flowyml 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

flowyml/stacks/azure.py

@@ -0,0 +1,295 @@
+"""Azure Stack Components and Preset Stack."""
+
+from __future__ import annotations
+
+import subprocess
+from pathlib import Path
+from typing import Any
+
+from flowyml.stacks.base import Stack
+from flowyml.stacks.components import ArtifactStore, ContainerRegistry, DockerConfig, ResourceConfig
+from flowyml.core.remote_orchestrator import RemoteOrchestrator
+from flowyml.stacks.plugins import register_component
+from flowyml.storage.metadata import SQLiteMetadataStore
+from flowyml.core.submission_result import SubmissionResult
+
+
+@register_component(name="azure_blob")
+class AzureBlobArtifactStore(ArtifactStore):
+    """Artifact store backed by Azure Blob Storage."""
+
+    def __init__(
+        self,
+        name: str = "azure_blob",
+        account_url: str | None = None,
+        container_name: str | None = None,
+        credential: Any | None = None,
+    ):
+        super().__init__(name)
+        self.account_url = account_url
+        self.container_name = container_name
+        self.credential = credential
+
+    def _client(self):
+        from azure.storage.blob import BlobServiceClient
+
+        credential = self.credential
+        if credential is None:
+            from azure.identity import DefaultAzureCredential
+
+            credential = DefaultAzureCredential(exclude_shared_token_cache_credential=True)
+
+        return BlobServiceClient(account_url=self.account_url, credential=credential)
+
+    def validate(self) -> bool:
+        if not self.account_url or not self.container_name:
+            raise ValueError("account_url and container_name are required for AzureBlobArtifactStore")
+        try:
+            client = self._client()
+            container_client = client.get_container_client(self.container_name)
+            container_client.get_container_properties()
+        except Exception as exc:
+            raise ValueError(f"Unable to access container '{self.container_name}': {exc}") from exc
+        return True
+
+    def save(self, artifact: Any, path: str) -> str:
+        blob_name = path.lstrip("/")
+        client = self._client().get_container_client(self.container_name)
+        if isinstance(artifact, (str, Path)) and Path(artifact).exists():
+            with open(Path(artifact), "rb") as f:
+                client.upload_blob(name=blob_name, data=f, overwrite=True)
+        else:
+            data = artifact if isinstance(artifact, bytes) else str(artifact).encode()
+            client.upload_blob(name=blob_name, data=data, overwrite=True)
+        return f"{self.account_url}/{self.container_name}/{blob_name}"
+
+    def load(self, path: str) -> bytes:
+        blob_name = path.lstrip("/")
+        client = self._client().get_blob_client(self.container_name, blob_name)
+        downloader = client.download_blob()
+        return downloader.readall()
+
+    def exists(self, path: str) -> bool:
+        blob_name = path.lstrip("/")
+        client = self._client().get_blob_client(self.container_name, blob_name)
+        return client.exists()
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "type": "azure_blob",
+            "account_url": self.account_url,
+            "container_name": self.container_name,
+        }
+
+
+@register_component(name="acr")
+class ACRContainerRegistry(ContainerRegistry):
+    """Azure Container Registry integration."""
+
+    def __init__(
+        self,
+        name: str = "acr",
+        registry_name: str | None = None,
+        login_server: str | None = None,
+        username: str | None = None,
+        password: str | None = None,
+    ):
+        super().__init__(name)
+        self.registry_name = registry_name
+        self.login_server = login_server or (f"{registry_name}.azurecr.io" if registry_name else None)
+        self.username = username
+        self.password = password
+
+    def validate(self) -> bool:
+        if not self.registry_name and not self.login_server:
+            raise ValueError("registry_name or login_server is required for ACRContainerRegistry")
+        return True
+
+    def _login(self) -> None:
+        if self.username and self.password:
+            subprocess.run(
+                ["docker", "login", self.login_server, "-u", self.username, "-p", self.password],
+                check=True,
+            )
+        else:
+            subprocess.run(["az", "acr", "login", "--name", self.registry_name], check=True)
+
+    def push_image(self, image_name: str, tag: str = "latest") -> str:
+        full_uri = self.get_image_uri(image_name, tag)
+        self._login()
+        subprocess.run(["docker", "tag", f"{image_name}:{tag}", full_uri], check=True)
+        subprocess.run(["docker", "push", full_uri], check=True)
+        return full_uri
+
+    def pull_image(self, image_name: str, tag: str = "latest") -> None:
+        full_uri = self.get_image_uri(image_name, tag)
+        self._login()
+        subprocess.run(["docker", "pull", full_uri], check=True)
+
+    def get_image_uri(self, image_name: str, tag: str = "latest") -> str:
+        return f"{self.login_server}/{image_name}:{tag}"
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "type": "acr",
+            "login_server": self.login_server,
+            "registry_name": self.registry_name,
+        }
+
+
+@register_component(name="azure_ml")
+class AzureMLOrchestrator(RemoteOrchestrator):
+    """Submit pipeline runs to Azure ML managed compute."""
+
+    def __init__(
+        self,
+        name: str = "azure_ml",
+        subscription_id: str | None = None,
+        resource_group: str | None = None,
+        workspace_name: str | None = None,
+        compute: str | None = None,
+        experiment_name: str = "flowyml",
+        credential: Any | None = None,
+    ):
+        super().__init__(name)
+        self.subscription_id = subscription_id
+        self.resource_group = resource_group
+        self.workspace_name = workspace_name
+        self.compute = compute
+        self.experiment_name = experiment_name
+        self.credential = credential
+
+    def _client(self):
+        from azure.ai.ml import MLClient
+        from azure.identity import DefaultAzureCredential
+
+        credential = self.credential or DefaultAzureCredential(exclude_shared_token_cache_credential=True)
+        return MLClient(
+            credential,
+            subscription_id=self.subscription_id,
+            resource_group_name=self.resource_group,
+            workspace_name=self.workspace_name,
+        )
+
+    def validate(self) -> bool:
+        if not all([self.subscription_id, self.resource_group, self.workspace_name, self.compute]):
+            raise ValueError(
+                "subscription_id, resource_group, workspace_name, and compute are required for AzureMLOrchestrator",
+            )
+        return True
+
+    def run_pipeline(
+        self,
+        pipeline: Any,
+        run_id: str,
+        resources: ResourceConfig | None = None,
+        docker_config: DockerConfig | None = None,
+        inputs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> SubmissionResult:
+        """Submit pipeline to Azure ML.
+
+        Returns:
+            SubmissionResult with job name and optional wait function.
+        """
+        from azure.ai.ml import command
+
+        ml_client = self._client()
+        job_name = kwargs.get("job_name") or f"{pipeline.name}-{run_id[:8]}"
+
+        # Build command job
+        job = command(
+            code=".",  # Use current directory or specify path
+            command="python -m flowyml.cli.run",
+            environment=docker_config.image if docker_config else self.environment_name,
+            compute=self.compute,
+            display_name=job_name,
+            experiment_name=self.experiment_name,
+        )
+
+        submitted = ml_client.jobs.create_or_update(job)
+
+        def wait_for_completion():
+            ml_client.jobs.stream(submitted.name)
+
+        return SubmissionResult(
+            job_id=submitted.name,
+            wait_for_completion=wait_for_completion,
+            metadata={"studio_url": submitted.studio_url},
+        )
+
+    def get_run_status(self, run_id: str) -> str:
+        client = self._client()
+        job = client.jobs.get(run_id)
+        return job.status
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "type": "azure_ml",
+            "subscription_id": self.subscription_id,
+            "resource_group": self.resource_group,
+            "workspace_name": self.workspace_name,
+            "compute": self.compute,
+        }
+
+
+class AzureMLStack(Stack):
+    """Managed Azure ML stack built from AzureML orchestrator, Blob storage, and ACR."""
+
+    def __init__(
+        self,
+        name: str = "azure",
+        subscription_id: str | None = None,
+        resource_group: str | None = None,
+        workspace_name: str | None = None,
+        compute: str | None = None,
+        account_url: str | None = None,
+        container_name: str | None = None,
+        registry_name: str | None = None,
+        login_server: str | None = None,
+        metadata_store: Any | None = None,
+    ):
+        orchestrator = AzureMLOrchestrator(
+            subscription_id=subscription_id,
+            resource_group=resource_group,
+            workspace_name=workspace_name,
+            compute=compute,
+        )
+        artifact_store = AzureBlobArtifactStore(account_url=account_url, container_name=container_name)
+        container_registry = ACRContainerRegistry(registry_name=registry_name, login_server=login_server)
+
+        if metadata_store is None:
+            metadata_store = SQLiteMetadataStore()
+
+        super().__init__(
+            name=name,
+            executor=None,
+            artifact_store=artifact_store,
+            metadata_store=metadata_store,
+            container_registry=container_registry,
+            orchestrator=orchestrator,
+        )
+
+        self.subscription_id = subscription_id
+        self.workspace_name = workspace_name
+
+    def validate(self) -> bool:
+        self.orchestrator.validate()
+        self.artifact_store.validate()
+        self.container_registry.validate()
+        return True
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "type": "azure",
+            "subscription_id": self.subscription_id,
+            "workspace_name": self.workspace_name,
+            "orchestrator": self.orchestrator.to_dict(),
+            "artifact_store": self.artifact_store.to_dict(),
+            "container_registry": self.container_registry.to_dict(),
+        }

flowyml/stacks/components.py

@@ -99,8 +99,30 @@ class Orchestrator(StackComponent):
         return ComponentType.ORCHESTRATOR
 
     @abstractmethod
-    def run_pipeline(self, pipeline: Any, **kwargs) -> Any:
-        """Run a pipeline on this orchestrator."""
+    def run_pipeline(
+        self,
+        pipeline: Any,
+        run_id: str,
+        resources: "ResourceConfig | None" = None,
+        docker_config: "DockerConfig | None" = None,
+        inputs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        **kwargs,
+    ) -> Any:
+        """Run a pipeline on this orchestrator.
+
+        Args:
+            pipeline: The pipeline to run.
+            run_id: The unique run identifier.
+            resources: Resource configuration.
+            docker_config: Docker configuration.
+            inputs: Input data.
+            context: Context variables.
+            **kwargs: Additional arguments.
+
+        Returns:
+            The run result or job ID.
+        """
         pass
 
     @abstractmethod

flowyml/stacks/gcp.py

@@ -6,17 +6,22 @@ on Google Cloud Platform using Vertex AI, Cloud Storage, and Container Registry.
 
 from typing import Any
 
+import subprocess
 from flowyml.stacks.base import Stack
 from flowyml.stacks.components import (
-    Orchestrator,
     ArtifactStore,
     ContainerRegistry,
     ResourceConfig,
     DockerConfig,
 )
+from flowyml.core.remote_orchestrator import RemoteOrchestrator
+from flowyml.stacks.plugins import register_component
+from flowyml.core.submission_result import SubmissionResult
+from flowyml.core.execution_status import ExecutionStatus
 
 
-class VertexAIOrchestrator(Orchestrator):
+@register_component(name="vertex_ai")
+class VertexAIOrchestrator(RemoteOrchestrator):
     """Vertex AI orchestrator for running pipelines on Google Cloud.
 
     This orchestrator submits pipeline jobs to Vertex AI Pipelines,
@@ -76,28 +81,35 @@ class VertexAIOrchestrator(Orchestrator):
     def run_pipeline(
         self,
         pipeline: Any,
+        run_id: str,
         resources: ResourceConfig | None = None,
         docker_config: DockerConfig | None = None,
+        inputs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
         **kwargs,
-    ) -> str:
+    ) -> "SubmissionResult":
         """Run pipeline on Vertex AI.
 
         Args:
             pipeline: Pipeline to run
+            run_id: Run identifier
             resources: Resource configuration
             docker_config: Docker configuration
+            inputs: Input data
+            context: Context variables
             **kwargs: Additional arguments
 
         Returns:
-            Job ID
+            SubmissionResult with job resource name
         """
         from google.cloud import aiplatform
+        import time
 
         # Initialize Vertex AI
         aiplatform.init(project=self.project_id, location=self.region)
 
         # Create custom job
-        job_display_name = f"{pipeline.name}-{pipeline.run_id}"
+        job_display_name = f"{pipeline.name}-{run_id[:8]}"
 
         # Build worker pool specs
         worker_pool_specs = self._build_worker_pool_specs(
@@ -114,16 +126,66 @@ class VertexAIOrchestrator(Orchestrator):
             encryption_spec_key_name=self.encryption_key,
         )
 
-        job.run(sync=False)
-
-        return job.resource_name
+        # Submit job asynchronously
+        job.submit()
+
+        job_id = job.resource_name
+
+        # Create wait function
+        def wait_for_completion():
+            """Poll job status until completion."""
+            while True:
+                status = self.get_run_status(job_id)
+                if status.is_finished:
+                    if not status.is_successful:
+                        raise RuntimeError(f"Vertex AI job {job_id} failed with status: {status}")
+                    break
+                time.sleep(15)  # Poll every 15 seconds
+
+        return SubmissionResult(
+            job_id=job_id,
+            wait_for_completion=wait_for_completion,
+            metadata={
+                "platform": "vertex_ai",
+                "project": self.project_id,
+                "region": self.region,
+                "job_name": job_display_name,
+            },
+        )
 
-    def get_run_status(self, run_id: str) -> str:
+    def get_run_status(self, job_id: str) -> "ExecutionStatus":
         """Get status of a Vertex AI job."""
         from google.cloud import aiplatform
 
-        job = aiplatform.CustomJob(run_id)
-        return job.state.name
+        try:
+            job = aiplatform.CustomJob(job_id)
+            state = job.state.name
+
+            # Map Vertex AI states to ExecutionStatus
+            status_map = {
+                "JOB_STATE_QUEUED": ExecutionStatus.PROVISIONING,
+                "JOB_STATE_PENDING": ExecutionStatus.PROVISIONING,
+                "JOB_STATE_RUNNING": ExecutionStatus.RUNNING,
+                "JOB_STATE_SUCCEEDED": ExecutionStatus.COMPLETED,
+                "JOB_STATE_FAILED": ExecutionStatus.FAILED,
+                "JOB_STATE_CANCELLING": ExecutionStatus.STOPPING,
+                "JOB_STATE_CANCELLED": ExecutionStatus.CANCELLED,
+            }
+            return status_map.get(state, ExecutionStatus.RUNNING)
+        except Exception as e:
+            print(f"Error fetching job status: {e}")
+            return ExecutionStatus.FAILED
+
+    def stop_run(self, job_id: str, graceful: bool = True) -> None:
+        """Cancel a Vertex AI job."""
+        from google.cloud import aiplatform
+
+        try:
+            job = aiplatform.CustomJob(job_id)
+            job.cancel()
+        except Exception as e:
+            print(f"Error cancelling job {job_id}: {e}")
+            raise
 
     def _build_worker_pool_specs(
         self,
@@ -171,6 +233,7 @@ class VertexAIOrchestrator(Orchestrator):
         }
 
 
+@register_component(name="gcs")
 class GCSArtifactStore(ArtifactStore):
     """Google Cloud Storage artifact store.
 
@@ -280,6 +343,7 @@ class GCSArtifactStore(ArtifactStore):
         }
 
 
+@register_component(name="gcr")
 class GCRContainerRegistry(ContainerRegistry):
     """Google Container Registry integration.
 
@@ -478,6 +542,8 @@ class GCPStack(Stack):
 
         self.project_id = project_id
         self.region = region
+        self.vertex_endpoints = VertexEndpointManager(project_id=project_id, region=region)
+        self.cloud_run = CloudRunDeployer(project_id=project_id, region=region)
 
     def validate(self) -> bool:
         """Validate all GCP stack components."""
@@ -497,3 +563,84 @@ class GCPStack(Stack):
             "artifact_store": self.artifact_store.to_dict(),
             "container_registry": self.container_registry.to_dict(),
         }
+
+
+class VertexEndpointManager:
+    """Deploy trained models as Vertex AI endpoints."""
+
+    def __init__(self, project_id: str | None, region: str = "us-central1"):
+        self.project_id = project_id
+        self.region = region
+
+    def deploy_model(
+        self,
+        model_display_name: str,
+        artifact_uri: str,
+        serving_image: str,
+        endpoint_display_name: str | None = None,
+        machine_type: str = "n1-standard-4",
+    ) -> str:
+        from google.cloud import aiplatform
+
+        aiplatform.init(project=self.project_id, location=self.region)
+        model = aiplatform.Model.upload(
+            display_name=model_display_name,
+            artifact_uri=artifact_uri,
+            serving_container_image_uri=serving_image,
+        )
+        endpoint = model.deploy(
+            machine_type=machine_type,
+            endpoint=aiplatform.Endpoint.create(
+                display_name=endpoint_display_name or f"{model_display_name}-endpoint",
+            ),
+        )
+        return endpoint.resource_name
+
+
+class CloudRunDeployer:
+    """Deploy container images to Cloud Run."""
+
+    def __init__(self, project_id: str | None, region: str = "us-central1"):
+        self.project_id = project_id
+        self.region = region
+
+    def deploy_service(
+        self,
+        service_name: str,
+        image: str,
+        env: dict[str, str] | None = None,
+        allow_unauthenticated: bool = True,
+    ) -> str:
+        command = [
+            "gcloud",
+            "run",
+            "deploy",
+            service_name,
+            f"--image={image}",
+            f"--region={self.region}",
+            f"--project={self.project_id}",
+        ]
+        if allow_unauthenticated:
+            command.append("--allow-unauthenticated")
+
+        env = env or {}
+        for key, value in env.items():
+            command.append(f"--set-env-vars={key}={value}")
+
+        subprocess.run(command, check=True)
+        url_result = subprocess.run(
+            [
+                "gcloud",
+                "run",
+                "services",
+                "describe",
+                service_name,
+                f"--region={self.region}",
+                f"--project={self.project_id}",
+                "--format=value(status.url)",
+            ],
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        return url_result.stdout.strip()

flowyml/stacks/local.py

@@ -46,12 +46,17 @@ class LocalStack(Stack):
         artifact_store = LocalArtifactStore(artifact_path)
         metadata_store = SQLiteMetadataStore(metadata_path)
 
+        from flowyml.core.orchestrator import LocalOrchestrator
+
+        orchestrator = LocalOrchestrator()
+
         # Initialize base stack
         super().__init__(
             name=name,
             executor=executor,
             artifact_store=artifact_store,
             metadata_store=metadata_store,
+            orchestrator=orchestrator,
         )
 
         # Ensure directories exist

flowyml/storage/artifacts.py

@@ -215,7 +215,7 @@ class LocalArtifactStore(ArtifactStore):
         from datetime import datetime
         from flowyml.storage.materializers.base import get_materializer
         import shutil
-        import pickle
+        import cloudpickle
         import json
 
         date_str = datetime.now().strftime("%Y-%m-%d")
@@ -236,11 +236,21 @@ class LocalArtifactStore(ArtifactStore):
         if materializer:
             materializer.save(obj, full_path)
         else:
-            # Fallback to pickle
-            with open(full_path / "data.pkl", "wb") as f:
-                pickle.dump(obj, f)
+            # Fallback to cloudpickle (more robust than pickle)
+            fallback_file = full_path / "data.pkl"
+            with open(fallback_file, "wb") as f:
+                cloudpickle.dump(obj, f)
             # Save metadata
             with open(full_path / "metadata.json", "w") as f:
-                json.dump({"type": "pickle", "format": "pickle"}, f, indent=2)
+                json.dump(
+                    {
+                        "type": type(obj).__name__,
+                        "serializer": "cloudpickle",
+                        "format": "pickle",
+                        "file": fallback_file.name,
+                    },
+                    f,
+                    indent=2,
+                )
 
         return str(full_path)

flowyml/storage/materializers/__init__.py

@@ -7,6 +7,7 @@ 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
+from flowyml.storage.materializers.cloudpickle import CloudpickleMaterializer
 
 __all__ = [
     "BaseMaterializer",
@@ -17,4 +18,5 @@ __all__ = [
     "PandasMaterializer",
     "NumPyMaterializer",
     "KerasMaterializer",
+    "CloudpickleMaterializer",
 ]