perf(serverless): dispatch queued jobs without polling

runpod/serverless/modules/rp_scale.py

@@ -56,6 +56,8 @@ class JobScaler:
 
     def __init__(self, config: dict[str, Any]):
         self._shutdown_event = asyncio.Event()
+        # Set by get_jobs after queuing so run_jobs dispatches without polling.
+        self._jobs_available = asyncio.Event()
         self._prestart_ready = asyncio.Event()
         self._prestart_succeeded = False
         self._prestart_error: dict[str, Any] | None = None
@@ -157,7 +159,14 @@ def handle_shutdown(self, signum, frame):
             frame: The current stack frame.
         """
         log.debug(f"Received shutdown signal: {signum}.")
-        self.kill_worker()
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            self.kill_worker()
+            return
+        # A signal.signal handler does not wake a loop blocked in its selector;
+        # call_soon_threadsafe writes to the loop's self-pipe so it does.
+        loop.call_soon_threadsafe(self.kill_worker)
 
     async def run(self):
         """Run prestart and the three persistent request loops concurrently."""
@@ -338,6 +347,7 @@ async def get_jobs(self, session: ClientSession):
                     await self.jobs_queue.put(job)
                     self.job_progress.add(job)
                     log.debug("Job Queued", job["id"])
+                self._jobs_available.set()
 
                 log.info(f"Jobs in queue: {self.jobs_queue.qsize()}")
 
@@ -371,6 +381,9 @@ async def run_jobs(self, session: ClientSession):
 
         last_task_count = 0
         while self.is_alive() or not self.jobs_queue.empty():
+            # Clear before draining: a job queued after this point re-arms it.
+            self._jobs_available.clear()
+
             # Fetch as many jobs as the concurrency allows
             while len(tasks) < self.current_concurrency and not self.jobs_queue.empty():
                 job = await self.jobs_queue.get()
@@ -379,23 +392,14 @@ async def run_jobs(self, session: ClientSession):
                 tasks.add(task)
                 self.jobs_tasks[job["id"]] = task
 
-            # Wait for any job to finish
             if tasks:
                 current_task_count = len(tasks)
                 if current_task_count != last_task_count:
                     log.info(f"Jobs in progress: {current_task_count}")
                     last_task_count = current_task_count
 
-                done, pending = await asyncio.wait(
-                    tasks, return_when=asyncio.FIRST_COMPLETED, timeout=0.1
-                )
-
-                # Remove completed tasks from the list
-                tasks.difference_update(done)
-
-            else:
-                # don't busy wait
-                await asyncio.sleep(0.1)
+            finished = await self._wait_for_dispatch_event(tasks)
+            tasks.difference_update(finished)
 
         # Ensure all remaining tasks finish before stopping. Stopped jobs raise
         # CancelledError during this drain, which is expected, but a genuine
@@ -409,6 +413,35 @@ async def run_jobs(self, session: ClientSession):
                     f"JobScaler.run_jobs | Task failed during shutdown drain: {result}"
                 )
 
+    async def _wait_for_dispatch_event(
+        self, tasks: set[asyncio.Task[Any]]
+    ) -> set[asyncio.Task[Any]]:
+        """Block until a job finishes, a job is queued, or shutdown begins.
+
+        Waiting on events instead of polling starts a fetched job as soon as
+        it is queued; a poll interval delayed every cold start's first job.
+
+        Returns the job tasks that finished.
+        """
+        # Snapshot once: a signal handler can stop the worker at any bytecode.
+        alive = self.is_alive()
+        if not alive and not tasks:
+            # Intake has stopped, so nothing would ever set _jobs_available.
+            return set()
+
+        waiters = {asyncio.create_task(self._jobs_available.wait())}
+        if alive:
+            # Once shutting down, the set event would wake this loop instantly.
+            waiters.add(asyncio.create_task(self._shutdown_event.wait()))
+        try:
+            done, _ = await asyncio.wait(
+                tasks | waiters, return_when=asyncio.FIRST_COMPLETED
+            )
+        finally:
+            for waiter in waiters:
+                waiter.cancel()
+        return done & tasks
+
     async def monitor_stop_signals(self, session: ClientSession):
         """
         Long-polls the dedicated stop channel and stops signalled jobs.

tests/test_serverless/test_rp_scale.py

@@ -1,4 +1,8 @@
 import asyncio
+import os
+import signal
+import threading
+import time
 from unittest.mock import AsyncMock, Mock, patch
 from dataclasses import dataclass
 
@@ -414,6 +418,148 @@ async def test_monitor_stop_signals_disabled_without_stop_url(job_scaler: PatchS
         await asyncio.wait_for(scaler.monitor_stop_signals(AsyncMock()), timeout=1)
 
 
+# Well under the old 100ms dispatch poll, so a polling loop fails these tests.
+DISPATCH_DEADLINE = 0.05  # seconds
+
+
+def _gated_fetcher(release: asyncio.Event, job_ids: list[str]):
+    """Fetcher that hands out the given jobs once `release` is set, then none."""
+    pending = list(job_ids)
+
+    async def fetcher(_session, jobs_needed):
+        await release.wait()
+        taken, pending[:] = pending[:jobs_needed], pending[jobs_needed:]
+        if not taken:
+            await asyncio.sleep(1)
+        return [generate_job(job_id) for job_id in taken]
+
+    return fetcher
+
+
+@pytest.mark.asyncio
+async def test_idle_dispatcher_starts_fetched_job_immediately(job_scaler: PatchScaler):
+    scaler = job_scaler.scaler
+    release = asyncio.Event()
+    handler_called = asyncio.Event()
+
+    async def handler(_session, _config, _job):
+        handler_called.set()
+
+    scaler.jobs_handler = handler
+    scaler.jobs_fetcher = _gated_fetcher(release, ["job-0"])
+    get_task = asyncio.create_task(scaler.get_jobs(AsyncMock()))
+    run_task = asyncio.create_task(scaler.run_jobs(None))
+    await asyncio.sleep(0.01)  # dispatcher is now idle
+
+    release.set()
+
+    await asyncio.wait_for(handler_called.wait(), timeout=DISPATCH_DEADLINE)
+    scaler.kill_worker()
+    get_task.cancel()
+    await asyncio.wait_for(run_task, timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_busy_dispatcher_starts_second_job_immediately(job_scaler: PatchScaler):
+    scaler = job_scaler.scaler
+    release_second = asyncio.Event()
+    finish_first = asyncio.Event()
+    started = {}
+
+    async def handler(_session, _config, job):
+        started[job["id"]] = True
+        if job["id"] == "job-0":
+            await finish_first.wait()
+
+    scaler.jobs_handler = handler
+    scaler.current_concurrency = 2
+    scaler.jobs_queue = asyncio.Queue(maxsize=2)
+    await scaler.jobs_queue.put(generate_job("job-0"))
+    job_scaler.progress.add(None)
+    run_task = asyncio.create_task(scaler.run_jobs(None))
+    await asyncio.sleep(0.01)  # job-0 running, dispatcher waiting on it
+    assert started == {"job-0": True}
+
+    scaler.jobs_fetcher = _gated_fetcher(release_second, ["job-1"])
+    get_task = asyncio.create_task(scaler.get_jobs(AsyncMock()))
+    await asyncio.sleep(0)
+    release_second.set()
+
+    async def second_started():
+        while "job-1" not in started:
+            await asyncio.sleep(0)
+
+    await asyncio.wait_for(second_started(), timeout=DISPATCH_DEADLINE)
+    finish_first.set()
+    scaler.kill_worker()
+    get_task.cancel()
+    await asyncio.wait_for(run_task, timeout=1)
+
+
+@pytest.mark.asyncio
+async def test_idle_dispatcher_exits_promptly_on_shutdown(job_scaler: PatchScaler):
+    scaler = job_scaler.scaler
+    run_task = asyncio.create_task(scaler.run_jobs(None))
+    await asyncio.sleep(0.01)  # dispatcher is now idle
+
+    scaler.kill_worker()
+
+    await asyncio.wait_for(run_task, timeout=DISPATCH_DEADLINE)
+
+
+@pytest.mark.asyncio
+async def test_idle_dispatcher_exits_when_shutdown_follows_loop_check(
+    job_scaler: PatchScaler,
+):
+    scaler = job_scaler.scaler
+    real_is_alive = scaler.is_alive
+    checks = 0
+
+    def is_alive_then_signalled():
+        # A signal handler can run between run_jobs' loop check and its wait.
+        nonlocal checks
+        checks += 1
+        alive = real_is_alive()
+        if checks == 1:
+            scaler.kill_worker()
+        return alive
+
+    scaler.is_alive = is_alive_then_signalled
+
+    await asyncio.wait_for(scaler.run_jobs(None), timeout=DISPATCH_DEADLINE)
+
+
+SIGNAL_DELAY = 0.1  # seconds
+SIGNAL_WAKE_TIMEOUT = 2  # seconds; the only timer that would wake a stalled loop
+
+
+@pytest.mark.skipif(not hasattr(signal, "SIGUSR1"), reason="needs POSIX signals")
+def test_idle_dispatcher_exits_promptly_on_os_signal(job_scaler: PatchScaler):
+    scaler = job_scaler.scaler
+    signalled_at = []
+
+    def handler(signum, frame):
+        signalled_at.append(time.monotonic())
+        scaler.handle_shutdown(signum, frame)
+
+    async def run_until_signalled():
+        run_task = asyncio.create_task(scaler.run_jobs(None))
+        # Deliver the signal once the loop is blocked in the selector.
+        threading.Timer(
+            SIGNAL_DELAY, os.kill, (os.getpid(), signal.SIGUSR1)
+        ).start()
+        await asyncio.wait_for(run_task, timeout=SIGNAL_WAKE_TIMEOUT)
+        return time.monotonic()
+
+    previous = signal.signal(signal.SIGUSR1, handler)
+    try:
+        finished_at = asyncio.run(run_until_signalled())
+    finally:
+        signal.signal(signal.SIGUSR1, previous)
+
+    assert finished_at - signalled_at[0] < DISPATCH_DEADLINE
+
+
 @pytest.mark.asyncio
 async def test_get_jobs_feeds_workers_end_to_end(job_scaler: PatchScaler):
     scaler = job_scaler.scaler