app/backend/fluksio/flow/artifacts.py

"""The artifact store: bytes a node produced, addressed by their content.

A typed message carries JSON, which is what lets the same value pass through
Redis, the work queue and the worker protocol unchanged. A model checkpoint is
not that, so it does not travel as a message — it is written here and the
message carries a reference to it.

Addressed by digest rather than by run, for three reasons. A sweep whose fifty
configs share one preprocessed input stores it once. A reference stays valid
however it is passed around, because it names content instead of a location.
And the digest is what a stage cache will compare, so building it in now is
what keeps that from being a change to the message contract later.
"""

from __future__ import annotations

import contextlib
import hashlib
import logging
import os
import tempfile
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

#: How much is read at a time when hashing or serving.
CHUNK = 1024 * 1024
DIGEST_PREFIX = "sha256:"


def is_reference(value: Any) -> bool:
    """Whether a message payload is an artifact reference."""
    return isinstance(value, dict) and str(value.get("digest", "")).startswith(
        DIGEST_PREFIX
    )


def valid_digest(digest: str) -> bool:
    """Guard for anything that reaches the filesystem from outside."""
    if not digest.startswith(DIGEST_PREFIX):
        return False
    body = digest[len(DIGEST_PREFIX) :]
    return len(body) == 64 and all(c in "0123456789abcdef" for c in body)


class ArtifactStore:
    """Content-addressed files under one directory."""

    def __init__(self, root: Path) -> None:
        self.root = root
        self.root.mkdir(parents=True, exist_ok=True)

    def _path(self, digest: str) -> Path:
        body = digest[len(DIGEST_PREFIX) :]
        # Two levels, so a directory listing stays usable at a hundred thousand
        # artifacts.
        return self.root / body[:2] / body

    def put(
        self, chunks: Iterable[bytes], name: str = "", media_type: str = ""
    ) -> dict[str, Any]:
        """Store a stream and return the reference to it.

        Written to a temporary file first and moved into place once the digest
        is known, so a half-written artifact never has a name anyone can find.
        A file already there is left alone: identical content is identical.
        """
        digester = hashlib.sha256()
        size = 0
        handle = tempfile.NamedTemporaryFile(dir=self.root, delete=False)
        try:
            with handle:
                for chunk in chunks:
                    digester.update(chunk)
                    size += len(chunk)
                    handle.write(chunk)
            digest = DIGEST_PREFIX + digester.hexdigest()
            target = self._path(digest)
            target.parent.mkdir(parents=True, exist_ok=True)
            if target.exists():
                os.unlink(handle.name)
            else:
                # Same content from two nodes at once is one rename winning and
                # the other replacing a byte-identical file.
                os.replace(handle.name, target)
        except BaseException:
            with contextlib.suppress(OSError):
                os.unlink(handle.name)
            raise
        return {
            "digest": digest,
            "size": size,
            "media_type": media_type or "application/octet-stream",
            "name": name,
        }

    def put_file(self, path: Path, media_type: str = "") -> dict[str, Any]:
        with path.open("rb") as handle:
            return self.put(
                iter(lambda: handle.read(CHUNK), b""),
                name=path.name,
                media_type=media_type,
            )

    def path(self, digest: str) -> Path | None:
        """Where the bytes are, or None if this store does not have them."""
        if not valid_digest(digest):
            return None
        target = self._path(digest)
        return target if target.exists() else None

    def read(self, digest: str) -> Iterator[bytes]:
        target = self.path(digest)
        if target is None:
            raise FileNotFoundError(digest)
        with target.open("rb") as handle:
            while chunk := handle.read(CHUNK):
                yield chunk

    def collect(self, keep: set[str]) -> int:
        """Delete what no run refers to any more. Returns how many went.

        The caller passes every digest still recorded; anything else in the
        store was produced by a run that has since been pruned, or never got a
        row at all because the run failed between writing and recording.
        """
        removed = 0
        for entry in self.root.glob("*/*"):
            if not entry.is_file():
                continue
            if DIGEST_PREFIX + entry.name in keep:
                continue
            try:
                entry.unlink()
                removed += 1
            except OSError:
                logger.warning("Could not remove artifact %s", entry.name)
        return removed