mirror of
https://github.com/langflow-ai/langflow.git
synced 2026-07-25 10:06:43 +08:00
feat: Redis-backed job queue for multi-worker deployments (#12588)
* feat: implement Redis job queue and fakeredis support
- Refactored the job queue service to support Redis-backed management for cross-worker scaling.
- Added environment variables for configuration:
- `LANGFLOW_JOB_QUEUE_TYPE=redis`
- `LANGFLOW_REDIS_QUEUE_DB=1`
- Updated job ownership methods to be asynchronous for improved concurrency handling.
- Enhanced Redis cache service with namespacing via key prefixes.
- Introduced `fakeredis` for in-memory Redis simulation in testin>
- Added comprehensive unit tests for Redis job queue components.
* fix: run experimental warning for RedisCache usage only once
- Introduced a mechanism to emit a one-time warning for the RedisCache experimental feature during server runtime.
- The warning is logged only if no other worker has already emitted it, ensuring clarity for users regarding the experimental status of RedisCache.
- The implementation includes a temporary file check to prevent multiple warnings across different processes.
* docs: document environment variables for worker management
- Added documentation for LANGFLOW_GUNICORN_PRELOAD to explain preloading for better performance.
- Detailed the use of LANGFLOW_JOB_QUEUE_TYPE for specifying backends (e.g., Redis).
- Included LANGFLOW_REDIS_QUEUE_DB to define the database index for job queues.
- Updated the "High-Load Environments" guide with these optimal configurations.
* docs: updated 'High-load and multi-worker environments' section
* feat: enhance RedisJobQueueService with consumer wrapper management
- Introduced a caching mechanism for Redis stream consumers to optimize job data retrieval.
- Added methods to manage consumer wrappers, ensuring they are reused across sequential polls.
- Implemented cleanup logic to cancel and clear consumer wrappers during job cleanup and service stop.
- Expanded unit tests to verify consumer wrapper reuse and cleanup behavior.
* fix: ensure Redis keys are deleted during job cleanup even on cancellation
- Updated the cleanup_job method in RedisJobQueueService to guarantee Redis keys are removed even if the job cleanup is interrupted by a CancelledError.
- Added a new unit test to verify that Redis keys are deleted correctly when cleanup is called during task cancellation.
* fix: manage connection check task in RedisJobQueueService
- Added handling for the connection check task in the stop method to ensure it is properly cancelled and awaited if still running.
- This change improves resource management and prevents potential issues during service shutdown.
* fix: handle unpublished sentinel requeue on cancellation in RedisJobQueueService
- Updated the job processing logic to ensure that if a job is cancelled during the xadd operation, the unpublished sentinel is requeued instead of being dropped.
- Introduced a new unit test to verify this behavior, ensuring robustness in job handling during cancellations.
* fix: improve Redis job queue service with enhanced configuration and cleanup
- Added atexit cleanup to remove stale temporary files for RedisCache.
- Refactored Redis job queue service to use shared constants for stream prefixes, improving maintainability.
- Updated type hints for better clarity and consistency in RedisQueueWrapper and RedisJobQueueService.
- Enhanced error handling with configurable backoff for transient read failures.
* fix: enhance Redis job queue service with maxlen configuration for xadd
- Updated the xadd method in RedisJobQueueService to include maxlen and approximate parameters, improving stream management and preventing excessive memory usage.
* fix: enhance job ownership retrieval in RedisJobQueueService
- Updated the get_job_owner method to refresh the Redis key TTL on successful lookups, ensuring long-running jobs maintain their ownership anchor.
- Improved code clarity by extracting the owner key into a variable and adding detailed docstring explanations for better understanding of the TTL management.
* fix: improve job ownership handling in cleanup_job method
- Enhanced the cleanup_job method in RedisJobQueueService to accurately capture job ownership before deleting Redis keys, preventing potential data corruption in multi-worker scenarios.
- Added comments for clarity on ownership logic and its implications during job cleanup.
* fix: optimize TTL management in RedisJobQueueService
- Introduced periodic TTL refresh logic in the _bridge_to_redis method to enhance Redis stream management, reducing round-trips and improving throughput.
- Added constants for TTL refresh events and seconds to maintain clarity and configurability.
- Updated event handling to ensure TTL is refreshed appropriately based on event count and time elapsed.
* fix: improve event handling in flow response management
- Removed unnecessary error handling for missing event tasks in get_flow_events_response, allowing for smoother operation when no task exists.
- Updated create_flow_response to handle optional event_task parameter, ensuring proper cleanup during disconnections.
- Added unit tests to verify behavior when event tasks are missing, enhancing robustness in streaming scenarios.
* fix: enhance RedisQueueWrapper with startup grace period and stream observation
- Added a startup grace period to prevent premature end-of-stream signals when the producer has not yet issued its first XADD.
- Introduced a flag to track whether the stream has been observed, improving the handling of early polling scenarios.
- Updated logic to ensure proper handling of stream existence checks and logging for better debugging during job processing.
* fix: implement cleanup for old cross-worker job queues in RedisJobQueueService
- Added a new method to clean up done cross-worker consumer wrappers that are not owned by the current worker, ensuring proper resource management.
- Enhanced the existing cleanup logic to prevent memory leaks by explicitly pruning stale entries from the consumer wrappers dictionary.
- Improved logging to provide better visibility into the cleanup process for cross-worker jobs.
* fix: enhance job handling in JobQueueService with guarded task execution
- Updated the start_job method to accept a Coroutine type for task_coro, ensuring type safety.
- Introduced a new _guarded_task method to wrap job coroutines, guaranteeing that unhandled exceptions emit an error event and write a sentinel to the Redis Stream, improving reliability in job processing.
- Enhanced documentation to clarify the behavior of the new task handling mechanism and its implications for cross-worker consumers.
* fix: improve client disconnection handling in create_flow_response
- Added logging for scenarios where a client disconnects without an associated event_task, clarifying that the producer will continue running until the build completes.
- Documented the limitation regarding cross-worker passive disconnects and the need for a Redis side-channel for proper cancellation, enhancing observability in the logs.
* fix: enhance RedisQueueWrapper to manage initial read state and buffer behavior
- Introduced a flag to track the completion of the first XREAD call, ensuring proper buffer management during the initial read phase.
- Updated the empty method to reflect the state of the buffer accurately, preventing premature exits from the drain loop until the first read is complete.
- Improved documentation to clarify the behavior of the new flag and its impact on job processing.
* fix: clarify RedisQueueWrapper behavior in tests for buffer state and no-op operations
- Updated test documentation to explain the behavior of the empty() method in relation to the first XREAD completion, ensuring accurate understanding of buffer state during job processing.
- Adjusted assertions in tests to reflect the intended behavior of the RedisQueueWrapper, specifically regarding the no-op nature of put_nowait and its impact on the internal buffer.
* fix: update dependencies in pyproject.toml and uv.lock
- Added fakeredis dependency with a minimum version of 2.0.0 to both pyproject.toml and uv.lock files.
- Ensured proper formatting and comments in pyproject.toml for clarity on onnxruntime version constraints.
* fix: enhance RedisQueueWrapper and job queue handling for improved cancellation and buffer management
- Introduced a structural protocol `_CancellableQueue` to ensure queues can handle cancellation properly during client disconnects.
- Updated `RedisQueueWrapper` to implement this protocol, allowing for graceful cancellation of background tasks.
- Added a maximum size limit to the internal buffer to prevent unbounded memory usage and ensure backpressure on slow consumers.
- Implemented a done callback to handle unexpected fill task crashes, ensuring consumers are not left hanging indefinitely.
- Enhanced unit tests to verify compliance with the new protocol and the behavior of the buffer under various conditions.
* fix: RedisQueueWrapper sentinel delivery, error time-bound, and cancel response
- Deliver end-of-stream sentinel on fill-task cancellation (_on_fill_done now
handles both cancelled and exception paths so consumers are never left hanging)
- Add _error_start time-bound to xread and exists() error loops: after
_STARTUP_GRACE_S seconds of continuous Redis errors the sentinel is delivered
instead of retrying forever
- Advance _last_id cursor only after buffer.put() succeeds so cancellation mid-put
does not silently skip that message in the Redis cursor
- Return False from cancel_flow_build when event_task is None (cross-worker path)
so the HTTP response correctly reports success=False instead of false success
* [autofix.ci] apply automated fixes
---------
Co-authored-by: ogabrielluiz <gabriel@langflow.org>
Co-authored-by: Jordan Frazier <jordan.frazier@datastax.com>
Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com>
This commit is contained in:
@ -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
|
||||
|
||||
@ -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
|
||||
|
||||
@ -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]]
|
||||
|
||||
@ -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:<job_id> 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")
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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:
|
||||
|
||||
@ -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()
|
||||
|
||||
@ -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 <str> data <bytes> ts <float>``
|
||||
* End-of-stream → ``XADD key * event_id __sentinel__ data __sentinel__ ts <float>``
|
||||
|
||||
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
|
||||
|
||||
@ -135,6 +135,7 @@ dev = [
|
||||
"pytest-flakefinder>=1.1.0",
|
||||
"types-markdown>=3.7.0.20240822",
|
||||
"codeflash>=0.8.4",
|
||||
"fakeredis>=2.0.0",
|
||||
]
|
||||
|
||||
|
||||
|
||||
@ -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:
|
||||
|
||||
652
src/backend/tests/unit/test_redis_job_queue_service.py
Normal file
652
src/backend/tests/unit/test_redis_job_queue_service.py
Normal file
@ -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()
|
||||
@ -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
|
||||
|
||||
18
uv.lock
generated
18
uv.lock
generated
@ -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" },
|
||||
|
||||
Reference in New Issue
Block a user