flowyml 1.1.0__py3-none-any.whl

flowyml/stacks/gcp.py

@@ -0,0 +1,499 @@
+"""GCP Stack - Google Cloud Platform integration for flowyml.
+
+This module provides GCP-specific implementations for running pipelines
+on Google Cloud Platform using Vertex AI, Cloud Storage, and Container Registry.
+"""
+
+from typing import Any
+
+from flowyml.stacks.base import Stack
+from flowyml.stacks.components import (
+    Orchestrator,
+    ArtifactStore,
+    ContainerRegistry,
+    ResourceConfig,
+    DockerConfig,
+)
+
+
+class VertexAIOrchestrator(Orchestrator):
+    """Vertex AI orchestrator for running pipelines on Google Cloud.
+
+    This orchestrator submits pipeline jobs to Vertex AI Pipelines,
+    allowing for scalable, managed execution in the cloud.
+
+    Example:
+        ```python
+        from flowyml.stacks.gcp import VertexAIOrchestrator
+
+        orchestrator = VertexAIOrchestrator(
+            project_id="my-gcp-project", region="us-central1", service_account="my-sa@my-project.iam.gserviceaccount.com"
+        )
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str = "vertex_ai",
+        project_id: str | None = None,
+        region: str = "us-central1",
+        service_account: str | None = None,
+        network: str | None = None,
+        encryption_key: str | None = None,
+    ):
+        """Initialize Vertex AI orchestrator.
+
+        Args:
+            name: Name of the orchestrator
+            project_id: GCP project ID
+            region: GCP region for Vertex AI
+            service_account: Service account email for job execution
+            network: VPC network for jobs
+            encryption_key: Customer-managed encryption key
+        """
+        super().__init__(name)
+        self.project_id = project_id
+        self.region = region
+        self.service_account = service_account
+        self.network = network
+        self.encryption_key = encryption_key
+
+    def validate(self) -> bool:
+        """Validate Vertex AI configuration."""
+        if not self.project_id:
+            raise ValueError("project_id is required for VertexAIOrchestrator")
+
+        # Check if google-cloud-aiplatform is installed
+        import importlib.util
+
+        if importlib.util.find_spec("google.cloud.aiplatform") is not None:
+            return True
+        raise ImportError(
+            "google-cloud-aiplatform is required for VertexAIOrchestrator. "
+            "Install with: pip install google-cloud-aiplatform",
+        )
+
+    def run_pipeline(
+        self,
+        pipeline: Any,
+        resources: ResourceConfig | None = None,
+        docker_config: DockerConfig | None = None,
+        **kwargs,
+    ) -> str:
+        """Run pipeline on Vertex AI.
+
+        Args:
+            pipeline: Pipeline to run
+            resources: Resource configuration
+            docker_config: Docker configuration
+            **kwargs: Additional arguments
+
+        Returns:
+            Job ID
+        """
+        from google.cloud import aiplatform
+
+        # Initialize Vertex AI
+        aiplatform.init(project=self.project_id, location=self.region)
+
+        # Create custom job
+        job_display_name = f"{pipeline.name}-{pipeline.run_id}"
+
+        # Build worker pool specs
+        worker_pool_specs = self._build_worker_pool_specs(
+            docker_config=docker_config,
+            resources=resources,
+        )
+
+        # Create and run custom job
+        job = aiplatform.CustomJob(
+            display_name=job_display_name,
+            worker_pool_specs=worker_pool_specs,
+            service_account=self.service_account,
+            network=self.network,
+            encryption_spec_key_name=self.encryption_key,
+        )
+
+        job.run(sync=False)
+
+        return job.resource_name
+
+    def get_run_status(self, run_id: str) -> str:
+        """Get status of a Vertex AI job."""
+        from google.cloud import aiplatform
+
+        job = aiplatform.CustomJob(run_id)
+        return job.state.name
+
+    def _build_worker_pool_specs(
+        self,
+        docker_config: DockerConfig | None,
+        resources: ResourceConfig | None,
+    ) -> list[dict]:
+        """Build worker pool specifications for Vertex AI."""
+        # Default resources
+        if resources is None:
+            resources = ResourceConfig()
+
+        machine_spec = {
+            "machine_type": resources.machine_type or "n1-standard-4",
+        }
+
+        if resources.gpu:
+            machine_spec["accelerator_type"] = resources.gpu
+            machine_spec["accelerator_count"] = resources.gpu_count
+
+        container_spec = {
+            "image_uri": (docker_config.image if docker_config else "gcr.io/flowyml/flowyml:latest"),
+            "command": ["python", "-m", "flowyml.cli.run"],
+        }
+
+        if docker_config and docker_config.env_vars:
+            container_spec["env"] = [{"name": k, "value": v} for k, v in docker_config.env_vars.items()]
+
+        return [
+            {
+                "replica_count": 1,
+                "machine_spec": machine_spec,
+                "container_spec": container_spec,
+            },
+        ]
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "name": self.name,
+            "type": "vertex_ai",
+            "project_id": self.project_id,
+            "region": self.region,
+            "service_account": self.service_account,
+            "network": self.network,
+        }
+
+
+class GCSArtifactStore(ArtifactStore):
+    """Google Cloud Storage artifact store.
+
+    Stores pipeline artifacts in Google Cloud Storage buckets.
+
+    Example:
+        ```python
+        from flowyml.stacks.gcp import GCSArtifactStore
+
+        artifact_store = GCSArtifactStore(bucket_name="my-flowyml-artifacts", project_id="my-gcp-project")
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str = "gcs",
+        bucket_name: str | None = None,
+        project_id: str | None = None,
+        prefix: str = "flowyml",
+    ):
+        """Initialize GCS artifact store.
+
+        Args:
+            name: Name of the artifact store
+            bucket_name: GCS bucket name
+            project_id: GCP project ID
+            prefix: Prefix for all artifacts in bucket
+        """
+        super().__init__(name)
+        self.bucket_name = bucket_name
+        self.project_id = project_id
+        self.prefix = prefix
+
+    def validate(self) -> bool:
+        """Validate GCS configuration."""
+        if not self.bucket_name:
+            raise ValueError("bucket_name is required for GCSArtifactStore")
+
+        # Check if google-cloud-storage is installed
+        import importlib.util
+
+        if importlib.util.find_spec("google.cloud.storage") is not None:
+            return True
+        raise ImportError(
+            "google-cloud-storage is required for GCSArtifactStore. Install with: pip install google-cloud-storage",
+        )
+
+    def save(self, artifact: Any, path: str) -> str:
+        """Save artifact to GCS."""
+        from google.cloud import storage
+        import pickle
+
+        client = storage.Client(project=self.project_id)
+        bucket = client.bucket(self.bucket_name)
+
+        # Full path with prefix
+        full_path = f"{self.prefix}/{path}"
+        blob = bucket.blob(full_path)
+
+        # Serialize and upload
+        data = pickle.dumps(artifact)
+        blob.upload_from_string(data)
+
+        return f"gs://{self.bucket_name}/{full_path}"
+
+    def load(self, path: str) -> Any:
+        """Load artifact from GCS."""
+        from google.cloud import storage
+        import pickle
+
+        client = storage.Client(project=self.project_id)
+        bucket = client.bucket(self.bucket_name)
+
+        # Handle both full gs:// URIs and relative paths
+        if path.startswith("gs://"):
+            # Extract bucket and path from URI
+            parts = path.replace("gs://", "").split("/", 1)
+            blob_path = parts[1] if len(parts) > 1 else ""
+        else:
+            blob_path = f"{self.prefix}/{path}"
+
+        blob = bucket.blob(blob_path)
+        data = blob.download_as_bytes()
+
+        return pickle.loads(data)
+
+    def exists(self, path: str) -> bool:
+        """Check if artifact exists in GCS."""
+        from google.cloud import storage
+
+        client = storage.Client(project=self.project_id)
+        bucket = client.bucket(self.bucket_name)
+
+        full_path = f"{self.prefix}/{path}"
+        blob = bucket.blob(full_path)
+
+        return blob.exists()
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "name": self.name,
+            "type": "gcs",
+            "bucket_name": self.bucket_name,
+            "project_id": self.project_id,
+            "prefix": self.prefix,
+        }
+
+
+class GCRContainerRegistry(ContainerRegistry):
+    """Google Container Registry integration.
+
+    Manages Docker images in Google Container Registry or Artifact Registry.
+
+    Example:
+        ```python
+        from flowyml.stacks.gcp import GCRContainerRegistry
+
+        registry = GCRContainerRegistry(project_id="my-gcp-project", registry_uri="gcr.io/my-gcp-project")
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str = "gcr",
+        project_id: str | None = None,
+        registry_uri: str | None = None,
+        region: str | None = None,  # For Artifact Registry
+    ):
+        """Initialize GCR container registry.
+
+        Args:
+            name: Name of the registry
+            project_id: GCP project ID
+            registry_uri: Full registry URI (e.g., gcr.io/project-id)
+            region: Region for Artifact Registry (e.g., us-central1)
+        """
+        super().__init__(name)
+        self.project_id = project_id
+        self.registry_uri = registry_uri or f"gcr.io/{project_id}"
+        self.region = region
+
+    def validate(self) -> bool:
+        """Validate registry configuration."""
+        if not self.project_id:
+            raise ValueError("project_id is required for GCRContainerRegistry")
+        return True
+
+    def push_image(self, image_name: str, tag: str = "latest") -> str:
+        """Push Docker image to GCR.
+
+        Args:
+            image_name: Name of the image
+            tag: Image tag
+
+        Returns:
+            Full image URI
+        """
+        import subprocess
+
+        full_uri = self.get_image_uri(image_name, tag)
+
+        # Tag image
+        subprocess.run(
+            ["docker", "tag", f"{image_name}:{tag}", full_uri],
+            check=True,
+        )
+
+        # Push to registry
+        subprocess.run(
+            ["docker", "push", full_uri],
+            check=True,
+        )
+
+        return full_uri
+
+    def pull_image(self, image_name: str, tag: str = "latest") -> None:
+        """Pull Docker image from GCR."""
+        import subprocess
+
+        full_uri = self.get_image_uri(image_name, tag)
+        subprocess.run(
+            ["docker", "pull", full_uri],
+            check=True,
+        )
+
+    def get_image_uri(self, image_name: str, tag: str = "latest") -> str:
+        """Get full URI for an image."""
+        return f"{self.registry_uri}/{image_name}:{tag}"
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary."""
+        return {
+            "name": self.name,
+            "type": "gcr",
+            "project_id": self.project_id,
+            "registry_uri": self.registry_uri,
+            "region": self.region,
+        }
+
+
+class GCPStack(Stack):
+    """Complete GCP stack for running flowyml pipelines on Google Cloud Platform.
+
+    This stack integrates:
+    - Vertex AI for orchestration
+    - Google Cloud Storage for artifact storage
+    - Google Container Registry for Docker images
+    - Cloud SQL or Firestore for metadata storage
+
+    Example:
+        ```python
+        from flowyml.stacks.gcp import GCPStack
+        from flowyml.stacks.components import ResourceConfig, DockerConfig
+        from flowyml import Pipeline, step
+
+        # Create GCP stack
+        stack = GCPStack(
+            name="production",
+            project_id="my-gcp-project",
+            region="us-central1",
+            bucket_name="my-flowyml-artifacts",
+            registry_uri="gcr.io/my-gcp-project",
+        )
+
+        # Define resource requirements
+        resources = ResourceConfig(cpu="4", memory="16Gi", gpu="nvidia-tesla-t4", gpu_count=1, machine_type="n1-highmem-4")
+
+        # Define Docker configuration
+        docker_config = DockerConfig(
+            image="gcr.io/my-gcp-project/ml-pipeline:v1",
+            requirements=["tensorflow>=2.12.0", "scikit-learn>=1.0.0"],
+            env_vars={"PYTHONUNBUFFERED": "1"},
+        )
+
+
+        # Run pipeline on GCP
+        @step
+        def train_model():
+            # Your training code
+            pass
+
+
+        pipeline = Pipeline("training", stack=stack)
+        pipeline.add_step(train_model)
+
+        result = pipeline.run(resources=resources, docker_config=docker_config)
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str = "gcp",
+        project_id: str | None = None,
+        region: str = "us-central1",
+        bucket_name: str | None = None,
+        registry_uri: str | None = None,
+        service_account: str | None = None,
+        metadata_store: Any | None = None,
+    ):
+        """Initialize GCP stack.
+
+        Args:
+            name: Stack name
+            project_id: GCP project ID
+            region: GCP region
+            bucket_name: GCS bucket for artifacts
+            registry_uri: Container registry URI
+            service_account: Service account for job execution
+            metadata_store: Metadata store (optional, defaults to local SQLite)
+        """
+        # Create GCP components
+        orchestrator = VertexAIOrchestrator(
+            project_id=project_id,
+            region=region,
+            service_account=service_account,
+        )
+
+        artifact_store = GCSArtifactStore(
+            bucket_name=bucket_name,
+            project_id=project_id,
+        )
+
+        container_registry = GCRContainerRegistry(
+            project_id=project_id,
+            registry_uri=registry_uri,
+            region=region,
+        )
+
+        # Use local metadata store if not provided
+        if metadata_store is None:
+            from flowyml.storage.metadata import SQLiteMetadataStore
+
+            metadata_store = SQLiteMetadataStore()
+
+        # Initialize base stack
+        super().__init__(
+            name=name,
+            executor=None,  # Vertex AI handles execution
+            artifact_store=artifact_store,
+            metadata_store=metadata_store,
+            container_registry=container_registry,
+            orchestrator=orchestrator,
+        )
+
+        self.project_id = project_id
+        self.region = region
+
+    def validate(self) -> bool:
+        """Validate all GCP stack components."""
+        self.orchestrator.validate()
+        self.artifact_store.validate()
+        self.container_registry.validate()
+        return True
+
+    def to_dict(self) -> dict[str, Any]:
+        """Convert stack configuration to dictionary."""
+        return {
+            "name": self.name,
+            "type": "gcp",
+            "project_id": self.project_id,
+            "region": self.region,
+            "orchestrator": self.orchestrator.to_dict(),
+            "artifact_store": self.artifact_store.to_dict(),
+            "container_registry": self.container_registry.to_dict(),
+        }

flowyml/stacks/local.py

@@ -0,0 +1,112 @@
+"""Local Stack - For local development and testing."""
+
+from pathlib import Path
+from flowyml.stacks.base import Stack
+from flowyml.core.executor import LocalExecutor
+from flowyml.storage.artifacts import LocalArtifactStore
+from flowyml.storage.metadata import SQLiteMetadataStore
+
+
+class LocalStack(Stack):
+    """Local stack for development and testing.
+
+    This stack uses:
+    - LocalExecutor for running steps locally
+    - LocalArtifactStore for filesystem-based artifact storage
+    - SQLiteMetadataStore for metadata storage
+
+    Example:
+        ```python
+        from flowyml import LocalStack, Pipeline
+
+        # Create local stack
+        stack = LocalStack(name="local", artifact_path=".flowyml/artifacts", metadata_path=".flowyml/metadata.db")
+
+        # Use with pipeline
+        pipeline = Pipeline("my_pipeline", stack=stack)
+        result = pipeline.run()
+        ```
+    """
+
+    def __init__(
+        self,
+        name: str = "local",
+        artifact_path: str = ".flowyml/artifacts",
+        metadata_path: str = ".flowyml/metadata.db",
+    ):
+        """Initialize LocalStack.
+
+        Args:
+            name: Stack name
+            artifact_path: Path for artifact storage
+            metadata_path: Path for metadata database
+        """
+        # Create storage backends
+        executor = LocalExecutor()
+        artifact_store = LocalArtifactStore(artifact_path)
+        metadata_store = SQLiteMetadataStore(metadata_path)
+
+        # Initialize base stack
+        super().__init__(
+            name=name,
+            executor=executor,
+            artifact_store=artifact_store,
+            metadata_store=metadata_store,
+        )
+
+        # Ensure directories exist
+        Path(artifact_path).mkdir(parents=True, exist_ok=True)
+        Path(metadata_path).parent.mkdir(parents=True, exist_ok=True)
+
+    def validate(self) -> bool:
+        """Validate stack configuration.
+
+        Returns:
+            True if stack is valid
+        """
+        # Check artifact store path exists and is writable
+        artifact_path = Path(self.artifact_store.base_path)
+        if not artifact_path.exists():
+            artifact_path.mkdir(parents=True, exist_ok=True)
+
+        # Check metadata store path parent exists
+        metadata_path = Path(self.metadata_store.db_path)
+        if not metadata_path.parent.exists():
+            metadata_path.parent.mkdir(parents=True, exist_ok=True)
+
+        return True
+
+    def cleanup(self) -> None:
+        """Clean up stack resources."""
+        # Could implement cache cleanup, temp file removal, etc.
+        pass
+
+    def get_stats(self) -> dict:
+        """Get stack usage statistics.
+
+        Returns:
+            Dictionary with stack statistics
+        """
+        from pathlib import Path
+
+        artifact_path = Path(self.artifact_store.base_path)
+        metadata_path = Path(self.metadata_store.db_path)
+
+        # Calculate artifact storage size
+        artifact_size = sum(f.stat().st_size for f in artifact_path.rglob("*") if f.is_file())
+
+        # Get metadata size
+        metadata_size = metadata_path.stat().st_size if metadata_path.exists() else 0
+
+        # Get metadata stats from store
+        metadata_stats = self.metadata_store.get_statistics()
+
+        return {
+            "name": self.name,
+            "artifact_storage_mb": artifact_size / (1024 * 1024),
+            "metadata_storage_mb": metadata_size / (1024 * 1024),
+            "total_runs": metadata_stats.get("total_runs", 0),
+            "total_artifacts": metadata_stats.get("total_artifacts", 0),
+            "total_metrics": metadata_stats.get("total_metrics", 0),
+            "total_pipelines": metadata_stats.get("total_pipelines", 0),
+        }

flowyml/stacks/migration.py

@@ -0,0 +1,97 @@
+"""Stack Migration Tools.
+
+This module provides tools to migrate stacks from external systems (e.g., ZenML)
+to flowyml configuration.
+"""
+
+from typing import Any
+import yaml
+
+from flowyml.stacks.plugin_config import PluginConfig
+
+
+class StackMigrator:
+    """Migrates stacks from external systems to flowyml."""
+
+    def migrate_zenml_stack(self, stack_name: str) -> dict[str, Any]:
+        """Migrate a ZenML stack to flowyml configuration."""
+        try:
+            from zenml.client import Client
+
+            client = Client()
+            stack = client.get_stack(stack_name)
+        except ImportError:
+            raise ImportError("ZenML is not installed. Please install it with: pip install zenml")
+        except KeyError:
+            raise ValueError(f"ZenML stack '{stack_name}' not found.")
+
+        # Generate plugin configs for each component
+        plugins = []
+        stack_config = {
+            "name": stack.name,
+            "components": {},
+        }
+
+        for component_type, components in stack.components.items():
+            # Handle list of components (ZenML specific)
+            if not isinstance(components, list):
+                components = [components]
+
+            for component in components:
+                # Create plugin config
+                plugin_name = f"zenml_{component.name}"
+
+                # Determine source class path
+                # This relies on ZenML component having a valid module path
+                # We might need to inspect the object class
+                source_class = component.__class__.__module__ + "." + component.__class__.__name__
+
+                plugin = PluginConfig(
+                    name=plugin_name,
+                    source=source_class,
+                    component_type=self._map_zenml_type(component.type),
+                    adaptation={
+                        "method_mapping": {"run_pipeline": "run"} if component.type == "orchestrator" else {},
+                    },
+                )
+                plugins.append(plugin)
+
+                # Add to stack config
+                # We use the plugin name as the component type in flowyml stack config
+                # This assumes the plugin is registered with this name
+                stack_config["components"][component_type] = {
+                    "type": plugin_name,
+                    "config": component.configuration,
+                }
+
+        return {
+            "plugins": [
+                {
+                    "name": p.name,
+                    "source": p.source,
+                    "type": p.component_type,  # Serialize as 'type' for compatibility
+                    "adaptation": p.adaptation,
+                }
+                for p in plugins
+            ],
+            "stack": stack_config,
+        }
+
+    def _map_zenml_type(self, zenml_type: str) -> str:
+        """Map ZenML component type to flowyml type."""
+        type_mapping = {
+            "orchestrator": "orchestrator",
+            "artifact_store": "artifact_store",
+            "container_registry": "container_registry",
+        }
+        return type_mapping.get(zenml_type, "custom")
+
+    def generate_yaml(self, migration_data: dict[str, Any]) -> str:
+        """Generate YAML configuration from migration data."""
+        output = {
+            "plugins": migration_data["plugins"],
+            "stacks": {
+                migration_data["stack"]["name"]: migration_data["stack"]["components"],
+            },
+        }
+        return yaml.dump(output, sort_keys=False)