feat: add Experiment.export() to write runnable scripts without submitting

[ko3n1g]

Summary

• Adds Experiment.export(output_dir, exist_ok=False) which writes one script per job to a self-contained directory plus a submit_all.sh launcher — no jobs are submitted
• Each scheduler's _submit_dryrun() now writes its script to executor.experiment_dir when set:
  • LocalExecutor → <task>.sh (executable bash)
  • DockerExecutor → <task>.yaml
  • SkypilotExecutor / SkypilotJobsExecutor → <task>.yaml
  • LeptonExecutor → <task>.sh (executable bash)
  • DGXCloudExecutor → <task>_torchrun_job.sh (was hardcoded, now uses job_name)
• Experiment.export() redirects all executor experiment_dirs to output_dir, runs dryrun to trigger script writing, then generates submit_all.sh with the correct submit command per executor type (sbatch, sky launch, etc.)

Example — SLURM

with run.Experiment("my-exp") as exp:
    exp.add(
        run.Script("scripts/pretrain.sh"),
        executor=SlurmExecutor(account="my_account", partition="gpu", nodes=4, ...),
        name="pretrain",
    )
    exp.add(
        run.Script("scripts/finetune.sh"),
        executor=SlurmExecutor(account="my_account", partition="gpu", nodes=1, ...),
        name="finetune",
    )
    exp.export("/tmp/my_exp_scripts")

Output directory:

/tmp/my_exp_scripts/
├── finetune_sbatch.sh
├── pretrain_sbatch.sh
└── submit_all.sh

pretrain_sbatch.sh

#!/bin/bash
#
# Generated by NeMo Run
# Run with: sbatch --requeue --parsable
#

# Parameters
#SBATCH --account=my_account
#SBATCH --gpus-per-node=8
#SBATCH --job-name=my_account-account.pretrain
#SBATCH --nodes=4
#SBATCH --ntasks-per-node=8
#SBATCH --open-mode=append
#SBATCH --output=/scratch/myuser/nemo_jobs/my-exp/my-exp_.../pretrain/sbatch_my_account-account.pretrain_%j.out
#SBATCH --partition=gpu
#SBATCH --time=04:00:00

set -evx

export PYTHONUNBUFFERED=1
export SLURM_UNBUFFEREDIO=1
export TORCHX_MAX_RETRIES=0

set +e

# setup

nodes=( $( scontrol show hostnames $SLURM_JOB_NODELIST ) )
nodes_array=($nodes)
head_node=${nodes_array[0]}
head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname --ip-address)

# Command 1

srun \
  --output .../pretrain/log-my_account-account.pretrain_%j_${SLURM_RESTART_COUNT:-0}.out \
  --container-image nvcr.io/nvidia/nemo:latest \
  --container-mounts .../pretrain:/nemo_run \
  --container-workdir /nemo_run/code \
  --wait=60 --kill-on-bad-exit=1 \
  bash scripts/pretrain.sh

exitcode=$?

set -e

echo "job exited with code $exitcode"
if [ $exitcode -ne 0 ]; then
    if [ "$TORCHX_MAX_RETRIES" -gt "${SLURM_RESTART_COUNT:-0}" ]; then
        scontrol requeue "$SLURM_JOB_ID"
    fi
    exit $exitcode
fi

submit_all.sh

#!/bin/bash
# Submit all jobs for experiment: my-exp
# Generated by NeMo Run

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

sbatch "$SCRIPT_DIR/pretrain_sbatch.sh"
sbatch "$SCRIPT_DIR/finetune_sbatch.sh"

Test plan

• uv run -- pytest test/run/test_experiment.py -k "export" -v — 4 new tests pass
• uv run -- pytest test/run/test_experiment.py — all 66 tests pass
• ruff check --fix . && ruff format . — clean

🤖 Generated with Claude Code