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/__init__.py +3 -0
- simrun/argo.py +449 -0
- simrun/build_constraints.txt +20 -0
- simrun/cgroup.py +77 -0
- simrun/cli.py +350 -0
- simrun/client.py +407 -0
- simrun/config.py +63 -0
- simrun/cwl.py +287 -0
- simrun/demo.py +106 -0
- simrun/edit.py +408 -0
- simrun/environment.py +370 -0
- simrun/ledger.py +286 -0
- simrun/mcp.py +161 -0
- simrun/planner.py +303 -0
- simrun/plugins.py +149 -0
- simrun/plugins_builtin.py +33 -0
- simrun/podspec.py +166 -0
- simrun/resources.py +270 -0
- simrun/service/__init__.py +1 -0
- simrun/service/app.py +1721 -0
- simrun/service/argo_backend.py +842 -0
- simrun/service/auth.py +93 -0
- simrun/service/backend.py +170 -0
- simrun/service/compare.py +77 -0
- simrun/service/controller.py +788 -0
- simrun/service/edits.py +406 -0
- simrun/service/index.py +292 -0
- simrun/service/logs.py +91 -0
- simrun/service/submission.py +349 -0
- simrun/snapshot.py +196 -0
- simrun/store.py +190 -0
- simrun/testing.py +263 -0
- simrun/worker.py +1097 -0
- simrun/workflow.py +188 -0
- simrun-0.1.0.dist-info/METADATA +107 -0
- simrun-0.1.0.dist-info/RECORD +39 -0
- simrun-0.1.0.dist-info/WHEEL +4 -0
- simrun-0.1.0.dist-info/entry_points.txt +10 -0
- simrun-0.1.0.dist-info/licenses/LICENSE +7 -0
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)())
|