mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-27 21:05:34 +00:00
Andrew Park
e6c47e9532
The orchestrator SFT pipeline was living entirely in the working tree and never got committed. One script in that lineage (clean_sft_data.py) has already been lost with no way to recover it, so track the rest before the same thing happens again: run_sft_fsdp.py full-parameter FSDP trainer (launched via accelerate) sft_tokenize.py conversation -> tokens + assistant-only masking render_sft_data.py record -> chat-template rendering make_splits.py train / holdout / overfit splits chat_orchestrator.py REPL for eyeballing a trained checkpoint's routing format_eval_sample.py, upload_to_braintrust.py scripts/train/ fsdp4.yaml, fsdp8.yaml, eval sbatch toolorchestra/tracing.py was also untracked despite being imported by unified.py and rollout.py, which left the branch broken for a fresh checkout. Add it along with the rest of the ToolOrchestra package split. Also here: mmlu_pro / supergpqa scorers, expert pricing for the OpenRouter Qwen builds (estimates, flagged in-file), an expanded calculator/web_search tool surface, and the reject-sampling clean gate that run_sft_fsdp reads via --require-clean. GRPO is left at its upstream state -- the RL stage hasn't been run yet, so it stays out of this change. Paths in the sbatch are env-driven rather than hardcoded to one cluster home.
110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
#!/usr/bin/env python
|
|
"""Interactive REPL to chat with a trained orchestrator checkpoint.
|
|
|
|
Reuses the SAME routing loop the eval uses (OrchestratorBackend.generate_full),
|
|
so what you see here is exactly how it behaves on the benchmarks: it reads your
|
|
question, decides whether to answer itself / call a tool / delegate to an expert
|
|
model, runs that rollout, and prints the final answer + a routing summary.
|
|
|
|
Prereqs:
|
|
1. Serve the checkpoint you want on a vLLM port (see --endpoint). e.g. 8k:
|
|
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True CUDA_VISIBLE_DEVICES=0 \
|
|
.venv/bin/python -m vllm.entrypoints.openai.api_server \
|
|
--model "$OJ_WORK/qwen_eval_ckpts/8k_served" \
|
|
--served-model-name sft-qwen-8k --port 8020 \
|
|
--gpu-memory-utilization 0.88 --trust-remote-code --max-model-len 32768 \
|
|
--enforce-eager --enable-auto-tool-choice --tool-call-parser qwen3_xml
|
|
2. source .env (cloud experts GPT-5.5 / Opus need the API keys)
|
|
|
|
Usage:
|
|
.venv/bin/python scripts/orchestrator/chat_orchestrator.py \
|
|
--endpoint http://localhost:8020/v1 --model sft-qwen-8k
|
|
|
|
# point local expert tiers at their own vLLMs if you have them up (optional):
|
|
# --local-endpoint Qwen/Qwen3.6-27B=http://localhost:8002/v1
|
|
"""
|
|
import argparse
|
|
import sys
|
|
import time
|
|
|
|
|
|
def _parse_args(argv=None):
|
|
p = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
p.add_argument("--endpoint", default="http://localhost:8020/v1",
|
|
help="vLLM endpoint serving the orchestrator checkpoint.")
|
|
p.add_argument("--model", default="sft-qwen-8k",
|
|
help="Served-model-name of the checkpoint.")
|
|
p.add_argument("--api-key", default="EMPTY", help="API key for the endpoint (local vLLM = EMPTY).")
|
|
p.add_argument("--temperature", type=float, default=0.0)
|
|
p.add_argument("--max-turns", type=int, default=8, help="Same default as the eval.")
|
|
p.add_argument("--max-tokens", type=int, default=3000,
|
|
help="Keep high: the model reasons a lot before it emits the delegation.")
|
|
p.add_argument("--local-endpoint", action="append", default=[],
|
|
help="MODEL_ID=URL for a local expert tier. Repeatable.")
|
|
return p.parse_args(argv)
|
|
|
|
|
|
def main(argv=None):
|
|
args = _parse_args(argv)
|
|
|
|
local_endpoints = {}
|
|
for pair in args.local_endpoint:
|
|
if "=" not in pair:
|
|
raise SystemExit(f"--local-endpoint expects MODEL_ID=URL, got: {pair!r}")
|
|
mid, url = pair.split("=", 1)
|
|
local_endpoints[mid.strip()] = url.strip()
|
|
|
|
from openjarvis.learning.intelligence.orchestrator.eval_backend import (
|
|
OrchestratorBackend,
|
|
)
|
|
|
|
backend = OrchestratorBackend(
|
|
orchestrator_endpoint=args.endpoint,
|
|
orchestrator_model=args.model,
|
|
api_key=args.api_key,
|
|
local_endpoints=local_endpoints,
|
|
max_turns=args.max_turns,
|
|
temperature=args.temperature,
|
|
)
|
|
|
|
print(f"orchestrator: {args.model} @ {args.endpoint} "
|
|
f"(temp={args.temperature}, max_turns={args.max_turns})")
|
|
print("type a question and hit enter. Ctrl-D or 'quit' to exit.\n")
|
|
|
|
try:
|
|
while True:
|
|
try:
|
|
prompt = input("you> ").strip()
|
|
except EOFError:
|
|
print()
|
|
break
|
|
if not prompt:
|
|
continue
|
|
if prompt.lower() in {"quit", "exit"}:
|
|
break
|
|
|
|
started = time.time()
|
|
full = backend.generate_full(
|
|
prompt, model=args.model, temperature=args.temperature,
|
|
max_tokens=args.max_tokens,
|
|
)
|
|
elapsed = time.time() - started
|
|
|
|
if full.get("error"):
|
|
print(f"\n[ERROR] {full['error']}\n", file=sys.stderr)
|
|
continue
|
|
|
|
print(f"\norch> {full.get('content', '').strip()}\n")
|
|
print(f" [turns={full.get('turn_count', '?')} "
|
|
f"tool_calls={full.get('tool_calls', '?')} "
|
|
f"cost=${full.get('cost_usd', 0.0):.4f} "
|
|
f"{elapsed:.1f}s]\n")
|
|
finally:
|
|
backend.close()
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|