diff --git a/runpod/serverless/modules/rp_scale.py b/runpod/serverless/modules/rp_scale.py index 0f0c9db0c..117026731 100644 --- a/runpod/serverless/modules/rp_scale.py +++ b/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. diff --git a/tests/test_serverless/test_rp_scale.py b/tests/test_serverless/test_rp_scale.py index 30f3f2f3b..ec6a752af 100644 --- a/tests/test_serverless/test_rp_scale.py +++ b/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