meds-algorithm 0.4.0a4__py3-none-any.whl

meds_algorithm/__init__.py

@@ -0,0 +1,94 @@
+"""Shared MEDS algorithm infrastructure; data dependencies are imported only when used."""
+
+from .algorithms import (
+    Algorithm,
+    FinetuneAlgorithm,
+    PredictOnlyAlgorithm,
+    PretrainedAlgorithm,
+    ProbeAlgorithm,
+    SupervisedAlgorithm,
+    ZeroShotDirectAlgorithm,
+    ZeroShotMaterializedAlgorithm,
+)
+from .artifacts.manifest import digest_config, digest_file, measured_artifact, read_manifest
+from .artifacts.references import (
+    ArtifactDestination,
+    ArtifactReference,
+    EmbeddingsArtifact,
+    FinetunedAlgorithmArtifact,
+    PayloadMetadata,
+    PredictionsArtifact,
+    PreparedDataArtifact,
+    PretrainedAlgorithmArtifact,
+    ProbeAlgorithmArtifact,
+    SupervisedAlgorithmArtifact,
+    TrajectoriesArtifact,
+)
+from .artifacts.results import write_run_result
+from .configuration import AlgorithmConfig, ExecutionOptions, NNConfig
+from .contracts.commands import (
+    ALGORITHM_PROFILES,
+    STEP_SPECS,
+    ZERO_SHOT_TASK_COMMANDS,
+    ArtifactInput,
+    ArtifactType,
+    Command,
+    CommandName,
+    LegacyAlgorithmRepository,
+    StepSpec,
+)
+from .data.index import read_index, split_files
+from .data.labels import read_context_labels
+from .evaluation.predictions import attach_ground_truth, check_decoded
+from .runner import AlgorithmRepository, StepRunner
+from .tasks.spec import RelaxationLedger, TaskSpec
+from .version import CONTRACT_VERSION
+
+__all__ = [
+    "ALGORITHM_PROFILES",
+    "CONTRACT_VERSION",
+    "STEP_SPECS",
+    "ZERO_SHOT_TASK_COMMANDS",
+    "Algorithm",
+    "AlgorithmConfig",
+    "AlgorithmRepository",
+    "ArtifactDestination",
+    "ArtifactInput",
+    "ArtifactReference",
+    "ArtifactType",
+    "Command",
+    "CommandName",
+    "EmbeddingsArtifact",
+    "ExecutionOptions",
+    "FinetuneAlgorithm",
+    "FinetunedAlgorithmArtifact",
+    "LegacyAlgorithmRepository",
+    "NNConfig",
+    "PayloadMetadata",
+    "PredictOnlyAlgorithm",
+    "PredictionsArtifact",
+    "PreparedDataArtifact",
+    "PretrainedAlgorithm",
+    "PretrainedAlgorithmArtifact",
+    "ProbeAlgorithm",
+    "ProbeAlgorithmArtifact",
+    "RelaxationLedger",
+    "StepRunner",
+    "StepSpec",
+    "SupervisedAlgorithm",
+    "SupervisedAlgorithmArtifact",
+    "TaskSpec",
+    "TrajectoriesArtifact",
+    "ZeroShotDirectAlgorithm",
+    "ZeroShotMaterializedAlgorithm",
+    "attach_ground_truth",
+    "check_decoded",
+    "digest_config",
+    "digest_file",
+    "measured_artifact",
+    "read_context_labels",
+    "read_index",
+    "read_manifest",
+    "split_files",
+    "write_run_result",
+]

meds_algorithm/algorithms.py

@@ -0,0 +1,157 @@
+"""Profile ABCs are the authority for lifecycle order and typed operation signatures."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import ClassVar, get_type_hints
+
+from .artifacts.references import PayloadMetadata
+from .configuration import AlgorithmConfig
+from .contracts.enums import CommandName
+from .requests import (
+    InferEmbeddingsRequest,
+    InferTrajectoriesRequest,
+    PredictFinetuneRequest,
+    PredictOnlyRequest,
+    PredictProbeRequest,
+    PredictSupervisedRequest,
+    PredictZeroShotDirectRequest,
+    PredictZeroShotMaterializedRequest,
+    PreprocessDataRequest,
+    PretrainRequest,
+    StepRequest,
+    TrainFinetuneRequest,
+    TrainProbeRequest,
+    TrainSupervisedRequest,
+)
+
+
+class Algorithm(ABC):
+    name: ClassVar[str]
+    config_type: ClassVar[type[AlgorithmConfig]] = AlgorithmConfig
+
+    def __init__(self, config: AlgorithmConfig | None = None):
+        self.config = self.config_type() if config is None else config
+        if type(self.config) is not self.config_type:
+            raise TypeError(f"Expected {self.config_type.__name__}")
+        if not getattr(self, "name", "").strip():
+            raise ValueError("Algorithm.name must be nonempty")
+        self.config.validate()
+
+    @abstractmethod
+    def preprocess_data(self, request: PreprocessDataRequest) -> PayloadMetadata:
+        """Write task-independent prepared data into request.output.directory."""
+
+    def validate_inputs(self, request: StepRequest) -> None:
+        """Default: fitted inputs must match this algorithm and resolved configuration.
+
+        Extend for payload/vocabulary/cutoff checks. Override deliberately if a documented family
+        supports compatible external pretrained algorithms or different training/inference settings.
+        """
+        for _, artifact in request.artifact_inputs():
+            if artifact.kind.value.endswith("algorithm"):
+                producer = artifact.read_manifest().get("algorithm", {})
+                recorded = producer.get("config")
+                if isinstance(recorded, dict):
+                    # Artifacts predating this optional field mean no feature additions.
+                    recorded = {"additional_predicates_file": None, **recorded}
+                if producer.get("name") != self.name or recorded != self.config.to_dict():
+                    raise ValueError(
+                        "Algorithm algorithm/configuration mismatch; load its recorded configuration"
+                    )
+
+
+class SupervisedAlgorithm(Algorithm):
+    profile = "supervised"
+    steps = ("preprocess_data", "train_supervised", "predict_supervised")
+
+    @abstractmethod
+    def train_supervised(self, request: TrainSupervisedRequest) -> PayloadMetadata: ...
+
+    @abstractmethod
+    def predict_supervised(self, request: PredictSupervisedRequest) -> PayloadMetadata: ...
+
+
+class PretrainedAlgorithm(Algorithm):
+    @abstractmethod
+    def pretrain(self, request: PretrainRequest) -> PayloadMetadata: ...
+
+
+class FinetuneAlgorithm(PretrainedAlgorithm):
+    profile = "finetune"
+    steps = ("preprocess_data", "pretrain", "train_finetune", "predict_finetune")
+
+    @abstractmethod
+    def train_finetune(self, request: TrainFinetuneRequest) -> PayloadMetadata: ...
+
+    @abstractmethod
+    def predict_finetune(self, request: PredictFinetuneRequest) -> PayloadMetadata: ...
+
+
+class ProbeAlgorithm(PretrainedAlgorithm):
+    profile = "probe"
+    steps = ("preprocess_data", "pretrain", "infer_embeddings", "train_probe", "predict_probe")
+
+    @abstractmethod
+    def infer_embeddings(self, request: InferEmbeddingsRequest) -> PayloadMetadata: ...
+
+    @abstractmethod
+    def train_probe(self, request: TrainProbeRequest) -> PayloadMetadata: ...
+
+    @abstractmethod
+    def predict_probe(self, request: PredictProbeRequest) -> PayloadMetadata: ...
+
+
+class ZeroShotDirectAlgorithm(PretrainedAlgorithm):
+    profile = "zero_shot_direct"
+    steps = ("preprocess_data", "pretrain", "predict_zero_shot_direct")
+
+    @abstractmethod
+    def predict_zero_shot_direct(self, request: PredictZeroShotDirectRequest) -> PayloadMetadata: ...
+
+
+class ZeroShotMaterializedAlgorithm(PretrainedAlgorithm):
+    profile = "zero_shot_materialized"
+    steps = ("preprocess_data", "pretrain", "infer_trajectories", "predict_zero_shot_materialized")
+
+    @abstractmethod
+    def infer_trajectories(self, request: InferTrajectoriesRequest) -> PayloadMetadata: ...
+
+    @abstractmethod
+    def predict_zero_shot_materialized(
+        self, request: PredictZeroShotMaterializedRequest
+    ) -> PayloadMetadata: ...
+
+
+class PredictOnlyAlgorithm(Algorithm):
+    profile = "predict_only"
+    steps = ("preprocess_data", "predict_only")
+
+    @abstractmethod
+    def predict_only(self, request: PredictOnlyRequest) -> PayloadMetadata: ...
+
+
+def profile_classes() -> tuple[type[Algorithm], ...]:
+    """Explicit built-ins; no plugin discovery or duplicated command specification."""
+    return (
+        SupervisedAlgorithm,
+        FinetuneAlgorithm,
+        ProbeAlgorithm,
+        ZeroShotDirectAlgorithm,
+        ZeroShotMaterializedAlgorithm,
+        PredictOnlyAlgorithm,
+    )
+
+
+def profile_class(profile: str) -> type[Algorithm]:
+    for cls in profile_classes():
+        if cls.profile == profile:
+            return cls
+    raise ValueError(f"Unknown algorithm profile: {profile}")
+
+
+def request_type(profile: str, step: str | CommandName) -> type[StepRequest]:
+    cls = profile_class(profile)
+    if step not in cls.steps:
+        raise ValueError(f"{step} is not part of the {profile} algorithm profile")
+    return get_type_hints(getattr(cls, step))["request"]

meds_algorithm/artifacts/__init__.py

@@ -0,0 +1 @@
+"""Artifact manifests, publication, input resolution and run evidence."""

meds_algorithm/artifacts/inputs.py

@@ -0,0 +1,38 @@
+"""Expose an upstream artifact reference as an explicit orchestration argument."""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from ..contracts.commands import STEP_SPECS, CommandName
+from .manifest import read_manifest
+
+
+def resolve_input(artifact: Path, argument: str) -> Path:
+    """Resolve a producer's recorded input; validate the referenced envelope type."""
+    manifest = read_manifest(artifact)
+    spec = STEP_SPECS[CommandName(manifest["command"]["name"])]
+    expected = next((item for item in spec.artifact_inputs if item.argument == argument), None)
+    if expected is None:
+        raise ValueError(f"{argument!r} is not an artifact input of this producer")
+    entry = manifest.get("inputs", {}).get(argument)
+    if not isinstance(entry, dict) or not entry.get("path"):
+        raise ValueError(f"Producer must record inputs.{argument}.path")
+    path = Path(entry["path"])
+    if not path.is_absolute():
+        path = artifact.resolve() / path
+    path = path.resolve()
+    upstream = read_manifest(path)
+    if upstream["artifact"]["type"] != expected.artifact_type.value:
+        raise ValueError(f"{argument} has the wrong artifact type")
+    # Content, vocabulary, and algorithm-specific compatibility are still checked by the consumer.
+    return path
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("artifact", type=Path)
+    parser.add_argument("argument")
+    args = parser.parse_args()
+    print(resolve_input(args.artifact, args.argument))

meds_algorithm/artifacts/manifest.py

@@ -0,0 +1,142 @@
+"""Measured artifacts with explicit algorithm provenance and rollback-safe replacement."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import time
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from ..contracts.commands import STEP_SPECS, CommandName
+from ..runtime.provenance import code_provenance
+from ..runtime.resources import _ProcessTreeMemorySampler
+from ..version import CONTRACT_VERSION
+from .publication import staged_directory
+
+
+def read_manifest(directory: Path | str) -> dict[str, Any]:
+    path = Path(directory) / "manifest.yaml"
+    if not path.is_file():
+        raise FileNotFoundError(f"Artifact {directory} has no manifest.yaml")
+    value = yaml.safe_load(path.read_text())
+    if not isinstance(value, dict) or value.get("contract_version") != CONTRACT_VERSION:
+        raise ValueError(f"{path} is not a version 4 MEDS algorithm artifact manifest")
+    for field in ("artifact", "command", "code", "inputs", "resources"):
+        if not isinstance(value.get(field), dict):
+            raise ValueError(f"{path} requires a {field} mapping")
+    try:
+        spec = STEP_SPECS[CommandName(value["command"]["name"])]
+    except (KeyError, ValueError) as error:
+        raise ValueError(f"{path} has an unknown or missing command") from error
+    if value["artifact"].get("type") != spec.artifact_type.value:
+        raise ValueError(f"{path} artifact type does not match its command")
+    for field in ("producer", "payload_format"):
+        if not isinstance(value["artifact"].get(field), str) or not value["artifact"][field].strip():
+            raise ValueError(f"{path} requires artifact.{field}")
+    return value
+
+
+def digest_file(path: Path | str) -> str:
+    with Path(path).open("rb") as stream:
+        digest = hashlib.file_digest(stream, "sha256").hexdigest()
+    return f"sha256:{digest}"
+
+
+def digest_config(config: Any) -> str:
+    """Digest a command's configuration: its arguments plus any algorithm-owned settings it resolved.
+
+    The value is canonical JSON (sorted keys; paths and other objects as strings), so the same arguments
+    always give the same digest regardless of order.
+    """
+    canonical = json.dumps(config, sort_keys=True, default=str, separators=(",", ":"))
+    return f"sha256:{hashlib.sha256(canonical.encode()).hexdigest()}"
+
+
+@contextmanager
+def measured_artifact(
+    destination: Path | str,
+    *,
+    artifact_type: str,
+    command: str,
+    producer: str,
+    payload_format: str | Callable[[], str],
+    inputs: dict[str, Any] | None = None,
+    config_digest: str | None = None,
+    semantics_digest: str | None = None,
+    external_predicates_file: Path | str | None = None,
+    overwrite: bool = False,
+    repo: Path | str | None = None,
+    manifest_fields: dict[str, Any] | None = None,
+) -> Iterator[Path]:
+    """Measure staged output and publish its manifest with rollback on replacement failure."""
+    if STEP_SPECS[CommandName(command)].artifact_type.value != artifact_type:
+        raise ValueError("Artifact type does not match its command")
+    if (
+        not isinstance(producer, str)
+        or not producer.strip()
+        or (
+            not callable(payload_format)
+            and (not isinstance(payload_format, str) or not payload_format.strip())
+        )
+    ):
+        raise ValueError("producer and payload_format must be nonempty strings")
+    with staged_directory(destination, overwrite=overwrite) as staging:
+        started = time.monotonic()
+        sampler = _ProcessTreeMemorySampler()
+        sampler.start()
+        try:
+            yield staging
+        finally:
+            sampler.stop()
+        resolved_format = payload_format() if callable(payload_format) else payload_format
+        if not isinstance(resolved_format, str) or not resolved_format.strip():
+            raise ValueError("payload_format must resolve to a nonempty string")
+        predicates = None
+        if external_predicates_file is not None:
+            predicates = {
+                "path": str(Path(external_predicates_file).resolve()),
+                "digest": digest_file(external_predicates_file),
+            }
+        recorded_inputs = dict(inputs or {})
+        if predicates is not None:
+            recorded_inputs["external_predicates_file"] = predicates
+        manifest = {
+            "contract_version": CONTRACT_VERSION,
+            "artifact": {
+                "type": artifact_type,
+                "payload_format": resolved_format,
+                "producer": producer,
+            },
+            "command": {"name": command, "config_digest": config_digest},
+            "code": code_provenance(repo),
+            "inputs": recorded_inputs,
+            "semantics_digest": semantics_digest,
+            "resources": {
+                "wall_seconds": time.monotonic() - started,
+                "memory": {
+                    "peak_process_tree_pss_bytes": sampler.peak_pss_bytes,
+                    "peak_process_tree_rss_bytes": sampler.peak_rss_bytes,
+                    "sampling_interval_seconds": sampler.interval_seconds,
+                    "source": sampler.source,
+                    "includes_descendants": sampler.includes_descendants,
+                },
+                "gpu": {
+                    "peak_allocated_bytes": None,
+                    "peak_reserved_bytes": None,
+                    "source": None,
+                },
+                "measurement": {
+                    "schema_version": "1",
+                },
+            },
+        }
+        if manifest_fields:
+            if manifest.keys() & manifest_fields.keys():
+                raise ValueError("Extra manifest fields cannot replace the common envelope")
+            manifest.update(manifest_fields)
+        (staging / "manifest.yaml").write_text(yaml.safe_dump(manifest, sort_keys=False))

meds_algorithm/artifacts/publication.py

@@ -0,0 +1,68 @@
+"""Unique staging, exclusive writers, and recoverable directory replacement.
+
+New artifacts appear by one rename. Overwrite uses two renames and can briefly hide the
+destination from readers. A process crash may leave the backup/lock for manual recovery;
+we never silently delete another writer's staging area or a previous artifact.
+"""
+
+import os
+import shutil
+import tempfile
+import warnings
+from collections.abc import Iterator
+from contextlib import contextmanager
+from pathlib import Path
+from uuid import uuid4
+
+
+@contextmanager
+def staged_directory(destination: Path | str, *, overwrite: bool = False) -> Iterator[Path]:
+    destination = Path(destination).absolute()
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    lock = destination.with_name(f".{destination.name}.lock")
+    # Atomic acquisition also serializes writers in different threads/processes.
+    try:
+        fd = os.open(lock, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
+    except FileExistsError as error:
+        raise FileExistsError(
+            f"Artifact is locked: {lock}; inspect the writer before recovering a stale lock"
+        ) from error
+    staging = None
+    backup = None
+    try:
+        with os.fdopen(fd, "w") as stream:
+            stream.write(f"pid={os.getpid()}\n")
+        if destination.is_symlink() or (destination.exists() and not destination.is_dir()):
+            raise ValueError(f"Artifact destination must be a directory, not a file/symlink: {destination}")
+        if destination.exists() and not overwrite:
+            raise FileExistsError(f"{destination} already exists; pass overwrite=true to replace it")
+        staging = Path(tempfile.mkdtemp(prefix=f".{destination.name}.staging-", dir=destination.parent))
+        yield staging
+        if destination.exists():
+            backup = destination.with_name(f".{destination.name}.backup-{uuid4().hex}")
+            destination.replace(backup)
+        try:
+            staging.replace(destination)
+        except BaseException as publish_error:
+            if backup is not None:
+                try:
+                    backup.replace(destination)
+                except BaseException as restore_error:
+                    raise RuntimeError(
+                        f"Publication failed; previous artifact is preserved at {backup}; "
+                        f"restore also failed: {restore_error}"
+                    ) from publish_error
+            raise
+        if backup is not None:
+            try:
+                shutil.rmtree(backup)
+            except OSError as error:
+                warnings.warn(
+                    f"Published {destination}; retained old backup {backup}: {error}",
+                    RuntimeWarning,
+                    stacklevel=2,
+                )
+    finally:
+        if staging is not None and staging.exists():
+            shutil.rmtree(staging, ignore_errors=True)
+        lock.unlink(missing_ok=True)

meds_algorithm/artifacts/references.py

@@ -0,0 +1,79 @@
+"""Immutable, lazy references to persisted artifacts; no payload loading at construction."""
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import ClassVar
+
+from ..contracts.enums import ArtifactType
+from ..serialization import Serializable
+
+
+@dataclass(frozen=True)
+class ArtifactReference(Serializable):
+    directory: Path
+    kind: ClassVar[ArtifactType]
+
+    @classmethod
+    def open(cls, directory: Path | str):
+        reference = cls(Path(directory))
+        reference.read_manifest()
+        return reference
+
+    def read_manifest(self) -> dict:
+        from .manifest import read_manifest
+
+        manifest = read_manifest(self.directory)
+        if manifest["artifact"]["type"] != self.kind.value:
+            raise ValueError(f"Expected {self.kind.value} artifact: {self.directory}")
+        return manifest
+
+
+class PreparedDataArtifact(ArtifactReference):
+    kind = ArtifactType.patients
+
+
+class PretrainedAlgorithmArtifact(ArtifactReference):
+    kind = ArtifactType.pretrained_algorithm
+
+
+class EmbeddingsArtifact(ArtifactReference):
+    kind = ArtifactType.embeddings
+
+
+class TrajectoriesArtifact(ArtifactReference):
+    kind = ArtifactType.trajectories
+
+
+class SupervisedAlgorithmArtifact(ArtifactReference):
+    kind = ArtifactType.supervised_algorithm
+
+
+class FinetunedAlgorithmArtifact(ArtifactReference):
+    kind = ArtifactType.finetuned_algorithm
+
+
+class ProbeAlgorithmArtifact(ArtifactReference):
+    kind = ArtifactType.probe_algorithm
+
+
+class PredictionsArtifact(ArtifactReference):
+    kind = ArtifactType.predictions
+
+
+@dataclass(frozen=True)
+class ArtifactDestination(Serializable):
+    """Final destination for callers; runner substitutes its private staging path for algorithm code."""
+
+    directory: Path
+
+
+@dataclass(frozen=True)
+class PayloadMetadata(Serializable):
+    """Algorithm methods describe their payload; only the runner publishes the final artifact."""
+
+    format: str
+    version: int = 1
+
+    def validate(self):
+        if not self.format.strip() or self.version < 1:
+            raise ValueError("Payload format must be nonempty and version must be positive")

meds_algorithm/artifacts/results.py

@@ -0,0 +1,76 @@
+"""Structured, commit-safe run evidence assembled from command manifests."""
+
+from __future__ import annotations
+
+import os
+import platform
+import shutil
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from ..contracts.commands import ALGORITHM_PROFILES, STEP_SPECS, CommandName
+from ..runtime.provenance import code_provenance
+from .manifest import digest_file, read_manifest
+
+
+def write_run_result(
+    run_dir: Path | str,
+    *,
+    repo: Path | str,
+    profile: str,
+    artifacts: dict[str, Path | str],
+    dataset: dict[str, Any],
+    task: dict[str, Any],
+    external_predicates_file: Path | str | None = None,
+    status: str = "complete",
+) -> Path:
+    """Write `spec.yaml`, copied manifests, and a run-level `result.yaml`."""
+    if profile not in ALGORITHM_PROFILES or set(artifacts) - set(ALGORITHM_PROFILES[profile]):
+        raise ValueError("Run artifacts must belong to a declared contract profile")
+    run_dir = Path(run_dir)
+    repo = Path(repo).resolve()
+    manifests_dir = run_dir / "manifests"
+    manifests_dir.mkdir(parents=True, exist_ok=True)
+    manifest_refs = {}
+    for command, artifact in artifacts.items():
+        manifest = read_manifest(artifact)
+        if (
+            manifest["command"]["name"] != command
+            or manifest["artifact"]["type"] != STEP_SPECS[CommandName(command)].artifact_type.value
+        ):
+            raise ValueError(f"Run entry {command} disagrees with its artifact manifest")
+        destination = manifests_dir / f"{command}.yaml"
+        shutil.copy2(Path(artifact) / "manifest.yaml", destination)
+        manifest_refs[command] = {
+            "path": str(destination.relative_to(run_dir)),
+            "artifact_type": manifest["artifact"]["type"],
+            "resources": manifest["resources"],
+        }
+
+    predicates = None
+    if external_predicates_file is not None:
+        predicates = {
+            "path": str(Path(external_predicates_file).resolve()),
+            "digest": digest_file(external_predicates_file),
+        }
+    provenance = code_provenance(repo)
+    spec = {"profile": profile, "dataset": dataset, "task": task, "external_predicates_file": predicates}
+    result = {
+        "profile": profile,
+        "status": status,
+        "source": provenance,
+        "environment": {
+            "hostname": platform.node(),
+            "platform": platform.platform(),
+            "python": platform.python_version(),
+            "cpu_count": os.cpu_count(),
+            "slurm_job_id": os.environ.get("SLURM_JOB_ID"),
+        },
+        "commands": manifest_refs,
+    }
+    (run_dir / "spec.yaml").write_text(yaml.safe_dump(spec, sort_keys=False))
+    output = run_dir / "result.yaml"
+    output.write_text(yaml.safe_dump(result, sort_keys=False))
+    return output