refactor: add backend support + cancellation and batched PBS query - #14
Merged
Conversation
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.
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.
Includes .env.example for SSH configuration. PBSPY_SSH_HOST and PBSPY_SSH_USER are also read from environment as argparse defaults.
This fixes a bug in the Dockerfile caused by the pbspy_server module not being installed by `hatchling`.
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.
Required so that `pbs_get_result` gets output and error files correctly if their paths are modified by the JobDescription.
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.
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.
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).
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.
# Conflicts: # src/pbspy/__init__.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This is a stripped down version of the
simple-serverbranch, retaining the backend support and onlyLocalBackend.LLM found a few issues and I think it is much better to put the actual job tracking in the prefect worker + persistent prefect variables rather than an independent process.
Added
Backendinterface and defaultLocalBackendimplementationJobDescription.submit()accepts a backend supplied by the callerJobretains its backend for waiting, result retrieval, and cancellationJob.cancel()Changed
Job.wait_all()andJob.result_all()calls by backend so each backend can wait for its jobs togetherLocalBackendfor jobs submitted without an explicit backendNonewhen PBS has not reported a completed job's exit statusFixed
JobDescription.output_pathanderror_pathon submitted jobs and use them when retrieving results