nv-ingest 2025.10.24.dev20251024__tar.gz → 2025.11.9.dev20251109__tar.gz

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nv-ingest
-Version: 2025.10.24.dev20251024
+Version: 2025.11.9.dev20251109
 Summary: Python module for multimodal document ingestion
 Author-email: Jeremy Dyer <jdyer@nvidia.com>
 License:                                  Apache License

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/api/v2/README.md

@@ -11,15 +11,6 @@ The V2 API introduces automatic PDF splitting at the REST layer to improve proce
 3. **Transparent Aggregation**: Results are automatically aggregated when fetching parent jobs
 4. **Backward Compatible**: PDFs with page counts ≤ `PDF_SPLIT_PAGE_COUNT` behave identical to V1
 
-## Tracing & Aggregated Metadata
-
-- V2 endpoints open an OpenTelemetry span using the shared `traced_endpoint` decorator. The span name defaults to the function name, or can be overridden when applying the decorator.
-- `submit_job_v2` records the parent span's `trace_id` into each subjob's `tracing_options`, enabling downstream Ray stages (e.g., the message broker sink) to attach chunk-level telemetry consistently.
-- Response headers still return `x-trace-id` derived from the active span context, allowing clients to correlate downstream work.
-- When `/v2/fetch_job/{parent_id}` aggregates completed chunks, it captures any `trace` / `annotations` dictionaries emitted by the sink for each subjob and includes them in the response payload (see "Aggregated response" below).
-
-This behaviour matches the V1 tracing model and sets the foundation for adding W3C `traceparent` propagation in future changes.
-
 ## How It Works
 
 1. **Submit**: When a PDF with pages exceeding `PDF_SPLIT_PAGE_COUNT` is submitted to `/v2/submit_job`:
@@ -36,6 +27,33 @@ This behaviour matches the V1 tracing model and sets the foundation for adding W
    - Pending work returns 202 (processing)
    - Failed chunks are noted without failing the entire job; metadata records which chunks failed
 
+
+## Client Library Features
+
+### Accessing Trace Metrics
+
+The Python client library provides convenient access to trace metrics via the `return_traces` parameter:
+
+```python
+from nv_ingest_client.client import Ingestor
+
+ingestor = Ingestor(
+    message_client_hostname="localhost",
+    message_client_port=7670,
+    message_client_kwargs={"api_version": "v2"}
+).files("/path/to/pdfs").extract().embed()
+
+# Get results with trace metrics
+results, traces = ingestor.ingest(return_traces=True)
+
+# Access timing for first document
+pdf_time = traces[0]["trace::resident_time::pdf_extractor"] / 1e9
+table_time = traces[0]["trace::resident_time::table_extractor"] / 1e9
+print(f"PDF: {pdf_time:.2f}s, Tables: {table_time:.2f}s")
+```
+
+**Note:** For split PDFs, `resident_time` represents aggregated compute time across all chunks. For non-split PDFs, it is computed client-side from entry/exit pairs.
+
 ### Aggregated response
 
 The fetch endpoint returns a JSON body shaped like the following:
@@ -163,15 +181,23 @@ For split PDFs, parent-level metrics are automatically computed for each stage (
 - Failed chunk entries remain in `failed_subjobs`; missing chunks indicate the sink did not emit telemetry
 - **To access chunk traces:** Use `metadata.trace_segments[]` - each segment contains the full trace dict for that chunk
 
-## Testing
+### Advanced: Accessing Full Metadata
 
-Use the V2 test script with environment variable:
-```bash
-# Run with V2 endpoints
-DATASET_DIR=/data/splits python scripts/tests/cases/dc20_v2_e2e.py
-```
+For advanced use cases requiring per-chunk trace breakdown or full metadata, use `include_parent_trace_ids`:
+
+```python
+results, traces, parent_trace_ids = ingestor.ingest(
+    return_traces=True,
+    include_parent_trace_ids=True
+)
+
+# Fetch full parent job metadata (including trace_segments)
+import requests
+response = requests.get(f"http://localhost:7670/v2/fetch_job/{parent_trace_ids[0]}")
+metadata = response.json()["metadata"]
 
-Or set the API version for any existing code:
-```bash
-export NV_INGEST_API_VERSION=v2
+# Access per-chunk traces
+for segment in metadata["trace_segments"]:
+    print(f"Chunk {segment['chunk_index']}: pages {segment['start_page']}-{segment['end_page']}")
+    print(f"  Traces: {len(segment['trace'])} entries")
 ```

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/api/v2/ingest.py

@@ -12,6 +12,7 @@ import logging
 import os
 import time
 import uuid
+import random
 
 from fastapi import APIRouter, Request, Response
 from fastapi import HTTPException
@@ -44,6 +45,42 @@ router = APIRouter()
 
 DEFAULT_PDF_SPLIT_PAGE_COUNT = 32
 
+# Default QoS thresholds (pages). Tunable via environment variables:
+# QOS_MAX_PAGES_MICRO, QOS_MAX_PAGES_SMALL, QOS_MAX_PAGES_MEDIUM
+_QOS_DEFAULTS = {
+    "micro": 8,
+    "small": 64,
+    "medium": 256,
+}
+
+
+def get_qos_tier_for_page_count(page_count: int) -> str:
+    """
+    Select QoS tier for a document based on its total page count.
+    Tiers: 'micro', 'small', 'medium', 'large', 'default'
+    Thresholds can be tuned via environment variables:
+      - QOS_MAX_PAGES_MICRO (default: 4)
+      - QOS_MAX_PAGES_SMALL (default: 16)
+      - QOS_MAX_PAGES_MEDIUM (default: 64)
+    Anything above MEDIUM is 'large'. Non-positive page_count returns 'default'.
+    """
+    try:
+        micro_max = int(os.getenv("QOS_MAX_PAGES_MICRO", str(_QOS_DEFAULTS["micro"])))
+        small_max = int(os.getenv("QOS_MAX_PAGES_SMALL", str(_QOS_DEFAULTS["small"])))
+        medium_max = int(os.getenv("QOS_MAX_PAGES_MEDIUM", str(_QOS_DEFAULTS["medium"])))
+    except ValueError:
+        micro_max, small_max, medium_max = _QOS_DEFAULTS["micro"], _QOS_DEFAULTS["small"], _QOS_DEFAULTS["medium"]
+
+    if page_count <= 0:
+        return "default"
+    if page_count <= micro_max:
+        return "micro"
+    if page_count <= small_max:
+        return "small"
+    if page_count <= medium_max:
+        return "medium"
+    return "large"
+
 
 def get_pdf_split_page_count(client_override: Optional[int] = None) -> int:
     """
@@ -432,6 +469,76 @@ def _extract_ray_telemetry(result: Dict[str, Any]) -> Tuple[Optional[Dict[str, A
     return trace_dict, annotations_dict
 
 
+def _normalize_chunk_records(
+    records: Optional[List[Any]],
+    descriptor: Dict[str, Any],
+    parent_metadata: Dict[str, Any],
+) -> List[Any]:
+    """Re-map chunk-local metadata to document-level context for aggregation."""
+
+    if not isinstance(records, list):
+        return []
+
+    total_pages = parent_metadata.get("total_pages")
+    original_source_id = parent_metadata.get("original_source_id")
+    original_source_name = parent_metadata.get("original_source_name")
+
+    start_page = descriptor.get("start_page")
+    page_offset = start_page - 1 if isinstance(start_page, int) and start_page > 0 else 0
+
+    normalized_entries: List[Any] = []
+
+    for entry in records:
+        if not isinstance(entry, dict):
+            normalized_entries.append(entry)
+            continue
+
+        normalized_entry = entry.copy()
+        original_metadata = entry.get("metadata")
+
+        if isinstance(original_metadata, dict):
+            normalized_metadata = original_metadata.copy()
+            normalized_entry["metadata"] = normalized_metadata
+
+            original_source_meta = original_metadata.get("source_metadata")
+            if isinstance(original_source_meta, dict):
+                normalized_source_meta = original_source_meta.copy()
+                normalized_metadata["source_metadata"] = normalized_source_meta
+
+                if original_source_id:
+                    normalized_source_meta["source_id"] = original_source_id
+                if original_source_name:
+                    normalized_source_meta["source_name"] = original_source_name
+
+            original_content_meta = original_metadata.get("content_metadata")
+            if isinstance(original_content_meta, dict):
+                normalized_content_meta = original_content_meta.copy()
+                normalized_metadata["content_metadata"] = normalized_content_meta
+
+                page_number = normalized_content_meta.get("page_number")
+                if isinstance(page_number, int) and page_number >= 0:
+                    normalized_content_meta["page_number"] = page_number + page_offset
+
+                if isinstance(total_pages, int) and isinstance(normalized_content_meta.get("page_count"), int):
+                    # Ensure optional per-record page count reflects the full document
+                    normalized_content_meta["page_count"] = total_pages
+
+                original_hierarchy = original_content_meta.get("hierarchy")
+                if isinstance(original_hierarchy, dict):
+                    normalized_hierarchy = original_hierarchy.copy()
+                    normalized_content_meta["hierarchy"] = normalized_hierarchy
+
+                    hierarchy_page = normalized_hierarchy.get("page")
+                    if isinstance(hierarchy_page, int) and hierarchy_page >= 0:
+                        normalized_hierarchy["page"] = hierarchy_page + page_offset
+                    if isinstance(total_pages, int):
+                        normalized_hierarchy["page_count"] = total_pages
+
+        normalized_entries.append(normalized_entry)
+
+    return normalized_entries
+
+
 def _aggregate_parent_traces(chunk_traces: Dict[str, Any]) -> Dict[str, Any]:
     """
     Aggregate chunk-level traces into parent-level metrics.
@@ -574,7 +681,8 @@ def _build_aggregated_response(
         if result is not None:
             # Add page data to aggregated result
             if "data" in result:
-                aggregated_result["data"].extend(result["data"])
+                normalized_records = _normalize_chunk_records(result.get("data"), descriptor, metadata)
+                aggregated_result["data"].extend(normalized_records)
             chunk_entry = dict(descriptor)
             aggregated_result["metadata"]["chunks"].append(chunk_entry)
 
@@ -631,6 +739,51 @@ def _build_aggregated_response(
     return aggregated_result
 
 
+# ---------------------------------------------------------------------------
+# Bursty submission helpers (fairness without long-lived in-flight tasks)
+# ---------------------------------------------------------------------------
+
+
+def _get_submit_burst_params() -> Tuple[int, int, int]:
+    """
+    Returns (burst_size, pause_ms, jitter_ms) from environment with sane defaults.
+    - V2_SUBMIT_BURST_SIZE (default: 16)
+    - V2_SUBMIT_BURST_PAUSE_MS (default: 25)
+    - V2_SUBMIT_BURST_JITTER_MS (default: 10)
+    """
+    burst_size = int(os.getenv("V2_SUBMIT_BURST_SIZE", "16"))
+    pause_ms = int(os.getenv("V2_SUBMIT_BURST_PAUSE_MS", "50"))
+    jitter_ms = int(os.getenv("V2_SUBMIT_BURST_JITTER_MS", "15"))
+
+    return max(1, burst_size), max(0, pause_ms), max(0, jitter_ms)
+
+
+async def _submit_subjobs_in_bursts(
+    items: List[Tuple[str, MessageWrapper]],
+    ingest_service: "INGEST_SERVICE_T",
+    *,
+    burst_size: int,
+    pause_ms: int,
+    jitter_ms: int,
+) -> None:
+    """
+    Submit subjobs in sequential bursts and await each burst to completion.
+    This avoids keeping a large number of pending tasks in the REST handler
+    and allows other concurrent requests to interleave enqueue work between bursts.
+    """
+    for offset in range(0, len(items), burst_size):
+        burst = items[offset : offset + burst_size]
+        tasks = [ingest_service.submit_job(wrapper, subjob_id) for (subjob_id, wrapper) in burst]
+        # Propagate any errors from this burst
+        await asyncio.gather(*tasks)
+
+        # Pause with jitter to yield to other request handlers before next burst
+        if offset + burst_size < len(items):
+            delay_ms = pause_ms + (random.randint(0, jitter_ms) if jitter_ms > 0 else 0)
+            if delay_ms > 0:
+                await asyncio.sleep(delay_ms / 1000.0)
+
+
 # POST /v2/submit_job
 @router.post(
     "/submit_job",
@@ -681,22 +834,24 @@ async def submit_job_v2(
             pdf_content = base64.b64decode(payloads[0])
             page_count = get_pdf_page_count(pdf_content)
             pdf_page_count_cache = page_count  # Cache for later use
+            qos_tier = get_qos_tier_for_page_count(page_count)
             pages_per_chunk = get_pdf_split_page_count(client_override=client_split_page_count)
 
             # Split if the document has more pages than our chunk size
             if page_count > pages_per_chunk:
                 logger.warning(
-                    "Splitting PDF %s into %s-page chunks (total pages: %s)",
+                    "Splitting PDF %s into %s-page chunks (total pages: %s) -> (qos_tier: %s)",
                     original_source_name,
                     pages_per_chunk,
                     page_count,
+                    qos_tier,
                 )
 
                 chunks = split_pdf_to_chunks(pdf_content, pages_per_chunk)
 
                 subjob_ids: List[str] = []
                 subjob_descriptors: List[Dict[str, Any]] = []
-                submission_tasks = []
+                submission_items: List[Tuple[str, MessageWrapper]] = []
 
                 try:
                     parent_uuid = uuid.UUID(parent_job_id)
@@ -717,7 +872,19 @@ async def submit_job_v2(
                         original_source_id=original_source_id,
                         original_source_name=original_source_name,
                     )
-                    submission_tasks.append(ingest_service.submit_job(subjob_wrapper, subjob_id))
+
+                    # Inject QoS routing hint into subjob routing_options (keeps API and service loosely coupled)
+                    try:
+                        sub_spec = json.loads(subjob_wrapper.payload)
+                        routing_opts = sub_spec.get("routing_options") or {}
+                        routing_opts["queue_hint"] = qos_tier
+                        sub_spec["routing_options"] = routing_opts
+                        subjob_wrapper = MessageWrapper(payload=json.dumps(sub_spec))
+                    except Exception:
+                        # Best-effort; if we cannot inject, fall back to default routing
+                        pass
+
+                    submission_items.append((subjob_id, subjob_wrapper))
                     subjob_ids.append(subjob_id)
                     subjob_descriptors.append(
                         {
@@ -729,8 +896,15 @@ async def submit_job_v2(
                         }
                     )
 
-                if submission_tasks:
-                    await asyncio.gather(*submission_tasks)
+                if submission_items:
+                    burst_size, pause_ms, jitter_ms = _get_submit_burst_params()
+                    await _submit_subjobs_in_bursts(
+                        submission_items,
+                        ingest_service,
+                        burst_size=burst_size,
+                        pause_ms=pause_ms,
+                        jitter_ms=jitter_ms,
+                    )
 
                 parent_metadata: Dict[str, Any] = {
                     "total_pages": page_count,
@@ -758,6 +932,16 @@ async def submit_job_v2(
         if "tracing_options" not in job_spec_dict:
             job_spec_dict["tracing_options"] = {"trace": True}
         job_spec_dict["tracing_options"]["trace_id"] = str(current_trace_id)
+        # If this was a PDF and we computed page_count, route the single job using the same QoS tier
+        try:
+            if (
+                document_types
+                and document_types[0].lower() == "pdf"
+                and "queue_hint" not in (job_spec_dict.get("routing_options") or {})
+            ):
+                job_spec_dict.setdefault("routing_options", {})["queue_hint"] = qos_tier
+        except Exception:
+            pass
         updated_job_spec = MessageWrapper(payload=json.dumps(job_spec_dict))
 
         span.add_event("Submitting as single job (no split needed)")

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/framework/orchestration/process/dependent_services.py

@@ -18,6 +18,18 @@ from nv_ingest_api.util.message_brokers.simple_message_broker.broker import Simp
 logger = logging.getLogger(__name__)
 
 
+def _broker_server_target(host, port, max_queue_size):
+    """
+    Target function to be run in a separate process for the SimpleMessageBroker.
+    """
+    server = SimpleMessageBroker(host, port, max_queue_size)
+    try:
+        server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    except Exception:
+        pass
+    server.serve_forever()
+
+
 def start_simple_message_broker(broker_client: dict) -> multiprocessing.Process:
     """
     Starts a SimpleMessageBroker server in a separate process.
@@ -58,16 +70,11 @@ def start_simple_message_broker(broker_client: dict) -> multiprocessing.Process:
             f"continuing to spawn a broker process (tests expect a Process to be returned)"
         )
 
-    def broker_server():
-        # Optionally, set socket options here for reuse (note: binding occurs in server __init__).
-        server = SimpleMessageBroker(server_host, server_port, max_queue_size)
-        try:
-            server.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
-        except Exception:
-            pass
-        server.serve_forever()
-
-    p = multiprocessing.Process(target=broker_server)
+    p = multiprocessing.Process(
+        target=_broker_server_target,
+        args=(server_host, server_port, max_queue_size),
+        daemon=True,
+    )
     # If we're launching from inside the pipeline subprocess, mark daemon so the
     # broker dies automatically when the subprocess exits.
     p.daemon = os.environ.get("NV_INGEST_BROKER_IN_SUBPROCESS") == "1"

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/framework/orchestration/process/strategies.py

@@ -11,9 +11,10 @@ Strategy pattern for clean separation of execution concerns.
 """
 
 import atexit
-import os
 import logging
 import multiprocessing
+import os
+import sys
 import time
 from abc import ABC, abstractmethod
 
@@ -132,7 +133,10 @@ class SubprocessStrategy(ProcessExecutionStrategy):
         logger.info("Launching pipeline in Python subprocess using multiprocessing.")
 
         # Create subprocess using fork context
-        ctx = multiprocessing.get_context("fork")
+        start_method = "fork"
+        if sys.platform.lower() == "darwin":
+            start_method = "spawn"
+        ctx = multiprocessing.get_context(start_method)
         process = ctx.Process(
             target=run_pipeline_process,
             args=(

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/framework/orchestration/process/termination.py

@@ -19,20 +19,45 @@ logger = logging.getLogger(__name__)
 
 
 def _safe_log(level: int, msg: str) -> None:
-    """Best-effort logging that won't crash during interpreter shutdown."""
+    """Best-effort logging that won't emit handler tracebacks on closed streams.
+
+    Temporarily disables logging.raiseExceptions to prevent the logging module
+    from printing "--- Logging error ---" to stderr if a handler's stream is
+    already closed (common during process teardown). Falls back to writing to
+    sys.__stderr__ if available.
+    """
     try:
-        logger.log(level, msg)
+        import logging as _logging
+
+        prev = getattr(_logging, "raiseExceptions", True)
+        # Suppress handler errors being printed to stderr
+        _logging.raiseExceptions = False
+
+        # If there are no handlers, skip and use stderr fallback
+        if logger.handlers:
+            logger.log(level, msg)
+            return
     except Exception:
+        # Intentionally ignore and try stderr fallback
+        pass
+    finally:
         try:
-            # Fallback to stderr if available
-            import sys
+            import logging as _logging  # re-import safe even if earlier failed
 
-            if hasattr(sys, "__stderr__") and sys.__stderr__:
-                sys.__stderr__.write(msg + "\n")
-                sys.__stderr__.flush()
+            _logging.raiseExceptions = prev  # type: ignore[name-defined]
         except Exception:
             pass
 
+    # Fallback to stderr if available
+    try:
+        import sys
+
+        if hasattr(sys, "__stderr__") and sys.__stderr__:
+            sys.__stderr__.write(msg + "\n")
+            sys.__stderr__.flush()
+    except Exception:
+        pass
+
 
 def kill_pipeline_process_group(process) -> None:
     """
@@ -74,7 +99,17 @@ def kill_pipeline_process_group(process) -> None:
 
     try:
         # Send graceful termination to the entire process group
-        os.killpg(os.getpgid(pid), signal.SIGTERM)
+        try:
+            pgid = os.getpgid(pid)
+        except Exception:
+            # Process already gone
+            _safe_log(logging.DEBUG, f"Process group for PID {pid} not found during SIGTERM phase")
+            return
+        try:
+            os.killpg(pgid, signal.SIGTERM)
+        except ProcessLookupError:
+            _safe_log(logging.DEBUG, f"Process group for PID {pid} no longer exists (SIGTERM)")
+            return
 
         # If we have a Process handle, give it a chance to exit cleanly
         if proc is not None and hasattr(proc, "join"):
@@ -95,7 +130,12 @@ def kill_pipeline_process_group(process) -> None:
         if still_alive:
             _safe_log(logging.WARNING, "Process group did not terminate gracefully, using SIGKILL")
             try:
-                os.killpg(os.getpgid(pid), signal.SIGKILL)
+                try:
+                    pgid2 = os.getpgid(pid)
+                except Exception:
+                    _safe_log(logging.DEBUG, f"Process group for PID {pid} vanished before SIGKILL")
+                    return
+                os.killpg(pgid2, signal.SIGKILL)
             finally:
                 if proc is not None and hasattr(proc, "join"):
                     try:

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/framework/orchestration/ray/examples/pipeline_test_harness.py

@@ -152,11 +152,11 @@ if __name__ == "__main__":
     os.environ["OCR_MODEL_NAME"] = "paddle"
     os.environ["NEMORETRIEVER_PARSE_HTTP_ENDPOINT"] = "https://integrate.api.nvidia.com/v1/chat/completions"
     os.environ["VLM_CAPTION_ENDPOINT"] = "https://integrate.api.nvidia.com/v1/chat/completions"
-    os.environ["VLM_CAPTION_MODEL_NAME"] = "nvidia/llama-3.1-nemotron-nano-vl-8b-v1"
+    os.environ["VLM_CAPTION_MODEL_NAME"] = "nvidia/nemotron-nano-12b-v2-vl"
     logger.info("Environment variables set.")
 
     image_caption_endpoint_url = "https://integrate.api.nvidia.com/v1/chat/completions"
-    model_name = "nvidia/llama-3.1-nemotron-nano-vl-8b-v1"
+    model_name = "nvidia/nemotron-nano-12b-v2-vl"
     yolox_grpc, yolox_http, yolox_auth, yolox_protocol = get_nim_service("yolox")
     (
         yolox_table_structure_grpc,

{nv_ingest-2025.10.24.dev20251024 → nv_ingest-2025.11.9.dev20251109}/nv_ingest/framework/orchestration/ray/stages/sources/message_broker_task_source.py

@@ -30,6 +30,7 @@ from nv_ingest_api.internal.schemas.meta.ingest_job_schema import validate_inges
 from nv_ingest_api.util.message_brokers.simple_message_broker.simple_client import SimpleClient
 from nv_ingest_api.util.service_clients.redis.redis_client import RedisClient
 from nv_ingest_api.util.logging.sanitize import sanitize_for_logging
+from nv_ingest_api.util.message_brokers.qos_scheduler import QosScheduler
 
 logger = logging.getLogger(__name__)
 
@@ -89,8 +90,10 @@ class MessageBrokerTaskSourceConfig(BaseModel):
 
     # Use the discriminated union for broker_client
     broker_client: Union[RedisClientConfig, SimpleClientConfig] = Field(..., discriminator="client_type")
-    task_queue: str = Field(..., description="The name of the queue to fetch tasks from.")
-    poll_interval: float = Field(default=0.1, gt=0, description="Polling interval in seconds.")
+    task_queue: str = Field(
+        ..., description="The base name of the queue to fetch tasks from. Derives sub-queues for fair scheduling."
+    )
+    poll_interval: float = Field(default=0.0, gt=0, description="Polling interval in seconds.")
 
 
 @ray.remote
@@ -134,7 +137,29 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         self._current_backoff_sleep: float = 0.0
         self._last_backoff_log_time: float = 0.0
 
-        self._logger.debug("MessageBrokerTaskSourceStage initialized. Task queue: %s", self.task_queue)
+        # Initialize QoS scheduler. Use a simple base-queue strategy for SimpleClient.
+        strategy = "simple" if isinstance(self.client, SimpleClient) else "lottery"
+        self.scheduler = QosScheduler(
+            self.task_queue,
+            num_prefetch_threads=6,  # one per category (no-op for simple strategy)
+            total_buffer_capacity=96,  # e.g., ~16 per thread
+            prefetch_poll_interval=0.002,  # faster polling for responsiveness
+            prefetch_non_immediate=True,  # enable prefetch for non-immediate categories
+            strategy=strategy,
+        )
+
+        self._logger.info(
+            "MessageBrokerTaskSourceStage initialized. Base task queue: %s | Derived queues: %s",
+            self.task_queue,
+            {
+                "immediate": f"{self.task_queue}_immediate",
+                "micro": f"{self.task_queue}_micro",
+                "small": f"{self.task_queue}_small",
+                "medium": f"{self.task_queue}_medium",
+                "large": f"{self.task_queue}_large",
+                "default": f"{self.task_queue}",
+            },
+        )
 
     # --- Private helper methods ---
     def _create_client(self):
@@ -265,14 +290,21 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
 
         return control_message
 
-    def _fetch_message(self, timeout=100):
+    def _fetch_message(self, timeout=0):
         """
-        Fetch a message from the message broker.
+        Fetch a message from the message broker using fair scheduling across derived queues.
+        This is a non-blocking sweep across all queues for the current scheduling cycle. If no
+        message is found across any queue, return None so the caller can sleep briefly.
         """
         try:
-            job = self.client.fetch_message(self.task_queue, timeout)
+            # Use scheduler to fetch next. In simple strategy this will block up to poll_interval on base queue.
+            job = self.scheduler.fetch_next(self.client, timeout=self.config.poll_interval)
             if job is None:
-                self._logger.debug("No message received from '%s'", self.task_queue)
+                self._logger.debug(
+                    "No message received from derived queues for base "
+                    "'%s' (immediate, micro, small, medium, large, default)",
+                    self.task_queue,
+                )
                 # Do not treat normal empty polls as failures
                 self._fetch_failure_count = 0
                 self._current_backoff_sleep = 0.0
@@ -336,7 +368,8 @@ class MessageBrokerTaskSourceStage(RayActorSourceStage):
         Instead of reading from an input edge, fetch a message from the broker.
         """
         self._logger.debug("read_input: calling _fetch_message()")
-        job = self._fetch_message(timeout=100)
+        # Perform a non-blocking sweep across all queues for this cycle
+        job = self._fetch_message(timeout=0)
         if job is None:
             # Sleep for either the configured poll interval or the current backoff, whichever is larger
             sleep_time = max(self.config.poll_interval, getattr(self, "_current_backoff_sleep", 0.0))