---
id: train-actions
title: Train Actions
sidebar_position: 1
---

# Train Actions

Train actions provide a base class for building model training workflows. Extend `BaseTrainAction` to create custom training plugins with built-in support for datasets, checkpoints, and model uploads.

## Overview

`BaseTrainAction` simplifies training workflows by handling common operations: fetching datasets from the backend, managing checkpoints, and uploading trained models.

It supports two execution modes—simple and step-based—to match your workflow complexity.

### At a Glance

| Feature | Description |
|---------|-------------|
| Dataset loading | Fetch datasets via `get_dataset()` |
| Checkpoint support | Resume training with `get_checkpoint()` |
| Model upload | Upload trained models with `create_model()` |
| Progress tracking | Built-in categories for DATASET, TRAIN, MODEL_UPLOAD |
| Execution modes | Simple `execute()` or step-based workflows |

## BaseTrainAction

### How it works

```mermaid
flowchart LR
    subgraph Pipeline
        D["Dataset<br/>Loading"] --> C["Checkpoint<br/>Loading"] --> T["Training<br/>Process"] --> U["Upload<br/>Model"]
    end
    D -.-> GD["get_dataset()"]
    C -.-> GC["get_checkpoint()"]
    T -.-> EX["execute()"]
    U -.-> CM["create_model()"]

    style Pipeline fill:#e0e7ff,stroke:#6366f1,color:#1f2937
    style D fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style C fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style T fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style U fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style GD fill:#8b5cf6,stroke:#7c3aed,color:#ffffff
    style GC fill:#8b5cf6,stroke:#7c3aed,color:#ffffff
    style EX fill:#8b5cf6,stroke:#7c3aed,color:#ffffff
    style CM fill:#8b5cf6,stroke:#7c3aed,color:#ffffff
```

1. **Dataset Loading**: Fetch training data from the backend using `get_dataset()`
2. **Checkpoint Loading**: Optionally load a pre-trained model via `get_checkpoint()`
3. **Training Process**: Run your training logic in `execute()` or step-based workflow
4. **Model Upload**: Upload the trained model using `create_model()`

### Class Signature

```python filename="synapse_sdk/plugins/actions/train/action.py"
class BaseTrainAction(BaseAction[P]):
    """Base class for training actions."""

    category = PluginCategory.NEURAL_NET
    progress = TrainProgressCategories()
```

The generic parameter `P` represents your custom params model extending `BaseTrainParams`.

### Helper Methods

| Method | Description | Returns |
|--------|-------------|---------|
| `get_dataset()` | Fetch dataset from backend by `params.dataset_id` | `dict[str, Any]` |
| `create_model(path, **kwargs)` | Upload trained model to backend | `dict[str, Any]` |
| `get_model(model_id)` | Retrieve existing model metadata | `dict[str, Any]` |
| `get_checkpoint()` | Get checkpoint model for transfer learning | `dict[str, Any] \| None` |

#### get_dataset

Fetches the training dataset from the backend. Override this method for custom data sources.

**Default implementation:**

```python filename="synapse_sdk/plugins/actions/train/action.py"
def get_dataset(self) -> dict[str, Any]:
    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)
```

**Custom override example (S3):**

```python filename="plugins/my_train/action.py"
def get_dataset(self) -> dict[str, Any]:
    return download_from_s3(self.params.s3_path)
```

#### get_checkpoint

Resolves checkpoints in two ways:
1. From `ctx.checkpoint` if already set (remote mode)
2. From `params.checkpoint` model ID (fetches and extracts)

```python filename="plugins/my_train/action.py"
checkpoint = self.get_checkpoint()
if checkpoint:
    model_path = checkpoint['path']      # Path object to extracted model
    category = checkpoint['category']    # 'base' or fine-tuned category
```

> **Good to know**: The checkpoint dict contains `category`, `path` (as `pathlib.Path`), `id`, and `name` fields. Use `checkpoint['path'] / 'model.pt'` for path operations.

### Progress Categories

Use `TrainProgressCategories` to report progress at each training phase:

| Category | Constant | Purpose |
|----------|----------|---------|
| Dataset | `self.progress.DATASET` | Data loading and preprocessing |
| Training | `self.progress.TRAIN` | Model training iterations |
| Model Upload | `self.progress.MODEL_UPLOAD` | Final model upload |

```python filename="plugins/my_train/action.py"
# Report dataset loading progress
self.set_progress(1, 10, self.progress.DATASET)

# Report training progress
for epoch in range(100):
    train_epoch()
    self.set_progress(epoch + 1, 100, self.progress.TRAIN)
```

## Execution Modes

Choose between two execution modes based on your workflow complexity.

### Simple Mode

Override `execute()` directly for straightforward training workflows.

```python filename="plugins/my_train/action.py"
class MyTrainAction(BaseTrainAction[MyParams]):
    action_name = 'my_train'
    params_model = MyParams
    result_model = TrainResult

    def execute(self) -> TrainResult:
        # 1. Load dataset
        dataset = self.get_dataset()
        self.set_progress(1, 1, self.progress.DATASET)

        # 2. Load checkpoint (optional)
        checkpoint = self.get_checkpoint()

        # 3. Train model
        model = train_model(dataset, checkpoint)
        self.set_progress(100, 100, self.progress.TRAIN)

        # 4. Upload model
        result = self.create_model(model.save_path)

        return TrainResult(
            weights_path=model.save_path,
            final_epoch=100,
            train_metrics={'loss': 0.01},
        )
```

### Step-Based Mode

Override `setup_steps()` for complex workflows requiring multiple reusable steps. Steps execute sequentially with automatic progress tracking and rollback on failure.

```mermaid
flowchart LR
    subgraph Steps
        DS["DatasetStep<br/>weight=0.1"] --> TS["TrainingStep<br/>weight=0.8"] --> US["UploadStep<br/>weight=0.1"]
    end
    DS -.-> D["ctx.dataset"]
    TS -.-> M["ctx.model_path"]
    US -.-> U["ctx.model"]

    style Steps fill:#e0e7ff,stroke:#6366f1,color:#1f2937
    style DS fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style TS fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style US fill:#3b82f6,stroke:#1e40af,color:#ffffff
    style D fill:#06b6d4,stroke:#0891b2,color:#ffffff
    style M fill:#06b6d4,stroke:#0891b2,color:#ffffff
    style U fill:#06b6d4,stroke:#0891b2,color:#ffffff
```

#### Define Steps

```python filename="plugins/my_train/steps.py"
from synapse_sdk.plugins.steps import BaseStep, StepResult
from synapse_sdk.plugins.actions.train import TrainContext

class DatasetStep(BaseStep[TrainContext]):
    @property
    def name(self) -> str:
        return 'dataset'

    @property
    def progress_weight(self) -> float:
        return 0.1  # 10% of total progress

    def execute(self, ctx: TrainContext) -> StepResult:
        dataset = ctx.client.get_data_collection(ctx.params['dataset_id'])
        ctx.dataset = dataset
        return StepResult(success=True, data={'count': len(dataset)})


class TrainingStep(BaseStep[TrainContext]):
    @property
    def name(self) -> str:
        return 'training'

    @property
    def progress_weight(self) -> float:
        return 0.8  # 80% of total progress

    def execute(self, ctx: TrainContext) -> StepResult:
        model = train(ctx.dataset, epochs=ctx.params['epochs'])
        ctx.model_path = model.save()
        return StepResult(success=True, data={'model_path': ctx.model_path})

    def rollback(self, ctx: TrainContext, result: StepResult) -> None:
        # Clean up on failure
        if ctx.model_path:
            Path(ctx.model_path).unlink(missing_ok=True)
```

> **Good to know**: `ctx.client` raises `RuntimeError` if the backend client is not configured in `RuntimeContext`. Ensure your action is executed with a valid client, or handle the exception in your step.

#### Register Steps

```python filename="plugins/my_train/action.py"
from synapse_sdk.plugins.steps import StepRegistry
from synapse_sdk.plugins.actions.train import BaseTrainAction, TrainContext

class MyTrainAction(BaseTrainAction[MyParams]):
    action_name = 'my_train'

    def setup_steps(self, registry: StepRegistry[TrainContext]) -> None:
        registry.register(DatasetStep())
        registry.register(TrainingStep())
        registry.register(UploadStep())
```

> **Good to know**: When `setup_steps()` registers steps, the step-based workflow takes precedence and `execute()` is not called.

## TrainContext

`TrainContext` carries state between steps in step-based workflows. It extends `BaseStepContext` with training-specific fields.

### Context Fields

| Field | Type | Description |
|-------|------|-------------|
| `params` | `dict[str, Any]` | Training parameters from action |
| `dataset` | `Any` | Loaded dataset (populated by step) |
| `model_path` | `str \| None` | Path to trained model (populated by step) |
| `model` | `dict[str, Any] \| None` | Uploaded model metadata (populated by step) |
| `client` | `BackendClient` | Backend client for API calls (raises `RuntimeError` if not set) |
| `runtime_ctx` | `RuntimeContext` | Parent runtime context |

**Inherited from `BaseStepContext`:**

| Field | Type | Description |
|-------|------|-------------|
| `step_results` | `list[StepResult]` | Results from each executed step |
| `errors` | `list[str]` | Accumulated error messages |
| `current_step` | `str \| None` | Name of the currently executing step |

### Inherited Methods

From `BaseStepContext`:

```python filename="plugins/my_train/steps.py"
# Log an event (optional file parameter for associated file path)
ctx.log('training_started', {'epochs': 100})
ctx.log('model_saved', {'path': '/models/best.pt'}, file='/models/best.pt')

# Set progress (category defaults to current step name)
ctx.set_progress(50, 100)

# Set metrics
ctx.set_metrics({'loss': 0.01, 'accuracy': 0.95}, 'train')
```

## Example: Complete Training Plugin

A complete example showing params, action, and result models working together.

```python filename="plugins/yolo_train/params.py"
from pydantic import Field
from synapse_sdk.plugins.actions.train import BaseTrainParams

class YOLOTrainParams(BaseTrainParams):
    """Parameters for YOLO training."""

    dataset_id: int = Field(description='Dataset ID to train on')
    epochs: int = Field(default=100, description='Number of training epochs')
    batch_size: int = Field(default=16, description='Batch size')
    imgsz: int = Field(default=640, description='Input image size')
    # checkpoint field inherited from BaseTrainParams
```

```python filename="plugins/yolo_train/action.py"
from ultralytics import YOLO
from synapse_sdk.plugins.actions.train import BaseTrainAction
from synapse_sdk.plugins import TrainResult  # or: from synapse_sdk.plugins.schemas.results import TrainResult

class YOLOTrainAction(BaseTrainAction[YOLOTrainParams]):
    """YOLO model training action."""

    action_name = 'yolo_train'
    params_model = YOLOTrainParams
    result_model = TrainResult

    def execute(self) -> TrainResult:
        # Enable autologging for Ultralytics
        self.autolog('ultralytics')

        # Load dataset
        dataset = self.get_dataset()
        self.set_progress(1, 1, self.progress.DATASET)

        # Initialize model (from checkpoint or pretrained)
        checkpoint = self.get_checkpoint()
        if checkpoint:
            model = YOLO(checkpoint['path'] / 'best.pt')
        else:
            model = YOLO('yolov8n.pt')

        # Train
        results = model.train(
            data=dataset['config_path'],
            epochs=self.params.epochs,
            batch=self.params.batch_size,
            imgsz=self.params.imgsz,
        )

        # Upload model
        weights_path = str(results.save_dir / 'weights' / 'best.pt')
        self.create_model(weights_path, name='yolo-trained')

        return TrainResult(
            weights_path=weights_path,
            final_epoch=self.params.epochs,
            train_metrics=results.results_dict,
        )
```

> **Good to know**: The `autolog('ultralytics')` call automatically logs training metrics, validation samples, and model artifacts.

## API Reference

### BaseTrainParams

Base parameters for training actions. Extend this class to add plugin-specific parameters.

| Field | Type | Required | Default | Description |
|-------|------|----------|---------|-------------|
| `checkpoint` | `int \| None` | No | `None` | Model ID to use as starting checkpoint |

### TrainProgressCategories

Standard progress category names for training workflows.

| Category | Value | Purpose |
|----------|-------|---------|
| `DATASET` | `'dataset'` | Data loading and preprocessing |
| `TRAIN` | `'train'` | Model training iterations |
| `MODEL_UPLOAD` | `'model_upload'` | Final model upload to backend |

### TrainResult

Combined result type for training actions.

| Field | Type | Required | Description |
|-------|------|----------|-------------|
| `weights_path` | `str` | Yes | Path to trained model weights |
| `final_epoch` | `int` | Yes | Last completed epoch |
| `best_epoch` | `int \| None` | No | Best epoch by validation metric |
| `train_metrics` | `dict[str, float]` | No | Final training metrics |
| `val_metrics` | `dict[str, float]` | No | Final validation metrics |

## Related

- [BaseAction](/plugins/defining-actions) - Parent class for all actions
- [Step-Based Workflows](/plugins/steps-workflow) - Step orchestration system
