ezpz.submit⚓︎

def detect_env_setup() -> str:
    """Return shell commands that reproduce the current Python environment.

    Checks (in order):

    1. ``EZPZ_SETUP_ENV`` — if it points to a file, sources it; otherwise
       used as inline shell commands.
    2. The installed ``ezpz/bin/utils.sh``, sourced by path.
    3. Only if that is missing, the network copy fetched via curl --
       which cannot work on a compute node.

    Returns:
        A (possibly multi-line) string of shell commands.
    """
    setup_env = os.environ.get("EZPZ_SETUP_ENV")
    if setup_env:
        if Path(setup_env).is_file():
            return f"source {shlex.quote(setup_env)}"
        return setup_env

    # Prefer the installed copy of utils.sh over fetching it.
    #
    # Compute nodes on Aurora / Polaris / Sunspot have no outbound
    # route, so `source <(curl ...)` hangs for ~270 s and then leaves
    # the environment UNCONFIGURED -- the failure surfaces much later as
    # `CONDA_PREFIX still not set`, or as every ezpz_* function being
    # undefined. SKILL.md lists this first under "inside a batch script,
    # three things change"; generating the curl form by default
    # contradicted that advice.
    local_utils = Path(__file__).parent / "bin" / "utils.sh"
    if local_utils.is_file():
        return f"source {shlex.quote(str(local_utils))} && ezpz_setup_env"

    # No installed copy (zip import, odd packaging): the network form
    # still works from a login node.
    return "source <(curl -fsSL https://ezpz.cool/utils.sh) && ezpz_setup_env"

def generate_pbs_script(
    command: str,
    *,
    nodes: int = 1,
    time: str = "01:00:00",
    queue: str = "debug",
    account: str | None = None,
    filesystems: str = "home",
    job_name: str | None = None,
    working_dir: str | None = None,
    env_setup: str | None = None,
    wrap_with_launch: bool = True,
    strict: bool = True,
) -> str:
    """Generate a PBS job script.

    Args:
        command: The command to execute (already formatted as a string).
        nodes: Number of compute nodes.
        time: Walltime in ``HH:MM:SS`` format.
        queue: PBS queue name.
        account: PBS account/project.  Falls back to ``PBS_ACCOUNT`` or
            ``PROJECT`` environment variables.
        filesystems: Comma-separated or colon-separated filesystem list.
        job_name: Job name (defaults to ``"ezpz"``).
        working_dir: Directory to ``cd`` into (defaults to cwd).
        env_setup: Shell commands for environment activation.
        wrap_with_launch: If ``True``, prefix *command* with ``ezpz launch``.
        strict: When ``True`` (default) emit ``set -eo pipefail``. Pass
            ``False`` for multi-arm experiment scripts where one arm
            timing out must not abort the remaining arms. ``set -u`` is
            never emitted: the Lmod module system is not ``set -u``-clean
            and sourcing ``/etc/profile`` under it aborts the script
            before a single line is printed (see SKILL.md).

    Returns:
        The complete job script as a string.
    """
    account = account or os.environ.get(
        "PBS_ACCOUNT", os.environ.get("PROJECT", "")
    )
    job_name = job_name or "ezpz"
    # working_dir="" suppresses the `cd` line entirely. Remote jobs need
    # this: os.getcwd() is a LOCAL path that does not exist on the target,
    # and shlex.quote("$HOME") yields a literal directory named `$HOME`.
    # Both PBS and SLURM start the job in $HOME anyway.
    if working_dir == "":
        working_dir = None
    elif working_dir is None:
        working_dir = os.getcwd()
    if env_setup is None:
        env_setup = detect_env_setup()

    fs = filesystems.replace(",", ":")

    lines = [
        "#!/bin/bash --login",
        f"#PBS -l select={nodes}",
        f"#PBS -l walltime={time}",
        f"#PBS -l filesystems={fs}",
    ]
    if account:
        lines.append(f"#PBS -A {account}")
    lines += [
        "#PBS -k doe",
        "#PBS -j oe",
        f"#PBS -q {queue}",
        f"#PBS -N {job_name}",
        "",
        "set -eo pipefail" if strict else "set -o pipefail",
        *([f"cd {shlex.quote(working_dir)}"] if working_dir else []),
    ]
    if env_setup:
        lines += ["", "# ── Environment setup ──", env_setup]

    run_cmd = f"ezpz launch {command}" if wrap_with_launch else command
    lines += ["", run_cmd, ""]
    return "\n".join(lines)

def generate_slurm_script(
    command: str,
    *,
    nodes: int = 1,
    time: str = "01:00:00",
    queue: str = "debug",
    account: str | None = None,
    job_name: str | None = None,
    working_dir: str | None = None,
    env_setup: str | None = None,
    wrap_with_launch: bool = True,
    gpus_per_node: int | None = None,
    ntasks_per_node: int | None = None,
    constraint: str | None = None,
    strict: bool = True,
) -> str:
    """Generate a SLURM job script.

    Args:
        command: The command to execute (already formatted as a string).
        nodes: Number of compute nodes.
        time: Walltime in ``HH:MM:SS`` format.
        queue: SLURM partition name.
        account: SLURM account.  Falls back to ``SLURM_ACCOUNT`` or
            ``PROJECT`` environment variables.
        job_name: Job name (defaults to ``"ezpz"``).
        working_dir: Directory to ``cd`` into (defaults to cwd).
        env_setup: Shell commands for environment activation.
        wrap_with_launch: If ``True``, prefix *command* with ``ezpz launch``.
        gpus_per_node: ``--gpus-per-node``. Required on Perlmutter: without
            it a GPU job gets none allocated.
        ntasks_per_node: ``--ntasks-per-node``. Sets ranks per node; pair
            with *gpus_per_node* for one rank per GPU.
        constraint: ``-C/--constraint`` (e.g. ``"gpu"`` on Perlmutter,
            which selects the GPU partition's node type).
        strict: When ``True`` (default) emit ``set -eo pipefail``. Pass
            ``False`` for multi-arm experiment scripts where one arm
            timing out must not abort the remaining arms. ``set -u`` is
            never emitted: the Lmod module system is not ``set -u``-clean
            and sourcing ``/etc/profile`` under it aborts the script
            before a single line is printed (see SKILL.md).

    Returns:
        The complete job script as a string.
    """
    account = account or os.environ.get(
        "SLURM_ACCOUNT", os.environ.get("PROJECT", "")
    )
    job_name = job_name or "ezpz"
    if working_dir == "":
        working_dir = None  # remote: let the scheduler land in $HOME
    elif working_dir is None:
        working_dir = os.getcwd()
    if env_setup is None:
        env_setup = detect_env_setup()

    lines = [
        "#!/bin/bash --login",
        f"#SBATCH --nodes={nodes}",
        f"#SBATCH --time={time}",
    ]
    if account:
        lines.append(f"#SBATCH --account={account}")
    if constraint:
        lines.append(f"#SBATCH --constraint={constraint}")
    lines += [
        f"#SBATCH --partition={queue}",
        f"#SBATCH --job-name={job_name}",
    ]
    if gpus_per_node is not None:
        lines.append(f"#SBATCH --gpus-per-node={gpus_per_node}")
    if ntasks_per_node is not None:
        lines.append(f"#SBATCH --ntasks-per-node={ntasks_per_node}")
    lines += [
        "",
        "set -eo pipefail" if strict else "set -o pipefail",
        *([f"cd {shlex.quote(working_dir)}"] if working_dir else []),
    ]
    if env_setup:
        lines += ["", "# ── Environment setup ──", env_setup]

    run_cmd = f"ezpz launch {command}" if wrap_with_launch else command
    lines += ["", run_cmd, ""]
    return "\n".join(lines)

def submit(
    command: list[str] | None = None,
    script: str | Path | None = None,
    *,
    nodes: int = 1,
    time: str = "01:00:00",
    queue: str = "debug",
    account: str | None = None,
    filesystems: str = "home",
    job_name: str | None = None,
    working_dir: str | None = None,
    scheduler: str | None = None,
    wrap_with_launch: bool = True,
    dry_run: bool = False,
    env_setup: str | None = None,
    gpus_per_node: int | None = None,
    ntasks_per_node: int | None = None,
    constraint: str | None = None,
    strict: bool = True,
) -> str | None:
    """Submit a job to the active scheduler.

    Either *command* (a list of args to wrap) or *script* (path to an
    existing job script) must be provided.

    Args:
        command: Command to wrap in a generated job script.
        script: Path to an existing job script to submit directly.
        nodes: Number of compute nodes.
        time: Walltime in ``HH:MM:SS`` format.
        queue: Queue or partition name.
        account: Project/account for billing.
        filesystems: PBS filesystems directive (ignored for SLURM).
        job_name: Job name.
        working_dir: Working directory for the job.
        scheduler: Override scheduler detection (``"PBS"`` or ``"SLURM"``).
        wrap_with_launch: Wrap *command* with ``ezpz launch``.
        dry_run: Print the script but do not submit.
        gpus_per_node: SLURM ``--gpus-per-node`` (ignored for PBS).
        ntasks_per_node: SLURM ``--ntasks-per-node`` (ignored for PBS).
        constraint: SLURM ``--constraint`` (ignored for PBS).
        strict: Emit ``set -eo pipefail``; ``False`` drops ``-e`` for
            multi-arm experiment scripts.

    Returns:
        The job ID string, or ``None`` on failure / dry-run.
    """
    from ezpz.configs import get_scheduler

    if scheduler is None:
        scheduler = get_scheduler()

    if scheduler.upper() not in ("PBS", "SLURM"):
        print(
            f"No supported scheduler detected (got {scheduler!r}).\n"
            "Pass --scheduler PBS or --scheduler SLURM explicitly.",
            file=sys.stderr,
        )
        return None

    # ── Mode 2: submit existing script ───────────────────────────────────
    if script is not None:
        script_path = Path(script)
        if not script_path.is_file():
            print(f"Script not found: {script_path}", file=sys.stderr)
            return None
        if dry_run:
            print(script_path.read_text())
            print(f"[dry-run] Would submit {script_path} via {scheduler}")
            return None
        job_id = submit_job(script_path, scheduler)
        if job_id:
            print(f"Submitted job {job_id}")
        return job_id

    # ── Mode 1: generate script from command ─────────────────────────────
    if command is None or not command:
        print("No command or script provided.", file=sys.stderr)
        return None

    cmd_str = shlex.join(command)

    if job_name is None:
        # Derive from command: "python3 -m ezpz.examples.test" → "ezpz.examples.test"
        for i, arg in enumerate(command):
            if arg == "-m" and i + 1 < len(command):
                job_name = command[i + 1]
                break
        if job_name is None:
            job_name = Path(command[0]).stem

    if scheduler.upper() == "PBS":
        script_text = generate_pbs_script(
            cmd_str,
            nodes=nodes,
            time=time,
            queue=queue,
            account=account,
            filesystems=filesystems,
            job_name=job_name,
            working_dir=working_dir,
            wrap_with_launch=wrap_with_launch,
            env_setup=env_setup,
            strict=strict,
        )
    else:
        script_text = generate_slurm_script(
            cmd_str,
            nodes=nodes,
            time=time,
            queue=queue,
            account=account,
            job_name=job_name,
            working_dir=working_dir,
            wrap_with_launch=wrap_with_launch,
            env_setup=env_setup,
            gpus_per_node=gpus_per_node,
            ntasks_per_node=ntasks_per_node,
            constraint=constraint,
            strict=strict,
        )

    # Print for transparency
    print("Generated job script:")
    print("-" * 60)
    print(script_text)
    print("-" * 60)

    if dry_run:
        print(f"[dry-run] Would submit via {scheduler.lower()}")
        return None

    # Write to a file so the user can inspect/resubmit later
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    script_dir = Path(working_dir or os.getcwd())
    script_path = script_dir / f".ezpz-submit-{timestamp}.sh"
    script_path.write_text(script_text, encoding="utf-8")
    script_path.chmod(0o755)

    job_id = submit_job(script_path, scheduler)
    if job_id:
        print(f"Submitted job {job_id}")
        print(f"Script saved to {script_path}")
    return job_id

def submit_job(script_path: str | Path, scheduler: str) -> str | None:
    """Submit a job script and return the job ID.

    Args:
        script_path: Path to the job script.
        scheduler: ``"PBS"`` or ``"SLURM"``.

    Returns:
        The job ID string, or ``None`` if submission failed.
    """
    script_path = str(script_path)
    if scheduler.upper() == "PBS":
        cmd = ["qsub", script_path]
    elif scheduler.upper() == "SLURM":
        cmd = ["sbatch", script_path]
    else:
        logger.error("Unknown scheduler %r — cannot submit", scheduler)
        return None

    logger.info("Submitting: %s", " ".join(cmd))
    try:
        result = subprocess.run(
            cmd, capture_output=True, text=True, check=True
        )
        job_id = result.stdout.strip()
        return job_id
    except FileNotFoundError:
        logger.error(
            "%s not found — is the scheduler available on this system?",
            cmd[0],
        )
        return None
    except subprocess.CalledProcessError as exc:
        logger.error(
            "Submission failed (exit %d): %s",
            exc.returncode,
            exc.stderr.strip(),
        )
        return None