mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-24 02:47:59 +08:00
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.
This commit is contained in:
ogabrielluiz
@ -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()
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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
|
||||
@ -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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user