flowyml 1.7.1__py3-none-any.whl → 1.8.0__py3-none-any.whl

flowyml/registry/model_registry.py

@@ -1,12 +1,19 @@
-"""Model registry for version management and deployment."""
+"""Model registry for version management and deployment.
 
-import json
-from dataclasses import dataclass, field, asdict
+This module provides SQL-backed model registry capabilities for managing
+model versions, stages, and metadata in a production-safe manner.
+"""
+
+import os
+import shutil
+from dataclasses import asdict, dataclass, field
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import Any
-import shutil
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from flowyml.assets.base import Asset
 
 
 class ModelStage(str, Enum):
@@ -31,6 +38,7 @@ class ModelVersion:
     framework: str
     metrics: dict[str, float] = field(default_factory=dict)
     tags: dict[str, str] = field(default_factory=dict)
+    schema: dict[str, Any] = field(default_factory=dict)
     description: str = ""
     author: str | None = None
     parent_version: str | None = None
@@ -44,13 +52,21 @@ class ModelVersion:
     @classmethod
     def from_dict(cls, data: dict[str, Any]) -> "ModelVersion":
         """Create from dictionary."""
-        data["stage"] = ModelStage(data["stage"])
+        data = data.copy()
+        # Handle stage conversion
+        if isinstance(data.get("stage"), str):
+            data["stage"] = ModelStage(data["stage"])
+        # Remove SQL-specific fields
+        data.pop("id", None)
         return cls(**data)
 
 
 class ModelRegistry:
     """Registry for managing model versions and deployments.
 
+    This registry uses SQL storage for production-safe concurrent access.
+    Model files are still stored on the filesystem, but metadata is in the database.
+
     Example:
         ```python
         from flowyml import ModelRegistry
@@ -78,30 +94,32 @@ class ModelRegistry:
         ```
     """
 
-    def __init__(self, registry_path: str = ".flowyml/model_registry"):
+    def __init__(
+        self,
+        registry_path: str = ".flowyml/model_registry",
+        db_url: str | None = None,
+    ):
         """Initialize model registry.
 
         Args:
-            registry_path: Path to registry storage
+            registry_path: Path to model file storage
+            db_url: Database URL for metadata storage (uses env var if not provided)
         """
         self.registry_path = Path(registry_path)
         self.registry_path.mkdir(parents=True, exist_ok=True)
-        self.metadata_file = self.registry_path / "registry.json"
-        self._metadata: dict[str, list[dict]] = {}
-        self._load_metadata()
-
-    def _load_metadata(self) -> None:
-        """Load registry metadata from disk."""
-        if self.metadata_file.exists():
-            with open(self.metadata_file) as f:
-                self._metadata = json.load(f)
-        else:
-            self._metadata = {}
 
-    def _save_metadata(self) -> None:
-        """Save registry metadata to disk."""
-        with open(self.metadata_file, "w") as f:
-            json.dump(self._metadata, f, indent=2)
+        # Initialize SQL storage for metadata
+        self._db_url = db_url or os.getenv("FLOWYML_DATABASE_URL")
+        self._store = None
+
+    @property
+    def _metadata_store(self):
+        """Lazy-load the metadata store."""
+        if self._store is None:
+            from flowyml.storage.sql import SQLMetadataStore
+
+            self._store = SQLMetadataStore(db_url=self._db_url)
+        return self._store
 
     def register(
         self,
@@ -112,6 +130,7 @@ class ModelRegistry:
         stage: ModelStage = ModelStage.DEVELOPMENT,
         metrics: dict[str, float] | None = None,
         tags: dict[str, str] | None = None,
+        schema: dict[str, Any] | None = None,
         description: str = "",
         author: str | None = None,
         parent_version: str | None = None,
@@ -126,6 +145,7 @@ class ModelRegistry:
             stage: Deployment stage
             metrics: Model metrics
             tags: Model tags
+            schema: Optional explicit schema (overrides introspection)
             description: Model description
             author: Model author
             parent_version: Parent version if this is an update
@@ -137,10 +157,18 @@ class ModelRegistry:
             ValueError: If version already exists
         """
         # Check if version already exists
-        if name in self._metadata:
-            existing_versions = [v["version"] for v in self._metadata[name]]
-            if version in existing_versions:
-                raise ValueError(f"Version {version} already exists for model {name}")
+        existing = self._metadata_store.get_model_version(name, version)
+        if existing:
+            raise ValueError(f"Version {version} already exists for model {name}")
+
+        # Introspect model schema if not provided
+        from flowyml.utils.model_introspection import introspect_model
+
+        inferred_schema = introspect_model(model, framework)
+        # Merge inferred schema with provided schema (provided takes precedence)
+        final_schema = inferred_schema
+        if schema:
+            final_schema.update(schema)
 
         # Create model directory
         model_dir = self.registry_path / name / version
@@ -152,7 +180,23 @@ class ModelRegistry:
 
         # Create version metadata
         now = datetime.now().isoformat()
-        model_version = ModelVersion(
+
+        # Save to SQL database
+        self._metadata_store.save_model_version(
+            name=name,
+            version=version,
+            stage=stage.value,
+            framework=framework,
+            model_path=str(model_path),
+            metrics=metrics,
+            tags=tags,
+            schema=final_schema,
+            description=description,
+            author=author,
+            parent_version=parent_version,
+        )
+
+        return ModelVersion(
             name=name,
             version=version,
             stage=stage,
@@ -162,20 +206,12 @@ class ModelRegistry:
             framework=framework,
             metrics=metrics or {},
             tags=tags or {},
+            schema=final_schema,
             description=description,
             author=author,
             parent_version=parent_version,
         )
 
-        # Add to metadata
-        if name not in self._metadata:
-            self._metadata[name] = []
-
-        self._metadata[name].append(model_version.to_dict())
-        self._save_metadata()
-
-        return model_version
-
     def _save_model(self, model: Any, path: Path, framework: str) -> None:
         """Save model using appropriate method.
 
@@ -226,7 +262,7 @@ class ModelRegistry:
             import pickle
 
             with open(path, "rb") as f:
-                return pickle.load(f)
+                return pickle.load(f)  # noqa: S301
 
     def get_version(self, name: str, version: str) -> ModelVersion | None:
         """Get specific model version.
@@ -238,13 +274,9 @@ class ModelRegistry:
         Returns:
             ModelVersion or None if not found
         """
-        if name not in self._metadata:
-            return None
-
-        for v in self._metadata[name]:
-            if v["version"] == version:
-                return ModelVersion.from_dict(v)
-
+        data = self._metadata_store.get_model_version(name, version)
+        if data:
+            return ModelVersion.from_dict(data)
         return None
 
     def list_versions(self, name: str) -> list[ModelVersion]:
@@ -256,10 +288,8 @@ class ModelRegistry:
         Returns:
             List of ModelVersion instances
         """
-        if name not in self._metadata:
-            return []
-
-        return [ModelVersion.from_dict(v) for v in self._metadata[name]]
+        versions = self._metadata_store.list_model_versions(name=name)
+        return [ModelVersion.from_dict(v) for v in versions]
 
     def list_models(self) -> list[str]:
         """List all registered models.
@@ -267,7 +297,7 @@ class ModelRegistry:
         Returns:
             List of model names
         """
-        return list(self._metadata.keys())
+        return self._metadata_store.list_registered_models()
 
     def get_latest_version(self, name: str, stage: ModelStage | None = None) -> ModelVersion | None:
         """Get latest version of a model.
@@ -279,17 +309,11 @@ class ModelRegistry:
         Returns:
             Latest ModelVersion or None
         """
-        versions = self.list_versions(name)
-
-        if stage:
-            versions = [v for v in versions if v.stage == stage]
-
-        if not versions:
-            return None
-
-        # Sort by created_at
-        versions.sort(key=lambda v: v.created_at, reverse=True)
-        return versions[0]
+        stage_str = stage.value if stage else None
+        data = self._metadata_store.get_latest_model_version(name, stage=stage_str)
+        if data:
+            return ModelVersion.from_dict(data)
+        return None
 
     def load(
         self,
@@ -341,14 +365,10 @@ class ModelRegistry:
         if not model_version:
             raise ValueError(f"Model {name} version {version} not found")
 
-        # Update stage in metadata
-        for v in self._metadata[name]:
-            if v["version"] == version:
-                v["stage"] = to_stage.value
-                v["updated_at"] = datetime.now().isoformat()
-                break
-
-        self._save_metadata()
+        # Update stage in database
+        success = self._metadata_store.promote_model(name, version, to_stage.value)
+        if not success:
+            raise ValueError(f"Failed to promote model {name} version {version}")
 
         return self.get_version(name, version)
 
@@ -392,16 +412,14 @@ class ModelRegistry:
         if model_version.stage == ModelStage.PRODUCTION:
             raise ValueError("Cannot delete production model. Demote first.")
 
-        # Remove from metadata
-        self._metadata[name] = [v for v in self._metadata[name] if v["version"] != version]
+        # Delete from database
+        self._metadata_store.delete_model_version(name, version)
 
         # Delete model files
         model_dir = Path(model_version.model_path).parent
         if model_dir.exists():
             shutil.rmtree(model_dir)
 
-        self._save_metadata()
-
     def compare_versions(
         self,
         name: str,
@@ -447,26 +465,23 @@ class ModelRegistry:
         Returns:
             List of matching ModelVersion instances
         """
-        results = []
+        stage_str = stage.value if stage else None
+        all_versions = self._metadata_store.list_model_versions(stage=stage_str)
 
-        for name in self._metadata:
-            for version_dict in self._metadata[name]:
-                version = ModelVersion.from_dict(version_dict)
+        results = []
+        for version_dict in all_versions:
+            version = ModelVersion.from_dict(version_dict)
 
-                # Check stage
-                if stage and version.stage != stage:
-                    continue
+            # Check tags
+            if tags and not all(version.tags.get(k) == v for k, v in tags.items()):
+                continue
 
-                # Check tags
-                if tags and not all(version.tags.get(k) == v for k, v in tags.items()):
+            # Check metrics
+            if min_metrics:
+                if not all(version.metrics.get(k, float("-inf")) >= v for k, v in min_metrics.items()):
                     continue
 
-                # Check metrics
-                if min_metrics:
-                    if not all(version.metrics.get(k, float("-inf")) >= v for k, v in min_metrics.items()):
-                        continue
-
-                results.append(version)
+            results.append(version)
 
         return results
 
@@ -476,16 +491,146 @@ class ModelRegistry:
         Returns:
             Dictionary with statistics
         """
-        total_models = len(self._metadata)
-        total_versions = sum(len(versions) for versions in self._metadata.values())
-
+        all_versions = self._metadata_store.list_model_versions()
+        models = set()
         stage_counts = {stage.value: 0 for stage in ModelStage}
-        for versions in self._metadata.values():
-            for v in versions:
-                stage_counts[v["stage"]] += 1
+
+        for v in all_versions:
+            models.add(v["name"])
+            stage_counts[v["stage"]] += 1
 
         return {
-            "total_models": total_models,
-            "total_versions": total_versions,
+            "total_models": len(models),
+            "total_versions": len(all_versions),
             "by_stage": stage_counts,
         }
+
+    # ========== Asset Integration Methods ==========
+
+    def register_asset(
+        self,
+        model_asset: "Asset",
+        version: str | None = None,
+        stage: ModelStage = ModelStage.DEVELOPMENT,
+        description: str = "",
+        author: str | None = None,
+        capture_environment: bool = False,
+    ) -> ModelVersion:
+        """Register a Model asset directly to the registry.
+
+        This method leverages the auto-extracted metadata from Model assets,
+        making registration simpler and more consistent.
+
+        Args:
+            model_asset: A Model asset (from flowyml.assets.model)
+            version: Version string (defaults to asset's version)
+            stage: Deployment stage
+            description: Model description
+            author: Model author
+            capture_environment: Whether to capture Python environment
+
+        Returns:
+            ModelVersion instance
+
+        Example:
+            >>> from flowyml import Model, ModelRegistry
+            >>> model_asset = Model.create(data=trained_model, name="classifier")
+            >>> registry = ModelRegistry()
+            >>> registry.register_asset(model_asset, version="v1.0.0")
+        """
+        # Get version from asset if not provided
+        version = version or model_asset.version
+
+        # Extract metadata from asset
+        properties = model_asset.properties if hasattr(model_asset, "properties") else {}
+        tags = model_asset.tags if hasattr(model_asset, "tags") else {}
+
+        # Get framework from auto-extracted properties
+        framework = properties.get("framework", "unknown")
+
+        # Extract metrics from properties (common keys)
+        metrics = {}
+        for key in ["accuracy", "f1", "precision", "recall", "loss", "auc", "mse", "mae"]:
+            if key in properties:
+                metrics[key] = properties[key]
+
+        # Capture environment if requested
+        if capture_environment:
+            from flowyml.registry.model_environment import ModelEnvironment
+
+            env = ModelEnvironment.from_current()
+            tags["python_version"] = env.python_version
+            tags["platform"] = env.platform
+            properties["environment"] = env.to_dict()
+
+        return self.register(
+            model=model_asset.data,
+            name=model_asset.name,
+            version=version,
+            framework=framework,
+            stage=stage,
+            metrics=metrics,
+            tags=tags,
+            description=description or properties.get("description", ""),
+            author=author,
+            parent_version=None,
+        )
+
+    def to_asset(
+        self,
+        name: str,
+        version: str | None = None,
+        stage: ModelStage | None = None,
+    ) -> "Asset":
+        """Load a model version as a Model asset.
+
+        This creates a Model asset with all the stored metadata,
+        enabling seamless integration with FlowyML pipelines.
+
+        Args:
+            name: Model name
+            version: Specific version (if None, loads latest)
+            stage: Stage filter (if version is None)
+
+        Returns:
+            Model asset with loaded model and metadata
+
+        Example:
+            >>> registry = ModelRegistry()
+            >>> model_asset = registry.to_asset("classifier", version="v1.0.0")
+            >>> print(model_asset.properties)
+        """
+        from flowyml.assets.model import Model
+
+        model_version = self.get_version(name, version) if version else self.get_latest_version(name, stage)
+
+        if not model_version:
+            raise ValueError(f"Model {name} not found")
+
+        model_data = self.load(name, version or model_version.version)
+
+        properties = {
+            "framework": model_version.framework,
+            "stage": model_version.stage.value,
+            "created_at": model_version.created_at,
+            "updated_at": model_version.updated_at,
+            **model_version.metrics,
+        }
+
+        return Model(
+            name=name,
+            version=model_version.version,
+            data=model_data,
+            tags=model_version.tags,
+            properties=properties,
+        )
+
+    def capture_environment(self) -> dict[str, Any]:
+        """Capture current Python environment.
+
+        Returns:
+            Dictionary with environment info
+        """
+        from flowyml.registry.model_environment import ModelEnvironment
+
+        return ModelEnvironment.from_current().to_dict()

flowyml/serving/__init__.py

@@ -0,0 +1,17 @@
+"""Model serving utilities for FlowyML deployments."""
+
+from .model_server import (
+    ModelServer,
+    ServerConfig,
+    start_model_server,
+    stop_model_server,
+    load_and_predict,
+)
+
+__all__ = [
+    "ModelServer",
+    "ServerConfig",
+    "start_model_server",
+    "stop_model_server",
+    "load_and_predict",
+]