Frozen Slurm launch plans

trade_study.slurm

Frozen Python task submission and node packing for existing study runners.

Adapted from the FLU_USA, denv-ade-v2 and mcvf frozen-worktree/packing scripts. Tasks are explicit Python argv lists; no live evaluators are pickled. Existing grid ledgers and adaptive journals remain responsible for within-task recovery.

FrozenCheckout(path, commit) dataclass

A detached checkout's location and full commit identity.

Construction alone does not verify files. verify checks the current commit and tracked worktree/index before submission and each task boundary. Keep data, logs, virtual environments and outputs outside the tracked tree.

verify()

Refuse a moved commit or modified tracked source/index.

Raises:

Type	Description
ValueError	The current commit or tracked source differs.

Source code in src/trade_study/slurm.py

def verify(self) -> None:
    """Refuse a moved commit or modified tracked source/index.

    Raises:
        ValueError: The current commit or tracked source differs.
    """
    if not re.fullmatch(r"[0-9a-f]{40}|[0-9a-f]{64}", self.commit):
        msg = "A frozen checkout requires a full commit hash"
        raise ValueError(msg)
    if _git(self.path, "rev-parse", "--abbrev-ref", "HEAD") != "HEAD":
        msg = "Frozen execution requires a detached worktree"
        raise ValueError(msg)
    if _git(self.path, "rev-parse", "HEAD") != self.commit or _git(
        self.path, "status", "--porcelain", "--untracked-files=no"
    ):
        msg = "Frozen checkout commit or tracked source/index mismatch"
        raise ValueError(msg)

SlurmResources(width=1, cpus_per_task=1, time_limit='01:00:00', memory=None, partition=None, account=None) dataclass

One node's CPU packing and site-specific submission choices.

width * cpus_per_task CPUs are requested in one allocation. Child tasks share that allocation's CPU set; no per-child CPU affinity is imposed. Memory and wall time apply to the whole allocation, rather than each child.

Examples:

>>> resource = SlurmResources(width=4, cpus_per_task=2)
>>> resource.width * resource.cpus_per_task
8

__post_init__()

Reject invalid packing sizes and malformed resource tokens.

Raises:

Type	Description
ValueError	Counts are not positive integers or a token is invalid.

Source code in src/trade_study/slurm.py

def __post_init__(self) -> None:
    """Reject invalid packing sizes and malformed resource tokens.

    Raises:
        ValueError: Counts are not positive integers or a token is invalid.
    """
    if any(
        type(value) is not int or value < 1
        for value in (self.width, self.cpus_per_task)
    ):
        msg = "Slurm width and CPUs per task must be positive integers"
        raise ValueError(msg)
    for value in (self.time_limit, self.memory, self.partition, self.account):
        if value is not None and not re.fullmatch(r"[A-Za-z0-9_.:+/-]+", value):
            msg = "Slurm resource values must be nonempty single tokens"
            raise ValueError(msg)

SlurmTask(args, outputs=()) dataclass

One Python argv vector and optional declared output files.

Output paths are absolute and outside the frozen checkout. Completion receipts hash these files. Missing files cause reevaluation; changed content is refused. Without declared outputs, a receipt attests successful exit only.

__post_init__()

Validate immutable task arguments and explicit output locations.

Raises:

Type	Description
ValueError	Arguments or output paths do not meet the task contract.

Source code in src/trade_study/slurm.py

def __post_init__(self) -> None:
    """Validate immutable task arguments and explicit output locations.

    Raises:
        ValueError: Arguments or output paths do not meet the task contract.
    """
    if (
        not isinstance(self.args, tuple)
        or not self.args
        or not isinstance(self.outputs, tuple)
        or any(
            not isinstance(arg, str) or not arg or "\0" in arg for arg in self.args
        )
        or any(
            not isinstance(path, str) or not Path(path).is_absolute()
            for path in self.outputs
        )
    ):
        msg = "Use immutable Python argv and absolute declared output paths"
        raise ValueError(msg)

SlurmPlan(path, identity) dataclass

A reviewed launch manifest and its independently retained content identity.

submit explicitly launches a batch job. run_local exercises packing without a scheduler. Both refuse changed source/runtime and reuse only task completions matching this exact plan. Keep the output directory unchanged.

run_local()

Run the same guarded pack locally, propagating any child failure.

Native child failures raise subprocess.CalledProcessError.

Raises:

Type	Description
ValueError	Plan, launcher, source or runtime differs.

Source code in src/trade_study/slurm.py

def run_local(self) -> None:
    """Run the same guarded pack locally, propagating any child failure.

    Native child failures raise ``subprocess.CalledProcessError``.

    Raises:
        ValueError: Plan, launcher, source or runtime differs.
    """
    plan = _read_plan(self.path, self.identity)
    launcher = self.path.parent / "launcher.py"
    if _hash(launcher.read_bytes()) != plan["launcher_digest"]:
        msg = "Slurm launcher integrity mismatch"
        raise ValueError(msg)
    _capture([plan["python"], str(launcher), str(self.path), self.identity])

submit(*, after=())

Submit one packed job, checking finished dependencies before afterok.

Scheduler errors propagate as subprocess.CalledProcessError; failed submission commands are not silently retried.

Parameters:

Name	Type	Description	Default
after	Sequence[str]	Numeric job IDs on the connected Slurm cluster. Active jobs become afterok dependencies; confirmed successful jobs are omitted.	()

Returns:

Type	Description
str	Native parsable sbatch job reference, retaining any cluster suffix.

Raises:

Type	Description
ValueError	Plan/source/runtime or dependency status is incompatible.

Source code in src/trade_study/slurm.py

def submit(self, *, after: Sequence[str] = ()) -> str:
    """Submit one packed job, checking finished dependencies before afterok.

    Scheduler errors propagate as ``subprocess.CalledProcessError``; failed
    submission commands are not silently retried.

    Args:
        after: Numeric job IDs on the connected Slurm cluster. Active jobs
            become afterok dependencies; confirmed successful jobs are omitted.

    Returns:
        Native parsable sbatch job reference, retaining any cluster suffix.

    Raises:
        ValueError: Plan/source/runtime or dependency status is incompatible.
    """
    plan = _read_plan(self.path, self.identity)
    checkout = _verify_plan(plan)
    launcher = self.path.parent / "launcher.py"
    if _hash(launcher.read_bytes()) != plan["launcher_digest"]:
        msg = "Slurm launcher integrity mismatch"
        raise ValueError(msg)
    # The fixed bootstrap checks copied code before executing it. Job args
    # are shell-quoted only here; all task commands themselves use argv.
    import shlex

    bootstrap = (
        "import hashlib,pathlib,runpy,sys; "
        "source=pathlib.Path(sys.argv.pop(1)); expected=sys.argv.pop(1); "
        "actual=hashlib.sha256(source.read_bytes()).hexdigest(); "
        "sys.exit('Slurm launcher integrity mismatch') "
        "if actual!=expected else None; "
        "sys.argv[0]=str(source); runpy.run_path(str(source),run_name='__main__')"
    )
    script = (
        "#!/bin/bash\nset -euo pipefail\nexec "
        + shlex.join([
            plan["python"],
            "-c",
            bootstrap,
            str(launcher),
            plan["launcher_digest"],
            str(self.path),
            self.identity,
        ])
        + "\n"
    )
    batch = self.path.parent / "batch.sh"
    _write_once(batch, script.encode())
    resource = SlurmResources(**plan["resources"])
    args = [
        _program("sbatch"),
        "--parsable",
        "--nodes=1",
        "--ntasks=1",
        f"--cpus-per-task={resource.width * resource.cpus_per_task}",
        f"--time={resource.time_limit}",
        f"--chdir={checkout.path}",
        f"--output={self.path.parent / 'slurm-%j.log'}",
    ]
    for flag, value in (
        ("mem", resource.memory),
        ("partition", resource.partition),
        ("account", resource.account),
    ):
        if value is not None:
            args.append(f"--{flag}={value}")
    pending = _dependencies(after)
    if pending:
        args.append("--dependency=afterok:" + ":".join(pending))
    submitted = _capture([*args, str(batch)])
    if not re.fullmatch(r"[0-9]+(?:;[A-Za-z0-9_.-]+)?", submitted):
        msg = (
            "Unexpected sbatch job reference; inspect the scheduler before retrying"
        )
        raise ValueError(msg)
    return submitted

freeze_checkout(repository, destination, *, ref='HEAD')

Create or reuse a detached worktree at an explicitly resolved commit.

Parameters:

Name	Type	Description	Default
repository	str | Path	Existing local Git repository; no fetch or branch change occurs.	required
destination	str | Path	New worktree location, or the same previously frozen checkout.	required
ref	str	Commit or ref resolved once, before any job is submitted.	'HEAD'

Returns:

Type	Description
FrozenCheckout	Verified full-commit checkout. Queued jobs do not read the submitter's tree.

Git errors propagate as subprocess.CalledProcessError; verification raises ValueError for a moved commit or changed tracked source.

Source code in src/trade_study/slurm.py

def freeze_checkout(
    repository: str | Path, destination: str | Path, *, ref: str = "HEAD"
) -> FrozenCheckout:
    """Create or reuse a detached worktree at an explicitly resolved commit.

    Args:
        repository: Existing local Git repository; no fetch or branch change occurs.
        destination: New worktree location, or the same previously frozen checkout.
        ref: Commit or ref resolved once, before any job is submitted.

    Returns:
        Verified full-commit checkout. Queued jobs do not read the submitter's tree.

    Git errors propagate as ``subprocess.CalledProcessError``; verification
    raises ``ValueError`` for a moved commit or changed tracked source.
    """
    repository = Path(repository).resolve()
    destination = Path(destination).resolve()
    commit = _git(
        repository, "rev-parse", "--verify", "--end-of-options", f"{ref}^{{commit}}"
    )
    if not destination.exists():
        destination.parent.mkdir(parents=True, exist_ok=True)
        _git(repository, "worktree", "add", "--detach", str(destination), commit)
    frozen = FrozenCheckout(destination, commit)
    frozen.verify()
    return frozen

prepare_slurm_plan(checkout, tasks, *, directory, revision, imports, python=sys.executable, resources=None)

Freeze explicit Python task argv, runtime, code and external-data revision.

Parameters:

Name	Type	Description	Default
checkout	FrozenCheckout	Detached worktree created by freeze_checkout.	required
tasks	Sequence[SlurmTask]	Ordered tasks, e.g. SlurmTask(('-m', 'project.run', '0')). The selected interpreter is prepended; strings are not shell programs.	required
directory	str | Path	Plan/log/completion directory outside the frozen source tree.	required
revision	str	Nonempty caller-managed identity for external data, configuration and environment behavior not captured by commits/package versions.	required
imports	Sequence[str]	Project modules whose tracked import origins must be in the frozen checkout; at least one guard is required.	required
python	str | Path	Exact shared interpreter path, preserving virtualenv symlinks.	executable
resources	SlurmResources | None	Single-node CPU packing and site-specific Slurm choices.	None

Returns:

Type	Description
SlurmPlan	Immutable plan that can be reviewed, run locally, or explicitly submitted.

Raises:

Type	Description
ValueError	Commands, revision, imports, source, runtime or reuse are invalid.

Source code in src/trade_study/slurm.py

def prepare_slurm_plan(
    checkout: FrozenCheckout,
    tasks: Sequence[SlurmTask],
    *,
    directory: str | Path,
    revision: str,
    imports: Sequence[str],
    python: str | Path = sys.executable,
    resources: SlurmResources | None = None,
) -> SlurmPlan:
    """Freeze explicit Python task argv, runtime, code and external-data revision.

    Args:
        checkout: Detached worktree created by ``freeze_checkout``.
        tasks: Ordered tasks, e.g. ``SlurmTask(('-m', 'project.run', '0'))``.
            The selected interpreter is prepended; strings are not shell programs.
        directory: Plan/log/completion directory outside the frozen source tree.
        revision: Nonempty caller-managed identity for external data, configuration
            and environment behavior not captured by commits/package versions.
        imports: Project modules whose tracked import origins must be in the frozen
            checkout; at least one guard is required.
        python: Exact shared interpreter path, preserving virtualenv symlinks.
        resources: Single-node CPU packing and site-specific Slurm choices.

    Returns:
        Immutable plan that can be reviewed, run locally, or explicitly submitted.

    Raises:
        ValueError: Commands, revision, imports, source, runtime or reuse are invalid.
    """
    checkout = FrozenCheckout(checkout.path.resolve(), checkout.commit)
    checkout.verify()
    directory = Path(directory).resolve()
    python = str(Path(python).absolute())
    if directory.is_relative_to(checkout.path) or not revision.strip() or not imports:
        msg = "Use external outputs, a nonempty revision and project import guards"
        raise ValueError(msg)
    if not tasks or any(not isinstance(task, SlurmTask) for task in tasks):
        msg = "Tasks require explicit SlurmTask Python argv records"
        raise ValueError(msg)
    outputs = [str(Path(path).resolve()) for task in tasks for path in task.outputs]
    if len(set(outputs)) != len(outputs) or any(
        Path(path).is_relative_to(checkout.path) or Path(path).is_relative_to(directory)
        for path in outputs
    ):
        msg = "Declared outputs must be distinct files outside source and plan"
        raise ValueError(msg)
    imports = _task_sources(checkout, tasks, imports)
    runtime = _runtime(python, checkout, imports)
    launcher = Path(__file__).read_bytes()
    payload = {
        "schema": "trade-study-slurm-v1",
        "checkout": str(checkout.path),
        "commit": checkout.commit,
        "python": python,
        "tasks": [
            {
                "args": list(task.args),
                "outputs": [str(Path(p).resolve()) for p in task.outputs],
            }
            for task in tasks
        ],
        "imports": list(imports),
        "revision": revision,
        "resources": asdict(resources or SlurmResources()),
        "runtime": runtime,
        "launcher_digest": _hash(launcher),
    }
    directory.mkdir(parents=True, exist_ok=True)
    _write_once(directory / "launcher.py", launcher)
    _write_once(directory / "plan.json", _json(payload))
    return SlurmPlan(directory / "plan.json", _hash(_json(payload)))