podstack 1.2.0__py3-none-any.whl

podstack/registry/__init__.py

@@ -0,0 +1,402 @@
+"""
+Podstack Registry - Experiment Tracking and Model Management
+
+Usage:
+    from podstack import registry
+
+    # Initialize
+    registry.init(api_key="your-api-key", project_id="your-project-id")
+
+    # Set experiment
+    registry.set_experiment("my-experiment")
+
+    # Track a run
+    with registry.start_run(name="training-v1") as run:
+        registry.log_params({"learning_rate": 0.001, "batch_size": 32})
+        for epoch in range(10):
+            loss = train_epoch()
+            registry.log_metrics({"loss": loss}, step=epoch)
+
+    # Register a model
+    registry.register_model(
+        name="my-model",
+        run_id=run.id,
+        description="My trained model"
+    )
+"""
+
+from .client import RegistryClient
+from .experiment import Experiment, Run
+from .model import RegisteredModel, ModelVersion, ModelAlias, StageTransition
+from .exceptions import (
+    RegistryError,
+    ExperimentNotFoundError,
+    RunNotFoundError,
+    ModelNotFoundError,
+    NoActiveRunError,
+    NoExperimentSetError,
+    ArtifactNotFoundError,
+    ModelSerializationError,
+    FrameworkNotInstalledError,
+)
+
+__all__ = [
+    # Client
+    "RegistryClient",
+    "init",
+    "set_experiment",
+    "get_experiment",
+    "list_experiments",
+    "start_run",
+    "end_run",
+    "log_params",
+    "log_metrics",
+    "log_artifact",
+    "set_tag",
+    "register_model",
+    "get_model",
+    "list_models",
+    "set_model_stage",
+    "set_model_alias",
+    # New methods
+    "log_model",
+    "load_model",
+    "log_dataset",
+    "compare_runs",
+    "get_metric_history",
+    "download_artifact",
+    "search_runs",
+    # Classes
+    "Experiment",
+    "Run",
+    "RegisteredModel",
+    "ModelVersion",
+    "ModelAlias",
+    "StageTransition",
+    # Exceptions
+    "RegistryError",
+    "ExperimentNotFoundError",
+    "RunNotFoundError",
+    "ModelNotFoundError",
+    "NoActiveRunError",
+    "NoExperimentSetError",
+    "ArtifactNotFoundError",
+    "ModelSerializationError",
+    "FrameworkNotInstalledError",
+]
+
+# Global client instance
+_client: RegistryClient = None
+
+
+def init(
+    api_url: str = None,
+    api_key: str = None,
+    project_id: str = None
+) -> RegistryClient:
+    """
+    Initialize the registry client.
+
+    Args:
+        api_url: Registry service URL. Defaults to PODSTACK_REGISTRY_URL env var.
+        api_key: API key for authentication. Defaults to PODSTACK_API_KEY env var.
+        project_id: Project ID for multi-tenant isolation. Defaults to PODSTACK_PROJECT_ID env var.
+
+    Returns:
+        RegistryClient instance
+    """
+    global _client
+    _client = RegistryClient(
+        api_url=api_url,
+        api_key=api_key,
+        project_id=project_id
+    )
+    return _client
+
+
+def _get_client() -> RegistryClient:
+    """Get the global client instance, initializing if needed."""
+    global _client
+    if _client is None:
+        _client = RegistryClient()
+    return _client
+
+
+def set_experiment(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
+    """
+    return _get_client().set_experiment(name, description)
+
+
+def get_experiment(experiment_id: str) -> Experiment:
+    """Get an experiment by ID."""
+    return _get_client().get_experiment(experiment_id)
+
+
+def list_experiments(limit: int = 20, offset: int = 0) -> list:
+    """List experiments in the current project."""
+    return _get_client().list_experiments(limit, offset)
+
+
+def start_run(name: str = None, tags: dict = None) -> Run:
+    """
+    Start a new run in the active experiment.
+
+    Use as a context manager:
+        with registry.start_run(name="training") as run:
+            registry.log_params({"lr": 0.001})
+            registry.log_metrics({"loss": 0.5})
+
+    Args:
+        name: Optional run name
+        tags: Optional tags dict
+
+    Returns:
+        Run object (context manager)
+    """
+    return _get_client().start_run(name, tags)
+
+
+def end_run(status: str = "completed"):
+    """
+    End the active run.
+
+    Args:
+        status: Run status (completed, failed, killed)
+    """
+    _get_client().end_run(status)
+
+
+def log_params(params: dict):
+    """
+    Log parameters for the active run.
+
+    Args:
+        params: Dict of parameter names to values
+    """
+    _get_client().log_params(params)
+
+
+def log_metrics(metrics: dict, step: int = None):
+    """
+    Log metrics for the active run.
+
+    Args:
+        metrics: Dict of metric names to values
+        step: Optional step number
+    """
+    _get_client().log_metrics(metrics, step)
+
+
+def log_artifact(local_path: str, artifact_path: str = None):
+    """
+    Log an artifact file for the active run.
+
+    Args:
+        local_path: Local path to the file
+        artifact_path: Optional path within the artifact store
+    """
+    _get_client().log_artifact(local_path, artifact_path)
+
+
+def set_tag(key: str, value: str):
+    """
+    Set a tag on the active run.
+
+    Args:
+        key: Tag key
+        value: Tag value
+    """
+    _get_client().set_tag(key, value)
+
+
+def register_model(
+    name: str,
+    run_id: str = None,
+    description: str = None,
+    tags: dict = None
+) -> RegisteredModel:
+    """
+    Register a new model.
+
+    Args:
+        name: Model name
+        run_id: Optional run ID to link the model to
+        description: Optional description
+        tags: Optional tags dict
+
+    Returns:
+        RegisteredModel object
+    """
+    return _get_client().register_model(name, run_id, description, tags)
+
+
+def get_model(name_or_id: str) -> RegisteredModel:
+    """Get a registered model by name or ID."""
+    return _get_client().get_model(name_or_id)
+
+
+def list_models(limit: int = 20, offset: int = 0) -> list:
+    """List registered models in the current project."""
+    return _get_client().list_models(limit, offset)
+
+
+def set_model_stage(
+    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
+    """
+    return _get_client().set_model_stage(model_name, version, stage, comment)
+
+
+def set_model_alias(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 to alias
+
+    Returns:
+        ModelAlias object
+    """
+    return _get_client().set_model_alias(model_name, alias, version)
+
+
+def log_model(model, artifact_path: str = "model", framework: str = None, metadata: dict = None):
+    """
+    Serialize and upload a model as an artifact for the active run.
+
+    Args:
+        model: The model object to save.
+        artifact_path: Path within the artifact store (default: "model").
+        framework: Framework name. Auto-detected if not provided.
+        metadata: Optional metadata dict.
+    """
+    _get_client().log_model(model, artifact_path, framework, metadata)
+
+
+def load_model(model_name: str, version: int = None, stage: str = None, framework: str = None):
+    """
+    Download and deserialize a model from the registry.
+
+    Args:
+        model_name: Registered model name.
+        version: Version number to load.
+        stage: Stage to load from (e.g., "production").
+        framework: Framework name for deserialization.
+
+    Returns:
+        The loaded model object.
+    """
+    return _get_client().load_model(model_name, version, stage, framework)
+
+
+def log_dataset(
+    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.
+
+    Args:
+        name: Dataset name.
+        path: Dataset path or URI.
+        version: Dataset version string.
+        description: Dataset description.
+        digest: Hash/digest of the dataset.
+        num_rows: Number of rows/samples.
+        num_features: Number of features/columns.
+    """
+    _get_client().log_dataset(name, path, version, description, digest, num_rows, num_features)
+
+
+def compare_runs(run_ids: list, metric_keys: list = None) -> dict:
+    """
+    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 comparison data.
+    """
+    return _get_client().compare_runs(run_ids, metric_keys)
+
+
+def get_metric_history(run_id: str, metric_key: str) -> list:
+    """
+    Get the full history of a metric across all steps for a run.
+
+    Args:
+        run_id: Run ID.
+        metric_key: Metric key.
+
+    Returns:
+        List of Metric objects.
+    """
+    return _get_client().get_metric_history(run_id, metric_key)
+
+
+def download_artifact(run_id: str, artifact_path: str, local_path: str) -> str:
+    """
+    Download an artifact file from a run.
+
+    Args:
+        run_id: Run ID.
+        artifact_path: Artifact path within the run.
+        local_path: Local directory to save to.
+
+    Returns:
+        Local path to the downloaded file.
+    """
+    return _get_client().download_artifact(run_id, artifact_path, local_path)
+
+
+def search_runs(
+    experiment_id: str = None,
+    status: str = None,
+    max_results: int = 100,
+    offset: int = 0
+) -> list:
+    """
+    Search / list runs.
+
+    Args:
+        experiment_id: Filter by experiment ID.
+        status: Filter by status (running, completed, failed, cancelled).
+        max_results: Maximum number of results.
+        offset: Pagination offset.
+
+    Returns:
+        List of matching Run objects.
+    """
+    return _get_client().search_runs(experiment_id, status, max_results, offset)