feat: copy scripts to output dir#46
Conversation
There was a problem hiding this comment.
Pull request overview
This PR aims to make TRL containerized runs more reproducible by “freezing” the code into the run output directory and executing from that frozen copy, while also normalizing output paths so a frozen config can re-derive the same run directory from any working directory.
Changes:
- Resolve
paths.output_base/paths.debug_baseto absolute paths during run directory setup and persist them back onto the config. - Add guardrail output indicating when a run directory already contains frozen
src/post_trainingand/orscripts. - Update the TRL container SLURM template to copy
src/post_trainingandscriptsintorun_dirand run withPYTHONPATH/cwdpointing atrun_dir.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
src/post_training/utils/paths.py |
Resolves output/debug base paths to absolute and persists them into the config before computing run_dir. |
src/post_training/utils/guardrails.py |
Adds a warning row when frozen source directories already exist under the run directory (TRL container path). |
src/post_training/slurm/job_trl_container.sh.jinja |
Adds a runtime copy step to “freeze” code into run_dir, and switches execution to use run_dir for PYTHONPATH and working directory. |
scripts/submit.py |
Ensures the frozen config path passed into SLURM generation/submission is an absolute resolved path. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| 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 |
| # ── 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" |
| 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, | ||
| ) |
| # ── 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 |
| frozen_src = run_dir / "src" / "post_training" | ||
| frozen_scripts = run_dir / "scripts" | ||
|
|
||
| if frozen_src.exists(): |
There was a problem hiding this comment.
I think this is still relevant to what copilot was suggesting for the stricter guardrails.
For example- if there already exists a incomplete/corrupted frozen_src then it would create problems because this is just a directory level match. And if theres a file inside which is outdated- it will not be overwritten and the rest of the copy will be skipped as well.
There was a problem hiding this comment.
And if theres a file inside which is outdated- it will not be overwritten
That's the point. Changes in the post-training codebase after the submission of the job should not be reflected in future (continued) runs. The source code should not change between runs.
The output directory is either specified in the config file, or generated on the fly using the timestamp (so every submission gets a unique output directory). The recommended approach is to have a fixed output directory for a given run.
There was a problem hiding this comment.
Maybe I will explain a bit better:
Two cases I can think of. First, if we set a run_name in the config for an experiment and a stale copy already exists for that folder, the copy is completely skipped - not just the outdated files, the whole thing. Second, if a copy failed halfway (system issue, process killed, etc) then the next submit also skips over whatever folder is already there. So you end up with a stale or partial frozen copy and nothing tells you.
The part that makes this worse than it looks- since this PR the job cds into the run dir and launches from there:
cd "{{ run_dir }}"
accelerate launch ... scripts/train.py --config {{ config_path }}So its not just that the frozen copy is a bad record of the run, its that the run silently executes the stale code. I tested this locally- with an old example.py sitting in run_dir/src/post_training, post_freeze keeps the old one and doesn't copy any of the new files in either.
Maybe we could copy to a temp dir and swap it in? That way it refreshes the content instead of just checking if the folder is there, so it covers both the stale and the half-written case:
tmp = run_dir / ".src.tmp"
shutil.copytree(repo_dir / "src" / "post_training", tmp)
if frozen_src.exists():
shutil.rmtree(frozen_src)
os.replace(tmp, frozen_src)I also thought about always applying a timestamp to the run name so we never reuse a folder, but I dont think it can go in generate_run_name as train.py calls setup_run_directory too at job runtime, so it would derive a different run_dir than the one submit.py copied into. Would have to be in submit.py, and then resubmitting a frozen config would keep making new dirs. So refreshing the copy seems like the cleaner fix to me.
There was a problem hiding this comment.
Maybe we could copy to a temp dir and swap it in?
Good suggestion. Went with that.
|
I like the slurmpilot like idea here. I left a comment. Apart from that in this PR, I checked functionality and its working well! |
Summary
The PR copies the
scriptsandsrc/post_trainingto the output directory and consumes those during the run. This has two benefits:Guardrails have also been added to warn the user if the run directory already contains
scriptsandsrc/post_training. In that case, the newer code is not copied over — the existing copies are left in place. This is based on the assumption that a given run should continue using the same codebase it started with.Type of change
Validation
Validated on LUMI with a short run with a checkpointing frequency of 1.