nat.finetuning.interfaces.trainer_adapter

Classes

TrainerAdapter

Adapter to send Trajectories to remote training cluster for weights updates.

Module Contents

class TrainerAdapter(

adapter_config: nat.data_models.finetuning.TrainerAdapterConfig,

)

Bases: abc.ABC

Adapter to send Trajectories to remote training cluster for weights updates.

adapter_config

run_config: nat.data_models.finetuning.FinetuneConfig = None

async initialize(

run_config: nat.data_models.finetuning.FinetuneConfig,

) → None

Asynchronously initialize any resources needed for the trainer adapter.

abstractmethod is_healthy() → bool

Async:

Check the health of the remote training backend.

Returns:

bool: True if the backend is healthy, False otherwise.

abstractmethod submit(

trajectories: nat.data_models.finetuning.TrajectoryCollection,

) → nat.data_models.finetuning.TrainingJobRef

Async:

Submit trajectories to remote training backend.

Args:

trajectories (list[Trajectory]): The list of trajectories to submit.

Returns:

TrainingJobRef: Reference to the submitted training job.

abstractmethod status(

ref: nat.data_models.finetuning.TrainingJobRef,

) → nat.data_models.finetuning.TrainingJobStatus

Async:

Get the status of a submitted training job.

Args:

ref (TrainingJobRef): Reference to the training job.

Returns:

TrainingJobStatus: The current status of the training job.

abstractmethod wait_until_complete(

ref: nat.data_models.finetuning.TrainingJobRef,

poll_interval: float = 10.0,

) → nat.data_models.finetuning.TrainingJobStatus

Async:

Wait until the training job is complete.

Args:

ref (TrainingJobRef): Reference to the training job. poll_interval (float): Time in seconds between status checks.

Returns:

TrainingJobStatus: The final status of the training job.

abstractmethod log_progress(

ref: nat.data_models.finetuning.TrainingJobRef,

metrics: dict[str, Any],

output_dir: str | None = None,

) → None

Log training adapter progress.

Args:

ref: Training job reference metrics: Dictionary of metrics to log output_dir: Optional output directory override