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

synapse_sdk/plugins/actions/export/action.py

@@ -0,0 +1,183 @@
+"""Export action base class with optional step support."""
+
+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.export.context import ExportContext
+from synapse_sdk.plugins.enums import PluginCategory
+from synapse_sdk.plugins.steps import Orchestrator, StepRegistry
+
+P = TypeVar('P', bound=BaseModel)
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.backend import BackendClient
+
+
+class ExportProgressCategories:
+    """Standard progress category names for export workflows.
+
+    Use these constants with set_progress() to track export phases:
+        - DATASET_CONVERSION: Data conversion and file generation
+
+    Example:
+        >>> self.set_progress(50, 100, self.progress.DATASET_CONVERSION)
+    """
+
+    DATASET_CONVERSION: str = 'dataset_conversion'
+
+
+class BaseExportAction(BaseAction[P]):
+    """Base class for export actions.
+
+    Provides helper methods for export workflows.
+    Override get_filtered_results() for your specific target type.
+
+    Supports two execution modes:
+    1. Simple execute: Override execute() directly for simple workflows
+    2. Step-based: Override setup_steps() to register workflow steps
+
+    If setup_steps() registers any steps, the step-based workflow
+    takes precedence and execute() is not called directly.
+
+    Attributes:
+        category: Plugin category (defaults to EXPORT).
+        progress: Standard progress category names.
+
+    Example (simple execute):
+        >>> class MyExportAction(BaseExportAction[MyParams]):
+        ...     action_name = 'export'
+        ...     params_model = MyParams
+        ...
+        ...     def get_filtered_results(self, filters: dict) -> tuple[Any, int]:
+        ...         return self.client.get_assignments(filters)
+        ...
+        ...     def execute(self) -> dict[str, Any]:
+        ...         results, count = self.get_filtered_results(self.params.filter)
+        ...         self.set_progress(0, count, self.progress.DATASET_CONVERSION)
+        ...         for i, item in enumerate(results, 1):
+        ...             # Process and export item
+        ...             self.set_progress(i, count, self.progress.DATASET_CONVERSION)
+        ...         return {'exported': count}
+
+    Example (step-based):
+        >>> class MyExportAction(BaseExportAction[MyParams]):
+        ...     def setup_steps(self, registry: StepRegistry[ExportContext]) -> None:
+        ...         registry.register(FetchResultsStep())
+        ...         registry.register(ProcessStep())
+        ...         registry.register(FinalizeStep())
+    """
+
+    category = PluginCategory.EXPORT
+    progress = ExportProgressCategories()
+
+    @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. Either provide a client via RuntimeContext or override get_filtered_results().'
+            )
+        return self.ctx.client
+
+    def setup_steps(self, registry: StepRegistry[ExportContext]) -> None:
+        """Register workflow steps for step-based execution.
+
+        Override this method to register custom steps for your export workflow.
+        If steps are registered, step-based execution takes precedence.
+
+        Args:
+            registry: StepRegistry to register steps with.
+
+        Example:
+            >>> def setup_steps(self, registry: StepRegistry[ExportContext]) -> None:
+            ...     registry.register(FetchResultsStep())
+            ...     registry.register(ProcessStep())
+            ...     registry.register(FinalizeStep())
+        """
+        pass  # Default: no steps, uses simple execute()
+
+    def create_context(self) -> ExportContext:
+        """Create export context for step-based workflow.
+
+        Override to customize context creation or add additional state.
+
+        Returns:
+            ExportContext instance with params and runtime context.
+        """
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+        return ExportContext(
+            runtime_ctx=self.ctx,
+            params=params_dict,
+        )
+
+    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[ExportContext] = StepRegistry()
+        self.setup_steps(registry)
+
+        if registry:
+            # Step-based execution
+            context = self.create_context()
+            orchestrator: Orchestrator[ExportContext] = Orchestrator(
+                registry=registry,
+                context=context,
+                progress_callback=lambda curr, total: self.set_progress(curr, total),
+            )
+            result = orchestrator.execute()
+
+            # Add context data to result
+            result['exported_count'] = context.exported_count
+            if context.output_path:
+                result['output_path'] = context.output_path
+
+            return result
+
+        # Simple execute mode
+        return self.execute()
+
+    def get_filtered_results(self, filters: dict[str, Any]) -> tuple[Any, int]:
+        """Fetch filtered results for export.
+
+        Override this method for your specific target type
+        (assignments, ground_truth, tasks, or custom).
+
+        Args:
+            filters: Filter criteria dict.
+
+        Returns:
+            Tuple of (results_iterator, total_count).
+
+        Raises:
+            NotImplementedError: Must be overridden by subclass.
+
+        Example:
+            >>> # Override for assignments:
+            >>> def get_filtered_results(self, filters: dict) -> tuple[Any, int]:
+            ...     return self.client.get_assignments(filters)
+            >>>
+            >>> # Override for ground truth:
+            >>> def get_filtered_results(self, filters: dict) -> tuple[Any, int]:
+            ...     return self.client.get_ground_truth(filters)
+        """
+        raise NotImplementedError(
+            'Override get_filtered_results() for your target type. Example: return self.client.get_assignments(filters)'
+        )

synapse_sdk/plugins/actions/export/context.py

@@ -0,0 +1,59 @@
+"""Export context for sharing state between workflow steps."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+from synapse_sdk.plugins.steps import BaseStepContext
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.backend import BackendClient
+
+
+@dataclass
+class ExportContext(BaseStepContext):
+    """Shared context passed between export workflow steps.
+
+    Extends BaseStepContext with export-specific state fields.
+    Carries parameters and accumulated state as the workflow
+    progresses through steps.
+
+    Attributes:
+        params: Export parameters (from action params).
+        results: Fetched results to export (populated by fetch step).
+        total_count: Total number of items to export.
+        exported_count: Number of items successfully exported.
+        output_path: Path to export output file/directory.
+
+    Example:
+        >>> context = ExportContext(
+        ...     runtime_ctx=runtime_ctx,
+        ...     params={'format': 'coco', 'filter': {}},
+        ... )
+        >>> # Steps populate state as they execute
+        >>> context.results = fetched_data
+    """
+
+    # Export parameters
+    params: dict[str, Any] = field(default_factory=dict)
+
+    # Processing state (populated by steps)
+    results: Any | None = None
+    total_count: int = 0
+    exported_count: int = 0
+    output_path: str | None = None
+
+    @property
+    def client(self) -> BackendClient:
+        """Backend client from runtime context.
+
+        Returns:
+            BackendClient instance.
+
+        Raises:
+            RuntimeError: If no client in runtime context.
+        """
+        if self.runtime_ctx.client is None:
+            raise RuntimeError('No client in runtime context')
+        return self.runtime_ctx.client

synapse_sdk/plugins/actions/inference/__init__.py

@@ -0,0 +1,84 @@
+"""Inference action module with Ray Serve integration.
+
+This module provides base classes for inference and deployment actions:
+
+- **BaseInferenceAction**: For batch inference or REST API inference
+- **BaseDeploymentAction**: For deploying models to Ray Serve
+- **BaseServeDeployment**: Ray Serve deployment class for model serving
+- **InferenceContext**: Step-based workflow context for inference
+- **DeploymentContext**: Step-based workflow context for deployment
+
+Example (Simple Inference):
+    >>> from synapse_sdk.plugins.actions.inference import BaseInferenceAction
+    >>>
+    >>> class MyInferenceAction(BaseInferenceAction[MyParams]):
+    ...     action_name = 'inference'
+    ...     category = 'neural_net'
+    ...     params_model = MyParams
+    ...
+    ...     def execute(self) -> dict:
+    ...         model_info = self.load_model(self.params.model_id)
+    ...         results = self.infer(model_info, self.params.inputs)
+    ...         return {'results': results}
+
+Example (Ray Serve Deployment):
+    >>> from synapse_sdk.plugins.actions.inference import (
+    ...     BaseDeploymentAction,
+    ...     BaseServeDeployment,
+    ... )
+    >>>
+    >>> class MyServeDeployment(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 [{'pred': model(inp)} for inp in inputs]
+    >>>
+    >>> class MyDeploymentAction(BaseDeploymentAction[MyParams]):
+    ...     action_name = 'deployment'
+    ...     category = 'neural_net'
+    ...     params_model = MyParams
+    ...     entrypoint = MyServeDeployment
+    ...
+    ...     def execute(self) -> dict:
+    ...         self.ray_init()
+    ...         self.deploy()
+    ...         app_id = self.register_serve_application()
+    ...         return {'serve_application': app_id}
+"""
+
+from __future__ import annotations
+
+from synapse_sdk.plugins.actions.inference.action import (
+    BaseInferenceAction,
+    InferenceProgressCategories,
+)
+from synapse_sdk.plugins.actions.inference.context import (
+    DeploymentContext,
+    InferenceContext,
+)
+from synapse_sdk.plugins.actions.inference.deployment import (
+    BaseDeploymentAction,
+    DeploymentProgressCategories,
+)
+from synapse_sdk.plugins.actions.inference.serve import (
+    BaseServeDeployment,
+    create_serve_multiplexed_model_id,
+)
+
+__all__ = [
+    # Actions
+    'BaseInferenceAction',
+    'BaseDeploymentAction',
+    # Serve
+    'BaseServeDeployment',
+    'create_serve_multiplexed_model_id',
+    # Contexts
+    'InferenceContext',
+    'DeploymentContext',
+    # Progress Categories
+    'InferenceProgressCategories',
+    'DeploymentProgressCategories',
+]

synapse_sdk/plugins/actions/inference/action.py

@@ -0,0 +1,285 @@
+"""Inference action base class with optional step support."""
+
+from __future__ import annotations
+
+import tempfile
+from pathlib import Path
+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 InferenceContext
+from synapse_sdk.plugins.enums import PluginCategory
+from synapse_sdk.plugins.steps import Orchestrator, StepRegistry
+from synapse_sdk.utils.file.archive import extract_archive
+
+P = TypeVar('P', bound=BaseModel)
+
+if TYPE_CHECKING:
+    from synapse_sdk.clients.backend import BackendClient
+
+
+class InferenceProgressCategories:
+    """Standard progress category names for inference workflows.
+
+    Use these constants with set_progress() to track inference phases:
+        - MODEL_LOAD: Model loading and initialization
+        - INFERENCE: Running inference on inputs
+        - POSTPROCESS: Post-processing results
+
+    Example:
+        >>> self.set_progress(1, 3, self.progress.MODEL_LOAD)
+        >>> self.set_progress(2, 3, self.progress.INFERENCE)
+    """
+
+    MODEL_LOAD: str = 'model_load'
+    INFERENCE: str = 'inference'
+    POSTPROCESS: str = 'postprocess'
+
+
+class BaseInferenceAction(BaseAction[P]):
+    """Base class for inference actions.
+
+    Provides helper methods for model loading and inference workflows.
+    Supports both REST API-based inference (via Ray Serve) and batch
+    inference with step-based workflows.
+
+    Supports two execution modes:
+    1. Simple execute: Override execute() directly for simple workflows
+    2. Step-based: Override setup_steps() to register workflow steps
+
+    If setup_steps() registers any steps, the step-based workflow
+    takes precedence and execute() is not called directly.
+
+    Attributes:
+        category: Plugin category (defaults to NEURAL_NET).
+        progress: Standard progress category names.
+
+    Example (simple execute):
+        >>> class MyInferenceAction(BaseInferenceAction[MyParams]):
+        ...     action_name = 'inference'
+        ...     params_model = MyParams
+        ...
+        ...     def execute(self) -> dict[str, Any]:
+        ...         model = self.load_model(self.params.model_id)
+        ...         self.set_progress(1, 3, self.progress.MODEL_LOAD)
+        ...         results = self.infer(model, self.params.inputs)
+        ...         self.set_progress(2, 3, self.progress.INFERENCE)
+        ...         return {'results': results}
+
+    Example (step-based):
+        >>> class MyInferenceAction(BaseInferenceAction[MyParams]):
+        ...     def setup_steps(self, registry: StepRegistry[InferenceContext]) -> None:
+        ...         registry.register(LoadModelStep())
+        ...         registry.register(InferenceStep())
+        ...         registry.register(PostProcessStep())
+    """
+
+    category = PluginCategory.NEURAL_NET
+    progress = InferenceProgressCategories()
+
+    @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. Either provide a client via RuntimeContext '
+                'or override the helper methods (load_model, get_model).'
+            )
+        return self.ctx.client
+
+    def setup_steps(self, registry: StepRegistry[InferenceContext]) -> None:
+        """Register workflow steps for step-based execution.
+
+        Override this method to register custom steps for your inference workflow.
+        If steps are registered, step-based execution takes precedence.
+
+        Args:
+            registry: StepRegistry to register steps with.
+
+        Example:
+            >>> def setup_steps(self, registry: StepRegistry[InferenceContext]) -> None:
+            ...     registry.register(LoadModelStep())
+            ...     registry.register(InferenceStep())
+            ...     registry.register(PostProcessStep())
+        """
+        pass  # Default: no steps, uses simple execute()
+
+    def create_context(self) -> InferenceContext:
+        """Create inference context for step-based workflow.
+
+        Override to customize context creation or add additional state.
+
+        Returns:
+            InferenceContext instance with params and runtime context.
+        """
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+        return InferenceContext(
+            runtime_ctx=self.ctx,
+            params=params_dict,
+            model_id=params_dict.get('model_id'),
+        )
+
+    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[InferenceContext] = StepRegistry()
+        self.setup_steps(registry)
+
+        if registry:
+            # Step-based execution
+            context = self.create_context()
+            orchestrator: Orchestrator[InferenceContext] = Orchestrator(
+                registry=registry,
+                context=context,
+                progress_callback=lambda curr, total: self.set_progress(curr, total),
+            )
+            result = orchestrator.execute()
+
+            # Add context data to result
+            result['processed_count'] = context.processed_count
+            if context.results:
+                result['results'] = context.results
+
+            return result
+
+        # Simple execute mode
+        return self.execute()
+
+    def get_model(self, model_id: int) -> dict[str, Any]:
+        """Retrieve model metadata by ID.
+
+        Args:
+            model_id: Model identifier.
+
+        Returns:
+            Model metadata dictionary including file URL.
+
+        Raises:
+            RuntimeError: If no client in context.
+
+        Example:
+            >>> model = self.get_model(123)
+            >>> print(model['name'], model['file'])
+        """
+        return self.client.get_model(model_id)
+
+    def download_model(
+        self,
+        model_id: int,
+        output_dir: str | Path | None = None,
+    ) -> Path:
+        """Download and extract model artifacts.
+
+        Fetches model metadata, downloads the model archive, and extracts
+        it to the specified directory (or a temp directory if not specified).
+
+        Args:
+            model_id: Model identifier.
+            output_dir: Directory to extract model to. If None, uses tempdir.
+
+        Returns:
+            Path to extracted model directory.
+
+        Raises:
+            RuntimeError: If no client in context.
+            ValueError: If model has no file URL.
+
+        Example:
+            >>> model_path = self.download_model(123)
+            >>> # Load model from model_path
+        """
+        model = self.get_model(model_id)
+
+        if not model.get('file'):
+            raise ValueError(f'Model {model_id} has no file URL')
+
+        # Determine output directory
+        if output_dir is None:
+            output_dir = Path(tempfile.mkdtemp(prefix='synapse_model_'))
+        else:
+            output_dir = Path(output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+
+        # Download and extract
+        from synapse_sdk.utils.file.download import download_file
+
+        archive_path = output_dir / 'model.zip'
+        download_file(model['file'], archive_path)
+        extract_archive(archive_path, output_dir)
+
+        # Remove archive after extraction
+        archive_path.unlink(missing_ok=True)
+
+        return output_dir
+
+    def load_model(self, model_id: int) -> dict[str, Any]:
+        """Load model for inference.
+
+        Downloads model artifacts and returns model info with local path.
+        Override this method for custom model loading (e.g., loading into
+        specific framework like PyTorch, TensorFlow).
+
+        Args:
+            model_id: Model identifier.
+
+        Returns:
+            Model metadata dict with 'path' key for local artifacts.
+
+        Example:
+            >>> model_info = self.load_model(123)
+            >>> model_path = model_info['path']
+            >>> # Load your model framework here:
+            >>> # model = torch.load(model_path / 'model.pt')
+        """
+        model = self.get_model(model_id)
+        model_path = self.download_model(model_id)
+        model['path'] = str(model_path)
+        return model
+
+    def infer(
+        self,
+        model: Any,
+        inputs: list[dict[str, Any]],
+    ) -> list[dict[str, Any]]:
+        """Run inference on inputs.
+
+        Override this method to implement your inference logic.
+        This is called by execute() in simple mode.
+
+        Args:
+            model: Loaded model (framework-specific).
+            inputs: List of input dictionaries.
+
+        Returns:
+            List of result dictionaries.
+
+        Raises:
+            NotImplementedError: Must be overridden by subclass.
+
+        Example:
+            >>> def infer(self, model, inputs):
+            ...     results = []
+            ...     for inp in inputs:
+            ...         prediction = model.predict(inp['image'])
+            ...         results.append({'prediction': prediction})
+            ...     return results
+        """
+        raise NotImplementedError(
+            'Override infer() to implement inference logic. Example: return model.predict(inputs)'
+        )

synapse_sdk/plugins/actions/inference/context.py

@@ -0,0 +1,81 @@
+"""Inference context for step-based workflows."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+from synapse_sdk.plugins.steps import BaseStepContext
+
+
+@dataclass
+class InferenceContext(BaseStepContext):
+    """Context for inference action step-based workflows.
+
+    Extends BaseStepContext with inference-specific state including
+    model information, request/response tracking, and batch processing.
+
+    Attributes:
+        params: Action parameters dict.
+        model_id: ID of the model being used for inference.
+        model: Loaded model information from backend.
+        model_path: Local path to downloaded/extracted model.
+        requests: Input requests to process.
+        results: Inference results.
+        batch_size: Batch size for processing.
+        processed_count: Number of processed items.
+
+    Example:
+        >>> context = InferenceContext(
+        ...     runtime_ctx=self.ctx,
+        ...     params={'model_id': 123},
+        ...     model_id=123,
+        ... )
+        >>> context.results.append({'prediction': 0.95})
+    """
+
+    params: dict[str, Any] = field(default_factory=dict)
+    model_id: int | None = None
+    model: dict[str, Any] | None = None
+    model_path: str | None = None
+    requests: list[dict[str, Any]] = field(default_factory=list)
+    results: list[dict[str, Any]] = field(default_factory=list)
+    batch_size: int = 1
+    processed_count: int = 0
+
+
+@dataclass
+class DeploymentContext(BaseStepContext):
+    """Context for deployment action step-based workflows.
+
+    Extends BaseStepContext with deployment-specific state including
+    model information, serve application configuration, and deployment status.
+
+    Attributes:
+        params: Action parameters dict.
+        model_id: ID of the model to deploy.
+        model: Model information from backend.
+        model_path: Local path to model artifacts.
+        serve_app_name: Name of the Ray Serve application.
+        serve_app_id: ID of the created serve application.
+        route_prefix: URL route prefix for the deployment.
+        ray_actor_options: Ray actor configuration options.
+        deployed: Whether deployment succeeded.
+
+    Example:
+        >>> context = DeploymentContext(
+        ...     runtime_ctx=self.ctx,
+        ...     params={'model_id': 123},
+        ...     serve_app_name='my-model-v1',
+        ... )
+    """
+
+    params: dict[str, Any] = field(default_factory=dict)
+    model_id: int | None = None
+    model: dict[str, Any] | None = None
+    model_path: str | None = None
+    serve_app_name: str | None = None
+    serve_app_id: int | None = None
+    route_prefix: str | None = None
+    ray_actor_options: dict[str, Any] = field(default_factory=dict)
+    deployed: bool = False