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

flowyml/plugins/deployers/sagemaker.py

@@ -0,0 +1,306 @@
+"""SageMaker Endpoint Deployer - Native FlowyML Plugin.
+
+This plugin provides direct integration with AWS SageMaker
+endpoints for model serving and inference.
+
+Usage:
+    from flowyml.plugins import get_plugin
+
+    deployer = get_plugin("sagemaker_endpoint", region="us-east-1")
+    endpoint = deployer.deploy(
+        model_name="my-model",
+        endpoint_name="my-endpoint",
+        instance_type="ml.m5.large"
+    )
+"""
+
+import logging
+from typing import Any
+import time
+
+from flowyml.plugins.base import ModelDeployerPlugin, PluginMetadata, PluginType
+
+logger = logging.getLogger(__name__)
+
+
+class SageMakerEndpointDeployer(ModelDeployerPlugin):
+    """Native SageMaker Endpoint deployer for FlowyML.
+
+    This plugin deploys models to SageMaker Endpoints for real-time inference.
+
+    Args:
+        region: AWS region.
+        role_arn: IAM role ARN for SageMaker.
+    """
+
+    metadata = PluginMetadata(
+        name="sagemaker_endpoint",
+        version="1.0.0",
+        description="AWS SageMaker Endpoint Deployer",
+        author="FlowyML Team",
+        plugin_type=PluginType.CUSTOM,
+    )
+
+    def __init__(
+        self,
+        region: str = None,
+        role_arn: str = None,
+        **kwargs,
+    ):
+        """Initialize the SageMaker Endpoint deployer.
+
+        Args:
+            region: AWS region (uses default if not specified).
+            role_arn: IAM role ARN for SageMaker operations.
+            **kwargs: Additional plugin arguments.
+        """
+        super().__init__(**kwargs)
+        self.region = region
+        self.role_arn = role_arn
+        self._boto_session = None
+        self._sm_client = None
+        self._runtime_client = None
+        self._initialized = False
+
+    @property
+    def plugin_type(self) -> PluginType:
+        return PluginType.CUSTOM
+
+    def initialize(self) -> None:
+        """Initialize connection to SageMaker."""
+        if self._initialized:
+            return
+
+        try:
+            import boto3
+
+            self._boto_session = boto3.Session(region_name=self.region)
+            self._sm_client = self._boto_session.client("sagemaker")
+            self._runtime_client = self._boto_session.client("sagemaker-runtime")
+
+            self._initialized = True
+            logger.info(f"SageMaker Endpoint Deployer initialized in region {self.region}")
+        except ImportError:
+            raise ImportError(
+                "boto3 is required. Install with: pip install boto3",
+            )
+
+    def _ensure_initialized(self) -> None:
+        """Ensure SageMaker is initialized."""
+        if not self._initialized:
+            self.initialize()
+
+    def deploy(
+        self,
+        model_uri: str,
+        endpoint_name: str,
+        instance_type: str = "ml.m5.large",
+        instance_count: int = 1,
+        inference_image_uri: str = None,
+        wait: bool = False,
+        **kwargs,
+    ) -> str:
+        """Deploy a model to a SageMaker endpoint.
+
+        Args:
+            model_uri: S3 URI to model artifacts or model package ARN.
+            endpoint_name: Name for the endpoint.
+            instance_type: Instance type (e.g., ml.m5.large).
+            instance_count: Number of instances.
+            inference_image_uri: Docker image for inference.
+            wait: Whether to wait for deployment to complete.
+            **kwargs: Additional deployment arguments.
+
+        Returns:
+            Endpoint ARN.
+        """
+        self._ensure_initialized()
+
+        try:
+            model_name = f"{endpoint_name}-model"
+            config_name = f"{endpoint_name}-config"
+
+            # Create model (if model_uri is S3 path, need image)
+            if model_uri.startswith("s3://"):
+                if not inference_image_uri:
+                    raise ValueError("inference_image_uri required for S3 model")
+
+                try:
+                    self._sm_client.create_model(
+                        ModelName=model_name,
+                        PrimaryContainer={
+                            "Image": inference_image_uri,
+                            "ModelDataUrl": model_uri,
+                        },
+                        ExecutionRoleArn=self.role_arn,
+                    )
+                except self._sm_client.exceptions.ResourceInUse:
+                    logger.info(f"Model {model_name} already exists")
+            elif model_uri.startswith("arn:"):
+                # Model package ARN
+                try:
+                    self._sm_client.create_model(
+                        ModelName=model_name,
+                        PrimaryContainer={
+                            "ModelPackageName": model_uri,
+                        },
+                        ExecutionRoleArn=self.role_arn,
+                    )
+                except self._sm_client.exceptions.ResourceInUse:
+                    logger.info(f"Model {model_name} already exists")
+            else:
+                raise ValueError(f"Invalid model_uri format: {model_uri}")
+
+            # Create endpoint config
+            try:
+                self._sm_client.create_endpoint_config(
+                    EndpointConfigName=config_name,
+                    ProductionVariants=[
+                        {
+                            "VariantName": "primary",
+                            "ModelName": model_name,
+                            "InstanceType": instance_type,
+                            "InitialInstanceCount": instance_count,
+                        },
+                    ],
+                )
+            except self._sm_client.exceptions.ResourceInUse:
+                logger.info(f"Endpoint config {config_name} already exists")
+
+            # Create or update endpoint
+            try:
+                self._sm_client.create_endpoint(
+                    EndpointName=endpoint_name,
+                    EndpointConfigName=config_name,
+                )
+                logger.info(f"Creating endpoint: {endpoint_name}")
+            except self._sm_client.exceptions.ResourceInUse:
+                self._sm_client.update_endpoint(
+                    EndpointName=endpoint_name,
+                    EndpointConfigName=config_name,
+                )
+                logger.info(f"Updating endpoint: {endpoint_name}")
+
+            if wait:
+                self._wait_for_endpoint(endpoint_name)
+
+            # Get endpoint ARN
+            endpoint_desc = self._sm_client.describe_endpoint(
+                EndpointName=endpoint_name,
+            )
+
+            logger.info(f"Deployed to endpoint: {endpoint_name}")
+            return endpoint_desc["EndpointArn"]
+
+        except Exception as e:
+            logger.error(f"Failed to deploy model: {e}")
+            raise
+
+    def _wait_for_endpoint(self, endpoint_name: str, timeout: int = 600) -> None:
+        """Wait for endpoint to be in service."""
+        start = time.time()
+        while time.time() - start < timeout:
+            desc = self._sm_client.describe_endpoint(EndpointName=endpoint_name)
+            status = desc["EndpointStatus"]
+
+            if status == "InService":
+                logger.info(f"Endpoint {endpoint_name} is InService")
+                return
+            elif status == "Failed":
+                raise RuntimeError(f"Endpoint deployment failed: {desc.get('FailureReason')}")
+
+            logger.debug(f"Waiting for endpoint, status: {status}")
+            time.sleep(30)
+
+        raise TimeoutError(f"Endpoint deployment timed out after {timeout}s")
+
+    def undeploy(
+        self,
+        endpoint_name: str,
+    ) -> bool:
+        """Delete an endpoint.
+
+        Args:
+            endpoint_name: Endpoint name.
+
+        Returns:
+            True if successful.
+        """
+        self._ensure_initialized()
+
+        try:
+            self._sm_client.delete_endpoint(EndpointName=endpoint_name)
+            logger.info(f"Deleted endpoint: {endpoint_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to delete endpoint: {e}")
+            return False
+
+    def get_endpoint(
+        self,
+        endpoint_name: str,
+    ) -> dict | None:
+        """Get endpoint details.
+
+        Args:
+            endpoint_name: Endpoint name.
+
+        Returns:
+            Endpoint details dictionary.
+        """
+        self._ensure_initialized()
+
+        try:
+            desc = self._sm_client.describe_endpoint(EndpointName=endpoint_name)
+
+            return {
+                "name": endpoint_name,
+                "arn": desc["EndpointArn"],
+                "status": desc["EndpointStatus"],
+                "config_name": desc["EndpointConfigName"],
+                "created": str(desc.get("CreationTime", "")),
+                "last_modified": str(desc.get("LastModifiedTime", "")),
+            }
+
+        except self._sm_client.exceptions.ResourceNotFoundException:
+            return None
+        except Exception as e:
+            logger.error(f"Failed to get endpoint: {e}")
+            return None
+
+    def predict(
+        self,
+        endpoint_name: str,
+        data: Any,
+        content_type: str = "application/json",
+    ) -> Any:
+        """Make predictions using a deployed endpoint.
+
+        Args:
+            endpoint_name: Endpoint name.
+            data: Input data (will be JSON serialized if dict/list).
+            content_type: Content type of the request.
+
+        Returns:
+            Prediction result.
+        """
+        self._ensure_initialized()
+
+        try:
+            import json
+
+            body = json.dumps(data) if isinstance(data, (dict, list)) else str(data)
+
+            response = self._runtime_client.invoke_endpoint(
+                EndpointName=endpoint_name,
+                ContentType=content_type,
+                Body=body.encode("utf-8"),
+            )
+
+            result = json.loads(response["Body"].read().decode("utf-8"))
+            return result
+
+        except Exception as e:
+            logger.error(f"Failed to predict: {e}")
+            raise

flowyml/plugins/deployers/vertex.py

@@ -0,0 +1,290 @@
+"""Vertex AI Endpoint Deployer - Native FlowyML Plugin.
+
+This plugin provides direct integration with Google Cloud Vertex AI
+endpoints for model serving and inference.
+
+Usage:
+    from flowyml.plugins import get_plugin
+
+    deployer = get_plugin("vertex_endpoint", project="my-project")
+    endpoint = deployer.deploy(
+        model_name="my-model",
+        endpoint_name="my-endpoint",
+        machine_type="n1-standard-4"
+    )
+"""
+
+import logging
+
+from flowyml.plugins.base import ModelDeployerPlugin, PluginMetadata, PluginType
+
+logger = logging.getLogger(__name__)
+
+
+class VertexEndpointDeployer(ModelDeployerPlugin):
+    """Native Vertex AI Endpoint deployer for FlowyML.
+
+    This plugin deploys models to Vertex AI Endpoints for real-time inference.
+
+    Args:
+        project: GCP project ID.
+        location: GCP region (default: us-central1).
+        staging_bucket: GCS bucket for staging artifacts.
+    """
+
+    metadata = PluginMetadata(
+        name="vertex_endpoint",
+        version="1.0.0",
+        description="Google Cloud Vertex AI Endpoint Deployer",
+        author="FlowyML Team",
+        plugin_type=PluginType.CUSTOM,
+    )
+
+    def __init__(
+        self,
+        project: str,
+        location: str = "us-central1",
+        staging_bucket: str = None,
+        **kwargs,
+    ):
+        """Initialize the Vertex Endpoint deployer.
+
+        Args:
+            project: GCP project ID.
+            location: GCP region.
+            staging_bucket: GCS bucket for artifacts.
+            **kwargs: Additional plugin arguments.
+        """
+        super().__init__(**kwargs)
+        self.project = project
+        self.location = location
+        self.staging_bucket = staging_bucket
+        self._aiplatform = None
+        self._initialized = False
+
+    @property
+    def plugin_type(self) -> PluginType:
+        return PluginType.CUSTOM
+
+    def initialize(self) -> None:
+        """Initialize connection to Vertex AI."""
+        if self._initialized:
+            return
+
+        try:
+            from google.cloud import aiplatform
+
+            aiplatform.init(
+                project=self.project,
+                location=self.location,
+                staging_bucket=self.staging_bucket,
+            )
+            self._aiplatform = aiplatform
+            self._initialized = True
+            logger.info(f"Vertex Endpoint Deployer initialized for project {self.project}")
+        except ImportError:
+            raise ImportError(
+                "google-cloud-aiplatform is required. " "Install with: pip install google-cloud-aiplatform",
+            )
+
+    def _ensure_initialized(self) -> None:
+        """Ensure Vertex AI is initialized."""
+        if not self._initialized:
+            self.initialize()
+
+    def deploy(
+        self,
+        model_uri: str,
+        endpoint_name: str,
+        machine_type: str = "n1-standard-4",
+        min_replica_count: int = 1,
+        max_replica_count: int = 1,
+        accelerator_type: str = None,
+        accelerator_count: int = 0,
+        service_account: str = None,
+        **kwargs,
+    ) -> str:
+        """Deploy a model to a Vertex AI endpoint.
+
+        Args:
+            model_uri: Model resource name or GCS URI.
+            endpoint_name: Name for the endpoint.
+            machine_type: Machine type (e.g., n1-standard-4).
+            min_replica_count: Minimum replicas.
+            max_replica_count: Maximum replicas.
+            accelerator_type: GPU type (NVIDIA_TESLA_T4, etc.).
+            accelerator_count: Number of GPUs.
+            service_account: Service account for endpoint.
+            **kwargs: Additional deployment arguments.
+
+        Returns:
+            Endpoint resource name.
+        """
+        self._ensure_initialized()
+
+        try:
+            # Get or create endpoint
+            endpoints = self._aiplatform.Endpoint.list(
+                filter=f'display_name="{endpoint_name}"',
+            )
+
+            if endpoints:
+                endpoint = endpoints[0]
+                logger.info(f"Using existing endpoint: {endpoint_name}")
+            else:
+                endpoint = self._aiplatform.Endpoint.create(
+                    display_name=endpoint_name,
+                )
+                logger.info(f"Created endpoint: {endpoint_name}")
+
+            # Get the model
+            if model_uri.startswith("projects/"):
+                model = self._aiplatform.Model(model_name=model_uri)
+            else:
+                # Search by display name
+                models = self._aiplatform.Model.list(
+                    filter=f'display_name="{model_uri}"',
+                    order_by="create_time desc",
+                )
+                if not models:
+                    raise ValueError(f"Model '{model_uri}' not found")
+                model = models[0]
+
+            # Build deployment config
+            deploy_config = {
+                "machine_type": machine_type,
+                "min_replica_count": min_replica_count,
+                "max_replica_count": max_replica_count,
+            }
+
+            if accelerator_type and accelerator_count > 0:
+                deploy_config["accelerator_type"] = accelerator_type
+                deploy_config["accelerator_count"] = accelerator_count
+
+            if service_account:
+                deploy_config["service_account"] = service_account
+
+            # Deploy model to endpoint
+            model.deploy(endpoint=endpoint, **deploy_config)
+
+            logger.info(f"Deployed model to endpoint: {endpoint.resource_name}")
+            return endpoint.resource_name
+
+        except Exception as e:
+            logger.error(f"Failed to deploy model: {e}")
+            raise
+
+    def undeploy(
+        self,
+        endpoint_name: str,
+        deployed_model_id: str = None,
+    ) -> bool:
+        """Undeploy a model from an endpoint.
+
+        Args:
+            endpoint_name: Endpoint name or resource name.
+            deployed_model_id: Specific deployed model ID (undeployes all if None).
+
+        Returns:
+            True if successful.
+        """
+        self._ensure_initialized()
+
+        try:
+            if endpoint_name.startswith("projects/"):
+                endpoint = self._aiplatform.Endpoint(endpoint_name=endpoint_name)
+            else:
+                endpoints = self._aiplatform.Endpoint.list(
+                    filter=f'display_name="{endpoint_name}"',
+                )
+                if not endpoints:
+                    logger.warning(f"Endpoint '{endpoint_name}' not found")
+                    return False
+                endpoint = endpoints[0]
+
+            if deployed_model_id:
+                endpoint.undeploy(deployed_model_id=deployed_model_id)
+            else:
+                endpoint.undeploy_all()
+
+            logger.info(f"Undeployed model(s) from endpoint: {endpoint_name}")
+            return True
+
+        except Exception as e:
+            logger.error(f"Failed to undeploy: {e}")
+            return False
+
+    def get_endpoint(
+        self,
+        endpoint_name: str,
+    ) -> dict | None:
+        """Get endpoint details.
+
+        Args:
+            endpoint_name: Endpoint name or resource name.
+
+        Returns:
+            Endpoint details dictionary.
+        """
+        self._ensure_initialized()
+
+        try:
+            if endpoint_name.startswith("projects/"):
+                endpoint = self._aiplatform.Endpoint(endpoint_name=endpoint_name)
+            else:
+                endpoints = self._aiplatform.Endpoint.list(
+                    filter=f'display_name="{endpoint_name}"',
+                )
+                if not endpoints:
+                    return None
+                endpoint = endpoints[0]
+
+            return {
+                "name": endpoint.display_name,
+                "resource_name": endpoint.resource_name,
+                "deployed_models": [
+                    {
+                        "id": dm.id,
+                        "model": dm.model,
+                    }
+                    for dm in endpoint.deployed_models
+                ],
+            }
+
+        except Exception as e:
+            logger.error(f"Failed to get endpoint: {e}")
+            return None
+
+    def predict(
+        self,
+        endpoint_name: str,
+        instances: list[dict],
+    ) -> list[dict]:
+        """Make predictions using a deployed endpoint.
+
+        Args:
+            endpoint_name: Endpoint name or resource name.
+            instances: List of input instances.
+
+        Returns:
+            List of predictions.
+        """
+        self._ensure_initialized()
+
+        try:
+            if endpoint_name.startswith("projects/"):
+                endpoint = self._aiplatform.Endpoint(endpoint_name=endpoint_name)
+            else:
+                endpoints = self._aiplatform.Endpoint.list(
+                    filter=f'display_name="{endpoint_name}"',
+                )
+                if not endpoints:
+                    raise ValueError(f"Endpoint '{endpoint_name}' not found")
+                endpoint = endpoints[0]
+
+            response = endpoint.predict(instances=instances)
+            return response.predictions
+
+        except Exception as e:
+            logger.error(f"Failed to predict: {e}")
+            raise