Repository: skypilot-org/skypilot
Pull request: [API server] performance improvement in local mode
<!-- Describe the changes in this PR -->

Given same resources, the max parallelism cannot be **exactly** as good as 0.8.0 due to the overhead of uvicorn server and status refresh daemon (see the test result). But this branch can already achieve similar concurrency if user provide `--async` to the CLI (so there is no overhead of client process), so I think this close  https://github.com/skypilot-org/skypilot/issues/4979

## Test

All the following tests are done with aws, azure and gcp enabled, but set `--cloud=aws` to launch instance only on AWS.

### Low resources: 0.5c 750MB

> Setup: launch an AWS c6i.large instance, manually start a docker container with 0.5c 750MB resources limit

For low resources env, we focus on that `sky` can provide similar performance compared to `0.8.0` 

- **This Branch**: support 3 concurrent `sky launch --async`, 4 concurrent `sky.launch` cause one of the worker process being OOM killed, extra `sky status` or `sky logs` cause the launch worker to be OOMKilled
- **0.8.0**: support 3 concurrent `sky launch &`, each takes about 210 MB peak RSS, 4 concurrent `sky launch` cause OOM, there is still memory to run `sky logs` and `sky status` without blocking;
- **master**: support 1 long request (e.g. `sky launch`) in parallel and 1 short request (e.g. `sky logs`)

Compared to 0.8.0, the overhead of this branch is about 200MB, including a single uvicorn process and a status refresh daemon process.

### Common cases: 4c16g

> Setup: launch an AWS c6i.4xlarge instance, manually start a docker container with 4c16g resources limit

For environment that has relatively enough resources for local usage, we focus on:
1. whether the parallelism can be burst to a high number if needed;
2. the resource consumption after the burst peak;

- **This Branch**: up to 30 parallel launch (each with worker process with about ~250MB peak memory usage), the resident memory after peak is about 3GB (4 long workers + 17 short workers)
- **0.8.0**:  similar to this branch, but no resident memory after peak
- **master**: support 4 long requests and 17 short requests in parallel

For common cases, current branch has about 500MB memory overhead compared to 0.8.0, the extra 300MB plus to uvicorn process and status refresh daemon process including:

- long worker process
- short worker process
- Queue server process

Screenshot of `htop`:

<img width="1920" alt="image" src="https://github.com/user-attachments/assets/06d8e582-ec68-4b1d-9686-ead9251f723e" />

## Future works

1. Better OOM handling: in 0.8.0, the visibility of CLI processing being OOM killed is good: user just see it in the terminal or `ps`, but the API server hides workers from user, which need some UX improvement;
2. More cooperative operations: the overhead of a dedicated process for each request is significant, make operation cooperative improves the performance in all scenarios.


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

## Tests

Tested (run the relevant ones):

- [x] Code formatting: install pre-commit (auto-check on commit) or `bash format.sh`
- [ ] Any manual or new tests for this PR (please specify below)
- [ ] All smoke tests: `/smoke-test` (CI) or `pytest tests/test_smoke.py` (local)
- [ ] Relevant individual tests: `/smoke-test -k test_name` (CI) or `pytest tests/test_smoke.py::test_name` (local)
- [ ] Backward compatibility: `/quicktest-core` (CI) or `pytest tests/smoke_tests/test_backward_compat.py` (local)

<!-- CI commands (/-prefixed) can only be triggered by repo members -->

### Queue Scheduling Benchmark

Run benchmark:

```
python tests/load_tests/test_queue_dispatcher.py
```

Result:

```bash
Benchmark Results:
--------------------------------------------------
Process queue + 1 process dispatcher:
  Requests/sec:  451.57
Local thread queue + 1 thread dispatcher:
  Requests/sec:  502.04
Process queue + 1 thread dispatcher:
  Requests/sec:  468.69
Process queue + 10 process dispatcher:
  Requests/sec:  411.62
Local thread queue + 10 thread dispatcher (in one process):
  Requests/sec:  410.38
```

A simple benchmark is added to see whether it is okay to switch to thread-based queue-server and dispatcher in low resource env.

It is interesting that multiprocess dispatchers does not outperform single-process/thread dispatcher on my laptop, worth digging deeper when we do comprehensive queue optimization.


Problem statement:
[API server] Local API server has lower concurrency compared to 0.8.0
One of our users runs concurrent `sky.launch` (via Python API) in kubernetes Pod with relatively low resource limit:

```
Requests - CPU: 500m, Memory: 1Gi
Limits - CPU: 750m, Memory: 4Gi
```

After upgrading to recent versions with local API server enabled, the max concurrency long-running `sky` operations (e.g. `sky.launch`) is limited to 1, which breaks the existing workflow.

Hints from issue discussion:

Commit message: Address review comments

Signed-off-by: Aylei <rayingecho@gmail.com>

Patch to review:
diff --git a/sky/server/requests/executor.py b/sky/server/requests/executor.py
index 9ec75f5a68f..18464320d86 100644
--- a/sky/server/requests/executor.py
+++ b/sky/server/requests/executor.py
@@ -18,9 +18,7 @@
 
 See the [README.md](../README.md) for detailed architecture of the executor.
 """
-import concurrent.futures
 import contextlib
-import dataclasses
 import enum
 import multiprocessing
 import os
@@ -42,12 +40,13 @@
 from sky.server import constants as server_constants
 from sky.server.requests import payloads
 from sky.server.requests import preconditions
+from sky.server.requests import process
 from sky.server.requests import requests as api_requests
+from sky.server.requests.queues import local_queue
 from sky.server.requests.queues import mp_queue
 from sky.skylet import constants
 from sky.utils import annotations
 from sky.utils import common_utils
-from sky.utils import subprocess_utils
 from sky.utils import timeline
 
 if typing.TYPE_CHECKING:
@@ -101,22 +100,23 @@
 # workers so at least 2 workers are needed to ensure responsiveness.
 _MIN_SHORT_WORKERS = 2
 
+# Default number of burstable workers for local API server. A heuristic number
+# that is large enough for most local cases.
+# TODO(aylei): the number of burstable workers should be auto-tuned based on the
+# system usage stats.
+_BURSTABLE_WORKERS_FOR_LOCAL = 1024
+
 
 class QueueBackend(enum.Enum):
+    # Local queue backend serves queues in each process locally, which has
+    # lower resource usage but the consumer must be in the same process, i.e.
+    # this only works in single-process mode.
+    LOCAL = 'local'
+    # Multi-process queue backend starts a dedicated process for serving queues.
     MULTIPROCESSING = 'multiprocessing'
     # TODO(zhwu): we can add redis backend in the future.
 
 
-@dataclasses.dataclass
-class RequestWorker:
-    id: int
-    # The type of queue this worker works on.
-    schedule_type: api_requests.ScheduleType
-
-    def __str__(self) -> str:
-        return f'Worker(id={self.id}, schedule_type={self.schedule_type.value})'
-
-
 class RequestQueue:
     """The queue for the requests, either redis or multiprocessing.
 
@@ -128,9 +128,12 @@ def __init__(self,
                  backend: Optional[QueueBackend] = None) -> None:
         self.name = schedule_type.value
         self.backend = backend
-        assert (backend is None or
-                backend == QueueBackend.MULTIPROCESSING), backend
-        self.queue = mp_queue.get_queue(self.name)
+        if backend == QueueBackend.MULTIPROCESSING:
+            self.queue = mp_queue.get_queue(self.name)
+        elif backend == QueueBackend.LOCAL:
+            self.queue = local_queue.get_queue(self.name)
+        else:
+            raise RuntimeError(f'Invalid queue backend: {backend}')
 
     def put(self, request: Tuple[str, bool]) -> None:
         """Put and request to the queue.
@@ -161,6 +164,104 @@ def __len__(self) -> int:
 queue_backend = QueueBackend.MULTIPROCESSING
 
 
+def executor_initializer(proc_group: str):
+    setproctitle.setproctitle(f'SkyPilot:executor:{proc_group}:'
+                              f'{multiprocessing.current_process().pid}')
+
+
+class RequestWorker:
+    """A worker that polls requests from the queue and runs them.
+
+    The worker can run at least `garanteed_parallelism` requests in parallel.
+    If there are more resources available, it can spin up extra workers up to
+    `garanteed_parallelism + burstable_parallelism`.
+    """
+
+    # The type of queue this worker works on.
+    schedule_type: api_requests.ScheduleType
+    # The least number of requests that this worker can run in parallel.
+    garanteed_parallelism: int
+    # The extra number of requests that this worker can run in parallel
+    # if there are available CPU/memory resources.
+    burstable_parallelism: int = 0
+
+    def __init__(self,
+                 schedule_type: api_requests.ScheduleType,
+                 garanteed_parallelism: int,
+                 burstable_parallelism: int = 0) -> None:
+        self.schedule_type = schedule_type
+        self.garanteed_parallelism = garanteed_parallelism
+        self.burstable_parallelism = burstable_parallelism
+
+    def __str__(self) -> str:
+        return f'Worker(schedule_type={self.schedule_type.value})'
+
+    def process_request(self, executor: process.BurstableExecutor,
+                        queue: RequestQueue) -> None:
+        try:
+            request_element = queue.get()
+            if request_element is None:
+                time.sleep(0.1)
+                return
+            request_id, ignore_return_value = request_element
+            request = api_requests.get_request(request_id)
+            assert request is not None, f'Request with ID {request_id} is None'
+            if request.status == api_requests.RequestStatus.CANCELLED:
+                return
+            logger.info(f'[{self}] Submitting request: {request_id}')
+            # Start additional process to run the request, so that it can be
+            # cancelled when requested by a user.
+            # TODO(zhwu): since the executor is reusing the request process,
+            # multiple requests can share the same process pid, which may cause
+            # issues with SkyPilot core functions if they rely on the exit of
+            # the process, such as subprocess_daemon.py.
+            executor.submit_until_success(_request_execution_wrapper,
+                                          request_id, ignore_return_value)
+
+            logger.info(f'[{self}] Submitted request: {request_id}')
+        except (Exception, SystemExit) as e:  # pylint: disable=broad-except
+            # Catch any other exceptions to avoid crashing the worker process.
+            logger.error(
+                f'[{self}] Error processing request: '
+                f'{request_id if "request_id" in locals() else ""} '
+                f'{common_utils.format_exception(e, use_bracket=True)}')
+
+    def run(self) -> None:
+        # Handle the SIGTERM signal to abort the executor process gracefully.
+        proc_group = f'{self.schedule_type.value}'
+        if threading.current_thread() is threading.main_thread():
+            signal.signal(signal.SIGTERM, _sigterm_handler)
+            setproctitle.setproctitle(f'SkyPilot:worker:{proc_group}')
+        queue = _get_queue(self.schedule_type)
+
+        # Use concurrent.futures.ProcessPoolExecutor instead of
+        # multiprocessing.Pool because the former is more efficient with the
+        # support of lazy creation of worker processes.
+        # We use executor instead of individual multiprocessing.Process to avoid
+        # the overhead of forking a new process for each request, which can be
+        # about 1s delay.
+        try:
+            executor = process.BurstableExecutor(
+                garanteed_workers=self.garanteed_parallelism,
+                burst_workers=self.burstable_parallelism,
+                initializer=executor_initializer,
+                initargs=(proc_group,))
+            while True:
+                self.process_request(executor, queue)
+        # TODO(aylei): better to distinct between KeyboardInterrupt and SIGTERM.
+        except KeyboardInterrupt:
+            pass
+        finally:
+            # In most cases, here we receive either ctrl-c in foreground
+            # execution or SIGTERM on server exiting. Gracefully exit the
+            # worker process and the executor.
+            # TODO(aylei): worker may also be killed by system daemons like
+            # OOM killer, crash the API server or recreate the worker process
+            # to avoid broken state in such cases.
+            logger.info(f'[{self}] Worker process interrupted')
+            executor.shutdown()
+
+
 @annotations.lru_cache(scope='global', maxsize=None)
 def _get_queue(schedule_type: api_requests.ScheduleType) -> RequestQueue:
     return RequestQueue(schedule_type, backend=queue_backend)
@@ -343,110 +444,77 @@ def enqueue():
         enqueue()
 
 
-def executor_initializer(proc_group: str):
-    setproctitle.setproctitle(f'SkyPilot:executor:{proc_group}:'
-                              f'{multiprocessing.current_process().pid}')
-
-
-def request_worker(worker: RequestWorker, max_parallel_size: int) -> None:
-    """Worker for the requests.
-
-    Args:
-        max_parallel_size: Maximum number of parallel jobs this worker can run.
-    """
-    # Handle the SIGTERM signal to abort the executor process gracefully.
-    signal.signal(signal.SIGTERM, _sigterm_handler)
-    proc_group = f'{worker.schedule_type.value}-{worker.id}'
-    setproctitle.setproctitle(f'SkyPilot:worker:{proc_group}')
-    queue = _get_queue(worker.schedule_type)
-
-    def process_request(executor: concurrent.futures.ProcessPoolExecutor):
-        try:
-            request_element = queue.get()
-            if request_element is None:
-                time.sleep(0.1)
-                return
-            request_id, ignore_return_value = request_element
-            request = api_requests.get_request(request_id)
-            assert request is not None, f'Request with ID {request_id} is None'
-            if request.status == api_requests.RequestStatus.CANCELLED:
-                return
-            logger.info(f'[{worker}] Submitting request: {request_id}')
-            # Start additional process to run the request, so that it can be
-            # cancelled when requested by a user.
-            # TODO(zhwu): since the executor is reusing the request process,
-            # multiple requests can share the same process pid, which may cause
-            # issues with SkyPilot core functions if they rely on the exit of
-            # the process, such as subprocess_daemon.py.
-            future = executor.submit(_request_execution_wrapper, request_id,
-                                     ignore_return_value)
-
-            if worker.schedule_type == api_requests.ScheduleType.LONG:
-                try:
-                    future.result(timeout=None)
-                except Exception as e:  # pylint: disable=broad-except
-                    logger.error(f'[{worker}] Request {request_id} failed: {e}')
-                logger.info(f'[{worker}] Finished request: {request_id}')
-            else:
-                logger.info(f'[{worker}] Submitted request: {request_id}')
-        except (Exception, SystemExit) as e:  # pylint: disable=broad-except
-            # Catch any other exceptions to avoid crashing the worker process.
-            logger.error(
-                f'[{worker}] Error processing request: '
-                f'{request_id if "request_id" in locals() else ""} '
-                f'{common_utils.format_exception(e, use_bracket=True)}')
-
-    # Use concurrent.futures.ProcessPoolExecutor instead of multiprocessing.Pool
-    # because the former is more efficient with the support of lazy creation of
-    # worker processes.
-    # We use executor instead of individual multiprocessing.Process to avoid
-    # the overhead of forking a new process for each request, which can be about
-    # 1s delay.
-    try:
-        executor = concurrent.futures.ProcessPoolExecutor(
-            max_workers=max_parallel_size,
-            initializer=executor_initializer,
-            initargs=(proc_group,))
-        while True:
-            process_request(executor)
-    # TODO(aylei): better to distinct between KeyboardInterrupt and SIGTERM.
-    except KeyboardInterrupt:
-        pass
-    finally:
-        # In most cases, here we receive either ctrl-c in foreground execution
-        # or SIGTERM on server exiting. Gracefully exit the worker process and
-        # the executor.
-        # TODO(aylei): worker may also be killed by system daemons like OOM
-        # killer, crash the API server or recreate the worker process to avoid
-        # broken state in such cases.
-        logger.info(f'[{worker}] Worker process interrupted')
-        executor_processes = list(executor._processes.values())  # pylint: disable=protected-access,line-too-long
-        # Shutdown the executor so that executor process can exit once the
-        # running task is finished or interrupted.
-        executor.shutdown(wait=False)
-        # Proactively interrupt the running task to avoid indefinite waiting.
-        subprocess_utils.run_in_parallel(
-            subprocess_utils.kill_process_with_grace_period,
-            executor_processes,
-            num_threads=len(executor_processes))
-
-
 def start(deploy: bool) -> List[multiprocessing.Process]:
-    """Start the request workers."""
+    """Start the request workers.
+
+    Request workers run in background, schedule the requests and delegate the
+    request execution to executor processes. We have different assumptions for
+    the resources in different deployment modes, which leads to different
+    worker setups:
+
+    - Deployment mode (deploy=True), we assume the resources are dedicated to
+      the API server and the resources will be tuned for serious use cases, so:
+      - Use multiprocessing queue backend and dedicated workers processes to
+        avoid GIL contention.
+      - Parallelism (number of executor processes) is fixed and executor
+        processes have same lifecycle with the server, which ensures
+        best-effort cache reusing and stable resources consumption.
+      - Reject to start in low resource environments, to avoid flaky
+        deployments.
+    - Local mode (deploy=False), we assume the server is running in a shared
+      environment (e.g. laptop) and users typically do not pay attention to
+      the resource setup of the server. Moreover, existing users may expect
+      some consistent behaviors with old versions, i.e. before API server was
+      introduced, so:
+      - The max number of long-running executor processes are limited, to avoid
+        high memory consumption when the server is idle.
+      - Allow burstable workers to handle requests when all long-running
+        workers are busy, which mimics the behavior of local sky CLI before
+        API server was introduced.
+      - Works in low resources environments, and further reduce the memory
+        consumption in low resource environments.
+
+    Note that there is still significant overhead for SDK users when migrate to
+    local API server. Since the users are free to run sky operations in Threads
+    when using SDK but all client operations will occupy at least one worker
+    process after API server was introduced.
+    """
     # Determine the job capacity of the workers based on the system resources.
     cpu_count = common_utils.get_cpu_count()
     mem_size_gb = common_utils.get_mem_size_gb()
     mem_size_gb = max(0, mem_size_gb - server_constants.MIN_AVAIL_MEM_GB)
+    # Runs in low resource mode if the available memory is less than
+    # server_constants.MIN_AVAIL_MEM_GB.
     max_parallel_for_long = _max_long_worker_parallism(cpu_count,
                                                        mem_size_gb,
                                                        local=not deploy)
     max_parallel_for_short = _max_short_worker_parallism(
         mem_size_gb, max_parallel_for_long)
-    logger.info(
-        f'SkyPilot API server will start {max_parallel_for_long} workers for '
-        f'long requests and will allow at max '
-        f'{max_parallel_for_short} short requests in parallel.')
-
+    if mem_size_gb < server_constants.MIN_AVAIL_MEM_GB:
+        # Permanent worker process may have significant memory consumption
+        # (~350MB per worker) after running commands like `sky check`, so we
+        # don't start any permanent workers in low resource local mode. This
+        # mimics the behavior of local sky CLI before API server was
+        # introduced, where the CLI will start new process everytime and
+        # never reject to start due to resource constraints.
+        # Note that the refresh daemon will still occupy one worker
+        # permanently because it never exits.
+        max_parallel_for_long = 0
+        max_parallel_for_short = 0
+        logger.warning(
+            'SkyPilot API server will run in low resource mode because '
+            'the available memory is less than '
+            f'{server_constants.MIN_AVAIL_MEM_GB}GB.')
+    else:
+        logger.info(
+            f'SkyPilot API server will start {max_parallel_for_long} workers '
+            f'for long requests and will allow at max '
+            f'{max_parallel_for_short} short requests in parallel.')
+    if not deploy:
+        # For local mode, use local queue backend since we only run 1 uvicorn
+        # worker in local mode.
+        global queue_backend
+        queue_backend = QueueBackend.LOCAL
     sub_procs = []
     # Setup the queues.
     if queue_backend == QueueBackend.MULTIPROCESSING:
@@ -465,28 +533,37 @@ def start(deploy: bool) -> List[multiprocessing.Process]:
             target=mp_queue.start_queue_manager, args=(queue_names, port))
         queue_server.start()
         sub_procs.append(queue_server)
-        mp_queue.wait_for_queues_to_be_ready(queue_names, queue_server, port)
+        mp_queue.wait_for_queues_to_be_ready(queue_names,
+                                             queue_server,
+                                             port=port)
+    elif queue_backend == QueueBackend.LOCAL:
+        # No setup is needed for local queue backend.
+        pass
+    else:
+        # Should be checked earlier, but just in case.
+        raise RuntimeError(f'Invalid queue backend: {queue_backend}')
 
     logger.info('Request queues created')
 
-    long_workers = []
-    for worker_id in range(max_parallel_for_long):
-        worker = RequestWorker(id=worker_id,
-                               schedule_type=api_requests.ScheduleType.LONG)
-        worker_proc = multiprocessing.Process(target=request_worker,
-                                              args=(worker, 1))
-        long_workers.append(worker_proc)
-        sub_procs.append(worker_proc)
-    threading.Thread(target=subprocess_utils.slow_start_processes,
-                     args=(long_workers,),
-                     daemon=True).start()
+    def run_worker_in_background(worker: RequestWorker):
+        # Thread dispatcher is sufficient for current scale, refer to
+        # tests/load_tests/test_queue_dispatcher.py for more details.
+        # Use daemon thread for automatic cleanup.
+        thread = threading.Thread(target=worker.run, daemon=True)
+        thread.start()
+
+    burstable_parallelism = _BURSTABLE_WORKERS_FOR_LOCAL if not deploy else 0
+    # Start a worker for long requests.
+    long_worker = RequestWorker(schedule_type=api_requests.ScheduleType.LONG,
+                                garanteed_parallelism=max_parallel_for_long,
+                                burstable_parallelism=burstable_parallelism)
+    run_worker_in_background(long_worker)
 
     # Start a worker for short requests.
-    worker = RequestWorker(id=1, schedule_type=api_requests.ScheduleType.SHORT)
-    worker_proc = multiprocessing.Process(target=request_worker,
-                                          args=(worker, max_parallel_for_short))
-    worker_proc.start()
-    sub_procs.append(worker_proc)
+    short_worker = RequestWorker(schedule_type=api_requests.ScheduleType.SHORT,
+                                 garanteed_parallelism=max_parallel_for_short,
+                                 burstable_parallelism=burstable_parallelism)
+    run_worker_in_background(short_worker)
     return sub_procs
 
 
diff --git a/sky/server/requests/process.py b/sky/server/requests/process.py
new file mode 100644
index 00000000000..a70c75f1740
--- /dev/null
+++ b/sky/server/requests/process.py
@@ -0,0 +1,212 @@
+"""ProcessPoolExecutor with additional supports for skypilot."""
+import concurrent.futures
+import logging
+import multiprocessing
+import threading
+import time
+from typing import Callable, Dict, Optional, Tuple
+
+from sky.utils import atomic
+from sky.utils import subprocess_utils
+
+logger = logging.getLogger(__name__)
+
+
+class PoolExecutor(concurrent.futures.ProcessPoolExecutor):
+    """A custom ProcessPoolExecutor with additional supports for skypilot.
+
+    The additional supports include:
+    1. Disposable workers: support control whether the worker process should
+       exit after complete a task.
+    2. Idle check: support check if there are any idle workers.
+    3. Proactive shutdown: SIGTERM worker processes when the executor is
+       shutting down instead of indefinitely waiting.
+    """
+
+    def __init__(self, max_workers: int, **kwargs):
+        super().__init__(max_workers=max_workers, **kwargs)
+        self.max_workers: int = max_workers
+        # The number of workers that are handling tasks, atomicity across
+        # multiple threads is sufficient since the idleness check is
+        # best-effort and does not affect the correctness.
+        # E.g. the following case is totally fine:
+        # 1. Thread 1 checks running == max_workers
+        # 2. Thread 2 decrements running
+        # 3. Thread 1 schedules the task to other pool even if the pool is
+        #    currently idle.
+        self.running: atomic.AtomicInt = atomic.AtomicInt(0)
+
+    def submit(self, fn, *args, **kwargs) -> concurrent.futures.Future:
+        """Submit a task for execution.
+
+        If reuse_worker is False, wraps the function to exit after completion.
+        """
+        self.running.increment()
+        future = super().submit(fn, *args, **kwargs)
+        future.add_done_callback(lambda _: self.running.decrement())
+        return future
+
+    def has_idle_workers(self) -> bool:
+        """Check if there are any idle workers."""
+        return self.running.get() < self.max_workers
+
+    def shutdown(self, wait: bool = True) -> None:
+        """Shutdown the executor."""
+        # Here wait means wait for the proactive cancellation complete.
+        # TODO(aylei): we may support wait=True in the future if needed.
+        assert wait is True, 'wait=False is not supported'
+        executor_processes = list(self._processes.values())
+        # Shutdown the executor so that executor process can exit once the
+        # running task is finished or interrupted.
+        super().shutdown(wait=False)
+        # Proactively interrupt the running task to avoid indefinite waiting.
+        subprocess_utils.run_in_parallel(
+            subprocess_utils.kill_process_with_grace_period,
+            executor_processes,
+            num_threads=len(executor_processes))
+
+
+# Define the worker function outside of the class to avoid pickling self
+def _disposable_worker(fn, initializer: Optional[Callable], initargs: Tuple,
+                       args, kwargs):
+    try:
+        if initializer is not None:
+            initializer(*initargs)
+        fn(*args, **kwargs)
+    except BaseException as e:  # pylint: disable=broad-except
+        return e
+
+
+class DisposableExecutor:
+    """A simple wrapper that creates a new process for each task.
+
+    This is a workaround for Python 3.10 since `max_tasks_per_child` of
+    ProcessPoolExecutor was introduced in 3.11. There is no way to control
+    the worker lifetime in 3.10.
+    Ref: https://docs.python.org/3/library/concurrent.futures.html#concurrent.futures.ProcessPoolExecutor # pylint: disable=line-too-long
+    TODO(aylei): use the official `max_tasks_per_child` when upgrade to 3.11
+    """
+
+    def __init__(self,
+                 max_workers: Optional[int] = None,
+                 initializer: Optional[Callable] = None,
+                 initargs: Tuple = ()):
+        self.max_workers: Optional[int] = max_workers
+        self.workers: Dict[int, multiprocessing.Process] = {}
+        self._shutdown: bool = False
+        self._lock: threading.Lock = threading.Lock()
+        self._initializer: Optional[Callable] = initializer
+        self._initargs: Tuple = initargs
+
+    def _monitor_worker(self, process: multiprocessing.Process) -> None:
+        """Monitor the worker process and cleanup when it's done."""
+        process.join()
+        if process.pid:
+            with self._lock:
+                if process.pid in self.workers:
+                    del self.workers[process.pid]
+
+    # Submit is not compatible with ProcessPoolExecutor because we does not
+    # bother to return a Future. Can be improved if needed.
+    def submit(self, fn, *args, **kwargs) -> bool:
+        """Submit a task for execution."""
+        if self._shutdown:
+            return False
+        with self._lock:
+            if (self.max_workers is not None and
+                    len(self.workers) >= self.max_workers):
+                return False
+
+        process = multiprocessing.Process(target=_disposable_worker,
+                                          args=(fn, self._initializer,
+                                                self._initargs, args, kwargs))
+        process.start()
+
+        with self._lock:
+            pid = process.pid or 0
+            if pid == 0:
+                raise RuntimeError('Failed to start process')
+            self.workers[pid] = process
+
+        # Start monitor thread to cleanup the worker process when it's done.
+        monitor_thread = threading.Thread(target=self._monitor_worker,
+                                          args=(process,),
+                                          daemon=True)
+        monitor_thread.start()
+
+        return True
+
+    def has_idle_workers(self) -> bool:
+        """Check if there are any idle workers."""
+        if self.max_workers is None:
+            return True
+        with self._lock:
+            return len(self.workers) < self.max_workers
+
+    def shutdown(self):
+        """Shutdown the executor."""
+        with self._lock:
+            self._shutdown = True
+        subprocess_utils.run_in_parallel(
+            subprocess_utils.kill_process_with_grace_period,
+            list(self.workers.values()),  # Convert dict values to list
+            num_threads=len(self.workers))
+
+
+class BurstableExecutor:
+    """An multiprocessing executor that supports bursting worker processes."""
+
+    # _executor is a PoolExecutor that is used to run guaranteed requests.
+    _executor: Optional[PoolExecutor] = None
+    # _burst_executor is a ProcessPoolExecutor that is used to run burst
+    # requests.
+    _burst_executor: Optional[DisposableExecutor] = None
+
+    def __init__(self,
+                 garanteed_workers: int,
+                 burst_workers: int = 0,
+                 **kwargs):
+        if garanteed_workers > 0:
+            self._executor = PoolExecutor(max_workers=garanteed_workers,
+                                          **kwargs)
+        if burst_workers > 0:
+            self._burst_executor = DisposableExecutor(max_workers=burst_workers,
+                                                      **kwargs)
+
+    def submit_until_success(self, fn, *args, **kwargs):
+        """Submit a task for execution until success.
+
+        Prioritizes submitting to the guaranteed pool. If no idle workers
+        are available in the guaranteed pool, it will submit to the burst
+        pool.
+        TODO(aylei): this is coupled with executor.RequestWorker since we
+        know the worker is dedicated to request scheduling and it either
+        blocks on request polling or request submitting. So it is no harm
+        to make submit blocking here. But for general cases, we need an
+        internal queue to decouple submit and run.
+        """
+
+        while True:
+            if self._executor is not None and self._executor.has_idle_workers():
+                self._executor.submit(fn, *args, **kwargs)
+                break
+            if (self._burst_executor is not None and
+                    self._burst_executor.has_idle_workers()):
+                self._burst_executor.submit(fn, *args, **kwargs)
+                break
+            if self._executor is not None:
+                # No idle workers in either pool, still queue the request
+                # to the guaranteed pool to keep behavior consistent.
+                self._executor.submit(fn, *args, **kwargs)
+                break
+            logger.debug('No guaranteed pool set and the burst pool is full, '
+                         'retry later.')
+            time.sleep(0.1)
+
+    def shutdown(self) -> None:
+        """Shutdown the executor."""
+
+        if self._burst_executor is not None:
+            self._burst_executor.shutdown()
+        if self._executor is not None:
+            self._executor.shutdown(wait=True)
diff --git a/sky/server/requests/queues/local_queue.py b/sky/server/requests/queues/local_queue.py
new file mode 100644
index 00000000000..5082e97f0e3
--- /dev/null
+++ b/sky/server/requests/queues/local_queue.py
@@ -0,0 +1,16 @@
+"""Process-local queue implementation."""
+import queue
+import threading
+from typing import Dict
+
+# Global dict to store queues
+_queues: Dict[str, queue.Queue] = {}
+_lock = threading.Lock()
+
+
+def get_queue(queue_name: str) -> queue.Queue:
+    """Get or create a queue by name."""
+    with _lock:
+        if queue_name not in _queues:
+            _queues[queue_name] = queue.Queue()
+        return _queues[queue_name]
diff --git a/sky/utils/atomic.py b/sky/utils/atomic.py
new file mode 100644
index 00000000000..e7870458daa
--- /dev/null
+++ b/sky/utils/atomic.py
@@ -0,0 +1,52 @@
+"""Atomic structures and utilties."""
+
+import threading
+
+
+class AtomicInt:
+    """A thread-safe atomic integer implementation."""
+
+    def __init__(self, initial_value: int = 0):
+        self._value = initial_value
+        self._lock = threading.Lock()
+
+    def get(self) -> int:
+        """Get the current value atomically.
+
+        Returns:
+            The current integer value.
+        """
+        with self._lock:
+            return self._value
+
+    def increment(self, delta: int = 1) -> int:
+        """Atomically increment by delta and return new value.
+
+        Args:
+            delta: Amount to increment by (default: 1)
+
+        Returns:
+            The new value after incrementing.
+        """
+        with self._lock:
+            self._value += delta
+            return self._value
+
+    def decrement(self, delta: int = 1) -> int:
+        """Atomically decrement by delta and return new value.
+
+        Args:
+            delta: Amount to decrement by (default: 1)
+
+        Returns:
+            The new value after decrementing.
+        """
+        with self._lock:
+            self._value -= delta
+            return self._value
+
+    def __str__(self) -> str:
+        return str(self.get())
+
+    def __repr__(self) -> str:
+        return f'AtomicInt({self.get()})'
diff --git a/tests/load_tests/test_queue_dispatcher.py b/tests/load_tests/test_queue_dispatcher.py
new file mode 100644
index 00000000000..fdc511b5e4e
--- /dev/null
+++ b/tests/load_tests/test_queue_dispatcher.py
@@ -0,0 +1,138 @@
+import multiprocessing
+import threading
+import time
+import uuid
+
+from sky import sky_logging
+from sky.server.requests import executor
+from sky.server.requests import payloads
+from sky.server.requests import process
+from sky.server.requests import requests as api_requests
+from sky.server.requests.queues import mp_queue
+from sky.skylet import constants
+
+logger = sky_logging.init_logger(__name__)
+
+
+def disable_logger() -> None:
+    server_logger = sky_logging.init_logger('sky.server')
+    server_logger.handlers.clear()
+    server_logger.propagate = False
+
+
+disable_logger()
+
+
+class MockExecutor(process.BurstableExecutor):
+
+    def submit_until_success(self, *args, **kwargs):
+        return None
+
+
+def dummy_func(*args, **kwargs):
+    pass
+
+
+def schedule():
+    worker = executor.RequestWorker(
+        schedule_type=api_requests.ScheduleType.LONG,
+        garanteed_parallelism=1,
+        burstable_parallelism=1)
+    mock_executor = MockExecutor(garanteed_workers=1, burst_workers=1)
+    queue = executor._get_queue(api_requests.ScheduleType.LONG)  # pylint: disable=protected-access
+    while not queue.queue.empty():
+        worker.process_request(mock_executor, queue)
+
+
+def benchmark_queue_dispatcher(queue_backend: str,
+                               worker_mode: str,
+                               count=1) -> float:
+    api_requests.reset_db_and_logs()
+    num_requests = 10000
+    executor.queue_backend = queue_backend
+
+    queue_server = None
+    if queue_backend == executor.QueueBackend.MULTIPROCESSING:
+        port = mp_queue.DEFAULT_QUEUE_MANAGER_PORT
+        queue_server = multiprocessing.Process(
+            target=mp_queue.start_queue_manager,
+            args=([api_requests.ScheduleType.LONG.value], port))
+        queue_server.start()
+        mp_queue.wait_for_queues_to_be_ready(
+            [api_requests.ScheduleType.LONG.value], queue_server, port=port)
+    elif queue_backend == executor.QueueBackend.LOCAL:
+        pass
+    else:
+        raise RuntimeError(f'Invalid queue backend: {queue_backend}')
+
+    def enqueue_requests():
+        for _ in range(num_requests):
+            body = payloads.StatusBody()
+            body.env_vars = dict({
+                constants.USER_ID_ENV_VAR: 'dummy',
+            })
+            executor.schedule_request(
+                request_id=str(uuid.uuid4()),
+                request_name='status',
+                is_skypilot_system=False,
+                request_body=body,
+                func=dummy_func,
+                schedule_type=api_requests.ScheduleType.LONG)
+
+    enqueue_thread = threading.Thread(target=enqueue_requests)
+    enqueue_thread.start()
+    enqueue_thread.join()
+
+    start_time = time.time()
+    if worker_mode == 'thread':
+        threads = []
+        for _ in range(count):
+            schedule_thread = threading.Thread(target=schedule)
+            schedule_thread.start()
+            threads.append(schedule_thread)
+        for t in threads:
+            t.join()
+    else:
+        procs = []
+        for _ in range(count):
+            schedule_proc = multiprocessing.Process(target=schedule)
+            schedule_proc.start()
+            procs.append(schedule_proc)
+        for p in procs:
+            p.join()
+    end_time = time.time()
+    duration = end_time - start_time
+    requests_per_second = num_requests / duration
+    if queue_server:
+        queue_server.terminate()
+        queue_server.join()
+    # Clear queue cache
+    executor._get_queue.cache_clear()  # pylint: disable=protected-access
+    return requests_per_second
+
+
+if __name__ == '__main__':
+    t1 = benchmark_queue_dispatcher(executor.QueueBackend.MULTIPROCESSING,
+                                    'process')
+    print('\nBenchmark Results:')
+    print('-' * 50)
+    print(f'Process queue + 1 process dispatcher:')
+    print(f'  Requests/sec:  {t1:,.2f}')
+
+    t2 = benchmark_queue_dispatcher(executor.QueueBackend.LOCAL, 'thread')
+    print(f'\nLocal thread queue + 1 thread dispatcher:')
+    print(f'  Requests/sec:  {t2:,.2f}')
+
+    t3 = benchmark_queue_dispatcher(executor.QueueBackend.MULTIPROCESSING,
+                                    'thread')
+    print(f'\nProcess queue + 1 thread dispatcher:')
+    print(f'  Requests/sec:  {t3:,.2f}')
+
+    t4 = benchmark_queue_dispatcher(executor.QueueBackend.MULTIPROCESSING,
+                                    'process', 10)
+    print(f'\nProcess queue + 10 process dispatcher:')
+    print(f'  Requests/sec:  {t4:,.2f}')
+
+    t5 = benchmark_queue_dispatcher(executor.QueueBackend.LOCAL, 'thread', 10)
+    print(f'\nLocal thread queue + 10 thread dispatcher (in one process):')
+    print(f'  Requests/sec:  {t5:,.2f}')
diff --git a/tests/unit_tests/sky/adaptors/test_oci.py b/tests/unit_tests/test_sky/adaptors/test_oci.py
similarity index 100%
rename from tests/unit_tests/sky/adaptors/test_oci.py
rename to tests/unit_tests/test_sky/adaptors/test_oci.py
diff --git a/tests/unit_tests/sky/api/requests/queues/test_mp_queue.py b/tests/unit_tests/test_sky/server/requests/queues/test_mp_queue.py
similarity index 100%
rename from tests/unit_tests/sky/api/requests/queues/test_mp_queue.py
rename to tests/unit_tests/test_sky/server/requests/queues/test_mp_queue.py
diff --git a/tests/unit_tests/sky/api/requests/test_executor.py b/tests/unit_tests/test_sky/server/requests/test_executor.py
similarity index 100%
rename from tests/unit_tests/sky/api/requests/test_executor.py
rename to tests/unit_tests/test_sky/server/requests/test_executor.py
diff --git a/tests/unit_tests/sky/api/requests/test_precond.py b/tests/unit_tests/test_sky/server/requests/test_precond.py
similarity index 100%
rename from tests/unit_tests/sky/api/requests/test_precond.py
rename to tests/unit_tests/test_sky/server/requests/test_precond.py
diff --git a/tests/unit_tests/test_sky/server/requests/test_process.py b/tests/unit_tests/test_sky/server/requests/test_process.py
new file mode 100644
index 00000000000..494cc810d9a
--- /dev/null
+++ b/tests/unit_tests/test_sky/server/requests/test_process.py
@@ -0,0 +1,155 @@
+"""Unit tests for sky/server/requests/process.py."""
+from concurrent.futures import Future
+import time
+
+import pytest
+
+from sky.server.requests.process import BurstableExecutor
+from sky.server.requests.process import DisposableExecutor
+from sky.server.requests.process import PoolExecutor
+
+
+def dummy_task(sleep_time=0.1):
+    """A dummy task that sleeps for a given time."""
+    time.sleep(sleep_time)
+    return True
+
+
+def failing_task():
+    """A task that raises an exception."""
+    raise ValueError('Task failed')
+
+
+def wait_for_workers_cleanup(executor, timeout=5):
+    """Wait for workers to be cleaned up.
+    
+    Args:
+        executor: The DisposableExecutor instance
+        timeout: Maximum time to wait in seconds
+    
+    Returns:
+        bool: True if workers are cleaned up, False if timeout
+    """
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        with executor._lock:
+            if len(executor.workers) == 0:
+                return True
+        time.sleep(0.1)
+    return False
+
+
+def wait_for_futures(futures, timeout=5):
+    """Wait for futures to complete.
+    
+    Args:
+        futures: List of futures to wait for
+        timeout: Maximum time to wait in seconds
+    
+    Returns:
+        bool: True if all futures completed, False if timeout
+    """
+    start_time = time.time()
+    try:
+        for future in futures:
+            remaining = max(0, timeout - (time.time() - start_time))
+            future.result(timeout=remaining)
+        return True
+    except TimeoutError:
+        return False
+
+
+def test_pool_executor():
+    """Test PoolExecutor functionality."""
+    executor = PoolExecutor(max_workers=2)
+    futures = []
+    try:
+        # Test submit and has_idle_workers
+        assert executor.has_idle_workers()
+        future = executor.submit(dummy_task, sleep_time=0.1)
+        futures.append(future)
+        assert isinstance(future, Future)
+
+        # Test multiple tasks
+        for _ in range(2):
+            futures.append(executor.submit(dummy_task, sleep_time=0.1))
+        # Should have no idle workers when both are running
+        assert not executor.has_idle_workers()
+
+        # Wait for all futures to complete before shutdown
+        assert wait_for_futures(futures), "Tasks did not complete in time"
+        assert all(f.done() for f in futures), "Not all tasks completed"
+        assert all(f.result() for f in futures), "Some tasks failed"
+
+        # Should have idle workers after completion
+        assert executor.has_idle_workers()
+    finally:
+        # Wait a bit to ensure all tasks are truly done
+        time.sleep(0.1)
+        executor.shutdown()
+
+
+def test_disposable_executor():
+    """Test DisposableExecutor functionality."""
+    executor = DisposableExecutor(max_workers=2)
+    try:
+        # Test submit and has_idle_workers
+        assert executor.has_idle_workers()
+        assert executor.submit(dummy_task)
+
+        # Test multiple tasks
+        assert executor.submit(dummy_task)
+        assert not executor.has_idle_workers()  # No idle workers when full
+
+        # Wait for tasks to complete and workers to be cleaned up
+        assert wait_for_workers_cleanup(executor), "Workers not cleaned up"
+        assert executor.has_idle_workers()  # Should have idle workers now
+
+        # Test with failing task
+        assert executor.submit(failing_task)
+        assert wait_for_workers_cleanup(
+            executor), "Failed task worker not cleaned up"
+        assert executor.has_idle_workers()  # Worker should be cleaned up
+    finally:
+        executor.shutdown()
+
+
+def test_burstable_executor():
+    """Test BurstableExecutor functionality."""
+    executor = BurstableExecutor(garanteed_workers=1, burst_workers=1)
+    try:
+        # Submit tasks that should go to guaranteed pool first
+        executor.submit_until_success(dummy_task)
+        # Submit another task that should go to burst pool
+        executor.submit_until_success(dummy_task)
+        # Submit one more task that should wait and go to guaranteed pool
+        executor.submit_until_success(dummy_task)
+        # Wait for tasks to complete
+        time.sleep(0.3)
+    finally:
+        executor.shutdown()
+
+
+def test_burstable_executor_no_guaranteed():
+    """Test BurstableExecutor with only burst workers."""
+    executor = BurstableExecutor(garanteed_workers=0, burst_workers=1)
+    try:
+        # Should use burst pool
+        executor.submit_until_success(dummy_task)
+        time.sleep(0.2)
+        # Should be able to submit another task after first one completes
+        executor.submit_until_success(dummy_task)
+    finally:
+        executor.shutdown()
+
+
+def test_burstable_executor_no_burst():
+    """Test BurstableExecutor with only guaranteed workers."""
+    executor = BurstableExecutor(garanteed_workers=1, burst_workers=0)
+    try:
+        # Should use guaranteed pool
+        executor.submit_until_success(dummy_task)
+        # Should queue to guaranteed pool even when busy
+        executor.submit_until_success(dummy_task)
+    finally:
+        executor.shutdown()
diff --git a/tests/unit_tests/sky/api/requests/test_requests.py b/tests/unit_tests/test_sky/server/requests/test_requests.py
similarity index 100%
rename from tests/unit_tests/sky/api/requests/test_requests.py
rename to tests/unit_tests/test_sky/server/requests/test_requests.py
diff --git a/tests/unit_tests/sky/server/test_server.py b/tests/unit_tests/test_sky/server/test_server.py
similarity index 100%
rename from tests/unit_tests/sky/server/test_server.py
rename to tests/unit_tests/test_sky/server/test_server.py
diff --git a/tests/unit_tests/sky/test_task.py b/tests/unit_tests/test_sky/test_task.py
similarity index 100%
rename from tests/unit_tests/sky/test_task.py
rename to tests/unit_tests/test_sky/test_task.py
diff --git a/tests/unit_tests/sky/utils/test_config_utils.py b/tests/unit_tests/test_sky/utils/test_config_utils.py
similarity index 100%
rename from tests/unit_tests/sky/utils/test_config_utils.py
rename to tests/unit_tests/test_sky/utils/test_config_utils.py
diff --git a/tests/unit_tests/sky/utils/test_rich_utils.py b/tests/unit_tests/test_sky/utils/test_rich_utils.py
similarity index 100%
rename from tests/unit_tests/sky/utils/test_rich_utils.py
rename to tests/unit_tests/test_sky/utils/test_rich_utils.py
diff --git a/tests/unit_tests/sky/utils/test_subprocess_utils.py b/tests/unit_tests/test_sky/utils/test_subprocess_utils.py
similarity index 100%
rename from tests/unit_tests/sky/utils/test_subprocess_utils.py
rename to tests/unit_tests/test_sky/utils/test_subprocess_utils.py


Code review comments:
diff_hunk: @@ -67,10 +68,13 @@ def benchmark_queue_dispatcher(queue_backend: str,
     def enqueue_requests():
         for _ in range(num_requests):
             body = payloads.StatusBody()
+            body.env_vars = dict({
+                constants.USER_ID_ENV_VAR: 'dummy',
path: tests/load_tests/test_queue_dispatcher.py
text: @user1:
USER_ID may have to be a valid hex with 8 characters to be correctly picked up. Please check https://github.com/skypilot-org/skypilot/blob/7e9586c1d43e46332d9ef2a2f61fd33fc2cb3991/sky/utils/common_utils.py#L96

@author:
Thanks! I think the test succeed since it does not actually touch the real execution path, will update!