Skip to main content

Generate Slurm Job Scripts with AI

Level: Intermediate · Time: 35–45 min · Category: Command line

Tags: Python · OpenAI SDK · Slurm · HPC

Build a CLI that turns a plain-English job description into a draft sbatch script. You will ground the model with your cluster's real constraints so it does not invent partitions or options, then add checks that catch obvious mistakes before you submit.

Always review generated scripts

An AI draft is a starting point, not a submission. Read every #SBATCH line and run a test before submitting a real job. This tutorial never submits jobs for you.

Prerequisites​

Requires an active HPC account on Sol

To actually validate or submit a generated script you need an HPC account on the Sol supercomputer — the sbatch steps cannot be completed without one. HPC accounts require a faculty sponsor — only sponsored users (or faculty themselves) can be approved. Request one through Getting Access and Requesting an Account before you start.

Confirm Python 3 is available and that you can reach the Slurm client on a login node:

python3 --version    # expect 3.9 or newer
squeue --version     # confirms the Slurm client is installed

If python3 is missing, load a Python module (for example module load python). If squeue is not found, connect to a Slurm login node before continuing.

What you will do​

  • Create a Python CLI workspace with the OpenAI SDK.
  • Write a small file of real cluster constraints to ground the model.
  • Ask for the first job-script generator.
  • Add validation that rejects unknown partitions and missing directives.
  • Save a reviewed script for a test submission.

Build it step by step​

1. Create the CLI workspace​

mkdir slurm-script-generator
cd slurm-script-generator
python3 -m venv .venv
source .venv/bin/activate
pip install openai python-dotenv rich
Keep this out of Git

If you track this project with Git, do not commit your virtual environment or any secrets. Create a .gitignore file in the project folder with at least these lines:

.gitignore
.venv/
.env
__pycache__/

.venv/ is large and specific to your machine, and .env holds your API key — neither belongs in a shared repository. New to Git? See Track a New Project with Git.

2. Add API connection values​

Your program needs three values to reach the API: your API key, the gateway URL, and a model ID. You keep them in a small text file named .env so they stay out of your code — and out of Git.

What is a .env file, and where does it go?

A .env ("dot-env") file is a plain text file that stores settings as NAME=value lines, one per line. Programs read it at startup so you never have to paste secrets like API keys directly into your code.

  • The filename is literally .env — a leading dot with nothing before it (not config.env or .env.txt).
  • It goes in your project folder: the same directory you just created, where the .venv lives and where you run your commands.
  • The python-dotenv package loads it with load_dotenv(), which looks for .env in the folder you run the program from.
  • Because it holds a secret, never commit it to Git (see the note below).

Create a new file named .env in your project folder — in VS Code use File → New File, or run code .env in the terminal — then paste these three lines and replace the placeholders:

.env
OPENAI_API_KEY="paste-your-key-here"
OPENAI_BASE_URL="https://openai.rc.asu.edu/v1"
RC_LLM_MODEL="paste-model-id-here"
Where do the values come from?

Get your API key and a valid model ID from the LLM API guide. The OPENAI_BASE_URL above is already correct for the Research Computing gateway.

  • Use the /v1 base URL, not the full /chat/completions URL, for the Python SDK.
  • Use a model ID that appears in the portal model list for your account.
  • Keep .env out of Git: add a line containing .env to your .gitignore (see Track a New Project with Git).

3. Write your cluster constraints​

Grounding is the most important step. Give the model the real, allowed options so it cannot guess. Edit these values to match your account and cluster.

cluster-constraints.md
# Cluster constraints

- Allowed partitions: general, htc, gpu
- Default time limit: 04:00:00, maximum: 7-00:00:00
- GPUs: request with `#SBATCH --gres=gpu:1` on the gpu partition only
- Always include: --job-name, --partition, --time, --nodes, --ntasks, --output
- Load software with `module load <name>` in the script body
- Do not use partitions or options not listed here

4. Ask for the first generator​

Tell the assistant to obey the constraints file exactly.

You are helping me build a beginner Python CLI that drafts Slurm sbatch scripts using an OpenAI-compatible API.

Project goal:
- Command: python generate_job.py "run a 4-hour single-GPU PyTorch training job"
- Read cluster-constraints.md and treat it as the only source of allowed options.
- Send the description plus the constraints to the chat completions API.
- Print a complete sbatch script to the terminal.
- Read OPENAI_API_KEY, OPENAI_BASE_URL, and RC_LLM_MODEL from a .env file.

Constraints:
- The generated script must only use partitions and options listed in cluster-constraints.md.
- Include a comment reminding the user to review before submitting.
- Keep the first version in one file named generate_job.py.

Please:
1. Create generate_job.py.
2. Explain how the constraints are passed to the model.
3. Show an example command and its output.

5. Add validation guardrails​

Do not trust the draft blindly. Ask for local checks that catch the most common mistakes.

Please add validation to generate_job.py before printing the script:

Checks:
- The script starts with #!/bin/bash.
- It contains #SBATCH lines for --job-name, --partition, --time, --nodes, --ntasks, and --output.
- The partition is one of the allowed partitions from cluster-constraints.md.
- --gres=gpu is only present when the partition is gpu.

Behavior:
- If a check fails, print a clear warning listing what is wrong instead of pretending the script is ready.
- Explain how to test the checks with a deliberately wrong request.

6. Save a script to test​

Once a draft passes the checks, save it and test it with Slurm's own dry run.

Please add a --save option that writes the script to a file (for example job.sh).
Keep terminal output unchanged when --save is not used, and print a reminder to run:

    sbatch --test-only job.sh

so I can validate the script without actually queuing a job.

More prompts to try​

  • Add a --gpus N option that the model must respect within the allowed limits.
  • Add an explanation block that describes what each #SBATCH line does.
  • Support an interactive mode that asks follow-up questions when the description is vague.

Troubleshooting​

  • If the model invents a partition, make cluster-constraints.md more explicit and repeat the "only use listed options" rule.
  • If validation is too strict or too loose, ask the assistant to adjust the specific check.
  • If sbatch --test-only reports an error, fix the script; do not submit it until it passes.
  • Never let the tool submit jobs automatically. Keep a human in the loop.

Next steps​