diff --git a/scripts/submit.py b/scripts/submit.py index 9d201bd..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, str(frozen), 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/backend.py b/src/post_training/backend.py index 0f883eb..1ee6452 100644 --- a/src/post_training/backend.py +++ b/src/post_training/backend.py @@ -4,8 +4,11 @@ import hashlib import json +import logging import math +import os import re +import shutil from abc import ABC, abstractmethod from pathlib import Path from typing import TYPE_CHECKING @@ -15,6 +18,8 @@ if TYPE_CHECKING: from post_training.config import PostTrainingConfig +logger = logging.getLogger(__name__) + # --------------------------------------------------------------------------- # Interface @@ -188,7 +193,33 @@ 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 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() + self._freeze_dir(repo_dir / "src" / "post_training", run_dir / "src" / "post_training") + self._freeze_dir(repo_dir / "scripts", run_dir / "scripts") + + @staticmethod + def _freeze_dir(src: Path, dst: Path) -> None: + if dst.exists(): + logger.warning("%s already exists — leaving frozen source in place.", dst) + return + + 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) # --------------------------------------------------------------------------- diff --git a/src/post_training/slurm/job_trl_container.sh.jinja b/src/post_training/slurm/job_trl_container.sh.jinja index 06ad9cc..86699b0 100644 --- a/src/post_training/slurm/job_trl_container.sh.jinja +++ b/src/post_training/slurm/job_trl_container.sh.jinja @@ -81,7 +81,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,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 {{ repo_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 a8327cd..1977685 100644 --- a/src/post_training/utils/guardrails.py +++ b/src/post_training/utils/guardrails.py @@ -252,6 +252,28 @@ 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 and frozen_scripts_exists: + _row( + "Frozen source", + _yellow( + "*** 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) 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 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()