datadoom 0.1.0.dev0__py3-none-any.whl

datadoom/jobs/worker.py

@@ -0,0 +1,162 @@
+"""In-process generation worker (03 §3.3, 17 step 8).
+
+A run is submitted to a thread pool; the worker loads the immutable spec snapshot
+from ``store``, drives the single ``engine.pipeline`` entry point, streams
+progress to the :class:`EventHub`, persists artifacts + report, and flips the
+``GenerationRun`` (and its dataset) status. Cancellation is cooperative — checked
+at each pipeline stage boundary by :class:`HubProgressEmitter`.
+
+This is the only code path that turns a queued run into artifacts; the CLI and
+``datadoom.generate()`` share the same engine underneath, never a fork of it.
+"""
+
+from __future__ import annotations
+
+import traceback
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+from datadoom.engine import generate, parse_spec
+from datadoom.store import (
+    ArtifactRepository,
+    ArtifactStore,
+    Database,
+    DatasetRepository,
+    ReportRepository,
+    RunRepository,
+    SpecRow,
+    utcnow_iso,
+)
+
+from .progress import EventHub, HubProgressEmitter, RunCancelled
+
+
+class WorkerPool:
+    def __init__(
+        self,
+        db: Database,
+        artifacts: ArtifactStore,
+        hub: EventHub,
+        package_version: str,
+        max_workers: int = 2,
+    ) -> None:
+        self.db = db
+        self.artifacts = artifacts
+        self.hub = hub
+        self.package_version = package_version
+        self._pool = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="dd-run")
+
+    def submit(self, run_id: str) -> None:
+        """Schedule a queued run for execution (returns immediately)."""
+        self._pool.submit(self._execute, run_id)
+
+    def shutdown(self) -> None:
+        self._pool.shutdown(wait=False, cancel_futures=True)
+
+    # --- execution ----------------------------------------------------------------
+    def _execute(self, run_id: str) -> None:
+        try:
+            self._run_pipeline(run_id)
+        except RunCancelled:
+            self._mark_cancelled(run_id)
+        except Exception as exc:  # noqa: BLE001 — persist any failure, never crash the pool
+            self._mark_failed(run_id, exc)
+
+    def _run_pipeline(self, run_id: str) -> None:
+        # Load run + spec; flip to running.
+        with self.db.session() as s:
+            run = RunRepository(s).get(run_id)
+            if run is None:
+                return
+            if self.hub.is_cancelled(run_id):
+                raise RunCancelled("queued")
+            spec_row = s.get(SpecRow, run.spec_id)
+            assert spec_row is not None
+            spec_body = dict(spec_row.body)
+            dataset_id = run.dataset_id
+            seed = run.seed
+            run.status = "running"
+            run.stage = "intake"
+            run.started_at = utcnow_iso()
+            ds = DatasetRepository(s).get(dataset_id)
+            if ds is not None:
+                ds.status = "running"
+                ds.updated_at = utcnow_iso()
+
+        spec = parse_spec(spec_body)
+        out_dir: Path = self.artifacts.run_dir(dataset_id, run_id)
+        emitter = HubProgressEmitter(self.hub, run_id)
+
+        result = generate(spec, seed=seed, out_dir=out_dir, progress=emitter)
+        emitter.finish()
+
+        # Persist artifacts + report + final status (one transaction).
+        with self.db.session() as s:
+            arts = ArtifactRepository(s)
+            for art in result.artifacts:
+                abs_path = out_dir / art.path
+                arts.add(
+                    run_id=run_id,
+                    version=art.version,
+                    fmt=art.format,
+                    storage_uri=self.artifacts.to_uri(abs_path),
+                    checksum_sha256=art.checksum_sha256,
+                    size_bytes=art.size_bytes,
+                    split="full",
+                )
+            report = ReportRepository(s).upsert(run_id, result.report.to_dict())
+            run = RunRepository(s).get(run_id)
+            assert run is not None
+            run.status = "completed"
+            run.stage = "packaging"
+            run.progress_pct = 100
+            run.finished_at = utcnow_iso()
+            run.metrics = {"compliance_score": result.compliance.score}
+            ds = DatasetRepository(s).get(dataset_id)
+            if ds is not None:
+                ds.status = "completed"
+                ds.latest_run_id = run_id
+                ds.updated_at = utcnow_iso()
+            report_id = report.report_id
+
+        self.hub.publish(
+            run_id,
+            {
+                "type": "completed",
+                "run_id": run_id,
+                "compliance_score": result.compliance.score,
+                "report_id": report_id,
+            },
+        )
+
+    # --- terminal states ----------------------------------------------------------
+    def _mark_cancelled(self, run_id: str) -> None:
+        with self.db.session() as s:
+            run = RunRepository(s).get(run_id)
+            if run is not None:
+                run.status = "cancelled"
+                run.finished_at = utcnow_iso()
+                ds = DatasetRepository(s).get(run.dataset_id)
+                if ds is not None and ds.status == "running":
+                    ds.status = "draft"
+                    ds.updated_at = utcnow_iso()
+        self.hub.publish(run_id, {"type": "cancelled", "run_id": run_id})
+
+    def _mark_failed(self, run_id: str, exc: Exception) -> None:
+        stage = "unknown"
+        tb = traceback.format_exc()
+        with self.db.session() as s:
+            run = RunRepository(s).get(run_id)
+            if run is not None:
+                stage = run.stage or "unknown"
+                run.status = "failed"
+                run.error = {"message": str(exc), "traceback": tb, "stage": stage}
+                run.finished_at = utcnow_iso()
+                ds = DatasetRepository(s).get(run.dataset_id)
+                if ds is not None:
+                    ds.status = "failed"
+                    ds.updated_at = utcnow_iso()
+        self.hub.publish(
+            run_id,
+            {"type": "failed", "stage": stage, "message": str(exc), "traceback": tb},
+        )

datadoom/plugin.py

@@ -0,0 +1,35 @@
+"""Public, stable contract surface for plugin authors (09 §4).
+
+Import the base classes and the ``schema`` helper from here::
+
+    from datadoom.plugin import Distribution, schema
+
+    class WeibullDistribution(Distribution):
+        name = "weibull"
+        param_schema = schema({"k": {"type": "number", "minimum": 0}})
+        def sample(self, rng, n, params): ...
+        def cdf(self, x, params): ...
+
+These names are re-exported from ``datadoom.plugins.contracts`` and stay stable
+even as the engine's internal layout evolves.
+"""
+
+from __future__ import annotations
+
+from datadoom.plugins.contracts import (
+    Distribution,
+    Exporter,
+    FailureMode,
+    ProbeModel,
+    StructuralFn,
+    schema,
+)
+
+__all__ = [
+    "Distribution",
+    "StructuralFn",
+    "FailureMode",
+    "Exporter",
+    "ProbeModel",
+    "schema",
+]

datadoom/plugins/__init__.py

@@ -0,0 +1,47 @@
+"""DataDoom plugin system — registry + loader + scaffolder (09, 17 step 17).
+
+Plugins extend DataDoom *without forking core*: a small class implementing one of
+the engine ABCs (re-exported as ``datadoom.plugin``) is discovered at startup and
+inserted into the engine's lookup tables, so it works in the CLI, the API, and the
+web UI with no core change. This package depends only on the engine (10 §4).
+"""
+
+from __future__ import annotations
+
+from .contracts import (
+    Distribution,
+    Exporter,
+    FailureMode,
+    ProbeModel,
+    StructuralFn,
+    schema,
+)
+from .loader import default_plugins_dir, load_plugins
+from .registry import (
+    PluginConflictError,
+    PluginError,
+    PluginRecord,
+    PluginRegistry,
+    get_registry,
+)
+from .scaffold import ObjectCheck, check_object, check_plugin, scaffold_plugin
+
+__all__ = [
+    "Distribution",
+    "StructuralFn",
+    "FailureMode",
+    "Exporter",
+    "ProbeModel",
+    "schema",
+    "load_plugins",
+    "default_plugins_dir",
+    "get_registry",
+    "PluginRegistry",
+    "PluginRecord",
+    "PluginError",
+    "PluginConflictError",
+    "scaffold_plugin",
+    "check_plugin",
+    "check_object",
+    "ObjectCheck",
+]

datadoom/plugins/contracts.py

@@ -0,0 +1,72 @@
+"""Plugin author contract surface — re-exported as ``datadoom.plugin`` (09 §4).
+
+A plugin is a small class implementing one of the engine's typed base classes.
+This module re-exports those ABCs (so authors import them from a stable name that
+does not change as the engine's internal layout evolves) and ships the tiny
+``schema`` helper for declaring a plugin's ``param_schema`` JSON-schema fragment.
+
+The canonical ABCs live in ``datadoom.engine`` — the engine *consumes* registered
+plugin instances through the registry (it never imports ``datadoom.plugins``); the
+registry mutates the engine's lookup tables in place (10 §4). Plugins therefore
+depend only on these ABCs.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from typing import Any
+
+from datadoom.engine.causal.functions import StructuralFn
+from datadoom.engine.difficulty.probes import ProbeModel
+from datadoom.engine.dist.base import Distribution
+from datadoom.engine.export.base import Exporter
+from datadoom.engine.failure.base import FailureMode
+
+__all__ = [
+    "Distribution",
+    "StructuralFn",
+    "FailureMode",
+    "Exporter",
+    "ProbeModel",
+    "schema",
+    "PLUGIN_BASES",
+    "KEY_ATTR",
+]
+
+# Plugin kind -> the base class an instance must subclass. The order is the
+# precedence used when resolving an object's kind (kinds are disjoint in practice).
+PLUGIN_BASES: dict[str, type] = {
+    "distribution": Distribution,
+    "structural_fn": StructuralFn,
+    "failure_mode": FailureMode,
+    "exporter": Exporter,
+    "probe_model": ProbeModel,
+}
+
+# The attribute that names each kind inside its registry. Exporters key on
+# ``format`` (it matches ``export.formats[]``); everything else keys on ``name``.
+KEY_ATTR: dict[str, str] = {
+    "distribution": "name",
+    "structural_fn": "name",
+    "failure_mode": "name",
+    "exporter": "format",
+    "probe_model": "name",
+}
+
+
+def schema(
+    properties: Mapping[str, Mapping[str, Any]],
+    *,
+    required: Sequence[str] | None = None,
+) -> dict[str, Any]:
+    """Wrap a property map into a JSON-schema ``object`` the Canvas can render.
+
+    ``param_schema = schema({"k": {"type": "number", "minimum": 0}})`` →
+    ``{"type": "object", "properties": {...}, "required": [...]}``. The UI reads
+    this fragment and renders form controls (number inputs with min/max, enums as
+    dropdowns…) wherever the plugin is selectable, with no frontend changes (09 §6).
+    """
+    obj: dict[str, Any] = {"type": "object", "properties": dict(properties)}
+    if required:
+        obj["required"] = list(required)
+    return obj

datadoom/plugins/loader.py

@@ -0,0 +1,125 @@
+"""Plugin discovery + loading (09 §3, 17 step 17).
+
+Two mechanisms resolve into the in-memory :class:`PluginRegistry` at server/CLI
+startup:
+
+1. **Python entry points** under the ``datadoom.plugins`` group (preferred for
+   published ``datadoom-plugin-*`` packages).
+2. **A local plugins directory** (``$DATADOOM_HOME/plugins/*.py``) auto-imported
+   for prototyping and per-project plugins.
+
+Each discovered object is registered (validated + inserted into the engine
+tables). Conflicts and broken plugins fail loudly with a clear message rather
+than silently degrading a run.
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import inspect
+import os
+import sys
+from importlib.metadata import entry_points
+from pathlib import Path
+from types import ModuleType
+
+from .registry import PluginError, PluginRegistry, get_registry, resolve_kind_class
+
+ENTRY_POINT_GROUP = "datadoom.plugins"
+
+
+def default_plugins_dir() -> Path:
+    """``$DATADOOM_HOME/plugins`` (or ``~/.datadoom/plugins``).
+
+    Resolved from the environment directly so this package depends only on the
+    engine (it does not import ``datadoom.config``).
+    """
+    home_env = os.environ.get("DATADOOM_HOME")
+    home = Path(home_env).expanduser() if home_env else Path.home() / ".datadoom"
+    return home / "plugins"
+
+
+def _instantiate(target: object, *, where: str) -> object:
+    """Coerce an entry-point/discovered target to a plugin *instance*."""
+    if inspect.isclass(target):
+        try:
+            return target()
+        except Exception as exc:  # noqa: BLE001 - surface as a clean plugin error
+            raise PluginError(f"could not instantiate {where}: {exc}") from exc
+    return target
+
+
+def load_entry_points(registry: PluginRegistry) -> list[str]:
+    """Register every plugin advertised under the ``datadoom.plugins`` group."""
+    loaded: list[str] = []
+    for ep in entry_points(group=ENTRY_POINT_GROUP):
+        try:
+            target = ep.load()
+        except Exception as exc:  # noqa: BLE001
+            raise PluginError(f"entry point {ep.name!r} failed to import: {exc}") from exc
+        obj = _instantiate(target, where=f"entry point {ep.name!r}")
+        version = _dist_version(ep)
+        record = registry.register(obj, source="entrypoint", version=version)
+        loaded.append(f"{record.kind}:{record.name}")
+    return loaded
+
+
+def load_local_dir(registry: PluginRegistry, directory: Path) -> list[str]:
+    """Import every ``*.py`` in ``directory`` and register the plugin classes it defines."""
+    loaded: list[str] = []
+    if not directory.is_dir():
+        return loaded
+    for path in sorted(directory.glob("*.py")):
+        if path.name.startswith("_"):
+            continue
+        module = _import_path(path)
+        for _, member in inspect.getmembers(module, inspect.isclass):
+            # Only classes *defined here* (skip the imported ABCs themselves).
+            if member.__module__ != module.__name__:
+                continue
+            if resolve_kind_class(member) is None:
+                continue
+            record = registry.register(member(), source="local", module=str(path))
+            loaded.append(f"{record.kind}:{record.name}")
+    return loaded
+
+
+def _import_path(path: Path) -> ModuleType:
+    mod_name = f"datadoom_local_plugin_{path.stem}"
+    spec = importlib.util.spec_from_file_location(mod_name, path)
+    if spec is None or spec.loader is None:
+        raise PluginError(f"could not load local plugin {path}")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[mod_name] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception as exc:  # noqa: BLE001
+        sys.modules.pop(mod_name, None)
+        raise PluginError(f"local plugin {path.name} failed to import: {exc}") from exc
+    return module
+
+
+def _dist_version(ep: object) -> str | None:
+    dist = getattr(ep, "dist", None)
+    return getattr(dist, "version", None) if dist is not None else None
+
+
+def load_plugins(
+    registry: PluginRegistry | None = None,
+    *,
+    use_entry_points: bool = True,
+    local_dir: Path | None = None,
+    use_local: bool = True,
+) -> PluginRegistry:
+    """Seed built-ins, then discover entry-point and local-directory plugins.
+
+    Idempotent for built-ins; intended to be called once at startup. ``local_dir``
+    defaults to :func:`default_plugins_dir` when ``use_local`` is set.
+    """
+    registry = registry or get_registry()
+    registry.seed_builtins()
+    if use_entry_points:
+        load_entry_points(registry)
+    if use_local:
+        load_local_dir(registry, local_dir or default_plugins_dir())
+    return registry

datadoom/plugins/registry.py

@@ -0,0 +1,214 @@
+"""In-memory plugin registry (09 §3, 17 step 17).
+
+The registry is the single place that knows *every* capability available to a
+run — both core built-ins and discovered plugins. Registering a plugin inserts
+its instance into the matching **engine lookup table** (the same dict the
+pipeline reads by name), so a registered distribution/fn/failure/exporter/probe
+is picked up with no engine change. The engine never imports this module; the
+dependency points the other way (``engine ← plugins``, 10 §4).
+
+Conflicts (a plugin reusing a built-in or another plugin's name) fail loudly so
+an install never silently shadows a capability.
+"""
+
+from __future__ import annotations
+
+from collections.abc import MutableMapping
+from dataclasses import dataclass
+from typing import Any
+
+from .contracts import KEY_ATTR, PLUGIN_BASES
+
+
+class PluginError(Exception):
+    """A plugin failed validation (wrong base class, bad schema, missing name)."""
+
+
+class PluginConflictError(PluginError):
+    """Two capabilities of the same kind claim the same name."""
+
+
+_ALLOWED_SCHEMA_TYPES = {"object", "string", "number", "integer", "boolean", "array", "null"}
+
+
+@dataclass(frozen=True)
+class PluginRecord:
+    """Metadata for one registered capability (built-in or plugin)."""
+
+    name: str
+    kind: str
+    source: str  # "builtin" | "entrypoint" | "local"
+    module: str | None = None
+    version: str | None = None
+    schema: dict[str, Any] | None = None
+
+    @property
+    def builtin(self) -> bool:
+        return self.source == "builtin"
+
+    def to_info(self) -> dict[str, Any]:
+        """Shape consumed by ``GET /api/plugins`` (08 §10)."""
+        return {
+            "name": self.name,
+            "kind": self.kind,
+            "version": self.version,
+            "schema": self.schema,
+            "source": self.source,
+            "builtin": self.builtin,
+            "enabled": True,
+        }
+
+
+def resolve_kind(obj: object) -> str | None:
+    """Return the plugin kind ``obj`` (an instance) implements, or ``None``."""
+    for kind, base in PLUGIN_BASES.items():
+        if isinstance(obj, base):
+            return kind
+    return None
+
+
+def resolve_kind_class(cls: object) -> str | None:
+    """Return the plugin kind a *class* implements, or ``None`` (skips the ABCs)."""
+    if not isinstance(cls, type):
+        return None
+    for kind, base in PLUGIN_BASES.items():
+        if cls is not base and issubclass(cls, base):
+            return kind
+    return None
+
+
+def validate_param_schema(schema: object, *, where: str = "param_schema") -> None:
+    """Lightweight sanity check that ``schema`` is a renderable JSON-schema fragment."""
+    if not isinstance(schema, dict):
+        raise PluginError(f"{where} must be a dict (use the schema() helper)")
+    stype = schema.get("type")
+    if stype is not None and stype not in _ALLOWED_SCHEMA_TYPES:
+        raise PluginError(f"{where} has unsupported type {stype!r}")
+    props = schema.get("properties")
+    if stype == "object" and not isinstance(props, dict):
+        raise PluginError(f"{where} of type 'object' needs a 'properties' dict")
+    if isinstance(props, dict):
+        for pname, pspec in props.items():
+            if not isinstance(pspec, dict):
+                raise PluginError(f"{where}.properties[{pname!r}] must be a dict")
+            ptype = pspec.get("type")
+            if ptype is not None and ptype not in _ALLOWED_SCHEMA_TYPES:
+                raise PluginError(f"{where}.properties[{pname!r}] has unsupported type {ptype!r}")
+
+
+def _engine_registries() -> dict[str, MutableMapping[str, Any]]:
+    """The five live engine lookup tables, keyed by plugin kind.
+
+    Imported lazily so this module loads without forcing the whole engine, and so
+    the dicts are the *canonical* objects the pipeline reads (mutating them in
+    place propagates to every ``from ... import REGISTRY`` reference).
+    """
+    from datadoom.engine.causal.functions import STRUCTURAL_FNS
+    from datadoom.engine.difficulty.probes import PROBES
+    from datadoom.engine.dist.builtins import REGISTRY
+    from datadoom.engine.export import EXPORTERS
+    from datadoom.engine.failure.modes import FAILURE_MODES
+
+    return {
+        "distribution": REGISTRY,
+        "structural_fn": STRUCTURAL_FNS,
+        "failure_mode": FAILURE_MODES,
+        "exporter": EXPORTERS,
+        "probe_model": PROBES,
+    }
+
+
+class PluginRegistry:
+    """Tracks records and keeps the engine lookup tables in sync."""
+
+    def __init__(self) -> None:
+        self._records: dict[tuple[str, str], PluginRecord] = {}
+        self._builtins_seeded = False
+
+    def seed_builtins(self) -> None:
+        """Record the core built-ins already present in the engine tables (idempotent)."""
+        if self._builtins_seeded:
+            return
+        for kind, reg in _engine_registries().items():
+            for key, obj in reg.items():
+                self._records[(kind, key)] = PluginRecord(
+                    name=key,
+                    kind=kind,
+                    source="builtin",
+                    module=type(obj).__module__,
+                    schema=_schema_of(obj),
+                )
+        self._builtins_seeded = True
+
+    def register(
+        self,
+        obj: object,
+        *,
+        source: str,
+        module: str | None = None,
+        version: str | None = None,
+    ) -> PluginRecord:
+        """Validate ``obj`` and insert it into its engine table. Raise on conflict."""
+        kind = resolve_kind(obj)
+        if kind is None:
+            bases = ", ".join(PLUGIN_BASES)
+            raise PluginError(
+                f"{type(obj).__name__} is not a plugin: it must subclass one of [{bases}]"
+            )
+        key_attr = KEY_ATTR[kind]
+        key = getattr(obj, key_attr, None)
+        if not isinstance(key, str) or not key:
+            raise PluginError(
+                f"{kind} plugin {type(obj).__name__} must set a non-empty '{key_attr}'"
+            )
+        if (kind, key) in self._records:
+            existing = self._records[(kind, key)]
+            raise PluginConflictError(
+                f"{kind} {key!r} is already registered (source={existing.source}); "
+                "capability names must be unique within an install"
+            )
+        schema = _schema_of(obj)
+        if schema is not None:
+            validate_param_schema(schema, where=f"{kind} {key!r} param_schema")
+
+        _engine_registries()[kind][key] = obj
+        record = PluginRecord(
+            name=key,
+            kind=kind,
+            source=source,
+            module=module or type(obj).__module__,
+            version=version,
+            schema=schema,
+        )
+        self._records[(kind, key)] = record
+        return record
+
+    def records(self) -> list[PluginRecord]:
+        return sorted(self._records.values(), key=lambda r: (r.kind, r.name))
+
+    def to_info(self) -> list[dict[str, Any]]:
+        return [r.to_info() for r in self.records()]
+
+    def get(self, kind: str, name: str) -> PluginRecord | None:
+        return self._records.get((kind, name))
+
+    def reset(self) -> None:
+        """Remove every non-built-in from the engine tables and records (test aid)."""
+        regs = _engine_registries()
+        for (kind, key), rec in list(self._records.items()):
+            if rec.source != "builtin":
+                regs[kind].pop(key, None)
+                del self._records[(kind, key)]
+
+
+def _schema_of(obj: object) -> dict[str, Any] | None:
+    schema = getattr(obj, "param_schema", None)
+    return dict(schema) if isinstance(schema, dict) else None
+
+
+_DEFAULT = PluginRegistry()
+
+
+def get_registry() -> PluginRegistry:
+    """The process-wide registry (shared by the API, CLI, and engine tables)."""
+    return _DEFAULT