geronimo.artifacts

 25class ArtifactStore:
 26    """Versioned storage for ML artifacts.
 27
 28    Supports local filesystem, S3, and Geronimo Cloud backends.
 29    Reads defaults from ~/.geronimo/config.yaml (run `geronimo config show`).
 30
 31    Example:
 32        ```python
 33        from geronimo.artifacts import ArtifactStore
 34
 35        # Save during training (uses global config defaults)
 36        store = ArtifactStore(project="credit-risk", version="1.2.0")
 37        store.save("model", trained_model)
 38        store.save("encoder", fitted_encoder)
 39
 40        # Load in production
 41        store = ArtifactStore.load(project="credit-risk", version="1.2.0")
 42        model = store.get("model")
 43        encoder = store.get("encoder")
 44        ```
 45    
 46    Configuration:
 47        Run `geronimo config init` to set up your preferred backend.
 48        Run `geronimo config show` to view current settings.
 49    """
 50
 51    project: str
 52    """The project name."""
 53
 54    version: str
 55    """The version string (e.g., "1.2.0")."""
 56
 57    backend: str
 58    """The active storage backend ("local", "s3", "gdc", "mlflow", or "custom")."""
 59
 60    base_path: Optional[str]
 61    """Base path for local storage (if using local backend)."""
 62
 63    s3_bucket: Optional[str]
 64    """S3 bucket name (if using s3 backend)."""
 65
 66    def __init__(
 67        self,
 68        project: str,
 69        version: str,
 70        backend: Optional[Union[Literal["local", "s3", "gdc", "mlflow"], ArtifactBackend]] = None,
 71        base_path: Optional[str] = None,
 72        s3_bucket: Optional[str] = None,
 73    ):
 74        """Initialize artifact store.
 75
 76        Args:
 77            project: Project name.
 78            version: Version string (e.g., "1.2.0").
 79            backend: Storage backend ("local", "s3", "mlflow", "gdc") or custom ArtifactBackend instance.
 80                     Defaults to value from ~/.geronimo/config.yaml.
 81            base_path: Base path for local storage.
 82                       Defaults to ~/.geronimo/artifacts.
 83            s3_bucket: S3 bucket for s3 backend.
 84                       Defaults to value from config or GERONIMO_ARTIFACT_BUCKET env var.
 85        """
 86        # Load user config for defaults
 87        from geronimo.config.user_config import load_user_config
 88        user_config = load_user_config()
 89        
 90        self.project = project
 91        self.version = version
 92        
 93        # Handle custom backend instance
 94        if isinstance(backend, ArtifactBackend):
 95            self.backend = "custom"
 96            self._backend_impl = backend
 97        else:
 98            self.backend = backend or user_config.artifacts.backend
 99            
100            # Resolve s3_bucket with fallback chain: param -> config -> env -> default
101            self.s3_bucket = (
102                s3_bucket 
103                or user_config.artifacts.s3_bucket 
104                or os.getenv("GERONIMO_ARTIFACT_BUCKET", "ml-artifacts")
105            )
106            
107            # Resolve base_path
108            self.base_path = base_path or os.path.expanduser(user_config.artifacts.base_path)
109            
110            # Create backend instance using factory
111            self._backend_impl = self._create_backend()
112
113        self._metadata: dict[str, ArtifactMetadata] = {}
114
115    def _create_backend(self) -> ArtifactBackend:
116        """Factory method to create backend instance.
117        
118        Returns:
119            ArtifactBackend implementation based on self.backend type.
120        """
121        if self.backend == "local":
122            from geronimo.artifacts.local_backend import LocalArtifactBackend
123            return LocalArtifactBackend(
124                project=self.project,
125                version=self.version,
126                base_path=self.base_path,
127            )
128        elif self.backend == "s3":
129            from geronimo.artifacts.s3_backend import S3ArtifactBackend
130            return S3ArtifactBackend(
131                project=self.project,
132                version=self.version,
133                bucket=self.s3_bucket,
134            )
135        elif self.backend == "gdc":
136            from geronimo.artifacts.gdc_backend import GeronimoDeployCloudArtifactBackend
137            return GeronimoDeployCloudArtifactBackend(
138                project=self.project,
139                version=self.version,
140            )
141        elif self.backend == "mlflow":
142            from geronimo.artifacts.mlflow_backend import MLFlowArtifactBackend
143            return MLFlowArtifactBackend(
144                project=self.project,
145                version=self.version,
146            )
147        else:
148            raise ValueError(f"Unknown backend type: {self.backend}")
149
150    @classmethod
151    def load(
152        cls,
153        project: str,
154        version: str,
155        backend: Optional[Literal["local", "s3", "gdc", "mlflow"]] = None,
156        **kwargs,
157    ) -> "ArtifactStore":
158        """Load existing artifact store.
159
160        Args:
161            project: Project name.
162            version: Version string.
163            backend: Storage backend. Defaults to value from config.
164            **kwargs: Additional backend options.
165
166        Returns:
167            ArtifactStore instance with loaded metadata.
168        """
169        store = cls(project=project, version=version, backend=backend, **kwargs)
170        store._load_metadata()
171        return store
172
173    def save(
174        self,
175        name: str,
176        artifact: Any,
177        artifact_type: Optional[str] = None,
178        tags: Optional[dict[str, str]] = None,
179    ) -> str:
180        """Save an artifact.
181
182        Args:
183            name: Artifact name (e.g., "model", "encoder").
184            artifact: Python object to serialize.
185            artifact_type: Optional type hint (auto-detected if not provided).
186            tags: Optional metadata tags.
187
188        Returns:
189            Path or URI where artifact was saved.
190        """
191        artifact_type = artifact_type or type(artifact).__name__
192        tags = tags or {}
193        
194        meta_dict = {
195            "artifact_type": artifact_type,
196            "tags": tags,
197        }
198
199        # Delegate to backend
200        uri = self._backend_impl.save(name, artifact, meta_dict)
201        
202        # Keep local metadata cache in sync
203        self._metadata[name] = ArtifactMetadata(
204            name=name,
205            version=self.version,
206            artifact_type=artifact_type,
207            created_at=datetime.utcnow(),
208            size_bytes=0,  # Backend may have actual size
209            tags=tags,
210        )
211        
212        return uri
213
214    def get(self, name: str) -> Any:
215        """Load an artifact by name.
216
217        Args:
218            name: Artifact name.
219
220        Returns:
221            Deserialized artifact.
222
223        Raises:
224            KeyError: If artifact not found.
225        """
226        return self._backend_impl.load(name)
227
228    def list(self) -> list[ArtifactMetadata]:
229        """List all artifacts in this store.
230
231        Returns:
232            List of artifact metadata.
233        """
234        self._load_metadata()
235        return list(self._metadata.values())
236        
237    def delete(self, name: str) -> None:
238        """Delete an artifact.
239        
240        Args:
241           name: Artifact name to delete
242        """
243        self._backend_impl.delete(name)
244        if name in self._metadata:
245            del self._metadata[name]
246
247    def _load_metadata(self) -> None:
248        """Load metadata from backend."""
249        # Try to get metadata from backend if it supports it
250        if hasattr(self._backend_impl, 'get_metadata_index'):
251            data = self._backend_impl.get_metadata_index()
252            self._metadata = {
253                k: ArtifactMetadata.model_validate(v) for k, v in data.items()
254            }
255
256    def __repr__(self) -> str:
257        return f"ArtifactStore({self.project}@{self.version}, backend={self.backend})"

 8@runtime_checkable
 9class ArtifactBackend(Protocol):
10    """Protocol for custom artifact storage backends.
11    
12    Implement this protocol to use your own storage system
13    (internal object store, custom S3, etc.) with Geronimo.
14    
15    Example:
16        class MyCompanyStore(ArtifactBackend):
17            def save(self, name, artifact, metadata):
18                return self._internal_api.upload(artifact)
19            
20            def load(self, uri):
21                return self._internal_api.download(uri)
22            
23            def list(self, prefix):
24                return self._internal_api.list_objects(prefix)
25            
26            def delete(self, uri):
27                self._internal_api.delete(uri)
28    
29    Usage:
30        from geronimo.artifacts import ArtifactStore
31        
32        custom_backend = MyCompanyStore(config=my_config)
33        store = ArtifactStore(
34            project="my-model",
35            version="1.0.0",
36            backend=custom_backend,
37        )
38    """
39    
40    def save(self, name: str, artifact: Any, metadata: dict) -> str:
41        """Save an artifact.
42        
43        Args:
44            name: Artifact name (e.g., "model", "encoder")
45            artifact: Serialized artifact bytes
46            metadata: Artifact metadata dict
47        
48        Returns:
49            URI or path where artifact was saved
50        """
51        ...
52    
53    def load(self, uri: str) -> Any:
54        """Load an artifact by URI.
55        
56        Args:
57            uri: URI returned from save()
58        
59        Returns:
60            Deserialized artifact
61        """
62        ...
63    
64    def list(self, prefix: str) -> list[str]:
65        """List artifacts matching prefix.
66        
67        Args:
68            prefix: Prefix to filter artifacts
69        
70        Returns:
71            List of artifact URIs
72        """
73        ...
74    
75    def delete(self, uri: str) -> None:
76        """Delete an artifact.
77        
78        Args:
79            uri: URI of artifact to delete
80        """
81        ...

 17class LocalArtifactBackend(ArtifactBackend):
 18    """Artifact backend for local filesystem storage.
 19    
 20    Stores artifacts as pickle files with JSON metadata index.
 21    """
 22
 23    project: str
 24    """The project name."""
 25
 26    version: str
 27    """The version string."""
 28
 29    base_path: Path
 30    """Base directory for artifacts."""
 31
 32    artifact_path: Path
 33    """Directory for the specific project version."""
 34
 35    _metadata_file: Path
 36    """Path to the metadata index file."""
 37    
 38    def __init__(
 39        self,
 40        project: str,
 41        version: str,
 42        base_path: Optional[str] = None,
 43    ):
 44        """Initialize local backend.
 45        
 46        Args:
 47            project: Project name.
 48            version: Version string.
 49            base_path: Base directory for artifacts. 
 50                       Defaults to ~/.geronimo/artifacts.
 51        """
 52        self.project = project
 53        self.version = version
 54        self.base_path = Path(base_path or os.path.expanduser("~/.geronimo/artifacts"))
 55        self.artifact_path = self.base_path / project / version
 56        self.artifact_path.mkdir(parents=True, exist_ok=True)
 57        self._metadata_file = self.artifact_path / "metadata.json"
 58    
 59    def save(self, name: str, artifact: Any, metadata: dict) -> str:
 60        """Save an artifact to local filesystem.
 61        
 62        Args:
 63            name: Artifact name.
 64            artifact: Python object to serialize.
 65            metadata: Artifact metadata dict.
 66            
 67        Returns:
 68            Path to saved artifact file.
 69        """
 70        artifact_file = self.artifact_path / f"{name}.pkl"
 71        
 72        with open(artifact_file, "wb") as f:
 73            pickle.dump(artifact, f)
 74        
 75        logger.debug(f"Saved artifact '{name}' to {artifact_file}")
 76        
 77        # Update metadata index
 78        self._save_to_metadata_index(name, artifact_file.stat().st_size, metadata)
 79        
 80        return str(artifact_file)
 81    
 82    def load(self, uri: str) -> Any:
 83        """Load an artifact by name or path.
 84        
 85        Args:
 86            uri: Artifact name or file path.
 87            
 88        Returns:
 89            Deserialized artifact.
 90            
 91        Raises:
 92            KeyError: If artifact not found.
 93        """
 94        # If path, use directly; otherwise treat as name
 95        if os.path.isabs(uri) or os.path.exists(uri):
 96            artifact_file = Path(uri)
 97        else:
 98            artifact_file = self.artifact_path / f"{uri}.pkl"
 99        
100        if not artifact_file.exists():
101            raise KeyError(f"Artifact not found: {uri}")
102        
103        logger.debug(f"Loading artifact from {artifact_file}")
104        
105        with open(artifact_file, "rb") as f:
106            return pickle.load(f)
107    
108    def list(self, prefix: Optional[str] = None) -> list[str]:
109        """List artifacts in the store.
110        
111        Args:
112            prefix: Optional name prefix to filter.
113            
114        Returns:
115            List of artifact file paths.
116        """
117        paths = []
118        for pkl_file in self.artifact_path.glob("*.pkl"):
119            name = pkl_file.stem
120            if prefix is None or name.startswith(prefix):
121                paths.append(str(pkl_file))
122        return paths
123    
124    def delete(self, uri: str) -> None:
125        """Delete an artifact.
126        
127        Args:
128            uri: Artifact name or file path.
129        """
130        # If path, use directly; otherwise treat as name
131        if os.path.isabs(uri) or os.path.exists(uri):
132            artifact_file = Path(uri)
133            name = artifact_file.stem
134        else:
135            name = uri
136            artifact_file = self.artifact_path / f"{uri}.pkl"
137        
138        if artifact_file.exists():
139            artifact_file.unlink()
140            logger.info(f"Deleted artifact '{name}'")
141            self._remove_from_metadata_index(name)
142        else:
143            logger.warning(f"Artifact not found for deletion: {uri}")
144    
145    def _save_to_metadata_index(self, name: str, size_bytes: int, metadata: dict) -> None:
146        """Add artifact to metadata index."""
147        index = self._load_metadata_index()
148        
149        index[name] = {
150            "name": name,
151            "version": self.version,
152            "artifact_type": metadata.get("artifact_type", "unknown"),
153            "created_at": datetime.utcnow().isoformat(),
154            "size_bytes": size_bytes,
155            "tags": metadata.get("tags", {}),
156        }
157        
158        self._metadata_file.write_text(json.dumps(index, indent=2))
159    
160    def _remove_from_metadata_index(self, name: str) -> None:
161        """Remove artifact from metadata index."""
162        index = self._load_metadata_index()
163        if name in index:
164            del index[name]
165            self._metadata_file.write_text(json.dumps(index, indent=2))
166    
167    def _load_metadata_index(self) -> dict:
168        """Load metadata index from file."""
169        if self._metadata_file.exists():
170            return json.loads(self._metadata_file.read_text())
171        return {}
172    
173    def get_metadata_index(self) -> dict:
174        """Get the full metadata index for this store.
175        
176        Returns:
177            Dict mapping artifact names to metadata dicts.
178        """
179        return self._load_metadata_index()

 17class S3ArtifactBackend(ArtifactBackend):
 18    """Artifact backend for S3 storage.
 19    
 20    Stores artifacts as pickle files in S3 with JSON metadata index.
 21    """
 22
 23    project: str
 24    """The project name."""
 25
 26    version: str
 27    """The version string."""
 28
 29    bucket: str
 30    """The S3 bucket name."""
 31
 32    _prefix: str
 33    """The S3 key prefix for this project version."""
 34    
 35    def __init__(
 36        self,
 37        project: str,
 38        version: str,
 39        bucket: Optional[str] = None,
 40    ):
 41        """Initialize S3 backend.
 42        
 43        Args:
 44            project: Project name.
 45            version: Version string.
 46            bucket: S3 bucket name. Defaults to GERONIMO_ARTIFACT_BUCKET env var.
 47        """
 48        self.project = project
 49        self.version = version
 50        self.bucket = bucket or os.getenv("GERONIMO_ARTIFACT_BUCKET", "ml-artifacts")
 51        self._prefix = f"{project}/{version}"
 52    
 53    def _get_s3_client(self):
 54        """Get boto3 S3 client (lazy import)."""
 55        import boto3
 56        return boto3.client("s3")
 57    
 58    def save(self, name: str, artifact: Any, metadata: dict) -> str:
 59        """Save an artifact to S3.
 60        
 61        Args:
 62            name: Artifact name.
 63            artifact: Python object to serialize.
 64            metadata: Artifact metadata dict.
 65            
 66        Returns:
 67            S3 URI of saved artifact.
 68        """
 69        s3 = self._get_s3_client()
 70        key = f"{self._prefix}/{name}.pkl"
 71        
 72        # Serialize to temp file
 73        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
 74            pickle.dump(artifact, f)
 75            temp_path = f.name
 76        
 77        try:
 78            s3.upload_file(temp_path, self.bucket, key)
 79            size_bytes = os.path.getsize(temp_path)
 80            logger.debug(f"Uploaded artifact '{name}' to s3://{self.bucket}/{key}")
 81        finally:
 82            os.unlink(temp_path)
 83        
 84        # Update metadata index in S3
 85        self._save_to_metadata_index(name, size_bytes, metadata)
 86        
 87        return f"s3://{self.bucket}/{key}"
 88    
 89    def load(self, uri: str) -> Any:
 90        """Load an artifact by name or S3 URI.
 91        
 92        Args:
 93            uri: Artifact name or S3 URI.
 94            
 95        Returns:
 96            Deserialized artifact.
 97            
 98        Raises:
 99            KeyError: If artifact not found.
100        """
101        s3 = self._get_s3_client()
102        
103        # Parse URI or construct from name
104        if uri.startswith("s3://"):
105            # Parse s3://bucket/path/to/artifact.pkl
106            path = uri.replace("s3://", "")
107            parts = path.split("/", 1)
108            bucket = parts[0]
109            key = parts[1] if len(parts) > 1 else ""
110        else:
111            bucket = self.bucket
112            key = f"{self._prefix}/{uri}.pkl"
113        
114        with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f:
115            temp_path = f.name
116        
117        try:
118            s3.download_file(bucket, key, temp_path)
119            logger.debug(f"Downloaded artifact from s3://{bucket}/{key}")
120            
121            with open(temp_path, "rb") as f:
122                return pickle.load(f)
123        except s3.exceptions.NoSuchKey:
124            raise KeyError(f"Artifact not found: {uri}")
125        finally:
126            if os.path.exists(temp_path):
127                os.unlink(temp_path)
128    
129    def list(self, prefix: Optional[str] = None) -> list[str]:
130        """List artifacts in the store.
131        
132        Args:
133            prefix: Optional name prefix to filter.
134            
135        Returns:
136            List of S3 URIs.
137        """
138        s3 = self._get_s3_client()
139        
140        search_prefix = self._prefix
141        if prefix:
142            search_prefix = f"{search_prefix}/{prefix}"
143        
144        uris = []
145        paginator = s3.get_paginator("list_objects_v2")
146        
147        for page in paginator.paginate(Bucket=self.bucket, Prefix=search_prefix):
148            for obj in page.get("Contents", []):
149                key = obj["Key"]
150                if key.endswith(".pkl"):
151                    uris.append(f"s3://{self.bucket}/{key}")
152        
153        return uris
154    
155    def delete(self, uri: str) -> None:
156        """Delete an artifact.
157        
158        Args:
159            uri: Artifact name or S3 URI.
160        """
161        s3 = self._get_s3_client()
162        
163        # Parse URI or construct from name
164        if uri.startswith("s3://"):
165            path = uri.replace("s3://", "")
166            parts = path.split("/", 1)
167            bucket = parts[0]
168            key = parts[1] if len(parts) > 1 else ""
169            # Extract name from key for metadata
170            name = key.split("/")[-1].replace(".pkl", "")
171        else:
172            name = uri
173            bucket = self.bucket
174            key = f"{self._prefix}/{uri}.pkl"
175        
176        try:
177            s3.delete_object(Bucket=bucket, Key=key)
178            logger.info(f"Deleted artifact from s3://{bucket}/{key}")
179            self._remove_from_metadata_index(name)
180        except Exception as e:
181            logger.warning(f"Failed to delete artifact: {e}")
182    
183    def _save_to_metadata_index(self, name: str, size_bytes: int, metadata: dict) -> None:
184        """Add artifact to metadata index in S3."""
185        index = self._load_metadata_index()
186        
187        index[name] = {
188            "name": name,
189            "version": self.version,
190            "artifact_type": metadata.get("artifact_type", "unknown"),
191            "created_at": datetime.utcnow().isoformat(),
192            "size_bytes": size_bytes,
193            "tags": metadata.get("tags", {}),
194        }
195        
196        s3 = self._get_s3_client()
197        key = f"{self._prefix}/metadata.json"
198        s3.put_object(
199            Bucket=self.bucket,
200            Key=key,
201            Body=json.dumps(index, indent=2),
202            ContentType="application/json",
203        )
204    
205    def _remove_from_metadata_index(self, name: str) -> None:
206        """Remove artifact from metadata index."""
207        index = self._load_metadata_index()
208        if name in index:
209            del index[name]
210            s3 = self._get_s3_client()
211            key = f"{self._prefix}/metadata.json"
212            s3.put_object(
213                Bucket=self.bucket,
214                Key=key,
215                Body=json.dumps(index, indent=2),
216                ContentType="application/json",
217            )
218    
219    def _load_metadata_index(self) -> dict:
220        """Load metadata index from S3."""
221        s3 = self._get_s3_client()
222        key = f"{self._prefix}/metadata.json"
223        
224        try:
225            response = s3.get_object(Bucket=self.bucket, Key=key)
226            return json.loads(response["Body"].read())
227        except Exception:
228            return {}
229    
230    def get_metadata_index(self) -> dict:
231        """Get the full metadata index for this store.
232        
233        Returns:
234            Dict mapping artifact names to metadata dicts.
235        """
236        return self._load_metadata_index()