Explicit Queue Listening (#537)

By default, a process running DBOS dequeues from all declared queues.
Now, you can instead use `DBOS.listen_queues` to explicitly tell a
process running DBOS to only dequeue workflows from a specific set of
queues.

```python
DBOS.listen_queues(
  queues: List[Queue]
)
```

For example, you can use this to manage heterogeneous workers, so
workers of type A only dequeue and execute workflows from queue A while
workers of type B only dequeue and execute workflows from queue B. For
example:

```python
cpu_queue = Queue("queue_one")
gpu_queue = Queue("queue_two")

if __name__ == "__main__":
    config = ...
    worker_type = ... # "cpu' or 'gpu'
    DBOS(config=config)
    if worker_type = "gpu":
        # GPU workers will only dequeue and execute workflows from the GPU queue
        DBOS.listen_queues([gpu_queue])
    elif worker_type == "cpu":
        # CPU workers will only dequeue and execute workflows from the CPU queue
        DBOS.listen_queues([cpu_queue])
    DBOS.launch()
```

Also add logging of all queues being listened to.

Author: kraftp

dbos/_dbos.py

@@ -326,6 +326,7 @@ def __init__(
         self._app_db_field: Optional[ApplicationDatabase] = None
         self._registry: DBOSRegistry = _get_or_create_dbos_registry()
         self._registry.dbos = self
+        self._listening_queues: Optional[List[Queue]] = None
         self._admin_server_field: Optional[AdminServer] = None
         # Stop internal background threads (queue thread, timeout threads, etc.)
         self.background_thread_stop_events: List[threading.Event] = []
@@ -1576,6 +1577,21 @@ def tracer(self) -> DBOSTracer:
         """Return the DBOS OpenTelemetry tracer."""
         return dbos_tracer
 
+    @classmethod
+    def listen_queues(cls, queues: List[Queue]) -> None:
+        """
+        Configure this DBOS process to only listen to (dequeue workflows from) specific queues.
+
+        Args:
+            queues(List[Queue]): The list of queues to listen to
+        """
+        dbos = _get_dbos_instance()
+        if dbos._launched:
+            raise DBOSException("listen_queues called after DBOS is launched")
+        if dbos._listening_queues is not None:
+            raise DBOSException("listen_queues called more than once")
+        dbos._listening_queues = queues
+
 
 class WorkflowHandle(Generic[R], Protocol):
     """

dbos/_queue.py

@@ -175,9 +175,25 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
     queue_threads: dict[str, threading.Thread] = {}
     check_interval = 1.0  # Check for new queues every second
 
+    if dbos._listening_queues is not None:
+        listening_queues = dbos._listening_queues
+    else:
+        listening_queues = list(dbos._registry.queue_info_map.values())
+        listening_queues = [
+            q for q in listening_queues if q.name != INTERNAL_QUEUE_NAME
+        ]
+    dbos.logger.info(f"Listening to {len(listening_queues)} queues:")
+    log_queues(listening_queues)
+
     while not stop_event.is_set():
-        # Check for new queues
-        current_queues = dict(dbos._registry.queue_info_map)
+        if dbos._listening_queues is not None:
+            # If explicitly listening for queues, only use those queues
+            current_queues = {queue.name: queue for queue in dbos._listening_queues}
+            # Always listen to the internal queue
+            current_queues[INTERNAL_QUEUE_NAME] = dbos._registry.get_internal_queue()
+        else:
+            # Else, check all declared queues
+            current_queues = dict(dbos._registry.queue_info_map)
 
         # Start threads for new queues
         for queue_name, queue in current_queues.items():
@@ -212,3 +228,21 @@ def queue_thread(stop_event: threading.Event, dbos: "DBOS") -> None:
                 dbos.logger.debug(
                     f"Queue worker thread for {queue_name} stopped successfully"
                 )
+
+
+def log_queues(queues: list[Queue]) -> None:
+    """Helper function to log queues on DBOS launch."""
+    for q in queues:
+        opts = []
+        if q.concurrency is not None:
+            opts.append(f"concurrency={q.concurrency}")
+        if q.worker_concurrency is not None:
+            opts.append(f"worker_concurrency={q.worker_concurrency}")
+        if q.limiter is not None:
+            opts.append(f"limit={q.limiter['limit']}/{q.limiter['period']}s")
+        if q.priority_enabled:
+            opts.append("priority")
+        if q.partition_queue:
+            opts.append("partitioned")
+        opts_str = f" ({', '.join(opts)})" if opts else ""
+        dbos_logger.info(f"Queue: {q.name}{opts_str}")

tests/test_queue.py

@@ -1752,3 +1752,35 @@ def workflow() -> str:
         start_time = time.time()
         assert queue.enqueue(workflow).get_result(polling_interval_sec=0.1)
         assert time.time() - start_time < 1.0
+
+
+def test_listen_queue(dbos: DBOS, config: DBOSConfig) -> None:
+    DBOS.destroy(destroy_registry=True)
+    DBOS(config=config)
+
+    queue_one = Queue("queue_one")
+    queue_two = Queue("queue_two")
+
+    @DBOS.workflow()
+    def workflow() -> str:
+        assert DBOS.workflow_id
+        return DBOS.workflow_id
+
+    DBOS.listen_queues([queue_one])
+    DBOS.launch()
+
+    # While only listening to queue one, only workflows enqueued there execute
+    handle_one = queue_one.enqueue(workflow)
+    handle_two = queue_two.enqueue(workflow)
+    assert handle_one.get_result()
+    assert handle_two.get_status().status == "ENQUEUED"
+
+    DBOS.destroy()
+    DBOS(config=config)
+    DBOS.listen_queues([queue_two])
+    DBOS.launch()
+
+    # Listening to queue two completes its workflows
+    assert DBOS.retrieve_workflow(handle_two.workflow_id).get_result()
+    # Verify the internal queue works
+    assert DBOS.fork_workflow(handle_two.workflow_id, 0).get_result()