diff --git a/docs/docs/Develop/environment-variables.mdx b/docs/docs/Develop/environment-variables.mdx index 0688031314..59d442f3d8 100644 --- a/docs/docs/Develop/environment-variables.mdx +++ b/docs/docs/Develop/environment-variables.mdx @@ -412,6 +412,8 @@ The following environment variables set base Langflow server configuration, such | `LANGFLOW_WORKERS` | Integer | `1` | Number of worker processes. | | `LANGFLOW_WORKER_TIMEOUT` | Integer | `300` | Worker timeout in seconds. | | `LANGFLOW_GUNICORN_PRELOAD` | Boolean | `False` | **Experimental.** When `true`, enables Gunicorn `preload_app` (non-Windows): the app is loaded in the master process before worker processes fork. Can reduce per-worker startup cost; behavior and compatibility may change. | +| `LANGFLOW_JOB_QUEUE_TYPE` | String | `asyncio` | Job queue backend. Use `redis` to share queue events across workers. | +| `LANGFLOW_REDIS_QUEUE_DB` | Integer | `1` | Redis database index used by the Redis job queue backend. | | `LANGFLOW_SSL_CERT_FILE` | String | Not set | Path to the SSL certificate file for enabling HTTPS on the Langflow web server. This is separate from [database SSL connections](/configuration-custom-database#connect-langflow-to-a-local-postgresql-database). | | `LANGFLOW_SSL_KEY_FILE` | String | Not set | Path to the SSL key file for enabling HTTPS on the Langflow web server. This is separate from [database SSL connections](/configuration-custom-database#connect-langflow-to-a-local-postgresql-database). | | `LANGFLOW_DEACTIVATE_TRACING` | Boolean | `False` | Deactivate tracing functionality. | @@ -432,6 +434,19 @@ Here is what these arguments do: - `--max-requests 100`: Automatically restarts a worker after it has processed 100 requests. This helps prevent memory leaks from accumulating over time. - `--max-requests-jitter 20`: Adds a random jitter of up to 20 requests to the `--max-requests` value. This ensures that all workers don't restart simultaneously, which could cause a brief downtime or spikes in latency. +#### High-load and multi-worker environments + +For high-concurrency deployments (such as when increasing `LANGFLOW_WORKERS`), consider the following best practices: +- **Enable Gunicorn preload (experimental):** + Set `LANGFLOW_GUNICORN_PRELOAD=true` to enable Gunicorn’s `preload_app` mode, which can reduce per-worker startup overhead (non-Windows only). +- **Use Redis-backed job queue:** + Set `LANGFLOW_JOB_QUEUE_TYPE=redis` to share queue events across workers. + **Note:** Always configure a dedicated Redis database for the job queue, separate from the cache database. +- **Disable tracing:** + Set `LANGFLOW_DEACTIVATE_TRACING=True` to turn off tracing, which may cause concurrency bottlenecks under high load or in multi-worker environments. +- **Use PostgreSQL instead of SQLite:** + SQLite can experience database locks and deadlocks with concurrent, write-heavy usage. For production or multi-worker environments, configure Langflow to use PostgreSQL by setting the `LANGFLOW_DATABASE_URL` environment variable (see [Memory management options](/memory#configure-external-memory)). This is strongly recommended to avoid operational issues. + For more information about deploying Langflow servers, see [Langflow deployment overview](/deployment-overview). ### Storage diff --git a/docs/versioned_docs/version-1.8.0/Develop/environment-variables.mdx b/docs/versioned_docs/version-1.8.0/Develop/environment-variables.mdx index 34cd4e6f78..f7bd2ca79b 100644 --- a/docs/versioned_docs/version-1.8.0/Develop/environment-variables.mdx +++ b/docs/versioned_docs/version-1.8.0/Develop/environment-variables.mdx @@ -452,12 +452,30 @@ The following environment variables set base Langflow server configuration, such | `LANGFLOW_HEALTH_CHECK_MAX_RETRIES` | Integer | `5` | Set the maximum number of retries for Langflow's server status health checks. | | `LANGFLOW_WORKERS` | Integer | `1` | Number of worker processes. | | `LANGFLOW_WORKER_TIMEOUT` | Integer | `300` | Worker timeout in seconds. | +| `LANGFLOW_GUNICORN_PRELOAD` | Boolean | `False` | **Experimental.** When `true`, enables Gunicorn `preload_app` (non-Windows): the app is loaded in the master process before worker processes fork. Can reduce per-worker startup cost; behavior and compatibility may change. | +| `LANGFLOW_JOB_QUEUE_TYPE` | String | `asyncio` | Job queue backend. Use `redis` to share queue events across workers. | +| `LANGFLOW_REDIS_QUEUE_DB` | Integer | `1` | Redis database index used by the Redis job queue backend. | | `LANGFLOW_SSL_CERT_FILE` | String | Not set | Path to the SSL certificate file for enabling HTTPS on the Langflow web server. This is separate from [database SSL connections](/configuration-custom-database#connect-langflow-to-a-local-postgresql-database). | | `LANGFLOW_SSL_KEY_FILE` | String | Not set | Path to the SSL key file for enabling HTTPS on the Langflow web server. This is separate from [database SSL connections](/configuration-custom-database#connect-langflow-to-a-local-postgresql-database). | | `LANGFLOW_DEACTIVATE_TRACING` | Boolean | `False` | Deactivate tracing functionality. | | `LANGFLOW_CELERY_ENABLED` | Boolean | `False` | Enable Celery for distributed task processing. | | `LANGFLOW_ALEMBIC_LOG_TO_STDOUT` | Boolean | `False` | Whether to log Alembic database migration output to stdout instead of a log file. If `true`, Alembic logs to `stdout` and the default log file is ignored. | +### High-load and multi-worker environments + +### High-load and multi-worker environments + +For high-concurrency deployments (such as when increasing `LANGFLOW_WORKERS`), consider the following best practices: +- **Enable Gunicorn preload (experimental):** + Set `LANGFLOW_GUNICORN_PRELOAD=true` to enable Gunicorn’s `preload_app` mode, which can reduce per-worker startup overhead (non-Windows only). +- **Use Redis-backed job queue:** + Set `LANGFLOW_JOB_QUEUE_TYPE=redis` to share queue events across workers. + **Note:** Always configure a dedicated Redis database for the job queue, separate from the cache database. +- **Disable tracing:** + Set `LANGFLOW_DEACTIVATE_TRACING=True` to turn off tracing, which may cause concurrency bottlenecks under high load or in multi-worker environments. +- **Use PostgreSQL instead of SQLite:** + SQLite can experience database locks and deadlocks with concurrent, write-heavy usage. For production or multi-worker environments, configure Langflow to use PostgreSQL by setting the `LANGFLOW_DATABASE_URL` environment variable (see [Memory management options](/memory#configure-external-memory)). This is strongly recommended to avoid operational issues. + For more information about deploying Langflow servers, see [Langflow deployment overview](/deployment-overview). ### Storage diff --git a/pyproject.toml b/pyproject.toml index a833b61515..2ca61cdc22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,7 +58,8 @@ dev = [ "pyyaml>=6.0.2", "pyleak>=0.1.14", "mcp-server-fetch>=2025.1.17", - "onnxruntime>=1.20,<1.24" # >=1.24 does not support Python 3.10; <1.24 allows 1.23.x for agent-lifecycle-toolkit + "onnxruntime>=1.20,<1.24", # >=1.24 does not support Python 3.10; <1.24 allows 1.23.x for agent-lifecycle-toolkit + "fakeredis>=2.0.0", ] [[tool.uv.index]] diff --git a/src/backend/base/langflow/api/build.py b/src/backend/base/langflow/api/build.py index a35c69c6d4..331a7c62f9 100644 --- a/src/backend/base/langflow/api/build.py +++ b/src/backend/base/langflow/api/build.py @@ -4,6 +4,7 @@ import time import traceback import uuid from collections.abc import AsyncIterator +from typing import Protocol, runtime_checkable from fastapi import BackgroundTasks, HTTPException, Response from lfx.graph.graph.base import Graph @@ -42,6 +43,21 @@ from langflow.services.job_queue.service import JobQueueNotFoundError, JobQueueS from langflow.services.telemetry.schema import ComponentInputsPayload, ComponentPayload, PlaygroundPayload +@runtime_checkable +class _CancellableQueue(Protocol): + """Structural protocol for queues that expose an async ``cancel()`` hook. + + Used by ``create_flow_response.on_disconnect`` to terminate background work + owned by the queue itself. :class:`~langflow.services.job_queue.service.RedisQueueWrapper` + implements this so the wrapper's background fill task is cancelled on client + disconnect. Plain ``asyncio.Queue`` does not have a ``cancel`` method and + so does not satisfy the protocol — those cases are covered by the separate + ``event_task.cancel()`` call in ``on_disconnect``. + """ + + async def cancel(self) -> None: ... + + def _log_component_input_telemetry( vertex, vertex_id: str, @@ -134,9 +150,6 @@ async def get_flow_events_response( try: main_queue, event_manager, event_task, _ = queue_service.get_queue_data(job_id) if event_delivery in (EventDeliveryType.STREAMING, EventDeliveryType.DIRECT): - if event_task is None: - await logger.aerror(f"No event task found for job {job_id}") - raise HTTPException(status_code=404, detail="No event task found for job") return await create_flow_response( queue=main_queue, event_manager=event_manager, @@ -193,7 +206,7 @@ async def get_flow_events_response( async def create_flow_response( queue: asyncio.Queue, event_manager: EventManager, - event_task: asyncio.Task, + event_task: asyncio.Task | None, ) -> DisconnectHandlerStreamingResponse: """Create a streaming response for the flow build process.""" @@ -210,9 +223,27 @@ async def create_flow_response( await logger.aexception(f"Error consuming event: {exc}") break - def on_disconnect() -> None: + async def on_disconnect() -> None: logger.debug("Client disconnected, closing tasks") - event_task.cancel() + if event_task is not None: + event_task.cancel() + else: + # Known limitation: cross-worker passive disconnect cannot be propagated. + # When this worker does not own the build task (event_task is None), there + # is no in-process handle to cancel the producer. The producer worker will + # continue emitting events into the queue until the build completes naturally. + # Proper cross-worker cancellation would require a Redis side-channel + # (e.g. pubsub or a langflow:cancel: key) that the build loop polls + # periodically. Until that is implemented, log a warning so the silent + # no-op is at least observable in logs. + logger.warning( + "Client disconnected but no local event_task found — " + "this worker does not own the build task. " + "The producer will keep running until the build finishes naturally. " + "Cross-worker passive-disconnect cancellation is not yet implemented." + ) + if isinstance(queue, _CancellableQueue): + await queue.cancel() event_manager.on_end(data={}) return DisconnectHandlerStreamingResponse( @@ -636,8 +667,14 @@ async def cancel_flow_build( _, _, event_task, _ = queue_service.get_queue_data(job_id) if event_task is None: - await logger.awarning(f"No event task found for job_id {job_id}") - return True # Nothing to cancel is still a success + # Cross-worker path: the job is owned by another process. We have no local task + # to cancel, so the build continues unaffected. Return False so callers know the + # cancellation did not take effect rather than reporting a false success. + await logger.awarning( + f"No event task found for job_id {job_id} — likely owned by another worker. " + "Cross-worker cancellation is not supported; the build will continue." + ) + return False if event_task.done(): await logger.ainfo(f"Task for job_id {job_id} is already completed") diff --git a/src/backend/base/langflow/api/v1/chat.py b/src/backend/base/langflow/api/v1/chat.py index 881b3eac1d..626db507c3 100644 --- a/src/backend/base/langflow/api/v1/chat.py +++ b/src/backend/base/langflow/api/v1/chat.py @@ -69,7 +69,7 @@ async def _verify_job_ownership(job_id: str, current_user: CurrentActiveUser, qu Jobs with no registered owner (build_public_tmp) are accessible to any authenticated user. """ - job_owner = queue_service.get_job_owner(job_id) + job_owner = await queue_service.get_job_owner(job_id) if job_owner is not None and job_owner != current_user.id: await logger.awarning( "Ownership check failed: user %s tried to access job %s owned by %s", @@ -241,7 +241,7 @@ async def build_flow( queue_service=queue_service, flow_name=flow_name, ) - queue_service.register_job_owner(job_id, current_user.id) + await queue_service.register_job_owner(job_id, current_user.id) # This is required to support FE tests - we need to be able to set the event delivery to direct if event_delivery != EventDeliveryType.DIRECT: diff --git a/src/backend/base/langflow/services/cache/service.py b/src/backend/base/langflow/services/cache/service.py index 06cbeb2960..c0240b970d 100644 --- a/src/backend/base/langflow/services/cache/service.py +++ b/src/backend/base/langflow/services/cache/service.py @@ -1,8 +1,12 @@ import asyncio +import atexit +import os import pickle +import tempfile import threading import time from collections import OrderedDict +from pathlib import Path from typing import Generic, Union import dill @@ -18,6 +22,36 @@ from langflow.services.cache.base import ( LockType, ) +_redis_cache_experimental_warning_lock = threading.Lock() +_redis_cache_experimental_warning_emitted = False + + +def _warn_redis_experimental_once() -> None: + """Emit the RedisCache experimental warning only once per server run.""" + global _redis_cache_experimental_warning_emitted # noqa: PLW0603 + + with _redis_cache_experimental_warning_lock: + if _redis_cache_experimental_warning_emitted: + return + _redis_cache_experimental_warning_emitted = True + + # Cross-process deduplication: all workers forked from the same master + # share the same getppid() value, so they all target the same sentinel. + sentinel = Path(tempfile.gettempdir()) / f"langflow_redis_cache_warned_{os.getppid()}.sentinel" + try: + fd = os.open(sentinel, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + os.close(fd) + except FileExistsError: + return # Another worker already logged the warning + + # Best-effort cleanup so we don't leave a stale file in /tmp after every restart. + atexit.register(sentinel.unlink, missing_ok=True) + + logger.warning( + "RedisCache is an experimental feature and may not work as expected." + " Please report any issues to our GitHub repository." + ) + class ThreadingInMemoryCache(CacheService, Generic[LockType]): """A simple in-memory cache using an OrderedDict. @@ -196,6 +230,8 @@ class RedisCache(ExternalAsyncBaseCacheService, Generic[LockType]): b = cache["b"] """ + KEY_PREFIX = "langflow:cache:" + def __init__(self, host="localhost", port=6379, db=0, url=None, expiration_time=60 * 60) -> None: """Initialize a new RedisCache instance. @@ -210,16 +246,17 @@ class RedisCache(ExternalAsyncBaseCacheService, Generic[LockType]): # Redis is a main dependency, no need to import check from redis.asyncio import StrictRedis - logger.warning( - "RedisCache is an experimental feature and may not work as expected." - " Please report any issues to our GitHub repository." - ) + _warn_redis_experimental_once() if url: self._client = StrictRedis.from_url(url) else: self._client = StrictRedis(host=host, port=port, db=db) self.expiration_time = expiration_time + def _key(self, key) -> str: + """Return the namespaced Redis key.""" + return f"{self.KEY_PREFIX}{key}" + async def is_connected(self) -> bool: """Check if the Redis client is connected.""" import redis @@ -236,14 +273,14 @@ class RedisCache(ExternalAsyncBaseCacheService, Generic[LockType]): async def get(self, key, lock=None): if key is None: return CACHE_MISS - value = await self._client.get(str(key)) + value = await self._client.get(self._key(key)) return dill.loads(value) if value else CACHE_MISS @override async def set(self, key, value, lock=None) -> None: try: if pickled := dill.dumps(value, recurse=True): - result = await self._client.setex(str(key), self.expiration_time, pickled) + result = await self._client.setex(self._key(key), self.expiration_time, pickled) if not result: msg = "RedisCache could not set the value." raise ValueError(msg) @@ -273,18 +310,25 @@ class RedisCache(ExternalAsyncBaseCacheService, Generic[LockType]): @override async def delete(self, key, lock=None) -> None: - await self._client.delete(key) + await self._client.delete(self._key(key)) @override async def clear(self, lock=None) -> None: - """Clear all items from the cache.""" - await self._client.flushdb() + """Clear all items from the cache using a key-prefix scan to avoid nuking unrelated data.""" + cursor = 0 + pattern = f"{self.KEY_PREFIX}*" + while True: + cursor, keys = await self._client.scan(cursor, match=pattern, count=100) + if keys: + await self._client.delete(*keys) + if cursor == 0: + break async def contains(self, key) -> bool: """Check if the key is in the cache.""" if key is None: return False - return bool(await self._client.exists(str(key))) + return bool(await self._client.exists(self._key(key))) @override async def teardown(self) -> None: diff --git a/src/backend/base/langflow/services/job_queue/factory.py b/src/backend/base/langflow/services/job_queue/factory.py index 71d629fdc1..38589c3079 100644 --- a/src/backend/base/langflow/services/job_queue/factory.py +++ b/src/backend/base/langflow/services/job_queue/factory.py @@ -1,11 +1,31 @@ -from langflow.services.base import Service +from __future__ import annotations + +from typing import TYPE_CHECKING + +from typing_extensions import override + from langflow.services.factory import ServiceFactory -from langflow.services.job_queue.service import JobQueueService +from langflow.services.job_queue.service import JobQueueService, RedisJobQueueService + +if TYPE_CHECKING: + from lfx.services.settings.service import SettingsService class JobQueueServiceFactory(ServiceFactory): def __init__(self): super().__init__(JobQueueService) - def create(self) -> Service: + @override + def create(self, settings_service: SettingsService): + settings = settings_service.settings + if settings.job_queue_type == "redis": + host = settings.redis_queue_host or settings.redis_host + port = settings.redis_queue_port or settings.redis_port + return RedisJobQueueService( + host=host, + port=port, + db=settings.redis_queue_db, + url=settings.redis_queue_url, + ttl=settings.redis_queue_ttl, + ) return JobQueueService() diff --git a/src/backend/base/langflow/services/job_queue/service.py b/src/backend/base/langflow/services/job_queue/service.py index 79b93491b3..b6262283b2 100644 --- a/src/backend/base/langflow/services/job_queue/service.py +++ b/src/backend/base/langflow/services/job_queue/service.py @@ -1,15 +1,26 @@ from __future__ import annotations import asyncio -from typing import TYPE_CHECKING +import contextlib +import time +from typing import TYPE_CHECKING, Any +from uuid import UUID + +if TYPE_CHECKING: + from collections.abc import Coroutine from lfx.log.logger import logger from langflow.events.event_manager import EventManager from langflow.services.base import Service -if TYPE_CHECKING: - from uuid import UUID +# Sentinel value written to Redis Streams to signal end-of-stream to consumers. +_STREAM_SENTINEL_DATA = b"__sentinel__" + +# Shared Redis key prefix for job event streams. Producer (RedisJobQueueService) and +# consumer (RedisQueueWrapper) MUST agree on this — keep a single source of truth. +_STREAM_PREFIX = "langflow:queue:" +_OWNER_PREFIX = "langflow:owner:" class JobQueueNotFoundError(Exception): @@ -155,7 +166,7 @@ class JobQueueService(Service): logger.debug(f"Queue and event manager successfully created for job_id {job_id}") return main_queue, event_manager - def start_job(self, job_id: str, task_coro) -> None: + def start_job(self, job_id: str, task_coro: Coroutine) -> None: """Start an asynchronous task for a given job, replacing any existing active task. The method performs the following: @@ -164,6 +175,13 @@ class JobQueueService(Service): - Launches a new asynchronous task using the provided coroutine. - Updates the internal registry with the new task. + The coroutine is wrapped with :meth:`_guarded_task` so that any unhandled + exception causes an ``on_error`` event to be emitted and the end-of-stream + sentinel to be written before the task exits. This guarantees that + cross-worker consumers can always distinguish a clean end from a crash — + both paths terminate with the sentinel in the Redis Stream, but a crash + will be preceded by an ``error`` event. + Args: job_id (str): Unique identifier for the job. task_coro: A coroutine representing the job's asynchronous task. @@ -183,11 +201,52 @@ class JobQueueService(Service): logger.debug(f"Existing task for job_id {job_id} detected; cancelling it.") existing_task.cancel() - # Initiate the new asynchronous task. - task = asyncio.create_task(task_coro) + # Wrap the coroutine so that any crash emits on_error + sentinel before exit. + task = asyncio.create_task(self._guarded_task(job_id, task_coro, event_manager, main_queue)) self._queues[job_id] = (main_queue, event_manager, task, None) logger.debug(f"New task started for job_id {job_id}") + @staticmethod + async def _guarded_task( + job_id: str, + task_coro: Coroutine, + event_manager: EventManager, + main_queue: asyncio.Queue, + ) -> None: + """Run *task_coro* and guarantee the end-of-stream sentinel is written on crash. + + A well-behaved build coroutine (``generate_flow_events``) writes the sentinel + itself after emitting ``on_end``. If the coroutine raises an unexpected + exception before doing so, this wrapper: + + 1. Emits an ``on_error`` event so consumers can distinguish a crash from a + clean end — both cases deliver ``(None, None, ts)`` as the terminal item, + but a crash is always preceded by an error event in the stream. + 2. Puts the raw ``(None, None, ts)`` sentinel so the bridge flushes and + writes ``_STREAM_SENTINEL_DATA`` to the Redis Stream before cleanup + deletes the key, preventing consumers from hanging indefinitely. + + ``asyncio.CancelledError`` is not caught here; the caller (``cleanup_job`` / + ``cancel_flow_build``) is responsible for those paths and already handles + sentinel delivery via bridge cancellation + stream-key deletion. + """ + try: + await task_coro + except asyncio.CancelledError: + raise + except Exception as exc: + await logger.aerror( + f"Unhandled exception in build task for job_id {job_id}: {exc}", + exc_info=True, + ) + # 1. Emit an error event so the stream carries the failure record. + with contextlib.suppress(Exception): + event_manager.on_error(data={"error": str(exc)}) + # 2. Write the sentinel so the bridge terminates and flushes to Redis. + with contextlib.suppress(Exception): + main_queue.put_nowait((None, None, time.time())) + raise + def get_queue_data(self, job_id: str) -> tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]: """Retrieve the complete data structure associated with a job's queue. @@ -212,11 +271,11 @@ class JobQueueService(Service): except KeyError as exc: raise JobQueueNotFoundError(job_id) from exc - def register_job_owner(self, job_id: str, user_id: UUID) -> None: + async def register_job_owner(self, job_id: str, user_id: UUID) -> None: """Register the authenticated user who initiated a build job.""" self._job_owners[job_id] = user_id - def get_job_owner(self, job_id: str) -> UUID | None: + async def get_job_owner(self, job_id: str) -> UUID | None: """Return the user ID that owns a job, or None if not tracked.""" return self._job_owners.get(job_id) @@ -364,3 +423,529 @@ class JobQueueService(Service): for name, event_type in event_names_types: manager.register_event(name, event_type) return manager + + +class RedisQueueWrapper: + """Consumer-side asyncio.Queue interface backed by a Redis Stream. + + Created by :class:`RedisJobQueueService` when :meth:`get_queue_data` is called + for a job that was started on a different worker process. A background + ``_fill_task`` reads from the Redis Stream and populates a local buffer so that + the rest of ``build.py`` can use the familiar ``asyncio.Queue`` interface. + + Stream protocol + --------------- + * Normal event → ``XADD key * event_id data ts `` + * End-of-stream → ``XADD key * event_id __sentinel__ data __sentinel__ ts `` + + Self-termination + ---------------- + The fill task exits when it: + 1. Receives the end-of-stream sentinel from the stream, **or** + 2. Detects that the stream key no longer exists (job was cleaned up). + In both cases it puts ``(None, None, timestamp)`` into the local buffer so + that consumers in ``build.py`` see the normal end-of-stream signal. + """ + + STREAM_PREFIX = _STREAM_PREFIX + + # Tunables for the background fill task. Kept as class-level constants so + # subclasses or tests can override without touching the loop body. + _XREAD_BLOCK_MS = 1000 # how long XREAD blocks waiting for new entries + _XREAD_BATCH_COUNT = 100 # max entries fetched per XREAD call + _READ_ERROR_BACKOFF_S = 0.5 # backoff after a transient XREAD failure + # How long to keep polling before giving up on a stream that has never appeared. + # Protects against the early-poll race where the consumer wrapper is created + # before the producer worker has issued its first XADD. + _STARTUP_GRACE_S = 30.0 + # Upper bound on the local buffer. Provides backpressure so the fill task + # doesn't read arbitrarily far ahead of a slow consumer and accumulate events + # in memory unboundedly. 1 000 events x ~200 B each ~= 200 KB per wrapper — + # well within reason even with many concurrent jobs. + _BUFFER_MAXSIZE = 1000 + + def __init__(self, job_id: str, client: Any, ttl: int) -> None: + self._job_id = job_id + self._client = client + self._ttl = ttl + self._buffer: asyncio.Queue = asyncio.Queue(maxsize=self._BUFFER_MAXSIZE) + self._last_id = "0-0" # read from the beginning of the stream + # Flips to True the first time XREAD returns messages for this stream. + # Until then, "stream key does not exist" is NOT treated as end-of-stream + # — it just means the producer hasn't written its first event yet. + self._observed_stream: bool = False + self._created_at: float = time.monotonic() + # Flips to True after the very first XREAD call returns (regardless of + # whether it had results). empty() returns False until this flag is set + # so that the while-not-empty drain loop in build.py suspends on get() + # and lets the fill task populate the buffer before the loop exits. + self._first_read_done: bool = False + self._fill_task: asyncio.Task = asyncio.create_task(self._fill_from_redis()) + self._fill_task.add_done_callback(self._on_fill_done) + + @property + def _stream_key(self) -> str: + return f"{self.STREAM_PREFIX}{self._job_id}" + + def _on_fill_done(self, task: asyncio.Task) -> None: + """Ensure the consumer is unblocked if the fill task exits unexpectedly. + + A clean exit (sentinel received) already puts the end-of-stream sentinel + into the buffer. Cancellation and unhandled exceptions skip that path — + this callback catches both gaps so the consumer's ``await queue.get()`` + doesn't block forever. + + Done callbacks run synchronously, so we use the non-blocking + ``put_nowait``. If the buffer happens to be at capacity (slow consumer + + bounded buffer), evict the oldest item to make room: losing one event + is strictly preferable to a stuck consumer. + """ + exc = task.exception() if not task.cancelled() else None + if not task.cancelled() and exc is None: + # Clean exit: _fill_from_redis already put the sentinel in the buffer. + return + if exc is not None: + logger.error( + f"RedisQueueWrapper fill task raised for job {self._job_id}: {exc!r} " + "— delivering end-of-stream sentinel so the consumer is not left hanging." + ) + if self._buffer.full(): + with contextlib.suppress(asyncio.QueueEmpty): + self._buffer.get_nowait() + with contextlib.suppress(asyncio.QueueFull): + self._buffer.put_nowait((None, None, time.time())) + + async def _fill_from_redis(self) -> None: + """Read events from the Redis Stream and forward them to the local buffer.""" + _error_start: float | None = None + try: + while True: + try: + results = await self._client.xread( + {self._stream_key: self._last_id}, + block=self._XREAD_BLOCK_MS, + count=self._XREAD_BATCH_COUNT, + ) + except Exception as exc: # noqa: BLE001 + now = time.monotonic() + if _error_start is None: + _error_start = now + elapsed = now - _error_start + await logger.awarning( + f"RedisQueueWrapper read error for {self._job_id} (elapsed {elapsed:.1f}s): {exc}" + ) + if elapsed >= self._STARTUP_GRACE_S: + await logger.aerror( + f"RedisQueueWrapper: persistent Redis error for {self._job_id} " + f"after {elapsed:.1f}s; delivering end-of-stream sentinel." + ) + await self._buffer.put((None, None, time.time())) + return + await asyncio.sleep(self._READ_ERROR_BACKOFF_S) + continue + _error_start = None # reset on successful XREAD + + self._first_read_done = True + if results: + self._observed_stream = True + for _, messages in results: + for msg_id, fields in messages: + data = fields.get(b"data") + ts = float(fields.get(b"ts") or b"0") + if data == _STREAM_SENTINEL_DATA: + await self._buffer.put((None, None, ts)) + self._last_id = msg_id.decode() if isinstance(msg_id, bytes) else msg_id + return + event_id = (fields.get(b"event_id") or b"").decode() + await self._buffer.put((event_id, data, ts)) + # Advance cursor only after the item is safely in the buffer. + # Advancing before the await would skip this message on cancellation. + self._last_id = msg_id.decode() if isinstance(msg_id, bytes) else msg_id + # No results within the block timeout. + elif not self._observed_stream: + # Stream hasn't appeared yet — the producer may not have issued its + # first XADD (early-poll race between workers). Keep blocking until + # the startup grace period expires to avoid a false end-of-stream. + elapsed = time.monotonic() - self._created_at + if elapsed > self._STARTUP_GRACE_S: + await logger.awarning( + f"RedisQueueWrapper: stream for {self._job_id} never appeared " + f"after {elapsed:.1f}s; treating as end-of-stream." + ) + await self._buffer.put((None, None, time.time())) + return + # Otherwise keep looping — next XREAD will block again. + else: + # Stream was observed before; check whether the key still exists. + # Wrap in try/except so a transient Redis error here uses the same + # backoff path as an XREAD error rather than crashing the fill task. + try: + key_exists = await self._client.exists(self._stream_key) + except Exception as exc: # noqa: BLE001 + now = time.monotonic() + if _error_start is None: + _error_start = now + elapsed = now - _error_start + await logger.awarning( + f"RedisQueueWrapper exists() check failed for {self._job_id} " + f"(elapsed {elapsed:.1f}s): {exc}" + ) + if elapsed >= self._STARTUP_GRACE_S: + await logger.aerror( + f"RedisQueueWrapper: persistent Redis error for {self._job_id} " + f"after {elapsed:.1f}s; delivering end-of-stream sentinel." + ) + await self._buffer.put((None, None, time.time())) + return + await asyncio.sleep(self._READ_ERROR_BACKOFF_S) + continue + if not key_exists: + # Key gone — the job was cleaned up on the producer side. + await self._buffer.put((None, None, time.time())) + return + except asyncio.CancelledError: + return + + # ------------------------------------------------------------------ + # asyncio.Queue-compatible interface used by build.py + # ------------------------------------------------------------------ + + def empty(self) -> bool: + # Before the first XREAD completes the local buffer is empty even if + # Redis already has events queued. Returning False here causes the + # while-not-empty drain loop in build.py to suspend on await get(), + # which yields to the event loop so the fill task can run its first + # XREAD and populate the buffer. After warm-up, delegate to the + # actual buffer state. + return self._first_read_done and self._buffer.empty() + + async def get(self): + return await self._buffer.get() + + def get_nowait(self): + return self._buffer.get_nowait() + + def put_nowait(self, item) -> None: + """No-op: this wrapper is consumer-only; producers write via the bridge.""" + + def fill_done(self) -> bool: + """Return True if the background fill task has finished.""" + return self._fill_task.done() + + async def cancel(self) -> None: + """Cancel the background fill task.""" + if not self._fill_task.done(): + self._fill_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._fill_task + + +class RedisJobQueueService(JobQueueService): + """Redis-backed job queue service for multi-worker deployments. + + Replaces the in-memory :class:`JobQueueService` with one that uses Redis + Streams as a shared event bus, so that build events published by the worker + that started the job can be consumed by any other worker that receives the + subsequent HTTP poll / streaming request. + + Architecture + ------------ + Producer (build worker):: + + EventManager → local asyncio.Queue → bridge coroutine → Redis Stream + + Consumer (poll worker):: + + Redis Stream → RedisQueueWrapper fill task → local buffer → HTTP response + + Configuration + ------------- + Set ``LANGFLOW_JOB_QUEUE_TYPE=redis`` and, optionally: + + * ``LANGFLOW_REDIS_QUEUE_DB`` (default ``1``, separate from cache DB ``0``) + * ``LANGFLOW_REDIS_QUEUE_URL`` (full URL, overrides host/port/db) + * ``LANGFLOW_REDIS_QUEUE_HOST`` / ``LANGFLOW_REDIS_QUEUE_PORT`` + + Known limitations + ----------------- + * Cross-worker *cancel*: cancelling a build running on Worker A from Worker B + silently no-ops (returns success). True cross-worker cancel would require an + additional Redis signal channel checked inside the build loop. + """ + + STREAM_PREFIX = _STREAM_PREFIX + OWNER_PREFIX = _OWNER_PREFIX + + # Bridge retry / TTL tunables — class-level so tests can override without + # patching the loop body. + _MAX_BRIDGE_RETRY_DELAY_S = 4.0 + # Refresh the stream TTL on the first event, then every _TTL_REFRESH_EVENTS + # events *or* every _TTL_REFRESH_SECS seconds, whichever comes first. + # Calling expire() on every XADD doubles Redis round-trips and caps + # single-job throughput; periodic refresh preserves semantics at ~1/100 the cost. + _TTL_REFRESH_EVENTS = 100 + _TTL_REFRESH_SECS = 30.0 + + def __init__( + self, + host: str = "localhost", + port: int = 6379, + db: int = 1, + url: str | None = None, + ttl: int = 3600, + ) -> None: + super().__init__() + self._redis_host = host + self._redis_port = port + self._redis_db = db + self._redis_url = url + self._ttl = ttl + self._client: Any = None + self._connection_check_task: asyncio.Task | None = None + self._bridge_tasks: dict[str, asyncio.Task] = {} + self._consumer_wrappers: dict[str, RedisQueueWrapper] = {} + + def _stream_key(self, job_id: str) -> str: + return f"{self.STREAM_PREFIX}{job_id}" + + def _owner_key(self, job_id: str) -> str: + return f"{self.OWNER_PREFIX}{job_id}" + + def start(self) -> None: + """Create the Redis client and start the periodic cleanup routine.""" + from redis.asyncio import StrictRedis + + if self._redis_url: + self._client = StrictRedis.from_url(self._redis_url) + else: + self._client = StrictRedis(host=self._redis_host, port=self._redis_port, db=self._redis_db) + super().start() + # Schedule a connectivity check so startup logs a clear error if Redis is unreachable. + self._connection_check_task = asyncio.create_task(self._check_connection()) + logger.debug("RedisJobQueueService started.") + + async def _check_connection(self) -> None: + """Ping Redis and log a prominent error if the connection is unavailable.""" + try: + await self._client.ping() + await logger.adebug("RedisJobQueueService: Redis connection OK.") + except Exception as exc: # noqa: BLE001 + await logger.aerror( + f"RedisJobQueueService: cannot reach Redis at " + f"{self._redis_url or f'{self._redis_host}:{self._redis_port} db={self._redis_db}'} — {exc}. " + "Build events will NOT be delivered. " + "Set LANGFLOW_JOB_QUEUE_TYPE=asyncio or start Redis before running Langflow." + ) + + async def stop(self) -> None: + """Stop the service, cancel all bridge tasks, and close the Redis client.""" + if self._connection_check_task and not self._connection_check_task.done(): + self._connection_check_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._connection_check_task + self._connection_check_task = None + + for bridge in list(self._bridge_tasks.values()): + if not bridge.done(): + bridge.cancel() + with contextlib.suppress(asyncio.CancelledError): + await bridge + self._bridge_tasks.clear() + for wrapper in list(self._consumer_wrappers.values()): + await wrapper.cancel() + self._consumer_wrappers.clear() + await super().stop() + if self._client: + await self._client.aclose() + self._client = None + await logger.adebug("RedisJobQueueService stopped.") + + def create_queue(self, job_id: str) -> tuple[asyncio.Queue, EventManager]: + """Create a local queue + EventManager and start the producer bridge to Redis.""" + local_queue, event_manager = super().create_queue(job_id) + bridge = asyncio.create_task(self._bridge_to_redis(job_id, local_queue)) + self._bridge_tasks[job_id] = bridge + return local_queue, event_manager + + async def _bridge_to_redis(self, job_id: str, local_queue: asyncio.Queue) -> None: + """Drain the local queue and publish each event to the Redis Stream. + + Items are read from the local asyncio.Queue (written by EventManager) and + forwarded to a Redis Stream via XADD so that any worker can consume them. + If Redis is temporarily unavailable the item is re-queued and the bridge + backs off before retrying, preventing event loss. + """ + stream_key = self._stream_key(job_id) + _retry_delay = 0.1 + in_flight_item = None + published = False + event_count = 0 + last_ttl_refresh = time.monotonic() + try: + while True: + item = await local_queue.get() + in_flight_item = item + published = False + event_id, data, ts = item + is_sentinel = data is None + fields = ( + {"event_id": "__sentinel__", "data": _STREAM_SENTINEL_DATA, "ts": str(ts)} + if is_sentinel + else {"event_id": event_id or "", "data": data, "ts": str(ts)} + ) + # Refresh TTL on first event, every N events, every M seconds, or on sentinel. + now = time.monotonic() + needs_ttl_refresh = ( + event_count == 0 + or event_count % self._TTL_REFRESH_EVENTS == 0 + or (now - last_ttl_refresh) >= self._TTL_REFRESH_SECS + or is_sentinel + ) + while True: + try: + if not published: + await self._client.xadd(stream_key, fields, maxlen=10_000, approximate=True) + published = True + if needs_ttl_refresh: + await self._client.expire(stream_key, self._ttl) + last_ttl_refresh = time.monotonic() + in_flight_item = None + _retry_delay = 0.1 + break + except asyncio.CancelledError: + raise + except Exception as exc: # noqa: BLE001 + await logger.awarning( + f"Bridge XADD failed for job_id {job_id} (retrying in {_retry_delay}s): {exc}" + ) + await asyncio.sleep(_retry_delay) + _retry_delay = min(_retry_delay * 2, self._MAX_BRIDGE_RETRY_DELAY_S) + event_count += 1 + if is_sentinel: + return + except asyncio.CancelledError: + if in_flight_item is not None and not published: + local_queue.put_nowait(in_flight_item) + return + + def _get_consumer_wrapper(self, job_id: str) -> RedisQueueWrapper: + """Return the cached Redis stream consumer for a job, creating it if needed.""" + wrapper = self._consumer_wrappers.get(job_id) + if wrapper is None: + wrapper = RedisQueueWrapper(job_id, self._client, self._ttl) + self._consumer_wrappers[job_id] = wrapper + return wrapper + + def get_queue_data(self, job_id: str) -> tuple[asyncio.Queue, EventManager, asyncio.Task | None, float | None]: # type: ignore[override] + """Return queue data for a job, always backed by a Redis Stream consumer. + + The queue returned is always a :class:`RedisQueueWrapper` that reads from the + Redis Stream, regardless of whether the job was started on this worker. This + avoids the race condition that would occur if the bridge coroutine and the HTTP + consumer both read from the same local ``asyncio.Queue``. + + * **Same-worker path**: the bridge is the sole reader of the local queue; the + HTTP consumer reads from Redis via the wrapper. The real ``asyncio.Task`` and + ``EventManager`` are returned from the local registry so that disconnect + handling and ownership checks work normally. + * **Cross-worker path**: no local entry exists; a null ``EventManager`` and + ``None`` task are returned (cross-worker cancel is a known limitation). + """ + if self._closed: + msg = f"Queue service is closed for job_id: {job_id}" + raise RuntimeError(msg) + + if job_id in self._queues: + # Same-worker: keep task + event_manager from local registry; give a Redis + # stream wrapper to the consumer so the bridge is the only local-queue reader. + _, event_manager, task, cleanup_time = self._queues[job_id] + return ( # type: ignore[return-value] + self._get_consumer_wrapper(job_id), + event_manager, + task, + cleanup_time, + ) + + # Cross-worker path: create a Redis-backed consumer for the stream. + # EventManager(None) is a null manager — send_event is a no-op (queue is None-guarded). + return ( # type: ignore[return-value] + self._get_consumer_wrapper(job_id), + EventManager(None), + None, + None, + ) + + async def _cleanup_old_queues(self) -> None: + """Run base queue cleanup then sweep done cross-worker consumer wrappers. + + Cross-worker jobs are never inserted into ``self._queues``, so the base + sweep never sees them. Their ``RedisQueueWrapper._fill_task`` exits on + its own (sentinel or ``exists()=False``), but the dict entry in + ``_consumer_wrappers`` stays forever unless we explicitly prune it here. + """ + await super()._cleanup_old_queues() + + # Only prune wrappers that are NOT owned by this worker (i.e. absent from + # self._queues). Same-worker wrappers are removed by cleanup_job(); touching + # them here would race with the grace-period logic in the base class. + done_cross_worker = [ + job_id + for job_id, wrapper in self._consumer_wrappers.items() + if job_id not in self._queues and wrapper.fill_done() + ] + for job_id in done_cross_worker: + wrapper = self._consumer_wrappers.pop(job_id, None) + if wrapper is not None: + await logger.adebug(f"Swept done cross-worker consumer wrapper for job_id {job_id}") + + async def cleanup_job(self, job_id: str) -> None: + """Cancel local task and bridge, then delete the Redis Stream and owner keys.""" + # Capture ownership before super() pops the entry from self._queues so that + # the Redis-key deletion below is scoped to the owning worker only. A + # cross-worker poll populates _consumer_wrappers but not _queues; deleting + # the stream/owner keys from that worker would corrupt an in-flight build on + # the true owner. + is_owner = job_id in self._queues + + wrapper = self._consumer_wrappers.pop(job_id, None) + if wrapper is not None: + await wrapper.cancel() + + bridge = self._bridge_tasks.pop(job_id, None) + if bridge and not bridge.done(): + bridge.cancel() + with contextlib.suppress(asyncio.CancelledError): + await bridge + + try: + await super().cleanup_job(job_id) + finally: + if is_owner and self._client: + await self._client.delete(self._stream_key(job_id), self._owner_key(job_id)) + await logger.adebug(f"Redis keys deleted for job_id {job_id}") + + async def register_job_owner(self, job_id: str, user_id: UUID) -> None: + """Store the job owner in Redis for cross-worker ownership checks.""" + self._job_owners[job_id] = user_id + if self._client: + await self._client.set(self._owner_key(job_id), str(user_id), ex=self._ttl) + + async def get_job_owner(self, job_id: str) -> UUID | None: + """Retrieve the job owner, checking Redis when not found locally. + + The Redis key TTL is refreshed on every successful cross-worker lookup so + that long-running builds (agent loops, large RAG ingests, etc.) do not lose + their ownership anchor mid-flight. The in-memory path is not TTL-bound. + """ + local = self._job_owners.get(job_id) + if local is not None: + return local + if self._client: + owner_key = self._owner_key(job_id) + value = await self._client.get(owner_key) + if value: + # Slide the TTL forward so builds longer than the initial TTL + # continue to pass ownership checks as long as they are polled. + await self._client.expire(owner_key, self._ttl) + return UUID(value.decode()) + return None diff --git a/src/backend/base/pyproject.toml b/src/backend/base/pyproject.toml index 91c5813320..e811f02478 100644 --- a/src/backend/base/pyproject.toml +++ b/src/backend/base/pyproject.toml @@ -135,6 +135,7 @@ dev = [ "pytest-flakefinder>=1.1.0", "types-markdown>=3.7.0.20240822", "codeflash>=0.8.4", + "fakeredis>=2.0.0", ] diff --git a/src/backend/tests/unit/test_chat_endpoint.py b/src/backend/tests/unit/test_chat_endpoint.py index 193d72a882..343c5c907a 100644 --- a/src/backend/tests/unit/test_chat_endpoint.py +++ b/src/backend/tests/unit/test_chat_endpoint.py @@ -903,13 +903,13 @@ async def test_job_owner_cleaned_up_after_cleanup_job(): service.start_job(job_id, _noop()) await asyncio.sleep(0.05) - service.register_job_owner(job_id, user_id) + await service.register_job_owner(job_id, user_id) - assert service.get_job_owner(job_id) == user_id + assert await service.get_job_owner(job_id) == user_id await service.cleanup_job(job_id) - assert service.get_job_owner(job_id) is None + assert await service.get_job_owner(job_id) is None finally: service._closed = True if service._cleanup_task: diff --git a/src/backend/tests/unit/test_redis_job_queue_service.py b/src/backend/tests/unit/test_redis_job_queue_service.py new file mode 100644 index 0000000000..3029f20ec5 --- /dev/null +++ b/src/backend/tests/unit/test_redis_job_queue_service.py @@ -0,0 +1,652 @@ +"""Unit tests for RedisJobQueueService and RedisQueueWrapper. + +Uses fakeredis so no real Redis instance is required. +""" + +from __future__ import annotations + +import asyncio +import contextlib +import time +import uuid +from typing import Any + +import fakeredis.aioredis as fakeredis_aio +import pytest +from langflow.api.build import _CancellableQueue, create_flow_response, get_flow_events_response +from langflow.api.utils import EventDeliveryType +from langflow.events.event_manager import EventManager +from langflow.services.job_queue.service import ( + _STREAM_SENTINEL_DATA, + RedisJobQueueService, + RedisQueueWrapper, +) + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +async def _make_service(ttl: int = 60) -> tuple[RedisJobQueueService, fakeredis_aio.FakeRedis]: + """Return a started RedisJobQueueService backed by a FakeRedis client.""" + fake_client = fakeredis_aio.FakeRedis() + service = RedisJobQueueService(ttl=ttl) + # Inject the fake client directly so no real Redis connection is attempted. + service._client = fake_client + service._closed = False + service._cleanup_task = asyncio.create_task(service._periodic_cleanup()) + service.ready = True + return service, fake_client + + +async def _stop_service(service: RedisJobQueueService) -> None: + service._closed = True + if service._cleanup_task: + service._cleanup_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await service._cleanup_task + for bridge in list(service._bridge_tasks.values()): + if not bridge.done(): + bridge.cancel() + with contextlib.suppress(asyncio.CancelledError): + await bridge + service._bridge_tasks.clear() + for wrapper in list(service._consumer_wrappers.values()): + await wrapper.cancel() + service._consumer_wrappers.clear() + if service._client: + await service._client.aclose() + service._client = None + + +class _BlockingXaddClient: + """Redis client wrapper that blocks inside xadd until the caller cancels it.""" + + def __init__(self, client: fakeredis_aio.FakeRedis) -> None: + self._client = client + self.xadd_started = asyncio.Event() + + def __getattr__(self, name: str) -> Any: + return getattr(self._client, name) + + async def xadd(self, *_args: Any, **_kwargs: Any) -> None: + self.xadd_started.set() + await asyncio.Event().wait() + + +# --------------------------------------------------------------------------- +# RedisQueueWrapper tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_reads_events(): + """RedisQueueWrapper delivers events published to the Redis Stream.""" + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + + # Publish two events followed by the sentinel. + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"hello", "ts": "1.0"}) + await fake_client.xadd(stream_key, {"event_id": "e2", "data": b"world", "ts": "2.0"}) + await fake_client.xadd( + stream_key, + {"event_id": "__sentinel__", "data": _STREAM_SENTINEL_DATA, "ts": "3.0"}, + ) + + event1 = await asyncio.wait_for(wrapper.get(), timeout=5) + event2 = await asyncio.wait_for(wrapper.get(), timeout=5) + sentinel = await asyncio.wait_for(wrapper.get(), timeout=5) + + assert event1 == ("e1", b"hello", 1.0) + assert event2 == ("e2", b"world", 2.0) + assert sentinel == (None, None, 3.0) + + await wrapper.cancel() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_self_terminates_on_key_deletion(): + """RedisQueueWrapper sends the end-of-stream sentinel when the key is deleted.""" + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + + # Publish one event then delete the key (simulates cleanup_job on another worker). + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"data", "ts": "1.0"}) + event1 = await asyncio.wait_for(wrapper.get(), timeout=5) + assert event1[1] == b"data" + + await fake_client.delete(stream_key) + + # The wrapper's fill task should detect the missing key and put the sentinel. + sentinel = await asyncio.wait_for(wrapper.get(), timeout=5) + assert sentinel[0] is None + assert sentinel[1] is None + + await wrapper.cancel() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_empty_reflects_buffer(): + """empty() reports the local buffer state after the first XREAD completes. + + Before the first XREAD returns, empty() always returns False so that the + while-not-empty drain loop in build.py suspends on get() and yields to the + event loop, giving the fill task a chance to populate the buffer. + """ + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + # Before the first XREAD completes, empty() returns False regardless of + # the actual buffer state (warm-up guard for the build.py drain loop). + assert not wrapper.empty() + + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"x", "ts": "1.0"}) + await fake_client.xadd( + stream_key, + {"event_id": "__sentinel__", "data": _STREAM_SENTINEL_DATA, "ts": "2.0"}, + ) + + # Allow the fill task to process the events. + await asyncio.sleep(0.1) + assert not wrapper.empty() + + await wrapper.cancel() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_put_nowait_is_noop(): + """put_nowait on the consumer-side wrapper is a no-op (does not raise). + + The wrapper is consumer-only; producers write via the bridge task. Calling + put_nowait must not add anything to the internal buffer. We check the + underlying _buffer directly so the test is independent of the _first_read_done + warm-up guard that empty() applies. + """ + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + + wrapper.put_nowait(("e1", b"data", 1.0)) + # Check the raw buffer — put_nowait must not have added any item to it. + assert wrapper._buffer.empty() + + await wrapper.cancel() + await fake_client.aclose() + + +# --------------------------------------------------------------------------- +# RedisJobQueueService tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_redis_service_create_and_publish(): + """Events written via EventManager appear in the Redis Stream.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + _queue, event_manager = service.create_queue(job_id) + + async def _build(): + event_manager.on_token(data={"chunk": "hi"}) + await event_manager.queue.put((None, None, time.time())) + + service.start_job(job_id, _build()) + await asyncio.sleep(0.2) + + stream_key = f"langflow:queue:{job_id}" + messages = await fake_client.xrange(stream_key) + assert len(messages) >= 2 # at least one event + sentinel + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_cross_worker_get_queue_data(): + """get_queue_data returns a RedisQueueWrapper for jobs not on this worker.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + + # Simulate another worker having published events to this stream. + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"bytes", "ts": "1.0"}) + await fake_client.xadd( + stream_key, + {"event_id": "__sentinel__", "data": _STREAM_SENTINEL_DATA, "ts": "2.0"}, + ) + + queue, _event_manager, task, _ = service.get_queue_data(job_id) + + assert isinstance(queue, RedisQueueWrapper) + assert task is None + + event = await asyncio.wait_for(queue.get(), timeout=5) + assert event[1] == b"bytes" + sentinel = await asyncio.wait_for(queue.get(), timeout=5) + assert sentinel[0] is None + + await queue.cancel() + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_local_job_returns_redis_wrapper(): + """get_queue_data returns a RedisQueueWrapper even for same-worker jobs. + + The bridge coroutine is the sole reader of the local asyncio.Queue; giving the + HTTP consumer a RedisQueueWrapper prevents the race condition where both would + compete to drain the same queue and events would be lost. + """ + service, _ = await _make_service() + try: + job_id = str(uuid.uuid4()) + service.create_queue(job_id) + + async def _noop(): + await asyncio.sleep(0) + + service.start_job(job_id, _noop()) + + queue, _, task, _ = service.get_queue_data(job_id) + # Consumer always gets a RedisQueueWrapper, never the raw asyncio.Queue + assert isinstance(queue, RedisQueueWrapper) + # Task reference is still available from the local registry + assert task is not None + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_reuses_consumer_wrapper_for_sequential_polls(): + """Repeated get_queue_data calls for a job continue from the last Redis Stream ID.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"first", "ts": "1.0"}) + await fake_client.xadd(stream_key, {"event_id": "e2", "data": b"second", "ts": "2.0"}) + + first_queue, _, _, _ = service.get_queue_data(job_id) + first_event = await asyncio.wait_for(first_queue.get(), timeout=5) + + second_queue, _, _, _ = service.get_queue_data(job_id) + second_event = await asyncio.wait_for(second_queue.get(), timeout=5) + + assert second_queue is first_queue + assert first_event == ("e1", b"first", 1.0) + assert second_event == ("e2", b"second", 2.0) + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_cleanup_cancels_cached_consumer_wrapper(): + """cleanup_job cancels the cached Redis consumer fill task for the job.""" + service, _ = await _make_service() + try: + job_id = str(uuid.uuid4()) + queue, _, _, _ = service.get_queue_data(job_id) + + assert isinstance(queue, RedisQueueWrapper) + assert job_id in service._consumer_wrappers + + await service.cleanup_job(job_id) + + assert job_id not in service._consumer_wrappers + assert queue._fill_task.done() + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_raises_when_closed(): + """get_queue_data raises RuntimeError when the service is closed.""" + service, _ = await _make_service() + service._closed = True + + with pytest.raises(RuntimeError, match="closed"): + service.get_queue_data("some-job-id") + + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_cleanup_deletes_redis_keys(): + """cleanup_job removes the Redis Stream and owner keys.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + service.create_queue(job_id) + + async def _noop(): + await asyncio.sleep(0) + + service.start_job(job_id, _noop()) + await asyncio.sleep(0.05) + + stream_key = f"langflow:queue:{job_id}" + owner_key = f"langflow:owner:{job_id}" + + # Manually create the keys to verify deletion. + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"x", "ts": "1.0"}) + await fake_client.set(owner_key, "some-user") + + await service.cleanup_job(job_id) + + assert not await fake_client.exists(stream_key) + assert not await fake_client.exists(owner_key) + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_cleanup_deletes_redis_keys_when_cancelled(): + """cleanup_job removes Redis keys even when local task cancellation is re-raised.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + service.create_queue(job_id) + + async def _long_running(): + await asyncio.Event().wait() + + service.start_job(job_id, _long_running()) + await asyncio.sleep(0.05) + + stream_key = f"langflow:queue:{job_id}" + owner_key = f"langflow:owner:{job_id}" + + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"x", "ts": "1.0"}) + await fake_client.set(owner_key, "some-user") + + with pytest.raises(asyncio.CancelledError): + await service.cleanup_job(job_id) + + assert not await fake_client.exists(stream_key) + assert not await fake_client.exists(owner_key) + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_owner_stored_in_redis(): + """register_job_owner writes to Redis; get_job_owner reads it cross-worker.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + user_id = uuid.uuid4() + + await service.register_job_owner(job_id, user_id) + + # Verify Redis key was written. + raw = await fake_client.get(f"langflow:owner:{job_id}") + assert raw is not None + assert str(user_id) == raw.decode() + + # Simulate cross-worker lookup: clear in-memory dict. + service._job_owners.clear() + retrieved = await service.get_job_owner(job_id) + assert retrieved == user_id + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_owner_cleaned_up_after_cleanup_job(): + """cleanup_job removes the _job_owners entry and the Redis owner key.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + user_id = uuid.uuid4() + + service.create_queue(job_id) + + async def _noop(): + await asyncio.sleep(0) + + service.start_job(job_id, _noop()) + await asyncio.sleep(0.05) + await service.register_job_owner(job_id, user_id) + + assert await service.get_job_owner(job_id) == user_id + + await service.cleanup_job(job_id) + + assert await service.get_job_owner(job_id) is None + assert not await fake_client.exists(f"langflow:owner:{job_id}") + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_get_job_owner_returns_none_for_unknown_job(): + """get_job_owner returns None for a job_id that was never registered.""" + service, _ = await _make_service() + try: + result = await service.get_job_owner("nonexistent-job-id") + assert result is None + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_bridge_publishes_sentinel_on_end(): + """The bridge task publishes the end-of-stream sentinel when the build ends.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + _queue, event_manager = service.create_queue(job_id) + + async def _build(): + await event_manager.queue.put((None, None, time.time())) + + service.start_job(job_id, _build()) + + # Wait long enough for the bridge to publish. + await asyncio.sleep(0.2) + + stream_key = f"langflow:queue:{job_id}" + messages = await fake_client.xrange(stream_key) + assert messages, "Expected at least one message in the stream" + last_fields = messages[-1][1] + assert last_fields.get(b"data") == _STREAM_SENTINEL_DATA + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_service_bridge_requeues_sentinel_when_cancelled_during_xadd(): + """A cancelled bridge restores an unpublished sentinel instead of dropping it.""" + service, fake_client = await _make_service() + blocking_client = _BlockingXaddClient(fake_client) + service._client = blocking_client + try: + job_id = str(uuid.uuid4()) + local_queue, _event_manager = service.create_queue(job_id) + sentinel = (None, None, time.time()) + + await local_queue.put(sentinel) + await asyncio.wait_for(blocking_client.xadd_started.wait(), timeout=1) + + bridge = service._bridge_tasks[job_id] + bridge.cancel() + with contextlib.suppress(asyncio.CancelledError): + await bridge + + assert local_queue.get_nowait() == sentinel + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_redis_cross_worker_streaming_response_allows_missing_event_task(): + """Streaming cross-worker reads from Redis even when no local build task exists.""" + service, fake_client = await _make_service() + try: + job_id = str(uuid.uuid4()) + stream_key = f"langflow:queue:{job_id}" + await fake_client.xadd(stream_key, {"event_id": "e1", "data": b"payload\n", "ts": "1.0"}) + await fake_client.xadd( + stream_key, + {"event_id": "__sentinel__", "data": _STREAM_SENTINEL_DATA, "ts": "2.0"}, + ) + + response = await get_flow_events_response( + job_id=job_id, + queue_service=service, + event_delivery=EventDeliveryType.STREAMING, + ) + + chunks = [ + chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk async for chunk in response.body_iterator + ] + + assert response.status_code == 200 + assert chunks == ["payload\n"] + finally: + await _stop_service(service) + + +@pytest.mark.asyncio +async def test_streaming_disconnect_cancels_queue_wrapper_without_event_task(): + """Streaming disconnect cleanup uses the Redis wrapper when no local task exists.""" + + class _CancelableQueue(asyncio.Queue): + def __init__(self) -> None: + super().__init__() + self.cancelled = False + + async def cancel(self) -> None: + self.cancelled = True + + queue = _CancelableQueue() + response = await create_flow_response( + queue=queue, + event_manager=EventManager(None), + event_task=None, + ) + + await response.on_disconnect() + + assert queue.cancelled + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_satisfies_cancellable_queue_protocol(): + """RedisQueueWrapper satisfies the _CancellableQueue structural Protocol.""" + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + assert isinstance(wrapper, _CancellableQueue), ( + "RedisQueueWrapper must satisfy the _CancellableQueue Protocol so that " + "on_disconnect() can call wrapper.cancel() via isinstance check." + ) + await wrapper.cancel() + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_done_callback_delivers_sentinel_on_fill_crash(): + """If the fill task crashes, the done callback puts the sentinel in the buffer. + + This guards against consumers hanging indefinitely on ``await queue.get()`` + when an unexpected exception escapes ``_fill_from_redis``. + """ + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + + # Forcibly inject a RuntimeError into the fill task by cancelling it first, + # then directly invoking the done-callback with a fake failed task. + await wrapper.cancel() + + class _FailedTask: + def cancelled(self) -> bool: + return False + + def exception(self) -> BaseException: + return RuntimeError("simulated fill crash") + + # Drain anything already in the buffer so the sentinel is the next item. + while not wrapper._buffer.empty(): + wrapper._buffer.get_nowait() + + wrapper._on_fill_done(_FailedTask()) # type: ignore[arg-type] + + sentinel = wrapper._buffer.get_nowait() + assert sentinel[0] is None + assert sentinel[1] is None + + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_buffer_bounded(): + """The internal buffer respects _BUFFER_MAXSIZE to bound memory usage.""" + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + await wrapper.cancel() + + assert wrapper._buffer.maxsize == RedisQueueWrapper._BUFFER_MAXSIZE + assert wrapper._buffer.maxsize > 0 + + await fake_client.aclose() + + +@pytest.mark.asyncio +async def test_redis_queue_wrapper_done_callback_evicts_to_make_room_for_sentinel(): + """If the buffer is full when the fill task crashes, an oldest item is evicted. + + Losing one event is strictly preferable to leaving the consumer blocked on + ``await get()`` indefinitely. + """ + fake_client = fakeredis_aio.FakeRedis() + job_id = str(uuid.uuid4()) + wrapper = RedisQueueWrapper(job_id, fake_client, ttl=60) + await wrapper.cancel() + + # cancel() delivers the sentinel; drain it so we can test the exception-path + # eviction independently with a fully saturated buffer. + while not wrapper._buffer.empty(): + wrapper._buffer.get_nowait() + + # Saturate the buffer so put_nowait(sentinel) would otherwise fail. + for i in range(wrapper._buffer.maxsize): + wrapper._buffer.put_nowait((f"e{i}", b"data", float(i))) + assert wrapper._buffer.full() + + class _FailedTask: + def cancelled(self) -> bool: + return False + + def exception(self) -> BaseException: + return RuntimeError("simulated fill crash") + + wrapper._on_fill_done(_FailedTask()) # type: ignore[arg-type] + + # The buffer should still be at capacity (one evicted, one sentinel added). + assert wrapper._buffer.full() + + # Drain everything; the sentinel must be present. + items = [] + while not wrapper._buffer.empty(): + items.append(wrapper._buffer.get_nowait()) + sentinels = [item for item in items if item[0] is None and item[1] is None] + assert len(sentinels) == 1, "Exactly one sentinel must reach the consumer" + + await fake_client.aclose() diff --git a/src/lfx/src/lfx/services/settings/base.py b/src/lfx/src/lfx/services/settings/base.py index c9513491ea..6cc1efe8cd 100644 --- a/src/lfx/src/lfx/services/settings/base.py +++ b/src/lfx/src/lfx/services/settings/base.py @@ -247,6 +247,20 @@ class Settings(BaseSettings): redis_url: str | None = None redis_cache_expire: int = 3600 + # Job Queue + job_queue_type: Literal["asyncio", "redis"] = "asyncio" + """The job queue backend. Use 'redis' for multi-worker deployments to solve cross-worker JobQueueNotFoundError.""" + redis_queue_host: str | None = None + """Redis host for the job queue. Falls back to redis_host if not set.""" + redis_queue_port: int | None = None + """Redis port for the job queue. Falls back to redis_port if not set.""" + redis_queue_db: int = 1 + """Redis DB number for the job queue. Defaults to 1 to avoid conflict with the cache (DB 0).""" + redis_queue_url: str | None = None + """Full Redis URL for the job queue. Takes priority over host/port/db if set.""" + redis_queue_ttl: int = 3600 + """TTL in seconds for job stream keys in Redis.""" + # Sentry sentry_dsn: str | None = None sentry_traces_sample_rate: float | None = 1.0 @@ -496,8 +510,9 @@ class Settings(BaseSettings): def set_event_delivery(cls, value, info): # If workers > 1, we need to use direct delivery # because polling and streaming are not supported - # in multi-worker environments - if info.data.get("workers", 1) > 1: + # in multi-worker environments — unless a Redis-backed job queue is configured, + # which shares state across workers and supports all delivery modes. + if info.data.get("workers", 1) > 1 and info.data.get("job_queue_type", "asyncio") != "redis": logger.warning("Multi-worker environment detected, using direct event delivery") return "direct" return value diff --git a/uv.lock b/uv.lock index 52688ba7da..0fc2ea4994 100644 --- a/uv.lock +++ b/uv.lock @@ -3066,6 +3066,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/da/8a/708103325edff16a0b0e004de0d37db8ba216a32713948c64d71f6d4a4c2/faker-40.13.0-py3-none-any.whl", hash = "sha256:c1298fd0d819b3688fb5fd358c4ba8f56c7c8c740b411fd3dbd8e30bf2c05019", size = 1994597, upload-time = "2026-04-06T16:44:53.698Z" }, ] +[[package]] +name = "fakeredis" +version = "2.35.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "redis" }, + { name = "sortedcontainers" }, + { name = "typing-extensions", marker = "python_full_version < '3.11'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/50/b748233c02fa77e5105238190cc9bb58b852eb1c8b1d0763230d3a5b745a/fakeredis-2.35.1.tar.gz", hash = "sha256:5bae5eba7b9d93cb968944ac40936373cf2397ff71667d4b595df65c3d2e413f", size = 189118, upload-time = "2026-04-12T17:05:58.539Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/27/b8b057a23f7777177e92d3a602fd866751b6b45014964548997e92e048fd/fakeredis-2.35.1-py3-none-any.whl", hash = "sha256:67d97e11f562b7870e11e5c30cf182270bfb2dd37f6707dba47cc6d91628d1b9", size = 129678, upload-time = "2026-04-12T17:05:56.86Z" }, +] + [[package]] name = "farama-notifications" version = "0.0.4" @@ -6444,6 +6458,7 @@ dev = [ { name = "elevenlabs", version = "1.58.1", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version == '3.12.*'" }, { name = "elevenlabs", version = "1.59.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version != '3.12.*'" }, { name = "faker" }, + { name = "fakeredis" }, { name = "httpx" }, { name = "hypothesis" }, { name = "ipykernel" }, @@ -6506,6 +6521,7 @@ dev = [ { name = "elevenlabs", marker = "python_full_version != '3.12.*'", specifier = ">=1.52.0" }, { name = "elevenlabs", marker = "python_full_version == '3.12.*'", specifier = "==1.58.1" }, { name = "faker", specifier = ">=37.0.0" }, + { name = "fakeredis", specifier = ">=2.0.0" }, { name = "httpx", specifier = ">=0.28.1" }, { name = "hypothesis", specifier = ">=6.123.17" }, { name = "ipykernel", specifier = ">=6.29.0" }, @@ -7202,6 +7218,7 @@ dev = [ { name = "codeflash" }, { name = "devtools" }, { name = "dictdiffer" }, + { name = "fakeredis" }, { name = "httpx", extra = ["http2"] }, { name = "ipykernel" }, { name = "pre-commit" }, @@ -7542,6 +7559,7 @@ dev = [ { name = "codeflash", specifier = ">=0.8.4" }, { name = "devtools", specifier = ">=0.12.2" }, { name = "dictdiffer", specifier = ">=0.9.0" }, + { name = "fakeredis", specifier = ">=2.0.0" }, { name = "httpx", extras = ["http2"], specifier = ">=0.27" }, { name = "ipykernel", specifier = ">=6.29.0" }, { name = "pre-commit", specifier = ">=3.7.0" },