feat: copy scripts to output dir

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)
 
 

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)
 
 
 # ---------------------------------------------------------------------------

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 \

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"

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)

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
 

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()