From af1321e03b73968429802a1a409656114e3616d7 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Tue, 7 Jul 2026 16:32:26 +0200 Subject: [PATCH 1/8] feat: copy scripts to output dir --- src/post_training/slurm/job_trl_container.sh.jinja | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 06ad9cc..a51d650 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -71,6 +71,12 @@ handle_signal() { } trap handle_signal USR1 +# ── Freeze source code for this run ─────────────────────────────────────── +echo "Copying src/post_training and scripts to run directory..." +mkdir -p "{{ run_dir }}/src" +cp -r "{{ repo_dir }}/src/post_training" "{{ run_dir }}/src/" +cp -r "{{ repo_dir }}/scripts" "{{ run_dir }}/scripts" + # ── Launch inside Singularity container ─────────────────────────────────── srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ singularity exec --nv \ @@ -81,7 +87,7 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ # Prevent host Python/PATH from interfering with container export PATH=\"{{ container_path }}\" - export PYTHONPATH=\"{{ repo_dir }}/src\" + export PYTHONPATH=\"{{ run_dir }}/src\" export PYTHONNOUSERSITE=1 export NODE_RANK=\"\$SLURM_NODEID\" @@ -95,7 +101,7 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-${HF_HOME}/datasets} export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt - cd {{ repo_dir }} + cd {{ run_dir }} export WANDB_DIR={{ run_dir }} From 68eb4f2e5960c0f396a6d69dccdd2ad7e6028f3a Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 14:34:12 +0200 Subject: [PATCH 2/8] fix: fix path to config.yaml --- scripts/submit.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/submit.py b/scripts/submit.py index 9d201bd..c6edd17 100644 --- a/scripts/submit.py +++ b/scripts/submit.py @@ -147,7 +147,7 @@ def main() -> None: get_backend(config.backend).post_freeze(config, run_dir) - job_id = generate_and_submit(config, run_dir, str(frozen), tokenize_only=tokenize_only) + job_id = generate_and_submit(config, run_dir, "config.yaml", tokenize_only=tokenize_only) logger.info("SLURM job submitted: %s", job_id) From fbee0ae22b27eea54893f52ea6eef5ff4eddac2b Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 16:56:55 +0200 Subject: [PATCH 3/8] fix: fix nested output directories --- scripts/submit.py | 4 +++- src/post_training/utils/paths.py | 9 +++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/scripts/submit.py b/scripts/submit.py index c6edd17..6a31ea9 100644 --- a/scripts/submit.py +++ b/scripts/submit.py @@ -147,7 +147,9 @@ def main() -> None: get_backend(config.backend).post_freeze(config, run_dir) - job_id = generate_and_submit(config, run_dir, "config.yaml", tokenize_only=tokenize_only) + job_id = generate_and_submit( + config, run_dir, str(frozen.resolve()), tokenize_only=tokenize_only + ) logger.info("SLURM job submitted: %s", job_id) diff --git a/src/post_training/utils/paths.py b/src/post_training/utils/paths.py index 9a0277f..5ae985b 100644 --- a/src/post_training/utils/paths.py +++ b/src/post_training/utils/paths.py @@ -51,10 +51,15 @@ def setup_run_directory(config: PostTrainingConfig, allow_override: bool = False """ run_name = generate_run_name(config) + # Resolve to an absolute path and persist it back onto the config so that + # re-deriving the run directory from a frozen config.yaml later (e.g. from + # scripts/train.py, possibly running from a different cwd) is idempotent. if config.debug.enabled: - base = Path(config.paths.debug_base) + base = Path(config.paths.debug_base).resolve() + config.paths.debug_base = str(base) else: - base = Path(config.paths.output_base) + base = Path(config.paths.output_base).resolve() + config.paths.output_base = str(base) run_dir = base / run_name From 2f59af0c49ff94929a0e656b0e6f28830333fbcf Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 16:58:22 +0200 Subject: [PATCH 4/8] feat: warn user when scripts already present in run directory --- src/post_training/slurm/job_trl_container.sh.jinja | 14 ++++++++++---- src/post_training/utils/guardrails.py | 12 ++++++++++++ 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index a51d650..ff21ae1 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -72,10 +72,16 @@ handle_signal() { trap handle_signal USR1 # ── Freeze source code for this run ─────────────────────────────────────── -echo "Copying src/post_training and scripts to run directory..." -mkdir -p "{{ run_dir }}/src" -cp -r "{{ repo_dir }}/src/post_training" "{{ run_dir }}/src/" -cp -r "{{ repo_dir }}/scripts" "{{ run_dir }}/scripts" +if [ -d "{{ run_dir }}/src/post_training" ] || [ -d "{{ run_dir }}/scripts" ]; then + echo "WARNING: {{ run_dir }}/src/post_training or {{ run_dir }}/scripts already exists." + echo "WARNING: Skipping copy — keeping the source frozen at original submission time." + echo "WARNING: Delete these directories manually if you intentionally want to refresh the frozen source." +else + echo "Copying src/post_training and scripts to run directory..." + mkdir -p "{{ run_dir }}/src" + cp -r "{{ repo_dir }}/src/post_training" "{{ run_dir }}/src/" + cp -r "{{ repo_dir }}/scripts" "{{ run_dir }}/scripts" +fi # ── Launch inside Singularity container ─────────────────────────────────── srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py index a8327cd..fe639fc 100644 --- a/src/post_training/utils/guardrails.py +++ b/src/post_training/utils/guardrails.py @@ -252,6 +252,18 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo # ------------------------------------------------------------------ _section("Output") _row("Run directory", str(run_dir)) + if config.backend == "trl" and config.container and config.container.image: + frozen_src_exists = (run_dir / "src" / "post_training").exists() + frozen_scripts_exists = (run_dir / "scripts").exists() + if frozen_src_exists or frozen_scripts_exists: + _row( + "Frozen source", + _yellow( + "*** src/post_training and/or scripts/ already exist in the run " + "directory — they will NOT be replaced ***" + ), + warn=True, + ) _row("Logging", ", ".join(config.logging.report_to)) if "wandb" in config.logging.report_to: _row("WandB project", config.logging.wandb_project) From 545ef1d900b83782241f546a59504e0dde3ac596 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Thu, 9 Jul 2026 18:32:12 +0200 Subject: [PATCH 5/8] fix: copy files at submission time --- src/post_training/backend.py | 29 ++++++++++++++++++- .../slurm/job_trl_container.sh.jinja | 16 ++-------- src/post_training/slurm/launcher.py | 1 - src/post_training/utils/guardrails.py | 14 +++++++-- 4 files changed, 42 insertions(+), 18 deletions(-) diff --git a/src/post_training/backend.py b/src/post_training/backend.py index 0f883eb..fac8427 100644 --- a/src/post_training/backend.py +++ b/src/post_training/backend.py @@ -4,8 +4,10 @@ import hashlib import json +import logging import math import re +import shutil from abc import ABC, abstractmethod from pathlib import Path from typing import TYPE_CHECKING @@ -15,6 +17,8 @@ if TYPE_CHECKING: from post_training.config import PostTrainingConfig +logger = logging.getLogger(__name__) + # --------------------------------------------------------------------------- # Interface @@ -188,7 +192,30 @@ def render_slurm_script(self, config, run_dir, frozen_config_path, *, tokenize_o ) def post_freeze(self, config, run_dir): - pass + """Freeze src/post_training and scripts/ into run_dir for containerized runs. + + Done here (at submission time) rather than inside the SLURM job script, + so a job that sits queued isn't affected by later repo changes. Each + directory is guarded independently so a partial freeze left behind by + an interrupted previous attempt is completed rather than left broken. + """ + if config.container is None or not config.container.image: + return + + repo_dir = Path.cwd() + frozen_src = run_dir / "src" / "post_training" + frozen_scripts = run_dir / "scripts" + + if frozen_src.exists(): + logger.warning("%s already exists — leaving frozen source in place.", frozen_src) + else: + frozen_src.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(repo_dir / "src" / "post_training", frozen_src) + + if frozen_scripts.exists(): + logger.warning("%s already exists — leaving frozen scripts in place.", frozen_scripts) + else: + shutil.copytree(repo_dir / "scripts", frozen_scripts) # --------------------------------------------------------------------------- diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index ff21ae1..99b8b58 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -71,18 +71,6 @@ handle_signal() { } trap handle_signal USR1 -# ── Freeze source code for this run ─────────────────────────────────────── -if [ -d "{{ run_dir }}/src/post_training" ] || [ -d "{{ run_dir }}/scripts" ]; then - echo "WARNING: {{ run_dir }}/src/post_training or {{ run_dir }}/scripts already exists." - echo "WARNING: Skipping copy — keeping the source frozen at original submission time." - echo "WARNING: Delete these directories manually if you intentionally want to refresh the frozen source." -else - echo "Copying src/post_training and scripts to run directory..." - mkdir -p "{{ run_dir }}/src" - cp -r "{{ repo_dir }}/src/post_training" "{{ run_dir }}/src/" - cp -r "{{ repo_dir }}/scripts" "{{ run_dir }}/scripts" -fi - # ── Launch inside Singularity container ─────────────────────────────────── srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ singularity exec --nv \ @@ -107,9 +95,9 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-${HF_HOME}/datasets} export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt - cd {{ run_dir }} + cd "{{ run_dir }}" - export WANDB_DIR={{ run_dir }} + export WANDB_DIR="{{ run_dir }}" accelerate launch \ --num_machines \$NNODES \ diff --git a/src/post_training/slurm/launcher.py b/src/post_training/slurm/launcher.py index 1c88e1e..000bce4 100644 --- a/src/post_training/slurm/launcher.py +++ b/src/post_training/slurm/launcher.py @@ -143,7 +143,6 @@ def render_trl_container_slurm_script( bind_mounts=config.container.bind_mounts, env_file=config.container.env_file, container_path=config.container.path, - repo_dir=str(Path.cwd()), ) slurm_dir = run_dir / "slurm" diff --git a/src/post_training/utils/guardrails.py b/src/post_training/utils/guardrails.py index fe639fc..1977685 100644 --- a/src/post_training/utils/guardrails.py +++ b/src/post_training/utils/guardrails.py @@ -255,15 +255,25 @@ def run_guardrails(config: PostTrainingConfig, run_dir: Path, tokenize_only: boo if config.backend == "trl" and config.container and config.container.image: frozen_src_exists = (run_dir / "src" / "post_training").exists() frozen_scripts_exists = (run_dir / "scripts").exists() - if frozen_src_exists or frozen_scripts_exists: + if frozen_src_exists and frozen_scripts_exists: _row( "Frozen source", _yellow( - "*** src/post_training and/or scripts/ already exist in the run " + "*** src/post_training and scripts/ already exist in the run " "directory — they will NOT be replaced ***" ), warn=True, ) + elif frozen_src_exists or frozen_scripts_exists: + missing = "scripts/" if frozen_src_exists else "src/post_training" + _row( + "Frozen source", + _yellow( + f"*** partially frozen — {missing} is missing and will be copied; " + "the existing copy will NOT be replaced ***" + ), + warn=True, + ) _row("Logging", ", ".join(config.logging.report_to)) if "wandb" in config.logging.report_to: _row("WandB project", config.logging.wandb_project) From 50811ae0d7982c5146cf914d791bc6669d0316c4 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 15 Jul 2026 16:39:17 +0200 Subject: [PATCH 6/8] fix(slurm): escape run_dir quotes for cd/WANDB_DIR in TRL container template --- src/post_training/slurm/job_trl_container.sh.jinja | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 99b8b58..86699b0 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -95,9 +95,9 @@ srun --export=ALL --wait=60 --kill-on-bad-exit=1 \ export HF_DATASETS_CACHE=${HF_DATASETS_CACHE:-${HF_HOME}/datasets} export SSL_CERT_FILE=/etc/ssl/certs/ca-certificates.crt - cd "{{ run_dir }}" + cd \"{{ run_dir }}\" - export WANDB_DIR="{{ run_dir }}" + export WANDB_DIR=\"{{ run_dir }}\" accelerate launch \ --num_machines \$NNODES \ From b08b6a983574325a4c865444fee3a9acbcc2e5bd Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 15 Jul 2026 16:39:24 +0200 Subject: [PATCH 7/8] fix(backend): atomically swap frozen source into place --- src/post_training/backend.py | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/src/post_training/backend.py b/src/post_training/backend.py index fac8427..1ee6452 100644 --- a/src/post_training/backend.py +++ b/src/post_training/backend.py @@ -6,6 +6,7 @@ import json import logging import math +import os import re import shutil from abc import ABC, abstractmethod @@ -196,26 +197,29 @@ def post_freeze(self, config, run_dir): Done here (at submission time) rather than inside the SLURM job script, so a job that sits queued isn't affected by later repo changes. Each - directory is guarded independently so a partial freeze left behind by - an interrupted previous attempt is completed rather than left broken. + directory is copied into a temp path and atomically swapped into place, + so an interrupted previous attempt can never leave a partial freeze that + gets mistaken for a completed one. """ if config.container is None or not config.container.image: return repo_dir = Path.cwd() - frozen_src = run_dir / "src" / "post_training" - frozen_scripts = run_dir / "scripts" + self._freeze_dir(repo_dir / "src" / "post_training", run_dir / "src" / "post_training") + self._freeze_dir(repo_dir / "scripts", run_dir / "scripts") - if frozen_src.exists(): - logger.warning("%s already exists — leaving frozen source in place.", frozen_src) - else: - frozen_src.parent.mkdir(parents=True, exist_ok=True) - shutil.copytree(repo_dir / "src" / "post_training", frozen_src) + @staticmethod + def _freeze_dir(src: Path, dst: Path) -> None: + if dst.exists(): + logger.warning("%s already exists — leaving frozen source in place.", dst) + return - if frozen_scripts.exists(): - logger.warning("%s already exists — leaving frozen scripts in place.", frozen_scripts) - else: - shutil.copytree(repo_dir / "scripts", frozen_scripts) + tmp = dst.with_name(dst.name + ".tmp") + if tmp.exists(): + shutil.rmtree(tmp) + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copytree(src, tmp) + os.replace(tmp, dst) # --------------------------------------------------------------------------- From b8ae877dfa7b7abcc932fb60b63ddc861c7fe640 Mon Sep 17 00:00:00 2001 From: Arjun Krishnakumar Date: Wed, 15 Jul 2026 16:39:29 +0200 Subject: [PATCH 8/8] test(backend): add regression test for interrupted freeze copy --- tests/test_backend.py | 125 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 125 insertions(+) create mode 100644 tests/test_backend.py diff --git a/tests/test_backend.py b/tests/test_backend.py new file mode 100644 index 0000000..eebe4d3 --- /dev/null +++ b/tests/test_backend.py @@ -0,0 +1,125 @@ +"""Tests for TRLBackend.post_freeze() — the submission-time source freeze.""" + +import logging + +from post_training.backend import TRLBackend +from post_training.config import ContainerConfig, PostTrainingConfig + + +def _make_fake_repo(repo_dir, marker: str = "live") -> None: + (repo_dir / "src" / "post_training").mkdir(parents=True) + (repo_dir / "src" / "post_training" / "marker.txt").write_text(marker) + (repo_dir / "scripts").mkdir(parents=True) + (repo_dir / "scripts" / "marker.txt").write_text(marker) + + +def _config_with_container() -> PostTrainingConfig: + cfg = PostTrainingConfig() + cfg.container = ContainerConfig( + image="/shared/containers/trl.sif", + bind_mounts=["/scratch:/scratch"], + env_file="/shared/env/cluster.env", + ) + return cfg + + +def test_post_freeze_noop_without_container(tmp_path, monkeypatch): + """No container configured — nothing is copied, run_dir stays untouched.""" + repo_dir = tmp_path / "repo" + _make_fake_repo(repo_dir) + monkeypatch.chdir(repo_dir) + + cfg = PostTrainingConfig() + cfg.container = None + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + TRLBackend().post_freeze(cfg, run_dir) + + assert not (run_dir / "src").exists() + assert not (run_dir / "scripts").exists() + + +def test_post_freeze_copies_when_fresh(tmp_path, monkeypatch): + """Neither directory exists yet — both get copied from the live repo.""" + repo_dir = tmp_path / "repo" + _make_fake_repo(repo_dir, marker="live") + monkeypatch.chdir(repo_dir) + + run_dir = tmp_path / "outputs" / "my-run" + run_dir.mkdir(parents=True) + + TRLBackend().post_freeze(_config_with_container(), run_dir) + + assert (run_dir / "src" / "post_training" / "marker.txt").read_text() == "live" + assert (run_dir / "scripts" / "marker.txt").read_text() == "live" + + +def test_post_freeze_skips_when_fully_frozen(tmp_path, monkeypatch, caplog): + """Both already frozen — existing copies are preserved, not overwritten.""" + repo_dir = tmp_path / "repo" + _make_fake_repo(repo_dir, marker="live-later") + monkeypatch.chdir(repo_dir) + + run_dir = tmp_path / "outputs" / "my-run" + (run_dir / "src" / "post_training").mkdir(parents=True) + (run_dir / "src" / "post_training" / "marker.txt").write_text("frozen") + (run_dir / "scripts").mkdir(parents=True) + (run_dir / "scripts" / "marker.txt").write_text("frozen") + + with caplog.at_level(logging.WARNING): + TRLBackend().post_freeze(_config_with_container(), run_dir) + + assert (run_dir / "src" / "post_training" / "marker.txt").read_text() == "frozen" + assert (run_dir / "scripts" / "marker.txt").read_text() == "frozen" + assert "already exists" in caplog.text + + +def test_post_freeze_completes_partial_state(tmp_path, monkeypatch, caplog): + """Only one directory exists (interrupted prior attempt) — the missing one + is copied to complete the freeze; the existing one is left untouched.""" + repo_dir = tmp_path / "repo" + _make_fake_repo(repo_dir, marker="live") + monkeypatch.chdir(repo_dir) + + run_dir = tmp_path / "outputs" / "my-run" + (run_dir / "scripts").mkdir(parents=True) + (run_dir / "scripts" / "marker.txt").write_text("frozen") + # src/post_training deliberately left missing. + + with caplog.at_level(logging.WARNING): + TRLBackend().post_freeze(_config_with_container(), run_dir) + + assert (run_dir / "src" / "post_training" / "marker.txt").read_text() == "live" + assert (run_dir / "scripts" / "marker.txt").read_text() == "frozen" + assert "already exists" in caplog.text + + +def test_post_freeze_recovers_from_interrupted_copy(tmp_path, monkeypatch): + """A killed process mid-copytree leaves only a .tmp dir, not a corrupted + frozen dir — the next post_freeze must redo the copy in full rather than + silently adopting the half-written content. + + Regression test for the bug reported on PR #46: since post_freeze's + existence check was the only signal used to decide "already frozen, + leave it", a directory that existed but was only partially copied got + permanently treated as a completed freeze, and the job silently ran + with incomplete/stale code (`cd`'d into run_dir and imported from there). + """ + repo_dir = tmp_path / "repo" + _make_fake_repo(repo_dir, marker="live") + monkeypatch.chdir(repo_dir) + + run_dir = tmp_path / "outputs" / "my-run" + frozen_src = run_dir / "src" / "post_training" + tmp_src = frozen_src.with_name(frozen_src.name + ".tmp") + tmp_src.mkdir(parents=True) + (tmp_src / "half-written.py").write_text("partial") + # frozen_src itself deliberately does not exist yet — simulates the + # process dying after shutil.copytree started writing but before it finished. + + TRLBackend().post_freeze(_config_with_container(), run_dir) + + assert (frozen_src / "marker.txt").read_text() == "live" + assert not (frozen_src / "half-written.py").exists() + assert not tmp_src.exists()