feat(tinker): add TinkerNativeBackend (#532)

* feat: add TinkerNativeBackend for native training

Separate native Tinker training/inference from LocalBackend to keep the API
clear while enabling explicit loss/checkpoint behavior and config.

* fix: address pre-commit type and format issues

Align tinker native types with OpenAI tooling and update tests to avoid
invalid type expressions under pyright.

* feat: add safer state merge and policy tracking

Use merge_state for backend persistence to avoid clobbering model state, and
fail fast on trajectories without Choice objects to prevent no-op training.
Expose policy version fields on trajectories for off-policy tracking.

* feat(pipeline): add PipelineTrainer for async 3-stage training

Add a new PipelineTrainer module that implements an asynchronous
3-stage pipeline (rollout, training, eval) for efficient RL training:

- PipelineTrainer: Main trainer class with configurable workers,
  batch sizes, and off-policy limits
- StatusReporter: Live progress reporting with tqdm and periodic
  logging
- PipelineState: Shared state dataclass for stage coordination
- Type definitions for RolloutFn, SingleRolloutFn, EvalFn

Key features:
- Async rollout workers with policy version tracking
- Stale sample detection and automatic discard
- Zero-variance group handling with collapse detection
- Graceful signal handling (SIGINT/SIGTERM)
- State persistence for training resumption
- Eval scheduling with configurable intervals

Also includes:
- yes_no_maybe_pipeline.py: Simple example showing basic usage
- binary_prefix_tool_pipeline.py: Complex example with tool calls

Updates to tinker_native backend:
- Add debug logging via ART_TINKER_TRAIN_LOG/ART_TINKER_SAMPLE_LOG
- Add fallback for create_conversation_prefix_with_tools
- Fix tool_call id handling in OpenAI server responses

* fix: resolve type errors after rebasing on main with ty

- Fix import path for get_free_port (moved from service to server)
- Add cast for merge_state return type
- Fix test to use async function for TrajectoryGroup creation
- Move tinker deps to separate dependency group
- Add tinker to allowed-unresolved-imports for ty

---------

Co-authored-by: Cursor Bot <bot@cursor.com>

Author: corbt

pyproject.toml

@@ -9,8 +9,6 @@ dependencies = [
     "typer>=0.15.2",
     "litellm>=1.71.1",
     "weave>=0.52.23",
-    "tinker>=0.8.1",
-    "tinker-cookbook>=0.1.0",
     "polars>=1.26.0",
     "tblib>=3.0.0",
     "nest-asyncio>=1.6.0",
@@ -115,6 +113,9 @@ unused-ignore-comment = "ignore"
 # Allow unresolved imports for optional dependencies that may not be installed locally.
 # In CI, we install all optional deps so these will be resolved and type-checked.
 allowed-unresolved-imports = [
+    # tinker deps
+    "tinker.**",
+    "tinker_cookbook.**",
     # backend deps
     "accelerate.**",
     "awscli.**",
@@ -165,6 +166,12 @@ dev = [
     "pyarrow>=15.0.0",
     "prek>=0.2.29",
 ]
+tinker = [
+    "fastapi>=0.128.0",
+    "tinker>=0.8.1",
+    "tinker-cookbook>=0.1.0",
+    "uvicorn>=0.35.0",
+]
 
 [tool.uv.sources]
 panza = { git = "https://github.com/corbt/panza.git" }

src/art/__init__.py

@@ -57,7 +57,13 @@ def __init__(self, **kwargs):
 from .local import LocalBackend
 from .model import Model, TrainableModel
 from .serverless import ServerlessBackend
-from .tinker import TinkerBackend
+
+try:
+    from .tinker import TinkerBackend
+    from .tinker_native import TinkerNativeBackend
+except ModuleNotFoundError:
+    TinkerBackend = None  # type: ignore[assignment]
+    TinkerNativeBackend = None  # type: ignore[assignment]
 from .trajectories import Trajectory, TrajectoryGroup
 from .types import (
     LocalTrainResult,
@@ -91,9 +97,10 @@ def __init__(self, **kwargs):
     "retry",
     "TrainConfig",
     "TrainResult",
-    "TinkerBackend",
     "Trajectory",
     "TrajectoryGroup",
     "capture_yielded_trajectory",
     "yield_trajectory",
 ]
+if TinkerBackend is not None:
+    __all__.extend(["TinkerBackend", "TinkerNativeBackend"])

src/art/dev/__init__.py

@@ -4,6 +4,7 @@
     InternalModelConfig,
     PeftArgs,
     TinkerArgs,
+    TinkerNativeArgs,
     TinkerTrainingClientArgs,
     TrainerArgs,
 )
@@ -16,6 +17,7 @@
     "InitArgs",
     "PeftArgs",
     "TinkerArgs",
+    "TinkerNativeArgs",
     "TinkerTrainingClientArgs",
     "TrainerArgs",
     "get_openai_server_config",

src/art/dev/model.py

@@ -121,6 +121,7 @@ class InternalModelConfig(TypedDict, total=False):
     engine_args: "EngineArgs"
     peft_args: "PeftArgs"
     tinker_args: "TinkerArgs | None"
+    tinker_native_args: "TinkerNativeArgs | None"
     trainer_args: "TrainerArgs"
 
 
@@ -129,6 +130,11 @@ class TinkerArgs(TypedDict, total=False):
     training_client_args: "TinkerTrainingClientArgs"
 
 
+class TinkerNativeArgs(TypedDict, total=False):
+    renderer_name: Required[str]
+    training_client_args: "TinkerTrainingClientArgs"
+
+
 class TinkerTrainingClientArgs(TypedDict, total=False):
     rank: int
     seed: int | None

src/art/local/backend.py

@@ -102,6 +102,8 @@ async def register(
         Args:
             model: An art.Model instance.
         """
+        # Ensure model state/logging uses the backend path
+        model.base_path = self._path
         output_dir = get_model_dir(model=model, art_path=self._path)
         os.makedirs(output_dir, exist_ok=True)
         with open(f"{output_dir}/model.json", "w") as f:

src/art/model.py

@@ -261,18 +261,22 @@ def _get_output_dir(self) -> str:
         """Get the output directory for this model."""
         return f"{self.base_path}/{self.project}/models/{self.name}"
 
-    def write_state(self, state: StateType) -> None:
-        """Write persistent state to the model directory as JSON.
+    def overwrite_state(self, state: StateType) -> None:
+        """Overwrite persistent state in the model directory as JSON.
 
         This state is stored in `state.json` within the model's output directory
         and can be used to track training progress, dataset position, or any
         other information that should persist across runs.
 
+        Warning:
+            This overwrites the entire state file. Prefer `merge_state()` unless
+            you intentionally want to replace all existing keys.
+
         Args:
             state: A dictionary of JSON-serializable values to persist.
 
         Example:
-            model.write_state({
+            model.overwrite_state({
                 "step": 5,
                 "dataset_offset": 100,
                 "last_checkpoint_time": "2024-01-15T10:30:00",
@@ -283,6 +287,45 @@ def write_state(self, state: StateType) -> None:
         with open(f"{output_dir}/state.json", "w") as f:
             json.dump(state, f, indent=2)
 
+    def write_state(self, state: StateType) -> None:
+        """Deprecated: use `overwrite_state()` or `merge_state()` instead."""
+        warnings.warn(
+            "write_state() is deprecated. Use overwrite_state() or merge_state() instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        self.overwrite_state(state)
+
+    def merge_state(self, state: StateType) -> StateType:
+        """Deep-merge state into the existing state and persist it.
+
+        Args:
+            state: A dictionary of JSON-serializable values to merge.
+
+        Returns:
+            The merged state dictionary that was persisted.
+        """
+        existing = self.read_state() or {}
+        merged = self._deep_merge_dicts(existing, state)
+        self.overwrite_state(merged)
+        return cast(StateType, merged)
+
+    @staticmethod
+    def _deep_merge_dicts(
+        base: dict[str, Any], updates: dict[str, Any]
+    ) -> dict[str, Any]:
+        merged = dict(base)
+        for key, value in updates.items():
+            if (
+                key in merged
+                and isinstance(merged[key], dict)
+                and isinstance(value, dict)
+            ):
+                merged[key] = Model._deep_merge_dicts(merged[key], value)
+            else:
+                merged[key] = value
+        return merged
+
     def read_state(self) -> StateType | None:
         """Read persistent state from the model directory.
 

src/art/pipeline_trainer/__init__.py

@@ -0,0 +1,13 @@
+from .status import StatusReporter
+from .trainer import PipelineTrainer, make_group_rollout_fn
+from .types import EvalFn, RolloutFn, ScenarioT, SingleRolloutFn
+
+__all__ = [
+    "PipelineTrainer",
+    "make_group_rollout_fn",
+    "StatusReporter",
+    "RolloutFn",
+    "SingleRolloutFn",
+    "EvalFn",
+    "ScenarioT",
+]