haps-plugin-scheduler-runtime 0.3.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (20) hide show
  1. haps_plugin_scheduler_runtime-0.3.0/LICENSE +21 -0
  2. haps_plugin_scheduler_runtime-0.3.0/PKG-INFO +31 -0
  3. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/__init__.py +30 -0
  4. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/config.py +183 -0
  5. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/job.py +336 -0
  6. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/metadata.py +74 -0
  7. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/profiles/__init__.py +33 -0
  8. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/profiles/base.py +207 -0
  9. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/profiles/cines.py +19 -0
  10. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/profiles/idris.py +14 -0
  11. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/profiles/tgcc.py +46 -0
  12. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/remote.py +214 -0
  13. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime/states.py +145 -0
  14. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime.egg-info/PKG-INFO +31 -0
  15. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime.egg-info/SOURCES.txt +18 -0
  16. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime.egg-info/dependency_links.txt +1 -0
  17. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime.egg-info/requires.txt +6 -0
  18. haps_plugin_scheduler_runtime-0.3.0/haps_plugin_scheduler_runtime.egg-info/top_level.txt +1 -0
  19. haps_plugin_scheduler_runtime-0.3.0/pyproject.toml +59 -0
  20. haps_plugin_scheduler_runtime-0.3.0/setup.cfg +4 -0
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Benjamin De Zordo
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,31 @@
1
+ Metadata-Version: 2.4
2
+ Name: haps-plugin-scheduler-runtime
3
+ Version: 0.3.0
4
+ Summary: Run HAPS jobs on a French HPC center by submitting them to its scheduler over SSH
5
+ Author: Logica Team, Universite de Rennes / IRISA
6
+ Author-email: Benjamin De Zordo <benjamin.de-zordo@irisa.fr>
7
+ License-Expression: MIT
8
+ Project-URL: Homepage, https://www.irisa.fr/equipes/logica
9
+ Keywords: haps,hpc,plugin,distributed-computing
10
+ Classifier: Development Status :: 3 - Alpha
11
+ Classifier: Intended Audience :: Science/Research
12
+ Classifier: Intended Audience :: System Administrators
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: Operating System :: OS Independent
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3 :: Only
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Scientific/Engineering
20
+ Classifier: Topic :: System :: Clustering
21
+ Classifier: Topic :: System :: Distributed Computing
22
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
23
+ Classifier: Typing :: Typed
24
+ Requires-Python: >=3.11
25
+ License-File: LICENSE
26
+ Requires-Dist: haps-runtime==0.3.*
27
+ Requires-Dist: haps-runtime-executor==0.3.*
28
+ Requires-Dist: remotemanager
29
+ Provides-Extra: dev
30
+ Requires-Dist: ruff>=0.6; extra == "dev"
31
+ Dynamic: license-file
@@ -0,0 +1,30 @@
1
+ from .config import BatchResources, HpcConfig, SchedulerConfig
2
+ from .job import SchedulerJob
3
+ from .metadata import RuntimeInfo
4
+ from .profiles import (
5
+ CinesProfile,
6
+ IdrisProfile,
7
+ ProfileError,
8
+ SchedulerProfile,
9
+ TgccProfile,
10
+ )
11
+ from .remote import CommandResult, RemoteError, RemoteExecutor
12
+
13
+ __version__ = "0.3.0"
14
+
15
+ __all__ = [
16
+ "RuntimeInfo",
17
+ "BatchResources",
18
+ "CinesProfile",
19
+ "CommandResult",
20
+ "HpcConfig",
21
+ "IdrisProfile",
22
+ "ProfileError",
23
+ "RemoteError",
24
+ "RemoteExecutor",
25
+ "SchedulerConfig",
26
+ "SchedulerJob",
27
+ "SchedulerProfile",
28
+ "TgccProfile",
29
+ "__version__",
30
+ ]
@@ -0,0 +1,183 @@
1
+ from __future__ import annotations
2
+
3
+ import stat
4
+ import tomllib
5
+ from pathlib import Path
6
+ from typing import Any, ClassVar
7
+
8
+ from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
9
+
10
+ from .profiles import ProfileError, SchedulerProfile, to_seconds
11
+ from .remote import SAFE_ARGUMENT, check_argument
12
+
13
+
14
+ class HpcConfig(BaseModel):
15
+ """The HPC centre, the account used on it, and how to reach it."""
16
+
17
+ model_config = ConfigDict(extra="forbid")
18
+
19
+ host: str
20
+ user: str
21
+ workdir_root: str
22
+
23
+ passfile: Path | None = Field(default=None, repr=False)
24
+ keyfile: Path | None = Field(default=None, repr=False)
25
+
26
+ quick_timeout: int = 60
27
+ submit_timeout: int = 120
28
+ poll_interval: float = 60.0
29
+
30
+ @field_validator("host", "user")
31
+ @classmethod
32
+ def _sendable(cls, value: str) -> str:
33
+ return check_argument(value)
34
+
35
+ @field_validator("workdir_root")
36
+ @classmethod
37
+ def _sendable_path(cls, value: str) -> str:
38
+ """Refuse a work directory path the transport could not send safely."""
39
+ rest = value
40
+ if value.startswith("$"):
41
+ head, _, tail = value[1:].partition("/")
42
+ if not SAFE_ARGUMENT.match(head.strip("{}")):
43
+ raise ValueError(f"{value!r}: the variable name must be a plain name")
44
+ rest = tail
45
+ if rest and not SAFE_ARGUMENT.match(rest):
46
+ raise ValueError(
47
+ f"{value!r}: a path sent to the centre may only hold letters, digits and _@%+=:,./-"
48
+ )
49
+ return value
50
+
51
+ @field_validator("passfile", "keyfile")
52
+ @classmethod
53
+ def _readable_by_nobody_else(cls, value: Path | None) -> Path | None:
54
+ if value is None:
55
+ return None
56
+ path = value.expanduser()
57
+ if not path.is_file():
58
+ raise ValueError(f"file not found: {path}")
59
+ if stat.S_IMODE(path.stat().st_mode) & 0o077:
60
+ raise ValueError(f"{path} is readable by others - chmod 600 it before starting")
61
+ return path
62
+
63
+
64
+ class BatchResources(BaseModel):
65
+ """What to ask the scheduler for, or the script on the centre that already asks."""
66
+
67
+ model_config = ConfigDict(extra="forbid")
68
+
69
+ script_path: str | None = None
70
+
71
+ account: str | None = None
72
+ target: str | None = None
73
+ qos: str | None = None
74
+ time: str | int | None = None
75
+ nodes: int | None = None
76
+ ntasks: int | None = None
77
+ ntasks_per_node: int | None = None
78
+ cpus_per_task: int | None = None
79
+ gpus_per_node: int | None = None
80
+ hint: str | None = None
81
+ threads_per_core: int | None = None
82
+ filesystems: list[str] = Field(default_factory=list)
83
+
84
+ modules: list[str] = Field(default_factory=list)
85
+ module_purge: bool = False
86
+ environment: dict[str, str] = Field(default_factory=dict)
87
+
88
+ #: Fields a profile may turn into a directive. The rest describe the script
89
+ #: body, and `directives` must not leak them into the header.
90
+ DIRECTIVE_FIELDS: ClassVar[tuple[str, ...]] = (
91
+ "account",
92
+ "target",
93
+ "qos",
94
+ "time",
95
+ "nodes",
96
+ "ntasks",
97
+ "ntasks_per_node",
98
+ "cpus_per_task",
99
+ "gpus_per_node",
100
+ "hint",
101
+ "threads_per_core",
102
+ "filesystems",
103
+ )
104
+
105
+ @field_validator("script_path")
106
+ @classmethod
107
+ def _sendable(cls, value: str | None) -> str | None:
108
+ return None if value is None else check_argument(value)
109
+
110
+ @model_validator(mode="after")
111
+ def _one_mode_only(self) -> BatchResources:
112
+ if self.script_path is None:
113
+ return self
114
+ set_anyway = [
115
+ name
116
+ for name in (*self.DIRECTIVE_FIELDS, "modules", "module_purge", "environment")
117
+ if getattr(self, name) not in (None, False, [], {})
118
+ ]
119
+ if set_anyway:
120
+ raise ValueError(
121
+ "script_path means the script on the centre already carries its own "
122
+ f"directives; remove {', '.join(set_anyway)} from [resources]"
123
+ )
124
+ return self
125
+
126
+ def directives(self) -> dict[str, Any]:
127
+ """The resource fields that are set, for a profile to turn into directives."""
128
+ return {
129
+ name: getattr(self, name)
130
+ for name in self.DIRECTIVE_FIELDS
131
+ if getattr(self, name) not in (None, False, [], "")
132
+ }
133
+
134
+
135
+ class SchedulerConfig(BaseModel):
136
+ """Everything the maintainer decides, read once at startup and never changed."""
137
+
138
+ # arbitrary_types_allowed because `profile` carries a class, not TOML data.
139
+ model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
140
+
141
+ hpc: HpcConfig
142
+ resources: BatchResources
143
+ profile: type[SchedulerProfile]
144
+
145
+ @field_validator("profile")
146
+ @classmethod
147
+ def _is_a_profile_class(cls, value: Any) -> type[SchedulerProfile]:
148
+ if not (isinstance(value, type) and issubclass(value, SchedulerProfile)):
149
+ raise ValueError(
150
+ "profile must be a SchedulerProfile subclass (the class itself, not an instance)"
151
+ )
152
+ return value
153
+
154
+ @model_validator(mode="after")
155
+ def _profile_accepts_these_resources(self) -> SchedulerConfig:
156
+ """Refuse at startup the resources the profile would reject on every job."""
157
+ if self.resources.script_path is not None:
158
+ return self # the script on the centre carries its own header
159
+ directives = self.resources.directives()
160
+ missing = [
161
+ name
162
+ for name in self.profile.REQUIRED
163
+ # A REQUIRED name outside DIRECTIVE_FIELDS is derived by the profile
164
+ # from the job itself - job_name, chdir, output - and is never absent.
165
+ if name in BatchResources.DIRECTIVE_FIELDS and name not in directives
166
+ ]
167
+ if missing:
168
+ raise ValueError(
169
+ f"{self.profile.__name__} requires {', '.join(missing)} in [resources]"
170
+ )
171
+ if self.resources.time is not None:
172
+ try:
173
+ to_seconds(self.resources.time)
174
+ except ProfileError as exc:
175
+ raise ValueError(str(exc)) from exc
176
+ return self
177
+
178
+ @classmethod
179
+ def from_toml(cls, path: str | Path, *, profile: type[SchedulerProfile]) -> SchedulerConfig:
180
+ """Read a TOML file and validate it into a configuration."""
181
+ with open(path, "rb") as handle:
182
+ document = tomllib.load(handle)
183
+ return cls(profile=profile, **document)
@@ -0,0 +1,336 @@
1
+ from __future__ import annotations
2
+
3
+ import logging
4
+ from dataclasses import replace
5
+ from pathlib import Path, PurePosixPath
6
+ from tempfile import TemporaryDirectory
7
+ from typing import TYPE_CHECKING, Any
8
+
9
+ from haps_runtime_executor import RuntimeJob, RuntimeJobError, RuntimeServiceError
10
+
11
+ from .metadata import NO_TRACE, RuntimeInfo
12
+ from .profiles import ProfileError, SchedulerProfile
13
+ from .remote import RemoteError, RemoteExecutor
14
+ from .states import (
15
+ LOST,
16
+ SACCT_FIELDS,
17
+ SQUEUE_FIELDS,
18
+ is_terminal,
19
+ reason,
20
+ slurm_infos,
21
+ succeeded,
22
+ )
23
+
24
+ if TYPE_CHECKING:
25
+ from collections.abc import Callable
26
+
27
+ logger = logging.getLogger(__name__)
28
+
29
+ #: How long the scheduler's stdout and stderr stay readable after a job ends.
30
+ TRACE_LIFESPAN = 12 * 3600
31
+
32
+
33
+ class SchedulerJob(RuntimeJob):
34
+ """One HAPS job, run as one scheduler job on an HPC centre."""
35
+
36
+ def __init__(self, **kwargs) -> None:
37
+ super().__init__(**kwargs)
38
+ # The centre, as an instance: the configuration carries the class (D-09)
39
+ # and render / submit_argv / trace_paths are instance methods.
40
+ self.profile: SchedulerProfile = self.config.profile()
41
+ self.remote: RemoteExecutor | None = None
42
+ self.slurm_job_id: str | None = None
43
+ # This job's own directory on the centre.
44
+ # example: ``workdir_root / job_uuid``.
45
+ self.workdir: PurePosixPath | None = None
46
+ self.script: PurePosixPath | None = None
47
+ self.last_state: str | None = None
48
+ self._cancel_sent = False
49
+ # What this job has published so far. Rewritten whole on every call, so
50
+ # nothing depends on how deeply the server merges an update.
51
+ self.published = RuntimeInfo()
52
+
53
+ def __repr__(self) -> str:
54
+ return f"<{type(self).__name__} job={self.job_uuid} slurm={self.slurm_job_id or '-'}>"
55
+
56
+ # --- The sequence --------------------------------------------------------
57
+
58
+ def steps(self) -> tuple[tuple[str, Callable[[], None]], ...]:
59
+ """The ordered sequence of steps, which a subclass may reorder."""
60
+ return (
61
+ ("connect", self.connect),
62
+ ("setup directories", self.setup_directories),
63
+ ("fetch arguments", self.fetch_arguments),
64
+ ("fetch ebuffer inputs", self.fetch_ebinputs),
65
+ ("stage script", self.stage_script),
66
+ ("execute", self.execute),
67
+ ("push results", self.push_results),
68
+ ("push ebuffer outputs", self.push_eboutputs),
69
+ )
70
+
71
+ # --- Getting there -------------------------------------------------------
72
+
73
+ def connect(self) -> None:
74
+ """Open the SSH connection to the centre."""
75
+ executor = RemoteExecutor(self.config.hpc)
76
+ try:
77
+ executor.connect()
78
+ except RemoteError as exc:
79
+ raise RuntimeServiceError(f"HPC centre unreachable: {exc}") from exc
80
+ self.remote = executor
81
+
82
+ def setup_directories(self, path: Path | str | None = None) -> None:
83
+ """Create this job's own directory on the centre, named after its uuid."""
84
+ try:
85
+ root = path if path else self.remote.expand(self.config.hpc.workdir_root)
86
+ self.workdir = root / self.job_uuid
87
+ self.remote.mkdir(self.workdir)
88
+ except RemoteError as exc:
89
+ raise RuntimeJobError(
90
+ "the working directory could not be created on the centre",
91
+ detail=str(exc),
92
+ ) from exc
93
+
94
+ def stage_script(self) -> None:
95
+ """Write the batch script on the centre, or point at the one already there."""
96
+ resources = self.config.resources
97
+
98
+ # Slurm script already exists on HPC centre (arg script_path set)
99
+ if resources.script_path is not None:
100
+ self.script = PurePosixPath(resources.script_path)
101
+ if not self.remote.exists(self.script):
102
+ raise RuntimeJobError(
103
+ f"the configured script does not exist on the centre: {self.script}",
104
+ detail="[resources] script_path",
105
+ )
106
+ return
107
+
108
+ # There is no script, build slurm script
109
+ try:
110
+ text = self.profile.render(
111
+ resources,
112
+ job_uuid=self.job_uuid,
113
+ workdir=self.workdir,
114
+ core=self.script_core(),
115
+ )
116
+ except ProfileError as exc:
117
+ # The profile knows what is missing, not that a submitter is waiting
118
+ # for an answer. Translated here, where the blast radius is known -
119
+ # the same rule that sends `connect` to RuntimeServiceError and
120
+ # `mkdir` to RuntimeJobError from the one RemoteError.
121
+ raise RuntimeJobError(f"the batch script cannot be built: {exc}") from exc
122
+ self.script = self.workdir / "haps_job.sh"
123
+ try:
124
+ self.remote.write_text(self.script, text)
125
+ except RemoteError as exc:
126
+ raise RuntimeJobError(
127
+ "the batch script could not be written on the centre", detail=str(exc)
128
+ ) from exc
129
+
130
+ # Nothing a submitter sent belongs in this text: with
131
+ # arguments = ["data.in; rm -rf $WORK"] the second half would run under the
132
+ # maintainer's service account. Write such values to a file in self.workdir
133
+ # and have the script read it. Use self.profile.launch_cmd rather than a
134
+ # hard-coded srun, and the same class runs on the three centres.
135
+ def script_core(self) -> str:
136
+ """The core of the batch script, which every subclass must write."""
137
+ raise NotImplementedError(
138
+ f"{type(self).__name__} must implement script_core(): what the batch "
139
+ "script runs on the compute nodes"
140
+ )
141
+
142
+ # --- Doing the work ------------------------------------------------------
143
+
144
+ def execute(self) -> None:
145
+ """Submit the job, then watch it until the scheduler says it is over."""
146
+ self.slurm_job_id = self.submit_slurm_job()
147
+ self.publish(slurm={"job_id": self.slurm_job_id})
148
+ state: str | None = None
149
+ while True:
150
+ found = self.get_slurm_job_infos()
151
+ state = found.get("state")
152
+ if state is not None:
153
+ if state != self.last_state:
154
+ self.last_state = state
155
+ self.publish(slurm={"job_id": self.slurm_job_id, **found})
156
+ if is_terminal(state):
157
+ break
158
+ if self.wait(self.config.hpc.poll_interval):
159
+ return
160
+ if not succeeded(state):
161
+ raise RuntimeJobError(
162
+ f"HPC job {self.slurm_job_id} ended in {state}: {reason(state)} - "
163
+ f"SchedulerClient.download_slurm_trace() has the output, for "
164
+ f"{TRACE_LIFESPAN // 3600} h",
165
+ detail=f"scheduler job {self.slurm_job_id} state={state}",
166
+ )
167
+
168
+ def submit_slurm_job(self) -> str:
169
+ try:
170
+ result = self.remote.run(
171
+ self.profile.submit_argv(self.script),
172
+ timeout=self.config.hpc.submit_timeout,
173
+ )
174
+ except RemoteError as exc:
175
+ raise RuntimeJobError(
176
+ "the job could not be submitted to the scheduler", detail=str(exc)
177
+ ) from exc
178
+ if not result.ok:
179
+ raise RuntimeJobError(
180
+ "the scheduler refused this job - check the resources asked for",
181
+ detail=f"{self.profile.submit_cmd} rc={result.returncode}: {result.stderr}",
182
+ )
183
+ try:
184
+ slurm_job_id = self.profile.parse_submit(result.stdout)
185
+ except ProfileError as exc:
186
+ # Accepted, then unreadable: the job may well be running and nothing
187
+ # can be cancelled without its id. The output goes to `detail` - it
188
+ # is the maintainer's clue and means nothing to a submitter.
189
+ raise RuntimeJobError(
190
+ "the scheduler accepted the job but printed no job id",
191
+ detail=str(exc),
192
+ ) from exc
193
+ logger.info("%s | submitted as scheduler job %s", self.log_head, slurm_job_id)
194
+ return slurm_job_id
195
+
196
+ def get_slurm_job_infos(self) -> dict[str, str]:
197
+ """Ask the scheduler about this job, once. Empty when we could not ask."""
198
+ # Try a squeue to slurm
199
+ try:
200
+ queued = self.remote.run(self.profile.queue_argv(self.slurm_job_id))
201
+ except RemoteError as exc:
202
+ logger.warning("%s | queue not readable, will ask again: %s", self.log_head, exc)
203
+ return {}
204
+ if queued.ok:
205
+ found = slurm_infos(queued.stdout, SQUEUE_FIELDS)
206
+ if found:
207
+ return found
208
+ # Try a sacct to slurm
209
+ try:
210
+ acct = self.remote.run(self.profile.acct_argv(self.slurm_job_id))
211
+ except RemoteError as exc:
212
+ logger.warning("%s | accounting not readable, will ask again: %s", self.log_head, exc)
213
+ return {}
214
+ if not acct.ok:
215
+ return {}
216
+ # An empty answer here is a verdict, not a silence: accounting replied and
217
+ # knows no such job.
218
+ return slurm_infos(acct.stdout, SACCT_FIELDS) or {"state": LOST}
219
+
220
+ # --- Stopping ------------------------------------------------------------
221
+
222
+ def on_stop(self) -> None:
223
+ """Cancel the scheduler job when this one is stopped or cancelled."""
224
+ self.cancel_slurm_job()
225
+
226
+ def cancel_slurm_job(self) -> None:
227
+ """Ask the scheduler to stop this job, at most once."""
228
+ if self.remote is None or self.slurm_job_id is None:
229
+ return # nothing was submitted, nothing to cancel
230
+ if self._cancel_sent:
231
+ # `on_stop` runs before the terminal status, `cleanup` after: on a
232
+ # stop both would fire. A second scancel answers "invalid job id"
233
+ # and would log an error about a cancellation that worked.
234
+ return
235
+ self._cancel_sent = True
236
+ logger.info("%s | cancelling scheduler job %s", self.log_head, self.slurm_job_id)
237
+ try:
238
+ result = self.remote.run(self.profile.cancel_argv(self.slurm_job_id))
239
+ except RemoteError as exc:
240
+ logger.error("%s | %s not cancelled: %s", self.log_head, self.slurm_job_id, exc)
241
+ return
242
+ if not result.ok:
243
+ logger.error(
244
+ "%s | %s not cancelled: %s", self.log_head, self.slurm_job_id, result.stderr
245
+ )
246
+
247
+ def cleanup(self) -> None:
248
+ """Stop what is left running, save the trace, clear the centre, disconnect."""
249
+ if self.remote is None:
250
+ return
251
+ try:
252
+ # Anything that is not a verdict from the scheduler leaves the
253
+ # calculation running: an unexpected error, a refused metadata write,
254
+ # a failure in a later step. It would burn quota to the walltime for
255
+ # a result nobody collects, and `remove_workdir` below would pull the
256
+ # ground from under it. Cancelling first makes both go away.
257
+ if not is_terminal(self.last_state):
258
+ self.cancel_slurm_job()
259
+ self.push_slurm_trace()
260
+ self.remove_workdir()
261
+ finally:
262
+ self.remote.close()
263
+
264
+ def publish(self, **changes: Any) -> None:
265
+ """Add to what this job recorded, and rewrite the whole runtime section."""
266
+ self.published = replace(self.published, **changes)
267
+ self.set_runtime_metadata(self.published.to_section())
268
+
269
+ def push_slurm_trace(self) -> None:
270
+ """Put the scheduler's stdout and stderr in a buffer each and record where."""
271
+ # `script_path` set means the header is the maintainer's, so we do not
272
+ # know where the scheduler was told to write - probing our own naming
273
+ # would report "absent" when the truth is "we do not know where".
274
+ if self.config.resources.script_path is not None:
275
+ return
276
+ if self.slurm_job_id is None:
277
+ return
278
+ out, err = self.profile.trace_paths(self.workdir, self.job_uuid, self.slurm_job_id)
279
+ recorded: dict[str, str] = {}
280
+ for name, path in (("stdout_buffer", out), ("stderr_buffer", err)):
281
+ try:
282
+ recorded[name] = self._push_to_buffer(path)
283
+ except Exception:
284
+ logger.exception("%s | %s not saved", self.log_head, path.name)
285
+ if recorded:
286
+ try:
287
+ self.publish(**recorded)
288
+ except Exception:
289
+ logger.exception("%s | log buffers not recorded", self.log_head)
290
+ # The maintainer's log carries nothing of what the calculation
291
+ # printed; this is what points at the buffers that do - while they
292
+ # live, which is TRACE_LIFESPAN and not a minute longer.
293
+ logger.info("%s | trace saved: %s", self.log_head, recorded)
294
+
295
+ def _push_to_buffer(self, path: PurePosixPath) -> str:
296
+ """Copy one remote file into a buffer sized to it and return its uuid."""
297
+ with TemporaryDirectory() as folder:
298
+ local = Path(folder)
299
+ if self.remote.size_of(path) in (None, 0):
300
+ return NO_TRACE
301
+ self.remote.get(path, local)
302
+ staged = local / path.name
303
+ size = staged.stat().st_size
304
+ if size == 0:
305
+ return NO_TRACE
306
+ buffer = self.client.buffers.create(
307
+ reserved_size=size,
308
+ lifespan=TRACE_LIFESPAN,
309
+ tags=["haps-scheduler:trace", f"job::{self.job_uuid}"],
310
+ )
311
+ buffer.fill(staged)
312
+ # A reservation above the server's cap is trimmed in silence; only
313
+ # reading it back says so.
314
+ if buffer.reserved_size is not None and buffer.reserved_size < size:
315
+ logger.warning(
316
+ "%s | %s truncated by the server: %d asked, %d given",
317
+ self.log_head,
318
+ path.name,
319
+ size,
320
+ buffer.reserved_size,
321
+ )
322
+ return buffer.uuid
323
+
324
+ def remove_workdir(self) -> None:
325
+ """Delete this job's directory on the centre."""
326
+ if self.remote is None or self.workdir is None:
327
+ return
328
+ # A guard on the one destructive call of this plugin: the directory is
329
+ # named after the job, and nothing else may be removed by accident.
330
+ if self.job_uuid not in str(self.workdir):
331
+ logger.error("%s | refusing to remove %s", self.log_head, self.workdir)
332
+ return
333
+ try:
334
+ self.remote.remove_tree(self.workdir)
335
+ except RemoteError as exc:
336
+ logger.error("%s | %s not removed: %s", self.log_head, self.workdir, exc)
@@ -0,0 +1,74 @@
1
+ from __future__ import annotations
2
+
3
+ from dataclasses import asdict, dataclass, field, fields
4
+ from datetime import datetime
5
+ from typing import Any
6
+
7
+ #: The server keeps its own event log beside what the writers post, and puts
8
+ #: everything they write under this key. Measured, not assumed: ``extra_metadata``
9
+ #: comes back as {"events": [...], "user_data": {"runtime": {...}, "user": {...}}}.
10
+ USER_DATA = "user_data"
11
+ RUNTIME = "runtime"
12
+ NO_TRACE = "absent"
13
+
14
+ #: What a microservice must carry to be usable with this plugin, and what the
15
+ #: client half copies onto every job it submits. Tags can never be removed, so a
16
+ #: marker posed here cannot be erased by mistake.
17
+ PLUGIN_TAG = "plugin::scheduler"
18
+
19
+
20
+ @dataclass(frozen=True)
21
+ class RuntimeInfo:
22
+ """What a scheduler runtime records on a job, and the contract between the two halves."""
23
+
24
+ # What the scheduler said, under its own key: job_id, state, and whatever
25
+ # states.py is told to collect. A plain dict on purpose - adding a field there
26
+ # must not mean editing the contract here.
27
+ slurm: dict[str, Any] = field(default_factory=dict)
28
+ #: uuid of an ebuffer, or NO_TRACE. HAPS objects, not scheduler facts, so they
29
+ #: stay one level up.
30
+ stdout_buffer: str | None = None
31
+ stderr_buffer: str | None = None
32
+
33
+ def to_section(self) -> dict[str, Any]:
34
+ """What is set, as the sub-document to hand to set_runtime_metadata."""
35
+ return {k: v for k, v in asdict(self).items() if v not in (None, {})}
36
+
37
+ # Reads only the fields declared above and ignores everything else, so a job
38
+ # run by another runtime plugin - which writes its own keys under the same
39
+ # section - comes back empty rather than raising.
40
+ @classmethod
41
+ def read(cls, extra_metadata: Any) -> RuntimeInfo:
42
+ """What the runtime recorded on this job, empty when there is nothing."""
43
+ section = cls._section(extra_metadata)
44
+ known = {f.name for f in fields(cls)}
45
+ found = {k: v for k, v in section.items() if k in known and v is not None}
46
+ if not isinstance(found.get("slurm"), dict):
47
+ found.pop("slurm", None)
48
+ return cls(**found)
49
+
50
+ @staticmethod
51
+ def _section(extra_metadata: Any) -> dict[str, Any]:
52
+ section = written_by(extra_metadata).get(RUNTIME)
53
+ return section if isinstance(section, dict) else {}
54
+
55
+
56
+
57
+ def from_utc(text: str | None) -> datetime | None:
58
+ """Read an instant the runtime wrote, or None when it is absent or damaged."""
59
+ if not text:
60
+ return None
61
+ try:
62
+ return datetime.fromisoformat(text)
63
+ except ValueError:
64
+ return None
65
+
66
+
67
+ def written_by(extra_metadata: Any) -> dict[str, Any]:
68
+ """The document the writers own, under the server's own event log."""
69
+ # Falls back to the top level so a deployment that does not wrap keeps
70
+ # working: one lookup either way, and no version to keep track of.
71
+ if not isinstance(extra_metadata, dict):
72
+ return {}
73
+ inner = extra_metadata.get(USER_DATA)
74
+ return inner if isinstance(inner, dict) else extra_metadata