methodic-research 0.1.2__py3-none-any.whl

methodic/types.py

@@ -0,0 +1,344 @@
+"""Dataclasses for Chronicle response payloads.
+
+Wire types use plain strings for UUIDs and ISO 8601 datetimes — same as
+the JSON they originate from. Callers that want strong typing can parse
+on their side; we don't bake `uuid.UUID` / `datetime.datetime` parsing
+into the SDK to keep dependencies minimal and round-tripping explicit.
+
+Each dataclass exposes a `from_dict` classmethod that ignores unknown
+keys, so server-side additions don't break older clients.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+T = TypeVar("T")
+
+
+def _from_dict(cls: type[T], data: dict[str, Any]) -> T:
+    """Build a dataclass from a dict, dropping keys the dataclass doesn't declare."""
+    field_names = {f.name for f in dataclasses.fields(cls)}  # type: ignore[arg-type]
+    return cls(**{k: v for k, v in data.items() if k in field_names})  # type: ignore[arg-type]
+
+
+@dataclass
+class Experiment:
+    """Mirror of the server's `Experiment` struct.
+
+    `git_repo_state` defaults to `"pending"` so older server payloads (which
+    may not include the field yet) deserialize cleanly.
+    """
+
+    id: str
+    owner_subject: str
+    hypothesis_summary: str
+    created_at: str
+    created_by: str
+    state: str
+    rationale: str | None = None
+    description: str | None = None
+    committed_at: str | None = None
+    concluded_at: str | None = None
+    retracted_at: str | None = None
+    retraction_reason: str | None = None
+    git_repo_state: str = "pending"
+    git_repo_url: str | None = None
+    git_repo_failure_reason: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Experiment:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class VariationSummary:
+    """One variation as it appears in `ExperimentDetail.variations`."""
+
+    variation: int
+    created_at: str
+    run_count: int
+    state: str
+    description: str | None = None
+    latest_status: str | None = None
+    committed_at: str | None = None
+    retracted_at: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> VariationSummary:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class Variation:
+    """Mirror of the server's `Variation` struct.
+
+    `config_json`, `accelerate_config_json`, and `launch_config` arrive as
+    arbitrary JSON — kept as `dict[str, Any]` since the schema is open.
+    """
+
+    experiment_id: str
+    variation: int
+    config_json: dict[str, Any]
+    config_yaml: str
+    created_at: str
+    created_by: str
+    state: str
+    accelerate_config_json: dict[str, Any] | None = None
+    accelerate_config_yaml: str | None = None
+    launch_config: dict[str, Any] | None = None
+    description: str | None = None
+    committed_at: str | None = None
+    retracted_at: str | None = None
+    retraction_reason: str | None = None
+    git_ref: str | None = None
+    git_sha: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> Variation:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class ExperimentDetail:
+    """`GET /experiments/{id}` response: experiment + parents + variation summaries."""
+
+    experiment: Experiment
+    parent_ids: list[str]
+    variations: list[VariationSummary]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ExperimentDetail:
+        return cls(
+            experiment=Experiment.from_dict(data["experiment"]),
+            parent_ids=list(data.get("parent_ids", [])),
+            variations=[VariationSummary.from_dict(v) for v in data.get("variations", [])],
+        )
+
+
+@dataclass
+class ExperimentSummary:
+    """One row in the experiments list."""
+
+    id: str
+    hypothesis_summary: str
+    variation_count: int
+    created_at: str
+    created_by: str
+    state: str
+    status: str | None = None
+    committed_at: str | None = None
+    concluded_at: str | None = None
+    retracted_at: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> ExperimentSummary:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class LineageResponse:
+    experiment_id: str
+    ancestors: list[Experiment]
+    descendants: list[Experiment]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> LineageResponse:
+        return cls(
+            experiment_id=data["experiment_id"],
+            ancestors=[Experiment.from_dict(e) for e in data.get("ancestors", [])],
+            descendants=[Experiment.from_dict(e) for e in data.get("descendants", [])],
+        )
+
+
+@dataclass
+class UpstreamRetraction:
+    experiment_id: str
+    retracted_at: str
+    reason: str
+    depth: int
+    variation: int | None = None
+    document_asset_id: str | None = None
+    chain: list[str] | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> UpstreamRetraction:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class UpstreamRetractionsResponse:
+    has_retractions: bool
+    retractions: list[UpstreamRetraction]
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> UpstreamRetractionsResponse:
+        return cls(
+            has_retractions=bool(data.get("has_retractions", False)),
+            retractions=[UpstreamRetraction.from_dict(r) for r in data.get("retractions", [])],
+        )
+
+
+@dataclass
+class CreateExperimentResponse:
+    """`POST /experiments` response: the new experiment plus the always-created variation 0 / run 0."""
+
+    experiment_id: str
+    variation: int
+    run: int
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> CreateExperimentResponse:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class GitBranch:
+    """One branch on the experiment repo, returned by `experiments.git_status`."""
+
+    name: str
+    head_sha: str
+    variation: int | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> GitBranch:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class GitToken:
+    """1-hour GitHub installation access token returned by `experiments.mint_git_token`.
+
+    `token` is what callers paste into git as
+    `https://x:<token>@github.com/<org>/<repo>`. Per the design, the
+    token has Administration permission stripped server-side, so it cannot
+    push to `agent/*` branches or modify branch protection — only the
+    Chronicle App can do those.
+    """
+
+    token: str
+    expires_at: str
+    repo_url: str
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> GitToken:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class GitStatus:
+    """Response from `GET /experiments/{id}/git`.
+
+    `state` mirrors `Experiment.git_repo_state`. When `state == "ready"`,
+    `repo_url` is populated; when `state == "failed"`, `failure_reason` is
+    populated. `branches` is empty until branch enumeration is wired up
+    server-side (Phase 2 follow-up).
+    """
+
+    state: str
+    default_branch: str = "main"
+    repo_url: str | None = None
+    failure_reason: str | None = None
+    branches: list[GitBranch] = dataclasses.field(default_factory=list)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> GitStatus:
+        return cls(
+            state=data["state"],
+            default_branch=data.get("default_branch", "main"),
+            repo_url=data.get("repo_url"),
+            failure_reason=data.get("failure_reason"),
+            branches=[GitBranch.from_dict(b) for b in data.get("branches", [])],
+        )
+
+
+@dataclass
+class SearchFilters:
+    """Filters layered on top of the RBAC + namespace filters the server adds."""
+
+    asset_types: list[str] | None = None
+    organization_id: str | None = None
+    team_id: str | None = None
+    created_after: str | None = None
+    created_before: str | None = None
+    created_by: str | None = None
+    source_type: str | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        out: dict[str, Any] = {}
+        if self.asset_types is not None:
+            out["asset_types"] = self.asset_types
+        if self.organization_id is not None:
+            out["organization_id"] = self.organization_id
+        if self.team_id is not None:
+            out["team_id"] = self.team_id
+        if self.created_after is not None:
+            out["created_after"] = self.created_after
+        if self.created_before is not None:
+            out["created_before"] = self.created_before
+        if self.created_by is not None:
+            out["created_by"] = self.created_by
+        if self.source_type is not None:
+            out["source_type"] = self.source_type
+        return out
+
+
+@dataclass
+class SearchResult:
+    """One hit from the Vertex-backed search."""
+
+    document_id: str
+    source_type: str
+    relevance_score: float
+    lineage_boost: bool
+    asset_type: str | None = None
+    title: str | None = None
+    snippet: str | None = None
+    experiment_ids: list[str] = dataclasses.field(default_factory=list)
+    created_at: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SearchResult:
+        return _from_dict(cls, data)
+
+
+@dataclass
+class SearchResponse:
+    results: list[SearchResult]
+    total_size: int
+    next_page_token: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> SearchResponse:
+        return cls(
+            results=[SearchResult.from_dict(r) for r in data.get("results", [])],
+            total_size=int(data.get("total_size", 0)),
+            next_page_token=data.get("next_page_token"),
+        )
+
+
+@dataclass
+class ExperimentListPage:
+    """One page of `experiments.list` results plus a cursor for the next page.
+
+    The current server returns a flat array; we normalize that into a
+    single-page response with `next_page_token=None`. When the server
+    grows pagination, the same dataclass keeps working.
+    """
+
+    results: list[ExperimentSummary]
+    next_page_token: str | None = None
+
+    @classmethod
+    def from_dict(cls, data: Any) -> ExperimentListPage:
+        if isinstance(data, list):
+            return cls(
+                results=[ExperimentSummary.from_dict(e) for e in data],
+                next_page_token=None,
+            )
+        return cls(
+            results=[ExperimentSummary.from_dict(e) for e in data.get("results", [])],
+            next_page_token=data.get("next_page_token"),
+        )

methodic/upload_tracker.py

@@ -0,0 +1,181 @@
+"""Local SQLite tracker for binary asset upload state.
+
+Shared by all producers (checkpoint manager, renderer). Enables:
+1. Crash recovery — on restart, find the latest checkpoint for training resume
+2. Upload coordination — background thread drains the queue
+Uses WAL mode for safe concurrent access. See docs/design.md for full design.
+"""
+
+import logging
+import sqlite3
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+_CREATE_TABLE = """
+CREATE TABLE IF NOT EXISTS upload_state (
+    asset_id     TEXT NOT NULL,
+    asset_type   TEXT NOT NULL,
+    component    TEXT NOT NULL,
+    local_path   TEXT NOT NULL,
+    size         INTEGER,
+    state        TEXT NOT NULL,
+    created_at   TEXT NOT NULL,
+    PRIMARY KEY (asset_id, component)
+)
+"""
+
+
+@dataclass
+class PendingUpload:
+    """A component that needs uploading."""
+
+    asset_id: str
+    asset_type: str
+    component: str
+    local_path: str
+    state: str
+    created_at: str
+
+
+class UploadTracker:
+    """Tracks upload state in a local SQLite database.
+
+    Thread-safe via WAL mode. Multiple producers can register files
+    concurrently; the background upload thread drains the queue.
+
+    The tracker is passive storage — it doesn't gate work. Atomicity at
+    the asset level (an asset is invisible until all components ship) is
+    enforced by the Chronicle backend's `pending → ready` finalize state
+    machine, not here. The tracker provides two things: batch-atomic
+    registration via SQLite transactions, and persistent state so a
+    restarted process can resume incomplete uploads.
+
+    Caller requirement: register all components for a given asset_id in a
+    SINGLE `register_components()` batch. Splitting registration across
+    calls races against `all_components_completed()` — that check returns
+    true once every row currently in the table is COMPLETED, with no
+    notion of "components still to come". A second batch registered after
+    the first has already finalized lands too late: the asset is already
+    `ready` with only the first batch's components. `register_and_upload_async()`
+    on `Client` handles this for you by registering the full component
+    list in one shot.
+    """
+
+    def __init__(self, db_path: str | Path = "upload_state.db") -> None:
+        self._conn = sqlite3.connect(str(db_path), check_same_thread=False)
+        self._conn.execute("PRAGMA journal_mode=WAL")
+        self._conn.execute(_CREATE_TABLE)
+        self._conn.commit()
+
+    def _now(self) -> str:
+        return datetime.now(timezone.utc).isoformat()
+
+    def register_components(
+        self,
+        asset_id: str,
+        asset_type: str,
+        components: list[tuple[str, str, int | None]],
+    ) -> None:
+        """Register all components for an asset as PENDING in a single batch.
+
+        Args:
+            asset_id: Chronicle asset UUID.
+            asset_type: 'checkpoint', 'snapshot', 'visualization', etc.
+            components: List of (component_name, local_path, size) tuples.
+        """
+        now = self._now()
+        self._conn.executemany(
+            "INSERT OR REPLACE INTO upload_state "
+            "(asset_id, asset_type, component, local_path, size, state, created_at) "
+            "VALUES (?, ?, ?, ?, ?, 'PENDING', ?)",
+            [(asset_id, asset_type, name, path, size, now) for name, path, size in components],
+        )
+        self._conn.commit()
+
+    def mark_pending(
+        self,
+        asset_id: str,
+        component: str,
+        local_path: str,
+        size: int | None = None,
+        asset_type: str = "unknown",
+    ) -> None:
+        """Register a single component as PENDING."""
+        self._conn.execute(
+            "INSERT OR REPLACE INTO upload_state "
+            "(asset_id, asset_type, component, local_path, size, state, created_at) "
+            "VALUES (?, ?, ?, ?, ?, 'PENDING', ?)",
+            (asset_id, asset_type, component, local_path, size, self._now()),
+        )
+        self._conn.commit()
+
+    def mark_uploading(self, asset_id: str, component: str) -> None:
+        self._conn.execute(
+            "UPDATE upload_state SET state = 'UPLOADING' "
+            "WHERE asset_id = ? AND component = ?",
+            (asset_id, component),
+        )
+        self._conn.commit()
+
+    def mark_completed(self, asset_id: str, component: str) -> None:
+        self._conn.execute(
+            "UPDATE upload_state SET state = 'COMPLETED' "
+            "WHERE asset_id = ? AND component = ?",
+            (asset_id, component),
+        )
+        self._conn.commit()
+
+    def get_incomplete(self) -> list[PendingUpload]:
+        """Return all rows that are not COMPLETED."""
+        cursor = self._conn.execute(
+            "SELECT asset_id, asset_type, component, local_path, state, created_at "
+            "FROM upload_state WHERE state != 'COMPLETED'"
+        )
+        return [
+            PendingUpload(
+                asset_id=r[0], asset_type=r[1], component=r[2],
+                local_path=r[3], state=r[4], created_at=r[5],
+            )
+            for r in cursor.fetchall()
+        ]
+
+    def get_latest_checkpoint_path(self) -> str | None:
+        """Get the local_path of the most recently registered checkpoint asset.
+
+        Returns the local_path of any component from the latest checkpoint
+        asset (by created_at). The checkpoint directory is the parent of
+        any component file. Returns None if no checkpoints are registered.
+
+        Used for in-container crash recovery — the latest registered checkpoint
+        is the last one fully flushed to disk.
+        """
+        cursor = self._conn.execute(
+            "SELECT local_path FROM upload_state "
+            "WHERE asset_type = 'checkpoint' "
+            "ORDER BY created_at DESC LIMIT 1"
+        )
+        row = cursor.fetchone()
+        if row is None:
+            return None
+        # The checkpoint directory is the parent of the component file
+        return str(Path(row[0]).parent)
+
+    def has_registered_assets(self) -> bool:
+        """Check if there are any registered assets (indicates a restart)."""
+        cursor = self._conn.execute("SELECT COUNT(*) FROM upload_state")
+        return cursor.fetchone()[0] > 0
+
+    def all_components_completed(self, asset_id: str) -> bool:
+        """Check if all components for an asset are in COMPLETED state."""
+        cursor = self._conn.execute(
+            "SELECT COUNT(*) FROM upload_state "
+            "WHERE asset_id = ? AND state != 'COMPLETED'",
+            (asset_id,),
+        )
+        return cursor.fetchone()[0] == 0
+
+    def close(self) -> None:
+        self._conn.close()

methodic/variations.py

@@ -0,0 +1,166 @@
+"""Variations namespace and resource handle."""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from methodic.transport import Transport
+from methodic.types import Variation as VariationData
+
+if TYPE_CHECKING:
+    from methodic.chronicle import Chronicle
+
+logger = logging.getLogger(__name__)
+
+
+class VariationsAPI:
+    """Variations namespace. Keys every operation on (experiment_id, variation)."""
+
+    def __init__(self, transport: Transport, chronicle: Chronicle) -> None:
+        self._t = transport
+        self._chronicle = chronicle
+
+    @staticmethod
+    def _path(experiment_id: str, variation: int | None = None) -> str:
+        base = f"/experiments/{experiment_id}/variations"
+        return f"{base}/{variation}" if variation is not None else base
+
+    def create(
+        self,
+        experiment_id: str,
+        *,
+        config_yaml: str,
+        accelerate_config_yaml: str | None = None,
+        launch_config: dict[str, Any] | None = None,
+        description: str | None = None,
+        input_asset_ids: list[str] | None = None,
+        git_ref: str | None = None,
+    ) -> Variation:
+        """Create a new variation under `experiment_id`. Returns a `Variation` handle.
+
+        `git_ref` optionally associates the variation with a branch on the
+        experiment's GitHub repo. Server captures the branch name now;
+        SHA resolution + the branch-rename-to-`agent/...` flow happens at
+        variation commit (Phase 3). Pre-Phase-3, registering with `git_ref`
+        is informational only.
+        """
+        payload: dict[str, Any] = {"config_yaml": config_yaml}
+        if accelerate_config_yaml is not None:
+            payload["accelerate_config_yaml"] = accelerate_config_yaml
+        if launch_config is not None:
+            payload["launch_config"] = launch_config
+        if description is not None:
+            payload["description"] = description
+        if input_asset_ids is not None:
+            payload["input_asset_ids"] = input_asset_ids
+        if git_ref is not None:
+            payload["git_ref"] = git_ref
+
+        resp = self._t.post(self._path(experiment_id), json=payload)
+        return Variation(self._chronicle, experiment_id, resp["variation"])
+
+    def get(self, experiment_id: str, variation: int) -> VariationData:
+        return VariationData.from_dict(self._t.get(self._path(experiment_id, variation)))
+
+    def commit(self, experiment_id: str, variation: int) -> dict[str, Any]:
+        return self._t.put(f"{self._path(experiment_id, variation)}/commit")
+
+    def retract(
+        self,
+        experiment_id: str,
+        variation: int,
+        *,
+        reason: str,
+        document_asset_id: str | None = None,
+    ) -> dict[str, Any]:
+        payload: dict[str, Any] = {"reason": reason}
+        if document_asset_id is not None:
+            payload["document_asset_id"] = document_asset_id
+        return self._t.put(f"{self._path(experiment_id, variation)}/retract", json=payload)
+
+
+class _BoundVariations:
+    """`exp.variations` — VariationsAPI pre-bound to one experiment id."""
+
+    def __init__(self, api: VariationsAPI, experiment_id: str) -> None:
+        self._api = api
+        self._experiment_id = experiment_id
+
+    def create(self, **kwargs: Any) -> Variation:
+        return self._api.create(self._experiment_id, **kwargs)
+
+    def get(self, variation: int) -> VariationData:
+        return self._api.get(self._experiment_id, variation)
+
+    def commit(self, variation: int) -> dict[str, Any]:
+        return self._api.commit(self._experiment_id, variation)
+
+    def retract(
+        self, variation: int, *, reason: str, document_asset_id: str | None = None
+    ) -> dict[str, Any]:
+        return self._api.retract(
+            self._experiment_id,
+            variation,
+            reason=reason,
+            document_asset_id=document_asset_id,
+        )
+
+
+class Variation:
+    """Handle for one variation. Holds (experiment_id, variation) and lazy-loaded data."""
+
+    def __init__(
+        self,
+        chronicle: Chronicle,
+        experiment_id: str,
+        variation: int,
+        *,
+        _data: VariationData | None = None,
+    ) -> None:
+        self._chronicle = chronicle
+        self.experiment_id = experiment_id
+        self.variation = variation
+        self._data = _data
+
+    @property
+    def data(self) -> VariationData:
+        """Server-side variation record. Auto-fetched on first access; refetched after mutations."""
+        if self._data is None:
+            self._data = self._chronicle.variations.get(self.experiment_id, self.variation)
+        return self._data
+
+    @property
+    def state(self) -> str:
+        return self.data.state
+
+    @property
+    def committed_at(self) -> str | None:
+        return self.data.committed_at
+
+    @property
+    def retracted_at(self) -> str | None:
+        return self.data.retracted_at
+
+    @property
+    def config_yaml(self) -> str:
+        return self.data.config_yaml
+
+    @property
+    def launch_config(self) -> dict[str, Any] | None:
+        return self.data.launch_config
+
+    def commit(self) -> Variation:
+        self._chronicle.variations.commit(self.experiment_id, self.variation)
+        self._data = None
+        return self
+
+    def retract(self, *, reason: str, document_asset_id: str | None = None) -> Variation:
+        self._chronicle.variations.retract(
+            self.experiment_id,
+            self.variation,
+            reason=reason,
+            document_asset_id=document_asset_id,
+        )
+        self._data = None
+        return self