feat: wire in FFT worker (#114)

* feat: add full worker mode support

Add OPEN_RL_WORKER_MODE=full as the gateway-selected path for launching one FFT training worker process per created model.

The gateway owns the FFTWorkerLauncher lifecycle, starts a run-scoped clock_cycle process before enqueueing create_model/create_model_from_state, and shuts down launched workers with the FastAPI lifespan.

The worker loop now chooses FFTTrainingWorker in full mode, drains only its model queue through Redis, saves full checkpoints for sampler exports, and keeps public Tinker SDK metadata LoRA-shaped until the client has native full fine-tuning support.

Also moves create_model loading into the trainer workers and covers the LoRA/FFT create_model, checkpoint path, sampler-save, and optimizer behavior in tests.

* docs: add gsm8k fft example

* fix: decouple fft worker launch

Author: ShubyM

examples/sft/gsm8k/README.md

@@ -0,0 +1,73 @@
+# GSM8K full fine-tuning
+
+Full-parameter SFT of a small model on GSM8K, driven through the OpenRL gateway
+with the Tinker SDK.
+
+## Why full fine-tuning goes through dedicated workers
+
+The public Tinker SDK entrypoint is still `create_lora_training_client()`. For
+now, OpenRL routes that same client flow to a full fine-tuning worker when the
+gateway is started with `OPEN_RL_ENABLE_FFT=true`.
+
+## Run
+
+This branch launches one full fine-tuning worker process per created model. That
+worker shares requests and futures with the gateway through Redis.
+
+Start from the repository root in separate terminals.
+
+### Terminal 1: Redis
+
+```bash
+redis-server --port 6379 --save "" --appendonly no
+```
+
+### Terminal 2: Gateway
+
+```bash
+cd src/server
+REDIS_URL=redis://127.0.0.1:6379 \
+OPEN_RL_ENABLE_FFT=true \
+BASE_MODEL=Qwen/Qwen2.5-0.5B \
+SAMPLING_BACKEND=torch \
+uv run --extra gpu python -m uvicorn gateway:app --host 127.0.0.1 --port 9003
+```
+
+### Terminal 3: SFT Job
+
+```bash
+uv --project examples run python examples/sft/gsm8k/gsm8k_sft.py \
+  --log-path=examples/sft/gsm8k/artifacts/job_a \
+  --max-steps=20 \
+  --base-model=Qwen/Qwen2.5-0.5B
+```
+
+Training uses `tinker_cookbook.supervised.train`, so batching, LR scheduling,
+metric logging, and final checkpoint export are handled by the cookbook loop.
+The example deletes an existing log directory by default so stale checkpoint
+metadata does not trigger resume. The training script prints
+`eval_model_path=...` when it can resolve a final checkpoint path.
+
+## Eval
+
+Eval is decoupled from OpenRL. Point vLLM at the saved Hugging Face checkpoint:
+
+```bash
+python examples/sft/gsm8k/vllm_eval.py \
+  --path <eval_model_path> \
+  --data gsm8k_test.json
+```
+
+## Result
+
+Single-job result from the original FFT prototype run:
+
+| Setup | GSM8K |
+| --- | --- |
+| Qwen2.5-0.5B base, 0-shot exact match on 250 examples | ~1.5% |
+| Qwen2.5-0.5B after full-FT SFT, 1 epoch, lr 2e-5 | ~36% |
+
+Files:
+
+- `gsm8k_sft.py`: training via the OpenRL/Tinker server.
+- `vllm_eval.py`: fast eval of the saved checkpoint directory.

examples/sft/gsm8k/gsm8k_sft.py

@@ -0,0 +1,97 @@
+import asyncio
+import os
+from pathlib import Path
+from typing import Any, cast
+
+import chz
+import tinker
+from datasets import load_dataset
+from tinker import types
+from tinker_cookbook import checkpoint_utils, cli_utils
+from tinker_cookbook.supervised.data import SupervisedDatasetFromHFDataset
+from tinker_cookbook.supervised.train import Config as TrainConfig
+from tinker_cookbook.supervised.train import main as train
+from tinker_cookbook.supervised.types import SupervisedDatasetBuilder
+from tinker_cookbook.tokenizer_utils import get_tokenizer
+
+os.environ.setdefault("TINKER_API_KEY", "tml-dummy-key")
+
+
+@chz.chz
+class GSM8KDataset(SupervisedDatasetBuilder):
+  model_name: str
+  batch_size: int = 16
+  max_length: int = 640
+  seed: int = 0
+
+  def __call__(self):
+    tokenizer = get_tokenizer(self.model_name)
+    dataset = load_dataset("openai/gsm8k", "main", split="train").shuffle(seed=self.seed)
+
+    def make_datum(row: dict) -> tinker.Datum:
+      prompt = tokenizer.encode(f"Question: {row['question']}\nAnswer:", add_special_tokens=False)
+      completion = tokenizer.encode(" " + row["answer"].strip(), add_special_tokens=False) + [tokenizer.eos_token_id]
+      tokens = (prompt + completion)[: self.max_length]
+      weights = ([0] * len(prompt) + [1] * len(completion))[: self.max_length]
+      return types.Datum(
+        model_input=types.ModelInput.from_ints(tokens=tokens[:-1]),
+        loss_fn_inputs=cast(Any, {"target_tokens": tokens[1:], "weights": [float(w) for w in weights[1:]]}),
+      )
+
+    return SupervisedDatasetFromHFDataset(dataset, self.batch_size, map_fn=make_datum), None
+
+
+@chz.chz
+class Config:
+  base_model: str = "Qwen/Qwen2.5-0.5B"
+  base_url: str = os.getenv("TINKER_BASE_URL", os.getenv("BASE_URL", "http://127.0.0.1:9003"))
+  log_path: str = str(Path(__file__).with_name("artifacts") / "gsm8k_sft")
+  epochs: int = 1
+  batch: int = 16
+  lr: float = 2e-5
+  rank: int = 32
+  max_len: int = 640
+  seed: int = 0
+  max_steps: int | None = None
+  save_every: int = 0
+  behavior_if_log_dir_exists: cli_utils.LogdirBehavior = "delete"
+
+
+def main(config: Config) -> None:
+  cli_utils.check_log_dir(config.log_path, behavior_if_exists=config.behavior_if_log_dir_exists)
+  asyncio.run(
+    train(
+      TrainConfig(
+        log_path=config.log_path,
+        model_name=config.base_model,
+        dataset_builder=GSM8KDataset(
+          model_name=config.base_model,
+          batch_size=config.batch,
+          max_length=config.max_len,
+          seed=config.seed,
+        ),
+        learning_rate=config.lr,
+        lr_schedule="cosine",
+        num_epochs=config.epochs,
+        lora_rank=config.rank,
+        base_url=config.base_url,
+        save_every=config.save_every,
+        eval_every=0,
+        infrequent_eval_every=0,
+        max_steps=config.max_steps,
+      )
+    )
+  )
+  checkpoint = checkpoint_utils.get_last_checkpoint(config.log_path, required_key="sampler_path")
+  if checkpoint is None:
+    checkpoint = checkpoint_utils.get_last_checkpoint(config.log_path, required_key="state_path")
+  if checkpoint is not None:
+    path = checkpoint.sampler_path or checkpoint.state_path
+    if path and path.startswith("tinker://"):
+      path = str(Path(os.getenv("OPEN_RL_TMP_DIR", "/tmp/open-rl")) / "sampler_full" / path.removeprefix("tinker://"))
+    if path:
+      print(f"eval_model_path={path}")
+
+
+if __name__ == "__main__":
+  chz.nested_entrypoint(main, allow_hyphens=True)

examples/sft/gsm8k/vllm_eval.py

@@ -0,0 +1,43 @@
+import argparse
+import json
+import re
+import time
+
+ANS_RE = re.compile(r"-?\d[\d,]*")
+
+
+def extract(text: str) -> str | None:
+  text = re.split(r"\n\s*Question:", text)[0]
+  if "####" in text:
+    match = ANS_RE.search(text.split("####")[-1])
+    if match:
+      return match.group(0).replace(",", "")
+  numbers = ANS_RE.findall(text)
+  return numbers[-1].replace(",", "") if numbers else None
+
+
+def main() -> None:
+  from vllm import LLM, SamplingParams
+
+  parser = argparse.ArgumentParser()
+  parser.add_argument("--path", required=True)
+  parser.add_argument("--data", default="gsm8k_test.json")
+  args = parser.parse_args()
+
+  with open(args.data) as f:
+    data = json.load(f)
+
+  llm = LLM(model=args.path, dtype="bfloat16", gpu_memory_utilization=0.85, max_model_len=1024, enforce_eager=True)
+  sampling_params = SamplingParams(temperature=0.0, max_tokens=256, stop=["\nQuestion:"])
+  start = time.time()
+  outputs = llm.generate([datum["prompt"] for datum in data], sampling_params)
+  elapsed = time.time() - start
+  correct = sum(int(extract(output.outputs[0].text) == datum["gold"]) for datum, output in zip(data, outputs, strict=True))
+
+  print("***************************************************************")
+  print(f"[VLLM] {args.path} 0-shot GSM8K acc = {correct / len(data):.1%} on {len(data)} problems in {elapsed:.1f}s")
+  print("***************************************************************")
+
+
+if __name__ == "__main__":
+  main()