submit

ClusterSyncFailed

Bases: Exception

Raised when syncing with a cluster fails.

Source code in cluv/cli/submit.py

class ClusterSyncFailed(Exception):
    """Raised when syncing with a cluster fails."""

JobSubmissionFailed

Bases: Exception

Raised when a job submission fails.

Source code in cluv/cli/submit.py

class JobSubmissionFailed(Exception):
    """Raised when a job submission fails."""

SubmissionProgress

Bases: Generic[JobSubmission]

Live, mutable tracking of one Submission's progress into becoming a Job.

Tracks a submission from before it's even synced (state="SYNCING", no job yet), through submission (job known, state polled from sacct), to running and, possibly, cancellation.

A single flat list of these -- covering every submission, on every cluster, for one submit() call -- is all a live display needs to render the whole picture, from a plain rich.Live table (see render_job_table below) up to, eventually, a table shared across several concurrent submit/submit_first calls (rich.Live only supports one live region per console, so that would fuse several such lists together instead of replacing this one).

Returned by:

• Reference CLI submit
  • sync_and_submit_jobs_to_cluster
  • wait_for_first_running_job

Used by:

• Reference CLI submit
  • render_job_table
  • run_scancel
  • sync_and_submit_jobs_to_cluster
  • wait_for_first_running_job
  • wait_for_jobs_to_cancel

Source code in cluv/cli/submit.py

@dataclasses.dataclass
class SubmissionProgress(Generic[JobSubmission]):
    """Live, mutable tracking of one `Submission`'s progress into becoming a `Job`.

    Tracks a submission from before it's even synced (``state="SYNCING"``, no `job` yet),
    through submission (``job`` known, state polled from `sacct`), to running and, possibly,
    cancellation.

    A single flat list of these -- covering every submission, on every cluster, for one
    `submit()` call -- is all a live display needs to render the whole picture, from a plain
    `rich.Live` table (see `render_job_table` below) up to, eventually, a table shared across
    several concurrent `submit`/`submit_first` calls (`rich.Live` only supports one live region
    per console, so that would fuse several such lists together instead of replacing this one).
    """

    job: JobSubmission
    log_path: Path
    """Where this submission's `sbatch` output will be written."""
    state: JobState = "SYNCING"
    error: ClusterSyncFailed | JobSubmissionFailed | None = None

    @property
    def cluster(self) -> str:
        return self.job.cluster

    @property
    def job_id(self) -> int | None:
        return self.job.job_id if isinstance(self.job, Job) else None

log_path

log_path: Path

Where this submission's sbatch output will be written.

render_job_table

render_job_table(
    cluster_to_job_submissions: dict[
        str, list[SubmissionProgress]
    ],
    *,
    cancelling: bool = False,
) -> Table

Render the current state of every submission as a single table.

A plain rich.Live for now; the natural place to plug in a registry that fuses several concurrent submit/submit_first calls' rows into one shared live region later on.

Source code in cluv/cli/submit.py

def render_job_table(
    cluster_to_job_submissions: dict[str, list[SubmissionProgress]], *, cancelling: bool = False
) -> rich.table.Table:
    """Render the current state of every submission as a single table.

    A plain `rich.Live` for now; the natural place to plug in a registry that fuses several
    concurrent `submit`/`submit_first` calls' rows into one shared live region later on.
    """
    title = "Waiting for jobs to cancel..." if cancelling else "Submitting jobs..."
    table = rich.table.Table(
        "Cluster",
        "Job ID",
        "Status",
        "Command",
        title=title,
        box=rich.box.ROUNDED,
        show_lines=True,
        header_style="bold white on #1a1a2e",
        title_style="bold cyan",
        expand=True,
    )
    for cluster_name, cluster_jobs in cluster_to_job_submissions.items():
        for job_row in cluster_jobs:
            assert job_row.cluster == cluster_name
            table.add_row(
                job_row.cluster,
                str(job_row.job_id) if job_row.job_id is not None else "-",
                rich.text.Text(job_row.state, style=_state_style(job_row.state)),
                _command_and_log_cell(job_row),
            )
    return table

submit

submit(
    cluster: str,
    job_script: Path | None,
    sbatch_args: list[str],
    program_args: list[str],
    autocommit: bool = False,
    chunking: int | None = None,
    vram: str | None = None,
    _skip_sync: bool = False,
    sync_datasets: bool = True,
    parsable: bool = False,
) -> Job | None

Submit a job to the given cluster (or all clusters if cluster=="first"), and return the Job object if successful.

When vram is set, one job is submitted per GPU type of each cluster that has at least that much VRAM (including MIG slices), racing them the same way as multiple allocations.

If parsable is True, print only the job ID (or ':' when cluster is 'first') to stdout, for programmatic use, instead of the usual human-readable summary. Everything else (logs, the live jobs table, command outputs) is silenced, as with --quiet.

Returns None if the submission failed.

Source code in cluv/cli/submit.py

async def submit(
    cluster: str,
    job_script: Path | None,
    sbatch_args: list[str],
    program_args: list[str],
    autocommit: bool = False,
    chunking: int | None = None,
    vram: str | None = None,
    _skip_sync: bool = False,
    sync_datasets: bool = True,
    parsable: bool = False,
) -> Job | None:
    """Submit a job to the given cluster (or all clusters if `cluster=="first"`),
    and return the Job object if successful.

    When `vram` is set, one job is submitted per GPU type of each cluster that has at least
    that much VRAM (including MIG slices), racing them the same way as multiple allocations.

    If `parsable` is True, print only the job ID (or '<cluster>:<job_id>' when `cluster`
    is 'first') to stdout, for programmatic use, instead of the usual human-readable summary.
    Everything else (logs, the live jobs table, command outputs) is silenced, as with `--quiet`.

    Returns None if the submission failed.
    """
    # `--parsable` promises that stdout carries nothing but the job id. `--quiet` already suppresses
    # everything the submission writes through the shared console (and the raw command outputs in
    # `cluv.remote.run`, which check `console.quiet` too), so it is all this needs to do.
    if parsable:
        console.quiet = True
    submit_command = build_submit_command(
        cluster=cluster, job_script=job_script, sbatch_args=sbatch_args, program_args=program_args
    )
    git_commit = ensure_clean_git_state(autocommit=autocommit, submit_command=submit_command)
    cluster_to_remote = await get_cluster_to_remote(cluster)

    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    log_dir = get_submission_log_dir()

    # Dict from cluster_name to *potential* job submissions.
    # They get actually converted into `Jobs` later when we actually sbatch them.
    _cluster_to_submissions = await gather_dict(
        {
            cluster_name: get_submissions(
                cluster=cluster_name,
                remote=remote,
                job_script=job_script,
                sbatch_args=sbatch_args,
                program_args=program_args,
                chunking=chunking,
                vram=vram,
                git_commit=git_commit,
            )
            for cluster_name, remote in cluster_to_remote.items()
        }
    )

    # Wrap the submissions in a mutable dataclass where we will modify the `state` and maybe set the `job` fields.
    cluster_to_job_submissions = {
        cluster_name: [
            SubmissionProgress(
                job=submission, log_path=log_dir / f"{timestamp}_{cluster_name}_{i}.txt"
            )
            for i, submission in enumerate(cluster_jobs)
        ]
        for cluster_name, cluster_jobs in _cluster_to_submissions.items()
    }
    # Create the log files right away, so the links in the jobs table always point somewhere.
    all_log_paths = tuple(
        row.log_path for rows in cluster_to_job_submissions.values() for row in rows
    )
    for rows in cluster_to_job_submissions.values():
        for row in rows:
            row.log_path.write_text(
                f"cluster: {row.cluster}\ncommand: {row.job.sbatch_command}\n\n--- sync ---\n"
            )

    if not _skip_sync:
        remotes = [r for r in cluster_to_remote.values() if r]
        with logging_commands_to(all_log_paths):
            await sync_common_part(remotes, sync_datasets=sync_datasets)

    found_running_job = asyncio.Event()

    cancelling = False

    def _render() -> rich.table.Table:
        return render_job_table(cluster_to_job_submissions, cancelling=cancelling)

    try:
        with Live(get_renderable=_render, console=console, refresh_per_second=1) as live:
            winning_job = await wait_for_first_running_job(
                cluster_to_job_submissions,
                cluster_to_remote=cluster_to_remote,
                found_running_job=found_running_job,
                _skip_sync=_skip_sync,
                sync_datasets=sync_datasets,
            )
            if winning_job is None:
                console.log("All job submissions have failed! Exiting.")
                return None

            for cluster_jobs in cluster_to_job_submissions.values():
                for job in cluster_jobs:
                    if job is not winning_job and not has_job(job) and job.state == "SYNCING":
                        job.state = "SKIPPED"

            cancelling = True

            jobs_to_cancel = [
                job_submission
                for cluster_jobs in cluster_to_job_submissions.values()
                for job_submission in cluster_jobs
                if job_submission is not winning_job and has_job(job_submission)
            ]
            if jobs_to_cancel:
                console.log("Jobs to cancel:")
                for job in jobs_to_cancel:
                    console.log(f"  {job.job_id} on {job.cluster}")
            await wait_for_jobs_to_cancel(jobs_to_cancel, cluster_to_remote)
            live.refresh()
    except (KeyboardInterrupt, asyncio.CancelledError):
        # The user stopped `cluv submit` while jobs were still in flight -- cancel everything.
        console.log("Interrupted by user. Cancelling all submitted jobs...")
        all_jobs = list(itertools.chain.from_iterable(cluster_to_job_submissions.values()))
        await run_scancel(all_jobs)
        raise

    job = winning_job.job
    assert job is not None

    if parsable:
        print(f"{job.cluster}:{job.job_id}" if cluster == "first" else job.job_id)
    else:
        console.print(
            f"Successfully submitted job {job.job_id} on cluster {job.cluster}.\n"
            f"Use `ssh {job.cluster} sacct -j {job.job_id}` to view its status, and `cluv sync"
            f" {job.cluster}` to fetch results once it is complete.",
            style="green",
        )

    save_job(job)
    return job

wait_for_first_running_job

wait_for_first_running_job(
    cluster_to_job_submissions: dict[
        str, list[SubmissionProgress[Submission]]
    ],
    *,
    cluster_to_remote: dict[str, Remote | None],
    found_running_job: Event,
    _skip_sync: bool,
    sync_datasets: bool,
    initial_delay: int = 10,
    max_wait_time_seconds: int = 60,
) -> SubmissionProgress[Job] | None

Poll sacct on each cluster until one submitted job starts running, or every submission has failed.

Mutates the job submissions in place with the latest known job id / state, so a live display can render them at any point during this wait. Sets found_running_job the moment a job starts, so that clusters which haven't submitted their own jobs yet can skip doing so.

Returns the submission for the job that started, or None if every submission ended up failing.

Source code in cluv/cli/submit.py

async def wait_for_first_running_job(
    cluster_to_job_submissions: dict[str, list[SubmissionProgress[Submission]]],
    *,
    cluster_to_remote: dict[str, Remote | None],
    found_running_job: asyncio.Event,
    _skip_sync: bool,
    sync_datasets: bool,
    initial_delay: int = 10,
    max_wait_time_seconds: int = 60,
) -> SubmissionProgress[Job] | None:
    """Poll `sacct` on each cluster until one submitted job starts running, or every submission has failed.

    Mutates the job submissions in place with the latest known job id / state, so a live display
    can render them at any point during this wait. Sets `found_running_job` the moment a job starts,
    so that clusters which haven't submitted their own jobs yet can skip doing so.

    Returns the submission for the job that started, or None if every submission ended up failing.
    """
    # Need one remote for each cluster that has job submissions.
    assert set(cluster_to_remote.keys()) >= set(cluster_to_job_submissions.keys())
    all_submissions = list(itertools.chain.from_iterable(cluster_to_job_submissions.values()))

    # Errors are reported (and stored on the submissions) by `sync_and_submit_jobs_to_cluster`.
    sync_and_submit = asyncio.gather(
        *(
            sync_and_submit_jobs_to_cluster(
                cluster_name,
                remote=cluster_to_remote[cluster_name],
                job_submissions=cluster_job_submissions,
                found_running_job=found_running_job,
                _skip_sync=_skip_sync,
                sync_datasets=sync_datasets,
            )
            for cluster_name, cluster_job_submissions in cluster_to_job_submissions.items()
        ),
        return_exceptions=True,
    )
    delay = initial_delay
    try:
        while True:
            # Read both before awaiting, so no job can get submitted in between.
            submitted_everywhere = sync_and_submit.done()
            queued_jobs = [s for s in all_submissions if has_job(s)]

            await asyncio.gather(
                *(
                    update_job_states_with_sacct(cluster_to_remote[cluster], cluster_jobs)
                    for cluster, cluster_jobs in group_by_cluster(queued_jobs).items()
                )
            )
            logger.debug("Job states: %s", {(j.cluster, j.job_id): j.state for j in queued_jobs})

            if started_jobs := (
                [j for j in queued_jobs if j.state == "COMPLETED"]
                or [j for j in queued_jobs if j.state == "RUNNING"]
            ):
                found_running_job.set()
                logger.debug(f"Found {len(started_jobs)} running (or completed) jobs.")
                return started_jobs[0]

            if submitted_everywhere:
                if all(j.state.startswith(tuple(FAILED_JOB_STATES)) for j in queued_jobs):
                    console.print(
                        f"Submitted {len(all_submissions)} jobs but none succeeded!", style="red"
                    )
                    return None
                # Skip the wait if we'd only be waiting on one job to start.
                pending_jobs = [j for j in queued_jobs if j.state == "PENDING"]
                if len(pending_jobs) == 1:
                    console.log("Only one job pending. Skipping wait for a running job.")
                    return pending_jobs[0]

            await asyncio.sleep(delay)
            delay = min(delay * 2, max_wait_time_seconds)
    finally:
        sync_and_submit.cancel()
        # Retrieve the gather's `CancelledError`, otherwise asyncio complains that it was never
        # retrieved when the future gets garbage-collected.
        with contextlib.suppress(asyncio.CancelledError):
            await sync_and_submit

wait_for_jobs_to_cancel

wait_for_jobs_to_cancel(
    job_submissions: list[SubmissionProgress[Job]],
    cluster_to_remote: dict[str, Remote | None],
    max_wait_time_seconds: int = 60,
    initial_delay: int = 1,
) -> None

Cancel every (already-submitted) job in rows, and wait until they're all done.

Source code in cluv/cli/submit.py

async def wait_for_jobs_to_cancel(
    job_submissions: list[SubmissionProgress[Job]],
    cluster_to_remote: dict[str, Remote | None],
    max_wait_time_seconds: int = 60,
    initial_delay: int = 1,
) -> None:
    """Cancel every (already-submitted) job in `rows`, and wait until they're all done."""
    cluster_to_job_submissions = group_by_cluster(job_submissions)

    assert all(job.job_id is not None for job in job_submissions)
    cluster_to_job_ids = {
        cluster: [job.job_id for job in cluster_jobs if job.job_id is not None]
        for cluster, cluster_jobs in cluster_to_job_submissions.items()
    }
    delay = initial_delay

    while not all(
        job.state.startswith(("CANCELLED", "COMPLETED", "FAILED")) for job in job_submissions
    ):
        for job in job_submissions:
            try:
                await run_scancel([job])
            except Exception as err:
                logging.debug(f"Error running scancel for job {job.job_id}: {err}")
        job_states = await gather_dict(
            {
                cluster: get_job_states_with_sacct(cluster_to_remote[cluster], cluster_job_ids)
                for cluster, cluster_job_ids in cluster_to_job_ids.items()
            }
        )
        for cluster, cluster_jobs in cluster_to_job_submissions.items():
            for job, state in zip(cluster_jobs, job_states[cluster]):
                job.state = state.strip().split()[0]  # keep only the first word of the state.

        await asyncio.sleep(delay)
        delay = min(delay * 2, max_wait_time_seconds)
    console.log(f"Cancelled {len(job_submissions)} job submission(s).")
    for job_submission in job_submissions:
        console.log(
            f"Job {job_submission.job_id} on custer {job_submission.cluster} that was in state: {job_submission.state}"
        )

run_scancel

run_scancel(jobs: list[SubmissionProgress]) -> None

Cancel the (already-submitted) jobs behind rows, grouped by remote.

Source code in cluv/cli/submit.py

async def run_scancel(jobs: list[SubmissionProgress]) -> None:
    """Cancel the (already-submitted) jobs behind `rows`, grouped by remote."""
    if not jobs:
        return
    by_remote: dict[Remote | None, list[SubmissionProgress]] = {}
    for job in jobs:
        if job.job_id is not None:
            by_remote.setdefault(job.job.remote, []).append(job)

    async def cancel(remote: Remote | None, cluster_rows: list[SubmissionProgress]) -> None:
        job_ids = [job.job_id for job in cluster_rows if job.job_id is not None]
        scancel_command = f"scancel {' '.join(map(str, job_ids))}"
        if remote is not None:
            await remote.get_output(scancel_command, hide=True)
        else:
            await run(tuple(shlex.split(scancel_command)), hide=True)

    await asyncio.gather(
        *(
            cancel(remote, cluster_rows)
            for remote, cluster_rows in by_remote.items()
            if cluster_rows
        ),
        return_exceptions=True,  # so we don't stop all the cancels if one fails.
    )

sync_and_submit_jobs_to_cluster

sync_and_submit_jobs_to_cluster(
    cluster: str,
    remote: Remote | None,
    job_submissions: list[SubmissionProgress[Submission]],
    found_running_job: Event,
    _skip_sync: bool = False,
    sync_datasets: bool = True,
) -> list[SubmissionProgress[Job]]

Sync then submit every submission for one cluster, in parallel.

NOTE: This modifies the job submissions in-place.

Returns the job submissions that were successfully submitted and now have a job id.

Source code in cluv/cli/submit.py

async def sync_and_submit_jobs_to_cluster(
    cluster: str,
    remote: Remote | None,
    job_submissions: list[SubmissionProgress[Submission]],
    found_running_job: asyncio.Event,
    _skip_sync: bool = False,
    sync_datasets: bool = True,
) -> list[SubmissionProgress[Job]]:
    """Sync then submit every submission for one cluster, in parallel.

    NOTE: This modifies the job submissions in-place.

    Returns the job submissions that were successfully submitted and now have a job id.
    """
    if found_running_job.is_set():
        # If a job has already started on another cluster, we don't need to submit more jobs.
        # NOTE: (@lebrice) I'm not sure if this case can be encountered in practice, given my limited
        # knowledge of how asyncio actually works.
        console.log(
            f"Skipping syncing with cluster {cluster} because a job "
            f"has already started on another cluster."
        )
        for job_submission in job_submissions:
            job_submission.state = "SKIPPED"
        return []

    if not _skip_sync:
        try:
            with logging_commands_to(tuple(row.log_path for row in job_submissions)):
                await sync_per_cluster_part(remote, sync_datasets=sync_datasets)
        except Exception as exc:
            console.log(f"Failed to sync with cluster {cluster}: {exc}")
            for job_submission in job_submissions:
                # todo: Not sure if this is the right way to do this.
                job_submission.error = ClusterSyncFailed().with_traceback(exc.__traceback__)
                job_submission.state = "FAILED (unable to sync)"
            raise ClusterSyncFailed() from exc

    if found_running_job.is_set():
        # If a job has already started on another cluster, we don't need to submit more jobs.
        console.log(
            f"Skipping submission of jobs to cluster {cluster} because a job "
            f"has already started on another cluster."
        )
        for job_submission in job_submissions:
            job_submission.state = "SKIPPED"
        return []

    for job_submission in job_submissions:
        job_submission.state = "SUBMITTING"

    results = await asyncio.gather(
        *(submit_job(row.job, row.log_path) for row in job_submissions),
        return_exceptions=True,
    )

    assert len(results) == len(job_submissions)
    successful_job_submissions: list[SubmissionProgress[Job]] = []
    for job_submission, result in zip(job_submissions, results):
        if isinstance(result, Job):
            job_submission.job = result
            job_submission.state = "PENDING"
            assert has_job(job_submission)
            successful_job_submissions.append(job_submission)
        elif isinstance(result, JobSubmissionFailed):
            job_submission.error = result
            job_submission.state = "FAILED"
            console.log(f"[red]{result}[/red]")
        else:
            assert isinstance(result, BaseException)
            logger.error(f"Unexpected exception: {result}")
            raise result
    return successful_job_submissions

get_submissions

get_submissions(
    cluster: str,
    remote: Remote | None,
    *,
    job_script: Path | None,
    sbatch_args: list[str],
    program_args: list[str],
    chunking: int | None,
    git_commit: str,
    vram: str | None = None,
) -> list[Submission]

Expand the possible job configurations for a cluster. Returns a list of Submission objects.

One Submission is produced per allocation configured for cluster (see [tool.cluv.clusters.<name>].sbatch_args), times one per compatible GPU type when vram is set (see expand_for_vram).

Does not do the actual job submission with sbatch.

Source code in cluv/cli/submit.py

async def get_submissions(
    cluster: str,
    remote: Remote | None,
    *,
    job_script: Path | None,
    sbatch_args: list[str],
    program_args: list[str],
    chunking: int | None,
    git_commit: str,
    vram: str | None = None,
) -> list[Submission]:
    """Expand the possible job configurations for a cluster. Returns a list of `Submission` objects.

    One `Submission` is produced per allocation configured for `cluster` (see
    `[tool.cluv.clusters.<name>].sbatch_args`), times one per compatible GPU type when `vram`
    is set (see `expand_for_vram`).

    Does *not* do the actual job submission with `sbatch`.
    """
    submissions: list[Submission] = []
    config = get_cluv_config()
    cluster_config = config.get_cluster_config(cluster)
    job_resources_options = cluster_config.sbatch_args

    # Get job vars
    job_env_vars = get_job_env_vars(
        cluster=cluster, git_commit=git_commit, cluster_config=cluster_config
    )

    # Get local job script
    local_job_script_path = get_local_job_script(
        job_script, cluster=cluster, cluster_config=cluster_config
    )

    # Get the job script path on the remote cluster
    cluster_job_script_path = get_cluster_job_script_path(
        local_job_script_path=local_job_script_path, cluster=cluster, cluster_config=cluster_config
    )

    for job_resources in job_resources_options:
        job_resources = merge_sbatch_args(from_config=job_resources, from_cli=sbatch_args)
        n_chunks, job_resources = apply_chunking(
            job_resources,
            job_script=local_job_script_path,
            chunking=chunking,
            env_vars=job_env_vars,
        )

        for expanded_resources in await expand_for_vram(
            cluster,
            remote,
            job_resources,
            job_script=local_job_script_path,
            vram=vram,
            env_vars=job_env_vars,
        ):
            expanded_resources = add_cluv_sbatch_args(
                expanded_resources,
                job_script=local_job_script_path,
                cluster=cluster,
                cluster_config=cluster_config,
            )
            sbatch_command = get_sbatch_command(
                env_vars=job_env_vars,
                job_script=cluster_job_script_path,
                sbatch_args=expanded_resources,
                program_args=program_args,
                project_dir_on_cluster=cluster_config.project_dir,
            )
            submissions.append(
                Submission(
                    cluster=cluster,
                    remote=remote,
                    job_script=local_job_script_path,
                    sbatch_args=expanded_resources,
                    program_args=program_args,
                    sbatch_command=sbatch_command,
                    n_chunks=n_chunks,
                    git_commit=git_commit,
                )
            )
    return submissions

merge_sbatch_args

merge_sbatch_args(
    from_config: SbatchArgs, from_cli: list[str]
) -> SbatchArgs

Merge the sbatch args from the config and from the CLI, with CLI args taking precedence.

-t is normalized to time (its long-flag alias) as it's merged in, so --time=1:00:00 -t=2:00:00 -- config or CLI, either order -- resolves to a single time value (the last one written) instead of leaving two separate keys for what's really the same sbatch option.

Source code in cluv/cli/submit.py

def merge_sbatch_args(from_config: SbatchArgs, from_cli: list[str]) -> SbatchArgs:
    """Merge the sbatch args from the config and from the CLI, with CLI args taking precedence.

    `-t` is normalized to `time` (its long-flag alias) as it's merged in, so `--time=1:00:00
    -t=2:00:00` -- config or CLI, either order -- resolves to a single `time` value (the last
    one written) instead of leaving two separate keys for what's really the same sbatch option.
    """
    sbatch_args_from_config = sbatch_args_to_list(from_config)
    return sbatch_args_from_list(sbatch_args_from_config + from_cli)

check_path_is_safe_to_interpolate

check_path_is_safe_to_interpolate(
    path: PurePosixPath | str, setting: str
) -> None

Raise a ValueError if path can't be interpolated into the sbatch command as-is.

Values that may contain environment variables are written into the command unquoted, so that the cluster's login shell is the one that expands them (see get_sbatch_command). Escaping them with shlex.quote first would stop exactly that, so anything else that shell treats specially has to be rejected up front instead: a space would split the value into two sbatch arguments, and a ; would end the command and start another one.

check_path_is_safe_to_interpolate("$SCRATCH/logs/imagenet", "results_path") check_path_is_safe_to_interpolate("/home/me/my logs", "results_path") ... # doctest: +ELLIPSIS Traceback (most recent call last): ... ValueError: The results_path '/home/me/my logs' contains a ' ', ...

Source code in cluv/cli/submit.py

def check_path_is_safe_to_interpolate(path: PurePosixPath | str, setting: str) -> None:
    """Raise a `ValueError` if `path` can't be interpolated into the sbatch command as-is.

    Values that may contain environment variables are written into the command *unquoted*, so that
    the cluster's login shell is the one that expands them (see `get_sbatch_command`). Escaping them
    with `shlex.quote` first would stop exactly that, so anything else that shell treats specially
    has to be rejected up front instead: a space would split the value into two `sbatch` arguments,
    and a `;` would end the command and start another one.

    >>> check_path_is_safe_to_interpolate("$SCRATCH/logs/imagenet", "results_path")
    >>> check_path_is_safe_to_interpolate("/home/me/my logs", "results_path")
    ... # doctest: +ELLIPSIS
    Traceback (most recent call last):
        ...
    ValueError: The results_path '/home/me/my logs' contains a ' ', ...
    """
    if match := _UNSAFE_PATH_CHARS.search(str(path)):
        raise ValueError(
            f"The {setting} {str(path)!r} contains a {match.group()!r}, which cluv can't pass to "
            f"sbatch safely: it goes into a `bash --login -c ...` command unquoted, so that the "
            f"cluster's login shell expands variables like $SCRATCH in it. Please use a value "
            f"without whitespace or shell metacharacters."
        )

add_cluv_sbatch_args

add_cluv_sbatch_args(
    sbatch_args: SbatchArgs,
    job_script: Path,
    cluster: str,
    cluster_config: ClusterConfig,
) -> SbatchArgs

• Add the --output flag (So that outputs are created in the results_path for the run prescribed by Cluv)
• Add the --job-name flag (So that we can identify the cluv jobs later)
• Add the --export=ALL flag (since trillium and trillium-gpu apparently have --export=None as default).
• Add the --chdir flag to move to the project folder when running the command.

Returns a new dict; the one passed in is left alone.

Source code in cluv/cli/submit.py

def add_cluv_sbatch_args(
    sbatch_args: SbatchArgs,
    job_script: Path,
    cluster: str,
    cluster_config: ClusterConfig,
) -> SbatchArgs:
    """
    - Add the --output flag (So that outputs are created in the `results_path` for the run prescribed by Cluv)
    - Add the --job-name flag (So that we can identify the cluv jobs later)
    - Add the --export=ALL flag (since trillium and trillium-gpu apparently have `--export=None` as default).
    - Add the --chdir flag to move to the project folder when running the command.

    Returns a new dict; the one passed in is left alone.
    """
    sbatch_args = sbatch_args.copy()

    base_name = sbatch_args.get("job-name") or Path(job_script).stem
    sbatch_args["job-name"] = f"cluv-{base_name}"

    if "output" not in sbatch_args and (
        _header_output := next(
            (
                line
                for line in job_script.read_text().splitlines()
                if line.strip().startswith("#SBATCH") and "--output" in line
            ),
            None,
        )
    ):
        logger.warning(
            f"[yellow]The job script {job_script} sets {_header_output.strip()!r}, which "
            f"will be overridden by cluv's --output so that results can be synced "
            f"back. Consider using cluv in your Python script to decide where to store "
            f"results instead.[/yellow]"
        )

    # Chunked (job array) jobs need `%A`/`%a` (array job id / task id) instead of `%j`.
    if "array" in sbatch_args:
        sbatch_args["output"] = str(cluster_config.results_path / f"{cluster}_%A/slurm-%A_%a.out")
    else:
        sbatch_args["output"] = str(cluster_config.results_path / f"{cluster}_%j/slurm-%j.out")

    # NOTE: `cluster_config.project_dir` already falls back to the global `project_dir`, so the
    # `$HOME/<project>` default below is only used when neither is set.
    local_project_dir = find_pyproject().parent
    remote_project_dir = cluster_config.project_dir or (
        PurePosixPath("$HOME") / local_project_dir.relative_to(Path.home())
    )
    sbatch_args["chdir"] = str(remote_project_dir)
    # Some clusters (trillium, trillium-gpu) have a wrapper around `sbatch` that sets `--export=NONE` by default,
    # which would discard the environment variables.
    sbatch_args["export"] = "ALL"
    return sbatch_args

get_sbatch_command

get_sbatch_command(
    job_script: PurePosixPath,
    sbatch_args: SbatchArgs,
    program_args: list[str],
    env_vars: dict[str, str],
    project_dir_on_cluster: PurePosixPath | None = None,
) -> str

Generate the command to submit the job via sbatch on the cluster.

Source code in cluv/cli/submit.py

def get_sbatch_command(
    job_script: PurePosixPath,
    sbatch_args: SbatchArgs,
    program_args: list[str],
    env_vars: dict[str, str],
    project_dir_on_cluster: PurePosixPath | None = None,
) -> str:
    """Generate the command to submit the job via `sbatch` on the cluster."""
    if job_script.is_absolute():
        raise RuntimeError(
            f"The job script path {str(job_script)!r} is an absolute path on this machine, but it "
            f"has to be the path of the script *on the cluster* (relative to the project there, or "
            f"starting with a variable like $HOME that the cluster's shell expands). This is what "
            f"`get_cluster_job_script_path` returns."
        )

    sbatch_flags = sbatch_args_to_list(sbatch_args)
    env_vars_prefix = " ".join(f"{k}={v}" for k, v in env_vars.items())

    # These are interpolated unquoted, so the cluster's login shell expands any env vars in them
    # (`$SCRATCH`, `$HOME`); they can't be `shlex`-escaped first, since quoting would both stop
    # that expansion and close the surrounding single-quoted string.
    check_path_is_safe_to_interpolate(job_script, "job_script")
    if isinstance(chdir := sbatch_args.get("chdir"), str):
        check_path_is_safe_to_interpolate(chdir, "chdir")
    if isinstance(output := sbatch_args.get("output"), str):
        check_path_is_safe_to_interpolate(output, "output")
    for name, value in env_vars.items():
        # Same deal: `UV_CACHE_DIR=$SCRATCH/.cache/uv` has to stay unquoted to be expanded on the
        # cluster, so a value with a space in it would make the login shell read the second word
        # as the command to run, and `sbatch` would never be reached.
        check_path_is_safe_to_interpolate(value, f"{name} environment variable")

    # `program_args` is the one part that is *not* meant to be expanded here: it is whatever the
    # user wrote after `--`, so it gets escaped, and a `$SLURM_TMPDIR` in it survives for the job
    # itself to expand. That only holds together because the whole inner command is quoted in one
    # go below - `shlex.join`'s quotes would otherwise close a hand-written `'...'` around it, and
    # an argument containing a space would break apart (POSIX single quotes don't nest).
    # `export` is required: with a bare `K=V K2=V2; sbatch ...` the assignments are plain shell
    # variables that the `sbatch` child process never sees, so `--export=ALL` (which only forwards
    # the exported environment) would carry none of them, `GIT_COMMIT` included.
    if env_vars_prefix:
        env_vars_prefix = f"export {env_vars_prefix}; "
    cd_command = ""
    if project_dir_on_cluster:
        cd_command = f"cd {project_dir_on_cluster} && "
    inner_command = (
        f"{env_vars_prefix}{cd_command}sbatch --parsable {' '.join(sbatch_flags)} {job_script} "
        f"{shlex.join(program_args)}"
    )
    return f"bash --login -c {shlex.quote(inner_command)}"

submit_job

submit_job(submission: Submission, log_path: Path) -> Job

Does the actual sbatch call, and writes its outcome to log_path.

Raise a JobSubmissionFailed if the job submission fails for some reason.

Source code in cluv/cli/submit.py

async def submit_job(submission: Submission, log_path: Path) -> Job:
    """Does the actual sbatch call, and writes its outcome to `log_path`.

    Raise a `JobSubmissionFailed` if the job submission fails for some reason.
    """
    display = display_commands.get()
    hide = not display
    warn = not raise_on_command_error.get()

    if submission.remote is not None:
        result = await submission.remote.run(
            submission.sbatch_command, display=display, warn=warn, hide=hide
        )
    else:
        result = await run(
            tuple(shlex.split(submission.sbatch_command)), _display=display, warn=warn, hide=hide
        )

    job_id = int(result.stdout.strip()) if result.returncode == 0 else None
    error = None
    if result.returncode != 0:
        error = (
            f"Failed to submit job on cluster {submission.cluster}: "
            f"{result.stderr or result.stdout}"
        )
    write_submission_log(log_path, job_id=job_id, error=error, result=result)

    if error is not None:
        raise JobSubmissionFailed(error)
    assert job_id is not None

    return Job(
        cluster=submission.cluster,
        remote=submission.remote,
        job_script=submission.job_script,
        sbatch_args=submission.sbatch_args,
        program_args=submission.program_args,
        sbatch_command=submission.sbatch_command,
        n_chunks=submission.n_chunks,
        git_commit=submission.git_commit,
        job_id=job_id,
        submitted_at=datetime.datetime.now(),
    )

logging_commands_to

logging_commands_to(log_paths: tuple[Path, ...])

Append the commands run in this context (and their outputs) to log_paths, as well as the traceback of any exception raised.

Source code in cluv/cli/submit.py

@contextlib.contextmanager
def logging_commands_to(log_paths: tuple[Path, ...]):
    """Append the commands run in this context (and their outputs) to `log_paths`, as well as
    the traceback of any exception raised."""
    try:
        with set_context(command_log_files, log_paths):
            yield
    except Exception:
        for log_path in log_paths:
            with log_path.open("a") as f:
                f.write(traceback.format_exc())
        raise

write_submission_log

write_submission_log(
    log_path: Path,
    *,
    job_id: int | None,
    error: str | None,
    result: CompletedProcess[str],
) -> None

Append the outcome of one sbatch call to log_path, so it can be found on its own even when many submissions' console output is interleaved.

Source code in cluv/cli/submit.py

def write_submission_log(
    log_path: Path,
    *,
    job_id: int | None,
    error: str | None,
    result: subprocess.CompletedProcess[str],
) -> None:
    """Append the outcome of one `sbatch` call to `log_path`, so it can be found on its own even
    when many submissions' console output is interleaved.
    """
    lines = [
        "",
        "--- submission ---",
        f"status: {'FAILED' if error is not None else 'SUBMITTED'}",
        f"job_id: {job_id if job_id is not None else '-'}",
        f"returncode: {result.returncode}",
        "",
        "--- stdout ---",
        result.stdout,
        "--- stderr ---",
        result.stderr,
    ]
    with log_path.open("a") as f:
        f.write("\n".join(lines))

build_submit_command

build_submit_command(
    cluster: str,
    job_script: str | Path | PurePosixPath | None,
    sbatch_args: list[str],
    program_args: list[str],
) -> str

Build the local cluv submit command line used to launch the job.

Source code in cluv/cli/submit.py

def build_submit_command(
    cluster: str,
    job_script: str | Path | PurePosixPath | None,
    sbatch_args: list[str],
    program_args: list[str],
) -> str:
    """Build the local `cluv submit` command line used to launch the job."""
    command_parts = ["cluv", "submit", cluster]
    if job_script is not None:
        command_parts.append(str(job_script))
    command_parts.extend(sbatch_args)
    if program_args:
        command_parts.extend(["--", *program_args])
    return shlex.join(command_parts)

create_submit_commit

create_submit_commit(submit_command: str) -> None

Create a commit with tracked changes and include the launched job command in the body.

Source code in cluv/cli/submit.py

def create_submit_commit(submit_command: str) -> None:
    """Create a commit with tracked changes and include the launched job command in the body."""
    try:
        subprocess.run(["git", "add", "-u"], check=True, capture_output=True, text=True)
        subprocess.run(
            [
                "git",
                "commit",
                "-m",
                "cluv submit: auto-commit tracked changes",
                "-m",
                f"Launched job command:\n\n{submit_command}",
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    except subprocess.CalledProcessError as err:
        error_text = (err.stderr or err.stdout or str(err)).strip()
        console.print(
            "[red]Failed to create automatic submit commit before job submission:[/red] "
            f"{error_text}"
        )
        raise

ensure_clean_git_state

ensure_clean_git_state(
    autocommit: bool = False,
    submit_command: str | None = None,
) -> str

Check git is clean locally and return the current commit hash.

Source code in cluv/cli/submit.py

def ensure_clean_git_state(autocommit: bool = False, submit_command: str | None = None) -> str:
    """
    Check git is clean locally and return the current commit hash.
    """
    git_status = subprocess.run(["git", "status", "--porcelain"], capture_output=True, text=True)
    dirty_lines = [line for line in git_status.stdout.splitlines() if not line.startswith("??")]
    if dirty_lines:
        if autocommit:
            if submit_command is None:
                raise ValueError("submit_command is required when autocommit=True")
            create_submit_commit(submit_command)
        elif not (os.environ.get("SKIP_CLEAN_GIT_CHECK", "0") == "1"):
            console.print(
                "Working directory is dirty. Please commit your changes before submitting, "
                "or use `--autocommit` (`hydra.launcher.autocommit=True` when using Hydra).",
                style="red",
            )
            sys.exit(1)

    # In GitHub Actions PR jobs we can be on a detached merge commit that doesn't exist on
    # the synced remote checkout. Prefer the branch tip commit in that case.
    current_branch = subprocess.check_output(
        ["git", "rev-parse", "--abbrev-ref", "HEAD"], text=True
    ).strip()
    if current_branch == "HEAD" and os.environ.get("GITHUB_ACTIONS"):
        github_head_ref = os.environ.get("GITHUB_HEAD_REF", "").strip()
        if github_head_ref:
            remote_head_ref = f"origin/{github_head_ref}"
            remote_head_result = subprocess.run(
                ["git", "rev-parse", "--verify", remote_head_ref],
                capture_output=True,
                text=True,
            )
            if remote_head_result.returncode == 0:
                return remote_head_result.stdout.strip()
            console.log(
                f"Could not resolve {remote_head_ref}. Falling back to local HEAD commit.",
                style="yellow",
            )

    # Capture current commit hash.
    return subprocess.check_output(["git", "rev-parse", "HEAD"], text=True).strip()