flowyml 1.1.0__py3-none-any.whl

flowyml/core/project.py

@@ -0,0 +1,253 @@
+"""Multi-tenancy and project organization."""
+
+import json
+from pathlib import Path
+from typing import Any, TYPE_CHECKING
+from datetime import datetime
+
+if TYPE_CHECKING:
+    from flowyml.core.pipeline import Pipeline
+
+
+class Project:
+    """Project for organizing pipelines, runs, and artifacts.
+
+    Provides multi-tenancy and better organization.
+
+    Examples:
+        >>> from flowyml import Project, Pipeline, step
+        >>> # Create project
+        >>> project = Project("recommendation_system")
+        >>> # Create pipeline in project
+        >>> pipeline = project.create_pipeline("training")
+        >>> pipeline.add_step(...)
+        >>> result = pipeline.run()
+        >>> # List all runs in project
+        >>> runs = project.list_runs()
+        >>> # Get artifacts
+        >>> artifacts = project.get_artifacts()
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: str = "",
+        projects_dir: str = ".flowyml/projects",
+    ):
+        self.name = name
+        self.description = description
+
+        # Project directory
+        self.project_dir = Path(projects_dir) / name
+        self.project_dir.mkdir(parents=True, exist_ok=True)
+
+        # Sub-directories
+        self.pipelines_dir = self.project_dir / "pipelines"
+        self.pipelines_dir.mkdir(exist_ok=True)
+
+        self.runs_dir = self.project_dir / "runs"
+        self.runs_dir.mkdir(exist_ok=True)
+
+        self.artifacts_dir = self.project_dir / "artifacts"
+        self.artifacts_dir.mkdir(exist_ok=True)
+
+        # Metadata
+        self.metadata_file = self.project_dir / "project.json"
+        self._load_or_create_metadata()
+
+        # Storage
+        from flowyml.storage.metadata import SQLiteMetadataStore
+
+        db_path = str(self.project_dir / "metadata.db")
+        self.metadata_store = SQLiteMetadataStore(db_path)
+
+    def _load_or_create_metadata(self) -> None:
+        """Load or create project metadata."""
+        if self.metadata_file.exists():
+            with open(self.metadata_file) as f:
+                self.metadata = json.load(f)
+        else:
+            self.metadata = {
+                "name": self.name,
+                "description": self.description,
+                "created_at": datetime.now().isoformat(),
+                "pipelines": [],
+                "tags": {},
+            }
+            self._save_metadata()
+
+    def _save_metadata(self) -> None:
+        """Save project metadata."""
+        with open(self.metadata_file, "w") as f:
+            json.dump(self.metadata, f, indent=2)
+
+    def create_pipeline(self, name: str, **kwargs) -> "Pipeline":
+        """Create a pipeline in this project.
+
+        Args:
+            name: Pipeline name
+            **kwargs: Additional arguments for Pipeline
+
+        Returns:
+            Pipeline instance
+        """
+        from flowyml.core.pipeline import Pipeline
+
+        # Create pipeline with project-specific settings
+        pipeline = Pipeline(
+            name=name,
+            cache_dir=str(self.artifacts_dir / "cache"),
+            **kwargs,
+        )
+
+        # Override runs directory to project runs
+        pipeline.runs_dir = self.runs_dir
+
+        # Use project metadata store
+        pipeline.metadata_store = self.metadata_store
+
+        # Register pipeline
+        if name not in self.metadata["pipelines"]:
+            self.metadata["pipelines"].append(name)
+            self._save_metadata()
+
+        return pipeline
+
+    def list_runs(
+        self,
+        pipeline_name: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """List all runs in this project.
+
+        Args:
+            pipeline_name: Filter by pipeline name
+            limit: Maximum number of runs
+
+        Returns:
+            List of run metadata
+        """
+        if pipeline_name:
+            runs = self.metadata_store.query(pipeline_name=pipeline_name)
+        else:
+            runs = self.metadata_store.list_runs(limit=limit)
+
+        return runs
+
+    def get_artifacts(
+        self,
+        artifact_type: str | None = None,
+        limit: int = 100,
+    ) -> list[dict[str, Any]]:
+        """Get artifacts in this project.
+
+        Args:
+            artifact_type: Filter by type
+            limit: Maximum number of artifacts
+
+        Returns:
+            List of artifact metadata
+        """
+        filters = {}
+        if artifact_type:
+            filters["type"] = artifact_type
+
+        return self.metadata_store.list_assets(limit=limit, **filters)
+
+    def get_pipelines(self) -> list[str]:
+        """Get list of pipelines in this project."""
+        return self.metadata["pipelines"]
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get project statistics."""
+        stats = self.metadata_store.get_statistics()
+        stats["project_name"] = self.name
+        stats["pipelines"] = len(self.metadata["pipelines"])
+        return stats
+
+    def add_tag(self, key: str, value: str) -> None:
+        """Add a tag to the project."""
+        self.metadata["tags"][key] = value
+        self._save_metadata()
+
+    def get_tags(self) -> dict[str, str]:
+        """Get project tags."""
+        return self.metadata["tags"]
+
+    def export_metadata(self, output_file: str) -> None:
+        """Export project metadata."""
+        export_data = {
+            "project": self.metadata,
+            "runs": self.list_runs(),
+            "artifacts": self.get_artifacts(),
+            "stats": self.get_stats(),
+        }
+
+        with open(output_file, "w") as f:
+            json.dump(export_data, f, indent=2)
+
+    def __repr__(self) -> str:
+        return f"Project(name='{self.name}', pipelines={len(self.metadata['pipelines'])})"
+
+
+class ProjectManager:
+    """Manage multiple projects.
+
+    Examples:
+        >>> from flowyml import ProjectManager
+        >>> manager = ProjectManager()
+        >>> # Create projects
+        >>> rec_sys = manager.create_project("recommendation_system")
+        >>> fraud = manager.create_project("fraud_detection")
+        >>> # List all projects
+        >>> projects = manager.list_projects()
+        >>> # Get project
+        >>> project = manager.get_project("recommendation_system")
+    """
+
+    def __init__(self, projects_dir: str = ".flowyml/projects"):
+        self.projects_dir = Path(projects_dir)
+        self.projects_dir.mkdir(parents=True, exist_ok=True)
+
+    def create_project(self, name: str, description: str = "") -> Project:
+        """Create a new project."""
+        # Fix: Pass the projects directory itself, not the parent
+        project = Project(name, description, str(self.projects_dir))
+        return project
+
+    def get_project(self, name: str) -> Project | None:
+        """Get an existing project."""
+        project_dir = self.projects_dir / name
+        if not project_dir.exists():
+            return None
+        # Fix: Pass the projects directory itself
+        return Project(name, projects_dir=str(self.projects_dir))
+
+    def list_projects(self) -> list[dict[str, Any]]:
+        """List all projects."""
+        projects = []
+        if not self.projects_dir.exists():
+            return projects
+
+        for project_dir in self.projects_dir.iterdir():
+            if project_dir.is_dir():
+                metadata_file = project_dir / "project.json"
+                if metadata_file.exists():
+                    try:
+                        with open(metadata_file) as f:
+                            metadata = json.load(f)
+                            projects.append(metadata)
+                    except Exception:
+                        pass
+        return projects
+
+    def delete_project(self, name: str, confirm: bool = False) -> None:
+        """Delete a project (requires confirmation)."""
+        if not confirm:
+            return
+
+        project_dir = self.projects_dir / name
+        if project_dir.exists():
+            import shutil
+
+            shutil.rmtree(project_dir)

flowyml/core/resources.py

@@ -0,0 +1,424 @@
+"""Resource specification for pipeline steps.
+
+This module provides orchestrator-agnostic resource specification for flowyml pipeline steps,
+including CPU, memory, GPU, storage, and node affinity requirements.
+"""
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+import re
+
+
+@dataclass
+class GPUConfig:
+    """GPU configuration specification.
+
+    Args:
+        gpu_type: GPU type/model (e.g., 'nvidia-tesla-v100', 'nvidia-a100')
+        count: Number of GPUs required
+        memory: GPU memory per device (e.g., '16Gi', '32Gi')
+
+    Examples:
+        >>> gpu = GPUConfig(gpu_type="nvidia-tesla-v100", count=2, memory="16Gi")
+        >>> gpu = GPUConfig(gpu_type="nvidia-a100", count=4)
+    """
+
+    gpu_type: str
+    count: int = 1
+    memory: Optional[str] = None
+
+    def __post_init__(self):
+        """Validate GPU configuration."""
+        if self.count < 1:
+            msg = f"GPU count must be >= 1, got {self.count}"
+            raise ValueError(msg)
+        if self.memory and not self._is_valid_memory(self.memory):
+            msg = f"Invalid GPU memory format: {self.memory}"
+            raise ValueError(msg)
+
+    @staticmethod
+    def _is_valid_memory(memory: str) -> bool:
+        """Check if memory string is valid (e.g., '16Gi', '32768Mi')."""
+        return bool(re.match(r"^\d+(\.\d+)?(Ki|Mi|Gi|Ti|K|M|G|T)$", memory))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "type": self.gpu_type,
+            "count": self.count,
+            "memory": self.memory,
+        }
+
+    def merge_with(self, other: "GPUConfig") -> "GPUConfig":
+        """Merge with another GPU config, taking max count and best GPU type.
+
+        Args:
+            other: Another GPUConfig to merge with
+
+        Returns:
+            New GPUConfig with merged specifications
+        """
+        # Prefer A100 > V100 > T4 > other, or just take first if unknown
+        gpu_hierarchy = ["nvidia-a100", "nvidia-tesla-a100", "nvidia-tesla-v100", "nvidia-v100", "nvidia-t4"]
+
+        best_type = self.gpu_type
+        for gpu_type in gpu_hierarchy:
+            if gpu_type in self.gpu_type.lower():
+                self_rank = gpu_hierarchy.index(gpu_type)
+                break
+        else:
+            self_rank = 999
+
+        for gpu_type in gpu_hierarchy:
+            if gpu_type in other.gpu_type.lower():
+                other_rank = gpu_hierarchy.index(gpu_type)
+                break
+        else:
+            other_rank = 999
+
+        if other_rank < self_rank:
+            best_type = other.gpu_type
+
+        # Take max count
+        max_count = max(self.count, other.count)
+
+        # Take max memory if both specified
+        max_memory = None
+        if self.memory and other.memory:
+            max_memory = self._compare_memory(self.memory, other.memory)
+        elif self.memory:
+            max_memory = self.memory
+        elif other.memory:
+            max_memory = other.memory
+
+        return GPUConfig(
+            gpu_type=best_type,
+            count=max_count,
+            memory=max_memory,
+        )
+
+    @staticmethod
+    def _compare_memory(mem1: str, mem2: str) -> str:
+        """Return the larger memory specification."""
+
+        # Simple comparison - convert to bytes and compare
+        def to_bytes(mem: str) -> int:
+            import re
+
+            match = re.match(r"^(\d+(?:\.\d+)?)(Ki|Mi|Gi|Ti|K|M|G|T)?$", mem)
+            if not match:
+                return 0
+            value, unit = float(match.group(1)), match.group(2) or ""
+            multipliers = {
+                "Ki": 1024,
+                "Mi": 1024**2,
+                "Gi": 1024**3,
+                "Ti": 1024**4,
+                "K": 1000,
+                "M": 1000**2,
+                "G": 1000**3,
+                "T": 1000**4,
+                "": 1,
+            }
+            return int(value * multipliers.get(unit, 1))
+
+        return mem1 if to_bytes(mem1) >= to_bytes(mem2) else mem2
+
+
+@dataclass
+class NodeAffinity:
+    """Node affinity and anti-affinity rules.
+
+    Args:
+        required: Required node labels (hard constraints)
+        preferred: Preferred node labels (soft constraints)
+        tolerations: Tolerations for node taints
+
+    Examples:
+        >>> affinity = NodeAffinity(
+        ...     required={"cloud.google.com/gke-nodepool": "gpu-pool"}, preferred={"instance-type": "n1-standard-8"}
+        ... )
+    """
+
+    required: dict[str, str] = field(default_factory=dict)
+    preferred: dict[str, str] = field(default_factory=dict)
+    tolerations: list[dict[str, str]] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        return {
+            "required": self.required,
+            "preferred": self.preferred,
+            "tolerations": self.tolerations,
+        }
+
+    def merge_with(self, other: "NodeAffinity") -> "NodeAffinity":
+        """Merge with another node affinity, combining constraints.
+
+        Args:
+            other: Another NodeAffinity to merge with
+
+        Returns:
+            New NodeAffinity with merged constraints
+        """
+        # Merge required labels (intersection - both must be satisfied)
+        merged_required = {**self.required, **other.required}
+
+        # Merge preferred labels (union - prefer either)
+        merged_preferred = {**self.preferred, **other.preferred}
+
+        # Merge tolerations (union - tolerate all)
+        merged_tolerations = list(self.tolerations)
+        for tol in other.tolerations:
+            if tol not in merged_tolerations:
+                merged_tolerations.append(tol)
+
+        return NodeAffinity(
+            required=merged_required,
+            preferred=merged_preferred,
+            tolerations=merged_tolerations,
+        )
+
+
+@dataclass
+class ResourceRequirements:
+    """Resource requirements for a pipeline step.
+
+    Orchestrator-agnostic resource specification that can be translated to
+    platform-specific formats (Kubernetes, Vertex AI, SageMaker, etc.).
+
+    Args:
+        cpu: CPU cores (e.g., "2", "500m", "2.5")
+        memory: RAM amount (e.g., "4Gi", "8192Mi", "16G")
+        storage: Ephemeral storage (e.g., "100Gi", "50G")
+        gpu: GPU configuration
+        node_affinity: Node selection rules
+
+    Examples:
+        >>> # Simple CPU/memory
+        >>> resources = ResourceRequirements(cpu="2", memory="4Gi")
+
+        >>> # With GPU
+        >>> resources = ResourceRequirements(cpu="4", memory="16Gi", gpu=GPUConfig(gpu_type="nvidia-tesla-v100", count=2))
+
+        >>> # With node affinity
+        >>> resources = ResourceRequirements(
+        ...     cpu="8",
+        ...     memory="32Gi",
+        ...     node_affinity=NodeAffinity(
+        ...         required={"gpu": "true"}, tolerations=[{"key": "nvidia.com/gpu", "operator": "Exists"}]
+        ...     ),
+        ... )
+    """
+
+    cpu: Optional[str] = None
+    memory: Optional[str] = None
+    storage: Optional[str] = None
+    gpu: Optional[GPUConfig] = None
+    node_affinity: Optional[NodeAffinity] = None
+
+    def __post_init__(self):
+        """Validate resource specifications."""
+        if self.cpu and not self._is_valid_cpu(self.cpu):
+            msg = f"Invalid CPU format: {self.cpu}"
+            raise ValueError(msg)
+        if self.memory and not self._is_valid_memory(self.memory):
+            msg = f"Invalid memory format: {self.memory}"
+            raise ValueError(msg)
+        if self.storage and not self._is_valid_memory(self.storage):
+            msg = f"Invalid storage format: {self.storage}"
+            raise ValueError(msg)
+
+    @staticmethod
+    def _is_valid_cpu(cpu: str) -> bool:
+        """Check if CPU string is valid (e.g., '2', '500m', '2.5')."""
+        return bool(re.match(r"^\d+(\.\d+)?m?$", cpu))
+
+    @staticmethod
+    def _is_valid_memory(memory: str) -> bool:
+        """Check if memory string is valid (e.g., '16Gi', '32768Mi')."""
+        return bool(re.match(r"^\d+(\.\d+)?(Ki|Mi|Gi|Ti|K|M|G|T|B)?$", memory))
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary representation."""
+        result = {}
+        if self.cpu:
+            result["cpu"] = self.cpu
+        if self.memory:
+            result["memory"] = self.memory
+        if self.storage:
+            result["storage"] = self.storage
+        if self.gpu:
+            result["gpu"] = self.gpu.to_dict()
+        if self.node_affinity:
+            result["node_affinity"] = self.node_affinity.to_dict()
+        return result
+
+    def has_gpu(self) -> bool:
+        """Check if GPU resources are requested."""
+        return self.gpu is not None
+
+    def get_gpu_count(self) -> int:
+        """Get total number of GPUs requested."""
+        return self.gpu.count if self.gpu else 0
+
+    @staticmethod
+    def _compare_cpu(cpu1: str, cpu2: str) -> str:
+        """Return the larger CPU requirement.
+
+        Args:
+            cpu1: First CPU specification (e.g., "2", "500m")
+            cpu2: Second CPU specification
+
+        Returns:
+            The larger CPU specification
+        """
+
+        def to_millicores(cpu: str) -> int:
+            if cpu.endswith("m"):
+                return int(cpu[:-1])
+            return int(float(cpu) * 1000)
+
+        return cpu1 if to_millicores(cpu1) >= to_millicores(cpu2) else cpu2
+
+    @staticmethod
+    def _compare_memory(mem1: str, mem2: str) -> str:
+        """Return the larger memory requirement.
+
+        Args:
+            mem1: First memory specification (e.g., "4Gi", "8192Mi")
+            mem2: Second memory specification
+
+        Returns:
+            The larger memory specification (in original format)
+        """
+        import re
+
+        def to_bytes(mem: str) -> int:
+            match = re.match(r"^(\d+(?:\.\d+)?)(Ki|Mi|Gi|Ti|K|M|G|T|B)?$", mem)
+            if not match:
+                return 0
+            value, unit = float(match.group(1)), match.group(2) or "B"
+            multipliers = {
+                "Ki": 1024,
+                "Mi": 1024**2,
+                "Gi": 1024**3,
+                "Ti": 1024**4,
+                "K": 1000,
+                "M": 1000**2,
+                "G": 1000**3,
+                "T": 1000**4,
+                "B": 1,
+                "": 1,
+            }
+            return int(value * multipliers.get(unit, 1))
+
+        bytes1 = to_bytes(mem1)
+        bytes2 = to_bytes(mem2)
+
+        # Return whichever is larger, but keep original format
+        return mem1 if bytes1 >= bytes2 else mem2
+
+    def merge_with(self, other: "ResourceRequirements") -> "ResourceRequirements":
+        """Merge with another ResourceRequirements, taking maximum of each.
+
+        This is used when grouping steps to aggregate their resource needs.
+        Strategy:
+        - CPU: Take maximum
+        - Memory: Take maximum
+        - Storage: Take maximum
+        - GPU: Merge configs (max count, best type)
+        - Node affinity: Merge constraints
+
+        Args:
+            other: Another ResourceRequirements to merge with
+
+        Returns:
+            New ResourceRequirements with merged specifications
+        """
+        # Merge CPU
+        merged_cpu = None
+        if self.cpu and other.cpu:
+            merged_cpu = self._compare_cpu(self.cpu, other.cpu)
+        elif self.cpu:
+            merged_cpu = self.cpu
+        elif other.cpu:
+            merged_cpu = other.cpu
+
+        # Merge memory
+        merged_memory = None
+        if self.memory and other.memory:
+            merged_memory = self._compare_memory(self.memory, other.memory)
+        elif self.memory:
+            merged_memory = self.memory
+        elif other.memory:
+            merged_memory = other.memory
+
+        # Merge storage
+        merged_storage = None
+        if self.storage and other.storage:
+            merged_storage = self._compare_memory(self.storage, other.storage)
+        elif self.storage:
+            merged_storage = self.storage
+        elif other.storage:
+            merged_storage = other.storage
+
+        # Merge GPU
+        merged_gpu = None
+        if self.gpu and other.gpu:
+            merged_gpu = self.gpu.merge_with(other.gpu)
+        elif self.gpu:
+            merged_gpu = self.gpu
+        elif other.gpu:
+            merged_gpu = other.gpu
+
+        # Merge node affinity
+        merged_affinity = None
+        if self.node_affinity and other.node_affinity:
+            merged_affinity = self.node_affinity.merge_with(other.node_affinity)
+        elif self.node_affinity:
+            merged_affinity = self.node_affinity
+        elif other.node_affinity:
+            merged_affinity = other.node_affinity
+
+        return ResourceRequirements(
+            cpu=merged_cpu,
+            memory=merged_memory,
+            storage=merged_storage,
+            gpu=merged_gpu,
+            node_affinity=merged_affinity,
+        )
+
+
+def resources(
+    cpu: Optional[str] = None,
+    memory: Optional[str] = None,
+    storage: Optional[str] = None,
+    gpu: Optional[GPUConfig] = None,
+    node_affinity: Optional[NodeAffinity] = None,
+) -> ResourceRequirements:
+    """Create a ResourceRequirements object with validation.
+
+    Convenience function for creating resource specifications with cleaner syntax.
+
+    Args:
+        cpu: CPU cores (e.g., "2", "500m")
+        memory: RAM amount (e.g., "4Gi", "8192Mi")
+        storage: Ephemeral storage (e.g., "100Gi")
+        gpu: GPU configuration
+        node_affinity: Node selection rules
+
+    Returns:
+        Validated ResourceRequirements object
+
+    Examples:
+        >>> req = resources(cpu="2", memory="4Gi")
+        >>> req = resources(cpu="4", memory="16Gi", gpu=GPUConfig(gpu_type="nvidia-v100", count=2))
+    """
+    return ResourceRequirements(
+        cpu=cpu,
+        memory=memory,
+        storage=storage,
+        gpu=gpu,
+        node_affinity=node_affinity,
+    )