podstack 1.2.0__py3-none-any.whl

podstack/registry/client.py

@@ -0,0 +1,957 @@
+"""
+Registry Client
+
+Client for interacting with the Podstack Registry API.
+"""
+
+import os
+import time
+import tempfile
+import shutil
+from typing import Optional, Dict, Any, List
+import requests
+
+from .experiment import Experiment, Run, Metric, Param
+from .model import RegisteredModel, ModelVersion, ModelAlias, StageTransition
+from .exceptions import (
+    RegistryError,
+    ExperimentNotFoundError,
+    RunNotFoundError,
+    ModelNotFoundError,
+    NoActiveRunError,
+    NoExperimentSetError,
+    InvalidStageError,
+    ModelVersionNotFoundError,
+    ArtifactNotFoundError,
+    ModelSerializationError,
+)
+from ..exceptions import AuthenticationError
+
+
+class RegistryClient:
+    """
+    Podstack Registry Client for experiment tracking and model management.
+
+    Usage:
+        client = RegistryClient(
+            api_url="https://cloud.podstack.ai/registry",
+            api_key="your-api-key",
+            project_id="your-project-id"
+        )
+
+        # Set experiment and track runs
+        client.set_experiment("my-experiment")
+        with client.start_run(name="training") as run:
+            client.log_params({"lr": 0.001})
+            client.log_metrics({"loss": 0.5}, step=1)
+
+        # Register and manage models
+        model = client.register_model("my-model", run_id=run.id)
+        client.set_model_stage("my-model", version=1, stage="production")
+    """
+
+    DEFAULT_API_URL = "https://cloud.podstack.ai/registry"
+
+    def __init__(
+        self,
+        api_url: str = None,
+        api_key: str = None,
+        project_id: str = None,
+        timeout: int = 30
+    ):
+        """
+        Initialize the registry client.
+
+        Args:
+            api_url: Registry API URL. Defaults to PODSTACK_REGISTRY_URL env var.
+            api_key: API key (psk_* token). Defaults to PODSTACK_API_KEY env var.
+            project_id: Project ID. Defaults to PODSTACK_PROJECT_ID env var.
+            timeout: Request timeout in seconds.
+        """
+        self.api_url = api_url or os.getenv("PODSTACK_REGISTRY_URL", self.DEFAULT_API_URL)
+        self.api_key = api_key or os.getenv("PODSTACK_API_KEY")
+        self.project_id = project_id or os.getenv("PODSTACK_PROJECT_ID")
+        self.timeout = timeout
+
+        # State
+        self._experiment_id: Optional[str] = None
+        self._experiment: Optional[Experiment] = None
+        self._active_run: Optional[Run] = None
+
+    def _get_headers(self) -> Dict[str, str]:
+        """Get request headers with auth and project ID."""
+        headers = {
+            "Content-Type": "application/json",
+        }
+        if self.api_key:
+            headers["Authorization"] = f"Bearer {self.api_key}"
+        if self.project_id:
+            headers["X-Project-ID"] = self.project_id
+        return headers
+
+    def _request(
+        self,
+        method: str,
+        endpoint: str,
+        json: Dict[str, Any] = None,
+        params: Dict[str, Any] = None,
+        files: Dict = None
+    ) -> Dict[str, Any]:
+        """
+        Make an API request.
+
+        Args:
+            method: HTTP method
+            endpoint: API endpoint path
+            json: JSON body
+            params: Query parameters
+            files: Files to upload
+
+        Returns:
+            Response data as dict
+        """
+        url = f"{self.api_url}/api/v1{endpoint}"
+        headers = self._get_headers()
+
+        # Remove Content-Type for file uploads
+        if files:
+            del headers["Content-Type"]
+
+        try:
+            response = requests.request(
+                method=method,
+                url=url,
+                headers=headers,
+                json=json,
+                params=params,
+                files=files,
+                timeout=self.timeout
+            )
+
+            if response.status_code == 401:
+                raise AuthenticationError("Invalid or expired API key")
+            elif response.status_code == 404:
+                error_msg = response.json().get("error", "Not found")
+                raise RegistryError(error_msg, code="not_found")
+            elif response.status_code >= 400:
+                error_data = response.json() if response.content else {}
+                error_msg = error_data.get("error", f"Request failed: {response.status_code}")
+                raise RegistryError(error_msg)
+
+            if response.status_code == 204 or not response.content:
+                return {}
+            return response.json()
+
+        except requests.exceptions.Timeout:
+            raise RegistryError("Request timed out", code="timeout")
+        except requests.exceptions.ConnectionError as e:
+            raise RegistryError(f"Connection error: {e}", code="connection_error")
+
+    # ==================== Experiment Methods ====================
+
+    def set_experiment(self, name: str, description: str = None) -> Experiment:
+        """
+        Set the active experiment. Creates if doesn't exist.
+
+        Args:
+            name: Experiment name
+            description: Optional description
+
+        Returns:
+            Experiment object
+        """
+        # Try to get existing experiment
+        try:
+            data = self._request("GET", f"/experiments/name/{name}")
+            self._experiment = Experiment.from_dict(data)
+        except RegistryError:
+            # Create new experiment
+            data = self._request("POST", "/experiments", json={
+                "name": name,
+                "description": description or ""
+            })
+            exp_data = data.get("experiment", data)
+            self._experiment = Experiment.from_dict(exp_data)
+
+        self._experiment_id = self._experiment.id
+        return self._experiment
+
+    def get_experiment(self, experiment_id: str) -> Experiment:
+        """Get an experiment by ID."""
+        try:
+            data = self._request("GET", f"/experiments/{experiment_id}")
+            return Experiment.from_dict(data)
+        except RegistryError as e:
+            if e.code == "not_found":
+                raise ExperimentNotFoundError(experiment_id)
+            raise
+
+    def list_experiments(self, limit: int = 20, offset: int = 0) -> List[Experiment]:
+        """List experiments in the current project."""
+        data = self._request("GET", "/experiments", params={
+            "limit": limit,
+            "offset": offset
+        })
+        return [Experiment.from_dict(e) for e in data.get("experiments", [])]
+
+    def archive_experiment(self, experiment_id: str):
+        """Archive an experiment."""
+        self._request("DELETE", f"/experiments/{experiment_id}")
+
+    # ==================== Run Methods ====================
+
+    def start_run(self, name: str = None, tags: dict = None) -> Run:
+        """
+        Start a new run in the active experiment.
+
+        Args:
+            name: Optional run name
+            tags: Optional tags dict
+
+        Returns:
+            Run object (can be used as context manager)
+        """
+        if not self._experiment_id:
+            raise NoExperimentSetError()
+
+        data = self._request("POST", "/runs", json={
+            "experiment_id": self._experiment_id,
+            "name": name or f"run-{int(time.time())}",
+            "tags": tags or {}
+        })
+        run_data = data.get("run", data)
+        self._active_run = Run.from_dict(run_data, client=self)
+        return self._active_run
+
+    def get_run(self, run_id: str) -> Run:
+        """Get a run by ID."""
+        try:
+            data = self._request("GET", f"/runs/{run_id}")
+            return Run.from_dict(data, client=self)
+        except RegistryError as e:
+            if e.code == "not_found":
+                raise RunNotFoundError(run_id)
+            raise
+
+    def list_runs(
+        self,
+        experiment_id: str = None,
+        status: str = None,
+        limit: int = 20,
+        offset: int = 0
+    ) -> List[Run]:
+        """
+        List runs.
+
+        Args:
+            experiment_id: Filter by experiment ID
+            status: Filter by status (running, completed, failed, killed)
+            limit: Max results
+            offset: Offset for pagination
+
+        Returns:
+            List of Run objects
+        """
+        params = {"limit": limit, "offset": offset}
+        if experiment_id:
+            params["experiment_id"] = experiment_id
+        if status:
+            params["status"] = status
+
+        data = self._request("GET", "/runs", params=params)
+        return [Run.from_dict(r, client=self) for r in data.get("runs", [])]
+
+    def end_run(self, status: str = "completed"):
+        """
+        End the active run.
+
+        Args:
+            status: Run status (completed, failed, killed)
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        self._request("POST", f"/runs/{self._active_run.id}/end", json={
+            "status": status
+        })
+        self._active_run.status = status
+        self._active_run = None
+
+    def log_params(self, params: Dict[str, Any]):
+        """
+        Log parameters for the active run.
+
+        Args:
+            params: Dict of parameter names to values
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        # Convert all values to strings
+        str_params = {k: str(v) for k, v in params.items()}
+        self._request("POST", f"/runs/{self._active_run.id}/params", json={
+            "params": str_params
+        })
+
+    def log_metrics(self, metrics: Dict[str, float], step: int = None):
+        """
+        Log metrics for the active run.
+
+        Args:
+            metrics: Dict of metric names to values
+            step: Optional step number
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        timestamp = int(time.time() * 1000)
+        metrics_list = [
+            {"key": k, "value": float(v), "step": step or 0, "timestamp": timestamp}
+            for k, v in metrics.items()
+        ]
+        self._request("POST", f"/runs/{self._active_run.id}/metrics", json={
+            "metrics": metrics_list
+        })
+
+    @staticmethod
+    def _get_artifact_dir(run_id: str) -> str:
+        """Return the local artifact directory for a run."""
+        base = os.getenv(
+            "PODSTACK_ARTIFACT_DIR",
+            os.path.join(os.path.expanduser("~"), ".podstack", "artifacts"),
+        )
+        return os.path.join(base, run_id)
+
+    def log_artifact(self, local_path: str, artifact_path: str = None):
+        """
+        Log an artifact file for the active run.
+
+        Copies the file into a local artifact directory
+        (``~/.podstack/artifacts/<run_id>/``) and records the path as a param.
+        The backend registry service does not expose an artifact-upload
+        endpoint; artifact files live on the local / shared filesystem.
+
+        Args:
+            local_path: Local path to the file
+            artifact_path: Optional relative path within the artifact store
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        artifact_path = artifact_path or os.path.basename(local_path)
+        dest_dir = self._get_artifact_dir(self._active_run.id)
+        dest = os.path.join(dest_dir, artifact_path)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+        shutil.copy2(local_path, dest)
+
+        # Record the artifact reference as a param
+        self._request("POST", f"/runs/{self._active_run.id}/params", json={
+            "params": {f"_artifact.{artifact_path}": dest}
+        })
+
+    def set_tag(self, key: str, value: str):
+        """
+        Set a tag on the active run.
+
+        Tags are persisted as params with a ``_tag.`` prefix via
+        ``POST /runs/:id/params`` (the backend has no dedicated tags endpoint).
+        The in-memory ``run.tags`` dict is also updated for local access.
+
+        Args:
+            key: Tag key
+            value: Tag value
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        # Backend has no /runs/:id/tags endpoint; persist via params
+        self._request("POST", f"/runs/{self._active_run.id}/params", json={
+            "params": {f"_tag.{key}": str(value)}
+        })
+        self._active_run.tags[key] = value
+
+    def get_run_metrics(self, run_id: str) -> List[Metric]:
+        """Get all metrics for a run."""
+        data = self._request("GET", f"/runs/{run_id}/metrics")
+        return [Metric.from_dict(m) for m in data.get("metrics", [])]
+
+    def get_run_params(self, run_id: str) -> List[Param]:
+        """Get all parameters for a run."""
+        data = self._request("GET", f"/runs/{run_id}/params")
+        return [Param.from_dict(p) for p in data.get("params", [])]
+
+    # ==================== Model Registry Methods ====================
+
+    def register_model(
+        self,
+        name: str,
+        run_id: str = None,
+        description: str = None,
+        tags: dict = None
+    ) -> RegisteredModel:
+        """
+        Register a new model.
+
+        If ``run_id`` is provided a first model version linked to that run is
+        created automatically (the backend ``POST /models`` only accepts
+        *name*, *description* and *tags* — ``run_id`` belongs on the version).
+
+        Args:
+            name: Model name
+            run_id: Optional run ID — creates version 1 linked to this run
+            description: Optional description
+            tags: Optional tags dict
+
+        Returns:
+            RegisteredModel object
+        """
+        body = {"name": name}
+        if description:
+            body["description"] = description
+        if tags:
+            body["tags"] = tags
+
+        data = self._request("POST", "/models", json=body)
+        model_data = data.get("model", data)
+        model = RegisteredModel.from_dict(model_data, client=self)
+
+        # Auto-create version 1 when run_id is provided
+        if run_id:
+            artifact_dir = self._get_artifact_dir(run_id)
+            source = f"runs/{run_id}/artifacts/model"
+            if os.path.isdir(artifact_dir):
+                source = artifact_dir
+            self.create_model_version(
+                model.id, run_id=run_id, source=source
+            )
+
+        return model
+
+    def get_model(self, name_or_id: str) -> RegisteredModel:
+        """
+        Get a model by name or ID.
+
+        Args:
+            name_or_id: Model name or ID
+
+        Returns:
+            RegisteredModel object
+        """
+        # Try by ID first
+        try:
+            data = self._request("GET", f"/models/{name_or_id}")
+            model_data = data.get("model", data)
+            return RegisteredModel.from_dict(model_data, client=self)
+        except RegistryError:
+            pass
+
+        # Try by name
+        try:
+            data = self._request("GET", f"/models/name/{name_or_id}")
+            model_data = data.get("model", data)
+            return RegisteredModel.from_dict(model_data, client=self)
+        except RegistryError as e:
+            if e.code == "not_found":
+                raise ModelNotFoundError(name_or_id)
+            raise
+
+    def list_models(self, limit: int = 20, offset: int = 0) -> List[RegisteredModel]:
+        """List registered models."""
+        data = self._request("GET", "/models", params={
+            "limit": limit,
+            "offset": offset
+        })
+        return [RegisteredModel.from_dict(m, client=self) for m in data.get("models", [])]
+
+    def delete_model(self, model_id: str):
+        """Delete a model."""
+        self._request("DELETE", f"/models/{model_id}")
+
+    def create_model_version(
+        self,
+        model_id: str,
+        run_id: str = None,
+        source: str = None,
+        description: str = None
+    ) -> ModelVersion:
+        """
+        Create a new version of a model.
+
+        Args:
+            model_id: Model ID
+            run_id: Optional run ID to link
+            source: Optional source path
+            description: Optional description
+
+        Returns:
+            ModelVersion object
+        """
+        body = {}
+        if run_id:
+            body["run_id"] = run_id
+        if source:
+            body["source"] = source
+        if description:
+            body["description"] = description
+
+        data = self._request("POST", f"/models/{model_id}/versions", json=body)
+        version_data = data.get("version", data)
+        return ModelVersion.from_dict(version_data, client=self)
+
+    def get_model_version(self, model_id: str, version: int) -> ModelVersion:
+        """Get a specific version of a model."""
+        try:
+            data = self._request("GET", f"/models/{model_id}/versions/{version}")
+            version_data = data.get("version", data)
+            return ModelVersion.from_dict(version_data, client=self)
+        except RegistryError as e:
+            if e.code == "not_found":
+                raise ModelVersionNotFoundError(model_id, version)
+            raise
+
+    def list_model_versions(
+        self,
+        model_id: str,
+        limit: int = 20,
+        offset: int = 0
+    ) -> List[ModelVersion]:
+        """List versions of a model."""
+        data = self._request("GET", f"/models/{model_id}/versions", params={
+            "limit": limit,
+            "offset": offset
+        })
+        return [ModelVersion.from_dict(v, client=self) for v in data.get("versions", [])]
+
+    def set_model_stage(
+        self,
+        model_name: str,
+        version: int,
+        stage: str,
+        comment: str = None
+    ) -> ModelVersion:
+        """
+        Transition a model version to a new stage.
+
+        Args:
+            model_name: Model name
+            version: Version number
+            stage: Target stage (development, staging, production, archived)
+            comment: Optional comment
+
+        Returns:
+            Updated ModelVersion
+        """
+        valid_stages = ["development", "staging", "production", "archived"]
+        if stage not in valid_stages:
+            raise InvalidStageError(stage)
+
+        # Get model by name to get ID
+        model = self.get_model(model_name)
+
+        body = {"stage": stage}
+        if comment:
+            body["comment"] = comment
+
+        data = self._request(
+            "PUT",
+            f"/models/{model.id}/versions/{version}/stage",
+            json=body
+        )
+        version_data = data.get("version", data)
+        return ModelVersion.from_dict(version_data, client=self)
+
+    def transition_model_stage(
+        self,
+        model_id: str,
+        version: int,
+        stage: str,
+        comment: str = None
+    ) -> ModelVersion:
+        """Transition a model version stage by model ID."""
+        valid_stages = ["development", "staging", "production", "archived"]
+        if stage not in valid_stages:
+            raise InvalidStageError(stage)
+
+        body = {"stage": stage}
+        if comment:
+            body["comment"] = comment
+
+        data = self._request(
+            "PUT",
+            f"/models/{model_id}/versions/{version}/stage",
+            json=body
+        )
+        version_data = data.get("version", data)
+        return ModelVersion.from_dict(version_data, client=self)
+
+    def set_model_alias(
+        self,
+        model_name: str,
+        alias: str,
+        version: int
+    ) -> ModelAlias:
+        """
+        Set an alias for a model version.
+
+        Args:
+            model_name: Model name
+            alias: Alias name (e.g., "champion", "challenger")
+            version: Version number
+
+        Returns:
+            ModelAlias object
+        """
+        model = self.get_model(model_name)
+        data = self._request("POST", f"/models/{model.id}/aliases", json={
+            "alias": alias,
+            "version": version
+        })
+        alias_data = data.get("alias", data)
+        return ModelAlias.from_dict(alias_data)
+
+    def get_model_aliases(self, model_id: str) -> List[ModelAlias]:
+        """Get all aliases for a model."""
+        data = self._request("GET", f"/models/{model_id}/aliases")
+        return [ModelAlias.from_dict(a) for a in data.get("aliases", [])]
+
+    def delete_model_alias(self, model_id: str, alias: str):
+        """Delete a model alias."""
+        self._request("DELETE", f"/models/{model_id}/aliases/{alias}")
+
+    def get_model_by_alias(self, model_id: str, alias: str) -> ModelVersion:
+        """Get a model version by alias."""
+        data = self._request("GET", f"/models/{model_id}/alias/{alias}")
+        version_data = data.get("version", data)
+        return ModelVersion.from_dict(version_data, client=self)
+
+    def get_stage_transitions(
+        self,
+        model_id: str,
+        version: int = None
+    ) -> List[StageTransition]:
+        """
+        Get stage transition history.
+
+        Args:
+            model_id: Model ID
+            version: Optional version number to filter
+
+        Returns:
+            List of StageTransition objects
+        """
+        if version:
+            endpoint = f"/models/{model_id}/versions/{version}/transitions"
+        else:
+            endpoint = f"/models/{model_id}/transitions"
+
+        data = self._request("GET", endpoint)
+        return [StageTransition.from_dict(t) for t in data.get("transitions", [])]
+
+    # ==================== Model Artifact Methods ====================
+
+    def log_model(
+        self,
+        model: Any,
+        artifact_path: str = "model",
+        framework: str = None,
+        metadata: Dict[str, str] = None
+    ):
+        """
+        Serialize a model to the local artifact directory for the active run.
+
+        The model is saved under ``~/.podstack/artifacts/<run_id>/<artifact_path>/``
+        and its metadata (framework, path) is recorded as run params so it can
+        be retrieved later with :meth:`load_model`.
+
+        Args:
+            model: The model object to save (PyTorch, TensorFlow, sklearn, HuggingFace, etc.)
+            artifact_path: Sub-path inside the artifact dir (default: "model")
+            framework: Framework name. Auto-detected if not provided.
+            metadata: Optional metadata dict stored as run params.
+
+        Raises:
+            NoActiveRunError: If no run is active.
+            ModelSerializationError: If model serialization fails.
+        """
+        from .model_utils import save_model, detect_framework
+
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        if framework is None:
+            framework = detect_framework(model)
+
+        # Serialize model into the persistent artifact directory
+        artifact_dir = self._get_artifact_dir(self._active_run.id)
+        model_dir = os.path.join(artifact_dir, artifact_path)
+        save_model(model, model_dir, framework)
+
+        # Record model metadata as params (backend supports POST /runs/:id/params)
+        model_params = {
+            "_model.framework": framework,
+            "_model.artifact_path": artifact_path,
+            "_model.local_dir": model_dir,
+        }
+        if metadata:
+            for key, value in metadata.items():
+                model_params[f"_model.{key}"] = str(value)
+
+        self.log_params(model_params)
+
+    def get_model_version_by_stage(
+        self,
+        model_id: str,
+        stage: str
+    ) -> ModelVersion:
+        """
+        Get the model version currently assigned to a stage.
+
+        Uses the dedicated ``GET /models/:id/stage/:stage`` backend endpoint.
+
+        Args:
+            model_id: Model ID.
+            stage: One of development, staging, production, archived.
+
+        Returns:
+            ModelVersion object.
+        """
+        data = self._request("GET", f"/models/{model_id}/stage/{stage}")
+        version_data = data.get("version", data)
+        return ModelVersion.from_dict(version_data, client=self)
+
+    def load_model(
+        self,
+        model_name: str,
+        version: int = None,
+        stage: str = None,
+        framework: str = None
+    ) -> Any:
+        """
+        Load a previously-saved model from the local artifact directory.
+
+        Resolves the model version (by explicit version number, stage via
+        ``GET /models/:id/stage/:stage``, or latest), reads framework metadata
+        from the run's params, then deserializes the model from the local
+        artifact directory.
+
+        Args:
+            model_name: Registered model name.
+            version: Version number to load. Mutually exclusive with *stage*.
+            stage: Stage to load from (e.g. "production"). Mutually exclusive with *version*.
+            framework: Framework name for deserialization.
+                       Read from run params if not provided.
+
+        Returns:
+            The loaded model object.
+
+        Raises:
+            ModelNotFoundError: If the model or version is not found.
+            ArtifactNotFoundError: If the local artifact directory is missing.
+        """
+        from .model_utils import load_model_from_path
+
+        # Resolve model and version
+        registered_model = self.get_model(model_name)
+
+        if stage:
+            # Use the dedicated backend endpoint
+            try:
+                model_version = self.get_model_version_by_stage(
+                    registered_model.id, stage
+                )
+            except RegistryError:
+                raise ModelNotFoundError(f"{model_name} (stage={stage})")
+        elif version is not None:
+            model_version = self.get_model_version(registered_model.id, version)
+        else:
+            model_version = self.get_model_version(
+                registered_model.id, registered_model.latest_version
+            )
+
+        # Read model metadata from run params
+        artifact_path = "model"
+        if model_version.run_id:
+            run_params = self.get_run_params(model_version.run_id)
+            params_dict = {p.key: p.value for p in run_params}
+            if not framework:
+                framework = params_dict.get("_model.framework", "pickle")
+            artifact_path = params_dict.get("_model.artifact_path", "model")
+
+        framework = framework or "pickle"
+
+        # Locate model on disk
+        if model_version.run_id:
+            model_dir = os.path.join(
+                self._get_artifact_dir(model_version.run_id), artifact_path
+            )
+        else:
+            model_dir = os.path.join(
+                self._get_artifact_dir(model_version.id), artifact_path
+            )
+
+        if not os.path.exists(model_dir):
+            raise ArtifactNotFoundError(
+                model_version.run_id or model_version.id, artifact_path
+            )
+
+        return load_model_from_path(model_dir, framework)
+
+    def log_dataset(
+        self,
+        name: str,
+        path: str = None,
+        version: str = None,
+        description: str = None,
+        digest: str = None,
+        num_rows: int = None,
+        num_features: int = None
+    ):
+        """
+        Log dataset metadata for the active run.
+
+        All metadata is stored as run params via ``POST /runs/:id/params``
+        using a ``dataset.`` prefix for easy retrieval.
+
+        Args:
+            name: Dataset name.
+            path: Dataset path or URI (e.g., "s3://bucket/data").
+            version: Dataset version string.
+            description: Dataset description.
+            digest: Hash/digest of the dataset for reproducibility.
+            num_rows: Number of rows/samples in the dataset.
+            num_features: Number of features/columns.
+
+        Raises:
+            NoActiveRunError: If no run is active.
+        """
+        if not self._active_run:
+            raise NoActiveRunError()
+
+        params = {"dataset.name": name}
+        if path:
+            params["dataset.path"] = path
+        if version:
+            params["dataset.version"] = version
+        if description:
+            params["dataset.description"] = description
+        if digest:
+            params["dataset.digest"] = digest
+        if num_rows is not None:
+            params["dataset.num_rows"] = str(num_rows)
+        if num_features is not None:
+            params["dataset.num_features"] = str(num_features)
+
+        self.log_params(params)
+
+    def compare_runs(
+        self,
+        run_ids: List[str],
+        metric_keys: List[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Compare multiple runs side by side.
+
+        Args:
+            run_ids: List of run IDs to compare.
+            metric_keys: Optional list of metric keys to include.
+
+        Returns:
+            Dict with structured comparison data including params and metrics
+            for each run.
+        """
+        body = {"run_ids": run_ids}
+        if metric_keys:
+            body["metric_keys"] = metric_keys
+
+        return self._request("POST", "/runs/compare", json=body)
+
+    def get_metric_history(
+        self,
+        run_id: str,
+        metric_key: str
+    ) -> List[Metric]:
+        """
+        Get the full history of a metric across all steps for a run.
+
+        Args:
+            run_id: Run ID.
+            metric_key: Metric key to retrieve history for.
+
+        Returns:
+            List of Metric objects ordered by step.
+        """
+        data = self._request("GET", f"/runs/{run_id}/metrics/{metric_key}/history")
+        return [Metric.from_dict(m) for m in data.get("metrics", [])]
+
+    def download_artifact(
+        self,
+        run_id: str,
+        artifact_path: str,
+        local_path: str
+    ) -> str:
+        """
+        Copy an artifact from the local artifact store to *local_path*.
+
+        Artifacts are stored on the local / shared filesystem under
+        ``~/.podstack/artifacts/<run_id>/``.  The backend registry service
+        does not expose an artifact-download endpoint, so this method reads
+        from the same directory that :meth:`log_artifact` / :meth:`log_model`
+        wrote to.
+
+        Args:
+            run_id: Run ID.
+            artifact_path: Relative path within the run's artifact directory.
+            local_path: Destination directory.
+
+        Returns:
+            Absolute path to the copied artifact.
+
+        Raises:
+            ArtifactNotFoundError: If the artifact is not on disk.
+        """
+        src = os.path.join(self._get_artifact_dir(run_id), artifact_path)
+
+        if not os.path.exists(src):
+            raise ArtifactNotFoundError(run_id, artifact_path)
+
+        os.makedirs(local_path, exist_ok=True)
+        dest = os.path.join(local_path, artifact_path)
+        os.makedirs(os.path.dirname(dest), exist_ok=True)
+
+        if os.path.isdir(src):
+            shutil.copytree(src, dest, dirs_exist_ok=True)
+        else:
+            shutil.copy2(src, dest)
+
+        return dest
+
+    def search_runs(
+        self,
+        experiment_id: str = None,
+        status: str = None,
+        max_results: int = 100,
+        offset: int = 0
+    ) -> List[Run]:
+        """
+        Search / list runs.
+
+        The backend ``GET /runs`` supports filtering by *experiment_id* and
+        *status* with *limit* / *offset* pagination.
+
+        Args:
+            experiment_id: Filter by experiment ID.
+            status: Filter by status (running, completed, failed, cancelled).
+            max_results: Maximum number of results (default 100).
+            offset: Pagination offset.
+
+        Returns:
+            List of matching Run objects.
+        """
+        params: Dict[str, Any] = {"limit": max_results, "offset": offset}
+        if experiment_id:
+            params["experiment_id"] = experiment_id
+        if status:
+            params["status"] = status
+
+        data = self._request("GET", "/runs", params=params)
+        return [Run.from_dict(r, client=self) for r in data.get("runs", [])]