Files
Giacomo Vercesi 4f9e811110 pypeline: implement Pypeline Data Viewer
Implement the Pypeline Data Viewer, an HTTP server that allows browsing
the data stored by PRSS.
2026-06-17 17:48:13 +02:00

320 lines
8.5 KiB
Python

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#
from __future__ import annotations
import asyncio
import enum
from abc import ABC, abstractmethod
from contextlib import AbstractAsyncContextManager
from dataclasses import dataclass
from typing import TypedDict
class LockRequest(enum.Enum):
ARTIFACT = enum.auto()
ANALYSIS = enum.auto()
class LockCheckType(enum.Enum):
# No lock required, check will always return True
NONE = enum.auto()
# Any lock, as long as it's valid, will return True
ANY = enum.auto()
# Requires an Artifact/Analysis lock
ARTIFACT = enum.auto()
# Requires an Analysis lock
ANALYSIS = enum.auto()
class CheckLockInvalid(Exception):
"""Exception throw by hold_lock if the lock is invalid"""
@dataclass(frozen=True, slots=True)
class InvalidatedObject:
"""A single invalidated object row returned after invalidation."""
object_id: str
container_id: str
savepoint_id: int
configuration_hash: str
class ProjectMetadataRow(TypedDict):
version: str | None
pipeline_description_hash: str | None
last_fetch: float
last_object_save: float
last_model_save: float
@dataclass(frozen=True, slots=True)
class DependencyEntry:
"""A single dependency to store."""
savepoint_id_start: int
savepoint_id_end: int
container_id: str
configuration_hash: str
object_id: bytes
model_path: str
@dataclass(frozen=True, slots=True)
class AdditionalObjectEntry:
savepoint_id_start: int
savepoint_id_end: int
container_id: str
configuration_hash: str
object_id: bytes
@dataclass(frozen=True, slots=True)
class CustomDependency:
"""A single custom dependency to store."""
pipe_id: int
configuration_hash: str
argument_index: int
object_id: bytes
data: bytes
@dataclass(frozen=True, slots=True)
class CustomDependencyEntry:
argument_index: int
object_id: bytes
data: bytes
@dataclass(frozen=True, slots=True)
class ObjectEntry:
"""A single object to store."""
savepoint_id: int
container_id: str
configuration_hash: str
object_id: bytes
object_id_string: str
content: bytes
@dataclass(frozen=True, slots=True)
class ModelSetResult:
"""Result of an invalidate-and-set-model operation."""
new_epoch: int
invalidated: list[InvalidatedObject]
class CheckMetadataResult(enum.Enum):
OK = enum.auto()
PIPELINE_DESCRIPTION_HASH_MISSING = enum.auto()
PRUNE_DONE = enum.auto()
class RSSStorage(ABC):
@staticmethod
@abstractmethod
def migrate(connection_string: str):
"""Apply migrations if needed"""
@staticmethod
@abstractmethod
def background_tasks(connection_string: str) -> list[asyncio.Task]:
"""Create the background tasks needed for housekeeping"""
@staticmethod
@abstractmethod
async def make_locked(connection_string: str) -> RSSStorage:
"""
Constructor method for the class, will return an instance which will
acquire a lock when creating a RSSLockedProjectStorage instance
"""
@staticmethod
@abstractmethod
async def make_unlocked(connection_string: str) -> RSSStorage:
"""
Constructor method for the class, will return an instance which **will
not** acquire a lock when creating a RSSLockedProjectStorage instance
"""
@abstractmethod
async def close(self):
"""Release any resources held by the storage backend."""
@abstractmethod
def get_project_storage(self, project_id: str) -> RSSProjectStorage:
"""Get the project storage for the specified project ID"""
class RSSProjectStorage(ABC):
# Locking primitives
@abstractmethod
async def make_lock(self, lock_type: LockRequest) -> str:
"""
Request a new lock of `lock_type`. This function will wait until the
lock is ready to be used.
"""
@abstractmethod
async def renew_lock(self, lock_id: str) -> bool:
"""
Renew the lock's refresh timestamp.
Returns False if the lock does not exist or is expired.
"""
@abstractmethod
async def release_lock(self, lock_id: str):
"""
Release the lock. This will also trigger the DB to unlock as many new
waiting locks as possible.
"""
# Metadata
@abstractmethod
async def get_metadata(self) -> ProjectMetadataRow:
"""Return the metadata of a project"""
@abstractmethod
async def create_project_if_missing(self):
"""Check if the project ID is present in the DB, if not add it"""
@abstractmethod
async def upgrade_project(
self, version: str, pipeline_description_hash: str
) -> CheckMetadataResult:
"""
Check that the version and has of the pipeline description match those
stored on the DB. If not all the objects will be pruned.
"""
@abstractmethod
async def put_pipeline_description(self, hash_: str, content: bytes):
"""Store a pipeline description blob"""
@abstractmethod
async def get_pipeline_description(self):
"""Get the current pipeline description"""
# Locked instance accessor
@abstractmethod
def get_locked(
self, lock_id: str, initial_type: LockCheckType
) -> AbstractAsyncContextManager[RSSLockedProjectStorage]:
"""
Get the LockedProjectStorage instance, the lock used will be held until
the context manager is closed.
"""
class RSSLockedProjectStorage(ABC):
# Savepoints
@abstractmethod
async def has_objects(
self,
savepoint_id: int,
container_id: str,
configuration_hash: str,
object_ids: list[bytes],
) -> list[bytes]:
"""
Given a list of wanted `object_ids` return a list of the subset that is
present in storage.
"""
@abstractmethod
async def get_objects(
self,
savepoint_id: int,
container_id: str,
configuration_hash: str,
object_ids: list[bytes],
) -> list[tuple[bytes, bytes]]:
"""
Given a list of `object_ids` that *need* to be present in the DB,
retrieve them and return their contents.
"""
@abstractmethod
async def get_custom_invalidation_data(
self,
pipe_id: int,
configuration_hash: str,
) -> list[CustomDependencyEntry]:
"""
Given a pipe_id and configuration_hash, return all the entries in the
DB that match.
"""
@abstractmethod
async def add_objects(
self,
dependencies: list[DependencyEntry],
custom_dependencies: list[CustomDependency],
objects: list[ObjectEntry],
):
"""
Store a group of objects and their dependencies into the DB
"""
@abstractmethod
async def prune_objects(self):
"""
Delete all objects, dependencies, and custom dependencies for a project.
"""
@abstractmethod
async def list_objects(self, savepoint_id: int, container_id: str) -> list[str]:
"""Return a list of available objects in storage"""
@abstractmethod
async def get_object(
self, savepoint_id: int, container_id: str, object_id: str
) -> bytes | None:
"""Return the contents of the specified object"""
# Model
@abstractmethod
async def get_epoch(self) -> int:
"""Return the current epoch for the project"""
@abstractmethod
async def get_model(self) -> tuple[bytes | None, int]:
"""Return the model (in bytes) and the epoch"""
@abstractmethod
async def invalidate_and_set_model(
self,
model_paths: list[str],
additional_objects: list[AdditionalObjectEntry],
model_bytes: bytes,
) -> ModelSetResult:
"""
Given a new model and the invalidation it caused (in the form of
changed model paths and objects that need to be removed), do the
following:
1. Run invalidation so that any object related to the modified paths
is dropped (in addition to the explicitly-specified
`additional_objects`)
2. Set the new model, increment the epoch by 1
"""
# File hashmap
@abstractmethod
async def put_file(self, hash_: str, content: bytes):
"""Store a file blob by its hash"""
@abstractmethod
async def get_files(self, hashes: list[str]) -> dict[str, bytes]:
"""Retrieve one or more files given a list of wanted hashes"""