methodic-research 0.1.2__py3-none-any.whl

methodic/runs.py

@@ -0,0 +1,306 @@
+"""Run lifecycle and worker-side asset uploads.
+
+`RunsAPI` is the stateless namespace; `Run` is a resource handle bound to a
+specific `(experiment_id, variation, run)` triple. Workers should use the
+handle (`chronicle.run(exp_id, v, r)`) — the convenience methods auto-populate
+the `output_of` field on uploaded assets and track in-flight async uploads.
+"""
+
+from __future__ import annotations
+
+import logging
+from concurrent.futures import Future, ThreadPoolExecutor
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from methodic.assets import AssetsAPI, AssetUploadInfo
+from methodic.transport import Transport
+
+if TYPE_CHECKING:
+    from methodic.upload_tracker import UploadTracker
+
+logger = logging.getLogger(__name__)
+
+
+class RunsAPI:
+    """Run-lifecycle namespace. Stateless across calls; takes the run triple as args."""
+
+    def __init__(
+        self, transport: Transport, assets: AssetsAPI, executor: ThreadPoolExecutor
+    ) -> None:
+        self._t = transport
+        self._assets = assets
+        self._executor = executor
+
+    @staticmethod
+    def _variation_path(experiment_id: str, variation: int) -> str:
+        return f"/experiments/{experiment_id}/variations/{variation}"
+
+    def _run_path(self, experiment_id: str, variation: int, run: int) -> str:
+        return f"{self._variation_path(experiment_id, variation)}/runs/{run}"
+
+    def get_variation_config(self, experiment_id: str, variation: int) -> dict[str, Any]:
+        path = self._variation_path(experiment_id, variation)
+        logger.info("Pulling variation config from %s", path)
+        return self._t.get(path)
+
+    def get_status(self, experiment_id: str, variation: int, run: int) -> Any:
+        return self._t.get(f"{self._run_path(experiment_id, variation, run)}/status")
+
+    def start(self, experiment_id: str, variation: int, run: int) -> None:
+        logger.info(
+            "Starting run %s/v%s/r%s", experiment_id, variation, run
+        )
+        self._t.post(f"{self._run_path(experiment_id, variation, run)}/start")
+
+    def succeed(self, experiment_id: str, variation: int, run: int) -> None:
+        logger.info("Run %s/v%s/r%s succeeded", experiment_id, variation, run)
+        self._t.post(f"{self._run_path(experiment_id, variation, run)}/succeed")
+
+    def fail(
+        self, experiment_id: str, variation: int, run: int, *, reason: str = "crash"
+    ) -> None:
+        """Mark a run failed. `reason` is `crash` (worker error) or `abandoned` (cancel)."""
+        logger.info(
+            "Run %s/v%s/r%s failed (reason=%s)", experiment_id, variation, run, reason
+        )
+        self._t.post(
+            f"{self._run_path(experiment_id, variation, run)}/fail", json={"reason": reason}
+        )
+
+    def heartbeat(self, experiment_id: str, variation: int, run: int) -> None:
+        self._t.post(
+            "/heartbeat",
+            json={"experiment_id": experiment_id, "variation": variation, "run": run},
+            timeout=10,
+        )
+
+
+class Run:
+    """Handle for a specific `(experiment_id, variation, run)` run.
+
+    Mutators return `self` so worker code can chain (`run.start().heartbeat()`).
+    Asset-upload helpers auto-populate `output_of` from the bound triple.
+    """
+
+    def __init__(self, api: RunsAPI, experiment_id: str, variation: int, run: int) -> None:
+        self._api = api
+        self.experiment_id = experiment_id
+        self.variation = variation
+        self.run = run
+        self._pending_uploads: list[Future] = []
+
+    @property
+    def _output_of(self) -> dict[str, Any]:
+        return {
+            "experiment_id": self.experiment_id,
+            "variation": self.variation,
+            "run": self.run,
+        }
+
+    # --- Lifecycle ---
+
+    def get_variation_config(self) -> dict[str, Any]:
+        return self._api.get_variation_config(self.experiment_id, self.variation)
+
+    def get_status(self) -> Any:
+        return self._api.get_status(self.experiment_id, self.variation, self.run)
+
+    def start(self) -> Run:
+        self._api.start(self.experiment_id, self.variation, self.run)
+        return self
+
+    def succeed(self) -> Run:
+        """Mark the run succeeded after waiting for any pending async uploads."""
+        self.wait_for_uploads()
+        self._api.succeed(self.experiment_id, self.variation, self.run)
+        return self
+
+    def fail(self, reason: str = "crash") -> Run:
+        self._api.fail(self.experiment_id, self.variation, self.run, reason=reason)
+        return self
+
+    def heartbeat(self) -> Run:
+        self._api.heartbeat(self.experiment_id, self.variation, self.run)
+        return self
+
+    # --- Asset uploads (output_of bound to this run) ---
+
+    def upload_asset(
+        self,
+        asset_type: str,
+        content: Any,
+        name: str | None = None,
+        content_type: str = "application/json",
+        asset_config: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Upload a small inline asset (auto-finalized). Linked to this run as output."""
+        return self._api._assets.create_inline(
+            asset_type=asset_type,
+            content=content,
+            name=name,
+            content_type=content_type,
+            output_of=self._output_of,
+            asset_config=asset_config,
+        )
+
+    def upload_environment(self, environment: dict[str, Any]) -> None:
+        self.upload_asset(asset_type="environment", content=environment)
+
+    def create_asset_presigned(
+        self,
+        asset_type: str,
+        components: list[str],
+        name: str | None = None,
+        content_type: str = "application/octet-stream",
+    ) -> AssetUploadInfo:
+        """Register a new asset for component upload via presigned URLs."""
+        return self._api._assets.create_with_presigned(
+            asset_type=asset_type,
+            components=components,
+            name=name,
+            content_type=content_type,
+            output_of=self._output_of,
+        )
+
+    def upload_component(
+        self, upload_url: str, local_path: Path, content_type: str
+    ) -> None:
+        self._api._assets.upload_component(upload_url, local_path, content_type)
+
+    def finalize_asset(self, asset_id: str) -> None:
+        self._api._assets.finalize(asset_id)
+
+    def get_asset(self, asset_id: str, include_presigned: bool = False) -> dict[str, Any]:
+        return self._api._assets.get(asset_id, include_presigned=include_presigned)
+
+    def presign_asset(
+        self, asset_id: str, operation: str = "read", components: list[str] | None = None
+    ) -> dict[str, Any]:
+        return self._api._assets.presign(asset_id, operation=operation, components=components)
+
+    def download_asset(self, asset_id: str, local_dir: Path) -> Path:
+        return self._api._assets.download(asset_id, local_dir)
+
+    def upload_file(
+        self,
+        local_path: Path,
+        asset_type: str,
+        content_type: str = "application/octet-stream",
+    ) -> str:
+        info = self.create_asset_presigned(
+            asset_type=asset_type,
+            components=[local_path.name],
+            content_type=content_type,
+        )
+        upload_url = info.upload_urls.get(local_path.name)
+        if upload_url:
+            self.upload_component(upload_url, local_path, content_type)
+        self.finalize_asset(info.asset_id)
+        return info.asset_id
+
+    # --- Async directory upload ---
+
+    def register_and_upload_async(
+        self,
+        local_dir: Path,
+        asset_type: str,
+        upload_tracker: UploadTracker,
+        content_type: str = "application/octet-stream",
+    ) -> str:
+        """Register every file in a directory, then upload + finalize on a background thread."""
+        files = [f for f in local_dir.rglob("*") if f.is_file()]
+        if not files:
+            return ""
+
+        components = [str(f.relative_to(local_dir)) for f in files]
+        info = self.create_asset_presigned(
+            asset_type=asset_type, components=components, content_type=content_type
+        )
+
+        upload_tracker.register_components(
+            info.asset_id,
+            asset_type,
+            [
+                (comp, str(local_dir / comp), (local_dir / comp).stat().st_size)
+                for comp in components
+            ],
+        )
+
+        future = self._api._executor.submit(
+            self._upload_registered_components,
+            info.asset_id,
+            info.upload_urls,
+            local_dir,
+            components,
+            content_type,
+            upload_tracker,
+        )
+        self._pending_uploads.append(future)
+        return info.asset_id
+
+    def _upload_registered_components(
+        self,
+        asset_id: str,
+        upload_urls: dict[str, str],
+        local_dir: Path,
+        components: list[str],
+        content_type: str,
+        tracker: UploadTracker,
+    ) -> None:
+        for comp in components:
+            upload_url = upload_urls.get(comp)
+            if not upload_url:
+                logger.warning("No upload URL for component %s, skipping", comp)
+                continue
+            tracker.mark_uploading(asset_id, comp)
+            self.upload_component(upload_url, local_dir / comp, content_type)
+            tracker.mark_completed(asset_id, comp)
+
+        if tracker.all_components_completed(asset_id):
+            self.finalize_asset(asset_id)
+            logger.info("Finalized asset %s (%d components)", asset_id, len(components))
+
+    def upload_directory_async(
+        self,
+        local_dir: Path,
+        asset_type: str,
+        content_type: str = "application/octet-stream",
+        upload_tracker: UploadTracker | None = None,
+    ) -> None:
+        """Upload every file in a directory on a background thread.
+
+        With `upload_tracker`, uses the register-then-upload flow for crash
+        recovery. Without, a thinner path that uploads and finalizes inline.
+        """
+        if upload_tracker:
+            self.register_and_upload_async(
+                local_dir, asset_type, upload_tracker, content_type
+            )
+            return
+
+        future = self._api._executor.submit(
+            self._upload_directory_no_tracking, local_dir, asset_type, content_type
+        )
+        self._pending_uploads.append(future)
+
+    def _upload_directory_no_tracking(
+        self, local_dir: Path, asset_type: str, content_type: str
+    ) -> None:
+        files = [f for f in local_dir.rglob("*") if f.is_file()]
+        if not files:
+            return
+        components = [str(f.relative_to(local_dir)) for f in files]
+        info = self.create_asset_presigned(
+            asset_type=asset_type, components=components, content_type=content_type
+        )
+        for comp in components:
+            upload_url = info.upload_urls.get(comp)
+            if upload_url:
+                self.upload_component(upload_url, local_dir / comp, content_type)
+        self.finalize_asset(info.asset_id)
+
+    def wait_for_uploads(self) -> None:
+        for future in self._pending_uploads:
+            future.result()
+        self._pending_uploads.clear()

methodic/search.py

@@ -0,0 +1,78 @@
+"""Search namespace.
+
+Wraps `POST /search`, which is backed by Vertex AI Search server-side. The
+endpoint returns 503 in environments without Vertex configured (local dev,
+CI without credentials) — callers that want to gate against this should
+catch `methodic.errors.ServerError` and check `status_code == 503`.
+
+RBAC and storage-prefix isolation are applied server-side: callers don't
+need to filter for them. `SearchFilters` is layered on top of those server
+filters for narrowing by asset type, time window, source, etc.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Iterator
+
+from methodic.transport import Transport
+from methodic.types import SearchFilters, SearchResponse, SearchResult
+
+if TYPE_CHECKING:
+    from methodic.chronicle import Chronicle
+
+logger = logging.getLogger(__name__)
+
+
+class SearchAPI:
+    """Vertex-backed search across research docs, experiment metadata, and arxiv assets."""
+
+    def __init__(self, transport: Transport, chronicle: Chronicle) -> None:
+        self._t = transport
+        self._chronicle = chronicle
+
+    def query(
+        self,
+        query: str,
+        *,
+        filters: SearchFilters | dict[str, Any] | None = None,
+        experiment_context: list[str] | None = None,
+        page_size: int | None = None,
+        page_token: str | None = None,
+    ) -> SearchResponse:
+        """Run a single search request. Returns one page; use `iter` to walk pages."""
+        payload: dict[str, Any] = {"query": query}
+        if filters is not None:
+            payload["filters"] = (
+                filters.to_dict() if isinstance(filters, SearchFilters) else filters
+            )
+        if experiment_context is not None:
+            payload["experiment_context"] = experiment_context
+        if page_size is not None:
+            payload["page_size"] = page_size
+        if page_token is not None:
+            payload["page_token"] = page_token
+        return SearchResponse.from_dict(self._t.post("/search", json=payload))
+
+    def iter(
+        self,
+        query: str,
+        *,
+        filters: SearchFilters | dict[str, Any] | None = None,
+        experiment_context: list[str] | None = None,
+        page_size: int | None = None,
+    ) -> Iterator[SearchResult]:
+        """Yield every search hit, paging server-side as needed."""
+        token: str | None = None
+        while True:
+            page = self.query(
+                query,
+                filters=filters,
+                experiment_context=experiment_context,
+                page_size=page_size,
+                page_token=token,
+            )
+            yield from page.results
+            if page.next_page_token is None:
+                return
+            token = page.next_page_token

methodic/transport.py

@@ -0,0 +1,91 @@
+"""HTTP transport shared by every namespace.
+
+Namespaces never call `requests` directly — they go through `Transport`, which
+owns the connection pool, attaches the auth header, and translates non-2xx
+responses into the right `APIError` subclass.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import requests
+
+from methodic.errors import raise_for_response
+
+logger = logging.getLogger(__name__)
+
+
+class Transport:
+    """Tiny wrapper over `requests.Session` for the Chronicle REST API."""
+
+    def __init__(self, server_url: str, api_key: str, timeout: int = 30) -> None:
+        self.base = server_url.rstrip("/")
+        self._api_key = api_key
+        self._timeout = timeout
+        self._session = requests.Session()
+
+    @property
+    def auth_headers(self) -> dict[str, str]:
+        return {"Authorization": f"Bearer {self._api_key}"}
+
+    def _url(self, path: str) -> str:
+        return f"{self.base}{path}" if path.startswith("/") else f"{self.base}/{path}"
+
+    def get(self, path: str, *, params: dict | None = None, timeout: int | None = None) -> Any:
+        resp = self._session.get(
+            self._url(path),
+            params=params,
+            headers=self.auth_headers,
+            timeout=timeout or self._timeout,
+        )
+        raise_for_response(resp)
+        return resp.json() if resp.content else None
+
+    def post(self, path: str, *, json: Any = None, timeout: int | None = None) -> Any:
+        resp = self._session.post(
+            self._url(path),
+            json=json,
+            headers=self.auth_headers,
+            timeout=timeout or self._timeout,
+        )
+        raise_for_response(resp)
+        return resp.json() if resp.content else None
+
+    def put(self, path: str, *, json: Any = None, timeout: int | None = None) -> Any:
+        resp = self._session.put(
+            self._url(path),
+            json=json,
+            headers=self.auth_headers,
+            timeout=timeout or self._timeout,
+        )
+        raise_for_response(resp)
+        return resp.json() if resp.content else None
+
+    def delete(self, path: str, *, timeout: int | None = None) -> Any:
+        resp = self._session.delete(
+            self._url(path),
+            headers=self.auth_headers,
+            timeout=timeout or self._timeout,
+        )
+        raise_for_response(resp)
+        return resp.json() if resp.content else None
+
+    def put_to_presigned(
+        self, url: str, *, data: Any, content_type: str, timeout: int = 600
+    ) -> None:
+        """PUT to a presigned cloud-storage URL. No auth header — the URL itself carries it."""
+        resp = self._session.put(
+            url, data=data, headers={"Content-Type": content_type}, timeout=timeout
+        )
+        resp.raise_for_status()
+
+    def get_streaming(self, url: str, *, timeout: int = 600) -> requests.Response:
+        """GET a presigned URL and return the raw streaming response. Caller closes."""
+        resp = self._session.get(url, stream=True, timeout=timeout)
+        resp.raise_for_status()
+        return resp
+
+    def close(self) -> None:
+        self._session.close()