Skip to content

Commit 0a87a3a

Browse files
mihowclaudeCopilot
authored
fix(jobs): throttle and defer pipeline heartbeat update (#1258)
* fix(jobs): throttle + defer pipeline heartbeat update Move _mark_pipeline_pull_services_seen off the HTTP request path by dispatching a new Celery task (update_pipeline_pull_services_seen) via .delay() from the /tasks and /result endpoints. The task throttles DB writes to once per ~30s per job, cutting concurrent UPDATE pressure under async_api load. Adds 6 unit tests covering the dispatch, throttle, and no-op edge cases. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(jobs): harden throttled pipeline heartbeat updates (#1260) * fix heartbeat review feedback Agent-Logs-Url: https://github.com/RolnickLab/antenna/sessions/bc1907bb-7118-4133-abab-4c4dd852ecc0 Co-authored-by: mihow <158175+mihow@users.noreply.github.com> * refine heartbeat timestamp handling Agent-Logs-Url: https://github.com/RolnickLab/antenna/sessions/bc1907bb-7118-4133-abab-4c4dd852ecc0 Co-authored-by: mihow <158175+mihow@users.noreply.github.com> * align heartbeat timestamps with local time Agent-Logs-Url: https://github.com/RolnickLab/antenna/sessions/bc1907bb-7118-4133-abab-4c4dd852ecc0 Co-authored-by: mihow <158175+mihow@users.noreply.github.com> --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: mihow <158175+mihow@users.noreply.github.com> * fix(jobs): gate heartbeat dispatch with Redis cache to cut broker churn Add a view-level cache.add() gate in _mark_pipeline_pull_services_seen keyed on (pipeline_id, project_id) with a HEARTBEAT_THROTTLE_SECONDS timeout. Previously every /tasks and /result request enqueued a Celery task whose sole job was usually to check the throttle and return; now most requests skip the enqueue entirely. The task's own stale-row check remains as a safety net under cache eviction. Key layout is intentionally noted to move to heartbeat:service:<id> once per-service identification lands via application-token auth (PR #1117), so one service's poll cannot suppress another's heartbeat. Adds test_view_gate_suppresses_redundant_dispatches; clears cache in setUp so the gate does not leak state between tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(jobs): simplify heartbeat task to rely on view-level gate The view's Redis cache gate (HEARTBEAT_THROTTLE_SECONDS, added in fd8e379) is the load-bearing throttle — it skips the whole .delay() when a recent heartbeat has already fired for the same (pipeline, project). With dispatch already gated, the task doesn't need its own throttle/staleness machinery. - Drop seen_at_iso param, Q-filter, .exists() preflight, expires=. - Task is now just .update(last_seen=now(), last_seen_live=True). - Drop 5 task-body tests (staleness, skip-when-recent, no-op-for- missing-job, no-op-for-no-pipeline, regression-guard) that exercised logic we removed. Keep 4: two endpoint dispatch smokes, broker- failure tolerance, and the view-gate suppression test. Crash-safety is unchanged: the write is still off the gunicorn request path, which was the SIGSEGV mitigation. --------- Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: mihow <158175+mihow@users.noreply.github.com>
1 parent 7f35bc5 commit 0a87a3a

3 files changed

Lines changed: 206 additions & 18 deletions

File tree

‎ami/jobs/tasks.py‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,47 @@
3434
# "nobody's listening" signal.
3535
WORKER_AVAILABILITY_ONLINE_CUTOFF = datetime.timedelta(minutes=5)
3636

37+
# Minimum interval between heartbeat dispatches for a given (pipeline, project).
38+
# The view-level Redis cache gate uses this window to skip .delay() under
39+
# concurrent polling; the task itself does no throttling.
40+
HEARTBEAT_THROTTLE_SECONDS = 30
41+
42+
43+
@celery_app.task(
44+
soft_time_limit=10,
45+
time_limit=15,
46+
ignore_result=True,
47+
# No retries — a missed heartbeat is benign; retrying adds load for no gain.
48+
)
49+
def update_pipeline_pull_services_seen(job_id: int) -> None:
50+
"""
51+
Fire-and-forget heartbeat task: record last_seen/last_seen_live for async
52+
(pull-mode) processing services linked to a job's pipeline.
53+
54+
Throttling lives in the view (Redis cache gate over HEARTBEAT_THROTTLE_SECONDS),
55+
so this task is dispatched at most once per (pipeline, project) per window
56+
and can just write.
57+
58+
Scope: marks ALL async services on the pipeline within this project as live,
59+
not just the specific service that made the request. Once application-token
60+
auth is available (PR #1117), this should be scoped to the individual
61+
calling service instead.
62+
"""
63+
from ami.jobs.models import Job # avoid circular import
64+
65+
try:
66+
job = Job.objects.select_related("pipeline").get(pk=job_id)
67+
except Job.DoesNotExist:
68+
return
69+
70+
if not job.pipeline_id:
71+
return
72+
73+
job.pipeline.processing_services.async_services().filter(projects=job.project_id).update(
74+
last_seen=datetime.datetime.now(),
75+
last_seen_live=True,
76+
)
77+
3778

3879
@celery_app.task(bind=True, soft_time_limit=default_soft_time_limit, time_limit=default_time_limit)
3980
def run_job(self, job_id: int) -> None:

‎ami/jobs/tests/test_jobs.py‎

Lines changed: 141 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
from ami.jobs.models import Job, JobDispatchMode, JobProgress, JobState, MLJob, SourceImageCollectionPopulateJob
1212
from ami.main.models import Project, SourceImage, SourceImageCollection
1313
from ami.ml.models import Pipeline
14+
from ami.ml.models.processing_service import ProcessingService
1415
from ami.ml.orchestration.jobs import queue_images_to_nats
1516
from ami.users.models import User
1617

@@ -1016,3 +1017,143 @@ def test_tasks_endpoint_rejects_non_async_jobs(self):
10161017
resp = self.client.post(tasks_url, {"batch_size": 1}, format="json")
10171018
self.assertEqual(resp.status_code, 400)
10181019
self.assertIn("async_api", resp.json()[0].lower())
1020+
1021+
1022+
class TestPipelineHeartbeatTask(APITestCase):
1023+
"""
1024+
Unit tests for update_pipeline_pull_services_seen and the view-level
1025+
_mark_pipeline_pull_services_seen fire-and-forget dispatch.
1026+
"""
1027+
1028+
def setUp(self):
1029+
from django.core.cache import cache
1030+
1031+
# Cache-based gate in _mark_pipeline_pull_services_seen would otherwise
1032+
# carry over between tests and suppress the .delay() we want to assert.
1033+
cache.clear()
1034+
1035+
self.project = Project.objects.create(name="Heartbeat Test Project")
1036+
self.pipeline = Pipeline.objects.create(name="Heartbeat Pipeline", slug="heartbeat-pipeline")
1037+
self.pipeline.projects.add(self.project)
1038+
self.collection = SourceImageCollection.objects.create(name="HB Collection", project=self.project)
1039+
self.job = Job.objects.create(
1040+
job_type_key=MLJob.key,
1041+
project=self.project,
1042+
name="Heartbeat Test Job",
1043+
pipeline=self.pipeline,
1044+
source_image_collection=self.collection,
1045+
dispatch_mode=JobDispatchMode.ASYNC_API,
1046+
)
1047+
self.service = ProcessingService.objects.create(
1048+
name="Heartbeat Worker",
1049+
endpoint_url=None, # None = pull-mode / async service
1050+
)
1051+
self.service.pipelines.add(self.pipeline)
1052+
self.service.projects.add(self.project)
1053+
1054+
def test_tasks_endpoint_dispatches_heartbeat_task(self):
1055+
"""The /tasks endpoint calls update_pipeline_pull_services_seen.delay(), not the DB directly."""
1056+
from unittest.mock import patch
1057+
1058+
job = self.job
1059+
job.status = JobState.STARTED
1060+
job.save(update_fields=["status"])
1061+
1062+
images = [
1063+
SourceImage.objects.create(
1064+
path=f"hb_tasks_{i}.jpg",
1065+
public_base_url="http://example.com",
1066+
project=self.project,
1067+
)
1068+
for i in range(2)
1069+
]
1070+
queue_images_to_nats(job, images)
1071+
1072+
user = User.objects.create_user(email="hbtest@example.com", is_superuser=True, is_active=True)
1073+
self.client.force_authenticate(user=user)
1074+
1075+
with patch("ami.jobs.views.update_pipeline_pull_services_seen.delay") as mock_delay:
1076+
tasks_url = reverse_with_params("api:job-tasks", args=[job.pk], params={"project_id": self.project.pk})
1077+
resp = self.client.post(tasks_url, {"batch_size": 1}, format="json")
1078+
1079+
self.assertEqual(resp.status_code, 200)
1080+
mock_delay.assert_called_once_with(job.pk)
1081+
1082+
def test_result_endpoint_dispatches_heartbeat_task(self):
1083+
"""The /result endpoint calls update_pipeline_pull_services_seen.delay(), not the DB directly."""
1084+
from unittest.mock import MagicMock, patch
1085+
1086+
user = User.objects.create_user(email="hbresult@example.com", is_superuser=True, is_active=True)
1087+
self.client.force_authenticate(user=user)
1088+
1089+
result_data = {
1090+
"results": [
1091+
{
1092+
"reply_subject": "test.reply.hb",
1093+
"result": {
1094+
"pipeline": "heartbeat-pipeline",
1095+
"algorithms": {},
1096+
"total_time": 0.1,
1097+
"source_images": [],
1098+
"detections": [],
1099+
"errors": None,
1100+
},
1101+
}
1102+
]
1103+
}
1104+
1105+
mock_async_result = MagicMock()
1106+
mock_async_result.id = "hb-task-id"
1107+
with (
1108+
patch("ami.jobs.views.process_nats_pipeline_result.delay", return_value=mock_async_result),
1109+
patch("ami.jobs.views.update_pipeline_pull_services_seen.delay") as mock_delay,
1110+
):
1111+
result_url = reverse_with_params(
1112+
"api:job-result", args=[self.job.pk], params={"project_id": self.project.pk}
1113+
)
1114+
resp = self.client.post(result_url, result_data, format="json")
1115+
1116+
self.assertEqual(resp.status_code, 200)
1117+
mock_delay.assert_called_once_with(self.job.pk)
1118+
1119+
def test_tasks_endpoint_tolerates_heartbeat_dispatch_failure(self):
1120+
"""Heartbeat enqueue errors should not fail the /tasks response."""
1121+
from unittest.mock import patch
1122+
1123+
from kombu.exceptions import OperationalError
1124+
1125+
job = self.job
1126+
job.status = JobState.STARTED
1127+
job.save(update_fields=["status"])
1128+
1129+
image = SourceImage.objects.create(
1130+
path="hb_tasks_broker.jpg",
1131+
public_base_url="http://example.com",
1132+
project=self.project,
1133+
)
1134+
queue_images_to_nats(job, [image])
1135+
1136+
user = User.objects.create_user(email="hbbroker@example.com", is_superuser=True, is_active=True)
1137+
self.client.force_authenticate(user=user)
1138+
1139+
with patch(
1140+
"ami.jobs.views.update_pipeline_pull_services_seen.delay",
1141+
side_effect=OperationalError("broker unavailable"),
1142+
):
1143+
tasks_url = reverse_with_params("api:job-tasks", args=[job.pk], params={"project_id": self.project.pk})
1144+
resp = self.client.post(tasks_url, {"batch_size": 1}, format="json")
1145+
1146+
self.assertEqual(resp.status_code, 200)
1147+
self.assertEqual(len(resp.json()["tasks"]), 1)
1148+
1149+
def test_view_gate_suppresses_redundant_dispatches(self):
1150+
"""Rapid repeated calls to _mark_pipeline_pull_services_seen should only enqueue once per window."""
1151+
from unittest.mock import patch
1152+
1153+
from ami.jobs.views import _mark_pipeline_pull_services_seen
1154+
1155+
with patch("ami.jobs.views.update_pipeline_pull_services_seen.delay") as mock_delay:
1156+
for _ in range(5):
1157+
_mark_pipeline_pull_services_seen(self.job)
1158+
1159+
self.assertEqual(mock_delay.call_count, 1)

‎ami/jobs/views.py‎

Lines changed: 24 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import kombu.exceptions
55
import nats.errors
66
from asgiref.sync import async_to_sync
7+
from django.core.cache import cache
78
from django.db.models import Q
89
from django.db.models.query import QuerySet
910
from django.forms import IntegerField
@@ -24,7 +25,7 @@
2425
MLJobTasksRequestSerializer,
2526
MLJobTasksResponseSerializer,
2627
)
27-
from ami.jobs.tasks import process_nats_pipeline_result
28+
from ami.jobs.tasks import HEARTBEAT_THROTTLE_SECONDS, process_nats_pipeline_result, update_pipeline_pull_services_seen
2829
from ami.main.api.schemas import project_id_doc_param
2930
from ami.main.api.views import DefaultViewSet
3031
from ami.utils.fields import url_boolean_param
@@ -52,26 +53,31 @@ def _actor_log_context(request) -> tuple[str, str | None]:
5253

5354
def _mark_pipeline_pull_services_seen(job: "Job") -> None:
5455
"""
55-
Record a heartbeat for async (pull-mode) processing services linked to the job's pipeline.
56-
57-
Called on every task-fetch and result-submit request so that the worker's polling activity
58-
keeps last_seen/last_seen_live current. The periodic check_processing_services_online task
59-
will mark services offline if this heartbeat stops arriving within PROCESSING_SERVICE_LAST_SEEN_MAX.
60-
61-
IMPORTANT: This marks ALL async services on the pipeline within this project as live, not just
62-
the specific service that made the request. If multiple async services share the same pipeline
63-
within a project, a single worker polling will keep all of them appearing online.
64-
Once application-token auth is available (PR #1117), this should be scoped to the individual
65-
calling service instead.
56+
Enqueue a fire-and-forget heartbeat for async (pull-mode) processing services
57+
linked to the job's pipeline.
58+
59+
A Redis cache gate skips the dispatch when a heartbeat for the same
60+
(pipeline, project) has already fired within HEARTBEAT_THROTTLE_SECONDS,
61+
so under concurrent polling we avoid broker + task churn. The Celery task
62+
keeps the DB write off the HTTP request path.
63+
64+
Cache key scope: currently `heartbeat:pipeline:<pipeline_id>:project:<project_id>`
65+
because we cannot yet identify the specific calling service. Once
66+
application-token auth lands (PR #1117), the key should become
67+
`heartbeat:service:<service_id>` so each service gets its own throttle
68+
window and one service's poll does not suppress another's heartbeat.
6669
"""
67-
import datetime
68-
6970
if not job.pipeline_id:
7071
return
71-
job.pipeline.processing_services.async_services().filter(projects=job.project_id).update(
72-
last_seen=datetime.datetime.now(),
73-
last_seen_live=True,
74-
)
72+
cache_key = f"heartbeat:pipeline:{job.pipeline_id}:project:{job.project_id}"
73+
if not cache.add(cache_key, 1, timeout=HEARTBEAT_THROTTLE_SECONDS):
74+
return
75+
try:
76+
update_pipeline_pull_services_seen.delay(job.pk)
77+
except (kombu.exceptions.KombuError, ConnectionError, OSError) as exc:
78+
msg = f"Failed to enqueue non-critical pipeline heartbeat for job {job.pk}: {exc}"
79+
logger.warning(msg)
80+
job.logger.warning(msg)
7581

7682

7783
class JobFilterSet(filters.FilterSet):

0 commit comments

Comments
 (0)