podstack 1.2.0__py3-none-any.whl

podstack/registry/exceptions.py

@@ -0,0 +1,107 @@
+"""
+Registry Exception Classes
+
+Custom exceptions for the Podstack Registry SDK.
+"""
+
+from ..exceptions import PodstackError
+
+
+class RegistryError(PodstackError):
+    """Base exception for registry errors"""
+
+    def __init__(self, message: str, code: str = "registry_error"):
+        super().__init__(message, code=code)
+
+
+class ExperimentNotFoundError(RegistryError):
+    """Raised when an experiment is not found"""
+
+    def __init__(self, experiment_id: str):
+        message = f"Experiment '{experiment_id}' not found"
+        super().__init__(message, code="experiment_not_found")
+        self.experiment_id = experiment_id
+
+
+class RunNotFoundError(RegistryError):
+    """Raised when a run is not found"""
+
+    def __init__(self, run_id: str):
+        message = f"Run '{run_id}' not found"
+        super().__init__(message, code="run_not_found")
+        self.run_id = run_id
+
+
+class ModelNotFoundError(RegistryError):
+    """Raised when a model is not found"""
+
+    def __init__(self, model_id: str):
+        message = f"Model '{model_id}' not found"
+        super().__init__(message, code="model_not_found")
+        self.model_id = model_id
+
+
+class NoActiveRunError(RegistryError):
+    """Raised when trying to log without an active run"""
+
+    def __init__(self):
+        message = "No active run. Call start_run() first."
+        super().__init__(message, code="no_active_run")
+
+
+class NoExperimentSetError(RegistryError):
+    """Raised when trying to create a run without setting an experiment"""
+
+    def __init__(self):
+        message = "No experiment set. Call set_experiment() first."
+        super().__init__(message, code="no_experiment_set")
+
+
+class InvalidStageError(RegistryError):
+    """Raised when an invalid stage is specified"""
+
+    VALID_STAGES = ["development", "staging", "production", "archived"]
+
+    def __init__(self, stage: str):
+        message = f"Invalid stage '{stage}'. Must be one of: {', '.join(self.VALID_STAGES)}"
+        super().__init__(message, code="invalid_stage")
+        self.stage = stage
+
+
+class ModelVersionNotFoundError(RegistryError):
+    """Raised when a model version is not found"""
+
+    def __init__(self, model_name: str, version: int):
+        message = f"Version {version} of model '{model_name}' not found"
+        super().__init__(message, code="model_version_not_found")
+        self.model_name = model_name
+        self.version = version
+
+
+class ArtifactNotFoundError(RegistryError):
+    """Raised when a requested artifact cannot be found"""
+
+    def __init__(self, run_id: str, artifact_path: str):
+        message = f"Artifact '{artifact_path}' not found in run '{run_id}'"
+        super().__init__(message, code="artifact_not_found")
+        self.run_id = run_id
+        self.artifact_path = artifact_path
+
+
+class ModelSerializationError(RegistryError):
+    """Raised when model serialization or deserialization fails"""
+
+    def __init__(self, message: str):
+        super().__init__(message, code="model_serialization_error")
+
+
+class FrameworkNotInstalledError(RegistryError):
+    """Raised when a required ML framework is not installed"""
+
+    def __init__(self, framework: str):
+        message = (
+            f"Framework '{framework}' is not installed. "
+            f"Install it with: pip install podstack[{framework}]"
+        )
+        super().__init__(message, code="framework_not_installed")
+        self.framework = framework

podstack/registry/experiment.py

@@ -0,0 +1,227 @@
+"""
+Experiment and Run Classes
+
+Data classes for experiment tracking.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+
+
+@dataclass
+class Experiment:
+    """Represents an experiment in the registry."""
+
+    id: str
+    project_id: str
+    name: str
+    description: Optional[str] = None
+    artifact_location: Optional[str] = None
+    status: str = "active"
+    tags: Dict[str, str] = field(default_factory=dict)
+    run_count: int = 0
+    last_run_at: Optional[datetime] = None
+    created_by_user_id: Optional[str] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Experiment":
+        """Create an Experiment from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            project_id=data.get("project_id", ""),
+            name=data.get("name", ""),
+            description=data.get("description"),
+            artifact_location=data.get("artifact_location"),
+            status=data.get("status", "active"),
+            tags=data.get("tags") or {},
+            run_count=data.get("run_count", 0),
+            last_run_at=_parse_datetime(data.get("last_run_at")),
+            created_by_user_id=data.get("created_by_user_id"),
+            created_at=_parse_datetime(data.get("created_at")),
+            updated_at=_parse_datetime(data.get("updated_at")),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "project_id": self.project_id,
+            "name": self.name,
+            "description": self.description,
+            "artifact_location": self.artifact_location,
+            "status": self.status,
+            "tags": self.tags,
+            "run_count": self.run_count,
+            "last_run_at": self.last_run_at.isoformat() if self.last_run_at else None,
+            "created_by_user_id": self.created_by_user_id,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }
+
+
+@dataclass
+class Metric:
+    """Represents a logged metric."""
+
+    key: str
+    value: float
+    timestamp: Optional[int] = None
+    step: Optional[int] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Metric":
+        return cls(
+            key=data.get("key", ""),
+            value=data.get("value", 0.0),
+            timestamp=data.get("timestamp"),
+            step=data.get("step"),
+        )
+
+
+@dataclass
+class Param:
+    """Represents a logged parameter."""
+
+    key: str
+    value: str
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "Param":
+        return cls(
+            key=data.get("key", ""),
+            value=data.get("value", ""),
+        )
+
+
+@dataclass
+class Run:
+    """
+    Represents a run in an experiment.
+
+    Can be used as a context manager:
+        with registry.start_run(name="training") as run:
+            registry.log_params({"lr": 0.001})
+            registry.log_metrics({"loss": 0.5})
+    """
+
+    id: str
+    project_id: str
+    experiment_id: str
+    user_id: Optional[str] = None
+    name: Optional[str] = None
+    status: str = "running"
+    source_type: Optional[str] = None
+    source_name: Optional[str] = None
+    artifact_uri: Optional[str] = None
+    start_time: Optional[int] = None
+    end_time: Optional[int] = None
+    duration_ms: int = 0
+    tags: Dict[str, str] = field(default_factory=dict)
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    # Reference to client for context manager
+    _client: Any = field(default=None, repr=False)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], client: Any = None) -> "Run":
+        """Create a Run from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            project_id=data.get("project_id", ""),
+            experiment_id=data.get("experiment_id", ""),
+            user_id=data.get("user_id"),
+            name=data.get("name"),
+            status=data.get("status", "running"),
+            source_type=data.get("source_type"),
+            source_name=data.get("source_name"),
+            artifact_uri=data.get("artifact_uri"),
+            start_time=data.get("start_time"),
+            end_time=data.get("end_time"),
+            duration_ms=data.get("duration_ms", 0),
+            tags=data.get("tags") or {},
+            created_at=_parse_datetime(data.get("created_at")),
+            updated_at=_parse_datetime(data.get("updated_at")),
+            _client=client,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "project_id": self.project_id,
+            "experiment_id": self.experiment_id,
+            "user_id": self.user_id,
+            "name": self.name,
+            "status": self.status,
+            "source_type": self.source_type,
+            "source_name": self.source_name,
+            "artifact_uri": self.artifact_uri,
+            "start_time": self.start_time,
+            "end_time": self.end_time,
+            "duration_ms": self.duration_ms,
+            "tags": self.tags,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }
+
+    def __enter__(self) -> "Run":
+        """Enter context manager."""
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Exit context manager, ending the run."""
+        if self._client:
+            if exc_type is not None:
+                # Exception occurred, mark as failed
+                self._client.end_run(status="failed")
+            else:
+                self._client.end_run(status="completed")
+        return False  # Don't suppress exceptions
+
+    def log_params(self, params: Dict[str, Any]):
+        """Log parameters for this run."""
+        if self._client:
+            self._client.log_params(params)
+
+    def log_metrics(self, metrics: Dict[str, float], step: int = None):
+        """Log metrics for this run."""
+        if self._client:
+            self._client.log_metrics(metrics, step)
+
+    def log_artifact(self, local_path: str, artifact_path: str = None):
+        """Log an artifact for this run."""
+        if self._client:
+            self._client.log_artifact(local_path, artifact_path)
+
+    def set_tag(self, key: str, value: str):
+        """Set a tag on this run."""
+        if self._client:
+            self._client.set_tag(key, value)
+        self.tags[key] = value
+
+    def end(self, status: str = "completed"):
+        """End this run."""
+        if self._client:
+            self._client.end_run(status)
+        self.status = status
+
+
+def _parse_datetime(value: Any) -> Optional[datetime]:
+    """Parse a datetime from string or return None."""
+    if value is None:
+        return None
+    if isinstance(value, datetime):
+        return value
+    if isinstance(value, str):
+        try:
+            # Handle ISO format with Z suffix
+            if value.endswith("Z"):
+                value = value[:-1] + "+00:00"
+            return datetime.fromisoformat(value)
+        except ValueError:
+            return None
+    return None

podstack/registry/model.py

@@ -0,0 +1,273 @@
+"""
+Model Registry Classes
+
+Data classes for model registry.
+"""
+
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Any, List
+from datetime import datetime
+
+
+@dataclass
+class RegisteredModel:
+    """Represents a registered model in the registry."""
+
+    id: str
+    project_id: str
+    name: str
+    description: Optional[str] = None
+    tags: Dict[str, str] = field(default_factory=dict)
+    version_count: int = 0
+    latest_version: int = 0
+    created_by_user_id: Optional[str] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    # Loaded lazily
+    _versions: Optional[List["ModelVersion"]] = field(default=None, repr=False)
+    _aliases: Optional[List["ModelAlias"]] = field(default=None, repr=False)
+    _client: Any = field(default=None, repr=False)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], client: Any = None) -> "RegisteredModel":
+        """Create a RegisteredModel from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            project_id=data.get("project_id", ""),
+            name=data.get("name", ""),
+            description=data.get("description"),
+            tags=data.get("tags") or {},
+            version_count=data.get("version_count", 0),
+            latest_version=data.get("latest_version", 0),
+            created_by_user_id=data.get("created_by_user_id"),
+            created_at=_parse_datetime(data.get("created_at")),
+            updated_at=_parse_datetime(data.get("updated_at")),
+            _client=client,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "project_id": self.project_id,
+            "name": self.name,
+            "description": self.description,
+            "tags": self.tags,
+            "version_count": self.version_count,
+            "latest_version": self.latest_version,
+            "created_by_user_id": self.created_by_user_id,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }
+
+    def get_version(self, version: int) -> "ModelVersion":
+        """Get a specific version of this model."""
+        if self._client:
+            return self._client.get_model_version(self.id, version)
+        raise RuntimeError("Model not connected to client")
+
+    def list_versions(self, limit: int = 20, offset: int = 0) -> List["ModelVersion"]:
+        """List versions of this model."""
+        if self._client:
+            return self._client.list_model_versions(self.id, limit, offset)
+        raise RuntimeError("Model not connected to client")
+
+    def create_version(
+        self,
+        run_id: str = None,
+        source: str = None,
+        description: str = None
+    ) -> "ModelVersion":
+        """Create a new version of this model."""
+        if self._client:
+            return self._client.create_model_version(
+                self.id, run_id=run_id, source=source, description=description
+            )
+        raise RuntimeError("Model not connected to client")
+
+    def set_alias(self, alias: str, version: int) -> "ModelAlias":
+        """Set an alias for a version of this model."""
+        if self._client:
+            return self._client.set_model_alias(self.name, alias, version)
+        raise RuntimeError("Model not connected to client")
+
+    def get_aliases(self) -> List["ModelAlias"]:
+        """Get all aliases for this model."""
+        if self._client:
+            return self._client.get_model_aliases(self.id)
+        raise RuntimeError("Model not connected to client")
+
+
+@dataclass
+class ModelVersion:
+    """Represents a version of a registered model."""
+
+    id: str
+    model_id: str
+    project_id: str
+    version: int
+    run_id: Optional[str] = None
+    source: Optional[str] = None
+    status: str = "ready"
+    stage: str = "development"
+    description: Optional[str] = None
+    tags: Dict[str, str] = field(default_factory=dict)
+    created_by_user_id: Optional[str] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    _client: Any = field(default=None, repr=False)
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any], client: Any = None) -> "ModelVersion":
+        """Create a ModelVersion from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            model_id=data.get("model_id", ""),
+            project_id=data.get("project_id", ""),
+            version=data.get("version", 0),
+            run_id=data.get("run_id"),
+            source=data.get("source"),
+            status=data.get("status", "ready"),
+            stage=data.get("stage", "development"),
+            description=data.get("description"),
+            tags=data.get("tags") or {},
+            created_by_user_id=data.get("created_by_user_id"),
+            created_at=_parse_datetime(data.get("created_at")),
+            updated_at=_parse_datetime(data.get("updated_at")),
+            _client=client,
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "model_id": self.model_id,
+            "project_id": self.project_id,
+            "version": self.version,
+            "run_id": self.run_id,
+            "source": self.source,
+            "status": self.status,
+            "stage": self.stage,
+            "description": self.description,
+            "tags": self.tags,
+            "created_by_user_id": self.created_by_user_id,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }
+
+    def transition_stage(self, stage: str, comment: str = None) -> "ModelVersion":
+        """Transition this version to a new stage."""
+        if self._client:
+            return self._client.transition_model_stage(
+                self.model_id, self.version, stage, comment
+            )
+        raise RuntimeError("ModelVersion not connected to client")
+
+
+@dataclass
+class ModelAlias:
+    """Represents an alias for a model version."""
+
+    id: str
+    model_id: str
+    project_id: str
+    alias: str
+    model_version_id: str
+    version: int
+    created_by_user_id: Optional[str] = None
+    created_at: Optional[datetime] = None
+    updated_at: Optional[datetime] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ModelAlias":
+        """Create a ModelAlias from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            model_id=data.get("model_id", ""),
+            project_id=data.get("project_id", ""),
+            alias=data.get("alias", ""),
+            model_version_id=data.get("model_version_id", ""),
+            version=data.get("version", 0),
+            created_by_user_id=data.get("created_by_user_id"),
+            created_at=_parse_datetime(data.get("created_at")),
+            updated_at=_parse_datetime(data.get("updated_at")),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "model_id": self.model_id,
+            "project_id": self.project_id,
+            "alias": self.alias,
+            "model_version_id": self.model_version_id,
+            "version": self.version,
+            "created_by_user_id": self.created_by_user_id,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+            "updated_at": self.updated_at.isoformat() if self.updated_at else None,
+        }
+
+
+@dataclass
+class StageTransition:
+    """Represents a stage transition history record."""
+
+    id: str
+    model_id: str
+    model_version_id: str
+    project_id: str
+    version: int
+    from_stage: Optional[str] = None
+    to_stage: str = ""
+    transitioned_by: Optional[str] = None
+    comment: Optional[str] = None
+    created_at: Optional[datetime] = None
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "StageTransition":
+        """Create a StageTransition from a dict."""
+        return cls(
+            id=data.get("id", ""),
+            model_id=data.get("model_id", ""),
+            model_version_id=data.get("model_version_id", ""),
+            project_id=data.get("project_id", ""),
+            version=data.get("version", 0),
+            from_stage=data.get("from_stage"),
+            to_stage=data.get("to_stage", ""),
+            transitioned_by=data.get("transitioned_by"),
+            comment=data.get("comment"),
+            created_at=_parse_datetime(data.get("created_at")),
+        )
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dict."""
+        return {
+            "id": self.id,
+            "model_id": self.model_id,
+            "model_version_id": self.model_version_id,
+            "project_id": self.project_id,
+            "version": self.version,
+            "from_stage": self.from_stage,
+            "to_stage": self.to_stage,
+            "transitioned_by": self.transitioned_by,
+            "comment": self.comment,
+            "created_at": self.created_at.isoformat() if self.created_at else None,
+        }
+
+
+def _parse_datetime(value: Any) -> Optional[datetime]:
+    """Parse a datetime from string or return None."""
+    if value is None:
+        return None
+    if isinstance(value, datetime):
+        return value
+    if isinstance(value, str):
+        try:
+            if value.endswith("Z"):
+                value = value[:-1] + "+00:00"
+            return datetime.fromisoformat(value)
+        except ValueError:
+            return None
+    return None