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

synapse_sdk/plugins/actions/train/action.py

@@ -0,0 +1,326 @@
+"""Train action base class with optional step support."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, TypeVar
+
+from pydantic import BaseModel, Field
+
+from synapse_sdk.plugins.action import BaseAction
+from synapse_sdk.plugins.actions.train.context import TrainContext
+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 BaseTrainParams(BaseModel):
+    """Base parameters for training actions.
+
+    Provides common fields used across training workflows.
+    Extend this class to add plugin-specific training parameters.
+
+    Attributes:
+        checkpoint: Optional model ID to use as starting checkpoint.
+
+    Example:
+        >>> class MyTrainParams(BaseTrainParams):
+        ...     epochs: int = 100
+        ...     batch_size: int = 32
+    """
+
+    checkpoint: int | None = Field(default=None, description='Checkpoint model ID to resume from')
+
+
+class TrainProgressCategories:
+    """Standard progress category names for training workflows.
+
+    Use these constants with set_progress() to track multi-phase training:
+        - DATASET: Data loading and preprocessing
+        - TRAIN: Model training iterations
+        - MODEL_UPLOAD: Final model upload to backend
+
+    Example:
+        >>> self.set_progress(1, 10, TrainProgressCategories.DATASET)
+        >>> self.set_progress(50, 100, TrainProgressCategories.TRAIN)
+    """
+
+    DATASET: str = 'dataset'
+    TRAIN: str = 'train'
+    MODEL_UPLOAD: str = 'model_upload'
+
+
+class BaseTrainAction(BaseAction[P]):
+    """Base class for training actions.
+
+    Provides helper methods for common training workflows:
+    dataset fetching, model creation, and progress tracking.
+
+    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 (without result schema):
+        >>> class MyTrainAction(BaseTrainAction[MyParams]):
+        ...     def execute(self) -> dict[str, Any]:
+        ...         return {'weights_path': '/model.pt'}
+
+    Example (with result schema):
+        >>> class TrainResult(BaseModel):
+        ...     weights_path: str
+        ...     final_loss: float
+        >>>
+        >>> class MyTrainAction(BaseTrainAction[MyParams]):
+        ...     result_model = TrainResult
+        ...
+        ...     def execute(self) -> TrainResult:
+        ...         return TrainResult(weights_path='/model.pt', final_loss=0.1)
+
+    Example (step-based):
+        >>> class MyTrainAction(BaseTrainAction[MyParams]):
+        ...     def setup_steps(self, registry: StepRegistry[TrainContext]) -> None:
+        ...         registry.register(LoadDatasetStep())
+        ...         registry.register(TrainStep())
+        ...         registry.register(UploadModelStep())
+    """
+
+    category = PluginCategory.NEURAL_NET
+    progress = TrainProgressCategories()
+
+    def autolog(self, framework: str) -> None:
+        """Enable automatic logging for an ML framework.
+
+        Call this before creating model objects. The SDK will automatically
+        attach callbacks to log progress, metrics, and artifacts.
+
+        Supported frameworks:
+            - 'ultralytics': YOLO object detection models
+
+        Args:
+            framework: Framework name (e.g., 'ultralytics').
+
+        Raises:
+            ValueError: If framework is not recognized.
+            ImportError: If framework package is not installed.
+
+        Example:
+            >>> def execute(self):
+            ...     self.autolog('ultralytics')
+            ...     model = YOLO('yolov8n.pt')  # Callbacks auto-attached
+            ...     model.train(epochs=100)     # Metrics logged automatically
+        """
+        from synapse_sdk.integrations import autolog as _autolog
+
+        _autolog(framework, self)
+
+    @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 (get_dataset, create_model, get_model).'
+            )
+        return self.ctx.client
+
+    def setup_steps(self, registry: StepRegistry[TrainContext]) -> None:
+        """Register workflow steps for step-based execution.
+
+        Override this method to register custom steps for your training workflow.
+        If steps are registered, step-based execution takes precedence.
+
+        Args:
+            registry: StepRegistry to register steps with.
+
+        Example:
+            >>> def setup_steps(self, registry: StepRegistry[TrainContext]) -> None:
+            ...     registry.register(LoadDatasetStep())
+            ...     registry.register(TrainStep())
+            ...     registry.register(UploadModelStep())
+        """
+        pass  # Default: no steps, uses simple execute()
+
+    def create_context(self) -> TrainContext:
+        """Create training context for step-based workflow.
+
+        Override to customize context creation or add additional state.
+
+        Returns:
+            TrainContext instance with params and runtime context.
+        """
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+        return TrainContext(
+            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[TrainContext] = StepRegistry()
+        self.setup_steps(registry)
+
+        if registry:
+            # Step-based execution
+            context = self.create_context()
+            orchestrator: Orchestrator[TrainContext] = 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.model:
+                result['model'] = context.model
+
+            return result
+
+        # Simple execute mode
+        return self.execute()
+
+    def get_dataset(self) -> dict[str, Any]:
+        """Fetch training dataset from backend.
+
+        Default implementation uses params.dataset_id to fetch
+        via client.get_data_collection(). Override for custom behavior.
+
+        Returns:
+            Dataset metadata dictionary.
+
+        Raises:
+            ValueError: If params.dataset_id is not set.
+            RuntimeError: If no client in context.
+
+        Example:
+            >>> # Override for S3:
+            >>> def get_dataset(self) -> dict[str, Any]:
+            ...     return download_from_s3(self.params.s3_path)
+        """
+        dataset_id = getattr(self.params, 'dataset_id', None)
+        if dataset_id is None:
+            raise ValueError(
+                'params.dataset_id is required for default get_dataset(). '
+                'Either set dataset_id in your params model or override get_dataset().'
+            )
+        return self.client.get_data_collection(dataset_id)
+
+    def create_model(self, path: str, **kwargs: Any) -> dict[str, Any]:
+        """Upload trained model to backend.
+
+        Default implementation uploads via client.create_model().
+        Override for custom behavior (e.g., MLflow, S3).
+
+        Args:
+            path: Local path to model artifacts.
+            **kwargs: Additional fields for model creation.
+
+        Returns:
+            Created model metadata dictionary.
+
+        Raises:
+            RuntimeError: If no client in context.
+
+        Example:
+            >>> model = self.create_model('./model', name='my-model')
+        """
+        return self.client.create_model({
+            'file': path,
+            **kwargs,
+        })
+
+    def get_model(self, model_id: int) -> dict[str, Any]:
+        """Retrieve existing model by ID.
+
+        Args:
+            model_id: Model identifier.
+
+        Returns:
+            Model metadata dictionary.
+
+        Raises:
+            RuntimeError: If no client in context.
+        """
+        return self.client.get_model(model_id)
+
+    def get_checkpoint(self) -> dict[str, Any] | None:
+        """Get checkpoint for training, either from context or by fetching from backend.
+
+        This method handles checkpoint resolution in the following order:
+        1. If ctx.checkpoint is already set (remote mode), returns it directly
+        2. If params has a 'checkpoint' field (model ID), fetches and extracts the model
+
+        The returned checkpoint dict contains:
+            - category: 'base' or fine-tuned model category
+            - path: Local path to the extracted model files
+
+        Returns:
+            Checkpoint dict with 'category' and 'path', or None if no checkpoint.
+
+        Raises:
+            RuntimeError: If no client in context and checkpoint ID provided.
+            FileNotFoundError: If model file cannot be downloaded/extracted.
+
+        Example:
+            >>> checkpoint = self.get_checkpoint()
+            >>> if checkpoint:
+            ...     model_path = checkpoint['path']
+            ...     is_base = checkpoint['category'] == 'base'
+        """
+        from synapse_sdk.utils.file import extract_archive, get_temp_path
+
+        # If checkpoint is already in context (remote mode), return it
+        if self.ctx.checkpoint is not None:
+            return self.ctx.checkpoint
+
+        # Check if params has a checkpoint field (model ID)
+        checkpoint_id = getattr(self.params, 'checkpoint', None)
+        if checkpoint_id is None:
+            return None
+
+        # Fetch model from backend
+        model = self.get_model(checkpoint_id)
+
+        # The model['file'] is downloaded by the client's url_conversion
+        model_file = Path(model['file'])
+
+        # Extract to temp path
+        output_path = get_temp_path(f'models/{model_file.stem}')
+        if not output_path.exists():
+            output_path.mkdir(parents=True, exist_ok=True)
+            extract_archive(model_file, output_path)
+
+        # Determine category - base models vs fine-tuned
+        category = model.get('category') or 'base'
+
+        return {
+            'category': category,
+            'path': output_path,
+            'id': model.get('id'),
+            'name': model.get('name'),
+        }

synapse_sdk/plugins/actions/train/context.py

@@ -0,0 +1,57 @@
+"""Train 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 TrainContext(BaseStepContext):
+    """Shared context passed between training workflow steps.
+
+    Extends BaseStepContext with training-specific state fields.
+    Carries parameters and accumulated state as the workflow
+    progresses through steps.
+
+    Attributes:
+        params: Training parameters (from action params).
+        dataset: Loaded dataset (populated by dataset step).
+        model_path: Path to trained model (populated by training step).
+        model: Created model metadata (populated by upload step).
+
+    Example:
+        >>> context = TrainContext(
+        ...     runtime_ctx=runtime_ctx,
+        ...     params={'dataset_id': 1, 'epochs': 10},
+        ... )
+        >>> # Steps populate state as they execute
+        >>> context.dataset = loaded_dataset
+    """
+
+    # Training parameters
+    params: dict[str, Any] = field(default_factory=dict)
+
+    # Processing state (populated by steps)
+    dataset: Any | None = None
+    model_path: str | None = None
+    model: dict[str, Any] | 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/upload/__init__.py

@@ -0,0 +1,49 @@
+"""Upload action module with workflow step support.
+
+Provides a full step-based workflow system for upload actions:
+    - BaseUploadAction: Base class for upload workflows
+    - UploadContext: Upload-specific context extending BaseStepContext
+    - UploadProgressCategories: Standard progress category names
+
+For step infrastructure (BaseStep, StepRegistry, Orchestrator),
+use the steps module:
+    from synapse_sdk.plugins.steps import BaseStep, StepRegistry
+
+Example:
+    >>> from synapse_sdk.plugins.steps import BaseStep, StepResult
+    >>> from synapse_sdk.plugins.actions.upload import (
+    ...     BaseUploadAction,
+    ...     UploadContext,
+    ... )
+    >>>
+    >>> class InitStep(BaseStep[UploadContext]):
+    ...     @property
+    ...     def name(self) -> str:
+    ...         return 'initialize'
+    ...
+    ...     @property
+    ...     def progress_weight(self) -> float:
+    ...         return 0.1
+    ...
+    ...     def execute(self, context: UploadContext) -> StepResult:
+    ...         # Initialize storage, validate params
+    ...         return StepResult(success=True)
+    >>>
+    >>> class MyUploadAction(BaseUploadAction[MyParams]):
+    ...     def setup_steps(self, registry) -> None:
+    ...         registry.register(InitStep())
+"""
+
+from synapse_sdk.plugins.actions.upload.action import (
+    BaseUploadAction,
+    UploadProgressCategories,
+)
+from synapse_sdk.plugins.actions.upload.context import UploadContext
+
+__all__ = [
+    # Action
+    'BaseUploadAction',
+    'UploadProgressCategories',
+    # Context
+    'UploadContext',
+]

synapse_sdk/plugins/actions/upload/action.py

@@ -0,0 +1,165 @@
+"""Upload action base class with workflow 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.upload.context import UploadContext
+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 UploadProgressCategories:
+    """Standard progress category names for upload workflows.
+
+    .. deprecated::
+        With automatic category inference from step names (US2),
+        explicit category constants are less necessary. Step names
+        are now automatically used as progress/metrics categories.
+
+    Use these constants with set_progress() to track upload phases:
+        - INITIALIZE: Storage and path setup
+        - VALIDATE: File validation
+        - UPLOAD: File upload to storage
+        - CLEANUP: Post-upload cleanup
+
+    Example:
+        >>> self.set_progress(1, 4, self.progress.INITIALIZE)
+    """
+
+    INITIALIZE: str = 'initialize'
+    VALIDATE: str = 'validate'
+    UPLOAD: str = 'upload'
+    CLEANUP: str = 'cleanup'
+
+    def __init__(self) -> None:
+        """Initialize with deprecation warning."""
+        import warnings
+
+        warnings.warn(
+            'UploadProgressCategories is deprecated. '
+            'Step names are now automatically used as progress/metrics categories. '
+            'Consider using step.name directly or omitting the category parameter.',
+            DeprecationWarning,
+            stacklevel=2,
+        )
+
+
+class BaseUploadAction(BaseAction[P]):
+    """Base class for upload actions with workflow step support.
+
+    Provides a full step-based workflow system:
+    - Override setup_steps() to register custom steps
+    - Steps execute in order with automatic rollback on failure
+    - Progress tracked across all steps based on weights
+
+    Attributes:
+        category: Plugin category (defaults to UPLOAD).
+        progress: Standard progress category names.
+
+    Example:
+        >>> class MyUploadAction(BaseUploadAction[MyParams]):
+        ...     action_name = 'upload'
+        ...     params_model = MyParams
+        ...
+        ...     def setup_steps(self, registry: StepRegistry) -> None:
+        ...         registry.register(InitializeStep())
+        ...         registry.register(ValidateStep())
+        ...         registry.register(UploadFilesStep())
+        ...         registry.register(CleanupStep())
+        >>>
+        >>> # Steps are executed in order with automatic rollback on failure
+        >>> # Progress is tracked based on step weights
+    """
+
+    category = PluginCategory.UPLOAD
+    progress = UploadProgressCategories()
+
+    @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
+
+    def setup_steps(self, registry: StepRegistry[UploadContext]) -> None:
+        """Register workflow steps.
+
+        Override this method to register custom steps for your upload workflow.
+        Steps are executed in registration order.
+
+        Args:
+            registry: StepRegistry to register steps with.
+
+        Example:
+            >>> def setup_steps(self, registry: StepRegistry[UploadContext]) -> None:
+            ...     registry.register(InitializeStep())
+            ...     registry.register(ValidateStep())
+            ...     registry.register(UploadFilesStep())
+        """
+        pass  # Subclasses override to add steps
+
+    def create_context(self) -> UploadContext:
+        """Create upload context for the workflow.
+
+        Override to customize context creation or add additional state.
+
+        Returns:
+            UploadContext instance with params and runtime context.
+        """
+        params_dict = self.params.model_dump() if hasattr(self.params, 'model_dump') else dict(self.params)
+        return UploadContext(
+            params=params_dict,
+            runtime_ctx=self.ctx,
+        )
+
+    def execute(self) -> dict[str, Any]:
+        """Execute the upload workflow.
+
+        Creates registry, registers steps via setup_steps(), creates context,
+        and runs the orchestrator.
+
+        Returns:
+            Dict with success status and workflow results.
+
+        Raises:
+            RuntimeError: If no steps registered or a step fails.
+        """
+        # Setup
+        registry: StepRegistry[UploadContext] = StepRegistry()
+        self.setup_steps(registry)
+
+        if not registry:
+            raise RuntimeError('No steps registered. Override setup_steps() to register workflow steps.')
+
+        # Create context and orchestrator
+        context = self.create_context()
+        orchestrator: Orchestrator[UploadContext] = Orchestrator(
+            registry=registry,
+            context=context,
+            progress_callback=lambda curr, total: self.set_progress(curr, total),
+        )
+
+        # Execute workflow
+        result = orchestrator.execute()
+
+        # Add upload-specific result data
+        result['uploaded_files'] = len(context.uploaded_files)
+        result['data_units'] = len(context.data_units)
+
+        return result

synapse_sdk/plugins/actions/upload/context.py

@@ -0,0 +1,61 @@
+"""Upload 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 UploadContext(BaseStepContext):
+    """Shared context passed between upload workflow steps.
+
+    Extends BaseStepContext with upload-specific state fields.
+    Carries parameters and accumulated state as the workflow
+    progresses through steps.
+
+    Attributes:
+        params: Upload parameters (from action params).
+        storage: Storage configuration (populated by init step).
+        pathlib_cwd: Working directory path (populated by init step).
+        organized_files: Files organized for upload (populated by organize step).
+        uploaded_files: Successfully uploaded files (populated by upload step).
+        data_units: Created data units (populated by generate step).
+
+    Example:
+        >>> context = UploadContext(
+        ...     runtime_ctx=runtime_ctx,
+        ...     params={'storage': 1, 'path': '/data'},
+        ... )
+        >>> # Steps populate state as they execute
+        >>> context.organized_files.append({'path': 'file1.jpg'})
+    """
+
+    # Upload parameters
+    params: dict[str, Any] = field(default_factory=dict)
+
+    # Processing state (populated by steps)
+    storage: Any | None = None
+    pathlib_cwd: Any | None = None
+    organized_files: list[dict[str, Any]] = field(default_factory=list)
+    uploaded_files: list[dict[str, Any]] = field(default_factory=list)
+    data_units: list[dict[str, Any]] = field(default_factory=list)
+
+    @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