feat: Implement tau-bench minimal training script

dev/tau-bench-minimal.py

@@ -0,0 +1,94 @@
+import asyncio
+import itertools as it
+from typing import Annotated
+
+from dotenv import load_dotenv
+import more_itertools as mit
+import typer
+
+import art
+from art import tau_bench
+from art.tinker import TinkerBackend
+
+app = typer.Typer()
+
+DEFAULT_BASE_MODEL = "Qwen/Qwen3-30B-A3B-Instruct-2507"
+DEFAULT_NAME = "001"
+
+
+async def train(
+    base_model: str,
+    name: str,
+):
+    load_dotenv()
+
+    steps = 100
+    groups = 1
+    trajectories = 16
+
+    backend = TinkerBackend()
+    model = art.TrainableModel(
+        name=name,
+        project="tau-bench",
+        base_model=base_model,
+    )
+    await model.register(backend)
+    scenarios = await tau_bench.get_scenarios(domain="telecom", split="test")
+
+    async def val() -> None:
+        val_trajectories = await art.gather_trajectory_groups(
+            (
+                art.TrajectoryGroup(
+                    tau_bench.rollout(scenario, model, max_turns=30) for _ in range(2)
+                )
+                for scenario in scenarios
+            ),
+            pbar_desc="val",
+        )
+        await model.log(val_trajectories)
+
+    for step in range(await model.get_step(), steps):
+        if step == 0:
+            await val()
+        train_groups = await art.gather_trajectory_groups(
+            (
+                art.TrajectoryGroup(
+                    tau_bench.rollout(scenario, model, max_turns=30)
+                    for _ in range(trajectories)
+                )
+                for scenario in mit.take((step + 1) * groups, it.cycle(scenarios))[
+                    -groups:
+                ]
+            ),
+            pbar_desc=f"gather({step:03d})",
+        )
+        await model.delete_checkpoints()
+        await backend.train(model, train_groups)
+        if step + 1 == steps:
+            await val()
+
+
+@app.command()
+def main(
+    base_model: Annotated[
+        str,
+        typer.Option(
+            "--base-model",
+            "-bm",
+            help="Base model to train.",
+        ),
+    ] = DEFAULT_BASE_MODEL,
+    name: Annotated[
+        str,
+        typer.Option(
+            "--name",
+            "-n",
+            help="Trainable model name.",
+        ),
+    ] = DEFAULT_NAME,
+):
+    asyncio.run(train(base_model, name))
+
+
+if __name__ == "__main__":
+    app()

src/art/megatron/model_support/registry.py

@@ -1,10 +1,5 @@
-from art.megatron.model_support.handlers import (
-    DEFAULT_DENSE_HANDLER,
-    QWEN3_5_DENSE_HANDLER,
-    QWEN3_5_MOE_HANDLER,
-    QWEN3_DENSE_HANDLER,
-    QWEN3_MOE_HANDLER,
-)
+import importlib
+
 from art.megatron.model_support.spec import (
     DependencyFloor,
     ModelSupportHandler,
@@ -49,27 +44,27 @@
 
 DEFAULT_DENSE_SPEC = ModelSupportSpec(
     key="default_dense",
-    handler_key=DEFAULT_DENSE_HANDLER.key,
+    handler_key="default_dense",
     default_target_modules=_DENSE_TARGET_MODULES,
-    native_vllm_lora_status=DEFAULT_DENSE_HANDLER.native_vllm_lora_status,
+    native_vllm_lora_status="disabled",
 )
 
 QWEN3_MOE_SPEC = ModelSupportSpec(
     key="qwen3_moe",
-    handler_key=QWEN3_MOE_HANDLER.key,
+    handler_key="qwen3_moe",
     model_names=(
         "Qwen/Qwen3-30B-A3B",
         "Qwen/Qwen3-30B-A3B-Base",
         "Qwen/Qwen3-30B-A3B-Instruct-2507",
         "Qwen/Qwen3-235B-A22B-Instruct-2507",
     ),
     default_target_modules=_QWEN3_MOE_TARGET_MODULES,
-    native_vllm_lora_status=QWEN3_MOE_HANDLER.native_vllm_lora_status,
+    native_vllm_lora_status="validated",
 )
 
 QWEN3_DENSE_SPEC = ModelSupportSpec(
     key="qwen3_dense",
-    handler_key=QWEN3_DENSE_HANDLER.key,
+    handler_key="qwen3_dense",
     model_names=(
         "Qwen/Qwen3-0.6B",
         "Qwen/Qwen3-0.6B-Base",
@@ -87,34 +82,34 @@
         "Qwen/Qwen3-32B-Base",
     ),
     default_target_modules=_DENSE_TARGET_MODULES,
-    native_vllm_lora_status=QWEN3_DENSE_HANDLER.native_vllm_lora_status,
+    native_vllm_lora_status="validated",
 )
 
 QWEN3_5_DENSE_SPEC = ModelSupportSpec(
     key="qwen3_5_dense",
-    handler_key=QWEN3_5_DENSE_HANDLER.key,
+    handler_key="qwen3_5_dense",
     model_names=(
         "Qwen/Qwen3.5-4B",
         "Qwen/Qwen3.5-27B",
         "Qwen/Qwen3.6-27B",
     ),
     default_target_modules=_QWEN3_5_DENSE_TARGET_MODULES,
-    native_vllm_lora_status=QWEN3_5_DENSE_HANDLER.native_vllm_lora_status,
+    native_vllm_lora_status="validated",
     dependency_floor=DependencyFloor(
         megatron_bridge="e049cc00c24d03e2ae45d2608c7a44e2d2364e3d",
     ),
 )
 
 QWEN3_5_MOE_SPEC = ModelSupportSpec(
     key="qwen3_5_moe",
-    handler_key=QWEN3_5_MOE_HANDLER.key,
+    handler_key="qwen3_5_moe",
     model_names=(
         "Qwen/Qwen3.5-35B-A3B",
         "Qwen/Qwen3.5-397B-A17B",
         "Qwen/Qwen3.6-35B-A3B",
     ),
     default_target_modules=_QWEN3_5_MOE_TARGET_MODULES,
-    native_vllm_lora_status=QWEN3_5_MOE_HANDLER.native_vllm_lora_status,
+    native_vllm_lora_status="validated",
     dependency_floor=DependencyFloor(
         megatron_bridge="e049cc00c24d03e2ae45d2608c7a44e2d2364e3d",
     ),
@@ -143,13 +138,14 @@
     for spec in PROBE_ONLY_MODEL_SUPPORT_SPECS
     for model_name in spec.model_names
 }
-_HANDLERS_BY_KEY: dict[str, ModelSupportHandler] = {
-    DEFAULT_DENSE_HANDLER.key: DEFAULT_DENSE_HANDLER,
-    QWEN3_DENSE_HANDLER.key: QWEN3_DENSE_HANDLER,
-    QWEN3_MOE_HANDLER.key: QWEN3_MOE_HANDLER,
-    QWEN3_5_DENSE_HANDLER.key: QWEN3_5_DENSE_HANDLER,
-    QWEN3_5_MOE_HANDLER.key: QWEN3_5_MOE_HANDLER,
+_HANDLER_EXPORTS_BY_KEY = {
+    "default_dense": ("default_dense", "DEFAULT_DENSE_HANDLER"),
+    "qwen3_dense": ("qwen3_dense", "QWEN3_DENSE_HANDLER"),
+    "qwen3_moe": ("qwen3_moe", "QWEN3_MOE_HANDLER"),
+    "qwen3_5_dense": ("qwen3_5", "QWEN3_5_DENSE_HANDLER"),
+    "qwen3_5_moe": ("qwen3_5", "QWEN3_5_MOE_HANDLER"),
 }
+_MOE_HANDLER_KEYS = {"qwen3_moe", "qwen3_5_moe"}
 
 QWEN3_DENSE_MODELS = frozenset(QWEN3_DENSE_SPEC.model_names)
 QWEN3_MOE_MODELS = frozenset(QWEN3_MOE_SPEC.model_names)
@@ -195,7 +191,11 @@ def get_model_support_handler(
 def get_model_support_handler_for_spec(
     spec: ModelSupportSpec,
 ) -> ModelSupportHandler:
-    return _HANDLERS_BY_KEY[spec.handler_key]
+    module_name, export_name = _HANDLER_EXPORTS_BY_KEY[spec.handler_key]
+    return getattr(
+        importlib.import_module(f"art.megatron.model_support.handlers.{module_name}"),
+        export_name,
+    )
 
 
 def default_target_modules_for_model(
@@ -216,7 +216,7 @@ def native_vllm_lora_status_for_model(
     *,
     allow_unvalidated_arch: bool = False,
 ) -> str:
-    return get_model_support_handler(
+    return get_model_support_spec(
         base_model,
         allow_unvalidated_arch=allow_unvalidated_arch,
     ).native_vllm_lora_status
@@ -241,11 +241,12 @@ def model_uses_expert_parallel(
     *,
     allow_unvalidated_arch: bool = False,
 ) -> bool:
-    return bool(
-        get_model_support_handler(
+    return (
+        get_model_support_spec(
             base_model,
             allow_unvalidated_arch=allow_unvalidated_arch,
-        ).is_moe
+        ).handler_key
+        in _MOE_HANDLER_KEYS
     )