pyattacker 0.1.0__py3-none-any.whl

pyattacker/artifact.py

@@ -0,0 +1,255 @@
+"""Artifact —— the persistent carrier of a task's state.
+
+Design notes:
+* Artifacts are **content-addressed** (blake2b digest), which gives deduplication and integrity checking for free.
+* An artifact is persisted as soon as it is produced —— that is the task-level checkpoint and the entire source of resume capability.
+* Encoding goes through a pluggable codec, JSON by default (supports dataclass / primitives / Enum / bytes).
+* An artifact with ``seq == SEED_SEQ`` is the seed input of the whole pipeline (one row of the dataset).
+"""
+
+from __future__ import annotations
+
+import base64
+import dataclasses
+import hashlib
+import json
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Protocol, runtime_checkable
+
+from .errors import ArtifactCodecError
+
+__all__ = [
+    "SEED_SEQ",
+    "SEED_TASK",
+    "Encoded",
+    "Artifact",
+    "Codec",
+    "JsonCodec",
+    "BytesCodec",
+    "CodecRegistry",
+    "DEFAULT_REGISTRY",
+    "canonical_json",
+    "digest_of",
+]
+
+SEED_SEQ = -1
+SEED_TASK = "__seed__"
+
+
+def _json_default(obj: Any) -> Any:
+    if dataclasses.is_dataclass(obj) and not isinstance(obj, type):
+        return dataclasses.asdict(obj)
+    if isinstance(obj, Enum):
+        return obj.value
+    if isinstance(obj, (set, frozenset)):
+        return sorted(obj, key=repr)
+    if isinstance(obj, (bytes, bytearray, memoryview)):
+        return {"__bytes__": base64.b64encode(bytes(obj)).decode("ascii")}
+    to_json = getattr(obj, "__json__", None)
+    if callable(to_json):
+        return to_json()
+    raise ArtifactCodecError(
+        f"artifact payload is not JSON-encodable: {type(obj).__name__}; "
+        f"register a codec for it (CodecRegistry.register)"
+    )
+
+
+def canonical_json(obj: Any) -> str:
+    """Stable serialization: sorted keys, no extra whitespace. Used for digest computation, not for display."""
+    return json.dumps(
+        obj, sort_keys=True, ensure_ascii=False, separators=(",", ":"), default=_json_default
+    )
+
+
+def digest_of(data: bytes | str) -> str:
+    if isinstance(data, str):
+        data = data.encode("utf-8")
+    return hashlib.blake2b(data, digest_size=16).hexdigest()
+
+
+@dataclass(frozen=True, slots=True)
+class Encoded:
+    """The result of one encoding pass. ``data`` is always bytes, ready to be written straight into a BLOB."""
+
+    type_name: str
+    codec: str
+    data: bytes
+    digest: str
+    size: int
+
+
+@runtime_checkable
+class Codec(Protocol):
+    name: str
+
+    def can_encode(self, obj: Any) -> bool: ...
+
+    def dumps(self, obj: Any) -> bytes: ...
+
+    def loads(self, data: bytes) -> Any: ...
+
+
+class JsonCodec:
+    name = "json"
+
+    def can_encode(self, obj: Any) -> bool:
+        try:
+            canonical_json(obj)
+        except ArtifactCodecError:
+            return False
+        return True
+
+    def dumps(self, obj: Any) -> bytes:
+        return canonical_json(obj).encode("utf-8")
+
+    def loads(self, data: bytes) -> Any:
+        return json.loads(data.decode("utf-8"))
+
+
+class BytesCodec:
+    """Raw byte payloads (images, audio, log chunks, ...)."""
+
+    name = "bytes"
+
+    def can_encode(self, obj: Any) -> bool:
+        return isinstance(obj, (bytes, bytearray, memoryview))
+
+    def dumps(self, obj: Any) -> bytes:
+        return bytes(obj)
+
+    def loads(self, data: bytes) -> Any:
+        return data
+
+
+class CodecRegistry:
+    """Registry mapping types → codecs; also responsible for restoring JSON payloads back into user types."""
+
+    def __init__(self) -> None:
+        self._codecs: dict[str, Codec] = {"json": JsonCodec(), "bytes": BytesCodec()}
+        self._by_type: dict[type, str] = {
+            bytes: "bytes",
+            bytearray: "bytes",
+            memoryview: "bytes",
+        }
+        self._rebuild: dict[str, type] = {}
+
+    def register(self, codec: Codec, *, for_types: Iterable[type] = (), name: str | None = None) -> None:
+        """Register a codec. **Later registrations take precedence over earlier ones.**
+
+        ``codec_for`` falls back to asking each registered codec in turn, and the built-in JSON
+        codec accepts almost anything, so insertion order decides whether a specialised codec ever
+        gets a chance. Moving each new registration to the front makes the rule explicit: if you
+        register a codec that claims a payload, it wins. An explicit ``for_types`` mapping still
+        beats the scan, because naming the type is the strongest statement of intent.
+        """
+        codec_name = name or codec.name
+        self._codecs = {codec_name: codec, **{k: v for k, v in self._codecs.items() if k != codec_name}}
+        for tp in for_types:
+            self._by_type[tp] = codec_name
+
+    def register_type(self, cls: type) -> type:
+        """Register a (dataclass) type; on restore the object is rebuilt as ``cls(**payload)``."""
+        self._rebuild[cls.__name__] = cls
+        return cls
+
+    def type_name_of(self, obj: Any) -> str:
+        return type(obj).__name__
+
+    def codec_for(self, obj: Any) -> Codec:
+        codec_name = self._by_type.get(type(obj))
+        if codec_name is not None:
+            return self._codecs[codec_name]
+        for base in type(obj).__mro__[1:]:
+            codec_name = self._by_type.get(base)
+            if codec_name is not None:
+                return self._codecs[codec_name]
+        for codec in self._codecs.values():
+            if codec.can_encode(obj):
+                return codec
+        raise ArtifactCodecError(f"no codec available for {type(obj).__name__}")
+
+    def dump(self, obj: Any) -> Encoded:
+        codec = self.codec_for(obj)
+        data = codec.dumps(obj)
+        return Encoded(
+            type_name=self.type_name_of(obj),
+            codec=codec.name,
+            data=data,
+            digest=digest_of(data),
+            size=len(data),
+        )
+
+    def load_raw(self, encoded: Encoded) -> Any:
+        codec = self._codecs.get(encoded.codec)
+        if codec is None:
+            raise ArtifactCodecError(f"unknown codec: {encoded.codec}")
+        return codec.loads(encoded.data)
+
+    def load(self, encoded: Encoded) -> Any:
+        """Decode + best-effort restore of the user type (a registered dataclass)."""
+        value = self.load_raw(encoded)
+        cls = self._rebuild.get(encoded.type_name)
+        if cls is None:
+            return value
+        if dataclasses.is_dataclass(cls) and isinstance(value, dict):
+            try:
+                return cls(**value)
+            except TypeError as exc:  # on field mismatch fall back to the raw dict, staying diagnosable
+                raise ArtifactCodecError(
+                    f"cannot restore artifact as {cls.__name__}: {exc}"
+                ) from exc
+        if isinstance(cls, type) and not isinstance(value, cls):
+            try:
+                return cls(value)
+            except Exception:  # pragma: no cover - best effort
+                return value
+        return value
+
+
+DEFAULT_REGISTRY = CodecRegistry()
+
+
+@dataclass(frozen=True, slots=True)
+class Artifact:
+    """The persisted state produced by a task."""
+
+    id: str
+    pipeline_id: str
+    task_name: str
+    seq: int
+    type_name: str
+    codec: str
+    digest: str
+    size: int
+    payload: bytes | None
+    created_at: float
+    is_final: bool = False
+    blob_ref: str | None = None
+    meta: dict[str, Any] = field(default_factory=dict)
+
+    @property
+    def available(self) -> bool:
+        """Whether the bytes are reachable right now.
+
+        Stores hydrate ``payload`` from the artifact backend on read, so this is False only when
+        the payload was deliberately dropped (``journal=summary``, a ``null`` backend) or is
+        genuinely missing — which is exactly the signal resume needs to rerun the pipeline.
+        """
+        return self.payload is not None
+
+    def encoded(self) -> Encoded:
+        if self.payload is None:
+            raise ArtifactCodecError(f"artifact {self.id} has no stored payload (a limitation of journal mode)")
+        return Encoded(
+            type_name=self.type_name,
+            codec=self.codec,
+            data=self.payload,
+            digest=self.digest,
+            size=self.size,
+        )
+
+    @staticmethod
+    def build_id(pipeline_id: str, seq: int) -> str:
+        return f"{pipeline_id}:{seq}"

pyattacker/backends.py

@@ -0,0 +1,228 @@
+"""Artifact backends: where payload bytes live.
+
+By default a payload lives inline in the store (a SQLite BLOB), which is right for the common
+case — a model response is a few kilobytes. It stops being right when an artifact is a rendered
+image, an audio clip or a multi-megabyte transcript: the database doubles in size, every backup
+copies it, and `export` has to materialize it.
+
+A backend answers one question — *should these bytes go somewhere else?* — and then owns them:
+
+* :class:`InlineBackend` (default): never spills, behaviour is exactly as before;
+* :class:`FileBackend`: content-addressed files under a root directory, written atomically;
+* :class:`NullBackend`: keeps the digest and drops the bytes (hash-only journals).
+
+The contract with the store is deliberately small: the store still records `digest`/`size`/`codec`
+in its own row, plus an opaque `blob_ref`. On read the store fills `payload` back in from the
+backend, so an artifact stays a single object to everything upstream.
+
+**Content addressing is what makes this safe**: the file name is the digest, so identical payloads
+collapse into one file, a partially written file can never be mistaken for a complete one, and a
+backend can be shared by many runs and many stores.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import json
+import os
+import tempfile
+from dataclasses import dataclass
+from typing import Any, Protocol, runtime_checkable
+
+from .artifact import Artifact
+from .errors import ConfigError, PyAttackerError
+
+__all__ = [
+    "ArtifactBackend",
+    "InlineBackend",
+    "FileBackend",
+    "NullBackend",
+    "resolve_backend",
+    "BACKENDS",
+]
+
+DEFAULT_MIN_BYTES = 256 * 1024
+
+
+@runtime_checkable
+class ArtifactBackend(Protocol):
+    """Where artifact payloads are kept."""
+
+    name: str
+
+    def wants(self, artifact: Artifact) -> bool:
+        """Should this artifact's payload be stored outside the database?"""
+        ...
+
+    def put(self, artifact: Artifact) -> str:
+        """Store the payload, returning an opaque reference the store will keep."""
+        ...
+
+    def get(self, ref: str) -> bytes | None:
+        """Fetch a payload by reference (``None`` when it is gone)."""
+        ...
+
+    def delete(self, ref: str) -> bool:
+        """Best-effort removal. Never called by the kernel; provided for operators."""
+        ...
+
+
+@dataclass
+class InlineBackend:
+    """Keep everything in the store. The default, and the previous behaviour exactly."""
+
+    name: str = "inline"
+
+    def wants(self, artifact: Artifact) -> bool:
+        return False
+
+    def put(self, artifact: Artifact) -> str:  # pragma: no cover - never called
+        raise PyAttackerError("InlineBackend never spills payloads")
+
+    def get(self, ref: str) -> bytes | None:
+        return None
+
+    def delete(self, ref: str) -> bool:
+        return False
+
+
+@dataclass
+class FileBackend:
+    """Content-addressed files under ``root``: ``<root>/ab/cdef…``.
+
+    Writes go to a temporary file in the same directory and are then ``os.replace``-d into place,
+    so a crash mid-write cannot leave a truncated blob that looks valid.
+    """
+
+    root: str
+    min_bytes: int = DEFAULT_MIN_BYTES
+    name: str = "file"
+
+    def __post_init__(self) -> None:
+        self.root = os.path.abspath(os.path.expanduser(str(self.root)))
+        os.makedirs(self.root, exist_ok=True)
+
+    def wants(self, artifact: Artifact) -> bool:
+        return artifact.payload is not None and artifact.size >= self.min_bytes
+
+    def path_for(self, digest: str) -> str:
+        return os.path.join(self.root, digest[:2], digest[2:])
+
+    def put(self, artifact: Artifact) -> str:
+        if artifact.payload is None:
+            raise PyAttackerError(f"artifact {artifact.id} has no payload to spill")
+        target = self.path_for(artifact.digest)
+        if not os.path.exists(target):
+            os.makedirs(os.path.dirname(target), exist_ok=True)
+            # Not a context manager on purpose: the file must outlive the block so it can be
+            # fsync-ed and then atomically renamed into place.
+            handle = tempfile.NamedTemporaryFile(  # noqa: SIM115
+                dir=os.path.dirname(target), prefix=".tmp-", delete=False
+            )
+            try:
+                handle.write(artifact.payload)
+                handle.flush()
+                os.fsync(handle.fileno())
+                handle.close()
+                os.replace(handle.name, target)
+            except BaseException:
+                handle.close()
+                with contextlib.suppress(OSError):  # pragma: no cover - best effort
+                    os.unlink(handle.name)
+                raise
+        return f"file://{target}"
+
+    def resolve(self, ref: str) -> str:
+        """Turn a reference into a path: ``file:///abs/path`` or a path relative to ``root``."""
+        path = ref[len("file://") :] if ref.startswith("file://") else ref
+        if not os.path.isabs(path):
+            path = os.path.join(self.root, path)
+        return path
+
+    def get(self, ref: str) -> bytes | None:
+        try:
+            with open(self.resolve(ref), "rb") as handle:
+                return handle.read()
+        except FileNotFoundError:
+            return None
+
+    def delete(self, ref: str) -> bool:
+        try:
+            os.unlink(self.resolve(ref))
+            return True
+        except OSError:
+            return False
+
+    def stats(self) -> dict[str, Any]:
+        """How much the backend holds (for operators; walks the tree, so not for hot paths)."""
+        files = 0
+        total = 0
+        for dirpath, _, names in os.walk(self.root):
+            for name in names:
+                if name.startswith(".tmp-"):
+                    continue
+                files += 1
+                try:
+                    total += os.path.getsize(os.path.join(dirpath, name))
+                except OSError:  # pragma: no cover
+                    continue
+        return {"root": self.root, "files": files, "bytes": total}
+
+
+@dataclass
+class NullBackend:
+    """Keep the digest, drop the bytes: useful when you only care *that* an artifact existed."""
+
+    name: str = "null"
+    min_bytes: int = 0
+    _seen: int = 0
+
+    def wants(self, artifact: Artifact) -> bool:
+        return artifact.payload is not None
+
+    def put(self, artifact: Artifact) -> str:
+        self._seen += 1
+        return f"null:{artifact.digest}:{artifact.size}"
+
+    def get(self, ref: str) -> bytes | None:
+        return None
+
+    def delete(self, ref: str) -> bool:
+        return True
+
+
+BACKENDS: dict[str, Any] = {"inline": InlineBackend, "file": FileBackend, "null": NullBackend}
+
+
+def resolve_backend(spec: Any) -> ArtifactBackend:
+    """``None``/``"inline"``/``"file:///data/blobs"``/``{"kind": "file", "root": ...}``/an instance."""
+    if spec is None:
+        return InlineBackend()
+    if isinstance(spec, ArtifactBackend):
+        return spec
+    if isinstance(spec, str):
+        raw = spec.strip()
+        if raw in ("", "inline", "none"):
+            return InlineBackend()
+        if raw == "null":
+            return NullBackend()
+        if raw.startswith("file://"):
+            return FileBackend(root=raw[len("file://") :])
+        if raw.startswith("{"):
+            try:
+                return resolve_backend(json.loads(raw))
+            except json.JSONDecodeError as exc:
+                raise ConfigError(f"artifact backend is not valid JSON: {exc}") from exc
+        # a bare path is treated as a file backend root, which is what people mean by "--blobs /data"
+        return FileBackend(root=raw)
+    if isinstance(spec, dict):
+        params = dict(spec)
+        kind = str(params.pop("kind", params.pop("name", "file")))
+        factory = BACKENDS.get(kind)
+        if factory is None:
+            raise ConfigError(f"unknown artifact backend {kind!r}; available: {sorted(BACKENDS)}")
+        try:
+            return factory(**params)
+        except TypeError as exc:
+            raise ConfigError(f"cannot configure {kind!r} backend: {exc}") from exc
+    raise ConfigError(f"cannot resolve artifact backend from {spec!r}")