perf(serverless): opt-in job take-ahead for saturated workers

docs/serverless/worker.md

@@ -22,6 +22,7 @@ The `config` parameter is a dictionary containing the following keys:
 |-------------------|------------|--------------------------------------------------------------------|
 | `handler`         | `function` | The handler function called with each job input.                   |
 | `prestart_timeout`| `number`   | Optional deadline in seconds for all registered prestart hooks.    |
+| `take_ahead`      | `int`      | Optional. Extra jobs a busy worker may hold while the queue is deep. Defaults to `RUNPOD_TAKE_AHEAD`, or `0` (off). See [Take-ahead](#take-ahead). |
 
 ### handler
 
@@ -118,6 +119,25 @@ A worker can process more than one job concurrently. When a single request is ca
 
 No handler changes are required to support this. Async handlers that hold resources can perform cleanup by catching `asyncio.CancelledError`, but they **must re-raise** it after cleaning up. Swallowing the cancellation makes the worker report the job as completed instead of stopped.
 
+## Take-ahead
+
+A job-take sent while jobs are already queued costs a roughly flat round trip, however many jobs it asks for ([#647](https://github.com/runpod/runpod-python/issues/647)). By default a full worker waits for a slot to free before it sends the next take, so on a saturated high-concurrency worker the slots sit idle for that round trip on every cycle.
+
+Take-ahead is an opt-in way to hide that wait. When the last take came back full (the queue is deep), the worker asks for its free slots plus up to `take_ahead` extra jobs and holds the extras until slots free up, so the next take overlaps the jobs that are running.
+
+```python
+runpod.serverless.start({"handler": handler, "concurrency_modifier": lambda _: 32, "take_ahead": 32})
+```
+
+or set `RUNPOD_TAKE_AHEAD=32` in the endpoint's environment. A value around your concurrency is a good start.
+
+Trade-offs:
+
+- A worker holding jobs keeps them even if another worker frees up first, so it only takes ahead while takes are coming back full. With a shallow queue, takes ask only for free slots, as before.
+- A held job counts as in progress on this worker while it waits for a slot, for up to about one job's run time.
+- Stop signals still apply to held jobs: a stopped job is dropped before it starts.
+- It is disabled when `refresh_worker` is set in the config, and should not be combined with handlers that return `refresh_worker`, since held jobs would be lost when the worker restarts.
+
 ## See Also
 
 - [Worker Fitness Checks](./worker_fitness_checks.md) - Validate your worker environment at startup

runpod/serverless/modules/rp_scale.py

@@ -5,6 +5,7 @@
 
 import asyncio
 import json
+import os
 import signal
 import sys
 import traceback
@@ -32,13 +33,45 @@
 # wait so one long job can't hold the other slots idle, and so set_scale still
 # runs as often as it did under the old fixed sleep.
 FULL_WAIT_SECONDS = 1
+# Opt-in take-ahead (#647): while the platform queue is deep, a busy worker may
+# hold up to this many extra jobs, so the next job-take overlaps the jobs that
+# are running instead of waiting for them. Off (0) by default.
+TAKE_AHEAD_ENV = "RUNPOD_TAKE_AHEAD"
 
 
 def _handle_uncaught_exception(exc_type, exc_value, exc_traceback):
     exc = traceback.format_exception(exc_type, exc_value, exc_traceback)
     log.error(f"Uncaught exception | {exc}")
 
 
+def _take_ahead_from(config: dict[str, Any]) -> int:
+    """Read the take-ahead size from config["take_ahead"] or RUNPOD_TAKE_AHEAD."""
+    value = config.get("take_ahead")
+    if value is None:
+        value = os.environ.get(TAKE_AHEAD_ENV, "0")
+
+    try:
+        if isinstance(value, bool):
+            raise TypeError("take_ahead is a job count, not a flag")
+        take_ahead = int(value)
+    except (TypeError, ValueError):
+        log.warn(
+            f"JobScaler | Ignoring invalid take_ahead {value!r}; expected an integer >= 0."
+        )
+        return 0
+
+    if take_ahead < 0:
+        log.warn(f"JobScaler | Ignoring negative take_ahead {take_ahead}.")
+        return 0
+
+    if take_ahead and config.get("refresh_worker", False):
+        # The worker exits after its first job, so held jobs would be lost.
+        log.warn("JobScaler | take_ahead is disabled because refresh_worker is set.")
+        return 0
+
+    return take_ahead
+
+
 def _default_concurrency_modifier(current_concurrency: int) -> int:
     """
     Default concurrency modifier.
@@ -65,6 +98,9 @@ def __init__(self, config: dict[str, Any]):
         self._jobs_available = asyncio.Event()
         # Set when a job leaves the worker so a full get_jobs re-checks for room.
         self._slot_freed = asyncio.Event()
+        # Set when run_jobs takes a job off jobs_queue, so a take-ahead worker
+        # holding jobs knows when its buffer has drained.
+        self._job_dequeued = asyncio.Event()
         self._prestart_ready = asyncio.Event()
         self._prestart_succeeded = False
         self._prestart_error: dict[str, Any] | None = None
@@ -84,7 +120,18 @@ def __init__(self, config: dict[str, Any]):
         self.stop_signals_fetcher = get_stop_signals
         self.stop_signals_fetcher_timeout = 90
 
-        self.jobs_queue = asyncio.Queue(maxsize=self.current_concurrency)
+        # Extra jobs a busy worker may hold while the queue is deep (#647).
+        self.take_ahead = _take_ahead_from(config)
+        # Whether the last job-take was filled, i.e. the platform queue is deep.
+        # Take-ahead only asks for extra jobs then, so a shallow queue still
+        # spreads across workers as before.
+        self._queue_is_deep = False
+        # Ids of jobs waiting in jobs_queue, and those of them stopped before
+        # they started, so a stop signal still applies to a held job.
+        self._queued_job_ids: set[str] = set()
+        self._stopped_queued_job_ids: set[str] = set()
+
+        self.jobs_queue = asyncio.Queue(maxsize=self._queue_size())
 
         self.concurrency_modifier = _default_concurrency_modifier
         self.jobs_fetcher = get_job
@@ -117,10 +164,13 @@ def __init__(self, config: dict[str, Any]):
         ):
             self.stop_signals_fetcher_timeout = stop_signals_fetcher_timeout
 
+    def _queue_size(self) -> int:
+        return self.current_concurrency + self.take_ahead
+
     async def set_scale(self):
         self.current_concurrency = self.concurrency_modifier(self.current_concurrency)
 
-        if self.jobs_queue and (self.current_concurrency == self.jobs_queue.maxsize):
+        if self.jobs_queue and (self._queue_size() == self.jobs_queue.maxsize):
             # no need to resize
             return
 
@@ -129,7 +179,7 @@ async def set_scale(self):
             await asyncio.sleep(1)
             continue
 
-        self.jobs_queue = asyncio.Queue(maxsize=self.current_concurrency)
+        self.jobs_queue = asyncio.Queue(maxsize=self._queue_size())
         log.debug(
             f"JobScaler.set_scale | New concurrency set to: {self.current_concurrency}"
         )
@@ -320,7 +370,14 @@ async def get_jobs(self, session: ClientSession):
                 break
             await self.set_scale()
 
-            jobs_needed = self.current_concurrency - self.current_occupancy()
+            if self.take_ahead and self.jobs_queue.qsize() > 0:
+                # Held jobs are still waiting for a slot; a take now could
+                # only ask for fewer jobs than a full one.
+                log.debug("JobScaler.get_jobs | Holding jobs. Waiting for a slot.")
+                await self._wait_until_queue_drained()
+                continue
+
+            jobs_needed = self._jobs_needed()
             if jobs_needed <= 0:
                 log.debug("JobScaler.get_jobs | Queue is full. Waiting to drain.")
                 await self._wait_until_drained()
@@ -335,6 +392,9 @@ async def get_jobs(self, session: ClientSession):
                     timeout=self.jobs_fetcher_timeout,
                 )
 
+                # A filled take means more jobs are likely queued behind it.
+                self._queue_is_deep = len(acquired_jobs or []) >= jobs_needed
+
                 if not acquired_jobs:
                     log.debug("JobScaler.get_jobs | No jobs acquired.")
                     continue
@@ -352,6 +412,7 @@ async def get_jobs(self, session: ClientSession):
 
                 for job in acquired_jobs:
                     await self.jobs_queue.put(job)
+                    self._queued_job_ids.add(job["id"])
                     self.job_progress.add(job)
                     log.debug("Job Queued", job["id"])
                 self._jobs_available.set()
@@ -378,6 +439,35 @@ async def get_jobs(self, session: ClientSession):
                 # Yield control back to the event loop
                 await asyncio.sleep(0)
 
+    def _jobs_needed(self) -> int:
+        """How many jobs the next job-take should ask for."""
+        jobs_needed = self.current_concurrency - self.current_occupancy()
+        if not self.take_ahead or self.current_concurrency <= 0:
+            return jobs_needed
+
+        if self._queue_is_deep:
+            # Ask for the free slots plus up to take_ahead jobs to hold, so
+            # this take overlaps the jobs already running.
+            jobs_needed += self.take_ahead
+        return jobs_needed
+
+    async def _wait_until_queue_drained(self) -> None:
+        """Block until run_jobs empties jobs_queue or shutdown begins."""
+        while self.is_alive():
+            # Clear, then re-check: a dequeue after this point re-arms the event.
+            self._job_dequeued.clear()
+            if self.jobs_queue.empty():
+                return
+            dequeued = asyncio.create_task(self._job_dequeued.wait())
+            stopping = asyncio.create_task(self._shutdown_event.wait())
+            try:
+                await asyncio.wait(
+                    {dequeued, stopping}, return_when=asyncio.FIRST_COMPLETED
+                )
+            finally:
+                dequeued.cancel()
+                stopping.cancel()
+
     async def _wait_until_drained(self) -> None:
         """Block until no job is in progress, shutdown begins, or FULL_WAIT_SECONDS pass."""
         loop = asyncio.get_running_loop()
@@ -419,6 +509,11 @@ async def run_jobs(self, session: ClientSession):
             # 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()
+                self._queued_job_ids.discard(job["id"])
+                self._job_dequeued.set()
+                if job["id"] in self._stopped_queued_job_ids:
+                    self._drop_stopped_job(job)
+                    continue
                 # Create a new task for each job and track it by job id
                 task = asyncio.create_task(self.handle_job(session, job))
                 tasks.add(task)
@@ -535,13 +630,26 @@ async def stop_job(self, job_id: str) -> bool:
         """
         task = self.jobs_tasks.get(job_id)
         if task is None:
+            if job_id in self._queued_job_ids:
+                # Held in jobs_queue and not started: run_jobs drops it.
+                log.info("Stopping job before it starts.", job_id)
+                self._stopped_queued_job_ids.add(job_id)
+                return True
             log.debug(f"JobScaler.stop_job | No in-progress job for {job_id}.")
             return False
 
         log.info("Stopping job.", job_id)
         task.cancel()
         return True
 
+    def _drop_stopped_job(self, job: dict[str, Any]) -> None:
+        """Release a job that was stopped while it waited in jobs_queue."""
+        self._stopped_queued_job_ids.discard(job["id"])
+        self.jobs_queue.task_done()
+        self.job_progress.remove(job)
+        self._slot_freed.set()
+        log.info("Job stopped.", job["id"])
+
     async def handle_job(self, session: ClientSession, job: dict[str, Any]):
         """
         Process an individual job. This function is run concurrently for multiple jobs.