diff --git a/scripts/convert_inferact_to_trace.py b/scripts/convert_inferact_to_trace.py new file mode 100644 index 0000000..52b281f --- /dev/null +++ b/scripts/convert_inferact_to_trace.py @@ -0,0 +1,189 @@ +"""Convert Inferact codex_swebenchpro_traces (ShareGPT) to agentic-pd-hybrid trace JSONL. + +Output schema (one JSON object per line, matching src/agentic_pd_hybrid/trace.py): + chat_id, parent_chat_id, timestamp, input_length, output_length, type, turn, hash_ids + +Each trial in the input becomes one session. Each (human, gpt) pair within a trial +becomes one turn. The prefix at turn N is the concatenation of all (human, gpt) pairs +from turns 0..N-1 plus the current human message — this mirrors how agentic coding +agents grow context across calls. + +hash_ids are derived per 24-token block via sha256 of the block's text + previous hash, +which gives stable, deterministic, prefix-shared hashes across turns of the same session. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import sys +import time +from pathlib import Path + +BLOCK_TOKEN_BUDGET = 24 + + +def _block_hash(text: str, prev_hash: int) -> int: + h = hashlib.sha256(text.encode("utf-8") + prev_hash.to_bytes(8, "big")).digest() + return int.from_bytes(h[:8], "big") & 0x7FFFFFFFFFFFFFFF + + +def _build_hash_ids(token_ids: list[int]) -> list[int]: + out: list[int] = [] + prev = 0 + for start in range(0, len(token_ids), BLOCK_TOKEN_BUDGET): + block = token_ids[start : start + BLOCK_TOKEN_BUDGET] + block_repr = ",".join(str(t) for t in block) + prev = _block_hash(block_repr, prev) + out.append(prev) + return out + + +def _pair_turns(conv: list[dict]) -> list[tuple[str, str]]: + """Pair consecutive (human, gpt) messages. Skip malformed.""" + pairs: list[tuple[str, str]] = [] + i = 0 + while i + 1 < len(conv): + a, b = conv[i], conv[i + 1] + if ( + isinstance(a, dict) + and isinstance(b, dict) + and a.get("from") == "human" + and b.get("from") == "gpt" + ): + pairs.append((str(a.get("value", "")), str(b.get("value", "")))) + i += 2 + else: + i += 1 + return pairs + + +def convert( + input_path: Path, + output_path: Path, + *, + tokenizer_path: str, + max_trials: int | None, + inter_turn_gap_s: float, + session_stagger_s: float, + request_type: str, +) -> None: + from transformers import AutoTokenizer + + print(f"loading tokenizer from {tokenizer_path}", file=sys.stderr) + tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True) + + print(f"loading {input_path}", file=sys.stderr) + data = json.loads(input_path.read_text()) + if max_trials is not None: + data = data[:max_trials] + print(f"{len(data)} trials to process", file=sys.stderr) + + next_chat_id = 1_000_000 + written = 0 + skipped_trials = 0 + t0 = time.time() + + with output_path.open("w", encoding="utf-8") as out_f: + for trial_idx, trial in enumerate(data): + conv = trial.get("conversations") or [] + turns = _pair_turns(conv) + if not turns: + skipped_trials += 1 + continue + + base_ts = trial_idx * session_stagger_s + ts = base_ts + parent_chat_id = -1 + prefix_text = "" + + for turn_idx, (human, assistant) in enumerate(turns): + # Input at this turn = full prior context + current human message. + current_text = ( + prefix_text + ("\n\n[USER]\n" if prefix_text else "[USER]\n") + human + ) + input_ids = tokenizer.encode(current_text, add_special_tokens=False) + input_length = len(input_ids) + + output_ids = tokenizer.encode(assistant, add_special_tokens=False) + output_length = max(1, len(output_ids)) + + hash_ids = _build_hash_ids(input_ids) + + chat_id = next_chat_id + next_chat_id += 1 + record = { + "chat_id": chat_id, + "parent_chat_id": parent_chat_id, + "timestamp": round(ts, 6), + "input_length": input_length, + "output_length": output_length, + "type": request_type, + "turn": turn_idx, + "hash_ids": hash_ids, + } + out_f.write(json.dumps(record) + "\n") + written += 1 + + parent_chat_id = chat_id + ts += inter_turn_gap_s + prefix_text = current_text + "\n\n[ASSISTANT]\n" + assistant + + if (trial_idx + 1) % 20 == 0: + elapsed = time.time() - t0 + rate = (trial_idx + 1) / elapsed if elapsed > 0 else 0 + eta = (len(data) - trial_idx - 1) / rate if rate > 0 else 0 + print( + f" trial {trial_idx + 1}/{len(data)} reqs={written} " + f"rate={rate:.1f} trial/s eta={eta:.0f}s", + file=sys.stderr, + ) + + elapsed = time.time() - t0 + print( + f"done: wrote {written} requests across {len(data) - skipped_trials} sessions " + f"({skipped_trials} trials skipped, empty conversations) in {elapsed:.1f}s " + f"to {output_path}", + file=sys.stderr, + ) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument( + "--input", + type=Path, + default=Path("third_party/codex_swebenchpro_traces/codex_swebenchpro.json"), + ) + p.add_argument("--output", type=Path, required=True) + p.add_argument( + "--tokenizer", + default="/mnt/models/Qwen/Qwen3-30B-A3B-Instruct-2507", + help="Path or HF id for the tokenizer. Default matches v2 sweep model.", + ) + p.add_argument( + "--max-trials", + type=int, + default=None, + help="Cap number of trials processed (useful for smoke / quick tests).", + ) + p.add_argument("--inter-turn-gap-s", type=float, default=2.5) + p.add_argument("--session-stagger-s", type=float, default=1.0) + p.add_argument("--request-type", default="chat") + args = p.parse_args() + + args.output.parent.mkdir(parents=True, exist_ok=True) + convert( + input_path=args.input, + output_path=args.output, + tokenizer_path=args.tokenizer, + max_trials=args.max_trials, + inter_turn_gap_s=args.inter_turn_gap_s, + session_stagger_s=args.session_stagger_s, + request_type=args.request_type, + ) + + +if __name__ == "__main__": + main() diff --git a/scripts/setup_env.sh b/scripts/setup_env.sh new file mode 100755 index 0000000..897b75e --- /dev/null +++ b/scripts/setup_env.sh @@ -0,0 +1,35 @@ +#!/usr/bin/env bash +# Source this file in every shell that will run agentic-pd-hybrid. +# +# source scripts/setup_env.sh +# +# Why all three are needed: +# - CUDA_HOME / PATH point tvm_ffi (vendor sglang JIT compiler) at cu12.8 nvcc. +# Without this it falls back to /usr/local/cuda-13.0/bin/nvcc and the +# resulting .so links libcudart.so.13 which driver 570 (cu12.8 API) rejects +# with cudaErrorInsufficientDriver. +# - LD_LIBRARY_PATH must expose libcudart.so.12 for mooncake.engine (cu12 wheel) +# AND ~/cuda-12.8/lib64 for tvm_ffi compile-time linker searches. +# +# See docs/H200_DRIVER570_SETUP_ZH.md for the full rationale. + +REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + +if [ ! -x "$HOME/cuda-12.8/bin/nvcc" ]; then + echo "ERROR: $HOME/cuda-12.8/bin/nvcc not found." >&2 + echo "Install cu12.8 toolkit first (see docs/H200_DRIVER570_SETUP_ZH.md §3)." >&2 + return 1 2>/dev/null || exit 1 +fi + +if [ ! -f "$REPO_ROOT/.venv/lib/python3.12/site-packages/nvidia/cuda_runtime/lib/libcudart.so.12" ]; then + echo "ERROR: venv libcudart.so.12 missing. Run 'uv sync' from $REPO_ROOT." >&2 + return 1 2>/dev/null || exit 1 +fi + +export CUDA_HOME="$HOME/cuda-12.8" +export PATH="$HOME/cuda-12.8/bin:$PATH" +export LD_LIBRARY_PATH="$REPO_ROOT/.venv/lib/python3.12/site-packages/nvidia/cuda_runtime/lib:$HOME/cuda-12.8/lib64${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + +echo "agentic-pd-hybrid env ready:" +echo " CUDA_HOME=$CUDA_HOME ($(nvcc --version | grep release | sed 's/.*release //'))" +echo " libcudart.so.12 at $REPO_ROOT/.venv/lib/python3.12/site-packages/nvidia/cuda_runtime/lib"