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 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..29b004c8 --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-cobol-data.sh @@ -0,0 +1,253 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +set -uo pipefail + +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}" + +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" + 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 +} + +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 "ERROR: pip install failed"; exit 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") +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) +failed = [] +for split_name in ["question_answering", "multiple_choice_question", "COBOL_code_summarization"]: + try: + 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}: 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 + +collect_eval_benchmarks() { + python3 - <<'PYEOF' +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"), +]: + 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) + break + except Exception: + continue + if ds is None: + 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: + 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}") + 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/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 + 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 + +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}..." + 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 +} + +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 -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 + +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}" +fail_if_incomplete "${TOTAL_COBOL}" + +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..a55dd9c5 --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-gated.py @@ -0,0 +1,220 @@ +#!/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, + ) + 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, + ) + ) + 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..c6ef178f --- /dev/null +++ b/examples/cobol-cpt-grpo/01-collect/collect-rl-data.sh @@ -0,0 +1,362 @@ +#!/usr/bin/env bash +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +set -uo pipefail + +DATA_ROOT="${DATA_ROOT:-/data/datasets/cobol}" +export RAW="${DATA_ROOT}/raw" +LOGS="${DATA_ROOT}/logs" +mkdir -p "${RAW}" "${LOGS}" + +LOG_FILE="${LOGS}/collect-rl.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" + 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 +} + +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 "ERROR: pip install failed"; exit 1; } +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) +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) + 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) + +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") +except Exception: + try: + 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) + 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}..." + 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" +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" + return 1 + 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}" +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" +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..472b8d90 --- /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: + DATA_ROOT: /data/datasets/cobol + 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..17a252b9 --- /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: + DATA_ROOT: /data/datasets/cobol + 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..03e03245 --- /dev/null +++ b/examples/cobol-cpt-grpo/02-data-prep/prepare_data.py @@ -0,0 +1,619 @@ +#!/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}") + malformed = 0 + 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: + 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}") + 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 (OSError, UnicodeDecodeError) as e: + print(f" WARNING: skipped {jpath}: {e}") + 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..60955b9b --- /dev/null +++ b/examples/cobol-cpt-grpo/03-benchmarks/build_benchmarks.py @@ -0,0 +1,435 @@ +#!/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 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. +""" + +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 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 = [] + 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(), + "body_end": start + len(body), + "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 + Suffix: the untouched procedure text after the body, appended for compilation + """ + 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_before_procedure_division(text) + + 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, + "suffix": proc_text[para["body_end"] :], + "reference_lines": len( + [ln for ln in reference.split("\n") if ln.strip()] + ), + } + ) + + 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_before_procedure_division(text).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_text": text, + "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/03-benchmarks/test_build_benchmarks.py b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py new file mode 100644 index 00000000..56ae54c5 --- /dev/null +++ b/examples/cobol-cpt-grpo/03-benchmarks/test_build_benchmarks.py @@ -0,0 +1,125 @@ +# 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. 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 +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"] + + +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/04-cpt/continued_pretrain.py b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py new file mode 100644 index 00000000..0628db60 --- /dev/null +++ b/examples/cobol-cpt-grpo/04-cpt/continued_pretrain.py @@ -0,0 +1,443 @@ +# 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) + 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, + 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..4f6983b1 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/prepare_rl_prompts.py @@ -0,0 +1,454 @@ +#!/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. + +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 + +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")) +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) + +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) + +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() + + +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 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) + and not overlaps_val(prepare_text(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: + 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 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) + f.flush() + src = f.name + try: + result = subprocess.run( + [ + "cobc", + "-fsyntax-only", + "-free", + "-frelax-syntax-checks", + "-frelax-level-hierarchy", + "-fmissing-statement=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", + "-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): + 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): + 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_stripped = prepare_text(raw_text) + + 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}") + + 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}") + + 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)") + + 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 + 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"] + 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") + 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) + 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, + "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..c1c93286 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/tau.yaml @@ -0,0 +1,36 @@ +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 + VAL: /data/datasets/cobol/prepared/val.jsonl + VAL_OVERLAP_LINES: "8" + 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/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..5b5ae572 --- /dev/null +++ b/examples/cobol-cpt-grpo/05-rl-prompts/test_prepare_rl_prompts.py @@ -0,0 +1,146 @@ +# 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 subprocess +from unittest import mock + +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_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)]) + 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) + + +@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() 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..0b087817 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/driver.py @@ -0,0 +1,204 @@ +# 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 contextlib +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: + 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) + 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..ce6f4719 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/test_driver.py @@ -0,0 +1,275 @@ +# 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, 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) +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((host, 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((ray.util.get_node_ip_address(), 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), + ray.util.get_node_ip_address(), + ] + + +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..d102addf --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/test_reward.py @@ -0,0 +1,893 @@ +# 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 json +import os +import subprocess +import sys +import tempfile +import types +import unittest +import unittest.mock as mock +from concurrent.futures import ThreadPoolExecutor + + +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 + TauGRPOMetricsCallback, + _compilation_reward, + _count_cobol_verbs, + _count_noop_verbs, + _has_non_display_verb, + _parse_cobc, + _repeated_statement_ratio, + _strip_non_code, + _structural_reward, + _try_compile, + compile_format, + load_prompts, + reward_fn, +) + + +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) + + +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 new file mode 100755 index 00000000..f79ff926 --- /dev/null +++ b/examples/cobol-cpt-grpo/06-grpo/train_grpo.py @@ -0,0 +1,887 @@ +#!/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") + +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() + 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, 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, fmt) + 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, compile_format(rec.get("compile_fmt")))) + + if max_body_tokens and tokenizer: + before = len(candidates) + 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: + print( + f" Body token filter: {before - len(candidates)} dropped " + f"(>{max_body_tokens} tokens)", + flush=True, + ) + + if not preflight: + return [{"prompt": p, "compile_fmt": fmt} for p, _, fmt in candidates] + + prompts = [] + skipped = 0 + 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: + 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" + ) + return prompts + + +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", flag, 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, +) + +_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, fmt = args + full = prompt + completion + 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) + 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, 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)) + for c in completions + ] + else: + n_tokens = [len(c) // 4 for c in completions] + 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)) + + 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) + + +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): + """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 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 + + 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}) on rank " + f"{GLOBAL_RANK}. Stopping.", + flush=True, + ) + 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( + f"ALARM: Mean verb count < 3 for 20 steps on rank {GLOBAL_RANK}. " + "Reward hacking. Stopping.", + flush=True, + ) + stop = True + else: + self.short_alarm_count = 0 + return stop + + +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) + 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 = [ + 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(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) + 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..402d311a --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/run_eval_vllm.py @@ -0,0 +1,778 @@ +#!/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. + +Before loading the model, main() fails unless every suite in REQUIRED_SUITES +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, and with an empty generation, in +place of the generation. +""" + +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")) +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") + +os.makedirs(OUT, exist_ok=True) +os.environ["HF_HOME"] = HF_HOME + + +REQUIRED_SUITES = ( + "perplexity", + "paragraph_completion", + "program_completion", + "structural_validation", + "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, + "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: + 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 + 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( + "Benchmark suites are incomplete; run 03-benchmarks first: " + + "; ".join(problems) + ) + return suites + + +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) + 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, + ) + + 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, + errors="replace", + 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 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 + empty_compiled = 0 + fmt_counts = {"free": 0, "fixed": 0} + results = [] + + for task, gen_text, ref_text in zip(tasks, generated_texts, reference_texts): + suffix = task.get("suffix", "") + success, errors, fmt_used = try_compile(task["prompt"] + gen_text + suffix) + ref_success, _, _ = try_compile(ref_text) + 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 + 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, + "empty_compiled": empty_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, + ) + print( + f" Reference compiles: {reference_compiled}/{total}; generation compiles " + 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), + "compiled": compiled, + "failed": failed, + "compiled_free": fmt_counts["free"], + "compiled_fixed": fmt_counts["fixed"], + "reference_compiled": reference_compiled, + "compiled_where_reference_compiles": compiled_where_reference_compiles, + "empty_compiled": empty_compiled, + "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) + + print(f" Compiler: {check_compiler()}", flush=True) + suites = load_required_suites(EVAL_DIR) + from transformers import AutoTokenizer + + 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: + 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", + max_model_len=MAX_MODEL_LEN, + gpu_memory_utilization=0.90, + download_dir=HF_HOME, + distributed_executor_backend="mp", + ) + print(" vLLM model loaded.", flush=True) + + para_tasks = suites["paragraph_completion"] + if para_tasks: + print( + f"\n--- Paragraph completion ({len(para_tasks)} tasks, max_tokens={MAX_NEW_TOKENS}) ---", + flush=True, + ) + 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"], "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" + } + 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 = suites["program_completion"] + 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, + [t["reference_text"] for t in prog_tasks], + ) + 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 = suites["structural_validation"] + 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, + [t["reference_text"] for t in strv_tasks], + ) + 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") + + 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) + 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) + print(f" Copybook BLEU: {all_results['copybook_usage']['avg_bleu']}", 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/07-eval/test_run_eval.py b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py new file mode 100644 index 00000000..cb7ce9ec --- /dev/null +++ b/examples/cobol-cpt-grpo/07-eval/test_run_eval.py @@ -0,0 +1,285 @@ +# 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, 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 +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=(), 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: + 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): + 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_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" + + +@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"], ["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_and_empty(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", + "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): + 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): + 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/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..d3a4ca66 --- /dev/null +++ b/examples/cobol-cpt-grpo/README.md @@ -0,0 +1,316 @@ +# 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`. 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. | + +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. 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/`. + +## 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). 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` 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. 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). 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. 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 +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 +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. `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 +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 +if a suite has no task left. + +## 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 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_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. 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 + +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, 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. +- 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). + +| Metric | Base | CPT | GRPO | +| --- | --- | --- | --- | +| COBOL perplexity | 1.90 | 1.54 | 1.54 | +| 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% | + +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. diff --git a/examples/cobol-cpt-grpo/image-requirements.txt b/examples/cobol-cpt-grpo/image-requirements.txt new file mode 100644 index 00000000..3e4f6825 --- /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.9 +wrapt==2.5.0