flowyml 1.1.0__py3-none-any.whl

flowyml/storage/__init__.py

@@ -0,0 +1,26 @@
+"""Storage module for artifacts, metadata, and materializers."""
+
+from flowyml.storage.artifacts import ArtifactStore, LocalArtifactStore
+from flowyml.storage.metadata import MetadataStore, SQLiteMetadataStore
+from flowyml.storage.materializers.base import BaseMaterializer, materializer_registry
+from flowyml.storage.materializers.pytorch import PyTorchMaterializer
+from flowyml.storage.materializers.tensorflow import TensorFlowMaterializer
+from flowyml.storage.materializers.sklearn import SklearnMaterializer
+from flowyml.storage.materializers.pandas import PandasMaterializer
+from flowyml.storage.materializers.numpy import NumPyMaterializer
+
+__all__ = [
+    # Core Storage
+    "ArtifactStore",
+    "LocalArtifactStore",
+    "MetadataStore",
+    "SQLiteMetadataStore",
+    # Materializers
+    "BaseMaterializer",
+    "materializer_registry",
+    "PyTorchMaterializer",
+    "TensorFlowMaterializer",
+    "SklearnMaterializer",
+    "PandasMaterializer",
+    "NumPyMaterializer",
+]

flowyml/storage/artifacts.py

@@ -0,0 +1,246 @@
+"""Artifact storage backends for flowyml."""
+
+import shutil
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+import pickle
+
+
+class ArtifactStore(ABC):
+    """Base class for artifact storage backends."""
+
+    @abstractmethod
+    def save(self, artifact: Any, path: str, metadata: dict | None = None) -> str:
+        """Save an artifact to storage.
+
+        Args:
+            artifact: The artifact to save
+            path: Storage path for the artifact
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Full path where artifact was saved
+        """
+        pass
+
+    @abstractmethod
+    def load(self, path: str) -> Any:
+        """Load an artifact from storage.
+
+        Args:
+            path: Storage path of the artifact
+
+        Returns:
+            The loaded artifact
+        """
+        pass
+
+    @abstractmethod
+    def exists(self, path: str) -> bool:
+        """Check if artifact exists at path."""
+        pass
+
+    @abstractmethod
+    def delete(self, path: str) -> None:
+        """Delete artifact at path."""
+        pass
+
+    @abstractmethod
+    def list_artifacts(self, prefix: str = "") -> list[str]:
+        """List all artifacts with optional prefix filter."""
+        pass
+
+    def materialize(self, obj: Any, name: str, run_id: str, step_name: str, project_name: str = "default") -> str:
+        """Materialize artifact to structured storage.
+
+        Args:
+            obj: Object to materialize
+            name: Name of the artifact
+            run_id: ID of the current run
+            step_name: Name of the step producing the artifact
+            project_name: Name of the project
+
+        Returns:
+            Path where artifact was saved
+        """
+        pass
+
+
+class LocalArtifactStore(ArtifactStore):
+    """Local filesystem artifact storage."""
+
+    def __init__(self, base_path: str = ".flowyml/artifacts"):
+        """Initialize local artifact store.
+
+        Args:
+            base_path: Base directory for storing artifacts
+        """
+        self.base_path = Path(base_path)
+        self.base_path.mkdir(parents=True, exist_ok=True)
+
+    def save(self, artifact: Any, path: str, metadata: dict | None = None) -> str:
+        """Save artifact to local filesystem.
+
+        Args:
+            artifact: The artifact to save
+            path: Relative path for the artifact
+            metadata: Optional metadata dictionary
+
+        Returns:
+            Full path where artifact was saved
+        """
+        full_path = self.base_path / path
+        full_path.parent.mkdir(parents=True, exist_ok=True)
+
+        # Save artifact using pickle by default
+        with open(full_path, "wb") as f:
+            pickle.dump(artifact, f)
+
+        # Save metadata if provided
+        if metadata:
+            metadata_path = full_path.with_suffix(".meta.json")
+            import json
+
+            with open(metadata_path, "w") as f:
+                json.dump(metadata, f, indent=2)
+
+        return str(full_path)
+
+    def load(self, path: str) -> Any:
+        """Load artifact from local filesystem.
+
+        Args:
+            path: Relative or absolute path to the artifact
+
+        Returns:
+            The loaded artifact
+        """
+        full_path = Path(path) if Path(path).is_absolute() else self.base_path / path
+
+        if not full_path.exists():
+            raise FileNotFoundError(f"Artifact not found at {full_path}")
+
+        with open(full_path, "rb") as f:
+            return pickle.load(f)
+
+    def exists(self, path: str) -> bool:
+        """Check if artifact exists.
+
+        Args:
+            path: Relative path to check
+
+        Returns:
+            True if artifact exists, False otherwise
+        """
+        full_path = self.base_path / path
+        return full_path.exists()
+
+    def delete(self, path: str) -> None:
+        """Delete artifact from filesystem.
+
+        Args:
+            path: Relative path to delete
+        """
+        full_path = self.base_path / path
+        if full_path.exists():
+            if full_path.is_dir():
+                shutil.rmtree(full_path)
+            else:
+                full_path.unlink()
+
+            # Also delete metadata if exists
+            metadata_path = full_path.with_suffix(".meta.json")
+            if metadata_path.exists():
+                metadata_path.unlink()
+
+    def list_artifacts(self, prefix: str = "") -> list[str]:
+        """List all artifacts with optional prefix.
+
+        Args:
+            prefix: Optional prefix filter
+
+        Returns:
+            List of artifact paths
+        """
+        search_path = self.base_path / prefix if prefix else self.base_path
+
+        if not search_path.exists():
+            return []
+
+        artifacts = []
+        for item in search_path.rglob("*"):
+            if item.is_file() and not item.name.endswith(".meta.json"):
+                rel_path = item.relative_to(self.base_path)
+                artifacts.append(str(rel_path))
+
+        return sorted(artifacts)
+
+    def get_metadata(self, path: str) -> dict | None:
+        """Get metadata for an artifact.
+
+        Args:
+            path: Relative path to the artifact
+
+        Returns:
+            Metadata dictionary or None if no metadata exists
+        """
+        full_path = self.base_path / path
+        metadata_path = full_path.with_suffix(".meta.json")
+
+        if not metadata_path.exists():
+            return None
+
+        import json
+
+        with open(metadata_path) as f:
+            return json.load(f)
+
+    def size(self, path: str) -> int:
+        """Get size of artifact in bytes.
+
+        Args:
+            path: Relative path to the artifact
+
+        Returns:
+            Size in bytes
+        """
+        full_path = self.base_path / path
+        if full_path.exists():
+            return full_path.stat().st_size
+        return 0
+
+    def materialize(self, obj: Any, name: str, run_id: str, step_name: str, project_name: str = "default") -> str:
+        """Materialize artifact to structured storage."""
+        from datetime import datetime
+        from flowyml.storage.materializers.base import get_materializer
+        import shutil
+        import pickle
+        import json
+
+        date_str = datetime.now().strftime("%Y-%m-%d")
+        # Structure: project / date / run_id / data / step / name
+        rel_path = Path(project_name) / date_str / run_id / "data" / step_name / name
+        full_path = self.base_path / rel_path
+
+        # Clean up if exists
+        if full_path.exists():
+            if full_path.is_dir():
+                shutil.rmtree(full_path)
+            else:
+                full_path.unlink()
+
+        full_path.mkdir(parents=True, exist_ok=True)
+
+        materializer = get_materializer(obj)
+        if materializer:
+            materializer.save(obj, full_path)
+        else:
+            # Fallback to pickle
+            with open(full_path / "data.pkl", "wb") as f:
+                pickle.dump(obj, f)
+            # Save metadata
+            with open(full_path / "metadata.json", "w") as f:
+                json.dump({"type": "pickle", "format": "pickle"}, f, indent=2)
+
+        return str(full_path)

flowyml/storage/materializers/__init__.py

@@ -0,0 +1,20 @@
+"""Materializers for framework-specific serialization."""
+
+from flowyml.storage.materializers.base import BaseMaterializer, materializer_registry
+from flowyml.storage.materializers.pytorch import PyTorchMaterializer
+from flowyml.storage.materializers.tensorflow import TensorFlowMaterializer
+from flowyml.storage.materializers.sklearn import SklearnMaterializer
+from flowyml.storage.materializers.pandas import PandasMaterializer
+from flowyml.storage.materializers.numpy import NumPyMaterializer
+from flowyml.storage.materializers.keras import KerasMaterializer
+
+__all__ = [
+    "BaseMaterializer",
+    "materializer_registry",
+    "PyTorchMaterializer",
+    "TensorFlowMaterializer",
+    "SklearnMaterializer",
+    "PandasMaterializer",
+    "NumPyMaterializer",
+    "KerasMaterializer",
+]

flowyml/storage/materializers/base.py

@@ -0,0 +1,133 @@
+"""Base materializer class and registry."""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any
+
+
+class BaseMaterializer(ABC):
+    """Base class for all materializers.
+
+    Materializers handle serialization/deserialization of specific types.
+    """
+
+    @abstractmethod
+    def save(self, obj: Any, path: Path) -> None:
+        """Save object to path.
+
+        Args:
+            obj: Object to save
+            path: Directory path where object should be saved
+        """
+        pass
+
+    @abstractmethod
+    def load(self, path: Path) -> Any:
+        """Load object from path.
+
+        Args:
+            path: Directory path from which to load object
+
+        Returns:
+            Loaded object
+        """
+        pass
+
+    @classmethod
+    @abstractmethod
+    def supported_types(cls) -> list[type]:
+        """Return list of types this materializer supports."""
+        pass
+
+    @classmethod
+    def can_handle(cls, obj: Any) -> bool:
+        """Check if this materializer can handle the given object.
+
+        Args:
+            obj: Object to check
+
+        Returns:
+            True if this materializer can handle the object
+        """
+        obj_type = type(obj)
+        for supported_type in cls.supported_types():
+            if isinstance(obj, supported_type):
+                return True
+            # Check module name for cross-import compatibility
+            if hasattr(supported_type, "__module__") and hasattr(obj_type, "__module__"):
+                if obj_type.__name__ == supported_type.__name__ and obj_type.__module__ == supported_type.__module__:
+                    return True
+        return False
+
+
+class MaterializerRegistry:
+    """Registry for materializers."""
+
+    def __init__(self):
+        self._materializers: dict[str, type[BaseMaterializer]] = {}
+
+    def register(self, materializer: type[BaseMaterializer]) -> None:
+        """Register a materializer.
+
+        Args:
+            materializer: Materializer class to register
+        """
+        for supported_type in materializer.supported_types():
+            key = f"{supported_type.__module__}.{supported_type.__name__}"
+            self._materializers[key] = materializer
+
+    def get_materializer(self, obj: Any) -> BaseMaterializer | None:
+        """Get materializer for object.
+
+        Args:
+            obj: Object to find materializer for
+
+        Returns:
+            Materializer instance or None if no suitable materializer found
+        """
+        # First try exact match
+        obj_type = type(obj)
+        key = f"{obj_type.__module__}.{obj_type.__name__}"
+
+        if key in self._materializers:
+            return self._materializers[key]()
+
+        # Then try checking all materializers
+        for materializer_cls in set(self._materializers.values()):
+            if materializer_cls.can_handle(obj):
+                return materializer_cls()
+
+        return None
+
+    def list_materializers(self) -> dict[str, type[BaseMaterializer]]:
+        """List all registered materializers.
+
+        Returns:
+            Dictionary mapping type names to materializer classes
+        """
+        return self._materializers.copy()
+
+
+# Global registry instance
+materializer_registry = MaterializerRegistry()
+
+
+def register_materializer(materializer: type[BaseMaterializer]) -> None:
+    """Register a materializer with the global registry.
+
+    Args:
+        materializer: Materializer class to register
+    """
+    materializer_registry.register(materializer)
+
+
+def get_materializer(obj: Any) -> BaseMaterializer | None:
+    """Get materializer for object from global registry.
+
+    Args:
+        obj: Object to find materializer for
+
+    Returns:
+        Materializer instance or None
+    """
+    return materializer_registry.get_materializer(obj)

flowyml/storage/materializers/keras.py

@@ -0,0 +1,185 @@
+"""Keras materializer for model serialization."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from flowyml.storage.materializers.base import BaseMaterializer, register_materializer
+
+try:
+    import keras
+
+    # Verify Keras has expected attributes
+    _ = keras.Model
+    _ = keras.Sequential
+    KERAS_AVAILABLE = True
+except (ImportError, AttributeError):
+    KERAS_AVAILABLE = False
+    keras = None
+
+
+if KERAS_AVAILABLE:
+
+    class KerasMaterializer(BaseMaterializer):
+        """Materializer for Keras models with full support for Sequential and Functional API."""
+
+        def save(self, obj: Any, path: Path) -> None:
+            """Save Keras model to path.
+
+            Args:
+                obj: Keras model (Sequential or Functional)
+                path: Directory path where model should be saved
+            """
+            path.mkdir(parents=True, exist_ok=True)
+
+            if isinstance(obj, keras.Model):
+                # Save model in Keras native format (.keras)
+                model_path = path / "model.keras"
+                try:
+                    obj.save(model_path)
+                except ValueError as e:
+                    if "save_format" in str(e):
+                        # Fallback for Keras 3 compatibility - try H5
+                        h5_path = path / "model.h5"
+                        obj.save(h5_path)
+                    else:
+                        raise e
+
+                # Save comprehensive metadata
+                metadata = {
+                    "type": "keras_model",
+                    "class_name": obj.__class__.__name__,
+                    "model_type": "Sequential" if isinstance(obj, keras.Sequential) else "Functional",
+                }
+
+                # Get input/output shapes
+                try:
+                    if hasattr(obj, "input_shape"):
+                        metadata["input_shape"] = str(obj.input_shape)
+                    if hasattr(obj, "output_shape"):
+                        metadata["output_shape"] = str(obj.output_shape)
+                except Exception:
+                    pass
+
+                # Get model architecture summary
+                try:
+                    import io
+
+                    string_buffer = io.StringIO()
+                    obj.summary(print_fn=lambda x: string_buffer.write(x + "\n"))
+                    metadata["summary"] = string_buffer.getvalue()
+                except Exception:
+                    pass
+
+                # Get layer information
+                try:
+                    metadata["num_layers"] = len(obj.layers)
+                    metadata["layer_names"] = [layer.name for layer in obj.layers]
+                    metadata["trainable_params"] = obj.count_params()
+                except Exception:
+                    pass
+
+                # Get optimizer info if compiled
+                try:
+                    if obj.optimizer:
+                        metadata["optimizer"] = obj.optimizer.__class__.__name__
+                        metadata["learning_rate"] = float(keras.backend.get_value(obj.optimizer.learning_rate))
+                except Exception:
+                    pass
+
+                # Get loss function if compiled
+                try:
+                    if obj.loss:
+                        if hasattr(obj.loss, "__name__"):
+                            metadata["loss"] = obj.loss.__name__
+                        else:
+                            metadata["loss"] = str(obj.loss)
+                except Exception:
+                    pass
+
+                # Get metrics if compiled
+                try:
+                    if obj.metrics:
+                        metadata["metrics"] = [m.name for m in obj.metrics]
+                except Exception:
+                    pass
+
+                # Save model config
+                try:
+                    config = obj.get_config()
+                    config_path = path / "config.json"
+                    with open(config_path, "w") as f:
+                        json.dump(config, f, indent=2, default=str)
+                    metadata["has_config"] = True
+                except Exception:
+                    metadata["has_config"] = False
+
+                # Save weights separately for backup - REMOVED for Keras 3 compatibility
+                # try:
+                #     weights_path = path / "weights.h5"
+                #     obj.save_weights(weights_path)
+                #     metadata["has_weights_backup"] = True
+                # except:
+                #     metadata["has_weights_backup"] = False
+
+                # Save metadata
+                metadata_path = path / "metadata.json"
+                with open(metadata_path, "w") as f:
+                    json.dump(metadata, f, indent=2, default=str)
+
+            else:
+                raise TypeError(f"Expected keras.Model, got {type(obj)}")
+
+        def load(self, path: Path) -> Any:
+            """Load Keras model from path.
+
+            Args:
+                path: Directory path from which to load model
+
+            Returns:
+                Loaded Keras model
+            """
+            # Load metadata
+            metadata_path = path / "metadata.json"
+            if metadata_path.exists():
+                with open(metadata_path) as f:
+                    json.load(f)
+            else:
+                pass
+
+            # Load model from .keras file
+            model_path = path / "model.keras"
+            if model_path.exists():
+                model = keras.models.load_model(model_path)
+                return model
+
+            # Fallback: try to load from SavedModel format
+            saved_model_path = path / "saved_model"
+            if saved_model_path.exists():
+                model = keras.models.load_model(saved_model_path)
+                return model
+
+            raise FileNotFoundError(f"No Keras model found at {path}")
+
+        @classmethod
+        def supported_types(cls) -> list[type]:
+            """Return Keras types supported by this materializer."""
+            return [keras.Model, keras.Sequential]
+
+    # Auto-register
+    register_materializer(KerasMaterializer)
+
+else:
+    # Placeholder when Keras not available
+    class KerasMaterializer(BaseMaterializer):
+        """Placeholder materializer when Keras is not installed."""
+
+        def save(self, obj: Any, path: Path) -> None:
+            raise ImportError("Keras/TensorFlow is not installed. Install with: pip install tensorflow")
+
+        def load(self, path: Path) -> Any:
+            raise ImportError("Keras/TensorFlow is not installed. Install with: pip install tensorflow")
+
+        @classmethod
+        def supported_types(cls) -> list[type]:
+            return []

flowyml/storage/materializers/numpy.py

@@ -0,0 +1,94 @@
+"""NumPy materializer for array serialization."""
+
+import json
+from pathlib import Path
+from typing import Any
+
+from flowyml.storage.materializers.base import BaseMaterializer, register_materializer
+
+try:
+    import numpy as np
+
+    NUMPY_AVAILABLE = True
+except ImportError:
+    NUMPY_AVAILABLE = False
+
+
+if NUMPY_AVAILABLE:
+
+    class NumPyMaterializer(BaseMaterializer):
+        """Materializer for NumPy arrays."""
+
+        def save(self, obj: Any, path: Path) -> None:
+            """Save NumPy array to path.
+
+            Args:
+                obj: NumPy array
+                path: Directory path where object should be saved
+            """
+            path.mkdir(parents=True, exist_ok=True)
+
+            # Save array
+            array_path = path / "array.npy"
+            np.save(array_path, obj)
+
+            # Save metadata
+            metadata = {
+                "type": "numpy_array",
+                "shape": list(obj.shape),
+                "dtype": str(obj.dtype),
+                "size": int(obj.size),
+                "ndim": int(obj.ndim),
+            }
+
+            # Add statistics for numerical arrays
+            if np.issubdtype(obj.dtype, np.number):
+                try:
+                    metadata["min"] = float(np.min(obj))
+                    metadata["max"] = float(np.max(obj))
+                    metadata["mean"] = float(np.mean(obj))
+                    metadata["std"] = float(np.std(obj))
+                except Exception:
+                    pass
+
+            with open(path / "metadata.json", "w") as f:
+                json.dump(metadata, f, indent=2)
+
+        def load(self, path: Path) -> Any:
+            """Load NumPy array from path.
+
+            Args:
+                path: Directory path from which to load object
+
+            Returns:
+                Loaded NumPy array
+            """
+            array_path = path / "array.npy"
+
+            if not array_path.exists():
+                raise FileNotFoundError(f"Array file not found at {array_path}")
+
+            return np.load(array_path)
+
+        @classmethod
+        def supported_types(cls) -> list[type]:
+            """Return NumPy types supported by this materializer."""
+            return [np.ndarray]
+
+    # Auto-register
+    register_materializer(NumPyMaterializer)
+
+else:
+    # Placeholder when NumPy not available
+    class NumPyMaterializer(BaseMaterializer):
+        """Placeholder materializer when NumPy is not installed."""
+
+        def save(self, obj: Any, path: Path) -> None:
+            raise ImportError("NumPy is not installed. Install with: pip install numpy")
+
+        def load(self, path: Path) -> Any:
+            raise ImportError("NumPy is not installed. Install with: pip install numpy")
+
+        @classmethod
+        def supported_types(cls) -> list[type]:
+            return []