simrun 0.1.0__py3-none-any.whl

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.
simrun/service/app.py ADDED
@@ -0,0 +1,1721 @@
1
+ """The API of the service: people and agents submit, follow and compare runs.
2
+
3
+ Every route of the API needs a bearer token of the tokens file
4
+ (`simrun.service.auth`); only the files of the web app at `/ui/` (setting `web`)
5
+ answer without one. The identity of the token decides the priority of its runs. A
6
+ submission is validated, its resources resolved and its tasks planned on the nodes
7
+ of the backend, and the record (`submitting`) and the settings of every task
8
+ written; `POST /runs` returns then, and the environment, the image and the workflow
9
+ follow in the background (`simrun.service.submission`). The controller
10
+ (`simrun.service.controller`) runs the submissions and follows the workflows in a
11
+ background task.
12
+
13
+ Run it with `uvicorn --factory simrun.service.app:serve --timeout-graceful-shutdown 5`,
14
+ which reads the `ServiceSettings` from the YAML file named by `SIMRUN_CONFIG`; the
15
+ timeout ends the open event streams when the service stops.
16
+ """
17
+
18
+ import asyncio
19
+ import dataclasses
20
+ import json
21
+ import logging
22
+ import os
23
+ import shutil
24
+ import tempfile
25
+ import time
26
+ from collections.abc import AsyncGenerator, Callable, Iterable
27
+ from contextlib import asynccontextmanager, suppress
28
+ from datetime import UTC, datetime, timedelta
29
+ from pathlib import Path, PurePosixPath
30
+ from typing import Annotated, Any, Literal
31
+
32
+ import yaml
33
+ from fastapi import Body, Depends, FastAPI, HTTPException, Query, Request
34
+ from fastapi.responses import PlainTextResponse, StreamingResponse
35
+ from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
36
+ from fastapi.staticfiles import StaticFiles
37
+ from pydantic import (
38
+ AliasChoices,
39
+ BaseModel,
40
+ Field,
41
+ TypeAdapter,
42
+ ValidationError,
43
+ )
44
+
45
+ from simrun.argo import workflow_name
46
+ from simrun.config import ClusterConfig
47
+ from simrun.ledger import Ledger, LedgerCounts
48
+ from simrun.planner import Node, PlanError, TaskPlan, plan_task
49
+ from simrun.plugins import TaskType, load_task_types, resolve
50
+ from simrun.plugins_builtin import COMMAND
51
+ from simrun.resources import ResolvedResources, ResourceStore, resource_key
52
+ from simrun.service.auth import Identity, authenticate, load_tokens, priority_of
53
+ from simrun.service.backend import Backend, Placement
54
+ from simrun.service.compare import CompareTable, compare_runs
55
+ from simrun.service.controller import (
56
+ FINAL_TASK_STATES,
57
+ TERMINAL_RUN_STATES,
58
+ Controller,
59
+ attempt_of,
60
+ is_oom,
61
+ )
62
+ from simrun.service.edits import (
63
+ EDIT_WAIT_S,
64
+ EditRequest,
65
+ EditResult,
66
+ EditSession,
67
+ EditState,
68
+ EditStore,
69
+ OpenEdit,
70
+ RunBody,
71
+ SaveBody,
72
+ )
73
+ from simrun.service.index import IndexedRunStore, RunIndex
74
+ from simrun.service.logs import DEFAULT_TAIL_LINES, read_kept, render
75
+ from simrun.service.submission import Submitter, has_model_project
76
+ from simrun.snapshot import SourceProvenance
77
+ from simrun.store import RunRecord, RunState, RunSummary, TaskState
78
+ from simrun.worker import (
79
+ MERGED_MARKER,
80
+ PROGRESS_FILE,
81
+ PROVENANCE_FILE,
82
+ SETTINGS_FILE,
83
+ UNITS_DIR,
84
+ WORKERS_DIR,
85
+ WorkerReport,
86
+ )
87
+ from simrun.workflow import Priority, SplitMode, Workflow, workflow_yaml
88
+
89
+ logger = logging.getLogger(__name__)
90
+
91
+ CONFIG_ENV = "SIMRUN_CONFIG"
92
+ WAIT_POLL_S = 0.2
93
+ WAIT_MAX_S = 300.0
94
+ # the longest long poll of the command of an edit session
95
+ EDIT_POLL_MAX_S = 30.0
96
+ NOT_CONNECTED = "the machine of the session is not connected"
97
+ # a snapshot younger than this is kept by a cleanup: a client may be submitting it
98
+ SNAPSHOT_GRACE_S = 3600.0
99
+ LATEST_LINK = "latest"
100
+ # an ended run is moved to `runs/<prefix><id>-<pid>` under the lock, then removed
101
+ TRASH_PREFIX = ".trash-"
102
+ PROGRESS_TAIL_BYTES = 8192
103
+ # the most a task which is not over counts towards the progress of its run
104
+ UNFINISHED_TASK_MAX = 0.99
105
+
106
+
107
+ class ServiceSettings(BaseModel):
108
+ """The settings of the service.
109
+
110
+ Attributes:
111
+ store: Root of the run store.
112
+ db: SQLite file of the resource observations and the index of the runs, on
113
+ a local volume.
114
+ tokens: YAML file of the tokens, see `simrun.service.auth.load_tokens`.
115
+ owner: The person who owns the cluster; their runs have the top priority.
116
+ cluster: The cluster.
117
+ poll_s: Seconds between two reconciliations of the controller, and
118
+ between two checks of an event stream.
119
+ events_max_s: Seconds after which an event stream ends and the client
120
+ reconnects.
121
+ default_image: Image of a run whose model has no `pyproject.toml`.
122
+ simrun_wheel: The simrun wheel the images of the environments install.
123
+ python: Python version of the images.
124
+ build_constraints: The pinned versions of the build backends of the
125
+ environments, the file of simrun (`simrun.environment.BUILD_CONSTRAINTS`)
126
+ by default.
127
+ backend: `module:callable` returning the backend, for `serve`.
128
+ web: Directory of the built web app, served at `/ui/` without a token; the
129
+ app itself talks to the API with the token of the user.
130
+ """
131
+
132
+ store: Path
133
+ db: Path
134
+ tokens: Path
135
+ owner: str
136
+ cluster: ClusterConfig
137
+ poll_s: float = 5.0
138
+ events_max_s: float = 300.0
139
+ default_image: str | None = None
140
+ simrun_wheel: Path | None = None
141
+ python: str = "3.14"
142
+ build_constraints: Path | None = None
143
+ backend: str | None = None
144
+ web: Path | None = None
145
+
146
+
147
+ class SubmitRequest(BaseModel):
148
+ """A submission of a workflow.
149
+
150
+ Attributes:
151
+ workflow: The workflow.
152
+ sources: The snapshots of the sources, the model first.
153
+ group: Group of the run, e.g. a study.
154
+ split: Split mode of every task, else the one of the task or the workflow.
155
+ reserve: A node reserved for the run, else the one of the workflow.
156
+ """
157
+
158
+ workflow: Workflow
159
+ sources: list[SourceProvenance] = Field(default_factory=list)
160
+ group: str | None = None
161
+ split: SplitMode | None = None
162
+ reserve: str | None = None
163
+
164
+
165
+ class ValidationResult(BaseModel):
166
+ """The result of a validation: the errors of a workflow, none when valid."""
167
+
168
+ ok: bool
169
+ errors: list[str]
170
+
171
+
172
+ class TaskEstimate(BaseModel):
173
+ """The resources, the plan and the expected duration of a task.
174
+
175
+ Attributes:
176
+ resources: The resolved resources.
177
+ plan: The plan of the pods.
178
+ seconds: The expected seconds of a split task: the median time of a unit in
179
+ one process of past runs times the rounds of its units over the
180
+ processes of its pods, None without past runs.
181
+ """
182
+
183
+ resources: ResolvedResources
184
+ plan: TaskPlan
185
+ seconds: float | None
186
+
187
+
188
+ class TaskProgress(BaseModel):
189
+ """The progress of a task: its ledger and the first and last lines of its progress file.
190
+
191
+ The first line is the first unit done, whose time is when the task began to
192
+ deliver units.
193
+ """
194
+
195
+ counts: LedgerCounts | None
196
+ # a default: a client validates the answers of a service not yet redeployed
197
+ first: dict[str, Any] | None = None
198
+ last: dict[str, Any] | None
199
+
200
+
201
+ class RunItem(RunSummary):
202
+ """A run of the list of runs: its summary with its owner and workflow name.
203
+
204
+ Attributes:
205
+ owner: The owner of the run, of an agent the person it acts for.
206
+ workflow: The name of the workflow.
207
+ progress: The progress of an active run (submitting, queued, running) from
208
+ 0 to 1, from the states of its tasks and their ledgers (`run_progress`);
209
+ None for a run which ended.
210
+ """
211
+
212
+ owner: str
213
+ workflow: str
214
+ progress: float | None = None
215
+
216
+
217
+ class PodUsage(BaseModel):
218
+ """A pod of a run on a server.
219
+
220
+ Attributes:
221
+ run_id: The run of the pod, None when the pod belongs to no known run.
222
+ task_id: The task of the pod.
223
+ pod: The name of the pod.
224
+ phase: The phase of the pod.
225
+ cores: The cores it requests.
226
+ memory_gb: The GiB of memory it requests.
227
+ """
228
+
229
+ run_id: str | None
230
+ task_id: str
231
+ pod: str
232
+ phase: str
233
+ cores: float
234
+ memory_gb: float
235
+
236
+
237
+ class ServerUsage(BaseModel):
238
+ """A server of the cluster with the pods of runs on it.
239
+
240
+ Attributes:
241
+ name: The name of the server.
242
+ cores: The cores the scheduler gives the pods of the runs: the allocatable
243
+ ones less the requests of other pods, e.g. of DaemonSets.
244
+ memory_gb: The GiB of memory the scheduler gives the pods of the runs.
245
+ gpus: The GPUs of the server.
246
+ reserved_by: The run which reserved the server, None if free.
247
+ pods: The pods of runs on the server, with their requests, which count
248
+ against `cores` and `memory_gb`.
249
+ """
250
+
251
+ name: str
252
+ cores: float
253
+ memory_gb: float
254
+ gpus: int
255
+ reserved_by: str | None
256
+ pods: list[PodUsage]
257
+
258
+
259
+ class WorkerView(BaseModel):
260
+ """A worker of a task, from its reports, its pod and its claims.
261
+
262
+ Attributes:
263
+ worker: The worker id.
264
+ node: The server of the attempt shown: of the live pod while it runs,
265
+ None while that pod is not bound, else of the last report.
266
+ state: `running` while a pod of an attempt without a report lives (a
267
+ retry after an eviction, a preemption, a lost node or a kill for
268
+ memory, as a worker reports also when it is stopped), else from the
269
+ last report: `failed` when it was killed for memory, else `done`.
270
+ attempts: The attempts of the worker: one more than the highest attempt
271
+ of its reports and of its live pod.
272
+ processes: The processes of its pool, None before a report.
273
+ units_done: The distinct units its `unit_done` lines in the progress
274
+ file of the task name, live and over all its attempts.
275
+ units_claimed: The units it holds claimed now.
276
+ peak_memory_gb: The peak memory of the last report.
277
+ memory_limit_gb: The memory limit of the last report.
278
+ oom_killed: Whether the attempt of the last report was killed for memory.
279
+ """
280
+
281
+ worker: str
282
+ node: str | None
283
+ state: Literal["running", "done", "failed"]
284
+ attempts: int
285
+ processes: int | None
286
+ units_done: int
287
+ units_claimed: int
288
+ peak_memory_gb: float | None
289
+ memory_limit_gb: float | None
290
+ oom_killed: bool
291
+
292
+
293
+ class OutputFile(BaseModel):
294
+ """A file of the output of a task, its path relative to the store."""
295
+
296
+ path: str
297
+ bytes: int
298
+
299
+
300
+ class HintView(BaseModel):
301
+ """The default resources of a task type."""
302
+
303
+ memory_base_gb: float
304
+ memory_per_worker_gb: float
305
+ scratch_gb: float
306
+ gpus: int
307
+
308
+
309
+ class TaskTypeView(BaseModel):
310
+ """A task type with the JSON Schema of its settings, for the editor."""
311
+
312
+ name: str
313
+ settings_schema: dict[str, Any]
314
+ inputs: list[str]
315
+ outputs: list[str]
316
+ splittable: bool
317
+ uses_cores: bool
318
+ has_plan: bool
319
+ cost_setting: str | None
320
+ hint: HintView
321
+ description: str
322
+
323
+ @classmethod
324
+ def of(cls, task_type: TaskType) -> TaskTypeView:
325
+ """The view of a task type."""
326
+ return cls(
327
+ name=task_type.name,
328
+ settings_schema=task_type.settings.model_json_schema(),
329
+ inputs=list(task_type.inputs),
330
+ outputs=list(task_type.outputs),
331
+ splittable=task_type.splittable,
332
+ uses_cores=task_type.uses_cores,
333
+ has_plan=task_type.plan is not None,
334
+ cost_setting=task_type.cost_setting,
335
+ hint=HintView(**dataclasses.asdict(task_type.hint)),
336
+ description=task_type.description,
337
+ )
338
+
339
+
340
+ class CleanupRequest(BaseModel):
341
+ """Which ended runs a cleanup removes.
342
+
343
+ Attributes:
344
+ older_than_days: Runs which ended at least this many days ago.
345
+ group: Only the runs of this group.
346
+ failed_only: Only failed runs.
347
+ dry_run: List what would be removed, remove nothing.
348
+ """
349
+
350
+ older_than_days: float = Field(ge=0)
351
+ group: str | None = None
352
+ failed_only: bool = False
353
+ dry_run: bool = False
354
+
355
+
356
+ class CleanupResult(BaseModel):
357
+ """What a cleanup or a deletion removed, paths relative to the store."""
358
+
359
+ runs: list[str]
360
+ snapshots: list[str]
361
+ bytes_freed: int
362
+ dry_run: bool
363
+
364
+
365
+ class Usage(BaseModel):
366
+ """Bytes of the run store per run and per group (`""` without a group)."""
367
+
368
+ runs: dict[str, int]
369
+ groups: dict[str, int]
370
+
371
+
372
+ class CompareRequest(BaseModel):
373
+ """The runs to compare."""
374
+
375
+ runs: list[str]
376
+
377
+
378
+ def unknown_settings(model: type[BaseModel], settings: dict[str, Any]) -> list[str]:
379
+ """The keys of settings which are no field of their model, sorted.
380
+
381
+ A field counts by its name and its aliases; a model which allows extra keys
382
+ (`extra="allow"`) has none unknown.
383
+ """
384
+ if model.model_config.get("extra") == "allow":
385
+ return []
386
+ known: set[str] = set()
387
+ for name, field in model.model_fields.items():
388
+ known.add(name)
389
+ if field.alias is not None:
390
+ known.add(field.alias)
391
+ alias = field.validation_alias
392
+ if isinstance(alias, str):
393
+ known.add(alias)
394
+ elif isinstance(alias, AliasChoices):
395
+ known.update(a for a in alias.choices if isinstance(a, str))
396
+ return sorted(set(settings) - known)
397
+
398
+
399
+ def validate_workflow(workflow: Workflow, types: dict[str, TaskType]) -> list[str]:
400
+ """The errors of a workflow beyond its structure: types, settings and ports.
401
+
402
+ A setting which is no field of the settings model of its type is an error, so a
403
+ misspelled one is not silently dropped. An invalid setting is one error per
404
+ field, `task '<id>': <loc>: <msg>`, the location dotted.
405
+
406
+ Args:
407
+ workflow: The workflow.
408
+ types: The task types by name.
409
+
410
+ Returns:
411
+ The errors, empty when the workflow is valid.
412
+ """
413
+ errors: list[str] = []
414
+ for task in workflow.tasks:
415
+ task_type = types.get(task.type)
416
+ if task_type is None:
417
+ errors.append(f"task {task.id!r}: unknown task type {task.type!r}")
418
+ continue
419
+ errors.extend(
420
+ f"task {task.id!r}: unknown setting {key!r}"
421
+ for key in unknown_settings(task_type.settings, task.settings)
422
+ )
423
+ try:
424
+ task_type.settings.model_validate(task.settings)
425
+ except ValidationError as error:
426
+ # one message per field, `<loc>: <msg>`, which the Editor puts under the field
427
+ errors.extend(f"task {task.id!r}: {message}" for message in _errors(error))
428
+ for port in task.stage:
429
+ if port not in task.inputs:
430
+ errors.append(f"task {task.id!r} stages {port!r}, which is no input")
431
+ for port, ref in task.inputs.items():
432
+ source, _, output = ref.partition(".")
433
+ source_type = types.get(workflow.task(source).type)
434
+ if source_type is not None and output not in source_type.outputs:
435
+ errors.append(
436
+ f"input {port!r} of task {task.id!r}: task {source!r} has no "
437
+ f"output {output!r}"
438
+ )
439
+ return errors
440
+
441
+
442
+ def _errors(error: ValidationError) -> list[str]:
443
+ return [
444
+ f"{'.'.join(map(str, d['loc']))}: {d['msg']}" if d["loc"] else d["msg"]
445
+ for d in error.errors()
446
+ ]
447
+
448
+
449
+ def _files(root: Path) -> Iterable[os.stat_result]:
450
+ """The stat of every file below `root`, or of `root` itself, links not followed."""
451
+ if root.is_symlink() or root.is_file():
452
+ yield root.lstat()
453
+ return
454
+ for dirpath, _, filenames in os.walk(root):
455
+ for name in filenames:
456
+ with suppress(FileNotFoundError):
457
+ yield (Path(dirpath) / name).lstat()
458
+
459
+
460
+ def freed_bytes(paths: Iterable[Path]) -> int:
461
+ """The bytes removing `paths` frees: files of which no hard link stays."""
462
+ seen: dict[tuple[int, int], list[int]] = {}
463
+ for path in paths:
464
+ for st in _files(path):
465
+ entry = seen.setdefault(
466
+ (st.st_dev, st.st_ino), [st.st_size, st.st_nlink, 0]
467
+ )
468
+ entry[2] += 1
469
+ return sum(size for size, links, count in seen.values() if count >= links)
470
+
471
+
472
+ def disk_usage(path: Path) -> int:
473
+ """The bytes of the files below `path`, every hard linked file once."""
474
+ sizes = {(st.st_dev, st.st_ino): st.st_size for st in _files(path)}
475
+ return sum(sizes.values())
476
+
477
+
478
+ def _task_of(dag_task: str) -> str:
479
+ """The task id of a DAG task: `<task>-w0`, `<task>-extra`, `<task>-merge`."""
480
+ for suffix in ("-w0", "-extra", "-merge"):
481
+ if dag_task.endswith(suffix):
482
+ return dag_task.removesuffix(suffix)
483
+ return dag_task
484
+
485
+
486
+ def _gb(value: int | None) -> float | None:
487
+ """Bytes as GiB, None stays None."""
488
+ return None if value is None else value / 1024**3
489
+
490
+
491
+ def _newer(pod: Placement, than: Placement) -> bool:
492
+ """Whether a pod is of a later attempt; a pod of no known attempt is newest."""
493
+ if pod.attempt is None or than.attempt is None:
494
+ return than.attempt is not None
495
+ return pod.attempt > than.attempt
496
+
497
+
498
+ def run_progress(
499
+ tasks: dict[str, TaskState], counts: dict[str, LedgerCounts | None]
500
+ ) -> float:
501
+ """The progress of a run from 0 to 1, as the web app counts it.
502
+
503
+ The mean over its tasks: a task which is over (done or skipped) counts whole, a
504
+ task with a ledger the share of its units done, at most `UNFINISHED_TASK_MAX`
505
+ (its merge and outputs are still to come), any other task nothing; so the
506
+ progress is 1 only once every task is over.
507
+
508
+ Args:
509
+ tasks: The states of the tasks.
510
+ counts: The ledger counts by task, None or missing without a ledger.
511
+
512
+ Returns:
513
+ The progress, 0 for a run without tasks.
514
+ """
515
+ if not tasks:
516
+ return 0.0
517
+ total = 0.0
518
+ for task_id, state in tasks.items():
519
+ found = counts.get(task_id)
520
+ if state in (TaskState.DONE, TaskState.SKIPPED):
521
+ total += 1.0
522
+ elif found is not None and found.total > 0:
523
+ total += min(found.done / found.total, UNFINISHED_TASK_MAX)
524
+ return total / len(tasks)
525
+
526
+
527
+ def _pending(record: RunRecord) -> bool:
528
+ """Whether a run still has to release its node or cancel its workflow."""
529
+ return record.pending_release is not None or record.pending_cancel
530
+
531
+
532
+ def _units_by_worker(path: Path) -> dict[str, set[int]]:
533
+ """The distinct units of the `unit_done` lines of a progress file, per worker.
534
+
535
+ A worker appends a line per unit it completes, as it completes it, so the
536
+ count is live; a unit done twice (a lost claim, a retried attempt) counts
537
+ once. Other lines, and lines which are no JSON (a line being written),
538
+ are skipped.
539
+ """
540
+ units: dict[str, set[int]] = {}
541
+ try:
542
+ with path.open("rb") as handle:
543
+ for line in handle:
544
+ try:
545
+ value = json.loads(line)
546
+ except ValueError:
547
+ continue
548
+ if (
549
+ isinstance(value, dict)
550
+ and value.get("event") == "unit_done"
551
+ and isinstance(value.get("worker"), str)
552
+ and isinstance(value.get("unit"), int)
553
+ ):
554
+ units.setdefault(value["worker"], set()).add(value["unit"])
555
+ except FileNotFoundError:
556
+ pass
557
+ return units
558
+
559
+
560
+ def _json_object(line: bytes) -> dict[str, Any] | None:
561
+ """A line of a JSON lines file as an object, None for anything else."""
562
+ try:
563
+ value = json.loads(line)
564
+ except ValueError:
565
+ return None
566
+ return value if isinstance(value, dict) else None
567
+
568
+
569
+ def _first_line(path: Path) -> dict[str, Any] | None:
570
+ """The first line of a JSON lines file, None without one."""
571
+ try:
572
+ with path.open("rb") as handle:
573
+ line = handle.readline(PROGRESS_TAIL_BYTES)
574
+ except FileNotFoundError:
575
+ return None
576
+ return _json_object(line) if line.strip() else None
577
+
578
+
579
+ def _last_line(path: Path) -> dict[str, Any] | None:
580
+ """The last line of a JSON lines file, None without one."""
581
+ try:
582
+ with path.open("rb") as handle:
583
+ handle.seek(0, os.SEEK_END)
584
+ handle.seek(max(0, handle.tell() - PROGRESS_TAIL_BYTES))
585
+ lines = handle.read().splitlines()
586
+ except FileNotFoundError:
587
+ return None
588
+ for line in reversed(lines):
589
+ if line.strip():
590
+ return _json_object(line)
591
+ return None
592
+
593
+
594
+ class _Service:
595
+ """The work behind the routes."""
596
+
597
+ def __init__(
598
+ self,
599
+ settings: ServiceSettings,
600
+ backend: Backend,
601
+ types: dict[str, TaskType],
602
+ ) -> None:
603
+ self.settings = settings
604
+ self.backend = backend
605
+ self.types = types
606
+ # the index is rebuilt from the records, the store is the source of truth
607
+ self.store = IndexedRunStore(settings.store, RunIndex(settings.db))
608
+ self.store.rebuild()
609
+ self.resources = ResourceStore(settings.db)
610
+ self.controller = Controller(self.store, self.resources, backend, types)
611
+ self.submitter = Submitter(
612
+ self.store,
613
+ backend,
614
+ types,
615
+ self.controller,
616
+ settings.cluster,
617
+ settings.python,
618
+ settings.simrun_wheel,
619
+ settings.build_constraints,
620
+ )
621
+ self.controller.resume = self.submitter.run
622
+ self.tokens = load_tokens(settings.tokens)
623
+ self.edits = EditStore(settings.store)
624
+
625
+ # submission
626
+
627
+ def _nodes(self, reserve: str | None) -> list[Node]:
628
+ nodes = self.backend.nodes()
629
+ if reserve is None:
630
+ return nodes
631
+ node = self._free_node(nodes, reserve)
632
+ return [dataclasses.replace(node, reserved_by=None)]
633
+
634
+ @staticmethod
635
+ def _free_node(nodes: list[Node], name: str) -> Node:
636
+ node = next((n for n in nodes if n.name == name), None)
637
+ if node is None:
638
+ raise HTTPException(422, f"unknown node {name!r} to reserve")
639
+ if node.reserved_by is not None:
640
+ raise HTTPException(
641
+ 409, f"node {name!r} is reserved by run {node.reserved_by}"
642
+ )
643
+ return node
644
+
645
+ def plan(
646
+ self, request: SubmitRequest, priority: Priority
647
+ ) -> dict[str, tuple[ResolvedResources, TaskPlan]]:
648
+ """Resolve the resources and plan every task, in the order of the workflow.
649
+
650
+ Raises:
651
+ HTTPException: 422 for an invalid workflow or a task which fits no node,
652
+ 409 for a reserved node which is taken.
653
+ """
654
+ workflow = request.workflow
655
+ errors = validate_workflow(workflow, self.types)
656
+ if errors:
657
+ raise HTTPException(422, "; ".join(errors))
658
+ self._check_sources(request.sources)
659
+ nodes = self._nodes(request.reserve or workflow.reserve)
660
+ model = request.sources[0].name if request.sources else ""
661
+ planned: dict[str, tuple[ResolvedResources, TaskPlan]] = {}
662
+ for task in workflow.order():
663
+ task_type = self.types[task.type]
664
+ units = None
665
+ if task_type.splittable and task_type.count_units is not None:
666
+ settings = task_type.settings.model_validate(task.settings)
667
+ units = resolve(task_type.count_units)(settings)
668
+ key = resource_key(task_type, task.settings, model)
669
+ resources = self.resources.resolve(key, task_type.hint, task.resources)
670
+ mode = request.split or task.split or workflow.split
671
+ try:
672
+ plan = plan_task(
673
+ task, task_type, resources, nodes, mode, priority, units
674
+ )
675
+ except PlanError as error:
676
+ raise HTTPException(422, f"task {task.id!r}: {error}") from error
677
+ planned[task.id] = (resources, plan)
678
+ return planned
679
+
680
+ def _check_sources(self, sources: list[SourceProvenance]) -> None:
681
+ """Check that the snapshots are below the store and exist in it.
682
+
683
+ Raises:
684
+ HTTPException: 422 for a path outside the store or a snapshot which the
685
+ store of the service does not have, e.g. one a client wrote with a
686
+ `SIMRUN_STORE` other than the store of the service.
687
+ """
688
+ for source in sources:
689
+ for rel in (source.snapshot, source.patch):
690
+ if rel is None:
691
+ continue
692
+ path = PurePosixPath(rel)
693
+ if path.is_absolute() or ".." in path.parts:
694
+ raise HTTPException(
695
+ 422,
696
+ f"source {source.name!r}: {rel!r} is not relative to the store",
697
+ )
698
+ if not (self.settings.store / source.snapshot).is_dir():
699
+ raise HTTPException(
700
+ 422,
701
+ f"source {source.name!r}: the snapshot {source.snapshot!r} is not "
702
+ "in the store of the service; is SIMRUN_STORE the mount of its "
703
+ "store?",
704
+ )
705
+
706
+ def _needs_run_image(self, workflow: Workflow) -> bool:
707
+ return any(
708
+ self.types[t.type].image is None and t.type != COMMAND.name
709
+ for t in workflow.tasks
710
+ )
711
+
712
+ def _image(self, request: SubmitRequest) -> str | None:
713
+ """The image of a run without its own environment, the default image.
714
+
715
+ A run whose model has a `pyproject.toml` gets the image of its environment
716
+ in the background (None until then).
717
+
718
+ Raises:
719
+ HTTPException: 422 when the run needs an image and the service has no
720
+ default image.
721
+ """
722
+ if has_model_project(self.settings.store, request.sources):
723
+ return None
724
+ if self.settings.default_image is None and self._needs_run_image(
725
+ request.workflow
726
+ ):
727
+ raise HTTPException(
728
+ 422,
729
+ "the run needs an image: its model has no pyproject.toml and the "
730
+ "service has no default_image",
731
+ )
732
+ return self.settings.default_image
733
+
734
+ def submit(self, request: SubmitRequest, identity: Identity) -> RunRecord:
735
+ """Plan a run and write it `submitting`; the rest follows in the background.
736
+
737
+ The environment, the image and the workflow are resolved, built and
738
+ submitted by the `Submitter` in the pool of the controller.
739
+ """
740
+ priority = priority_of(identity, self.settings.owner)
741
+ planned = self.plan(request, priority)
742
+ image = self._image(request)
743
+ workflow = request.workflow
744
+ run_id = self.store.new_run_id()
745
+ record = RunRecord(
746
+ id=run_id,
747
+ submitted_by=identity.name,
748
+ owner=identity.owner,
749
+ agent=identity.name if identity.kind == "agent" else None,
750
+ priority=priority,
751
+ group=request.group,
752
+ workflow=workflow,
753
+ sources=request.sources,
754
+ image=image,
755
+ reserve=request.reserve or workflow.reserve,
756
+ state=RunState.SUBMITTING,
757
+ tasks=dict.fromkeys((t.id for t in workflow.tasks), TaskState.WAITING),
758
+ plans={k: dataclasses.asdict(plan) for k, (_, plan) in planned.items()},
759
+ # named before the submission, so a workflow a failed or an interrupted
760
+ # submission created is found and cancelled
761
+ backend_name=workflow_name(run_id),
762
+ created=datetime.now(UTC),
763
+ finished=None,
764
+ message=None,
765
+ )
766
+ self._write(record)
767
+ self.controller.start_submission(record.id, self.submitter.run)
768
+ logger.info("run %s accepted from %s", record.id, identity.name)
769
+ return record
770
+
771
+ def _write(self, record: RunRecord) -> None:
772
+ """Write the record, the settings and the provenance of every task.
773
+
774
+ The settings are written as submitted, not as the model of the service
775
+ validates them: the worker validates them with the plugin of the snapshot,
776
+ whose defaults may differ from the ones of the service. The provenance is
777
+ the list of the sources of the run, which a task reads through
778
+ `SIMRUN_PROVENANCE`.
779
+ """
780
+ with self.controller.lock:
781
+ self.store.write(record)
782
+ run_dir = self.store.run_dir(record.id)
783
+ provenance = json.dumps(
784
+ [source.model_dump(mode="json") for source in record.sources], indent=2
785
+ )
786
+ for task in record.workflow.tasks:
787
+ task_dir = run_dir / "tasks" / task.id
788
+ task_dir.mkdir(parents=True, exist_ok=True)
789
+ (task_dir / SETTINGS_FILE).write_text(json.dumps(task.settings, indent=2))
790
+ (task_dir / PROVENANCE_FILE).write_text(provenance)
791
+
792
+ # queries
793
+
794
+ def read(self, run_id: str) -> RunRecord:
795
+ """The record of a run; 404 when unknown."""
796
+ try:
797
+ return self.store.read(run_id)
798
+ except KeyError:
799
+ raise HTTPException(404, f"unknown run {run_id!r}") from None
800
+
801
+ def task_dir(self, record: RunRecord, task_id: str) -> Path:
802
+ """The directory of a task of a run; 404 when the run has no such task."""
803
+ if task_id not in record.tasks:
804
+ raise HTTPException(404, f"run {record.id} has no task {task_id!r}")
805
+ return self.store.run_dir(record.id) / "tasks" / task_id
806
+
807
+ def counts(self, record: RunRecord, task_id: str) -> LedgerCounts | None:
808
+ """The ledger counts of a task, None for a task without a ledger.
809
+
810
+ A ledger which vanishes while it is read (a run cleaned up, the store
811
+ unreachable for a moment) counts as none, so one run does not break the
812
+ list of runs or a stream.
813
+ """
814
+ task_dir = self.task_dir(record, task_id)
815
+ units = record.plans.get(task_id, {}).get("units")
816
+ if units is None or not (task_dir / UNITS_DIR).is_dir():
817
+ return None
818
+ try:
819
+ return Ledger(task_dir / UNITS_DIR, units).counts()
820
+ except OSError as e:
821
+ logger.warning("the ledger of %s/%s: %s", record.id, task_id, e)
822
+ return None
823
+
824
+ def progress(self, record: RunRecord) -> dict[str, TaskProgress]:
825
+ """The ledger counts and the first and the last progress line of every task."""
826
+ result: dict[str, TaskProgress] = {}
827
+ for task_id in record.tasks:
828
+ path = self.task_dir(record, task_id) / PROGRESS_FILE
829
+ result[task_id] = TaskProgress(
830
+ counts=self.counts(record, task_id),
831
+ first=_first_line(path),
832
+ last=_last_line(path),
833
+ )
834
+ return result
835
+
836
+ def outputs(self, record: RunRecord, task_id: str) -> list[OutputFile]:
837
+ """The files of the output of a task, relative to the store."""
838
+ out = self.task_dir(record, task_id) / "out"
839
+ files = (
840
+ [p for p in out.rglob("*") if p.is_file() and p.name != MERGED_MARKER]
841
+ if out.is_dir()
842
+ else []
843
+ )
844
+ return [
845
+ OutputFile(
846
+ path=p.relative_to(self.settings.store).as_posix(),
847
+ bytes=p.stat().st_size,
848
+ )
849
+ for p in sorted(files)
850
+ ]
851
+
852
+ def run_items(
853
+ self,
854
+ group: str | None,
855
+ owner: str | None,
856
+ agent: str | None,
857
+ state: RunState | None,
858
+ limit: int,
859
+ ) -> list[RunItem]:
860
+ """The most recent runs matching the filters, the newest first.
861
+
862
+ An active run carries its progress, from its record and its ledgers.
863
+ """
864
+ runs = self.store.index.runs(group, owner, agent, state)
865
+ runs.sort(key=lambda r: (r.summary.created, r.summary.id), reverse=True)
866
+ return [
867
+ RunItem(
868
+ **r.summary.model_dump(),
869
+ owner=r.owner,
870
+ workflow=r.workflow,
871
+ progress=self._run_progress(r.summary),
872
+ )
873
+ for r in runs[:limit]
874
+ ]
875
+
876
+ def _run_progress(self, summary: RunSummary) -> float | None:
877
+ """The progress of an active run, None for one which ended or is gone."""
878
+ if summary.state in TERMINAL_RUN_STATES:
879
+ return None
880
+ try:
881
+ record = self.store.read(summary.id)
882
+ except KeyError:
883
+ return None
884
+ counts = {t: self.counts(record, t) for t in record.tasks}
885
+ return run_progress(record.tasks, counts)
886
+
887
+ def servers(self) -> list[ServerUsage]:
888
+ """Every node with the pods of the runs on it."""
889
+ run_ids = {
890
+ workflow_name(r.summary.id): r.summary.id for r in self.store.index.runs()
891
+ }
892
+ pods: dict[str, list[PodUsage]] = {}
893
+ for p in self.backend.placements():
894
+ if p.node is None:
895
+ continue
896
+ pods.setdefault(p.node, []).append(
897
+ PodUsage(
898
+ run_id=run_ids.get(p.workflow),
899
+ task_id=_task_of(p.dag_task),
900
+ pod=p.pod,
901
+ phase=p.phase,
902
+ cores=p.cores,
903
+ memory_gb=p.memory_gb,
904
+ )
905
+ )
906
+ return [
907
+ ServerUsage(
908
+ name=n.name,
909
+ # what the scheduler counts the requests of the pods against, not
910
+ # the capacity of one worker pod the planner plans with
911
+ cores=n.allocatable_cores,
912
+ memory_gb=n.allocatable_memory_gb,
913
+ gpus=n.gpus,
914
+ reserved_by=n.reserved_by,
915
+ pods=pods.get(n.name, []),
916
+ )
917
+ for n in self.backend.nodes()
918
+ ]
919
+
920
+ def live_pods(self, record: RunRecord) -> list[Placement]:
921
+ """The pods of the workers of a run which have not terminated.
922
+
923
+ None before the workflow of the run was submitted and none once the run
924
+ ended: the backend is not asked then.
925
+ """
926
+ if record.backend_name is None or record.state in TERMINAL_RUN_STATES:
927
+ return []
928
+ return [
929
+ p
930
+ for p in self.backend.placements()
931
+ if p.workflow == record.backend_name and p.worker is not None
932
+ ]
933
+
934
+ def task_workers(self, record: RunRecord) -> dict[str, list[WorkerView]]:
935
+ """The workers of every task of a run, from one listing of the pods."""
936
+ pods = self.live_pods(record)
937
+ return {t: self.workers(record, t, pods) for t in record.tasks}
938
+
939
+ def run_state(self, record: RunRecord) -> tuple[dict[str, str], bool]:
940
+ """A tick of the stream of a run: its events and whether they are final.
941
+
942
+ Args:
943
+ record: The record of the run.
944
+
945
+ Returns:
946
+ The JSON text of the events `run`, `progress` and `workers`, and
947
+ whether the run is terminal with the states of its tasks final.
948
+ """
949
+ final = record.state in TERMINAL_RUN_STATES and all(
950
+ state in FINAL_TASK_STATES for state in record.tasks.values()
951
+ )
952
+ return {
953
+ "run": record.model_dump_json(),
954
+ "progress": _PROGRESS.dump_json(self.progress(record)).decode(),
955
+ "workers": _WORKERS.dump_json(self.task_workers(record)).decode(),
956
+ }, final
957
+
958
+ def workers(
959
+ self,
960
+ record: RunRecord,
961
+ task_id: str,
962
+ pods: list[Placement] | None = None,
963
+ ) -> list[WorkerView]:
964
+ """The workers of a task, sorted by id; 404 for an unknown task.
965
+
966
+ Args:
967
+ record: The record of the run.
968
+ task_id: The task.
969
+ pods: The live pods of the run (`live_pods`), listed when None.
970
+
971
+ Returns:
972
+ A view per worker of its reports or its pod.
973
+ """
974
+ task_dir = self.task_dir(record, task_id)
975
+ reports: dict[str, dict[int, WorkerReport]] = {}
976
+ for path in sorted((task_dir / WORKERS_DIR).glob("*.json")):
977
+ try:
978
+ report = WorkerReport(**json.loads(path.read_text()))
979
+ except (OSError, ValueError, TypeError) as error:
980
+ logger.warning("report %s of run %s: %s", path.name, record.id, error)
981
+ continue
982
+ reports.setdefault(report.worker, {})[attempt_of(path.name)] = report
983
+ live: dict[str, Placement] = {}
984
+ for p in self.live_pods(record) if pods is None else pods:
985
+ if p.worker is None or _task_of(p.dag_task) != task_id:
986
+ continue
987
+ # of two pods of a worker, the one of the later attempt
988
+ seen = live.get(p.worker)
989
+ if seen is None or _newer(p, seen):
990
+ live[p.worker] = p
991
+ units = record.plans.get(task_id, {}).get("units")
992
+ claims: dict[str, list[int]] = {}
993
+ if units is not None and (task_dir / UNITS_DIR).is_dir():
994
+ claims = Ledger(task_dir / UNITS_DIR, units).claims()
995
+ done = _units_by_worker(task_dir / PROGRESS_FILE)
996
+ views = []
997
+ for worker in sorted(reports.keys() | live.keys()):
998
+ attempts = reports.get(worker, {})
999
+ reported = max(attempts, default=-1)
1000
+ last = attempts.get(reported)
1001
+ oom = last is not None and is_oom(last, reported, record.oom_kills)
1002
+ pod = live.get(worker)
1003
+ # a worker reports in `finally`, also when it was evicted: a pod of an
1004
+ # attempt without a report is a retry which waits or runs, and a pod of
1005
+ # no known attempt counts as newer than every report
1006
+ state: Literal["running", "done", "failed"]
1007
+ if pod is not None and pod.attempt not in attempts:
1008
+ state = "running"
1009
+ node = pod.node # None while unbound, never an earlier node
1010
+ current = reported + 1 if pod.attempt is None else pod.attempt
1011
+ else:
1012
+ state = "failed" if oom else "done"
1013
+ node = last.node if last is not None and last.node else None
1014
+ current = reported
1015
+ views.append(
1016
+ WorkerView(
1017
+ worker=worker,
1018
+ node=node,
1019
+ state=state,
1020
+ attempts=max(current, reported) + 1,
1021
+ processes=last.processes if last else None,
1022
+ units_done=len(done.get(worker, ())),
1023
+ units_claimed=len(claims.get(worker, [])),
1024
+ peak_memory_gb=_gb(last.peak_memory_bytes if last else None),
1025
+ memory_limit_gb=_gb(last.memory_limit_bytes if last else None),
1026
+ oom_killed=oom,
1027
+ )
1028
+ )
1029
+ return views
1030
+
1031
+ # removal
1032
+
1033
+ def delete(self, run_id: str) -> CleanupResult:
1034
+ """Remove an ended run; 409 for a run which has not ended."""
1035
+ trash = self._trash(run_id, check=True)
1036
+ if trash is None:
1037
+ raise HTTPException(404, f"unknown run {run_id!r}")
1038
+ freed = freed_bytes([trash])
1039
+ shutil.rmtree(trash)
1040
+ logger.info("run %s deleted", run_id)
1041
+ return CleanupResult(
1042
+ runs=[run_id], snapshots=[], bytes_freed=freed, dry_run=False
1043
+ )
1044
+
1045
+ def _trash(self, run_id: str, check: bool = False) -> Path | None:
1046
+ """Move an ended run out of the runs, to remove it without the lock.
1047
+
1048
+ Returns:
1049
+ The directory it was moved to, None when the run is gone or has not
1050
+ ended and `check` is false.
1051
+
1052
+ Raises:
1053
+ HTTPException: 409 when `check` is set and the run has not ended.
1054
+ """
1055
+ with self.controller.lock:
1056
+ try:
1057
+ record = self.store.read(run_id)
1058
+ except KeyError:
1059
+ return None
1060
+ if record.state not in TERMINAL_RUN_STATES:
1061
+ if check:
1062
+ raise HTTPException(
1063
+ 409, f"run {run_id} is {record.state}, not ended"
1064
+ )
1065
+ return None
1066
+ if _pending(record):
1067
+ if check:
1068
+ raise HTTPException(
1069
+ 409,
1070
+ f"run {run_id} still has to release its node or cancel its "
1071
+ "workflow",
1072
+ )
1073
+ return None
1074
+ trash = self.store.root / "runs" / f"{TRASH_PREFIX}{run_id}-{os.getpid()}"
1075
+ os.rename(self.store.run_dir(run_id), trash)
1076
+ self.store.forget(run_id)
1077
+ return trash
1078
+
1079
+ def _unreferenced_snapshots(
1080
+ self, referenced: set[str], cutoff: datetime
1081
+ ) -> list[Path]:
1082
+ """The snapshots older than `cutoff` which no kept run references.
1083
+
1084
+ The target of `latest` and snapshots younger than `SNAPSHOT_GRACE_S` stay,
1085
+ since a client may be submitting them.
1086
+ """
1087
+ root = self.settings.store / "snapshots"
1088
+ if not root.is_dir():
1089
+ return []
1090
+ young = datetime.now(UTC).timestamp() - SNAPSHOT_GRACE_S
1091
+ found: list[Path] = []
1092
+ for parent in sorted(p for p in root.iterdir() if p.is_dir()):
1093
+ latest = parent / LATEST_LINK
1094
+ target = os.readlink(latest) if latest.is_symlink() else None
1095
+ for entry in sorted(parent.iterdir()):
1096
+ if (
1097
+ entry.name in (LATEST_LINK, target)
1098
+ or entry.name.startswith(".")
1099
+ or entry.is_symlink()
1100
+ or not entry.is_dir()
1101
+ ):
1102
+ continue
1103
+ st = entry.stat()
1104
+ if (
1105
+ entry.relative_to(self.settings.store).as_posix() in referenced
1106
+ or st.st_mtime >= cutoff.timestamp()
1107
+ or st.st_ctime >= young
1108
+ ):
1109
+ continue
1110
+ found.append(entry)
1111
+ return found
1112
+
1113
+ def cleanup(self, request: CleanupRequest) -> CleanupResult:
1114
+ """Remove ended runs and the snapshots no remaining run references."""
1115
+ cutoff = datetime.now(UTC) - timedelta(days=request.older_than_days)
1116
+ runs = self.store.index.runs()
1117
+ chosen = [
1118
+ r.summary.id
1119
+ for r in runs
1120
+ if r.summary.state in TERMINAL_RUN_STATES
1121
+ and (r.summary.finished or r.summary.created) <= cutoff
1122
+ and (request.group is None or r.summary.group == request.group)
1123
+ and (not request.failed_only or r.summary.state == RunState.FAILED)
1124
+ and not r.pending
1125
+ ]
1126
+ referenced = {
1127
+ s for r in runs if r.summary.id not in chosen for s in r.snapshots
1128
+ }
1129
+ snapshots = self._unreferenced_snapshots(referenced, cutoff)
1130
+ patches = [
1131
+ p for s in snapshots if (p := s.with_name(f"{s.name}.patch")).is_file()
1132
+ ]
1133
+ if request.dry_run:
1134
+ runs = [self.store.run_dir(run_id) for run_id in chosen]
1135
+ else:
1136
+ # an ended run does not change, so moving it under the lock suffices
1137
+ moved = {run_id: self._trash(run_id) for run_id in chosen}
1138
+ chosen = [run_id for run_id, path in moved.items() if path is not None]
1139
+ runs = [path for path in moved.values() if path is not None]
1140
+ runs += [ # left by an interrupted removal
1141
+ p
1142
+ for p in (self.store.root / "runs").glob(f"{TRASH_PREFIX}*")
1143
+ if p not in runs
1144
+ ]
1145
+ paths = runs + snapshots + patches
1146
+ freed = freed_bytes(paths)
1147
+ if not request.dry_run:
1148
+ for path in paths:
1149
+ if path.is_dir():
1150
+ shutil.rmtree(path)
1151
+ else:
1152
+ path.unlink(missing_ok=True)
1153
+ logger.info(
1154
+ "cleanup removed %d runs and %d snapshots", len(chosen), len(snapshots)
1155
+ )
1156
+ return CleanupResult(
1157
+ runs=sorted(chosen),
1158
+ snapshots=[
1159
+ s.relative_to(self.settings.store).as_posix() for s in snapshots
1160
+ ],
1161
+ bytes_freed=freed,
1162
+ dry_run=request.dry_run,
1163
+ )
1164
+
1165
+ def usage(self) -> Usage:
1166
+ """The bytes of every run and every group."""
1167
+ runs: dict[str, int] = {}
1168
+ groups: dict[str, int] = {}
1169
+ for run in self.store.index.runs():
1170
+ size = disk_usage(self.store.run_dir(run.summary.id))
1171
+ runs[run.summary.id] = size
1172
+ group = run.summary.group or ""
1173
+ groups[group] = groups.get(group, 0) + size
1174
+ return Usage(runs=runs, groups=groups)
1175
+
1176
+
1177
+ EVENTS_MAX_RUNS = 200
1178
+ KEEP_ALIVE_S = 15.0
1179
+ _RUN_ITEMS = TypeAdapter(list[RunItem])
1180
+ _SERVERS = TypeAdapter(list[ServerUsage])
1181
+ _PROGRESS = TypeAdapter(dict[str, TaskProgress])
1182
+ _WORKERS = TypeAdapter(dict[str, list[WorkerView]])
1183
+
1184
+
1185
+ def _sse(event: str, data: str) -> str:
1186
+ """One server-sent event; `data` is a single line of JSON."""
1187
+ return f"event: {event}\ndata: {data}\n\n"
1188
+
1189
+
1190
+ async def _stream(
1191
+ request: Request,
1192
+ settings: ServiceSettings,
1193
+ compute: Callable[[], tuple[dict[str, str], bool]],
1194
+ ) -> AsyncGenerator[str]:
1195
+ """Send the events `compute` returns at the start and whenever they change.
1196
+
1197
+ Args:
1198
+ request: The request, whose disconnect ends the stream.
1199
+ settings: The settings, `poll_s` and `events_max_s`.
1200
+ compute: Returns the JSON text of every event by name and whether the
1201
+ stream is complete after them; it runs in a thread.
1202
+
1203
+ Yields:
1204
+ The events changed since the last check, and a keep-alive comment after
1205
+ `KEEP_ALIVE_S` seconds without one. The stream ends when it is complete,
1206
+ after `events_max_s` seconds or on a disconnect.
1207
+ """
1208
+ started = last_sent = time.monotonic()
1209
+ sent: dict[str, str] = {}
1210
+ while True:
1211
+ try:
1212
+ payloads, complete = await asyncio.to_thread(compute)
1213
+ except HTTPException as error:
1214
+ logger.debug("an event stream ends: %s", error.detail)
1215
+ return
1216
+ except Exception:
1217
+ logger.warning(
1218
+ "an event stream ends, its state cannot be read", exc_info=True
1219
+ )
1220
+ return
1221
+ for event, data in payloads.items():
1222
+ if sent.get(event) != data:
1223
+ sent[event] = data
1224
+ last_sent = time.monotonic()
1225
+ yield _sse(event, data)
1226
+ if complete or time.monotonic() - started >= settings.events_max_s:
1227
+ return
1228
+ if time.monotonic() - last_sent >= KEEP_ALIVE_S:
1229
+ last_sent = time.monotonic()
1230
+ yield ": keep-alive\n\n"
1231
+ await asyncio.sleep(settings.poll_s)
1232
+ if await request.is_disconnected():
1233
+ return
1234
+
1235
+
1236
+ def _event_response(stream: AsyncGenerator[str]) -> StreamingResponse:
1237
+ """The response of an event stream: no caching, and no buffering by a proxy.
1238
+
1239
+ Args:
1240
+ stream: The events, e.g. of `_stream`.
1241
+
1242
+ Returns:
1243
+ A streaming response of `text/event-stream`.
1244
+ """
1245
+ return StreamingResponse(
1246
+ stream,
1247
+ media_type="text/event-stream",
1248
+ headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"},
1249
+ )
1250
+
1251
+
1252
+ async def _control(controller: Controller, poll_s: float) -> None:
1253
+ """Recover, then reconcile every `poll_s` seconds in a thread, off the event loop.
1254
+
1255
+ A failure is logged and the loop continues; a recovery which failed is tried
1256
+ again at the next poll.
1257
+ """
1258
+ recovered = False
1259
+ while True:
1260
+ if not recovered:
1261
+ try:
1262
+ await asyncio.to_thread(controller.recover)
1263
+ recovered = True
1264
+ except Exception:
1265
+ logger.exception("the recovery of the controller failed")
1266
+ try:
1267
+ await asyncio.to_thread(controller.reconcile)
1268
+ except Exception:
1269
+ logger.exception("the reconciliation failed")
1270
+ await asyncio.sleep(poll_s)
1271
+
1272
+
1273
+ def create_app(
1274
+ settings: ServiceSettings,
1275
+ backend: Backend,
1276
+ types: dict[str, TaskType] | None = None,
1277
+ ) -> FastAPI:
1278
+ """Create the app of the service.
1279
+
1280
+ Args:
1281
+ settings: The settings.
1282
+ backend: The backend the workflows run on.
1283
+ types: The task types by name, the ones of the entry points by default.
1284
+
1285
+ Returns:
1286
+ The app; its lifespan runs the controller every `settings.poll_s` seconds.
1287
+ """
1288
+ service = _Service(settings, backend, types or load_task_types())
1289
+ bearer = HTTPBearer(auto_error=False)
1290
+
1291
+ def identity(
1292
+ credentials: Annotated[HTTPAuthorizationCredentials | None, Depends(bearer)],
1293
+ ) -> Identity:
1294
+ found = (
1295
+ None
1296
+ if credentials is None
1297
+ else authenticate(service.tokens, credentials.credentials)
1298
+ )
1299
+ if found is None:
1300
+ raise HTTPException(
1301
+ 401, "a valid bearer token is required", {"WWW-Authenticate": "Bearer"}
1302
+ )
1303
+ return found
1304
+
1305
+ @asynccontextmanager
1306
+ async def lifespan(_: FastAPI) -> AsyncGenerator[None]:
1307
+ task = asyncio.create_task(_control(service.controller, settings.poll_s))
1308
+ try:
1309
+ yield
1310
+ finally:
1311
+ task.cancel()
1312
+ with suppress(asyncio.CancelledError):
1313
+ await task
1314
+ await asyncio.to_thread(service.controller.close)
1315
+
1316
+ # no documentation routes: they would answer without a token
1317
+ app = FastAPI(
1318
+ title="simrun",
1319
+ lifespan=lifespan,
1320
+ dependencies=[Depends(identity)],
1321
+ docs_url=None,
1322
+ redoc_url=None,
1323
+ openapi_url=None,
1324
+ )
1325
+ app.state.settings = settings
1326
+ app.state.backend = backend
1327
+ app.state.controller = service.controller
1328
+ app.state.service = service
1329
+ Caller = Annotated[Identity, Depends(identity)]
1330
+
1331
+ @app.post("/validate")
1332
+ def validate(body: Annotated[dict[str, Any], Body()]) -> ValidationResult:
1333
+ try:
1334
+ workflow = Workflow.model_validate(body)
1335
+ except ValidationError as error:
1336
+ return ValidationResult(ok=False, errors=_errors(error))
1337
+ errors = validate_workflow(workflow, service.types)
1338
+ return ValidationResult(ok=not errors, errors=errors)
1339
+
1340
+ @app.get("/task-types")
1341
+ def task_types() -> list[TaskTypeView]:
1342
+ return [TaskTypeView.of(service.types[name]) for name in sorted(service.types)]
1343
+
1344
+ @app.post("/workflows/yaml", response_class=PlainTextResponse)
1345
+ def workflow_to_yaml(workflow: Workflow) -> str:
1346
+ return workflow_yaml(workflow)
1347
+
1348
+ def own_edit(session_id: str, caller: Identity, owned: bool = True) -> EditSession:
1349
+ """The session, 404 when unknown, 403 when `owned` and not the caller's."""
1350
+ try:
1351
+ session = service.edits.read(session_id)
1352
+ except KeyError:
1353
+ raise HTTPException(404, f"unknown edit session {session_id}") from None
1354
+ if owned and session.owner != caller.owner:
1355
+ raise HTTPException(403, "only the owner of the session may do this")
1356
+ return session
1357
+
1358
+ def answer_or_session(
1359
+ session_id: str, request_id: str
1360
+ ) -> tuple[EditResult | None, EditSession]:
1361
+ """The answer to a request, none while there is none, and its session.
1362
+
1363
+ The session is read first: the command posts a result before it closes, so
1364
+ a session read as closed and no result means no answer will come.
1365
+ """
1366
+ try:
1367
+ session = service.edits.read(session_id)
1368
+ except KeyError:
1369
+ raise HTTPException(404, "the edit session is gone") from None
1370
+ return service.edits.result(session_id, request_id), session
1371
+
1372
+ async def ask(session_id: str, request_id: str) -> EditResult:
1373
+ """Wait for the answer of the command to a request, then drop the request.
1374
+
1375
+ A 503 when the command does not answer within `EDIT_WAIT_S`, and as soon as
1376
+ the session is no longer connected (the command closed it or went away).
1377
+ """
1378
+ loop = asyncio.get_running_loop()
1379
+ deadline = loop.time() + EDIT_WAIT_S
1380
+ try:
1381
+ while True:
1382
+ result, session = await asyncio.to_thread(
1383
+ answer_or_session, session_id, request_id
1384
+ )
1385
+ remaining = deadline - loop.time()
1386
+ if result is not None:
1387
+ return result
1388
+ if not session.connected:
1389
+ raise HTTPException(503, NOT_CONNECTED)
1390
+ if remaining <= 0:
1391
+ raise HTTPException(
1392
+ 503, "the machine of the session did not answer"
1393
+ )
1394
+ await asyncio.sleep(min(WAIT_POLL_S, remaining))
1395
+ finally:
1396
+ await asyncio.to_thread(service.edits.drop_request, session_id, request_id)
1397
+
1398
+ def refused(result: EditResult) -> HTTPException:
1399
+ """The error of a refusal of the command.
1400
+
1401
+ A conflict (409) and a refusal of the workflow (422) go to the browser as
1402
+ they are; any other status is the command's trouble with the service, a
1403
+ 502 with its detail.
1404
+ """
1405
+ status = result.status if result.status in (409, 422) else 502
1406
+ return HTTPException(status, result.detail)
1407
+
1408
+ @app.post("/edits", status_code=201)
1409
+ def open_edit(body: OpenEdit, caller: Caller) -> EditSession:
1410
+ return service.edits.open(caller, body)
1411
+
1412
+ @app.get("/edits")
1413
+ def list_edits(caller: Caller) -> list[EditSession]:
1414
+ return service.edits.list(caller.owner)
1415
+
1416
+ @app.get("/edits/{session_id}")
1417
+ def get_edit(session_id: str, caller: Caller) -> EditSession:
1418
+ return own_edit(session_id, caller, owned=False)
1419
+
1420
+ @app.put("/edits/{session_id}/state")
1421
+ def put_edit_state(
1422
+ session_id: str, state: EditState, caller: Caller
1423
+ ) -> EditSession:
1424
+ own_edit(session_id, caller)
1425
+ return service.edits.update(session_id, state)
1426
+
1427
+ @app.post("/edits/{session_id}/close")
1428
+ def close_edit(session_id: str, caller: Caller) -> EditSession:
1429
+ own_edit(session_id, caller)
1430
+ return service.edits.close(session_id)
1431
+
1432
+ @app.get("/edits/{session_id}/requests")
1433
+ async def edit_requests(
1434
+ session_id: str,
1435
+ request: Request,
1436
+ caller: Caller,
1437
+ timeout: Annotated[float, Query(ge=0, le=EDIT_POLL_MAX_S)] = 20.0,
1438
+ ) -> list[EditRequest]:
1439
+ await asyncio.to_thread(own_edit, session_id, caller)
1440
+ await asyncio.to_thread(service.edits.seen, session_id)
1441
+ loop = asyncio.get_running_loop()
1442
+ deadline = loop.time() + timeout
1443
+ while True:
1444
+ found = await asyncio.to_thread(service.edits.requests, session_id)
1445
+ remaining = deadline - loop.time()
1446
+ if await request.is_disconnected():
1447
+ return found # the command went away: no `seen`
1448
+ if found or remaining <= 0:
1449
+ await asyncio.to_thread(service.edits.seen, session_id)
1450
+ return found
1451
+ await asyncio.sleep(min(WAIT_POLL_S, remaining))
1452
+
1453
+ @app.post("/edits/{session_id}/requests/{request_id}/result", status_code=204)
1454
+ def edit_result(
1455
+ session_id: str, request_id: str, result: EditResult, caller: Caller
1456
+ ) -> None:
1457
+ own_edit(session_id, caller)
1458
+ try:
1459
+ service.edits.answer(session_id, request_id, result)
1460
+ except KeyError:
1461
+ raise HTTPException(404, f"unknown request {request_id}") from None
1462
+
1463
+ @app.post("/edits/{session_id}/save")
1464
+ async def save_edit(session_id: str, body: SaveBody, caller: Caller) -> EditSession:
1465
+ session = await asyncio.to_thread(own_edit, session_id, caller)
1466
+ if not session.connected:
1467
+ raise HTTPException(503, NOT_CONNECTED)
1468
+ request = await asyncio.to_thread(
1469
+ lambda: service.edits.add_request(
1470
+ session_id,
1471
+ "save",
1472
+ body.workflow,
1473
+ base=body.base,
1474
+ force=body.force,
1475
+ )
1476
+ )
1477
+ result = await ask(session_id, request.id)
1478
+ if not result.ok:
1479
+ raise refused(result)
1480
+ return await asyncio.to_thread(
1481
+ service.edits.saved,
1482
+ session_id,
1483
+ body.workflow,
1484
+ result.file_sha256,
1485
+ )
1486
+
1487
+ @app.post("/edits/{session_id}/run")
1488
+ async def run_edit(session_id: str, body: RunBody, caller: Caller) -> RunSummary:
1489
+ session = await asyncio.to_thread(own_edit, session_id, caller)
1490
+ if not session.connected:
1491
+ raise HTTPException(503, NOT_CONNECTED)
1492
+ request = await asyncio.to_thread(
1493
+ lambda: service.edits.add_request(
1494
+ session_id,
1495
+ "run",
1496
+ body.workflow,
1497
+ group=body.group,
1498
+ split=body.split,
1499
+ reserve=body.reserve,
1500
+ )
1501
+ )
1502
+ result = await ask(session_id, request.id)
1503
+ if not result.ok:
1504
+ raise refused(result)
1505
+ if result.run is None:
1506
+ raise HTTPException(502, "the machine of the session sent no run")
1507
+ return result.run
1508
+
1509
+ @app.get("/edits/{session_id}/events")
1510
+ def edit_events(
1511
+ session_id: str, request: Request, caller: Caller
1512
+ ) -> StreamingResponse:
1513
+ own_edit(session_id, caller, owned=False)
1514
+
1515
+ def compute() -> tuple[dict[str, str], bool]:
1516
+ try:
1517
+ session = service.edits.read(session_id)
1518
+ except KeyError:
1519
+ raise HTTPException(404, "the edit session is gone") from None
1520
+ return {"edit": session.model_dump_json()}, False
1521
+
1522
+ return _event_response(_stream(request, settings, compute))
1523
+
1524
+ @app.post("/estimate")
1525
+ def estimate(request: SubmitRequest, caller: Caller) -> dict[str, TaskEstimate]:
1526
+ priority = priority_of(caller, settings.owner)
1527
+ result: dict[str, TaskEstimate] = {}
1528
+ for task_id, (resources, plan) in service.plan(request, priority).items():
1529
+ seconds = None
1530
+ if resources.seconds_per_unit is not None and plan.units is not None:
1531
+ slots = plan.pods * plan.shape.processes
1532
+ rounds = -(-plan.units // slots)
1533
+ seconds = resources.seconds_per_unit * rounds
1534
+ result[task_id] = TaskEstimate(
1535
+ resources=resources, plan=plan, seconds=seconds
1536
+ )
1537
+ return result
1538
+
1539
+ @app.post("/runs", status_code=201)
1540
+ def submit(request: SubmitRequest, caller: Caller) -> RunSummary:
1541
+ return RunSummary.of(service.submit(request, caller))
1542
+
1543
+ @app.get("/me")
1544
+ def me(caller: Caller) -> Identity:
1545
+ return caller
1546
+
1547
+ @app.get("/servers")
1548
+ def servers() -> list[ServerUsage]:
1549
+ return service.servers()
1550
+
1551
+ @app.get("/events")
1552
+ def events(
1553
+ request: Request,
1554
+ group: str | None = None,
1555
+ owner: str | None = None,
1556
+ agent: str | None = None,
1557
+ state: RunState | None = None,
1558
+ ) -> StreamingResponse:
1559
+ def compute() -> tuple[dict[str, str], bool]:
1560
+ runs = service.run_items(group, owner, agent, state, EVENTS_MAX_RUNS)
1561
+ return {
1562
+ "runs": _RUN_ITEMS.dump_json(runs).decode(),
1563
+ "servers": _SERVERS.dump_json(service.servers()).decode(),
1564
+ }, False
1565
+
1566
+ return _event_response(_stream(request, settings, compute))
1567
+
1568
+ @app.get("/runs/{run_id}/events")
1569
+ def run_events(run_id: str, request: Request) -> StreamingResponse:
1570
+ service.read(run_id) # 404 before the stream starts
1571
+
1572
+ def compute() -> tuple[dict[str, str], bool]:
1573
+ return service.run_state(service.read(run_id))
1574
+
1575
+ return _event_response(_stream(request, settings, compute))
1576
+
1577
+ @app.get("/runs/{run_id}/tasks/{task_id}/workers")
1578
+ def workers(run_id: str, task_id: str) -> list[WorkerView]:
1579
+ return service.workers(service.read(run_id), task_id)
1580
+
1581
+ @app.get("/runs")
1582
+ def list_runs(
1583
+ group: str | None = None,
1584
+ owner: str | None = None,
1585
+ agent: str | None = None,
1586
+ state: RunState | None = None,
1587
+ ) -> list[RunSummary]:
1588
+ runs = service.store.index.runs(group, owner, agent, state)
1589
+ return [r.summary for r in runs]
1590
+
1591
+ @app.get("/runs/{run_id}")
1592
+ def get_run(run_id: str) -> RunRecord:
1593
+ return service.read(run_id)
1594
+
1595
+ @app.get("/runs/{run_id}/wait")
1596
+ async def wait(
1597
+ run_id: str,
1598
+ timeout: Annotated[float, Query(ge=0, le=WAIT_MAX_S)] = 60.0,
1599
+ ) -> RunSummary:
1600
+ loop = asyncio.get_running_loop()
1601
+ deadline = loop.time() + timeout
1602
+ while True:
1603
+ record = await asyncio.to_thread(service.read, run_id)
1604
+ remaining = deadline - loop.time()
1605
+ if record.state in TERMINAL_RUN_STATES or remaining <= 0:
1606
+ return RunSummary.of(record)
1607
+ await asyncio.sleep(min(WAIT_POLL_S, remaining))
1608
+
1609
+ @app.get("/runs/{run_id}/progress")
1610
+ def progress(run_id: str) -> dict[str, TaskProgress]:
1611
+ return service.progress(service.read(run_id))
1612
+
1613
+ @app.get("/runs/{run_id}/logs/{task_id}", response_class=PlainTextResponse)
1614
+ def logs(
1615
+ run_id: str,
1616
+ task_id: str,
1617
+ tail_lines: Annotated[int, Query(ge=0)] = DEFAULT_TAIL_LINES,
1618
+ ) -> str:
1619
+ """The logs of the pods of a task, the last `tail_lines` of each (0: all).
1620
+
1621
+ The logs kept in the store once the task was final, else the ones of its
1622
+ live pods.
1623
+ """
1624
+ record = service.read(run_id)
1625
+ task_dir = service.task_dir(record, task_id)
1626
+ lines = tail_lines or None
1627
+ found = read_kept(task_dir, lines)
1628
+ if found is None:
1629
+ if record.backend_name is None:
1630
+ raise HTTPException(409, f"run {run_id} has no workflow")
1631
+ found = backend.logs(record.backend_name, task_id, lines)
1632
+ return render(found)
1633
+
1634
+ @app.get("/runs/{run_id}/outputs/{task_id}")
1635
+ def outputs(run_id: str, task_id: str) -> list[OutputFile]:
1636
+ return service.outputs(service.read(run_id), task_id)
1637
+
1638
+ @app.post("/runs/{run_id}/cancel")
1639
+ def cancel(run_id: str) -> RunSummary:
1640
+ try:
1641
+ return RunSummary.of(service.controller.cancel(run_id))
1642
+ except KeyError:
1643
+ raise HTTPException(404, f"unknown run {run_id!r}") from None
1644
+ except ValueError as error:
1645
+ raise HTTPException(409, str(error)) from error
1646
+
1647
+ @app.delete("/runs/{run_id}")
1648
+ def delete(run_id: str) -> CleanupResult:
1649
+ return service.delete(run_id)
1650
+
1651
+ @app.post("/cleanup")
1652
+ def cleanup(request: CleanupRequest) -> CleanupResult:
1653
+ return service.cleanup(request)
1654
+
1655
+ @app.get("/usage")
1656
+ def usage() -> Usage:
1657
+ return service.usage()
1658
+
1659
+ @app.post("/compare")
1660
+ def compare(request: CompareRequest) -> CompareTable:
1661
+ try:
1662
+ return compare_runs(service.store, request.runs)
1663
+ except KeyError as error:
1664
+ raise HTTPException(404, f"unknown run {error.args[0]!r}") from None
1665
+
1666
+ if settings.web is not None:
1667
+ # the dependencies of the app apply to its routes only, not to a mount
1668
+ app.mount("/ui", StaticFiles(directory=settings.web, html=True), name="ui")
1669
+
1670
+ return app
1671
+
1672
+
1673
+ def openapi_document() -> dict[str, Any]:
1674
+ """The OpenAPI document of the service, the input of the types of the frontend.
1675
+
1676
+ The app is created on a fake backend in a temporary directory; the routes do not
1677
+ depend on either, so the document is the same on every machine.
1678
+ """
1679
+ from simrun.testing import FakeBackend
1680
+
1681
+ with tempfile.TemporaryDirectory() as directory:
1682
+ root = Path(directory)
1683
+ (root / "tokens.yaml").write_text("{}\n")
1684
+ settings = ServiceSettings(
1685
+ store=root / "store",
1686
+ db=root / "db" / "resources.sqlite",
1687
+ tokens=root / "tokens.yaml",
1688
+ owner="owner",
1689
+ cluster=ClusterConfig(registry="localhost:5000"),
1690
+ )
1691
+ app = create_app(settings, FakeBackend(), {"command": COMMAND})
1692
+ document = app.openapi()
1693
+ # the models only the event streams send have no response schema in the routes
1694
+ schemas = document.setdefault("components", {}).setdefault("schemas", {})
1695
+ for model in (RunItem,):
1696
+ schema = model.model_json_schema(ref_template="#/components/schemas/{model}")
1697
+ schemas.update(schema.pop("$defs", {}) | {model.__name__: schema})
1698
+ return json.loads(json.dumps(document))
1699
+
1700
+
1701
+ def load_settings(path: Path) -> ServiceSettings:
1702
+ """Read the settings of the service from a YAML file."""
1703
+ return ServiceSettings.model_validate(yaml.safe_load(path.read_text()))
1704
+
1705
+
1706
+ def serve() -> FastAPI:
1707
+ """The app of the settings in the YAML file named by `SIMRUN_CONFIG`.
1708
+
1709
+ The backend is the `backend` of the settings, a `module:callable` called without
1710
+ arguments.
1711
+
1712
+ Raises:
1713
+ RuntimeError: `SIMRUN_CONFIG` or the backend is not set.
1714
+ """
1715
+ path = os.environ.get(CONFIG_ENV)
1716
+ if not path:
1717
+ raise RuntimeError(f"{CONFIG_ENV} must name the YAML file of the settings")
1718
+ settings = load_settings(Path(path))
1719
+ if settings.backend is None:
1720
+ raise RuntimeError(f"{path}: no backend is configured")
1721
+ return create_app(settings, resolve(settings.backend)())