Conversation
There was a problem hiding this comment.
Pull request overview
This PR aims to support specifying a SLURM account when generating/submitting post-training SLURM job scripts.
Changes:
- Adds conditional
#SBATCH --account={{ account }}directives to the SLURM job script templates. - Adds an
accountfield toSlurmConfig.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
src/post_training/slurm/job.sh.jinja |
Adds optional --account SBATCH directive in the non-container TRL template. |
src/post_training/slurm/job_trl_container.sh.jinja |
Adds optional --account SBATCH directive in the containerized TRL template. |
src/post_training/slurm/job_llamafactory.sh.jinja |
Adds optional --account SBATCH directive in the LlamaFactory template. |
src/post_training/config.py |
Introduces SlurmConfig.account configuration field. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| {% if account %}#SBATCH --account={{ account }} | ||
| {% endif %}#SBATCH --nodes={{ num_nodes }} |
There was a problem hiding this comment.
account is referenced in the template but is not currently provided in the template.render(...) context in src/post_training/slurm/launcher.py (render_*_slurm_script functions). As a result the --account directive will never be emitted (Jinja treats missing vars as Undefined/falsy by default), so this feature won’t work as intended. Pass account=config.slurm.account into the render context for each script renderer (and consider keeping the {% if %} / #SBATCH directives on separate lines for consistency with the other conditional blocks in this template).
| {% if account %}#SBATCH --account={{ account }} | |
| {% endif %}#SBATCH --nodes={{ num_nodes }} | |
| {% if account %} | |
| #SBATCH --account={{ account }} | |
| {% endif %} | |
| #SBATCH --nodes={{ num_nodes }} |
| {% if account %}#SBATCH --account={{ account }} | ||
| {% endif %}#SBATCH --nodes={{ num_nodes }} |
There was a problem hiding this comment.
account is referenced in the template but is not currently provided in the template.render(...) context in src/post_training/slurm/launcher.py. This makes the {% if account %} block effectively dead code, so containerized TRL jobs will never set #SBATCH --account even when configured. Ensure account=config.slurm.account is passed when rendering this template.
| {% if account %}#SBATCH --account={{ account }} | |
| {% endif %}#SBATCH --nodes={{ num_nodes }} | |
| {% if account is defined and account %} | |
| #SBATCH --account={{ account }} | |
| {% endif %} | |
| #SBATCH --nodes={{ num_nodes }} |
| {% if account %}#SBATCH --account={{ account }} | ||
| {% endif %}#SBATCH --nodes={{ num_nodes }} |
There was a problem hiding this comment.
account is referenced in the template but is not currently provided in the template.render(...) context in src/post_training/slurm/launcher.py. Without passing account=config.slurm.account, the #SBATCH --account directive will never be written into the generated script for this backend.
| {% if account %}#SBATCH --account={{ account }} | |
| {% endif %}#SBATCH --nodes={{ num_nodes }} | |
| {% if account is defined and account %} | |
| #SBATCH --account={{ account }} | |
| {% endif %} | |
| #SBATCH --nodes={{ num_nodes }} |
| """SLURM job scheduler parameters.""" | ||
|
|
||
| partition: str = "gpu" | ||
| account: str | None = "OELLM_prod2026" |
There was a problem hiding this comment.
The PR description says the SLURM account is optional, but the new config default sets account to a concrete value ("OELLM_prod2026"). This makes the option effectively enabled by default (once the renderer passes it through) and also hard-codes an environment-specific account into the library default. Consider defaulting to None and letting users set it via YAML/CLI overrides when needed.
| account: str | None = "OELLM_prod2026" | |
| account: str | None = None |
Summary
What does this PR do?
This PR adds an optional account filed to the SLURM job that is submitted.
Type of change