Stage 5: Add TinkerTrainingAdapter for forward_backward/optim_step

[tyler-griggs]

Summary

Implements Tinker-compatible training operations (forward_backward and optim_step) through SkyRL's WorkerDispatch.

Components:

skyrl-train (core training logic):

• TinkerTrainingAdapter: Converts Tinker format → WorkerDispatch calls
  • Supports loss functions: cross_entropy, importance_sampling, ppo
  • Maps Tinker Datum to TrainingInputBatch with left-padding
  • Async wrappers around WorkerDispatch methods
• 16 unit tests covering all loss functions and edge cases

skyrl-tx (API integration):

• SkyRLTrainingClient: Thin wrapper for pydantic conversion + database storage
  • call_forward_backward_and_store(): Background task for async API
  • call_optim_step_and_store(): Background task for optimizer step
  • attach_skyrl_training(): Easy integration with FastAPI app

Architecture:

Tinker API (/api/v1/forward_backward, /api/v1/optim_step)
    ↓
SkyRLTrainingClient (skyrl-tx)
    • Tinker pydantic ↔ plain Python dicts
    • Database storage (FutureDB)
    ↓
TinkerTrainingAdapter (skyrl-train)
    • Datum list → TrainingInputBatch
    • Loss function mapping
    ↓
WorkerDispatch.forward_backward() / optim_step()

Loss Function Support:

• cross_entropy: Supervised learning (requires: target_tokens, weights)
• importance_sampling: REINFORCE with IS (requires: target_tokens, logprobs, advantages)
• ppo: PPO with clipping (mapped to SkyRL's "regular" loss)

Tests:
✅ 16/16 CPU unit tests passing

Next Steps (Stage 7-9):

• Stage 7: ServiceClient factory for creating training/sampling clients
• Stage 8: Checkpoint management
• Stage 9: End-to-end integration test with tinker-cookbook scripts

[gemini-code-assist]

Code Review

This pull request introduces a TinkerTrainingAdapter to bridge Tinker-style training operations with SkyRL's WorkerDispatch, along with the SkyRLTrainingClient for API integration and a comprehensive set of unit tests. A critical Denial of Service vulnerability was identified where the TinkerTrainingAdapter defines async methods that perform blocking synchronous calls to Ray, which will freeze the FastAPI event loop and make the server unresponsive during training operations. This should be remediated by offloading the blocking calls to a thread pool. Additionally, the review identified a critical issue regarding potential crashes in background tasks due to missing None checks when retrieving database records, a high-severity bug where important metadata is overwritten, and a medium-severity issue related to inconsistent data handling within the adapter.

[gemini-code-assist]

session.get(FutureDB, request_id) can return None if no record is found for the given request_id. Accessing attributes on future (e.g., future.result_data on line 90) would then raise an AttributeError, crashing the background task. You should handle the case where future is None.

[gemini-code-assist]

session.get(FutureDB, request_id) can return None if no record is found for the given request_id. Accessing attributes on future (e.g., future.result_data on line 119) would then raise an AttributeError, crashing the background task. You should handle the case where future is None.

[gemini-code-assist]

The forward_backward method is defined as asynchronous, but it performs a synchronous, blocking operation by calling self.worker_dispatch.forward_backward. In the provided context, WorkerDispatch.forward_backward internally calls ray.get(), which blocks the execution thread until the remote task completes. Since this is called within the FastAPI event loop (via the background task in SkyRLTrainingClient), it will block the entire event loop, making the API server unresponsive to other concurrent requests during the training operation. This is a significant Denial of Service (DoS) vulnerability.

Suggested change

	metrics = self.worker_dispatch.forward_backward(self.model_name, training_batch)
	metrics = await asyncio.to_thread(self.worker_dispatch.forward_backward, self.model_name, training_batch)

[gemini-code-assist]

Similar to the forward_backward method, optim_step is defined as async but calls the synchronous, blocking self.worker_dispatch.optim_step method. This will block the FastAPI event loop, leading to a Denial of Service. The blocking call should be offloaded to a separate thread.

Suggested change

	grad_norm = self.worker_dispatch.optim_step(self.model_name)
	grad_norm = await asyncio.to_thread(self.worker_dispatch.optim_step, self.model_name)

[gemini-code-assist]

The metadata from _convert_data_to_batch (which contains num_actions) is being overwritten here. You should update the existing metadata dictionary instead of replacing it to preserve all metadata.

Suggested change

	training_batch.metadata = {
	"loss_fn": self.LOSS_FN_MAP[loss_fn],
	"loss_fn_config": loss_fn_config or {},
	}
	training_batch.metadata.update({
	"loss_fn": self.LOSS_FN_MAP[loss_fn],
	"loss_fn_config": loss_fn_config or {},
	})

[gemini-code-assist]

The current implementation for calculating max_seq_len is not fully robust and duplicates logic from extract_tokens_from_model_input.

1. It uses direct key access d["model_input"] which can raise a KeyError. It's safer to use .get().
2. It only considers "tokens" and ignores the "chunks" format, which extract_tokens_from_model_input is designed to handle.

Using the extract_tokens_from_model_input static method here and on line 224 will improve consistency and robustness.

Suggested change

	max_seq_len = max(
	len(d["model_input"].get("tokens", []))
	for d in data
	)
	max_seq_len = max(
	len(TinkerTrainingAdapter.extract_tokens_from_model_input(d.get("model_input", {})))
	for d in data
	)