"""Auto-generated stub for module: trainer."""
from typing import Any, Dict, List, Optional

# Constants
logger: Any

# Classes
class TrainResult:
    # Output of a completed training run.
    #
    #     Captures the best metric achieved, the number of epochs that ran,
    #     saved checkpoint paths, and the evaluation payload that was (or will
    #     be) sent to ActionTracker.

    ...
class Trainer:
    # Abstract base for all BYOM training jobs.
    #
    #     Modelled after AWS SageMaker's ``Estimator``: call :meth:`fit` and
    #     the full lifecycle — data loading, model initialisation, training,
    #     checkpointing, evaluation — runs with ActionTracker status updates
    #     at every phase.
    #
    #     Subclass one of the concrete variants instead of this class directly:
    #
    #     * ``LoopTrainer`` — you own the PyTorch for-loop
    #       (classification, R2Plus1D).
    #     * ``DelegatedTrainer`` — the framework owns the loop
    #       (Ultralytics YOLO, D-FINE solver).
    #
    #     Each subclass only implements a handful of hooks; all boilerplate
    #     (ActionTracker wiring, error handling, checkpoint upload) is handled
    #     here.
    #
    #     Example::
    #
    #         class MyTrainer(LoopTrainer):
    #             def setup_data(self) -> None: ...
    #             def setup_model(self) -> None: ...
    #
    #             ...
    #
    #
    #         trainer = MyTrainer(action_id, model_info, config)
    #         result = trainer.fit()

    def __init__(self: Any, action_id: str | None, model_info: Any, config: Any) -> None:
        """
        Initialize the trainer.
        
                Args:
                    action_id: Platform action ID.  Pass ``None`` for local /
                        offline runs (a ``LocalActionTracker`` or no-op stub
                        will be created automatically).
                    model_info: Model identity and capabilities.
                    config: Training hyperparameters.
                    checkpoint_dir: Directory where checkpoints are written.
                    callbacks: Optional list of :class:`TrainerCallback`
                        instances.  When ``None`` the default set
                        (``ActionTrackerCallback``) is registered automatically.
        """
        ...

    def fit(self: Any) -> Any:
        """
        Run the full training lifecycle.
        
                This is the single entry-point, analogous to
                ``sagemaker.estimator.Estimator.fit()``.  The phases are:
        
                1. **Acknowledge** — ``MDL_TRN_ACK``
                2. **Data loading** — ``MDL_TRN_DTL``
                3. **Model init** — ``MDL_TRN_MDL``
                4. **Training loop** — ``MDL_TRN_CMPL``
                5. **Checkpoint upload** — ``MDL_TRN_BMS``
                6. **Evaluation** — ``MDL_TRN_EVL``
        
                Each phase is wrapped in ActionTracker status updates.  On
                failure the tracker receives an ``ERROR`` status and a
                :class:`TrainerError` is raised.
        
                Returns:
                    A :class:`TrainResult` summarising the completed run.
        
                Raises:
                    TrainerError: If any lifecycle phase fails.
        """
        ...

    def get_label_map(self: Any) -> dict[str, str]:
        """
        Return an index-to-label mapping for the ActionTracker.
        
                Returns:
                    Dict mapping **string** class indices to human-readable
                    class names, e.g. ``{"0": "cat", "1": "dog"}``.
        """
        ...

    def get_model_load_strategies(self: Any) -> Any | None:
        """
        Optionally provide shared model-loading strategies.
        
                Returns:
                    ``ModelLoadStrategies`` to enable automatic model-loading in
                    :meth:`_do_setup_model`, or ``None`` to use subclass
                    :meth:`setup_model` behavior.
        """
        ...

    def load_model_with_strategies(self: Any, strategies: Any) -> Any:
        """
        Load model using ActionTracker checkpoint/predefined strategy.
        
                Args:
                    strategies: Strategy callbacks for checkpoint/predefined/scratch branches.
        
                Returns:
                    Loaded model artifact selected from the first available branch.
        """
        ...

    def log_epoch(self: Any, epoch: int, metrics: list[dict[str, Any]]) -> None:
        """
        Forward per-epoch metrics to the ActionTracker.
        
                Called by ``LoopTrainer`` (or your own loop) at the end of each
                epoch.  Override to add side-channels such as TensorBoard or
                Weights & Biases logging.
        
                Args:
                    epoch: Zero-based epoch number.
                    metrics: List of metric dicts, each containing at minimum
                        ``splitType``, ``metricName``, and ``metricValue``.
        """
        ...

    def notify(self: Any, hook: str, *args: Any, **kwargs: Any) -> None:
        """
        Fire a named callback hook on every registered callback.
        
                Exceptions from individual callbacks are logged but do **not**
                interrupt the training loop.
        
                Args:
                    hook: Method name to invoke (e.g. ``"on_epoch_end"``).
                    *args: Positional arguments forwarded to the hook.
                    **kwargs: Keyword arguments forwarded to the hook.
        """
        ...

    def run_evaluation(self: Any) -> list[dict[str, Any]]:
        """
        Run post-training evaluation and return the results payload.
        
                The returned list is forwarded to
                ``action_tracker.save_evaluation_results()``.
        
                Returns:
                    Metric dicts in ActionTracker format::
        
                        [{"splitType": "val", "metricName": "acc@1", "metricValue": 0.95, "category": "all"}, ...]
        """
        ...

    def save_model(self: Any) -> None:
        """
        Save the best model and upload checkpoints to the platform.
        
                The default implementation uploads ``model_best.pth.tar`` and
                ``model_best.pt`` from :attr:`checkpoint_dir`.  If the ``.pt``
                file does not exist but ``self.model`` is set, it is serialised
                with ``torch.save`` first.
        
                Override for frameworks that use a different checkpoint layout
                (e.g. Ultralytics writes to ``runs/detect/train/weights/``).
        """
        ...

    def setup_data(self: Any) -> None:
        """
        Load datasets and create data loaders.
        
                Implementations should store loaders as instance attributes
                (e.g. ``self.train_loader``) so that :meth:`train_loop` and
                :meth:`run_evaluation` can access them.
        """
        ...

    def setup_model(self: Any) -> None:
        """
        Construct or load the model and assign it to ``self.model``.
        
                The compute device is resolved automatically *after* this hook
                returns.  Moving the model to the device is the responsibility
                of the subclass (in this hook or in :meth:`train_loop`), because
                different frameworks handle device placement differently.
        
                For checkpoint/predefined selection flows, subclasses can use
                :meth:`load_model_with_strategies`.
        """
        ...

    def train_loop(self: Any) -> None:
        """
        Execute the main training loop.
        
                ``LoopTrainer`` provides a standard PyTorch epoch-loop.
                ``DelegatedTrainer`` delegates to the framework's own API.
        
                **Contract:** implementations *must* populate ``self._result``
                (a :class:`TrainResult`) before returning.
        """
        ...

class TrainerError(Exception):
    # Raised when a training lifecycle phase fails.
    #
    #     Attributes:
    #         phase: The ActionTracker status code of the failed phase
    #             (e.g. ``"MDL_TRN_DTL"``).

    def __init__(self: Any, phase: str, message: str) -> None:
        """
        Initialize with the failing phase and error message.
        
                Args:
                    phase: ActionTracker status code (e.g. ``MDL_TRN_DTL``).
                    message: Human-readable description of the failure.
        """
        ...

