Function do_sync_training

examples/tinker/agl_tinker/train.py:111–224

Implements fully synchronous on-policy training. See `tinker_cookbook.rl.train.do_sync_training` for the original flow. The Agent-lightning adaptation diverges in a few places: * A LiteLLM proxy is restarted every batch so refreshed player checkpoints are immediately visible to r

(
    *,
    start_batch: int,
    end_batch: int,
    num_batches: int,
    cfg: Config,
    training_client: tinker.TrainingClient,
    service_client: tinker.ServiceClient,
    evaluators: list[AGLTestSetEvaluator[Any]],
    dataset: AGLDataset[Any],
    ml_logger: ml_log.Logger,
    tokenizer: Tokenizer,
    store: LightningStore,
    adapter: TraceToTripletBase,
    llm_proxy: LLMProxy,
)

@scope
async def do_sync_training(
    *,
    start_batch: int,
    end_batch: int,
    num_batches: int,
    cfg: Config,
    training_client: tinker.TrainingClient,
    service_client: tinker.ServiceClient,
    evaluators: list[AGLTestSetEvaluator[Any]],
    dataset: AGLDataset[Any],
    ml_logger: ml_log.Logger,
    tokenizer: Tokenizer,
    store: LightningStore,
    adapter: TraceToTripletBase,
    llm_proxy: LLMProxy,
):
    """Implements fully synchronous on-policy training.

    See `tinker_cookbook.rl.train.do_sync_training` for the original flow. The
    Agent-lightning adaptation diverges in a few places:

    * A LiteLLM proxy is restarted every batch so refreshed player checkpoints
      are immediately visible to rollout workers.
    * Trajectories are gathered via `do_group_of_group_rollouts`, which in turn
      dequeues tasks from the Agent-lightning store and rebuilds transitions from
      trace triplets.
    * Evaluation hooks call `AGLTestSetEvaluator` so validation samples reuse the
      same CrewAI-based agent rather than invoking a raw token completer.
    """

    # Initial sampling client
    logger.info(f"Creating sampling client with training client {training_client} and start batch {start_batch}")
    sampling_client, _ = await save_checkpoint_and_get_sampling_client(
        training_client, start_batch, cfg.log_path, cfg.save_every
    )
    logger.info(f"Creating renderer with name {cfg.renderer_name}")
    renderer = get_renderer(cfg.renderer_name, tokenizer)

    tinker_llm = TinkerLLM(
        model_name=cfg.model_name,
        sampling_client=sampling_client,
        renderer=renderer,
        tokenizer=tokenizer,
        max_tokens=cfg.max_tokens,
        temperature=cfg.train_temperature,
        top_k=cfg.top_k,
        top_p=cfg.top_p,
    ).rewrite_litellm_custom_providers()

    logger.info(f"Starting training from batch {start_batch} to {end_batch}")
    for i_batch in range(start_batch, end_batch):
        metrics = {
            "progress/batch": i_batch,
            "optim/lr": cfg.learning_rate,
            "progress/done_frac": (i_batch + 1) / num_batches,
        }
        logger.info(f"[Batch {i_batch}] Starting training step. Learning rate: {cfg.learning_rate}")
        t_start = time.time()