From 133c9b1431bdd624854ab16aced770adb7f3c7c0 Mon Sep 17 00:00:00 2001 From: June Liu Date: Wed, 21 Oct 2026 10:31:23 -0700 Subject: [PATCH 01/15] feat(examples): add COBOL CPT, GRPO, and eval pipeline Add examples/cobol-cpt-grpo, a production-shaped pipeline that adapts Qwen2.5-Coder-32B to COBOL on 2 x 8 H200 with RDMA: - 01-collect, 02-data-prep, 03-benchmarks and 05-rl-prompts run as CPU Jobs. Source repositories default to permissive licenses (SOURCE_LICENSES=all opts in to the rest), and gated Hugging Face datasets run only when HF_TOKEN is provided through a Secret. - 04-cpt runs DeepSpeed ZeRO-3 continued pretraining under Ray Train. - 06-grpo runs the trl vllm-serve server and the GRPO trainer in one RayJob with NCCL weight sync and a GnuCOBOL compile reward. - 07-eval scores five benchmarks with vLLM and format-aware cobc checks. The README documents data licensing, placeholders, stage commands, and the reference results (64/220 compile). Serving is not included. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/README.md | 1 + examples/cobol-cpt-grpo/.dockerignore | 2 + .../01-collect/collect-cobol-data.sh | 228 +++++ .../01-collect/collect-gated.py | 222 +++++ .../01-collect/collect-rl-data.sh | 342 +++++++ .../cobol-cpt-grpo/01-collect/tau-gated.yaml | 31 + .../cobol-cpt-grpo/01-collect/tau-rl.yaml | 29 + examples/cobol-cpt-grpo/01-collect/tau.yaml | 29 + .../02-data-prep/prepare_data.py | 616 +++++++++++++ examples/cobol-cpt-grpo/02-data-prep/tau.yaml | 30 + .../03-benchmarks/build_benchmarks.py | 428 +++++++++ .../cobol-cpt-grpo/03-benchmarks/tau.yaml | 30 + .../04-cpt/continued_pretrain.py | 444 +++++++++ examples/cobol-cpt-grpo/04-cpt/driver.py | 65 ++ examples/cobol-cpt-grpo/04-cpt/tau.yaml | 65 ++ .../04-cpt/test_continued_pretrain.py | 264 ++++++ examples/cobol-cpt-grpo/04-cpt/test_driver.py | 125 +++ .../05-rl-prompts/prepare_rl_prompts.py | 363 ++++++++ .../cobol-cpt-grpo/05-rl-prompts/tau.yaml | 34 + examples/cobol-cpt-grpo/06-grpo/driver.py | 202 +++++ examples/cobol-cpt-grpo/06-grpo/tau.yaml | 59 ++ .../cobol-cpt-grpo/06-grpo/test_driver.py | 272 ++++++ .../cobol-cpt-grpo/06-grpo/test_reward.py | 718 +++++++++++++++ examples/cobol-cpt-grpo/06-grpo/train_grpo.py | 851 ++++++++++++++++++ examples/cobol-cpt-grpo/07-eval/driver.py | 36 + .../cobol-cpt-grpo/07-eval/run_eval_vllm.py | 612 +++++++++++++ examples/cobol-cpt-grpo/07-eval/tau.yaml | 42 + .../cobol-cpt-grpo/07-eval/test_driver.py | 73 ++ examples/cobol-cpt-grpo/Dockerfile | 15 + examples/cobol-cpt-grpo/README.md | 236 +++++ .../cobol-cpt-grpo/image-requirements.txt | 32 + 31 files changed, 6496 insertions(+) create mode 100644 examples/cobol-cpt-grpo/.dockerignore create mode 100755 examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh create mode 100755 examples/cobol-cpt-grpo/01-collect/collect-gated.py create mode 100755 examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh create mode 100644 examples/cobol-cpt-grpo/01-collect/tau-gated.yaml create mode 100644 examples/cobol-cpt-grpo/01-collect/tau-rl.yaml create mode 100644 examples/cobol-cpt-grpo/01-collect/tau.yaml create mode 100755 examples/cobol-cpt-grpo/02-data-prep/prepare_data.py create mode 100644 examples/cobol-cpt-grpo/02-data-prep/tau.yaml create mode 100755 examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py create mode 100644 examples/cobol-cpt-grpo/03-benchmarks/tau.yaml create mode 100644 examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py create mode 100644 examples/cobol-cpt-grpo/04-cpt/driver.py create mode 100644 examples/cobol-cpt-grpo/04-cpt/tau.yaml create mode 100644 examples/cobol-cpt-grpo/04-cpt/test_continued_pretrain.py create mode 100644 examples/cobol-cpt-grpo/04-cpt/test_driver.py create mode 100755 examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py create mode 100644 examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml create mode 100644 examples/cobol-cpt-grpo/06-grpo/driver.py create mode 100644 examples/cobol-cpt-grpo/06-grpo/tau.yaml create mode 100644 examples/cobol-cpt-grpo/06-grpo/test_driver.py create mode 100644 examples/cobol-cpt-grpo/06-grpo/test_reward.py create mode 100755 examples/cobol-cpt-grpo/06-grpo/train_grpo.py create mode 100644 examples/cobol-cpt-grpo/07-eval/driver.py create mode 100755 examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py create mode 100644 examples/cobol-cpt-grpo/07-eval/tau.yaml create mode 100644 examples/cobol-cpt-grpo/07-eval/test_driver.py create mode 100644 examples/cobol-cpt-grpo/Dockerfile create mode 100644 examples/cobol-cpt-grpo/README.md create mode 100644 examples/cobol-cpt-grpo/image-requirements.txt diff --git a/examples/README.md b/examples/README.md index 293ff395..8f98109d 100644 --- a/examples/README.md +++ b/examples/README.md @@ -46,6 +46,7 @@ Standalone Tau example pages should include the same status block near the top: | [`aks-cpu-quickstart`](./aks-cpu-quickstart/) | `operator runbook` | Creates a new CPU-only AKS cluster and runs `az aks create`, `tau cluster install`, a real CPU PyTorch RayJob, and live Stellar metrics. Creates billable Azure resources. | | [`aks-gpu-quickstart`](./aks-gpu-quickstart/) | `operator runbook` | GPU counterpart to `aks-cpu-quickstart` on a single A100 node, including the AKS device plugin and MIG configuration issues. Verifies CUDA execution with evidence a CPU cannot produce. Creates expensive billable Azure resources. | | [`market-policy`](./market-policy/) | `production-shaped` | Trains the exact compact actor-critic rendered in the documentation, requires CUDA on TauGrid, and exports browser-ready policy/value weights. | +| [`cobol-cpt-grpo`](./cobol-cpt-grpo/) | `production-shaped` | Full COBOL pipeline on 2 x 8 H200 with RDMA: CPU data collection and prep Jobs, ZeRO-3 continued pretraining under Ray Train, GRPO with a GnuCOBOL compile reward and an in-RayJob vLLM server, and a compile-checked vLLM eval. Requires a self-built training image and a writable `/data` PVC. | | [`kind-smoke`](./kind-smoke/) | `smoke/debug` | Verify the local Tau -> Kueue -> Kubernetes Job path on a disposable Kind cluster before trying AKS/GPU-specific prepare and storage flows. | | [`nanogpt-ray`](./nanogpt-ray/) | `smoke/debug` for local and small GPU runs; `production-shaped` for H200 target routes | Start with the 1-GPU smoke for onboarding, then use the pre-tokenized H200 configs for FineWeb target/regression work. | | [`ray-tune-smoke`](./ray-tune-smoke/) | `smoke/debug` | Minimal Ray Tune HPO path: the researcher supplies `train_func(config)` and a search space, and Tau generates the Tuner/TorchTrainer wrapper. Use the [canonical AKS walkthrough](../site/content/en/docs/examples/gpu-ray-tune.md) for provider setup and the repository-first handoff. | diff --git a/examples/cobol-cpt-grpo/.dockerignore b/examples/cobol-cpt-grpo/.dockerignore new file mode 100644 index 00000000..d464b260 --- /dev/null +++ b/examples/cobol-cpt-grpo/.dockerignore @@ -0,0 +1,2 @@ +* +!image-requirements.txt diff --git a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh new file mode 100755 index 00000000..7d69f241 --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh @@ -0,0 +1,228 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +set -uo pipefail + +DATA_ROOT="/data/datasets/cobol" +RAW="${DATA_ROOT}/raw" +PROCESSED="${DATA_ROOT}/processed" +LOGS="${DATA_ROOT}/logs" +mkdir -p "${RAW}" "${PROCESSED}" "${LOGS}" + +log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*" | tee -a "${LOGS}/collect.log"; } +run_step() { + local name="$1"; shift + log "--- ${name} ---" + if "$@"; then + log "${name}: OK" + else + log "${name}: FAILED (continuing)" + fi +} + +log "=== COBOL data collection starting ===" +log "Target: ${DATA_ROOT}" +SOURCE_LICENSES="${SOURCE_LICENSES:-permissive}" +log "Source licenses: ${SOURCE_LICENSES}" + +pip install -q huggingface_hub==1.28.0 datasets==5.0.1 pandas==3.0.6 pyarrow==25.0.1 2>&1 | tail -1 +log "pip install done" + +collect_troy_cobol() { + python3 - <<'PYEOF' +import os, json +from datasets import load_dataset +out = os.environ.get("RAW", "/data/datasets/cobol/raw") + "/troy-cobol" +os.makedirs(out, exist_ok=True) +ds = load_dataset("TroyDoesAI/Cobol", split="train", trust_remote_code=True) +for i, row in enumerate(ds): + path = row.get("path", f"file_{i}.cbl") + fname = f"{i:06d}_{os.path.basename(path)}" + with open(os.path.join(out, fname), "w") as f: + f.write(row["content"]) +print(f"Saved {len(ds)} files to {out}") +with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "TroyDoesAI/Cobol", "count": len(ds)}, f) +PYEOF +} +run_step "TroyDoesAI/Cobol (2,980 files)" collect_troy_cobol + +collect_commitpack() { + python3 - <<'PYEOF' +import os, json, shutil +from huggingface_hub import hf_hub_download +out = os.environ.get("RAW", "/data/datasets/cobol/raw") + "/commitpack" +os.makedirs(out, exist_ok=True) +local = hf_hub_download("bigcode/commitpack", "data/cobol/cobol-0001.jsonl", repo_type="dataset") +shutil.copy(local, os.path.join(out, "commitpack-cobol.jsonl")) +count = sum(1 for _ in open(os.path.join(out, "commitpack-cobol.jsonl"))) +print(f"Saved {count} commit records to {out}") +PYEOF +} +run_step "bigcode/commitpack COBOL (24,953 commits)" collect_commitpack + +collect_mainframebench() { + python3 - <<'PYEOF' +import os, json +from datasets import load_dataset +out = os.environ.get("RAW", "/data/datasets/cobol/raw") + "/mainframebench" +os.makedirs(out, exist_ok=True) +for split_name in ["question_answering", "multiple_choice_question", "COBOL_code_summarization"]: + try: + ds = load_dataset("Fsoft-AIC/MainframeBench", split_name, split="train", trust_remote_code=True) + with open(os.path.join(out, f"{split_name}.jsonl"), "w") as f: + for row in ds: + f.write(json.dumps(dict(row)) + "\n") + print(f" {split_name}: {len(ds)} samples") + except Exception as e: + print(f" {split_name}: skipped ({e})") +PYEOF +} +run_step "MainframeBench (7,052 samples)" collect_mainframebench + +collect_eval_benchmarks() { + python3 - <<'PYEOF' +import os, json +from datasets import load_dataset +raw = os.environ.get("RAW", "/data/datasets/cobol/raw") +for name, hf_id in [ + ("coboleval", "msplits/human-eval-cobol"), + ("cobolcodebench", "harshini-kumar/CobolCodeBench"), +]: + out = os.path.join(raw, name) + os.makedirs(out, exist_ok=True) + try: + splits = ["default", "train"] + ds = None + for s in splits: + try: + ds = load_dataset(hf_id, split=s, trust_remote_code=True) + break + except Exception: + continue + if ds is None: + ds = load_dataset(hf_id, trust_remote_code=True) + first_split = list(ds.keys())[0] + ds = ds[first_split] + with open(os.path.join(out, f"{name}.jsonl"), "w") as f: + for row in ds: + f.write(json.dumps(dict(row)) + "\n") + print(f"{name}: {len(ds)} samples") + except Exception as e: + print(f"{name}: error - {e}") +PYEOF +} +run_step "Evaluation benchmarks" collect_eval_benchmarks + +collect_xcobol() { + mkdir -p "${RAW}/x-cobol" + curl -fsSL "https://zenodo.org/api/records/14269462/files-archive" -o "${RAW}/x-cobol/x-cobol-v2.zip" + cd "${RAW}/x-cobol" && unzip -qo x-cobol-v2.zip -d extracted + find "${RAW}/x-cobol/extracted" -name '*.zip' -exec sh -c 'unzip -qo "$1" -d "$(dirname "$1")/$(basename "$1" .zip)"' _ {} \; + XCOUNT=$(find ${RAW}/x-cobol/extracted -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) + log "X-COBOL: ${XCOUNT} COBOL files" + if [ "${XCOUNT}" -eq 0 ]; then + log "X-COBOL: listing top-level contents for debugging:" + find "${RAW}/x-cobol/extracted" -maxdepth 3 -type f 2>/dev/null | head -30 | tee -a "${LOGS}/collect.log" + fi +} +run_step "X-COBOL v2 from Zenodo" collect_xcobol + +log "--- GitHub repos ---" +GITHUB="${RAW}/github" +mkdir -p "${GITHUB}" + +clone_repo() { + local url="$1" dir="$2" + if [ ! -d "${GITHUB}/${dir}/.git" ]; then + log " cloning ${dir}..." + git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1 || log " WARN: failed to clone ${dir}" + else + log " ${dir} already exists, skipping" + fi +} + +clone_repo "https://github.com/aws-samples/aws-mainframe-modernization-carddemo.git" "carddemo" +clone_repo "https://github.com/meyfa/CobolCraft.git" "cobolcraft" +clone_repo "https://github.com/openmainframeproject/cobol-programming-course.git" "omp-course" +clone_repo "https://github.com/DillonDepeel/Cobol-Programming-Collection.git" "cobol-collection" +clone_repo "https://github.com/shamrice/COBOL-Examples.git" "cobol-examples" +clone_repo "https://github.com/krisds/koopa.git" "koopa" +clone_repo "https://github.com/bushidocodes/pc-pricer.git" "cms-pricer" +clone_repo "https://github.com/exercism/cobol.git" "exercism-cobol" + +if [ "${SOURCE_LICENSES}" = "all" ]; then + clone_repo "https://github.com/OCamlPro/gnucobol-contrib.git" "gnucobol-contrib" + clone_repo "https://github.com/opensourcecobol/opensourcecobol4j.git" "opensourcecobol4j" + clone_repo "https://github.com/Martinfx/Cobol.git" "martinfx-cobol" + clone_repo "https://github.com/etalab/taxe-fonciere.git" "taxe-fonciere" + clone_repo "https://github.com/josuerochadev/poei-cobol.git" "poei-cobol" + clone_repo "https://github.com/cchipman21804/EnterpriseCOBOLv6.3.git" "enterprise-cobol" + clone_repo "https://github.com/Apress/beg-cobol-for-programmers.git" "apress-cobol" + clone_repo "https://github.com/andrea-git/GESLUX.git" "geslux" + clone_repo "https://github.com/z390development/z390.git" "z390" + clone_repo "https://github.com/writ3it/cobol-examples.git" "writ3it-examples" + clone_repo "https://github.com/fgregg/tax_extension.git" "cook-county-tax" + clone_repo "https://github.com/sentientsergio/COBOL-Legacy-Benchmark-Suite.git" "legacy-benchmark" +else + log " skipping copyleft and unlicensed repos (set SOURCE_LICENSES=all to include them)" +fi + +TOTAL_GH=$(find "${GITHUB}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' -o -iname '*.CBL' -o -iname '*.COB' \) 2>/dev/null | wc -l) +log "GitHub repos: ${TOTAL_GH} COBOL files total" + +collect_nist() { + mkdir -p "${RAW}/nist" + if [ ! -f "${RAW}/nist/newcob.val" ]; then + curl -L --max-redirs 5 "https://sourceforge.net/projects/gnucobol/files/nist/nist_2013-10-10.tar.gz/download" \ + -o "${RAW}/nist/nist.tar.gz" + if file "${RAW}/nist/nist.tar.gz" | grep -q "gzip"; then + cd "${RAW}/nist" && tar xzf nist.tar.gz + else + log "NIST: downloaded file is not a valid gzip (likely redirect page), trying alternate URL" + curl -L "https://master.dl.sourceforge.net/project/gnucobol/nist/nist_2013-10-10.tar.gz" \ + -o "${RAW}/nist/nist.tar.gz" + cd "${RAW}/nist" && tar xzf nist.tar.gz + fi + fi +} +run_step "NIST COBOL85 test suite" collect_nist + +log "=== Collection complete ===" +log "Disk usage:" +du -sh "${RAW}"/* 2>/dev/null | tee -a "${LOGS}/collect.log" +log "Total raw:" +du -sh "${RAW}" | tee -a "${LOGS}/collect.log" + +TOTAL_COBOL=$(find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) | wc -l) +TOTAL_JSONL=$(find "${RAW}" -type f -name '*.jsonl' | wc -l) +log "Total COBOL source files: ${TOTAL_COBOL}" +log "Total JSONL data files: ${TOTAL_JSONL}" + +cat > "${DATA_ROOT}/manifest.json" < "${DATA_ROOT}/source-file-index.txt" +find "${RAW}" -type f -name '*.jsonl' >> "${DATA_ROOT}/source-file-index.txt" +INDEX_COUNT=$(wc -l < "${DATA_ROOT}/source-file-index.txt") +log "File index written: ${INDEX_COUNT} entries" +log "=== Done ===" diff --git a/examples/cobol-cpt-grpo/01-collect/collect-gated.py b/examples/cobol-cpt-grpo/01-collect/collect-gated.py new file mode 100755 index 00000000..4c1feaff --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-gated.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Download gated HF datasets: The Stack v1 COBOL, The Stack v2 COBOL, cobol-pretraining-corpus.""" + +import json +import os +import subprocess +import sys + +subprocess.check_call( + [ + sys.executable, + "-m", + "pip", + "install", + "-q", + "huggingface_hub==1.28.0", + "datasets==5.0.1", + "pandas==3.0.6", + "pyarrow==25.0.1", + ] +) + +from datasets import load_dataset # noqa: E402 +from huggingface_hub import HfApi # noqa: E402 + +DATA_ROOT = os.environ.get("DATA_ROOT", "/data/datasets/cobol") +RAW = os.path.join(DATA_ROOT, "raw") +os.makedirs(RAW, exist_ok=True) + +token = os.environ["HF_TOKEN"] + +api = HfApi(token=token) +print(f"Authenticated as: {api.whoami()['name']}", flush=True) + + +def collect_the_stack_v1(): + out = os.path.join(RAW, "the-stack-v1-cobol") + os.makedirs(out, exist_ok=True) + print("\n=== The Stack v1 — COBOL subset ===", flush=True) + + ds = load_dataset( + "bigcode/the-stack", data_dir="data/cobol", split="train", token=token + ) + print(f" Loaded {len(ds)} records", flush=True) + + jsonl_path = os.path.join(out, "the-stack-v1-cobol.jsonl") + count = 0 + with open(jsonl_path, "w") as f: + for row in ds: + f.write( + json.dumps( + { + "hexsha": row.get("hexsha", ""), + "path": row.get("max_stars_repo_path", row.get("path", "")), + "repo": row.get("max_stars_repo_name", ""), + "content": row.get("content", ""), + "size": row.get("size", 0), + "language": "COBOL", + "source": "bigcode/the-stack", + } + ) + + "\n" + ) + count += 1 + + print(f" Saved {count} records to {jsonl_path}", flush=True) + + with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump( + {"source": "bigcode/the-stack", "subset": "cobol", "count": count}, + f, + indent=2, + ) + + src_count = 0 + src_dir = os.path.join(out, "source") + os.makedirs(src_dir, exist_ok=True) + ds2 = load_dataset( + "bigcode/the-stack", data_dir="data/cobol", split="train", token=token + ) + for i, row in enumerate(ds2): + content = row.get("content", "") + if not content.strip(): + continue + path = row.get("max_stars_repo_path", row.get("path", f"file_{i}.cbl")) + ext = os.path.splitext(path)[1].lower() + if ext not in (".cbl", ".cob", ".cpy", ".CBL", ".COB", ".CPY"): + ext = ".cbl" + fname = f"{i:06d}_{os.path.basename(path)}" + if not fname.lower().endswith((".cbl", ".cob", ".cpy")): + fname += ext + with open(os.path.join(src_dir, fname), "w") as f: + f.write(content) + src_count += 1 + + print(f" Extracted {src_count} source files to {src_dir}", flush=True) + + +def collect_the_stack_v2(): + out = os.path.join(RAW, "the-stack-v2-cobol") + os.makedirs(out, exist_ok=True) + print("\n=== The Stack v2 — COBOL subset ===", flush=True) + + try: + ds = load_dataset( + "bigcode/the-stack-v2", + data_dir="data/cobol", + split="train", + token=token, + trust_remote_code=True, + ) + print(f" Loaded {len(ds)} records", flush=True) + + jsonl_path = os.path.join(out, "the-stack-v2-cobol.jsonl") + count = 0 + with open(jsonl_path, "w") as f: + for row in ds: + record = { + k: v + for k, v in row.items() + if isinstance(v, (str, int, float, bool)) + } + record["source"] = "bigcode/the-stack-v2" + record["language"] = "COBOL" + f.write(json.dumps(record) + "\n") + count += 1 + + print(f" Saved {count} records to {jsonl_path}", flush=True) + + has_content = any( + row.get("content") + for row in load_dataset( + "bigcode/the-stack-v2", + data_dir="data/cobol", + split="train[:5]", + token=token, + trust_remote_code=True, + ) + ) + if has_content: + src_dir = os.path.join(out, "source") + os.makedirs(src_dir, exist_ok=True) + src_count = 0 + for i, row in enumerate(ds): + content = row.get("content", "") + if not content or not content.strip(): + continue + fname = f"{i:06d}.cbl" + with open(os.path.join(src_dir, fname), "w") as f: + f.write(content) + src_count += 1 + print(f" Extracted {src_count} source files", flush=True) + else: + print( + " Note: v2 contains SWHIDs only, no source content (needs Software Heritage S3)", + flush=True, + ) + + with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump( + { + "source": "bigcode/the-stack-v2", + "subset": "cobol", + "count": count, + "has_content": has_content, + }, + f, + indent=2, + ) + + except Exception as e: + print(f" Failed: {e}", flush=True) + with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "bigcode/the-stack-v2", "error": str(e)}, f, indent=2) + + +def collect_pretraining_corpus(): + out = os.path.join(RAW, "cobol-pretraining-corpus") + os.makedirs(out, exist_ok=True) + print("\n=== 4sp1d3r2/cobol-pretraining-corpus ===", flush=True) + + ds = load_dataset("4sp1d3r2/cobol-pretraining-corpus", split="train", token=token) + print(f" Loaded {len(ds)} records", flush=True) + print(f" Columns: {ds.column_names}", flush=True) + + jsonl_path = os.path.join(out, "cobol-pretraining-corpus.jsonl") + count = 0 + with open(jsonl_path, "w") as f: + for row in ds: + record = { + k: v for k, v in row.items() if isinstance(v, (str, int, float, bool)) + } + record["source"] = "4sp1d3r2/cobol-pretraining-corpus" + f.write(json.dumps(record) + "\n") + count += 1 + + print(f" Saved {count} records to {jsonl_path}", flush=True) + + with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump( + { + "source": "4sp1d3r2/cobol-pretraining-corpus", + "count": count, + "columns": ds.column_names, + }, + f, + indent=2, + ) + + +if __name__ == "__main__": + print("=== Gated COBOL dataset collection ===", flush=True) + print(f" Output: {RAW}", flush=True) + + collect_the_stack_v1() + collect_the_stack_v2() + collect_pretraining_corpus() + + print("\n=== Collection complete ===", flush=True) diff --git a/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh new file mode 100755 index 00000000..cbcc0fa3 --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh @@ -0,0 +1,342 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +set -uo pipefail + +DATA_ROOT="/data/datasets/cobol" +RAW="${DATA_ROOT}/raw" +LOGS="${DATA_ROOT}/logs" +mkdir -p "${RAW}" "${LOGS}" + +log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*" | tee -a "${LOGS}/collect-rl.log"; } +run_step() { + local name="$1"; shift + log "--- ${name} ---" + if "$@"; then + log "${name}: OK" + else + log "${name}: FAILED (continuing)" + fi +} + +log "=== RL data collection starting ===" +log "HF_TOKEN set: $([ -n "${HF_TOKEN:-}" ] && echo yes || echo NO)" +SOURCE_LICENSES="${SOURCE_LICENSES:-permissive}" +log "Source licenses: ${SOURCE_LICENSES}" + +pip install -q huggingface_hub==1.28.0 datasets==5.0.1 pandas==3.0.6 pyarrow==25.0.1 2>&1 | tail -3 +log "pip install done" + +collect_the_stack_v1() { + local out="${RAW}/the-stack-v1-cobol" + [ -f "${out}/metadata.json" ] && { log "Stack v1 already collected, skipping"; return 0; } + python3 - <<'PYEOF' +import json, os +from datasets import load_dataset + +token = os.environ["HF_TOKEN"] +out = os.environ["RAW"] + "/the-stack-v1-cobol" +os.makedirs(out, exist_ok=True) +src_dir = os.path.join(out, "source") +os.makedirs(src_dir, exist_ok=True) + +print("Loading bigcode/the-stack COBOL subset...", flush=True) +ds = load_dataset("bigcode/the-stack", data_dir="data/cobol", split="train", token=token) +print(f" Loaded {len(ds)} records", flush=True) + +count = 0 +for i, row in enumerate(ds): + content = row.get("content", "") + if not content.strip(): + continue + path = row.get("max_stars_repo_path", row.get("path", f"file_{i}.cbl")) + ext = os.path.splitext(path)[1].lower() + if ext not in (".cbl", ".cob", ".cpy"): + ext = ".cbl" + fname = f"{i:06d}_{os.path.basename(path)}" + if not fname.lower().endswith((".cbl", ".cob", ".cpy")): + fname += ext + with open(os.path.join(src_dir, fname), "w") as f: + f.write(content) + count += 1 + +with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "bigcode/the-stack", "subset": "cobol", "count": count}, f, indent=2) +print(f" Saved {count} source files to {src_dir}", flush=True) +PYEOF +} +if [ -n "${HF_TOKEN:-}" ]; then + run_step "The Stack v1 COBOL (~4K files)" collect_the_stack_v1 +else + log "HF_TOKEN unset: skipping gated dataset the_stack_v1" +fi + +collect_the_stack_v2_dedup() { + local out="${RAW}/the-stack-v2-dedup-cobol" + [ -f "${out}/metadata.json" ] && { log "Stack v2-dedup already collected, skipping"; return 0; } + python3 - <<'PYEOF' +import json, os +from datasets import load_dataset + +token = os.environ["HF_TOKEN"] +out = os.environ["RAW"] + "/the-stack-v2-dedup-cobol" +os.makedirs(out, exist_ok=True) +src_dir = os.path.join(out, "source") +os.makedirs(src_dir, exist_ok=True) + +print("Loading bigcode/the-stack-v2-dedup COBOL...", flush=True) +try: + ds = load_dataset("bigcode/the-stack-v2-dedup", "COBOL", split="train", + token=token, trust_remote_code=True) +except Exception as e1: + print(f" Config 'COBOL' failed ({e1}), trying data_dir...", flush=True) + try: + ds = load_dataset("bigcode/the-stack-v2-dedup", data_dir="data/COBOL", + split="train", token=token, trust_remote_code=True) + except Exception as e2: + print(f" data_dir 'data/COBOL' failed ({e2}), trying lowercase...", flush=True) + ds = load_dataset("bigcode/the-stack-v2-dedup", data_dir="data/cobol", + split="train", token=token, trust_remote_code=True) + +print(f" Loaded {len(ds)} records", flush=True) +print(f" Columns: {ds.column_names}", flush=True) + +has_content = "content" in ds.column_names +count = 0 + +if has_content: + for i, row in enumerate(ds): + content = row.get("content", "") + if not content or not content.strip(): + continue + fname = f"{i:06d}.cbl" + with open(os.path.join(src_dir, fname), "w") as f: + f.write(content) + count += 1 + print(f" Extracted {count} source files", flush=True) +else: + jsonl_path = os.path.join(out, "the-stack-v2-dedup-cobol.jsonl") + with open(jsonl_path, "w") as f: + for row in ds: + record = {k: v for k, v in row.items() if isinstance(v, (str, int, float, bool))} + f.write(json.dumps(record) + "\n") + count += 1 + print(f" No 'content' column — saved {count} metadata records (SWHIDs only)", flush=True) + +with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "bigcode/the-stack-v2-dedup", "subset": "COBOL", + "count": count, "has_content": has_content}, f, indent=2) +PYEOF +} +if [ -n "${HF_TOKEN:-}" ]; then + run_step "The Stack v2-dedup COBOL (~36K files)" collect_the_stack_v2_dedup +else + log "HF_TOKEN unset: skipping gated dataset the_stack_v2_dedup" +fi + +collect_pretraining_corpus() { + local out="${RAW}/cobol-pretraining-corpus" + [ -f "${out}/metadata.json" ] && { log "Pretraining corpus already collected, skipping"; return 0; } + python3 - <<'PYEOF' +import json, os +from datasets import load_dataset + +token = os.environ["HF_TOKEN"] +out = os.environ["RAW"] + "/cobol-pretraining-corpus" +os.makedirs(out, exist_ok=True) +src_dir = os.path.join(out, "source") +os.makedirs(src_dir, exist_ok=True) + +print("Loading 4sp1d3r2/cobol-pretraining-corpus...", flush=True) +ds = load_dataset("4sp1d3r2/cobol-pretraining-corpus", split="train", token=token) +print(f" Loaded {len(ds)} records", flush=True) +print(f" Columns: {ds.column_names}", flush=True) + +code_count = 0 +web_count = 0 +for i, row in enumerate(ds): + text = row.get("text", "") + source = row.get("source", "") + if not text or not text.strip(): + continue + is_code = any(kw in text[:500].upper() for kw in + ["IDENTIFICATION DIVISION", "PROCEDURE DIVISION", "DATA DIVISION", + "PROGRAM-ID", "WORKING-STORAGE", "01 ", "PERFORM ", "MOVE "]) + if is_code: + fname = f"{i:06d}.cbl" + with open(os.path.join(src_dir, fname), "w") as f: + f.write(text) + code_count += 1 + else: + web_count += 1 + +print(f" Extracted {code_count} code files, skipped {web_count} web/educational text", flush=True) +with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "4sp1d3r2/cobol-pretraining-corpus", + "total": len(ds), "code_files": code_count, "web_skipped": web_count}, f, indent=2) +PYEOF +} +if [ -n "${HF_TOKEN:-}" ]; then + run_step "cobol-pretraining-corpus (~217K rows, extracting code only)" collect_pretraining_corpus +else + log "HF_TOKEN unset: skipping gated dataset pretraining_corpus" +fi + +collect_zorse() { + local out="${RAW}/zorse" + [ -f "${out}/metadata.json" ] && { log "zorse already collected, skipping"; return 0; } + python3 - <<'PYEOF' +import json, os +from datasets import load_dataset + +out = os.environ["RAW"] + "/zorse" +os.makedirs(out, exist_ok=True) +src_dir = os.path.join(out, "source") +os.makedirs(src_dir, exist_ok=True) + +print("Loading zorse/zorse...", flush=True) +try: + ds = load_dataset("zorse/zorse", split="train", trust_remote_code=True) +except Exception: + try: + ds_dict = load_dataset("zorse/zorse", trust_remote_code=True) + first_split = list(ds_dict.keys())[0] + ds = ds_dict[first_split] + except Exception as e: + print(f" Failed to load zorse/zorse: {e}", flush=True) + with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "zorse/zorse", "error": str(e)}, f, indent=2) + raise + +print(f" Loaded {len(ds)} records", flush=True) +print(f" Columns: {ds.column_names}", flush=True) + +content_col = None +for col in ["content", "text", "code", "source_code", "program"]: + if col in ds.column_names: + content_col = col + break + +count = 0 +if content_col: + for i, row in enumerate(ds): + text = row.get(content_col, "") + if not text or not text.strip(): + continue + fname = f"{i:06d}.cbl" + with open(os.path.join(src_dir, fname), "w") as f: + f.write(text) + count += 1 + print(f" Extracted {count} source files (column: {content_col})", flush=True) +else: + jsonl_path = os.path.join(out, "zorse.jsonl") + with open(jsonl_path, "w") as f: + for row in ds: + record = {k: v for k, v in row.items() if isinstance(v, (str, int, float, bool))} + f.write(json.dumps(record) + "\n") + count += 1 + print(f" No content column found, saved {count} records as JSONL", flush=True) + print(f" Available columns: {ds.column_names}", flush=True) + +with open(os.path.join(out, "metadata.json"), "w") as f: + json.dump({"source": "zorse/zorse", "count": count, + "content_column": content_col, "columns": ds.column_names}, f, indent=2) +PYEOF +} +if [ "${SOURCE_LICENSES}" = "all" ]; then + run_step "zorse/zorse (~40K COBOL rows)" collect_zorse +else + log "skipping zorse/zorse: no stated license (set SOURCE_LICENSES=all to include it)" +fi + +log "--- Large GitHub parser repos ---" +GITHUB="${RAW}/github" +mkdir -p "${GITHUB}" + +clone_repo() { + local url="$1" dir="$2" + if [ -d "${GITHUB}/${dir}" ] && [ "$(find "${GITHUB}/${dir}" -maxdepth 1 -type f 2>/dev/null | head -1)" ]; then + log " ${dir} already exists, skipping" + return 0 + fi + log " cloning ${dir}..." + git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1 || log " WARN: failed to clone ${dir}" +} + +clone_repo "https://github.com/uwol/proleap-cobol-parser.git" "proleap" +clone_repo "https://github.com/avishek-sen-gupta/cobol-rekt.git" "cobol-rekt" +clone_repo "https://github.com/mrm413/Ironclad-COBOL-to-Rust.git" "ironclad" +clone_repo "https://github.com/navikt/DSF.git" "dsf-norway" +clone_repo "https://github.com/IBM/Bank-of-Z.git" "bank-of-z" +clone_repo "https://github.com/icitry/FPS.cob.git" "fps-cobol" + +if [ "${SOURCE_LICENSES}" = "all" ]; then + clone_repo "https://github.com/eclipse-che4z/che-che4z-lsp-for-cobol.git" "che4z-lsp" + clone_repo "https://github.com/neopragma/cobol-unit-test.git" "cobol-unit-test" + clone_repo "https://github.com/mridoni/gixsql.git" "gixsql" + clone_repo "https://github.com/JohnDovey/GNUCobol-Samples.git" "gnucobol-samples" +else + log " skipping copyleft repos (set SOURCE_LICENSES=all to include them)" +fi + +NEW_GH=$(find "${GITHUB}/che4z-lsp" "${GITHUB}/proleap" "${GITHUB}/cobol-rekt" "${GITHUB}/ironclad" "${GITHUB}/cobol-unit-test" "${GITHUB}/dsf-norway" "${GITHUB}/bank-of-z" "${GITHUB}/gixsql" "${GITHUB}/gnucobol-samples" "${GITHUB}/fps-cobol" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' -o -iname '*.CBL' -o -iname '*.COB' \) 2>/dev/null | wc -l) +log "New GitHub parser repos: ${NEW_GH} COBOL files" + +collect_nist_from_oss4j() { + local nist_dir="${RAW}/nist-extracted" + [ -d "${nist_dir}" ] && [ "$(find "${nist_dir}" -name '*.CBL' -o -name '*.cbl' 2>/dev/null | head -1)" ] && { + log "NIST already extracted, skipping" + return 0 + } + mkdir -p "${nist_dir}" + + local oss4j="${GITHUB}/opensourcecobol4j" + if [ ! -d "${oss4j}" ]; then + log "opensourcecobol4j not found, cloning..." + git clone --depth 1 --single-branch "https://github.com/opensourcecobol/opensourcecobol4j.git" "${oss4j}" 2>&1 | tail -1 + fi + + local nist_src + nist_src=$(find "${oss4j}" -type d -name "CCVS85" -o -name "nist" -o -name "cobol85" 2>/dev/null | head -1) + if [ -z "${nist_src}" ]; then + nist_src=$(find "${oss4j}" -type f -iname '*.CBL' -path '*test*' 2>/dev/null | head -1) + nist_src=$(dirname "${nist_src:-/nonexistent}") + fi + + if [ -d "${nist_src}" ]; then + find "${nist_src}" -type f \( -iname '*.cbl' -o -iname '*.CBL' -o -iname '*.cob' -o -iname '*.cpy' \) -exec cp {} "${nist_dir}/" \; + local count + count=$(find "${nist_dir}" -type f 2>/dev/null | wc -l) + log "NIST extracted: ${count} files from ${nist_src}" + else + log "NIST: could not find test directory in opensourcecobol4j" + fi +} +if [ "${SOURCE_LICENSES}" = "all" ]; then + run_step "NIST test suite (from opensourcecobol4j)" collect_nist_from_oss4j +else + log "skipping NIST extraction from opensourcecobol4j (set SOURCE_LICENSES=all to include it)" +fi + +log "=== RL data collection complete ===" +log "New data disk usage:" +for d in the-stack-v1-cobol the-stack-v2-dedup-cobol cobol-pretraining-corpus zorse; do + du -sh "${RAW}/${d}" 2>/dev/null | tee -a "${LOGS}/collect-rl.log" || true +done +for d in che4z-lsp proleap cobol-rekt ironclad cobol-unit-test dsf-norway bank-of-z gixsql gnucobol-samples fps-cobol nist-extracted; do + du -sh "${GITHUB}/${d}" 2>/dev/null || du -sh "${RAW}/${d}" 2>/dev/null || true +done | tee -a "${LOGS}/collect-rl.log" + +TOTAL_NEW=$(find "${RAW}/the-stack-v1-cobol" "${RAW}/the-stack-v2-dedup-cobol" "${RAW}/cobol-pretraining-corpus" "${RAW}/zorse" "${GITHUB}/che4z-lsp" "${GITHUB}/proleap" "${GITHUB}/cobol-rekt" "${GITHUB}/ironclad" "${GITHUB}/cobol-unit-test" "${GITHUB}/dsf-norway" "${GITHUB}/bank-of-z" "${GITHUB}/gixsql" "${GITHUB}/gnucobol-samples" "${GITHUB}/fps-cobol" "${RAW}/nist-extracted" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) +log "Total NEW COBOL source files: ${TOTAL_NEW}" + +TOTAL_ALL=$(find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) +log "Total ALL COBOL source files (old + new): ${TOTAL_ALL}" + +log "Rebuilding file index..." +find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) > "${DATA_ROOT}/source-file-index.txt" +find "${RAW}" -type f -name '*.jsonl' >> "${DATA_ROOT}/source-file-index.txt" +INDEX_COUNT=$(wc -l < "${DATA_ROOT}/source-file-index.txt") +log "File index rebuilt: ${INDEX_COUNT} entries" +log "=== Done ===" diff --git a/examples/cobol-cpt-grpo/01-collect/tau-gated.yaml b/examples/cobol-cpt-grpo/01-collect/tau-gated.yaml new file mode 100644 index 00000000..be4b31ae --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/tau-gated.yaml @@ -0,0 +1,31 @@ +name: cobol-collect-gated +engine: job +entrypoint: collect-gated.py + +experiment: + project: cobol-cpt-grpo + name: collect-gated + group: data + +runtime: + image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a + env: + DATA_ROOT: /data/datasets/cobol + HF_HOME: /tmp/hf-cache + PYTHONUNBUFFERED: "1" + env_secret: + HF_TOKEN: hf-token:token + +compute: + gpus: 0 + cpu_request: "4" + memory_request: 16Gi + cpu_limit: "8" + memory_limit: 32Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/collect-gated + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml b/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml new file mode 100644 index 00000000..b7f95cb7 --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml @@ -0,0 +1,29 @@ +name: cobol-collect-rl +engine: job +entrypoint: collect-rl-data.sh + +experiment: + project: cobol-cpt-grpo + name: collect-rl + group: data + +runtime: + image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a + env: + RAW: /data/datasets/cobol/raw + HF_HOME: /tmp/hf-cache + SOURCE_LICENSES: permissive + +compute: + gpus: 0 + cpu_request: "4" + memory_request: 16Gi + cpu_limit: "8" + memory_limit: 32Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/collect-rl + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/01-collect/tau.yaml b/examples/cobol-cpt-grpo/01-collect/tau.yaml new file mode 100644 index 00000000..2d75bf5d --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/tau.yaml @@ -0,0 +1,29 @@ +name: cobol-collect +engine: job +entrypoint: collect-cobol-data.sh + +experiment: + project: cobol-cpt-grpo + name: collect + group: data + +runtime: + image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a + env: + RAW: /data/datasets/cobol/raw + HF_HOME: /tmp/hf-cache + SOURCE_LICENSES: permissive + +compute: + gpus: 0 + cpu_request: "2" + memory_request: 8Gi + cpu_limit: "4" + memory_limit: 16Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/collect + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py b/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py new file mode 100755 index 00000000..71b91c98 --- /dev/null +++ b/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py @@ -0,0 +1,616 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +COBOL training data preparation pipeline. + +Reads raw collected data from /data/datasets/cobol/raw/, +applies dedup, quality filtering, formatting, and train/val split. +Writes training-ready JSONL to /data/datasets/cobol/prepared/. + +Stages: + 1. Ingest — load all COBOL source files and JSONL records into a unified list + 2. Normalize — strip trailing whitespace, normalize line endings, detect format + 3. Dedup — SHA-256 exact dedup, then MinHash near-dedup (Jaccard > 0.7) + 4. Quality filter — remove stubs (<20 non-empty lines), binaries, generated listings + 5. Format — produce training JSONL in completion format with metadata + 6. Split — 95/5 train/val by file, write shards + +Uses Ray for parallel processing when available, falls back to multiprocessing. +""" + +import hashlib +import json +import os +import re +import random +from collections import defaultdict +from pathlib import Path + +RAW = os.environ.get("RAW", "/data/datasets/cobol/raw") +OUT = os.environ.get("OUT", "/data/datasets/cobol/prepared") +REPORT = os.environ.get("REPORT", "/data/datasets/cobol/prepared/prep-report.json") +SEED = 42 + +os.makedirs(OUT, exist_ok=True) +random.seed(SEED) + +COBOL_EXTS = {".cbl", ".cob", ".cobol", ".cpy", ".CBL", ".COB", ".CPY"} + + +def discover_files_from_index(raw_dir, data_root): + """BlobFuse with cancel-list-on-mount-seconds can't list large dirs. + Use the tokenizer per-file-stats.jsonl and known JSONL paths from + the manifest as an authoritative file index instead of os.walk.""" + per_file_stats = os.path.join( + data_root, "tokenizer-analysis", "per-file-stats.jsonl" + ) + + source_paths = [] + jsonl_paths = [] + + file_index = os.path.join(data_root, "source-file-index.txt") + + if os.path.isfile(per_file_stats): + print(f" Reading file index from {per_file_stats}") + with open(per_file_stats, "r") as f: + for line in f: + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + rel = rec.get("file", "") + if rel: + full = os.path.join(raw_dir, rel) + source_paths.append(full) + except (json.JSONDecodeError, KeyError): + pass + print(f" Found {len(source_paths)} source file paths from tokenizer index") + elif os.path.isfile(file_index): + print(f" Reading file index from {file_index}") + with open(file_index, "r") as f: + for line in f: + fpath = line.strip() + if not fpath: + continue + if any(fpath.lower().endswith(ext) for ext in COBOL_EXTS): + source_paths.append(fpath) + elif fpath.endswith(".jsonl"): + jsonl_paths.append(fpath) + print( + f" Found {len(source_paths)} source files and {len(jsonl_paths)} JSONL files from file index" + ) + else: + print(" WARNING: no file index found, walking known subdirectories") + known_subdirs = [ + "troy-cobol", + "github", + "x-cobol", + "nist", + "the-stack-v1-cobol", + ] + for subdir in known_subdirs: + subpath = os.path.join(raw_dir, subdir) + if not os.path.isdir(subpath): + continue + count = 0 + for dirpath, _, filenames in os.walk(subpath): + for fname in filenames: + if any(fname.endswith(ext) for ext in COBOL_EXTS): + source_paths.append(os.path.join(dirpath, fname)) + count += 1 + print(f" {subdir}: {count} source files") + + indexed_set = set(source_paths) + extra_source_dirs = [ + os.path.join(raw_dir, "the-stack-v1-cobol", "source"), + ] + extra_count = 0 + for src_dir in extra_source_dirs: + if not os.path.isdir(src_dir): + continue + for fname in os.listdir(src_dir): + fpath = os.path.join(src_dir, fname) + if fpath not in indexed_set and any( + fname.lower().endswith(ext) for ext in COBOL_EXTS + ): + source_paths.append(fpath) + extra_count += 1 + if extra_count: + print(f" Found {extra_count} additional source files from new datasets") + + known_jsonl = [ + os.path.join(raw_dir, "commitpack", "commitpack-cobol.jsonl"), + os.path.join(raw_dir, "mainframebench", "question_answering.jsonl"), + os.path.join(raw_dir, "mainframebench", "multiple_choice_question.jsonl"), + os.path.join(raw_dir, "mainframebench", "COBOL_code_summarization.jsonl"), + os.path.join(raw_dir, "coboleval", "coboleval.jsonl"), + os.path.join(raw_dir, "cobolcodebench", "cobolcodebench.jsonl"), + os.path.join(raw_dir, "the-stack-v1-cobol", "the-stack-v1-cobol.jsonl"), + os.path.join( + raw_dir, "cobol-pretraining-corpus", "cobol-pretraining-corpus.jsonl" + ), + ] + for jpath in known_jsonl: + if os.path.isfile(jpath): + jsonl_paths.append(jpath) + print(f" Found {len(jsonl_paths)} JSONL files from known paths") + + return source_paths, jsonl_paths + + +def ingest_source_files_from_paths(raw_dir, file_paths): + """Load COBOL source files from a known list of paths.""" + files = [] + errors = 0 + for fpath in file_paths: + rel = os.path.relpath(fpath, raw_dir) + try: + with open(fpath, "r", errors="replace") as f: + content = f.read() + files.append({"source": rel, "content": content, "type": "source"}) + except Exception: + errors += 1 + if errors: + print(f" WARNING: {errors} files could not be read") + return files + + +def ingest_jsonl_from_paths(raw_dir, jsonl_paths): + """Load structured records from known JSONL file paths.""" + records = [] + for jpath in jsonl_paths: + rel = os.path.relpath(jpath, raw_dir) + source_name = Path(jpath).stem + try: + with open(jpath, "r") as f: + for i, line in enumerate(f): + line = line.strip() + if not line: + continue + try: + rec = json.loads(line) + except json.JSONDecodeError: + continue + + content = extract_content_from_record(rec, source_name) + if content and len(content.strip()) > 50: + records.append( + { + "source": f"{rel}:{i}", + "content": content, + "type": "jsonl", + "dataset": source_name, + } + ) + except Exception: + pass + return records + + +def extract_content_from_record(rec, source_name): + """Extract COBOL-relevant text from a JSONL record.""" + if "commitpack" in source_name: + old = rec.get("old_contents", "") + new = rec.get("new_contents", "") + if new and looks_like_cobol(new): + return new + if old and looks_like_cobol(old): + return old + return None + + if "mainframebench" in source_name.lower(): + q = rec.get("question", rec.get("input", "")) + a = rec.get("answer", rec.get("output", "")) + if a and looks_like_cobol(a): + return a + if q and looks_like_cobol(q): + return q + return f"{q}\n{a}" if q and a else None + + if "coboleval" in source_name.lower() or "cobolcodebench" in source_name.lower(): + prompt = rec.get("prompt", rec.get("input", "")) + solution = rec.get( + "canonical_solution", rec.get("solution", rec.get("output", "")) + ) + if solution and looks_like_cobol(solution): + return solution + return prompt if prompt else None + + if "the-stack" in source_name.lower(): + content = rec.get("content", "") + if content and len(content.strip()) > 50: + return content + return None + + if "pretraining-corpus" in source_name.lower(): + for key in ["text", "content", "code"]: + val = rec.get(key, "") + if val and isinstance(val, str) and len(val.strip()) > 50: + return val + return None + + for key in ["content", "code", "text", "source", "output", "new_contents"]: + val = rec.get(key, "") + if val and isinstance(val, str) and len(val) > 50: + return val + return None + + +def looks_like_cobol(text): + """Heuristic: does this text look like COBOL code?""" + upper = text.upper() + indicators = [ + "IDENTIFICATION DIVISION", + "PROCEDURE DIVISION", + "DATA DIVISION", + "WORKING-STORAGE", + "PERFORM ", + "MOVE ", + "PIC ", + "PICTURE ", + "01 ", + "05 ", + "DISPLAY ", + "STOP RUN", + "GOBACK", + ] + return sum(1 for ind in indicators if ind in upper) >= 2 + + +def normalize(doc): + """Normalize whitespace and line endings.""" + content = doc["content"] + content = content.replace("\r\n", "\n").replace("\r", "\n") + lines = content.split("\n") + lines = [line.rstrip() for line in lines] + while lines and not lines[-1].strip(): + lines.pop() + doc["content"] = "\n".join(lines) + doc["line_count"] = len(lines) + doc["non_empty_lines"] = sum(1 for ln in lines if ln.strip()) + doc["byte_size"] = len(doc["content"].encode("utf-8")) + return doc + + +def content_hash(text): + return hashlib.sha256(text.encode("utf-8")).hexdigest() + + +def exact_dedup(docs): + """Remove exact duplicates by content hash.""" + seen = {} + unique = [] + dupes = 0 + for doc in docs: + h = content_hash(doc["content"]) + if h not in seen: + seen[h] = doc["source"] + doc["content_hash"] = h + unique.append(doc) + else: + dupes += 1 + return unique, dupes + + +def _line_fingerprint(text, max_lines=200): + """Create a set of normalized line hashes for near-dedup comparison.""" + lines = text.split("\n")[:max_lines] + fp = set() + for line in lines: + stripped = line.strip().lower() + if len(stripped) > 10: + fp.add(hashlib.md5(stripped.encode()).digest()[:8]) + return fp + + +def near_dedup(docs, threshold=0.7): + """Line-fingerprint near-dedup using bucket + Jaccard check. Fast in pure Python.""" + n = len(docs) + print(f" Computing line fingerprints for {n} documents...", flush=True) + + fps = [] + for idx, doc in enumerate(docs): + if idx % 5000 == 0: + print(f" Fingerprints: {idx}/{n}...", flush=True) + fps.append(_line_fingerprint(doc["content"])) + + print(" Building candidate buckets...", flush=True) + from collections import defaultdict + + buckets = defaultdict(list) + for i, fp in enumerate(fps): + if not fp: + continue + size_bucket = len(fp) // 10 + sample = tuple(sorted(fp)[:5]) + for h in sample: + buckets[(size_bucket, h)].append(i) + + candidates = set() + for members in buckets.values(): + if 1 < len(members) <= 500: + for a in range(len(members)): + for b in range(a + 1, len(members)): + i, j = members[a], members[b] + if i < j: + candidates.add((i, j)) + else: + candidates.add((j, i)) + + print( + f" Checking {len(candidates)} candidate pairs (from {n * (n - 1) // 2} total)...", + flush=True, + ) + + removed = set() + for i, j in candidates: + if i in removed or j in removed: + continue + fp_i, fp_j = fps[i], fps[j] + if not fp_i or not fp_j: + continue + intersection = len(fp_i & fp_j) + union = len(fp_i | fp_j) + jacc = intersection / union if union > 0 else 0 + if jacc >= threshold: + size_i = docs[i].get("byte_size", 0) + size_j = docs[j].get("byte_size", 0) + if size_j >= size_i: + removed.add(i) + else: + removed.add(j) + + print(f" Near dedup removed {len(removed)} documents", flush=True) + unique = [doc for i, doc in enumerate(docs) if i not in removed] + return unique, len(removed) + + +LISTING_PATTERNS = [ + re.compile(r"^1?\s*PP\s+\d+", re.MULTILINE), + re.compile(r"COMPILER LISTING", re.IGNORECASE), + re.compile(r"^.*LINE\s+SEQ\s+", re.MULTILINE), + re.compile(r"^\d{6}\s", re.MULTILINE), +] + +BINARY_THRESHOLD = 0.10 + + +def quality_filter(docs, min_non_empty_lines=20): + """Remove stubs, binary-looking files, compiler listings.""" + passed = [] + filtered = defaultdict(int) + + for doc in docs: + content = doc["content"] + + if doc.get("non_empty_lines", 0) < min_non_empty_lines: + filtered["too_short"] += 1 + continue + + null_ratio = content.count("\x00") / max(len(content), 1) + if null_ratio > BINARY_THRESHOLD: + filtered["binary"] += 1 + continue + + non_ascii = sum(1 for c in content if ord(c) > 127) + if non_ascii / max(len(content), 1) > 0.3: + filtered["non_ascii"] += 1 + continue + + is_listing = False + for pat in LISTING_PATTERNS: + matches = pat.findall(content[:2000]) + if len(matches) >= 3: + is_listing = True + break + if is_listing: + filtered["listing"] += 1 + continue + + lines = content.split("\n") + hex_lines = sum( + 1 + for ln in lines + if re.match(r"^[0-9a-fA-F\s:]+$", ln.strip()) and len(ln.strip()) > 40 + ) + if hex_lines > len(lines) * 0.3: + filtered["hex_data"] += 1 + continue + + passed.append(doc) + + return passed, dict(filtered) + + +def detect_cobol_structure(content): + """Detect which COBOL divisions are present.""" + upper = content.upper() + divisions = [] + for div in ["IDENTIFICATION", "ENVIRONMENT", "DATA", "PROCEDURE"]: + if f"{div} DIVISION" in upper: + divisions.append(div) + has_copybook_markers = ( + "COPY " in upper + or content.strip().startswith("01 ") + or content.strip().startswith("05 ") + ) + return { + "divisions": divisions, + "is_complete_program": "IDENTIFICATION" in divisions + and "PROCEDURE" in divisions, + "has_data_division": "DATA" in divisions, + "is_copybook": has_copybook_markers and "PROCEDURE" not in divisions, + } + + +def format_for_training(doc): + """Produce a training record in completion format.""" + content = doc["content"] + structure = detect_cobol_structure(content) + + record = { + "text": content, + "source": doc["source"], + "type": doc.get("type", "source"), + "dataset": doc.get("dataset", ""), + "line_count": doc.get("line_count", 0), + "byte_size": doc.get("byte_size", 0), + "content_hash": doc.get("content_hash", ""), + "is_complete_program": structure["is_complete_program"], + "is_copybook": structure["is_copybook"], + "divisions": structure["divisions"], + } + return record + + +def train_val_split(records, val_ratio=0.05): + """Split records into train/val sets.""" + random.shuffle(records) + val_size = max(1, int(len(records) * val_ratio)) + return records[val_size:], records[:val_size] + + +def write_jsonl(records, path): + with open(path, "w") as f: + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + + +def main(): + report = {"stages": {}} + + DATA_ROOT = os.path.dirname(RAW) + + print("=== Stage 0: Discover Files ===") + source_file_paths, jsonl_file_paths = discover_files_from_index(RAW, DATA_ROOT) + + print("\n=== Stage 1: Ingest ===") + source_files = ingest_source_files_from_paths(RAW, source_file_paths) + print(f" Source files: {len(source_files)}") + jsonl_records = ingest_jsonl_from_paths(RAW, jsonl_file_paths) + print(f" JSONL records: {len(jsonl_records)}") + all_docs = source_files + jsonl_records + print(f" Total documents: {len(all_docs)}") + report["stages"]["ingest"] = { + "source_files": len(source_files), + "jsonl_records": len(jsonl_records), + "total": len(all_docs), + } + + print("\n=== Stage 2: Normalize ===") + all_docs = [normalize(doc) for doc in all_docs] + total_lines = sum(d["line_count"] for d in all_docs) + total_bytes = sum(d["byte_size"] for d in all_docs) + print(f" Total lines: {total_lines:,}") + print(f" Total bytes: {total_bytes:,}") + report["stages"]["normalize"] = { + "total_lines": total_lines, + "total_bytes": total_bytes, + } + + print("\n=== Stage 3: Dedup ===") + docs_after_exact, exact_dupes = exact_dedup(all_docs) + print(f" Exact dedup: removed {exact_dupes}, remaining {len(docs_after_exact)}") + + if len(docs_after_exact) <= 50000: + docs_after_near, near_dupes = near_dedup(docs_after_exact, threshold=0.7) + print(f" Near dedup: removed {near_dupes}, remaining {len(docs_after_near)}") + else: + print(f" Near dedup: skipped (>{50000} docs, would be too slow)") + docs_after_near = docs_after_exact + near_dupes = 0 + + report["stages"]["dedup"] = { + "before": len(all_docs), + "exact_dupes_removed": exact_dupes, + "near_dupes_removed": near_dupes, + "after": len(docs_after_near), + } + + print("\n=== Stage 4: Quality Filter ===") + docs_filtered, filter_reasons = quality_filter(docs_after_near) + print(f" Passed: {len(docs_filtered)}") + print(f" Filtered: {sum(filter_reasons.values())} — {dict(filter_reasons)}") + report["stages"]["quality_filter"] = { + "before": len(docs_after_near), + "passed": len(docs_filtered), + "filtered_reasons": filter_reasons, + } + + print("\n=== Stage 5: Format ===") + training_records = [format_for_training(doc) for doc in docs_filtered] + + complete = sum(1 for r in training_records if r["is_complete_program"]) + copybooks = sum(1 for r in training_records if r["is_copybook"]) + fragments = len(training_records) - complete - copybooks + total_training_lines = sum(r["line_count"] for r in training_records) + total_training_bytes = sum(r["byte_size"] for r in training_records) + print(f" Complete programs: {complete}") + print(f" Copybooks: {copybooks}") + print(f" Fragments/other: {fragments}") + print(f" Total training lines: {total_training_lines:,}") + print( + f" Total training bytes: {total_training_bytes:,} ({total_training_bytes / 1024 / 1024:.1f} MB)" + ) + report["stages"]["format"] = { + "total_records": len(training_records), + "complete_programs": complete, + "copybooks": copybooks, + "fragments": fragments, + "total_lines": total_training_lines, + "total_bytes": total_training_bytes, + } + + print("\n=== Stage 6: Train/Val Split ===") + train_records, val_records = train_val_split(training_records) + print(f" Train: {len(train_records)} records") + print(f" Val: {len(val_records)} records") + + train_path = os.path.join(OUT, "train.jsonl") + val_path = os.path.join(OUT, "val.jsonl") + write_jsonl(train_records, train_path) + write_jsonl(val_records, val_path) + print(f" Written: {train_path}") + print(f" Written: {val_path}") + + train_bytes = os.path.getsize(train_path) + val_bytes = os.path.getsize(val_path) + print(f" Train file: {train_bytes / 1024 / 1024:.1f} MB") + print(f" Val file: {val_bytes / 1024 / 1024:.1f} MB") + + report["stages"]["split"] = { + "train_records": len(train_records), + "val_records": len(val_records), + "train_bytes": train_bytes, + "val_bytes": val_bytes, + } + + sizes = sorted([r["byte_size"] for r in training_records]) + if sizes: + import statistics + + report["size_distribution"] = { + "min": sizes[0], + "median": int(statistics.median(sizes)), + "mean": int(statistics.mean(sizes)), + "p90": sizes[int(len(sizes) * 0.9)], + "p99": sizes[int(len(sizes) * 0.99)], + "max": sizes[-1], + } + print("\n=== Size Distribution (bytes) ===") + print(f" Min: {sizes[0]:,}") + print(f" Median: {int(statistics.median(sizes)):,}") + print(f" P90: {sizes[int(len(sizes) * 0.9)]:,}") + print(f" Max: {sizes[-1]:,}") + + with open(REPORT, "w") as f: + json.dump(report, f, indent=2) + print(f"\nReport written to {REPORT}") + print("\n=== Data preparation complete ===") + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/02-data-prep/tau.yaml b/examples/cobol-cpt-grpo/02-data-prep/tau.yaml new file mode 100644 index 00000000..997acf0c --- /dev/null +++ b/examples/cobol-cpt-grpo/02-data-prep/tau.yaml @@ -0,0 +1,30 @@ +name: cobol-data-prep +engine: job +entrypoint: prepare_data.py + +experiment: + project: cobol-cpt-grpo + name: data-prep + group: data + +runtime: + image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a + env: + RAW: /data/datasets/cobol/raw + OUT: /data/datasets/cobol/prepared + REPORT: /data/datasets/cobol/prepared/prep-report.json + PYTHONUNBUFFERED: "1" + +compute: + gpus: 0 + cpu_request: "1" + memory_request: 4Gi + cpu_limit: "4" + memory_limit: 8Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/data-prep + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py new file mode 100755 index 00000000..d77e9f72 --- /dev/null +++ b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py @@ -0,0 +1,428 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +Build COBOL evaluation benchmarks from the prepared validation set. + +Reads val.jsonl and constructs four benchmark suites: + 1. paragraph_completion — mask a PROCEDURE DIVISION paragraph body, predict it + 2. copybook_usage — given a copybook + partial program, predict field references + 3. program_completion — mask PROCEDURE DIVISION of a complete program + 4. perplexity — held-out complete programs for perplexity measurement + +Each benchmark is a JSONL file with prompt/reference/metadata fields. +Writes to /data/datasets/cobol/eval/ on the PVC. +""" + +import json +import os +import re +import random +import hashlib + +VAL_PATH = os.environ.get("VAL", "/data/datasets/cobol/prepared/val.jsonl") +TRAIN_PATH = os.environ.get("TRAIN", "/data/datasets/cobol/prepared/train.jsonl") +OUT = os.environ.get("OUT", "/data/datasets/cobol/eval") +SEED = 42 + +os.makedirs(OUT, exist_ok=True) +random.seed(SEED) + +PARA_RE = re.compile(r"^ (\d{4}-[A-Z0-9-]+|[A-Z][A-Z0-9-]*)\.\s*$", re.MULTILINE) + +DIVISION_RE = re.compile( + r"(IDENTIFICATION|ENVIRONMENT|DATA|PROCEDURE)\s+DIVISION", re.IGNORECASE +) + + +def load_jsonl(path): + records = [] + with open(path, "r") as f: + for line in f: + line = line.strip() + if line: + records.append(json.loads(line)) + return records + + +def split_divisions(text): + """Split COBOL source into divisions.""" + divisions = {} + parts = DIVISION_RE.split(text) + if len(parts) < 2: + return divisions + i = 1 + while i < len(parts): + div_name = parts[i].upper() + if i + 1 < len(parts): + div_body = parts[i + 1] + next_div = DIVISION_RE.search(div_body) + if next_div: + div_body = div_body[: next_div.start()] + divisions[div_name] = div_body + i += 2 + return divisions + + +def extract_paragraphs(procedure_text): + """Extract named paragraphs from PROCEDURE DIVISION text.""" + paragraphs = [] + matches = list(PARA_RE.finditer(procedure_text)) + for idx, match in enumerate(matches): + name = match.group(1) + start = match.end() + end = ( + matches[idx + 1].start() if idx + 1 < len(matches) else len(procedure_text) + ) + body = procedure_text[start:end].rstrip() + if body.strip(): + paragraphs.append( + {"name": name, "body": body, "start": match.start(), "end": end} + ) + return paragraphs + + +def build_paragraph_completion(records, max_tasks=200): + """Build paragraph completion benchmark. + + Prompt: everything before the paragraph body (divisions + paragraph name) + Reference: the paragraph body + """ + tasks = [] + for rec in records: + if not rec.get("is_complete_program"): + continue + text = rec["text"] + divisions = split_divisions(text) + if "PROCEDURE" not in divisions: + continue + + proc_text = divisions["PROCEDURE"] + paragraphs = extract_paragraphs(proc_text) + + good_paras = [ + p + for p in paragraphs + if len([ln for ln in p["body"].split("\n") if ln.strip()]) >= 3 + ] + if not good_paras: + continue + + para = random.choice(good_paras) + + proc_div_header = " PROCEDURE DIVISION" + proc_prefix = proc_text[: para["start"]] + non_proc = ( + text.split("PROCEDURE DIVISION")[0] + if "PROCEDURE DIVISION" in text.upper() + else "" + ) + + prompt = non_proc.rstrip() + "\n" + proc_div_header + proc_prefix + prompt = prompt.rstrip() + "\n " + para["name"] + ".\n" + + reference = para["body"] + + tasks.append( + { + "id": f"para-{hashlib.md5(prompt.encode()).hexdigest()[:8]}", + "type": "paragraph_completion", + "source": rec.get("source", ""), + "paragraph_name": para["name"], + "prompt": prompt, + "reference": reference, + "reference_lines": len( + [ln for ln in reference.split("\n") if ln.strip()] + ), + "stop_at": [ + "\n " + p["name"] + "." + for p in paragraphs + if p["name"] != para["name"] + ][:5], + } + ) + + random.shuffle(tasks) + return tasks[:max_tasks] + + +def build_copybook_usage(records, max_tasks=100): + """Build copybook usage benchmark. + + Find programs that reference copybook-like data structures, then mask + statements that use those fields. + """ + copybooks = [r for r in records if r.get("is_copybook")] + programs = [r for r in records if r.get("is_complete_program")] + + if not copybooks or not programs: + return [] + + field_re = re.compile(r"^\s+\d{2}\s+([A-Z][A-Z0-9-]+)\s", re.MULTILINE) + copybook_fields = {} + for cb in copybooks: + fields = field_re.findall(cb["text"]) + if len(fields) >= 3: + copybook_fields[cb.get("source", "")] = { + "text": cb["text"], + "fields": fields, + } + + tasks = [] + for prog in programs: + text = prog["text"] + divisions = split_divisions(text) + if "PROCEDURE" not in divisions or "DATA" not in divisions: + continue + + proc_text = divisions["PROCEDURE"] + + for cb_src, cb_info in copybook_fields.items(): + used_fields = [f for f in cb_info["fields"] if f in proc_text] + if len(used_fields) < 2: + continue + + lines = proc_text.split("\n") + field_lines = [] + for i, line in enumerate(lines): + if any(f in line for f in used_fields): + field_lines.append((i, line)) + + if len(field_lines) < 2: + continue + + start_idx = field_lines[0][0] + end_idx = min(start_idx + 8, field_lines[-1][0] + 1) + masked_lines = lines[start_idx:end_idx] + reference = "\n".join(masked_lines) + + context_before = "\n".join(lines[:start_idx]) + context_after = "\n".join(lines[end_idx : end_idx + 5]) + + prompt = ( + "Given this copybook:\n" + + cb_info["text"].strip() + + "\n\nAnd this partial COBOL program:\n" + + text.split("PROCEDURE DIVISION")[0].rstrip() + + "\n PROCEDURE DIVISION.\n" + + context_before.rstrip() + + "\n [COMPLETE THE FOLLOWING LINES]\n" + ) + + tasks.append( + { + "id": f"cpyb-{hashlib.md5(prompt.encode()).hexdigest()[:8]}", + "type": "copybook_usage", + "source": prog.get("source", ""), + "copybook_source": cb_src, + "copybook_fields_used": used_fields, + "prompt": prompt, + "reference": reference, + "context_after": context_after, + "reference_lines": len(masked_lines), + } + ) + break + + random.shuffle(tasks) + return tasks[:max_tasks] + + +def build_program_completion(records, max_tasks=50): + """Build full PROCEDURE DIVISION completion benchmark. + + Prompt: IDENTIFICATION + ENVIRONMENT + DATA divisions + Reference: PROCEDURE DIVISION + """ + tasks = [] + for rec in records: + if not rec.get("is_complete_program"): + continue + text = rec["text"] + + divisions = split_divisions(text) + if "PROCEDURE" not in divisions: + continue + + proc_div = divisions["PROCEDURE"] + proc_lines = [ln for ln in proc_div.split("\n") if ln.strip()] + if len(proc_lines) < 10 or len(proc_lines) > 200: + continue + + upper = text.upper() + proc_idx = upper.find("PROCEDURE DIVISION") + if proc_idx < 0: + continue + + prompt = text[:proc_idx].rstrip() + "\n PROCEDURE DIVISION.\n" + reference = proc_div + + tasks.append( + { + "id": f"prog-{hashlib.md5(prompt.encode()).hexdigest()[:8]}", + "type": "program_completion", + "source": rec.get("source", ""), + "prompt": prompt, + "reference": reference, + "reference_lines": len(proc_lines), + "total_program_lines": rec.get("line_count", 0), + } + ) + + random.shuffle(tasks) + return tasks[:max_tasks] + + +def build_perplexity_set(records, max_docs=500): + """Build perplexity evaluation set — complete programs for loss measurement.""" + complete = [ + r + for r in records + if r.get("is_complete_program") and r.get("line_count", 0) >= 30 + ] + random.shuffle(complete) + tasks = [] + for rec in complete[:max_docs]: + tasks.append( + { + "id": f"ppl-{hashlib.md5(rec['text'].encode()).hexdigest()[:8]}", + "type": "perplexity", + "source": rec.get("source", ""), + "text": rec["text"], + "line_count": rec.get("line_count", 0), + "byte_size": rec.get("byte_size", 0), + } + ) + return tasks + + +def build_structural_validation_set(records, max_tasks=100): + """Build set for structural validation — programs to test compilation and formatting.""" + complete = [ + r + for r in records + if r.get("is_complete_program") and 30 <= r.get("line_count", 0) <= 500 + ] + random.shuffle(complete) + tasks = [] + for rec in complete[:max_tasks]: + text = rec["text"] + divisions = split_divisions(text) + if "PROCEDURE" not in divisions: + continue + + upper = text.upper() + proc_idx = upper.find("PROCEDURE DIVISION") + if proc_idx < 0: + continue + + prompt = text[:proc_idx].rstrip() + "\n PROCEDURE DIVISION.\n" + + tasks.append( + { + "id": f"strv-{hashlib.md5(prompt.encode()).hexdigest()[:8]}", + "type": "structural_validation", + "source": rec.get("source", ""), + "prompt": prompt, + "reference_text": text, + "reference_divisions": list(divisions.keys()), + "line_count": rec.get("line_count", 0), + } + ) + return tasks + + +def write_jsonl(records, path): + with open(path, "w") as f: + for rec in records: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + print(f" Written {len(records)} tasks to {path}", flush=True) + + +def main(): + print("=== Building COBOL Evaluation Benchmarks ===", flush=True) + + print(f"\nLoading validation set from {VAL_PATH}...", flush=True) + val_records = load_jsonl(VAL_PATH) + print(f" Loaded {len(val_records)} validation records", flush=True) + + print(f"Loading training set from {TRAIN_PATH}...", flush=True) + train_records = load_jsonl(TRAIN_PATH) + print(f" Loaded {len(train_records)} training records", flush=True) + + random.seed(SEED) + all_records = val_records + + print("\n--- Benchmark 1: Paragraph Completion ---", flush=True) + para_tasks = build_paragraph_completion(all_records, max_tasks=200) + print(f" Generated {len(para_tasks)} paragraph completion tasks", flush=True) + if para_tasks: + avg_lines = sum(t["reference_lines"] for t in para_tasks) / len(para_tasks) + print(f" Average reference lines: {avg_lines:.1f}", flush=True) + write_jsonl(para_tasks, os.path.join(OUT, "paragraph_completion.jsonl")) + + print("\n--- Benchmark 2: Copybook Usage ---", flush=True) + cpyb_pool = val_records + [r for r in train_records if r.get("is_copybook")] + cpyb_tasks = build_copybook_usage(cpyb_pool, max_tasks=100) + print(f" Generated {len(cpyb_tasks)} copybook usage tasks", flush=True) + write_jsonl(cpyb_tasks, os.path.join(OUT, "copybook_usage.jsonl")) + + print("\n--- Benchmark 3: Program Completion ---", flush=True) + prog_tasks = build_program_completion(all_records, max_tasks=50) + print(f" Generated {len(prog_tasks)} program completion tasks", flush=True) + if prog_tasks: + avg_lines = sum(t["reference_lines"] for t in prog_tasks) / len(prog_tasks) + print(f" Average PROCEDURE DIVISION lines: {avg_lines:.1f}", flush=True) + write_jsonl(prog_tasks, os.path.join(OUT, "program_completion.jsonl")) + + print("\n--- Benchmark 4: Perplexity ---", flush=True) + ppl_tasks = build_perplexity_set(all_records, max_docs=len(all_records)) + print(f" Generated {len(ppl_tasks)} perplexity evaluation docs", flush=True) + write_jsonl(ppl_tasks, os.path.join(OUT, "perplexity.jsonl")) + + print("\n--- Benchmark 5: Structural Validation ---", flush=True) + strv_tasks = build_structural_validation_set(all_records, max_tasks=100) + print(f" Generated {len(strv_tasks)} structural validation tasks", flush=True) + write_jsonl(strv_tasks, os.path.join(OUT, "structural_validation.jsonl")) + + total = ( + len(para_tasks) + + len(cpyb_tasks) + + len(prog_tasks) + + len(ppl_tasks) + + len(strv_tasks) + ) + summary = { + "benchmarks": { + "paragraph_completion": { + "count": len(para_tasks), + "file": "paragraph_completion.jsonl", + }, + "copybook_usage": { + "count": len(cpyb_tasks), + "file": "copybook_usage.jsonl", + }, + "program_completion": { + "count": len(prog_tasks), + "file": "program_completion.jsonl", + }, + "perplexity": {"count": len(ppl_tasks), "file": "perplexity.jsonl"}, + "structural_validation": { + "count": len(strv_tasks), + "file": "structural_validation.jsonl", + }, + }, + "total_tasks": total, + "source": VAL_PATH, + } + with open(os.path.join(OUT, "benchmark-summary.json"), "w") as f: + json.dump(summary, f, indent=2) + print( + f"\n=== Benchmark suite complete: {total} total tasks across 5 benchmarks ===", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/03-benchmarks/tau.yaml b/examples/cobol-cpt-grpo/03-benchmarks/tau.yaml new file mode 100644 index 00000000..0870c2d9 --- /dev/null +++ b/examples/cobol-cpt-grpo/03-benchmarks/tau.yaml @@ -0,0 +1,30 @@ +name: cobol-build-benchmarks +engine: job +entrypoint: build_benchmarks.py + +experiment: + project: cobol-cpt-grpo + name: build-benchmarks + group: data + +runtime: + image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a + env: + VAL: /data/datasets/cobol/prepared/val.jsonl + TRAIN: /data/datasets/cobol/prepared/train.jsonl + OUT: /data/datasets/cobol/eval + PYTHONUNBUFFERED: "1" + +compute: + gpus: 0 + cpu_request: "1" + memory_request: 4Gi + cpu_limit: "4" + memory_limit: 8Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/build-benchmarks + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py new file mode 100644 index 00000000..65220791 --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py @@ -0,0 +1,444 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +Continued pretraining of Qwen2.5-Coder-32B on COBOL corpus. + +Uses HuggingFace Trainer + DeepSpeed ZeRO-3 (config built inline) across 2 nodes x 8 GPUs. +Publishes Stellar-compatible metrics chunks for TauGrid experiment tracking. + +Launch: driver.py runs main() in 16 Ray Train workers (TorchTrainer, one GPU each). +Ray Train initializes torch.distributed and sets RANK, LOCAL_RANK, WORLD_SIZE and +MASTER_ADDR/PORT before main() runs; accelerate and DeepSpeed reuse that process group. +Importing this module has no side effects; all env reads happen in main(). + +MAX_STEPS (optional) stops training after that many optimizer steps for a short check, +through a callback, so the LR schedule and TrainingArguments stay those of the full run. +Unset means a full run. +""" + +import functools +import json +import math +import os +import time +from pathlib import Path + +import torch +from datasets import Dataset +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + Trainer, + TrainerCallback, + TrainingArguments, +) + +_orig_torch_load = torch.load + + +def _patched_torch_load(*args, **kwargs): + if "weights_only" not in kwargs: + kwargs["weights_only"] = False + return _orig_torch_load(*args, **kwargs) + + +def _save_file_direct(tensors, filename, metadata=None): + """Writes a safetensors file under its final name. + + safetensors 0.8 writes to a temp file and renames it before closing it, and blobfuse2 + fails that rename with EIO. Serialize in memory and write the file directly instead, + as train_grpo.py's _blobfuse_save_safetensors in 06-grpo does. + """ + import safetensors.torch + + data = safetensors.torch.save(tensors, metadata=metadata) + with open(filename, "wb") as f: + f.write(data) + f.flush() + os.fsync(f.fileno()) + + +class StellarMetricsCallback(TrainerCallback): + def __init__(self, metrics_dir): + self.metrics_dir = Path(metrics_dir) + self.metrics_dir.mkdir(parents=True, exist_ok=True) + + def _publish(self, step, metrics): + row = {"_step": step, "_timestamp": time.time()} + row.update(metrics) + token = f"{step:08d}-{time.time_ns()}" + final = self.metrics_dir / f"chunk-{token}.jsonl" + tmp = self.metrics_dir / f".{final.name}.tmp-{os.getpid()}" + with tmp.open("x") as f: + f.write(json.dumps(row) + "\n") + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, final) + + def on_log(self, args, state, control, logs=None, **kwargs): + if logs and state.global_step > 0: + metrics = {} + for k, v in logs.items(): + if isinstance(v, (int, float)) and not math.isnan(v): + metrics[k] = v + if metrics: + self._publish(state.global_step, metrics) + + def on_evaluate(self, args, state, control, metrics=None, **kwargs): + if metrics and state.global_step > 0: + eval_metrics = {} + for k, v in metrics.items(): + if isinstance(v, (int, float)) and not math.isnan(v): + eval_metrics[k] = v + if eval_metrics: + self._publish(state.global_step, eval_metrics) + + +class StopAtStepCallback(TrainerCallback): + """Stops after max_steps without changing the full-run LR schedule.""" + + def __init__(self, max_steps): + self.max_steps = max_steps + + def on_step_end(self, args, state, control, **kwargs): + if state.global_step >= self.max_steps: + control.should_training_stop = True + return control + + +def load_jsonl(path): + records = [] + with open(path) as f: + for line in f: + line = line.strip() + if line: + records.append(json.loads(line)) + return records + + +def tokenize_dataset(records, tokenizer, max_seq_len): + all_input_ids = [] + for rec in records: + text = rec.get("text", "") + if not text.strip(): + continue + tokens = tokenizer(text, truncation=False, add_special_tokens=False)[ + "input_ids" + ] + tokens.append(tokenizer.eos_token_id) + all_input_ids.extend(tokens) + + total_tokens = len(all_input_ids) + n_chunks = total_tokens // max_seq_len + print(f" Total tokens: {total_tokens:,}", flush=True) + print(f" Chunks of {max_seq_len}: {n_chunks:,}", flush=True) + print( + f" Discarded tail: {total_tokens - n_chunks * max_seq_len:,} tokens", + flush=True, + ) + + chunks = [] + for i in range(n_chunks): + start = i * max_seq_len + chunk = all_input_ids[start : start + max_seq_len] + chunks.append({"input_ids": chunk, "labels": chunk.copy()}) + + return Dataset.from_list(chunks) + + +def main(): + torch.load = _patched_torch_load + from transformers import modeling_utils + + if not hasattr(modeling_utils, "safe_save_file"): + raise RuntimeError( + "transformers.modeling_utils.safe_save_file not found; update _save_file_direct" + ) + modeling_utils.safe_save_file = _save_file_direct + + MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-Coder-32B") + TRAIN_FILE = os.environ.get( + "TRAIN_FILE", "/data/datasets/cobol/prepared/train.jsonl" + ) + VAL_FILE = os.environ.get("VAL_FILE", "/data/datasets/cobol/prepared/val.jsonl") + OUTPUT_DIR = os.environ.get( + "OUTPUT_DIR", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt" + ) + EPOCHS = int(os.environ.get("EPOCHS", "3")) + LR = float(os.environ.get("LR", "2e-5")) + BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1")) + GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "16")) + MAX_SEQ_LEN = int(os.environ.get("MAX_SEQ_LEN", "4096")) + SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "500")) + WARMUP_RATIO = float(os.environ.get("WARMUP_RATIO", "0.05")) + WEIGHT_DECAY = float(os.environ.get("WEIGHT_DECAY", "0.1")) + HF_HOME = os.environ.get("HF_HOME", "/tmp/hf-cache") + RESUME_FROM = os.environ.get("TAU_RESUME_FROM", None) + MAX_STEPS = os.environ.get("MAX_STEPS", "").strip() + stop_at_step = int(MAX_STEPS) if MAX_STEPS else None + + os.environ["HF_HOME"] = HF_HOME + os.makedirs(OUTPUT_DIR, exist_ok=True) + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + is_main = local_rank == 0 + global_rank = int(os.environ.get("RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + + if is_main: + print("=" * 60, flush=True) + print("COBOL Continued Pretraining (multi-node)", flush=True) + print("=" * 60, flush=True) + print(f" Model: {MODEL_ID}", flush=True) + print(f" Train file: {TRAIN_FILE}", flush=True) + print(f" Val file: {VAL_FILE}", flush=True) + print(f" Output: {OUTPUT_DIR}", flush=True) + print(f" Epochs: {EPOCHS}", flush=True) + print(f" LR: {LR}", flush=True) + print(f" Batch size: {BATCH_SIZE} (per device)", flush=True) + print(f" Grad accum: {GRAD_ACCUM}", flush=True) + print(f" Max seq len: {MAX_SEQ_LEN}", flush=True) + print(f" Save steps: {SAVE_STEPS}", flush=True) + print(f" Stop at step: {MAX_STEPS or 'None (full run)'}", flush=True) + print(f" World size: {world_size} processes", flush=True) + print(f" GPUs/node: {torch.cuda.device_count()}", flush=True) + print(f" Resume from: {RESUME_FROM or 'None (fresh start)'}", flush=True) + effective_batch = BATCH_SIZE * GRAD_ACCUM * world_size + print( + f" Effective batch: {effective_batch} sequences = {effective_batch * MAX_SEQ_LEN:,} tokens", + flush=True, + ) + + if is_main: + print("\nLoading tokenizer...", flush=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + + if is_main: + print("\nTokenizing training data...", flush=True) + train_records = load_jsonl(TRAIN_FILE) + if is_main: + print(f" Loaded {len(train_records)} training records", flush=True) + train_dataset = tokenize_dataset(train_records, tokenizer, MAX_SEQ_LEN) + del train_records + + if is_main: + print("\nTokenizing validation data...", flush=True) + val_records = load_jsonl(VAL_FILE) + if is_main: + print(f" Loaded {len(val_records)} validation records", flush=True) + val_dataset = tokenize_dataset(val_records, tokenizer, MAX_SEQ_LEN) + del val_records + + if is_main: + print(f"\n Train: {len(train_dataset)} chunks", flush=True) + print(f" Val: {len(val_dataset)} chunks", flush=True) + train_tokens = len(train_dataset) * MAX_SEQ_LEN + print(f" Train tokens per epoch: {train_tokens:,}", flush=True) + print( + f" Total tokens ({EPOCHS} epochs): {train_tokens * EPOCHS:,}", flush=True + ) + + if is_main: + print("\nLoading model...", flush=True) + model = AutoModelForCausalLM.from_pretrained( + MODEL_ID, + dtype=torch.bfloat16, + trust_remote_code=True, + attn_implementation="sdpa", + low_cpu_mem_usage=True, + ) + model.save_pretrained = functools.partial( + model.save_pretrained, max_shard_size="5GB" + ) + + if is_main: + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + total = sum(p.numel() for p in model.parameters()) + print( + f" Parameters: {total / 1e9:.1f}B total, {trainable / 1e9:.1f}B trainable", + flush=True, + ) + + resume_checkpoint = None + if RESUME_FROM and os.path.isdir(RESUME_FROM): + resume_checkpoint = RESUME_FROM + if is_main: + print(f"\n Resuming from TauGrid checkpoint: {RESUME_FROM}", flush=True) + else: + checkpoints = sorted( + [d for d in Path(OUTPUT_DIR).glob("checkpoint-*") if d.is_dir()], + key=lambda p: int(p.name.split("-")[-1]), + ) + if checkpoints: + resume_checkpoint = str(checkpoints[-1]) + if is_main: + print( + f"\n Resuming from latest checkpoint: {resume_checkpoint}", + flush=True, + ) + + metrics_dir = os.path.join(OUTPUT_DIR, "metrics-history-attempt-0") + stellar_cb = StellarMetricsCallback(metrics_dir) if global_rank == 0 else None + + ds_config = { + "bf16": {"enabled": True}, + "zero_optimization": { + "stage": 3, + "offload_optimizer": {"device": "none"}, + "offload_param": {"device": "none"}, + "overlap_comm": False, + "contiguous_gradients": True, + "reduce_bucket_size": 1e8, + "stage3_prefetch_bucket_size": 8.4e7, + "stage3_param_persistence_threshold": 1e6, + "stage3_gather_16bit_weights_on_model_save": True, + }, + "gradient_accumulation_steps": "auto", + "gradient_clipping": 1.0, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": False, + } + ds_config_path = os.path.join(OUTPUT_DIR, "ds_config.json") + if global_rank == 0: + with open(ds_config_path, "w") as f: + json.dump(ds_config, f, indent=2) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + training_args = TrainingArguments( + output_dir=OUTPUT_DIR, + num_train_epochs=EPOCHS, + per_device_train_batch_size=BATCH_SIZE, + per_device_eval_batch_size=BATCH_SIZE, + gradient_accumulation_steps=GRAD_ACCUM, + learning_rate=LR, + lr_scheduler_type="cosine", + warmup_steps=WARMUP_RATIO, + weight_decay=WEIGHT_DECAY, + bf16=True, + logging_steps=10, + eval_strategy="steps", + eval_steps=SAVE_STEPS, + save_strategy="steps", + save_steps=SAVE_STEPS, + save_total_limit=5, + load_best_model_at_end=False, + dataloader_num_workers=4, + dataloader_pin_memory=True, + gradient_checkpointing=True, + gradient_checkpointing_kwargs={"use_reentrant": False}, + deepspeed=ds_config_path, + report_to="none", + seed=42, + remove_unused_columns=False, + ) + + stop_cb = StopAtStepCallback(stop_at_step) if stop_at_step is not None else None + callbacks = [cb for cb in [stellar_cb, stop_cb] if cb is not None] + trainer = Trainer( + model=model, + args=training_args, + train_dataset=train_dataset, + eval_dataset=val_dataset, + callbacks=callbacks, + ) + + if is_main: + print("\n" + "=" * 60, flush=True) + print("Starting training...", flush=True) + print("=" * 60, flush=True) + + trainer.train(resume_from_checkpoint=resume_checkpoint) + + if is_main: + print("\nRunning final evaluation...", flush=True) + final_metrics = trainer.evaluate() + final_eval_loss = final_metrics["eval_loss"] + + best_eval_loss = trainer.state.best_metric + best_ckpt = trainer.state.best_model_checkpoint + if is_main: + print( + f" Final eval loss: {final_eval_loss:.4f} (ppl {math.exp(final_eval_loss):.2f})", + flush=True, + ) + if best_eval_loss is not None and best_ckpt: + print( + f" Best checkpoint: {best_ckpt} (eval loss {best_eval_loss:.4f}, ppl {math.exp(best_eval_loss):.2f})", + flush=True, + ) + else: + print(" No best checkpoint tracked — using final weights", flush=True) + + use_best = best_eval_loss is not None and best_eval_loss < final_eval_loss + chosen_loss = best_eval_loss if use_best else final_eval_loss + + if use_best: + if is_main: + print(f" Best checkpoint is better — copying from {best_ckpt}", flush=True) + else: + if is_main: + print(" Final weights are best — saving directly", flush=True) + + final_dir = os.path.join(OUTPUT_DIR, "final") + if use_best: + if global_rank == 0: + print(f"\nCopying best checkpoint to {final_dir}...", flush=True) + import shutil + + if os.path.exists(final_dir): + shutil.rmtree(final_dir) + shutil.copytree( + best_ckpt, + final_dir, + ignore=shutil.ignore_patterns( + "global_step*", "zero_pp*", "optimizer*", "*.pt" + ), + ) + tokenizer.save_pretrained(final_dir) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + else: + if global_rank == 0: + print("\nSaving final model...", flush=True) + trainer.save_model(final_dir) + tokenizer.save_pretrained(final_dir) + + if global_rank == 0: + print(f"\n Saved eval loss: {chosen_loss:.4f}", flush=True) + print(f" Saved eval perplexity: {math.exp(chosen_loss):.2f}", flush=True) + print(f" Model saved to: {final_dir}", flush=True) + + config = { + "model_id": MODEL_ID, + "epochs": EPOCHS, + "lr": LR, + "batch_size": BATCH_SIZE, + "grad_accum": GRAD_ACCUM, + "max_seq_len": MAX_SEQ_LEN, + "effective_batch_tokens": BATCH_SIZE + * GRAD_ACCUM + * world_size + * MAX_SEQ_LEN, + "train_chunks": len(train_dataset), + "val_chunks": len(val_dataset), + "final_eval_loss": final_eval_loss, + "best_checkpoint_eval_loss": best_eval_loss, + "saved_eval_loss": chosen_loss, + "saved_eval_ppl": math.exp(chosen_loss), + "used_best_checkpoint": use_best, + "world_size": world_size, + "gpus_per_node": torch.cuda.device_count(), + } + with open(os.path.join(OUTPUT_DIR, "train-config.json"), "w") as f: + json.dump(config, f, indent=2) + + print("\n=== Training complete ===", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/04-cpt/driver.py b/examples/cobol-cpt-grpo/04-cpt/driver.py new file mode 100644 index 00000000..63c031d5 --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/driver.py @@ -0,0 +1,65 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""RayJob entrypoint for COBOL continued pretraining. + +TauGrid runs this file as the Ray driver on a CPU-only head pod, so it must not import +transformers, deepspeed or the training module. A Ray Train TorchTrainer starts 16 +workers, one GPU each (2 worker pods x 8 GPUs), and each worker calls +continued_pretrain.main(). + +Ray Train initializes torch.distributed (NCCL on GPUs) and sets RANK, LOCAL_RANK, +WORLD_SIZE and MASTER_ADDR/PORT in every worker before train_fn runs, the same contract +torchrun gives. continued_pretrain is imported inside train_fn; TauGrid ships +this directory as the Ray working_dir, so the import resolves on the workers. +Configuration (MODEL_ID, OUTPUT_DIR, BATCH_SIZE, MAX_STEPS, ...) comes from the pod env, +which TauGrid sets on the head and on every worker. + +HF Trainer keeps its own checkpointing and resume under OUTPUT_DIR; Ray Train only +stores its run metadata under OUTPUT_DIR/ray. +""" + +import os + +from ray.train import RunConfig, ScalingConfig +from ray.train.torch import TorchTrainer + +NUM_WORKERS = 16 +CPUS_PER_WORKER = 6 + + +def train_fn(): + import continued_pretrain + + continued_pretrain.main() + + +def build_trainer( + num_workers=NUM_WORKERS, + use_gpu=True, + cpus_per_worker=CPUS_PER_WORKER, + storage_path=None, +): + resources = {"CPU": cpus_per_worker} + if use_gpu: + resources["GPU"] = 1 + if storage_path is None: + output_dir = os.environ.get( + "OUTPUT_DIR", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt" + ) + storage_path = os.path.join(output_dir, "ray") + return TorchTrainer( + train_loop_per_worker=train_fn, + scaling_config=ScalingConfig( + num_workers=num_workers, use_gpu=use_gpu, resources_per_worker=resources + ), + run_config=RunConfig(storage_path=storage_path), + ) + + +def main(): + build_trainer().fit() + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/04-cpt/tau.yaml b/examples/cobol-cpt-grpo/04-cpt/tau.yaml new file mode 100644 index 00000000..068ae6a3 --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/tau.yaml @@ -0,0 +1,65 @@ +name: cobol-cpt +engine: rayjob +run: + entrypoint: driver.py + working_dir: . + working_dir_excludes: + - __pycache__ + - .pytest_cache + - .ruff_cache + - test_*.py + +experiment: + project: cobol-cpt-grpo + name: cpt + group: cpt-ray-zero3 + +runtime: + image: registry.example.com/cobol-train:v0.28.0-ray2.56.0 + rdma: + enabled: true + env: + MODEL_ID: Qwen/Qwen2.5-Coder-32B + TRAIN_FILE: /data/datasets/cobol/prepared/train.jsonl + VAL_FILE: /data/datasets/cobol/prepared/val.jsonl + OUTPUT_DIR: /data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt + EPOCHS: "3" + LR: "2e-5" + BATCH_SIZE: "4" + GRAD_ACCUM: "4" + MAX_SEQ_LEN: "4096" + SAVE_STEPS: "25" + WARMUP_RATIO: "0.05" + WEIGHT_DECAY: "0.1" + PYTHONUNBUFFERED: "1" + HF_HOME: /tmp/hf-cache + NCCL_DEBUG: INFO + NCCL_SOCKET_IFNAME: eth0 + NCCL_IB_HCA: mlx5 + NCCL_IB_GID_INDEX: "3" + NCCL_NET_GDR_LEVEL: "5" + NCCL_ALGO: Ring + NCCL_IB_QPS_PER_CONNECTION: "4" + NCCL_IB_SPLIT_DATA_ON_QPS: "1" + NCCL_IB_TIMEOUT: "22" + NCCL_TIMEOUT: "1800000" + +execution: + launcher: ray-train + allow_nccl_override: true + +compute: + workers: 2 + gpus_per_worker: 8 + worker_cpu_request: "32" + worker_memory_request: 768Gi + worker_cpu_limit: "64" + worker_memory_limit: 1500Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt + +policy: + workspace: taugrid-default + topology: same-network-domain diff --git a/examples/cobol-cpt-grpo/04-cpt/test_continued_pretrain.py b/examples/cobol-cpt-grpo/04-cpt/test_continued_pretrain.py new file mode 100644 index 00000000..a7d6a2aa --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/test_continued_pretrain.py @@ -0,0 +1,264 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local tests for continued_pretrain's launch contract. + +transformers and datasets are replaced with fakes, so this needs CPU torch only: + python -m pytest examples/cobol-cpt-grpo/04-cpt/test_continued_pretrain.py -v + +Checks that importing the module has no side effects, that TrainingArguments are the +same with MAX_STEPS set or unset, that MAX_STEPS only adds a stop callback, and that +saves use 5GB shards written without a temp-file rename. +""" + +import dataclasses +import importlib.util +import json +import os +import subprocess +import sys +import types + +import pytest +import torch + +CPT_DIR = os.path.dirname(os.path.abspath(__file__)) + +EXPECTED_ARGS = { + "num_train_epochs": 3, + "per_device_train_batch_size": 4, + "per_device_eval_batch_size": 4, + "gradient_accumulation_steps": 4, + "learning_rate": 2e-5, + "lr_scheduler_type": "cosine", + "warmup_steps": 0.05, + "weight_decay": 0.1, + "bf16": True, + "logging_steps": 10, + "eval_strategy": "steps", + "eval_steps": 25, + "save_strategy": "steps", + "save_steps": 25, + "save_total_limit": 5, + "load_best_model_at_end": False, + "dataloader_num_workers": 4, + "dataloader_pin_memory": True, + "gradient_checkpointing": True, + "gradient_checkpointing_kwargs": {"use_reentrant": False}, + "report_to": "none", + "seed": 42, + "remove_unused_columns": False, +} + +CPT_ENV = { + "EPOCHS": "3", + "LR": "2e-5", + "BATCH_SIZE": "4", + "GRAD_ACCUM": "4", + "MAX_SEQ_LEN": "8", + "SAVE_STEPS": "25", + "WARMUP_RATIO": "0.05", + "WEIGHT_DECAY": "0.1", +} + + +class _StopBeforeTrainer(Exception): + pass + + +class _FakeTokenizer: + pad_token = None + eos_token = "" + eos_token_id = 0 + + def __call__(self, text, **kwargs): + return {"input_ids": list(range(1, len(text) + 1))} + + +class _FakeModel: + def __init__(self): + self.saves = [] + + def parameters(self): + return [torch.nn.Parameter(torch.zeros(2))] + + def save_pretrained(self, output_dir, **kwargs): + self.saves.append((output_dir, kwargs)) + + +def _fake_hf_modules(captured): + transformers = types.ModuleType("transformers") + transformers.AutoTokenizer = types.SimpleNamespace( + from_pretrained=lambda *a, **k: _FakeTokenizer() + ) + transformers.AutoModelForCausalLM = types.SimpleNamespace( + from_pretrained=lambda *a, **k: _FakeModel() + ) + transformers.TrainerCallback = type("TrainerCallback", (), {}) + + def training_arguments(**kwargs): + captured["args"] = kwargs + return types.SimpleNamespace(**kwargs) + + def trainer(**kwargs): + captured["callbacks"] = kwargs["callbacks"] + captured["model"] = kwargs["model"] + raise _StopBeforeTrainer + + transformers.TrainingArguments = training_arguments + transformers.modeling_utils = types.SimpleNamespace(safe_save_file=None) + transformers.Trainer = trainer + + datasets = types.ModuleType("datasets") + datasets.Dataset = types.SimpleNamespace(from_list=lambda rows: rows) + return {"transformers": transformers, "datasets": datasets} + + +@pytest.fixture +def cpt(tmp_path, monkeypatch): + """Imports continued_pretrain with fake HF modules and fails on any subprocess.""" + captured = {} + for name, module in _fake_hf_modules(captured).items(): + monkeypatch.setitem(sys.modules, name, module) + + def no_subprocess(*args, **kwargs): + raise AssertionError(f"unexpected subprocess: {args}") + + for fn in ("run", "call", "check_call", "check_output", "Popen"): + monkeypatch.setattr(subprocess, fn, no_subprocess) + + output_dir = tmp_path / "out" + monkeypatch.setenv("OUTPUT_DIR", str(output_dir)) + monkeypatch.setenv("HF_HOME", str(tmp_path / "hf")) + monkeypatch.delenv("MAX_STEPS", raising=False) + monkeypatch.delenv("TAU_RESUME_FROM", raising=False) + for key in ("RANK", "LOCAL_RANK", "WORLD_SIZE"): + monkeypatch.delenv(key, raising=False) + + orig_load = torch.load + spec = importlib.util.spec_from_file_location( + "continued_pretrain", os.path.join(CPT_DIR, "continued_pretrain.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + yield types.SimpleNamespace( + module=module, captured=captured, output_dir=output_dir, orig_load=orig_load + ) + torch.load = orig_load + + +def _run_main(cpt, tmp_path, monkeypatch): + data = tmp_path / "data.jsonl" + data.write_text("".join(json.dumps({"text": "x" * 30}) + "\n" for _ in range(4))) + for key, value in CPT_ENV.items(): + monkeypatch.setenv(key, value) + monkeypatch.setenv("TRAIN_FILE", str(data)) + monkeypatch.setenv("VAL_FILE", str(data)) + with pytest.raises(_StopBeforeTrainer): + cpt.module.main() + return cpt.captured + + +def test_import_has_no_side_effects(cpt): + assert torch.load is cpt.orig_load + assert not cpt.output_dir.exists() + assert not torch.distributed.is_initialized() + assert not torch.cuda.is_initialized() + for name in ("LOCAL_RANK", "MODEL_ID", "OUTPUT_DIR", "BATCH_SIZE"): + assert not hasattr(cpt.module, name) + + +def _expected_args(cpt): + output_dir = str(cpt.output_dir) + return { + "output_dir": output_dir, + "deepspeed": os.path.join(output_dir, "ds_config.json"), + **EXPECTED_ARGS, + } + + +def _stop_callbacks(cpt): + return [ + cb + for cb in cpt.captured["callbacks"] + if isinstance(cb, cpt.module.StopAtStepCallback) + ] + + +def test_max_steps_unset_keeps_training_arguments(cpt, tmp_path, monkeypatch): + captured = _run_main(cpt, tmp_path, monkeypatch) + assert captured["args"] == _expected_args(cpt) + assert _stop_callbacks(cpt) == [] + assert len(captured["callbacks"]) == 1 + assert torch.load is cpt.module._patched_torch_load + + +def test_max_steps_set_adds_stop_callback_only(cpt, tmp_path, monkeypatch): + monkeypatch.setenv("MAX_STEPS", "50") + captured = _run_main(cpt, tmp_path, monkeypatch) + assert captured["args"] == _expected_args(cpt) + stops = _stop_callbacks(cpt) + assert len(stops) == 1 + assert stops[0].max_steps == 50 + + +def test_stop_callback_stops_at_max_steps(cpt): + cb = cpt.module.StopAtStepCallback(50) + for step, should_stop in ((1, False), (49, False), (50, True), (51, True)): + control = types.SimpleNamespace(should_training_stop=False) + state = types.SimpleNamespace(global_step=step) + cb.on_step_end(None, state, control) + assert control.should_training_stop is should_stop, step + + +def test_training_arguments_exist_in_installed_transformers(): + """The fake TrainingArguments accepts any kwarg; check the names against the real one.""" + real = pytest.importorskip("transformers") + fields = {f.name for f in dataclasses.fields(real.TrainingArguments)} + assert set(EXPECTED_ARGS) - fields == set() + assert {"output_dir", "deepspeed"} <= fields + + +def test_saves_use_5gb_shards(cpt, tmp_path, monkeypatch): + model = _run_main(cpt, tmp_path, monkeypatch)["model"] + model.save_pretrained("ckpt") + model.save_pretrained("final", state_dict={}, max_shard_size="1GB") + assert model.saves == [ + ("ckpt", {"max_shard_size": "5GB"}), + ("final", {"state_dict": {}, "max_shard_size": "1GB"}), + ] + + +def test_main_replaces_safe_save_file(cpt, tmp_path, monkeypatch): + _run_main(cpt, tmp_path, monkeypatch) + modeling_utils = sys.modules["transformers"].modeling_utils + assert modeling_utils.safe_save_file is cpt.module._save_file_direct + + +def test_save_file_direct_matches_save_file_without_temp_file(cpt, tmp_path): + """Needs the real safetensors package; the fixture only fakes transformers and datasets.""" + st = pytest.importorskip("safetensors.torch") + tensors = { + "a": torch.arange(6, dtype=torch.bfloat16).reshape(2, 3), + "b": torch.ones(4), + } + metadata = {"format": "pt"} + st.save_file(tensors, str(tmp_path / "ref.safetensors"), metadata=metadata) + + opened = [] + real_open = open + + def tracking_open(path, *args, **kwargs): + opened.append(os.path.basename(path)) + return real_open(path, *args, **kwargs) + + renames = [] + out = tmp_path / "direct.safetensors" + with pytest.MonkeyPatch.context() as mp: + mp.setattr("builtins.open", tracking_open) + mp.setattr(os, "rename", lambda *a: renames.append(a)) + mp.setattr(os, "replace", lambda *a: renames.append(a)) + cpt.module._save_file_direct(tensors, str(out), metadata=metadata) + + assert opened == ["direct.safetensors"] and renames == [] + assert out.read_bytes() == (tmp_path / "ref.safetensors").read_bytes() diff --git a/examples/cobol-cpt-grpo/04-cpt/test_driver.py b/examples/cobol-cpt-grpo/04-cpt/test_driver.py new file mode 100644 index 00000000..a02da79b --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/test_driver.py @@ -0,0 +1,125 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local test for the CPT RayJob driver. + +Runs the driver's TorchTrainer on a local Ray cluster with 2 CPU workers (gloo) and a +stub continued_pretrain, then checks that every worker ran main() with the torchrun-style +env and an initialized process group, and that the driver process itself never imported +transformers or deepspeed. Needs ray[train] and CPU torch, not transformers or a GPU: + python -m pytest examples/cobol-cpt-grpo/04-cpt/test_driver.py -v +""" + +import importlib.util +import json +import os +import shutil +import sys + +import pytest +import ray + +CPT_DIR = os.path.dirname(os.path.abspath(__file__)) + +STUB_CONTINUED_PRETRAIN = """\ +import json +import os + +import deepspeed # noqa: F401 +import torch +import torch.distributed as dist +import transformers # noqa: F401 + + +def main(): + t = torch.ones(1) + dist.all_reduce(t) + rank = os.environ["RANK"] + with open(os.path.join(os.environ["STUB_OUT"], f"rank-{rank}.json"), "w") as f: + json.dump({ + "rank": int(rank), + "local_rank": int(os.environ["LOCAL_RANK"]), + "world_size": int(os.environ["WORLD_SIZE"]), + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "dist_initialized": dist.is_initialized(), + "dist_rank": dist.get_rank(), + "dist_world_size": dist.get_world_size(), + "backend": dist.get_backend(), + "all_reduce": t.item(), + "pid": os.getpid(), + }, f) +""" + + +@pytest.fixture +def work_dir(tmp_path, monkeypatch): + work = tmp_path / "work" + work.mkdir() + shutil.copy(os.path.join(CPT_DIR, "driver.py"), work / "driver.py") + (work / "continued_pretrain.py").write_text(STUB_CONTINUED_PRETRAIN) + (work / "transformers.py").write_text("") + (work / "deepspeed.py").write_text("") + monkeypatch.syspath_prepend(str(work)) + for name in ("driver", "continued_pretrain", "transformers", "deepspeed"): + monkeypatch.delitem(sys.modules, name, raising=False) + return work + + +def test_build_trainer_defaults_match_plan(work_dir, monkeypatch): + monkeypatch.setenv( + "OUTPUT_DIR", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt" + ) + import driver + + trainer = driver.build_trainer() + scaling = trainer.scaling_config + assert scaling.num_workers == 16 + assert scaling.use_gpu is True + assert scaling.resources_per_worker == {"CPU": 6, "GPU": 1} + assert ( + trainer.run_config.storage_path + == "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/ray" + ) + + +def test_driver_runs_main_on_every_worker_without_hf_imports(work_dir, tmp_path): + out = tmp_path / "out" + out.mkdir() + ray.init( + num_cpus=4, + include_dashboard=False, + runtime_env={"working_dir": str(work_dir), "env_vars": {"STUB_OUT": str(out)}}, + ) + try: + import driver + + driver.build_trainer( + num_workers=2, + use_gpu=False, + cpus_per_worker=1, + storage_path=str(tmp_path / "ray"), + ).fit() + finally: + ray.shutdown() + + results = sorted( + (json.loads(p.read_text()) for p in out.glob("rank-*.json")), + key=lambda r: r["rank"], + ) + assert [r["rank"] for r in results] == [0, 1] + for r in results: + assert r["world_size"] == 2 + assert r["dist_initialized"] is True + assert r["dist_rank"] == r["rank"] + assert r["dist_world_size"] == 2 + assert r["backend"] == "gloo" + assert r["all_reduce"] == 2.0 + assert r["master_addr"] and r["master_port"] + assert r["pid"] != os.getpid() + assert sorted(r["local_rank"] for r in results) == [0, 1] + + assert importlib.util.find_spec("transformers") is not None + assert importlib.util.find_spec("deepspeed") is not None + for name in ("continued_pretrain", "transformers", "deepspeed"): + assert name not in sys.modules diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py new file mode 100755 index 00000000..34496c84 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py @@ -0,0 +1,363 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +RL prompt extraction pipeline. + +Reads ~60K raw COBOL source files from the PVC, deduplicates, extracts +PROCEDURE DIVISION prompts, compile-tests each with cobc, and writes +an RL-ready JSONL. + +Input: /data/datasets/cobol/source-file-index.txt (rebuilt by 01-collect/collect-rl-data.sh) + /data/datasets/cobol/raw/**/*.{cbl,cob,cpy} +Output: /data/datasets/cobol/prepared/rl-prompts.jsonl + /data/datasets/cobol/prepared/rl-prep-report.json + +Compatible with 06-grpo/train_grpo.py's load_prompts() — each JSONL row has +{"text": ""} where the program contains both the prompt +prefix (up to PROCEDURE DIVISION header) and the body. + +Requires: cobc (GnuCOBOL) installed in the container. +""" + +import hashlib +import json +import os +import re +import subprocess +import tempfile +import time +from collections import Counter +from concurrent.futures import ProcessPoolExecutor, as_completed + +RAW = os.environ.get("RAW", "/data/datasets/cobol/raw") +DATA_ROOT = os.environ.get("DATA_ROOT", "/data/datasets/cobol") +OUT = os.environ.get("OUT", "/data/datasets/cobol/prepared") +MAX_WORKERS = int(os.environ.get("MAX_WORKERS", "8")) +MIN_BODY_LINES = int(os.environ.get("MIN_BODY_LINES", "5")) +MAX_BODY_LINES = int(os.environ.get("MAX_BODY_LINES", "200")) +MAX_PROMPT_CHARS = int(os.environ.get("MAX_PROMPT_CHARS", "50000")) + +os.makedirs(OUT, exist_ok=True) + +COBOL_EXTS = {".cbl", ".cob", ".cobol", ".cpy", ".CBL", ".COB", ".CPY"} + +PROC_DIV_RE = re.compile( + r"^(.*PROCEDURE\s+DIVISION[^\n]*\.?)$", re.IGNORECASE | re.MULTILINE +) + +EXEC_BLOCK_RE = re.compile( + r"EXEC\s+(SQL|CICS|DLI)\b.*?END-EXEC\s*\.?", re.IGNORECASE | re.DOTALL +) + +COPY_STMT_RE = re.compile(r"^\s*COPY\s+\S+.*?\.\s*$", re.IGNORECASE | re.MULTILINE) + +stats = Counter() + + +def log(msg): + print(f"[{time.strftime('%H:%M:%S')}] {msg}", flush=True) + + +def discover_files(): + index_path = os.path.join(DATA_ROOT, "source-file-index.txt") + source_files = [] + if os.path.isfile(index_path): + log(f"Reading file index from {index_path}") + with open(index_path) as f: + for line in f: + fpath = line.strip() + if not fpath: + continue + ext = os.path.splitext(fpath)[1].lower() + if ext in {e.lower() for e in COBOL_EXTS}: + source_files.append(fpath) + else: + log("No file index found, walking raw directory...") + for dirpath, _, filenames in os.walk(RAW): + for fname in filenames: + if os.path.splitext(fname)[1] in COBOL_EXTS: + source_files.append(os.path.join(dirpath, fname)) + log(f"Discovered {len(source_files)} source files") + return source_files + + +def exact_dedup(texts_with_paths): + seen = set() + unique = [] + for path, text in texts_with_paths: + h = hashlib.sha256(text.encode("utf-8", errors="replace")).hexdigest() + if h not in seen: + seen.add(h) + unique.append((path, text)) + duped = len(texts_with_paths) - len(unique) + return unique, duped + + +def extract_prompt(text): + upper = text.upper() + if "IDENTIFICATION DIVISION" not in upper and "PROGRAM-ID" not in upper: + return None, None, 0 + if "PROCEDURE DIVISION" not in upper: + return None, None, 0 + m = PROC_DIV_RE.search(text) + if not m: + return None, None, 0 + prompt = text[: m.end()] + body = text[m.end() :] + body_lines = [ln for ln in body.split("\n") if ln.strip()] + return prompt, body, len(body_lines) + + +def strip_exec_blocks(text): + return EXEC_BLOCK_RE.sub(" CONTINUE", text) + + +def strip_copy_stmts(text): + return COPY_STMT_RE.sub("", text) + + +def try_compile_free(cobol_text): + with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: + f.write(cobol_text) + f.flush() + src = f.name + try: + result = subprocess.run( + [ + "cobc", + "-fsyntax-only", + "-free", + "-frelax-syntax-checks", + "-frelax-level-hierarchy", + "-fmissing-statement=ok", + "-fmissing-period=ok", + "-w", + src, + ], + capture_output=True, + timeout=10, + ) + stderr = result.stderr.decode("utf-8", errors="replace") + return result.returncode, stderr + except subprocess.TimeoutExpired: + return -1, "timeout" + except Exception as e: + return -1, str(e) + finally: + try: + os.unlink(src) + except OSError: + pass + + +def try_compile_fixed(cobol_text): + with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: + f.write(cobol_text) + f.flush() + src = f.name + try: + result = subprocess.run( + [ + "cobc", + "-fsyntax-only", + "-fixed", + "-frelax-syntax-checks", + "-frelax-level-hierarchy", + "-fmissing-statement=ok", + "-fmissing-period=ok", + "-w", + src, + ], + capture_output=True, + timeout=10, + ) + stderr = result.stderr.decode("utf-8", errors="replace") + return result.returncode, stderr + except subprocess.TimeoutExpired: + return -1, "timeout" + except Exception as e: + return -1, str(e) + finally: + try: + os.unlink(src) + except OSError: + pass + + +def is_copybook_only_error(stderr): + if not stderr.strip(): + return False + for line in stderr.splitlines(): + low = line.lower() + if "error:" in low: + if "copy" not in low or "not found" not in low: + return False + return True + + +def compile_check(text): + rc, stderr = try_compile_free(text) + if rc == 0: + return "free", rc, stderr + if is_copybook_only_error(stderr): + return "free-copybook-only", 0, stderr + rc2, stderr2 = try_compile_fixed(text) + if rc2 == 0: + return "fixed", rc2, stderr2 + if is_copybook_only_error(stderr2): + return "fixed-copybook-only", 0, stderr2 + return "failed", rc, stderr + + +def process_one(path_and_text): + path, raw_text = path_and_text + result = {"path": path, "status": None} + + text = raw_text.replace("\r\n", "\n").replace("\r", "\n") + text = text.rstrip() + + text_stripped = strip_exec_blocks(text) + text_stripped = strip_copy_stmts(text_stripped) + + prompt, body, body_lines = extract_prompt(text_stripped) + if prompt is None: + result["status"] = "no_proc_div" + return result + if body_lines < MIN_BODY_LINES: + result["status"] = "body_too_short" + return result + if body_lines > MAX_BODY_LINES: + result["status"] = "body_too_long" + return result + if len(prompt) > MAX_PROMPT_CHARS: + result["status"] = "prompt_too_long" + return result + + full = prompt + body + fmt, rc, stderr = compile_check(full) + if rc != 0: + result["status"] = "compile_fail" + result["compile_fmt"] = fmt + return result + + result["status"] = "ok" + result["compile_fmt"] = fmt + result["body_lines"] = body_lines + result["text"] = text_stripped + return result + + +def main(): + log("=== RL Prompt Extraction Pipeline ===") + log(f"RAW={RAW}, OUT={OUT}, workers={MAX_WORKERS}") + log(f"Body lines: {MIN_BODY_LINES}-{MAX_BODY_LINES}") + + subprocess.run(["cobc", "--version"], capture_output=True, check=True) + log("cobc available") + + file_paths = discover_files() + + log("Reading source files...") + texts = [] + read_errors = 0 + for fpath in file_paths: + try: + with open(fpath, "r", errors="replace") as f: + content = f.read() + if len(content.strip()) > 50: + texts.append((fpath, content)) + except Exception: + read_errors += 1 + log(f"Read {len(texts)} files ({read_errors} errors)") + + texts, duped = exact_dedup(texts) + log(f"After exact dedup: {len(texts)} unique ({duped} duplicates removed)") + + log(f"Processing {len(texts)} files with {MAX_WORKERS} workers...") + results = [] + done = 0 + ok_count = 0 + with ProcessPoolExecutor(max_workers=MAX_WORKERS) as pool: + futures = {pool.submit(process_one, item): item[0] for item in texts} + for future in as_completed(futures): + try: + r = future.result() + results.append(r) + if r["status"] == "ok": + ok_count += 1 + except Exception as e: + results.append( + {"path": futures[future], "status": "error", "error": str(e)} + ) + done += 1 + if done % 5000 == 0: + log( + f" ... {done}/{len(texts)} processed, {ok_count} compilable so far" + ) + + for r in results: + stats[r["status"]] += 1 + + log("Processing complete:") + for status, count in stats.most_common(): + log(f" {status}: {count}") + + ok_results = [r for r in results if r["status"] == "ok"] + log(f"Writing {len(ok_results)} compilable programs to rl-prompts.jsonl...") + + out_path = os.path.join(OUT, "rl-prompts.jsonl") + with open(out_path, "w") as f: + for r in ok_results: + f.write( + json.dumps( + { + "text": r["text"], + "source": r["path"], + "compile_fmt": r["compile_fmt"], + "body_lines": r["body_lines"], + } + ) + + "\n" + ) + + out_size = os.path.getsize(out_path) + log(f"Wrote {out_path} ({out_size / 1024 / 1024:.1f} MB)") + + report = { + "total_files_discovered": len(file_paths), + "total_read": len(texts) + duped, + "exact_dupes_removed": duped, + "after_dedup": len(texts), + "status_counts": dict(stats.most_common()), + "compilable_programs": len(ok_results), + "output_file": out_path, + "output_size_bytes": out_size, + "body_lines_range": [MIN_BODY_LINES, MAX_BODY_LINES], + } + + body_lines = [r["body_lines"] for r in ok_results] + if body_lines: + body_lines.sort() + report["body_lines_stats"] = { + "min": body_lines[0], + "median": body_lines[len(body_lines) // 2], + "p90": body_lines[int(len(body_lines) * 0.9)], + "max": body_lines[-1], + "mean": sum(body_lines) / len(body_lines), + } + + fmt_counts = Counter(r["compile_fmt"] for r in ok_results) + report["compile_format_counts"] = dict(fmt_counts.most_common()) + + report_path = os.path.join(OUT, "rl-prep-report.json") + with open(report_path, "w") as f: + json.dump(report, f, indent=2) + log(f"Report: {report_path}") + + log(f"=== Done: {len(ok_results)} RL-ready prompts ===") + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml b/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml new file mode 100644 index 00000000..bd2c6363 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml @@ -0,0 +1,34 @@ +name: cobol-rl-prompts +engine: job +entrypoint: prepare_rl_prompts.py + +experiment: + project: cobol-cpt-grpo + name: rl-prompts + group: data + +runtime: + image: registry.example.com/cobol-train:v0.28.0-ray2.56.0 + env: + RAW: /data/datasets/cobol/raw + DATA_ROOT: /data/datasets/cobol + OUT: /data/datasets/cobol/prepared + MAX_WORKERS: "8" + MIN_BODY_LINES: "5" + MAX_BODY_LINES: "200" + MAX_PROMPT_CHARS: "50000" + PYTHONUNBUFFERED: "1" + +compute: + gpus: 0 + cpu_request: "4" + memory_request: 16Gi + cpu_limit: "8" + memory_limit: 32Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/rl-prompts + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/06-grpo/driver.py b/examples/cobol-cpt-grpo/06-grpo/driver.py new file mode 100644 index 00000000..a15d2bff --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/driver.py @@ -0,0 +1,202 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""RayJob entrypoint for COBOL GRPO training. + +TauGrid runs this file as the Ray driver on a CPU-only head pod, so it must not import +trl, vllm, transformers or the training module. The run uses the two GPU worker pods +of one Ray cluster: + +1. A VLLMServer actor reserves all 8 GPUs of one worker and runs `trl vllm-serve` + (TP=8) as a subprocess. The driver waits for + http://:/health. +2. A Ray Train TorchTrainer starts 8 workers, one GPU each. The server holds every GPU + of its node, so they land on the other worker. Each sets VLLM_SERVER_URL and calls + train_grpo.main(). + +Ray Train initializes torch.distributed (NCCL on GPUs) and sets RANK, LOCAL_RANK, +WORLD_SIZE and MASTER_ADDR/PORT in every worker before train_fn runs, the same contract +torchrun gives. train_grpo is imported inside train_fn; TauGrid ships this +directory as the Ray working_dir, so the import resolves on the workers. Configuration +(MODEL_ID, OUTPUT_DIR, K_SAMPLES, MAX_STEPS, ...) comes from the pod env, which TauGrid +sets on the head and on every worker. + +TRL's weight sync does not go through this driver: trainer rank 0 opens an NCCL group +on the group port (51216 in TRL 1.13.0) and the server's TP workers join it. + +HF Trainer keeps its own checkpointing and resume under OUTPUT_DIR; Ray Train only +stores its run metadata under OUTPUT_DIR/ray. The server is stopped when training +returns or fails. +""" + +import os +import signal +import subprocess +import time +import urllib.request + +import ray +from ray.train import RunConfig, ScalingConfig +from ray.train.torch import TorchTrainer + +SERVER_GPUS = 8 +NUM_WORKERS = 8 +HEALTH_TIMEOUT_S = 900 +STOP_TIMEOUT_S = 30 +SERVER_ACTOR_NAME = "vllm-server" + + +def vllm_serve_command(): + """The `trl vllm-serve` command line, built from the pod env.""" + model_id = os.environ.get( + "MODEL_ID", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/final" + ) + port = int(os.environ.get("VLLM_PORT", "8000")) + tp_size = int(os.environ.get("TP_SIZE", str(SERVER_GPUS))) + max_prompt = int(os.environ.get("MAX_PROMPT_LEN", "3072")) + max_gen = int(os.environ.get("MAX_GEN_LEN", "1024")) + max_len = max_prompt + max_gen + 256 + return [ + "trl", + "vllm-serve", + "--model", + model_id, + "--port", + str(port), + "--dtype", + "bfloat16", + "--tensor-parallel-size", + str(tp_size), + "--gpu-memory-utilization", + os.environ.get("GPU_MEM_UTIL", "0.90"), + "--max-model-len", + str(max_len), + "--trust-remote-code", + ] + + +@ray.remote(num_gpus=SERVER_GPUS) +class VLLMServer: + """Owns the vLLM server subprocess on a GPU worker.""" + + def __init__(self, cmd): + self._cmd = cmd + self._proc = None + + def start(self): + print(f"Starting vLLM server: {' '.join(self._cmd)}", flush=True) + self._proc = subprocess.Popen(self._cmd, start_new_session=True) + return ray.util.get_node_ip_address() + + def returncode(self): + return self._proc.poll() if self._proc else None + + def stop(self): + if self._proc is None: + return + if self._proc.poll() is None: + self._proc.terminate() + try: + self._proc.wait(timeout=STOP_TIMEOUT_S) + except subprocess.TimeoutExpired: + pass + try: + os.killpg(self._proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError): + pass + self._proc.wait() + + +def wait_for_health(server, base_url, timeout_s): + """Poll /health until it answers, the server exits, or timeout_s passes.""" + health_url = base_url + "/health" + start = time.monotonic() + next_log = 0 + while True: + rc = ray.get(server.returncode.remote()) + if rc is not None: + raise RuntimeError( + f"vLLM server exited with code {rc} before {health_url} was ready" + ) + try: + with urllib.request.urlopen(health_url, timeout=3): + print(f"vLLM ready at {base_url}", flush=True) + return + except OSError: + pass + elapsed = time.monotonic() - start + if elapsed >= timeout_s: + raise TimeoutError(f"vLLM at {health_url} not ready after {timeout_s:.0f}s") + if elapsed >= next_log: + print(f" Waiting for vLLM at {health_url}... ({elapsed:.0f}s)", flush=True) + next_log += 30 + time.sleep(1) + + +def stop_server(server): + try: + ray.get(server.stop.remote(), timeout=STOP_TIMEOUT_S + 30) + except ray.exceptions.RayError as e: + print(f"WARNING: stopping the vLLM server failed: {e}", flush=True) + ray.kill(server) + + +def train_fn(config): + os.environ["VLLM_SERVER_URL"] = config["vllm_server_url"] + import train_grpo + + train_grpo.main() + + +def build_trainer( + vllm_server_url, num_workers=NUM_WORKERS, use_gpu=True, storage_path=None +): + if storage_path is None: + output_dir = os.environ.get( + "OUTPUT_DIR", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo" + ) + storage_path = os.path.join(output_dir, "ray") + return TorchTrainer( + train_loop_per_worker=train_fn, + train_loop_config={"vllm_server_url": vllm_server_url}, + scaling_config=ScalingConfig(num_workers=num_workers, use_gpu=use_gpu), + run_config=RunConfig(storage_path=storage_path), + ) + + +def run( + server_cmd=None, + port=None, + num_workers=NUM_WORKERS, + use_gpu=True, + storage_path=None, + health_timeout_s=HEALTH_TIMEOUT_S, +): + if server_cmd is None: + server_cmd = vllm_serve_command() + if port is None: + port = int(os.environ.get("VLLM_PORT", "8000")) + + server = VLLMServer.options(name=SERVER_ACTOR_NAME).remote(server_cmd) + try: + deadline = time.monotonic() + health_timeout_s + ip = ray.get(server.start.remote(), timeout=health_timeout_s) + base_url = f"http://{ip}:{port}" + wait_for_health(server, base_url, max(0.0, deadline - time.monotonic())) + build_trainer( + base_url, + num_workers=num_workers, + use_gpu=use_gpu, + storage_path=storage_path, + ).fit() + finally: + stop_server(server) + + +def main(): + ray.init(address="auto", ignore_reinit_error=True) + run() + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/06-grpo/tau.yaml b/examples/cobol-cpt-grpo/06-grpo/tau.yaml new file mode 100644 index 00000000..077a01a5 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/tau.yaml @@ -0,0 +1,59 @@ +name: cobol-grpo +engine: rayjob +run: + entrypoint: driver.py + working_dir: . + working_dir_excludes: + - __pycache__ + - .pytest_cache + - .ruff_cache + - test_*.py + +experiment: + project: cobol-cpt-grpo + name: grpo + group: grpo-ray + +runtime: + image: registry.example.com/cobol-train:v0.28.0-ray2.56.0 + rdma: + enabled: true + env: + MODEL_ID: /data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/final + TRAIN_FILE: /data/datasets/cobol/prepared/rl-prompts.jsonl + OUTPUT_DIR: /data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo + NUM_EPOCHS: "3" + LR: "5e-7" + K_SAMPLES: "4" + TEMPERATURE: "0.8" + KL_COEFF: "0.05" + BATCH_SIZE: "4" + GRAD_ACCUM: "4" + MAX_PROMPT_LEN: "3072" + MAX_GEN_LEN: "1024" + SAVE_STEPS: "50" + PYTHONUNBUFFERED: "1" + HF_HOME: /tmp/hf-cache + PYTORCH_CUDA_ALLOC_CONF: expandable_segments:True + +compute: + workers: 2 + gpus_per_worker: 8 + worker_cpu_request: "32" + worker_memory_request: 768Gi + worker_cpu_limit: "64" + worker_memory_limit: 1500Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo + +metrics: + history: + - metrics-history-attempt-0/*.jsonl + offload: + enabled: true + +policy: + workspace: taugrid-default + topology: same-network-domain diff --git a/examples/cobol-cpt-grpo/06-grpo/test_driver.py b/examples/cobol-cpt-grpo/06-grpo/test_driver.py new file mode 100644 index 00000000..75ad05e0 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/test_driver.py @@ -0,0 +1,272 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local test for the GRPO RayJob driver. + +Runs the driver on a local Ray cluster with 8 logical GPUs. A tiny HTTP server stands +in for `trl vllm-serve` and a stub stands in for train_grpo on 2 CPU workers (gloo). +Checks the health wait, the fast failure when the server or its actor dies before it +is ready, the cleanup of the actor and its subprocess, the env each trainer worker sees, +and that the driver process never imported trl, vllm or transformers. Needs ray[train] +and CPU torch, not trl or a GPU: + python -m pytest examples/cobol-cpt-grpo/06-grpo/test_driver.py -v +""" + +import importlib.util +import json +import os +import shutil +import socket +import sys +import time + +import pytest +import ray + +RL_DIR = os.path.dirname(os.path.abspath(__file__)) + +FAKE_SERVER = """\ +import http.server +import json +import os +import signal +import sys +import time + +port, delay, mode, out = int(sys.argv[1]), float(sys.argv[2]), sys.argv[3], sys.argv[4] +with open(out, "w") as f: + json.dump({"pid": os.getpid(), "cuda": os.environ.get("CUDA_VISIBLE_DEVICES")}, f) +time.sleep(delay) +if mode == "kill-parent": + os.kill(os.getppid(), signal.SIGKILL) + sys.exit(0) +if mode.startswith("exit-"): + sys.exit(int(mode[len("exit-"):])) + + +class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + self.send_response(200 if self.path == "/health" else 404) + self.end_headers() + + def log_message(self, *args): + pass + + +http.server.HTTPServer(("0.0.0.0", port), Handler).serve_forever() +""" + +STUB_TRAIN_GRPO = """\ +import json +import os +import urllib.request + +import torch +import torch.distributed as dist +import transformers # noqa: F401 +import trl # noqa: F401 +import vllm # noqa: F401 + + +def main(): + t = torch.ones(1) + dist.all_reduce(t) + url = os.environ["VLLM_SERVER_URL"] + with urllib.request.urlopen(url + "/health", timeout=5) as resp: + health = resp.status + rank = os.environ["RANK"] + with open(os.path.join(os.environ["STUB_OUT"], f"rank-{rank}.json"), "w") as f: + json.dump({ + "rank": int(rank), + "local_rank": int(os.environ["LOCAL_RANK"]), + "world_size": int(os.environ["WORLD_SIZE"]), + "master_addr": os.environ.get("MASTER_ADDR"), + "master_port": os.environ.get("MASTER_PORT"), + "vllm_server_url": url, + "health": health, + "dist_initialized": dist.is_initialized(), + "backend": dist.get_backend(), + "all_reduce": t.item(), + "pid": os.getpid(), + }, f) +""" + +FAKE_MODULES = ("trl", "vllm", "transformers") + + +def _free_port(): + with socket.socket() as s: + s.bind(("", 0)) + return s.getsockname()[1] + + +def _pid_alive(pid): + try: + os.kill(pid, 0) + except ProcessLookupError: + return False + return True + + +@pytest.fixture +def work_dir(tmp_path, monkeypatch): + work = tmp_path / "work" + work.mkdir() + shutil.copy(os.path.join(RL_DIR, "driver.py"), work / "driver.py") + (work / "train_grpo.py").write_text(STUB_TRAIN_GRPO) + (work / "fake_server.py").write_text(FAKE_SERVER) + for name in FAKE_MODULES: + (work / f"{name}.py").write_text("") + monkeypatch.syspath_prepend(str(work)) + for name in ("driver", "train_grpo", *FAKE_MODULES): + monkeypatch.delitem(sys.modules, name, raising=False) + return work + + +@pytest.fixture +def ray_cluster(work_dir, tmp_path): + out = tmp_path / "out" + out.mkdir() + ray.init( + num_cpus=4, + num_gpus=8, + include_dashboard=False, + runtime_env={"working_dir": str(work_dir), "env_vars": {"STUB_OUT": str(out)}}, + ) + try: + yield out + finally: + ray.shutdown() + + +def _server_cmd(work_dir, port, delay, mode, record): + return [ + sys.executable, + str(work_dir / "fake_server.py"), + str(port), + str(delay), + mode, + str(record), + ] + + +def _assert_cleaned_up(record): + import driver + + with pytest.raises(ValueError): + ray.get_actor(driver.SERVER_ACTOR_NAME) + pid = json.loads(record.read_text())["pid"] + deadline = time.monotonic() + 10 + while _pid_alive(pid) and time.monotonic() < deadline: + time.sleep(0.2) + assert not _pid_alive(pid) + + +def test_defaults_match_plan(work_dir, monkeypatch): + env = { + "MODEL_ID": "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/final", + "MAX_PROMPT_LEN": "3072", + "MAX_GEN_LEN": "1024", + "OUTPUT_DIR": "/data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo", + } + for key in ("VLLM_PORT", "TP_SIZE", "GPU_MEM_UTIL"): + monkeypatch.delenv(key, raising=False) + for key, value in env.items(): + monkeypatch.setenv(key, value) + import driver + + assert driver.vllm_serve_command() == [ + "trl", + "vllm-serve", + "--model", + "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/final", + "--port", + "8000", + "--dtype", + "bfloat16", + "--tensor-parallel-size", + "8", + "--gpu-memory-utilization", + "0.90", + "--max-model-len", + "4352", + "--trust-remote-code", + ] + assert driver.VLLMServer._default_options["num_gpus"] == 8 + assert driver.HEALTH_TIMEOUT_S == 900 + + trainer = driver.build_trainer("http://10.0.0.1:8000") + assert trainer.scaling_config.num_workers == 8 + assert trainer.scaling_config.use_gpu is True + assert ( + trainer.run_config.storage_path + == "/data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo/ray" + ) + assert trainer.train_loop_config == {"vllm_server_url": "http://10.0.0.1:8000"} + + +def test_run_waits_for_health_trains_and_cleans_up(work_dir, ray_cluster, tmp_path): + import driver + + port = _free_port() + record = tmp_path / "server.json" + driver.run( + server_cmd=_server_cmd(work_dir, port, 3, "serve", record), + port=port, + num_workers=2, + use_gpu=False, + storage_path=str(tmp_path / "ray"), + ) + + server = json.loads(record.read_text()) + assert len(server["cuda"].split(",")) == 8 + + results = sorted( + (json.loads(p.read_text()) for p in ray_cluster.glob("rank-*.json")), + key=lambda r: r["rank"], + ) + assert [r["rank"] for r in results] == [0, 1] + for r in results: + assert r["world_size"] == 2 + assert r["dist_initialized"] is True + assert r["backend"] == "gloo" + assert r["all_reduce"] == 2.0 + assert r["master_addr"] and r["master_port"] + assert r["vllm_server_url"] == f"http://{ray.util.get_node_ip_address()}:{port}" + assert not r["vllm_server_url"].endswith("/v1") + assert r["health"] == 200 + assert r["pid"] != os.getpid() + assert sorted(r["local_rank"] for r in results) == [0, 1] + + _assert_cleaned_up(record) + + for name in FAKE_MODULES: + assert importlib.util.find_spec(name) is not None + for name in ("train_grpo", *FAKE_MODULES): + assert name not in sys.modules + + +@pytest.mark.parametrize( + ("mode", "error"), + [ + ("exit-3", RuntimeError), + ("kill-parent", ray.exceptions.RayActorError), + ], +) +def test_run_fails_fast_when_server_dies(work_dir, ray_cluster, tmp_path, mode, error): + import driver + + port = _free_port() + record = tmp_path / "server.json" + start = time.monotonic() + with pytest.raises(error): + driver.run( + server_cmd=_server_cmd(work_dir, port, 1, mode, record), + port=port, + num_workers=2, + use_gpu=False, + storage_path=str(tmp_path / "ray"), + ) + assert time.monotonic() - start < 60 + assert not list(ray_cluster.glob("rank-*.json")) + _assert_cleaned_up(record) diff --git a/examples/cobol-cpt-grpo/06-grpo/test_reward.py b/examples/cobol-cpt-grpo/06-grpo/test_reward.py new file mode 100644 index 00000000..dea08904 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/test_reward.py @@ -0,0 +1,718 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Table-driven reward function tests. + +Run inside the training container (has torch, trl, etc.): + python -m pytest examples/cobol-cpt-grpo/06-grpo/test_reward.py -v + +Run locally (no torch/trl needed): + python examples/cobol-cpt-grpo/06-grpo/test_reward.py +""" + +import os +import subprocess +import sys +import types +import unittest +from unittest import mock + + +os.environ.setdefault("WORLD_SIZE", "2") +os.environ.setdefault("RANK", "0") +os.environ.setdefault("LOCAL_RANK", "0") +os.environ.setdefault("OUTPUT_DIR", "/tmp/test_reward_out") +os.environ.setdefault("TRAIN_FILE", "/dev/null") + +for mod_name in [ + "torch", + "torch.cuda", + "torch.distributed", + "datasets", + "transformers", + "trl", + "safetensors", + "safetensors.torch", + "safetensors._safetensors_rust", + "deepspeed", +]: + if mod_name not in sys.modules: + fake = types.ModuleType(mod_name) + if mod_name == "torch": + fake.cuda = types.ModuleType("torch.cuda") + fake.cuda.device_count = lambda: 0 + fake.bfloat16 = "bfloat16" + if mod_name == "transformers": + fake.AutoTokenizer = type("AutoTokenizer", (), {}) + fake.TrainerCallback = type("TrainerCallback", (), {}) + if mod_name == "trl": + fake.GRPOConfig = type("GRPOConfig", (), {}) + fake.GRPOTrainer = type("GRPOTrainer", (), {}) + if mod_name == "datasets": + fake.Dataset = type("Dataset", (), {}) + sys.modules[mod_name] = fake + +sys.path.insert(0, os.path.dirname(__file__)) + +from train_grpo import ( # noqa: E402 + _compilation_reward, + _count_cobol_verbs, + _count_noop_verbs, + _has_non_display_verb, + _parse_cobc, + _repeated_statement_ratio, + _strip_non_code, + _structural_reward, + _try_compile, +) + + +class TestStripNonCode(unittest.TestCase): + def test_free_comment(self): + text = "*> DISPLAY MOVE IF\n DISPLAY X\n" + stripped = _strip_non_code(text) + self.assertNotIn("MOVE", stripped) + self.assertIn("DISPLAY", stripped) + + def test_multiply_operator_at_column_7(self): + text = " IF A * B > 0\n MOVE A TO C\n DISPLAY C\n END-IF.\n STOP RUN.\n" + stripped = _strip_non_code(text) + self.assertIn("IF", stripped) + self.assertIn("MOVE", stripped) + self.assertIn("DISPLAY", stripped) + + def test_divide_operator_at_column_7(self): + text = " IF A / B > 0\n MOVE A TO C\n DISPLAY C\n END-IF.\n STOP RUN.\n" + stripped = _strip_non_code(text) + self.assertIn("IF", stripped) + self.assertIn("MOVE", stripped) + self.assertIn("DISPLAY", stripped) + + def test_six_space_D_not_stripped_in_free_format(self): + text = " DISPLAY X\n" + stripped = _strip_non_code(text) + self.assertIn("DISPLAY", stripped) + + def test_string_literal(self): + text = """ DISPLAY "MOVE IF ADD"\n""" + stripped = _strip_non_code(text) + self.assertNotIn("MOVE", stripped) + self.assertNotIn("ADD", stripped) + self.assertIn("DISPLAY", stripped) + + def test_single_quoted_literal(self): + text = " DISPLAY 'MOVE IF'\n" + stripped = _strip_non_code(text) + self.assertNotIn("MOVE", stripped) + self.assertIn("DISPLAY", stripped) + + def test_end_terminator(self): + text = " END-IF\n END-PERFORM\n" + stripped = _strip_non_code(text) + self.assertNotIn("END-IF", stripped) + self.assertNotIn("END-PERFORM", stripped) + + def test_hyphenated_identifier(self): + text = " MOVE-RECORD.\n WS-COUNTER\n" + stripped = _strip_non_code(text) + self.assertNotIn("MOVE-RECORD", stripped) + self.assertNotIn("MOVE", stripped.replace("_ID_", "")) + + +class TestCobolVerbCount(unittest.TestCase): + def test_empty(self): + self.assertEqual(_count_cobol_verbs(""), 0) + + def test_comments_containing_verbs_are_zero(self): + text = "*> DISPLAY one\n*> MOVE two\n*> IF three\n" + self.assertEqual(_count_cobol_verbs(text), 0) + + def test_string_containing_verbs(self): + text = ' DISPLAY "MOVE IF"\n' + self.assertEqual(_count_cobol_verbs(text), 1) + + def test_hyphenated_identifier_not_counted(self): + text = " MOVE-RECORD.\n" + self.assertEqual(_count_cobol_verbs(text), 0) + + def test_end_if_not_counted(self): + text = " END-IF\n END-PERFORM\n" + self.assertEqual(_count_cobol_verbs(text), 0) + + def test_display_is_meaningful(self): + self.assertEqual(_count_cobol_verbs(' DISPLAY "HELLO"\n'), 1) + + def test_stop_not_meaningful(self): + self.assertEqual(_count_cobol_verbs(" STOP RUN.\n"), 0) + + def test_goback_not_meaningful(self): + self.assertEqual(_count_cobol_verbs(" GOBACK.\n"), 0) + + def test_exit_not_meaningful(self): + self.assertEqual(_count_cobol_verbs(" EXIT.\n"), 0) + + def test_continue_not_meaningful(self): + self.assertEqual(_count_cobol_verbs(" CONTINUE.\n"), 0) + + def test_go_to_is_meaningful(self): + self.assertEqual(_count_cobol_verbs(" GO TO PARA-1\n"), 1) + + def test_free_is_meaningful(self): + self.assertEqual(_count_cobol_verbs(" FREE WS-PTR\n"), 1) + + def test_six_space_display_counted(self): + text = " DISPLAY X\n DISPLAY Y\n DISPLAY Z\n" + self.assertEqual(_count_cobol_verbs(text), 3) + + def test_perform_move_if_display(self): + text = ( + " PERFORM INIT-PARA\n" + " MOVE 1 TO WS-X\n" + " IF WS-X > 0\n" + " DISPLAY WS-X\n" + " END-IF\n" + " STOP RUN.\n" + ) + self.assertEqual(_count_cobol_verbs(text), 4) + + def test_case_insensitive(self): + self.assertEqual(_count_cobol_verbs("display move perform"), 3) + + def test_multiply_at_column_7_not_stripped(self): + text = " IF A * B > 0\n MOVE A TO C\n DISPLAY C\n END-IF.\n STOP RUN.\n" + self.assertEqual(_count_cobol_verbs(text), 3) + + def test_divide_at_column_7_not_stripped(self): + text = " IF A / B > 0\n MOVE A TO C\n DISPLAY C\n END-IF.\n STOP RUN.\n" + self.assertEqual(_count_cobol_verbs(text), 3) + + +class TestNoopVerbCount(unittest.TestCase): + def test_continue_is_noop(self): + self.assertEqual(_count_noop_verbs(" CONTINUE.\n"), 1) + + def test_display_not_noop(self): + self.assertEqual(_count_noop_verbs(" DISPLAY X.\n"), 0) + + def test_stop_not_noop(self): + self.assertEqual(_count_noop_verbs(" STOP RUN.\n"), 0) + + def test_goback_not_noop(self): + self.assertEqual(_count_noop_verbs(" GOBACK.\n"), 0) + + def test_exit_not_noop(self): + self.assertEqual(_count_noop_verbs(" EXIT.\n"), 0) + + def test_go_not_noop(self): + self.assertEqual(_count_noop_verbs(" GO TO PARA-1.\n"), 0) + + +class TestHasNonDisplayVerb(unittest.TestCase): + def test_display_only(self): + text = ( + ' DISPLAY "A"\n DISPLAY "B"\n DISPLAY "C"\n' + ) + self.assertFalse(_has_non_display_verb(text)) + + def test_move_present(self): + text = " MOVE 1 TO X\n DISPLAY X\n" + self.assertTrue(_has_non_display_verb(text)) + + def test_go_to_present(self): + text = " GO TO PARA-1\n DISPLAY X\n" + self.assertTrue(_has_non_display_verb(text)) + + def test_empty(self): + self.assertFalse(_has_non_display_verb("")) + + def test_verb_in_literal_with_comment_marker(self): + text = ( + ' DISPLAY "MOVE *> A".\n' + ' DISPLAY "MOVE *> B".\n' + ' DISPLAY "MOVE *> C".\n' + ) + self.assertFalse(_has_non_display_verb(text)) + + +class TestVerbCountWithCommentInLiteral(unittest.TestCase): + def test_move_inside_literal_not_counted(self): + text = ( + ' DISPLAY "MOVE *> A".\n' + ' DISPLAY "MOVE *> B".\n' + ' DISPLAY "MOVE *> C".\n' + " STOP RUN.\n" + ) + self.assertEqual(_count_cobol_verbs(text), 3) + + def test_structural_penalty_for_display_only_with_literal_verbs(self): + completion = ( + ' DISPLAY "MOVE *> A".\n' + ' DISPLAY "MOVE *> B".\n' + ' DISPLAY "MOVE *> C".\n' + " STOP RUN.\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.20) + + +class TestRepeatedStatementRatio(unittest.TestCase): + def test_all_same(self): + text = ' DISPLAY "A".\n' * 5 + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_all_different(self): + text = ( + " MOVE 1 TO X.\n" + " ADD 1 TO Y.\n" + " DISPLAY X.\n" + " DISPLAY Y.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertLessEqual(ratio, 0.5) + + def test_short_text(self): + self.assertEqual(_repeated_statement_ratio(" DISPLAY X.\n"), 0.0) + + def test_distinct_moves_with_different_identifiers(self): + text = ( + " MOVE WS-A TO WS-B.\n" + " MOVE WS-C TO WS-D.\n" + " MOVE WS-E TO WS-F.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertLessEqual(ratio, 0.5) + + def test_same_statement_different_case_and_spacing(self): + text = ( + " MOVE 1 TO WS-X.\n" + " move 1 to ws-x.\n" + " MOVE\t1 TO WS-X.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_distinct_call_literals_not_repeated(self): + text = ( + ' CALL "PROG-A" USING WS-X.\n' + ' CALL "PROG-B" USING WS-X.\n' + ' CALL "PROG-C" USING WS-X.\n' + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertLessEqual(ratio, 0.5) + + def test_comment_marker_inside_literal_preserved(self): + text = ( + ' DISPLAY "PHASE *> A".\n' + ' DISPLAY "PHASE *> B".\n' + ' DISPLAY "PHASE *> C".\n' + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertLessEqual(ratio, 0.5) + + def test_same_statement_split_across_lines(self): + text = ( + " MOVE 1\n" + " TO WS-X.\n" + " MOVE 1\n" + " TO WS-X.\n" + " MOVE 1\n" + " TO WS-X.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_same_statement_different_line_breaks(self): + text = ( + " MOVE 1 TO WS-X.\n" + " MOVE 1\n" + " TO WS-X.\n" + " MOVE\n" + " 1 TO\n" + " WS-X.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_intra_sentence_repetition(self): + text = ( + " MOVE 1 TO WS-Y\n" + " MOVE 1 TO WS-Y\n" + " MOVE 1 TO WS-Y\n" + " DISPLAY WS-Y.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_decimal_literal_not_split(self): + text = ( + " MOVE 1.1 TO WS-X.\n" + " MOVE 1.1 TO WS-X.\n" + " MOVE 1.1 TO WS-X.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_attached_comment_distinct_code(self): + text = ( + ' DISPLAY "A".*> inline note\n' + ' DISPLAY "B".*> inline note\n' + ' DISPLAY "C".*> inline note\n' + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertLessEqual(ratio, 0.5) + + def test_attached_comment_same_code(self): + text = ( + " MOVE 1 TO WS-Y*> first\n" + " MOVE 1 TO WS-Y*> second\n" + " MOVE 1 TO WS-Y*> third\n" + " DISPLAY WS-Y.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + def test_spaced_comment_same_code(self): + text = ( + " MOVE 1 TO WS-Y *> comment A\n" + " MOVE 1 TO WS-Y *> comment B\n" + " MOVE 1 TO WS-Y *> comment C\n" + " DISPLAY WS-Y.\n" + " STOP RUN.\n" + ) + ratio = _repeated_statement_ratio(text) + self.assertGreater(ratio, 0.5) + + +class TestCompilationReward(unittest.TestCase): + def test_clean_compile(self): + self.assertEqual(_compilation_reward(0, ""), 1.0) + + def test_warnings_only(self): + stderr = "prog.cbl:10: warning: some warning" + self.assertEqual(_compilation_reward(0, stderr), 0.7) + + def test_timeout_raises(self): + with self.assertRaises(RuntimeError): + _compilation_reward(-1, "compilation timeout") + + def test_cobc_not_found_raises(self): + with self.assertRaises(RuntimeError): + _compilation_reward(-1, "cobc not found") + + def test_unknown_infra_failure_raises(self): + with self.assertRaises(RuntimeError): + _compilation_reward(-1, "something unexpected") + + def test_cobc_signal_crash_scores_as_failed_compile(self): + stderr = ( + "x.cbl:15: error: 'LINE-NUMBER' is not defined\n" + "x.cbl:17: error: cannot specify both COLUMN and AT screen-location\n" + "x.cbl:23: error: syntax error, unexpected DEPENDING\n" + "free(): invalid size\n" + ) + crashed = subprocess.CompletedProcess(["cobc"], -6, b"", stderr.encode()) + with mock.patch("train_grpo.subprocess.run", return_value=crashed): + rc, err = _try_compile("IDENTIFICATION DIVISION.") + self.assertEqual(rc, 1) + self.assertAlmostEqual(_compilation_reward(rc, err), -0.6) + + def test_copybook_only(self): + stderr = "prog.cbl:5: error: COPY-FILE.cpy: No such file or directory" + self.assertEqual(_compilation_reward(1, stderr), 0.0) + + def test_one_error(self): + stderr = "prog.cbl:10: error: 'WS-FOO' is not defined" + self.assertAlmostEqual(_compilation_reward(1, stderr), -0.2) + + def test_two_errors(self): + stderr = ( + "prog.cbl:10: error: 'WS-FOO' is not defined\n" + "prog.cbl:11: error: 'WS-BAR' is not defined" + ) + self.assertAlmostEqual(_compilation_reward(1, stderr), -0.4) + + def test_five_errors_clamped(self): + lines = [f"prog.cbl:{i}: error: err{i}" for i in range(5)] + self.assertAlmostEqual(_compilation_reward(1, "\n".join(lines)), -1.0) + + def test_six_errors_clamped(self): + lines = [f"prog.cbl:{i}: error: err{i}" for i in range(6)] + self.assertAlmostEqual(_compilation_reward(1, "\n".join(lines)), -1.0) + + def test_no_errors_no_warnings_nonzero_rc(self): + self.assertEqual(_compilation_reward(1, ""), -1.0) + + def test_mixed_copybook_and_errors(self): + stderr = ( + "prog.cbl:5: error: COPY-FILE.cpy: No such file or directory\n" + "prog.cbl:10: error: 'WS-FOO' is not defined" + ) + self.assertAlmostEqual(_compilation_reward(1, stderr), -0.2) + + +class TestStructuralReward(unittest.TestCase): + def test_comment_only_hard_reject(self): + completion = "*> filler\n" * 10 + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.50) + + def test_two_meaningful_verb_program(self): + completion = ' DISPLAY "HELLO"\n MOVE 1 TO X.\n' + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.50) + + def test_three_displays_only_penalized(self): + completion = ( + ' DISPLAY "A"\n DISPLAY "B"\n DISPLAY "C"\n' + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.20) + + def test_three_identical_displays_penalized(self): + completion = ' DISPLAY "A"\n' * 3 + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.20) + + def test_three_continues_degenerate(self): + completion = " CONTINUE.\n" * 3 + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.50) + + def test_noop_majority_penalty(self): + completion = ( + " MOVE 1 TO X\n" + " DISPLAY Y\n" + " ADD 1 TO Z\n" + " CONTINUE.\n" + " CONTINUE.\n" + " CONTINUE.\n" + " CONTINUE.\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.30) + + def test_truncation_penalty(self): + completion = ( + " MOVE 1 TO X\n" + " ADD 1 TO Y\n" + " DISPLAY X\n" + " DISPLAY Y\n" + ) + r_normal = _structural_reward(completion, 1024, 50) + r_truncated = _structural_reward(completion, 1024, 1020) + self.assertLess(r_truncated, r_normal) + self.assertEqual(r_truncated, r_normal - 0.15) + + def test_perform_bonus_defined(self): + completion = ( + " INIT-PARA.\n" + " MOVE 1 TO WS-X.\n" + " PERFORM INIT-PARA\n" + " DISPLAY WS-X\n" + " ADD 1 TO WS-Y\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, 0.15) + + def test_dangling_perform_no_bonus(self): + completion = ( + " PERFORM NONEXISTENT-PARA\n" + " DISPLAY WS-X\n" + " MOVE 1 TO WS-Y\n" + " ADD 1 TO WS-Z\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, 0.0) + + def test_perform_until_not_paragraph(self): + completion = ( + " PERFORM UNTIL WS-EOF\n" + " READ INPUT-FILE\n" + " DISPLAY WS-REC\n" + " END-PERFORM\n" + " CLOSE INPUT-FILE\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, 0.0) + + def test_max_gen_zero_no_truncation_penalty(self): + completion = ( + " MOVE 1 TO X\n" + " ADD 1 TO Y\n" + " DISPLAY X\n" + " DISPLAY Y\n" + ) + r = _structural_reward(completion, 0, 999) + self.assertGreaterEqual(r, 0.0) + + def test_comments_with_verbs_rejected(self): + completion = ( + "*> DISPLAY something\n" + "*> MOVE something\n" + "*> IF something\n" + "*> PERFORM something\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.50) + + def test_useful_short_program_not_penalized(self): + completion = ( + " MOVE 1 TO WS-CTR\n" + " DISPLAY WS-CTR\n" + " ADD 1 TO WS-CTR\n" + " DISPLAY WS-CTR\n" + " STOP RUN.\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertGreaterEqual(r, 0.0) + + def test_distinct_moves_not_penalized(self): + completion = ( + " MOVE WS-A TO WS-B.\n" + " MOVE WS-C TO WS-D.\n" + " MOVE WS-E TO WS-F.\n" + " DISPLAY WS-B.\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertGreaterEqual(r, 0.0) + + def test_go_to_program_not_penalized(self): + completion = ( + " GO TO PARA-1\n" + " DISPLAY X\n" + " MOVE 1 TO Y\n" + " STOP RUN.\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertGreaterEqual(r, 0.0) + + def test_six_space_displays_not_stripped(self): + completion = ( + " DISPLAY X\n MOVE 1 TO Y\n ADD 1 TO Z\n DISPLAY Z\n" + ) + r = _structural_reward(completion, 1024, 50) + self.assertGreaterEqual(r, 0.0) + + def test_repeated_statements_penalized(self): + completion = " MOVE 1 TO X.\n" * 5 + " DISPLAY X.\n" + r = _structural_reward(completion, 1024, 50) + self.assertEqual(r, -0.20) + + +class TestParseCobc(unittest.TestCase): + def test_empty(self): + errors, warnings, copybook = _parse_cobc("") + self.assertEqual(errors, []) + self.assertEqual(warnings, []) + self.assertEqual(copybook, []) + + def test_error(self): + stderr = "prog.cbl:10: error: 'WS-X' is not defined" + errors, warnings, copybook = _parse_cobc(stderr) + self.assertEqual(len(errors), 1) + self.assertEqual(len(copybook), 0) + + def test_copybook(self): + stderr = "prog.cbl:5: error: COPY-FILE.cpy: No such file or directory" + errors, warnings, copybook = _parse_cobc(stderr) + self.assertEqual(len(errors), 0) + self.assertEqual(len(copybook), 1) + + def test_warning(self): + stderr = "prog.cbl:10: warning: unreachable code" + errors, warnings, copybook = _parse_cobc(stderr) + self.assertEqual(len(warnings), 1) + self.assertEqual(len(errors), 0) + + def test_mixed(self): + stderr = ( + "prog.cbl:5: error: COPY-FILE.cpy: No such file or directory\n" + "prog.cbl:10: error: 'WS-X' is not defined\n" + "prog.cbl:15: warning: unreachable code" + ) + errors, warnings, copybook = _parse_cobc(stderr) + self.assertEqual(len(errors), 1) + self.assertEqual(len(copybook), 1) + self.assertEqual(len(warnings), 1) + + +class TestScoreIntegration(unittest.TestCase): + def test_clean_compile_plus_good_structure(self): + r_compile = _compilation_reward(0, "") + completion = ( + " INIT-PARA.\n" + " MOVE 1 TO WS-X.\n" + " PERFORM INIT-PARA\n" + " DISPLAY WS-X\n" + " ADD 1 TO WS-Y\n" + ) + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 1.15) + + def test_one_error_trivial_completion(self): + r_compile = _compilation_reward(1, "prog.cbl:10: error: 'X' is not defined") + r_struct = _structural_reward(" DISPLAY 1\n", 1024, 10) + total = r_compile + r_struct + self.assertEqual(total, -0.7) + + def test_comment_exploit_blocked(self): + r_compile = _compilation_reward(0, "") + completion = "*> DISPLAY MOVE IF ADD COMPUTE\n" * 5 + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 0.5) + + def test_three_continues_compile_but_degenerate(self): + r_compile = _compilation_reward(0, "") + completion = " CONTINUE.\n" * 3 + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 0.5) + + def test_three_displays_not_max_reward(self): + r_compile = _compilation_reward(0, "") + completion = ( + ' DISPLAY "A"\n DISPLAY "B"\n DISPLAY "C"\n' + ) + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 0.80) + + def test_useful_program_max_reward(self): + r_compile = _compilation_reward(0, "") + completion = ( + " MOVE 1 TO WS-CTR\n" + " ADD 1 TO WS-CTR\n" + " DISPLAY WS-CTR\n" + " STOP RUN.\n" + ) + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 1.0) + + def test_go_to_program_gets_full_reward(self): + r_compile = _compilation_reward(0, "") + completion = ( + " GO TO PARA-1\n" + " DISPLAY X\n" + " MOVE 1 TO Y\n" + " STOP RUN.\n" + ) + r_struct = _structural_reward(completion, 1024, 50) + total = r_compile + r_struct + self.assertEqual(total, 1.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/cobol-cpt-grpo/06-grpo/train_grpo.py b/examples/cobol-cpt-grpo/06-grpo/train_grpo.py new file mode 100755 index 00000000..ecbb962f --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/train_grpo.py @@ -0,0 +1,851 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +GRPO training with compiler-as-teacher reward for COBOL code generation. + +Runs under Ray Train: driver.py starts the `trl vllm-serve` server (TP=8) on one +node, then a TorchTrainer whose 8 workers, one per GPU on the other node, each call +main(). Ray Train initializes torch.distributed and sets RANK, LOCAL_RANK and +WORLD_SIZE before main() runs, and the driver sets VLLM_SERVER_URL. Importing this +module has no side effects. +""" + +import json +import math +import os +import re +import subprocess +import tempfile +import time +from concurrent.futures import ProcessPoolExecutor +from pathlib import Path + +import torch +from datasets import Dataset +from transformers import AutoTokenizer, TrainerCallback +from trl import GRPOConfig, GRPOTrainer + +VLLM_PORT = 8000 + +LOCAL_RANK = 0 +GLOBAL_RANK = 0 +WORLD_SIZE = 1 + +MODEL_ID = os.environ.get( + "MODEL_ID", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/cpt/final" +) +TRAIN_FILE = os.environ.get("TRAIN_FILE", "/data/datasets/cobol/prepared/train.jsonl") +OUTPUT_DIR = os.environ.get( + "OUTPUT_DIR", "/data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo" +) +NUM_EPOCHS = int(os.environ.get("NUM_EPOCHS", "3")) +LR = float(os.environ.get("LR", "5e-7")) +K_SAMPLES = int(os.environ.get("K_SAMPLES", "8")) +TEMPERATURE = float(os.environ.get("TEMPERATURE", "0.8")) +KL_COEFF = float(os.environ.get("KL_COEFF", "0.05")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "8")) +GRAD_ACCUM = int(os.environ.get("GRAD_ACCUM", "4")) +MAX_PROMPT_LEN = int(os.environ.get("MAX_PROMPT_LEN", "4096")) +MAX_GEN_LEN = int(os.environ.get("MAX_GEN_LEN", "1024")) +SAVE_STEPS = int(os.environ.get("SAVE_STEPS", "50")) +MAX_STEPS = os.environ.get("MAX_STEPS") +HF_HOME = os.environ.get("HF_HOME", "/tmp/hf-cache") + + +def extract_prompt(text): + upper = text.upper() + if "IDENTIFICATION DIVISION" not in upper or "PROCEDURE DIVISION" not in upper: + return None, None, 0 + proc_pat = re.compile( + r"^(.*PROCEDURE\s+DIVISION[^\n]*\.?)$", re.IGNORECASE | re.MULTILINE + ) + m = proc_pat.search(text) + if not m: + return None, None, 0 + prompt = text[: m.end()] + body = text[m.end() :] + body_lines = [ln for ln in body.split("\n") if ln.strip()] + return prompt, body, len(body_lines) + + +def _preflight_compilable(prompt, body): + """Check if prompt + reference body compiles cleanly (rc == 0). + Reject programs with missing copybooks — the compiler can't validate + the procedure body when dependencies are unresolved, so these produce + all-failure groups regardless of completion quality.""" + full = prompt + body + rc, _ = _try_compile(full) + return rc == 0 + + +def _check_compiler_available(): + try: + result = subprocess.run( + ["cobc", "--version"], + capture_output=True, + timeout=10, + ) + if result.returncode != 0: + raise RuntimeError(f"cobc exited with code {result.returncode}") + except FileNotFoundError: + raise RuntimeError("cobc not found — GnuCOBOL must be installed") + except subprocess.TimeoutExpired: + raise RuntimeError("cobc --version timed out") + + +def load_prompts( + path, min_body=5, max_body=200, preflight=True, max_body_tokens=None, tokenizer=None +): + _check_compiler_available() + + candidates = [] + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + rec = json.loads(line) + prompt, body, body_len = extract_prompt(rec.get("text", "")) + if prompt is None or body_len < min_body or body_len > max_body: + continue + candidates.append((prompt, body)) + + if max_body_tokens and tokenizer: + before = len(candidates) + candidates = [ + (p, b) + for p, b in candidates + if len(tokenizer.encode(b, add_special_tokens=False)) <= max_body_tokens + ] + if GLOBAL_RANK == 0 and len(candidates) < before: + print( + f" Body token filter: {before - len(candidates)} dropped " + f"(>{max_body_tokens} tokens)", + flush=True, + ) + + if not preflight: + return [p for p, _ in candidates] + + prompts = [] + skipped = 0 + for prompt, body in candidates: + if _preflight_compilable(prompt, body): + prompts.append(prompt) + else: + skipped += 1 + + if GLOBAL_RANK == 0: + print(f" Preflight: {len(prompts)} compilable, {skipped} skipped", flush=True) + if not prompts: + raise RuntimeError( + f"No compilable prompts found in {path} — check compiler and data" + ) + return prompts + + +def _try_compile(cobol_text): + with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: + f.write(cobol_text) + f.flush() + src = f.name + try: + result = subprocess.run( + ["cobc", "-fsyntax-only", "-free", src], + capture_output=True, + timeout=10, + ) + result = subprocess.CompletedProcess( + result.args, + result.returncode, + result.stdout.decode("utf-8", errors="replace"), + result.stderr.decode("utf-8", errors="replace"), + ) + if result.returncode < 0: + print( + f" cobc killed by signal {-result.returncode}; scoring as failed compile", + flush=True, + ) + return 1, result.stderr.strip() + return result.returncode, result.stderr.strip() + except subprocess.TimeoutExpired: + return -1, "compilation timeout" + except FileNotFoundError: + return -1, "cobc not found" + finally: + os.unlink(src) + + +def _parse_cobc(stderr): + errors, warnings, copybook = [], [], [] + for line in stderr.strip().split("\n"): + if not line.strip(): + continue + ll = line.lower() + if "error:" in ll: + if "no such file or directory" in ll: + copybook.append(line) + else: + errors.append(line) + elif "warning:" in ll: + warnings.append(line) + return errors, warnings, copybook + + +def _compilation_reward(returncode, stderr): + if returncode == 0: + errors, warnings, _ = _parse_cobc(stderr) + return 0.7 if warnings else 1.0 + if returncode < 0: + raise RuntimeError(f"Compiler infrastructure failure: {stderr}") + errors, warnings, copybook = _parse_cobc(stderr) + if not errors and not copybook: + return -1.0 + if copybook and not errors: + return 0.0 + n = len(errors) + return max(-1.0, -0.2 * n) + + +_STRING_LITERAL_RE = re.compile(r"""(?:"[^"]*"|'[^']*')""") +_END_TERMINATOR_RE = re.compile(r"\bEND-[A-Z]+\b", re.IGNORECASE) +_HYPHENATED_IDENT_RE = re.compile(r"\b[A-Z0-9]+-[A-Z0-9][-A-Z0-9]*\b", re.IGNORECASE) +_INLINE_COMMENT_RE = re.compile(r"""("[^"]*"|'[^']*')|\*>.*$""", re.MULTILINE) + + +def _strip_inline_comments(text): + def _repl(m): + if m.group(1) is not None: + return m.group(0) + return "" + + return _INLINE_COMMENT_RE.sub(_repl, text) + + +_MEANINGFUL_VERBS = re.compile( + r"\b(?:ACCEPT|ADD|ALLOCATE|ALTER|CALL|CANCEL|CLOSE|COMPUTE|" + r"DELETE|DISPLAY|DIVIDE|EVALUATE|FREE|GO|IF|INITIALIZE|INSPECT|" + r"MERGE|MOVE|MULTIPLY|OPEN|PERFORM|READ|RELEASE|RETURN|REWRITE|" + r"SEARCH|SET|SORT|START|STRING|SUBTRACT|UNSTRING|WRITE)\b", + re.IGNORECASE, +) + +_NOOP_VERBS = re.compile( + r"\b(?:CONTINUE)\b", + re.IGNORECASE, +) + +_DISPLAY_ONLY_RE = re.compile(r"\bDISPLAY\b", re.IGNORECASE) +_NON_DISPLAY_VERBS = re.compile( + r"\b(?:ACCEPT|ADD|ALLOCATE|ALTER|CALL|CANCEL|CLOSE|COMPUTE|" + r"DELETE|DIVIDE|EVALUATE|FREE|GO|IF|INITIALIZE|INSPECT|" + r"MERGE|MOVE|MULTIPLY|OPEN|PERFORM|READ|RELEASE|RETURN|REWRITE|" + r"SEARCH|SET|SORT|START|STRING|SUBTRACT|UNSTRING|WRITE)\b", + re.IGNORECASE, +) + + +def _strip_comments_and_strings(text): + text = _strip_inline_comments(text) + text = _STRING_LITERAL_RE.sub('""', text) + return text + + +def _strip_non_code(text): + text = _strip_comments_and_strings(text) + text = _END_TERMINATOR_RE.sub("", text) + text = _HYPHENATED_IDENT_RE.sub("_ID_", text) + return text + + +def _count_cobol_verbs(text): + cleaned = _strip_non_code(text) + return len(_MEANINGFUL_VERBS.findall(cleaned)) + + +def _count_noop_verbs(text): + cleaned = _strip_non_code(text) + return len(_NOOP_VERBS.findall(cleaned)) + + +def _has_non_display_verb(text): + cleaned = _strip_non_code(text) + return bool(_NON_DISPLAY_VERBS.search(cleaned)) + + +_COLLAPSE_WS_RE = re.compile(r"[ \t]+") +_COBOL_TOKEN_RE = re.compile( + r"""("[^"]*"|'[^']*'""" + r"""|\d+\.\d+""" + r"""|[^\s.]+""" + r"""|\.)\s*""", + re.MULTILINE, +) + + +def _tokenize_cobol(text): + text = _strip_inline_comments(text) + tokens = [] + for m in _COBOL_TOKEN_RE.finditer(text): + tok = m.group(1) + if tok.startswith('"') or tok.startswith("'"): + tokens.append(tok) + else: + tokens.append(tok.upper()) + return tokens + + +def _split_statements(tokens): + stmts = [] + cur = [] + for tok in tokens: + cur.append(tok) + if tok == ".": + normalized = " ".join(cur) + if normalized != ".": + stmts.append(normalized) + cur = [] + if cur: + normalized = " ".join(cur) + stmts.append(normalized) + return stmts + + +def _normalize_line(line): + s = line.strip() + if not s: + return None + parts = _STRING_LITERAL_RE.split(s) + literals = _STRING_LITERAL_RE.findall(s) + result = [] + for i, part in enumerate(parts): + result.append(_COLLAPSE_WS_RE.sub(" ", part).upper()) + if i < len(literals): + result.append(literals[i]) + return "".join(result) + + +def _line_repetition_ratio(text): + text = _strip_inline_comments(text) + lines = [] + for line in text.split("\n"): + n = _normalize_line(line) + if n: + lines.append(n) + if len(lines) < 3: + return 0.0 + from collections import Counter + + counts = Counter(lines) + return counts.most_common(1)[0][1] / len(lines) + + +def _stmt_repetition_ratio(text): + tokens = _tokenize_cobol(text) + stmts = _split_statements(tokens) + if len(stmts) < 3: + return 0.0 + from collections import Counter + + counts = Counter(stmts) + return counts.most_common(1)[0][1] / len(stmts) + + +def _repeated_statement_ratio(text): + return max(_line_repetition_ratio(text), _stmt_repetition_ratio(text)) + + +def _structural_reward(completion, max_gen_tokens, completion_n_tokens): + verb_count = _count_cobol_verbs(completion) + noop_count = _count_noop_verbs(completion) + if verb_count < 3: + return -0.50 + if verb_count + noop_count > 0 and noop_count / (verb_count + noop_count) > 0.5: + return -0.30 + if not _has_non_display_verb(completion): + return -0.20 + if _repeated_statement_ratio(completion) > 0.5: + return -0.20 + + bonus = 0.0 + + code = _strip_comments_and_strings(completion) + para_def_re = re.compile(r"^[ ]{7}([A-Z0-9][A-Z0-9-]*)\.\s*$", re.MULTILINE) + defined = {m.group(1).upper() for m in para_def_re.finditer(code)} + perform_re = re.compile(r"PERFORM[ \t]+([A-Z0-9][A-Z0-9-]+)", re.IGNORECASE) + skip = {"UNTIL", "VARYING", "TIMES", "TEST", "WITH"} + referenced, dangling = set(), set() + for m in perform_re.finditer(code): + t = m.group(1).upper() + if t in skip: + continue + referenced.add(t) + if t not in defined: + dangling.add(t) + if referenced and not dangling: + bonus += 0.15 + + if max_gen_tokens > 0 and completion_n_tokens >= max_gen_tokens - 5: + bonus -= 0.15 + + return bonus + + +def _score_one(args): + prompt, completion, max_gen_tokens, completion_n_tokens = args + full = prompt + completion + rc, stderr = _try_compile(full) + r_compile = _compilation_reward(rc, stderr) + r_struct = _structural_reward(completion, max_gen_tokens, completion_n_tokens) + errors, warnings, copybook = _parse_cobc(stderr) + verb_count = _count_cobol_verbs(completion) + noop_count = _count_noop_verbs(completion) + return { + "reward": r_compile + r_struct, + "compiled": rc == 0, + "copybook_only": bool(copybook) and not bool(errors), + "n_errors": len(errors), + "verb_count": verb_count, + "noop_count": noop_count, + "trivial": verb_count < 3, + } + + +_reward_stats = {} +_reward_tokenizer = None + + +def reward_fn(prompts, completions, **kwargs): + """Compute rewards for a batch of (prompt, completion) pairs.""" + if _reward_tokenizer is not None: + n_tokens = [ + len(_reward_tokenizer.encode(c, add_special_tokens=False)) + for c in completions + ] + else: + n_tokens = [len(c) // 4 for c in completions] + args = [(p, c, MAX_GEN_LEN, nt) for p, c, nt in zip(prompts, completions, n_tokens)] + with ProcessPoolExecutor(max_workers=min(8, len(args))) as pool: + results = list(pool.map(_score_one, args)) + + rewards = [r["reward"] for r in results] + + _reward_stats["compile_rate"] = sum(r["compiled"] for r in results) / len(results) + _reward_stats["copybook_only_rate"] = sum( + r["copybook_only"] for r in results + ) / len(results) + _reward_stats["trivial_rate"] = sum(r["trivial"] for r in results) / len(results) + _reward_stats["mean_verb_count"] = sum(r["verb_count"] for r in results) / len( + results + ) + _reward_stats["mean_noop_count"] = sum(r["noop_count"] for r in results) / len( + results + ) + _reward_stats["mean_errors"] = sum(r["n_errors"] for r in results) / len(results) + + return rewards + + +class TauMetricsPublisher: + def __init__(self, output_dir): + self.metrics_dir = Path(output_dir) / "metrics-history-attempt-0" + self.metrics_dir.mkdir(parents=True, exist_ok=True) + + def publish(self, step, metrics): + row = {"_step": step, "_timestamp": time.time()} + for k, v in metrics.items(): + if isinstance(v, (int, float)) and not math.isnan(v): + row[k] = v + token = f"{step:08d}-{time.time_ns()}" + final = self.metrics_dir / f"chunk-{token}.jsonl" + tmp = self.metrics_dir / f".{final.name}.tmp-{os.getpid()}" + with tmp.open("x", encoding="utf-8") as f: + f.write(json.dumps(row) + "\n") + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, final) + + +_REWARD_STAT_KEYS = [ + "compile_rate", + "copybook_only_rate", + "trivial_rate", + "mean_verb_count", + "mean_noop_count", + "mean_errors", +] + + +def _aggregate_reward_stats(): + return dict(_reward_stats) + + +class TauGRPOMetricsCallback(TrainerCallback): + def __init__(self, publisher): + self.publisher = publisher + self.kl_alarm_count = 0 + self.short_alarm_count = 0 + + def on_log(self, args, state, control, logs=None, **kwargs): + if not logs or state.global_step <= 0: + return + if GLOBAL_RANK != 0: + return + + agg = _aggregate_reward_stats() + + metrics = {} + for k, v in logs.items(): + if isinstance(v, (int, float)) and not math.isnan(v): + metrics[k] = v + + if agg: + metrics["reward/compile_rate"] = agg.get("compile_rate", 0) + metrics["reward/copybook_only_rate"] = agg.get("copybook_only_rate", 0) + metrics["reward/trivial_rate"] = agg.get("trivial_rate", 0) + metrics["gen/verb_count"] = agg.get("mean_verb_count", 0) + metrics["gen/noop_count"] = agg.get("mean_noop_count", 0) + metrics["gen/mean_errors"] = agg.get("mean_errors", 0) + + if metrics: + self.publisher.publish(state.global_step, metrics) + + kl = logs.get("kl", 0) + if kl > 5.0: + self.kl_alarm_count += 1 + if self.kl_alarm_count >= 10: + print( + f"ALARM: KL divergence > 5.0 for 10 steps ({kl:.2f}). Stopping.", + flush=True, + ) + control.should_training_stop = True + else: + self.kl_alarm_count = 0 + + mean_verbs = agg.get("mean_verb_count", 100) + if mean_verbs < 3: + self.short_alarm_count += 1 + if self.short_alarm_count >= 20: + print( + "ALARM: Mean verb count < 3 for 20 steps. Reward hacking. Stopping.", + flush=True, + ) + control.should_training_stop = True + else: + self.short_alarm_count = 0 + + +class MaxStepsCallback(TrainerCallback): + """Stops training once global_step reaches max_steps. + + GRPOConfig is left unchanged, so the run keeps the full-run LR schedule (warmup and + cosine decay over all epochs) and its first steps track a full run. Give a short run + its own OUTPUT_DIR: its checkpoints would otherwise be auto-resumed by a later full + run. + """ + + def __init__(self, max_steps): + self.max_steps = max_steps + + def on_step_end(self, args, state, control, **kwargs): + if state.global_step >= self.max_steps: + control.should_training_stop = True + + +def _parse_size_bytes(size_str): + size_str = size_str.strip().upper() + units = {"TB": 1024**4, "GB": 1024**3, "MB": 1024**2, "KB": 1024, "B": 1} + for suffix, mult in units.items(): + if size_str.endswith(suffix): + return int(float(size_str[: -len(suffix)]) * mult) + return int(size_str) + + +def _blobfuse_save_safetensors(tensors, filepath, metadata=None): + """safetensors.torch.save_file() uses Rust serialize_file() which is + completely broken on blobfuse2 FUSE mounts (I/O error on any file size). + Workaround: serialize to bytes in memory, then plain Python write().""" + from safetensors.torch import save as _st_save + + raw = _st_save(tensors, metadata=metadata) + with open(filepath, "wb") as f: + f.write(raw) + f.flush() + os.fsync(f.fileno()) + return len(raw) + + +def _blobfuse_save_model(state_dict, output_dir, max_shard_size="5GB"): + """Save sharded safetensors to blobfuse2, bypassing broken Rust serialize_file.""" + max_bytes = _parse_size_bytes(max_shard_size) + + shards, current, current_sz = [], {}, 0 + for key in sorted(state_dict.keys()): + t = state_dict[key] + tsz = t.numel() * t.element_size() + if current_sz + tsz > max_bytes and current: + shards.append(current) + current, current_sz = {}, 0 + current[key] = t + current_sz += tsz + if current: + shards.append(current) + + if len(shards) <= 1: + path = os.path.join(output_dir, "model.safetensors") + nbytes = _blobfuse_save_safetensors(shards[0] if shards else {}, path) + print( + f" [blobfuse-save] model.safetensors ({nbytes / 1e9:.2f} GB)", flush=True + ) + else: + weight_map, total_size = {}, 0 + for i, shard in enumerate(shards, 1): + fname = f"model-{i:05d}-of-{len(shards):05d}.safetensors" + nbytes = _blobfuse_save_safetensors(shard, os.path.join(output_dir, fname)) + for key in shard: + weight_map[key] = fname + total_size += sum(t.numel() * t.element_size() for t in shard.values()) + print( + f" [blobfuse-save] {fname} ({nbytes / 1e9:.2f} GB) [{i}/{len(shards)}]", + flush=True, + ) + index = {"metadata": {"total_size": total_size}, "weight_map": weight_map} + with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f: + json.dump(index, f, indent=2) + + +def _patch_trainer_save_shard_size(trainer, max_shard_size="5GB"): + """Monkey-patch Trainer._save to bypass safetensors Rust serialize_file + (broken on blobfuse2) and use in-memory serialize + Python write instead.""" + + def _patched_save(output_dir=None, state_dict=None): + output_dir = output_dir if output_dir is not None else trainer.args.output_dir + os.makedirs(output_dir, exist_ok=True) + + if state_dict is not None: + _blobfuse_save_model(state_dict, output_dir, max_shard_size=max_shard_size) + else: + unwrapped = trainer.accelerator.unwrap_model( + trainer.model, keep_torch_compile=False + ) + _blobfuse_save_model( + unwrapped.state_dict(), output_dir, max_shard_size=max_shard_size + ) + + unwrapped = trainer.accelerator.unwrap_model( + trainer.model, keep_torch_compile=False + ) + if hasattr(unwrapped, "config"): + unwrapped.config.save_pretrained(output_dir) + + if trainer.processing_class is not None: + trainer.processing_class.save_pretrained(output_dir) + + torch.save(trainer.args, os.path.join(output_dir, "training_args.bin")) + + trainer._save = _patched_save + + +def main(): + global LOCAL_RANK, GLOBAL_RANK, WORLD_SIZE + LOCAL_RANK = int(os.environ.get("LOCAL_RANK", "0")) + GLOBAL_RANK = int(os.environ.get("RANK", "0")) + WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) + + os.environ["HF_HOME"] = HF_HOME + os.makedirs(OUTPUT_DIR, exist_ok=True) + + is_main = GLOBAL_RANK == 0 + + if is_main: + print("=" * 60, flush=True) + print("COBOL GRPO Training (compiler-as-teacher)", flush=True) + print("=" * 60, flush=True) + print(f" Model: {MODEL_ID}", flush=True) + print(f" Train file: {TRAIN_FILE}", flush=True) + print(f" Output: {OUTPUT_DIR}", flush=True) + print(f" Epochs: {NUM_EPOCHS}", flush=True) + print(f" LR: {LR}", flush=True) + print(f" K samples: {K_SAMPLES}", flush=True) + print(f" Temperature: {TEMPERATURE}", flush=True) + print(f" KL coeff: {KL_COEFF}", flush=True) + print(f" Batch size: {BATCH_SIZE} (per device)", flush=True) + print(f" Grad accum: {GRAD_ACCUM}", flush=True) + print(f" GPUs: {torch.cuda.device_count()}", flush=True) + print(f" World size: {WORLD_SIZE}", flush=True) + if MAX_STEPS: + print(f" Max steps: {MAX_STEPS}", flush=True) + + global _reward_tokenizer + if is_main: + print("\nLoading tokenizer...", flush=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + if tokenizer.pad_token is None: + tokenizer.pad_token = tokenizer.eos_token + _reward_tokenizer = tokenizer + + if is_main: + print("\nExtracting prompts from training data...", flush=True) + prompts = load_prompts( + TRAIN_FILE, + max_body_tokens=MAX_GEN_LEN, + tokenizer=tokenizer, + ) + if is_main: + print(f" Extracted {len(prompts)} prompts", flush=True) + + orig_count = len(prompts) + prompts = [p for p in prompts if len(tokenizer.encode(p)) <= MAX_PROMPT_LEN] + if is_main and len(prompts) < orig_count: + print( + f" Filtered {orig_count - len(prompts)} prompts exceeding {MAX_PROMPT_LEN} tokens", + flush=True, + ) + + dataset = Dataset.from_list([{"prompt": p} for p in prompts]) + if is_main: + print(f" Dataset size: {len(dataset)}", flush=True) + + ds_config = { + "bf16": {"enabled": True}, + "zero_optimization": { + "stage": 3, + "offload_optimizer": {"device": "cpu", "pin_memory": True}, + "offload_param": {"device": "none"}, + "overlap_comm": True, + "contiguous_gradients": True, + "reduce_bucket_size": 5e8, + "stage3_prefetch_bucket_size": 5e8, + "stage3_param_persistence_threshold": 1e6, + "stage3_gather_16bit_weights_on_model_save": True, + }, + "gradient_accumulation_steps": "auto", + "gradient_clipping": 1.0, + "train_batch_size": "auto", + "train_micro_batch_size_per_gpu": "auto", + "wall_clock_breakdown": False, + } + ds_config_path = os.path.join(OUTPUT_DIR, "ds_config_rl.json") + if GLOBAL_RANK == 0: + with open(ds_config_path, "w") as f: + json.dump(ds_config, f, indent=2) + if torch.distributed.is_initialized(): + torch.distributed.barrier() + + vllm_base_url = os.environ.get("VLLM_SERVER_URL", f"http://localhost:{VLLM_PORT}") + if is_main: + print(f" vLLM server: {vllm_base_url}", flush=True) + + grpo_kwargs = dict( + output_dir=OUTPUT_DIR, + num_train_epochs=NUM_EPOCHS, + per_device_train_batch_size=BATCH_SIZE, + gradient_accumulation_steps=GRAD_ACCUM, + learning_rate=LR, + max_completion_length=MAX_GEN_LEN, + num_generations=K_SAMPLES, + temperature=TEMPERATURE, + top_p=0.95, + beta=KL_COEFF, + logging_steps=1, + save_steps=SAVE_STEPS, + save_total_limit=3, + gradient_checkpointing=True, + report_to="none", + seed=42, + ) + import inspect + + _sig = inspect.signature(GRPOConfig) + _optional = { + "lr_scheduler_type": "cosine", + "warmup_ratio": 0.1, + "bf16": True, + "max_prompt_length": MAX_PROMPT_LEN, + "gradient_checkpointing_kwargs": {"use_reentrant": False}, + "deepspeed": ds_config_path, + "save_strategy": "steps", + "use_vllm": True, + "vllm_server_base_url": vllm_base_url, + "vllm_mode": "server", + } + for k, v in _optional.items(): + if k in _sig.parameters: + grpo_kwargs[k] = v + elif is_main: + print(f" Skipping unsupported GRPOConfig param: {k}", flush=True) + training_args = GRPOConfig(**grpo_kwargs) + + publisher = TauMetricsPublisher(OUTPUT_DIR) if GLOBAL_RANK == 0 else None + metrics_cb = TauGRPOMetricsCallback(publisher) if publisher else None + max_steps_cb = MaxStepsCallback(int(MAX_STEPS)) if MAX_STEPS else None + callbacks = [cb for cb in [metrics_cb, max_steps_cb] if cb is not None] + + if is_main: + print("\nInitializing GRPOTrainer...", flush=True) + + trainer = GRPOTrainer( + model=MODEL_ID, + reward_funcs=reward_fn, + args=training_args, + train_dataset=dataset, + processing_class=tokenizer, + callbacks=callbacks, + ) + + _patch_trainer_save_shard_size(trainer, max_shard_size="5GB") + + resume_ckpt = None + ckpts = sorted( + [d for d in Path(OUTPUT_DIR).glob("checkpoint-*") if d.is_dir()], + key=lambda p: int(p.name.split("-")[-1]), + ) + if ckpts: + resume_ckpt = str(ckpts[-1]) + if is_main: + print(f"\n Resuming from {resume_ckpt}", flush=True) + + if is_main: + print("\n" + "=" * 60, flush=True) + if resume_ckpt: + print( + f"Resuming GRPO training from {Path(resume_ckpt).name}...", flush=True + ) + else: + print("Starting GRPO training from scratch...", flush=True) + print("=" * 60, flush=True) + + trainer.train(resume_from_checkpoint=resume_ckpt) + + if is_main: + print("\nSaving final model...", flush=True) + final_dir = os.path.join(OUTPUT_DIR, "final") + trainer.save_model(final_dir) + tokenizer.save_pretrained(final_dir) + + if GLOBAL_RANK == 0: + config = { + "model_id": MODEL_ID, + "algorithm": "GRPO", + "num_epochs": NUM_EPOCHS, + "lr": LR, + "k_samples": K_SAMPLES, + "temperature": TEMPERATURE, + "kl_coeff": KL_COEFF, + "batch_size": BATCH_SIZE, + "grad_accum": GRAD_ACCUM, + "max_prompt_len": MAX_PROMPT_LEN, + "max_gen_len": MAX_GEN_LEN, + "num_prompts": len(prompts), + "world_size": WORLD_SIZE, + "gpus_per_node": torch.cuda.device_count(), + } + with open(os.path.join(final_dir, "train-config.json"), "w") as f: + json.dump(config, f, indent=2) + print("\n=== GRPO training complete ===", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/07-eval/driver.py b/examples/cobol-cpt-grpo/07-eval/driver.py new file mode 100644 index 00000000..ca3d5b90 --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/driver.py @@ -0,0 +1,36 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""RayJob entrypoint for the vLLM eval harness. + +TauGrid runs this file as the Ray driver on a CPU-only head pod, so it must not import +vllm or torch. The eval runs in one Ray task that reserves a whole GPU worker and calls +run_eval_vllm.main() there; cobc also runs inside that task, on the worker. + +run_eval_vllm is imported inside the task because its module-level code imports vllm and +torch. TauGrid ships this directory as the Ray working_dir, so the import resolves on the +worker. Configuration (EVAL_DIR, MODEL_ID, OUT, MAX_NEW_TOKENS, TP_SIZE, ...) comes from +the pod env, which TauGrid sets on the head and on every worker. +""" + +import os + +import ray + +NUM_GPUS = int(os.environ.get("TP_SIZE", "8")) + + +@ray.remote(num_gpus=NUM_GPUS, max_retries=0) +def run_eval(): + import run_eval_vllm + + run_eval_vllm.main() + + +def main(): + ray.init(address="auto", ignore_reinit_error=True) + ray.get(run_eval.remote()) + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py new file mode 100755 index 00000000..5f8cc1aa --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -0,0 +1,612 @@ +#!/usr/bin/env python3 +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +""" +COBOL evaluation harness — vLLM backend for multi-GPU batch inference. + +Uses vLLM with tensor parallelism across all available GPUs on a node. +All benchmark tasks are batched and processed concurrently. +""" + +import json +import math +import os +import re +import subprocess +import tempfile +from collections import defaultdict + +import torch +from vllm import LLM, SamplingParams + +EVAL_DIR = os.environ.get("EVAL_DIR", "/data/datasets/cobol/eval") +MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-Coder-32B") +OUT = os.environ.get("OUT", "/data/datasets/cobol/eval/results/baseline") +MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "512")) +TP_SIZE = int(os.environ.get("TP_SIZE", str(torch.cuda.device_count()))) +HF_HOME = os.environ.get("HF_HOME", "/tmp/hf-cache") + +os.makedirs(OUT, exist_ok=True) +os.environ["HF_HOME"] = HF_HOME + + +def load_jsonl(path): + records = [] + if not os.path.isfile(path): + return records + with open(path) as f: + for line in f: + line = line.strip() + if line: + records.append(json.loads(line)) + return records + + +def eval_perplexity(tasks, max_length=4096): + """Compute perplexity via HuggingFace model (teacher-forcing, not generation).""" + from transformers import AutoTokenizer, AutoModelForCausalLM + + print(f"\n=== Perplexity Evaluation ({len(tasks)} docs) ===", flush=True) + print(" Loading HF model for perplexity (single GPU)...", flush=True) + + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + n_gpus = torch.cuda.device_count() + max_mem = {i: "17GiB" for i in range(n_gpus)} + model = AutoModelForCausalLM.from_pretrained( + MODEL_ID, + torch_dtype=torch.bfloat16, + device_map="auto", + max_memory=max_mem, + trust_remote_code=True, + ) + + total_nll = 0.0 + total_tokens = 0 + per_doc = [] + + model.eval() + with torch.no_grad(): + for idx, task in enumerate(tasks): + if idx % 20 == 0: + print(f" Progress: {idx}/{len(tasks)}...", flush=True) + + text = task["text"] + encodings = tokenizer( + text, return_tensors="pt", truncation=True, max_length=max_length + ) + input_ids = encodings.input_ids.to(model.device) + + if input_ids.shape[1] < 2: + continue + + outputs = model(input_ids, labels=input_ids) + nll = outputs.loss.item() * (input_ids.shape[1] - 1) + n_tokens = input_ids.shape[1] - 1 + + total_nll += nll + total_tokens += n_tokens + doc_ppl = math.exp(nll / n_tokens) if n_tokens > 0 else float("inf") + per_doc.append( + { + "id": task["id"], + "source": task.get("source", ""), + "perplexity": round(doc_ppl, 2), + "tokens": n_tokens, + } + ) + + avg_ppl = math.exp(total_nll / total_tokens) if total_tokens > 0 else float("inf") + print( + f" Aggregate perplexity: {avg_ppl:.2f} ({total_tokens:,} tokens)", flush=True + ) + + del model + del tokenizer + torch.cuda.empty_cache() + + return { + "aggregate_perplexity": round(avg_ppl, 2), + "total_tokens": total_tokens, + "num_docs": len(per_doc), + "per_doc": per_doc, + } + + +def validate_division_structure(text): + upper = text.upper() + errors = [] + if "IDENTIFICATION DIVISION" not in upper: + errors.append("missing_identification_division") + if "PROCEDURE DIVISION" not in upper: + errors.append("missing_procedure_division") + + positions = {} + for div in ["IDENTIFICATION", "ENVIRONMENT", "DATA", "PROCEDURE"]: + idx = upper.find(f"{div} DIVISION") + if idx >= 0: + positions[div] = idx + + sorted_divs = sorted(positions.items(), key=lambda x: x[1]) + expected_order = ["IDENTIFICATION", "ENVIRONMENT", "DATA", "PROCEDURE"] + actual_order = [d for d, _ in sorted_divs] + if actual_order != [d for d in expected_order if d in positions]: + errors.append("wrong_division_order") + + return errors + + +def validate_column_formatting(text): + errors = [] + lines = text.split("\n") + violations = 0 + total_code_lines = 0 + + for line in lines: + if len(line) < 7: + continue + total_code_lines += 1 + col7 = line[6] if len(line) > 6 else " " + if col7 not in " */-dD": + violations += 1 + if len(line) > 72: + content_past_72 = line[72:].strip() + if content_past_72 and not line[72:].strip().isdigit(): + violations += 1 + + violation_rate = violations / max(total_code_lines, 1) + if violation_rate > 0.1: + errors.append(f"column_violations_{violation_rate:.0%}") + + return errors, violations, total_code_lines + + +def validate_pic_clauses(text): + pic_re = re.compile( + r"PIC(?:TURE)?\s+IS\s+([^\s.]+)|PIC(?:TURE)?\s+([^\s.]+)", re.IGNORECASE + ) + matches = pic_re.findall(text) + valid = 0 + invalid = 0 + for m in matches: + pic_str = m[0] or m[1] + if re.match(r"^[X9AZSV()\d+\-.]+$", pic_str, re.IGNORECASE): + valid += 1 + else: + invalid += 1 + return valid, invalid + + +def validate_paragraph_refs(text): + para_def_re = re.compile(r"^ ([A-Z0-9][A-Z0-9-]*)\.\s*$", re.MULTILINE) + defined = set(m.group(1) for m in para_def_re.finditer(text)) + + perform_re = re.compile(r"PERFORM\s+([A-Z0-9][A-Z0-9-]+)", re.IGNORECASE) + referenced = set() + dangling = set() + for m in perform_re.finditer(text): + target = m.group(1).upper() + if target in ("UNTIL", "VARYING", "TIMES", "TEST", "WITH"): + continue + referenced.add(target) + if target not in {d.upper() for d in defined}: + dangling.add(target) + + return len(referenced), len(dangling), defined + + +def eval_structural(tasks, generated_texts): + print(f"\n=== Structural Validation ({len(tasks)} tasks) ===", flush=True) + results = [] + totals = defaultdict(int) + + for task, gen_text in zip(tasks, generated_texts): + full_text = task["prompt"] + gen_text + + div_errors = validate_division_structure(full_text) + col_errors, col_violations, col_total = validate_column_formatting(full_text) + pic_valid, pic_invalid = validate_pic_clauses(full_text) + ref_total, ref_dangling, _ = validate_paragraph_refs(full_text) + + totals["division_valid"] += 1 if not div_errors else 0 + totals["column_valid"] += 1 if not col_errors else 0 + totals["pic_valid"] += pic_valid + totals["pic_invalid"] += pic_invalid + totals["ref_total"] += ref_total + totals["ref_dangling"] += ref_dangling + totals["total"] += 1 + + results.append( + { + "id": task["id"], + "division_errors": div_errors, + "column_violations": col_violations, + "column_total_lines": col_total, + "pic_valid": pic_valid, + "pic_invalid": pic_invalid, + "ref_total": ref_total, + "ref_dangling": ref_dangling, + } + ) + + n = max(totals["total"], 1) + summary = { + "division_structure_valid_pct": round(totals["division_valid"] / n * 100, 1), + "column_format_valid_pct": round(totals["column_valid"] / n * 100, 1), + "pic_accuracy_pct": round( + totals["pic_valid"] + / max(totals["pic_valid"] + totals["pic_invalid"], 1) + * 100, + 1, + ), + "perform_ref_accuracy_pct": round( + (totals["ref_total"] - totals["ref_dangling"]) + / max(totals["ref_total"], 1) + * 100, + 1, + ), + "num_tasks": n, + } + print( + f" Division structure valid: {summary['division_structure_valid_pct']}%", + flush=True, + ) + print( + f" Column formatting valid: {summary['column_format_valid_pct']}%", flush=True + ) + print(f" PIC clause accuracy: {summary['pic_accuracy_pct']}%", flush=True) + print( + f" PERFORM reference accuracy: {summary['perform_ref_accuracy_pct']}%", + flush=True, + ) + + return {"summary": summary, "details": results} + + +def _compile_once(src_path, fmt_flag): + try: + result = subprocess.run( + ["cobc", "-fsyntax-only", fmt_flag, src_path], + capture_output=True, + text=True, + timeout=10, + ) + success = result.returncode == 0 + errors = result.stderr.strip() if not success else "" + return success, errors + except FileNotFoundError: + return None, "cobc not found" + except subprocess.TimeoutExpired: + return False, "compilation timeout" + + +def try_compile(cobol_text): + with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: + f.write(cobol_text) + f.flush() + src_path = f.name + + try: + success, errors = _compile_once(src_path, "-free") + fmt_used = "free" + if not success: + success_fixed, errors_fixed = _compile_once(src_path, "-fixed") + if success_fixed: + success, errors, fmt_used = True, "", "fixed" + elif errors_fixed and len(errors_fixed) < len(errors): + errors = errors_fixed + fmt_used = "fixed" + return success, errors, fmt_used + finally: + os.unlink(src_path) + + +def eval_compilation(tasks, generated_texts): + print(f"\n=== Compilation Evaluation ({len(tasks)} tasks) ===", flush=True) + compiled = 0 + failed = 0 + skipped = 0 + fmt_counts = {"free": 0, "fixed": 0} + results = [] + + for task, gen_text in zip(tasks, generated_texts): + full_text = task["prompt"] + gen_text + success, errors, fmt_used = try_compile(full_text) + if success is None: + skipped += 1 + continue + if success: + compiled += 1 + fmt_counts[fmt_used] += 1 + else: + failed += 1 + + results.append( + { + "id": task["id"], + "compiled": success, + "errors": errors[:500] if errors else "", + "format": fmt_used, + } + ) + + total = compiled + failed + rate = compiled / max(total, 1) * 100 + print(f" Compiled: {compiled}/{total} ({rate:.1f}%)", flush=True) + print( + f" Format breakdown: {fmt_counts['free']} free, {fmt_counts['fixed']} fixed", + flush=True, + ) + if skipped: + print(f" Skipped (cobc not found): {skipped}", flush=True) + + return { + "compilation_rate_pct": round(rate, 1), + "compiled": compiled, + "failed": failed, + "skipped": skipped, + "compiled_free": fmt_counts["free"], + "compiled_fixed": fmt_counts["fixed"], + "details": results, + } + + +def normalized_edit_distance(s1, s2): + if not s1 and not s2: + return 0.0 + if not s1 or not s2: + return 1.0 + + m, n = len(s1), len(s2) + if m > 2000 or n > 2000: + s1, s2 = s1[:2000], s2[:2000] + m, n = len(s1), len(s2) + + dp = list(range(n + 1)) + for i in range(1, m + 1): + prev = dp[0] + dp[0] = i + for j in range(1, n + 1): + temp = dp[j] + if s1[i - 1] == s2[j - 1]: + dp[j] = prev + else: + dp[j] = 1 + min(dp[j], dp[j - 1], prev) + prev = temp + + return dp[n] / max(m, n) + + +def compute_bleu(reference, hypothesis, max_n=4): + ref_tokens = reference.split() + hyp_tokens = hypothesis.split() + + if not hyp_tokens or not ref_tokens: + return 0.0 + + precisions = [] + for n in range(1, max_n + 1): + ref_ngrams = defaultdict(int) + for i in range(len(ref_tokens) - n + 1): + ngram = tuple(ref_tokens[i : i + n]) + ref_ngrams[ngram] += 1 + + matches = 0 + total = 0 + for i in range(len(hyp_tokens) - n + 1): + ngram = tuple(hyp_tokens[i : i + n]) + if ref_ngrams.get(ngram, 0) > 0: + matches += 1 + ref_ngrams[ngram] -= 1 + total += 1 + + if total == 0: + precisions.append(0) + else: + precisions.append(matches / total) + + if any(p == 0 for p in precisions): + return 0.0 + + log_avg = sum(math.log(p) for p in precisions) / max_n + + bp = 1.0 + if len(hyp_tokens) < len(ref_tokens): + bp = math.exp(1 - len(ref_tokens) / max(len(hyp_tokens), 1)) + + return bp * math.exp(log_avg) + + +def eval_generation_quality(tasks, generated_texts, benchmark_name): + print( + f"\n=== Generation Quality: {benchmark_name} ({len(tasks)} tasks) ===", + flush=True, + ) + bleu_scores = [] + edit_dists = [] + results = [] + + for task, gen_text in zip(tasks, generated_texts): + reference = task["reference"] + bleu = compute_bleu(reference, gen_text) + ed = normalized_edit_distance(reference.strip(), gen_text.strip()) + + bleu_scores.append(bleu) + edit_dists.append(ed) + results.append( + { + "id": task["id"], + "bleu": round(bleu, 4), + "edit_distance": round(ed, 4), + } + ) + + avg_bleu = sum(bleu_scores) / max(len(bleu_scores), 1) + avg_ed = sum(edit_dists) / max(len(edit_dists), 1) + print(f" Average BLEU: {avg_bleu:.4f}", flush=True) + print(f" Average normalized edit distance: {avg_ed:.4f}", flush=True) + + return { + "avg_bleu": round(avg_bleu, 4), + "avg_edit_distance": round(avg_ed, 4), + "num_tasks": len(tasks), + "details": results, + } + + +def batch_generate(llm, tasks, max_new_tokens): + """Generate completions for all tasks in one batched vLLM call.""" + prompts = [task["prompt"] for task in tasks] + sampling_params = SamplingParams( + max_tokens=max_new_tokens, + temperature=0, + top_p=1.0, + ) + print( + f" Generating {len(prompts)} completions (batch, TP={TP_SIZE})...", flush=True + ) + outputs = llm.generate(prompts, sampling_params) + generated = [output.outputs[0].text for output in outputs] + print(f" Done — {len(generated)} completions generated.", flush=True) + return generated + + +def main(): + print("=== COBOL Evaluation Harness (vLLM) ===", flush=True) + print(f" Model: {MODEL_ID}", flush=True) + print(f" GPUs: {torch.cuda.device_count()}, TP size: {TP_SIZE}", flush=True) + print(f" Eval dir: {EVAL_DIR}", flush=True) + + all_results = {"model": MODEL_ID, "backend": "vllm", "tp_size": TP_SIZE} + + ppl_tasks = load_jsonl(os.path.join(EVAL_DIR, "perplexity.jsonl")) + if ppl_tasks: + ppl_results = eval_perplexity(ppl_tasks) + all_results["perplexity"] = { + k: v for k, v in ppl_results.items() if k != "per_doc" + } + with open(os.path.join(OUT, "perplexity_details.jsonl"), "w") as f: + for doc in ppl_results.get("per_doc", []): + f.write(json.dumps(doc) + "\n") + + print(f"\nLoading vLLM model (TP={TP_SIZE})...", flush=True) + llm = LLM( + model=MODEL_ID, + tensor_parallel_size=TP_SIZE, + dtype="bfloat16", + trust_remote_code=True, + max_model_len=int(os.environ.get("MAX_MODEL_LEN", "32768")), + gpu_memory_utilization=0.90, + download_dir=HF_HOME, + distributed_executor_backend="mp", + ) + print(" vLLM model loaded.", flush=True) + + para_tasks = load_jsonl(os.path.join(EVAL_DIR, "paragraph_completion.jsonl")) + if para_tasks: + print( + f"\n--- Paragraph completion ({len(para_tasks)} tasks, max_tokens={MAX_NEW_TOKENS}) ---", + flush=True, + ) + para_gen = batch_generate(llm, para_tasks, max_new_tokens=MAX_NEW_TOKENS) + para_quality = eval_generation_quality( + para_tasks, para_gen, "paragraph_completion" + ) + all_results["paragraph_completion"] = para_quality + para_compile = eval_compilation( + [{"id": t["id"], "prompt": t["prompt"]} for t in para_tasks], para_gen + ) + all_results["paragraph_compilation"] = { + k: v for k, v in para_compile.items() if k != "details" + } + with open(os.path.join(OUT, "paragraph_compilation_details.jsonl"), "w") as f: + for d in para_compile.get("details", []): + f.write(json.dumps(d) + "\n") + + prog_tasks = load_jsonl(os.path.join(EVAL_DIR, "program_completion.jsonl")) + if prog_tasks: + print( + f"\n--- Program completion ({len(prog_tasks)} tasks, max_tokens=1024) ---", + flush=True, + ) + prog_gen = batch_generate(llm, prog_tasks, max_new_tokens=1024) + prog_quality = eval_generation_quality( + prog_tasks, prog_gen, "program_completion" + ) + all_results["program_completion"] = prog_quality + prog_compile = eval_compilation( + [{"id": t["id"], "prompt": t["prompt"]} for t in prog_tasks], prog_gen + ) + all_results["program_compilation"] = { + k: v for k, v in prog_compile.items() if k != "details" + } + with open(os.path.join(OUT, "program_compilation_details.jsonl"), "w") as f: + for d in prog_compile.get("details", []): + f.write(json.dumps(d) + "\n") + + strv_tasks = load_jsonl(os.path.join(EVAL_DIR, "structural_validation.jsonl")) + if strv_tasks: + print( + f"\n--- Structural validation ({len(strv_tasks)} tasks, max_tokens=1024) ---", + flush=True, + ) + strv_gen = batch_generate(llm, strv_tasks, max_new_tokens=1024) + strv_results = eval_structural(strv_tasks, strv_gen) + all_results["structural_validation"] = strv_results["summary"] + strv_compile = eval_compilation( + [{"id": t["id"], "prompt": t["prompt"]} for t in strv_tasks], strv_gen + ) + all_results["structural_compilation"] = { + k: v for k, v in strv_compile.items() if k != "details" + } + with open(os.path.join(OUT, "structural_compilation_details.jsonl"), "w") as f: + for d in strv_compile.get("details", []): + f.write(json.dumps(d) + "\n") + + summary_path = os.path.join(OUT, "eval-summary.json") + with open(summary_path, "w") as f: + json.dump(all_results, f, indent=2) + print(f"\n=== Evaluation complete. Results: {summary_path} ===", flush=True) + + print("\n=== Results Summary ===", flush=True) + if "perplexity" in all_results: + print( + f" COBOL Perplexity: {all_results['perplexity']['aggregate_perplexity']}", + flush=True, + ) + if "paragraph_completion" in all_results: + print( + f" Paragraph BLEU: {all_results['paragraph_completion']['avg_bleu']}", + flush=True, + ) + print( + f" Paragraph Edit Dist: {all_results['paragraph_completion']['avg_edit_distance']}", + flush=True, + ) + if "paragraph_compilation" in all_results: + print( + f" Paragraph Compile Rate: {all_results['paragraph_compilation']['compilation_rate_pct']}%", + flush=True, + ) + if "program_completion" in all_results: + print( + f" Program BLEU: {all_results['program_completion']['avg_bleu']}", + flush=True, + ) + if "program_compilation" in all_results: + print( + f" Program Compile Rate: {all_results['program_compilation']['compilation_rate_pct']}%", + flush=True, + ) + if "structural_validation" in all_results: + sv = all_results["structural_validation"] + print( + f" Division Structure: {sv['division_structure_valid_pct']}%", flush=True + ) + print(f" Column Format: {sv['column_format_valid_pct']}%", flush=True) + print(f" PIC Accuracy: {sv['pic_accuracy_pct']}%", flush=True) + print(f" PERFORM Ref Accuracy: {sv['perform_ref_accuracy_pct']}%", flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/cobol-cpt-grpo/07-eval/tau.yaml b/examples/cobol-cpt-grpo/07-eval/tau.yaml new file mode 100644 index 00000000..3d80520e --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/tau.yaml @@ -0,0 +1,42 @@ +name: cobol-eval +engine: rayjob +run: + entrypoint: driver.py + working_dir: . + working_dir_excludes: + - __pycache__ + - .pytest_cache + - .ruff_cache + - test_*.py + +experiment: + project: cobol-cpt-grpo + name: eval-grpo + group: eval + +runtime: + image: registry.example.com/cobol-train:v0.28.0-ray2.56.0 + env: + EVAL_DIR: /data/datasets/cobol/eval + MODEL_ID: /data/projects/taugrid-default/runs/cobol-cpt-grpo/grpo/final + OUT: /data/projects/taugrid-default/runs/cobol-cpt-grpo/eval/grpo + MAX_NEW_TOKENS: "512" + TP_SIZE: "8" + PYTHONUNBUFFERED: "1" + HF_HOME: /tmp/hf-cache + VLLM_WORKER_MULTIPROC_METHOD: spawn + +compute: + workers: 1 + gpus_per_worker: 8 + worker_cpu_request: "32" + worker_memory_request: 128Gi + worker_cpu_limit: "64" + worker_memory_limit: 256Gi + +storage: + data_pvc: workspace-data + output: /data/projects/taugrid-default/runs/cobol-cpt-grpo/eval/grpo + +policy: + workspace: taugrid-default diff --git a/examples/cobol-cpt-grpo/07-eval/test_driver.py b/examples/cobol-cpt-grpo/07-eval/test_driver.py new file mode 100644 index 00000000..e14521a7 --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/test_driver.py @@ -0,0 +1,73 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local test for the eval RayJob driver. + +Starts a local Ray cluster with 8 fake logical GPUs and replaces run_eval_vllm with a +stub, then checks that the driver runs main() in a task holding all 8 GPUs and that the +driver process itself never imports vllm or torch. Needs ray, not vllm or a GPU: + python -m pytest examples/cobol-cpt-grpo/07-eval/test_driver.py -v +""" + +import json +import os +import shutil +import sys + +import pytest +import ray + +EVAL_DIR = os.path.dirname(os.path.abspath(__file__)) + +STUB_RUN_EVAL = """\ +import json +import os + +import ray + + +def main(): + with open(os.environ["STUB_MARKER"], "w") as f: + json.dump({"gpu_ids": ray.get_gpu_ids(), "pid": os.getpid()}, f) +""" + + +@pytest.fixture +def ray_cluster(tmp_path, monkeypatch): + work = tmp_path / "work" + work.mkdir() + shutil.copy(os.path.join(EVAL_DIR, "driver.py"), work / "driver.py") + (work / "run_eval_vllm.py").write_text(STUB_RUN_EVAL) + marker = tmp_path / "marker.json" + + monkeypatch.delenv("TP_SIZE", raising=False) + monkeypatch.syspath_prepend(str(work)) + sys.modules.pop("driver", None) + ray.init( + num_cpus=2, + num_gpus=8, + include_dashboard=False, + runtime_env={ + "working_dir": str(work), + "env_vars": {"STUB_MARKER": str(marker)}, + }, + ) + yield marker + ray.shutdown() + sys.modules.pop("driver", None) + + +def test_driver_runs_eval_on_gpu_task_without_gpu_imports(ray_cluster): + assert "vllm" not in sys.modules + assert "torch" not in sys.modules + + import driver + + driver.main() + + result = json.loads(ray_cluster.read_text()) + assert len(result["gpu_ids"]) == 8 + assert result["pid"] != os.getpid() + assert "run_eval_vllm" not in sys.modules + assert "vllm" not in sys.modules + assert "torch" not in sys.modules diff --git a/examples/cobol-cpt-grpo/Dockerfile b/examples/cobol-cpt-grpo/Dockerfile new file mode 100644 index 00000000..3e7eb4d8 --- /dev/null +++ b/examples/cobol-cpt-grpo/Dockerfile @@ -0,0 +1,15 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +FROM vllm/vllm-openai:v0.28.0@sha256:61fc8a896b0a4fbbbdc063bc4b0dbc25ce98e02b5050c24aeb7830ac02039b14 + +RUN apt-get update && \ + apt-get install -y --no-install-recommends \ + gnucobol3=3.1.2-5.1ubuntu1 \ + wget=1.21.4-1ubuntu4.5 && \ + rm -rf /var/lib/apt/lists/* + +ARG PIP_INDEX_URL=https://pypi.org/simple +COPY image-requirements.txt /tmp/image-requirements.txt +RUN pip install --no-cache-dir -r /tmp/image-requirements.txt && \ + rm /tmp/image-requirements.txt diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md new file mode 100644 index 00000000..9a5a1e25 --- /dev/null +++ b/examples/cobol-cpt-grpo/README.md @@ -0,0 +1,236 @@ +# COBOL continued pretraining, GRPO, and eval + +> **Status:** `production-shaped` +> **Intended use:** adapt `Qwen/Qwen2.5-Coder-32B` to COBOL with continued +> pretraining (CPT), then GRPO with a compiler reward, and score the result on +> compile-checked COBOL benchmarks, on 2 x 8 H200 nodes with RDMA. +> **Not for:** first-time onboarding, CPU-only or single-GPU clusters, or +> serving. Start with [`nanogpt-ray`](../nanogpt-ray/) or +> [`ray-tune-smoke`](../ray-tune-smoke/) to validate a workspace first. + +The pipeline has seven stages. Each stage folder has its own `tau.yaml`. +CPU stages run as Jobs; GPU stages run as RayJobs that ship the stage +folder as the Ray `working_dir`. + +| Stage | Engine | What it does | +| --- | --- | --- | +| [`01-collect/`](./01-collect/) | Job (CPU) | Clones COBOL repositories and downloads public Hugging Face datasets into `/data/datasets/cobol/raw`. `tau-rl.yaml` collects the parser and test corpora used for RL prompts. `tau-gated.yaml` downloads gated datasets and needs `HF_TOKEN`. | +| [`02-data-prep/`](./02-data-prep/) | Job (CPU) | Extracts, deduplicates, filters, and splits the corpus into `prepared/train.jsonl` and `prepared/val.jsonl`. | +| [`03-benchmarks/`](./03-benchmarks/) | Job (CPU) | Builds five eval suites from the validation split into `/data/datasets/cobol/eval`. | +| [`04-cpt/`](./04-cpt/) | RayJob, 2 x 8 GPUs | Continued pretraining with DeepSpeed ZeRO-3 under Ray Train. | +| [`05-rl-prompts/`](./05-rl-prompts/) | Job (CPU) | Masks paragraph bodies of programs that compile with `cobc` to produce `prepared/rl-prompts.jsonl`. | +| [`06-grpo/`](./06-grpo/) | RayJob, 2 x 8 GPUs | GRPO from the CPT checkpoint with a compiler reward and an in-cluster vLLM generation server. | +| [`07-eval/`](./07-eval/) | RayJob, 1 x 8 GPUs | Runs the benchmarks with vLLM (TP=8) and compiles every generated program with GnuCOBOL. | + +Serving the trained model is not part of this example; no `08-serve` stage is +included. + +## Data sources and licensing + +`01-collect/collect-cobol-data.sh` and `01-collect/collect-rl-data.sh` read +`SOURCE_LICENSES`: + +- `permissive` (the checked-in default) clones only repositories whose GitHub + license is MIT, Apache-2.0, BSD, CC-BY-4.0, or US-government public domain: + - CPT corpus: `aws-samples/aws-mainframe-modernization-carddemo` + (Apache-2.0), CobolCraft (MIT), the openmainframeproject COBOL course + (CC-BY-4.0), Cobol-Programming-Collection (MIT), shamrice (MIT), koopa + (BSD), pc-pricer (US government, public domain), and the exercism COBOL + track (MIT). + - RL corpora: proleap (MIT), cobol-rekt (MIT), Ironclad (Apache-2.0), DSF + (MIT), Bank-of-Z (Apache-2.0), and FPS.cob (Apache-2.0). +- `all` also clones repositories that are copyleft (GPL-2.0, GPL-3.0, + EPL-2.0, CECILL-2.1) or have no stated license. It also enables the NIST + COBOL85 suite extracted from opensourcecobol4j and the zorse dataset. Use + `all` only if your use of the trained weights complies with those terms. + +The ungated Hugging Face datasets in `collect-cobol-data.sh` are +`TroyDoesAI/Cobol`, `bigcode/commitpack` (COBOL subset), +`Fsoft-AIC/MainframeBench`, `msplits/human-eval-cobol`, and +`harshini-kumar/CobolCodeBench`. Check each dataset card's terms before use. + +Gated datasets are opt-in. They are `bigcode/the-stack` and +`bigcode/the-stack-v2` (COBOL subsets) and `4sp1d3r2/cobol-pretraining-corpus`. +Accept their terms on Hugging Face, then store a read token in a Secret in the +workspace namespace: + +```bash +kubectl create secret generic hf-token --from-literal=token= \ + --namespace +``` + +`01-collect/tau-gated.yaml` maps that Secret to `HF_TOKEN` with +`env_secret: {HF_TOKEN: hf-token:token}`. `collect-rl-data.sh` skips its gated +steps when `HF_TOKEN` is unset. To include them, add the same `env_secret` +block to `01-collect/tau-rl.yaml`. `02-data-prep` picks up the gated outputs +when they are present under `raw/`. + +## Method + +**CPT (`04-cpt`).** `driver.py` runs on the CPU-only Ray head and starts a Ray +Train `TorchTrainer` with 16 workers, one per GPU. Each worker calls +`continued_pretrain.main()`, which runs HF Trainer with DeepSpeed ZeRO-3 on +packed 4096-token chunks. By default that is 3 epochs at LR 2e-5, an effective +batch of 4 x 4 x 16, and eval every 25 steps. Ray Train sets up +`torch.distributed` and the rank environment that torchrun would. HF Trainer +owns checkpointing and resume under `OUTPUT_DIR`; Ray Train stores only its +run metadata under `OUTPUT_DIR/ray`. `final/` holds the trained model as 5 GB +safetensors shards. + +**GRPO (`06-grpo`).** One RayJob runs both halves of TRL's server mode on the +two GPU workers: + +1. A `VLLMServer` actor reserves all 8 GPUs of one worker and runs + `trl vllm-serve` (TP=8). The driver waits for `/health` on `VLLM_PORT` + (8000). +2. A `TorchTrainer` with 8 workers lands on the other node. Each worker sets + `VLLM_SERVER_URL` and calls `train_grpo.main()`. + +TRL syncs updated weights from trainer rank 0 to the server over an NCCL group +on TRL's group port (51216 in TRL 1.13.0). The reward has two parts: + +- **Compile reward.** Each prompt plus completion is checked with + `cobc -fsyntax-only -free`. A clean compile scores 1.0 and a compile with + warnings scores 0.7. A `cobc` signal crash counts as a failed compile; a + timeout or a missing `cobc` raises. +- **Structural reward.** Penalizes degenerate completions (display-only + bodies, no-op majorities, repeated statements, truncation). It adds a bonus + when every `PERFORM` target is a paragraph defined in the completion. + +The defaults are K=4 samples per prompt at temperature 0.8, LR 5e-7, KL +coefficient 0.05, and 3 epochs. `06-grpo/test_reward.py` covers the reward +without a GPU. + +**Eval (`07-eval`).** One Ray task reserves a whole 8-GPU worker and runs +`run_eval_vllm.main()` with greedy decoding (512 new tokens). The five suites +are paragraph completion, copybook usage, program completion, perplexity, and +structural validation. Compile scoring is format-aware: each generation is +compiled with `cobc -free`, retried with `-fixed` if that fails, and both +counts are reported. The headline compile score is out of 220 tasks: 70 +paragraph, 50 program, and 100 structural. + +## Hardware + +- CPT and GRPO: 2 GPU workers x 8 H200, `rdma.enabled: true`, and + `policy.topology: same-network-domain`. The NCCL settings in + `04-cpt/tau.yaml` assume InfiniBand HCAs named `mlx5*` and a pod interface + `eth0`. +- Eval: 1 GPU worker x 8 H200, no RDMA. +- CPU stages: 1 to 8 vCPUs and 4 to 32 GiB memory. +- Storage: a writable workspace PVC mounted at `/data`. Datasets live under + `/data/datasets/cobol`; checkpoints and eval output live under + `/data/projects/taugrid-default/runs/cobol-cpt-grpo`. + +If `/data` is a blobfuse2 mount, `train_grpo.py` writes safetensors through a +Python writer because the Rust `serialize_file` path fails on blobfuse2. For +the same reason, `continued_pretrain.py` saves with `max_shard_size="5GB"`. + +## Substitute placeholders + +The checked-in values are placeholders for a platform-specific copy of the +configs: + +| Value | Where | Replace with | +| --- | --- | --- | +| `registry.example.com/cobol-train:v0.28.0-ray2.56.0` | `05-rl-prompts`, `04-cpt`, `06-grpo`, `07-eval` | An image you build from [`Dockerfile`](./Dockerfile) (vLLM 0.28.0, GnuCOBOL 3.1.2, and the pins in `image-requirements.txt`) and push to your registry. Reference it by digest. | +| `policy.workspace: taugrid-default` | every `tau.yaml` | Your TauWorkspace. | +| `storage.data_pvc: workspace-data` | every `tau.yaml` | The workspace's writable PVC. | +| `/data/projects/taugrid-default/runs/cobol-cpt-grpo/...` | `OUTPUT_DIR`, `MODEL_ID`, `OUT`, `storage.output` | Your workspace output root. | +| `policy.topology` | `04-cpt`, `06-grpo` | The placement of the workload profile you target. Set `policy.profile` if more than one ready profile matches; list them with `tau cluster profiles export`. | +| `TAU_METRICS_OFFLOAD_*` | environment of the `06-grpo` submitter | See below. | + +The CPU stages use `python:3.11.17-bookworm` pinned by digest, from Docker +Hub. It includes `git`, `curl`, `unzip`, and `file`, which the collect +scripts need. The training image is based on the upstream +`vllm/vllm-openai:v0.28.0` image from Docker Hub, pinned by digest. Both are +exceptions to the MCR base-image preference; mirror them into your registry +if your platform requires it. + +`06-grpo/tau.yaml` enables the Stellar metrics offload sidecar for +`metrics-history-attempt-0/*.jsonl`. Provide the collector image and ADX +settings, as in [`portal-ray-stellar`](../portal-ray-stellar/): + +```bash +export TAU_METRICS_OFFLOAD_IMAGE= +export TAU_METRICS_OFFLOAD_ADX_CLUSTER_URI=https://..kusto.windows.net +export TAU_METRICS_OFFLOAD_ADX_DATABASE= +export TAU_METRICS_OFFLOAD_ADX_CLIENT_ID= +export TAU_METRICS_OFFLOAD_OUT=/var/run/tau/metrics-offload +``` + +To run without Stellar, remove the `metrics.offload` block. + +## Run it + +Run from the repository root. Render each stage with a client dry-run before +submitting it, and wait for each stage to finish before starting the next. + +```bash +make install-tau-cli + +tau run --config examples/cobol-cpt-grpo/01-collect/tau.yaml --dry-run=client +tau run --config examples/cobol-cpt-grpo/01-collect/tau.yaml +tau run --config examples/cobol-cpt-grpo/01-collect/tau-rl.yaml +tau run --config examples/cobol-cpt-grpo/01-collect/tau-gated.yaml + +tau run --config examples/cobol-cpt-grpo/02-data-prep/tau.yaml +tau run --config examples/cobol-cpt-grpo/03-benchmarks/tau.yaml + +tau run --config examples/cobol-cpt-grpo/04-cpt/tau.yaml --dry-run=client +tau run --config examples/cobol-cpt-grpo/04-cpt/tau.yaml + +tau run --config examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml + +tau run --config examples/cobol-cpt-grpo/06-grpo/tau.yaml --dry-run=client +tau run --config examples/cobol-cpt-grpo/06-grpo/tau.yaml + +tau run --config examples/cobol-cpt-grpo/07-eval/tau.yaml +``` + +`tau-gated.yaml` is optional. `05-rl-prompts` only needs `01-collect` and +`02-data-prep`, so it can run alongside CPT. + +To evaluate the base model or the CPT checkpoint, copy `07-eval/tau.yaml` and +change `MODEL_ID` (`Qwen/Qwen2.5-Coder-32B` or `.../cpt/final`) and `OUT` +(`.../eval/base` or `.../eval/cpt`). + +For a short end-to-end check of the GPU stages, set `MAX_STEPS` in +`04-cpt/tau.yaml` or `06-grpo/tau.yaml` (for example `MAX_STEPS: "5"`). A +short CPT run produces a model that is not useful for GRPO. + +The driver tests run on a CPU-only machine with a local Ray install: + +```bash +cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_reward.py +cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_driver.py +cd examples/cobol-cpt-grpo/04-cpt && python -m pytest test_continued_pretrain.py test_driver.py +cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_driver.py +``` + +Run `test_reward.py` in its own process; it replaces `torch` and `trl` with +stubs. `06-grpo/test_driver.py` starts a local Ray cluster and takes one to +five minutes. + +## Results + +The GRPO model was evaluated with `07-eval` (greedy decoding, 512 new tokens, +TP=8 on 8 H200): + +| Metric | GRPO model | +| --- | --- | +| Compile, overall | **64/220** | +| Compile, paragraph completion | 18/70 (7 free + 11 fixed) | +| Compile, program completion | 18/50 | +| Compile, structural validation | 28/100 (16 free + 12 fixed) | +| COBOL perplexity | 1.49 | +| BLEU, paragraph / program | 0.155 / 0.190 | +| PERFORM reference accuracy | 41.8% | +| Column formatting | 72.0% | +| PIC clause accuracy | 98.9% | + +The checkpoint was trained with an earlier Job-based launch of the same CPT +and GRPO code. Running the `07-eval` RayJob on that checkpoint reproduced +every number above exactly. + +> **TODO:** add full RayJob rerun results (CPT, GRPO, and eval all launched +> from this example). diff --git a/examples/cobol-cpt-grpo/image-requirements.txt b/examples/cobol-cpt-grpo/image-requirements.txt new file mode 100644 index 00000000..25b34780 --- /dev/null +++ b/examples/cobol-cpt-grpo/image-requirements.txt @@ -0,0 +1,32 @@ +vllm==0.28.0 +torch==2.13.0+cu130 +transformers==5.15.1 +accelerate==1.14.0 +huggingface_hub==1.28.0 +sentencepiece==0.2.2 +protobuf==6.33.6 + +trl==1.13.0 +deepspeed==0.19.7 +datasets==5.0.1 +pyarrow==25.0.1 +fsspec==2026.6.0 +nvidia-nccl-cu13==2.29.7 +hjson==3.1.0 +multiprocess==0.70.19 +pandas==3.0.6 +xxhash==4.0.1 + +ray[default,train]==2.56.0 +aiohttp-cors==0.8.1 +colorful==0.5.8 +distlib==0.4.3 +opencensus==0.11.4 +opencensus-context==0.1.3 +platformdirs==4.12.4 +py-spy==0.4.2 +python-discovery==1.6.1 +smart_open==8.0.2 +tensorboardX==2.6.5 +virtualenv==21.14.6 +wrapt==2.5.0 From 31f1402893b3a8643706ed075d3ffbda578fb933 Mon Sep 17 00:00:00 2001 From: June Liu Date: Wed, 21 Oct 2026 10:40:19 -0700 Subject: [PATCH 02/15] docs(examples): add COBOL full RayJob rerun results Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/cobol-cpt-grpo/README.md | 20 ++++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index 9a5a1e25..a98dbe42 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -232,5 +232,21 @@ The checkpoint was trained with an earlier Job-based launch of the same CPT and GRPO code. Running the `07-eval` RayJob on that checkpoint reproduced every number above exactly. -> **TODO:** add full RayJob rerun results (CPT, GRPO, and eval all launched -> from this example). +A full rerun with every stage launched from this example scored lower. The +Job-based CPT column is the checkpoint that the GRPO model above started from. + + +| Metric | Job-based CPT | RayJob CPT | RayJob GRPO | +| --- | --- | --- | --- | +| Compile, overall | 58/220 | 53/220 | 57/220 | +| Compile, paragraph completion | 16/70 | 14/70 | 14/70 | +| Compile, program completion | 16/50 | 18/50 | 18/50 | +| Compile, structural validation | 26/100 | 21/100 | 25/100 | +| COBOL perplexity | 1.49 | 1.48 | 1.48 | +| PERFORM reference accuracy | 40.4% | 27.3% | 28.1% | + +All three columns were scored with the same `07-eval` RayJob. GRPO adds a +similar amount on either CPT checkpoint (+4 on the RayJob CPT, +6 on the +Job-based CPT), so the gap comes from CPT. The cause is not yet known. The +training image moved to newer `transformers` and DeepSpeed releases than the +Job-based run used. From 9f4264a7cc09afa0963a75d00063fd538ac11c54 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 12:48:11 -0700 Subject: [PATCH 03/15] fix(examples): address COBOL pipeline review feedback - 05-rl-prompts: exclude programs that share an 8-line window with the validation split and fail if any overlap remains. - 06-grpo: keep each prompt's compile_fmt and compile with the matching -free or -fixed flag in preflight and the reward. Check the KL and reward-hacking alarms on every rank and combine the stop decision. - 03-benchmarks: find PROCEDURE DIVISION case-insensitively so mixed-case programs no longer leak the procedure body into paragraph and copybook prompts. - 07-eval: score copybook_usage with BLEU and edit distance. Fail before loading the model if a suite is missing or empty or cobc is unusable, and fail if cobc disappears mid-run. - 01-collect: exit nonzero when a required step or clone fails or no COBOL files are collected. - Bump virtualenv to 21.14.7 (GHSA-vxj4-q4g9-crj7), drop trust_remote_code, bind test sockets to the node IP, and replace empty excepts with logged handling. - README: describe the changes and mark the results as coming from the earlier pipeline, with the overlap and the dropped fixed-format prompts. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../01-collect/collect-cobol-data.sh | 52 +++-- .../01-collect/collect-gated.py | 2 - .../01-collect/collect-rl-data.sh | 42 +++-- .../02-data-prep/prepare_data.py | 11 +- .../03-benchmarks/build_benchmarks.py | 17 +- .../03-benchmarks/test_build_benchmarks.py | 67 +++++++ .../04-cpt/continued_pretrain.py | 3 +- .../05-rl-prompts/prepare_rl_prompts.py | 79 +++++++- .../cobol-cpt-grpo/05-rl-prompts/tau.yaml | 2 + .../05-rl-prompts/test_prepare_rl_prompts.py | 101 ++++++++++ examples/cobol-cpt-grpo/06-grpo/driver.py | 10 +- .../cobol-cpt-grpo/06-grpo/test_driver.py | 9 +- .../cobol-cpt-grpo/06-grpo/test_reward.py | 177 +++++++++++++++++- .../cobol-cpt-grpo/06-grpo/test_stop_sync.py | 82 ++++++++ examples/cobol-cpt-grpo/06-grpo/train_grpo.py | 156 +++++++++------ .../cobol-cpt-grpo/07-eval/run_eval_vllm.py | 82 ++++++-- .../cobol-cpt-grpo/07-eval/test_run_eval.py | 102 ++++++++++ examples/cobol-cpt-grpo/README.md | 71 +++++-- .../cobol-cpt-grpo/image-requirements.txt | 2 +- 19 files changed, 929 insertions(+), 138 deletions(-) create mode 100644 examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py create mode 100644 examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py create mode 100644 examples/cobol-cpt-grpo/06-grpo/test_stop_sync.py create mode 100644 examples/cobol-cpt-grpo/07-eval/test_run_eval.py diff --git a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh index 7d69f241..82a9d93a 100755 --- a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh +++ b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh @@ -10,14 +10,30 @@ PROCESSED="${DATA_ROOT}/processed" LOGS="${DATA_ROOT}/logs" mkdir -p "${RAW}" "${PROCESSED}" "${LOGS}" -log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*" | tee -a "${LOGS}/collect.log"; } +LOG_FILE="${LOGS}/collect.log" +FAILED_STEPS=() + +log() { echo "[$(date -u +%Y-%m-%dT%H:%M:%SZ)] $*" | tee -a "${LOG_FILE}"; } run_step() { local name="$1"; shift log "--- ${name} ---" if "$@"; then log "${name}: OK" else - log "${name}: FAILED (continuing)" + log "${name}: FAILED" + FAILED_STEPS+=("${name}") + fi +} +fail_if_incomplete() { + local records="$1" + if [ "${#FAILED_STEPS[@]}" -gt 0 ]; then + log "ERROR: ${#FAILED_STEPS[@]} required step(s) failed:" + printf ' %s\n' "${FAILED_STEPS[@]}" | tee -a "${LOG_FILE}" + exit 1 + fi + if [ "${records}" -eq 0 ]; then + log "ERROR: collection produced no COBOL source files" + exit 1 fi } @@ -26,7 +42,8 @@ log "Target: ${DATA_ROOT}" SOURCE_LICENSES="${SOURCE_LICENSES:-permissive}" log "Source licenses: ${SOURCE_LICENSES}" -pip install -q huggingface_hub==1.28.0 datasets==5.0.1 pandas==3.0.6 pyarrow==25.0.1 2>&1 | tail -1 +pip install -q huggingface_hub==1.28.0 datasets==5.0.1 pandas==3.0.6 pyarrow==25.0.1 2>&1 | tail -1 \ + || { log "ERROR: pip install failed"; exit 1; } log "pip install done" collect_troy_cobol() { @@ -35,7 +52,7 @@ import os, json from datasets import load_dataset out = os.environ.get("RAW", "/data/datasets/cobol/raw") + "/troy-cobol" os.makedirs(out, exist_ok=True) -ds = load_dataset("TroyDoesAI/Cobol", split="train", trust_remote_code=True) +ds = load_dataset("TroyDoesAI/Cobol", split="train") for i, row in enumerate(ds): path = row.get("path", f"file_{i}.cbl") fname = f"{i:06d}_{os.path.basename(path)}" @@ -68,15 +85,19 @@ import os, json from datasets import load_dataset out = os.environ.get("RAW", "/data/datasets/cobol/raw") + "/mainframebench" os.makedirs(out, exist_ok=True) +failed = [] for split_name in ["question_answering", "multiple_choice_question", "COBOL_code_summarization"]: try: - ds = load_dataset("Fsoft-AIC/MainframeBench", split_name, split="train", trust_remote_code=True) + ds = load_dataset("Fsoft-AIC/MainframeBench", split_name, split="train") with open(os.path.join(out, f"{split_name}.jsonl"), "w") as f: for row in ds: f.write(json.dumps(dict(row)) + "\n") print(f" {split_name}: {len(ds)} samples") except Exception as e: - print(f" {split_name}: skipped ({e})") + print(f" {split_name}: failed ({e})") + failed.append(split_name) +if failed: + raise SystemExit(f"MainframeBench splits failed: {', '.join(failed)}") PYEOF } run_step "MainframeBench (7,052 samples)" collect_mainframebench @@ -86,6 +107,7 @@ collect_eval_benchmarks() { import os, json from datasets import load_dataset raw = os.environ.get("RAW", "/data/datasets/cobol/raw") +failed = [] for name, hf_id in [ ("coboleval", "msplits/human-eval-cobol"), ("cobolcodebench", "harshini-kumar/CobolCodeBench"), @@ -97,12 +119,12 @@ for name, hf_id in [ ds = None for s in splits: try: - ds = load_dataset(hf_id, split=s, trust_remote_code=True) + ds = load_dataset(hf_id, split=s) break except Exception: continue if ds is None: - ds = load_dataset(hf_id, trust_remote_code=True) + ds = load_dataset(hf_id) first_split = list(ds.keys())[0] ds = ds[first_split] with open(os.path.join(out, f"{name}.jsonl"), "w") as f: @@ -111,20 +133,24 @@ for name, hf_id in [ print(f"{name}: {len(ds)} samples") except Exception as e: print(f"{name}: error - {e}") + failed.append(name) +if failed: + raise SystemExit(f"Evaluation datasets failed: {', '.join(failed)}") PYEOF } run_step "Evaluation benchmarks" collect_eval_benchmarks collect_xcobol() { mkdir -p "${RAW}/x-cobol" - curl -fsSL "https://zenodo.org/api/records/14269462/files-archive" -o "${RAW}/x-cobol/x-cobol-v2.zip" - cd "${RAW}/x-cobol" && unzip -qo x-cobol-v2.zip -d extracted + curl -fsSL "https://zenodo.org/api/records/14269462/files-archive" -o "${RAW}/x-cobol/x-cobol-v2.zip" || return 1 + (cd "${RAW}/x-cobol" && unzip -qo x-cobol-v2.zip -d extracted) || return 1 find "${RAW}/x-cobol/extracted" -name '*.zip' -exec sh -c 'unzip -qo "$1" -d "$(dirname "$1")/$(basename "$1" .zip)"' _ {} \; XCOUNT=$(find ${RAW}/x-cobol/extracted -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) log "X-COBOL: ${XCOUNT} COBOL files" if [ "${XCOUNT}" -eq 0 ]; then log "X-COBOL: listing top-level contents for debugging:" find "${RAW}/x-cobol/extracted" -maxdepth 3 -type f 2>/dev/null | head -30 | tee -a "${LOGS}/collect.log" + return 1 fi } run_step "X-COBOL v2 from Zenodo" collect_xcobol @@ -137,7 +163,10 @@ clone_repo() { local url="$1" dir="$2" if [ ! -d "${GITHUB}/${dir}/.git" ]; then log " cloning ${dir}..." - git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1 || log " WARN: failed to clone ${dir}" + if ! git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1; then + log " FAILED to clone ${dir}" + FAILED_STEPS+=("clone ${dir}") + fi else log " ${dir} already exists, skipping" fi @@ -199,6 +228,7 @@ TOTAL_COBOL=$(find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -inam TOTAL_JSONL=$(find "${RAW}" -type f -name '*.jsonl' | wc -l) log "Total COBOL source files: ${TOTAL_COBOL}" log "Total JSONL data files: ${TOTAL_JSONL}" +fail_if_incomplete "${TOTAL_COBOL}" cat > "${DATA_ROOT}/manifest.json" <&1 | tail -3 +pip install -q huggingface_hub==1.28.0 datasets==5.0.1 pandas==3.0.6 pyarrow==25.0.1 2>&1 | tail -3 \ + || { log "ERROR: pip install failed"; exit 1; } log "pip install done" collect_the_stack_v1() { @@ -88,16 +105,16 @@ os.makedirs(src_dir, exist_ok=True) print("Loading bigcode/the-stack-v2-dedup COBOL...", flush=True) try: ds = load_dataset("bigcode/the-stack-v2-dedup", "COBOL", split="train", - token=token, trust_remote_code=True) + token=token) except Exception as e1: print(f" Config 'COBOL' failed ({e1}), trying data_dir...", flush=True) try: ds = load_dataset("bigcode/the-stack-v2-dedup", data_dir="data/COBOL", - split="train", token=token, trust_remote_code=True) + split="train", token=token) except Exception as e2: print(f" data_dir 'data/COBOL' failed ({e2}), trying lowercase...", flush=True) ds = load_dataset("bigcode/the-stack-v2-dedup", data_dir="data/cobol", - split="train", token=token, trust_remote_code=True) + split="train", token=token) print(f" Loaded {len(ds)} records", flush=True) print(f" Columns: {ds.column_names}", flush=True) @@ -197,16 +214,14 @@ os.makedirs(src_dir, exist_ok=True) print("Loading zorse/zorse...", flush=True) try: - ds = load_dataset("zorse/zorse", split="train", trust_remote_code=True) + ds = load_dataset("zorse/zorse", split="train") except Exception: try: - ds_dict = load_dataset("zorse/zorse", trust_remote_code=True) + ds_dict = load_dataset("zorse/zorse") first_split = list(ds_dict.keys())[0] ds = ds_dict[first_split] except Exception as e: print(f" Failed to load zorse/zorse: {e}", flush=True) - with open(os.path.join(out, "metadata.json"), "w") as f: - json.dump({"source": "zorse/zorse", "error": str(e)}, f, indent=2) raise print(f" Loaded {len(ds)} records", flush=True) @@ -261,7 +276,10 @@ clone_repo() { return 0 fi log " cloning ${dir}..." - git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1 || log " WARN: failed to clone ${dir}" + if ! git clone --depth 1 --single-branch "$url" "${GITHUB}/${dir}" 2>&1 | tail -1; then + log " FAILED to clone ${dir}" + FAILED_STEPS+=("clone ${dir}") + fi } clone_repo "https://github.com/uwol/proleap-cobol-parser.git" "proleap" @@ -311,6 +329,7 @@ collect_nist_from_oss4j() { log "NIST extracted: ${count} files from ${nist_src}" else log "NIST: could not find test directory in opensourcecobol4j" + return 1 fi } if [ "${SOURCE_LICENSES}" = "all" ]; then @@ -333,6 +352,7 @@ log "Total NEW COBOL source files: ${TOTAL_NEW}" TOTAL_ALL=$(find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) log "Total ALL COBOL source files (old + new): ${TOTAL_ALL}" +fail_if_incomplete "${TOTAL_NEW}" log "Rebuilding file index..." find "${RAW}" -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) > "${DATA_ROOT}/source-file-index.txt" diff --git a/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py b/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py index 71b91c98..03e03245 100755 --- a/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py +++ b/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py @@ -54,6 +54,7 @@ def discover_files_from_index(raw_dir, data_root): if os.path.isfile(per_file_stats): print(f" Reading file index from {per_file_stats}") + malformed = 0 with open(per_file_stats, "r") as f: for line in f: line = line.strip() @@ -65,8 +66,10 @@ def discover_files_from_index(raw_dir, data_root): if rel: full = os.path.join(raw_dir, rel) source_paths.append(full) - except (json.JSONDecodeError, KeyError): - pass + except json.JSONDecodeError: + malformed += 1 + if malformed: + print(f" WARNING: skipped {malformed} malformed lines in {per_file_stats}") print(f" Found {len(source_paths)} source file paths from tokenizer index") elif os.path.isfile(file_index): print(f" Reading file index from {file_index}") @@ -185,8 +188,8 @@ def ingest_jsonl_from_paths(raw_dir, jsonl_paths): "dataset": source_name, } ) - except Exception: - pass + except (OSError, UnicodeDecodeError) as e: + print(f" WARNING: skipped {jpath}: {e}") return records diff --git a/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py index d77e9f72..40e92ae8 100755 --- a/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py +++ b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py @@ -5,11 +5,12 @@ """ Build COBOL evaluation benchmarks from the prepared validation set. -Reads val.jsonl and constructs four benchmark suites: +Reads val.jsonl and constructs five benchmark suites: 1. paragraph_completion — mask a PROCEDURE DIVISION paragraph body, predict it 2. copybook_usage — given a copybook + partial program, predict field references 3. program_completion — mask PROCEDURE DIVISION of a complete program 4. perplexity — held-out complete programs for perplexity measurement + 5. structural_validation — complete a PROCEDURE DIVISION that must compile Each benchmark is a JSONL file with prompt/reference/metadata fields. Writes to /data/datasets/cobol/eval/ on the PVC. @@ -65,6 +66,12 @@ def split_divisions(text): return divisions +def text_before_procedure_division(text): + """Return the original text before the first PROCEDURE DIVISION, in any case.""" + proc_idx = text.upper().find("PROCEDURE DIVISION") + return text[:proc_idx] if proc_idx >= 0 else "" + + def extract_paragraphs(procedure_text): """Extract named paragraphs from PROCEDURE DIVISION text.""" paragraphs = [] @@ -113,11 +120,7 @@ def build_paragraph_completion(records, max_tasks=200): proc_div_header = " PROCEDURE DIVISION" proc_prefix = proc_text[: para["start"]] - non_proc = ( - text.split("PROCEDURE DIVISION")[0] - if "PROCEDURE DIVISION" in text.upper() - else "" - ) + non_proc = text_before_procedure_division(text) prompt = non_proc.rstrip() + "\n" + proc_div_header + proc_prefix prompt = prompt.rstrip() + "\n " + para["name"] + ".\n" @@ -204,7 +207,7 @@ def build_copybook_usage(records, max_tasks=100): "Given this copybook:\n" + cb_info["text"].strip() + "\n\nAnd this partial COBOL program:\n" - + text.split("PROCEDURE DIVISION")[0].rstrip() + + text_before_procedure_division(text).rstrip() + "\n PROCEDURE DIVISION.\n" + context_before.rstrip() + "\n [COMPLETE THE FOLLOWING LINES]\n" diff --git a/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py new file mode 100644 index 00000000..c126db26 --- /dev/null +++ b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py @@ -0,0 +1,67 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local tests for the prompt boundary in build_benchmarks. + +Needs no cobc and no GPU: + python -m pytest examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py -v + +Checks that a PROCEDURE DIVISION header in any case ends the prompt prefix, so the +procedure body that holds the reference never appears in a prompt. +""" + +import importlib.util +import os + +import pytest + +BENCH_DIR = os.path.dirname(os.path.abspath(__file__)) + +PROGRAM = ( + " IDENTIFICATION DIVISION.\n" + " PROGRAM-ID. MIXED.\n" + " DATA DIVISION.\n" + " WORKING-STORAGE SECTION.\n" + " 01 WS-COUNT PIC 9(4).\n" + " Procedure Division.\n" + " MAIN-PARA.\n" + " MOVE 1 TO WS-COUNT\n" + " ADD 2 TO WS-COUNT\n" + " DISPLAY 'SECRET REFERENCE LINE'\n" + " STOP RUN.\n" +) + + +@pytest.fixture +def bench(tmp_path, monkeypatch): + monkeypatch.setenv("OUT", str(tmp_path / "eval")) + spec = importlib.util.spec_from_file_location( + "build_benchmarks", os.path.join(BENCH_DIR, "build_benchmarks.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize( + "header", ["PROCEDURE DIVISION", "Procedure Division", "procedure division"] +) +def test_prefix_stops_at_procedure_division_in_any_case(bench, header): + text = PROGRAM.replace("Procedure Division", header) + prefix = bench.text_before_procedure_division(text) + assert prefix == text[: text.index(header)] + assert "SECRET REFERENCE LINE" not in prefix + + +def test_prefix_is_empty_without_procedure_division(bench): + assert bench.text_before_procedure_division(" DATA DIVISION.\n") == "" + + +def test_paragraph_prompt_does_not_leak_mixed_case_body(bench): + tasks = bench.build_paragraph_completion( + [{"text": PROGRAM, "is_complete_program": True, "source": "mixed.cbl"}] + ) + assert len(tasks) == 1 + assert "SECRET REFERENCE LINE" in tasks[0]["reference"] + assert "SECRET REFERENCE LINE" not in tasks[0]["prompt"] + assert "WS-COUNT PIC 9(4)" in tasks[0]["prompt"] diff --git a/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py index 65220791..0628db60 100644 --- a/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py +++ b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py @@ -212,7 +212,7 @@ def main(): if is_main: print("\nLoading tokenizer...", flush=True) - tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token @@ -246,7 +246,6 @@ def main(): model = AutoModelForCausalLM.from_pretrained( MODEL_ID, dtype=torch.bfloat16, - trust_remote_code=True, attn_implementation="sdpa", low_cpu_mem_usage=True, ) diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py index 34496c84..cab02852 100755 --- a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py +++ b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py @@ -9,8 +9,14 @@ PROCEDURE DIVISION prompts, compile-tests each with cobc, and writes an RL-ready JSONL. +Programs that share any VAL_OVERLAP_LINES consecutive normalized lines with a +val.jsonl program are excluded, because 03-benchmarks builds the eval suites +from val.jsonl. The stage fails if val.jsonl is missing or empty, or if any +written prompt still overlaps the validation split. + Input: /data/datasets/cobol/source-file-index.txt (rebuilt by 01-collect/collect-rl-data.sh) /data/datasets/cobol/raw/**/*.{cbl,cob,cpy} + /data/datasets/cobol/prepared/val.jsonl (from 02-data-prep) Output: /data/datasets/cobol/prepared/rl-prompts.jsonl /data/datasets/cobol/prepared/rl-prep-report.json @@ -38,6 +44,8 @@ MIN_BODY_LINES = int(os.environ.get("MIN_BODY_LINES", "5")) MAX_BODY_LINES = int(os.environ.get("MAX_BODY_LINES", "200")) MAX_PROMPT_CHARS = int(os.environ.get("MAX_PROMPT_CHARS", "50000")) +VAL_PATH = os.environ.get("VAL", os.path.join(OUT, "val.jsonl")) +VAL_OVERLAP_LINES = int(os.environ.get("VAL_OVERLAP_LINES", "8")) os.makedirs(OUT, exist_ok=True) @@ -53,6 +61,8 @@ COPY_STMT_RE = re.compile(r"^\s*COPY\s+\S+.*?\.\s*$", re.IGNORECASE | re.MULTILINE) +WHITESPACE_RE = re.compile(r"\s+") + stats = Counter() @@ -95,6 +105,60 @@ def exact_dedup(texts_with_paths): return unique, duped +def normalized_lines(text): + return [ + WHITESPACE_RE.sub(" ", ln).strip().upper() + for ln in text.replace("\r", "").split("\n") + if ln.strip() + ] + + +def line_windows(text, size): + lines = normalized_lines(text) + return { + hashlib.sha256("\n".join(lines[i : i + size]).encode()).hexdigest() + for i in range(max(0, len(lines) - size + 1)) + } + + +def load_val_windows(path, size): + if not os.path.isfile(path): + raise RuntimeError(f"{path} not found; run 02-data-prep first") + windows = set() + count = 0 + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + windows |= line_windows(json.loads(line).get("text", ""), size) + count += 1 + if not count or not windows: + raise RuntimeError(f"{path} has no validation programs") + return windows, count + + +def overlaps_val(text, val_windows, size): + return not line_windows(text, size).isdisjoint(val_windows) + + +def exclude_val_overlap(texts_with_paths, val_windows, size): + kept = [ + (path, text) + for path, text in texts_with_paths + if not overlaps_val(text, val_windows, size) + ] + return kept, len(texts_with_paths) - len(kept) + + +def assert_disjoint_from_val(records, val_windows, size): + leaked = [r["path"] for r in records if overlaps_val(r["text"], val_windows, size)] + if leaked: + raise RuntimeError( + f"{len(leaked)} RL prompts overlap the validation split, e.g. {leaked[0]}" + ) + + def extract_prompt(text): upper = text.upper() if "IDENTIFICATION DIVISION" not in upper and "PROGRAM-ID" not in upper: @@ -257,6 +321,9 @@ def main(): subprocess.run(["cobc", "--version"], capture_output=True, check=True) log("cobc available") + val_windows, val_count = load_val_windows(VAL_PATH, VAL_OVERLAP_LINES) + log(f"Loaded {val_count} validation programs from {VAL_PATH}") + file_paths = discover_files() log("Reading source files...") @@ -275,6 +342,12 @@ def main(): texts, duped = exact_dedup(texts) log(f"After exact dedup: {len(texts)} unique ({duped} duplicates removed)") + texts, val_overlap = exclude_val_overlap(texts, val_windows, VAL_OVERLAP_LINES) + log( + f"After validation-split exclusion: {len(texts)} remain " + f"({val_overlap} share {VAL_OVERLAP_LINES}+ lines with val.jsonl)" + ) + log(f"Processing {len(texts)} files with {MAX_WORKERS} workers...") results = [] done = 0 @@ -305,6 +378,7 @@ def main(): log(f" {status}: {count}") ok_results = [r for r in results if r["status"] == "ok"] + assert_disjoint_from_val(ok_results, val_windows, VAL_OVERLAP_LINES) log(f"Writing {len(ok_results)} compilable programs to rl-prompts.jsonl...") out_path = os.path.join(OUT, "rl-prompts.jsonl") @@ -329,7 +403,10 @@ def main(): "total_files_discovered": len(file_paths), "total_read": len(texts) + duped, "exact_dupes_removed": duped, - "after_dedup": len(texts), + "after_dedup": len(texts) + val_overlap, + "val_overlap_removed": val_overlap, + "val_overlap_lines": VAL_OVERLAP_LINES, + "after_val_exclusion": len(texts), "status_counts": dict(stats.most_common()), "compilable_programs": len(ok_results), "output_file": out_path, diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml b/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml index bd2c6363..c1c93286 100644 --- a/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml +++ b/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml @@ -13,6 +13,8 @@ runtime: RAW: /data/datasets/cobol/raw DATA_ROOT: /data/datasets/cobol OUT: /data/datasets/cobol/prepared + VAL: /data/datasets/cobol/prepared/val.jsonl + VAL_OVERLAP_LINES: "8" MAX_WORKERS: "8" MIN_BODY_LINES: "5" MAX_BODY_LINES: "200" diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py new file mode 100644 index 00000000..74ec1507 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py @@ -0,0 +1,101 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local tests for the validation-split exclusion in prepare_rl_prompts. + +Needs no cobc and no GPU: + python -m pytest examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py -v +""" + +import importlib.util +import json +import os + +import pytest + +RL_DIR = os.path.dirname(os.path.abspath(__file__)) + + +@pytest.fixture +def prep(tmp_path, monkeypatch): + monkeypatch.setenv("OUT", str(tmp_path / "prepared")) + spec = importlib.util.spec_from_file_location( + "prepare_rl_prompts", os.path.join(RL_DIR, "prepare_rl_prompts.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _program(name, body_lines): + head = [ + " IDENTIFICATION DIVISION.", + f" PROGRAM-ID. {name}.", + " PROCEDURE DIVISION.", + ] + return "\n".join(head + body_lines) + "\n" + + +SHARED = [f" DISPLAY 'SHARED LINE {i}'." for i in range(10)] +OTHER = [f" DISPLAY 'TRAIN ONLY {i}'." for i in range(10)] + + +def _write_val(path, texts): + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w") as f: + for text in texts: + f.write(json.dumps({"text": text}) + "\n") + + +def test_val_paths_default_to_prepared_dir(prep, tmp_path): + assert prep.VAL_PATH == str(tmp_path / "prepared" / "val.jsonl") + assert prep.VAL_OVERLAP_LINES == 8 + + +def test_windows_ignore_case_whitespace_and_blank_lines(prep): + a = "\n".join(SHARED) + b = "\n\n".join(" " + ln.lower().replace(" ", " ") for ln in SHARED) + assert prep.line_windows(a, 8) == prep.line_windows(b, 8) + assert len(prep.line_windows(a, 8)) == 3 + + +def test_exclude_val_overlap_drops_programs_sharing_a_window(prep, tmp_path): + val_path = tmp_path / "prepared" / "val.jsonl" + _write_val(val_path, [_program("VALPROG", SHARED)]) + windows, count = prep.load_val_windows(str(val_path), 8) + assert count == 1 + + copied = _program("RENAMED", OTHER[:2] + SHARED[1:9] + OTHER[2:4]) + short = _program("SHORT", OTHER[:3] + SHARED[:7] + OTHER[3:6]) + clean = _program("CLEAN", OTHER) + kept, removed = prep.exclude_val_overlap( + [("copied.cbl", copied), ("short.cbl", short), ("clean.cbl", clean)], + windows, + 8, + ) + assert removed == 1 + assert [p for p, _ in kept] == ["short.cbl", "clean.cbl"] + + +def test_assert_disjoint_raises_on_overlap(prep, tmp_path): + val_path = tmp_path / "prepared" / "val.jsonl" + _write_val(val_path, [_program("VALPROG", SHARED)]) + windows, _ = prep.load_val_windows(str(val_path), 8) + + prep.assert_disjoint_from_val( + [{"path": "clean.cbl", "text": _program("CLEAN", OTHER)}], windows, 8 + ) + with pytest.raises(RuntimeError, match="1 RL prompts overlap"): + prep.assert_disjoint_from_val( + [{"path": "val.cbl", "text": _program("VALPROG", SHARED)}], windows, 8 + ) + + +@pytest.mark.parametrize("contents", [None, "", "\n\n"]) +def test_load_val_windows_fails_without_validation_programs(prep, tmp_path, contents): + val_path = tmp_path / "prepared" / "val.jsonl" + if contents is not None: + val_path.parent.mkdir(parents=True, exist_ok=True) + val_path.write_text(contents) + with pytest.raises(RuntimeError): + prep.load_val_windows(str(val_path), 8) diff --git a/examples/cobol-cpt-grpo/06-grpo/driver.py b/examples/cobol-cpt-grpo/06-grpo/driver.py index a15d2bff..0b087817 100644 --- a/examples/cobol-cpt-grpo/06-grpo/driver.py +++ b/examples/cobol-cpt-grpo/06-grpo/driver.py @@ -29,6 +29,7 @@ returns or fails. """ +import contextlib import os import signal import subprocess @@ -99,11 +100,12 @@ def stop(self): try: self._proc.wait(timeout=STOP_TIMEOUT_S) except subprocess.TimeoutExpired: - pass - try: + print( + f"vLLM server still running after {STOP_TIMEOUT_S}s; killing it", + flush=True, + ) + with contextlib.suppress(ProcessLookupError, PermissionError): os.killpg(self._proc.pid, signal.SIGKILL) - except (ProcessLookupError, PermissionError): - pass self._proc.wait() diff --git a/examples/cobol-cpt-grpo/06-grpo/test_driver.py b/examples/cobol-cpt-grpo/06-grpo/test_driver.py index 75ad05e0..ce6f4719 100644 --- a/examples/cobol-cpt-grpo/06-grpo/test_driver.py +++ b/examples/cobol-cpt-grpo/06-grpo/test_driver.py @@ -33,7 +33,9 @@ import sys import time -port, delay, mode, out = int(sys.argv[1]), float(sys.argv[2]), sys.argv[3], sys.argv[4] +port, delay, mode, out, host = ( + int(sys.argv[1]), float(sys.argv[2]), sys.argv[3], sys.argv[4], sys.argv[5] +) with open(out, "w") as f: json.dump({"pid": os.getpid(), "cuda": os.environ.get("CUDA_VISIBLE_DEVICES")}, f) time.sleep(delay) @@ -53,7 +55,7 @@ def log_message(self, *args): pass -http.server.HTTPServer(("0.0.0.0", port), Handler).serve_forever() +http.server.HTTPServer((host, port), Handler).serve_forever() """ STUB_TRAIN_GRPO = """\ @@ -96,7 +98,7 @@ def main(): def _free_port(): with socket.socket() as s: - s.bind(("", 0)) + s.bind((ray.util.get_node_ip_address(), 0)) return s.getsockname()[1] @@ -147,6 +149,7 @@ def _server_cmd(work_dir, port, delay, mode, record): str(delay), mode, str(record), + ray.util.get_node_ip_address(), ] diff --git a/examples/cobol-cpt-grpo/06-grpo/test_reward.py b/examples/cobol-cpt-grpo/06-grpo/test_reward.py index dea08904..d102addf 100644 --- a/examples/cobol-cpt-grpo/06-grpo/test_reward.py +++ b/examples/cobol-cpt-grpo/06-grpo/test_reward.py @@ -10,12 +10,15 @@ python examples/cobol-cpt-grpo/06-grpo/test_reward.py """ +import json import os import subprocess import sys +import tempfile import types import unittest -from unittest import mock +import unittest.mock as mock +from concurrent.futures import ThreadPoolExecutor os.environ.setdefault("WORLD_SIZE", "2") @@ -55,6 +58,7 @@ sys.path.insert(0, os.path.dirname(__file__)) from train_grpo import ( # noqa: E402 + TauGRPOMetricsCallback, _compilation_reward, _count_cobol_verbs, _count_noop_verbs, @@ -64,6 +68,9 @@ _strip_non_code, _structural_reward, _try_compile, + compile_format, + load_prompts, + reward_fn, ) @@ -714,5 +721,173 @@ def test_go_to_program_gets_full_reward(self): self.assertEqual(total, 1.0) +FIXED_PROGRAM = ( + " IDENTIFICATION DIVISION.\n" + " PROGRAM-ID. FIXEDP.\n" + " PROCEDURE DIVISION.\n" + " DISPLAY 'X'.\n" * 6 +) + + +def _ok(args, **kwargs): + return subprocess.CompletedProcess(args, 0, b"", b"") + + +class TestCompileFormat(unittest.TestCase): + def test_normalizes_stage_05_values(self): + self.assertEqual(compile_format("free"), "free") + self.assertEqual(compile_format("fixed"), "fixed") + self.assertEqual(compile_format("fixed-copybook-only"), "fixed") + self.assertEqual(compile_format("free-copybook-only"), "free") + + def test_missing_defaults_to_free(self): + self.assertEqual(compile_format(None), "free") + self.assertEqual(compile_format(""), "free") + + def test_unknown_raises(self): + with self.assertRaises(ValueError): + compile_format("failed") + + def test_try_compile_passes_format_flag(self): + with mock.patch("train_grpo.subprocess.run", side_effect=_ok) as run: + _try_compile("IDENTIFICATION DIVISION.", "fixed") + _try_compile("IDENTIFICATION DIVISION.") + self.assertEqual(run.call_args_list[0].args[0][2], "-fixed") + self.assertEqual(run.call_args_list[1].args[0][2], "-free") + + +class TestLoadPromptsKeepsFormat(unittest.TestCase): + def _write(self, records): + f = tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False) + with f: + for rec in records: + f.write(json.dumps(rec) + "\n") + self.addCleanup(os.unlink, f.name) + return f.name + + def test_preflight_compiles_each_record_in_its_format(self): + path = self._write( + [ + {"text": FIXED_PROGRAM, "compile_fmt": "fixed"}, + { + "text": FIXED_PROGRAM.replace("FIXEDP", "FREEP"), + "compile_fmt": "free", + }, + {"text": FIXED_PROGRAM.replace("FIXEDP", "OLDREC")}, + ] + ) + + def fixed_only(args, **kwargs): + rc = 0 if args[1] == "--version" or args[2] == "-fixed" else 1 + return subprocess.CompletedProcess(args, rc, b"", b"") + + with mock.patch("train_grpo.subprocess.run", side_effect=fixed_only): + prompts = load_prompts(path) + self.assertEqual([r["compile_fmt"] for r in prompts], ["fixed"]) + self.assertIn("FIXEDP", prompts[0]["prompt"]) + + def test_without_preflight_returns_format(self): + path = self._write( + [{"text": FIXED_PROGRAM, "compile_fmt": "fixed-copybook-only"}] + ) + with mock.patch("train_grpo.subprocess.run", side_effect=_ok): + prompts = load_prompts(path, preflight=False) + self.assertEqual(prompts[0]["compile_fmt"], "fixed") + + +class TestRewardFnUsesFormat(unittest.TestCase): + def test_each_completion_compiled_in_its_prompt_format(self): + seen = [] + + def fake_compile(text, fmt="free"): + seen.append((text, fmt)) + return 0, "" + + with ( + mock.patch("train_grpo.ProcessPoolExecutor", ThreadPoolExecutor), + mock.patch("train_grpo._try_compile", side_effect=fake_compile), + ): + rewards = reward_fn( + ["A", "B"], + [" MOVE 1 TO X\n", " MOVE 2 TO Y\n"], + compile_fmt=["fixed", "free"], + ) + self.assertEqual(len(rewards), 2) + self.assertEqual( + sorted(seen), + [ + ("A MOVE 1 TO X\n", "fixed"), + ("B MOVE 2 TO Y\n", "free"), + ], + ) + + def test_missing_format_column_defaults_to_free(self): + seen = [] + + def fake_compile(text, fmt="free"): + seen.append(fmt) + return 0, "" + + with ( + mock.patch("train_grpo.ProcessPoolExecutor", ThreadPoolExecutor), + mock.patch("train_grpo._try_compile", side_effect=fake_compile), + ): + reward_fn(["A"], [" MOVE 1 TO X\n"]) + self.assertEqual(seen, ["free"]) + + +class TestAlarmStopIsShared(unittest.TestCase): + def _step(self, cb, logs, agg, any_rank): + control = types.SimpleNamespace(should_training_stop=False) + state = types.SimpleNamespace(global_step=5) + with ( + mock.patch("train_grpo._aggregate_reward_stats", return_value=agg), + mock.patch("train_grpo._any_rank", side_effect=any_rank) as synced, + ): + cb.on_log(None, state, control, logs=logs) + return control.should_training_stop, synced + + def test_non_zero_rank_without_publisher_stops_on_local_alarm(self): + cb = TauGRPOMetricsCallback() + cb.kl_alarm_count = 9 + stopped, synced = self._step(cb, {"kl": 6.0}, {}, lambda flag: flag) + self.assertTrue(stopped) + synced.assert_called_once_with(True) + + def test_stops_when_another_rank_raised_the_alarm(self): + cb = TauGRPOMetricsCallback() + stopped, synced = self._step( + cb, {"kl": 0.1}, {"mean_verb_count": 10}, lambda flag: True + ) + self.assertTrue(stopped) + synced.assert_called_once_with(False) + + def test_no_alarm_keeps_training(self): + cb = TauGRPOMetricsCallback() + stopped, _ = self._step( + cb, {"kl": 0.1}, {"mean_verb_count": 10}, lambda flag: flag + ) + self.assertFalse(stopped) + + def test_only_publisher_rank_publishes(self): + publisher = mock.Mock() + self._step( + TauGRPOMetricsCallback(publisher), {"kl": 0.1}, {}, lambda flag: flag + ) + publisher.publish.assert_called_once() + stopped, synced = self._step( + TauGRPOMetricsCallback(), {"kl": 0.1}, {}, lambda flag: flag + ) + synced.assert_called_once() + + def test_step_zero_skips_the_collective(self): + cb = TauGRPOMetricsCallback() + control = types.SimpleNamespace(should_training_stop=False) + with mock.patch("train_grpo._any_rank") as synced: + cb.on_log( + None, types.SimpleNamespace(global_step=0), control, logs={"kl": 9} + ) + synced.assert_not_called() + + if __name__ == "__main__": unittest.main() diff --git a/examples/cobol-cpt-grpo/06-grpo/test_stop_sync.py b/examples/cobol-cpt-grpo/06-grpo/test_stop_sync.py new file mode 100644 index 00000000..1d220088 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/test_stop_sync.py @@ -0,0 +1,82 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local test that a GRPO alarm on one rank stops every rank. + +Starts two CPU processes in a gloo group, each running TauGRPOMetricsCallback without +a publisher. Only rank 1 raises the KL alarm; both ranks must set +should_training_stop at the same step. trl, transformers and datasets are replaced +with stubs, so this needs CPU torch only: + python -m pytest examples/cobol-cpt-grpo/06-grpo/test_stop_sync.py -v +""" + +import json +import os +import socket +import subprocess +import sys + +RL_DIR = os.path.dirname(os.path.abspath(__file__)) + +WORKER = """\ +import json +import os +import sys +import types + +for name in ["datasets", "transformers", "trl"]: + sys.modules[name] = types.ModuleType(name) +sys.modules["datasets"].Dataset = type("Dataset", (), {}) +sys.modules["transformers"].AutoTokenizer = type("AutoTokenizer", (), {}) +sys.modules["transformers"].TrainerCallback = type("TrainerCallback", (), {}) +sys.modules["trl"].GRPOConfig = type("GRPOConfig", (), {}) +sys.modules["trl"].GRPOTrainer = type("GRPOTrainer", (), {}) + +rl_dir, port, rank, out = sys.argv[1], sys.argv[2], int(sys.argv[3]), sys.argv[4] +sys.path.insert(0, rl_dir) + +import torch.distributed as dist + +import train_grpo + +dist.init_process_group( + "gloo", init_method=f"tcp://127.0.0.1:{port}", rank=rank, world_size=2 +) +cb = train_grpo.TauGRPOMetricsCallback() +if rank == 1: + cb.kl_alarm_count = 8 +stops = [] +for step in (1, 2): + control = types.SimpleNamespace(should_training_stop=False) + kl = 6.0 if rank == 1 else 0.1 + cb.on_log(None, types.SimpleNamespace(global_step=step), control, logs={"kl": kl}) + stops.append(control.should_training_stop) +dist.destroy_process_group() +with open(out, "w") as f: + json.dump(stops, f) +""" + + +def _free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def test_alarm_on_one_rank_stops_both(tmp_path): + worker = tmp_path / "worker.py" + worker.write_text(WORKER) + port = str(_free_port()) + outs = [tmp_path / f"rank-{rank}.json" for rank in (0, 1)] + procs = [ + subprocess.Popen( + [sys.executable, str(worker), RL_DIR, port, str(rank), str(outs[rank])] + ) + for rank in (0, 1) + ] + for proc in procs: + assert proc.wait(timeout=120) == 0 + assert [json.loads(out.read_text()) for out in outs] == [ + [False, True], + [False, True], + ] diff --git a/examples/cobol-cpt-grpo/06-grpo/train_grpo.py b/examples/cobol-cpt-grpo/06-grpo/train_grpo.py index ecbb962f..f79ff926 100755 --- a/examples/cobol-cpt-grpo/06-grpo/train_grpo.py +++ b/examples/cobol-cpt-grpo/06-grpo/train_grpo.py @@ -53,6 +53,15 @@ MAX_STEPS = os.environ.get("MAX_STEPS") HF_HOME = os.environ.get("HF_HOME", "/tmp/hf-cache") +COMPILE_FORMAT_FLAGS = {"free": "-free", "fixed": "-fixed"} + + +def compile_format(record_fmt): + fmt = (record_fmt or "free").split("-", 1)[0] + if fmt not in COMPILE_FORMAT_FLAGS: + raise ValueError(f"Unknown compile_fmt {record_fmt!r}") + return fmt + def extract_prompt(text): upper = text.upper() @@ -70,13 +79,13 @@ def extract_prompt(text): return prompt, body, len(body_lines) -def _preflight_compilable(prompt, body): - """Check if prompt + reference body compiles cleanly (rc == 0). - Reject programs with missing copybooks — the compiler can't validate - the procedure body when dependencies are unresolved, so these produce - all-failure groups regardless of completion quality.""" +def _preflight_compilable(prompt, body, fmt="free"): + """Check if prompt + reference body compiles cleanly (rc == 0) in its + source format. Reject programs with missing copybooks — the compiler + can't validate the procedure body when dependencies are unresolved, so + these produce all-failure groups regardless of completion quality.""" full = prompt + body - rc, _ = _try_compile(full) + rc, _ = _try_compile(full, fmt) return rc == 0 @@ -110,13 +119,13 @@ def load_prompts( prompt, body, body_len = extract_prompt(rec.get("text", "")) if prompt is None or body_len < min_body or body_len > max_body: continue - candidates.append((prompt, body)) + candidates.append((prompt, body, compile_format(rec.get("compile_fmt")))) if max_body_tokens and tokenizer: before = len(candidates) candidates = [ - (p, b) - for p, b in candidates + (p, b, fmt) + for p, b, fmt in candidates if len(tokenizer.encode(b, add_special_tokens=False)) <= max_body_tokens ] if GLOBAL_RANK == 0 and len(candidates) < before: @@ -127,18 +136,23 @@ def load_prompts( ) if not preflight: - return [p for p, _ in candidates] + return [{"prompt": p, "compile_fmt": fmt} for p, _, fmt in candidates] prompts = [] skipped = 0 - for prompt, body in candidates: - if _preflight_compilable(prompt, body): - prompts.append(prompt) + for prompt, body, fmt in candidates: + if _preflight_compilable(prompt, body, fmt): + prompts.append({"prompt": prompt, "compile_fmt": fmt}) else: skipped += 1 if GLOBAL_RANK == 0: - print(f" Preflight: {len(prompts)} compilable, {skipped} skipped", flush=True) + n_fixed = sum(r["compile_fmt"] == "fixed" for r in prompts) + print( + f" Preflight: {len(prompts)} compilable ({len(prompts) - n_fixed} free, " + f"{n_fixed} fixed), {skipped} skipped", + flush=True, + ) if not prompts: raise RuntimeError( f"No compilable prompts found in {path} — check compiler and data" @@ -146,14 +160,15 @@ def load_prompts( return prompts -def _try_compile(cobol_text): +def _try_compile(cobol_text, fmt="free"): + flag = COMPILE_FORMAT_FLAGS[compile_format(fmt)] with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: f.write(cobol_text) f.flush() src = f.name try: result = subprocess.run( - ["cobc", "-fsyntax-only", "-free", src], + ["cobc", "-fsyntax-only", flag, src], capture_output=True, timeout=10, ) @@ -237,7 +252,6 @@ def _repl(m): re.IGNORECASE, ) -_DISPLAY_ONLY_RE = re.compile(r"\bDISPLAY\b", re.IGNORECASE) _NON_DISPLAY_VERBS = re.compile( r"\b(?:ACCEPT|ADD|ALLOCATE|ALTER|CALL|CANCEL|CLOSE|COMPUTE|" r"DELETE|DIVIDE|EVALUATE|FREE|GO|IF|INITIALIZE|INSPECT|" @@ -394,9 +408,9 @@ def _structural_reward(completion, max_gen_tokens, completion_n_tokens): def _score_one(args): - prompt, completion, max_gen_tokens, completion_n_tokens = args + prompt, completion, max_gen_tokens, completion_n_tokens, fmt = args full = prompt + completion - rc, stderr = _try_compile(full) + rc, stderr = _try_compile(full, fmt) r_compile = _compilation_reward(rc, stderr) r_struct = _structural_reward(completion, max_gen_tokens, completion_n_tokens) errors, warnings, copybook = _parse_cobc(stderr) @@ -417,8 +431,9 @@ def _score_one(args): _reward_tokenizer = None -def reward_fn(prompts, completions, **kwargs): - """Compute rewards for a batch of (prompt, completion) pairs.""" +def reward_fn(prompts, completions, compile_fmt=None, **kwargs): + """Compute rewards for a batch of (prompt, completion) pairs, compiling + each in the format recorded for its prompt.""" if _reward_tokenizer is not None: n_tokens = [ len(_reward_tokenizer.encode(c, add_special_tokens=False)) @@ -426,7 +441,11 @@ def reward_fn(prompts, completions, **kwargs): ] else: n_tokens = [len(c) // 4 for c in completions] - args = [(p, c, MAX_GEN_LEN, nt) for p, c, nt in zip(prompts, completions, n_tokens)] + fmts = compile_fmt if compile_fmt is not None else ["free"] * len(prompts) + args = [ + (p, c, MAX_GEN_LEN, nt, fmt) + for p, c, nt, fmt in zip(prompts, completions, n_tokens, fmts, strict=True) + ] with ProcessPoolExecutor(max_workers=min(8, len(args))) as pool: results = list(pool.map(_score_one, args)) @@ -468,59 +487,72 @@ def publish(self, step, metrics): os.replace(tmp, final) -_REWARD_STAT_KEYS = [ - "compile_rate", - "copybook_only_rate", - "trivial_rate", - "mean_verb_count", - "mean_noop_count", - "mean_errors", -] - - def _aggregate_reward_stats(): return dict(_reward_stats) +def _any_rank(flag): + dist = torch.distributed + if not (dist.is_available() and dist.is_initialized()): + return flag + if dist.get_backend() == "nccl": + device = torch.device("cuda", torch.cuda.current_device()) + else: + device = torch.device("cpu") + t = torch.tensor([1.0 if flag else 0.0], device=device) + dist.all_reduce(t, op=dist.ReduceOp.MAX) + return bool(t.item()) + + class TauGRPOMetricsCallback(TrainerCallback): - def __init__(self, publisher): + """Installed on every rank. Rank 0 publishes metrics; every rank evaluates + the alarms and the stop decision is shared across ranks, so all ranks + leave the training loop together.""" + + def __init__(self, publisher=None): self.publisher = publisher self.kl_alarm_count = 0 self.short_alarm_count = 0 def on_log(self, args, state, control, logs=None, **kwargs): - if not logs or state.global_step <= 0: - return - if GLOBAL_RANK != 0: + if state.global_step <= 0: return agg = _aggregate_reward_stats() + stop = self._check_alarms(logs, agg) if logs else False + + if self.publisher is not None and logs: + metrics = {} + for k, v in logs.items(): + if isinstance(v, (int, float)) and not math.isnan(v): + metrics[k] = v + + if agg: + metrics["reward/compile_rate"] = agg.get("compile_rate", 0) + metrics["reward/copybook_only_rate"] = agg.get("copybook_only_rate", 0) + metrics["reward/trivial_rate"] = agg.get("trivial_rate", 0) + metrics["gen/verb_count"] = agg.get("mean_verb_count", 0) + metrics["gen/noop_count"] = agg.get("mean_noop_count", 0) + metrics["gen/mean_errors"] = agg.get("mean_errors", 0) + + if metrics: + self.publisher.publish(state.global_step, metrics) + + if _any_rank(stop): + control.should_training_stop = True - metrics = {} - for k, v in logs.items(): - if isinstance(v, (int, float)) and not math.isnan(v): - metrics[k] = v - - if agg: - metrics["reward/compile_rate"] = agg.get("compile_rate", 0) - metrics["reward/copybook_only_rate"] = agg.get("copybook_only_rate", 0) - metrics["reward/trivial_rate"] = agg.get("trivial_rate", 0) - metrics["gen/verb_count"] = agg.get("mean_verb_count", 0) - metrics["gen/noop_count"] = agg.get("mean_noop_count", 0) - metrics["gen/mean_errors"] = agg.get("mean_errors", 0) - - if metrics: - self.publisher.publish(state.global_step, metrics) - + def _check_alarms(self, logs, agg): + stop = False kl = logs.get("kl", 0) if kl > 5.0: self.kl_alarm_count += 1 if self.kl_alarm_count >= 10: print( - f"ALARM: KL divergence > 5.0 for 10 steps ({kl:.2f}). Stopping.", + f"ALARM: KL divergence > 5.0 for 10 steps ({kl:.2f}) on rank " + f"{GLOBAL_RANK}. Stopping.", flush=True, ) - control.should_training_stop = True + stop = True else: self.kl_alarm_count = 0 @@ -529,12 +561,14 @@ def on_log(self, args, state, control, logs=None, **kwargs): self.short_alarm_count += 1 if self.short_alarm_count >= 20: print( - "ALARM: Mean verb count < 3 for 20 steps. Reward hacking. Stopping.", + f"ALARM: Mean verb count < 3 for 20 steps on rank {GLOBAL_RANK}. " + "Reward hacking. Stopping.", flush=True, ) - control.should_training_stop = True + stop = True else: self.short_alarm_count = 0 + return stop class MaxStepsCallback(TrainerCallback): @@ -681,7 +715,7 @@ def main(): global _reward_tokenizer if is_main: print("\nLoading tokenizer...", flush=True) - tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token _reward_tokenizer = tokenizer @@ -697,14 +731,16 @@ def main(): print(f" Extracted {len(prompts)} prompts", flush=True) orig_count = len(prompts) - prompts = [p for p in prompts if len(tokenizer.encode(p)) <= MAX_PROMPT_LEN] + prompts = [ + r for r in prompts if len(tokenizer.encode(r["prompt"])) <= MAX_PROMPT_LEN + ] if is_main and len(prompts) < orig_count: print( f" Filtered {orig_count - len(prompts)} prompts exceeding {MAX_PROMPT_LEN} tokens", flush=True, ) - dataset = Dataset.from_list([{"prompt": p} for p in prompts]) + dataset = Dataset.from_list(prompts) if is_main: print(f" Dataset size: {len(dataset)}", flush=True) @@ -779,7 +815,7 @@ def main(): training_args = GRPOConfig(**grpo_kwargs) publisher = TauMetricsPublisher(OUTPUT_DIR) if GLOBAL_RANK == 0 else None - metrics_cb = TauGRPOMetricsCallback(publisher) if publisher else None + metrics_cb = TauGRPOMetricsCallback(publisher) max_steps_cb = MaxStepsCallback(int(MAX_STEPS)) if MAX_STEPS else None callbacks = [cb for cb in [metrics_cb, max_steps_cb] if cb is not None] diff --git a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py index 5f8cc1aa..c4540c83 100755 --- a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -7,6 +7,9 @@ Uses vLLM with tensor parallelism across all available GPUs on a node. All benchmark tasks are batched and processed concurrently. + +Before loading the model, main() fails unless every suite in REQUIRED_SUITES +exists in EVAL_DIR with at least one task and `cobc --version` succeeds. """ import json @@ -31,6 +34,50 @@ os.environ["HF_HOME"] = HF_HOME +REQUIRED_SUITES = ( + "perplexity", + "paragraph_completion", + "program_completion", + "structural_validation", + "copybook_usage", +) + + +def check_compiler(): + try: + result = subprocess.run( + ["cobc", "--version"], capture_output=True, text=True, timeout=10 + ) + except FileNotFoundError: + raise RuntimeError("cobc not found; GnuCOBOL must be installed") from None + except subprocess.TimeoutExpired: + raise RuntimeError("cobc --version timed out") from None + if result.returncode != 0: + raise RuntimeError(f"cobc --version exited with code {result.returncode}") + return result.stdout.splitlines()[0] if result.stdout else "" + + +def load_required_suites(eval_dir): + suites = {} + problems = [] + for name in REQUIRED_SUITES: + path = os.path.join(eval_dir, f"{name}.jsonl") + if not os.path.isfile(path): + problems.append(f"{path} is missing") + continue + tasks = load_jsonl(path) + if not tasks: + problems.append(f"{path} has no tasks") + continue + suites[name] = tasks + if problems: + raise RuntimeError( + "Benchmark suites are incomplete; run 03-benchmarks first: " + + "; ".join(problems) + ) + return suites + + def load_jsonl(path): records = [] if not os.path.isfile(path): @@ -50,7 +97,7 @@ def eval_perplexity(tasks, max_length=4096): print(f"\n=== Perplexity Evaluation ({len(tasks)} docs) ===", flush=True) print(" Loading HF model for perplexity (single GPU)...", flush=True) - tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) n_gpus = torch.cuda.device_count() max_mem = {i: "17GiB" for i in range(n_gpus)} model = AutoModelForCausalLM.from_pretrained( @@ -58,7 +105,6 @@ def eval_perplexity(tasks, max_length=4096): torch_dtype=torch.bfloat16, device_map="auto", max_memory=max_mem, - trust_remote_code=True, ) total_nll = 0.0 @@ -305,7 +351,6 @@ def eval_compilation(tasks, generated_texts): print(f"\n=== Compilation Evaluation ({len(tasks)} tasks) ===", flush=True) compiled = 0 failed = 0 - skipped = 0 fmt_counts = {"free": 0, "fixed": 0} results = [] @@ -313,8 +358,7 @@ def eval_compilation(tasks, generated_texts): full_text = task["prompt"] + gen_text success, errors, fmt_used = try_compile(full_text) if success is None: - skipped += 1 - continue + raise RuntimeError(f"cobc not found while compiling task {task['id']}") if success: compiled += 1 fmt_counts[fmt_used] += 1 @@ -337,14 +381,11 @@ def eval_compilation(tasks, generated_texts): f" Format breakdown: {fmt_counts['free']} free, {fmt_counts['fixed']} fixed", flush=True, ) - if skipped: - print(f" Skipped (cobc not found): {skipped}", flush=True) return { "compilation_rate_pct": round(rate, 1), "compiled": compiled, "failed": failed, - "skipped": skipped, "compiled_free": fmt_counts["free"], "compiled_fixed": fmt_counts["fixed"], "details": results, @@ -477,9 +518,14 @@ def main(): print(f" GPUs: {torch.cuda.device_count()}, TP size: {TP_SIZE}", flush=True) print(f" Eval dir: {EVAL_DIR}", flush=True) + print(f" Compiler: {check_compiler()}", flush=True) + suites = load_required_suites(EVAL_DIR) + for name in REQUIRED_SUITES: + print(f" {name}: {len(suites[name])} tasks", flush=True) + all_results = {"model": MODEL_ID, "backend": "vllm", "tp_size": TP_SIZE} - ppl_tasks = load_jsonl(os.path.join(EVAL_DIR, "perplexity.jsonl")) + ppl_tasks = suites["perplexity"] if ppl_tasks: ppl_results = eval_perplexity(ppl_tasks) all_results["perplexity"] = { @@ -494,7 +540,6 @@ def main(): model=MODEL_ID, tensor_parallel_size=TP_SIZE, dtype="bfloat16", - trust_remote_code=True, max_model_len=int(os.environ.get("MAX_MODEL_LEN", "32768")), gpu_memory_utilization=0.90, download_dir=HF_HOME, @@ -502,7 +547,7 @@ def main(): ) print(" vLLM model loaded.", flush=True) - para_tasks = load_jsonl(os.path.join(EVAL_DIR, "paragraph_completion.jsonl")) + para_tasks = suites["paragraph_completion"] if para_tasks: print( f"\n--- Paragraph completion ({len(para_tasks)} tasks, max_tokens={MAX_NEW_TOKENS}) ---", @@ -523,7 +568,7 @@ def main(): for d in para_compile.get("details", []): f.write(json.dumps(d) + "\n") - prog_tasks = load_jsonl(os.path.join(EVAL_DIR, "program_completion.jsonl")) + prog_tasks = suites["program_completion"] if prog_tasks: print( f"\n--- Program completion ({len(prog_tasks)} tasks, max_tokens=1024) ---", @@ -544,7 +589,7 @@ def main(): for d in prog_compile.get("details", []): f.write(json.dumps(d) + "\n") - strv_tasks = load_jsonl(os.path.join(EVAL_DIR, "structural_validation.jsonl")) + strv_tasks = suites["structural_validation"] if strv_tasks: print( f"\n--- Structural validation ({len(strv_tasks)} tasks, max_tokens=1024) ---", @@ -563,6 +608,16 @@ def main(): for d in strv_compile.get("details", []): f.write(json.dumps(d) + "\n") + cpyb_tasks = suites["copybook_usage"] + print( + f"\n--- Copybook usage ({len(cpyb_tasks)} tasks, max_tokens={MAX_NEW_TOKENS}) ---", + flush=True, + ) + cpyb_gen = batch_generate(llm, cpyb_tasks, max_new_tokens=MAX_NEW_TOKENS) + all_results["copybook_usage"] = eval_generation_quality( + cpyb_tasks, cpyb_gen, "copybook_usage" + ) + summary_path = os.path.join(OUT, "eval-summary.json") with open(summary_path, "w") as f: json.dump(all_results, f, indent=2) @@ -606,6 +661,7 @@ def main(): print(f" Column Format: {sv['column_format_valid_pct']}%", flush=True) print(f" PIC Accuracy: {sv['pic_accuracy_pct']}%", flush=True) print(f" PERFORM Ref Accuracy: {sv['perform_ref_accuracy_pct']}%", flush=True) + print(f" Copybook BLEU: {all_results['copybook_usage']['avg_bleu']}", flush=True) if __name__ == "__main__": diff --git a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py new file mode 100644 index 00000000..e6755bca --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py @@ -0,0 +1,102 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Local tests for the eval preflight in run_eval_vllm. + +vllm is replaced with a stub, so this needs CPU torch only: + python -m pytest examples/cobol-cpt-grpo/07-eval/test_run_eval.py -v + +Checks that a missing or empty suite and a missing or broken cobc fail the run +before any model is loaded, and that a cobc that disappears mid-run fails instead +of counting tasks as skipped. +""" + +import importlib.util +import json +import os +import subprocess +import sys +import types +from unittest import mock + +import pytest + +EVAL_DIR = os.path.dirname(os.path.abspath(__file__)) + + +@pytest.fixture +def run_eval(tmp_path, monkeypatch): + monkeypatch.setenv("OUT", str(tmp_path / "out")) + fake_vllm = types.ModuleType("vllm") + fake_vllm.LLM = mock.Mock(side_effect=AssertionError("model loaded")) + fake_vllm.SamplingParams = object + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + spec = importlib.util.spec_from_file_location( + "run_eval_vllm", os.path.join(EVAL_DIR, "run_eval_vllm.py") + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _write_suites(eval_dir, names, empty=()): + eval_dir.mkdir(parents=True, exist_ok=True) + for name in names: + with open(eval_dir / f"{name}.jsonl", "w") as f: + if name not in empty: + f.write(json.dumps({"id": f"{name}-1", "prompt": "X"}) + "\n") + + +def _cobc_ok(args, **kwargs): + return subprocess.CompletedProcess(args, 0, "cobc (GnuCOBOL) 3.1.2\n", "") + + +def test_load_required_suites_returns_every_suite(run_eval, tmp_path): + _write_suites(tmp_path / "eval", run_eval.REQUIRED_SUITES) + suites = run_eval.load_required_suites(str(tmp_path / "eval")) + assert sorted(suites) == sorted(run_eval.REQUIRED_SUITES) + assert "copybook_usage" in run_eval.REQUIRED_SUITES + + +def test_missing_and_empty_suites_fail_together(run_eval, tmp_path): + present = [n for n in run_eval.REQUIRED_SUITES if n != "program_completion"] + _write_suites(tmp_path / "eval", present, empty=("perplexity",)) + with pytest.raises(RuntimeError) as err: + run_eval.load_required_suites(str(tmp_path / "eval")) + assert "program_completion.jsonl is missing" in str(err.value) + assert "perplexity.jsonl has no tasks" in str(err.value) + + +def test_check_compiler_reports_version(run_eval): + with mock.patch.object(run_eval.subprocess, "run", side_effect=_cobc_ok): + assert run_eval.check_compiler() == "cobc (GnuCOBOL) 3.1.2" + + +@pytest.mark.parametrize( + "side_effect", + [ + FileNotFoundError("cobc"), + subprocess.TimeoutExpired("cobc", 10), + lambda args, **kw: subprocess.CompletedProcess(args, 1, "", "boom"), + ], +) +def test_check_compiler_fails_closed(run_eval, side_effect): + with mock.patch.object(run_eval.subprocess, "run", side_effect=side_effect): + with pytest.raises(RuntimeError): + run_eval.check_compiler() + + +def test_main_fails_before_loading_the_model(run_eval, tmp_path, monkeypatch): + monkeypatch.setattr(run_eval, "EVAL_DIR", str(tmp_path / "missing")) + with mock.patch.object(run_eval.subprocess, "run", side_effect=_cobc_ok): + with pytest.raises(RuntimeError, match="run 03-benchmarks first"): + run_eval.main() + run_eval.LLM.assert_not_called() + + +def test_compilation_fails_when_cobc_disappears(run_eval): + with mock.patch.object( + run_eval.subprocess, "run", side_effect=FileNotFoundError("cobc") + ): + with pytest.raises(RuntimeError, match="cobc not found while compiling"): + run_eval.eval_compilation([{"id": "t1", "prompt": "X"}], ["Y"]) diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index a98dbe42..b49062f2 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -18,7 +18,7 @@ folder as the Ray `working_dir`. | [`02-data-prep/`](./02-data-prep/) | Job (CPU) | Extracts, deduplicates, filters, and splits the corpus into `prepared/train.jsonl` and `prepared/val.jsonl`. | | [`03-benchmarks/`](./03-benchmarks/) | Job (CPU) | Builds five eval suites from the validation split into `/data/datasets/cobol/eval`. | | [`04-cpt/`](./04-cpt/) | RayJob, 2 x 8 GPUs | Continued pretraining with DeepSpeed ZeRO-3 under Ray Train. | -| [`05-rl-prompts/`](./05-rl-prompts/) | Job (CPU) | Masks paragraph bodies of programs that compile with `cobc` to produce `prepared/rl-prompts.jsonl`. | +| [`05-rl-prompts/`](./05-rl-prompts/) | Job (CPU) | Masks paragraph bodies of programs that compile with `cobc` to produce `prepared/rl-prompts.jsonl`. Programs that share any 8-line window with the validation split are excluded. | | [`06-grpo/`](./06-grpo/) | RayJob, 2 x 8 GPUs | GRPO from the CPT checkpoint with a compiler reward and an in-cluster vLLM generation server. | | [`07-eval/`](./07-eval/) | RayJob, 1 x 8 GPUs | Runs the benchmarks with vLLM (TP=8) and compiles every generated program with GnuCOBOL. | @@ -61,7 +61,8 @@ kubectl create secret generic hf-token --from-literal=token= `01-collect/tau-gated.yaml` maps that Secret to `HF_TOKEN` with `env_secret: {HF_TOKEN: hf-token:token}`. `collect-rl-data.sh` skips its gated -steps when `HF_TOKEN` is unset. To include them, add the same `env_secret` +steps when `HF_TOKEN` is unset. Both collect scripts exit nonzero if any step +that runs fails, including a clone, or if no COBOL files are collected. To include them, add the same `env_secret` block to `01-collect/tau-rl.yaml`. `02-data-prep` picks up the gated outputs when they are present under `raw/`. @@ -87,27 +88,41 @@ two GPU workers: `VLLM_SERVER_URL` and calls `train_grpo.main()`. TRL syncs updated weights from trainer rank 0 to the server over an NCCL group -on TRL's group port (51216 in TRL 1.13.0). The reward has two parts: +on TRL's group port (51216 in TRL 1.13.0). Stage 05 records whether each prompt +compiles as free or fixed format; GRPO keeps that `compile_fmt` and uses the +matching `cobc` flag for both the prompt preflight and the reward. The reward +has two parts: - **Compile reward.** Each prompt plus completion is checked with - `cobc -fsyntax-only -free`. A clean compile scores 1.0 and a compile with - warnings scores 0.7. A `cobc` signal crash counts as a failed compile; a - timeout or a missing `cobc` raises. + `cobc -fsyntax-only` and `-free` or `-fixed`. A clean compile scores 1.0 and + a compile with warnings scores 0.7. A `cobc` signal crash counts as a failed + compile; a timeout or a missing `cobc` raises. - **Structural reward.** Penalizes degenerate completions (display-only bodies, no-op majorities, repeated statements, truncation). It adds a bonus when every `PERFORM` target is a paragraph defined in the completion. The defaults are K=4 samples per prompt at temperature 0.8, LR 5e-7, KL -coefficient 0.05, and 3 epochs. `06-grpo/test_reward.py` covers the reward -without a GPU. +coefficient 0.05, and 3 epochs. Every trainer rank checks the KL and +reward-hacking alarms, and the ranks combine their stop decisions so that all +of them stop at the same step. Only rank 0 publishes metrics. +`06-grpo/test_reward.py` covers the reward without a GPU. **Eval (`07-eval`).** One Ray task reserves a whole 8-GPU worker and runs -`run_eval_vllm.main()` with greedy decoding (512 new tokens). The five suites -are paragraph completion, copybook usage, program completion, perplexity, and -structural validation. Compile scoring is format-aware: each generation is -compiled with `cobc -free`, retried with `-fixed` if that fails, and both -counts are reported. The headline compile score is out of 220 tasks: 70 -paragraph, 50 program, and 100 structural. +`run_eval_vllm.main()` with greedy decoding (512 new tokens). It scores five +suites: + +- Perplexity on held-out programs. +- Paragraph completion and program completion: BLEU, edit distance, and + compilation. +- Structural validation: structure checks and compilation. +- Copybook usage: BLEU and edit distance only. + +Compile scoring is format-aware: each generation is compiled with `cobc -free`, +retried with `-fixed` if that fails, and both counts are reported. The headline +compile score is out of 220 tasks: 70 paragraph, 50 program, and 100 +structural. Before loading the model, the eval checks that `cobc --version` +works and that all five suite files exist and are not empty. It fails if +either check fails, and it fails if `cobc` disappears during the run. ## Hardware @@ -198,21 +213,41 @@ For a short end-to-end check of the GPU stages, set `MAX_STEPS` in `04-cpt/tau.yaml` or `06-grpo/tau.yaml` (for example `MAX_STEPS: "5"`). A short CPT run produces a model that is not useful for GRPO. -The driver tests run on a CPU-only machine with a local Ray install: +The tests run on a CPU-only machine with CPU `torch` and a local Ray install: ```bash +cd examples/cobol-cpt-grpo/03-benchmarks && python -m pytest test_build_benchmarks.py +cd examples/cobol-cpt-grpo/05-rl-prompts && python -m pytest test_prepare_rl_prompts.py cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_reward.py +cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_stop_sync.py cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_driver.py cd examples/cobol-cpt-grpo/04-cpt && python -m pytest test_continued_pretrain.py test_driver.py -cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_driver.py +cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_run_eval.py test_driver.py ``` Run `test_reward.py` in its own process; it replaces `torch` and `trl` with -stubs. `06-grpo/test_driver.py` starts a local Ray cluster and takes one to -five minutes. +stubs. `test_stop_sync.py` starts two CPU processes in a gloo group. +`06-grpo/test_driver.py` starts a local Ray cluster and takes one to five +minutes. ## Results +The numbers in this section come from an earlier version of the pipeline that +had two defects, both fixed in the code here: + +- Stage 05 did not exclude the validation split, so GRPO prompts included + validation programs. For up to 18/50 program-completion, 5/70 paragraph, and + 19/100 structural tasks, the reference answer appears in a GRPO training + prompt. The reward only checks compilation, so GRPO never trained on those + answers. CPT is not affected. +- GRPO compiled every prompt with `cobc -free`, so it dropped the 3,676 of + 6,852 prompts (54%) that compile only in fixed format. + +The GRPO results and the GRPO gains below are therefore partly contaminated, +and they come from half of the intended prompts. The CPT columns are not +affected. Copybook usage was not scored at the time. GRPO and eval have not +yet been rerun with the fixes. + The GRPO model was evaluated with `07-eval` (greedy decoding, 512 new tokens, TP=8 on 8 H200): diff --git a/examples/cobol-cpt-grpo/image-requirements.txt b/examples/cobol-cpt-grpo/image-requirements.txt index 25b34780..44fea346 100644 --- a/examples/cobol-cpt-grpo/image-requirements.txt +++ b/examples/cobol-cpt-grpo/image-requirements.txt @@ -28,5 +28,5 @@ py-spy==0.4.2 python-discovery==1.6.1 smart_open==8.0.2 tensorboardX==2.6.5 -virtualenv==21.14.6 +virtualenv==21.14.7 wrapt==2.5.0 From 230b54ef0fc679a199b7b9330e605f6f3c5570a0 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 13:05:35 -0700 Subject: [PATCH 04/15] fix(examples): let the COBOL collect scripts take DATA_ROOT from the environment collect-gated.py already reads DATA_ROOT; the two shell scripts hardcoded it. Both now default to /data/datasets/cobol, derive RAW from DATA_ROOT, and export it for the embedded Python steps. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh | 4 ++-- examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh | 4 ++-- examples/cobol-cpt-grpo/01-collect/tau-rl.yaml | 2 +- examples/cobol-cpt-grpo/01-collect/tau.yaml | 2 +- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh index 82a9d93a..0b8ef0ed 100755 --- a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh +++ b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh @@ -4,8 +4,8 @@ set -uo pipefail -DATA_ROOT="/data/datasets/cobol" -RAW="${DATA_ROOT}/raw" +DATA_ROOT="${DATA_ROOT:-/data/datasets/cobol}" +export RAW="${DATA_ROOT}/raw" PROCESSED="${DATA_ROOT}/processed" LOGS="${DATA_ROOT}/logs" mkdir -p "${RAW}" "${PROCESSED}" "${LOGS}" diff --git a/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh index 77138830..c6ef178f 100755 --- a/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh +++ b/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh @@ -4,8 +4,8 @@ set -uo pipefail -DATA_ROOT="/data/datasets/cobol" -RAW="${DATA_ROOT}/raw" +DATA_ROOT="${DATA_ROOT:-/data/datasets/cobol}" +export RAW="${DATA_ROOT}/raw" LOGS="${DATA_ROOT}/logs" mkdir -p "${RAW}" "${LOGS}" diff --git a/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml b/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml index b7f95cb7..472b8d90 100644 --- a/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml +++ b/examples/cobol-cpt-grpo/01-collect/tau-rl.yaml @@ -10,7 +10,7 @@ experiment: runtime: image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a env: - RAW: /data/datasets/cobol/raw + DATA_ROOT: /data/datasets/cobol HF_HOME: /tmp/hf-cache SOURCE_LICENSES: permissive diff --git a/examples/cobol-cpt-grpo/01-collect/tau.yaml b/examples/cobol-cpt-grpo/01-collect/tau.yaml index 2d75bf5d..17a252b9 100644 --- a/examples/cobol-cpt-grpo/01-collect/tau.yaml +++ b/examples/cobol-cpt-grpo/01-collect/tau.yaml @@ -10,7 +10,7 @@ experiment: runtime: image: python:3.11.17-bookworm@sha256:c255868cf24d6ca01031d262fa952f59b108b96776625f3d7f1af85d5f28e26a env: - RAW: /data/datasets/cobol/raw + DATA_ROOT: /data/datasets/cobol HF_HOME: /tmp/hf-cache SOURCE_LICENSES: permissive From 037e2585d5e896b826818bea86e01bf5b3f94d79 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 13:17:58 -0700 Subject: [PATCH 05/15] fix(examples): fetch X-COBOL and NIST from current pinned URLs Zenodo removed the files-archive endpoint for the X-COBOL record and the SourceForge nist_2013-10-10.tar.gz file no longer exists, so both required collection steps failed. Download COBOL_Files.zip and newcob.val.tar.gz directly, verify their sha256 digests, and keep the existing extracted layout. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../01-collect/collect-cobol-data.sh | 23 ++++++++----------- 1 file changed, 9 insertions(+), 14 deletions(-) diff --git a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh index 0b8ef0ed..29b004c8 100755 --- a/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh +++ b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh @@ -141,10 +141,10 @@ PYEOF run_step "Evaluation benchmarks" collect_eval_benchmarks collect_xcobol() { - mkdir -p "${RAW}/x-cobol" - curl -fsSL "https://zenodo.org/api/records/14269462/files-archive" -o "${RAW}/x-cobol/x-cobol-v2.zip" || return 1 - (cd "${RAW}/x-cobol" && unzip -qo x-cobol-v2.zip -d extracted) || return 1 - find "${RAW}/x-cobol/extracted" -name '*.zip' -exec sh -c 'unzip -qo "$1" -d "$(dirname "$1")/$(basename "$1" .zip)"' _ {} \; + mkdir -p "${RAW}/x-cobol/extracted" + curl -fsSL "https://zenodo.org/api/records/14269462/files/COBOL_Files.zip/content" -o "${RAW}/x-cobol/extracted/COBOL_Files.zip" || return 1 + echo "9b97b2b77f025a11895730580965bc79ab3d715f9922ab79c0a67e30068c6e6d ${RAW}/x-cobol/extracted/COBOL_Files.zip" | sha256sum -c - || return 1 + (cd "${RAW}/x-cobol/extracted" && unzip -qo COBOL_Files.zip -d COBOL_Files) || return 1 XCOUNT=$(find ${RAW}/x-cobol/extracted -type f \( -iname '*.cbl' -o -iname '*.cob' -o -iname '*.cobol' -o -iname '*.cpy' \) 2>/dev/null | wc -l) log "X-COBOL: ${XCOUNT} COBOL files" if [ "${XCOUNT}" -eq 0 ]; then @@ -204,17 +204,12 @@ log "GitHub repos: ${TOTAL_GH} COBOL files total" collect_nist() { mkdir -p "${RAW}/nist" if [ ! -f "${RAW}/nist/newcob.val" ]; then - curl -L --max-redirs 5 "https://sourceforge.net/projects/gnucobol/files/nist/nist_2013-10-10.tar.gz/download" \ - -o "${RAW}/nist/nist.tar.gz" - if file "${RAW}/nist/nist.tar.gz" | grep -q "gzip"; then - cd "${RAW}/nist" && tar xzf nist.tar.gz - else - log "NIST: downloaded file is not a valid gzip (likely redirect page), trying alternate URL" - curl -L "https://master.dl.sourceforge.net/project/gnucobol/nist/nist_2013-10-10.tar.gz" \ - -o "${RAW}/nist/nist.tar.gz" - cd "${RAW}/nist" && tar xzf nist.tar.gz - fi + curl -fsSL "https://downloads.sourceforge.net/project/gnucobol/nist/newcob.val.tar.gz" -o "${RAW}/nist/newcob.val.tar.gz" || return 1 + echo "e4513f26a9b38911f7bf882fe3d3339a80b45cabc2caf85eb3055b0f5ce87ee0 ${RAW}/nist/newcob.val.tar.gz" | sha256sum -c - || return 1 + tar xzf "${RAW}/nist/newcob.val.tar.gz" -C "${RAW}/nist" || return 1 fi + [ -s "${RAW}/nist/newcob.val" ] || return 1 + log "NIST: newcob.val $(wc -l < "${RAW}/nist/newcob.val") lines" } run_step "NIST COBOL85 test suite" collect_nist From 067765852f1d06df560fa7b079e9eee35fbf0be6 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 14:04:42 -0700 Subject: [PATCH 06/15] fix(examples): exclude RL prompts whose stripped text overlaps validation The validation exclusion checked raw source text, but prompts are built after EXEC blocks and COPY statements are stripped. Removing those lines can join shared lines into a window that matches the validation split, which the final disjointness check then rejects. Check the stripped text as well. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../05-rl-prompts/prepare_rl_prompts.py | 12 +++++++----- .../05-rl-prompts/test_prepare_rl_prompts.py | 15 +++++++++++++++ 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py index cab02852..95921196 100755 --- a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py +++ b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py @@ -147,6 +147,7 @@ def exclude_val_overlap(texts_with_paths, val_windows, size): (path, text) for path, text in texts_with_paths if not overlaps_val(text, val_windows, size) + and not overlaps_val(prepare_text(text), val_windows, size) ] return kept, len(texts_with_paths) - len(kept) @@ -182,6 +183,11 @@ def strip_copy_stmts(text): return COPY_STMT_RE.sub("", text) +def prepare_text(raw_text): + text = raw_text.replace("\r\n", "\n").replace("\r", "\n").rstrip() + return strip_copy_stmts(strip_exec_blocks(text)) + + def try_compile_free(cobol_text): with tempfile.NamedTemporaryFile(suffix=".cbl", mode="w", delete=False) as f: f.write(cobol_text) @@ -279,11 +285,7 @@ def process_one(path_and_text): path, raw_text = path_and_text result = {"path": path, "status": None} - text = raw_text.replace("\r\n", "\n").replace("\r", "\n") - text = text.rstrip() - - text_stripped = strip_exec_blocks(text) - text_stripped = strip_copy_stmts(text_stripped) + text_stripped = prepare_text(raw_text) prompt, body, body_lines = extract_prompt(text_stripped) if prompt is None: diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py index 74ec1507..21e5cc98 100644 --- a/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py +++ b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py @@ -77,6 +77,21 @@ def test_exclude_val_overlap_drops_programs_sharing_a_window(prep, tmp_path): assert [p for p, _ in kept] == ["short.cbl", "clean.cbl"] +def test_exclude_val_overlap_checks_text_after_copy_stripping(prep, tmp_path): + val_path = tmp_path / "prepared" / "val.jsonl" + _write_val(val_path, [_program("VALPROG", SHARED)]) + windows, _ = prep.load_val_windows(str(val_path), 8) + + split = _program( + "SPLIT", OTHER[:2] + SHARED[:4] + [" COPY FOO."] + SHARED[4:8] + ) + assert not prep.overlaps_val(split, windows, 8) + assert prep.overlaps_val(prep.prepare_text(split), windows, 8) + kept, removed = prep.exclude_val_overlap([("split.cbl", split)], windows, 8) + assert removed == 1 + assert kept == [] + + def test_assert_disjoint_raises_on_overlap(prep, tmp_path): val_path = tmp_path / "prepared" / "val.jsonl" _write_val(val_path, [_program("VALPROG", SHARED)]) From 8e9926dfdf7b7300967fb83c712fce9d15dc85a0 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 14:12:36 -0700 Subject: [PATCH 07/15] fix(examples): make the RL prompt compile filter fail closed GnuCOBOL 3.1.2 in the training image rejects -fmissing-period=ok with an "unrecognized option" message that has no error: line, and the copybook fallback counted that as a copybook-only failure, so every program passed without being compiled. Drop the unsupported flag, require at least one missing-copybook error before treating a failure as copybook-only, and check that cobc compiles a known-good program in both formats before processing. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../05-rl-prompts/prepare_rl_prompts.py | 36 ++++++++++++------- .../05-rl-prompts/test_prepare_rl_prompts.py | 30 ++++++++++++++++ 2 files changed, 54 insertions(+), 12 deletions(-) diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py index 95921196..4f6983b1 100755 --- a/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py +++ b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py @@ -61,6 +61,14 @@ COPY_STMT_RE = re.compile(r"^\s*COPY\s+\S+.*?\.\s*$", re.IGNORECASE | re.MULTILINE) +PREFLIGHT_PROGRAM = ( + " IDENTIFICATION DIVISION.\n" + " PROGRAM-ID. PREFLIGHT.\n" + " PROCEDURE DIVISION.\n" + " DISPLAY 'OK'.\n" + " STOP RUN.\n" +) + WHITESPACE_RE = re.compile(r"\s+") stats = Counter() @@ -202,7 +210,6 @@ def try_compile_free(cobol_text): "-frelax-syntax-checks", "-frelax-level-hierarchy", "-fmissing-statement=ok", - "-fmissing-period=ok", "-w", src, ], @@ -236,7 +243,6 @@ def try_compile_fixed(cobol_text): "-frelax-syntax-checks", "-frelax-level-hierarchy", "-fmissing-statement=ok", - "-fmissing-period=ok", "-w", src, ], @@ -257,14 +263,17 @@ def try_compile_fixed(cobol_text): def is_copybook_only_error(stderr): - if not stderr.strip(): - return False - for line in stderr.splitlines(): - low = line.lower() - if "error:" in low: - if "copy" not in low or "not found" not in low: - return False - return True + errors = [ln.lower() for ln in stderr.splitlines() if "error:" in ln.lower()] + return bool(errors) and all("copy" in ln and "not found" in ln for ln in errors) + + +def check_compiler(): + for name, compile_fn in [("free", try_compile_free), ("fixed", try_compile_fixed)]: + rc, stderr = compile_fn(PREFLIGHT_PROGRAM) + if rc != 0: + raise RuntimeError( + f"cobc {name}-format preflight failed (rc={rc}): {stderr.strip()[:500]}" + ) def compile_check(text): @@ -320,8 +329,11 @@ def main(): log(f"RAW={RAW}, OUT={OUT}, workers={MAX_WORKERS}") log(f"Body lines: {MIN_BODY_LINES}-{MAX_BODY_LINES}") - subprocess.run(["cobc", "--version"], capture_output=True, check=True) - log("cobc available") + version = subprocess.run( + ["cobc", "--version"], capture_output=True, text=True, check=True + ).stdout.splitlines()[0] + check_compiler() + log(f"{version} compiles the preflight program in free and fixed format") val_windows, val_count = load_val_windows(VAL_PATH, VAL_OVERLAP_LINES) log(f"Loaded {val_count} validation programs from {VAL_PATH}") diff --git a/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py index 21e5cc98..5b5ae572 100644 --- a/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py +++ b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py @@ -10,6 +10,8 @@ import importlib.util import json import os +import subprocess +from unittest import mock import pytest @@ -114,3 +116,31 @@ def test_load_val_windows_fails_without_validation_programs(prep, tmp_path, cont val_path.write_text(contents) with pytest.raises(RuntimeError): prep.load_val_windows(str(val_path), 8) + + +@pytest.mark.parametrize( + "stderr, expected", + [ + ("", False), + ("cobc: unrecognized option '-fmissing-period=ok'", False), + ("x.cbl:3: error: COPY book 'FOO' not found", True), + ( + "x.cbl:3: error: COPY book 'FOO' not found\nx.cbl:9: error: 'X' is not defined", + False, + ), + ], +) +def test_copybook_only_needs_only_missing_copybook_errors(prep, stderr, expected): + assert prep.is_copybook_only_error(stderr) is expected + + +def test_check_compiler_fails_on_rejected_flags(prep): + rejected = subprocess.CompletedProcess( + [], 1, b"", b"cobc: unrecognized option '-fmissing-period=ok'" + ) + with mock.patch.object(prep.subprocess, "run", return_value=rejected): + with pytest.raises(RuntimeError, match="free-format preflight failed"): + prep.check_compiler() + ok = subprocess.CompletedProcess([], 0, b"", b"") + with mock.patch.object(prep.subprocess, "run", return_value=ok): + prep.check_compiler() From cdb4c8b2647494c1fbc454909ae54a9c8c6fbec8 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 22:23:40 -0700 Subject: [PATCH 08/15] fix(examples): set aside eval prompts that exceed the model context Stage 03 keeps whole programs as prompt context, and a few programs in a larger corpus produce prompts longer than the 32768-token context. vLLM rejects the batch and the eval fails. Tokenize each generation prompt before loading the model, set aside those that cannot fit with their output budget, record them in eval-summary.json, and fail if a suite is left empty. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../cobol-cpt-grpo/07-eval/run_eval_vllm.py | 55 ++++++++++++++-- .../cobol-cpt-grpo/07-eval/test_run_eval.py | 65 ++++++++++++++++++- examples/cobol-cpt-grpo/README.md | 6 +- 3 files changed, 119 insertions(+), 7 deletions(-) diff --git a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py index c4540c83..1e0e460b 100755 --- a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -27,6 +27,7 @@ MODEL_ID = os.environ.get("MODEL_ID", "Qwen/Qwen2.5-Coder-32B") OUT = os.environ.get("OUT", "/data/datasets/cobol/eval/results/baseline") MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "512")) +MAX_MODEL_LEN = int(os.environ.get("MAX_MODEL_LEN", "32768")) TP_SIZE = int(os.environ.get("TP_SIZE", str(torch.cuda.device_count()))) HF_HOME = os.environ.get("HF_HOME", "/tmp/hf-cache") @@ -42,6 +43,39 @@ "copybook_usage", ) +GENERATION_BUDGETS = { + "paragraph_completion": MAX_NEW_TOKENS, + "program_completion": 1024, + "structural_validation": 1024, + "copybook_usage": MAX_NEW_TOKENS, +} + + +def fit_to_context(suites, tokenizer, max_model_len): + """Set aside generation tasks whose prompt plus output budget exceeds the context.""" + excluded = {} + emptied = [] + for name, budget in GENERATION_BUDGETS.items(): + limit = max_model_len - budget + kept = [] + too_long = [] + for task in suites[name]: + n_tokens = len(tokenizer(task["prompt"])["input_ids"]) + if n_tokens > limit: + too_long.append({"id": task["id"], "prompt_tokens": n_tokens}) + else: + kept.append(task) + if too_long: + excluded[name] = {"prompt_token_limit": limit, "tasks": too_long} + if not kept: + emptied.append(name) + suites[name] = kept + if emptied: + raise RuntimeError( + f"No task fits in max_model_len={max_model_len} for: {', '.join(emptied)}" + ) + return excluded + def check_compiler(): try: @@ -520,10 +554,23 @@ def main(): print(f" Compiler: {check_compiler()}", flush=True) suites = load_required_suites(EVAL_DIR) - for name in REQUIRED_SUITES: - print(f" {name}: {len(suites[name])} tasks", flush=True) + from transformers import AutoTokenizer - all_results = {"model": MODEL_ID, "backend": "vllm", "tp_size": TP_SIZE} + excluded = fit_to_context( + suites, AutoTokenizer.from_pretrained(MODEL_ID), MAX_MODEL_LEN + ) + for name in REQUIRED_SUITES: + dropped = len(excluded.get(name, {}).get("tasks", [])) + note = f" ({dropped} over the context limit set aside)" if dropped else "" + print(f" {name}: {len(suites[name])} tasks{note}", flush=True) + + all_results = { + "model": MODEL_ID, + "backend": "vllm", + "tp_size": TP_SIZE, + "max_model_len": MAX_MODEL_LEN, + "excluded_over_context": excluded, + } ppl_tasks = suites["perplexity"] if ppl_tasks: @@ -540,7 +587,7 @@ def main(): model=MODEL_ID, tensor_parallel_size=TP_SIZE, dtype="bfloat16", - max_model_len=int(os.environ.get("MAX_MODEL_LEN", "32768")), + max_model_len=MAX_MODEL_LEN, gpu_memory_utilization=0.90, download_dir=HF_HOME, distributed_executor_backend="mp", diff --git a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py index e6755bca..baac1983 100644 --- a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py +++ b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py @@ -7,8 +7,9 @@ python -m pytest examples/cobol-cpt-grpo/07-eval/test_run_eval.py -v Checks that a missing or empty suite and a missing or broken cobc fail the run -before any model is loaded, and that a cobc that disappears mid-run fails instead -of counting tasks as skipped. +before any model is loaded, that a cobc that disappears mid-run fails instead +of counting tasks as skipped, and that prompts too long for the model context are +set aside and recorded, failing the run if a suite is left empty. """ import importlib.util @@ -100,3 +101,63 @@ def test_compilation_fails_when_cobc_disappears(run_eval): ): with pytest.raises(RuntimeError, match="cobc not found while compiling"): run_eval.eval_compilation([{"id": "t1", "prompt": "X"}], ["Y"]) + + +def _fake_tokenizer(text): + return {"input_ids": text.split()} + + +def _suites_with_prompts(run_eval, lengths): + return { + name: [ + {"id": f"{name}-{n}", "prompt": " ".join(["w"] * n)} + for n in lengths.get(name, [10]) + ] + for name in run_eval.REQUIRED_SUITES + } + + +def test_fit_to_context_sets_aside_prompts_over_each_budget(run_eval): + budgets = run_eval.GENERATION_BUDGETS + max_len = 2048 + para_limit = max_len - budgets["paragraph_completion"] + prog_limit = max_len - budgets["program_completion"] + suites = _suites_with_prompts( + run_eval, + { + "paragraph_completion": [para_limit, para_limit + 1], + "program_completion": [prog_limit + 1, 10], + }, + ) + excluded = run_eval.fit_to_context(suites, _fake_tokenizer, max_len) + assert [t["id"] for t in suites["paragraph_completion"]] == [ + f"paragraph_completion-{para_limit}" + ] + assert [t["id"] for t in suites["program_completion"]] == ["program_completion-10"] + assert excluded == { + "paragraph_completion": { + "prompt_token_limit": para_limit, + "tasks": [ + { + "id": f"paragraph_completion-{para_limit + 1}", + "prompt_tokens": para_limit + 1, + } + ], + }, + "program_completion": { + "prompt_token_limit": prog_limit, + "tasks": [ + { + "id": f"program_completion-{prog_limit + 1}", + "prompt_tokens": prog_limit + 1, + } + ], + }, + } + assert len(suites["perplexity"]) == 1 + + +def test_fit_to_context_fails_when_a_suite_is_emptied(run_eval): + suites = _suites_with_prompts(run_eval, {"copybook_usage": [5000]}) + with pytest.raises(RuntimeError, match="copybook_usage"): + run_eval.fit_to_context(suites, _fake_tokenizer, 2048) diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index b49062f2..aa0b8276 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -122,7 +122,11 @@ retried with `-fixed` if that fails, and both counts are reported. The headline compile score is out of 220 tasks: 70 paragraph, 50 program, and 100 structural. Before loading the model, the eval checks that `cobc --version` works and that all five suite files exist and are not empty. It fails if -either check fails, and it fails if `cobc` disappears during the run. +either check fails, and it fails if `cobc` disappears during the run. It also +tokenizes every generation prompt and sets aside any prompt that cannot fit in +`MAX_MODEL_LEN` (32768) together with its output budget. Set-aside task IDs are +listed under `excluded_over_context` in `eval-summary.json`, and the run fails +if a suite has no task left. ## Hardware From 7fecbf066262f5f85a43d6ec65471ab4cc21c601 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 22:37:40 -0700 Subject: [PATCH 09/15] fix(examples): bump virtualenv to 21.14.9 for GHSA-c978-59m8-c77v virtualenv 21.14.7 is affected by GHSA-c978-59m8-c77v (high), which fails the dependency-review check. 21.14.9 is the first patched release and has the same dependency requirements. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/cobol-cpt-grpo/image-requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/cobol-cpt-grpo/image-requirements.txt b/examples/cobol-cpt-grpo/image-requirements.txt index 44fea346..3e4f6825 100644 --- a/examples/cobol-cpt-grpo/image-requirements.txt +++ b/examples/cobol-cpt-grpo/image-requirements.txt @@ -28,5 +28,5 @@ py-spy==0.4.2 python-discovery==1.6.1 smart_open==8.0.2 tensorboardX==2.6.5 -virtualenv==21.14.7 +virtualenv==21.14.9 wrapt==2.5.0 From 38a64a6c9cda14ee100b345d13b6b7f506f6b3f8 Mon Sep 17 00:00:00 2001 From: June Liu Date: Thu, 22 Oct 2026 22:58:50 -0700 Subject: [PATCH 10/15] docs(examples): report base, CPT, and GRPO results from the fixed pipeline Replace the results from the earlier, contaminated run with a base, CPT, and GRPO comparison scored on the same benchmarks built from the permissive corpus. Describe the compile score total as corpus dependent. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/cobol-cpt-grpo/README.md | 92 +++++++++++++------------------ 1 file changed, 38 insertions(+), 54 deletions(-) diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index aa0b8276..140fc597 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -119,8 +119,9 @@ suites: Compile scoring is format-aware: each generation is compiled with `cobc -free`, retried with `-fixed` if that fails, and both counts are reported. The headline -compile score is out of 220 tasks: 70 paragraph, 50 program, and 100 -structural. Before loading the model, the eval checks that `cobc --version` +compile score sums the paragraph, program, and structural tasks. Stage 03 caps +these at 200, 50, and 100, so the total depends on the corpus. Before loading +the model, the eval checks that `cobc --version` works and that all five suite files exist and are not empty. It fails if either check fails, and it fails if `cobc` disappears during the run. It also tokenizes every generation prompt and sets aside any prompt that cannot fit in @@ -236,56 +237,39 @@ minutes. ## Results -The numbers in this section come from an earlier version of the pipeline that -had two defects, both fixed in the code here: - -- Stage 05 did not exclude the validation split, so GRPO prompts included - validation programs. For up to 18/50 program-completion, 5/70 paragraph, and - 19/100 structural tasks, the reference answer appears in a GRPO training - prompt. The reward only checks compilation, so GRPO never trained on those - answers. CPT is not affected. -- GRPO compiled every prompt with `cobc -free`, so it dropped the 3,676 of - 6,852 prompts (54%) that compile only in fixed format. - -The GRPO results and the GRPO gains below are therefore partly contaminated, -and they come from half of the intended prompts. The CPT columns are not -affected. Copybook usage was not scored at the time. GRPO and eval have not -yet been rerun with the fixes. - -The GRPO model was evaluated with `07-eval` (greedy decoding, 512 new tokens, -TP=8 on 8 H200): - -| Metric | GRPO model | -| --- | --- | -| Compile, overall | **64/220** | -| Compile, paragraph completion | 18/70 (7 free + 11 fixed) | -| Compile, program completion | 18/50 | -| Compile, structural validation | 28/100 (16 free + 12 fixed) | -| COBOL perplexity | 1.49 | -| BLEU, paragraph / program | 0.155 / 0.190 | -| PERFORM reference accuracy | 41.8% | -| Column formatting | 72.0% | -| PIC clause accuracy | 98.9% | - -The checkpoint was trained with an earlier Job-based launch of the same CPT -and GRPO code. Running the `07-eval` RayJob on that checkpoint reproduced -every number above exactly. - -A full rerun with every stage launched from this example scored lower. The -Job-based CPT column is the checkpoint that the GRPO model above started from. - - -| Metric | Job-based CPT | RayJob CPT | RayJob GRPO | +All three models below were scored by one `07-eval` RayJob each (greedy +decoding, 512 new tokens, `MAX_MODEL_LEN` 32768, TP=8 on 8 H200) on the same +benchmarks. The pipeline ran end to end from this example with +`SOURCE_LICENSES=permissive` and no gated datasets: + +- CPT: 60 steps (3 epochs) from `Qwen/Qwen2.5-Coder-32B`. Eval loss was + 0.523 at step 25 and 0.507 at step 60. +- GRPO: 162 steps (3 epochs) from the CPT checkpoint. The compile preflight + kept 1,762 prompts, 918 free format and 844 fixed format. +- Benchmarks: 62 paragraph, 50 program, 100 structural, and 78 copybook tasks, + plus 158 held-out programs for perplexity. Three prompts did not fit in the + context and were set aside for every model: one paragraph task (51,129 + tokens) and two copybook tasks (55,731 and 59,490 tokens). The compile score + is out of the remaining 211 tasks. + +| Metric | Base | CPT | GRPO | | --- | --- | --- | --- | -| Compile, overall | 58/220 | 53/220 | 57/220 | -| Compile, paragraph completion | 16/70 | 14/70 | 14/70 | -| Compile, program completion | 16/50 | 18/50 | 18/50 | -| Compile, structural validation | 26/100 | 21/100 | 25/100 | -| COBOL perplexity | 1.49 | 1.48 | 1.48 | -| PERFORM reference accuracy | 40.4% | 27.3% | 28.1% | - -All three columns were scored with the same `07-eval` RayJob. GRPO adds a -similar amount on either CPT checkpoint (+4 on the RayJob CPT, +6 on the -Job-based CPT), so the gap comes from CPT. The cause is not yet known. The -training image moved to newer `transformers` and DeepSpeed releases than the -Job-based run used. +| Compile, overall | 50/211 | **89/211** | 88/211 | +| Compile, paragraph completion | 16/61 (6 free + 10 fixed) | 24/61 (9 + 15) | 25/61 (9 + 16) | +| Compile, program completion | 13/50 (8 + 5) | 20/50 (12 + 8) | 20/50 (11 + 9) | +| Compile, structural validation | 21/100 (11 + 10) | 45/100 (23 + 22) | 43/100 (22 + 21) | +| COBOL perplexity | 1.90 | 1.54 | 1.54 | +| BLEU, paragraph / program | 0.130 / 0.056 | 0.176 / 0.096 | 0.180 / 0.099 | +| Edit distance, paragraph / program | 0.677 / 0.661 | 0.632 / 0.632 | 0.632 / 0.623 | +| PERFORM reference accuracy | 15.9% | 40.7% | 34.6% | +| Column formatting | 74.0% | 80.0% | 78.0% | +| PIC clause accuracy | 98.5% | 98.5% | 98.5% | +| Division structure | 100% | 100% | 100% | +| BLEU / edit distance, copybook usage | 0.013 / 0.883 | 0.035 / 0.843 | 0.038 / 0.849 | + +CPT accounts for nearly all of the improvement over the base model: compiles +rise from 50 to 89, perplexity falls from 1.90 to 1.54, and PERFORM reference +accuracy more than doubles. GRPO does not improve on CPT on these benchmarks. +Its compile count and BLEU are within one task or a few thousandths of CPT, +and its PERFORM reference accuracy is 6.1 points lower even though the +structural reward includes a PERFORM bonus. From 85bd4bf1c2eb245c709a581e1204a2a412441a29 Mon Sep 17 00:00:00 2001 From: June Liu Date: Fri, 23 Oct 2026 09:40:00 -0700 Subject: [PATCH 11/15] fix(examples): compile paragraph completions with the rest of the procedure Stage 03 dropped every paragraph after the masked one, so a completion that performs a later paragraph could not compile even when it matched the reference. Keep the remaining procedure text as a suffix and append it when compiling. Cut paragraph completions at the next paragraph header, and report how many tasks compile with the reference in place of the generation. Decode compiler output with replacement so non-UTF-8 bytes in cobc errors do not abort the eval, require the current benchmark fields before loading the model, and run the stage 07 test files separately. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../03-benchmarks/build_benchmarks.py | 16 ++- .../03-benchmarks/test_build_benchmarks.py | 60 ++++++++- .../cobol-cpt-grpo/07-eval/run_eval_vllm.py | 73 ++++++++-- .../cobol-cpt-grpo/07-eval/test_run_eval.py | 127 +++++++++++++++++- examples/cobol-cpt-grpo/README.md | 39 ++++-- 5 files changed, 284 insertions(+), 31 deletions(-) diff --git a/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py index 40e92ae8..60955b9b 100755 --- a/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py +++ b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py @@ -85,7 +85,13 @@ def extract_paragraphs(procedure_text): body = procedure_text[start:end].rstrip() if body.strip(): paragraphs.append( - {"name": name, "body": body, "start": match.start(), "end": end} + { + "name": name, + "body": body, + "start": match.start(), + "body_end": start + len(body), + "end": end, + } ) return paragraphs @@ -95,6 +101,7 @@ def build_paragraph_completion(records, max_tasks=200): Prompt: everything before the paragraph body (divisions + paragraph name) Reference: the paragraph body + Suffix: the untouched procedure text after the body, appended for compilation """ tasks = [] for rec in records: @@ -135,14 +142,10 @@ def build_paragraph_completion(records, max_tasks=200): "paragraph_name": para["name"], "prompt": prompt, "reference": reference, + "suffix": proc_text[para["body_end"] :], "reference_lines": len( [ln for ln in reference.split("\n") if ln.strip()] ), - "stop_at": [ - "\n " + p["name"] + "." - for p in paragraphs - if p["name"] != para["name"] - ][:5], } ) @@ -268,6 +271,7 @@ def build_program_completion(records, max_tasks=50): "source": rec.get("source", ""), "prompt": prompt, "reference": reference, + "reference_text": text, "reference_lines": len(proc_lines), "total_program_lines": rec.get("line_count", 0), } diff --git a/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py index c126db26..56ae54c5 100644 --- a/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py +++ b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py @@ -7,7 +7,9 @@ python -m pytest examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py -v Checks that a PROCEDURE DIVISION header in any case ends the prompt prefix, so the -procedure body that holds the reference never appears in a prompt. +procedure body that holds the reference never appears in a prompt. Also checks +that a paragraph task keeps the procedure text after the masked paragraph as its +suffix, and that a program task keeps the original program. """ import importlib.util @@ -65,3 +67,59 @@ def test_paragraph_prompt_does_not_leak_mixed_case_body(bench): assert "SECRET REFERENCE LINE" in tasks[0]["reference"] assert "SECRET REFERENCE LINE" not in tasks[0]["prompt"] assert "WS-COUNT PIC 9(4)" in tasks[0]["prompt"] + + +MULTI_PARAGRAPH = ( + " IDENTIFICATION DIVISION.\n" + " PROGRAM-ID. MULTI.\n" + " DATA DIVISION.\n" + " WORKING-STORAGE SECTION.\n" + " 01 WS-COUNT PIC 9(4).\n" + " PROCEDURE DIVISION.\n" + " INIT-PARA.\n" + " MOVE 0 TO WS-COUNT.\n" + " MAIN-PARA.\n" + " PERFORM INIT-PARA\n" + " PERFORM ADD-PARA\n" + " DISPLAY WS-COUNT\n" + " STOP RUN.\n" + "\n" + " ADD-PARA.\n" + " ADD 2 TO WS-COUNT.\n" + " END-PARA.\n" + " EXIT.\n" +) + + +def _code_lines(text): + return [ln.rstrip() for ln in text.split("\n") if ln.strip()] + + +def test_paragraph_task_keeps_the_rest_of_the_procedure(bench): + tasks = bench.build_paragraph_completion( + [{"text": MULTI_PARAGRAPH, "is_complete_program": True, "source": "m.cbl"}] + ) + assert len(tasks) == 1 + task = tasks[0] + assert task["paragraph_name"] == "MAIN-PARA" + assert "ADD-PARA." in task["suffix"] + assert "END-PARA." in task["suffix"] + assert "ADD-PARA." not in task["prompt"] + assert "ADD-PARA." not in task["reference"] + assert _code_lines( + task["prompt"] + task["reference"] + task["suffix"] + ) == _code_lines(MULTI_PARAGRAPH) + + +def test_program_task_keeps_the_original_program(bench): + tasks = bench.build_program_completion( + [ + { + "text": MULTI_PARAGRAPH + + "".join(f" DISPLAY 'LINE {i}'.\n" for i in range(5)), + "is_complete_program": True, + } + ] + ) + assert len(tasks) == 1 + assert tasks[0]["reference_text"].startswith(MULTI_PARAGRAPH) diff --git a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py index 1e0e460b..534f2ea6 100755 --- a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -9,7 +9,12 @@ All benchmark tasks are batched and processed concurrently. Before loading the model, main() fails unless every suite in REQUIRED_SUITES -exists in EVAL_DIR with at least one task and `cobc --version` succeeds. +exists in EVAL_DIR with at least one task that has the fields in REQUIRED_FIELDS +and `cobc --version` succeeds. + +Paragraph completions are cut at the next paragraph header and compiled with +the rest of the original procedure appended. Each compile suite also reports +how many tasks compile with the reference in place of the generation. """ import json @@ -43,6 +48,16 @@ "copybook_usage", ) +REQUIRED_FIELDS = { + "perplexity": ("id", "text"), + "paragraph_completion": ("id", "prompt", "reference", "suffix"), + "program_completion": ("id", "prompt", "reference", "reference_text"), + "structural_validation": ("id", "prompt", "reference_text"), + "copybook_usage": ("id", "prompt", "reference"), +} + +PARA_RE = re.compile(r"^ (\d{4}-[A-Z0-9-]+|[A-Z][A-Z0-9-]*)\.\s*$", re.MULTILINE) + GENERATION_BUDGETS = { "paragraph_completion": MAX_NEW_TOKENS, "program_completion": 1024, @@ -103,6 +118,12 @@ def load_required_suites(eval_dir): if not tasks: problems.append(f"{path} has no tasks") continue + missing = sorted( + {f for task in tasks for f in REQUIRED_FIELDS[name] if f not in task} + ) + if missing: + problems.append(f"{path} tasks lack {', '.join(missing)}") + continue suites[name] = tasks if problems: raise RuntimeError( @@ -349,6 +370,7 @@ def _compile_once(src_path, fmt_flag): ["cobc", "-fsyntax-only", fmt_flag, src_path], capture_output=True, text=True, + errors="replace", timeout=10, ) success = result.returncode == 0 @@ -381,28 +403,42 @@ def try_compile(cobol_text): os.unlink(src_path) -def eval_compilation(tasks, generated_texts): +def cut_at_next_paragraph(text): + """Keep a paragraph completion up to the first paragraph header it starts.""" + match = PARA_RE.search(text) + return text[: match.start()] if match else text + + +def eval_compilation(tasks, generated_texts, reference_texts): print(f"\n=== Compilation Evaluation ({len(tasks)} tasks) ===", flush=True) compiled = 0 failed = 0 + reference_compiled = 0 + compiled_where_reference_compiles = 0 fmt_counts = {"free": 0, "fixed": 0} results = [] - for task, gen_text in zip(tasks, generated_texts): - full_text = task["prompt"] + gen_text + for task, gen_text, ref_text in zip(tasks, generated_texts, reference_texts): + full_text = task["prompt"] + gen_text + task.get("suffix", "") success, errors, fmt_used = try_compile(full_text) - if success is None: + ref_success, _, _ = try_compile(ref_text) + if success is None or ref_success is None: raise RuntimeError(f"cobc not found while compiling task {task['id']}") if success: compiled += 1 fmt_counts[fmt_used] += 1 else: failed += 1 + if ref_success: + reference_compiled += 1 + if success: + compiled_where_reference_compiles += 1 results.append( { "id": task["id"], "compiled": success, + "reference_compiled": ref_success, "errors": errors[:500] if errors else "", "format": fmt_used, } @@ -415,6 +451,11 @@ def eval_compilation(tasks, generated_texts): f" Format breakdown: {fmt_counts['free']} free, {fmt_counts['fixed']} fixed", flush=True, ) + print( + f" Reference compiles: {reference_compiled}/{total}; generation compiles " + f"on {compiled_where_reference_compiles} of those", + flush=True, + ) return { "compilation_rate_pct": round(rate, 1), @@ -422,6 +463,8 @@ def eval_compilation(tasks, generated_texts): "failed": failed, "compiled_free": fmt_counts["free"], "compiled_fixed": fmt_counts["fixed"], + "reference_compiled": reference_compiled, + "compiled_where_reference_compiles": compiled_where_reference_compiles, "details": results, } @@ -600,13 +643,21 @@ def main(): f"\n--- Paragraph completion ({len(para_tasks)} tasks, max_tokens={MAX_NEW_TOKENS}) ---", flush=True, ) - para_gen = batch_generate(llm, para_tasks, max_new_tokens=MAX_NEW_TOKENS) + para_gen = [ + cut_at_next_paragraph(g) + for g in batch_generate(llm, para_tasks, max_new_tokens=MAX_NEW_TOKENS) + ] para_quality = eval_generation_quality( para_tasks, para_gen, "paragraph_completion" ) all_results["paragraph_completion"] = para_quality para_compile = eval_compilation( - [{"id": t["id"], "prompt": t["prompt"]} for t in para_tasks], para_gen + [ + {"id": t["id"], "prompt": t["prompt"], "suffix": t["suffix"]} + for t in para_tasks + ], + para_gen, + [t["prompt"] + t["reference"] + t["suffix"] for t in para_tasks], ) all_results["paragraph_compilation"] = { k: v for k, v in para_compile.items() if k != "details" @@ -627,7 +678,9 @@ def main(): ) all_results["program_completion"] = prog_quality prog_compile = eval_compilation( - [{"id": t["id"], "prompt": t["prompt"]} for t in prog_tasks], prog_gen + [{"id": t["id"], "prompt": t["prompt"]} for t in prog_tasks], + prog_gen, + [t["reference_text"] for t in prog_tasks], ) all_results["program_compilation"] = { k: v for k, v in prog_compile.items() if k != "details" @@ -646,7 +699,9 @@ def main(): strv_results = eval_structural(strv_tasks, strv_gen) all_results["structural_validation"] = strv_results["summary"] strv_compile = eval_compilation( - [{"id": t["id"], "prompt": t["prompt"]} for t in strv_tasks], strv_gen + [{"id": t["id"], "prompt": t["prompt"]} for t in strv_tasks], + strv_gen, + [t["reference_text"] for t in strv_tasks], ) all_results["structural_compilation"] = { k: v for k, v in strv_compile.items() if k != "details" diff --git a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py index baac1983..ac56763f 100644 --- a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py +++ b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py @@ -6,15 +6,21 @@ vllm is replaced with a stub, so this needs CPU torch only: python -m pytest examples/cobol-cpt-grpo/07-eval/test_run_eval.py -v -Checks that a missing or empty suite and a missing or broken cobc fail the run -before any model is loaded, that a cobc that disappears mid-run fails instead -of counting tasks as skipped, and that prompts too long for the model context are -set aside and recorded, failing the run if a suite is left empty. +Checks that a missing or empty suite, a suite built without the current task +fields, and a missing or broken cobc fail the run before any model is loaded, +that a cobc that disappears mid-run fails instead of counting tasks as skipped, +and that prompts too long for the model context are set aside and recorded, +failing the run if a suite is left empty. + +Also checks that paragraph completions are cut at the next paragraph and +compiled with the rest of the original procedure. With GnuCOBOL on PATH, a +reference completion for a paragraph that PERFORMs a later one must compile. """ import importlib.util import json import os +import shutil import subprocess import sys import types @@ -40,12 +46,16 @@ def run_eval(tmp_path, monkeypatch): return module -def _write_suites(eval_dir, names, empty=()): +def _write_suites(eval_dir, names, empty=(), drop=None): eval_dir.mkdir(parents=True, exist_ok=True) for name in names: with open(eval_dir / f"{name}.jsonl", "w") as f: if name not in empty: - f.write(json.dumps({"id": f"{name}-1", "prompt": "X"}) + "\n") + task = {"id": f"{name}-1", "prompt": "X", "reference": "Y"} + task.update({"text": "X", "suffix": "", "reference_text": "XY"}) + if drop and drop[0] == name: + del task[drop[1]] + f.write(json.dumps(task) + "\n") def _cobc_ok(args, **kwargs): @@ -68,6 +78,18 @@ def test_missing_and_empty_suites_fail_together(run_eval, tmp_path): assert "perplexity.jsonl has no tasks" in str(err.value) +def test_paragraph_tasks_without_suffix_fail(run_eval, tmp_path): + _write_suites( + tmp_path / "eval", + run_eval.REQUIRED_SUITES, + drop=("paragraph_completion", "suffix"), + ) + with pytest.raises( + RuntimeError, match="paragraph_completion.jsonl tasks lack suffix" + ): + run_eval.load_required_suites(str(tmp_path / "eval")) + + def test_check_compiler_reports_version(run_eval): with mock.patch.object(run_eval.subprocess, "run", side_effect=_cobc_ok): assert run_eval.check_compiler() == "cobc (GnuCOBOL) 3.1.2" @@ -100,7 +122,98 @@ def test_compilation_fails_when_cobc_disappears(run_eval): run_eval.subprocess, "run", side_effect=FileNotFoundError("cobc") ): with pytest.raises(RuntimeError, match="cobc not found while compiling"): - run_eval.eval_compilation([{"id": "t1", "prompt": "X"}], ["Y"]) + run_eval.eval_compilation([{"id": "t1", "prompt": "X"}], ["Y"], ["XY"]) + + +def test_compile_tolerates_non_utf8_compiler_output(run_eval, tmp_path): + fake_cobc = tmp_path / "cobc" + fake_cobc.write_text("#!/bin/sh\nprintf 'bad \\200 byte' >&2\nexit 1\n") + fake_cobc.chmod(0o755) + with mock.patch.dict(os.environ, {"PATH": str(tmp_path)}): + success, errors, _ = run_eval.try_compile("X") + assert success is False + assert "bad" in errors + + +def test_cut_at_next_paragraph_keeps_only_the_target_body(run_eval): + body = " MOVE 1 TO WS-A.\n DISPLAY WS-A.\n" + later = " NEXT-PARA.\n STOP RUN.\n" + assert run_eval.cut_at_next_paragraph(body + later) == body + assert run_eval.cut_at_next_paragraph(body) == body + assert run_eval.cut_at_next_paragraph(" 0100-START.\n" + body) == "" + + +def test_eval_compilation_appends_suffix_and_counts_reference(run_eval): + compiled_texts = [] + + def fake_compile(text): + compiled_texts.append(text) + return text.endswith("SUFFIX") or text == "REF-OK", "", "fixed" + + tasks = [ + {"id": "a", "prompt": "P1-", "suffix": "SUFFIX"}, + {"id": "b", "prompt": "P2-"}, + ] + with mock.patch.object(run_eval, "try_compile", side_effect=fake_compile): + result = run_eval.eval_compilation(tasks, ["G1-", "G2"], ["REF-OK", "REF-BAD"]) + assert compiled_texts == ["P1-G1-SUFFIX", "REF-OK", "P2-G2", "REF-BAD"] + assert result["compiled"] == 1 + assert result["reference_compiled"] == 1 + assert result["compiled_where_reference_compiles"] == 1 + assert [d["reference_compiled"] for d in result["details"]] == [True, False] + + +def _load_benchmarks(tmp_path, monkeypatch): + monkeypatch.setenv("OUT", str(tmp_path / "bench")) + path = os.path.join(EVAL_DIR, "..", "03-benchmarks", "build_benchmarks.py") + spec = importlib.util.spec_from_file_location("build_benchmarks", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +CALLS_LATER_PARAGRAPH = ( + " IDENTIFICATION DIVISION.\n" + " PROGRAM-ID. CALLER.\n" + " DATA DIVISION.\n" + " WORKING-STORAGE SECTION.\n" + " 01 WS-COUNT PIC 9(4) VALUE 0.\n" + " PROCEDURE DIVISION.\n" + " MAIN-PARA.\n" + " MOVE 1 TO WS-COUNT\n" + " PERFORM ADD-PARA\n" + " DISPLAY WS-COUNT\n" + " STOP RUN.\n" + " ADD-PARA.\n" + " ADD 2 TO WS-COUNT.\n" +) + + +def test_paragraph_regex_matches_stage_03(run_eval, tmp_path, monkeypatch): + bench = _load_benchmarks(tmp_path, monkeypatch) + assert run_eval.PARA_RE.pattern == bench.PARA_RE.pattern + assert run_eval.PARA_RE.flags == bench.PARA_RE.flags + + +@pytest.mark.skipif(shutil.which("cobc") is None, reason="needs GnuCOBOL") +def test_reference_completion_compiles_with_later_paragraphs( + run_eval, tmp_path, monkeypatch +): + bench = _load_benchmarks(tmp_path, monkeypatch) + tasks = bench.build_paragraph_completion( + [{"text": CALLS_LATER_PARAGRAPH, "is_complete_program": True}] + ) + assert [t["paragraph_name"] for t in tasks] == ["MAIN-PARA"] + task = tasks[0] + assert "ADD-PARA." in task["suffix"] + assert not run_eval.try_compile(task["prompt"] + task["reference"])[0] + result = run_eval.eval_compilation( + [task], + [task["reference"]], + [task["prompt"] + task["reference"] + task["suffix"]], + ) + assert result["compiled"] == 1 + assert result["reference_compiled"] == 1 def _fake_tokenizer(text): diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index 140fc597..2bf9e12f 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -118,12 +118,26 @@ suites: - Copybook usage: BLEU and edit distance only. Compile scoring is format-aware: each generation is compiled with `cobc -free`, -retried with `-fixed` if that fails, and both counts are reported. The headline -compile score sums the paragraph, program, and structural tasks. Stage 03 caps -these at 200, 50, and 100, so the total depends on the corpus. Before loading -the model, the eval checks that `cobc --version` -works and that all five suite files exist and are not empty. It fails if -either check fails, and it fails if `cobc` disappears during the run. It also +retried with `-fixed` if that fails, and both counts are reported. A paragraph +completion is cut at the first paragraph header it starts. It is then compiled +with the rest of the original procedure appended, so a `PERFORM` of a later +paragraph resolves; BLEU and edit distance score only the cut completion. The +headline compile score sums the paragraph, program, and structural tasks. +Stage 03 caps these at 200, 50, and 100, so the total depends on the corpus. + +Many source programs do not compile on their own, for example because a +copybook is missing. Each compile suite therefore also compiles the reference +in place of the generation: the paragraph prompt, reference, and appended +procedure, or the original program for program and structural tasks. +`reference_compiled` counts the tasks that compile that way, and +`compiled_where_reference_compiles` counts the generations that compile on +those tasks. + +Before loading the model, the eval checks that `cobc --version` works and that +all five suite files exist, are not empty, and have the fields the current +stage 03 writes. It fails if any check fails, and it fails if `cobc` +disappears during the run. Rerun `03-benchmarks` after updating this example. +It also tokenizes every generation prompt and sets aside any prompt that cannot fit in `MAX_MODEL_LEN` (32768) together with its output budget. Set-aside task IDs are listed under `excluded_over_context` in `eval-summary.json`, and the run fails @@ -227,16 +241,25 @@ cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_reward.py cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_stop_sync.py cd examples/cobol-cpt-grpo/06-grpo && python -m pytest test_driver.py cd examples/cobol-cpt-grpo/04-cpt && python -m pytest test_continued_pretrain.py test_driver.py -cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_run_eval.py test_driver.py +cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_run_eval.py +cd examples/cobol-cpt-grpo/07-eval && python -m pytest test_driver.py ``` Run `test_reward.py` in its own process; it replaces `torch` and `trl` with -stubs. `test_stop_sync.py` starts two CPU processes in a gloo group. +stubs. Run `07-eval/test_driver.py` apart from `test_run_eval.py`; it checks +that the driver process never imports `torch`, which `test_run_eval.py` does. +With GnuCOBOL on `PATH`, `test_run_eval.py` also compiles a reference +completion for a paragraph that performs a later one. `test_stop_sync.py` starts two CPU processes in a gloo group. `06-grpo/test_driver.py` starts a local Ray cluster and takes one to five minutes. ## Results +These results predate the paragraph compile fix described under Eval: paragraph +completions were compiled without the rest of the procedure, and no reference +compile counts were recorded. The paragraph and overall compile rows will +change when the eval is rerun. + All three models below were scored by one `07-eval` RayJob each (greedy decoding, 512 new tokens, `MAX_MODEL_LEN` 32768, TP=8 on 8 H200) on the same benchmarks. The pipeline ran end to end from this example with From 114a7a0cd3168fefa564dd4a661f18b3960b5c4a Mon Sep 17 00:00:00 2001 From: June Liu Date: Fri, 23 Oct 2026 09:49:37 -0700 Subject: [PATCH 12/15] fix(examples): report the compile count of an empty generation cobc -fsyntax-only accepts an empty paragraph or an empty PROCEDURE DIVISION, so an empty generation compiles on many tasks. Report that floor next to each compile score and the reference ceiling. Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../build_benchmarks.cpython-312.pyc | Bin 0 -> 19499 bytes .../cobol-cpt-grpo/07-eval/run_eval_vllm.py | 16 ++++++++++++---- .../cobol-cpt-grpo/07-eval/test_run_eval.py | 13 +++++++++++-- examples/cobol-cpt-grpo/README.md | 5 ++++- 4 files changed, 27 insertions(+), 7 deletions(-) create mode 100644 examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc diff --git a/examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc b/examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69507617d721e03def6a967f12fb79b709869aac GIT binary patch literal 19499 zcmb_^X>=Rcm1Y%C*a{$aZr}zkA_cD6D9Tz*aT6t))M~e7%M8SVNKhm}SAklPLWb_{ zw82rz>p%O%pb~(`<&5_nIlS0r|my8 z_r5|EKnjwbc9*SJg{t@N+wQyfyZ7D4|DKn}QgHdc|LN#IK1Wgi3x6mVlYV%t(NNT7 zildHG9L;IQ=;JiGYmRHkU3*+h?#yw9+;zuw0nMmIA(Od}+A`0|6Q|{v_kccL$BlnP zIUT20%bEU&at6+*mSZ{7akGYk(nWF$veN*|vG1Feyfg**m{Cu54*5A?evg))NFMEI z>63XlOF%nX{^Waw`f)o@&s{PP-d8=Nq?5D$Id$B{*&uauc1ZI%2c!j@hjacpeY}t> z;aqSp;z~I;+>5y~E+6h5cwYeb60Q((mvTjrmIcbW;?wIXiZ?=g%HB&?OH=1H=QY0a zsec0##qOO7j&a_OfxQEL-oP3E*pxpU3{7}X1}26^#{K+h!F!4ijeEl*0q-OqnDp}j z&I{RsoSaVxgqvCRaKO(A__}#i2u+N6{S%ybC^R92`Kh5WPzat0WGH%tsbDxDY-hdR zjm=)5=pW|&lOuyeq4CMF0IB_dPQT+F_Y0@Je(%A<107wR#}0RSJ9}U3J<>bSpDE%! z8RE`2V!OEDP}my`1IbOzQ1j&ZlcCV*!70H%9PlEsVQ9ewP*Xm68<0vk=pTd9A*g>m zL*`U403*T&P6c=vgn$43;k!?=a`?A+r4zc9rI#sd@B(bi_KJib%Be{3*oO0d5wt(*fG9`T30<5PlM zXK*aQvR(e6k*p~S3ef!hM+W-)VDiQS-m}5*2=-!pGTfq$Ov`v6><2#ZLmuda(9FKX zydF6}$Rh>^(ft z{{p-eJKK-8iz=J=UKJV?ZXQZY8?rbJkZR%I40;Clak&)3BNdr zL8&3>#zG)Pe2m0Q`oklVfe#FYcutV$QxY>4oCpZ0!EjQ(Jk@ec2=GG7=+s1DEI0++ z#RXcz{;6R;$Tg)|?VKN1Q=0Rdh9<*J!~A5ZWkXvN;lk$TmJOSmGE%PPB+9|T%!}s9 z^Ad|K8YCi(cSGeMN-5z_Aep98R{M0%vMDdd&30UY{_6;irD!z)Z#nujr~zJsbH!

tQ`CvH$n1oVP!}D^lh?^p+eBXo;RyJcNY#)`Ds_$T}Vl4Y8A&Lw-b)Xg?E_}sZmuTVqW=D zW6a*`|%Jho-lG+Pwg6K6IrIU(U|q0p*@|^e>5MS*?7+*f4VCQX&~4 zqYkwuqgXP+1{BuMDC$+971c%=kr~ag+_R=j2RWXx%t0LB(uP)rVR_3}rY<6~3W%s7 zd5GZ3VX`Q!Gr?w2$p}q@NQrvU;GsqxN*$TIQZ{N7jib0C*IdeXP`^=7zxiuE!%q-X zM(E`;+9m3<+GSM67ke!lIL#GRID%sSfFUZ^r;TjwItTa%Ffqy!E`ylUX#Pp|B%I#N z7Xn$-(nJI?D1P`dcqJL}aS(O2780KsHL?(v3Mih7Kwd)8NG@4EI$;6fFbrns@w*_~0n8tV1Xhj1HNAh?>U^i|o~t_Hs$S?% zx;7`RTR^)PR4f>NQ?P!9jcH@YQa1O+Z_j*tZsS|ui7_dA{$vzX`W;^C8FZCoWrS~jV2}{)-OU>VR+_ok=UwWveOqPdw%3eMfSm=n` z*59+WBy257+eV;hwq4vlv;F%!QIFZJxfMF``x1HlIB7o12slXN?SNRQ3D#$ z2;aIGD#}KUK#PU^f0W4-vykra>EYn1^ z=hPJc+C1U9qRLkqR>+07nJmqc%i)vJHel>%fDZ_u7y<(bdF?Eu@T85JqQ>+ydFyjW z{*UZ-QN@yCu&;8xfMW$^*+F7*?Q>RV^&&i*s0QgRU^aT~7;13xS`A>ncCls*SMz3tiK|b_Z zp!Xp-MjS+hqT~t+isOVJ8FaYF$j79rA|8dfM6(Z2#zcx`F`ffz^%FG`F%rscytK*_n3_}Z0oZxxh2U&PTO+KTf34#F{J`YG8u@tBxWkh5gu7{3; z*F!R?ZV#YBMo6-$uLg;I>a$DQ;Xs&xaH3{pt%vMPw4J2q!7&yZm-Is;q2N$}064x4 zKcgO$3?qJFWGr}6(vEYjk{+B-VAo2fk-#}FI2=HAD~v$RV*#Iow_^d*xc}T>82n^{ zq?2!2Fq+B0@)aZl&e|ZVRQ@^qO;1Xiu?dMC4uv?P69qK3dB0>M8PWI{0aqTVaWi=V z&KF<>us}i+1c36rK!x8AiGWfrJx7_FtaUxIEu9(2jAxJk78D!8^8W~8VVZhqrSfcG zlGqAj+LXL` zl+OHS{~P_Yt#_E>l-WM(nf1lZrE#Y8i%0D=7y%&25h??SMj2@5nv2tB znN9jNAK|B>0;4+c1TLJUnVQB_==4uK7{xmL+9 z8ae$H4H_{yM=EOv1-&H~_A=;qG>HrWiAiV=E3#=wVa`A!QQaqrS~!zvIYSE!$KoA; zV3`Q-Q0_(Zz*F-K^eh7y1Sn${^Uh$ejG`Up(hPIyh}yII z?@&gCvk=3NzR)D*hDs@K-)@0*Qg4uebBXz$j2T?~w0vTLSirfl z{;b>{E+1x62J_DViUhV6U>ZhsMXn5AjaDk}GIs@Hl9Gj(ip5`{;t_Se zxRNWX34c1TwF$xtC}hJTOIgZ4xE089farW1fW(ZM*QIYa#w zeVK-Z&N9?_>StQm{-FDBP(H1sY2LsiaN?bi2uR55ot8n9Y3FbN9KDnL?ns@CpOkTv zw=tV{j)jK&W5VucHRG?42)6N`|2Or|(^R}*>y3tmL$j~Xb-r`{M#C3mzvR!X4w|?g z`3}hV;QzrzoIL#xK%-KT{2sC~3RZ=n&eL`_QpK|EFr(x(z&Nr`DsK@3EJX0Wo&B8+$yN;Xy;64O1 z*(6z4_t%s--<7^;OUhiyLLP{0ozxh z;DqeY!i6Uz`h)m%2$RE@9Koa?lcSLM$_RkxkKqeYZ~*0%o(RMljAIN>q*sH`(oE0g zj!K%LlQNIKh@}}aE$B)nxSvN?IanDd2L*lzdUtnbdn z@BEuk_zL_J^sr8Vy;m$$SxwTubG~Qx^@Z23PhFc@JaM;f#}6EHJ=48$`_A8%R3&X4 z^KG;K1^@NHwLrYF{ci2vAJB7c)BEDKj^CNx^ZQaZN6hq@&4GUMxM$}tPkuD^;aJjk z0KDV{MbiVHS_@;El%x2bqcY*BT&P-j@mhV-(G=4!o2{`kZ|#_?OqxsX7kTEk&a(il zS9ID|YiwX{@73=4?uBiOhm)lnKCx_kXrOFG5ACvHV=n^Vp3ODeJtO8aWEQvHEKW4+ zU1iVQuV}tku_aNlsI}+P0);f1 z+_N7Ho!K2pTg3v{Fx4r05!gnT2QCfBMpIYHQyM$4TvD@e>^CJ1mkhJ??3QI$K@z0} zBzIj~Q|^j;?%IUAc9CA(vT**cdn4FCrH!8(DZ6Vm-Jxm*?tarJfc+R)r zpKn?$Svq{9;-(|v+i}yM@a;(!JRj>xm3SA57u(_`n`2;{lwLK>n-&_AMJ=)3`(>3^ zU!H$?VQTUCVlZB}Em`(ltRGyRS9|As7XphtOT)M5pP!D`?oF0<#9jb@sl%6Y7SGjQ zYPsjEOE~M2PTz`_cGTZ@KDVONIJSSTr@-=ZC!JgGId>(TyOPd5D+aCOdHSJ|a&2BQ zF^$d{z9m)Hlxl2C zmDi*y>r?ACE!WjA?o8HgTXC@t=jU2Wp7oIjDtYt@O?lhEZ|f-jeBxQw)DaEk*!6{E z>-gf)CaUBBEqGx){6)e3dgjLrB-ab}w=n;Sf#g@unf7ln{=Jrgbd>>x{!3IdkADGX z)hh4rCeT(GIFL3bAdh-|bMG&JepD}OE(i@v*9P$3kmbV#zwooRlmS}1MT6|a1uv~3 zj1rSvtIQtgpB{EGb~^d+l1NlG2yi!>=vZD9OFWWexYVQ1(W z1`bF3v{AFD2T!pkr$J)Qm>{F>Os&e3XwC&&6>zDvY*aN7lyaFZdCtUH(f~B?DKl>6 zY>N3HnnjDq3F)`vkZg&Ui_mgf}DyL7g(Fd>)3Jz{Cf*h5Fz&{=`rbdCQkS z4UgcP1mM&}K$zrAgQzUeWCnyn&xQ1Vw2N;uA4<{3R&xf8n11G8FxXmapMI{4A&c&{aD3O2XxX z@P_J!>$|V*22WUZTQa{pQQbCoSoMURj90hC*m!<7LV(Cyd*%**Vy^h<=7q>l_ADLw zqXRn=I6mK3%*6mqwn9zr31|MD% z`tZ7+rFB}V0`S^dV30pH=6B`mKQ3n=U8Pp=BT*~=8D?RPR}R?etG#l3;Zyu}8cwSS zB92LK(V_RQQTE^yb--+)+t1z4ZLul4qoIWG$MsWqNxnNV% zfa@lgYIs5nKT8$pa8rwHiyDS9G0w+%)6p)!h%`)CWHw< z%u9Nk3M_t~b?N|gP&W0m!c(*2fgy^MM60TSrD+)>2dR#Q1i_mXp4u!+#|uD|2*C#^ zAo?TB7$-)`S^QYX z*!9126vf&hV(Mb+Ol#axd&gY2*#1G+jjpA}TXeE<@2!)cH130K1{0fd6amcEX2hCnZx z07>RRp9;-1l4Ns73vuR5+`%Xs-)k;#JgRQ&%qNAmiXwA3NR>RCJ_|+W7GK#Yl7+Y> z6_WzKS5e^8Yk!I$%i;jeenrLNAjm#_LXbJ8(7hrHG89?38?87SA)F6ly_JIt7$vXD zFo^DixOF8QUOw4?XBeP3#2bZS%RMd8-UjM?My!1f5}-c|1QOJD8uCOi!ud=Z@8xG& zkQJw&k$67CS~1s))qEck$(4b5>droCfX}*C_MxxpYa==^|G;f@I&fa_nX@1r@-%-L z6ICE)z@1FI0l#7c1v4aB;7t~n6*6?o*<~U*{~AhdM#+*a4H*!w~Mjs9dK zxD0l!me4CU2=w~BmIawyLfLB`Zb-}IJ*rIJ*7@9OnGD`&Stj3yhs#vfF# zC5AJ+uO9D$V@ppQFt|{gK3-sg;{|8C;RJ^zFo9=e2K~ZNFi4ED2$*YlE+q`x4Dd|i zB=CgF(Hpd+J2f^Xi~u-5*eB`E;_($8Ww4~Eal&Ex81Bs zxOd;R>{-!L?h8ei!84^_$uq`pG)iU4=qUG9qg1|Tln~>POU!w`;Qb1PO4Rg} zaZ?dEjGK4uxMj_~=^OQ`QfY^BU)8IsHNC3K^lBCVvQO!Q8i5RbsD4Ht*l*N_8l~;Z z9r}=wsH)7$dZ(@{L;+dvGu*ZMy*d|b@~i#B_fc~eCAC)0$*sGR%>!!PT;ZB8YVB}T zdUYAKey(!O7qtak4SwNjujoW`M%Qt5SF}+}*r&|fHyY>l%J_0V$k2LPuwaZIvDvUjr;{dwA zlu92J-?G#Tvbl&whNsCR+UDlwh{@~4Ob|32ab;&qAN-7X z6cXepFVe~oaCp_1O7?oTsj11Uw#K`0yZ4|H%j@ls4_9M+HRR&I2W|AxWdA*4@p^k8 z^a78pKubz8hh?(276@#7dKicr1Of&WIb}wn?VAn*foQe!G=_oP(2x*>e*P4Egdf9j zljv~vp)Y;}-bD(>@$B-G5g(I)k{+XtE(d%Rw*P!|$uURAJq>q@8h<|e#bXA({i)4C zLpQwb5HF7B-&G6JOZEe{h&}TyO}AkU-PY~iBjhhK;9m^zzL@caKoybFr}#NSK*mP z2u|1Woj5dOn7oC_JD9u+2}D4#YUr3CndRLoQUmdZL<-ceO~TYT{P~IVqzZlkDIhWh z3r`R6Iq(+{A{8tVON#K3|Eh$-BLU#@8nE(UH!N-O`G|Y6#veE{I(x6F+v~!hXsywb#?OJ!WUX-C4U{CG$DZ}mf>0W zh0)5C?4EfOpK#NPc#L^;eh_0*wlW#CGJRIb3_=pZeB|@)c|@|QKg|$~r2ZfWZ6JgQR_NelC^ROq31E;6TCg0(s1dbt|CR zljYq}OuazKgues@1hy?>inH5pF?X4+&*0Hg|Ngec`lSx#-LkXv%AX{i zwbKV6W32a`p*aXWsg38?#p|}m*&VmG+-3HWaO%}E%-R*E(kNwp^!hHnJ78vv5c5uur z4K418Td|?Weann(mc6p~YR`Pn)xP<@_g(Sk{mIe;cbGm>^TwO*yUgyCz2rtyvi#tU zriDW{H-6OmVe753yE_Jc^!j2`qWoaOeh|)OJJu}(Z$jhtJlVM2adyuw=s{nO$~u#! zU3ZvnQW>}b?lOCq;fz7UE$1Ca=aZE7#@UYBTkbLkzp|TKmR`OoEDgl1owqyh8V}{D zd0(=$_YSk4)ZDZLHE&?gA!oX> z>1x}2+tr=(JKxvFoAxA2p1;Gi6LK=|EIT$VmE6|ear8aek*+w~eH)nKNREabNS5~9 zVO{`ViOCY%|CVi;&5Mn`)dFEcCeLT-XPw66fef&X8+@x70)$M(5a0yAFrSGrd0@9g z$A0dnJY~otsiG|q-nGsO8#Bs_W_VGN{bCcmsI~%?b?X;RKk1L}>`!B;t zEf-BQrrEl=jSGkF8tc#t^gIoNlw#`_TNmr%8;`|}FQ%-{f6{4n`Y#@ZXsY56{Rh}j zgluuso-c(WSew6Q3l19ezpi8;{RemPL5ntS(H+b;NXEfIT)=~apj6T*=mW?Pzn0|N zkMLF@83u8mY;ceWD*l)7lApu`1tHmC!@$5ZK|VAg3%b9-_ct*i!tZb4NrK%r2pEJT z^ZbYS;dM-SOkh(8cK>>N`Uehob+jMp;<3rHK_{a zYh#jqY1(k#S`<5Zadc)h#q7Gte-!;N3TJO@nrZVR)mQD2D8>?te4;J>Eo1y{H}E*) IAtUtv0a#O68~^|S literal 0 HcmV?d00001 diff --git a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py index 534f2ea6..402d311a 100755 --- a/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -14,7 +14,8 @@ Paragraph completions are cut at the next paragraph header and compiled with the rest of the original procedure appended. Each compile suite also reports -how many tasks compile with the reference in place of the generation. +how many tasks compile with the reference, and with an empty generation, in +place of the generation. """ import json @@ -415,15 +416,19 @@ def eval_compilation(tasks, generated_texts, reference_texts): failed = 0 reference_compiled = 0 compiled_where_reference_compiles = 0 + empty_compiled = 0 fmt_counts = {"free": 0, "fixed": 0} results = [] for task, gen_text, ref_text in zip(tasks, generated_texts, reference_texts): - full_text = task["prompt"] + gen_text + task.get("suffix", "") - success, errors, fmt_used = try_compile(full_text) + suffix = task.get("suffix", "") + success, errors, fmt_used = try_compile(task["prompt"] + gen_text + suffix) ref_success, _, _ = try_compile(ref_text) - if success is None or ref_success is None: + empty_success, _, _ = try_compile(task["prompt"] + suffix) + if None in (success, ref_success, empty_success): raise RuntimeError(f"cobc not found while compiling task {task['id']}") + if empty_success: + empty_compiled += 1 if success: compiled += 1 fmt_counts[fmt_used] += 1 @@ -439,6 +444,7 @@ def eval_compilation(tasks, generated_texts, reference_texts): "id": task["id"], "compiled": success, "reference_compiled": ref_success, + "empty_compiled": empty_success, "errors": errors[:500] if errors else "", "format": fmt_used, } @@ -456,6 +462,7 @@ def eval_compilation(tasks, generated_texts, reference_texts): f"on {compiled_where_reference_compiles} of those", flush=True, ) + print(f" Empty generation compiles: {empty_compiled}/{total}", flush=True) return { "compilation_rate_pct": round(rate, 1), @@ -465,6 +472,7 @@ def eval_compilation(tasks, generated_texts, reference_texts): "compiled_fixed": fmt_counts["fixed"], "reference_compiled": reference_compiled, "compiled_where_reference_compiles": compiled_where_reference_compiles, + "empty_compiled": empty_compiled, "details": results, } diff --git a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py index ac56763f..cb7ce9ec 100644 --- a/examples/cobol-cpt-grpo/07-eval/test_run_eval.py +++ b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py @@ -143,7 +143,7 @@ def test_cut_at_next_paragraph_keeps_only_the_target_body(run_eval): assert run_eval.cut_at_next_paragraph(" 0100-START.\n" + body) == "" -def test_eval_compilation_appends_suffix_and_counts_reference(run_eval): +def test_eval_compilation_appends_suffix_and_counts_reference_and_empty(run_eval): compiled_texts = [] def fake_compile(text): @@ -156,11 +156,20 @@ def fake_compile(text): ] with mock.patch.object(run_eval, "try_compile", side_effect=fake_compile): result = run_eval.eval_compilation(tasks, ["G1-", "G2"], ["REF-OK", "REF-BAD"]) - assert compiled_texts == ["P1-G1-SUFFIX", "REF-OK", "P2-G2", "REF-BAD"] + assert compiled_texts == [ + "P1-G1-SUFFIX", + "REF-OK", + "P1-SUFFIX", + "P2-G2", + "REF-BAD", + "P2-", + ] assert result["compiled"] == 1 assert result["reference_compiled"] == 1 assert result["compiled_where_reference_compiles"] == 1 + assert result["empty_compiled"] == 1 assert [d["reference_compiled"] for d in result["details"]] == [True, False] + assert [d["empty_compiled"] for d in result["details"]] == [True, False] def _load_benchmarks(tmp_path, monkeypatch): diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index 2bf9e12f..dd0e90b8 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -131,7 +131,10 @@ in place of the generation: the paragraph prompt, reference, and appended procedure, or the original program for program and structural tasks. `reference_compiled` counts the tasks that compile that way, and `compiled_where_reference_compiles` counts the generations that compile on -those tasks. +those tasks. `empty_compiled` counts the tasks that compile with an empty +generation. `cobc -fsyntax-only` accepts an empty paragraph or an empty +PROCEDURE DIVISION, so a compile count is meaningful only above this floor and +alongside BLEU, edit distance, and the structural checks. Before loading the model, the eval checks that `cobc --version` works and that all five suite files exist, are not empty, and have the fields the current From f64b18b6edb1ec6d71e783d08e4954135b008091 Mon Sep 17 00:00:00 2001 From: June Liu Date: Fri, 23 Oct 2026 10:12:44 -0700 Subject: [PATCH 13/15] docs(examples): report eval results with reference and empty-generation compile counts Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- examples/cobol-cpt-grpo/README.md | 69 +++++++++++++++++++------------ 1 file changed, 42 insertions(+), 27 deletions(-) diff --git a/examples/cobol-cpt-grpo/README.md b/examples/cobol-cpt-grpo/README.md index dd0e90b8..d3a4ca66 100644 --- a/examples/cobol-cpt-grpo/README.md +++ b/examples/cobol-cpt-grpo/README.md @@ -122,7 +122,7 @@ retried with `-fixed` if that fails, and both counts are reported. A paragraph completion is cut at the first paragraph header it starts. It is then compiled with the rest of the original procedure appended, so a `PERFORM` of a later paragraph resolves; BLEU and edit distance score only the cut completion. The -headline compile score sums the paragraph, program, and structural tasks. +total compile count sums the paragraph, program, and structural tasks. Stage 03 caps these at 200, 50, and 100, so the total depends on the corpus. Many source programs do not compile on their own, for example because a @@ -258,15 +258,10 @@ minutes. ## Results -These results predate the paragraph compile fix described under Eval: paragraph -completions were compiled without the rest of the procedure, and no reference -compile counts were recorded. The paragraph and overall compile rows will -change when the eval is rerun. - All three models below were scored by one `07-eval` RayJob each (greedy -decoding, 512 new tokens, `MAX_MODEL_LEN` 32768, TP=8 on 8 H200) on the same -benchmarks. The pipeline ran end to end from this example with -`SOURCE_LICENSES=permissive` and no gated datasets: +decoding, 512 new tokens, `MAX_MODEL_LEN` 32768, TP=8 on 8 H200, GnuCOBOL +3.1.2) on the same benchmarks. The pipeline ran end to end from this example +with `SOURCE_LICENSES=permissive` and no gated datasets: - CPT: 60 steps (3 epochs) from `Qwen/Qwen2.5-Coder-32B`. Eval loss was 0.523 at step 25 and 0.507 at step 60. @@ -275,27 +270,47 @@ benchmarks. The pipeline ran end to end from this example with - Benchmarks: 62 paragraph, 50 program, 100 structural, and 78 copybook tasks, plus 158 held-out programs for perplexity. Three prompts did not fit in the context and were set aside for every model: one paragraph task (51,129 - tokens) and two copybook tasks (55,731 and 59,490 tokens). The compile score - is out of the remaining 211 tasks. + tokens) and two copybook tasks (55,731 and 59,490 tokens). | Metric | Base | CPT | GRPO | | --- | --- | --- | --- | -| Compile, overall | 50/211 | **89/211** | 88/211 | -| Compile, paragraph completion | 16/61 (6 free + 10 fixed) | 24/61 (9 + 15) | 25/61 (9 + 16) | -| Compile, program completion | 13/50 (8 + 5) | 20/50 (12 + 8) | 20/50 (11 + 9) | -| Compile, structural validation | 21/100 (11 + 10) | 45/100 (23 + 22) | 43/100 (22 + 21) | | COBOL perplexity | 1.90 | 1.54 | 1.54 | -| BLEU, paragraph / program | 0.130 / 0.056 | 0.176 / 0.096 | 0.180 / 0.099 | -| Edit distance, paragraph / program | 0.677 / 0.661 | 0.632 / 0.632 | 0.632 / 0.623 | -| PERFORM reference accuracy | 15.9% | 40.7% | 34.6% | -| Column formatting | 74.0% | 80.0% | 78.0% | +| BLEU, paragraph / program | 0.196 / 0.060 | 0.246 / 0.096 | 0.249 / 0.098 | +| Edit distance, paragraph / program | 0.566 / 0.667 | 0.497 / 0.632 | 0.502 / 0.623 | +| BLEU / edit distance, copybook usage | 0.013 / 0.884 | 0.036 / 0.842 | 0.038 / 0.848 | +| PERFORM reference accuracy | 27.1% | 39.8% | 35.4% | +| Column formatting | 75.0% | 78.0% | 79.0% | | PIC clause accuracy | 98.5% | 98.5% | 98.5% | | Division structure | 100% | 100% | 100% | -| BLEU / edit distance, copybook usage | 0.013 / 0.883 | 0.035 / 0.843 | 0.038 / 0.849 | - -CPT accounts for nearly all of the improvement over the base model: compiles -rise from 50 to 89, perplexity falls from 1.90 to 1.54, and PERFORM reference -accuracy more than doubles. GRPO does not improve on CPT on these benchmarks. -Its compile count and BLEU are within one task or a few thousandths of CPT, -and its PERFORM reference accuracy is 6.1 points lower even though the -structural reward includes a PERFORM bonus. + +Compile counts, with the generations that compile on tasks whose reference +compiles in parentheses: + +| Suite | Tasks | Reference | Empty generation | Base | CPT | GRPO | +| --- | --- | --- | --- | --- | --- | --- | +| Paragraph completion | 61 | 30 | 31 | 17 (16) | 22 (20) | 23 (22) | +| Program completion | 50 | 28 | 34 | 11 (10) | 20 (17) | 20 (17) | +| Structural validation | 100 | 52 | 68 | 26 (23) | 44 (36) | 42 (33) | +| Total | 211 | 110 | 133 | 54 (49) | 86 (73) | 85 (72) | + +An empty generation compiles on more tasks than any model and than the +reference, because `cobc -fsyntax-only` accepts an empty body while a +reference can contain statements that fail on their own. No model reaches +that floor in any suite. These compile counts show how often a generation +adds a compile error; they do not show that a model writes code that +compiles. Compare models on perplexity, BLEU, edit distance, and the +structural checks. + +CPT accounts for nearly all of the improvement over the base model: +perplexity falls from 1.90 to 1.54, paragraph BLEU rises from 0.196 to 0.246, +and PERFORM reference accuracy rises from 27.1% to 39.8%. GRPO does not +improve on CPT on these benchmarks. Its BLEU and edit distance are within a +few thousandths of CPT, and its PERFORM reference accuracy is 4.4 points +lower even though the structural reward includes a PERFORM bonus. + +Greedy decoding at TP=8 is not exactly reproducible between runs. An earlier +run of the same base model on the same program and structural tasks measured +13/50 program compiles and 15.9% PERFORM reference accuracy, against 11/50 +and 27.1% here, and 13 of its 50 program generations scored a different +BLEU. Treat differences of a few tasks or a few points between models as +within run-to-run variation. From 2300002e8a9e5106a9b3efd61c04d63d5c435cc2 Mon Sep 17 00:00:00 2001 From: June Liu Date: Fri, 23 Oct 2026 11:17:37 -0700 Subject: [PATCH 14/15] fix(examples): remove committed Python bytecode from 03-benchmarks Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .../build_benchmarks.cpython-312.pyc | Bin 19499 -> 0 bytes 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc diff --git a/examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc b/examples/cobol-cpt-grpo/03-benchmarks/__pycache__/build_benchmarks.cpython-312.pyc deleted file mode 100644 index 69507617d721e03def6a967f12fb79b709869aac..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 19499 zcmb_^X>=Rcm1Y%C*a{$aZr}zkA_cD6D9Tz*aT6t))M~e7%M8SVNKhm}SAklPLWb_{ zw82rz>p%O%pb~(`<&5_nIlS0r|my8 z_r5|EKnjwbc9*SJg{t@N+wQyfyZ7D4|DKn}QgHdc|LN#IK1Wgi3x6mVlYV%t(NNT7 zildHG9L;IQ=;JiGYmRHkU3*+h?#yw9+;zuw0nMmIA(Od}+A`0|6Q|{v_kccL$BlnP zIUT20%bEU&at6+*mSZ{7akGYk(nWF$veN*|vG1Feyfg**m{Cu54*5A?evg))NFMEI z>63XlOF%nX{^Waw`f)o@&s{PP-d8=Nq?5D$Id$B{*&uauc1ZI%2c!j@hjacpeY}t> z;aqSp;z~I;+>5y~E+6h5cwYeb60Q((mvTjrmIcbW;?wIXiZ?=g%HB&?OH=1H=QY0a zsec0##qOO7j&a_OfxQEL-oP3E*pxpU3{7}X1}26^#{K+h!F!4ijeEl*0q-OqnDp}j z&I{RsoSaVxgqvCRaKO(A__}#i2u+N6{S%ybC^R92`Kh5WPzat0WGH%tsbDxDY-hdR zjm=)5=pW|&lOuyeq4CMF0IB_dPQT+F_Y0@Je(%A<107wR#}0RSJ9}U3J<>bSpDE%! z8RE`2V!OEDP}my`1IbOzQ1j&ZlcCV*!70H%9PlEsVQ9ewP*Xm68<0vk=pTd9A*g>m zL*`U403*T&P6c=vgn$43;k!?=a`?A+r4zc9rI#sd@B(bi_KJib%Be{3*oO0d5wt(*fG9`T30<5PlM zXK*aQvR(e6k*p~S3ef!hM+W-)VDiQS-m}5*2=-!pGTfq$Ov`v6><2#ZLmuda(9FKX zydF6}$Rh>^(ft z{{p-eJKK-8iz=J=UKJV?ZXQZY8?rbJkZR%I40;Clak&)3BNdr zL8&3>#zG)Pe2m0Q`oklVfe#FYcutV$QxY>4oCpZ0!EjQ(Jk@ec2=GG7=+s1DEI0++ z#RXcz{;6R;$Tg)|?VKN1Q=0Rdh9<*J!~A5ZWkXvN;lk$TmJOSmGE%PPB+9|T%!}s9 z^Ad|K8YCi(cSGeMN-5z_Aep98R{M0%vMDdd&30UY{_6;irD!z)Z#nujr~zJsbH!

tQ`CvH$n1oVP!}D^lh?^p+eBXo;RyJcNY#)`Ds_$T}Vl4Y8A&Lw-b)Xg?E_}sZmuTVqW=D zW6a*`|%Jho-lG+Pwg6K6IrIU(U|q0p*@|^e>5MS*?7+*f4VCQX&~4 zqYkwuqgXP+1{BuMDC$+971c%=kr~ag+_R=j2RWXx%t0LB(uP)rVR_3}rY<6~3W%s7 zd5GZ3VX`Q!Gr?w2$p}q@NQrvU;GsqxN*$TIQZ{N7jib0C*IdeXP`^=7zxiuE!%q-X zM(E`;+9m3<+GSM67ke!lIL#GRID%sSfFUZ^r;TjwItTa%Ffqy!E`ylUX#Pp|B%I#N z7Xn$-(nJI?D1P`dcqJL}aS(O2780KsHL?(v3Mih7Kwd)8NG@4EI$;6fFbrns@w*_~0n8tV1Xhj1HNAh?>U^i|o~t_Hs$S?% zx;7`RTR^)PR4f>NQ?P!9jcH@YQa1O+Z_j*tZsS|ui7_dA{$vzX`W;^C8FZCoWrS~jV2}{)-OU>VR+_ok=UwWveOqPdw%3eMfSm=n` z*59+WBy257+eV;hwq4vlv;F%!QIFZJxfMF``x1HlIB7o12slXN?SNRQ3D#$ z2;aIGD#}KUK#PU^f0W4-vykra>EYn1^ z=hPJc+C1U9qRLkqR>+07nJmqc%i)vJHel>%fDZ_u7y<(bdF?Eu@T85JqQ>+ydFyjW z{*UZ-QN@yCu&;8xfMW$^*+F7*?Q>RV^&&i*s0QgRU^aT~7;13xS`A>ncCls*SMz3tiK|b_Z zp!Xp-MjS+hqT~t+isOVJ8FaYF$j79rA|8dfM6(Z2#zcx`F`ffz^%FG`F%rscytK*_n3_}Z0oZxxh2U&PTO+KTf34#F{J`YG8u@tBxWkh5gu7{3; z*F!R?ZV#YBMo6-$uLg;I>a$DQ;Xs&xaH3{pt%vMPw4J2q!7&yZm-Is;q2N$}064x4 zKcgO$3?qJFWGr}6(vEYjk{+B-VAo2fk-#}FI2=HAD~v$RV*#Iow_^d*xc}T>82n^{ zq?2!2Fq+B0@)aZl&e|ZVRQ@^qO;1Xiu?dMC4uv?P69qK3dB0>M8PWI{0aqTVaWi=V z&KF<>us}i+1c36rK!x8AiGWfrJx7_FtaUxIEu9(2jAxJk78D!8^8W~8VVZhqrSfcG zlGqAj+LXL` zl+OHS{~P_Yt#_E>l-WM(nf1lZrE#Y8i%0D=7y%&25h??SMj2@5nv2tB znN9jNAK|B>0;4+c1TLJUnVQB_==4uK7{xmL+9 z8ae$H4H_{yM=EOv1-&H~_A=;qG>HrWiAiV=E3#=wVa`A!QQaqrS~!zvIYSE!$KoA; zV3`Q-Q0_(Zz*F-K^eh7y1Sn${^Uh$ejG`Up(hPIyh}yII z?@&gCvk=3NzR)D*hDs@K-)@0*Qg4uebBXz$j2T?~w0vTLSirfl z{;b>{E+1x62J_DViUhV6U>ZhsMXn5AjaDk}GIs@Hl9Gj(ip5`{;t_Se zxRNWX34c1TwF$xtC}hJTOIgZ4xE089farW1fW(ZM*QIYa#w zeVK-Z&N9?_>StQm{-FDBP(H1sY2LsiaN?bi2uR55ot8n9Y3FbN9KDnL?ns@CpOkTv zw=tV{j)jK&W5VucHRG?42)6N`|2Or|(^R}*>y3tmL$j~Xb-r`{M#C3mzvR!X4w|?g z`3}hV;QzrzoIL#xK%-KT{2sC~3RZ=n&eL`_QpK|EFr(x(z&Nr`DsK@3EJX0Wo&B8+$yN;Xy;64O1 z*(6z4_t%s--<7^;OUhiyLLP{0ozxh z;DqeY!i6Uz`h)m%2$RE@9Koa?lcSLM$_RkxkKqeYZ~*0%o(RMljAIN>q*sH`(oE0g zj!K%LlQNIKh@}}aE$B)nxSvN?IanDd2L*lzdUtnbdn z@BEuk_zL_J^sr8Vy;m$$SxwTubG~Qx^@Z23PhFc@JaM;f#}6EHJ=48$`_A8%R3&X4 z^KG;K1^@NHwLrYF{ci2vAJB7c)BEDKj^CNx^ZQaZN6hq@&4GUMxM$}tPkuD^;aJjk z0KDV{MbiVHS_@;El%x2bqcY*BT&P-j@mhV-(G=4!o2{`kZ|#_?OqxsX7kTEk&a(il zS9ID|YiwX{@73=4?uBiOhm)lnKCx_kXrOFG5ACvHV=n^Vp3ODeJtO8aWEQvHEKW4+ zU1iVQuV}tku_aNlsI}+P0);f1 z+_N7Ho!K2pTg3v{Fx4r05!gnT2QCfBMpIYHQyM$4TvD@e>^CJ1mkhJ??3QI$K@z0} zBzIj~Q|^j;?%IUAc9CA(vT**cdn4FCrH!8(DZ6Vm-Jxm*?tarJfc+R)r zpKn?$Svq{9;-(|v+i}yM@a;(!JRj>xm3SA57u(_`n`2;{lwLK>n-&_AMJ=)3`(>3^ zU!H$?VQTUCVlZB}Em`(ltRGyRS9|As7XphtOT)M5pP!D`?oF0<#9jb@sl%6Y7SGjQ zYPsjEOE~M2PTz`_cGTZ@KDVONIJSSTr@-=ZC!JgGId>(TyOPd5D+aCOdHSJ|a&2BQ zF^$d{z9m)Hlxl2C zmDi*y>r?ACE!WjA?o8HgTXC@t=jU2Wp7oIjDtYt@O?lhEZ|f-jeBxQw)DaEk*!6{E z>-gf)CaUBBEqGx){6)e3dgjLrB-ab}w=n;Sf#g@unf7ln{=Jrgbd>>x{!3IdkADGX z)hh4rCeT(GIFL3bAdh-|bMG&JepD}OE(i@v*9P$3kmbV#zwooRlmS}1MT6|a1uv~3 zj1rSvtIQtgpB{EGb~^d+l1NlG2yi!>=vZD9OFWWexYVQ1(W z1`bF3v{AFD2T!pkr$J)Qm>{F>Os&e3XwC&&6>zDvY*aN7lyaFZdCtUH(f~B?DKl>6 zY>N3HnnjDq3F)`vkZg&Ui_mgf}DyL7g(Fd>)3Jz{Cf*h5Fz&{=`rbdCQkS z4UgcP1mM&}K$zrAgQzUeWCnyn&xQ1Vw2N;uA4<{3R&xf8n11G8FxXmapMI{4A&c&{aD3O2XxX z@P_J!>$|V*22WUZTQa{pQQbCoSoMURj90hC*m!<7LV(Cyd*%**Vy^h<=7q>l_ADLw zqXRn=I6mK3%*6mqwn9zr31|MD% z`tZ7+rFB}V0`S^dV30pH=6B`mKQ3n=U8Pp=BT*~=8D?RPR}R?etG#l3;Zyu}8cwSS zB92LK(V_RQQTE^yb--+)+t1z4ZLul4qoIWG$MsWqNxnNV% zfa@lgYIs5nKT8$pa8rwHiyDS9G0w+%)6p)!h%`)CWHw< z%u9Nk3M_t~b?N|gP&W0m!c(*2fgy^MM60TSrD+)>2dR#Q1i_mXp4u!+#|uD|2*C#^ zAo?TB7$-)`S^QYX z*!9126vf&hV(Mb+Ol#axd&gY2*#1G+jjpA}TXeE<@2!)cH130K1{0fd6amcEX2hCnZx z07>RRp9;-1l4Ns73vuR5+`%Xs-)k;#JgRQ&%qNAmiXwA3NR>RCJ_|+W7GK#Yl7+Y> z6_WzKS5e^8Yk!I$%i;jeenrLNAjm#_LXbJ8(7hrHG89?38?87SA)F6ly_JIt7$vXD zFo^DixOF8QUOw4?XBeP3#2bZS%RMd8-UjM?My!1f5}-c|1QOJD8uCOi!ud=Z@8xG& zkQJw&k$67CS~1s))qEck$(4b5>droCfX}*C_MxxpYa==^|G;f@I&fa_nX@1r@-%-L z6ICE)z@1FI0l#7c1v4aB;7t~n6*6?o*<~U*{~AhdM#+*a4H*!w~Mjs9dK zxD0l!me4CU2=w~BmIawyLfLB`Zb-}IJ*rIJ*7@9OnGD`&Stj3yhs#vfF# zC5AJ+uO9D$V@ppQFt|{gK3-sg;{|8C;RJ^zFo9=e2K~ZNFi4ED2$*YlE+q`x4Dd|i zB=CgF(Hpd+J2f^Xi~u-5*eB`E;_($8Ww4~Eal&Ex81Bs zxOd;R>{-!L?h8ei!84^_$uq`pG)iU4=qUG9qg1|Tln~>POU!w`;Qb1PO4Rg} zaZ?dEjGK4uxMj_~=^OQ`QfY^BU)8IsHNC3K^lBCVvQO!Q8i5RbsD4Ht*l*N_8l~;Z z9r}=wsH)7$dZ(@{L;+dvGu*ZMy*d|b@~i#B_fc~eCAC)0$*sGR%>!!PT;ZB8YVB}T zdUYAKey(!O7qtak4SwNjujoW`M%Qt5SF}+}*r&|fHyY>l%J_0V$k2LPuwaZIvDvUjr;{dwA zlu92J-?G#Tvbl&whNsCR+UDlwh{@~4Ob|32ab;&qAN-7X z6cXepFVe~oaCp_1O7?oTsj11Uw#K`0yZ4|H%j@ls4_9M+HRR&I2W|AxWdA*4@p^k8 z^a78pKubz8hh?(276@#7dKicr1Of&WIb}wn?VAn*foQe!G=_oP(2x*>e*P4Egdf9j zljv~vp)Y;}-bD(>@$B-G5g(I)k{+XtE(d%Rw*P!|$uURAJq>q@8h<|e#bXA({i)4C zLpQwb5HF7B-&G6JOZEe{h&}TyO}AkU-PY~iBjhhK;9m^zzL@caKoybFr}#NSK*mP z2u|1Woj5dOn7oC_JD9u+2}D4#YUr3CndRLoQUmdZL<-ceO~TYT{P~IVqzZlkDIhWh z3r`R6Iq(+{A{8tVON#K3|Eh$-BLU#@8nE(UH!N-O`G|Y6#veE{I(x6F+v~!hXsywb#?OJ!WUX-C4U{CG$DZ}mf>0W zh0)5C?4EfOpK#NPc#L^;eh_0*wlW#CGJRIb3_=pZeB|@)c|@|QKg|$~r2ZfWZ6JgQR_NelC^ROq31E;6TCg0(s1dbt|CR zljYq}OuazKgues@1hy?>inH5pF?X4+&*0Hg|Ngec`lSx#-LkXv%AX{i zwbKV6W32a`p*aXWsg38?#p|}m*&VmG+-3HWaO%}E%-R*E(kNwp^!hHnJ78vv5c5uur z4K418Td|?Weann(mc6p~YR`Pn)xP<@_g(Sk{mIe;cbGm>^TwO*yUgyCz2rtyvi#tU zriDW{H-6OmVe753yE_Jc^!j2`qWoaOeh|)OJJu}(Z$jhtJlVM2adyuw=s{nO$~u#! zU3ZvnQW>}b?lOCq;fz7UE$1Ca=aZE7#@UYBTkbLkzp|TKmR`OoEDgl1owqyh8V}{D zd0(=$_YSk4)ZDZLHE&?gA!oX> z>1x}2+tr=(JKxvFoAxA2p1;Gi6LK=|EIT$VmE6|ear8aek*+w~eH)nKNREabNS5~9 zVO{`ViOCY%|CVi;&5Mn`)dFEcCeLT-XPw66fef&X8+@x70)$M(5a0yAFrSGrd0@9g z$A0dnJY~otsiG|q-nGsO8#Bs_W_VGN{bCcmsI~%?b?X;RKk1L}>`!B;t zEf-BQrrEl=jSGkF8tc#t^gIoNlw#`_TNmr%8;`|}FQ%-{f6{4n`Y#@ZXsY56{Rh}j zgluuso-c(WSew6Q3l19ezpi8;{RemPL5ntS(H+b;NXEfIT)=~apj6T*=mW?Pzn0|N zkMLF@83u8mY;ceWD*l)7lApu`1tHmC!@$5ZK|VAg3%b9-_ct*i!tZb4NrK%r2pEJT z^ZbYS;dM-SOkh(8cK>>N`Uehob+jMp;<3rHK_{a zYh#jqY1(k#S`<5Zadc)h#q7Gte-!;N3TJO@nrZVR)mQD2D8>?te4;J>Eo1y{H}E*) IAtUtv0a#O68~^|S From d2cc73ea09ab117df1329e7407244698d0c8bd41 Mon Sep 17 00:00:00 2001 From: June Liu Date: Fri, 23 Oct 2026 11:18:55 -0700 Subject: [PATCH 15/15] chore: ignore Python bytecode caches Co-Authored-By: Claude Opus 5.5 Signed-off-by: June Liu --- .gitignore | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.gitignore b/.gitignore index e31ca608..00a96521 100644 --- a/.gitignore +++ b/.gitignore @@ -25,3 +25,6 @@ portal/frontend/*.tsbuildinfo metrics/experiment-metrics-collector/bin/ metrics/experiment-metrics-collector/taugrid-metrics-collector + +__pycache__/ +*.pyc