Skip to content

Commit fbd607e

Browse files
author
mezzeddine
committed
jobs/tests/docs: Add heartbeat command tests and lifecycle documentation
Add router-level tests covering JobCommand lifecycle: - Kill command created when status → KILLED or DELETED - Delivered exactly once via heartbeat - No duplicate delivery on subsequent heartbeats - Multiple jobs handled independently - Non-terminal transitions do not create commands Add developer documentation describing one-shot delivery semantics. Ignore local virtualenv (.venv).
1 parent bbb34dd commit fbd607e

4 files changed

Lines changed: 393 additions & 1 deletion

File tree

‎.gitignore‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,5 +101,8 @@ pixi.lock
101101
*.egg-info
102102
docs/templates/_builtin_markdown.jinja
103103

104+
# virtualenv
105+
.venv
106+
104107
# docs site
105108
site

‎diracx-routers/tests/jobs/test_heartbeat_commands.py‎

Lines changed: 306 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,306 @@
1+
from __future__ import annotations
2+
3+
from datetime import datetime, timezone
4+
from time import sleep
5+
6+
import pytest
7+
from fastapi.testclient import TestClient
8+
9+
from diracx.core.models.job import JobStatus
10+
11+
pytestmark = pytest.mark.enabled_dependencies(
12+
[
13+
"AuthSettings",
14+
"JobDB",
15+
"JobLoggingDB",
16+
"ConfigSource",
17+
"TaskQueueDB",
18+
"SandboxMetadataDB",
19+
"WMSAccessPolicy",
20+
"DevelopmentSettings",
21+
"JobParametersDB",
22+
]
23+
)
24+
25+
26+
def test_kill_command_created_and_delivered_once(
27+
normal_user_client: TestClient,
28+
valid_job_id: int,
29+
):
30+
"""Verify lifecycle of a Kill command.
31+
32+
1. Job initially has no commands.
33+
2. Setting Status=KILLED creates a Kill command.
34+
3. Command is delivered on next heartbeat.
35+
4. Command is not re-delivered.
36+
"""
37+
# ------------------------------------------------------------------
38+
# 1️⃣ Initial heartbeat → no commands
39+
# ------------------------------------------------------------------
40+
r = normal_user_client.patch(
41+
"/api/jobs/heartbeat",
42+
json={valid_job_id: {"Vsize": 1000}},
43+
)
44+
r.raise_for_status()
45+
assert r.json() == []
46+
47+
# ------------------------------------------------------------------
48+
# 2️⃣ Set job to KILLED (creates Kill command internally)
49+
# ------------------------------------------------------------------
50+
r = normal_user_client.patch(
51+
"/api/jobs/status",
52+
json={
53+
valid_job_id: {
54+
datetime.now(timezone.utc).isoformat(): {
55+
"Status": JobStatus.KILLED,
56+
"MinorStatus": "Marked for termination",
57+
}
58+
}
59+
},
60+
)
61+
r.raise_for_status()
62+
63+
# Avoid heartbeat timestamp collision
64+
sleep(1)
65+
66+
# ------------------------------------------------------------------
67+
# 3️⃣ First heartbeat → command delivered
68+
# ------------------------------------------------------------------
69+
r = normal_user_client.patch(
70+
"/api/jobs/heartbeat",
71+
json={valid_job_id: {"Vsize": 1001}},
72+
)
73+
r.raise_for_status()
74+
75+
commands = r.json()
76+
77+
assert len(commands) == 1
78+
assert commands[0]["job_id"] == valid_job_id
79+
assert commands[0]["command"] == "Kill"
80+
81+
sleep(1)
82+
83+
# ------------------------------------------------------------------
84+
# 4️⃣ Second heartbeat → command NOT delivered again
85+
# ------------------------------------------------------------------
86+
r = normal_user_client.patch(
87+
"/api/jobs/heartbeat",
88+
json={valid_job_id: {"Vsize": 1002}},
89+
)
90+
r.raise_for_status()
91+
92+
assert r.json() == []
93+
94+
95+
def test_multiple_jobs_receive_independent_kill_commands(
96+
normal_user_client: TestClient,
97+
valid_job_ids: list[int],
98+
):
99+
"""Verify that multiple jobs each receive their own Kill command.
100+
101+
Ensure there is no cross-contamination.
102+
"""
103+
job_ids = valid_job_ids
104+
105+
# ------------------------------------------------------------------
106+
# Kill all jobs
107+
# ------------------------------------------------------------------
108+
r = normal_user_client.patch(
109+
"/api/jobs/status",
110+
json={
111+
job_id: {
112+
datetime.now(timezone.utc).isoformat(): {
113+
"Status": JobStatus.KILLED,
114+
"MinorStatus": "Marked for termination",
115+
}
116+
}
117+
for job_id in job_ids
118+
},
119+
)
120+
r.raise_for_status()
121+
122+
sleep(1)
123+
124+
# ------------------------------------------------------------------
125+
# Heartbeat all jobs at once
126+
# ------------------------------------------------------------------
127+
r = normal_user_client.patch(
128+
"/api/jobs/heartbeat",
129+
json={job_id: {"Vsize": 2000} for job_id in job_ids},
130+
)
131+
r.raise_for_status()
132+
133+
commands = r.json()
134+
135+
assert len(commands) == len(job_ids)
136+
137+
returned_ids = {cmd["job_id"] for cmd in commands}
138+
assert returned_ids == set(job_ids)
139+
140+
for cmd in commands:
141+
assert cmd["command"] == "Kill"
142+
143+
sleep(1)
144+
145+
# ------------------------------------------------------------------
146+
# Second heartbeat → no commands anymore
147+
# ------------------------------------------------------------------
148+
r = normal_user_client.patch(
149+
"/api/jobs/heartbeat",
150+
json={job_id: {"Vsize": 2001} for job_id in job_ids},
151+
)
152+
r.raise_for_status()
153+
154+
assert r.json() == []
155+
156+
157+
def test_non_killed_status_does_not_create_command(
158+
normal_user_client: TestClient,
159+
valid_job_id: int,
160+
):
161+
"""Verify that updating a job to a non-KILLED status.
162+
163+
Ensure it does not create any JobCommand.
164+
"""
165+
# Set job to RUNNING
166+
r = normal_user_client.patch(
167+
"/api/jobs/status",
168+
json={
169+
valid_job_id: {
170+
datetime.now(timezone.utc).isoformat(): {
171+
"Status": "Running",
172+
"MinorStatus": "Normal transition",
173+
}
174+
}
175+
},
176+
)
177+
r.raise_for_status()
178+
179+
sleep(1)
180+
181+
# Heartbeat → should return no commands
182+
r = normal_user_client.patch(
183+
"/api/jobs/heartbeat",
184+
json={valid_job_id: {"Vsize": 500}},
185+
)
186+
r.raise_for_status()
187+
188+
assert r.json() == []
189+
190+
191+
def test_non_terminal_status_does_not_create_command(
192+
normal_user_client: TestClient,
193+
valid_job_id: int,
194+
):
195+
r = normal_user_client.patch(
196+
"/api/jobs/status",
197+
json={
198+
valid_job_id: {
199+
str(datetime.now(timezone.utc)): {
200+
"Status": JobStatus.RUNNING,
201+
"MinorStatus": "Normal transition",
202+
}
203+
}
204+
},
205+
)
206+
r.raise_for_status()
207+
208+
sleep(1)
209+
r = normal_user_client.patch(
210+
"/api/jobs/heartbeat",
211+
json={valid_job_id: {"Vsize": 1234}},
212+
)
213+
r.raise_for_status()
214+
assert r.json() == []
215+
216+
217+
def test_command_delivered_exactly_once(
218+
normal_user_client: TestClient,
219+
valid_job_id: int,
220+
):
221+
"""Explicitly verify that a Kill command is delivered.
222+
223+
Ensure it is delivered exactly once and removed after delivery.
224+
"""
225+
# Kill job
226+
normal_user_client.patch(
227+
"/api/jobs/status",
228+
json={
229+
valid_job_id: {
230+
datetime.now(timezone.utc).isoformat(): {
231+
"Status": "Killed",
232+
"MinorStatus": "Termination requested",
233+
}
234+
}
235+
},
236+
).raise_for_status()
237+
238+
sleep(1)
239+
240+
# First heartbeat → get Kill
241+
r = normal_user_client.patch(
242+
"/api/jobs/heartbeat",
243+
json={valid_job_id: {"Vsize": 111}},
244+
)
245+
r.raise_for_status()
246+
247+
assert len(r.json()) == 1
248+
249+
sleep(1)
250+
251+
# Second heartbeat → no more commands
252+
r = normal_user_client.patch(
253+
"/api/jobs/heartbeat",
254+
json={valid_job_id: {"Vsize": 112}},
255+
)
256+
r.raise_for_status()
257+
258+
assert r.json() == []
259+
260+
261+
def test_deleted_creates_kill_command(
262+
normal_user_client: TestClient,
263+
valid_job_id: int,
264+
):
265+
"""Verify that transitioning a job to DELETED.
266+
267+
Ensure it creates a Kill JobCommand (same behavior as KILLED).
268+
"""
269+
# Transition to DELETED
270+
r = normal_user_client.patch(
271+
"/api/jobs/status",
272+
json={
273+
valid_job_id: {
274+
datetime.now(timezone.utc).isoformat(): {
275+
"Status": "Deleted",
276+
"MinorStatus": "User removed job",
277+
}
278+
}
279+
},
280+
)
281+
r.raise_for_status()
282+
283+
sleep(1)
284+
285+
# Heartbeat should deliver Kill command
286+
r = normal_user_client.patch(
287+
"/api/jobs/heartbeat",
288+
json={valid_job_id: {"Vsize": 123}},
289+
)
290+
r.raise_for_status()
291+
292+
commands = r.json()
293+
294+
assert len(commands) == 1
295+
assert commands[0]["job_id"] == valid_job_id
296+
assert commands[0]["command"] == "Kill"
297+
298+
# Second heartbeat → no command anymore
299+
sleep(1)
300+
r = normal_user_client.patch(
301+
"/api/jobs/heartbeat",
302+
json={valid_job_id: {"Vsize": 124}},
303+
)
304+
r.raise_for_status()
305+
306+
assert r.json() == []

‎docs/dev/explanations/job_commands.md‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
## What Are Job Commands?
2+
3+
Job commands implement a **per-job control queue** between the WMS and the pilot (JobAgent).
4+
5+
They allow the WMS to send control instructions to running jobs via the heartbeat mechanism.
6+
7+
Currently used for:
8+
9+
- Remote job termination (`Kill` command)
10+
11+
The mechanism is generic and could support additional commands in the future (e.g. debug actions, core dump requests, etc.).
12+
13+
______________________________________________________________________
14+
15+
## Behaviour Summary
16+
17+
### Command Creation
18+
19+
When a job transitions to a terminal state (`KILLED` or `DELETED`),
20+
`set_job_statuses()` enqueues a `Kill` command in the JobDB.
21+
22+
### Command Delivery
23+
24+
Commands are delivered during:
25+
26+
```text
27+
PATCH /api/jobs/heartbeat
28+
```
29+
30+
Flow:
31+
32+
1. Heartbeat updates job state.
33+
2. `get_job_commands()` retrieves pending commands.
34+
3. Commands are marked as sent.
35+
4. Commands are returned to the pilot.
36+
37+
This guarantees **one-shot delivery semantics**:
38+
39+
- A command is delivered exactly once.
40+
- It is not re-delivered on subsequent heartbeats.
41+
42+
______________________________________________________________________
43+
44+
## Sequence Diagram (Kill Command Lifecycle)
45+
46+
```mermaid
47+
sequenceDiagram
48+
autonumber
49+
participant User
50+
participant Router
51+
participant JobDB
52+
participant WatchDog
53+
54+
User->>Router: PATCH /api/jobs/status (Killed)
55+
Router->>JobDB: update status → KILLED
56+
Router->>JobDB: enqueue "Kill"
57+
Router-->>User: 200 OK
58+
59+
WatchDog->>Router: PATCH /api/jobs/heartbeat
60+
Router->>JobDB: fetch + mark sent
61+
Router-->>WatchDog: [Kill]
62+
63+
WatchDog->>Router: PATCH /api/jobs/heartbeat
64+
Router-->>WatchDog: []
65+
```
66+
67+
______________________________________________________________________
68+
69+
## Activity View
70+
71+
```mermaid
72+
flowchart TD
73+
A[Status change → KILLED] --> B[Enqueue Kill command]
74+
B --> C[Stored in JobDB]
75+
76+
D[Heartbeat] --> E[Fetch commands]
77+
E --> F[Mark as sent]
78+
F --> G[Return to pilot]
79+
G --> D
80+
```
81+
82+
______________________________________________________________________

0 commit comments

Comments
 (0)