From 210d0ef615cb77d0d0f0f56675e8100e4c2437df Mon Sep 17 00:00:00 2001 From: omdowley Date: Fri, 29 May 2026 13:55:26 +1000 Subject: [PATCH 01/16] feat: add SSH backend and pbspy-server daemon Previously pbspy could only submit jobs when running directly on a supercomputer login node. This adds a client/server architecture so jobs can be submitted from any machine with SSH access. src/pbspy/__init__.py has been split into discrete modules: _backend.py - Backend Abstract Base Class _local_backend.py - LocalBackend: direct qsub/qstat calls (default) _ssh_backend.py - SSHBackend: talks to the daemon over SSH _pbs_core.py - Core PBS operations extracted from __init__.py _protocol.py - Length-prefixed pickle framing (send_frame / recv_frame) and request/response dataclasses A new top-level package pbspy_server contains: server.py - TCP daemon: accepts connections, dispatches requests, polls qstat in a background thread, streams status updates to waiting clients, shuts down automatically when no jobs are pending and no clients are connected proxy.py - SSH proxy mode: acquires a file lock, starts the daemon if needed, then splices stdin/stdout with the daemon socket so the remote client communicates transparently __main__.py - CLI entry point (--daemon / --proxy / --status) pyproject.toml adds the `pbspy-server` console script. `pbspy-server` binds a TCP port on all interfaces so that proxy processes on other login nodes can reach it. To prevent unauthorized access (pickle deserialization is remote code exec), every connection must authenticate with a shared secret token: - On startup the server generates secrets.token_hex(32) and writes the addr file on the login-node as JSON: {"host": ..., "port": ..., "token": ...} - Clients send AuthRequest(token=...) as the first frame; the server responds with AuthOkResponse or ErrorResponse and closes the connection - The proxy reads the token from the JSON addr file and authenticates using it - pbspy-server --status and _daemon_alive() also authenticate before sending PingRequest JobDescription.submit() now accepts an optional `backend` parameter (defaults to LocalBackend). Pass an SSHBackend instance to submit to a remote supercomputer: ```python backend = SSHBackend("user@gadi.nci.org.au") job = JobDescription(...).add_command([...]).submit(backend=backend) ``` Backend, LocalBackend, and SSHBackend are exported from pbspy.__all__. Added tests to new code: tests/test_protocol.py tests/test_server.py tests/test_default.py Updated pre-commit: .pre-commit-config.yaml: added pytest to mypy's additional_dependencies so the type stubs for pytest decorators are available during the hook run. --- .pre-commit-config.yaml | 1 + README.md | 43 +++- examples/run.py | 10 +- pyproject.toml | 3 + src/pbspy/__init__.py | 315 +++++++---------------------- src/pbspy/_backend.py | 62 ++++++ src/pbspy/_local_backend.py | 105 ++++++++++ src/pbspy/_pbs_core.py | 195 ++++++++++++++++++ src/pbspy/_protocol.py | 190 ++++++++++++++++++ src/pbspy/_ssh_backend.py | 226 +++++++++++++++++++++ src/pbspy_server/__init__.py | 1 + src/pbspy_server/__main__.py | 100 ++++++++++ src/pbspy_server/proxy.py | 174 ++++++++++++++++ src/pbspy_server/server.py | 305 ++++++++++++++++++++++++++++ tests/test_default.py | 114 ++++++++++- tests/test_protocol.py | 185 +++++++++++++++++ tests/test_server.py | 373 +++++++++++++++++++++++++++++++++++ 17 files changed, 2150 insertions(+), 252 deletions(-) create mode 100644 src/pbspy/_backend.py create mode 100644 src/pbspy/_local_backend.py create mode 100644 src/pbspy/_pbs_core.py create mode 100644 src/pbspy/_protocol.py create mode 100644 src/pbspy/_ssh_backend.py create mode 100644 src/pbspy_server/__init__.py create mode 100644 src/pbspy_server/__main__.py create mode 100644 src/pbspy_server/proxy.py create mode 100644 src/pbspy_server/server.py create mode 100644 tests/test_protocol.py create mode 100644 tests/test_server.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 18ee742..659ce06 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -50,6 +50,7 @@ repos: additional_dependencies: - rich - typer + - pytest - repo: https://github.com/scientific-python/cookie rev: 2024.04.23 hooks: diff --git a/README.md b/README.md index 9bfbe0f..91dcc8a 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,14 +40,49 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -Output (partially executed): +### Submitting from a remote machine via SSH + +Install `pbspy` on your supercomputer via the mechanism of your choice, for instance: + +```bash +uv tool install pbspy +``` + +Then use `SSHBackend` on your local machine — the server daemon starts automatically on first use +and is reused by subsequent connections: + +```python +from pbspy import Job, JobDescription, SSHBackend + +backend = SSHBackend("user@gadi.nci.org.au") + +job_a = ( + JobDescription(name="job_a", ncpus=4, mem="192GB", walltime="00:05:00") + .add_command(["echo", "A"]) + .submit(backend=backend) +) + +job_b = ( + JobDescription(name="job_b", ncpus=1, walltime="00:05:00", afterok=[job_a]) + .add_command(["echo", "B"]) + .submit(backend=backend) +) + +(result_a, result_b) = Job.result_all([job_a, job_b]) +print("job_a:", result_a.output.strip()) +print("job_b:", result_b.output.strip()) +``` + +The `SSHBackend` uses your existing SSH configuration (keys, `~/.ssh/config` aliases, ssh-agent). + +## Output (partially executed) ```text ✓ 124397435.gadi-pbs job_a 0:01:15 124397436.gadi-pbs job_b ``` -Output (completed): +## Output (completed) ```text ✓ 124397435.gadi-pbs job_a @@ -55,6 +92,6 @@ job_a: A job_b: B ``` -## Licence +## License `pbspy` is licensed under the MIT License [LICENSE](./LICENSE) or . diff --git a/examples/run.py b/examples/run.py index 5d8c612..35b23a1 100644 --- a/examples/run.py +++ b/examples/run.py @@ -1,10 +1,12 @@ import typer from rich import print -from pbspy import Job, JobDescription, QueueLimits +from pbspy import Job, JobDescription, QueueLimits, SSHBackend -def main() -> None: +def main(remote: str = "") -> None: + backend = SSHBackend(remote) if remote else None + # Run a job with some explicit parameters job_a = ( JobDescription( @@ -14,7 +16,7 @@ def main() -> None: walltime="00:05:00", ) .add_command(["echo", "A"]) - .submit() + .submit(backend=backend) ) # Submit another job that waits for job_a to finish @@ -30,7 +32,7 @@ def main() -> None: afterok=[job_a], ) .add_command(["echo", "B"]) - .submit() + .submit(backend=backend) ) # Get the result of the jobs diff --git a/pyproject.toml b/pyproject.toml index 9d54d1f..5369551 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,9 @@ dependencies = [ "rich~=13.8", ] +[project.scripts] +pbspy-server = "pbspy_server.__main__:main" + [dependency-groups] dev = [ "mypy>=1.14.1", diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index e057e38..f6e98af 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -4,161 +4,28 @@ from __future__ import annotations -import re -import shlex -import subprocess -import time from dataclasses import dataclass, field from typing import Any, Self, TypeAlias -from rich.progress import Progress, TextColumn, TimeElapsedColumn +from pbspy._backend import Backend +from pbspy._local_backend import LocalBackend +from pbspy._ssh_backend import SSHBackend -__all__ = ["JobDescription", "Job", "JobResult", "QueueLimits", "QueueLimitsMap", "gadi"] +# 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], progress: bool = True) -> None: - """ - Waits for the PBS jobs to complete. - - Args: - jobs (list[Job]): A list of Job objects representing the PBS jobs. - progress (bool): Whether or not to print the state of the job and status on exit. - - Returns: - None - """ - if progress: - _pbs_wait_for_jobs_with_progress(jobs) - else: - _pbs_wait_for_jobs_quiet(jobs) - - -def _pbs_wait_for_jobs_quiet(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')}") - - -def _pbs_wait_for_jobs_with_progress(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 - """ - last_check = time.time() - with Progress( - TimeElapsedColumn(), - TextColumn("[progress.description]{task.description}"), - auto_refresh=False, - ) as progress: - tasks = [progress.add_task(f"{job.job_id} {job.job_name} {job.description or ''}", total=1) for job in jobs] - task_done = [False] * len(jobs) - while not all(task_done): - time.sleep(1) - if time.time() - last_check > 60: - last_check = time.time() - for i, (job, task) in enumerate(zip(jobs, tasks, strict=False)): - 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 - progress.advance(task) - progress.update(task, visible=False) - - # Try and get the exit code - exit_code = _try_get_exit_code(job.job_id) - if exit_code is None: - output_status = "[yellow]?" - else: - output_status = "[green]✓" if exit_code == 0 else "[red]✗" - progress.console.print( - f'{output_status} {job.job_id} {job.job_name} {job.description or ""}' - ) - # else: - # raise RuntimeError(f"Failed to check job status: {process.stderr.decode('utf-8')}") - progress.refresh() +__all__ = [ + "JobDescription", + "Job", + "JobResult", + "QueueLimits", + "QueueLimitsMap", + "Backend", + "LocalBackend", + "SSHBackend", + "gadi", +] @dataclass @@ -196,6 +63,9 @@ class Job: description: str | None = None """A description of the job for progress updates. Unused by PBS.""" + backend: Backend = field(default=_DEFAULT_LOCAL_BACKEND, repr=False, compare=False) + """The backend used to submit and track this job.""" + def wait(self, progress: bool = True) -> None: """ Wait for the job to complete. @@ -203,56 +73,7 @@ def wait(self, progress: bool = True) -> None: Args: progress (bool): Whether or not to print the state of the job and status on exit. """ - _pbs_wait_for_jobs([self], progress=progress) - - 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], progress=progress) def result(self, progress: bool = True) -> JobResult: """ @@ -262,7 +83,7 @@ def result(self, progress: bool = True) -> JobResult: progress (bool): Whether or not to print the state of the job and status on exit. """ self.wait(progress=progress) - return self._result_no_wait() + return self.backend.get_result(self) # type: ignore[return-value] @staticmethod def wait_all(jobs: list[Job], progress: bool = True) -> None: @@ -272,7 +93,10 @@ def wait_all(jobs: list[Job], progress: bool = True) -> None: Args: progress (bool): Whether or not to print the state of the job and status on exit. """ - _pbs_wait_for_jobs(jobs, progress=progress) + if not jobs: + return + # Group by backend so each backend can wait for its own jobs efficiently + _wait_all_grouped(jobs, progress=progress) @staticmethod def result_all(jobs: list[Job], progress: bool = True) -> list[JobResult]: @@ -282,8 +106,20 @@ def result_all(jobs: list[Job], progress: bool = True) -> list[JobResult]: Args: progress (bool): Whether or not to print the state of the job and status on exit. """ - _pbs_wait_for_jobs(jobs, progress=progress) - return [job._result_no_wait() for job in jobs] + Job.wait_all(jobs, progress=progress) + return [job.backend.get_result(job) for job in jobs] # type: ignore[misc] + + +def _wait_all_grouped(jobs: list[Job], progress: bool) -> 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, progress=progress) @dataclass(kw_only=True) @@ -387,50 +223,41 @@ 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: + def submit(self, backend: Backend | None = None) -> Job: """ - Submits the job to the PBS queue using ``qsub``. - """ - 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). + Pass a :class:`~pbspy.SSHBackend` instance to submit to a + remote supercomputer over SSH. + """ + if backend is None: + backend = LocalBackend() - 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) + self.name = job_name - return Job(job_name=self.name, job_id=job_id, description=self.description) + return Job(job_name=job_name, job_id=job_id, description=self.description, backend=backend) diff --git a/src/pbspy/_backend.py b/src/pbspy/_backend.py new file mode 100644 index 0000000..c7f4a85 --- /dev/null +++ b/src/pbspy/_backend.py @@ -0,0 +1,62 @@ +""" +Backend abstract base class. + +Concrete backends: :class:`~pbspy._local_backend.LocalBackend` (default, +in-process) and :class:`~pbspy._ssh_backend.SSHBackend` (remote via SSH). +""" + +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 + +__all__ = ["Backend"] + + +class Backend(ABC): + """ + Abstract base class for PBS operation backends. + + All methods that touch PBS (qsub, qstat, file reading) go through a + Backend so the same :class:`~pbspy.JobDescription` / :class:`~pbspy.Job` + API works both locally and over SSH. + """ + + @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, + progress: bool = True, + ) -> None: + """ + Wait for *jobs* to finish. + + Args: + jobs: List of :class:`~pbspy.Job` objects. + on_update: Optional callback ``(job_id, state)`` on each status + change. *state* is ``None`` when the job has finished. + progress: Whether to show a progress display. + """ + ... + + @abstractmethod + def get_result(self, job: object) -> object: + """ + Return the :class:`~pbspy.JobResult` for a completed job. + """ + ... diff --git a/src/pbspy/_local_backend.py b/src/pbspy/_local_backend.py new file mode 100644 index 0000000..eacf2ac --- /dev/null +++ b/src/pbspy/_local_backend.py @@ -0,0 +1,105 @@ +""" +LocalBackend: direct in-process calls to :mod:`pbspy._pbs_core`. + +No sockets, no IPC, no serialisation. This is the default backend when +running directly on the supercomputer — it behaves identically to the +original pbspy code but delegates through the :class:`Backend` interface. +""" + +from __future__ import annotations + +import subprocess +import time +from collections.abc import Callable +from typing import TYPE_CHECKING + +from rich.progress import Progress, TextColumn, TimeElapsedColumn + +import pbspy._pbs_core as core +from pbspy._backend import Backend + +if TYPE_CHECKING: + from pbspy import Job + +__all__ = ["LocalBackend"] + +_PROGRESS_REFRESH_HZ = 1.0 + + +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, + progress: bool = True, + ) -> None: + """Wait for *jobs* to finish, optionally showing a progress display.""" + if progress: + self._wait_with_progress(jobs, on_update) + else: + core.pbs_wait_for_jobs(jobs, on_update) + + def get_result(self, job: object) -> object: # job: Job -> JobResult + """Return the :class:`~pbspy.JobResult` for a completed job.""" + return core.pbs_get_result(job) # type: ignore[arg-type] + + def _wait_with_progress( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None, + ) -> None: + last_check = time.time() - core._POLL_INTERVAL_SECONDS # poll immediately + + with Progress( + TimeElapsedColumn(), + TextColumn("[progress.description]{task.description}"), + auto_refresh=False, + ) as progress_bar: + tasks = { + job.job_id: progress_bar.add_task( + f"{job.job_id} {job.job_name or ''} {job.description or ''}", + total=1, + ) + for job in jobs + } + task_done: dict[str, bool] = {job.job_id: False for job in jobs} + + while not all(task_done.values()): + time.sleep(1.0 / _PROGRESS_REFRESH_HZ) + progress_bar.refresh() + + if time.time() - last_check < core._POLL_INTERVAL_SECONDS: + continue + last_check = time.time() + + for job in jobs: + if task_done[job.job_id]: + continue + process = subprocess.run(["qstat", job.job_id], capture_output=True) + if process.returncode != 0: + stdout = process.stdout.decode("utf-8") + if job.job_id not in stdout or "has finished" in process.stderr.decode("utf-8"): + task_done[job.job_id] = True + progress_bar.advance(tasks[job.job_id]) + progress_bar.update(tasks[job.job_id], visible=False) + + exit_code = core.try_get_exit_code(job.job_id) + status = ( + "[green]✓" if exit_code == 0 else "[red]✗" if exit_code is not None else "[yellow]?" + ) + progress_bar.console.print( + f"{status} {job.job_id} {job.job_name or ''} {job.description or ''}" + ) + if on_update: + on_update(job.job_id, None) diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py new file mode 100644 index 0000000..47c949f --- /dev/null +++ b/src/pbspy/_pbs_core.py @@ -0,0 +1,195 @@ +""" +Core PBS operations: qsub submission, qstat polling, and result file reading. + +These functions are called directly by LocalBackend (in-process) and by the +server daemon when handling pickled requests from SSHBackend. +""" + +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__ = ["pbs_submit", "pbs_wait_for_jobs", "pbs_get_result", "try_get_exit_code"] + +_POLL_INTERVAL_SECONDS = 60 +_PROGRESS_REFRESH_SECONDS = 1 + + +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") + 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) -> int | None: + """Return the exit status of a finished PBS job, or ``None`` if unavailable.""" + 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") + 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) -> tuple[str, str]: + """ + Submit a PBS job script via ``qsub``. + + Returns: + A tuple of ``(job_id, job_name)``. + """ + process = subprocess.Popen( + ["qsub"], + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + stdout, stderr = process.communicate(input=script.encode("utf-8")) + if process.returncode != 0: + raise RuntimeError(f'Failed to submit job: {stderr.strip().decode("utf-8")}') + job_id = stdout.strip().decode("utf-8") + if name is None: + name = _get_job_name(job_id) + 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. + + Args: + jobs: The jobs to wait for. + 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) + last_check = time.time() - _POLL_INTERVAL_SECONDS # poll immediately on first iteration + + while not all(task_done): + time.sleep(_PROGRESS_REFRESH_SECONDS) + if time.time() - last_check >= _POLL_INTERVAL_SECONDS: + last_check = time.time() + 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_result(job: Job) -> 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 = f"{job.job_name}.o{job_id_num}" + error_file = f"{job.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 = "" + + 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) + + 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/src/pbspy/_protocol.py b/src/pbspy/_protocol.py new file mode 100644 index 0000000..a3135c5 --- /dev/null +++ b/src/pbspy/_protocol.py @@ -0,0 +1,190 @@ +""" +Request and response dataclasses for the pbspy client-server protocol. + +Messages are exchanged as length-prefixed pickle frames (see :func:`send_frame` +and :func:`recv_frame`). :class:`~pbspy._local_backend.LocalBackend` uses the +same dataclasses but never pickles them — it calls :mod:`pbspy._pbs_core` +directly and passes Python objects in-process. +""" + +from __future__ import annotations + +import io +import pickle +import struct +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, BinaryIO + +if TYPE_CHECKING: + from pbspy import Job, JobResult + +__all__ = [ + # Requests + "AuthRequest", + "SubmitRequest", + "WaitRequest", + "ResultRequest", + "PingRequest", + # Responses + "AuthOkResponse", + "SubmittedResponse", + "StatusUpdateResponse", + "WaitDoneResponse", + "ResultResponse", + "PongResponse", + "ErrorResponse", + # Framing + "send_frame", + "recv_frame", +] + +_FRAME_HEADER = struct.Struct("!I") # 4-byte big-endian unsigned int + + +# --------------------------------------------------------------------------- +# Requests +# --------------------------------------------------------------------------- + + +@dataclass +class AuthRequest: + """ + Authentication handshake — must be the first frame sent on every connection. + + The server responds with :class:`AuthOkResponse` on success or + :class:`ErrorResponse` (and closes the connection) on failure. + """ + + token: str + + +@dataclass +class SubmitRequest: + """Ask the server to submit a PBS job script.""" + + script: str + name: str | None = None + + +@dataclass +class WaitRequest: + """ + Ask the server to wait for *jobs* and stream :class:`StatusUpdateResponse` + messages until all jobs finish, then send :class:`WaitDoneResponse`. + """ + + jobs: list[Job] = field(default_factory=list) + + +@dataclass +class ResultRequest: + """Ask the server to return the result of a completed job.""" + + job: Job + + +@dataclass +class PingRequest: + """Liveness check; server responds with :class:`PongResponse`.""" + + +# --------------------------------------------------------------------------- +# Responses +# --------------------------------------------------------------------------- + + +@dataclass +class AuthOkResponse: + """Returned after a successful :class:`AuthRequest`.""" + + +@dataclass +class SubmittedResponse: + """Returned after a successful :class:`SubmitRequest`.""" + + job: Job + + +@dataclass +class StatusUpdateResponse: + """ + Streamed periodically during a :class:`WaitRequest`. + + *state* is one of ``"Q"`` (queued), ``"R"`` (running), ``"E"`` (exiting), + ``"F"`` (finished), or ``None`` when the job has left the qstat queue. + """ + + job: Job + state: str | None + + +@dataclass +class WaitDoneResponse: + """Sent after all jobs in a :class:`WaitRequest` have finished.""" + + jobs: list[Job] = field(default_factory=list) + + +@dataclass +class ResultResponse: + """Returned after a successful :class:`ResultRequest`.""" + + job: Job + result: JobResult + + +@dataclass +class PongResponse: + """Response to :class:`PingRequest`.""" + + +@dataclass +class ErrorResponse: + """Returned when the server encounters an error handling a request.""" + + message: str + + +# --------------------------------------------------------------------------- +# Framing helpers +# --------------------------------------------------------------------------- + + +def send_frame(stream: BinaryIO, obj: object) -> None: + """ + Serialise *obj* with :mod:`pickle` and write it to *stream* as a + length-prefixed frame. + + Frame layout: ``[4-byte big-endian length][pickle bytes]`` + """ + data = pickle.dumps(obj) + stream.write(_FRAME_HEADER.pack(len(data))) + stream.write(data) + if hasattr(stream, "flush"): + stream.flush() + + +def recv_frame(stream: BinaryIO) -> object: + """ + Read one length-prefixed frame from *stream* and return the unpickled + object. + + Raises :class:`EOFError` if the stream is closed before a complete frame + is received. + """ + header = _read_exactly(stream, _FRAME_HEADER.size) + (length,) = _FRAME_HEADER.unpack(header) + data = _read_exactly(stream, length) + return pickle.loads(data) # noqa: S301 + + +def _read_exactly(stream: BinaryIO, n: int) -> bytes: + buf = io.BytesIO() + remaining = n + while remaining > 0: + chunk = stream.read(remaining) + if not chunk: + raise EOFError("Stream closed before a complete frame was received") + buf.write(chunk) + remaining -= len(chunk) + return buf.getvalue() diff --git a/src/pbspy/_ssh_backend.py b/src/pbspy/_ssh_backend.py new file mode 100644 index 0000000..037d8a1 --- /dev/null +++ b/src/pbspy/_ssh_backend.py @@ -0,0 +1,226 @@ +""" +SSHBackend: communicates with a remote pbspy-server daemon over SSH. + +The backend opens an SSH subprocess running ``pbspy-server --proxy`` on the +remote host. Messages are exchanged as length-prefixed pickle frames (see +:mod:`pbspy._protocol`). The proxy on the remote end automatically starts the +daemon if it is not already running (emacsclient-style). +""" + +from __future__ import annotations + +import select +import subprocess +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING, BinaryIO, cast + +import pbspy._protocol as proto +from pbspy._backend import Backend + +if TYPE_CHECKING: + from pbspy import Job + +__all__ = ["SSHBackend"] + + +class SSHBackend(Backend): + """ + Backend that submits and tracks PBS jobs on a remote supercomputer via SSH. + + Args: + host: SSH destination (e.g. ``"user@gadi.nci.org.au"``). + Any option accepted by ``ssh`` (host aliases, ``-i``, etc.) works + because the system ``ssh`` binary is used. + ssh_args: Extra arguments forwarded to the ``ssh`` command (e.g. + ``["-i", "/path/to/key"]``). + """ + + def __init__(self, host: str, ssh_args: list[str] | None = None) -> None: + self._host = host + self._ssh_args = ssh_args or [] + self._proc: subprocess.Popen[bytes] | None = None + self._stdin: BinaryIO | None = None + self._stdout: BinaryIO | None = None + self._lock = threading.Lock() + self._connect() + + def submit(self, script: str, name: str | None = None) -> tuple[str, str]: + response = self._rpc(proto.SubmitRequest(script=script, name=name)) + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.SubmittedResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.job.job_id, response.job.job_name or "" + + def wait( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None = None, + progress: bool = True, + ) -> None: + """ + Send a WaitRequest and stream StatusUpdateResponse frames until + WaitDoneResponse is received. Optionally shows a progress display. + """ + + if progress: + self._wait_with_progress(jobs, on_update) + else: + self._wait_quiet(jobs, on_update) + + def get_result(self, job: object) -> object: + response = self._rpc(proto.ResultRequest(job=job)) # type: ignore[arg-type] + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.ResultResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.result + + def _connect(self) -> None: + """Start the SSH subprocess and verify the proxy is alive.""" + cmd = ["ssh", *self._ssh_args, self._host, "pbspy-server", "--proxy"] + self._proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + assert self._proc.stdin is not None + assert self._proc.stdout is not None + self._stdin = cast(BinaryIO, self._proc.stdin) + self._stdout = cast(BinaryIO, self._proc.stdout) + + try: + pong = self._rpc(proto.PingRequest()) + except (OSError, EOFError) as exc: + stderr_text = self._read_stderr() + raise RuntimeError( + f"pbspy-server --proxy did not respond to ping on {self._host!r}: {exc}" + + (f"\nSSH stderr: {stderr_text}" if stderr_text else "") + ) from exc + if not isinstance(pong, proto.PongResponse): + stderr_text = self._read_stderr() + raise RuntimeError( + f"pbspy-server --proxy returned unexpected response: {pong!r}" + + (f"\nSSH stderr: {stderr_text}" if stderr_text else "") + ) + + def _read_stderr(self) -> str: + """Non-blocking read of any immediately available stderr from the SSH process.""" + if self._proc is None or self._proc.stderr is None: + return "" + try: + rlist, _, _ = select.select([self._proc.stderr], [], [], 0.0) + if rlist: + return self._proc.stderr.read(4096).decode(errors="replace") + except OSError: + pass + return "" + + def _reconnect(self) -> None: + """Terminate the current SSH subprocess and establish a fresh connection.""" + if self._proc is not None: + try: + self._proc.terminate() + except OSError: + pass + self._proc = None + self._connect() + + def _rpc(self, request: object) -> object: + """Send one request frame and return the first response frame.""" + with self._lock: + try: + assert self._stdin is not None + assert self._stdout is not None + proto.send_frame(self._stdin, request) + return proto.recv_frame(self._stdout) + except (OSError, EOFError): + # Connection lost — attempt one reconnect then retry. + self._reconnect() + assert self._stdin is not None + assert self._stdout is not None + proto.send_frame(self._stdin, request) + return proto.recv_frame(self._stdout) + + def _send(self, request: object) -> None: + with self._lock: + assert self._stdin is not None + proto.send_frame(self._stdin, request) + + def _recv(self) -> object: + assert self._stdout is not None + return proto.recv_frame(self._stdout) + + def _wait_quiet( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None, + ) -> None: + self._send(proto.WaitRequest(jobs=jobs)) + while True: + response = self._recv() + if isinstance(response, proto.StatusUpdateResponse): + if on_update: + on_update(response.job.job_id, response.state) + elif isinstance(response, proto.WaitDoneResponse): + return + elif isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error while waiting: {response.message}") + + def _wait_with_progress( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None, + ) -> None: + from rich.progress import Progress, TextColumn, TimeElapsedColumn + + self._send(proto.WaitRequest(jobs=jobs)) + + with Progress( + TimeElapsedColumn(), + TextColumn("[progress.description]{task.description}"), + auto_refresh=False, + ) as progress_bar: + tasks = { + job.job_id: progress_bar.add_task( + f"{job.job_id} {job.job_name or ''} {job.description or ''}", + total=1, + ) + for job in jobs + } + + while True: + response = self._recv() + progress_bar.refresh() + + if isinstance(response, proto.StatusUpdateResponse): + job = response.job + if response.state is None: + task = tasks.get(job.job_id) + if task is not None: + progress_bar.advance(task) + progress_bar.update(task, visible=False) + if on_update: + on_update(job.job_id, None) + elif on_update: + on_update(job.job_id, response.state) + + elif isinstance(response, proto.WaitDoneResponse): + return + + elif isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error while waiting: {response.message}") + + def close(self) -> None: + """Terminate the SSH subprocess.""" + if self._proc is not None: + self._proc.terminate() + self._proc = None + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass diff --git a/src/pbspy_server/__init__.py b/src/pbspy_server/__init__.py new file mode 100644 index 0000000..cf731a8 --- /dev/null +++ b/src/pbspy_server/__init__.py @@ -0,0 +1 @@ +"""pbspy-server: PBS job scheduler server daemon for pbspy.""" diff --git a/src/pbspy_server/__main__.py b/src/pbspy_server/__main__.py new file mode 100644 index 0000000..41206c1 --- /dev/null +++ b/src/pbspy_server/__main__.py @@ -0,0 +1,100 @@ +""" +pbspy-server command-line entry point. + +Usage:: + + pbspy-server --daemon # start server daemon in the foreground + pbspy-server --proxy # SSH proxy mode (stdio ↔ socket splice) + pbspy-server --status # print server status and exit +""" + +from __future__ import annotations + +import argparse +import logging +import socket +import sys +from pathlib import Path +from typing import BinaryIO, cast + +_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" + + +def _cmd_daemon() -> None: + from pbspy_server.server import run_server + + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + run_server() + + +def _cmd_proxy() -> None: + from pbspy_server.proxy import run_proxy + + run_proxy() + + +def _cmd_status() -> None: + if not _ADDR_PATH.exists(): + print("pbspy-server: not running (address file not found)") + sys.exit(1) + + try: + import json + + data = json.loads(_ADDR_PATH.read_text()) + host = data["host"] + port = int(data["port"]) + token = data["token"] + except (ValueError, KeyError, OSError) as exc: + print(f"pbspy-server: malformed address file ({exc})") + sys.exit(1) + + try: + import pbspy._protocol as proto + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(3.0) + sock.connect((host, port)) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + sock.close() + proto.send_frame(f, proto.AuthRequest(token=token)) + auth_resp = proto.recv_frame(f) + if not isinstance(auth_resp, proto.AuthOkResponse): + print(f"pbspy-server: auth failed: {auth_resp!r}") + f.close() + sys.exit(1) + proto.send_frame(f, proto.PingRequest()) + response = proto.recv_frame(f) + f.close() + if isinstance(response, proto.PongResponse): + print(f"pbspy-server: running on {host}:{port}") + else: + print(f"pbspy-server: unexpected response: {response!r}") + sys.exit(1) + except OSError as exc: + print(f"pbspy-server: not responding ({exc})") + sys.exit(1) + + +def main() -> None: + parser = argparse.ArgumentParser( + prog="pbspy-server", + description="PBS job scheduler server daemon for pbspy.", + ) + group = parser.add_mutually_exclusive_group(required=True) + group.add_argument("--daemon", action="store_true", help="Start the server daemon (runs in foreground)") + group.add_argument("--proxy", action="store_true", help="SSH proxy mode: splice stdin/stdout with server socket") + group.add_argument("--status", action="store_true", help="Print server status and exit") + + args = parser.parse_args() + + if args.daemon: + _cmd_daemon() + elif args.proxy: + _cmd_proxy() + elif args.status: + _cmd_status() + + +if __name__ == "__main__": + main() diff --git a/src/pbspy_server/proxy.py b/src/pbspy_server/proxy.py new file mode 100644 index 0000000..a817d61 --- /dev/null +++ b/src/pbspy_server/proxy.py @@ -0,0 +1,174 @@ +""" +SSH proxy mode for pbspy-server. + +When invoked as ``pbspy-server --proxy`` over an SSH connection, this module: + +1. Acquires an exclusive file lock (``~/.pbspy/server.lock``) to prevent two + proxy processes on different login nodes from racing to start the daemon. +2. Checks whether the server daemon is already running by reading + ``~/.pbspy/server.addr`` and attempting an authenticated TCP connection. +3. If the daemon is not running, starts it in the background and waits for the + address file to appear. +4. Splices ``sys.stdin`` / ``sys.stdout`` bytes bidirectionally with the server + TCP socket so the remote client can talk to the daemon transparently. + +Because the address file lives on the shared home filesystem, this works +correctly regardless of which login node the SSH session lands on — the daemon +may be running on a completely different node. +""" + +from __future__ import annotations + +import fcntl +import json +import os +import select +import socket +import subprocess +import sys +import time +from pathlib import Path +from typing import BinaryIO, cast + +import pbspy._protocol as proto + +__all__ = ["run_proxy"] + +_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" +_LOCK_PATH = Path.home() / ".pbspy" / "server.lock" +_STARTUP_TIMEOUT = 10.0 # seconds to wait for daemon to start +_POLL_INTERVAL = 0.1 +_BUF_SIZE = 65536 + + +def _read_server_addr() -> tuple[str, int, str] | None: + """Return ``(hostname, port, token)`` from the JSON addr file, or ``None`` if absent/invalid.""" + if not _ADDR_PATH.exists(): + return None + try: + data = json.loads(_ADDR_PATH.read_text()) + return data["host"], int(data["port"]), data["token"] + except (ValueError, KeyError, OSError): + return None + + +def _daemon_alive() -> bool: + """Return True if the addr file exists and the daemon responds to an auth+ping.""" + addr = _read_server_addr() + if addr is None: + return False + host, port, token = addr + try: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(2.0) + sock.connect((host, port)) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + sock.close() + proto.send_frame(f, proto.AuthRequest(token=token)) + resp = proto.recv_frame(f) + if not isinstance(resp, proto.AuthOkResponse): + f.close() + return False + proto.send_frame(f, proto.PingRequest()) + resp2 = proto.recv_frame(f) + f.close() + return isinstance(resp2, proto.PongResponse) + except OSError: + # Stale addr file — remove it so the next caller starts a fresh daemon. + try: + _ADDR_PATH.unlink(missing_ok=True) + except OSError: + pass + return False + + +def _start_daemon() -> None: + """Fork the server daemon in the background.""" + subprocess.Popen( + [sys.executable, "-m", "pbspy_server", "--daemon"], + stdin=subprocess.DEVNULL, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + start_new_session=True, + ) + + +def _wait_for_daemon(timeout: float = _STARTUP_TIMEOUT) -> None: + """Block until the daemon is alive or *timeout* seconds have elapsed.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if _daemon_alive(): + return + time.sleep(_POLL_INTERVAL) + raise TimeoutError(f"pbspy-server daemon did not start within {timeout:.0f}s") + + +def run_proxy() -> None: + """ + Entry point for ``--proxy`` mode. + + Ensures the daemon is running (starting it if necessary), then authenticates + and splices stdin/stdout with the server TCP socket until one side closes. + """ + _ADDR_PATH.parent.mkdir(parents=True, exist_ok=True) + + # Acquire exclusive lock while checking / starting the daemon to prevent + # two proxy processes on different login nodes from both deciding to start it. + lock_file = open(_LOCK_PATH, "w") # noqa: SIM115 + try: + fcntl.flock(lock_file, fcntl.LOCK_EX) + if not _daemon_alive(): + _start_daemon() + _wait_for_daemon() + addr = _read_server_addr() + finally: + fcntl.flock(lock_file, fcntl.LOCK_UN) + lock_file.close() + + if addr is None: + raise RuntimeError("pbspy-server daemon started but address file not found") + + host, port, token = addr + srv_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv_sock.connect((host, port)) + + # Authenticate with the daemon before entering the byte-splice loop. + srv_file = cast(BinaryIO, srv_sock.makefile("rwb", buffering=0)) + proto.send_frame(srv_file, proto.AuthRequest(token=token)) + auth_resp = proto.recv_frame(srv_file) + srv_file.close() # File object done; raw socket takes over below + if not isinstance(auth_resp, proto.AuthOkResponse): + srv_sock.close() + raise RuntimeError(f"pbspy-server rejected auth token: {auth_resp!r}") + + stdin_fd = sys.stdin.buffer.fileno() + srv_fd = srv_sock.fileno() + + # Put stdin in non-blocking mode; stdout remains blocking so writes + # don't need special partial-write handling. + _set_nonblocking(stdin_fd) + + try: + while True: + readable, _, exceptional = select.select([stdin_fd, srv_fd], [], [stdin_fd, srv_fd], 1.0) + if exceptional: + break + for fd in readable: + if fd == stdin_fd: + data = os.read(stdin_fd, _BUF_SIZE) + if not data: + return + srv_sock.sendall(data) + elif fd == srv_fd: + data = srv_sock.recv(_BUF_SIZE) + if not data: + return + sys.stdout.buffer.write(data) + sys.stdout.buffer.flush() + finally: + srv_sock.close() + + +def _set_nonblocking(fd: int) -> None: + flags = fcntl.fcntl(fd, fcntl.F_GETFL) + fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) diff --git a/src/pbspy_server/server.py b/src/pbspy_server/server.py new file mode 100644 index 0000000..b70572c --- /dev/null +++ b/src/pbspy_server/server.py @@ -0,0 +1,305 @@ +""" +pbspy-server daemon. + +Listens on a TCP socket (hostname:random-port), writes the address to +``~/.pbspy/server.addr`` so that any login node can find and connect to it, +accepts one connection per client, and dispatches pickled requests from +:mod:`pbspy._protocol`. + +A background thread polls ``qstat`` every 60 seconds for all unfinished jobs +and pushes :class:`StatusUpdateResponse` frames to any clients waiting on +those jobs. Job state is held in memory for the daemon's lifetime; it is not +persisted across restarts. +""" + +from __future__ import annotations + +import json +import logging +import secrets +import signal +import socket +import subprocess +import threading +from dataclasses import dataclass +from pathlib import Path +from typing import BinaryIO, cast + +import pbspy._pbs_core as core +import pbspy._protocol as proto +from pbspy import Job, JobResult + +__all__ = ["run_server"] + +logger = logging.getLogger(__name__) + +_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" +_POLL_INTERVAL = 60.0 # seconds between qstat polls + + +@dataclass +class _JobRecord: + """In-memory record of a submitted job's current state.""" + + job_id: str + job_name: str | None = None + state: str | None = None # "Q", "R", "E", "F" + exit_code: int | None = None + + +class _WaitSubscription: + """Tracks a client waiting for a set of jobs to finish.""" + + def __init__(self, jobs: list[Job], stream: BinaryIO) -> None: + self.pending: set[str] = {j.job_id for j in jobs} + self.jobs_by_id: dict[str, Job] = {j.job_id: j for j in jobs} + self.stream = stream + self.lock = threading.Lock() + self.done = threading.Event() + + +class _ServerState: + """Shared mutable state accessed by both connection threads and the poll thread.""" + + def __init__(self, token: str) -> None: + self.token = token + self.lock = threading.Lock() + self.active_connections: int = 0 + self._ever_had_connection: bool = False + self.jobs: dict[str, _JobRecord] = {} + self.subscriptions: dict[str, list[_WaitSubscription]] = {} + + def add_subscription(self, sub: _WaitSubscription) -> None: + with self.lock: + for job_id in sub.pending: + self.subscriptions.setdefault(job_id, []).append(sub) + + def notify_finished(self, job_id: str, job: Job) -> None: + """Mark job_id as done in all waiting subscriptions, sending frames as needed.""" + with self.lock: + subs = self.subscriptions.pop(job_id, []) + for sub in subs: + with sub.lock: + try: + proto.send_frame(sub.stream, proto.StatusUpdateResponse(job=job, state=None)) + except OSError: + pass + sub.pending.discard(job_id) + if not sub.pending: + try: + proto.send_frame(sub.stream, proto.WaitDoneResponse(jobs=list(sub.jobs_by_id.values()))) + except OSError: + pass + sub.done.set() + + +def _cancel_all_subscriptions(state: _ServerState) -> None: + """Signal all waiting subscriptions to unblock their connection threads (for shutdown).""" + with state.lock: + subs_by_id = dict(state.subscriptions) + state.subscriptions.clear() + seen: set[int] = set() + for subs in subs_by_id.values(): + for sub in subs: + if id(sub) not in seen: + seen.add(id(sub)) + sub.done.set() + + +def _should_shutdown(state: _ServerState) -> bool: + """Return True when the daemon has nothing left to do and can exit.""" + with state.lock: + if not state._ever_had_connection: + return False + if state.active_connections > 0: + return False + return not any(r.state != "F" for r in state.jobs.values()) + + +def _poll_loop(state: _ServerState, stop_event: threading.Event, poll_interval: float = _POLL_INTERVAL) -> None: + """Background thread: polls qstat for all unfinished jobs.""" + while not stop_event.wait(timeout=poll_interval): + try: + with state.lock: + unfinished = [r for r in state.jobs.values() if r.state != "F"] + for record in unfinished: + job = Job(job_id=record.job_id, job_name=record.job_name) + result = _check_job_finished(record.job_id) + if result is not None: + with state.lock: + record.state = "F" + record.exit_code = result + state.notify_finished(record.job_id, job) + if _should_shutdown(state): + logger.info("No active connections and no unfinished jobs; shutting down") + stop_event.set() + except Exception: + logger.exception("Error in poll loop") + + +def _check_job_finished(job_id: str) -> int | None: + """Return exit code if job has finished, else None.""" + process = subprocess.run(["qstat", job_id], capture_output=True) + if process.returncode != 0: + stdout = process.stdout.decode("utf-8") + if job_id not in stdout or "has finished" in process.stderr.decode("utf-8"): + return core.try_get_exit_code(job_id) or 0 + return None + + +def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: threading.Event) -> None: + """Handle a single client connection in its own thread.""" + with state.lock: + state.active_connections += 1 + state._ever_had_connection = True + stream = cast(BinaryIO, conn.makefile("rwb", buffering=0)) + try: + # First frame must be a valid AuthRequest. + try: + auth = proto.recv_frame(stream) + except EOFError: + return + if not isinstance(auth, proto.AuthRequest) or auth.token != state.token: + logger.warning("Connection rejected: invalid or missing auth token") + try: + proto.send_frame(stream, proto.ErrorResponse(message="Authentication failed")) + except OSError: + pass + return + proto.send_frame(stream, proto.AuthOkResponse()) + + while True: + try: + request = proto.recv_frame(stream) + except EOFError: + break + except Exception: + logger.exception("Error receiving frame") + break + + try: + if isinstance(request, proto.PingRequest): + proto.send_frame(stream, proto.PongResponse()) + + elif isinstance(request, proto.SubmitRequest): + job_id, job_name = core.pbs_submit(request.script, request.name) + job = Job(job_id=job_id, job_name=job_name) + with state.lock: + state.jobs[job_id] = _JobRecord(job_id=job_id, job_name=job_name, state="Q") + proto.send_frame(stream, proto.SubmittedResponse(job=job)) + + elif isinstance(request, proto.WaitRequest): + sub = _WaitSubscription(request.jobs, stream) + # Check state and register subscription atomically so notify_finished() + # cannot fire between the two steps and leave us waiting on a job + # that is already done. + with state.lock: + for job in request.jobs: + record = state.jobs.get(job.job_id) + if record and record.state == "F": + sub.pending.discard(job.job_id) + if sub.pending: + for job_id in sub.pending: + state.subscriptions.setdefault(job_id, []).append(sub) + if sub.pending: + sub.done.wait() + else: + proto.send_frame(stream, proto.WaitDoneResponse(jobs=request.jobs)) + + elif isinstance(request, proto.ResultRequest): + result: JobResult = core.pbs_get_result(request.job) + proto.send_frame(stream, proto.ResultResponse(job=request.job, result=result)) + + else: + proto.send_frame(stream, proto.ErrorResponse(message=f"Unknown request type: {type(request)}")) + + except Exception as exc: + logger.exception("Error handling request") + try: + proto.send_frame(stream, proto.ErrorResponse(message=str(exc))) + except OSError: + pass + finally: + stream.close() + conn.close() + with state.lock: + state.active_connections -= 1 + if _should_shutdown(state): + logger.info("Last client disconnected with no unfinished jobs; shutting down") + stop_event.set() + + +def run_server( + addr_path: Path = _ADDR_PATH, + poll_interval: float = _POLL_INTERVAL, +) -> None: + """ + Start the pbspy-server daemon. + + Runs in the foreground; the caller is responsible for daemonising the + process (e.g. ``--daemon`` in :mod:`pbspy_server.__main__` forks and + calls this in the child). + + The server binds a TCP port on ``0.0.0.0`` (OS chooses a free port) and + writes a JSON address file to *addr_path* containing the hostname, port, + and a random auth token. Clients must send the token as the first frame + on every connection; connections with an incorrect or missing token are + rejected immediately. + + Job state is held in memory for the daemon's lifetime and is not + persisted across restarts. + + Args: + addr_path: Path for the address file (JSON). Defaults to + ``~/.pbspy/server.addr``. + poll_interval: Seconds between qstat polls (default 60). Override in tests. + """ + addr_path.parent.mkdir(parents=True, exist_ok=True) + if addr_path.exists(): + addr_path.unlink() + + token = secrets.token_hex(32) + state = _ServerState(token=token) + + stop_event = threading.Event() + poll_thread = threading.Thread( + target=_poll_loop, args=(state, stop_event), kwargs={"poll_interval": poll_interval}, daemon=True + ) + poll_thread.start() + + # Signal handlers can only be registered from the main thread. + if threading.current_thread() is threading.main_thread(): + + def _handle_signal(signum: int, frame: object) -> None: + logger.info("Received signal %d, shutting down", signum) + stop_event.set() + + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) + + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind(("", 0)) + srv.listen(16) + srv.settimeout(1.0) + + hostname = socket.gethostname() + port = srv.getsockname()[1] + addr_path.write_text(json.dumps({"host": hostname, "port": port, "token": token})) + + logger.info("pbspy-server listening on %s:%d", hostname, port) + + try: + while not stop_event.is_set(): + try: + conn, _ = srv.accept() + except TimeoutError: + continue + t = threading.Thread(target=_handle_connection, args=(conn, state, stop_event), daemon=True) + t.start() + finally: + srv.close() + if addr_path.exists(): + addr_path.unlink() + _cancel_all_subscriptions(state) + logger.info("pbspy-server stopped") diff --git a/tests/test_default.py b/tests/test_default.py index 48d568b..a406310 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -1,6 +1,12 @@ -from pbspy import JobDescription +import shutil +from collections.abc import Callable +import pytest +from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend, SSHBackend + + +@pytest.mark.skipif(shutil.which("qsub") is None, reason="qsub not available") def test_default() -> None: job_description = ( JobDescription() @@ -15,3 +21,109 @@ 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] = {} + + 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, progress: bool = True + ) -> None: + self.waited.append(list(jobs)) + + def get_result(self, job: object) -> object: + assert isinstance(job, Job) + return self.results.get(job.job_id, JobResult(exit_code=None)) + + +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(progress=False) + + 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], progress=False) + + 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, LocalBackend, SSHBackend are importable from pbspy.""" + assert issubclass(LocalBackend, Backend) + assert issubclass(SSHBackend, Backend) diff --git a/tests/test_protocol.py b/tests/test_protocol.py new file mode 100644 index 0000000..0dac621 --- /dev/null +++ b/tests/test_protocol.py @@ -0,0 +1,185 @@ +"""Tests for the pickle framing helpers in pbspy._protocol.""" + +from __future__ import annotations + +import io +import pickle + +import pytest + +import pbspy._protocol as proto +from pbspy import Job, JobResult + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def roundtrip(obj: object) -> object: + """Serialise *obj* into a BytesIO buffer, then deserialise and return it.""" + buf = io.BytesIO() + proto.send_frame(buf, obj) + buf.seek(0) + return proto.recv_frame(buf) + + +# --------------------------------------------------------------------------- +# Request types +# --------------------------------------------------------------------------- + + +def test_roundtrip_auth_request() -> None: + req = proto.AuthRequest(token="abc123secret") + result = roundtrip(req) + assert isinstance(result, proto.AuthRequest) + assert result.token == "abc123secret" + + +def test_roundtrip_ping() -> None: + result = roundtrip(proto.PingRequest()) + assert isinstance(result, proto.PingRequest) + + +def test_roundtrip_submit_request() -> None: + req = proto.SubmitRequest(script="#!/bin/bash\necho hi", name="myjob") + result = roundtrip(req) + assert isinstance(result, proto.SubmitRequest) + assert result.script == req.script + assert result.name == req.name + + +def test_roundtrip_submit_request_no_name() -> None: + req = proto.SubmitRequest(script="#!/bin/bash") + result = roundtrip(req) + assert isinstance(result, proto.SubmitRequest) + assert result.name is None + + +def test_roundtrip_wait_request() -> None: + jobs = [Job(job_id="1.gadi", job_name="a"), Job(job_id="2.gadi", job_name="b")] + req = proto.WaitRequest(jobs=jobs) + result = roundtrip(req) + assert isinstance(result, proto.WaitRequest) + assert [j.job_id for j in result.jobs] == ["1.gadi", "2.gadi"] + + +def test_roundtrip_result_request() -> None: + job = Job(job_id="42.gadi", job_name="myjob") + req = proto.ResultRequest(job=job) + result = roundtrip(req) + assert isinstance(result, proto.ResultRequest) + assert result.job.job_id == "42.gadi" + + +# --------------------------------------------------------------------------- +# Response types +# --------------------------------------------------------------------------- + + +def test_roundtrip_auth_ok() -> None: + assert isinstance(roundtrip(proto.AuthOkResponse()), proto.AuthOkResponse) + + +def test_roundtrip_pong() -> None: + assert isinstance(roundtrip(proto.PongResponse()), proto.PongResponse) + + +def test_roundtrip_submitted_response() -> None: + job = Job(job_id="99.gadi", job_name="submitted") + resp = proto.SubmittedResponse(job=job) + result = roundtrip(resp) + assert isinstance(result, proto.SubmittedResponse) + assert result.job.job_id == "99.gadi" + + +def test_roundtrip_status_update() -> None: + job = Job(job_id="1.gadi") + resp = proto.StatusUpdateResponse(job=job, state="R") + result = roundtrip(resp) + assert isinstance(result, proto.StatusUpdateResponse) + assert result.state == "R" + + +def test_roundtrip_status_update_finished() -> None: + job = Job(job_id="1.gadi") + resp = proto.StatusUpdateResponse(job=job, state=None) + result = roundtrip(resp) + assert isinstance(result, proto.StatusUpdateResponse) + assert result.state is None + + +def test_roundtrip_wait_done() -> None: + jobs = [Job(job_id="1.gadi"), Job(job_id="2.gadi")] + resp = proto.WaitDoneResponse(jobs=jobs) + result = roundtrip(resp) + assert isinstance(result, proto.WaitDoneResponse) + assert len(result.jobs) == 2 + + +def test_roundtrip_result_response() -> None: + job = Job(job_id="5.gadi", job_name="myjob") + job_result = JobResult(exit_code=0, output="hello\n", error="", stats={"CPU": "1s"}) + resp = proto.ResultResponse(job=job, result=job_result) + result = roundtrip(resp) + assert isinstance(result, proto.ResultResponse) + assert result.result.exit_code == 0 + assert result.result.output == "hello\n" + assert result.result.stats == {"CPU": "1s"} + + +def test_roundtrip_error_response() -> None: + resp = proto.ErrorResponse(message="something went wrong") + result = roundtrip(resp) + assert isinstance(result, proto.ErrorResponse) + assert result.message == "something went wrong" + + +# --------------------------------------------------------------------------- +# Framing edge cases +# --------------------------------------------------------------------------- + + +def test_eof_on_empty_stream() -> None: + buf = io.BytesIO(b"") + with pytest.raises(EOFError): + proto.recv_frame(buf) + + +def test_eof_on_truncated_header() -> None: + buf = io.BytesIO(b"\x00\x00") # Only 2 bytes of a 4-byte header + with pytest.raises(EOFError): + proto.recv_frame(buf) + + +def test_eof_on_truncated_body() -> None: + data = pickle.dumps(proto.PingRequest()) + # Write a header claiming a longer body than we provide + header = len(data).to_bytes(4, "big") + buf = io.BytesIO(header + data[:2]) # body truncated + with pytest.raises(EOFError): + proto.recv_frame(buf) + + +def test_multiple_frames_in_sequence() -> None: + buf = io.BytesIO() + proto.send_frame(buf, proto.PingRequest()) + proto.send_frame(buf, proto.PongResponse()) + proto.send_frame(buf, proto.ErrorResponse(message="test")) + buf.seek(0) + + assert isinstance(proto.recv_frame(buf), proto.PingRequest) + assert isinstance(proto.recv_frame(buf), proto.PongResponse) + err = proto.recv_frame(buf) + assert isinstance(err, proto.ErrorResponse) + assert err.message == "test" + + +def test_large_object_roundtrip() -> None: + """Frames larger than a typical network buffer must reassemble correctly.""" + big_output = "x" * 500_000 + result_obj = JobResult(exit_code=0, output=big_output) + job = Job(job_id="big.gadi", job_name="bigtest") + resp = proto.ResultResponse(job=job, result=result_obj) + result = roundtrip(resp) + assert isinstance(result, proto.ResultResponse) + assert len(result.result.output) == 500_000 diff --git a/tests/test_server.py b/tests/test_server.py new file mode 100644 index 0000000..ef3e864 --- /dev/null +++ b/tests/test_server.py @@ -0,0 +1,373 @@ +""" +Integration tests for pbspy_server.server. + +Spins up a real server (in a background thread) with mocked PBS operations so +that no ``qsub``/``qstat`` installation is needed. +""" + +from __future__ import annotations + +import json +import socket +import threading +import time +from collections.abc import Generator +from pathlib import Path +from typing import BinaryIO, cast +from unittest.mock import MagicMock, patch + +import pytest + +import pbspy._protocol as proto +from pbspy import Job, JobResult +from pbspy_server.server import run_server + +# --------------------------------------------------------------------------- +# Fixture: running server +# --------------------------------------------------------------------------- + +_WAIT_TIMEOUT = 5.0 # seconds to wait for server address file to appear + + +def _wait_for_server(addr_path: Path, timeout: float = _WAIT_TIMEOUT) -> None: + """Wait until the server address file appears (server is bound and ready).""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if addr_path.exists(): + return + time.sleep(0.05) + raise TimeoutError(f"Server addr file {addr_path} did not appear within {timeout}s") + + +def _read_addr(addr_path: Path) -> tuple[str, int, str]: + data = json.loads(addr_path.read_text()) + return data["host"], int(data["port"]), data["token"] + + +class ServerHandle: + """Handle to a running test server.""" + + def __init__(self, addr_path: Path) -> None: + self.addr_path = addr_path + + def connect(self, timeout: float | None = None) -> BinaryIO: + """Open an authenticated TCP connection and return a binary file-like object.""" + host, port, token = _read_addr(self.addr_path) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + if timeout is not None: + sock.settimeout(timeout) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + # Close the socket object so that when the SocketIO is GC'd it properly + # calls _real_close() (via _decref_socketios). Without this, the socket + # FD leaks and Python emits a ResourceWarning. + sock.close() + proto.send_frame(f, proto.AuthRequest(token=token)) + auth_resp = proto.recv_frame(f) + assert isinstance(auth_resp, proto.AuthOkResponse), f"Auth failed: {auth_resp!r}" + return f + + def rpc(self, request: object) -> object: + """Send one request, return the first response, then close.""" + stream = self.connect() + try: + proto.send_frame(stream, request) + return proto.recv_frame(stream) + finally: + stream.close() + + +@pytest.fixture() +def mock_pbs_submit() -> Generator[MagicMock, None, None]: + """Patch pbs_submit to return a fake job without calling qsub.""" + with patch("pbspy_server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: + yield m + + +@pytest.fixture() +def mock_check_finished() -> Generator[MagicMock, None, None]: + """Patch _check_job_finished to report jobs as never finished (default).""" + with patch("pbspy_server.server._check_job_finished", return_value=None) as m: + yield m + + +@pytest.fixture() +def mock_pbs_get_result() -> Generator[MagicMock, None, None]: + with patch( + "pbspy_server.server.core.pbs_get_result", + return_value=JobResult(exit_code=0, output="hello\n"), + ) as m: + yield m + + +@pytest.fixture() +def server( + tmp_path: Path, + mock_pbs_submit: MagicMock, + mock_check_finished: MagicMock, + mock_pbs_get_result: MagicMock, +) -> Generator[ServerHandle, None, None]: + addr_path = tmp_path / "test.addr" + + t = threading.Thread( + target=run_server, + kwargs={"addr_path": addr_path, "poll_interval": 0.1}, + daemon=True, + ) + t.start() + _wait_for_server(addr_path) + yield ServerHandle(addr_path) + # Server exits on its own once no clients + no jobs (shutdown test). + # For other tests, give it a moment to clean up. + t.join(timeout=2.0) + + +# --------------------------------------------------------------------------- +# Ping / pong +# --------------------------------------------------------------------------- + + +def test_ping(server: ServerHandle) -> None: + response = server.rpc(proto.PingRequest()) + assert isinstance(response, proto.PongResponse) + + +# --------------------------------------------------------------------------- +# Auth — rejection cases +# --------------------------------------------------------------------------- + + +def test_auth_rejected_wrong_token(server: ServerHandle) -> None: + """Connecting with a wrong token must return ErrorResponse and close the connection.""" + host, port, _token = _read_addr(server.addr_path) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + sock.close() + + proto.send_frame(f, proto.AuthRequest(token="wrong-token")) + resp = proto.recv_frame(f) + assert isinstance(resp, proto.ErrorResponse) + # Server should close the connection after rejecting auth. + try: + proto.recv_frame(f) + received_eof = False + except EOFError: + received_eof = True + f.close() + assert received_eof, "Expected server to close connection after auth failure" + + +def test_auth_rejected_non_auth_first_frame(server: ServerHandle) -> None: + """Sending a non-AuthRequest as the first frame must be rejected.""" + host, port, _token = _read_addr(server.addr_path) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + sock.close() + + proto.send_frame(f, proto.PingRequest()) # Wrong first frame + resp = proto.recv_frame(f) + assert isinstance(resp, proto.ErrorResponse) + f.close() + + +# --------------------------------------------------------------------------- +# Submit +# --------------------------------------------------------------------------- + + +def test_submit_returns_job(server: ServerHandle) -> None: + response = server.rpc(proto.SubmitRequest(script="#!/bin/bash\necho hi", name="test_job")) + assert isinstance(response, proto.SubmittedResponse) + assert response.job.job_id == "100.mock" + assert response.job.job_name == "test_job" + + +# --------------------------------------------------------------------------- +# Result +# --------------------------------------------------------------------------- + + +def test_result_returns_job_result(server: ServerHandle) -> None: + job = Job(job_id="100.mock", job_name="test_job") + response = server.rpc(proto.ResultRequest(job=job)) + assert isinstance(response, proto.ResultResponse) + assert response.result.exit_code == 0 + assert response.result.output == "hello\n" + + +# --------------------------------------------------------------------------- +# Wait — already finished +# --------------------------------------------------------------------------- + + +def test_wait_already_finished_returns_immediately(tmp_path: Path) -> None: + """If the job is already marked F in state, WaitRequest should return without blocking.""" + addr_path = tmp_path / "wait-done.addr" + server_done = threading.Event() + + def _run() -> None: + with ( + patch("pbspy_server.server.core.pbs_submit", return_value=("done.mock", "done_job")), + # Return 0 immediately so the very first poll marks the job finished. + patch("pbspy_server.server._check_job_finished", return_value=0), + patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), + ): + run_server(addr_path=addr_path, poll_interval=0.1) + server_done.set() + + t = threading.Thread(target=_run, daemon=True) + t.start() + _wait_for_server(addr_path) + + # Submit the job, then wait long enough for the poll loop (0.1s) to mark it done. + stream = _raw_connect(addr_path) + proto.send_frame(stream, proto.SubmitRequest(script="#!/bin/bash", name="done_job")) + proto.recv_frame(stream) + stream.close() + + time.sleep(0.5) # poll interval is 0.1s; 0.5s is more than enough + + job = Job(job_id="done.mock", job_name="done_job") + stream = _raw_connect(addr_path) + proto.send_frame(stream, proto.WaitRequest(jobs=[job])) + + response = proto.recv_frame(stream) + stream.close() + + assert isinstance(response, proto.WaitDoneResponse) + assert response.jobs[0].job_id == "done.mock" + + t.join(timeout=2.0) + + +# --------------------------------------------------------------------------- +# Wait — job finishes during wait +# --------------------------------------------------------------------------- + + +def test_wait_receives_status_update_then_done(server: ServerHandle, mock_check_finished: MagicMock) -> None: + """ + Submit a job, start waiting, then let the poll loop (0.1s interval) detect + the job as finished and push WaitDoneResponse back to the client. + """ + server.rpc(proto.SubmitRequest(script="#!/bin/bash", name="test_job")) + + job = Job(job_id="100.mock", job_name="test_job") + stream = server.connect(timeout=5.0) + proto.send_frame(stream, proto.WaitRequest(jobs=[job])) + + # Allow the poll loop to mark the job as finished. + mock_check_finished.return_value = 0 + + # Collect responses until WaitDoneResponse (or timeout). + responses = [] + try: + while True: + r = proto.recv_frame(stream) + responses.append(r) + if isinstance(r, proto.WaitDoneResponse): + break + except (EOFError, OSError): + pass + + stream.close() + + assert any(isinstance(r, proto.WaitDoneResponse) for r in responses), f"Expected WaitDoneResponse; got: {responses}" + + +# --------------------------------------------------------------------------- +# Shutdown: no jobs + last client disconnects +# --------------------------------------------------------------------------- + + +def test_shutdown_when_no_jobs_and_no_clients(tmp_path: Path) -> None: + """ + The server should exit automatically when the last client disconnects + and there are no unfinished jobs. + """ + addr_path = tmp_path / "shutdown.addr" + + server_done = threading.Event() + + def _run() -> None: + with ( + patch("pbspy_server.server.core.pbs_submit", return_value=("1.mock", "j")), + patch("pbspy_server.server._check_job_finished", return_value=None), + patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), + ): + run_server(addr_path=addr_path) + server_done.set() + + t = threading.Thread(target=_run, daemon=True) + t.start() + _wait_for_server(addr_path) + + # Connect, ping, and disconnect — no jobs submitted + stream = _raw_connect(addr_path) + proto.send_frame(stream, proto.PingRequest()) + proto.recv_frame(stream) + stream.close() + + assert server_done.wait(timeout=5.0), "Server did not exit after last client disconnected with no jobs" + + +# --------------------------------------------------------------------------- +# Shutdown: stays alive while jobs pending +# --------------------------------------------------------------------------- + + +def test_server_stays_alive_while_jobs_pending(tmp_path: Path) -> None: + """ + If jobs are still unfinished, the server must NOT exit when the last + client disconnects; it should exit once the poll loop detects them done. + """ + addr_path = tmp_path / "pending.addr" + server_done = threading.Event() + + with ( + patch("pbspy_server.server.core.pbs_submit", return_value=("2.mock", "pending_job")) as _, + patch("pbspy_server.server._check_job_finished", return_value=None) as mock_check, + patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), + ): + + def _run() -> None: + run_server(addr_path=addr_path, poll_interval=0.1) + server_done.set() + + t = threading.Thread(target=_run, daemon=True) + t.start() + _wait_for_server(addr_path) + + # Submit a job then disconnect + stream = _raw_connect(addr_path) + proto.send_frame(stream, proto.SubmitRequest(script="#!/bin/bash", name="pending_job")) + proto.recv_frame(stream) + stream.close() + + # Server should NOT have shut down yet (jobs still pending) + assert not server_done.wait(timeout=1.0), "Server exited prematurely while jobs were still pending" + + # Let the poll loop see the job as finished → server should then exit. + mock_check.return_value = 0 + assert server_done.wait(timeout=5.0), "Server did not exit after all jobs finished" + + t.join(timeout=2.0) + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _raw_connect(addr_path: Path) -> BinaryIO: + host, port, token = _read_addr(addr_path) + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((host, port)) + f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) + sock.close() # Mark closed; SocketIO keeps FD open via _io_refs + proto.send_frame(f, proto.AuthRequest(token=token)) + auth_resp = proto.recv_frame(f) + assert isinstance(auth_resp, proto.AuthOkResponse), f"Auth failed: {auth_resp!r}" + return f From 623837f7746358789b0b0919bbcf61af40e3eb24 Mon Sep 17 00:00:00 2001 From: omdowley Date: Tue, 2 Jun 2026 14:24:27 +1000 Subject: [PATCH 02/16] refactor: move pbspy-server off the supercomputer The server now runs on a machine with SSH access to the supercomputer rather than on the login node itself. PBS commands (qsub, qstat, file reads) are issued by the server over SSH via the new PBSRunner. Clients connect to the server over a plain TCP socket. A thourough cleanup has removed no-longer required parts of the last commit. --- examples/run.py | 4 +- src/pbspy/__init__.py | 8 +- src/pbspy/_backend.py | 2 +- src/pbspy/_pbs_core.py | 63 ++++-- src/pbspy/_protocol.py | 25 +-- .../{_ssh_backend.py => _server_backend.py} | 155 +++++++------- src/pbspy_server/__main__.py | 112 ++++------ src/pbspy_server/proxy.py | 174 --------------- src/pbspy_server/server.py | 160 +++++--------- tests/test_default.py | 6 +- tests/test_protocol.py | 11 - tests/test_server.py | 202 ++---------------- 12 files changed, 244 insertions(+), 678 deletions(-) rename src/pbspy/{_ssh_backend.py => _server_backend.py} (54%) delete mode 100644 src/pbspy_server/proxy.py diff --git a/examples/run.py b/examples/run.py index 35b23a1..76a38d6 100644 --- a/examples/run.py +++ b/examples/run.py @@ -1,11 +1,11 @@ import typer from rich import print -from pbspy import Job, JobDescription, QueueLimits, SSHBackend +from pbspy import Job, JobDescription, QueueLimits, ServerBackend def main(remote: str = "") -> None: - backend = SSHBackend(remote) if remote else None + backend = ServerBackend(remote) if remote else None # Run a job with some explicit parameters job_a = ( diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index f6e98af..0ec8671 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -9,7 +9,7 @@ from pbspy._backend import Backend from pbspy._local_backend import LocalBackend -from pbspy._ssh_backend import SSHBackend +from pbspy._server_backend import ServerBackend # Shared LocalBackend instance so all locally-submitted jobs are grouped together # when calling Job.wait_all() / Job.result_all(), restoring concurrent polling. @@ -23,7 +23,7 @@ "QueueLimitsMap", "Backend", "LocalBackend", - "SSHBackend", + "ServerBackend", "gadi", ] @@ -248,8 +248,8 @@ def submit(self, backend: Backend | None = None) -> Job: Args: backend: The backend to use for submission. Defaults to :class:`~pbspy.LocalBackend` (runs ``qsub`` locally). - Pass a :class:`~pbspy.SSHBackend` instance to submit to a - remote supercomputer over SSH. + Pass a :class:`~pbspy.ServerBackend` instance to submit via + a pbspy-server daemon. """ if backend is None: backend = LocalBackend() diff --git a/src/pbspy/_backend.py b/src/pbspy/_backend.py index c7f4a85..e3df70f 100644 --- a/src/pbspy/_backend.py +++ b/src/pbspy/_backend.py @@ -2,7 +2,7 @@ Backend abstract base class. Concrete backends: :class:`~pbspy._local_backend.LocalBackend` (default, -in-process) and :class:`~pbspy._ssh_backend.SSHBackend` (remote via SSH). +in-process) and :class:`~pbspy._server_backend.ServerBackend` (remote via TCP). """ from __future__ import annotations diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py index 47c949f..51bb858 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -2,7 +2,7 @@ Core PBS operations: qsub submission, qstat polling, and result file reading. These functions are called directly by LocalBackend (in-process) and by the -server daemon when handling pickled requests from SSHBackend. +server daemon when handling pickled requests from ServerBackend clients. """ from __future__ import annotations @@ -17,14 +17,39 @@ if TYPE_CHECKING: from pbspy import Job, JobResult -__all__ = ["pbs_submit", "pbs_wait_for_jobs", "pbs_get_result", "try_get_exit_code"] +__all__ = ["PBSRunner", "pbs_submit", "pbs_wait_for_jobs", "pbs_get_result", "try_get_exit_code"] _POLL_INTERVAL_SECONDS = 60 _PROGRESS_REFRESH_SECONDS = 1 -def _get_job_name(job_id: str) -> str: - process = subprocess.run(["qstat", "-fx", job_id], capture_output=True) +class PBSRunner: + """ + Abstracts command execution and file reading for PBS operations. + + The default instance runs commands locally. Pass ``ssh_prefix`` to run + commands on a remote host via SSH (e.g. ``["ssh", "user@host"]``). + """ + + def __init__(self, ssh_prefix: list[str] | None = None) -> None: + self._prefix: list[str] = ssh_prefix or [] + + def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.CompletedProcess[bytes]: + return subprocess.run(self._prefix + cmd, input=input, capture_output=True) + + def read_file(self, path: str) -> str: + """Read a file (possibly on a remote host via ``ssh … cat``) 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") @@ -36,9 +61,9 @@ def _get_job_name(job_id: str) -> str: raise RuntimeError("Could not retrieve job name") -def try_get_exit_code(job_id: str) -> int | None: +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 = subprocess.run(["qstat", "-fx", job_id], capture_output=True) + 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") @@ -50,25 +75,19 @@ def try_get_exit_code(job_id: str) -> int | None: return None -def pbs_submit(script: str, name: str | None) -> tuple[str, str]: +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 = subprocess.Popen( - ["qsub"], - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - stdout, stderr = process.communicate(input=script.encode("utf-8")) + process = runner.run(["qsub"], input=script.encode("utf-8")) if process.returncode != 0: - raise RuntimeError(f'Failed to submit job: {stderr.strip().decode("utf-8")}') - job_id = stdout.strip().decode("utf-8") + 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) + name = _get_job_name(job_id, runner) return job_id, name @@ -105,7 +124,7 @@ def pbs_wait_for_jobs( on_update(job.job_id, None) -def pbs_get_result(job: Job) -> JobResult: +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`. @@ -118,13 +137,11 @@ def pbs_get_result(job: Job) -> JobResult: error_file = f"{job.job_name}.e{job_id_num}" try: - with open(output_file) as f: - output = f.read() + output = runner.read_file(output_file) except FileNotFoundError: output = "" try: - with open(error_file) as f: - error = f.read() + error = runner.read_file(error_file) except FileNotFoundError: error = "" @@ -149,7 +166,7 @@ def pbs_get_result(job: Job) -> JobResult: exit_code = int(exit_code_match.group()) if exit_code is None: - exit_code = try_get_exit_code(job.job_id) + exit_code = try_get_exit_code(job.job_id, runner) return JobResult(exit_code=exit_code, output=output, error=error, stats=pbs_stats) diff --git a/src/pbspy/_protocol.py b/src/pbspy/_protocol.py index a3135c5..e70debd 100644 --- a/src/pbspy/_protocol.py +++ b/src/pbspy/_protocol.py @@ -2,9 +2,9 @@ Request and response dataclasses for the pbspy client-server protocol. Messages are exchanged as length-prefixed pickle frames (see :func:`send_frame` -and :func:`recv_frame`). :class:`~pbspy._local_backend.LocalBackend` uses the -same dataclasses but never pickles them — it calls :mod:`pbspy._pbs_core` -directly and passes Python objects in-process. +and :func:`recv_frame`). Connections require no authentication — the server +is expected to run on a trusted machine with network access restricted to +authorised clients. """ from __future__ import annotations @@ -20,13 +20,11 @@ __all__ = [ # Requests - "AuthRequest", "SubmitRequest", "WaitRequest", "ResultRequest", "PingRequest", # Responses - "AuthOkResponse", "SubmittedResponse", "StatusUpdateResponse", "WaitDoneResponse", @@ -46,18 +44,6 @@ # --------------------------------------------------------------------------- -@dataclass -class AuthRequest: - """ - Authentication handshake — must be the first frame sent on every connection. - - The server responds with :class:`AuthOkResponse` on success or - :class:`ErrorResponse` (and closes the connection) on failure. - """ - - token: str - - @dataclass class SubmitRequest: """Ask the server to submit a PBS job script.""" @@ -93,11 +79,6 @@ class PingRequest: # --------------------------------------------------------------------------- -@dataclass -class AuthOkResponse: - """Returned after a successful :class:`AuthRequest`.""" - - @dataclass class SubmittedResponse: """Returned after a successful :class:`SubmitRequest`.""" diff --git a/src/pbspy/_ssh_backend.py b/src/pbspy/_server_backend.py similarity index 54% rename from src/pbspy/_ssh_backend.py rename to src/pbspy/_server_backend.py index 037d8a1..6006d3c 100644 --- a/src/pbspy/_ssh_backend.py +++ b/src/pbspy/_server_backend.py @@ -1,16 +1,13 @@ """ -SSHBackend: communicates with a remote pbspy-server daemon over SSH. +ServerBackend: communicates with a pbspy-server daemon over a plain TCP connection. -The backend opens an SSH subprocess running ``pbspy-server --proxy`` on the -remote host. Messages are exchanged as length-prefixed pickle frames (see -:mod:`pbspy._protocol`). The proxy on the remote end automatically starts the -daemon if it is not already running (emacsclient-style). +Messages are exchanged as length-prefixed pickle frames (see :mod:`pbspy._protocol`). +The server handles all SSH communication to the supercomputer internally. """ from __future__ import annotations -import select -import subprocess +import socket import threading from collections.abc import Callable from typing import TYPE_CHECKING, BinaryIO, cast @@ -21,27 +18,33 @@ if TYPE_CHECKING: from pbspy import Job -__all__ = ["SSHBackend"] +__all__ = ["ServerBackend"] -class SSHBackend(Backend): +class ServerBackend(Backend): """ - Backend that submits and tracks PBS jobs on a remote supercomputer via SSH. + Backend that connects to a pbspy-server daemon over TCP. + + The server (started with ``pbspy-server``) runs on any machine with SSH + access to the supercomputer; this client connects to it directly. Args: - host: SSH destination (e.g. ``"user@gadi.nci.org.au"``). - Any option accepted by ``ssh`` (host aliases, ``-i``, etc.) works - because the system ``ssh`` binary is used. - ssh_args: Extra arguments forwarded to the ``ssh`` command (e.g. - ``["-i", "/path/to/key"]``). + host: Hostname or IP address of the machine running pbspy-server. + port: TCP port the server is listening on (default: 9876). + connect_timeout: Seconds to wait for the initial TCP connection. """ - def __init__(self, host: str, ssh_args: list[str] | None = None) -> None: + def __init__( + self, + host: str, + port: int = 9876, + connect_timeout: float = 10.0, + ) -> None: self._host = host - self._ssh_args = ssh_args or [] - self._proc: subprocess.Popen[bytes] | None = None - self._stdin: BinaryIO | None = None - self._stdout: BinaryIO | None = None + self._port = port + self._connect_timeout = connect_timeout + self._sock: socket.socket | None = None + self._stream: BinaryIO | None = None self._lock = threading.Lock() self._connect() @@ -59,11 +62,6 @@ def wait( on_update: Callable[[str, str | None], None] | None = None, progress: bool = True, ) -> None: - """ - Send a WaitRequest and stream StatusUpdateResponse frames until - WaitDoneResponse is received. Optionally shows a progress display. - """ - if progress: self._wait_with_progress(jobs, on_update) else: @@ -78,80 +76,67 @@ def get_result(self, job: object) -> object: return response.result def _connect(self) -> None: - """Start the SSH subprocess and verify the proxy is alive.""" - cmd = ["ssh", *self._ssh_args, self._host, "pbspy-server", "--proxy"] - self._proc = subprocess.Popen( - cmd, - stdin=subprocess.PIPE, - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - ) - assert self._proc.stdin is not None - assert self._proc.stdout is not None - self._stdin = cast(BinaryIO, self._proc.stdin) - self._stdout = cast(BinaryIO, self._proc.stdout) + """Open a TCP connection to the server and verify liveness with a ping.""" + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.settimeout(self._connect_timeout) + try: + sock.connect((self._host, self._port)) + except OSError as exc: + sock.close() + raise RuntimeError(f"Could not connect to pbspy-server at {self._host}:{self._port}: {exc}") from exc + sock.settimeout(None) + self._sock = sock + self._stream = cast(BinaryIO, sock.makefile("rwb", buffering=0)) try: - pong = self._rpc(proto.PingRequest()) + proto.send_frame(self._stream, proto.PingRequest()) + pong = proto.recv_frame(self._stream) except (OSError, EOFError) as exc: - stderr_text = self._read_stderr() - raise RuntimeError( - f"pbspy-server --proxy did not respond to ping on {self._host!r}: {exc}" - + (f"\nSSH stderr: {stderr_text}" if stderr_text else "") - ) from exc + self._stream.close() + sock.close() + raise RuntimeError(f"pbspy-server at {self._host}:{self._port} did not respond to ping: {exc}") from exc if not isinstance(pong, proto.PongResponse): - stderr_text = self._read_stderr() - raise RuntimeError( - f"pbspy-server --proxy returned unexpected response: {pong!r}" - + (f"\nSSH stderr: {stderr_text}" if stderr_text else "") - ) - - def _read_stderr(self) -> str: - """Non-blocking read of any immediately available stderr from the SSH process.""" - if self._proc is None or self._proc.stderr is None: - return "" - try: - rlist, _, _ = select.select([self._proc.stderr], [], [], 0.0) - if rlist: - return self._proc.stderr.read(4096).decode(errors="replace") - except OSError: - pass - return "" + self._stream.close() + sock.close() + raise RuntimeError(f"pbspy-server returned unexpected response to ping: {pong!r}") def _reconnect(self) -> None: - """Terminate the current SSH subprocess and establish a fresh connection.""" - if self._proc is not None: + """Close the current connection and establish a fresh one.""" + if self._stream is not None: + try: + self._stream.close() + except OSError: + pass + self._stream = None + if self._sock is not None: try: - self._proc.terminate() + self._sock.close() except OSError: pass - self._proc = None + self._sock = None self._connect() def _rpc(self, request: object) -> object: """Send one request frame and return the first response frame.""" with self._lock: try: - assert self._stdin is not None - assert self._stdout is not None - proto.send_frame(self._stdin, request) - return proto.recv_frame(self._stdout) + assert self._stream is not None + proto.send_frame(self._stream, request) + return proto.recv_frame(self._stream) except (OSError, EOFError): - # Connection lost — attempt one reconnect then retry. self._reconnect() - assert self._stdin is not None - assert self._stdout is not None - proto.send_frame(self._stdin, request) - return proto.recv_frame(self._stdout) + assert self._stream is not None + proto.send_frame(self._stream, request) + return proto.recv_frame(self._stream) def _send(self, request: object) -> None: with self._lock: - assert self._stdin is not None - proto.send_frame(self._stdin, request) + assert self._stream is not None + proto.send_frame(self._stream, request) def _recv(self) -> object: - assert self._stdout is not None - return proto.recv_frame(self._stdout) + assert self._stream is not None + return proto.recv_frame(self._stream) def _wait_quiet( self, @@ -214,10 +199,18 @@ def _wait_with_progress( raise RuntimeError(f"Server error while waiting: {response.message}") def close(self) -> None: - """Terminate the SSH subprocess.""" - if self._proc is not None: - self._proc.terminate() - self._proc = None + if self._stream is not None: + try: + self._stream.close() + except OSError: + pass + self._stream = None + if self._sock is not None: + try: + self._sock.close() + except OSError: + pass + self._sock = None def __del__(self) -> None: try: diff --git a/src/pbspy_server/__main__.py b/src/pbspy_server/__main__.py index 41206c1..322fb78 100644 --- a/src/pbspy_server/__main__.py +++ b/src/pbspy_server/__main__.py @@ -3,97 +3,57 @@ Usage:: - pbspy-server --daemon # start server daemon in the foreground - pbspy-server --proxy # SSH proxy mode (stdio ↔ socket splice) - pbspy-server --status # print server status and exit + pbspy-server [--ssh-host HOST] [--ssh-user USER] [--port PORT] [--ssh-arg ARG ...] """ from __future__ import annotations import argparse import logging -import socket -import sys -from pathlib import Path -from typing import BinaryIO, cast - -_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" - - -def _cmd_daemon() -> None: - from pbspy_server.server import run_server - - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - run_server() - - -def _cmd_proxy() -> None: - from pbspy_server.proxy import run_proxy - - run_proxy() - - -def _cmd_status() -> None: - if not _ADDR_PATH.exists(): - print("pbspy-server: not running (address file not found)") - sys.exit(1) - - try: - import json - - data = json.loads(_ADDR_PATH.read_text()) - host = data["host"] - port = int(data["port"]) - token = data["token"] - except (ValueError, KeyError, OSError) as exc: - print(f"pbspy-server: malformed address file ({exc})") - sys.exit(1) - - try: - import pbspy._protocol as proto - - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(3.0) - sock.connect((host, port)) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - sock.close() - proto.send_frame(f, proto.AuthRequest(token=token)) - auth_resp = proto.recv_frame(f) - if not isinstance(auth_resp, proto.AuthOkResponse): - print(f"pbspy-server: auth failed: {auth_resp!r}") - f.close() - sys.exit(1) - proto.send_frame(f, proto.PingRequest()) - response = proto.recv_frame(f) - f.close() - if isinstance(response, proto.PongResponse): - print(f"pbspy-server: running on {host}:{port}") - else: - print(f"pbspy-server: unexpected response: {response!r}") - sys.exit(1) - except OSError as exc: - print(f"pbspy-server: not responding ({exc})") - sys.exit(1) def main() -> None: parser = argparse.ArgumentParser( prog="pbspy-server", - description="PBS job scheduler server daemon for pbspy.", + description="PBS job server daemon for pbspy.", + ) + parser.add_argument( + "--ssh-host", + metavar="HOST", + help="Supercomputer hostname to SSH into for PBS commands.", + ) + parser.add_argument( + "--ssh-user", + metavar="USER", + help="SSH username on the supercomputer.", + ) + parser.add_argument( + "--port", + type=int, + default=9876, + metavar="PORT", + help="TCP port to listen on (default: 9876).", + ) + parser.add_argument( + "--ssh-arg", + dest="ssh_args", + action="append", + metavar="ARG", + help="Extra SSH argument (may be repeated, e.g. --ssh-arg=-i --ssh-arg=/path/to/key).", ) - group = parser.add_mutually_exclusive_group(required=True) - group.add_argument("--daemon", action="store_true", help="Start the server daemon (runs in foreground)") - group.add_argument("--proxy", action="store_true", help="SSH proxy mode: splice stdin/stdout with server socket") - group.add_argument("--status", action="store_true", help="Print server status and exit") args = parser.parse_args() - if args.daemon: - _cmd_daemon() - elif args.proxy: - _cmd_proxy() - elif args.status: - _cmd_status() + logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") + + from pbspy_server.server import run_server + + run_server( + port=args.port, + ssh_host=args.ssh_host, + ssh_user=args.ssh_user, + ssh_args=args.ssh_args, + ) if __name__ == "__main__": diff --git a/src/pbspy_server/proxy.py b/src/pbspy_server/proxy.py deleted file mode 100644 index a817d61..0000000 --- a/src/pbspy_server/proxy.py +++ /dev/null @@ -1,174 +0,0 @@ -""" -SSH proxy mode for pbspy-server. - -When invoked as ``pbspy-server --proxy`` over an SSH connection, this module: - -1. Acquires an exclusive file lock (``~/.pbspy/server.lock``) to prevent two - proxy processes on different login nodes from racing to start the daemon. -2. Checks whether the server daemon is already running by reading - ``~/.pbspy/server.addr`` and attempting an authenticated TCP connection. -3. If the daemon is not running, starts it in the background and waits for the - address file to appear. -4. Splices ``sys.stdin`` / ``sys.stdout`` bytes bidirectionally with the server - TCP socket so the remote client can talk to the daemon transparently. - -Because the address file lives on the shared home filesystem, this works -correctly regardless of which login node the SSH session lands on — the daemon -may be running on a completely different node. -""" - -from __future__ import annotations - -import fcntl -import json -import os -import select -import socket -import subprocess -import sys -import time -from pathlib import Path -from typing import BinaryIO, cast - -import pbspy._protocol as proto - -__all__ = ["run_proxy"] - -_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" -_LOCK_PATH = Path.home() / ".pbspy" / "server.lock" -_STARTUP_TIMEOUT = 10.0 # seconds to wait for daemon to start -_POLL_INTERVAL = 0.1 -_BUF_SIZE = 65536 - - -def _read_server_addr() -> tuple[str, int, str] | None: - """Return ``(hostname, port, token)`` from the JSON addr file, or ``None`` if absent/invalid.""" - if not _ADDR_PATH.exists(): - return None - try: - data = json.loads(_ADDR_PATH.read_text()) - return data["host"], int(data["port"]), data["token"] - except (ValueError, KeyError, OSError): - return None - - -def _daemon_alive() -> bool: - """Return True if the addr file exists and the daemon responds to an auth+ping.""" - addr = _read_server_addr() - if addr is None: - return False - host, port, token = addr - try: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(2.0) - sock.connect((host, port)) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - sock.close() - proto.send_frame(f, proto.AuthRequest(token=token)) - resp = proto.recv_frame(f) - if not isinstance(resp, proto.AuthOkResponse): - f.close() - return False - proto.send_frame(f, proto.PingRequest()) - resp2 = proto.recv_frame(f) - f.close() - return isinstance(resp2, proto.PongResponse) - except OSError: - # Stale addr file — remove it so the next caller starts a fresh daemon. - try: - _ADDR_PATH.unlink(missing_ok=True) - except OSError: - pass - return False - - -def _start_daemon() -> None: - """Fork the server daemon in the background.""" - subprocess.Popen( - [sys.executable, "-m", "pbspy_server", "--daemon"], - stdin=subprocess.DEVNULL, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, - start_new_session=True, - ) - - -def _wait_for_daemon(timeout: float = _STARTUP_TIMEOUT) -> None: - """Block until the daemon is alive or *timeout* seconds have elapsed.""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if _daemon_alive(): - return - time.sleep(_POLL_INTERVAL) - raise TimeoutError(f"pbspy-server daemon did not start within {timeout:.0f}s") - - -def run_proxy() -> None: - """ - Entry point for ``--proxy`` mode. - - Ensures the daemon is running (starting it if necessary), then authenticates - and splices stdin/stdout with the server TCP socket until one side closes. - """ - _ADDR_PATH.parent.mkdir(parents=True, exist_ok=True) - - # Acquire exclusive lock while checking / starting the daemon to prevent - # two proxy processes on different login nodes from both deciding to start it. - lock_file = open(_LOCK_PATH, "w") # noqa: SIM115 - try: - fcntl.flock(lock_file, fcntl.LOCK_EX) - if not _daemon_alive(): - _start_daemon() - _wait_for_daemon() - addr = _read_server_addr() - finally: - fcntl.flock(lock_file, fcntl.LOCK_UN) - lock_file.close() - - if addr is None: - raise RuntimeError("pbspy-server daemon started but address file not found") - - host, port, token = addr - srv_sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - srv_sock.connect((host, port)) - - # Authenticate with the daemon before entering the byte-splice loop. - srv_file = cast(BinaryIO, srv_sock.makefile("rwb", buffering=0)) - proto.send_frame(srv_file, proto.AuthRequest(token=token)) - auth_resp = proto.recv_frame(srv_file) - srv_file.close() # File object done; raw socket takes over below - if not isinstance(auth_resp, proto.AuthOkResponse): - srv_sock.close() - raise RuntimeError(f"pbspy-server rejected auth token: {auth_resp!r}") - - stdin_fd = sys.stdin.buffer.fileno() - srv_fd = srv_sock.fileno() - - # Put stdin in non-blocking mode; stdout remains blocking so writes - # don't need special partial-write handling. - _set_nonblocking(stdin_fd) - - try: - while True: - readable, _, exceptional = select.select([stdin_fd, srv_fd], [], [stdin_fd, srv_fd], 1.0) - if exceptional: - break - for fd in readable: - if fd == stdin_fd: - data = os.read(stdin_fd, _BUF_SIZE) - if not data: - return - srv_sock.sendall(data) - elif fd == srv_fd: - data = srv_sock.recv(_BUF_SIZE) - if not data: - return - sys.stdout.buffer.write(data) - sys.stdout.buffer.flush() - finally: - srv_sock.close() - - -def _set_nonblocking(fd: int) -> None: - flags = fcntl.fcntl(fd, fcntl.F_GETFL) - fcntl.fcntl(fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) diff --git a/src/pbspy_server/server.py b/src/pbspy_server/server.py index b70572c..a60aad9 100644 --- a/src/pbspy_server/server.py +++ b/src/pbspy_server/server.py @@ -1,28 +1,25 @@ """ pbspy-server daemon. -Listens on a TCP socket (hostname:random-port), writes the address to -``~/.pbspy/server.addr`` so that any login node can find and connect to it, -accepts one connection per client, and dispatches pickled requests from -:mod:`pbspy._protocol`. +Listens on a configured TCP port, accepts connections from :class:`~pbspy.ServerBackend` +clients, and dispatches pickled requests from :mod:`pbspy._protocol`. + +PBS operations (qsub, qstat, file reads) are executed via SSH to a remote +supercomputer when ``ssh_host`` is provided, or locally when omitted. A background thread polls ``qstat`` every 60 seconds for all unfinished jobs and pushes :class:`StatusUpdateResponse` frames to any clients waiting on -those jobs. Job state is held in memory for the daemon's lifetime; it is not -persisted across restarts. +those jobs. Job state is held in memory and is not persisted across restarts. """ from __future__ import annotations -import json import logging -import secrets import signal import socket -import subprocess import threading +from collections.abc import Callable from dataclasses import dataclass -from pathlib import Path from typing import BinaryIO, cast import pbspy._pbs_core as core @@ -33,7 +30,6 @@ logger = logging.getLogger(__name__) -_ADDR_PATH = Path.home() / ".pbspy" / "server.addr" _POLL_INTERVAL = 60.0 # seconds between qstat polls @@ -43,8 +39,7 @@ class _JobRecord: job_id: str job_name: str | None = None - state: str | None = None # "Q", "R", "E", "F" - exit_code: int | None = None + finished: bool = False class _WaitSubscription: @@ -61,19 +56,12 @@ def __init__(self, jobs: list[Job], stream: BinaryIO) -> None: class _ServerState: """Shared mutable state accessed by both connection threads and the poll thread.""" - def __init__(self, token: str) -> None: - self.token = token + def __init__(self, runner: core.PBSRunner) -> None: + self.runner = runner self.lock = threading.Lock() - self.active_connections: int = 0 - self._ever_had_connection: bool = False self.jobs: dict[str, _JobRecord] = {} self.subscriptions: dict[str, list[_WaitSubscription]] = {} - def add_subscription(self, sub: _WaitSubscription) -> None: - with self.lock: - for job_id in sub.pending: - self.subscriptions.setdefault(job_id, []).append(sub) - def notify_finished(self, job_id: str, job: Job) -> None: """Mark job_id as done in all waiting subscriptions, sending frames as needed.""" with self.lock: @@ -106,68 +94,40 @@ def _cancel_all_subscriptions(state: _ServerState) -> None: sub.done.set() -def _should_shutdown(state: _ServerState) -> bool: - """Return True when the daemon has nothing left to do and can exit.""" - with state.lock: - if not state._ever_had_connection: - return False - if state.active_connections > 0: - return False - return not any(r.state != "F" for r in state.jobs.values()) - - -def _poll_loop(state: _ServerState, stop_event: threading.Event, poll_interval: float = _POLL_INTERVAL) -> None: +def _poll_loop( + state: _ServerState, + stop_event: threading.Event, + poll_interval: float = _POLL_INTERVAL, +) -> None: """Background thread: polls qstat for all unfinished jobs.""" while not stop_event.wait(timeout=poll_interval): try: with state.lock: - unfinished = [r for r in state.jobs.values() if r.state != "F"] + unfinished = [r for r in state.jobs.values() if not r.finished] for record in unfinished: job = Job(job_id=record.job_id, job_name=record.job_name) - result = _check_job_finished(record.job_id) - if result is not None: + if _check_job_finished(record.job_id, state.runner) is not None: with state.lock: - record.state = "F" - record.exit_code = result + record.finished = True state.notify_finished(record.job_id, job) - if _should_shutdown(state): - logger.info("No active connections and no unfinished jobs; shutting down") - stop_event.set() except Exception: logger.exception("Error in poll loop") -def _check_job_finished(job_id: str) -> int | None: +def _check_job_finished(job_id: str, runner: core.PBSRunner) -> int | None: """Return exit code if job has finished, else None.""" - process = subprocess.run(["qstat", job_id], capture_output=True) + process = runner.run(["qstat", job_id]) if process.returncode != 0: stdout = process.stdout.decode("utf-8") if job_id not in stdout or "has finished" in process.stderr.decode("utf-8"): - return core.try_get_exit_code(job_id) or 0 + return core.try_get_exit_code(job_id, runner) or 0 return None -def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: threading.Event) -> None: +def _handle_connection(conn: socket.socket, state: _ServerState) -> None: """Handle a single client connection in its own thread.""" - with state.lock: - state.active_connections += 1 - state._ever_had_connection = True stream = cast(BinaryIO, conn.makefile("rwb", buffering=0)) try: - # First frame must be a valid AuthRequest. - try: - auth = proto.recv_frame(stream) - except EOFError: - return - if not isinstance(auth, proto.AuthRequest) or auth.token != state.token: - logger.warning("Connection rejected: invalid or missing auth token") - try: - proto.send_frame(stream, proto.ErrorResponse(message="Authentication failed")) - except OSError: - pass - return - proto.send_frame(stream, proto.AuthOkResponse()) - while True: try: request = proto.recv_frame(stream) @@ -182,10 +142,10 @@ def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: thr proto.send_frame(stream, proto.PongResponse()) elif isinstance(request, proto.SubmitRequest): - job_id, job_name = core.pbs_submit(request.script, request.name) + job_id, job_name = core.pbs_submit(request.script, request.name, runner=state.runner) job = Job(job_id=job_id, job_name=job_name) with state.lock: - state.jobs[job_id] = _JobRecord(job_id=job_id, job_name=job_name, state="Q") + state.jobs[job_id] = _JobRecord(job_id=job_id, job_name=job_name) proto.send_frame(stream, proto.SubmittedResponse(job=job)) elif isinstance(request, proto.WaitRequest): @@ -196,7 +156,7 @@ def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: thr with state.lock: for job in request.jobs: record = state.jobs.get(job.job_id) - if record and record.state == "F": + if record and record.finished: sub.pending.discard(job.job_id) if sub.pending: for job_id in sub.pending: @@ -207,7 +167,7 @@ def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: thr proto.send_frame(stream, proto.WaitDoneResponse(jobs=request.jobs)) elif isinstance(request, proto.ResultRequest): - result: JobResult = core.pbs_get_result(request.job) + result: JobResult = core.pbs_get_result(request.job, runner=state.runner) proto.send_frame(stream, proto.ResultResponse(job=request.job, result=result)) else: @@ -222,52 +182,50 @@ def _handle_connection(conn: socket.socket, state: _ServerState, stop_event: thr finally: stream.close() conn.close() - with state.lock: - state.active_connections -= 1 - if _should_shutdown(state): - logger.info("Last client disconnected with no unfinished jobs; shutting down") - stop_event.set() def run_server( - addr_path: Path = _ADDR_PATH, + host: str = "0.0.0.0", + port: int = 9876, + ssh_host: str | None = None, + ssh_user: str | None = None, + ssh_args: list[str] | None = None, poll_interval: float = _POLL_INTERVAL, + _stop_event: threading.Event | None = None, + _ready_callback: Callable[[int], None] | None = None, ) -> None: """ Start the pbspy-server daemon. - Runs in the foreground; the caller is responsible for daemonising the - process (e.g. ``--daemon`` in :mod:`pbspy_server.__main__` forks and - calls this in the child). - - The server binds a TCP port on ``0.0.0.0`` (OS chooses a free port) and - writes a JSON address file to *addr_path* containing the hostname, port, - and a random auth token. Clients must send the token as the first frame - on every connection; connections with an incorrect or missing token are - rejected immediately. + Runs in the foreground; use your OS's service manager (systemd, nohup, + etc.) to run it as a background service. - Job state is held in memory for the daemon's lifetime and is not - persisted across restarts. + When *ssh_host* is provided, all PBS operations (qsub, qstat, file reads) + are executed on that host via SSH. Otherwise they run locally. Args: - addr_path: Path for the address file (JSON). Defaults to - ``~/.pbspy/server.addr``. - poll_interval: Seconds between qstat polls (default 60). Override in tests. + host: Local address to bind (default ``"0.0.0.0"``). + port: TCP port to listen on (default 9876; pass 0 for OS-assigned). + ssh_host: Supercomputer hostname to SSH into for PBS commands. + ssh_user: SSH username (combined with *ssh_host* as ``user@host``). + ssh_args: Extra arguments forwarded to the ``ssh`` command. + poll_interval: Seconds between qstat polls (default 60). """ - addr_path.parent.mkdir(parents=True, exist_ok=True) - if addr_path.exists(): - addr_path.unlink() + ssh_prefix: list[str] | None = None + if ssh_host is not None: + destination = f"{ssh_user}@{ssh_host}" if ssh_user else ssh_host + ssh_prefix = ["ssh", *(ssh_args or []), destination] + + runner = core.PBSRunner(ssh_prefix) + state = _ServerState(runner=runner) - token = secrets.token_hex(32) - state = _ServerState(token=token) + stop_event = _stop_event or threading.Event() - stop_event = threading.Event() poll_thread = threading.Thread( target=_poll_loop, args=(state, stop_event), kwargs={"poll_interval": poll_interval}, daemon=True ) poll_thread.start() - # Signal handlers can only be registered from the main thread. if threading.current_thread() is threading.main_thread(): def _handle_signal(signum: int, frame: object) -> None: @@ -279,15 +237,17 @@ def _handle_signal(signum: int, frame: object) -> None: srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - srv.bind(("", 0)) + srv.bind((host, port)) srv.listen(16) srv.settimeout(1.0) - hostname = socket.gethostname() - port = srv.getsockname()[1] - addr_path.write_text(json.dumps({"host": hostname, "port": port, "token": token})) + actual_port = srv.getsockname()[1] + logger.info("pbspy-server listening on %s:%d", host, actual_port) + if ssh_host: + logger.info("PBS commands will run via SSH on %s", destination) - logger.info("pbspy-server listening on %s:%d", hostname, port) + if _ready_callback is not None: + _ready_callback(actual_port) try: while not stop_event.is_set(): @@ -295,11 +255,9 @@ def _handle_signal(signum: int, frame: object) -> None: conn, _ = srv.accept() except TimeoutError: continue - t = threading.Thread(target=_handle_connection, args=(conn, state, stop_event), daemon=True) + t = threading.Thread(target=_handle_connection, args=(conn, state), daemon=True) t.start() finally: srv.close() - if addr_path.exists(): - addr_path.unlink() _cancel_all_subscriptions(state) logger.info("pbspy-server stopped") diff --git a/tests/test_default.py b/tests/test_default.py index a406310..7a94c98 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -3,7 +3,7 @@ import pytest -from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend, SSHBackend +from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend, ServerBackend @pytest.mark.skipif(shutil.which("qsub") is None, reason="qsub not available") @@ -124,6 +124,6 @@ def test_script_generation_unchanged() -> None: def test_backend_classes_importable() -> None: - """Backend, LocalBackend, SSHBackend are importable from pbspy.""" + """Backend, LocalBackend, ServerBackend are importable from pbspy.""" assert issubclass(LocalBackend, Backend) - assert issubclass(SSHBackend, Backend) + assert issubclass(ServerBackend, Backend) diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 0dac621..3c80be6 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -28,13 +28,6 @@ def roundtrip(obj: object) -> object: # --------------------------------------------------------------------------- -def test_roundtrip_auth_request() -> None: - req = proto.AuthRequest(token="abc123secret") - result = roundtrip(req) - assert isinstance(result, proto.AuthRequest) - assert result.token == "abc123secret" - - def test_roundtrip_ping() -> None: result = roundtrip(proto.PingRequest()) assert isinstance(result, proto.PingRequest) @@ -76,10 +69,6 @@ def test_roundtrip_result_request() -> None: # --------------------------------------------------------------------------- -def test_roundtrip_auth_ok() -> None: - assert isinstance(roundtrip(proto.AuthOkResponse()), proto.AuthOkResponse) - - def test_roundtrip_pong() -> None: assert isinstance(roundtrip(proto.PongResponse()), proto.PongResponse) diff --git a/tests/test_server.py b/tests/test_server.py index ef3e864..fdc933d 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -7,12 +7,11 @@ from __future__ import annotations -import json +import queue import socket import threading import time from collections.abc import Generator -from pathlib import Path from typing import BinaryIO, cast from unittest.mock import MagicMock, patch @@ -26,35 +25,20 @@ # Fixture: running server # --------------------------------------------------------------------------- -_WAIT_TIMEOUT = 5.0 # seconds to wait for server address file to appear - - -def _wait_for_server(addr_path: Path, timeout: float = _WAIT_TIMEOUT) -> None: - """Wait until the server address file appears (server is bound and ready).""" - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if addr_path.exists(): - return - time.sleep(0.05) - raise TimeoutError(f"Server addr file {addr_path} did not appear within {timeout}s") - - -def _read_addr(addr_path: Path) -> tuple[str, int, str]: - data = json.loads(addr_path.read_text()) - return data["host"], int(data["port"]), data["token"] +_WAIT_TIMEOUT = 5.0 # seconds to wait for the server to become ready class ServerHandle: """Handle to a running test server.""" - def __init__(self, addr_path: Path) -> None: - self.addr_path = addr_path + def __init__(self, host: str, port: int) -> None: + self.host = host + self.port = port def connect(self, timeout: float | None = None) -> BinaryIO: - """Open an authenticated TCP connection and return a binary file-like object.""" - host, port, token = _read_addr(self.addr_path) + """Open a TCP connection and return a binary file-like object.""" sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((host, port)) + sock.connect((self.host, self.port)) if timeout is not None: sock.settimeout(timeout) f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) @@ -62,9 +46,6 @@ def connect(self, timeout: float | None = None) -> BinaryIO: # calls _real_close() (via _decref_socketios). Without this, the socket # FD leaks and Python emits a ResourceWarning. sock.close() - proto.send_frame(f, proto.AuthRequest(token=token)) - auth_resp = proto.recv_frame(f) - assert isinstance(auth_resp, proto.AuthOkResponse), f"Auth failed: {auth_resp!r}" return f def rpc(self, request: object) -> object: @@ -102,23 +83,22 @@ def mock_pbs_get_result() -> Generator[MagicMock, None, None]: @pytest.fixture() def server( - tmp_path: Path, mock_pbs_submit: MagicMock, mock_check_finished: MagicMock, mock_pbs_get_result: MagicMock, ) -> Generator[ServerHandle, None, None]: - addr_path = tmp_path / "test.addr" + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() t = threading.Thread( target=run_server, - kwargs={"addr_path": addr_path, "poll_interval": 0.1}, + kwargs={"port": 0, "poll_interval": 0.1, "_ready_callback": port_q.put, "_stop_event": stop}, daemon=True, ) t.start() - _wait_for_server(addr_path) - yield ServerHandle(addr_path) - # Server exits on its own once no clients + no jobs (shutdown test). - # For other tests, give it a moment to clean up. + port = port_q.get(timeout=_WAIT_TIMEOUT) + yield ServerHandle("127.0.0.1", port) + stop.set() t.join(timeout=2.0) @@ -132,46 +112,6 @@ def test_ping(server: ServerHandle) -> None: assert isinstance(response, proto.PongResponse) -# --------------------------------------------------------------------------- -# Auth — rejection cases -# --------------------------------------------------------------------------- - - -def test_auth_rejected_wrong_token(server: ServerHandle) -> None: - """Connecting with a wrong token must return ErrorResponse and close the connection.""" - host, port, _token = _read_addr(server.addr_path) - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((host, port)) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - sock.close() - - proto.send_frame(f, proto.AuthRequest(token="wrong-token")) - resp = proto.recv_frame(f) - assert isinstance(resp, proto.ErrorResponse) - # Server should close the connection after rejecting auth. - try: - proto.recv_frame(f) - received_eof = False - except EOFError: - received_eof = True - f.close() - assert received_eof, "Expected server to close connection after auth failure" - - -def test_auth_rejected_non_auth_first_frame(server: ServerHandle) -> None: - """Sending a non-AuthRequest as the first frame must be rejected.""" - host, port, _token = _read_addr(server.addr_path) - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((host, port)) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - sock.close() - - proto.send_frame(f, proto.PingRequest()) # Wrong first frame - resp = proto.recv_frame(f) - assert isinstance(resp, proto.ErrorResponse) - f.close() - - # --------------------------------------------------------------------------- # Submit # --------------------------------------------------------------------------- @@ -202,10 +142,10 @@ def test_result_returns_job_result(server: ServerHandle) -> None: # --------------------------------------------------------------------------- -def test_wait_already_finished_returns_immediately(tmp_path: Path) -> None: +def test_wait_already_finished_returns_immediately() -> None: """If the job is already marked F in state, WaitRequest should return without blocking.""" - addr_path = tmp_path / "wait-done.addr" - server_done = threading.Event() + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() def _run() -> None: with ( @@ -214,23 +154,20 @@ def _run() -> None: patch("pbspy_server.server._check_job_finished", return_value=0), patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), ): - run_server(addr_path=addr_path, poll_interval=0.1) - server_done.set() + run_server(port=0, poll_interval=0.1, _ready_callback=port_q.put, _stop_event=stop) t = threading.Thread(target=_run, daemon=True) t.start() - _wait_for_server(addr_path) + port = port_q.get(timeout=_WAIT_TIMEOUT) + handle = ServerHandle("127.0.0.1", port) # Submit the job, then wait long enough for the poll loop (0.1s) to mark it done. - stream = _raw_connect(addr_path) - proto.send_frame(stream, proto.SubmitRequest(script="#!/bin/bash", name="done_job")) - proto.recv_frame(stream) - stream.close() + handle.rpc(proto.SubmitRequest(script="#!/bin/bash", name="done_job")) time.sleep(0.5) # poll interval is 0.1s; 0.5s is more than enough job = Job(job_id="done.mock", job_name="done_job") - stream = _raw_connect(addr_path) + stream = handle.connect() proto.send_frame(stream, proto.WaitRequest(jobs=[job])) response = proto.recv_frame(stream) @@ -239,6 +176,7 @@ def _run() -> None: assert isinstance(response, proto.WaitDoneResponse) assert response.jobs[0].job_id == "done.mock" + stop.set() t.join(timeout=2.0) @@ -275,99 +213,3 @@ def test_wait_receives_status_update_then_done(server: ServerHandle, mock_check_ stream.close() assert any(isinstance(r, proto.WaitDoneResponse) for r in responses), f"Expected WaitDoneResponse; got: {responses}" - - -# --------------------------------------------------------------------------- -# Shutdown: no jobs + last client disconnects -# --------------------------------------------------------------------------- - - -def test_shutdown_when_no_jobs_and_no_clients(tmp_path: Path) -> None: - """ - The server should exit automatically when the last client disconnects - and there are no unfinished jobs. - """ - addr_path = tmp_path / "shutdown.addr" - - server_done = threading.Event() - - def _run() -> None: - with ( - patch("pbspy_server.server.core.pbs_submit", return_value=("1.mock", "j")), - patch("pbspy_server.server._check_job_finished", return_value=None), - patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), - ): - run_server(addr_path=addr_path) - server_done.set() - - t = threading.Thread(target=_run, daemon=True) - t.start() - _wait_for_server(addr_path) - - # Connect, ping, and disconnect — no jobs submitted - stream = _raw_connect(addr_path) - proto.send_frame(stream, proto.PingRequest()) - proto.recv_frame(stream) - stream.close() - - assert server_done.wait(timeout=5.0), "Server did not exit after last client disconnected with no jobs" - - -# --------------------------------------------------------------------------- -# Shutdown: stays alive while jobs pending -# --------------------------------------------------------------------------- - - -def test_server_stays_alive_while_jobs_pending(tmp_path: Path) -> None: - """ - If jobs are still unfinished, the server must NOT exit when the last - client disconnects; it should exit once the poll loop detects them done. - """ - addr_path = tmp_path / "pending.addr" - server_done = threading.Event() - - with ( - patch("pbspy_server.server.core.pbs_submit", return_value=("2.mock", "pending_job")) as _, - patch("pbspy_server.server._check_job_finished", return_value=None) as mock_check, - patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), - ): - - def _run() -> None: - run_server(addr_path=addr_path, poll_interval=0.1) - server_done.set() - - t = threading.Thread(target=_run, daemon=True) - t.start() - _wait_for_server(addr_path) - - # Submit a job then disconnect - stream = _raw_connect(addr_path) - proto.send_frame(stream, proto.SubmitRequest(script="#!/bin/bash", name="pending_job")) - proto.recv_frame(stream) - stream.close() - - # Server should NOT have shut down yet (jobs still pending) - assert not server_done.wait(timeout=1.0), "Server exited prematurely while jobs were still pending" - - # Let the poll loop see the job as finished → server should then exit. - mock_check.return_value = 0 - assert server_done.wait(timeout=5.0), "Server did not exit after all jobs finished" - - t.join(timeout=2.0) - - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def _raw_connect(addr_path: Path) -> BinaryIO: - host, port, token = _read_addr(addr_path) - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((host, port)) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - sock.close() # Mark closed; SocketIO keeps FD open via _io_refs - proto.send_frame(f, proto.AuthRequest(token=token)) - auth_resp = proto.recv_frame(f) - assert isinstance(auth_resp, proto.AuthOkResponse), f"Auth failed: {auth_resp!r}" - return f From f3c8273c950e731e2e04342265ac6090e21eb718 Mon Sep 17 00:00:00 2001 From: omdowley Date: Tue, 2 Jun 2026 16:16:37 +1000 Subject: [PATCH 03/16] feat: add docker-compose for pbspy-server Includes .env.example for SSH configuration. PBSPY_SSH_HOST and PBSPY_SSH_USER are also read from environment as argparse defaults. --- .env.example | 3 +++ .gitignore | 1 + Dockerfile | 15 +++++++++++++++ docker-compose.yml | 10 ++++++++++ src/pbspy_server/__main__.py | 3 +++ 5 files changed, 32 insertions(+) create mode 100644 .env.example create mode 100644 Dockerfile create mode 100644 docker-compose.yml diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..53e4203 --- /dev/null +++ b/.env.example @@ -0,0 +1,3 @@ +PBSPY_SSH_HOST=gadi.nci.org.au +# PBSPY_SSH_USER=ab1234 # omit if your local username matches or ~/.ssh/config handles it +# PBSPY_PORT=9876 # override if 9876 is already in use on the host diff --git a/.gitignore b/.gitignore index f6c4123..bcde509 100644 --- a/.gitignore +++ b/.gitignore @@ -174,3 +174,4 @@ cython_debug/ *.e* *.o* +!.env.example diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..530f548 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,15 @@ +FROM python:3.12-slim + +RUN apt-get update \ + && apt-get install -y --no-install-recommends openssh-client \ + && rm -rf /var/lib/apt/lists/* + +RUN useradd -m -u 1000 pbspy +WORKDIR /home/pbspy + +COPY --chown=pbspy:pbspy . . +RUN pip install --no-cache-dir . + +USER pbspy +EXPOSE 9876 +ENTRYPOINT ["pbspy-server"] diff --git a/docker-compose.yml b/docker-compose.yml new file mode 100644 index 0000000..36f9995 --- /dev/null +++ b/docker-compose.yml @@ -0,0 +1,10 @@ +services: + pbspy-server: + build: . + ports: + - "${PBSPY_PORT:-9876}:9876" + volumes: + - "${HOME}/.ssh:/home/pbspy/.ssh:ro" + environment: + - PBSPY_SSH_HOST=${SSH_HOST} + - PBSPY_SSH_USER=${SSH_USER:-} diff --git a/src/pbspy_server/__main__.py b/src/pbspy_server/__main__.py index 322fb78..9dbc0d5 100644 --- a/src/pbspy_server/__main__.py +++ b/src/pbspy_server/__main__.py @@ -10,6 +10,7 @@ import argparse import logging +import os def main() -> None: @@ -20,11 +21,13 @@ def main() -> None: parser.add_argument( "--ssh-host", metavar="HOST", + default=os.environ.get("PBSPY_SSH_HOST"), help="Supercomputer hostname to SSH into for PBS commands.", ) parser.add_argument( "--ssh-user", metavar="USER", + default=os.environ.get("PBSPY_SSH_USER"), help="SSH username on the supercomputer.", ) parser.add_argument( From df1a25e18fc4ab3512d2779c22d0c30e410b7fb2 Mon Sep 17 00:00:00 2001 From: omdowley Date: Thu, 4 Jun 2026 14:22:13 +1000 Subject: [PATCH 04/16] fix: move pbspy_server module to pbspy.server This fixes a bug in the Dockerfile caused by the pbspy_server module not being installed by `hatchling`. --- pyproject.toml | 2 +- src/pbspy/server/__init__.py | 7 +++++++ src/{pbspy_server => pbspy/server}/__main__.py | 2 +- src/{pbspy_server => pbspy/server}/server.py | 1 + src/pbspy_server/__init__.py | 1 - tests/test_server.py | 16 ++++++++-------- 6 files changed, 18 insertions(+), 11 deletions(-) create mode 100644 src/pbspy/server/__init__.py rename src/{pbspy_server => pbspy/server}/__main__.py (96%) rename src/{pbspy_server => pbspy/server}/server.py (99%) delete mode 100644 src/pbspy_server/__init__.py diff --git a/pyproject.toml b/pyproject.toml index 5369551..0dee784 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ dependencies = [ ] [project.scripts] -pbspy-server = "pbspy_server.__main__:main" +pbspy-server = "pbspy.server.__main__:main" [dependency-groups] dev = [ diff --git a/src/pbspy/server/__init__.py b/src/pbspy/server/__init__.py new file mode 100644 index 0000000..bc2aee6 --- /dev/null +++ b/src/pbspy/server/__init__.py @@ -0,0 +1,7 @@ +"""pbspy-server: PBS job scheduler server daemon for pbspy.""" + +from __future__ import annotations + +from pbspy.server.server import run_server + +__all__ = ["run_server"] diff --git a/src/pbspy_server/__main__.py b/src/pbspy/server/__main__.py similarity index 96% rename from src/pbspy_server/__main__.py rename to src/pbspy/server/__main__.py index 9dbc0d5..18f51d2 100644 --- a/src/pbspy_server/__main__.py +++ b/src/pbspy/server/__main__.py @@ -49,7 +49,7 @@ def main() -> None: logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - from pbspy_server.server import run_server + from pbspy.server import run_server run_server( port=args.port, diff --git a/src/pbspy_server/server.py b/src/pbspy/server/server.py similarity index 99% rename from src/pbspy_server/server.py rename to src/pbspy/server/server.py index a60aad9..82306db 100644 --- a/src/pbspy_server/server.py +++ b/src/pbspy/server/server.py @@ -211,6 +211,7 @@ def run_server( ssh_args: Extra arguments forwarded to the ``ssh`` command. poll_interval: Seconds between qstat polls (default 60). """ + destination: str | None = None ssh_prefix: list[str] | None = None if ssh_host is not None: destination = f"{ssh_user}@{ssh_host}" if ssh_user else ssh_host diff --git a/src/pbspy_server/__init__.py b/src/pbspy_server/__init__.py deleted file mode 100644 index cf731a8..0000000 --- a/src/pbspy_server/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""pbspy-server: PBS job scheduler server daemon for pbspy.""" diff --git a/tests/test_server.py b/tests/test_server.py index fdc933d..7b7e16a 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -1,5 +1,5 @@ """ -Integration tests for pbspy_server.server. +Integration tests for pbspy.server. Spins up a real server (in a background thread) with mocked PBS operations so that no ``qsub``/``qstat`` installation is needed. @@ -19,7 +19,7 @@ import pbspy._protocol as proto from pbspy import Job, JobResult -from pbspy_server.server import run_server +from pbspy.server import run_server # --------------------------------------------------------------------------- # Fixture: running server @@ -61,21 +61,21 @@ def rpc(self, request: object) -> object: @pytest.fixture() def mock_pbs_submit() -> Generator[MagicMock, None, None]: """Patch pbs_submit to return a fake job without calling qsub.""" - with patch("pbspy_server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: + with patch("pbspy.server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: yield m @pytest.fixture() def mock_check_finished() -> Generator[MagicMock, None, None]: """Patch _check_job_finished to report jobs as never finished (default).""" - with patch("pbspy_server.server._check_job_finished", return_value=None) as m: + with patch("pbspy.server.server._check_job_finished", return_value=None) as m: yield m @pytest.fixture() def mock_pbs_get_result() -> Generator[MagicMock, None, None]: with patch( - "pbspy_server.server.core.pbs_get_result", + "pbspy.server.server.core.pbs_get_result", return_value=JobResult(exit_code=0, output="hello\n"), ) as m: yield m @@ -149,10 +149,10 @@ def test_wait_already_finished_returns_immediately() -> None: def _run() -> None: with ( - patch("pbspy_server.server.core.pbs_submit", return_value=("done.mock", "done_job")), + patch("pbspy.server.server.core.pbs_submit", return_value=("done.mock", "done_job")), # Return 0 immediately so the very first poll marks the job finished. - patch("pbspy_server.server._check_job_finished", return_value=0), - patch("pbspy_server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), + patch("pbspy.server.server._check_job_finished", return_value=0), + patch("pbspy.server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), ): run_server(port=0, poll_interval=0.1, _ready_callback=port_q.put, _stop_event=stop) From 1a98d566094e2423ab0829d97a2de2e4ef2b3e94 Mon Sep 17 00:00:00 2001 From: omdowley Date: Thu, 4 Jun 2026 14:56:10 +1000 Subject: [PATCH 05/16] fix: pickled socket error The Backend was being pickled when passing jobs to the server, causing crashes because Sockets aren't pickleable. This resolves this and adds a regression test ensuring we can pickle and unpickle a Job. --- src/pbspy/__init__.py | 11 +++++++++++ tests/test_default.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 0ec8671..4698301 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -66,6 +66,17 @@ class Job: backend: Backend = field(default=_DEFAULT_LOCAL_BACKEND, repr=False, compare=False) """The backend used to submit and track this job.""" + def __getstate__(self) -> dict[str, Any]: + """Exclude backend from pickling (it carries live sockets).""" + state = self.__dict__.copy() + del state["backend"] + return state + + def __setstate__(self, state: dict[str, Any]) -> None: + """Restore job state, assigning the default local backend.""" + self.__dict__.update(state) + self.backend = _DEFAULT_LOCAL_BACKEND + def wait(self, progress: bool = True) -> None: """ Wait for the job to complete. diff --git a/tests/test_default.py b/tests/test_default.py index 7a94c98..51b4dc7 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -1,3 +1,4 @@ +import pickle import shutil from collections.abc import Callable @@ -127,3 +128,36 @@ def test_backend_classes_importable() -> None: """Backend, LocalBackend, ServerBackend are importable from pbspy.""" assert issubclass(LocalBackend, Backend) assert issubclass(ServerBackend, Backend) + + +def test_job_pickle_round_trip() -> None: + """Job pickle round-trip excludes backend and restores all other fields. + + Regression test: Job.backend carries live sockets (ServerBackend) and + must not be pickled. See __getstate__/__setstate__ on Job. + """ + job = Job( + job_id="42.mock", + job_name="my_job", + description="some description", + backend=LocalBackend(), + ) + + # Pickle and unpickle + data = pickle.dumps(job) + restored = pickle.loads(data) + + # All non-backend fields preserved + assert restored.job_id == "42.mock" + assert restored.job_name == "my_job" + assert restored.description == "some description" + # backend is restored to the default local backend (never pickled) + assert isinstance(restored.backend, LocalBackend) + + +def test_job_pickle_excludes_backend_from_state() -> None: + """__getstate__ does not include the backend attribute.""" + job = Job(job_id="1.mock", backend=LocalBackend()) + state = job.__getstate__() + assert "backend" not in state + assert state == {"job_name": None, "job_id": "1.mock", "description": None} From 512706446573672f4d0c9494b74e372efe92b93e Mon Sep 17 00:00:00 2001 From: omdowley Date: Thu, 4 Jun 2026 16:27:19 +1000 Subject: [PATCH 06/16] fix: propagate output_path and error_path Required so that `pbs_get_result` gets output and error files correctly if their paths are modified by the JobDescription. --- src/pbspy/__init__.py | 15 ++++- src/pbspy/_pbs_core.py | 6 +- tests/test_default.py | 124 ++++++++++++++++++++++++++++++++++++++++- 3 files changed, 140 insertions(+), 5 deletions(-) diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 4698301..b1002c5 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -66,6 +66,12 @@ class Job: 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 __getstate__(self) -> dict[str, Any]: """Exclude backend from pickling (it carries live sockets).""" state = self.__dict__.copy() @@ -271,4 +277,11 @@ def submit(self, backend: Backend | None = None) -> Job: if self.name is None: self.name = job_name - return Job(job_name=job_name, job_id=job_id, description=self.description, backend=backend) + 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/_pbs_core.py b/src/pbspy/_pbs_core.py index 51bb858..0b26f2f 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -84,7 +84,7 @@ def pbs_submit(script: str, name: str | None, runner: PBSRunner = _LOCAL_RUNNER) """ 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")}') + 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) @@ -133,8 +133,8 @@ def pbs_get_result(job: Job, runner: PBSRunner = _LOCAL_RUNNER) -> JobResult: from pbspy import JobResult job_id_num = job.job_id.split(".")[0] - output_file = f"{job.job_name}.o{job_id_num}" - error_file = f"{job.job_name}.e{job_id_num}" + 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) diff --git a/tests/test_default.py b/tests/test_default.py index 51b4dc7..9da00ec 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -1,5 +1,6 @@ import pickle import shutil +import subprocess from collections.abc import Callable import pytest @@ -160,4 +161,125 @@ def test_job_pickle_excludes_backend_from_state() -> None: job = Job(job_id="1.mock", backend=LocalBackend()) state = job.__getstate__() assert "backend" not in state - assert state == {"job_name": None, "job_id": "1.mock", "description": None} + assert state == { + "job_name": None, + "job_id": "1.mock", + "description": None, + "output_path": None, + "error_path": None, + } + + +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 == "" + + +def test_job_pickle_preserves_output_and_error_paths() -> None: + """Job pickle round-trip preserves output_path and error_path.""" + job = Job( + job_id="789.mock", + job_name="my_job", + output_path="/scratch/project/job.out", + error_path="/scratch/project/job.err", + backend=LocalBackend(), + ) + + data = pickle.dumps(job) + restored = pickle.loads(data) + + assert restored.output_path == "/scratch/project/job.out" + assert restored.error_path == "/scratch/project/job.err" + assert restored.job_id == "789.mock" + assert restored.job_name == "my_job" From 56f5054ccd4b592f6b681561da24277e45a5c58e Mon Sep 17 00:00:00 2001 From: omdowley Date: Wed, 17 Jun 2026 15:57:05 +1000 Subject: [PATCH 07/16] fix: resolve ssh key/hostfile permission errors --- Dockerfile | 9 ++++++--- entrypoint.sh | 28 ++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+), 3 deletions(-) create mode 100644 entrypoint.sh diff --git a/Dockerfile b/Dockerfile index 530f548..28cfaec 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,7 +1,7 @@ FROM python:3.12-slim RUN apt-get update \ - && apt-get install -y --no-install-recommends openssh-client \ + && apt-get install -y --no-install-recommends openssh-client gosu \ && rm -rf /var/lib/apt/lists/* RUN useradd -m -u 1000 pbspy @@ -10,6 +10,9 @@ WORKDIR /home/pbspy COPY --chown=pbspy:pbspy . . RUN pip install --no-cache-dir . -USER pbspy +COPY entrypoint.sh /entrypoint.sh +RUN chmod +x /entrypoint.sh + EXPOSE 9876 -ENTRYPOINT ["pbspy-server"] +ENTRYPOINT ["/entrypoint.sh"] +CMD ["pbspy-server"] diff --git a/entrypoint.sh b/entrypoint.sh new file mode 100644 index 0000000..02ad13f --- /dev/null +++ b/entrypoint.sh @@ -0,0 +1,28 @@ +#!/bin/sh +set -e + +mkdir -p /home/pbspy/.ssh +chown pbspy:pbspy /home/pbspy/.ssh +chmod 700 /home/pbspy/.ssh + +# Copy known_hosts from the staged host .ssh directory so SSH can verify +# Gadi's host key regardless of the host file's ownership/permissions. +if [ -f /tmp/ssh_host/known_hosts ]; then + cp /tmp/ssh_host/known_hosts /home/pbspy/.ssh/known_hosts + chown pbspy:pbspy /home/pbspy/.ssh/known_hosts + chmod 600 /home/pbspy/.ssh/known_hosts +fi + +# If PBSPY_SSH_KEY names a key file (basename only — tilde already stripped by +# ${HOME} expansion in the mount), copy it and tell pbspy-server to use it. +if [ -n "${PBSPY_SSH_KEY}" ]; then + key_name=$(basename "${PBSPY_SSH_KEY}") + if [ -f "/tmp/ssh_host/${key_name}" ]; then + cp "/tmp/ssh_host/${key_name}" /home/pbspy/.ssh/pbspy_key + chown pbspy:pbspy /home/pbspy/.ssh/pbspy_key + chmod 600 /home/pbspy/.ssh/pbspy_key + set -- "$@" --ssh-arg=-i --ssh-arg=/home/pbspy/.ssh/pbspy_key + fi +fi + +exec gosu pbspy "$@" From 451faeb7f8475bfe543e3757f6d6b75ec4a4425e Mon Sep 17 00:00:00 2001 From: omdowley Date: Mon, 22 Jun 2026 15:19:20 +1000 Subject: [PATCH 08/16] feat: add API key auth and generic SSH exec server Clients connecting to a server started with --api-key must send AuthRequest as the first frame; wrong or missing key closes the connection with ErrorResponse. Auth is opt-in: servers without --api-key accept all connections as before. Adds ExecRequest/ExecResponse to the protocol so clients can run arbitrary commands through the server's SSH runner. Disabled by default (--allow-exec / PBSPY_ALLOW_EXEC=1); a startup warning is logged when enabled. --- docker-compose.yml | 2 + src/pbspy/_protocol.py | 40 ++++++++++- src/pbspy/_server_backend.py | 42 ++++++++++- src/pbspy/server/__main__.py | 15 ++++ src/pbspy/server/server.py | 45 +++++++++++- tests/test_server.py | 132 +++++++++++++++++++++++++++++++++++ 6 files changed, 270 insertions(+), 6 deletions(-) diff --git a/docker-compose.yml b/docker-compose.yml index 36f9995..12d664c 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -8,3 +8,5 @@ services: environment: - PBSPY_SSH_HOST=${SSH_HOST} - PBSPY_SSH_USER=${SSH_USER:-} + - PBSPY_API_KEY=${PBSPY_API_KEY:-} + - PBSPY_ALLOW_EXEC=${PBSPY_ALLOW_EXEC:-} diff --git a/src/pbspy/_protocol.py b/src/pbspy/_protocol.py index e70debd..012009b 100644 --- a/src/pbspy/_protocol.py +++ b/src/pbspy/_protocol.py @@ -2,9 +2,10 @@ Request and response dataclasses for the pbspy client-server protocol. Messages are exchanged as length-prefixed pickle frames (see :func:`send_frame` -and :func:`recv_frame`). Connections require no authentication — the server -is expected to run on a trusted machine with network access restricted to -authorised clients. +and :func:`recv_frame`). Connections optionally require authentication via an +API key (see :class:`AuthRequest`). When enabled, the server sends an +:class:`AuthOkResponse` on success or an :class:`ErrorResponse` on failure +before accepting any other requests. """ from __future__ import annotations @@ -20,17 +21,21 @@ __all__ = [ # Requests + "AuthRequest", "SubmitRequest", "WaitRequest", "ResultRequest", "PingRequest", + "ExecRequest", # Responses + "AuthOkResponse", "SubmittedResponse", "StatusUpdateResponse", "WaitDoneResponse", "ResultResponse", "PongResponse", "ErrorResponse", + "ExecResponse", # Framing "send_frame", "recv_frame", @@ -69,6 +74,13 @@ class ResultRequest: job: Job +@dataclass +class AuthRequest: + """Sent as the first frame on connection when the server requires API key auth.""" + + api_key: str | None = None + + @dataclass class PingRequest: """Liveness check; server responds with :class:`PongResponse`.""" @@ -119,6 +131,19 @@ class PongResponse: """Response to :class:`PingRequest`.""" +@dataclass +class AuthOkResponse: + """Sent after a successful :class:`AuthRequest`.""" + + +@dataclass +class ExecRequest: + """Ask the server to execute an arbitrary command via SSH.""" + + command: list[str] + stdin: bytes | None = None + + @dataclass class ErrorResponse: """Returned when the server encounters an error handling a request.""" @@ -126,6 +151,15 @@ class ErrorResponse: message: str +@dataclass +class ExecResponse: + """Returned after a successful :class:`ExecRequest`.""" + + returncode: int + stdout: bytes + stderr: bytes + + # --------------------------------------------------------------------------- # Framing helpers # --------------------------------------------------------------------------- diff --git a/src/pbspy/_server_backend.py b/src/pbspy/_server_backend.py index 6006d3c..f375141 100644 --- a/src/pbspy/_server_backend.py +++ b/src/pbspy/_server_backend.py @@ -32,6 +32,7 @@ class ServerBackend(Backend): host: Hostname or IP address of the machine running pbspy-server. port: TCP port the server is listening on (default: 9876). connect_timeout: Seconds to wait for the initial TCP connection. + api_key: API key for authentication (required if the server requires it). """ def __init__( @@ -39,10 +40,12 @@ def __init__( host: str, port: int = 9876, connect_timeout: float = 10.0, + api_key: str | None = None, ) -> None: self._host = host self._port = port self._connect_timeout = connect_timeout + self._api_key = api_key self._sock: socket.socket | None = None self._stream: BinaryIO | None = None self._lock = threading.Lock() @@ -75,8 +78,31 @@ def get_result(self, job: object) -> object: raise RuntimeError(f"Unexpected response: {response!r}") return response.result + def exec(self, command: list[str], stdin: bytes | None = None) -> tuple[int, bytes, bytes]: + """ + Execute an arbitrary command on the server via SSH. + + Requires the server to be started with ``--allow-exec``. + + Args: + command: Command and arguments to execute. + stdin: Optional data to pass as standard input. + + Returns: + A tuple of ``(returncode, stdout, stderr)``. + + Raises: + RuntimeError: If the server rejects the request. + """ + response = self._rpc(proto.ExecRequest(command=command, stdin=stdin)) + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.ExecResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.returncode, response.stdout, response.stderr + def _connect(self) -> None: - """Open a TCP connection to the server and verify liveness with a ping.""" + """Open a TCP connection to the server, authenticate if required, and verify liveness.""" sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(self._connect_timeout) try: @@ -89,6 +115,20 @@ def _connect(self) -> None: self._stream = cast(BinaryIO, sock.makefile("rwb", buffering=0)) try: + # Auth step first — server expects AuthRequest before any other frame when + # an API key is configured. + if self._api_key is not None: + proto.send_frame(self._stream, proto.AuthRequest(api_key=self._api_key)) + auth_resp = proto.recv_frame(self._stream) + if isinstance(auth_resp, proto.ErrorResponse): + self._stream.close() + sock.close() + raise RuntimeError(f"pbspy-server auth failed: {auth_resp.message}") + if not isinstance(auth_resp, proto.AuthOkResponse): + self._stream.close() + sock.close() + raise RuntimeError(f"pbspy-server returned unexpected auth response: {auth_resp!r}") + proto.send_frame(self._stream, proto.PingRequest()) pong = proto.recv_frame(self._stream) except (OSError, EOFError) as exc: diff --git a/src/pbspy/server/__main__.py b/src/pbspy/server/__main__.py index 18f51d2..05a1d04 100644 --- a/src/pbspy/server/__main__.py +++ b/src/pbspy/server/__main__.py @@ -4,6 +4,7 @@ Usage:: pbspy-server [--ssh-host HOST] [--ssh-user USER] [--port PORT] [--ssh-arg ARG ...] + [--api-key KEY] [--allow-exec] """ from __future__ import annotations @@ -44,6 +45,18 @@ def main() -> None: metavar="ARG", help="Extra SSH argument (may be repeated, e.g. --ssh-arg=-i --ssh-arg=/path/to/key).", ) + parser.add_argument( + "--api-key", + metavar="KEY", + default=os.environ.get("PBSPY_API_KEY"), + help="Require clients to authenticate with this key (env: PBSPY_API_KEY).", + ) + parser.add_argument( + "--allow-exec", + action="store_true", + default=os.environ.get("PBSPY_ALLOW_EXEC", "").lower() in ("1", "true", "yes"), + help="Allow clients to execute arbitrary SSH commands (env: PBSPY_ALLOW_EXEC; disabled by default).", + ) args = parser.parse_args() @@ -56,6 +69,8 @@ def main() -> None: ssh_host=args.ssh_host, ssh_user=args.ssh_user, ssh_args=args.ssh_args, + api_key=args.api_key, + allow_exec=args.allow_exec, ) diff --git a/src/pbspy/server/server.py b/src/pbspy/server/server.py index 82306db..c16f6b3 100644 --- a/src/pbspy/server/server.py +++ b/src/pbspy/server/server.py @@ -56,8 +56,10 @@ def __init__(self, jobs: list[Job], stream: BinaryIO) -> None: class _ServerState: """Shared mutable state accessed by both connection threads and the poll thread.""" - def __init__(self, runner: core.PBSRunner) -> None: + def __init__(self, runner: core.PBSRunner, api_key: str | None = None, allow_exec: bool = False) -> None: self.runner = runner + self.api_key = api_key + self.allow_exec = allow_exec self.lock = threading.Lock() self.jobs: dict[str, _JobRecord] = {} self.subscriptions: dict[str, list[_WaitSubscription]] = {} @@ -128,6 +130,17 @@ def _handle_connection(conn: socket.socket, state: _ServerState) -> None: """Handle a single client connection in its own thread.""" stream = cast(BinaryIO, conn.makefile("rwb", buffering=0)) try: + # Auth pre-check (runs before any other request) + if state.api_key is not None: + try: + first = proto.recv_frame(stream) + except EOFError: + return + if not isinstance(first, proto.AuthRequest) or first.api_key != state.api_key: + proto.send_frame(stream, proto.ErrorResponse(message="Authentication failed")) + return + proto.send_frame(stream, proto.AuthOkResponse()) + while True: try: request = proto.recv_frame(stream) @@ -170,6 +183,25 @@ def _handle_connection(conn: socket.socket, state: _ServerState) -> None: result: JobResult = core.pbs_get_result(request.job, runner=state.runner) proto.send_frame(stream, proto.ResultResponse(job=request.job, result=result)) + elif isinstance(request, proto.AuthRequest): + proto.send_frame(stream, proto.AuthOkResponse()) + + elif isinstance(request, proto.ExecRequest): + if not state.allow_exec: + proto.send_frame( + stream, proto.ErrorResponse(message="SSH command execution is disabled on this server") + ) + else: + exec_result = state.runner.run(request.command, input=request.stdin) + proto.send_frame( + stream, + proto.ExecResponse( + returncode=exec_result.returncode, + stdout=exec_result.stdout, + stderr=exec_result.stderr, + ), + ) + else: proto.send_frame(stream, proto.ErrorResponse(message=f"Unknown request type: {type(request)}")) @@ -191,6 +223,8 @@ def run_server( ssh_user: str | None = None, ssh_args: list[str] | None = None, poll_interval: float = _POLL_INTERVAL, + api_key: str | None = None, + allow_exec: bool = False, _stop_event: threading.Event | None = None, _ready_callback: Callable[[int], None] | None = None, ) -> None: @@ -210,6 +244,10 @@ def run_server( ssh_user: SSH username (combined with *ssh_host* as ``user@host``). ssh_args: Extra arguments forwarded to the ``ssh`` command. poll_interval: Seconds between qstat polls (default 60). + api_key: When set, clients must send a matching :class:`~pbspy._protocol.AuthRequest` + as the first frame or the connection is rejected. + allow_exec: When ``True``, clients may send :class:`~pbspy._protocol.ExecRequest` + to run arbitrary commands via the server's SSH runner. Disabled by default. """ destination: str | None = None ssh_prefix: list[str] | None = None @@ -218,7 +256,7 @@ def run_server( ssh_prefix = ["ssh", *(ssh_args or []), destination] runner = core.PBSRunner(ssh_prefix) - state = _ServerState(runner=runner) + state = _ServerState(runner=runner, api_key=api_key, allow_exec=allow_exec) stop_event = _stop_event or threading.Event() @@ -247,6 +285,9 @@ def _handle_signal(signum: int, frame: object) -> None: if ssh_host: logger.info("PBS commands will run via SSH on %s", destination) + if allow_exec: + logger.warning("SSH command execution is enabled (--allow-exec)") + if _ready_callback is not None: _ready_callback(actual_port) diff --git a/tests/test_server.py b/tests/test_server.py index 7b7e16a..c835752 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -102,6 +102,60 @@ def server( t.join(timeout=2.0) +@pytest.fixture() +def server_with_key( + mock_pbs_submit: MagicMock, + mock_check_finished: MagicMock, + mock_pbs_get_result: MagicMock, +) -> Generator[ServerHandle, None, None]: + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() + + t = threading.Thread( + target=run_server, + kwargs={ + "port": 0, + "poll_interval": 0.1, + "api_key": "secret", + "_ready_callback": port_q.put, + "_stop_event": stop, + }, + daemon=True, + ) + t.start() + port = port_q.get(timeout=_WAIT_TIMEOUT) + yield ServerHandle("127.0.0.1", port) + stop.set() + t.join(timeout=2.0) + + +@pytest.fixture() +def server_with_exec( + mock_pbs_submit: MagicMock, + mock_check_finished: MagicMock, + mock_pbs_get_result: MagicMock, +) -> Generator[ServerHandle, None, None]: + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() + + t = threading.Thread( + target=run_server, + kwargs={ + "port": 0, + "poll_interval": 0.1, + "allow_exec": True, + "_ready_callback": port_q.put, + "_stop_event": stop, + }, + daemon=True, + ) + t.start() + port = port_q.get(timeout=_WAIT_TIMEOUT) + yield ServerHandle("127.0.0.1", port) + stop.set() + t.join(timeout=2.0) + + # --------------------------------------------------------------------------- # Ping / pong # --------------------------------------------------------------------------- @@ -213,3 +267,81 @@ def test_wait_receives_status_update_then_done(server: ServerHandle, mock_check_ stream.close() assert any(isinstance(r, proto.WaitDoneResponse) for r in responses), f"Expected WaitDoneResponse; got: {responses}" + + +# --------------------------------------------------------------------------- +# Auth tests +# --------------------------------------------------------------------------- + + +def test_auth_required_with_api_key(server_with_key: ServerHandle) -> None: + """Connecting without sending AuthRequest first should fail.""" + stream = server_with_key.connect() + # Send a PingRequest directly (no auth) + proto.send_frame(stream, proto.PingRequest()) + response = proto.recv_frame(stream) + stream.close() + assert isinstance(response, proto.ErrorResponse) + assert "Authentication failed" in response.message + + +def test_auth_wrong_key(server_with_key: ServerHandle) -> None: + """Sending the wrong API key should fail.""" + stream = server_with_key.connect() + proto.send_frame(stream, proto.AuthRequest(api_key="wrong")) + response = proto.recv_frame(stream) + stream.close() + assert isinstance(response, proto.ErrorResponse) + assert "Authentication failed" in response.message + + +def test_auth_correct_key(server_with_key: ServerHandle) -> None: + """Correct API key followed by a ping should work.""" + stream = server_with_key.connect() + proto.send_frame(stream, proto.AuthRequest(api_key="secret")) + auth_resp = proto.recv_frame(stream) + assert isinstance(auth_resp, proto.AuthOkResponse) + + proto.send_frame(stream, proto.PingRequest()) + pong = proto.recv_frame(stream) + stream.close() + assert isinstance(pong, proto.PongResponse) + + +def test_server_backend_auth_end_to_end(server_with_key: ServerHandle) -> None: + """ServerBackend with the correct api_key should connect and submit successfully.""" + from pbspy._server_backend import ServerBackend + + backend = ServerBackend(server_with_key.host, server_with_key.port, api_key="secret") + job_id, job_name = backend.submit("#!/bin/bash\necho hi", name="test_job") + backend.close() + assert job_id == "100.mock" + assert job_name == "test_job" + + +def test_server_backend_auth_wrong_key(server_with_key: ServerHandle) -> None: + """ServerBackend with a wrong api_key should raise RuntimeError on connect.""" + from pbspy._server_backend import ServerBackend + + with pytest.raises(RuntimeError, match="auth failed"): + ServerBackend(server_with_key.host, server_with_key.port, api_key="wrong") + + +# --------------------------------------------------------------------------- +# Exec tests +# --------------------------------------------------------------------------- + + +def test_exec_disabled_by_default(server: ServerHandle) -> None: + """ExecRequest on a server without --allow-exec should return ErrorResponse.""" + response = server.rpc(proto.ExecRequest(command=["echo", "hi"])) + assert isinstance(response, proto.ErrorResponse) + assert "disabled" in response.message + + +def test_exec_runs_command(server_with_exec: ServerHandle) -> None: + """ExecRequest with allow_exec=True should return ExecResponse.""" + response = server_with_exec.rpc(proto.ExecRequest(command=["echo", "hi"])) + assert isinstance(response, proto.ExecResponse) + assert response.returncode == 0 + assert response.stdout == b"hi\n" From 21455cb01c5b60c22db854cc0ef921d3e4b1e112 Mon Sep 17 00:00:00 2001 From: omdowley Date: Fri, 26 Jun 2026 11:51:51 +1000 Subject: [PATCH 09/16] feat: add ProxyServer mode Introduces `pbspy-server --proxy-to HOST` (ProxyServer), a pure byte-relay that forwards client connections to a pbspy-server running in an Gadi persistent session via `ssh -W`. Existing `ServerBackend` clients connect to the proxy without any changes. Refactors shared RPC logic (auth, ping, submit/wait/result, reconnect) out of `ServerBackend` into a new `StreamBackend` base class, leaving only the transport (socket open/close) in the subclass. --- .env.example | 5 + README.md | 49 +++++-- docker-compose.yml | 5 + src/pbspy/_server_backend.py | 215 ++-------------------------- src/pbspy/_stream_backend.py | 236 +++++++++++++++++++++++++++++++ src/pbspy/server/__init__.py | 3 +- src/pbspy/server/__main__.py | 59 ++++++-- src/pbspy/server/proxy_server.py | 155 ++++++++++++++++++++ tests/test_proxy_server.py | 221 +++++++++++++++++++++++++++++ 9 files changed, 719 insertions(+), 229 deletions(-) create mode 100644 src/pbspy/_stream_backend.py create mode 100644 src/pbspy/server/proxy_server.py create mode 100644 tests/test_proxy_server.py diff --git a/.env.example b/.env.example index 53e4203..44b5e43 100644 --- a/.env.example +++ b/.env.example @@ -1,3 +1,8 @@ PBSPY_SSH_HOST=gadi.nci.org.au # PBSPY_SSH_USER=ab1234 # omit if your local username matches or ~/.ssh/config handles it # PBSPY_PORT=9876 # override if 9876 is already in use on the host + +# ProxyServer mode (relays to a Server running in a Gadi persistent session) instead of the +# direct-SSH-per-command mode above: comment out PBSPY_SSH_HOST and set these instead. +# PBSPY_PROXY_TO=mysession.ab1234.proj123.ps.gadi.nci.org.au +# PBSPY_REMOTE_PORT=9876 diff --git a/README.md b/README.md index 91dcc8a..65e6fe7 100644 --- a/README.md +++ b/README.md @@ -40,21 +40,17 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -### Submitting from a remote machine via SSH +### Submitting from a remote machine via a `pbspy-server` daemon -Install `pbspy` on your supercomputer via the mechanism of your choice, for instance: - -```bash -uv tool install pbspy -``` - -Then use `SSHBackend` on your local machine — the server daemon starts automatically on first use -and is reused by subsequent connections: +Run `pbspy-server` somewhere with SSH access to the supercomputer (see [docker-compose.yml](./docker-compose.yml) +for a containerised example), then connect to it with `ServerBackend`: ```python -from pbspy import Job, JobDescription, SSHBackend +from pbspy import Job, JobDescription, ServerBackend -backend = SSHBackend("user@gadi.nci.org.au") +backend = ServerBackend( + "pbspy-server.example.com", api_key="..." +) # api_key only if the server requires it job_a = ( JobDescription(name="job_a", ncpus=4, mem="192GB", walltime="00:05:00") @@ -73,7 +69,36 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -The `SSHBackend` uses your existing SSH configuration (keys, `~/.ssh/config` aliases, ssh-agent). +By default `pbspy-server` runs every PBS command (`qsub`, `qstat`, file reads) over a fresh SSH +connection to the supercomputer (`--ssh-host gadi.nci.org.au`). `ServerBackend` doesn't know or +care which of the deployments below is on the other end — the wire protocol is identical either way. + +### Running the server on the supercomputer itself (persistent sessions) + +If your supercomputer supports long-running background processes (e.g. NCI Gadi's +[persistent sessions](https://opus.nci.org.au/display/Help/Persistent+Sessions)), you can instead +run `pbspy-server` directly there — PBS commands then run in-process with no per-command SSH +round trip: + +```bash +persistent-sessions start pbspy +ssh pbspy...ps.gadi.nci.org.au +uv tool install pbspy +pbspy-server --host 127.0.0.1 # no --ssh-host: PBS commands run locally +``` + +Persistent-session hostnames are only resolvable from inside the supercomputer's network, so a +client outside it can't connect directly. Run a `ProxyServer` wherever your existing +`ServerBackend`-based code already points (e.g. in place of the direct-SSH daemon above) — it +relays each connection to the real server via `ssh -W`, so existing client code needs no changes at +all: + +```bash +pbspy-server --proxy-to pbspy...ps.gadi.nci.org.au --remote-port 9876 +``` + +(Add `ProxyJump`/`-J` to your SSH config for that host alias first, since the persistent-session +hostname needs to be reached via a login node.) ## Output (partially executed) diff --git a/docker-compose.yml b/docker-compose.yml index 12d664c..7d77709 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -6,7 +6,12 @@ services: volumes: - "${HOME}/.ssh:/home/pbspy/.ssh:ro" environment: + # Direct-SSH-per-command mode (Server issues qsub/qstat over SSH for every call): - PBSPY_SSH_HOST=${SSH_HOST} - PBSPY_SSH_USER=${SSH_USER:-} - PBSPY_API_KEY=${PBSPY_API_KEY:-} - PBSPY_ALLOW_EXEC=${PBSPY_ALLOW_EXEC:-} + # ProxyServer mode instead (relays to a Server running in a Gadi persistent session): + # leave PBSPY_SSH_HOST unset and set these instead. + - PBSPY_PROXY_TO=${PBSPY_PROXY_TO:-} + - PBSPY_REMOTE_PORT=${PBSPY_REMOTE_PORT:-} diff --git a/src/pbspy/_server_backend.py b/src/pbspy/_server_backend.py index f375141..21c714c 100644 --- a/src/pbspy/_server_backend.py +++ b/src/pbspy/_server_backend.py @@ -8,25 +8,20 @@ from __future__ import annotations import socket -import threading -from collections.abc import Callable -from typing import TYPE_CHECKING, BinaryIO, cast +from typing import BinaryIO, cast -import pbspy._protocol as proto -from pbspy._backend import Backend - -if TYPE_CHECKING: - from pbspy import Job +from pbspy._stream_backend import StreamBackend __all__ = ["ServerBackend"] -class ServerBackend(Backend): +class ServerBackend(StreamBackend): """ Backend that connects to a pbspy-server daemon over TCP. - The server (started with ``pbspy-server``) runs on any machine with SSH - access to the supercomputer; this client connects to it directly. + The server (started with ``pbspy-server``) runs on any machine with SSH access to the + supercomputer; this client connects to it directly. It also works unmodified against a + ProxyServer (``pbspy-server --proxy-to ...``), since the wire protocol is identical. Args: host: Hostname or IP address of the machine running pbspy-server. @@ -45,64 +40,10 @@ def __init__( self._host = host self._port = port self._connect_timeout = connect_timeout - self._api_key = api_key self._sock: socket.socket | None = None - self._stream: BinaryIO | None = None - self._lock = threading.Lock() - self._connect() - - def submit(self, script: str, name: str | None = None) -> tuple[str, str]: - response = self._rpc(proto.SubmitRequest(script=script, name=name)) - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.SubmittedResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - return response.job.job_id, response.job.job_name or "" - - def wait( - self, - jobs: list[Job], - on_update: Callable[[str, str | None], None] | None = None, - progress: bool = True, - ) -> None: - if progress: - self._wait_with_progress(jobs, on_update) - else: - self._wait_quiet(jobs, on_update) - - def get_result(self, job: object) -> object: - response = self._rpc(proto.ResultRequest(job=job)) # type: ignore[arg-type] - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.ResultResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - return response.result + super().__init__(api_key=api_key) - def exec(self, command: list[str], stdin: bytes | None = None) -> tuple[int, bytes, bytes]: - """ - Execute an arbitrary command on the server via SSH. - - Requires the server to be started with ``--allow-exec``. - - Args: - command: Command and arguments to execute. - stdin: Optional data to pass as standard input. - - Returns: - A tuple of ``(returncode, stdout, stderr)``. - - Raises: - RuntimeError: If the server rejects the request. - """ - response = self._rpc(proto.ExecRequest(command=command, stdin=stdin)) - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.ExecResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - return response.returncode, response.stdout, response.stderr - - def _connect(self) -> None: - """Open a TCP connection to the server, authenticate if required, and verify liveness.""" + def _open_stream(self) -> BinaryIO: sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) sock.settimeout(self._connect_timeout) try: @@ -112,148 +53,12 @@ def _connect(self) -> None: raise RuntimeError(f"Could not connect to pbspy-server at {self._host}:{self._port}: {exc}") from exc sock.settimeout(None) self._sock = sock - self._stream = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - - try: - # Auth step first — server expects AuthRequest before any other frame when - # an API key is configured. - if self._api_key is not None: - proto.send_frame(self._stream, proto.AuthRequest(api_key=self._api_key)) - auth_resp = proto.recv_frame(self._stream) - if isinstance(auth_resp, proto.ErrorResponse): - self._stream.close() - sock.close() - raise RuntimeError(f"pbspy-server auth failed: {auth_resp.message}") - if not isinstance(auth_resp, proto.AuthOkResponse): - self._stream.close() - sock.close() - raise RuntimeError(f"pbspy-server returned unexpected auth response: {auth_resp!r}") - - proto.send_frame(self._stream, proto.PingRequest()) - pong = proto.recv_frame(self._stream) - except (OSError, EOFError) as exc: - self._stream.close() - sock.close() - raise RuntimeError(f"pbspy-server at {self._host}:{self._port} did not respond to ping: {exc}") from exc - if not isinstance(pong, proto.PongResponse): - self._stream.close() - sock.close() - raise RuntimeError(f"pbspy-server returned unexpected response to ping: {pong!r}") + return cast(BinaryIO, sock.makefile("rwb", buffering=0)) - def _reconnect(self) -> None: - """Close the current connection and establish a fresh one.""" - if self._stream is not None: - try: - self._stream.close() - except OSError: - pass - self._stream = None + def _close_stream(self) -> None: if self._sock is not None: try: self._sock.close() except OSError: pass self._sock = None - self._connect() - - def _rpc(self, request: object) -> object: - """Send one request frame and return the first response frame.""" - with self._lock: - try: - assert self._stream is not None - proto.send_frame(self._stream, request) - return proto.recv_frame(self._stream) - except (OSError, EOFError): - self._reconnect() - assert self._stream is not None - proto.send_frame(self._stream, request) - return proto.recv_frame(self._stream) - - def _send(self, request: object) -> None: - with self._lock: - assert self._stream is not None - proto.send_frame(self._stream, request) - - def _recv(self) -> object: - assert self._stream is not None - return proto.recv_frame(self._stream) - - def _wait_quiet( - self, - jobs: list[Job], - on_update: Callable[[str, str | None], None] | None, - ) -> None: - self._send(proto.WaitRequest(jobs=jobs)) - while True: - response = self._recv() - if isinstance(response, proto.StatusUpdateResponse): - if on_update: - on_update(response.job.job_id, response.state) - elif isinstance(response, proto.WaitDoneResponse): - return - elif isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error while waiting: {response.message}") - - def _wait_with_progress( - self, - jobs: list[Job], - on_update: Callable[[str, str | None], None] | None, - ) -> None: - from rich.progress import Progress, TextColumn, TimeElapsedColumn - - self._send(proto.WaitRequest(jobs=jobs)) - - with Progress( - TimeElapsedColumn(), - TextColumn("[progress.description]{task.description}"), - auto_refresh=False, - ) as progress_bar: - tasks = { - job.job_id: progress_bar.add_task( - f"{job.job_id} {job.job_name or ''} {job.description or ''}", - total=1, - ) - for job in jobs - } - - while True: - response = self._recv() - progress_bar.refresh() - - if isinstance(response, proto.StatusUpdateResponse): - job = response.job - if response.state is None: - task = tasks.get(job.job_id) - if task is not None: - progress_bar.advance(task) - progress_bar.update(task, visible=False) - if on_update: - on_update(job.job_id, None) - elif on_update: - on_update(job.job_id, response.state) - - elif isinstance(response, proto.WaitDoneResponse): - return - - elif isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error while waiting: {response.message}") - - def close(self) -> None: - if self._stream is not None: - try: - self._stream.close() - except OSError: - pass - self._stream = None - if self._sock is not None: - try: - self._sock.close() - except OSError: - pass - self._sock = None - - def __del__(self) -> None: - try: - self.close() - except Exception: - pass diff --git a/src/pbspy/_stream_backend.py b/src/pbspy/_stream_backend.py new file mode 100644 index 0000000..0bb7838 --- /dev/null +++ b/src/pbspy/_stream_backend.py @@ -0,0 +1,236 @@ +""" +StreamBackend: shared RPC logic for backends that talk the pbspy wire protocol +over some duplex byte stream. + +:class:`~pbspy._server_backend.ServerBackend` connects over a plain TCP socket. A future +backend that talks directly to a Server over an SSH subprocess's stdin/stdout pipes (the +way the pre-0.0.10 ``SSHBackend`` did) would reuse this same RPC machinery — only how the +stream is opened differs between transports. +""" + +from __future__ import annotations + +import threading +from collections.abc import Callable +from typing import TYPE_CHECKING, BinaryIO + +import pbspy._protocol as proto +from pbspy._backend import Backend + +if TYPE_CHECKING: + from pbspy import Job + +__all__ = ["StreamBackend"] + + +class StreamBackend(Backend): + """ + Backend that exchanges length-prefixed pickle frames over a duplex byte stream. + + Subclasses implement :meth:`_open_stream` (and optionally :meth:`_close_stream`) to + provide the transport; everything else (auth handshake, ping, submit/wait/result RPCs, + reconnect-on-failure) is shared. + """ + + def __init__(self, api_key: str | None = None) -> None: + self._api_key = api_key + self._stream: BinaryIO | None = None + self._lock = threading.Lock() + self._connect() + + def _open_stream(self) -> BinaryIO: + """Open the transport and return a readable+writable stream. Implemented by subclasses.""" + raise NotImplementedError + + def _close_stream(self) -> None: + """Release any transport resources beyond the stream itself. Implemented by subclasses.""" + + def submit(self, script: str, name: str | None = None) -> tuple[str, str]: + response = self._rpc(proto.SubmitRequest(script=script, name=name)) + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.SubmittedResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.job.job_id, response.job.job_name or "" + + def wait( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None = None, + progress: bool = True, + ) -> None: + if progress: + self._wait_with_progress(jobs, on_update) + else: + self._wait_quiet(jobs, on_update) + + def get_result(self, job: object) -> object: + response = self._rpc(proto.ResultRequest(job=job)) # type: ignore[arg-type] + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.ResultResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.result + + def exec(self, command: list[str], stdin: bytes | None = None) -> tuple[int, bytes, bytes]: + """ + Execute an arbitrary command on the server via SSH. + + Requires the server to be started with ``--allow-exec``. + + Args: + command: Command and arguments to execute. + stdin: Optional data to pass as standard input. + + Returns: + A tuple of ``(returncode, stdout, stderr)``. + + Raises: + RuntimeError: If the server rejects the request. + """ + response = self._rpc(proto.ExecRequest(command=command, stdin=stdin)) + if isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error: {response.message}") + if not isinstance(response, proto.ExecResponse): + raise RuntimeError(f"Unexpected response: {response!r}") + return response.returncode, response.stdout, response.stderr + + def _connect(self) -> None: + """Open the transport, authenticate if required, and verify liveness with a ping.""" + stream = self._open_stream() + try: + if self._api_key is not None: + proto.send_frame(stream, proto.AuthRequest(api_key=self._api_key)) + auth_resp = proto.recv_frame(stream) + if isinstance(auth_resp, proto.ErrorResponse): + raise RuntimeError(f"pbspy-server auth failed: {auth_resp.message}") + if not isinstance(auth_resp, proto.AuthOkResponse): + raise RuntimeError(f"pbspy-server returned unexpected auth response: {auth_resp!r}") + + proto.send_frame(stream, proto.PingRequest()) + pong = proto.recv_frame(stream) + if not isinstance(pong, proto.PongResponse): + raise RuntimeError(f"pbspy-server returned unexpected response to ping: {pong!r}") + except (OSError, EOFError) as exc: + try: + stream.close() + except OSError: + pass + self._close_stream() + raise RuntimeError(f"pbspy-server did not respond to ping: {exc}") from exc + except RuntimeError: + try: + stream.close() + except OSError: + pass + self._close_stream() + raise + self._stream = stream + + def _reconnect(self) -> None: + """Close the current connection and establish a fresh one.""" + if self._stream is not None: + try: + self._stream.close() + except OSError: + pass + self._stream = None + self._close_stream() + self._connect() + + def _rpc(self, request: object) -> object: + """Send one request frame and return the first response frame.""" + with self._lock: + try: + assert self._stream is not None + proto.send_frame(self._stream, request) + return proto.recv_frame(self._stream) + except (OSError, EOFError): + self._reconnect() + assert self._stream is not None + proto.send_frame(self._stream, request) + return proto.recv_frame(self._stream) + + def _send(self, request: object) -> None: + with self._lock: + assert self._stream is not None + proto.send_frame(self._stream, request) + + def _recv(self) -> object: + assert self._stream is not None + return proto.recv_frame(self._stream) + + def _wait_quiet( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None, + ) -> None: + self._send(proto.WaitRequest(jobs=jobs)) + while True: + response = self._recv() + if isinstance(response, proto.StatusUpdateResponse): + if on_update: + on_update(response.job.job_id, response.state) + elif isinstance(response, proto.WaitDoneResponse): + return + elif isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error while waiting: {response.message}") + + def _wait_with_progress( + self, + jobs: list[Job], + on_update: Callable[[str, str | None], None] | None, + ) -> None: + from rich.progress import Progress, TextColumn, TimeElapsedColumn + + self._send(proto.WaitRequest(jobs=jobs)) + + with Progress( + TimeElapsedColumn(), + TextColumn("[progress.description]{task.description}"), + auto_refresh=False, + ) as progress_bar: + tasks = { + job.job_id: progress_bar.add_task( + f"{job.job_id} {job.job_name or ''} {job.description or ''}", + total=1, + ) + for job in jobs + } + + while True: + response = self._recv() + progress_bar.refresh() + + if isinstance(response, proto.StatusUpdateResponse): + job = response.job + if response.state is None: + task = tasks.get(job.job_id) + if task is not None: + progress_bar.advance(task) + progress_bar.update(task, visible=False) + if on_update: + on_update(job.job_id, None) + elif on_update: + on_update(job.job_id, response.state) + + elif isinstance(response, proto.WaitDoneResponse): + return + + elif isinstance(response, proto.ErrorResponse): + raise RuntimeError(f"Server error while waiting: {response.message}") + + def close(self) -> None: + if self._stream is not None: + try: + self._stream.close() + except OSError: + pass + self._stream = None + self._close_stream() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass diff --git a/src/pbspy/server/__init__.py b/src/pbspy/server/__init__.py index bc2aee6..eb19c24 100644 --- a/src/pbspy/server/__init__.py +++ b/src/pbspy/server/__init__.py @@ -2,6 +2,7 @@ from __future__ import annotations +from pbspy.server.proxy_server import run_proxy_server from pbspy.server.server import run_server -__all__ = ["run_server"] +__all__ = ["run_server", "run_proxy_server"] diff --git a/src/pbspy/server/__main__.py b/src/pbspy/server/__main__.py index 05a1d04..26a2097 100644 --- a/src/pbspy/server/__main__.py +++ b/src/pbspy/server/__main__.py @@ -3,8 +3,11 @@ Usage:: - pbspy-server [--ssh-host HOST] [--ssh-user USER] [--port PORT] [--ssh-arg ARG ...] - [--api-key KEY] [--allow-exec] + pbspy-server [--host HOST] [--port PORT] [--ssh-host HOST] [--ssh-user USER] + [--ssh-arg ARG ...] [--api-key KEY] [--allow-exec] + + pbspy-server --proxy-to HOST [--remote-port PORT] [--host HOST] [--port PORT] + [--ssh-arg ARG ...] """ from __future__ import annotations @@ -31,6 +34,12 @@ def main() -> None: default=os.environ.get("PBSPY_SSH_USER"), help="SSH username on the supercomputer.", ) + parser.add_argument( + "--host", + default="0.0.0.0", + metavar="HOST", + help="Local address to bind (default: 0.0.0.0).", + ) parser.add_argument( "--port", type=int, @@ -57,21 +66,49 @@ def main() -> None: default=os.environ.get("PBSPY_ALLOW_EXEC", "").lower() in ("1", "true", "yes"), help="Allow clients to execute arbitrary SSH commands (env: PBSPY_ALLOW_EXEC; disabled by default).", ) + parser.add_argument( + "--proxy-to", + metavar="HOST", + default=os.environ.get("PBSPY_PROXY_TO"), + help=( + "Run as a ProxyServer instead of a Server: relay each client connection to a remote " + "Server via `ssh -W localhost: HOST` (env: PBSPY_PROXY_TO)." + ), + ) + parser.add_argument( + "--remote-port", + type=int, + default=int(os.environ.get("PBSPY_REMOTE_PORT", "9876")), + metavar="PORT", + help="Port the remote Server is listening on, used with --proxy-to (default: 9876).", + ) args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - from pbspy.server import run_server + if args.proxy_to is not None: + from pbspy.server import run_proxy_server - run_server( - port=args.port, - ssh_host=args.ssh_host, - ssh_user=args.ssh_user, - ssh_args=args.ssh_args, - api_key=args.api_key, - allow_exec=args.allow_exec, - ) + run_proxy_server( + ssh_host=args.proxy_to, + host=args.host, + port=args.port, + remote_port=args.remote_port, + ssh_args=args.ssh_args, + ) + else: + from pbspy.server import run_server + + run_server( + host=args.host, + port=args.port, + ssh_host=args.ssh_host, + ssh_user=args.ssh_user, + ssh_args=args.ssh_args, + api_key=args.api_key, + allow_exec=args.allow_exec, + ) if __name__ == "__main__": diff --git a/src/pbspy/server/proxy_server.py b/src/pbspy/server/proxy_server.py new file mode 100644 index 0000000..a8e00b4 --- /dev/null +++ b/src/pbspy/server/proxy_server.py @@ -0,0 +1,155 @@ +""" +ProxyServer: relays pbspy client connections to a Server running elsewhere via SSH. + +For each accepted client connection, opens an ``ssh -W localhost: `` +subprocess and splices raw bytes bidirectionally between the client socket and the +subprocess's stdin/stdout. The pbspy wire protocol (see :mod:`pbspy._protocol`) is never +parsed here -- frames pass through untouched and are handled by the real Server, so any +client that works against a Server works identically against a ProxyServer. +""" + +from __future__ import annotations + +import logging +import os +import select +import signal +import socket +import subprocess +import threading +from collections.abc import Callable + +__all__ = ["run_proxy_server"] + +logger = logging.getLogger(__name__) + +_BUF_SIZE = 65536 + + +def _splice(client: socket.socket, proc: subprocess.Popen[bytes]) -> None: + """Bidirectionally copy bytes between *client* and *proc*'s stdin/stdout until either closes.""" + assert proc.stdin is not None + assert proc.stdout is not None + client_fd = client.fileno() + stdin_fd = proc.stdin.fileno() + stdout_fd = proc.stdout.fileno() + try: + while True: + readable, _, exceptional = select.select([client_fd, stdout_fd], [], [client_fd, stdout_fd], 1.0) + if exceptional or proc.poll() is not None: + return + for fd in readable: + if fd == client_fd: + data = client.recv(_BUF_SIZE) + if not data: + return + os.write(stdin_fd, data) + elif fd == stdout_fd: + data = os.read(stdout_fd, _BUF_SIZE) + if not data: + return + client.sendall(data) + except OSError: + return + + +def _handle_connection(conn: socket.socket, ssh_command: list[str]) -> None: + proc = subprocess.Popen( # noqa: S603 + ssh_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) + stderr_chunks: list[bytes] = [] + stderr_pipe = proc.stderr + assert stderr_pipe is not None + + def _drain_stderr() -> None: + for chunk in iter(lambda: stderr_pipe.read(4096), b""): + stderr_chunks.append(chunk) + + stderr_thread = threading.Thread(target=_drain_stderr, daemon=True) + stderr_thread.start() + try: + _splice(conn, proc) + finally: + conn.close() + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=5.0) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait() + stderr_thread.join(timeout=1.0) + stderr = b"".join(stderr_chunks) + if stderr: + logger.warning("ssh subprocess stderr: %s", stderr.decode(errors="replace").strip()) + assert proc.stdin is not None + assert proc.stdout is not None + assert proc.stderr is not None + proc.stdin.close() + proc.stdout.close() + proc.stderr.close() + + +def run_proxy_server( + ssh_host: str, + host: str = "0.0.0.0", + port: int = 9876, + remote_port: int = 9876, + ssh_args: list[str] | None = None, + _stop_event: threading.Event | None = None, + _ready_callback: Callable[[int], None] | None = None, + _ssh_command: list[str] | None = None, +) -> None: + """ + Start the pbspy ProxyServer. + + Listens on *host*:*port* and, for each client connection, relays bytes to a Server + reachable at *remote_port* on *ssh_host* via ``ssh -W``. The pbspy wire protocol is never + inspected -- this is a pure byte relay, so clients connecting to a ProxyServer behave + identically to clients connecting directly to a Server. + + Args: + ssh_host: SSH destination of the remote Server (e.g. a persistent-session host + alias). Combine with *ssh_args* (e.g. ``-J``/``ProxyJump``) if it isn't directly + reachable from this machine's default SSH configuration. + host: Local address to bind (default ``"0.0.0.0"``). + port: TCP port to listen on (default 9876; pass 0 for OS-assigned). + remote_port: Port the remote Server is listening on (default 9876). + ssh_args: Extra arguments forwarded to the ``ssh`` command. + """ + ssh_command = _ssh_command or ["ssh", *(ssh_args or []), "-W", f"localhost:{remote_port}", ssh_host] + + stop_event = _stop_event or threading.Event() + + if threading.current_thread() is threading.main_thread(): + + def _handle_signal(signum: int, frame: object) -> None: + logger.info("Received signal %d, shutting down", signum) + stop_event.set() + + signal.signal(signal.SIGTERM, _handle_signal) + signal.signal(signal.SIGINT, _handle_signal) + + srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + srv.bind((host, port)) + srv.listen(16) + srv.settimeout(1.0) + + actual_port = srv.getsockname()[1] + logger.info("pbspy ProxyServer listening on %s:%d, relaying to %s:%d", host, actual_port, ssh_host, remote_port) + + if _ready_callback is not None: + _ready_callback(actual_port) + + try: + while not stop_event.is_set(): + try: + conn, _ = srv.accept() + except TimeoutError: + continue + t = threading.Thread(target=_handle_connection, args=(conn, ssh_command), daemon=True) + t.start() + finally: + srv.close() + logger.info("pbspy ProxyServer stopped") diff --git a/tests/test_proxy_server.py b/tests/test_proxy_server.py new file mode 100644 index 0000000..13e8464 --- /dev/null +++ b/tests/test_proxy_server.py @@ -0,0 +1,221 @@ +""" +Integration tests for pbspy.server.proxy_server. + +ProxyServer never parses the pbspy wire protocol -- it just relays bytes to a real Server +over an SSH-like subprocess. These tests spin up a real `run_server()` (with mocked PBS +operations, as in test_server.py) plus a `run_proxy_server()` pointed at it, substituting a +small Python script for `_ssh_command` so no real `ssh` binary or network access is needed. +The same protocol-level assertions used against a direct Server connection are re-run against +the ProxyServer's port to verify the relay is transparent. +""" + +from __future__ import annotations + +import queue +import socket +import sys +import threading +import time +from collections.abc import Callable, Generator +from typing import BinaryIO, cast +from unittest.mock import MagicMock, patch + +import pytest + +import pbspy._protocol as proto +from pbspy import Job, JobResult +from pbspy._server_backend import ServerBackend +from pbspy.server import run_proxy_server, run_server +from tests.test_server import ServerHandle + +_WAIT_TIMEOUT = 5.0 + + +def _wait_for(condition: Callable[[], bool], timeout: float = _WAIT_TIMEOUT) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if condition(): + return True + time.sleep(0.02) + return False + + +# A stand-in for `ssh -W localhost: `: connects to the given local port and +# splices stdin/stdout with that socket, exactly as `ssh -W` would over an SSH channel. +_FAKE_SSH_SOURCE = """ +import os, select, socket, sys + +port = int(sys.argv[1]) +sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) +sock.connect(("127.0.0.1", port)) +stdin_fd = sys.stdin.buffer.fileno() +stdout_fd = sys.stdout.buffer.fileno() +sock_fd = sock.fileno() + +while True: + readable, _, exceptional = select.select([stdin_fd, sock_fd], [], [stdin_fd, sock_fd]) + if exceptional: + break + done = False + for fd in readable: + if fd == stdin_fd: + data = os.read(stdin_fd, 65536) + if not data: + done = True + break + sock.sendall(data) + elif fd == sock_fd: + data = sock.recv(65536) + if not data: + done = True + break + os.write(stdout_fd, data) + if done: + break +""" + + +def _fake_ssh_command(remote_port: int) -> list[str]: + return [sys.executable, "-c", _FAKE_SSH_SOURCE, str(remote_port)] + + +@pytest.fixture() +def mock_pbs_submit() -> Generator[MagicMock, None, None]: + with patch("pbspy.server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: + yield m + + +@pytest.fixture() +def mock_check_finished() -> Generator[MagicMock, None, None]: + with patch("pbspy.server.server._check_job_finished", return_value=None) as m: + yield m + + +@pytest.fixture() +def mock_pbs_get_result() -> Generator[MagicMock, None, None]: + with patch( + "pbspy.server.server.core.pbs_get_result", + return_value=JobResult(exit_code=0, output="hello\n"), + ) as m: + yield m + + +@pytest.fixture() +def real_server( + mock_pbs_submit: MagicMock, + mock_check_finished: MagicMock, + mock_pbs_get_result: MagicMock, +) -> Generator[ServerHandle, None, None]: + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() + + t = threading.Thread( + target=run_server, + kwargs={"port": 0, "poll_interval": 0.1, "_ready_callback": port_q.put, "_stop_event": stop}, + daemon=True, + ) + t.start() + port = port_q.get(timeout=_WAIT_TIMEOUT) + yield ServerHandle("127.0.0.1", port) + stop.set() + t.join(timeout=2.0) + + +@pytest.fixture() +def proxy_server(real_server: ServerHandle) -> Generator[ServerHandle, None, None]: + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() + + t = threading.Thread( + target=run_proxy_server, + kwargs={ + "ssh_host": "unused", + "port": 0, + "remote_port": real_server.port, + "_ssh_command": _fake_ssh_command(real_server.port), + "_ready_callback": port_q.put, + "_stop_event": stop, + }, + daemon=True, + ) + t.start() + port = port_q.get(timeout=_WAIT_TIMEOUT) + yield ServerHandle("127.0.0.1", port) + stop.set() + t.join(timeout=2.0) + + +def test_ping_through_proxy(proxy_server: ServerHandle) -> None: + response = proxy_server.rpc(proto.PingRequest()) + assert isinstance(response, proto.PongResponse) + + +def test_submit_through_proxy(proxy_server: ServerHandle) -> None: + response = proxy_server.rpc(proto.SubmitRequest(script="#!/bin/bash\necho hi", name="test_job")) + assert isinstance(response, proto.SubmittedResponse) + assert response.job.job_id == "100.mock" + assert response.job.job_name == "test_job" + + +def test_result_through_proxy(proxy_server: ServerHandle) -> None: + job = Job(job_id="100.mock", job_name="test_job") + response = proxy_server.rpc(proto.ResultRequest(job=job)) + assert isinstance(response, proto.ResultResponse) + assert response.result.exit_code == 0 + assert response.result.output == "hello\n" + + +def test_server_backend_works_unmodified_through_proxy(proxy_server: ServerHandle) -> None: + """ServerBackend, with no awareness of ProxyServer, should work exactly as against a Server.""" + backend = ServerBackend(proxy_server.host, proxy_server.port) + job_id, job_name = backend.submit("#!/bin/bash\necho hi", name="test_job") + backend.close() + assert job_id == "100.mock" + assert job_name == "test_job" + + +def test_closing_client_terminates_relay_subprocess(real_server: ServerHandle) -> None: + """When a client disconnects, the spawned ssh-stand-in subprocess should be cleaned up.""" + import subprocess + + spawned: list[subprocess.Popen[bytes]] = [] + real_popen = subprocess.Popen + + def _tracking_popen(*args: object, **kwargs: object) -> subprocess.Popen[bytes]: + proc: subprocess.Popen[bytes] = real_popen(*args, **kwargs) # type: ignore[call-overload] + spawned.append(proc) + return proc + + port_q: queue.SimpleQueue[int] = queue.SimpleQueue() + stop = threading.Event() + + with patch("pbspy.server.proxy_server.subprocess.Popen", side_effect=_tracking_popen): + t = threading.Thread( + target=run_proxy_server, + kwargs={ + "ssh_host": "unused", + "port": 0, + "remote_port": real_server.port, + "_ssh_command": _fake_ssh_command(real_server.port), + "_ready_callback": port_q.put, + "_stop_event": stop, + }, + daemon=True, + ) + t.start() + port = port_q.get(timeout=_WAIT_TIMEOUT) + + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect(("127.0.0.1", port)) + stream = cast("BinaryIO", sock.makefile("rwb", buffering=0)) + sock.close() + proto.send_frame(stream, proto.PingRequest()) + assert _wait_for(lambda: len(spawned) > 0), "subprocess not spawned in time" + stream.close() + assert _wait_for(lambda: spawned[0].poll() is not None), "subprocess not cleaned up after client disconnect" + + stop.set() + t.join(timeout=2.0) + + assert spawned, "expected a subprocess to have been spawned" + assert spawned[0].poll() is not None, "subprocess should have been terminated after client disconnect" From 4fd4b80a865e5019d38798398df04be77d991315 Mon Sep 17 00:00:00 2001 From: omdowley Date: Thu, 2 Jul 2026 14:03:18 +1000 Subject: [PATCH 10/16] feat: simplify server This commit entirely removes the SSH layer from pbspy -- setting up the transport layer is a user's concern and not this library's. The server is simplified with the idea that it should run on a node with direct access to PBS (e.g., an NCI Gadi persistent-session). --- .env.example | 10 +- Dockerfile | 18 --- README.md | 48 ++----- docker-compose.yml | 17 --- entrypoint.sh | 28 ---- src/pbspy/__init__.py | 12 +- src/pbspy/_backend.py | 10 ++ src/pbspy/_local_backend.py | 4 + src/pbspy/_pbs_core.py | 72 ++++++++-- src/pbspy/_protocol.py | 27 ++-- src/pbspy/_server_backend.py | 10 +- src/pbspy/_stream_backend.py | 17 +-- src/pbspy/server/__init__.py | 3 +- src/pbspy/server/__main__.py | 78 ++--------- src/pbspy/server/proxy_server.py | 155 ---------------------- src/pbspy/server/server.py | 71 +++------- tests/test_default.py | 113 ++++++++++++++++ tests/test_protocol.py | 12 ++ tests/test_proxy_server.py | 221 ------------------------------- tests/test_server.py | 87 ++++++------ 20 files changed, 310 insertions(+), 703 deletions(-) delete mode 100644 Dockerfile delete mode 100644 docker-compose.yml delete mode 100644 entrypoint.sh delete mode 100644 src/pbspy/server/proxy_server.py delete mode 100644 tests/test_proxy_server.py diff --git a/.env.example b/.env.example index 44b5e43..bba8e26 100644 --- a/.env.example +++ b/.env.example @@ -1,8 +1,2 @@ -PBSPY_SSH_HOST=gadi.nci.org.au -# PBSPY_SSH_USER=ab1234 # omit if your local username matches or ~/.ssh/config handles it -# PBSPY_PORT=9876 # override if 9876 is already in use on the host - -# ProxyServer mode (relays to a Server running in a Gadi persistent session) instead of the -# direct-SSH-per-command mode above: comment out PBSPY_SSH_HOST and set these instead. -# PBSPY_PROXY_TO=mysession.ab1234.proj123.ps.gadi.nci.org.au -# PBSPY_REMOTE_PORT=9876 +PBSPY_API_KEY= # optional: require clients to authenticate with this key +# PBSPY_PORT=9876 # override if 9876 is already in use on the host diff --git a/Dockerfile b/Dockerfile deleted file mode 100644 index 28cfaec..0000000 --- a/Dockerfile +++ /dev/null @@ -1,18 +0,0 @@ -FROM python:3.12-slim - -RUN apt-get update \ - && apt-get install -y --no-install-recommends openssh-client gosu \ - && rm -rf /var/lib/apt/lists/* - -RUN useradd -m -u 1000 pbspy -WORKDIR /home/pbspy - -COPY --chown=pbspy:pbspy . . -RUN pip install --no-cache-dir . - -COPY entrypoint.sh /entrypoint.sh -RUN chmod +x /entrypoint.sh - -EXPOSE 9876 -ENTRYPOINT ["/entrypoint.sh"] -CMD ["pbspy-server"] diff --git a/README.md b/README.md index 65e6fe7..d6a7c97 100644 --- a/README.md +++ b/README.md @@ -42,14 +42,25 @@ print("job_b:", result_b.output.strip()) ### Submitting from a remote machine via a `pbspy-server` daemon -Run `pbspy-server` somewhere with SSH access to the supercomputer (see [docker-compose.yml](./docker-compose.yml) -for a containerised example), then connect to it with `ServerBackend`: +`pbspy-server` runs every PBS command (`qsub`, `qstat`, `qdel`, file reads) locally, +so it must run somewhere with direct PBS access. On NCI Gadi that means a +[persistent session](https://opus.nci.org.au/display/Help/Persistent+Sessions): + +```bash +persistent-sessions start pbspy +ssh pbspy...ps.gadi.nci.org.au +uv tool install pbspy +pbspy-server --host 0.0.0.0 --api-key ... +``` + +Then connect to it with `ServerBackend` (from inside the supercomputer's network, or via an SSH +tunnel/port-forward if you need to reach it from elsewhere): ```python from pbspy import Job, JobDescription, ServerBackend backend = ServerBackend( - "pbspy-server.example.com", api_key="..." + "pbspy...ps.gadi.nci.org.au", api_key="..." ) # api_key only if the server requires it job_a = ( @@ -69,37 +80,6 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -By default `pbspy-server` runs every PBS command (`qsub`, `qstat`, file reads) over a fresh SSH -connection to the supercomputer (`--ssh-host gadi.nci.org.au`). `ServerBackend` doesn't know or -care which of the deployments below is on the other end — the wire protocol is identical either way. - -### Running the server on the supercomputer itself (persistent sessions) - -If your supercomputer supports long-running background processes (e.g. NCI Gadi's -[persistent sessions](https://opus.nci.org.au/display/Help/Persistent+Sessions)), you can instead -run `pbspy-server` directly there — PBS commands then run in-process with no per-command SSH -round trip: - -```bash -persistent-sessions start pbspy -ssh pbspy...ps.gadi.nci.org.au -uv tool install pbspy -pbspy-server --host 127.0.0.1 # no --ssh-host: PBS commands run locally -``` - -Persistent-session hostnames are only resolvable from inside the supercomputer's network, so a -client outside it can't connect directly. Run a `ProxyServer` wherever your existing -`ServerBackend`-based code already points (e.g. in place of the direct-SSH daemon above) — it -relays each connection to the real server via `ssh -W`, so existing client code needs no changes at -all: - -```bash -pbspy-server --proxy-to pbspy...ps.gadi.nci.org.au --remote-port 9876 -``` - -(Add `ProxyJump`/`-J` to your SSH config for that host alias first, since the persistent-session -hostname needs to be reached via a login node.) - ## Output (partially executed) ```text diff --git a/docker-compose.yml b/docker-compose.yml deleted file mode 100644 index 7d77709..0000000 --- a/docker-compose.yml +++ /dev/null @@ -1,17 +0,0 @@ -services: - pbspy-server: - build: . - ports: - - "${PBSPY_PORT:-9876}:9876" - volumes: - - "${HOME}/.ssh:/home/pbspy/.ssh:ro" - environment: - # Direct-SSH-per-command mode (Server issues qsub/qstat over SSH for every call): - - PBSPY_SSH_HOST=${SSH_HOST} - - PBSPY_SSH_USER=${SSH_USER:-} - - PBSPY_API_KEY=${PBSPY_API_KEY:-} - - PBSPY_ALLOW_EXEC=${PBSPY_ALLOW_EXEC:-} - # ProxyServer mode instead (relays to a Server running in a Gadi persistent session): - # leave PBSPY_SSH_HOST unset and set these instead. - - PBSPY_PROXY_TO=${PBSPY_PROXY_TO:-} - - PBSPY_REMOTE_PORT=${PBSPY_REMOTE_PORT:-} diff --git a/entrypoint.sh b/entrypoint.sh deleted file mode 100644 index 02ad13f..0000000 --- a/entrypoint.sh +++ /dev/null @@ -1,28 +0,0 @@ -#!/bin/sh -set -e - -mkdir -p /home/pbspy/.ssh -chown pbspy:pbspy /home/pbspy/.ssh -chmod 700 /home/pbspy/.ssh - -# Copy known_hosts from the staged host .ssh directory so SSH can verify -# Gadi's host key regardless of the host file's ownership/permissions. -if [ -f /tmp/ssh_host/known_hosts ]; then - cp /tmp/ssh_host/known_hosts /home/pbspy/.ssh/known_hosts - chown pbspy:pbspy /home/pbspy/.ssh/known_hosts - chmod 600 /home/pbspy/.ssh/known_hosts -fi - -# If PBSPY_SSH_KEY names a key file (basename only — tilde already stripped by -# ${HOME} expansion in the mount), copy it and tell pbspy-server to use it. -if [ -n "${PBSPY_SSH_KEY}" ]; then - key_name=$(basename "${PBSPY_SSH_KEY}") - if [ -f "/tmp/ssh_host/${key_name}" ]; then - cp "/tmp/ssh_host/${key_name}" /home/pbspy/.ssh/pbspy_key - chown pbspy:pbspy /home/pbspy/.ssh/pbspy_key - chmod 600 /home/pbspy/.ssh/pbspy_key - set -- "$@" --ssh-arg=-i --ssh-arg=/home/pbspy/.ssh/pbspy_key - fi -fi - -exec gosu pbspy "$@" diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 13250cc..901745e 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -83,21 +83,25 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.__dict__.update(state) self.backend = _DEFAULT_LOCAL_BACKEND - def wait(self, **kwargs: dict[str, Any]) -> None: + def wait(self, **kwargs: Any) -> None: """ Wait for the job to complete. """ self.backend.wait([self]) - def result(self, **kwargs: dict[str, Any]) -> JobResult: + def result(self, **kwargs: Any) -> JobResult: """ Waits for the job to complete and returns the result. """ self.wait() return self.backend.get_result(self) # type: ignore[return-value] + def cancel(self) -> None: + """Cancel (``qdel``) this job.""" + self.backend.delete([self.job_id]) + @staticmethod - def wait_all(jobs: list[Job], **kwargs: dict[str, Any]) -> None: + def wait_all(jobs: list[Job], **kwargs: Any) -> None: """ Waits for multiple jobs to complete. """ @@ -107,7 +111,7 @@ def wait_all(jobs: list[Job], **kwargs: dict[str, Any]) -> None: _wait_all_grouped(jobs) @staticmethod - def result_all(jobs: list[Job], **kwargs: dict[str, Any]) -> list[JobResult]: + def result_all(jobs: list[Job], **kwargs: Any) -> list[JobResult]: """ Waits for multiple jobs to complete and returns their results. """ diff --git a/src/pbspy/_backend.py b/src/pbspy/_backend.py index 5b484c0..c1e0902 100644 --- a/src/pbspy/_backend.py +++ b/src/pbspy/_backend.py @@ -58,3 +58,13 @@ def get_result(self, job: object) -> object: Return the :class:`~pbspy.JobResult` for a completed job. """ ... + + @abstractmethod + def delete(self, job_ids: list[str]) -> None: + """ + Cancel (``qdel``) the given jobs. + + Args: + job_ids: Job ids to cancel. + """ + ... diff --git a/src/pbspy/_local_backend.py b/src/pbspy/_local_backend.py index 7c007e7..2b30410 100644 --- a/src/pbspy/_local_backend.py +++ b/src/pbspy/_local_backend.py @@ -43,3 +43,7 @@ def wait( def get_result(self, job: object) -> object: # job: Job -> JobResult """Return the :class:`~pbspy.JobResult` for a completed job.""" return core.pbs_get_result(job) # type: ignore[arg-type] + + 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 index 0b26f2f..bb69818 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -17,28 +17,28 @@ if TYPE_CHECKING: from pbspy import Job, JobResult -__all__ = ["PBSRunner", "pbs_submit", "pbs_wait_for_jobs", "pbs_get_result", "try_get_exit_code"] +__all__ = [ + "PBSRunner", + "pbs_submit", + "pbs_wait_for_jobs", + "pbs_get_result", + "pbs_get_states", + "pbs_delete", + "try_get_exit_code", +] _POLL_INTERVAL_SECONDS = 60 _PROGRESS_REFRESH_SECONDS = 1 class PBSRunner: - """ - Abstracts command execution and file reading for PBS operations. - - The default instance runs commands locally. Pass ``ssh_prefix`` to run - commands on a remote host via SSH (e.g. ``["ssh", "user@host"]``). - """ - - def __init__(self, ssh_prefix: list[str] | None = None) -> None: - self._prefix: list[str] = ssh_prefix or [] + """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(self._prefix + cmd, input=input, capture_output=True) + return subprocess.run(cmd, input=input, capture_output=True) def read_file(self, path: str) -> str: - """Read a file (possibly on a remote host via ``ssh … cat``) and return its text.""" + """Read a local file and return its text.""" result = self.run(["cat", path]) if result.returncode != 0: raise FileNotFoundError(path) @@ -124,6 +124,52 @@ def pbs_wait_for_jobs( 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`. @@ -207,6 +253,6 @@ def format_job_script( {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 ""} +{f"#PBS -W depend=afterok:{":".join(afterok_ids)}" if afterok_ids else ""} {commands_str} """ diff --git a/src/pbspy/_protocol.py b/src/pbspy/_protocol.py index 012009b..bb3e9f0 100644 --- a/src/pbspy/_protocol.py +++ b/src/pbspy/_protocol.py @@ -26,7 +26,7 @@ "WaitRequest", "ResultRequest", "PingRequest", - "ExecRequest", + "DeleteRequest", # Responses "AuthOkResponse", "SubmittedResponse", @@ -35,7 +35,7 @@ "ResultResponse", "PongResponse", "ErrorResponse", - "ExecResponse", + "DeleteResponse", # Framing "send_frame", "recv_frame", @@ -86,6 +86,13 @@ class PingRequest: """Liveness check; server responds with :class:`PongResponse`.""" +@dataclass +class DeleteRequest: + """Ask the server to cancel (``qdel``) one or more jobs.""" + + job_ids: list[str] + + # --------------------------------------------------------------------------- # Responses # --------------------------------------------------------------------------- @@ -136,14 +143,6 @@ class AuthOkResponse: """Sent after a successful :class:`AuthRequest`.""" -@dataclass -class ExecRequest: - """Ask the server to execute an arbitrary command via SSH.""" - - command: list[str] - stdin: bytes | None = None - - @dataclass class ErrorResponse: """Returned when the server encounters an error handling a request.""" @@ -152,12 +151,8 @@ class ErrorResponse: @dataclass -class ExecResponse: - """Returned after a successful :class:`ExecRequest`.""" - - returncode: int - stdout: bytes - stderr: bytes +class DeleteResponse: + """Returned after a successful :class:`DeleteRequest`.""" # --------------------------------------------------------------------------- diff --git a/src/pbspy/_server_backend.py b/src/pbspy/_server_backend.py index 21c714c..939d028 100644 --- a/src/pbspy/_server_backend.py +++ b/src/pbspy/_server_backend.py @@ -2,7 +2,9 @@ ServerBackend: communicates with a pbspy-server daemon over a plain TCP connection. Messages are exchanged as length-prefixed pickle frames (see :mod:`pbspy._protocol`). -The server handles all SSH communication to the supercomputer internally. +The server runs PBS commands (qsub, qstat, qdel) locally on the machine it's started on, +so it's typically started on a machine with direct PBS access (e.g. a Gadi persistent +session). """ from __future__ import annotations @@ -19,9 +21,9 @@ class ServerBackend(StreamBackend): """ Backend that connects to a pbspy-server daemon over TCP. - The server (started with ``pbspy-server``) runs on any machine with SSH access to the - supercomputer; this client connects to it directly. It also works unmodified against a - ProxyServer (``pbspy-server --proxy-to ...``), since the wire protocol is identical. + The server (started with ``pbspy-server``) runs locally on a machine with direct PBS + access (e.g. a Gadi persistent session with ``/g/data`` mounted); this client connects + to it directly over TCP. Args: host: Hostname or IP address of the machine running pbspy-server. diff --git a/src/pbspy/_stream_backend.py b/src/pbspy/_stream_backend.py index 9baef02..c8de9de 100644 --- a/src/pbspy/_stream_backend.py +++ b/src/pbspy/_stream_backend.py @@ -68,28 +68,21 @@ def get_result(self, job: object) -> object: raise RuntimeError(f"Unexpected response: {response!r}") return response.result - def exec(self, command: list[str], stdin: bytes | None = None) -> tuple[int, bytes, bytes]: + def delete(self, job_ids: list[str]) -> None: """ - Execute an arbitrary command on the server via SSH. - - Requires the server to be started with ``--allow-exec``. + Cancel (``qdel``) the given jobs on the server. Args: - command: Command and arguments to execute. - stdin: Optional data to pass as standard input. - - Returns: - A tuple of ``(returncode, stdout, stderr)``. + job_ids: Job ids to cancel. Raises: RuntimeError: If the server rejects the request. """ - response = self._rpc(proto.ExecRequest(command=command, stdin=stdin)) + response = self._rpc(proto.DeleteRequest(job_ids=job_ids)) if isinstance(response, proto.ErrorResponse): raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.ExecResponse): + if not isinstance(response, proto.DeleteResponse): raise RuntimeError(f"Unexpected response: {response!r}") - return response.returncode, response.stdout, response.stderr def _connect(self) -> None: """Open the transport, authenticate if required, and verify liveness with a ping.""" diff --git a/src/pbspy/server/__init__.py b/src/pbspy/server/__init__.py index eb19c24..bc2aee6 100644 --- a/src/pbspy/server/__init__.py +++ b/src/pbspy/server/__init__.py @@ -2,7 +2,6 @@ from __future__ import annotations -from pbspy.server.proxy_server import run_proxy_server from pbspy.server.server import run_server -__all__ = ["run_server", "run_proxy_server"] +__all__ = ["run_server"] diff --git a/src/pbspy/server/__main__.py b/src/pbspy/server/__main__.py index 26a2097..29b9353 100644 --- a/src/pbspy/server/__main__.py +++ b/src/pbspy/server/__main__.py @@ -3,11 +3,7 @@ Usage:: - pbspy-server [--host HOST] [--port PORT] [--ssh-host HOST] [--ssh-user USER] - [--ssh-arg ARG ...] [--api-key KEY] [--allow-exec] - - pbspy-server --proxy-to HOST [--remote-port PORT] [--host HOST] [--port PORT] - [--ssh-arg ARG ...] + pbspy-server [--host HOST] [--port PORT] [--api-key KEY] """ from __future__ import annotations @@ -22,23 +18,11 @@ def main() -> None: prog="pbspy-server", description="PBS job server daemon for pbspy.", ) - parser.add_argument( - "--ssh-host", - metavar="HOST", - default=os.environ.get("PBSPY_SSH_HOST"), - help="Supercomputer hostname to SSH into for PBS commands.", - ) - parser.add_argument( - "--ssh-user", - metavar="USER", - default=os.environ.get("PBSPY_SSH_USER"), - help="SSH username on the supercomputer.", - ) parser.add_argument( "--host", - default="0.0.0.0", + default="127.0.0.1", metavar="HOST", - help="Local address to bind (default: 0.0.0.0).", + help="Local address to bind (default: 127.0.0.1).", ) parser.add_argument( "--port", @@ -47,68 +31,24 @@ def main() -> None: metavar="PORT", help="TCP port to listen on (default: 9876).", ) - parser.add_argument( - "--ssh-arg", - dest="ssh_args", - action="append", - metavar="ARG", - help="Extra SSH argument (may be repeated, e.g. --ssh-arg=-i --ssh-arg=/path/to/key).", - ) parser.add_argument( "--api-key", metavar="KEY", default=os.environ.get("PBSPY_API_KEY"), help="Require clients to authenticate with this key (env: PBSPY_API_KEY).", ) - parser.add_argument( - "--allow-exec", - action="store_true", - default=os.environ.get("PBSPY_ALLOW_EXEC", "").lower() in ("1", "true", "yes"), - help="Allow clients to execute arbitrary SSH commands (env: PBSPY_ALLOW_EXEC; disabled by default).", - ) - parser.add_argument( - "--proxy-to", - metavar="HOST", - default=os.environ.get("PBSPY_PROXY_TO"), - help=( - "Run as a ProxyServer instead of a Server: relay each client connection to a remote " - "Server via `ssh -W localhost: HOST` (env: PBSPY_PROXY_TO)." - ), - ) - parser.add_argument( - "--remote-port", - type=int, - default=int(os.environ.get("PBSPY_REMOTE_PORT", "9876")), - metavar="PORT", - help="Port the remote Server is listening on, used with --proxy-to (default: 9876).", - ) args = parser.parse_args() logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - if args.proxy_to is not None: - from pbspy.server import run_proxy_server + from pbspy.server import run_server - run_proxy_server( - ssh_host=args.proxy_to, - host=args.host, - port=args.port, - remote_port=args.remote_port, - ssh_args=args.ssh_args, - ) - else: - from pbspy.server import run_server - - run_server( - host=args.host, - port=args.port, - ssh_host=args.ssh_host, - ssh_user=args.ssh_user, - ssh_args=args.ssh_args, - api_key=args.api_key, - allow_exec=args.allow_exec, - ) + run_server( + host=args.host, + port=args.port, + api_key=args.api_key, + ) if __name__ == "__main__": diff --git a/src/pbspy/server/proxy_server.py b/src/pbspy/server/proxy_server.py deleted file mode 100644 index a8e00b4..0000000 --- a/src/pbspy/server/proxy_server.py +++ /dev/null @@ -1,155 +0,0 @@ -""" -ProxyServer: relays pbspy client connections to a Server running elsewhere via SSH. - -For each accepted client connection, opens an ``ssh -W localhost: `` -subprocess and splices raw bytes bidirectionally between the client socket and the -subprocess's stdin/stdout. The pbspy wire protocol (see :mod:`pbspy._protocol`) is never -parsed here -- frames pass through untouched and are handled by the real Server, so any -client that works against a Server works identically against a ProxyServer. -""" - -from __future__ import annotations - -import logging -import os -import select -import signal -import socket -import subprocess -import threading -from collections.abc import Callable - -__all__ = ["run_proxy_server"] - -logger = logging.getLogger(__name__) - -_BUF_SIZE = 65536 - - -def _splice(client: socket.socket, proc: subprocess.Popen[bytes]) -> None: - """Bidirectionally copy bytes between *client* and *proc*'s stdin/stdout until either closes.""" - assert proc.stdin is not None - assert proc.stdout is not None - client_fd = client.fileno() - stdin_fd = proc.stdin.fileno() - stdout_fd = proc.stdout.fileno() - try: - while True: - readable, _, exceptional = select.select([client_fd, stdout_fd], [], [client_fd, stdout_fd], 1.0) - if exceptional or proc.poll() is not None: - return - for fd in readable: - if fd == client_fd: - data = client.recv(_BUF_SIZE) - if not data: - return - os.write(stdin_fd, data) - elif fd == stdout_fd: - data = os.read(stdout_fd, _BUF_SIZE) - if not data: - return - client.sendall(data) - except OSError: - return - - -def _handle_connection(conn: socket.socket, ssh_command: list[str]) -> None: - proc = subprocess.Popen( # noqa: S603 - ssh_command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE - ) - stderr_chunks: list[bytes] = [] - stderr_pipe = proc.stderr - assert stderr_pipe is not None - - def _drain_stderr() -> None: - for chunk in iter(lambda: stderr_pipe.read(4096), b""): - stderr_chunks.append(chunk) - - stderr_thread = threading.Thread(target=_drain_stderr, daemon=True) - stderr_thread.start() - try: - _splice(conn, proc) - finally: - conn.close() - if proc.poll() is None: - proc.terminate() - try: - proc.wait(timeout=5.0) - except subprocess.TimeoutExpired: - proc.kill() - proc.wait() - stderr_thread.join(timeout=1.0) - stderr = b"".join(stderr_chunks) - if stderr: - logger.warning("ssh subprocess stderr: %s", stderr.decode(errors="replace").strip()) - assert proc.stdin is not None - assert proc.stdout is not None - assert proc.stderr is not None - proc.stdin.close() - proc.stdout.close() - proc.stderr.close() - - -def run_proxy_server( - ssh_host: str, - host: str = "0.0.0.0", - port: int = 9876, - remote_port: int = 9876, - ssh_args: list[str] | None = None, - _stop_event: threading.Event | None = None, - _ready_callback: Callable[[int], None] | None = None, - _ssh_command: list[str] | None = None, -) -> None: - """ - Start the pbspy ProxyServer. - - Listens on *host*:*port* and, for each client connection, relays bytes to a Server - reachable at *remote_port* on *ssh_host* via ``ssh -W``. The pbspy wire protocol is never - inspected -- this is a pure byte relay, so clients connecting to a ProxyServer behave - identically to clients connecting directly to a Server. - - Args: - ssh_host: SSH destination of the remote Server (e.g. a persistent-session host - alias). Combine with *ssh_args* (e.g. ``-J``/``ProxyJump``) if it isn't directly - reachable from this machine's default SSH configuration. - host: Local address to bind (default ``"0.0.0.0"``). - port: TCP port to listen on (default 9876; pass 0 for OS-assigned). - remote_port: Port the remote Server is listening on (default 9876). - ssh_args: Extra arguments forwarded to the ``ssh`` command. - """ - ssh_command = _ssh_command or ["ssh", *(ssh_args or []), "-W", f"localhost:{remote_port}", ssh_host] - - stop_event = _stop_event or threading.Event() - - if threading.current_thread() is threading.main_thread(): - - def _handle_signal(signum: int, frame: object) -> None: - logger.info("Received signal %d, shutting down", signum) - stop_event.set() - - signal.signal(signal.SIGTERM, _handle_signal) - signal.signal(signal.SIGINT, _handle_signal) - - srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - srv.bind((host, port)) - srv.listen(16) - srv.settimeout(1.0) - - actual_port = srv.getsockname()[1] - logger.info("pbspy ProxyServer listening on %s:%d, relaying to %s:%d", host, actual_port, ssh_host, remote_port) - - if _ready_callback is not None: - _ready_callback(actual_port) - - try: - while not stop_event.is_set(): - try: - conn, _ = srv.accept() - except TimeoutError: - continue - t = threading.Thread(target=_handle_connection, args=(conn, ssh_command), daemon=True) - t.start() - finally: - srv.close() - logger.info("pbspy ProxyServer stopped") diff --git a/src/pbspy/server/server.py b/src/pbspy/server/server.py index c16f6b3..ab796e7 100644 --- a/src/pbspy/server/server.py +++ b/src/pbspy/server/server.py @@ -4,8 +4,9 @@ Listens on a configured TCP port, accepts connections from :class:`~pbspy.ServerBackend` clients, and dispatches pickled requests from :mod:`pbspy._protocol`. -PBS operations (qsub, qstat, file reads) are executed via SSH to a remote -supercomputer when ``ssh_host`` is provided, or locally when omitted. +PBS operations (qsub, qstat, qdel, file reads) are executed locally, so this server is +intended to run natively on the machine with PBS access (e.g. a Gadi persistent session +with ``/g/data`` mounted) rather than behind SSH. A background thread polls ``qstat`` every 60 seconds for all unfinished jobs and pushes :class:`StatusUpdateResponse` frames to any clients waiting on @@ -56,10 +57,9 @@ def __init__(self, jobs: list[Job], stream: BinaryIO) -> None: class _ServerState: """Shared mutable state accessed by both connection threads and the poll thread.""" - def __init__(self, runner: core.PBSRunner, api_key: str | None = None, allow_exec: bool = False) -> None: + def __init__(self, runner: core.PBSRunner, api_key: str | None = None) -> None: self.runner = runner self.api_key = api_key - self.allow_exec = allow_exec self.lock = threading.Lock() self.jobs: dict[str, _JobRecord] = {} self.subscriptions: dict[str, list[_WaitSubscription]] = {} @@ -106,9 +106,12 @@ def _poll_loop( try: with state.lock: unfinished = [r for r in state.jobs.values() if not r.finished] + if not unfinished: + continue + states = core.pbs_get_states([r.job_id for r in unfinished], runner=state.runner) for record in unfinished: - job = Job(job_id=record.job_id, job_name=record.job_name) - if _check_job_finished(record.job_id, state.runner) is not None: + if states.get(record.job_id) is None: + job = Job(job_id=record.job_id, job_name=record.job_name) with state.lock: record.finished = True state.notify_finished(record.job_id, job) @@ -116,16 +119,6 @@ def _poll_loop( logger.exception("Error in poll loop") -def _check_job_finished(job_id: str, runner: core.PBSRunner) -> int | None: - """Return exit code if job has finished, else None.""" - process = runner.run(["qstat", job_id]) - if process.returncode != 0: - stdout = process.stdout.decode("utf-8") - if job_id not in stdout or "has finished" in process.stderr.decode("utf-8"): - return core.try_get_exit_code(job_id, runner) or 0 - return None - - def _handle_connection(conn: socket.socket, state: _ServerState) -> None: """Handle a single client connection in its own thread.""" stream = cast(BinaryIO, conn.makefile("rwb", buffering=0)) @@ -186,21 +179,9 @@ def _handle_connection(conn: socket.socket, state: _ServerState) -> None: elif isinstance(request, proto.AuthRequest): proto.send_frame(stream, proto.AuthOkResponse()) - elif isinstance(request, proto.ExecRequest): - if not state.allow_exec: - proto.send_frame( - stream, proto.ErrorResponse(message="SSH command execution is disabled on this server") - ) - else: - exec_result = state.runner.run(request.command, input=request.stdin) - proto.send_frame( - stream, - proto.ExecResponse( - returncode=exec_result.returncode, - stdout=exec_result.stdout, - stderr=exec_result.stderr, - ), - ) + elif isinstance(request, proto.DeleteRequest): + core.pbs_delete(request.job_ids, runner=state.runner) + proto.send_frame(stream, proto.DeleteResponse()) else: proto.send_frame(stream, proto.ErrorResponse(message=f"Unknown request type: {type(request)}")) @@ -219,12 +200,8 @@ def _handle_connection(conn: socket.socket, state: _ServerState) -> None: def run_server( host: str = "0.0.0.0", port: int = 9876, - ssh_host: str | None = None, - ssh_user: str | None = None, - ssh_args: list[str] | None = None, poll_interval: float = _POLL_INTERVAL, api_key: str | None = None, - allow_exec: bool = False, _stop_event: threading.Event | None = None, _ready_callback: Callable[[int], None] | None = None, ) -> None: @@ -234,29 +211,18 @@ def run_server( Runs in the foreground; use your OS's service manager (systemd, nohup, etc.) to run it as a background service. - When *ssh_host* is provided, all PBS operations (qsub, qstat, file reads) - are executed on that host via SSH. Otherwise they run locally. + All PBS operations (qsub, qstat, qdel, file reads) run locally, so this should be started + on a machine with PBS access (e.g. a Gadi persistent session with ``/g/data`` mounted). Args: host: Local address to bind (default ``"0.0.0.0"``). port: TCP port to listen on (default 9876; pass 0 for OS-assigned). - ssh_host: Supercomputer hostname to SSH into for PBS commands. - ssh_user: SSH username (combined with *ssh_host* as ``user@host``). - ssh_args: Extra arguments forwarded to the ``ssh`` command. poll_interval: Seconds between qstat polls (default 60). api_key: When set, clients must send a matching :class:`~pbspy._protocol.AuthRequest` as the first frame or the connection is rejected. - allow_exec: When ``True``, clients may send :class:`~pbspy._protocol.ExecRequest` - to run arbitrary commands via the server's SSH runner. Disabled by default. """ - destination: str | None = None - ssh_prefix: list[str] | None = None - if ssh_host is not None: - destination = f"{ssh_user}@{ssh_host}" if ssh_user else ssh_host - ssh_prefix = ["ssh", *(ssh_args or []), destination] - - runner = core.PBSRunner(ssh_prefix) - state = _ServerState(runner=runner, api_key=api_key, allow_exec=allow_exec) + runner = core.PBSRunner() + state = _ServerState(runner=runner, api_key=api_key) stop_event = _stop_event or threading.Event() @@ -282,11 +248,6 @@ def _handle_signal(signum: int, frame: object) -> None: actual_port = srv.getsockname()[1] logger.info("pbspy-server listening on %s:%d", host, actual_port) - if ssh_host: - logger.info("PBS commands will run via SSH on %s", destination) - - if allow_exec: - logger.warning("SSH command execution is enabled (--allow-exec)") if _ready_callback is not None: _ready_callback(actual_port) diff --git a/tests/test_default.py b/tests/test_default.py index 9da00ec..2b99664 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -37,6 +37,7 @@ 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)) @@ -54,6 +55,9 @@ def get_result(self, job: object) -> object: assert isinstance(job, Job) 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.""" @@ -283,3 +287,112 @@ def test_job_pickle_preserves_output_and_error_paths() -> None: assert restored.error_path == "/scratch/project/job.err" assert restored.job_id == "789.mock" assert restored.job_name == "my_job" + + +# --------------------------------------------------------------------------- +# 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"]] diff --git a/tests/test_protocol.py b/tests/test_protocol.py index 3c80be6..8a7e2ce 100644 --- a/tests/test_protocol.py +++ b/tests/test_protocol.py @@ -64,6 +64,13 @@ def test_roundtrip_result_request() -> None: assert result.job.job_id == "42.gadi" +def test_roundtrip_delete_request() -> None: + req = proto.DeleteRequest(job_ids=["1.gadi", "2.gadi"]) + result = roundtrip(req) + assert isinstance(result, proto.DeleteRequest) + assert result.job_ids == ["1.gadi", "2.gadi"] + + # --------------------------------------------------------------------------- # Response types # --------------------------------------------------------------------------- @@ -123,6 +130,11 @@ def test_roundtrip_error_response() -> None: assert result.message == "something went wrong" +def test_roundtrip_delete_response() -> None: + result = roundtrip(proto.DeleteResponse()) + assert isinstance(result, proto.DeleteResponse) + + # --------------------------------------------------------------------------- # Framing edge cases # --------------------------------------------------------------------------- diff --git a/tests/test_proxy_server.py b/tests/test_proxy_server.py deleted file mode 100644 index 13e8464..0000000 --- a/tests/test_proxy_server.py +++ /dev/null @@ -1,221 +0,0 @@ -""" -Integration tests for pbspy.server.proxy_server. - -ProxyServer never parses the pbspy wire protocol -- it just relays bytes to a real Server -over an SSH-like subprocess. These tests spin up a real `run_server()` (with mocked PBS -operations, as in test_server.py) plus a `run_proxy_server()` pointed at it, substituting a -small Python script for `_ssh_command` so no real `ssh` binary or network access is needed. -The same protocol-level assertions used against a direct Server connection are re-run against -the ProxyServer's port to verify the relay is transparent. -""" - -from __future__ import annotations - -import queue -import socket -import sys -import threading -import time -from collections.abc import Callable, Generator -from typing import BinaryIO, cast -from unittest.mock import MagicMock, patch - -import pytest - -import pbspy._protocol as proto -from pbspy import Job, JobResult -from pbspy._server_backend import ServerBackend -from pbspy.server import run_proxy_server, run_server -from tests.test_server import ServerHandle - -_WAIT_TIMEOUT = 5.0 - - -def _wait_for(condition: Callable[[], bool], timeout: float = _WAIT_TIMEOUT) -> bool: - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - if condition(): - return True - time.sleep(0.02) - return False - - -# A stand-in for `ssh -W localhost: `: connects to the given local port and -# splices stdin/stdout with that socket, exactly as `ssh -W` would over an SSH channel. -_FAKE_SSH_SOURCE = """ -import os, select, socket, sys - -port = int(sys.argv[1]) -sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) -sock.connect(("127.0.0.1", port)) -stdin_fd = sys.stdin.buffer.fileno() -stdout_fd = sys.stdout.buffer.fileno() -sock_fd = sock.fileno() - -while True: - readable, _, exceptional = select.select([stdin_fd, sock_fd], [], [stdin_fd, sock_fd]) - if exceptional: - break - done = False - for fd in readable: - if fd == stdin_fd: - data = os.read(stdin_fd, 65536) - if not data: - done = True - break - sock.sendall(data) - elif fd == sock_fd: - data = sock.recv(65536) - if not data: - done = True - break - os.write(stdout_fd, data) - if done: - break -""" - - -def _fake_ssh_command(remote_port: int) -> list[str]: - return [sys.executable, "-c", _FAKE_SSH_SOURCE, str(remote_port)] - - -@pytest.fixture() -def mock_pbs_submit() -> Generator[MagicMock, None, None]: - with patch("pbspy.server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: - yield m - - -@pytest.fixture() -def mock_check_finished() -> Generator[MagicMock, None, None]: - with patch("pbspy.server.server._check_job_finished", return_value=None) as m: - yield m - - -@pytest.fixture() -def mock_pbs_get_result() -> Generator[MagicMock, None, None]: - with patch( - "pbspy.server.server.core.pbs_get_result", - return_value=JobResult(exit_code=0, output="hello\n"), - ) as m: - yield m - - -@pytest.fixture() -def real_server( - mock_pbs_submit: MagicMock, - mock_check_finished: MagicMock, - mock_pbs_get_result: MagicMock, -) -> Generator[ServerHandle, None, None]: - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - t = threading.Thread( - target=run_server, - kwargs={"port": 0, "poll_interval": 0.1, "_ready_callback": port_q.put, "_stop_event": stop}, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - yield ServerHandle("127.0.0.1", port) - stop.set() - t.join(timeout=2.0) - - -@pytest.fixture() -def proxy_server(real_server: ServerHandle) -> Generator[ServerHandle, None, None]: - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - t = threading.Thread( - target=run_proxy_server, - kwargs={ - "ssh_host": "unused", - "port": 0, - "remote_port": real_server.port, - "_ssh_command": _fake_ssh_command(real_server.port), - "_ready_callback": port_q.put, - "_stop_event": stop, - }, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - yield ServerHandle("127.0.0.1", port) - stop.set() - t.join(timeout=2.0) - - -def test_ping_through_proxy(proxy_server: ServerHandle) -> None: - response = proxy_server.rpc(proto.PingRequest()) - assert isinstance(response, proto.PongResponse) - - -def test_submit_through_proxy(proxy_server: ServerHandle) -> None: - response = proxy_server.rpc(proto.SubmitRequest(script="#!/bin/bash\necho hi", name="test_job")) - assert isinstance(response, proto.SubmittedResponse) - assert response.job.job_id == "100.mock" - assert response.job.job_name == "test_job" - - -def test_result_through_proxy(proxy_server: ServerHandle) -> None: - job = Job(job_id="100.mock", job_name="test_job") - response = proxy_server.rpc(proto.ResultRequest(job=job)) - assert isinstance(response, proto.ResultResponse) - assert response.result.exit_code == 0 - assert response.result.output == "hello\n" - - -def test_server_backend_works_unmodified_through_proxy(proxy_server: ServerHandle) -> None: - """ServerBackend, with no awareness of ProxyServer, should work exactly as against a Server.""" - backend = ServerBackend(proxy_server.host, proxy_server.port) - job_id, job_name = backend.submit("#!/bin/bash\necho hi", name="test_job") - backend.close() - assert job_id == "100.mock" - assert job_name == "test_job" - - -def test_closing_client_terminates_relay_subprocess(real_server: ServerHandle) -> None: - """When a client disconnects, the spawned ssh-stand-in subprocess should be cleaned up.""" - import subprocess - - spawned: list[subprocess.Popen[bytes]] = [] - real_popen = subprocess.Popen - - def _tracking_popen(*args: object, **kwargs: object) -> subprocess.Popen[bytes]: - proc: subprocess.Popen[bytes] = real_popen(*args, **kwargs) # type: ignore[call-overload] - spawned.append(proc) - return proc - - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - with patch("pbspy.server.proxy_server.subprocess.Popen", side_effect=_tracking_popen): - t = threading.Thread( - target=run_proxy_server, - kwargs={ - "ssh_host": "unused", - "port": 0, - "remote_port": real_server.port, - "_ssh_command": _fake_ssh_command(real_server.port), - "_ready_callback": port_q.put, - "_stop_event": stop, - }, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect(("127.0.0.1", port)) - stream = cast("BinaryIO", sock.makefile("rwb", buffering=0)) - sock.close() - proto.send_frame(stream, proto.PingRequest()) - assert _wait_for(lambda: len(spawned) > 0), "subprocess not spawned in time" - stream.close() - assert _wait_for(lambda: spawned[0].poll() is not None), "subprocess not cleaned up after client disconnect" - - stop.set() - t.join(timeout=2.0) - - assert spawned, "expected a subprocess to have been spawned" - assert spawned[0].poll() is not None, "subprocess should have been terminated after client disconnect" diff --git a/tests/test_server.py b/tests/test_server.py index c835752..4915eb5 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -66,9 +66,15 @@ def mock_pbs_submit() -> Generator[MagicMock, None, None]: @pytest.fixture() -def mock_check_finished() -> Generator[MagicMock, None, None]: - """Patch _check_job_finished to report jobs as never finished (default).""" - with patch("pbspy.server.server._check_job_finished", return_value=None) as m: +def mock_get_states() -> Generator[MagicMock, None, None]: + """Patch core.pbs_get_states to report jobs as still running (never finished) by default.""" + state = {"value": "R"} + + def _get_states(job_ids: list[str], runner: object = None) -> dict[str, str | None]: + return dict.fromkeys(job_ids, state["value"]) + + with patch("pbspy.server.server.core.pbs_get_states", side_effect=_get_states) as m: + m.state = state yield m @@ -84,7 +90,7 @@ def mock_pbs_get_result() -> Generator[MagicMock, None, None]: @pytest.fixture() def server( mock_pbs_submit: MagicMock, - mock_check_finished: MagicMock, + mock_get_states: MagicMock, mock_pbs_get_result: MagicMock, ) -> Generator[ServerHandle, None, None]: port_q: queue.SimpleQueue[int] = queue.SimpleQueue() @@ -105,7 +111,7 @@ def server( @pytest.fixture() def server_with_key( mock_pbs_submit: MagicMock, - mock_check_finished: MagicMock, + mock_get_states: MagicMock, mock_pbs_get_result: MagicMock, ) -> Generator[ServerHandle, None, None]: port_q: queue.SimpleQueue[int] = queue.SimpleQueue() @@ -129,33 +135,6 @@ def server_with_key( t.join(timeout=2.0) -@pytest.fixture() -def server_with_exec( - mock_pbs_submit: MagicMock, - mock_check_finished: MagicMock, - mock_pbs_get_result: MagicMock, -) -> Generator[ServerHandle, None, None]: - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - t = threading.Thread( - target=run_server, - kwargs={ - "port": 0, - "poll_interval": 0.1, - "allow_exec": True, - "_ready_callback": port_q.put, - "_stop_event": stop, - }, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - yield ServerHandle("127.0.0.1", port) - stop.set() - t.join(timeout=2.0) - - # --------------------------------------------------------------------------- # Ping / pong # --------------------------------------------------------------------------- @@ -204,8 +183,8 @@ def test_wait_already_finished_returns_immediately() -> None: def _run() -> None: with ( patch("pbspy.server.server.core.pbs_submit", return_value=("done.mock", "done_job")), - # Return 0 immediately so the very first poll marks the job finished. - patch("pbspy.server.server._check_job_finished", return_value=0), + # Return None immediately so the very first poll marks the job finished. + patch("pbspy.server.server.core.pbs_get_states", return_value={"done.mock": None}), patch("pbspy.server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), ): run_server(port=0, poll_interval=0.1, _ready_callback=port_q.put, _stop_event=stop) @@ -239,7 +218,7 @@ def _run() -> None: # --------------------------------------------------------------------------- -def test_wait_receives_status_update_then_done(server: ServerHandle, mock_check_finished: MagicMock) -> None: +def test_wait_receives_status_update_then_done(server: ServerHandle, mock_get_states: MagicMock) -> None: """ Submit a job, start waiting, then let the poll loop (0.1s interval) detect the job as finished and push WaitDoneResponse back to the client. @@ -251,7 +230,7 @@ def test_wait_receives_status_update_then_done(server: ServerHandle, mock_check_ proto.send_frame(stream, proto.WaitRequest(jobs=[job])) # Allow the poll loop to mark the job as finished. - mock_check_finished.return_value = 0 + mock_get_states.state["value"] = None # Collect responses until WaitDoneResponse (or timeout). responses = [] @@ -328,20 +307,34 @@ def test_server_backend_auth_wrong_key(server_with_key: ServerHandle) -> None: # --------------------------------------------------------------------------- -# Exec tests +# Delete # --------------------------------------------------------------------------- -def test_exec_disabled_by_default(server: ServerHandle) -> None: - """ExecRequest on a server without --allow-exec should return ErrorResponse.""" - response = server.rpc(proto.ExecRequest(command=["echo", "hi"])) +def test_delete_returns_delete_response(server: ServerHandle) -> None: + """DeleteRequest should call core.pbs_delete once and return DeleteResponse.""" + with patch("pbspy.server.server.core.pbs_delete") as mock_delete: + response = server.rpc(proto.DeleteRequest(job_ids=["100.mock", "101.mock"])) + assert isinstance(response, proto.DeleteResponse) + mock_delete.assert_called_once() + args, kwargs = mock_delete.call_args + assert args[0] == ["100.mock", "101.mock"] + + +def test_delete_propagates_errors(server: ServerHandle) -> None: + """If core.pbs_delete raises, the server should send back an ErrorResponse.""" + with patch("pbspy.server.server.core.pbs_delete", side_effect=RuntimeError("boom")): + response = server.rpc(proto.DeleteRequest(job_ids=["100.mock"])) assert isinstance(response, proto.ErrorResponse) - assert "disabled" in response.message + assert "boom" in response.message + +def test_server_backend_delete_end_to_end(server: ServerHandle) -> None: + """ServerBackend.delete() sends one DeleteRequest and returns without error.""" + from pbspy._server_backend import ServerBackend -def test_exec_runs_command(server_with_exec: ServerHandle) -> None: - """ExecRequest with allow_exec=True should return ExecResponse.""" - response = server_with_exec.rpc(proto.ExecRequest(command=["echo", "hi"])) - assert isinstance(response, proto.ExecResponse) - assert response.returncode == 0 - assert response.stdout == b"hi\n" + with patch("pbspy.server.server.core.pbs_delete") as mock_delete: + backend = ServerBackend(server.host, server.port) + backend.delete(["100.mock"]) + backend.close() + mock_delete.assert_called_once() From 6bc4f2231c2a41d30f4cb51194bbb3e1535ce7d6 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Wed, 29 Jul 2026 15:24:21 +1000 Subject: [PATCH 11/16] refactor!: remove built-in server stack Retain the backend abstraction and local PBS implementation so alternate backends can be supplied by separate packages. Remove the bundled server, wire protocol, transport code, CLI, and associated documentation and tests. Also tighten backend result typing, share the default local backend, retain legacy Mango progress-call compatibility, and document the net changes since v0.0.9. --- .env.example | 2 - CHANGELOG.md | 23 ++- README.md | 42 +---- examples/run.py | 10 +- pyproject.toml | 3 - src/pbspy/__init__.py | 26 +-- src/pbspy/_backend.py | 24 ++- src/pbspy/_local_backend.py | 10 +- src/pbspy/_pbs_core.py | 15 +- src/pbspy/_protocol.py | 200 --------------------- src/pbspy/_server_backend.py | 66 ------- src/pbspy/_stream_backend.py | 181 ------------------- src/pbspy/server/__init__.py | 7 - src/pbspy/server/__main__.py | 55 ------ src/pbspy/server/server.py | 266 --------------------------- tests/test_default.py | 75 ++------ tests/test_protocol.py | 186 ------------------- tests/test_server.py | 340 ----------------------------------- 18 files changed, 73 insertions(+), 1458 deletions(-) delete mode 100644 .env.example delete mode 100644 src/pbspy/_protocol.py delete mode 100644 src/pbspy/_server_backend.py delete mode 100644 src/pbspy/_stream_backend.py delete mode 100644 src/pbspy/server/__init__.py delete mode 100644 src/pbspy/server/__main__.py delete mode 100644 src/pbspy/server/server.py delete mode 100644 tests/test_protocol.py delete mode 100644 tests/test_server.py diff --git a/.env.example b/.env.example deleted file mode 100644 index bba8e26..0000000 --- a/.env.example +++ /dev/null @@ -1,2 +0,0 @@ -PBSPY_API_KEY= # optional: require clients to authenticate with this key -# PBSPY_PORT=9876 # override if 9876 is already in use on the host diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d81445..7bbcd6a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,10 +7,29 @@ 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) +### Added + +- 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 + ### Removed -- Remove `progress` parameter from the `wait[_all]` and `result[_all]` methods of `Job` - - These methods no longer print job IDs or progress +- Remove progress display from the `Job` wait and result methods + - Legacy progress arguments remain accepted and are ignored while consumers migrate - Remove `rich` dependency ## [0.0.9](https://github.com/MaterialsPhysicsANU/pbspy/releases/tag/v0.0.9) - 2026-05-18 diff --git a/README.md b/README.md index d6a7c97..4704213 100644 --- a/README.md +++ b/README.md @@ -40,45 +40,11 @@ print("job_a:", result_a.output.strip()) print("job_b:", result_b.output.strip()) ``` -### Submitting from a remote machine via a `pbspy-server` daemon +### Backends -`pbspy-server` runs every PBS command (`qsub`, `qstat`, `qdel`, file reads) locally, -so it must run somewhere with direct PBS access. On NCI Gadi that means a -[persistent session](https://opus.nci.org.au/display/Help/Persistent+Sessions): - -```bash -persistent-sessions start pbspy -ssh pbspy...ps.gadi.nci.org.au -uv tool install pbspy -pbspy-server --host 0.0.0.0 --api-key ... -``` - -Then connect to it with `ServerBackend` (from inside the supercomputer's network, or via an SSH -tunnel/port-forward if you need to reach it from elsewhere): - -```python -from pbspy import Job, JobDescription, ServerBackend - -backend = ServerBackend( - "pbspy...ps.gadi.nci.org.au", api_key="..." -) # api_key only if the server requires it - -job_a = ( - JobDescription(name="job_a", ncpus=4, mem="192GB", walltime="00:05:00") - .add_command(["echo", "A"]) - .submit(backend=backend) -) - -job_b = ( - JobDescription(name="job_b", ncpus=1, walltime="00:05:00", afterok=[job_a]) - .add_command(["echo", "B"]) - .submit(backend=backend) -) - -(result_a, result_b) = Job.result_all([job_a, job_b]) -print("job_a:", result_a.output.strip()) -print("job_b:", result_b.output.strip()) -``` +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. ## Output (partially executed) diff --git a/examples/run.py b/examples/run.py index 3d9bd6c..3db9e21 100644 --- a/examples/run.py +++ b/examples/run.py @@ -1,11 +1,9 @@ import typer -from pbspy import Job, JobDescription, QueueLimits, ServerBackend +from pbspy import Job, JobDescription, QueueLimits -def main(remote: str = "") -> None: - backend = ServerBackend(remote) if remote else None - +def main() -> None: # Run a job with some explicit parameters job_a = ( JobDescription( @@ -15,7 +13,7 @@ def main(remote: str = "") -> None: walltime="00:05:00", ) .add_command(["echo", "A"]) - .submit(backend=backend) + .submit() ) # Submit another job that waits for job_a to finish @@ -31,7 +29,7 @@ def main(remote: str = "") -> None: afterok=[job_a], ) .add_command(["echo", "B"]) - .submit(backend=backend) + .submit() ) # Get the result of the jobs diff --git a/pyproject.toml b/pyproject.toml index a992d2a..30f158e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -21,9 +21,6 @@ license = {text = "MIT License"} dependencies = [ ] -[project.scripts] -pbspy-server = "pbspy.server.__main__:main" - [dependency-groups] dev = [ "mypy>=1.14.1", diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 901745e..957ab1e 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -9,7 +9,6 @@ from pbspy._backend import Backend from pbspy._local_backend import LocalBackend -from pbspy._server_backend import ServerBackend # Shared LocalBackend instance so all locally-submitted jobs are grouped together # when calling Job.wait_all() / Job.result_all(), restoring concurrent polling. @@ -23,7 +22,6 @@ "QueueLimitsMap", "Backend", "LocalBackend", - "ServerBackend", "gadi", ] @@ -72,17 +70,6 @@ class Job: error_path: str | None = None """Custom path for the job's stderr error file (``#PBS -e``), if any.""" - def __getstate__(self) -> dict[str, Any]: - """Exclude backend from pickling (it carries live sockets).""" - state = self.__dict__.copy() - del state["backend"] - return state - - def __setstate__(self, state: dict[str, Any]) -> None: - """Restore job state, assigning the default local backend.""" - self.__dict__.update(state) - self.backend = _DEFAULT_LOCAL_BACKEND - def wait(self, **kwargs: Any) -> None: """ Wait for the job to complete. @@ -94,7 +81,7 @@ def result(self, **kwargs: Any) -> JobResult: Waits for the job to complete and returns the result. """ self.wait() - return self.backend.get_result(self) # type: ignore[return-value] + return self.backend.get_result(self) def cancel(self) -> None: """Cancel (``qdel``) this job.""" @@ -111,12 +98,15 @@ def wait_all(jobs: list[Job], **kwargs: Any) -> None: _wait_all_grouped(jobs) @staticmethod - def result_all(jobs: list[Job], **kwargs: Any) -> list[JobResult]: + def result_all(jobs: list[Job], *_args: Any, **_kwargs: Any) -> list[JobResult]: """ Waits for multiple jobs to complete and returns their results. + + Additional arguments are accepted for compatibility with callers that + still pass the removed ``progress`` option; they have no effect. """ Job.wait_all(jobs) - return [job.backend.get_result(job) for job in jobs] # type: ignore[misc] + return [job.backend.get_result(job) for job in jobs] def _wait_all_grouped(jobs: list[Job]) -> None: @@ -257,11 +247,9 @@ def submit(self, backend: Backend | None = None) -> Job: Args: backend: The backend to use for submission. Defaults to :class:`~pbspy.LocalBackend` (runs ``qsub`` locally). - Pass a :class:`~pbspy.ServerBackend` instance to submit via - a pbspy-server daemon. """ if backend is None: - backend = LocalBackend() + backend = _DEFAULT_LOCAL_BACKEND job_script = self.script() job_id, job_name = backend.submit(job_script, self.name) diff --git a/src/pbspy/_backend.py b/src/pbspy/_backend.py index c1e0902..98ed9e0 100644 --- a/src/pbspy/_backend.py +++ b/src/pbspy/_backend.py @@ -1,8 +1,8 @@ """ Backend abstract base class. -Concrete backends: :class:`~pbspy._local_backend.LocalBackend` (default, -in-process) and :class:`~pbspy._server_backend.ServerBackend` (remote via TCP). +The built-in :class:`~pbspy._local_backend.LocalBackend` runs PBS operations +in-process. Alternate implementations can provide other execution mechanisms. """ from __future__ import annotations @@ -12,7 +12,7 @@ from typing import TYPE_CHECKING if TYPE_CHECKING: - from pbspy import Job + from pbspy import Job, JobResult __all__ = ["Backend"] @@ -21,9 +21,9 @@ class Backend(ABC): """ Abstract base class for PBS operation backends. - All methods that touch PBS (qsub, qstat, file reading) go through a - Backend so the same :class:`~pbspy.JobDescription` / :class:`~pbspy.Job` - API works both locally and over SSH. + 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 @@ -45,15 +45,14 @@ def wait( """ Wait for *jobs* to finish. - Args: - jobs: List of :class:`~pbspy.Job` objects. - on_update: Optional callback ``(job_id, state)`` on each status - change. *state* is ``None`` when the job has finished. + :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: object) -> object: + def get_result(self, job: Job) -> JobResult: """ Return the :class:`~pbspy.JobResult` for a completed job. """ @@ -64,7 +63,6 @@ def delete(self, job_ids: list[str]) -> None: """ Cancel (``qdel``) the given jobs. - Args: - job_ids: Job ids to cancel. + :param job_ids: Job ids to cancel. """ ... diff --git a/src/pbspy/_local_backend.py b/src/pbspy/_local_backend.py index 2b30410..4b1beaf 100644 --- a/src/pbspy/_local_backend.py +++ b/src/pbspy/_local_backend.py @@ -1,9 +1,7 @@ """ LocalBackend: direct in-process calls to :mod:`pbspy._pbs_core`. -No sockets, no IPC, no serialisation. This is the default backend when -running directly on the supercomputer — it behaves identically to the -original pbspy code but delegates through the :class:`Backend` interface. +This is the default backend when running on a host with PBS access. """ from __future__ import annotations @@ -15,7 +13,7 @@ from pbspy._backend import Backend if TYPE_CHECKING: - from pbspy import Job + from pbspy import Job, JobResult __all__ = ["LocalBackend"] @@ -40,9 +38,9 @@ def wait( """Wait for *jobs* to finish.""" core.pbs_wait_for_jobs(jobs, on_update) - def get_result(self, job: object) -> object: # job: Job -> JobResult + def get_result(self, job: Job) -> JobResult: """Return the :class:`~pbspy.JobResult` for a completed job.""" - return core.pbs_get_result(job) # type: ignore[arg-type] + return core.pbs_get_result(job) def delete(self, job_ids: list[str]) -> None: """Cancel (``qdel``) the given jobs.""" diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py index bb69818..ce5b5c5 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -1,8 +1,8 @@ """ Core PBS operations: qsub submission, qstat polling, and result file reading. -These functions are called directly by LocalBackend (in-process) and by the -server daemon when handling pickled requests from ServerBackend clients. +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 @@ -98,12 +98,11 @@ def pbs_wait_for_jobs( """ Poll qstat until all jobs in *jobs* have finished. - Args: - jobs: The jobs to wait for. - 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). + :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) last_check = time.time() - _POLL_INTERVAL_SECONDS # poll immediately on first iteration diff --git a/src/pbspy/_protocol.py b/src/pbspy/_protocol.py deleted file mode 100644 index bb3e9f0..0000000 --- a/src/pbspy/_protocol.py +++ /dev/null @@ -1,200 +0,0 @@ -""" -Request and response dataclasses for the pbspy client-server protocol. - -Messages are exchanged as length-prefixed pickle frames (see :func:`send_frame` -and :func:`recv_frame`). Connections optionally require authentication via an -API key (see :class:`AuthRequest`). When enabled, the server sends an -:class:`AuthOkResponse` on success or an :class:`ErrorResponse` on failure -before accepting any other requests. -""" - -from __future__ import annotations - -import io -import pickle -import struct -from dataclasses import dataclass, field -from typing import TYPE_CHECKING, BinaryIO - -if TYPE_CHECKING: - from pbspy import Job, JobResult - -__all__ = [ - # Requests - "AuthRequest", - "SubmitRequest", - "WaitRequest", - "ResultRequest", - "PingRequest", - "DeleteRequest", - # Responses - "AuthOkResponse", - "SubmittedResponse", - "StatusUpdateResponse", - "WaitDoneResponse", - "ResultResponse", - "PongResponse", - "ErrorResponse", - "DeleteResponse", - # Framing - "send_frame", - "recv_frame", -] - -_FRAME_HEADER = struct.Struct("!I") # 4-byte big-endian unsigned int - - -# --------------------------------------------------------------------------- -# Requests -# --------------------------------------------------------------------------- - - -@dataclass -class SubmitRequest: - """Ask the server to submit a PBS job script.""" - - script: str - name: str | None = None - - -@dataclass -class WaitRequest: - """ - Ask the server to wait for *jobs* and stream :class:`StatusUpdateResponse` - messages until all jobs finish, then send :class:`WaitDoneResponse`. - """ - - jobs: list[Job] = field(default_factory=list) - - -@dataclass -class ResultRequest: - """Ask the server to return the result of a completed job.""" - - job: Job - - -@dataclass -class AuthRequest: - """Sent as the first frame on connection when the server requires API key auth.""" - - api_key: str | None = None - - -@dataclass -class PingRequest: - """Liveness check; server responds with :class:`PongResponse`.""" - - -@dataclass -class DeleteRequest: - """Ask the server to cancel (``qdel``) one or more jobs.""" - - job_ids: list[str] - - -# --------------------------------------------------------------------------- -# Responses -# --------------------------------------------------------------------------- - - -@dataclass -class SubmittedResponse: - """Returned after a successful :class:`SubmitRequest`.""" - - job: Job - - -@dataclass -class StatusUpdateResponse: - """ - Streamed periodically during a :class:`WaitRequest`. - - *state* is one of ``"Q"`` (queued), ``"R"`` (running), ``"E"`` (exiting), - ``"F"`` (finished), or ``None`` when the job has left the qstat queue. - """ - - job: Job - state: str | None - - -@dataclass -class WaitDoneResponse: - """Sent after all jobs in a :class:`WaitRequest` have finished.""" - - jobs: list[Job] = field(default_factory=list) - - -@dataclass -class ResultResponse: - """Returned after a successful :class:`ResultRequest`.""" - - job: Job - result: JobResult - - -@dataclass -class PongResponse: - """Response to :class:`PingRequest`.""" - - -@dataclass -class AuthOkResponse: - """Sent after a successful :class:`AuthRequest`.""" - - -@dataclass -class ErrorResponse: - """Returned when the server encounters an error handling a request.""" - - message: str - - -@dataclass -class DeleteResponse: - """Returned after a successful :class:`DeleteRequest`.""" - - -# --------------------------------------------------------------------------- -# Framing helpers -# --------------------------------------------------------------------------- - - -def send_frame(stream: BinaryIO, obj: object) -> None: - """ - Serialise *obj* with :mod:`pickle` and write it to *stream* as a - length-prefixed frame. - - Frame layout: ``[4-byte big-endian length][pickle bytes]`` - """ - data = pickle.dumps(obj) - stream.write(_FRAME_HEADER.pack(len(data))) - stream.write(data) - if hasattr(stream, "flush"): - stream.flush() - - -def recv_frame(stream: BinaryIO) -> object: - """ - Read one length-prefixed frame from *stream* and return the unpickled - object. - - Raises :class:`EOFError` if the stream is closed before a complete frame - is received. - """ - header = _read_exactly(stream, _FRAME_HEADER.size) - (length,) = _FRAME_HEADER.unpack(header) - data = _read_exactly(stream, length) - return pickle.loads(data) # noqa: S301 - - -def _read_exactly(stream: BinaryIO, n: int) -> bytes: - buf = io.BytesIO() - remaining = n - while remaining > 0: - chunk = stream.read(remaining) - if not chunk: - raise EOFError("Stream closed before a complete frame was received") - buf.write(chunk) - remaining -= len(chunk) - return buf.getvalue() diff --git a/src/pbspy/_server_backend.py b/src/pbspy/_server_backend.py deleted file mode 100644 index 939d028..0000000 --- a/src/pbspy/_server_backend.py +++ /dev/null @@ -1,66 +0,0 @@ -""" -ServerBackend: communicates with a pbspy-server daemon over a plain TCP connection. - -Messages are exchanged as length-prefixed pickle frames (see :mod:`pbspy._protocol`). -The server runs PBS commands (qsub, qstat, qdel) locally on the machine it's started on, -so it's typically started on a machine with direct PBS access (e.g. a Gadi persistent -session). -""" - -from __future__ import annotations - -import socket -from typing import BinaryIO, cast - -from pbspy._stream_backend import StreamBackend - -__all__ = ["ServerBackend"] - - -class ServerBackend(StreamBackend): - """ - Backend that connects to a pbspy-server daemon over TCP. - - The server (started with ``pbspy-server``) runs locally on a machine with direct PBS - access (e.g. a Gadi persistent session with ``/g/data`` mounted); this client connects - to it directly over TCP. - - Args: - host: Hostname or IP address of the machine running pbspy-server. - port: TCP port the server is listening on (default: 9876). - connect_timeout: Seconds to wait for the initial TCP connection. - api_key: API key for authentication (required if the server requires it). - """ - - def __init__( - self, - host: str, - port: int = 9876, - connect_timeout: float = 10.0, - api_key: str | None = None, - ) -> None: - self._host = host - self._port = port - self._connect_timeout = connect_timeout - self._sock: socket.socket | None = None - super().__init__(api_key=api_key) - - def _open_stream(self) -> BinaryIO: - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.settimeout(self._connect_timeout) - try: - sock.connect((self._host, self._port)) - except OSError as exc: - sock.close() - raise RuntimeError(f"Could not connect to pbspy-server at {self._host}:{self._port}: {exc}") from exc - sock.settimeout(None) - self._sock = sock - return cast(BinaryIO, sock.makefile("rwb", buffering=0)) - - def _close_stream(self) -> None: - if self._sock is not None: - try: - self._sock.close() - except OSError: - pass - self._sock = None diff --git a/src/pbspy/_stream_backend.py b/src/pbspy/_stream_backend.py deleted file mode 100644 index c8de9de..0000000 --- a/src/pbspy/_stream_backend.py +++ /dev/null @@ -1,181 +0,0 @@ -""" -StreamBackend: shared RPC logic for backends that talk the pbspy wire protocol -over some duplex byte stream. - -:class:`~pbspy._server_backend.ServerBackend` connects over a plain TCP socket. A future -backend that talks directly to a Server over an SSH subprocess's stdin/stdout pipes (the -way the pre-0.0.10 ``SSHBackend`` did) would reuse this same RPC machinery — only how the -stream is opened differs between transports. -""" - -from __future__ import annotations - -import threading -from collections.abc import Callable -from typing import TYPE_CHECKING, BinaryIO - -import pbspy._protocol as proto -from pbspy._backend import Backend - -if TYPE_CHECKING: - from pbspy import Job - -__all__ = ["StreamBackend"] - - -class StreamBackend(Backend): - """ - Backend that exchanges length-prefixed pickle frames over a duplex byte stream. - - Subclasses implement :meth:`_open_stream` (and optionally :meth:`_close_stream`) to - provide the transport; everything else (auth handshake, ping, submit/wait/result RPCs, - reconnect-on-failure) is shared. - """ - - def __init__(self, api_key: str | None = None) -> None: - self._api_key = api_key - self._stream: BinaryIO | None = None - self._lock = threading.Lock() - self._connect() - - def _open_stream(self) -> BinaryIO: - """Open the transport and return a readable+writable stream. Implemented by subclasses.""" - raise NotImplementedError - - def _close_stream(self) -> None: - """Release any transport resources beyond the stream itself. Implemented by subclasses.""" - - def submit(self, script: str, name: str | None = None) -> tuple[str, str]: - response = self._rpc(proto.SubmitRequest(script=script, name=name)) - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.SubmittedResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - return response.job.job_id, response.job.job_name or "" - - def wait( - self, - jobs: list[Job], - on_update: Callable[[str, str | None], None] | None = None, - ) -> None: - self._wait_quiet(jobs, on_update) - - def get_result(self, job: object) -> object: - response = self._rpc(proto.ResultRequest(job=job)) # type: ignore[arg-type] - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.ResultResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - return response.result - - def delete(self, job_ids: list[str]) -> None: - """ - Cancel (``qdel``) the given jobs on the server. - - Args: - job_ids: Job ids to cancel. - - Raises: - RuntimeError: If the server rejects the request. - """ - response = self._rpc(proto.DeleteRequest(job_ids=job_ids)) - if isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error: {response.message}") - if not isinstance(response, proto.DeleteResponse): - raise RuntimeError(f"Unexpected response: {response!r}") - - def _connect(self) -> None: - """Open the transport, authenticate if required, and verify liveness with a ping.""" - stream = self._open_stream() - try: - if self._api_key is not None: - proto.send_frame(stream, proto.AuthRequest(api_key=self._api_key)) - auth_resp = proto.recv_frame(stream) - if isinstance(auth_resp, proto.ErrorResponse): - raise RuntimeError(f"pbspy-server auth failed: {auth_resp.message}") - if not isinstance(auth_resp, proto.AuthOkResponse): - raise RuntimeError(f"pbspy-server returned unexpected auth response: {auth_resp!r}") - - proto.send_frame(stream, proto.PingRequest()) - pong = proto.recv_frame(stream) - if not isinstance(pong, proto.PongResponse): - raise RuntimeError(f"pbspy-server returned unexpected response to ping: {pong!r}") - except (OSError, EOFError) as exc: - try: - stream.close() - except OSError: - pass - self._close_stream() - raise RuntimeError(f"pbspy-server did not respond to ping: {exc}") from exc - except RuntimeError: - try: - stream.close() - except OSError: - pass - self._close_stream() - raise - self._stream = stream - - def _reconnect(self) -> None: - """Close the current connection and establish a fresh one.""" - if self._stream is not None: - try: - self._stream.close() - except OSError: - pass - self._stream = None - self._close_stream() - self._connect() - - def _rpc(self, request: object) -> object: - """Send one request frame and return the first response frame.""" - with self._lock: - try: - assert self._stream is not None - proto.send_frame(self._stream, request) - return proto.recv_frame(self._stream) - except (OSError, EOFError): - self._reconnect() - assert self._stream is not None - proto.send_frame(self._stream, request) - return proto.recv_frame(self._stream) - - def _send(self, request: object) -> None: - with self._lock: - assert self._stream is not None - proto.send_frame(self._stream, request) - - def _recv(self) -> object: - assert self._stream is not None - return proto.recv_frame(self._stream) - - def _wait_quiet( - self, - jobs: list[Job], - on_update: Callable[[str, str | None], None] | None, - ) -> None: - self._send(proto.WaitRequest(jobs=jobs)) - while True: - response = self._recv() - if isinstance(response, proto.StatusUpdateResponse): - if on_update: - on_update(response.job.job_id, response.state) - elif isinstance(response, proto.WaitDoneResponse): - return - elif isinstance(response, proto.ErrorResponse): - raise RuntimeError(f"Server error while waiting: {response.message}") - - def close(self) -> None: - if self._stream is not None: - try: - self._stream.close() - except OSError: - pass - self._stream = None - self._close_stream() - - def __del__(self) -> None: - try: - self.close() - except Exception: - pass diff --git a/src/pbspy/server/__init__.py b/src/pbspy/server/__init__.py deleted file mode 100644 index bc2aee6..0000000 --- a/src/pbspy/server/__init__.py +++ /dev/null @@ -1,7 +0,0 @@ -"""pbspy-server: PBS job scheduler server daemon for pbspy.""" - -from __future__ import annotations - -from pbspy.server.server import run_server - -__all__ = ["run_server"] diff --git a/src/pbspy/server/__main__.py b/src/pbspy/server/__main__.py deleted file mode 100644 index 29b9353..0000000 --- a/src/pbspy/server/__main__.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -pbspy-server command-line entry point. - -Usage:: - - pbspy-server [--host HOST] [--port PORT] [--api-key KEY] -""" - -from __future__ import annotations - -import argparse -import logging -import os - - -def main() -> None: - parser = argparse.ArgumentParser( - prog="pbspy-server", - description="PBS job server daemon for pbspy.", - ) - parser.add_argument( - "--host", - default="127.0.0.1", - metavar="HOST", - help="Local address to bind (default: 127.0.0.1).", - ) - parser.add_argument( - "--port", - type=int, - default=9876, - metavar="PORT", - help="TCP port to listen on (default: 9876).", - ) - parser.add_argument( - "--api-key", - metavar="KEY", - default=os.environ.get("PBSPY_API_KEY"), - help="Require clients to authenticate with this key (env: PBSPY_API_KEY).", - ) - - args = parser.parse_args() - - logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") - - from pbspy.server import run_server - - run_server( - host=args.host, - port=args.port, - api_key=args.api_key, - ) - - -if __name__ == "__main__": - main() diff --git a/src/pbspy/server/server.py b/src/pbspy/server/server.py deleted file mode 100644 index ab796e7..0000000 --- a/src/pbspy/server/server.py +++ /dev/null @@ -1,266 +0,0 @@ -""" -pbspy-server daemon. - -Listens on a configured TCP port, accepts connections from :class:`~pbspy.ServerBackend` -clients, and dispatches pickled requests from :mod:`pbspy._protocol`. - -PBS operations (qsub, qstat, qdel, file reads) are executed locally, so this server is -intended to run natively on the machine with PBS access (e.g. a Gadi persistent session -with ``/g/data`` mounted) rather than behind SSH. - -A background thread polls ``qstat`` every 60 seconds for all unfinished jobs -and pushes :class:`StatusUpdateResponse` frames to any clients waiting on -those jobs. Job state is held in memory and is not persisted across restarts. -""" - -from __future__ import annotations - -import logging -import signal -import socket -import threading -from collections.abc import Callable -from dataclasses import dataclass -from typing import BinaryIO, cast - -import pbspy._pbs_core as core -import pbspy._protocol as proto -from pbspy import Job, JobResult - -__all__ = ["run_server"] - -logger = logging.getLogger(__name__) - -_POLL_INTERVAL = 60.0 # seconds between qstat polls - - -@dataclass -class _JobRecord: - """In-memory record of a submitted job's current state.""" - - job_id: str - job_name: str | None = None - finished: bool = False - - -class _WaitSubscription: - """Tracks a client waiting for a set of jobs to finish.""" - - def __init__(self, jobs: list[Job], stream: BinaryIO) -> None: - self.pending: set[str] = {j.job_id for j in jobs} - self.jobs_by_id: dict[str, Job] = {j.job_id: j for j in jobs} - self.stream = stream - self.lock = threading.Lock() - self.done = threading.Event() - - -class _ServerState: - """Shared mutable state accessed by both connection threads and the poll thread.""" - - def __init__(self, runner: core.PBSRunner, api_key: str | None = None) -> None: - self.runner = runner - self.api_key = api_key - self.lock = threading.Lock() - self.jobs: dict[str, _JobRecord] = {} - self.subscriptions: dict[str, list[_WaitSubscription]] = {} - - def notify_finished(self, job_id: str, job: Job) -> None: - """Mark job_id as done in all waiting subscriptions, sending frames as needed.""" - with self.lock: - subs = self.subscriptions.pop(job_id, []) - for sub in subs: - with sub.lock: - try: - proto.send_frame(sub.stream, proto.StatusUpdateResponse(job=job, state=None)) - except OSError: - pass - sub.pending.discard(job_id) - if not sub.pending: - try: - proto.send_frame(sub.stream, proto.WaitDoneResponse(jobs=list(sub.jobs_by_id.values()))) - except OSError: - pass - sub.done.set() - - -def _cancel_all_subscriptions(state: _ServerState) -> None: - """Signal all waiting subscriptions to unblock their connection threads (for shutdown).""" - with state.lock: - subs_by_id = dict(state.subscriptions) - state.subscriptions.clear() - seen: set[int] = set() - for subs in subs_by_id.values(): - for sub in subs: - if id(sub) not in seen: - seen.add(id(sub)) - sub.done.set() - - -def _poll_loop( - state: _ServerState, - stop_event: threading.Event, - poll_interval: float = _POLL_INTERVAL, -) -> None: - """Background thread: polls qstat for all unfinished jobs.""" - while not stop_event.wait(timeout=poll_interval): - try: - with state.lock: - unfinished = [r for r in state.jobs.values() if not r.finished] - if not unfinished: - continue - states = core.pbs_get_states([r.job_id for r in unfinished], runner=state.runner) - for record in unfinished: - if states.get(record.job_id) is None: - job = Job(job_id=record.job_id, job_name=record.job_name) - with state.lock: - record.finished = True - state.notify_finished(record.job_id, job) - except Exception: - logger.exception("Error in poll loop") - - -def _handle_connection(conn: socket.socket, state: _ServerState) -> None: - """Handle a single client connection in its own thread.""" - stream = cast(BinaryIO, conn.makefile("rwb", buffering=0)) - try: - # Auth pre-check (runs before any other request) - if state.api_key is not None: - try: - first = proto.recv_frame(stream) - except EOFError: - return - if not isinstance(first, proto.AuthRequest) or first.api_key != state.api_key: - proto.send_frame(stream, proto.ErrorResponse(message="Authentication failed")) - return - proto.send_frame(stream, proto.AuthOkResponse()) - - while True: - try: - request = proto.recv_frame(stream) - except EOFError: - break - except Exception: - logger.exception("Error receiving frame") - break - - try: - if isinstance(request, proto.PingRequest): - proto.send_frame(stream, proto.PongResponse()) - - elif isinstance(request, proto.SubmitRequest): - job_id, job_name = core.pbs_submit(request.script, request.name, runner=state.runner) - job = Job(job_id=job_id, job_name=job_name) - with state.lock: - state.jobs[job_id] = _JobRecord(job_id=job_id, job_name=job_name) - proto.send_frame(stream, proto.SubmittedResponse(job=job)) - - elif isinstance(request, proto.WaitRequest): - sub = _WaitSubscription(request.jobs, stream) - # Check state and register subscription atomically so notify_finished() - # cannot fire between the two steps and leave us waiting on a job - # that is already done. - with state.lock: - for job in request.jobs: - record = state.jobs.get(job.job_id) - if record and record.finished: - sub.pending.discard(job.job_id) - if sub.pending: - for job_id in sub.pending: - state.subscriptions.setdefault(job_id, []).append(sub) - if sub.pending: - sub.done.wait() - else: - proto.send_frame(stream, proto.WaitDoneResponse(jobs=request.jobs)) - - elif isinstance(request, proto.ResultRequest): - result: JobResult = core.pbs_get_result(request.job, runner=state.runner) - proto.send_frame(stream, proto.ResultResponse(job=request.job, result=result)) - - elif isinstance(request, proto.AuthRequest): - proto.send_frame(stream, proto.AuthOkResponse()) - - elif isinstance(request, proto.DeleteRequest): - core.pbs_delete(request.job_ids, runner=state.runner) - proto.send_frame(stream, proto.DeleteResponse()) - - else: - proto.send_frame(stream, proto.ErrorResponse(message=f"Unknown request type: {type(request)}")) - - except Exception as exc: - logger.exception("Error handling request") - try: - proto.send_frame(stream, proto.ErrorResponse(message=str(exc))) - except OSError: - pass - finally: - stream.close() - conn.close() - - -def run_server( - host: str = "0.0.0.0", - port: int = 9876, - poll_interval: float = _POLL_INTERVAL, - api_key: str | None = None, - _stop_event: threading.Event | None = None, - _ready_callback: Callable[[int], None] | None = None, -) -> None: - """ - Start the pbspy-server daemon. - - Runs in the foreground; use your OS's service manager (systemd, nohup, - etc.) to run it as a background service. - - All PBS operations (qsub, qstat, qdel, file reads) run locally, so this should be started - on a machine with PBS access (e.g. a Gadi persistent session with ``/g/data`` mounted). - - Args: - host: Local address to bind (default ``"0.0.0.0"``). - port: TCP port to listen on (default 9876; pass 0 for OS-assigned). - poll_interval: Seconds between qstat polls (default 60). - api_key: When set, clients must send a matching :class:`~pbspy._protocol.AuthRequest` - as the first frame or the connection is rejected. - """ - runner = core.PBSRunner() - state = _ServerState(runner=runner, api_key=api_key) - - stop_event = _stop_event or threading.Event() - - poll_thread = threading.Thread( - target=_poll_loop, args=(state, stop_event), kwargs={"poll_interval": poll_interval}, daemon=True - ) - poll_thread.start() - - if threading.current_thread() is threading.main_thread(): - - def _handle_signal(signum: int, frame: object) -> None: - logger.info("Received signal %d, shutting down", signum) - stop_event.set() - - signal.signal(signal.SIGTERM, _handle_signal) - signal.signal(signal.SIGINT, _handle_signal) - - srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - srv.bind((host, port)) - srv.listen(16) - srv.settimeout(1.0) - - actual_port = srv.getsockname()[1] - logger.info("pbspy-server listening on %s:%d", host, actual_port) - - if _ready_callback is not None: - _ready_callback(actual_port) - - try: - while not stop_event.is_set(): - try: - conn, _ = srv.accept() - except TimeoutError: - continue - t = threading.Thread(target=_handle_connection, args=(conn, state), daemon=True) - t.start() - finally: - srv.close() - _cancel_all_subscriptions(state) - logger.info("pbspy-server stopped") diff --git a/tests/test_default.py b/tests/test_default.py index 2b99664..71df64e 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -1,11 +1,10 @@ -import pickle import shutil import subprocess from collections.abc import Callable import pytest -from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend, ServerBackend +from pbspy import Backend, Job, JobDescription, JobResult, LocalBackend @pytest.mark.skipif(shutil.which("qsub") is None, reason="qsub not available") @@ -51,8 +50,7 @@ def wait( ) -> None: self.waited.append(list(jobs)) - def get_result(self, job: object) -> object: - assert isinstance(job, Job) + 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: @@ -105,6 +103,11 @@ def test_mock_backend_result_all() -> None: assert results[1].output == "B\n" +def test_result_all_accepts_legacy_positional_progress() -> None: + """Existing consumers may still pass the removed progress flag positionally.""" + assert Job.result_all([], False) == [] + + def test_script_generation_unchanged() -> None: """script() output is unchanged from the original implementation.""" jd = JobDescription( @@ -130,48 +133,19 @@ def test_script_generation_unchanged() -> None: def test_backend_classes_importable() -> None: - """Backend, LocalBackend, ServerBackend are importable from pbspy.""" + """Backend and LocalBackend are importable from pbspy.""" assert issubclass(LocalBackend, Backend) - assert issubclass(ServerBackend, Backend) -def test_job_pickle_round_trip() -> None: - """Job pickle round-trip excludes backend and restores all other fields. +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)) - Regression test: Job.backend carries live sockets (ServerBackend) and - must not be pickled. See __getstate__/__setstate__ on Job. - """ - job = Job( - job_id="42.mock", - job_name="my_job", - description="some description", - backend=LocalBackend(), - ) + job_a = JobDescription(name="job_a").submit() + job_b = JobDescription(name="job_b").submit() - # Pickle and unpickle - data = pickle.dumps(job) - restored = pickle.loads(data) - - # All non-backend fields preserved - assert restored.job_id == "42.mock" - assert restored.job_name == "my_job" - assert restored.description == "some description" - # backend is restored to the default local backend (never pickled) - assert isinstance(restored.backend, LocalBackend) - - -def test_job_pickle_excludes_backend_from_state() -> None: - """__getstate__ does not include the backend attribute.""" - job = Job(job_id="1.mock", backend=LocalBackend()) - state = job.__getstate__() - assert "backend" not in state - assert state == { - "job_name": None, - "job_id": "1.mock", - "description": None, - "output_path": None, - "error_path": None, - } + assert job_a.backend is job_b.backend def test_submit_carries_output_and_error_paths() -> None: @@ -270,25 +244,6 @@ def run(self, cmd: list[str], *, input: bytes | None = None) -> subprocess.Compl assert result.error == "" -def test_job_pickle_preserves_output_and_error_paths() -> None: - """Job pickle round-trip preserves output_path and error_path.""" - job = Job( - job_id="789.mock", - job_name="my_job", - output_path="/scratch/project/job.out", - error_path="/scratch/project/job.err", - backend=LocalBackend(), - ) - - data = pickle.dumps(job) - restored = pickle.loads(data) - - assert restored.output_path == "/scratch/project/job.out" - assert restored.error_path == "/scratch/project/job.err" - assert restored.job_id == "789.mock" - assert restored.job_name == "my_job" - - # --------------------------------------------------------------------------- # pbs_get_states / pbs_delete / Job.cancel # --------------------------------------------------------------------------- diff --git a/tests/test_protocol.py b/tests/test_protocol.py deleted file mode 100644 index 8a7e2ce..0000000 --- a/tests/test_protocol.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Tests for the pickle framing helpers in pbspy._protocol.""" - -from __future__ import annotations - -import io -import pickle - -import pytest - -import pbspy._protocol as proto -from pbspy import Job, JobResult - -# --------------------------------------------------------------------------- -# Helpers -# --------------------------------------------------------------------------- - - -def roundtrip(obj: object) -> object: - """Serialise *obj* into a BytesIO buffer, then deserialise and return it.""" - buf = io.BytesIO() - proto.send_frame(buf, obj) - buf.seek(0) - return proto.recv_frame(buf) - - -# --------------------------------------------------------------------------- -# Request types -# --------------------------------------------------------------------------- - - -def test_roundtrip_ping() -> None: - result = roundtrip(proto.PingRequest()) - assert isinstance(result, proto.PingRequest) - - -def test_roundtrip_submit_request() -> None: - req = proto.SubmitRequest(script="#!/bin/bash\necho hi", name="myjob") - result = roundtrip(req) - assert isinstance(result, proto.SubmitRequest) - assert result.script == req.script - assert result.name == req.name - - -def test_roundtrip_submit_request_no_name() -> None: - req = proto.SubmitRequest(script="#!/bin/bash") - result = roundtrip(req) - assert isinstance(result, proto.SubmitRequest) - assert result.name is None - - -def test_roundtrip_wait_request() -> None: - jobs = [Job(job_id="1.gadi", job_name="a"), Job(job_id="2.gadi", job_name="b")] - req = proto.WaitRequest(jobs=jobs) - result = roundtrip(req) - assert isinstance(result, proto.WaitRequest) - assert [j.job_id for j in result.jobs] == ["1.gadi", "2.gadi"] - - -def test_roundtrip_result_request() -> None: - job = Job(job_id="42.gadi", job_name="myjob") - req = proto.ResultRequest(job=job) - result = roundtrip(req) - assert isinstance(result, proto.ResultRequest) - assert result.job.job_id == "42.gadi" - - -def test_roundtrip_delete_request() -> None: - req = proto.DeleteRequest(job_ids=["1.gadi", "2.gadi"]) - result = roundtrip(req) - assert isinstance(result, proto.DeleteRequest) - assert result.job_ids == ["1.gadi", "2.gadi"] - - -# --------------------------------------------------------------------------- -# Response types -# --------------------------------------------------------------------------- - - -def test_roundtrip_pong() -> None: - assert isinstance(roundtrip(proto.PongResponse()), proto.PongResponse) - - -def test_roundtrip_submitted_response() -> None: - job = Job(job_id="99.gadi", job_name="submitted") - resp = proto.SubmittedResponse(job=job) - result = roundtrip(resp) - assert isinstance(result, proto.SubmittedResponse) - assert result.job.job_id == "99.gadi" - - -def test_roundtrip_status_update() -> None: - job = Job(job_id="1.gadi") - resp = proto.StatusUpdateResponse(job=job, state="R") - result = roundtrip(resp) - assert isinstance(result, proto.StatusUpdateResponse) - assert result.state == "R" - - -def test_roundtrip_status_update_finished() -> None: - job = Job(job_id="1.gadi") - resp = proto.StatusUpdateResponse(job=job, state=None) - result = roundtrip(resp) - assert isinstance(result, proto.StatusUpdateResponse) - assert result.state is None - - -def test_roundtrip_wait_done() -> None: - jobs = [Job(job_id="1.gadi"), Job(job_id="2.gadi")] - resp = proto.WaitDoneResponse(jobs=jobs) - result = roundtrip(resp) - assert isinstance(result, proto.WaitDoneResponse) - assert len(result.jobs) == 2 - - -def test_roundtrip_result_response() -> None: - job = Job(job_id="5.gadi", job_name="myjob") - job_result = JobResult(exit_code=0, output="hello\n", error="", stats={"CPU": "1s"}) - resp = proto.ResultResponse(job=job, result=job_result) - result = roundtrip(resp) - assert isinstance(result, proto.ResultResponse) - assert result.result.exit_code == 0 - assert result.result.output == "hello\n" - assert result.result.stats == {"CPU": "1s"} - - -def test_roundtrip_error_response() -> None: - resp = proto.ErrorResponse(message="something went wrong") - result = roundtrip(resp) - assert isinstance(result, proto.ErrorResponse) - assert result.message == "something went wrong" - - -def test_roundtrip_delete_response() -> None: - result = roundtrip(proto.DeleteResponse()) - assert isinstance(result, proto.DeleteResponse) - - -# --------------------------------------------------------------------------- -# Framing edge cases -# --------------------------------------------------------------------------- - - -def test_eof_on_empty_stream() -> None: - buf = io.BytesIO(b"") - with pytest.raises(EOFError): - proto.recv_frame(buf) - - -def test_eof_on_truncated_header() -> None: - buf = io.BytesIO(b"\x00\x00") # Only 2 bytes of a 4-byte header - with pytest.raises(EOFError): - proto.recv_frame(buf) - - -def test_eof_on_truncated_body() -> None: - data = pickle.dumps(proto.PingRequest()) - # Write a header claiming a longer body than we provide - header = len(data).to_bytes(4, "big") - buf = io.BytesIO(header + data[:2]) # body truncated - with pytest.raises(EOFError): - proto.recv_frame(buf) - - -def test_multiple_frames_in_sequence() -> None: - buf = io.BytesIO() - proto.send_frame(buf, proto.PingRequest()) - proto.send_frame(buf, proto.PongResponse()) - proto.send_frame(buf, proto.ErrorResponse(message="test")) - buf.seek(0) - - assert isinstance(proto.recv_frame(buf), proto.PingRequest) - assert isinstance(proto.recv_frame(buf), proto.PongResponse) - err = proto.recv_frame(buf) - assert isinstance(err, proto.ErrorResponse) - assert err.message == "test" - - -def test_large_object_roundtrip() -> None: - """Frames larger than a typical network buffer must reassemble correctly.""" - big_output = "x" * 500_000 - result_obj = JobResult(exit_code=0, output=big_output) - job = Job(job_id="big.gadi", job_name="bigtest") - resp = proto.ResultResponse(job=job, result=result_obj) - result = roundtrip(resp) - assert isinstance(result, proto.ResultResponse) - assert len(result.result.output) == 500_000 diff --git a/tests/test_server.py b/tests/test_server.py deleted file mode 100644 index 4915eb5..0000000 --- a/tests/test_server.py +++ /dev/null @@ -1,340 +0,0 @@ -""" -Integration tests for pbspy.server. - -Spins up a real server (in a background thread) with mocked PBS operations so -that no ``qsub``/``qstat`` installation is needed. -""" - -from __future__ import annotations - -import queue -import socket -import threading -import time -from collections.abc import Generator -from typing import BinaryIO, cast -from unittest.mock import MagicMock, patch - -import pytest - -import pbspy._protocol as proto -from pbspy import Job, JobResult -from pbspy.server import run_server - -# --------------------------------------------------------------------------- -# Fixture: running server -# --------------------------------------------------------------------------- - -_WAIT_TIMEOUT = 5.0 # seconds to wait for the server to become ready - - -class ServerHandle: - """Handle to a running test server.""" - - def __init__(self, host: str, port: int) -> None: - self.host = host - self.port = port - - def connect(self, timeout: float | None = None) -> BinaryIO: - """Open a TCP connection and return a binary file-like object.""" - sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - sock.connect((self.host, self.port)) - if timeout is not None: - sock.settimeout(timeout) - f = cast(BinaryIO, sock.makefile("rwb", buffering=0)) - # Close the socket object so that when the SocketIO is GC'd it properly - # calls _real_close() (via _decref_socketios). Without this, the socket - # FD leaks and Python emits a ResourceWarning. - sock.close() - return f - - def rpc(self, request: object) -> object: - """Send one request, return the first response, then close.""" - stream = self.connect() - try: - proto.send_frame(stream, request) - return proto.recv_frame(stream) - finally: - stream.close() - - -@pytest.fixture() -def mock_pbs_submit() -> Generator[MagicMock, None, None]: - """Patch pbs_submit to return a fake job without calling qsub.""" - with patch("pbspy.server.server.core.pbs_submit", return_value=("100.mock", "test_job")) as m: - yield m - - -@pytest.fixture() -def mock_get_states() -> Generator[MagicMock, None, None]: - """Patch core.pbs_get_states to report jobs as still running (never finished) by default.""" - state = {"value": "R"} - - def _get_states(job_ids: list[str], runner: object = None) -> dict[str, str | None]: - return dict.fromkeys(job_ids, state["value"]) - - with patch("pbspy.server.server.core.pbs_get_states", side_effect=_get_states) as m: - m.state = state - yield m - - -@pytest.fixture() -def mock_pbs_get_result() -> Generator[MagicMock, None, None]: - with patch( - "pbspy.server.server.core.pbs_get_result", - return_value=JobResult(exit_code=0, output="hello\n"), - ) as m: - yield m - - -@pytest.fixture() -def server( - mock_pbs_submit: MagicMock, - mock_get_states: MagicMock, - mock_pbs_get_result: MagicMock, -) -> Generator[ServerHandle, None, None]: - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - t = threading.Thread( - target=run_server, - kwargs={"port": 0, "poll_interval": 0.1, "_ready_callback": port_q.put, "_stop_event": stop}, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - yield ServerHandle("127.0.0.1", port) - stop.set() - t.join(timeout=2.0) - - -@pytest.fixture() -def server_with_key( - mock_pbs_submit: MagicMock, - mock_get_states: MagicMock, - mock_pbs_get_result: MagicMock, -) -> Generator[ServerHandle, None, None]: - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - t = threading.Thread( - target=run_server, - kwargs={ - "port": 0, - "poll_interval": 0.1, - "api_key": "secret", - "_ready_callback": port_q.put, - "_stop_event": stop, - }, - daemon=True, - ) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - yield ServerHandle("127.0.0.1", port) - stop.set() - t.join(timeout=2.0) - - -# --------------------------------------------------------------------------- -# Ping / pong -# --------------------------------------------------------------------------- - - -def test_ping(server: ServerHandle) -> None: - response = server.rpc(proto.PingRequest()) - assert isinstance(response, proto.PongResponse) - - -# --------------------------------------------------------------------------- -# Submit -# --------------------------------------------------------------------------- - - -def test_submit_returns_job(server: ServerHandle) -> None: - response = server.rpc(proto.SubmitRequest(script="#!/bin/bash\necho hi", name="test_job")) - assert isinstance(response, proto.SubmittedResponse) - assert response.job.job_id == "100.mock" - assert response.job.job_name == "test_job" - - -# --------------------------------------------------------------------------- -# Result -# --------------------------------------------------------------------------- - - -def test_result_returns_job_result(server: ServerHandle) -> None: - job = Job(job_id="100.mock", job_name="test_job") - response = server.rpc(proto.ResultRequest(job=job)) - assert isinstance(response, proto.ResultResponse) - assert response.result.exit_code == 0 - assert response.result.output == "hello\n" - - -# --------------------------------------------------------------------------- -# Wait — already finished -# --------------------------------------------------------------------------- - - -def test_wait_already_finished_returns_immediately() -> None: - """If the job is already marked F in state, WaitRequest should return without blocking.""" - port_q: queue.SimpleQueue[int] = queue.SimpleQueue() - stop = threading.Event() - - def _run() -> None: - with ( - patch("pbspy.server.server.core.pbs_submit", return_value=("done.mock", "done_job")), - # Return None immediately so the very first poll marks the job finished. - patch("pbspy.server.server.core.pbs_get_states", return_value={"done.mock": None}), - patch("pbspy.server.server.core.pbs_get_result", return_value=JobResult(exit_code=0)), - ): - run_server(port=0, poll_interval=0.1, _ready_callback=port_q.put, _stop_event=stop) - - t = threading.Thread(target=_run, daemon=True) - t.start() - port = port_q.get(timeout=_WAIT_TIMEOUT) - handle = ServerHandle("127.0.0.1", port) - - # Submit the job, then wait long enough for the poll loop (0.1s) to mark it done. - handle.rpc(proto.SubmitRequest(script="#!/bin/bash", name="done_job")) - - time.sleep(0.5) # poll interval is 0.1s; 0.5s is more than enough - - job = Job(job_id="done.mock", job_name="done_job") - stream = handle.connect() - proto.send_frame(stream, proto.WaitRequest(jobs=[job])) - - response = proto.recv_frame(stream) - stream.close() - - assert isinstance(response, proto.WaitDoneResponse) - assert response.jobs[0].job_id == "done.mock" - - stop.set() - t.join(timeout=2.0) - - -# --------------------------------------------------------------------------- -# Wait — job finishes during wait -# --------------------------------------------------------------------------- - - -def test_wait_receives_status_update_then_done(server: ServerHandle, mock_get_states: MagicMock) -> None: - """ - Submit a job, start waiting, then let the poll loop (0.1s interval) detect - the job as finished and push WaitDoneResponse back to the client. - """ - server.rpc(proto.SubmitRequest(script="#!/bin/bash", name="test_job")) - - job = Job(job_id="100.mock", job_name="test_job") - stream = server.connect(timeout=5.0) - proto.send_frame(stream, proto.WaitRequest(jobs=[job])) - - # Allow the poll loop to mark the job as finished. - mock_get_states.state["value"] = None - - # Collect responses until WaitDoneResponse (or timeout). - responses = [] - try: - while True: - r = proto.recv_frame(stream) - responses.append(r) - if isinstance(r, proto.WaitDoneResponse): - break - except (EOFError, OSError): - pass - - stream.close() - - assert any(isinstance(r, proto.WaitDoneResponse) for r in responses), f"Expected WaitDoneResponse; got: {responses}" - - -# --------------------------------------------------------------------------- -# Auth tests -# --------------------------------------------------------------------------- - - -def test_auth_required_with_api_key(server_with_key: ServerHandle) -> None: - """Connecting without sending AuthRequest first should fail.""" - stream = server_with_key.connect() - # Send a PingRequest directly (no auth) - proto.send_frame(stream, proto.PingRequest()) - response = proto.recv_frame(stream) - stream.close() - assert isinstance(response, proto.ErrorResponse) - assert "Authentication failed" in response.message - - -def test_auth_wrong_key(server_with_key: ServerHandle) -> None: - """Sending the wrong API key should fail.""" - stream = server_with_key.connect() - proto.send_frame(stream, proto.AuthRequest(api_key="wrong")) - response = proto.recv_frame(stream) - stream.close() - assert isinstance(response, proto.ErrorResponse) - assert "Authentication failed" in response.message - - -def test_auth_correct_key(server_with_key: ServerHandle) -> None: - """Correct API key followed by a ping should work.""" - stream = server_with_key.connect() - proto.send_frame(stream, proto.AuthRequest(api_key="secret")) - auth_resp = proto.recv_frame(stream) - assert isinstance(auth_resp, proto.AuthOkResponse) - - proto.send_frame(stream, proto.PingRequest()) - pong = proto.recv_frame(stream) - stream.close() - assert isinstance(pong, proto.PongResponse) - - -def test_server_backend_auth_end_to_end(server_with_key: ServerHandle) -> None: - """ServerBackend with the correct api_key should connect and submit successfully.""" - from pbspy._server_backend import ServerBackend - - backend = ServerBackend(server_with_key.host, server_with_key.port, api_key="secret") - job_id, job_name = backend.submit("#!/bin/bash\necho hi", name="test_job") - backend.close() - assert job_id == "100.mock" - assert job_name == "test_job" - - -def test_server_backend_auth_wrong_key(server_with_key: ServerHandle) -> None: - """ServerBackend with a wrong api_key should raise RuntimeError on connect.""" - from pbspy._server_backend import ServerBackend - - with pytest.raises(RuntimeError, match="auth failed"): - ServerBackend(server_with_key.host, server_with_key.port, api_key="wrong") - - -# --------------------------------------------------------------------------- -# Delete -# --------------------------------------------------------------------------- - - -def test_delete_returns_delete_response(server: ServerHandle) -> None: - """DeleteRequest should call core.pbs_delete once and return DeleteResponse.""" - with patch("pbspy.server.server.core.pbs_delete") as mock_delete: - response = server.rpc(proto.DeleteRequest(job_ids=["100.mock", "101.mock"])) - assert isinstance(response, proto.DeleteResponse) - mock_delete.assert_called_once() - args, kwargs = mock_delete.call_args - assert args[0] == ["100.mock", "101.mock"] - - -def test_delete_propagates_errors(server: ServerHandle) -> None: - """If core.pbs_delete raises, the server should send back an ErrorResponse.""" - with patch("pbspy.server.server.core.pbs_delete", side_effect=RuntimeError("boom")): - response = server.rpc(proto.DeleteRequest(job_ids=["100.mock"])) - assert isinstance(response, proto.ErrorResponse) - assert "boom" in response.message - - -def test_server_backend_delete_end_to_end(server: ServerHandle) -> None: - """ServerBackend.delete() sends one DeleteRequest and returns without error.""" - from pbspy._server_backend import ServerBackend - - with patch("pbspy.server.server.core.pbs_delete") as mock_delete: - backend = ServerBackend(server.host, server.port) - backend.delete(["100.mock"]) - backend.close() - mock_delete.assert_called_once() From f3f6c0acb26becf52d3250bc0216b46fa2ac4741 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 30 Jul 2026 09:52:37 +1000 Subject: [PATCH 12/16] chore: remove outdated output docs --- README.md | 17 ----------------- 1 file changed, 17 deletions(-) diff --git a/README.md b/README.md index 4704213..6025190 100644 --- a/README.md +++ b/README.md @@ -46,23 +46,6 @@ 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. -## 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 - -job_a: A -job_b: B -``` - ## License `pbspy` is licensed under the MIT License [LICENSE](./LICENSE) or . From 3aff0662ebd38e459524a914b8e92fd2382a4a76 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 30 Jul 2026 10:05:01 +1000 Subject: [PATCH 13/16] chore: remove vestigial .env.example gitignore entry --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index bcde509..f6c4123 100644 --- a/.gitignore +++ b/.gitignore @@ -174,4 +174,3 @@ cython_debug/ *.e* *.o* -!.env.example From a73bb02b395a09337d11ef1fa8d18bfc981974d1 Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 30 Jul 2026 10:05:06 +1000 Subject: [PATCH 14/16] refactor: drop legacy progress kwargs from Job wait/result methods --- CHANGELOG.md | 6 +----- src/pbspy/__init__.py | 11 ++++------- tests/test_default.py | 13 +++---------- 3 files changed, 8 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index a4ace19..56e4b87 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,17 +25,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - 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 - - Legacy progress arguments remain accepted and are ignored while consumers migrate - Remove `rich` dependency -### Fixed - -- Update `JOBFS` limit for jobs on the `copyq` queue of the `Gadi` supercomputer - ## [0.0.9](https://github.com/MaterialsPhysicsANU/pbspy/releases/tag/v0.0.9) - 2026-05-18 ### Added diff --git a/src/pbspy/__init__.py b/src/pbspy/__init__.py index 8d19616..2824846 100644 --- a/src/pbspy/__init__.py +++ b/src/pbspy/__init__.py @@ -70,13 +70,13 @@ class Job: error_path: str | None = None """Custom path for the job's stderr error file (``#PBS -e``), if any.""" - def wait(self, **kwargs: Any) -> None: + def wait(self) -> None: """ Wait for the job to complete. """ self.backend.wait([self]) - def result(self, **kwargs: Any) -> JobResult: + def result(self) -> JobResult: """ Waits for the job to complete and returns the result. """ @@ -88,7 +88,7 @@ def cancel(self) -> None: self.backend.delete([self.job_id]) @staticmethod - def wait_all(jobs: list[Job], **kwargs: Any) -> None: + def wait_all(jobs: list[Job]) -> None: """ Waits for multiple jobs to complete. """ @@ -98,12 +98,9 @@ def wait_all(jobs: list[Job], **kwargs: Any) -> None: _wait_all_grouped(jobs) @staticmethod - def result_all(jobs: list[Job], *_args: Any, **_kwargs: Any) -> list[JobResult]: + def result_all(jobs: list[Job]) -> list[JobResult]: """ Waits for multiple jobs to complete and returns their results. - - Additional arguments are accepted for compatibility with callers that - still pass the removed ``progress`` option; they have no effect. """ Job.wait_all(jobs) return [job.backend.get_result(job) for job in jobs] diff --git a/tests/test_default.py b/tests/test_default.py index 71df64e..2a01149 100644 --- a/tests/test_default.py +++ b/tests/test_default.py @@ -45,9 +45,7 @@ def submit(self, script: str, name: str | None = None) -> tuple[str, str]: 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, progress: bool = True - ) -> None: + 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: @@ -80,7 +78,7 @@ def test_mock_backend_result() -> None: jd.add_command(["echo", "hello"]) job = jd.submit(backend=mock) - result = job.result(progress=False) + result = job.result() assert len(mock.waited) == 1 assert result.exit_code == 0 @@ -96,18 +94,13 @@ def test_mock_backend_result_all() -> None: 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], progress=False) + 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_result_all_accepts_legacy_positional_progress() -> None: - """Existing consumers may still pass the removed progress flag positionally.""" - assert Job.result_all([], False) == [] - - def test_script_generation_unchanged() -> None: """script() output is unchanged from the original implementation.""" jd = JobDescription( From b86d843424b178e228362d47f80c943ee0979abd Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 30 Jul 2026 10:05:12 +1000 Subject: [PATCH 15/16] refactor: simplify pbs_wait_for_jobs poll loop to fixed interval --- src/pbspy/_pbs_core.py | 27 +++++++++++---------------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py index ce5b5c5..8922f94 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -28,7 +28,6 @@ ] _POLL_INTERVAL_SECONDS = 60 -_PROGRESS_REFRESH_SECONDS = 1 class PBSRunner: @@ -105,22 +104,18 @@ def pbs_wait_for_jobs( from qstat (finished). """ task_done = [False] * len(jobs) - last_check = time.time() - _POLL_INTERVAL_SECONDS # poll immediately on first iteration - while not all(task_done): - time.sleep(_PROGRESS_REFRESH_SECONDS) - if time.time() - last_check >= _POLL_INTERVAL_SECONDS: - last_check = time.time() - 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) + 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]: From 8cc3fc1d4fda9db0e67f8f98069016319d6b4aaf Mon Sep 17 00:00:00 2001 From: Lachlan Deakin Date: Thu, 30 Jul 2026 10:19:34 +1000 Subject: [PATCH 16/16] fix: mypy lint --- src/pbspy/_pbs_core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pbspy/_pbs_core.py b/src/pbspy/_pbs_core.py index 8922f94..978770f 100644 --- a/src/pbspy/_pbs_core.py +++ b/src/pbspy/_pbs_core.py @@ -247,6 +247,6 @@ def format_job_script( {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 ""} +{f'#PBS -W depend=afterok:{":".join(afterok_ids)}' if afterok_ids else ""} {commands_str} """