SLURM

The live SLURM surface is gx exp submit <experiment.yml>.

Implementation lives in graphids.exp.slurm and is exposed through graphids.cli.exp.submit. It validates an ExperimentConfig, renders a Ray allocation sbatch script under the configured SLURM log directory, and submits it with sbatch.

Useful commands:

gx exp submit configs/experiments/gat_snapshot_sequence_real.yml -C pitzer
gx exp submit configs/experiments/gat_snapshot_sequence_real.yml -C pitzer --dry-run

graphids.exp.slurm

slurm

SLURM submission for experiment YAML configs.

build_slurm_script

build_slurm_script(cfg: ExperimentConfig, yaml_path: str | Path, *, cluster: str | None = None, partition: str | None = None, time_limit: str | None = None, gres: str | None = None, nodes: int = 1, ray_port: int = 6379) -> tuple[Path, str]

Build a SLURM script that starts Ray inside one allocation.

Source code in graphids/exp/slurm.py

def build_slurm_script(
    cfg: ExperimentConfig,
    yaml_path: str | Path,
    *,
    cluster: str | None = None,
    partition: str | None = None,
    time_limit: str | None = None,
    gres: str | None = None,
    nodes: int = 1,
    ray_port: int = 6379,
) -> tuple[Path, str]:
    """Build a SLURM script that starts Ray inside one allocation."""
    yaml_abs = Path(yaml_path).resolve()
    if not yaml_abs.is_file():
        raise FileNotFoundError(f"experiment YAML not found: {yaml_abs}")

    job_name = _slug(f"ray-{cfg.experiment_name}")
    log_dir = _log_dir()
    script_dir = _script_dir()
    script_path = script_dir / f"{job_name}.sbatch"
    stdout_path = log_dir / f"{job_name}_%j.out"
    stderr_path = log_dir / f"{job_name}_%j.err"
    directives = _sbatch_directives(
        cfg,
        job_name=job_name,
        cluster=cluster,
        partition=partition,
        time_limit=time_limit,
        gres=gres,
    )
    directives[0] = f"#SBATCH --nodes={max(1, int(nodes))}"
    directives[1] = "#SBATCH --ntasks-per-node=1"
    directives.extend(
        [
            f"#SBATCH --output={stdout_path}",
            f"#SBATCH --error={stderr_path}",
        ]
    )
    gpus = _gpus(cfg)
    cpus = max(1, int(cfg.resources.cpus_per_worker))
    script = "\n".join(
        [
            "#!/usr/bin/env bash",
            *directives,
            "",
            "set -euo pipefail",
            f"cd {PROJECT_ROOT}",
            "source scripts/slurm/_preamble.sh",
            "",
            f"export SLURM_CPUS_PER_TASK={cpus}",
            "RAY_BASE_PORT=${GRAPHIDS_RAY_BASE_PORT:-$((20000 + (SLURM_JOB_ID % 200) * 200))}",
            "RAY_PORT=${GRAPHIDS_RAY_PORT:-${RAY_BASE_PORT}}",
            'RAY_OBJECT_MANAGER_PORT="$((RAY_BASE_PORT + 1))"',
            'RAY_NODE_MANAGER_PORT="$((RAY_BASE_PORT + 2))"',
            'RAY_DASHBOARD_AGENT_PORT="$((RAY_BASE_PORT + 4))"',
            'RAY_DASHBOARD_AGENT_GRPC_PORT="$((RAY_BASE_PORT + 5))"',
            'RAY_RUNTIME_ENV_AGENT_PORT="$((RAY_BASE_PORT + 6))"',
            'RAY_METRICS_EXPORT_PORT="$((RAY_BASE_PORT + 7))"',
            'RAY_MIN_WORKER_PORT="$((RAY_BASE_PORT + 20))"',
            'RAY_MAX_WORKER_PORT="$((RAY_BASE_PORT + 119))"',
            f"RAY_LOG_DIR={log_dir}/ray_${{SLURM_JOB_ID}}",
            'RAY_TMP_DIR="${TMPDIR:-/tmp}/graphids-ray-${SLURM_JOB_ID}"',
            'mkdir -p "${RAY_LOG_DIR}" "${RAY_TMP_DIR}"',
            "RAY_START_PIDS=()",
            'mapfile -t RAY_NODES < <(scontrol show hostnames "${SLURM_JOB_NODELIST}")',
            'HEAD_NODE="${RAY_NODES[0]}"',
            'HEAD_IP=$(srun --nodes=1 --ntasks=1 -w "${HEAD_NODE}" hostname --ip-address | awk \'{print $1}\')',
            'RAY_ADDRESS="${HEAD_IP}:${RAY_PORT}"',
            "export RAY_ADDRESS",
            "cleanup_ray() {",
            '  if ((${#RAY_START_PIDS[@]})); then',
            '    kill "${RAY_START_PIDS[@]}" >/dev/null 2>&1 || true',
            '    wait "${RAY_START_PIDS[@]}" >/dev/null 2>&1 || true',
            "  fi",
            '  rm -rf "${RAY_TMP_DIR}"',
            "}",
            "trap cleanup_ray EXIT",
            "",
            'mkdir -p "${RAY_LOG_DIR}" "${RAY_TMP_DIR}"',
            'srun --nodes=1 --ntasks=1 -w "${HEAD_NODE}" \\',
            '  ray start --head --node-ip-address="${HEAD_IP}" --port="${RAY_PORT}" \\',
            '    --include-dashboard=false \\',
            '    --object-manager-port="${RAY_OBJECT_MANAGER_PORT}" \\',
            '    --node-manager-port="${RAY_NODE_MANAGER_PORT}" \\',
            '    --dashboard-agent-listen-port="${RAY_DASHBOARD_AGENT_PORT}" \\',
            '    --dashboard-agent-grpc-port="${RAY_DASHBOARD_AGENT_GRPC_PORT}" \\',
            '    --runtime-env-agent-port="${RAY_RUNTIME_ENV_AGENT_PORT}" \\',
            '    --metrics-export-port="${RAY_METRICS_EXPORT_PORT}" \\',
            '    --min-worker-port="${RAY_MIN_WORKER_PORT}" \\',
            '    --max-worker-port="${RAY_MAX_WORKER_PORT}" \\',
            f'    --num-cpus={cpus} --num-gpus={gpus} --temp-dir="${{RAY_TMP_DIR}}/head" \\',
            '    --block >"${RAY_LOG_DIR}/head.log" 2>&1 &',
            'RAY_START_PIDS+=("$!")',
            "sleep 10",
            "",
            'for NODE in "${RAY_NODES[@]:1}"; do',
            '  srun --nodes=1 --ntasks=1 -w "${NODE}" \\',
            '    ray start --address="${RAY_ADDRESS}" \\',
            f'      --num-cpus={cpus} --num-gpus={gpus} --temp-dir="${{RAY_TMP_DIR}}/${{NODE}}" \\',
            '      --block >"${RAY_LOG_DIR}/${NODE}.log" 2>&1 &',
            '  RAY_START_PIDS+=("$!")',
            "done",
            "sleep 10",
            "",
            f'python -m graphids exp launch {yaml_abs} --address "${{RAY_ADDRESS}}"',
            "source scripts/slurm/_epilog.sh",
            "",
        ]
    )
    return script_path, script