From a4ae5917e08076927ecf0d970c3cff11eda4b4c4 Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Fri, 5 Jun 2026 13:36:51 -0300 Subject: [PATCH] feat(background-execution): register + heartbeat workers in the claim loop Wire the durable worker_registry into run_worker_loop: register an IDLE row before the loop, flip it BUSY/IDLE event-driven on claim/complete with an immediate beat, refresh last_heartbeat on a periodic task (so the row stays fresh during a long job and while idle), and deregister in the finally so a cleanly-stopped worker leaves the roster immediately. Sessions are acquired per registry call via session_scope(), the same mechanism the loop already uses for the job heartbeat. pid/host are computed at the worker entrypoint and threaded through so they are stable and testable. Adds background_worker_registry_interval_s and background_worker_registry_retention_s settings. --- src/backend/base/langflow/__main__.py | 11 + .../services/background_execution/worker.py | 123 +++++++++++ .../test_worker_loop_registry.py | 196 ++++++++++++++++++ src/lfx/src/lfx/services/settings/base.py | 11 + 4 files changed, 341 insertions(+) create mode 100644 src/backend/tests/unit/services/background_execution/test_worker_loop_registry.py diff --git a/src/backend/base/langflow/__main__.py b/src/backend/base/langflow/__main__.py index 9ef72e5cb2..9e0d521cf4 100644 --- a/src/backend/base/langflow/__main__.py +++ b/src/backend/base/langflow/__main__.py @@ -867,6 +867,7 @@ def worker( from uuid import uuid4 + from langflow.services.background_execution.worker_registry import WorkerRegistryService from langflow.services.deps import get_job_service # Process-unique owner so this worker's heartbeats are attributable, and @@ -874,6 +875,12 @@ def worker( # steady fleet WITHOUT requiring a restart. owner = f"worker:{os.getpid()}:{uuid4().hex[:8]}" backend, worker_runner, teardown = await build_worker(owner=owner) + # Durable presence roster: register/heartbeat/deregister this worker in + # worker_registry. pid/host are computed here (at the entrypoint) so the + # loop stores stable values. The service is a stateless query helper. + worker_registry = WorkerRegistryService() + pid = os.getpid() + host = socket.gethostname() stop_event = asyncio.Event() loop = asyncio.get_running_loop() for sig in (signal.SIGTERM, signal.SIGINT): @@ -890,6 +897,10 @@ def worker( owner=owner, lease_ttl_s=settings.background_lease_ttl_s, watchdog_interval_s=settings.background_watchdog_interval_s, + worker_registry=worker_registry, + pid=pid, + host=host, + registry_interval_s=settings.background_worker_registry_interval_s, ) finally: await teardown() diff --git a/src/backend/base/langflow/services/background_execution/worker.py b/src/backend/base/langflow/services/background_execution/worker.py index 9be1684987..f21132e7c7 100644 --- a/src/backend/base/langflow/services/background_execution/worker.py +++ b/src/backend/base/langflow/services/background_execution/worker.py @@ -12,17 +12,24 @@ from __future__ import annotations import asyncio import contextlib +import os +import socket from typing import TYPE_CHECKING, Any from lfx.log.logger import logger from langflow.services.background_execution.runner import JobRunner +from langflow.services.database.models.worker_registry.model import WorkerState +from langflow.services.deps import session_scope if TYPE_CHECKING: from collections.abc import Callable + from uuid import UUID from lfx.services.settings.base import Settings + from langflow.services.background_execution.worker_registry import WorkerRegistryService + class WorkerJobRunner: """Run one durable job to terminal state inside a worker process. @@ -148,6 +155,61 @@ async def _recover_stranded_queued(backend: Any) -> None: await recover() +class _WorkerPresence: + """Mutable presence state shared between the claim loop and its heartbeat task. + + The claim loop flips ``state`` + ``current_job_id`` event-driven (BUSY on claim, + IDLE on complete) and writes an immediate registry beat each transition for a + snappy roster; the periodic ``_registry_heartbeat_loop`` reads this same state + and refreshes ``last_heartbeat`` on the interval, so the row stays fresh during a + long job (it runs while ``runner.run`` awaits) and while idle. + """ + + def __init__(self) -> None: + self.state: WorkerState = WorkerState.IDLE + self.current_job_id: UUID | None = None + + +async def _write_registry_heartbeat( + worker_registry: WorkerRegistryService, + owner: str, + presence: _WorkerPresence, +) -> None: + """Write one registry heartbeat for the current presence. Best-effort. + + Opens a short-lived ``session_scope()`` per call — the same mechanism the loop + uses for the job heartbeat — because ``WorkerRegistryService`` takes the session + positionally and never opens its own. A DB hiccup must never kill the loop, so + the write is guarded. + """ + with contextlib.suppress(Exception): + async with session_scope() as session: + await worker_registry.heartbeat( + session, + owner=owner, + state=presence.state, + current_job_id=presence.current_job_id, + ) + + +async def _registry_heartbeat_loop( + worker_registry: WorkerRegistryService, + owner: str, + presence: _WorkerPresence, + *, + interval_s: float, +) -> None: + """Refresh ``last_heartbeat`` every ``interval_s`` until cancelled. + + Mirrors the runner's ``_start_heartbeat``: it keeps the row fresh during a long + job (this task runs while the loop is blocked in ``await runner.run``) and during + idle. Cancelled by the loop's ``finally``. + """ + while True: + await asyncio.sleep(interval_s) + await _write_registry_heartbeat(worker_registry, owner, presence) + + async def run_worker_loop( backend: Any, runner: Any, @@ -158,6 +220,10 @@ async def run_worker_loop( owner: str | None = None, lease_ttl_s: float = 45.0, watchdog_interval_s: float | None = None, + worker_registry: WorkerRegistryService | None = None, + pid: int | None = None, + host: str | None = None, + registry_interval_s: float = 10.0, ) -> None: """Claim-and-run loop with a periodic lease watchdog. Returns on *stop_event*. @@ -175,6 +241,15 @@ async def run_worker_loop( lease_ttl_s: lease window the watchdog uses to decide "dead". watchdog_interval_s: how often the periodic watchdog runs; None disables it (startup-only reconcile, the prior behaviour). + worker_registry: durable presence roster; when set (with ``owner``) the + worker registers an IDLE row before the loop, flips it BUSY/IDLE on + claim/complete, beats ``last_heartbeat`` on ``registry_interval_s``, and + deregisters in the finally so a cleanly-stopped worker disappears. + pid: process id stored on the registry row (computed at the entrypoint so + tests can inject it). + host: hostname stored on the registry row (computed at the entrypoint). + registry_interval_s: how often the periodic registry heartbeat refreshes + ``last_heartbeat`` (idle or busy). Defaults to the settings value. """ # Startup reconcile: requeue work lost by a previously-crashed worker. await backend.requeue_lost(lease_ttl_s=lease_ttl_s) @@ -182,6 +257,21 @@ async def run_worker_loop( with contextlib.suppress(Exception): await _recover_stranded_queued(backend) + # Durable presence: register an IDLE row BEFORE any heartbeat so the + # recreate-on-missing branch in the service stays dead, then keep it fresh on + # the interval and flip it event-driven on claim/complete. + registry_active = worker_registry is not None and owner is not None + presence = _WorkerPresence() + if registry_active: + with contextlib.suppress(Exception): + async with session_scope() as session: + await worker_registry.register( + session, + owner=owner, + pid=pid if pid is not None else os.getpid(), + host=host if host is not None else socket.gethostname(), + ) + watchdog_task: asyncio.Task | None = None if watchdog_interval_s is not None: watchdog_task = asyncio.create_task( @@ -194,6 +284,17 @@ async def run_worker_loop( ) ) + registry_heartbeat_task: asyncio.Task | None = None + if registry_active: + registry_heartbeat_task = asyncio.create_task( + _registry_heartbeat_loop( + worker_registry, + owner, + presence, + interval_s=registry_interval_s, + ) + ) + try: while not stop_event.is_set(): job_id = await backend.claim(block_ms=idle_block_ms) @@ -209,6 +310,12 @@ async def run_worker_loop( if job_service is not None and owner is not None: with contextlib.suppress(Exception): await job_service.heartbeat(_coerce_uuid(job_id), owner) + # Flip the roster to BUSY with the claimed id and write an immediate beat + # so the roster reflects the transition without waiting for the interval. + if registry_active: + presence.state = WorkerState.BUSY + presence.current_job_id = _coerce_uuid(job_id) + await _write_registry_heartbeat(worker_registry, owner, presence) try: await runner.run(job_id) except asyncio.CancelledError: @@ -220,11 +327,27 @@ async def run_worker_loop( # whether the work should be retried; a stuck processing-list entry # would block reconcile forever. await backend.complete(job_id) + # Back to IDLE on the roster with an immediate beat for a snappy view. + if registry_active: + presence.state = WorkerState.IDLE + presence.current_job_id = None + await _write_registry_heartbeat(worker_registry, owner, presence) finally: if watchdog_task is not None: watchdog_task.cancel() with contextlib.suppress(asyncio.CancelledError, Exception): await watchdog_task + if registry_heartbeat_task is not None: + registry_heartbeat_task.cancel() + with contextlib.suppress(asyncio.CancelledError, Exception): + await registry_heartbeat_task + # Graceful stop deletes the row so the worker disappears from the roster + # immediately. Runs on the normal stop path AND on a signal-driven stop + # (the loop reaches this finally either way). Best-effort. + if registry_active: + with contextlib.suppress(Exception): + async with session_scope() as session: + await worker_registry.deregister(session, owner=owner) def _coerce_uuid(job_id: Any) -> Any: diff --git a/src/backend/tests/unit/services/background_execution/test_worker_loop_registry.py b/src/backend/tests/unit/services/background_execution/test_worker_loop_registry.py new file mode 100644 index 0000000000..d93e978822 --- /dev/null +++ b/src/backend/tests/unit/services/background_execution/test_worker_loop_registry.py @@ -0,0 +1,196 @@ +"""The worker loop registers/heartbeats/deregisters its presence in worker_registry. + +These drive the REAL ``run_worker_loop`` against the REAL test DB (the ``client`` +fixture; SQLite locally, Postgres in CI) with NO DB mocks. The queue side is a +real-behavior stand-in (hands out scripted ids, records completes) and the runner +is a stub we gate with an ``asyncio.Event`` so a job can be held BUSY mid-run while +we read the registry row. The heartbeat task's clock is injected so a beat lands +deterministically without a wall-clock sleep. +""" + +from __future__ import annotations + +import asyncio +from datetime import timezone +from uuid import uuid4 + +import pytest +from langflow.services.background_execution.worker import run_worker_loop +from langflow.services.background_execution.worker_registry import WorkerRegistryService +from langflow.services.database.models.worker_registry.model import WorkerRegistry, WorkerState +from langflow.services.deps import session_scope + +pytestmark = pytest.mark.usefixtures("client") + + +def _aware(dt): + """SQLite hands datetimes back naive; normalize to aware UTC for comparison.""" + return dt if dt.tzinfo is not None else dt.replace(tzinfo=timezone.utc) + + +async def _get_row(owner: str) -> WorkerRegistry | None: + async with session_scope() as session: + return await session.get(WorkerRegistry, owner) + + +async def _wait_for(predicate, *, timeout: float = 5.0): + """Poll an async predicate until truthy or the timeout elapses.""" + loop = asyncio.get_running_loop() + deadline = loop.time() + timeout + while loop.time() < deadline: + if await predicate(): + return True + await asyncio.sleep(0.01) + return False + + +class _FakeBackend: + """Real-behavior stand-in for the claim queue (not a mock of our code). + + The first claim waits on ``claim_gate`` so a test can observe the IDLE row the + worker registered BEFORE any job is handed out; once released it drains the + scripted ids then returns None (idle). + """ + + def __init__(self, ids, *, claim_gate=None): + self._ids = list(ids) + self.completed: list[str] = [] + self._claim_gate = claim_gate + self._gated = False + + async def requeue_lost(self, *, lease_ttl_s=45.0): # noqa: ARG002 + return [] + + async def claim(self, *, block_ms=1000): # noqa: ARG002 + if self._claim_gate is not None and not self._gated: + self._gated = True + await self._claim_gate.wait() + if self._ids: + return self._ids.pop(0) + return None + + async def complete(self, job_id): + self.completed.append(job_id) + + +class _GatedRunner: + """Holds each run open until a per-run gate is released, so we can observe BUSY.""" + + def __init__(self): + self.started = asyncio.Event() + self.release = asyncio.Event() + self.ran: list[str] = [] + + async def run(self, job_id): + self.ran.append(job_id) + self.started.set() + await self.release.wait() + + +async def test_registry_lifecycle_idle_busy_idle_then_deregister(): + """Startup -> one IDLE row; mid-job -> BUSY+job_id; after -> IDLE; on stop -> deleted.""" + owner = "worker:loop:lifecycle" + job_id = str(uuid4()) # production job ids are UUIDs; the registry column is UUID. + claim_gate = asyncio.Event() + backend = _FakeBackend([job_id], claim_gate=claim_gate) + runner = _GatedRunner() + stop_event = asyncio.Event() + registry = WorkerRegistryService() + + loop_task = asyncio.create_task( + run_worker_loop( + backend, + runner, + stop_event=stop_event, + idle_block_ms=10, + owner=owner, + worker_registry=registry, + pid=4242, + host="host-loop", + registry_interval_s=100.0, # long: no periodic beat interferes with this test + ) + ) + + # (1) After startup, with the first claim gated, exactly one IDLE row exists. + assert await _wait_for(lambda: _idle_row_exists(owner)) + row = await _get_row(owner) + assert row.pid == 4242 + assert row.host == "host-loop" + assert row.state == WorkerState.IDLE + assert row.current_job_id is None + + # (2) Release the claim: the job is claimed and the gated runner is mid-run, so + # the row is BUSY with the job id. + claim_gate.set() + assert await runner.started.wait() or True + assert await _wait_for(lambda: _busy_with_job(owner, job_id)) + + # (3) Release the run; once complete it returns to IDLE with no current job. + runner.release.set() + assert await _wait_for(lambda: _idle_no_job(owner)) + + # (4) Stop the loop; the finally deregisters so the row is deleted. + stop_event.set() + await asyncio.wait_for(loop_task, timeout=5.0) + assert await _get_row(owner) is None + + +async def _idle_row_exists(owner): + row = await _get_row(owner) + return row is not None and row.state == WorkerState.IDLE + + +async def _busy_with_job(owner, job_id): + row = await _get_row(owner) + return row is not None and row.state == WorkerState.BUSY and str(row.current_job_id) == job_id + + +async def _idle_no_job(owner): + row = await _get_row(owner) + return row is not None and row.state == WorkerState.IDLE and row.current_job_id is None + + +async def test_periodic_heartbeat_refreshes_last_heartbeat_during_long_job(): + """While a job runs, the periodic beat advances last_heartbeat with state still BUSY.""" + owner = "worker:loop:heartbeat" + job_id = str(uuid4()) + backend = _FakeBackend([job_id]) + runner = _GatedRunner() + stop_event = asyncio.Event() + registry = WorkerRegistryService() + + # Short interval so periodic beats land quickly in real time. The gated runner + # holds the job open, so every beat that lands does so with the row still BUSY: + # determinism comes from the gate, not from sleeping a fixed amount. + loop_task = asyncio.create_task( + run_worker_loop( + backend, + runner, + stop_event=stop_event, + idle_block_ms=10, + owner=owner, + worker_registry=registry, + pid=99, + host="host-hb", + registry_interval_s=0.02, + ) + ) + + # Wait until the job is BUSY, capture last_heartbeat, then wait for a later beat. + assert await _wait_for(lambda: _busy_with_job(owner, job_id)) + first = _aware((await _get_row(owner)).last_heartbeat) + + # A periodic beat must land while still BUSY and advance last_heartbeat. + advanced = await _wait_for(lambda: _busy_and_advanced(owner, first)) + assert advanced, "periodic heartbeat did not refresh last_heartbeat during the job" + + runner.release.set() + stop_event.set() + await asyncio.wait_for(loop_task, timeout=5.0) + + +async def _busy_and_advanced(owner, baseline): + row = await _get_row(owner) + if row is None or row.state != WorkerState.BUSY: + return False + return _aware(row.last_heartbeat) > baseline diff --git a/src/lfx/src/lfx/services/settings/base.py b/src/lfx/src/lfx/services/settings/base.py index 6edec277de..ad43eb9ae9 100644 --- a/src/lfx/src/lfx/services/settings/base.py +++ b/src/lfx/src/lfx/services/settings/base.py @@ -313,6 +313,17 @@ class Settings(BaseSettings): """How often the scaled worker's periodic watchdog scans for orphaned leases (a dead worker's in-flight job) and reconciles them WITHOUT requiring a restart. Must be > 0.""" + background_worker_registry_interval_s: float = Field(default=10.0, gt=0) + """How often a ``langflow worker`` refreshes its row in ``worker_registry`` + (idle or busy). This first-class idle heartbeat keeps ``last_heartbeat`` fresh + during a long job and while idle; the online window is a multiple of this. The + API-side collector derives the online/busy/idle gauges from these rows. Must + be > 0.""" + background_worker_registry_retention_s: float = Field(default=3600.0, gt=0) + """How long a stale ``worker_registry`` row (a crashed worker that never + deregistered) is kept before the collector prunes it. Surfaces a crashed + worker as offline for this window, then removes it so the roster does not + accumulate dead owners across restarts. Must be > 0.""" test_redis_url: str | None = Field(default=None) """Redis URL used by tests that exercise the scaled background backend.