synapse-sdk 1.0.0a11__py3-none-any.whl → 2026.1.1b2__py3-none-any.whl

synapse_sdk/plugins/actions/inference/deployment.py

@@ -0,0 +1,322 @@
+"""Deployment action base class for Ray Serve deployments."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from pydantic import BaseModel
+
+from synapse_sdk.plugins.action import BaseAction
+from synapse_sdk.plugins.actions.inference.context import DeploymentContext
+from synapse_sdk.plugins.steps import Orchestrator, StepRegistry
+
+P = TypeVar('P', bound=BaseModel)
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.agent import AgentClient
+    from synapse_sdk.clients.backend import BackendClient
+
+
+class DeploymentProgressCategories:
+    """Standard progress category names for deployment workflows.
+
+    Use these constants with set_progress() to track deployment phases:
+        - INITIALIZE: Ray cluster initialization
+        - DEPLOY: Deploying to Ray Serve
+        - REGISTER: Registering with backend
+
+    Example:
+        >>> self.set_progress(1, 3, self.progress.INITIALIZE)
+        >>> self.set_progress(2, 3, self.progress.DEPLOY)
+    """
+
+    INITIALIZE: str = 'initialize'
+    DEPLOY: str = 'deploy'
+    REGISTER: str = 'register'
+
+
+class BaseDeploymentAction(BaseAction[P]):
+    """Base class for Ray Serve deployment actions.
+
+    Provides helper methods for deploying inference endpoints to Ray Serve.
+    Handles Ray initialization, deployment creation, and backend registration.
+
+    Supports two execution modes:
+    1. Simple execute: Override execute() directly for simple deployments
+    2. Step-based: Override setup_steps() to register workflow steps
+
+    Attributes:
+        progress: Standard progress category names.
+        entrypoint: The serve deployment class to deploy (set in subclass).
+
+    Example (simple execute):
+        >>> class MyDeploymentAction(BaseDeploymentAction[MyParams]):
+        ...     action_name = 'deployment'
+        ...     category = 'neural_net'
+        ...     params_model = MyParams
+        ...     entrypoint = MyServeDeployment
+        ...
+        ...     def execute(self) -> dict[str, Any]:
+        ...         self.ray_init()
+        ...         self.set_progress(1, 3, self.progress.INITIALIZE)
+        ...         self.deploy()
+        ...         self.set_progress(2, 3, self.progress.DEPLOY)
+        ...         app_id = self.register_serve_application()
+        ...         self.set_progress(3, 3, self.progress.REGISTER)
+        ...         return {'serve_application': app_id}
+
+    Example (step-based):
+        >>> class MyDeploymentAction(BaseDeploymentAction[MyParams]):
+        ...     entrypoint = MyServeDeployment
+        ...
+        ...     def setup_steps(self, registry: StepRegistry[DeploymentContext]) -> None:
+        ...         registry.register(InitializeRayStep())
+        ...         registry.register(DeployStep())
+        ...         registry.register(RegisterStep())
+    """
+
+    progress = DeploymentProgressCategories()
+
+    # Override in subclass with your serve deployment class
+    entrypoint: type | None = None
+
+    @property
+    def client(self) -> BackendClient:
+        """Backend client from context.
+
+        Returns:
+            BackendClient instance.
+
+        Raises:
+            RuntimeError: If no client in context.
+        """
+        if self.ctx.client is None:
+            raise RuntimeError('No client in context. Provide a client via RuntimeContext.')
+        return self.ctx.client
+
+    @property
+    def agent_client(self) -> AgentClient:
+        """Agent client from context.
+
+        Returns:
+            AgentClient instance for Ray operations.
+
+        Raises:
+            RuntimeError: If no agent_client in context.
+        """
+        if self.ctx.agent_client is None:
+            raise RuntimeError('No agent_client in context. Provide an agent_client via RuntimeContext.')
+        return self.ctx.agent_client
+
+    def setup_steps(self, registry: StepRegistry[DeploymentContext]) -> None:
+        """Register workflow steps for step-based execution.
+
+        Override this method to register custom steps for deployment workflow.
+        If steps are registered, step-based execution takes precedence.
+
+        Args:
+            registry: StepRegistry to register steps with.
+
+        Example:
+            >>> def setup_steps(self, registry: StepRegistry[DeploymentContext]) -> None:
+            ...     registry.register(InitializeRayStep())
+            ...     registry.register(DeployStep())
+            ...     registry.register(RegisterStep())
+        """
+        pass  # Default: no steps, uses simple execute()
+
+    def create_context(self) -> DeploymentContext:
+        """Create deployment context for step-based workflow.
+
+        Override to customize context creation or add additional state.
+
+        Returns:
+            DeploymentContext instance with params and runtime context.
+        """
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+        return DeploymentContext(
+            runtime_ctx=self.ctx,
+            params=params_dict,
+            model_id=params_dict.get('model_id'),
+            serve_app_name=self.get_serve_app_name(),
+            route_prefix=self.get_route_prefix(),
+            ray_actor_options=self.get_ray_actor_options(),
+        )
+
+    def run(self) -> Any:
+        """Run the action, using steps if registered.
+
+        This method is called by executors. It checks if steps are
+        registered and uses step-based execution if so.
+
+        Returns:
+            Action result (dict or any return type).
+        """
+        # Check if steps are registered
+        registry: StepRegistry[DeploymentContext] = StepRegistry()
+        self.setup_steps(registry)
+
+        if registry:
+            # Step-based execution
+            context = self.create_context()
+            orchestrator: Orchestrator[DeploymentContext] = Orchestrator(
+                registry=registry,
+                context=context,
+                progress_callback=lambda curr, total: self.set_progress(curr, total),
+            )
+            result = orchestrator.execute()
+
+            # Add context data to result
+            if context.serve_app_id:
+                result['serve_application'] = context.serve_app_id
+            result['deployed'] = context.deployed
+
+            return result
+
+        # Simple execute mode
+        return self.execute()
+
+    def get_serve_app_name(self) -> str:
+        """Get the name for the Ray Serve application.
+
+        Default uses plugin release code from SYNAPSE_PLUGIN_RELEASE_CODE env var.
+        Override for custom naming.
+
+        Returns:
+            Serve application name.
+        """
+        return self.ctx.env.get_str('SYNAPSE_PLUGIN_RELEASE_CODE', 'synapse-serve-app') or 'synapse-serve-app'
+
+    def get_route_prefix(self) -> str:
+        """Get the route prefix for the deployment.
+
+        Default uses plugin release checksum from SYNAPSE_PLUGIN_RELEASE_CHECKSUM env var.
+        Override for custom routing.
+
+        Returns:
+            Route prefix string (e.g., '/abc123').
+        """
+        checksum = self.ctx.env.get_str('SYNAPSE_PLUGIN_RELEASE_CHECKSUM', 'default') or 'default'
+        return f'/{checksum}'
+
+    def get_ray_actor_options(self) -> dict[str, Any]:
+        """Get Ray actor options for the deployment.
+
+        Default extracts num_cpus and num_gpus from params.
+        Override for custom resource allocation.
+
+        Returns:
+            Dict with Ray actor options (num_cpus, num_gpus, etc.).
+        """
+        options: dict[str, Any] = {
+            'runtime_env': self.get_runtime_env(),
+        }
+
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+
+        for option in ['num_cpus', 'num_gpus', 'memory']:
+            if value := params_dict.get(option):
+                options[option] = value
+
+        return options
+
+    def get_runtime_env(self) -> dict[str, Any]:
+        """Get Ray runtime environment.
+
+        Override to customize the runtime environment for deployments.
+
+        Returns:
+            Dict with runtime environment configuration.
+        """
+        return {}
+
+    def ray_init(self, **kwargs: Any) -> None:
+        """Initialize Ray cluster connection.
+
+        Call this before deploying to ensure Ray is connected.
+
+        Args:
+            **kwargs: Additional arguments for ray.init().
+        """
+        try:
+            import ray
+        except ImportError:
+            raise ImportError("Ray is required for deployment actions. Install with: pip install 'synapse-sdk[ray]'")
+
+        if not ray.is_initialized():
+            ray.init(**kwargs)
+
+    def deploy(self) -> None:
+        """Deploy the inference endpoint to Ray Serve.
+
+        Uses the entrypoint class and current configuration to create
+        a Ray Serve deployment.
+
+        Raises:
+            RuntimeError: If entrypoint is not set.
+            ImportError: If Ray Serve is not installed.
+        """
+        if self.entrypoint is None:
+            raise RuntimeError(
+                'entrypoint must be set to a serve deployment class. Example: entrypoint = MyServeDeployment'
+            )
+
+        try:
+            from ray import serve
+        except ImportError:
+            raise ImportError(
+                "Ray Serve is required for deployment actions. Install with: pip install 'synapse-sdk[ray]'"
+            )
+
+        # Get deployment configuration
+        app_name = self.get_serve_app_name()
+        route_prefix = self.get_route_prefix()
+        ray_actor_options = self.get_ray_actor_options()
+
+        # Delete existing deployment if present
+        try:
+            serve.delete(app_name)
+        except Exception:
+            pass  # Ignore if not exists
+
+        # Get backend URL for the deployment
+        backend_url = self.ctx.env.get_str('SYNAPSE_PLUGIN_RUN_HOST', '') or ''
+
+        # Create and deploy
+        # The entrypoint should be a class that implements BaseServeDeployment
+        deployment = serve.deployment(ray_actor_options=ray_actor_options)(self.entrypoint).bind(backend_url)
+
+        serve.run(
+            deployment,
+            name=app_name,
+            route_prefix=route_prefix,
+        )
+
+    def register_serve_application(self) -> int | None:
+        """Register the serve application with the backend.
+
+        Creates a serve application record in the backend for tracking.
+
+        Returns:
+            Serve application ID if created, None otherwise.
+        """
+        job_id = self.ctx.job_id
+        if not job_id:
+            return None
+
+        app_name = self.get_serve_app_name()
+
+        # Get serve application status from Ray
+        try:
+            serve_app = self.agent_client.get_serve_application(app_name)
+        except Exception:
+            return None
+
+        # Register with backend
+        result = self.client.create_serve_application({
+            'job': job_id,
+            'status': serve_app.get('status'),
+            'data': serve_app,
+        })
+
+        return result.get('id')

synapse_sdk/plugins/actions/inference/serve.py

@@ -0,0 +1,252 @@
+"""Base Ray Serve deployment class for inference endpoints."""
+
+from __future__ import annotations
+
+import tempfile
+from abc import ABC, abstractmethod
+from typing import Any
+
+
+class BaseServeDeployment(ABC):
+    """Base class for Ray Serve inference deployments.
+
+    Provides model loading with multiplexing support. Subclasses implement
+    _get_model() to load their specific model format and infer() to run
+    inference.
+
+    This class is designed to be used with Ray Serve's @serve.deployment
+    decorator and supports model multiplexing via @serve.multiplexed().
+
+    Attributes:
+        backend_url: URL of the Synapse backend for model fetching.
+        _model_cache: Internal cache for loaded models.
+
+    Example:
+        >>> from ray import serve
+        >>> from fastapi import FastAPI
+        >>>
+        >>> app = FastAPI()
+        >>>
+        >>> @serve.deployment
+        >>> @serve.ingress(app)
+        >>> class MyInference(BaseServeDeployment):
+        ...     async def _get_model(self, model_info: dict) -> Any:
+        ...         import torch
+        ...         return torch.load(model_info['path'] / 'model.pt')
+        ...
+        ...     async def infer(self, inputs: list[dict]) -> list[dict]:
+        ...         model = await self.get_model()
+        ...         return [{'prediction': model(inp)} for inp in inputs]
+        >>>
+        >>> # Deploy with:
+        >>> deployment = MyInference.bind(backend_url='https://api.example.com')
+        >>> serve.run(deployment)
+    """
+
+    def __init__(self, backend_url: str) -> None:
+        """Initialize the serve deployment.
+
+        Args:
+            backend_url: URL of the Synapse backend for fetching models.
+        """
+        self.backend_url = backend_url
+        self._model_cache: dict[str, Any] = {}
+
+    async def _load_model_from_token(self, model_token: str) -> Any:
+        """Load model from an encoded token.
+
+        Decodes the JWT token containing model info, fetches model from
+        backend, downloads and extracts artifacts, then calls _get_model().
+
+        Args:
+            model_token: JWT-encoded token with model info.
+
+        Returns:
+            Loaded model object (format depends on _get_model implementation).
+
+        Raises:
+            ImportError: If jwt or required dependencies not installed.
+        """
+        try:
+            import jwt
+        except ImportError:
+            raise ImportError('PyJWT is required for model token decoding. Install with: pip install PyJWT')
+
+        # Decode token to get model info
+        model_info = jwt.decode(model_token, self.backend_url, algorithms=['HS256'])
+
+        # Create backend client with user credentials
+        from synapse_sdk.clients.backend import BackendClient
+
+        client = BackendClient(
+            base_url=self.backend_url,
+            access_token=model_info['token'],
+            tenant=model_info.get('tenant'),
+        )
+
+        # Fetch model metadata
+        model = client.get_model(int(model_info['model']))
+
+        if not model.get('file'):
+            raise ValueError(f'Model {model_info["model"]} has no file URL')
+
+        # Download and extract model
+        with tempfile.TemporaryDirectory() as temp_path:
+            from pathlib import Path
+
+            from synapse_sdk.utils.file.archive import extract_archive
+            from synapse_sdk.utils.file.download import download_file
+
+            archive_path = Path(temp_path) / 'model.zip'
+            download_file(model['file'], archive_path)
+            extract_archive(archive_path, temp_path)
+
+            model['path'] = Path(temp_path)
+            return await self._get_model(model)
+
+    async def get_model(self) -> Any:
+        """Get the current model for inference.
+
+        Uses Ray Serve's multiplexing to load the appropriate model
+        based on the request's multiplexed model ID header.
+
+        Returns:
+            Loaded model object.
+
+        Note:
+            This method uses Ray Serve's @serve.multiplexed() decorator
+            internally. Ensure requests include the appropriate header.
+        """
+        # Import here to avoid issues when Ray is not installed
+        try:
+            from ray import serve
+        except ImportError:
+            raise ImportError("Ray Serve is required. Install with: pip install 'synapse-sdk[ray]'")
+
+        model_id = serve.get_multiplexed_model_id()
+        return await self._load_model_multiplexed(model_id)
+
+    async def _load_model_multiplexed(self, model_id: str) -> Any:
+        """Load model with multiplexing support.
+
+        This method is decorated with @serve.multiplexed() to enable
+        model multiplexing in Ray Serve deployments.
+
+        Args:
+            model_id: The model token/ID from request header.
+
+        Returns:
+            Loaded model object.
+        """
+        # Check cache first
+        if model_id in self._model_cache:
+            return self._model_cache[model_id]
+
+        # Load and cache
+        model = await self._load_model_from_token(model_id)
+        self._model_cache[model_id] = model
+        return model
+
+    @abstractmethod
+    async def _get_model(self, model_info: dict[str, Any]) -> Any:
+        """Load model from extracted artifacts.
+
+        Override this method to implement your specific model loading logic.
+        Called after model artifacts are downloaded and extracted.
+
+        Args:
+            model_info: Model metadata dict with 'path' key for local artifacts.
+                       The path is a Path object pointing to extracted directory.
+
+        Returns:
+            Loaded model object (framework-specific).
+
+        Example (PyTorch):
+            >>> async def _get_model(self, model_info: dict) -> Any:
+            ...     import torch
+            ...     model_path = model_info['path'] / 'model.pt'
+            ...     return torch.load(model_path)
+
+        Example (ONNX):
+            >>> async def _get_model(self, model_info: dict) -> Any:
+            ...     import onnxruntime as ort
+            ...     model_path = model_info['path'] / 'model.onnx'
+            ...     return ort.InferenceSession(str(model_path))
+        """
+        raise NotImplementedError(
+            'Override _get_model() to load your model format. '
+            'Example: return torch.load(model_info["path"] / "model.pt")'
+        )
+
+    @abstractmethod
+    async def infer(self, *args: Any, **kwargs: Any) -> Any:
+        """Run inference on inputs.
+
+        Override this method to implement your inference logic.
+        Use self.get_model() to obtain the loaded model.
+
+        Args:
+            *args: Inference inputs (format depends on implementation).
+            **kwargs: Additional inference parameters.
+
+        Returns:
+            Inference results (format depends on implementation).
+
+        Example:
+            >>> async def infer(self, inputs: list[dict]) -> list[dict]:
+            ...     model = await self.get_model()
+            ...     results = []
+            ...     for inp in inputs:
+            ...         prediction = model.predict(inp['data'])
+            ...         results.append({'prediction': prediction.tolist()})
+            ...     return results
+        """
+        raise NotImplementedError(
+            'Override infer() to implement inference logic. Example: return model.predict(inputs)'
+        )
+
+
+def create_serve_multiplexed_model_id(
+    model_id: int | str,
+    token: str,
+    backend_url: str,
+    tenant: str | None = None,
+) -> str:
+    """Create a JWT-encoded model ID for serve multiplexing.
+
+    This helper creates the token that should be passed in the
+    'serve_multiplexed_model_id' header for inference requests.
+
+    Args:
+        model_id: The model ID to encode.
+        token: User access token for authentication.
+        backend_url: Backend URL (used as JWT secret).
+        tenant: Optional tenant identifier.
+
+    Returns:
+        JWT-encoded model token string.
+
+    Example:
+        >>> model_token = create_serve_multiplexed_model_id(
+        ...     model_id=123,
+        ...     token='user_access_token',
+        ...     backend_url='https://api.example.com',
+        ...     tenant='my-tenant',
+        ... )
+        >>> # Use in request headers:
+        >>> headers = {'serve_multiplexed_model_id': model_token}
+    """
+    try:
+        import jwt
+    except ImportError:
+        raise ImportError('PyJWT is required for model token encoding. Install with: pip install PyJWT')
+
+    payload: dict[str, Any] = {
+        'model': str(model_id),
+        'token': token,
+    }
+
+    if tenant:
+        payload['tenant'] = tenant
+
+    return jwt.encode(payload, backend_url, algorithm='HS256')

synapse_sdk/plugins/actions/train/__init__.py

@@ -0,0 +1,54 @@
+"""Train action module with optional workflow step support.
+
+Provides the training action base class:
+    - BaseTrainAction: Base class for training workflows
+    - TrainContext: Training-specific context extending BaseStepContext
+    - TrainProgressCategories: Standard progress category names
+
+For step infrastructure (BaseStep, StepRegistry, Orchestrator),
+use the steps module:
+    from synapse_sdk.plugins.steps import BaseStep, StepRegistry
+
+Example (simple execute):
+    >>> class MyTrainAction(BaseTrainAction[MyParams]):
+    ...     def execute(self) -> dict[str, Any]:
+    ...         dataset = self.get_dataset()
+    ...         # ... train model ...
+    ...         return {'model_id': model['id']}
+
+Example (step-based):
+    >>> from synapse_sdk.plugins.steps import BaseStep, StepResult
+    >>>
+    >>> class LoadDatasetStep(BaseStep[TrainContext]):
+    ...     @property
+    ...     def name(self) -> str:
+    ...         return 'load_dataset'
+    ...
+    ...     @property
+    ...     def progress_weight(self) -> float:
+    ...         return 0.2
+    ...
+    ...     def execute(self, context: TrainContext) -> StepResult:
+    ...         context.dataset = load_data(context.params['dataset_id'])
+    ...         return StepResult(success=True)
+    >>>
+    >>> class MyTrainAction(BaseTrainAction[MyParams]):
+    ...     def setup_steps(self, registry) -> None:
+    ...         registry.register(LoadDatasetStep())
+    ...         registry.register(TrainStep())
+    ...         registry.register(UploadModelStep())
+"""
+
+from synapse_sdk.plugins.actions.train.action import (
+    BaseTrainAction,
+    BaseTrainParams,
+    TrainProgressCategories,
+)
+from synapse_sdk.plugins.actions.train.context import TrainContext
+
+__all__ = [
+    'BaseTrainAction',
+    'BaseTrainParams',
+    'TrainContext',
+    'TrainProgressCategories',
+]