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

synapse_sdk/integrations/_base.py

@@ -0,0 +1,119 @@
+"""Base integration infrastructure for autolog."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from synapse_sdk.plugins.action import BaseAction
+
+
+class BaseIntegration(ABC):
+    """Base class for framework integrations.
+
+    Subclasses implement framework-specific autologging by patching
+    the framework's training methods to automatically log metrics,
+    progress, and artifacts.
+
+    Example:
+        >>> @register_integration('my_framework')
+        ... class MyFrameworkIntegration(BaseIntegration):
+        ...     name = 'my_framework'
+        ...
+        ...     def enable(self, action):
+        ...         # Patch framework methods
+        ...         ...
+        ...
+        ...     def disable(self):
+        ...         # Restore original methods
+        ...         ...
+        ...
+        ...     def is_available(self):
+        ...         try:
+        ...             import my_framework
+        ...             return True
+        ...         except ImportError:
+        ...             return False
+    """
+
+    name: str
+
+    @abstractmethod
+    def enable(self, action: BaseAction) -> None:
+        """Enable autologging for this framework.
+
+        Args:
+            action: The current action instance to log to.
+        """
+        ...
+
+    @abstractmethod
+    def disable(self) -> None:
+        """Disable autologging and restore original behavior."""
+        ...
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        """Check if the framework is installed.
+
+        Returns:
+            True if framework is available, False otherwise.
+        """
+        ...
+
+
+# Registry of available integrations
+_integrations: dict[str, type[BaseIntegration]] = {}
+
+
+def register_integration(name: str) -> Callable[[type[BaseIntegration]], type[BaseIntegration]]:
+    """Decorator to register an integration.
+
+    Args:
+        name: Integration name (e.g., 'ultralytics').
+
+    Returns:
+        Decorator function.
+
+    Example:
+        >>> @register_integration('ultralytics')
+        ... class UltralyticsIntegration(BaseIntegration):
+        ...     ...
+    """
+
+    def decorator(cls: type[BaseIntegration]) -> type[BaseIntegration]:
+        _integrations[name] = cls
+        return cls
+
+    return decorator
+
+
+def get_integration(name: str) -> BaseIntegration:
+    """Get integration instance by name.
+
+    Args:
+        name: Integration name (e.g., 'ultralytics').
+
+    Returns:
+        BaseIntegration instance.
+
+    Raises:
+        ValueError: If integration is not registered.
+    """
+    if name not in _integrations:
+        available = list(_integrations.keys()) or ['none']
+        raise ValueError(f"Unknown integration: '{name}'. Available: {available}")
+    return _integrations[name]()
+
+
+def list_integrations() -> list[str]:
+    """List all registered integration names.
+
+    Returns:
+        List of integration names.
+    """
+    return list(_integrations.keys())
+
+
+__all__ = ['BaseIntegration', 'register_integration', 'get_integration', 'list_integrations']

synapse_sdk/integrations/_context.py

@@ -0,0 +1,53 @@
+"""Context management for autolog integrations."""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from synapse_sdk.plugins.action import BaseAction
+
+
+@dataclass
+class AutologContext:
+    """Holds action context for autolog callbacks.
+
+    Attributes:
+        action: The current BaseAction instance.
+        total_epochs: Total epochs (auto-detected from train() kwargs).
+        extra: Additional framework-specific data.
+    """
+
+    action: BaseAction[Any]
+    total_epochs: int | None = None
+    extra: dict[str, Any] = field(default_factory=dict)
+
+
+# Thread/async safe context storage
+_autolog_context: ContextVar[AutologContext | None] = ContextVar(
+    'synapse_autolog_context',
+    default=None,
+)
+
+
+def get_autolog_context() -> AutologContext | None:
+    """Get current autolog context.
+
+    Returns:
+        AutologContext if autolog is active, None otherwise.
+    """
+    return _autolog_context.get()
+
+
+def set_autolog_context(ctx: AutologContext | None) -> None:
+    """Set autolog context.
+
+    Args:
+        ctx: AutologContext to set, or None to clear.
+    """
+    _autolog_context.set(ctx)
+
+
+__all__ = ['AutologContext', 'get_autolog_context', 'set_autolog_context']

synapse_sdk/integrations/ultralytics/__init__.py

@@ -0,0 +1,78 @@
+"""Ultralytics YOLO autolog integration."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from synapse_sdk.integrations._base import BaseIntegration, register_integration
+from synapse_sdk.integrations._context import AutologContext, set_autolog_context
+
+if TYPE_CHECKING:
+    from synapse_sdk.plugins.action import BaseAction
+
+
+@register_integration('ultralytics')
+class UltralyticsIntegration(BaseIntegration):
+    """Ultralytics YOLO autologging integration.
+
+    Automatically logs training progress, metrics, and artifacts
+    when using Ultralytics YOLO models.
+
+    Logged metrics:
+        - train: box_loss, cls_loss, dfl_loss (per epoch)
+        - validation: mAP50, mAP50_95 (per epoch)
+
+    Logged artifacts:
+        - validation_samples: Validation batch prediction images
+        - model_weights: best.pt weights file
+        - training_results: results.csv file
+
+    Example:
+        >>> class TrainAction(BaseTrainAction[TrainParams]):
+        ...     def execute(self):
+        ...         self.autolog('ultralytics')
+        ...         model = YOLO('yolov8n.pt')
+        ...         model.train(data='coco.yaml', epochs=100)
+    """
+
+    name = 'ultralytics'
+
+    def is_available(self) -> bool:
+        """Check if ultralytics is installed."""
+        try:
+            import ultralytics  # noqa: F401
+
+            return True
+        except ImportError:
+            return False
+
+    def enable(self, action: BaseAction) -> None:
+        """Enable autologging for Ultralytics.
+
+        Args:
+            action: The action instance to log to.
+
+        Raises:
+            ImportError: If ultralytics is not installed.
+        """
+        if not self.is_available():
+            raise ImportError('ultralytics is not installed. Install with: pip install ultralytics')
+
+        # Set context for callbacks
+        set_autolog_context(AutologContext(action=action))
+
+        # Apply patches (idempotent)
+        from synapse_sdk.integrations.ultralytics._patches import patch_yolo
+
+        patch_yolo()
+
+    def disable(self) -> None:
+        """Disable autologging for Ultralytics.
+
+        Clears the context but keeps patches in place for efficiency.
+        Callbacks check context before logging.
+        """
+        set_autolog_context(None)
+
+
+__all__ = ['UltralyticsIntegration']

synapse_sdk/integrations/ultralytics/_callbacks.py

@@ -0,0 +1,126 @@
+"""Ultralytics YOLO callback implementations for autolog."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+from synapse_sdk.integrations._context import get_autolog_context
+
+
+def on_train_epoch_end(trainer: Any) -> None:
+    """Log training metrics at end of each epoch.
+
+    Logs:
+        - Progress: epoch/total_epochs
+        - Metrics: box_loss, cls_loss, dfl_loss (category='train')
+
+    Args:
+        trainer: Ultralytics trainer instance.
+    """
+    ctx = get_autolog_context()
+    if ctx is None:
+        return
+
+    action = ctx.action
+    epoch = trainer.epoch + 1
+    total_epochs = ctx.total_epochs or getattr(trainer, 'epochs', 100)
+
+    # Update progress
+    action.set_progress(epoch, total_epochs, 'train')
+
+    # Log loss metrics
+    if hasattr(trainer, 'loss_items') and trainer.loss_items is not None:
+        loss_items = trainer.loss_items
+        if hasattr(loss_items, 'cpu'):
+            loss_items = loss_items.cpu().numpy()
+
+        action.set_metrics(
+            {
+                'epoch': epoch,
+                'box_loss': float(loss_items[0]),
+                'cls_loss': float(loss_items[1]),
+                'dfl_loss': float(loss_items[2]),
+            },
+            category='train',
+        )
+
+
+def on_fit_epoch_end(trainer: Any) -> None:
+    """Log validation metrics after validation pass.
+
+    Logs:
+        - Metrics: mAP50, mAP50_95 (category='validation')
+        - Files: validation batch prediction images
+
+    Args:
+        trainer: Ultralytics trainer instance.
+    """
+    ctx = get_autolog_context()
+    if ctx is None:
+        return
+
+    action = ctx.action
+    epoch = trainer.epoch + 1
+    metrics = trainer.metrics
+
+    if metrics:
+        action.set_metrics(
+            {
+                'epoch': epoch,
+                'mAP50': metrics.get('metrics/mAP50(B)', 0),
+                'mAP50_95': metrics.get('metrics/mAP50-95(B)', 0),
+            },
+            category='validation',
+        )
+
+        # Log validation sample images
+        save_dir = Path(trainer.save_dir)
+        for i in range(3):
+            img_path = save_dir / f'val_batch{i}_pred.jpg'
+            if img_path.exists():
+                action.log(
+                    'validation_samples',
+                    {'group': epoch, 'index': i},
+                    file=str(img_path),
+                )
+
+    # Ray Tune integration
+    is_tune = action.ctx.env.get_bool('IS_TUNE', default=False)
+    if is_tune and metrics:
+        try:
+            from ray import tune
+
+            tune.report(**metrics)
+        except ImportError:
+            pass
+
+
+def on_train_end(trainer: Any) -> None:
+    """Log final artifacts when training completes.
+
+    Logs:
+        - Files: best.pt weights, results.csv
+
+    Args:
+        trainer: Ultralytics trainer instance.
+    """
+    ctx = get_autolog_context()
+    if ctx is None:
+        return
+
+    action = ctx.action
+    save_dir = Path(trainer.save_dir)
+
+    # Log final model weights
+    best_pt = save_dir / 'weights' / 'best.pt'
+    if best_pt.exists():
+        action.log('model_weights', {'type': 'best'}, file=str(best_pt))
+
+    # Log training results CSV
+    results_csv = save_dir / 'results.csv'
+    if results_csv.exists():
+        action.log('training_results', {}, file=str(results_csv))
+
+
+__all__ = ['on_train_epoch_end', 'on_fit_epoch_end', 'on_train_end']

synapse_sdk/integrations/ultralytics/_patches.py

@@ -0,0 +1,124 @@
+"""Monkey patches for Ultralytics YOLO autolog."""
+
+from __future__ import annotations
+
+import functools
+from typing import Any
+
+from synapse_sdk.integrations._context import get_autolog_context
+from synapse_sdk.integrations.ultralytics._callbacks import (
+    on_fit_epoch_end,
+    on_train_end,
+    on_train_epoch_end,
+)
+
+# Store original methods to avoid re-patching
+_original_yolo_init: Any = None
+_original_yolo_train: Any = None
+_patched = False
+
+
+def patch_yolo() -> None:
+    """Apply all YOLO patches.
+
+    Patches:
+        - YOLO.__init__: Auto-attach callbacks when model is created
+        - YOLO.train: Capture epochs from kwargs
+
+    This function is idempotent - calling it multiple times has no effect.
+    """
+    global _patched
+    if _patched:
+        return
+
+    _patch_yolo_init()
+    _patch_yolo_train()
+    _patched = True
+
+
+def _patch_yolo_init() -> None:
+    """Patch YOLO.__init__ to auto-attach callbacks."""
+    global _original_yolo_init
+
+    from ultralytics import YOLO
+
+    if _original_yolo_init is not None:
+        return  # Already patched
+
+    _original_yolo_init = YOLO.__init__
+
+    @functools.wraps(_original_yolo_init)
+    def patched_init(self: Any, *args: Any, **kwargs: Any) -> None:
+        result = _original_yolo_init(self, *args, **kwargs)
+
+        # Check if autolog is active
+        ctx = get_autolog_context()
+        if ctx is not None:
+            # Attach Synapse callbacks
+            self.add_callback('on_train_epoch_end', on_train_epoch_end)
+            self.add_callback('on_fit_epoch_end', on_fit_epoch_end)
+            self.add_callback('on_train_end', on_train_end)
+
+        return result
+
+    YOLO.__init__ = patched_init
+
+
+def _patch_yolo_train() -> None:
+    """Patch YOLO.train to capture epochs from kwargs."""
+    global _original_yolo_train
+
+    from ultralytics import YOLO
+
+    if _original_yolo_train is not None:
+        return  # Already patched
+
+    _original_yolo_train = YOLO.train
+
+    @functools.wraps(_original_yolo_train)
+    def patched_train(self: Any, *args: Any, **kwargs: Any) -> Any:
+        ctx = get_autolog_context()
+        if ctx is not None:
+            # Auto-detect epochs from train() kwargs
+            epochs = kwargs.get('epochs')
+            if epochs is None:
+                # Check data config for epochs
+                data = kwargs.get('data')
+                if isinstance(data, dict):
+                    epochs = data.get('epochs')
+
+            # Fall back to ultralytics default
+            if epochs is None:
+                epochs = 100
+
+            ctx.total_epochs = epochs
+
+        return _original_yolo_train(self, *args, **kwargs)
+
+    YOLO.train = patched_train
+
+
+def unpatch_yolo() -> None:
+    """Restore original YOLO methods.
+
+    Used for testing or cleanup.
+    """
+    global _original_yolo_init, _original_yolo_train, _patched
+
+    if not _patched:
+        return
+
+    from ultralytics import YOLO
+
+    if _original_yolo_init is not None:
+        YOLO.__init__ = _original_yolo_init
+        _original_yolo_init = None
+
+    if _original_yolo_train is not None:
+        YOLO.train = _original_yolo_train
+        _original_yolo_train = None
+
+    _patched = False
+
+
+__all__ = ['patch_yolo', 'unpatch_yolo']