diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 169dab9..ac7de36 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -49,6 +49,7 @@ repos: args: [src, tests, examples] additional_dependencies: - typer + - pytest - repo: https://github.com/scientific-python/cookie rev: 2024.04.23 hooks: diff --git a/CHANGELOG.md b/CHANGELOG.md index fe0b143..56e4b87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,16 +7,31 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased](https://github.com/MaterialsPhysicsANU/pbspy/compare/v0.0.9...HEAD) -### Removed +### Added -- Remove `progress` parameter from the `wait[_all]` and `result[_all]` methods of `Job` - - These methods no longer print job IDs or progress -- Remove `rich` dependency +- Add the public `Backend` interface and default `LocalBackend` implementation + - `JobDescription.submit()` accepts a backend supplied by the caller + - Each `Job` retains its backend for waiting, result retrieval, and cancellation +- Add `Job.cancel()` +- Add batched PBS state-query and deletion helpers + +### Changed + +- Group `Job.wait_all()` and `Job.result_all()` calls by backend so each backend can wait for its jobs together +- Reuse a shared `LocalBackend` for jobs submitted without an explicit backend +- Return `None` when PBS has not reported a completed job's exit status ### Fixed +- Preserve `JobDescription.output_path` and `error_path` on submitted jobs and use them when retrieving results +- Treat cancellation of finished or unknown PBS jobs as successful - Update `JOBFS` limit for jobs on the `copyq` queue of the `Gadi` supercomputer +### Removed + +- Remove progress display from the `Job` wait and result methods +- Remove `rich` dependency + ## [0.0.9](https://github.com/MaterialsPhysicsANU/pbspy/releases/tag/v0.0.9) - 2026-05-18 ### Added diff --git a/README.md b/README.md index 9bfbe0f..6025190 100644 --- a/README.md +++ b/README.md @@ -12,6 +12,8 @@ See the [documentation](https://MaterialsPhysicsANU.github.io/pbspy/) for more i ## Example +### Running directly on a supercomputer login-node + ```python from pbspy import Job, JobDescription @@ -38,23 +40,12 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -Output (partially executed): - -```text -✓ 124397435.gadi-pbs job_a -0:01:15 124397436.gadi-pbs job_b -``` - -Output (completed): - -```text -✓ 124397435.gadi-pbs job_a -✓ 124397436.gadi-pbs job_b +### Backends -job_a: A -job_b: B -``` +PBS operations use `LocalBackend` by default and run on the current host. +`JobDescription.submit()` also accepts any implementation of the public `Backend` +interface, allowing alternate execution mechanisms to live in separate packages. -## Licence +## License `pbspy` is licensed under the MIT License [LICENSE](./LICENSE) or . diff --git a/pyproject.toml b/pyproject.toml index b81fa64..30f158e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,7 +21,6 @@ license = {text = "MIT License"} dependencies = [ ] - [dependency-groups] dev = [ "mypy>=1.14.1", diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 2bca0a5..2824846 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -4,92 +4,26 @@ from __future__ import annotations -import re -import shlex -import subprocess -import time from dataclasses import dataclass, field from typing import Any, Self, TypeAlias -__all__ = ["JobDescription", "Job", "JobResult", "QueueLimits", "QueueLimitsMap", "gadi"] +from pbspy._backend import Backend +from pbspy._local_backend import LocalBackend +# Shared LocalBackend instance so all locally-submitted jobs are grouped together +# when calling Job.wait_all() / Job.result_all(), restoring concurrent polling. +_DEFAULT_LOCAL_BACKEND = LocalBackend() -def _get_job_name(job_id: str) -> str: - process = subprocess.run( - ["qstat", "-fx", job_id], - capture_output=True, - ) - if process.returncode != 0: - raise RuntimeError(f"Failed to get job status: {process.stderr.decode('utf-8')}") - output = process.stdout.decode("utf-8") - - # Grab the exit code - # NOTE: This is not always immediately set when the job finishes... - re_match = re.search(r"Job_Name = ([^\s]+)", output) - if re_match is not None: - group = re_match.group(1) - assert group - return group - else: - raise RuntimeError("Could not retrieve job name") - - -def _try_get_exit_code(job_id: str) -> int | None: - """ - Retrieves the exit code of a job with the given job ID. - - Args: - job_id (str): The ID of the job. - - Returns: - int | None: The exit code of the job if available, None otherwise. - """ - # Get the exit code - process = subprocess.run( - ["qstat", "-fx", job_id], - capture_output=True, - ) - if process.returncode != 0: - raise RuntimeError(f"Failed to get job status: {process.stderr.decode('utf-8')}") - output = process.stdout.decode("utf-8") - - # Grab the exit code - # NOTE: This is not always immediately set when the job finishes... - exit_code_match = re.search(r"Exit_status = (\d+)", output) - if exit_code_match is not None: - exit_code_group = exit_code_match.group(1) - assert exit_code_group - return int(exit_code_group) - else: - raise RuntimeError("Could not get exit code") - - -def _pbs_wait_for_jobs(jobs: list[Job]) -> None: - """ - Waits for the PBS jobs to complete. - - Args: - jobs (list[Job]): A list of Job objects representing the PBS jobs. - - Returns: - None - """ - task_done = [False] * len(jobs) - while not all(task_done): - time.sleep(60) - for i, job in enumerate(jobs): - if task_done[i]: - continue - process = subprocess.run( - ["qstat", job.job_id], - capture_output=True, - ) - if process.returncode != 0: - output = process.stdout.decode("utf-8") - if job.job_id not in output or "has finished" in process.stderr.decode("utf-8"): - task_done[i] = True - # else: - # raise RuntimeError(f"Failed to check job status: {process.stderr.decode('utf-8')}") +__all__ = [ + "JobDescription", + "Job", + "JobResult", + "QueueLimits", + "QueueLimitsMap", + "Backend", + "LocalBackend", + "gadi", +] @dataclass @@ -127,82 +61,61 @@ class Job: description: str | None = None """A description of the job. Unused by PBS.""" + backend: Backend = field(default=_DEFAULT_LOCAL_BACKEND, repr=False, compare=False) + """The backend used to submit and track this job.""" + + output_path: str | None = None + """Custom path for the job's stdout output file (``#PBS -o``), if any.""" + + error_path: str | None = None + """Custom path for the job's stderr error file (``#PBS -e``), if any.""" + def wait(self) -> None: """ Wait for the job to complete. """ - _pbs_wait_for_jobs([self]) - - def _result_no_wait(self) -> JobResult: - """ - Return the result of a job without waiting. - - If the job has not completed, this method will raise an error. - """ - - # Get output and error file - job_id_num = self.job_id.split(".")[0] # Remove .gadi-pbspy suffix - output_file = f"{self.job_name}.o{job_id_num}" - error_file = f"{self.job_name}.e{job_id_num}" - try: - with open(output_file) as f: - output = f.read() - except FileNotFoundError: - output = "" - try: - with open(error_file) as f: - error = f.read() - except FileNotFoundError: - error = "" - - # Try and process the output - output_split = output.split( - "\n======================================================================================\n" - ) - exit_code = None - if len(output_split) == 3: - output = output_split[0] - stats = output_split[1].strip() - stats = "\n".join(stats.split("\n")[1:]) # Skip first line "Resource usage on .." - pbs_stats = dict() - pattern = re.compile(r"\s*([^:]+):\s*([^\s]+)") - for match in pattern.finditer(stats): - key = match.group(1).strip() - value = match.group(2).strip() - pbs_stats[key] = value - if key == "Exit Status": - exit_code_match = re.search(r"^\d+", value) - if exit_code_match is not None: - exit_code = int(exit_code_match.group()) - else: - pbs_stats = None - - if exit_code is None: - exit_code = _try_get_exit_code(self.job_id) - - return JobResult(exit_code=exit_code, output=output, error=error, stats=pbs_stats) + self.backend.wait([self]) def result(self) -> JobResult: """ Waits for the job to complete and returns the result. """ self.wait() - return self._result_no_wait() + return self.backend.get_result(self) + + def cancel(self) -> None: + """Cancel (``qdel``) this job.""" + self.backend.delete([self.job_id]) @staticmethod def wait_all(jobs: list[Job]) -> None: """ Waits for multiple jobs to complete. """ - _pbs_wait_for_jobs(jobs) + if not jobs: + return + # Group by backend so each backend can wait for its own jobs efficiently + _wait_all_grouped(jobs) @staticmethod def result_all(jobs: list[Job]) -> list[JobResult]: """ Waits for multiple jobs to complete and returns their results. """ - _pbs_wait_for_jobs(jobs) - return [job._result_no_wait() for job in jobs] + Job.wait_all(jobs) + return [job.backend.get_result(job) for job in jobs] + + +def _wait_all_grouped(jobs: list[Job]) -> None: + """Group jobs by backend identity and wait per group.""" + groups: dict[int, tuple[Backend, list[Job]]] = {} + for job in jobs: + bid = id(job.backend) + if bid not in groups: + groups[bid] = (job.backend, []) + groups[bid][1].append(job) + for backend, group in groups.values(): + backend.wait(group) @dataclass(kw_only=True) @@ -306,50 +219,46 @@ def script(self) -> str: """ Generate a PBS job script based on job description. """ - commands: list[str] = [] - for command in self.commands: - if isinstance(command, str): - commands.append(command) - elif isinstance(command, list): - commands.append(" ".join(shlex.quote(arg) for arg in command)) - commands_str = "\n".join(commands) - - job_script = f"""#!/bin/bash -{f"#PBS -P {self.project}" if self.project else ""} -{f"#PBS -N {self.name}" if self.name else ""} -{f"#PBS -q {self.queue}" if self.queue else ""} -{f"#PBS -l ncpus={self.ncpus}" if self.ncpus else ""} -{f"#PBS -l mem={self.mem}" if self.mem else ""} -{f"#PBS -l jobfs={self.jobfs}" if self.jobfs else ""} -{f"#PBS -l walltime={self.walltime}" if self.walltime else ""} -{f"#PBS -l storage={self.storage}" if self.storage else ""} -{f"#PBS -o {self.output_path}" if self.output_path else ""} -{f"#PBS -e {self.error_path}" if self.error_path else ""} -{"#PBS -l wd" if self.wd else ""} -{f'#PBS -W depend=afterok:{":".join([job.job_id for job in self.afterok])}' if len(self.afterok) > 0 else ""} -{commands_str} -""" - return job_script + from pbspy._pbs_core import format_job_script + + return format_job_script( + name=self.name, + project=self.project, + queue=self.queue, + ncpus=self.ncpus, + mem=self.mem, + jobfs=self.jobfs, + walltime=self.walltime, + storage=self.storage, + wd=self.wd, + output_path=self.output_path, + error_path=self.error_path, + afterok_ids=[job.job_id for job in self.afterok], + commands=self.commands, + ) - def submit(self) -> Job: - """ - Submits the job to the PBS queue using ``qsub``. + def submit(self, backend: Backend | None = None) -> Job: """ - job_script = self.script() + Submits the job to the PBS queue. - process = subprocess.Popen( - ["qsub"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - stdout, stderr = process.communicate(input=job_script.encode("utf-8")) + Args: + backend: The backend to use for submission. Defaults to + :class:`~pbspy.LocalBackend` (runs ``qsub`` locally). + """ + if backend is None: + backend = _DEFAULT_LOCAL_BACKEND - if process.returncode != 0: - raise RuntimeError(f'Failed to submit job: {stderr.strip().decode("utf-8")}') - job_id = stdout.strip().decode("utf-8") + job_script = self.script() + job_id, job_name = backend.submit(job_script, self.name) if self.name is None: - self.name = _get_job_name(job_id) - - return Job(job_name=self.name, job_id=job_id, description=self.description) + self.name = job_name + + return Job( + job_name=job_name, + job_id=job_id, + description=self.description, + backend=backend, + output_path=self.output_path, + error_path=self.error_path, + ) diff --git a/src/pbspy/_backend.py b/src/pbspy/_backend.py new file mode 100644 index 0000000..98ed9e0 --- /dev/null +++ b/src/pbspy/_backend.py @@ -0,0 +1,68 @@ +""" +Backend abstract base class. + +The built-in :class:`~pbspy._local_backend.LocalBackend` runs PBS operations +in-process. Alternate implementations can provide other execution mechanisms. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pbspy import Job, JobResult + +__all__ = ["Backend"] + + +class Backend(ABC): + """ + Abstract base class for PBS operation backends. + + All methods that touch PBS (qsub, qstat, qdel, and file reading) go through + a backend so callers can supply a different implementation without changing + the :class:`~pbspy.JobDescription` / :class:`~pbspy.Job` API. + """ + + @abstractmethod + def submit(self, script: str, name: str | None = None) -> tuple[str, str]: + """ + Submit a PBS job script. + + Returns: + ``(job_id, job_name)`` + """ + ... + + @abstractmethod + def wait( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None = None, + ) -> None: + """ + Wait for *jobs* to finish. + + :param jobs: List of :class:`~pbspy.Job` objects. + :param on_update: Optional callback ``(job_id, state)`` on each status + change. *state* is ``None`` when the job has finished. + """ + ... + + @abstractmethod + def get_result(self, job: Job) -> JobResult: + """ + Return the :class:`~pbspy.JobResult` for a completed job. + """ + ... + + @abstractmethod + def delete(self, job_ids: list[str]) -> None: + """ + Cancel (``qdel``) the given jobs. + + :param job_ids: Job ids to cancel. + """ + ... diff --git a/src/pbspy/_local_backend.py b/src/pbspy/_local_backend.py new file mode 100644 index 0000000..4b1beaf --- /dev/null +++ b/src/pbspy/_local_backend.py @@ -0,0 +1,47 @@ +""" +LocalBackend: direct in-process calls to :mod:`pbspy._pbs_core`. + +This is the default backend when running on a host with PBS access. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TYPE_CHECKING + +import pbspy._pbs_core as core +from pbspy._backend import Backend + +if TYPE_CHECKING: + from pbspy import Job, JobResult + +__all__ = ["LocalBackend"] + + +class LocalBackend(Backend): + """ + Backend that calls PBS commands directly in the current process. + + This is the default backend; it requires ``qsub`` and ``qstat`` to be + available on ``PATH``. + """ + + def submit(self, script: str, name: str | None = None) -> tuple[str, str]: + """Submit *script* via ``qsub`` and return ``(job_id, job_name)``.""" + return core.pbs_submit(script, name) + + def wait( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None = None, + ) -> None: + """Wait for *jobs* to finish.""" + core.pbs_wait_for_jobs(jobs, on_update) + + def get_result(self, job: Job) -> JobResult: + """Return the :class:`~pbspy.JobResult` for a completed job.""" + return core.pbs_get_result(job) + + def delete(self, job_ids: list[str]) -> None: + """Cancel (``qdel``) the given jobs.""" + core.pbs_delete(job_ids) diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py new file mode 100644 index 0000000..978770f --- /dev/null +++ b/src/pbspy/_pbs_core.py @@ -0,0 +1,252 @@ +""" +Core PBS operations: qsub submission, qstat polling, and result file reading. + +These functions are used by :class:`pbspy.LocalBackend` and are also available +to applications that need lower-level PBS state queries. +""" + +from __future__ import annotations + +import re +import shlex +import subprocess +import time +from collections.abc import Callable +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from pbspy import Job, JobResult + +__all__ = [ + "PBSRunner", + "pbs_submit", + "pbs_wait_for_jobs", + "pbs_get_result", + "pbs_get_states", + "pbs_delete", + "try_get_exit_code", +] + +_POLL_INTERVAL_SECONDS = 60 + + +class PBSRunner: + """Abstracts command execution and file reading for PBS operations. Runs commands locally.""" + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.run(cmd, input=input, capture_output=True) + + def read_file(self, path: str) -> str: + """Read a local file and return its text.""" + result = self.run(["cat", path]) + if result.returncode != 0: + raise FileNotFoundError(path) + return result.stdout.decode("utf-8") + + +_LOCAL_RUNNER = PBSRunner() + + +def _get_job_name(job_id: str, runner: PBSRunner = _LOCAL_RUNNER) -> str: + process = runner.run(["qstat", "-fx", job_id]) + if process.returncode != 0: + raise RuntimeError(f"Failed to get job status: {process.stderr.decode('utf-8')}") + output = process.stdout.decode("utf-8") + re_match = re.search(r"Job_Name = ([^\s]+)", output) + if re_match is not None: + group = re_match.group(1) + assert group + return group + raise RuntimeError("Could not retrieve job name") + + +def try_get_exit_code(job_id: str, runner: PBSRunner = _LOCAL_RUNNER) -> int | None: + """Return the exit status of a finished PBS job, or ``None`` if unavailable.""" + process = runner.run(["qstat", "-fx", job_id]) + if process.returncode != 0: + raise RuntimeError(f"Failed to get job status: {process.stderr.decode('utf-8')}") + output = process.stdout.decode("utf-8") + exit_code_match = re.search(r"Exit_status = (\d+)", output) + if exit_code_match is not None: + exit_code_group = exit_code_match.group(1) + assert exit_code_group + return int(exit_code_group) + return None + + +def pbs_submit(script: str, name: str | None, runner: PBSRunner = _LOCAL_RUNNER) -> tuple[str, str]: + """ + Submit a PBS job script via ``qsub``. + + Returns: + A tuple of ``(job_id, job_name)``. + """ + process = runner.run(["qsub"], input=script.encode("utf-8")) + if process.returncode != 0: + raise RuntimeError(f"Failed to submit job: {process.stderr.strip().decode('utf-8')}") + job_id = process.stdout.strip().decode("utf-8") + if name is None: + name = _get_job_name(job_id, runner) + return job_id, name + + +def pbs_wait_for_jobs( + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None = None, +) -> None: + """ + Poll qstat until all jobs in *jobs* have finished. + + :param jobs: The jobs to wait for. + :param on_update: Optional callback invoked on each status change. Called + as ``on_update(job_id, state)`` where *state* is one of ``"Q"``, + ``"R"``, ``"E"``, ``"F"`` or ``None`` when the job has disappeared + from qstat (finished). + """ + task_done = [False] * len(jobs) + while not all(task_done): + time.sleep(_POLL_INTERVAL_SECONDS) + for i, job in enumerate(jobs): + if task_done[i]: + continue + process = subprocess.run(["qstat", job.job_id], capture_output=True) + if process.returncode != 0: + output = process.stdout.decode("utf-8") + if job.job_id not in output or "has finished" in process.stderr.decode("utf-8"): + task_done[i] = True + if on_update is not None: + on_update(job.job_id, None) + + +def pbs_get_states(job_ids: list[str], runner: PBSRunner = _LOCAL_RUNNER) -> dict[str, str | None]: + """ + Look up the PBS state of each id in *job_ids* with a single ``qstat`` call. + + Returns: + A mapping of job id -> state (``"Q"``, ``"R"``, ``"E"``, etc.), or ``None`` when the + id is no longer listed by ``qstat`` (i.e. the job has finished). + """ + if not job_ids: + return {} + + states: dict[str, str | None] = dict.fromkeys(job_ids) + process = runner.run(["qstat", *job_ids]) + output = process.stdout.decode("utf-8") + + for line in output.splitlines(): + line = line.strip() + if not line or line.startswith("Job id") or line.startswith("---"): + continue + fields = line.split() + if len(fields) < 5: + continue + job_id = fields[0] + if job_id in states: + states[job_id] = fields[4] + + # Any id not found in the qstat table is left as None (finished / unknown), matching the + # "job_id not in output or has finished" convention used elsewhere in this module. + return states + + +def pbs_delete(job_ids: list[str], runner: PBSRunner = _LOCAL_RUNNER) -> None: + """ + Cancel jobs via a single ``qdel`` call. Idempotent: already-finished or unknown job ids + are treated as success rather than raising. + """ + if not job_ids: + return + + process = runner.run(["qdel", *job_ids]) + if process.returncode != 0: + stderr = process.stderr.decode("utf-8") + if "has finished" not in stderr and "Unknown Job Id" not in stderr: + raise RuntimeError(f"Failed to delete job(s): {stderr.strip()}") + + +def pbs_get_result(job: Job, runner: PBSRunner = _LOCAL_RUNNER) -> JobResult: + """ + Read the output/error files for a completed job and return a :class:`~pbspy.JobResult`. + + The job must have already finished; this function does **not** wait. + """ + from pbspy import JobResult + + job_id_num = job.job_id.split(".")[0] + output_file = job.output_path or f"{job.job_name}.o{job_id_num}" + error_file = job.error_path or f"{job.job_name}.e{job_id_num}" + + try: + output = runner.read_file(output_file) + except FileNotFoundError: + output = "" + try: + error = runner.read_file(error_file) + except FileNotFoundError: + error = "" + + output_split = output.split( + "\n======================================================================================\n" + ) + exit_code = None + pbs_stats: dict[str, str] | None = None + if len(output_split) == 3: + output = output_split[0] + stats_block = output_split[1].strip() + stats_block = "\n".join(stats_block.split("\n")[1:]) # Skip "Resource usage on .." line + pbs_stats = {} + pattern = re.compile(r"\s*([^:]+):\s*([^\s]+)") + for match in pattern.finditer(stats_block): + key = match.group(1).strip() + value = match.group(2).strip() + pbs_stats[key] = value + if key == "Exit Status": + exit_code_match = re.search(r"^\d+", value) + if exit_code_match is not None: + exit_code = int(exit_code_match.group()) + + if exit_code is None: + exit_code = try_get_exit_code(job.job_id, runner) + + return JobResult(exit_code=exit_code, output=output, error=error, stats=pbs_stats) + + +def format_job_script( + name: str | None, + project: str | None, + queue: str | None, + ncpus: int | None, + mem: str | None, + jobfs: str | None, + walltime: str | None, + storage: str | None, + wd: bool, + output_path: str | None, + error_path: str | None, + afterok_ids: list[str], + commands: list[str | list[str]], +) -> str: + """Format a PBS job script from the given parameters.""" + command_strs: list[str] = [] + for command in commands: + if isinstance(command, str): + command_strs.append(command) + elif isinstance(command, list): + command_strs.append(" ".join(shlex.quote(arg) for arg in command)) + commands_str = "\n".join(command_strs) + + return f"""#!/bin/bash +{f"#PBS -P {project}" if project else ""} +{f"#PBS -N {name}" if name else ""} +{f"#PBS -q {queue}" if queue else ""} +{f"#PBS -l ncpus={ncpus}" if ncpus else ""} +{f"#PBS -l mem={mem}" if mem else ""} +{f"#PBS -l jobfs={jobfs}" if jobfs else ""} +{f"#PBS -l walltime={walltime}" if walltime else ""} +{f"#PBS -l storage={storage}" if storage else ""} +{f"#PBS -o {output_path}" if output_path else ""} +{f"#PBS -e {error_path}" if error_path else ""} +{"#PBS -l wd" if wd else ""} +{f'#PBS -W depend=afterok:{":".join(afterok_ids)}' if afterok_ids else ""} +{commands_str} +""" diff --git a/tests/test_default.py b/tests/test_default.py index 48d568b..2a01149 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -1,6 +1,13 @@ -from pbspy import JobDescription +import shutil +import subprocess +from collections.abc import Callable +import pytest +from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend + + +@pytest.mark.skipif(shutil.which("qsub") is None, reason="qsub not available") def test_default() -> None: job_description = ( JobDescription() @@ -15,3 +22,325 @@ def test_default() -> None: assert result.exit_code == 0 assert result.output == "a\nb\nc\nd\ne\n" # assert result.error == "" + + +# --------------------------------------------------------------------------- +# Mock backend tests — do not require qsub/qstat +# --------------------------------------------------------------------------- + + +class _MockBackend(Backend): + """In-process backend that records calls without invoking PBS.""" + + def __init__(self) -> None: + self.submitted: list[tuple[str, str | None]] = [] + self.waited: list[list[Job]] = [] + self.results: dict[str, JobResult] = {} + self.deleted: list[list[str]] = [] + + def submit(self, script: str, name: str | None = None) -> tuple[str, str]: + self.submitted.append((script, name)) + job_name = name or "mock_job" + job_id = "99999.mock" + self.results[job_id] = JobResult(exit_code=0, output="mock output\n") + return job_id, job_name + + def wait(self, jobs: list[Job], on_update: Callable[[str, str | None], None] | None = None) -> None: + self.waited.append(list(jobs)) + + def get_result(self, job: Job) -> JobResult: + return self.results.get(job.job_id, JobResult(exit_code=None)) + + def delete(self, job_ids: list[str]) -> None: + self.deleted.append(list(job_ids)) + + +def test_mock_backend_submit() -> None: + """JobDescription.submit() uses the provided backend.""" + mock = _MockBackend() + jd = JobDescription(name="test_job", ncpus=1, walltime="00:01:00") + jd.add_command(["echo", "hello"]) + + job = jd.submit(backend=mock) + + assert len(mock.submitted) == 1 + script, _ = mock.submitted[0] + assert "echo hello" in script + assert job.job_id == "99999.mock" + assert job.job_name == "test_job" + assert job.backend is mock + + +def test_mock_backend_result() -> None: + """Job.result() delegates to the backend.""" + mock = _MockBackend() + jd = JobDescription(name="test_job") + jd.add_command(["echo", "hello"]) + job = jd.submit(backend=mock) + + result = job.result() + + assert len(mock.waited) == 1 + assert result.exit_code == 0 + assert result.output == "mock output\n" + + +def test_mock_backend_result_all() -> None: + """Job.result_all() groups jobs and returns all results.""" + mock = _MockBackend() + mock.results["1.mock"] = JobResult(exit_code=0, output="A\n") + mock.results["2.mock"] = JobResult(exit_code=0, output="B\n") + + job_a = Job(job_id="1.mock", job_name="job_a", backend=mock) + job_b = Job(job_id="2.mock", job_name="job_b", backend=mock) + + results = Job.result_all([job_a, job_b]) + + assert len(results) == 2 + assert results[0].output == "A\n" + assert results[1].output == "B\n" + + +def test_script_generation_unchanged() -> None: + """script() output is unchanged from the original implementation.""" + jd = JobDescription( + name="myjob", + project="ab01", + queue="normal", + ncpus=4, + mem="16GB", + walltime="01:00:00", + wd=True, + ) + jd.add_command(["echo", "hello world"]) + script = jd.script() + + assert "#PBS -N myjob" in script + assert "#PBS -P ab01" in script + assert "#PBS -q normal" in script + assert "#PBS -l ncpus=4" in script + assert "#PBS -l mem=16GB" in script + assert "#PBS -l walltime=01:00:00" in script + assert "#PBS -l wd" in script + assert "echo 'hello world'" in script + + +def test_backend_classes_importable() -> None: + """Backend and LocalBackend are importable from pbspy.""" + assert issubclass(LocalBackend, Backend) + + +def test_default_submissions_share_local_backend(monkeypatch: pytest.MonkeyPatch) -> None: + """Default submissions share a backend so wait_all can group their polling.""" + submitted = iter([("1.mock", "job_a"), ("2.mock", "job_b")]) + monkeypatch.setattr(LocalBackend, "submit", lambda *_args, **_kwargs: next(submitted)) + + job_a = JobDescription(name="job_a").submit() + job_b = JobDescription(name="job_b").submit() + + assert job_a.backend is job_b.backend + + +def test_submit_carries_output_and_error_paths() -> None: + """JobDescription.submit() propagates output_path/error_path to the returned Job.""" + mock = _MockBackend() + jd = JobDescription( + name="test_job", + output_path="/scratch/project/job.out", + error_path="/scratch/project/job.err", + ) + jd.add_command(["echo", "hello"]) + + job = jd.submit(backend=mock) + + assert job.output_path == "/scratch/project/job.out" + assert job.error_path == "/scratch/project/job.err" + + +def test_submit_default_paths_are_none() -> None: + """When output_path/error_path are not set, Job gets None for both.""" + mock = _MockBackend() + jd = JobDescription(name="test_job") + jd.add_command(["echo", "hello"]) + + job = jd.submit(backend=mock) + + assert job.output_path is None + assert job.error_path is None + + +def test_pbs_get_result_uses_custom_paths() -> None: + """pbs_get_result reads from custom output_path/error_path when set.""" + from pbspy._pbs_core import PBSRunner, pbs_get_result + + class _FileTrackingRunner(PBSRunner): + def __init__(self) -> None: + super().__init__() + self.read_paths: list[str] = [] + + def read_file(self, path: str) -> str: + self.read_paths.append(path) + if "out" in path: + return "custom stdout\n" + return "custom stderr\n" + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + # Override to avoid real qstat calls + return subprocess.CompletedProcess(cmd, 0, stdout=b"", stderr=b"") + + runner = _FileTrackingRunner() + job = Job( + job_id="123.mock", + job_name="myjob", + output_path="/scratch/job.out", + error_path="/scratch/job.err", + ) + + result = pbs_get_result(job, runner=runner) + + # Should read from custom paths, not default names + assert "/scratch/job.out" in runner.read_paths + assert "/scratch/job.err" in runner.read_paths + assert result.output == "custom stdout\n" + assert result.error == "custom stderr\n" + + +def test_pbs_get_result_uses_default_paths_when_none() -> None: + """pbs_get_result falls back to {job_name}.o{job_id} when paths are None.""" + from pbspy._pbs_core import PBSRunner, pbs_get_result + + class _FileTrackingRunner(PBSRunner): + def __init__(self) -> None: + super().__init__() + self.read_paths: list[str] = [] + + def read_file(self, path: str) -> str: + self.read_paths.append(path) + return "" + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.CompletedProcess(cmd, 0, stdout=b"", stderr=b"") + + runner = _FileTrackingRunner() + job = Job( + job_id="456.mock", + job_name="defaultjob", + output_path=None, + error_path=None, + ) + + result = pbs_get_result(job, runner=runner) + + assert "defaultjob.o456" in runner.read_paths + assert "defaultjob.e456" in runner.read_paths + assert result.output == "" + assert result.error == "" + + +# --------------------------------------------------------------------------- +# pbs_get_states / pbs_delete / Job.cancel +# --------------------------------------------------------------------------- + + +def test_pbs_get_states_single_qstat_call() -> None: + """pbs_get_states issues exactly one qstat call for multiple job ids.""" + from pbspy._pbs_core import PBSRunner, pbs_get_states + + class _RecordingRunner(PBSRunner): + def __init__(self) -> None: + self.calls: list[list[str]] = [] + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + self.calls.append(cmd) + stdout = ( + b"Job id Name User Time Use S Queue\n" + b"---------------- ---------------- ---------------- -------- - -----\n" + b"123.gadi-pbs job_a user00 00:00:12 R normal\n" + b"124.gadi-pbs job_b user00 00:00:00 Q normal\n" + ) + return subprocess.CompletedProcess(cmd, 0, stdout=stdout, stderr=b"") + + runner = _RecordingRunner() + states = pbs_get_states(["123.gadi-pbs", "124.gadi-pbs", "125.gadi-pbs"], runner=runner) + + assert len(runner.calls) == 1 + assert runner.calls[0] == ["qstat", "123.gadi-pbs", "124.gadi-pbs", "125.gadi-pbs"] + assert states == { + "123.gadi-pbs": "R", + "124.gadi-pbs": "Q", + "125.gadi-pbs": None, # not listed => finished + } + + +def test_pbs_get_states_empty_list() -> None: + """pbs_get_states with no job ids makes no qstat call and returns an empty dict.""" + from pbspy._pbs_core import PBSRunner, pbs_get_states + + class _FailIfCalledRunner(PBSRunner): + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + raise AssertionError("qstat should not be called for an empty job list") + + assert pbs_get_states([], runner=_FailIfCalledRunner()) == {} + + +def test_pbs_delete_single_qdel_call() -> None: + """pbs_delete issues exactly one qdel call for multiple job ids.""" + from pbspy._pbs_core import PBSRunner, pbs_delete + + class _RecordingRunner(PBSRunner): + def __init__(self) -> None: + self.calls: list[list[str]] = [] + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + self.calls.append(cmd) + return subprocess.CompletedProcess(cmd, 0, stdout=b"", stderr=b"") + + runner = _RecordingRunner() + pbs_delete(["123.gadi-pbs", "124.gadi-pbs"], runner=runner) + + assert len(runner.calls) == 1 + assert runner.calls[0] == ["qdel", "123.gadi-pbs", "124.gadi-pbs"] + + +def test_pbs_delete_idempotent_for_finished_jobs() -> None: + """pbs_delete does not raise when qdel reports the job has already finished.""" + from pbspy._pbs_core import PBSRunner, pbs_delete + + class _FinishedRunner(PBSRunner): + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.CompletedProcess(cmd, 1, stdout=b"", stderr=b"qdel: Job has finished 123.gadi-pbs") + + pbs_delete(["123.gadi-pbs"], runner=_FinishedRunner()) # should not raise + + +def test_pbs_delete_idempotent_for_unknown_jobs() -> None: + """pbs_delete does not raise when qdel reports an unknown job id.""" + from pbspy._pbs_core import PBSRunner, pbs_delete + + class _UnknownRunner(PBSRunner): + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.CompletedProcess(cmd, 1, stdout=b"", stderr=b"qdel: Unknown Job Id 123.gadi-pbs") + + pbs_delete(["123.gadi-pbs"], runner=_UnknownRunner()) # should not raise + + +def test_pbs_delete_raises_on_other_errors() -> None: + """pbs_delete raises for genuine errors (not the finished/unknown cases).""" + from pbspy._pbs_core import PBSRunner, pbs_delete + + class _BrokenRunner(PBSRunner): + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.CompletedProcess(cmd, 1, stdout=b"", stderr=b"qdel: permission denied") + + with pytest.raises(RuntimeError, match="permission denied"): + pbs_delete(["123.gadi-pbs"], runner=_BrokenRunner()) + + +def test_job_cancel_calls_backend_delete() -> None: + """Job.cancel() delegates to backend.delete([job_id]).""" + mock = _MockBackend() + job = Job(job_id="123.mock", job_name="job_a", backend=mock) + + job.cancel() + + assert mock.deleted == [["123.mock"]]