1"""Track D runner: hosted black-box receivers (Anthropic claude-haiku-4-5; OpenAI gpt-5-nano) on A0-A6 messages built by2the LOCAL sender (Qwen2.5-3B), plus provider prefix caching (documented cache_control). Only synthetic + exec + public dev/test3items (no private data). Hard budget cap enforced from configs/resources.yaml."""4from __future__ import annotations5import json, time, yaml6from pathlib import Path7import mlx.core as mx8from .util import ROOT, read_json, sha256_json, sha256_text, utc_now, read_jsonl9from .records import new_record, write_record, trial_key10from .data.prepare import load_doc, load_key11from .scoring import score_synthetic, score_public12from .data.exec_tasks import run_hidden_tests13from .prompts import SYSTEM_RECEIVER, SYSTEM_RECEIVER_NOMSG, receiver_user_text, receiver_question_only14from .backends.mlx_backend import MLXModel15from .methods import text as TM16from .hosted import anthropic_answer, openai_answer, Budget1718def run(config_path: str, resume: bool = False) -> int:19 cfg = yaml.safe_load(open(config_path)); res = yaml.safe_load(open(ROOT / "configs/resources.yaml"))20 budget = Budget(res["budgets"]["api"]["hard_cap_usd"], ROOT / "runs/logs/api_cost_ledger.jsonl")21 mcfg = yaml.safe_load(open(ROOT / "configs/models.yaml"))["models"]22 sender = MLXModel(mcfg[cfg["sender"]]["repo"])23 man = read_json(ROOT / "data/manifests/splits.json")24 items = [it for it in man["items"] if it["split"] in cfg["splits"] and it["family"] in cfg["families"] and (it.get("target_tokens") is None or it["target_tokens"] in cfg.get("lengths", [it.get("target_tokens")]))]25 lim = cfg.get("limit_docs_per_family")26 if lim:27 items = [it for fam in cfg["families"] for it in [x for x in items if x["family"] == fam][:lim]]28 out_path = ROOT / "runs/raw" / f"{cfg['experiment_id']}.jsonl"; log_path = ROOT / "runs/logs" / f"{cfg['experiment_id']}.log"29 done = {r["trial_id"] for r in read_jsonl(out_path) if r["status"] == "success"} if (resume and out_path.exists()) else set()30 protocol_hash = cfg.get("protocol_hash", "unfrozen"); config_hash = sha256_json(cfg)31 def log(m):32 line = f"{utc_now()} {m}"; print(line, flush=True)33 with open(log_path, "a") as f: f.write(line + "\n")34 log(f"hosted start {cfg['experiment_id']} items={len(items)} spent_usd={budget.spent:.3f} cap={budget.cap}")35 for it in items:36 fam = it["family"]; doc = load_doc(it["document_id"], fam); key = load_key(it["document_id"], fam)37 qkeys = {q["question_id"]: q for q in key["questions"]}38 qs = doc["questions"][: cfg.get("max_questions_per_doc") or None]39 for prov in cfg["providers"]:40 for cond in cfg["conditions"]:41 method = cond["method"]; access = cond.get("access", "none"); cid = f"{method}|{prov}|" + "|".join(f"{k}={v}" for k, v in cond.items() if k not in ("method", "access"))42 pending = [(q, trial_key(cfg["experiment_id"], it["document_id"], q["question_id"], method, {k: v for k, v in cond.items() if k != "method"}, cfg["sender"], prov, access, protocol_hash, 0, "hosted")) for q in qs]43 pending = [(q, tk) for q, tk in pending if tk not in done]44 if not pending: continue45 shared = None46 if access in ("blind", "none"):47 shared = _build(method, cond, doc["text"], None, sender)48 for q, tk in pending:49 if budget.exhausted():50 log("BUDGET EXHAUSTED"); return 051 msg = shared or _build(method, cond, doc["text"], q["question"], sender)52 material = msg.text; fetched = 053 if method == "A6":54 material, ri = TM.a6_retrieve(msg, q["question"], cond.get("max_fetch_bytes", 2000), cond.get("max_lookups", 6)); fetched = ri["fetched_bytes"]55 if method == "A0":56 system, user = SYSTEM_RECEIVER_NOMSG, receiver_question_only(q["question"], fam)57 else:58 system, user = SYSTEM_RECEIVER, receiver_user_text(msg.material_label, material, q["question"], fam)59 rec = new_record(experiment_id=cfg["experiment_id"], trial_id=tk, run_group_id=f"{cfg['experiment_id']}:{cid}", protocol_hash=protocol_hash, config_hash=config_hash, task_family=fam,60 document_id=it["document_id"], document_hash=it["document_hash"], question_id=q["question_id"], split=it["split"], cluster_id=it["cluster_id"],61 sender_model_revision=f"{sender.info.repo_id}@{sender.info.revision}", receiver_model_revision=prov, tokenizer_hashes={"sender": sender.info.tokenizer_hash},62 backend_versions={"anthropic_sdk": __import__("anthropic").__version__}, hardware_id="hosted:" + prov.split(":")[0], device_state=None, weight_precision="provider_undisclosed",63 method_id=method, implementation_fidelity="custom", compression_config={k: v for k, v in cond.items() if k != "method"}, question_access=access, topology="hosted_single_handoff",64 handoff_index=1 if method != "A0" else 0, reuse_count=len(pending) if access in ("blind", "none") else 1, cache_hit=False, cache_compatibility_hash=None, random_seed=0,65 source_tokens_sender=len(sender.encode(doc["text"])), source_tokens_receiver=None, query_tokens=None, message_tokens=None, message_utf8_bytes=msg.utf8_bytes,66 cache_tensor_bytes=None, serialized_payload_bytes=msg.payload_bytes, side_information_bytes=0, fetched_bytes=fetched, timing_spans={"encode_ms": msg.encode_ms},67 input_usage={**msg.sender_usage}, output_usage={}, provider_cached_usage=None, process_memory=None, backend_memory=None, peak_prefill_memory=None, peak_decode_memory=None,68 mapper_memory=None, cost_with_source=None, prediction_artifact=None, evidence_artifact=None, reference_key_id=f"keys/{fam}/{it['document_id']}.json", task_scores=None,69 status="implementation_error", error_code=None, log_artifact=str(log_path.relative_to(ROOT)), payload_hash=sha256_text(material), output_hash=None)70 try:71 max_out = 320 if fam == "exec" else 6472 if prov.startswith("anthropic"):73 use_cache = cond.get("provider_cache", False)74 blocks = [{"text": user}]75 if use_cache and method != "A0":76 # split material (cacheable) from question (fresh) so the provider can cache the document prefix77 blocks = [{"text": f"{msg.material_label}:\n{material}\n\n", "cache": True}, {"text": user.split(material, 1)[-1].lstrip("\n")}]78 r = anthropic_answer(system, blocks, model=res["budgets"]["api"]["providers"]["anthropic"]["model"], max_tokens=max_out, use_cache_control=use_cache)79 rec["provider_cached_usage"] = {"cache_creation_input_tokens": r["cache_creation_input_tokens"], "cache_read_input_tokens": r["cache_read_input_tokens"]}80 else:81 r = openai_answer(system, user, model=res["budgets"]["api"]["providers"]["openai"]["model"], max_tokens=max_out)82 rec["provider_cached_usage"] = {"cached_tokens": r["cached_tokens"]}83 budget.charge(r["cost_usd_estimated"], {"experiment": cfg["experiment_id"], "trial": tk, "provider": prov, "model": r["model"]})84 pred = r["text"]85 rec.update({"prediction_artifact": pred, "output_hash": sha256_text(pred), "ttft_ms": r["ttft_ms"], "total_wall_ms": r["total_ms"] + msg.encode_ms,86 "input_usage": {**rec["input_usage"], "provider_input_tokens": r["input_tokens"]}, "output_usage": {"provider_output_tokens": r["output_tokens"]},87 "cost_with_source": {"usd": r["cost_usd_estimated"], "kind": "estimated", "price_source": r["price_source"], "model": r["model"]}, "receiver_model_revision": r["model"]})88 rec["timing_spans"].update({"provider_total_ms": r["total_ms"], "ttft_ms": r["ttft_ms"]})89 qk = qkeys[q["question_id"]]90 if fam == "exec":91 t = run_hidden_tests(pred, key["hidden_tests"]); rec["task_scores"] = {"correct": t["correct"], "pass_rate": t["pass_rate"], "error": t.get("error")}92 elif fam.startswith("public_"):93 rec["task_scores"] = score_public(pred, [qk["answer"]] + list(qk.get("aliases", [])))94 else:95 rec["task_scores"] = score_synthetic(pred, qk)96 rec["status"] = "success" if pred.strip() else "invalid_output"97 except Exception as e:98 rec["error_code"] = f"{type(e).__name__}: {str(e)[:300]}"99 if "rate" in str(e).lower(): time.sleep(5)100 write_record(rec, out_path)101 sc = rec["task_scores"] or {}102 log(f"{cid:>36s} {it['document_id']:>26s} {q['question_id']:>16s} {rec['status']:>8s} corr={sc.get('correct', sc.get('f1'))} cost={(rec['cost_with_source'] or {}).get('usd', 0):.5f} cached={rec['provider_cached_usage']} ans={json.dumps((rec['prediction_artifact'] or '')[:40])}")103 mx.clear_cache()104 log(f"hosted end spent_usd={budget.spent:.3f}")105 return 0106107def _build(method, cond, doc_text, question, sender):108 if method == "A0": return TM.a0_no_message()109 if method == "A1": return TM.a1_full(doc_text)110 if method == "A2": return TM.a2_truncate(doc_text, cond["budget_tokens"], sender.tokenizer)111 if method == "A3": return TM.a3_extractive(doc_text, cond["budget_tokens"], sender.tokenizer, question)112 if method == "A4": return TM.a4_summary(doc_text, cond["budget_tokens"], sender, question)113 if method == "A5": return TM.a5_structured(doc_text, cond["budget_tokens"], sender, question)114 if method == "A6": return TM.a6_references(doc_text, cond["budget_tokens"], sender, question)115 raise NotImplementedError(method)116