Jobs API¶
smint.alignment.jobs runs a registration in a separate process: a local
subprocess, or an sbatch job on a SLURM cluster. The napari plugin and
smint-alignment both go through it. Keeping the worker out of process lets a
napari session (numpy ≥ 2) drive an STalign registration (numpy < 2), and lets a
failed job fail without taking the caller down.
sequenceDiagram
participant You as Your code / napari / CLI
participant Dir as output_dir
participant W as Worker (smint-register)
You->>Dir: submit(spec): job_spec.json, status.json "submitted"
alt backend = "sbatch"
You->>Dir: submit.sh, then status "queued" (+ job_id)
else backend = "local"
You->>Dir: status "running" (+ pid), worker.log
end
W->>Dir: status "running"
W->>Dir: *_transformed.csv
W->>Dir: status "completed" + outputs + metrics (or "failed" + traceback)
You->>Dir: poll(output_dir), reconciled with squeue / sacct
Submitting and polling a job¶
import time
from smint.alignment.jobs import JobSpec, SlurmResources, submit, poll
spec = JobSpec(
round="centroid",
inputs={"source_file": "nuclei_centroids.csv",
"target_file": "xenium_cells.csv"},
params={"method": "multiround",
"source_cols": ("x_centroid", "y_centroid"),
"target_cols": ("x_centroid", "y_centroid")},
output_dir="runs/centroid",
backend="local",
)
info = submit(spec) # {'backend': 'local', 'pid': ..., ...}
while (status := poll(spec.output_dir))["state"] not in ("completed", "failed", "cancelled"):
time.sleep(5)
print(status["state"], status.get("metrics", {}).get("tre_validation"))
from smint.alignment.jobs import JobSpec, SlurmResources, submit
spec = JobSpec(
round="st_sm",
inputs={"st_file": "st_cells.csv", "sm_file": "sm_matrix.csv",
"st_points": "landmarks_st_points.npy",
"sm_points": "landmarks_sm_points.npy"},
params={"dx": 30.0, "niter": 1000, "epV": 200.0, "orientation": "stalign",
"sm_kwargs": {"scale_xy": 10.0, "rotate_left_90": True}},
output_dir="runs/st_sm", # must be on storage the compute node sees
backend="sbatch",
resources=SlurmResources(partition="gpuq", gpus=1, time_limit="01:00:00"),
worker_python="/path/to/stalign-env/bin/python", # numpy < 2, STalign, torch
)
print(submit(spec)) # {'backend': 'sbatch', 'job_id': '...', ...}
params are passed as keyword arguments to
register_centroid_files
(centroid) or register_sm_to_st
(st_sm); params["output_path"] overrides the default output file name. The
files a job writes are described in
Data formats.
Job specification¶
smint.alignment.jobs.JobSpec
dataclass
¶
JobSpec(round, inputs, output_dir, params=dict(), backend='sbatch', resources=SlurmResources(), worker_python=None, created=(lambda: datetime.now().isoformat(timespec='seconds'))())
A registration job, serialisable to JSON.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
round
|
(st_sm, centroid)
|
Which registration to run. |
'st_sm'
|
inputs
|
dict
|
Input paths. For |
required |
params
|
dict
|
Keyword arguments forwarded to the registration function. |
dict()
|
output_dir
|
str
|
Directory for outputs, |
required |
backend
|
(sbatch, local)
|
How to launch the worker. |
'sbatch'
|
resources
|
SlurmResources
|
Ignored when |
SlurmResources()
|
worker_python
|
str
|
Interpreter for the worker; defaults to |
None
|
validate
¶
Check required inputs are present, exist, and are reachable by the worker.
Raises:
| Type | Description |
|---|---|
ValueError
|
If a required key is missing, or an sbatch job references node-local paths the compute node cannot see. |
FileNotFoundError
|
If a named input file does not exist. Catching this here means a typo surfaces immediately instead of after a queue wait. |
smint.alignment.jobs.SlurmResources
dataclass
¶
SlurmResources(partition='regular', cpus_per_task=8, memory='64G', time_limit='04:00:00', job_name='smint_register', gpus=0, gpu_type='A30')
SLURM resource request for a registration job.
Set gpus to request GPUs, and remember to point partition at one
that actually has them -- asking for a GPU on a CPU partition leaves the
job pending indefinitely rather than failing, which is easy to mistake for
a slow queue. check warns about that combination.
gpu_type pins which card the job gets. A bare --gres=gpu:N is a
lottery across every model the partition holds -- a GPU partition often
mixes generations (A30, A100 and P100, say) -- so a run can land on a card
generations older than the one it was timed on. Registration is validated
on the A30, which is why that is the default. Set it to None to accept
any card.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
partition
|
str
|
SLURM partition. Must offer GPUs when |
'regular'
|
cpus_per_task
|
int
|
|
8
|
memory
|
str
|
|
'64G'
|
time_limit
|
str
|
|
'04:00:00'
|
job_name
|
str
|
|
'smint_register'
|
gpus
|
int
|
GPUs to request; 0 for a CPU job. |
0
|
gpu_type
|
str or None
|
GPU model for |
'A30'
|
check
¶
Validate the resource request before anything reaches sbatch.
Raises:
| Type | Description |
|---|---|
ValueError
|
If the partition is not one this cluster offers. Catching it here turns SLURM's terse "Invalid partition name specified" into a message naming the valid options. |
sbatch_directives
¶
Render the #SBATCH lines; logs go to log_dir/<job_name>_<jobid>.out|.err.
Launching and monitoring¶
smint.alignment.jobs.submit
¶
Launch a registration job.
Validates the spec, writes it to disk, then launches the worker via the requested backend.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
spec
|
JobSpec
|
|
required |
Returns:
| Type | Description |
|---|---|
dict
|
|
Raises:
| Type | Description |
|---|---|
RuntimeError
|
If |
smint.alignment.jobs.poll
¶
Current job status, reconciling the status file with SLURM.
The worker's own status file is authoritative once it exists. SLURM is consulted to catch jobs killed before the worker could report -- an OOM or timeout leaves the status file saying "running" forever otherwise.
Returns:
| Type | Description |
|---|---|
dict
|
|
smint.alignment.jobs.read_status
¶
Read a job's status.
Returns a dict with at least state. unknown means no status file
has appeared yet -- normal while a job is queued.
smint.alignment.jobs.write_status
¶
Write the worker's status file. Called by the worker, read by the GUI.
smint.alignment.jobs.slurm_state
¶
Query SLURM for a job's state, or None if SLURM is unavailable.
Tries squeue first (job still in the queue), then sacct (finished).
Spec files and sbatch scripts¶
smint.alignment.jobs.write_spec
¶
Serialise a job spec to JSON; defaults to <output_dir>/job_spec.json.
smint.alignment.jobs.build_sbatch_script
¶
Render the sbatch script that runs the worker for spec.
Worker environment and cluster discovery¶
smint.alignment.jobs.default_worker_python
¶
Best guess at an interpreter that can run registration.
Order: SMINT_WORKER_PYTHON, then STalign_env/bin/python beside the
installed package, then the current interpreter.
smint.alignment.jobs.worker_env
¶
Return a copy of base (default os.environ) with LEAKY_ENV_VARS dropped.
smint.alignment.jobs.available_partitions
¶
Partition names this cluster accepts, via sinfo.
Returns an empty list when SLURM is unavailable, so callers can fall back rather than blocking on a check they cannot perform.
smint.alignment.jobs.available_gpu_types
¶
GPU types this cluster advertises, via sinfo.
sinfo reports gres as gpu:A30:4(S:0-1); the middle field is the
type name --gres=gpu:<type>:<n> expects. Pass partition to list
only the types that partition offers.
Returns an empty list when SLURM is unavailable, so callers can fall back rather than blocking on a check they cannot perform.
Constants¶
smint.alignment.jobs.TERMINAL_STATES
module-attribute
¶
smint.alignment.jobs.LEAKY_ENV_VARS
module-attribute
¶
The worker¶
smint.cli.register
¶
Registration worker: runs a job spec written by smint.alignment.jobs.
Invoked as python -m smint.cli.register <spec.json>, normally by sbatch or
a local subprocess rather than by hand. This is the only place STalign is
imported, which is what lets a napari front end in a numpy>=2 environment drive
a registration that requires numpy<2.
Progress is reported by rewriting <output_dir>/status.json so a GUI can
poll it without parsing logs.
Next steps¶
- Command line:
smint-alignment … --submit --watchwraps all of this. - Running on HPC: partitions, GPU types and shared storage.
- Alignment API: the functions the worker calls.