Repository: skypilot-org/skypilot
Pull request: [Client/Server] fixed server hung during smoke test
<!-- Describe the changes in this PR -->

close https://github.com/skypilot-org/skypilot/issues/4732

API server stability enhancements in resource-constrained environments:
1. Optimize memory estimation of API Server
  - Long worker: 0.25GB => 0.4GB
  - Short worker: 0.15GB => 0.25GB
  - For example, the former parallelism in a 7c14g container is 14xLong and 56xShort. After this PR, it would be 14xLong and 25xShort
2. Remove busy-waiting for request status, which may use up CPUs since the wait time is unbounded
3. Prevent worker exits by catching all exceptions

## Benchmarks

Mostly this PR only affects the parallelism of short workers since most of the environment has memory:cpu > 2:1, so we focus on the throughput and RTT of short requests. We collect results of both 1st (cold) and 2nd (warm) run on an newly started API server. All the tests are done on the PR branch to get rid of the influence of server process busy-waiting:

- server: a container(7c14g) on AWS c6i.2xlarge (8c16g)
- client: AWS m6i.4xlarge 16c64g (Note: use a client >= 16c if you want to repro, an 8c client was the bottleneck during my test)

1. `python tests/load_tests/test_load_on_server.py -n 100 --apis cli_status`, **cold**

```text
# 14 Long and 56 Short (old estimation)
----------------------------------------------------------------------------------------------------
Kind                 Count    Total(s)   Avg(s)     Min(s)     Max(s)     P95(s)     P99(s)
----------------------------------------------------------------------------------------------------
sky status via API   100      2228.69    22.29      20.58      23.03      22.91      22.96

# 14 Long and 25 Short (new estimation)
sky status via API   100      1022.67    10.23      7.13       11.06      10.78      10.90
```

2. `python tests/load_tests/test_load_on_server.py -n 100 --api cli_status`, **warm**

```
# old estimation
----------------------------------------------------------------------------------------------------
Command              Count    Total(s)   Avg(s)     Min(s)     Max(s)     P95(s)     P99(s)
----------------------------------------------------------------------------------------------------
sky status via API   100      211.68     2.12       1.38       2.56       2.47       2.55

# new estimation
sky status via API   100      202.96     2.03       1.46       2.45       2.32       2.43
```

3. `python tests/load_tests/test_load_on_server.py -n 100 -api tail_logs`, **cold**

```
# old estimation
----------------------------------------------------------------------------------------------------
Command              Count    Total(s)   Avg(s)     Min(s)     Max(s)     P95(s)     P99(s)
----------------------------------------------------------------------------------------------------
API /tail_logs       100      5754.58    57.55      3.77       76.09      75.96      76.03

# new estimation
API /tail_logs       100      4163.28    41.63      2.36       65.19      64.98      65.10
```

4. `python tests/load_tests/test_load_on_server.py -n 100 --api tail_logs`, **warm**

```
# old estimation
----------------------------------------------------------------------------------------------------
Command              Count    Total(s)   Avg(s)     Min(s)     Max(s)     P95(s)     P99(s)
----------------------------------------------------------------------------------------------------
API /tail_logs       100      3535.23    35.35      4.09       53.39      53.25      53.33

# new estimation
API /tail_logs       100      3987.35    39.87      1.83       65.31      64.12      65.27
```

5. `pytest tests/test_smoke.py --aws`
  - old estimation: server hung and OOM kill found, interrupted
  - new estimation: 
    - complete without seeing resource issues
    - rerun failed flaky cases, passes
    - rerun cases failed due to azure credential issue, passes

For this specific test case, the new estimation is:

- more robust under complex load (smoke tests);
- has better performance during cold start;
- has same performance after warming compared to the old estimation. It turns out **the bottleneck is CPU** under high concurrency.

This conclusion can be generalized, because a typical deployment environment usually has a higher memory:CPU ratio (usually 4:1) compared to the test case, so the CPU bottleneck would be more significant.

## Where is the new estimation from? 

It is based on sys profiling gathered during smoke tests on GCP and AWS. The memory usage in smoke tests is still bit higher than the refined load tests:

- smoke tests (AWS & GCP)
```txt
Peak Non-blocking Executor Memory: 0.40GB
Peak Non-blocking Executor Memory Average: 0.23GB # This value is used, rounded to 0.5GB
Peak Blocking Executor Memory: 0.44GB
Peak Blocking Executor Memory Average: 0.38GB # This value is used, rounded to 0.5GB
```

- load tests (AWS & GCP & Kubernetes)
```txt
Peak Non-blocking Executor Memory: 0.37GB
Peak Non-blocking Executor Memory Average: 0.23GB
Peak Blocking Executor Memory: 0.35GB
Peak Blocking Executor Memory Average: 0.30GB
```

Since the RSS of executor process is rather stable, this subtle difference can cause https://github.com/skypilot-org/skypilot/issues/4732 in real-world scenarios while passing the load test. So the more conservative estimation from smoke tests is adopted by this PR.

> One thing I haven't figured out yet is why the RSS of executor process accumulates when performing different operations on different clouds. It seems like not a memory leak since the memory does not grow after performing same operations for multiple times. It's a good follow-up and I plan to make a detailed memory profile then. 

## Known issues

1. `sky logs --follow` may occupy short workers;
2. The resource usage of executor is highly dependent on the workload, so the estimation can be off a lot in a real-world case. Setting env vars to override the parallelism can be a mitigation for now.

## About 2 and 3

Note that points 2 and 3 are not must-to-have for fixing https://github.com/skypilot-org/skypilot/issues/4732. I just fixed them when investigating the root cause of hanging. Point 2 seems relevant, but the busy-wait will not block client requests since it yields to other coroutines in each loop. 

![image](https://github.com/user-attachments/assets/c93359e8-49e2-4f46-8649-79fc67161923)


<!-- Describe the tests ran -->
<!-- Unit tests (tests/test_*.py) are part of GitHub CI; below are tests that launch on the cloud. -->

Tested (run the relevant ones):

- [x] Code formatting: `bash format.sh`
- [ ] All smoke tests: `pytest tests/test_smoke.py` 
- [x] AWS smoke tests `pytest  tests/test_smoke.py --aws`
- [ ] Relevant individual smoke tests: `pytest tests/test_smoke.py::test_fill_in_the_name` 
- [ ] Backward compatibility tests: `conda deactivate; bash -i tests/backward_compatibility_tests.sh`




Problem statement:
[Client/Server] server hung on memory starvation
# Actual Behavior

I launched an API server on an 8c16g VM and set resources limits of the API server container to 7c14g via docker and test the server with parallel requests:

```bash
host> docker update --cpus 7 <container-id>
host> docker update --memory 14g --memory-swap 14g <container-id>

client> python tests/load_tests/test_load_on_server.py -n 100 -r launch
client> python tests/load_tests/test_load_on_server.py -n 100 -r status
```

The server container hung after ~30 seconds, then I inspected the container resource usage on the host:

```
# memory limit reached
host> docker stats --no-stream
CONTAINER ID   NAME            CPU %     MEM USAGE / LIMIT   MEM %     NET I/O   BLOCK I/O       PIDS
7605840c6190   sky_container   706.07%   14GiB / 14GiB       99.98%    0B / 0B   150GB / 124MB   764

# check cgroup interface files, memory.current is thrashing
host> cat memory.max
15032385536
host> while true; do cat memory.current; sleep 1; done;
15032377344
15032385536
15032385536
15032385536
15032385536
15032385536
15032385536
15032385536
15032385536
15032385536
15032381440
15032385536
15032377344
```

Also, there are extensive major pagefaults and heavy disk read:

```
host> cat memory.stat
...
pgfault 18139199
pgmajfault 772436 # major page faults, involve disk reads
pgrefill 5108525745
pgactivate 32401953
host> iotop:
  Total DISK READ:   95.15 M/s ⣀⣀⣀⣀⣀⣀⣀⣀⣀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ |   Total DISK WRITE:    0.00 B/s ⠀⠀⠀⠀⠀⠀⢸⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
Current DISK READ:   66.41 M/s ⣿⣿⣿⣿⣿⣿⣿⣿⣿⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ | Current DISK WRITE:    0.00 B/s ⢸⠀
```

# Expected Behavior

API server should not hung, either fail fast or throttle / reject requests to avoid memory starvation. Besides, there is an interesting of cgroup v2 here: it looks like there is always **some** memory to reclaim when memory.max is reached so the API server is never get really killed, such reclamation cause extensive major pagefaults and leads to server hung.

Hints from issue discussion:

Commit message: relax resource constrain

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 561cb01466a..d3c801586c2 100644
--- a/sky/server/requests/executor.py
+++ b/sky/server/requests/executor.py
@@ -73,8 +73,13 @@
 # Constants based on profiling the peak memory usage of
 # various sky commands. See `tests/load_test/` for details.
 # Max memory consumption for each request.
-_PER_BLOCKING_REQUEST_MEM_GB = 0.25
-_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.15
+# NOTE(dev): update these constants for each release according to the load
+# test results.
+# TODO(aylei): maintaining these constants is error-prone, we may need to
+# automatically tune parallelism at runtime according to system usage stats
+# in the future.
+_PER_BLOCKING_REQUEST_MEM_GB = 0.4
+_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.25
 # To control the number of blocking workers.
 _CPU_MULTIPLIER_FOR_BLOCKING_WORKERS = 2
 _MAX_BLOCKING_WORKERS_LOCAL = 4
@@ -301,34 +306,33 @@ def schedule_request(request_id: str,
     _get_queue(schedule_type).put(input_tuple)
 
 
+def executor_initializer(sub_title: str):
+    setproctitle.setproctitle(
+        f'{sub_title}:{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.
     """
-    logger.info(f'Starting {worker} with pid '
-                f'{multiprocessing.current_process().pid}')
     setproctitle.setproctitle(
         f'SkyPilot:worker:{worker.schedule_type.value}-{worker.id}')
+    sub_title = f'SkyPilot:executor:{worker.schedule_type.value}-{worker.id}'
     queue = _get_queue(worker.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.
-    with concurrent.futures.ProcessPoolExecutor(
-            max_workers=max_parallel_size) as executor:
-        while True:
+
+    def process_request(executor: concurrent.futures.ProcessPoolExecutor):
+        try:
             request_element = queue.get()
             if request_element is None:
                 time.sleep(0.1)
-                continue
+                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:
-                continue
+                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.
@@ -347,6 +351,23 @@ def request_worker(worker: RequestWorker, max_parallel_size: int) -> None:
                 logger.info(f'[{worker}] Finished request: {request_id}')
             else:
                 logger.info(f'[{worker}] Submitted request: {request_id}')
+        except Exception as e:  # pylint: disable=broad-except
+            logger.error(
+                f'[{worker}] Error processing request {request_id}: '
+                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.
+    with concurrent.futures.ProcessPoolExecutor(
+            max_workers=max_parallel_size,
+            initializer=executor_initializer,
+            initargs=(sub_title,)) as executor:
+        while True:
+            process_request(executor)
 
 
 def _get_cpu_count() -> int:
@@ -390,11 +411,9 @@ def start(deploy: bool) -> List[multiprocessing.Process]:
     cpu_count = _get_cpu_count()
     mem_size_gb = _get_mem_size_gb()
     mem_size_gb = max(0, mem_size_gb - server_constants.MIN_AVAIL_MEM_GB)
-    parallel_for_blocking = _max_parallel_size_for_blocking(
-        cpu_count, mem_size_gb)
-    if not deploy:
-        parallel_for_blocking = min(parallel_for_blocking,
-                                    _MAX_BLOCKING_WORKERS_LOCAL)
+    parallel_for_blocking = _max_parallel_size_for_blocking(cpu_count,
+                                                            mem_size_gb,
+                                                            local=not deploy)
     max_parallel_for_non_blocking = _max_parallel_size_for_non_blocking(
         mem_size_gb, parallel_for_blocking)
     logger.info(
@@ -443,12 +462,28 @@ def start(deploy: bool) -> List[multiprocessing.Process]:
 
 
 @annotations.lru_cache(scope='global', maxsize=1)
-def _max_parallel_size_for_blocking(cpu_count: int, mem_size_gb: float) -> int:
+def _max_parallel_size_for_blocking(cpu_count: int,
+                                    mem_size_gb: float,
+                                    local=False) -> int:
     """Max parallelism for blocking requests."""
+    # since the estimation can be off by a lot, we allow users to override
+    # the default value by setting the environment variable.
+    env_parallel_size = os.getenv('SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS')
+    if env_parallel_size is not None:
+        try:
+            n = int(env_parallel_size)
+            return max(1, n)
+        except ValueError:
+            logger.warning(
+                'Invalid SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS value: '
+                f'{env_parallel_size}. Falling back to calculated value.')
+
     cpu_based_max_parallel = cpu_count * _CPU_MULTIPLIER_FOR_BLOCKING_WORKERS
     mem_based_max_parallel = int(mem_size_gb * _MAX_MEM_PERCENT_FOR_BLOCKING /
                                  _PER_BLOCKING_REQUEST_MEM_GB)
     n = max(1, min(cpu_based_max_parallel, mem_based_max_parallel))
+    if local:
+        return min(n, _MAX_BLOCKING_WORKERS_LOCAL)
     return n
 
 
@@ -456,6 +491,19 @@ def _max_parallel_size_for_blocking(cpu_count: int, mem_size_gb: float) -> int:
 def _max_parallel_size_for_non_blocking(mem_size_gb: float,
                                         parallel_size_for_blocking: int) -> int:
     """Max parallelism for non-blocking requests."""
+    # since the estimation can be off by a lot, we allow users to override
+    # the default value by setting the environment variable.
+    env_parallel_size = os.getenv('SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS')
+    if env_parallel_size is not None:
+        try:
+            n = int(env_parallel_size)
+            return max(1, n)
+        except ValueError:
+            logger.warning(
+                'Invalid SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS value: '
+                f'{env_parallel_size}. Falling back to calculated value.')
+
+    # Fall back to calculation if env var is not set or invalid
     available_mem = mem_size_gb - (parallel_size_for_blocking *
                                    _PER_BLOCKING_REQUEST_MEM_GB)
     n = max(1, int(available_mem / _PER_NON_BLOCKING_REQUEST_MEM_GB))
diff --git a/sky/server/server.py b/sky/server/server.py
index 8c00aecc6a1..9c315424f6f 100644
--- a/sky/server/server.py
+++ b/sky/server/server.py
@@ -57,7 +57,9 @@
 
 def _add_timestamp_prefix_for_server_logs() -> None:
     server_logger = sky_logging.init_logger('sky.server')
-    # Disable propagation to avoid the root logger of SkyPilot being affected.
+    # Clear existing handlers first to prevent duplicates
+    server_logger.handlers.clear()
+    # Disable propagation to avoid the root logger of SkyPilot being affected
     server_logger.propagate = False
     # Add date prefix to the log message printed by loggers under
     # server.
@@ -70,6 +72,8 @@ def _add_timestamp_prefix_for_server_logs() -> None:
         uvicorn_logger = logging.getLogger(name)
         uvicorn_logger.handlers.clear()
         uvicorn_logger.addHandler(stream_handler)
+        # Also disable propagation for uvicorn loggers
+        uvicorn_logger.propagate = False
 
 
 _add_timestamp_prefix_for_server_logs()
@@ -794,10 +798,9 @@ async def api_get(request_id: str) -> requests_lib.RequestPayload:
                                             detail=dataclasses.asdict(
                                                 request_task.encode()))
             return request_task.encode()
-        # Sleep 0 to yield, so other coroutines can run. This busy waiting
-        # loop is performance critical for short-running requests, so we do
-        # not want to yield too long.
-        await asyncio.sleep(0)
+        # yield control to allow other coroutines to run, sleep shortly
+        # to avoid storming the DB and CPU in the meantime
+        await asyncio.sleep(0.1)
 
 
 @app.get('/api/stream')
diff --git a/sky/server/stream_utils.py b/sky/server/stream_utils.py
index 58fbb8b51e7..3ec11ae314f 100644
--- a/sky/server/stream_utils.py
+++ b/sky/server/stream_utils.py
@@ -68,7 +68,7 @@ async def log_streamer(request_id: Optional[str],
             # Sleep 0 to yield, so other coroutines can run. This busy waiting
             # loop is performance critical for short-running requests, so we do
             # not want to yield too long.
-            await asyncio.sleep(0)
+            await asyncio.sleep(0.1)
             request_task = requests_lib.get_request(request_id)
             if not follow:
                 break
@@ -88,6 +88,9 @@ async def log_streamer(request_id: Optional[str],
                 yield line_str
 
         while True:
+            # Sleep 0 to yield control to allow other coroutines to run,
+            # while keeps the loop tight to make log stream responsive.
+            await asyncio.sleep(0)
             line: Optional[bytes] = await f.readline()
             if not line:
                 if request_id is not None:
@@ -100,24 +103,18 @@ async def log_streamer(request_id: Optional[str],
                         break
                 if not follow:
                     break
-
-                # Sleep 0 to yield, so other coroutines can run. This busy
-                # waiting loop is performance critical for short-running
-                # requests, so we do not want to yield too long.
-                await asyncio.sleep(0)
+                # Sleep shortly to avoid storming the DB and CPU, this has
+                # little impact on the responsivness here since we are waiting
+                # for a new line to come in.
+                await asyncio.sleep(0.1)
                 continue
             line_str = line.decode('utf-8')
             if plain_logs:
                 is_payload, line_str = message_utils.decode_payload(
                     line_str, raise_for_mismatch=False)
                 if is_payload:
-                    # Sleep 0 to yield, so other coroutines can run. This busy
-                    # waiting loop is performance critical for short-running
-                    # requests, so we do not want to yield too long.
-                    await asyncio.sleep(0)
                     continue
             yield line_str
-            await asyncio.sleep(0)  # Allow other tasks to run
 
 
 def stream_response(
diff --git a/tests/load_tests/README.md b/tests/load_tests/README.md
index 2be0d64c2ba..d5997c4d3b9 100644
--- a/tests/load_tests/README.md
+++ b/tests/load_tests/README.md
@@ -5,6 +5,9 @@ This folder contains tools for load testing and profiling the SkyPilot API serve
 * System Profiling (`sys_profiling.py`): monitors system resource usage (CPU and memory) over time
 * Load Testing (`test_load_on_server.py`): sends concurrent requests to stress test the SkyPilot API server
 
+> **Note**: The load testing workload is simple and may not reflect the usage of the SkyPilot API server in real-world scenarios.
+> You may consider running part of or all smoke tests to get a more accurate measurement.
+
 ## Instructions
 
 1. Start the SkyPilot API server on a remote machine
@@ -17,27 +20,31 @@ python tests/load_tests/sys_profiling.py
 ```
 3. Run the load testing script locally to send requests to the server
 ```bash
-python tests/load_tests/test_load_on_server.py -n 100 -r launch
+python tests/load_tests/test_load_on_server.py -n 50 -r all --cloud aws
 ```
 4. Once the load testing is done, terminate the profiling process to get a resource usage summary printed in the terminal, e.g.
 ```bash
 ==================================================
 MONITORING SUMMARY
 ==================================================
-Duration: 28.6 seconds (0.01 hours)
+Duration: 1481.7 seconds (0.41 hours)
 
 BASELINE USAGE:
-Baseline CPU: 2.5%
-Baseline Memory: 14.17GB (47.6%)
+Baseline CPU: 1.2%
+Baseline Memory: 2.99GB (21.6%)
 
 PEAK USAGE:
-Peak CPU: 100.0%
-Peak Memory: 15.88GB (53.2%)
-Memory Delta: 1.7GB
+Peak CPU: 96.9%
+Peak Memory: 11.78GB (79.0%)
+Memory Delta: 8.8GB
+Peak Non-blocking Executor Memory: 0.37GB
+Peak Non-blocking Executor Memory Average: 0.23GB
+Peak Blocking Executor Memory: 0.35GB
+Peak Blocking Executor Memory Average: 0.30GB
 
 AVERAGE USAGE:
-Average CPU: 19.3%
-Average Memory: 49.5%
+Average CPU: 11.5%
+Average Memory: 63.1%
 ==================================================
 ```
-
+5. Update the `_PER_BLOCKING_REQUEST_MEM_GB` and `_PER_NON_BLOCKING_REQUEST_MEM_GB` constants in `sky/server/requests/executor.py` based on the peak memory usage of the blocking and non-blocking workers.
diff --git a/tests/load_tests/serve.yaml b/tests/load_tests/serve.yaml
new file mode 100644
index 00000000000..883ff9a8db1
--- /dev/null
+++ b/tests/load_tests/serve.yaml
@@ -0,0 +1,12 @@
+service:
+  readiness_probe:
+    path: /health
+    initial_delay_seconds: 180
+  replicas: 2
+
+resources:
+  ports: 8080
+  cpus: 2+
+
+workdir: examples/serve/http_server
+run: python3 server.py --port 8080
\ No newline at end of file
diff --git a/tests/load_tests/sys_profiling.py b/tests/load_tests/sys_profiling.py
index 5c0ab2dd1bd..1d00a5efde9 100644
--- a/tests/load_tests/sys_profiling.py
+++ b/tests/load_tests/sys_profiling.py
@@ -16,12 +16,34 @@ def get_size_gb(bytes):
     return bytes / (1024**3)
 
 
+def get_worker_memory(process_keyword):
+    top = 0
+    acc = 0
+    count = 0
+    for proc in psutil.process_iter(['name', 'cmdline']):
+        try:
+            cmdline = ' '.join(proc.cmdline()) if proc.cmdline() else ''
+            if process_keyword in cmdline:
+                top = max(top, proc.memory_info().rss)
+                acc += proc.memory_info().rss
+                count += 1
+        except (psutil.NoSuchProcess, psutil.AccessDenied,
+                psutil.ZombieProcess):
+            pass
+    avg = acc / count if count > 0 else 0
+    return top, avg
+
+
 def monitor_system():
     # Initialize peak tracking variables
     interval = 0.2
     peak_cpu = 0
     peak_memory_percent = 0
     peak_memory_used = 0
+    nonblocking_worker_memory_peak = 0
+    blocking_worker_memory_peak = 0
+    nonblocking_worker_memory_avg_peak = 0
+    blocking_worker_memory_avg_peak = 0
     start_time = datetime.now()
     samples = 0
     total_cpu = 0
@@ -35,6 +57,17 @@ def monitor_system():
 
     try:
         while True:
+            nb_top, nb_avg = get_worker_memory('SkyPilot:executor:short')
+            nonblocking_worker_memory_peak = max(nonblocking_worker_memory_peak,
+                                                 nb_top)
+            nonblocking_worker_memory_avg_peak = max(
+                nonblocking_worker_memory_avg_peak, nb_avg)
+            b_top, b_avg = get_worker_memory('SkyPilot:executor:long')
+            blocking_worker_memory_peak = max(blocking_worker_memory_peak,
+                                              b_top)
+            blocking_worker_memory_avg_peak = max(
+                blocking_worker_memory_avg_peak, b_avg)
+
             # CPU Usage
             cpu_percent = psutil.cpu_percent(interval=interval)
 
@@ -91,6 +124,14 @@ def monitor_system():
         print(
             f"Memory Delta: {get_size_gb(peak_memory_used - baseline_memory.used):.1f}GB"
         )
+        print("Peak Non-blocking Executor Memory: "
+              f"{get_size_gb(nonblocking_worker_memory_peak):.2f}GB")
+        print("Peak Non-blocking Executor Memory Average: "
+              f"{get_size_gb(nonblocking_worker_memory_avg_peak):.2f}GB")
+        print("Peak Blocking Executor Memory: "
+              f"{get_size_gb(blocking_worker_memory_peak):.2f}GB")
+        print("Peak Blocking Executor Memory Average: "
+              f"{get_size_gb(blocking_worker_memory_avg_peak):.2f}GB")
         print("\nAVERAGE USAGE:")
         print(f"Average CPU: {avg_cpu:.1f}%")
         print(f"Average Memory: {avg_memory:.1f}%")
diff --git a/tests/load_tests/test_load_on_server.py b/tests/load_tests/test_load_on_server.py
index e0600eb52fe..ce6091dfa96 100644
--- a/tests/load_tests/test_load_on_server.py
+++ b/tests/load_tests/test_load_on_server.py
@@ -6,6 +6,7 @@
 """
 
 import argparse
+import os
 import subprocess
 import threading
 import time
@@ -39,21 +40,24 @@ def run_concurrent_requests(num_requests, cmd):
         thread.join()
 
 
-def test_launch_requests(num_requests, is_async=True):
-    cmd = 'sky launch --cloud=aws --cpus=2 -y'
+def test_launch_requests(num_requests, cloud, is_async=True):
+    print(f"Testing {num_requests} launch requests")
+    cmd = f'sky launch --cloud={cloud} --cpus=2 -y'
     if is_async:
         cmd += ' --async'
     run_concurrent_requests(num_requests, cmd)
 
 
 def test_status_requests(num_requests):
+    print(f"Testing {num_requests} status requests")
     run_concurrent_requests(num_requests, 'sky status')
 
 
-def test_logs_requests(num_requests):
+def test_logs_requests(num_requests, cloud):
+    print(f"Testing {num_requests} logs requests")
     cluster_name = 'test'
     job_id = 1
-    setup_cmd = f'sky launch -c {cluster_name} --cloud=aws --cpus=2 -y'
+    setup_cmd = f'sky launch -c {cluster_name} --cloud={cloud} --cpus=2 -y'
     run_single_request(0, setup_cmd)
     cmd = f'sky logs {cluster_name} {job_id}'
     run_concurrent_requests(num_requests, cmd)
@@ -61,25 +65,87 @@ def test_logs_requests(num_requests):
     run_single_request(0, cleanup_cmd)
 
 
+def test_jobs_requests(num_requests, cloud, is_async=True):
+    print(f"Testing {num_requests} jobs launch requests")
+    cmd = f'sky jobs launch --cloud={cloud} --cpus=2 "echo hello" -y'
+    if is_async:
+        cmd += ' --async'
+    run_concurrent_requests(num_requests, cmd)
+
+
+def test_serve_requests(num_requests, cloud):
+    print(f"Testing {num_requests} serve status requests")
+    svc_name = 'test-serve'
+    # Use serve.yaml from the same directory
+    serve_yaml_path = os.path.join(os.path.dirname(__file__), 'serve.yaml')
+    setup_cmd = f'sky serve up -n {svc_name} --cloud={cloud} {serve_yaml_path} -y'
+    run_single_request(0, setup_cmd)
+
+    cmd = f'sky serve status {svc_name}'
+    run_concurrent_requests(num_requests, cmd)
+
+    cleanup_cmd = f'sky serve down {svc_name} -y'
+    run_single_request(0, cleanup_cmd)
+
+
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
     parser.add_argument('-n',
                         type=int,
                         default=1,
                         help='Number of requests to make')
-    parser.add_argument('-r',
+    parser.add_argument(
+        '-r',
+        '--requests',
+        type=str,
+        nargs='+',
+        choices=['launch', 'status', 'logs', 'jobs', 'serve', 'all'],
+        help='List of SkyPilot requests to test (e.g., launch status logs)')
+    # Test worker memory consumption when operates multiple clouds.
+    parser.add_argument('-c',
+                        '--clouds',
                         type=str,
-                        choices=['launch', 'status', 'logs'],
-                        help='Sky request name')
+                        nargs='+',
+                        default=['aws'],
+                        help='List of clouds to test')
+
     args = parser.parse_args()
 
+    if 'all' in args.requests:
+        args.requests = ['launch', 'status', 'logs', 'jobs', 'serve']
+
     print("Starting concurrent sky requests...")
     start_time = time.time()
-    if args.r == 'launch':
-        test_launch_requests(args.n, is_async=True)
-    elif args.r == 'status':
-        test_status_requests(args.n)
-    elif args.r == 'logs':
-        test_logs_requests(args.n)
+
+    # Create threads for each test type
+    test_threads = []
+    for cloud in args.clouds:
+        print(f"\nTesting cloud: {cloud}")
+        for request_type in args.requests:
+            thread = None
+            if request_type == 'launch':
+                thread = threading.Thread(target=test_launch_requests,
+                                          args=(args.n, cloud, True))
+            elif request_type == 'status':
+                thread = threading.Thread(target=test_status_requests,
+                                          args=(args.n,))
+            elif request_type == 'logs':
+                thread = threading.Thread(target=test_logs_requests,
+                                          args=(args.n, cloud))
+            elif request_type == 'jobs':
+                thread = threading.Thread(target=test_jobs_requests,
+                                          args=(args.n, cloud))
+            elif request_type == 'serve':
+                thread = threading.Thread(target=test_serve_requests,
+                                          args=(args.n, cloud))
+
+            if thread:
+                test_threads.append(thread)
+                thread.start()
+
+    # Wait for all test threads to complete
+    for thread in test_threads:
+        thread.join()
+
     end_time = time.time()
     print(f"All requests completed in {end_time - start_time:.2f} seconds")
diff --git a/tests/smoke_tests/test_basic.py b/tests/smoke_tests/test_basic.py
index 202a0afee81..b63360f11a4 100644
--- a/tests/smoke_tests/test_basic.py
+++ b/tests/smoke_tests/test_basic.py
@@ -145,7 +145,8 @@ def test_launch_fast_with_autostop(generic_cloud: str):
                 timeout=autostop_timeout),
             # Even the cluster is stopped, cloud platform may take a while to
             # delete the VM.
-            f'sleep 35',
+            # FIXME(aylei): this can be flaky, sleep longer for now.
+            f'sleep 60',
             # Launch again. Do full output validation - we expect the cluster to re-launch
             f'unset SKYPILOT_DEBUG; s=$(sky launch -y -c {name} --fast -i 1 tests/test_yamls/minimal.yaml) && {smoke_tests_utils.VALIDATE_LAUNCH_OUTPUT}',
             f'sky logs {name} 2 --status',
diff --git a/tests/smoke_tests/test_cluster_job.py b/tests/smoke_tests/test_cluster_job.py
index 549a5d1ee1b..1515495f54d 100644
--- a/tests/smoke_tests/test_cluster_job.py
+++ b/tests/smoke_tests/test_cluster_job.py
@@ -419,7 +419,7 @@ def test_multi_echo(generic_cloud: str):
             f's=$(sky queue {name}); echo "$s"; echo; echo; echo "$s" | grep "FAILED" && exit 1 || true',
             'sleep 30',
             # Make sure that our job scheduler is fast enough to have at least
-            # 15 RUNNING jobs in parallel.
+            # 18 RUNNING jobs in parallel.
             f's=$(sky queue {name}); echo "$s"; echo; echo; echo "$s" | grep "RUNNING" | wc -l | awk \'{{if ($1 < 18) exit 1}}\'',
             'sleep 30',
             f's=$(sky queue {name}); echo "$s"; echo; echo; echo "$s" | grep "FAILED" && exit 1 || true',
@@ -1265,13 +1265,13 @@ def test_autodown(generic_cloud: str):
     # the VM is terminated.
     autodown_timeout = 900 if generic_cloud == 'azure' else 240
     total_timeout_minutes = 90 if generic_cloud == 'azure' else 20
+    check_autostop_set = f's=$(sky status) && echo "$s" && echo "==check autostop set==" && echo "$s" | grep {name} | grep "1m (down)"'
     test = smoke_tests_utils.Test(
         'autodown',
         [
             f'sky launch -y -d -c {name} --num-nodes 2 --cloud {generic_cloud} tests/test_yamls/minimal.yaml',
             f'sky autostop -y {name} --down -i 1',
-            # Ensure autostop is set.
-            f'sky status | grep {name} | grep "1m (down)"',
+            check_autostop_set,
             # Ensure the cluster is not terminated early.
             'sleep 40',
             f's=$(sky status {name} --refresh); echo "$s"; echo; echo; echo "$s"  | grep {name} | grep UP',
@@ -1281,7 +1281,7 @@ def test_autodown(generic_cloud: str):
             f'sky launch -y -d -c {name} --cloud {generic_cloud} --num-nodes 2 --down tests/test_yamls/minimal.yaml',
             f'sky status | grep {name} | grep UP',  # Ensure the cluster is UP.
             f'sky exec {name} --cloud {generic_cloud} tests/test_yamls/minimal.yaml',
-            f'sky status | grep {name} | grep "1m (down)"',
+            check_autostop_set,
             f'sleep {autodown_timeout}',
             # Ensure the cluster is terminated.
             f's=$(SKYPILOT_DEBUG=0 sky status {name} --refresh) && echo "$s" && {{ echo "$s" | grep {name} | grep "Autodowned cluster\|Cluster \'{name}\' not found"; }} || {{ echo "$s" | grep {name} && exit 1 || exit 0; }}',
diff --git a/tests/smoke_tests/test_managed_job.py b/tests/smoke_tests/test_managed_job.py
index e98252337f1..e2d8808d324 100644
--- a/tests/smoke_tests/test_managed_job.py
+++ b/tests/smoke_tests/test_managed_job.py
@@ -856,7 +856,9 @@ def test_managed_jobs_storage(generic_cloud: str):
                 *smoke_tests_utils.STORAGE_SETUP_COMMANDS,
                 smoke_tests_utils.launch_cluster_for_cloud_cmd(
                     generic_cloud, name),
-                f'sky jobs launch -n {name}{use_spot} --cloud {generic_cloud}{region_flag} {file_path} -y -d',
+                # Override CPU/memory requirements to relax resource constraints
+                # and reduce the chance of out-of-stock
+                f'sky jobs launch -n {name}{use_spot} --cpus 2+ --memory 4+ --cloud {generic_cloud}{region_flag} {file_path} -y -d',
                 region_validation_cmd,  # Check if the bucket is created in the correct region
                 smoke_tests_utils.
                 get_cmd_wait_until_managed_job_status_contains_matching_job_name(
diff --git a/tests/unit_tests/sky/api/requests/test_executor.py b/tests/unit_tests/sky/api/requests/test_executor.py
index e83bb95b645..1afcc6d383c 100644
--- a/tests/unit_tests/sky/api/requests/test_executor.py
+++ b/tests/unit_tests/sky/api/requests/test_executor.py
@@ -1,8 +1,18 @@
+import os
+
 import pytest
 
 from sky.server.requests import executor
 
 
+@pytest.fixture(autouse=True)
+def clear_lru_cache():
+    """Clear the LRU cache before each test."""
+    executor._max_parallel_size_for_blocking.cache_clear()
+    executor._max_parallel_size_for_non_blocking.cache_clear()
+    yield
+
+
 def test_parallel_size_blocking():
     # Test with insufficient memory
     cpu_count = 4
@@ -25,6 +35,20 @@ def test_parallel_size_blocking():
     assert executor._max_parallel_size_for_blocking(cpu_count,
                                                     mem_size_gb) == expected
 
+    # Test with environment variable override
+    os.environ['SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS'] = '3'
+    executor._max_parallel_size_for_blocking.cache_clear()
+    assert executor._max_parallel_size_for_blocking(cpu_count, mem_size_gb) == 3
+    del os.environ['SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS']
+    executor._max_parallel_size_for_blocking.cache_clear()
+
+    # Test with invalid environment variable
+    os.environ['SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS'] = 'invalid'
+    executor._max_parallel_size_for_blocking.cache_clear()
+    assert executor._max_parallel_size_for_blocking(cpu_count, mem_size_gb) == 1
+
+    del os.environ['SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS']
+
 
 def test_parallel_size_non_blocking():
     # Test with insufficient memory
@@ -37,13 +61,28 @@ def test_parallel_size_non_blocking():
     # Test with sufficient memory
     blocking_size = 8
     mem_size_gb = 10.5
-    expected = 56
+    expected = 29
     assert executor._max_parallel_size_for_non_blocking(
         mem_size_gb, blocking_size) == expected
 
     # Test with limited memory
     blocking_size = 1
     mem_size_gb = 1
-    expected = 5
+    expected = 2
     assert executor._max_parallel_size_for_non_blocking(
         mem_size_gb, blocking_size) == expected
+
+    # Test with environment variable override
+    os.environ['SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS'] = '10'
+    executor._max_parallel_size_for_non_blocking.cache_clear()
+    assert executor._max_parallel_size_for_non_blocking(mem_size_gb,
+                                                        blocking_size) == 10
+    del os.environ['SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS']
+    executor._max_parallel_size_for_non_blocking.cache_clear()
+
+    # Test with invalid environment variable
+    os.environ['SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS'] = 'invalid'
+    executor._max_parallel_size_for_non_blocking.cache_clear()
+    assert executor._max_parallel_size_for_non_blocking(mem_size_gb,
+                                                        blocking_size) == 2
+    del os.environ['SKYPILOT_MAX_PARALLEL_NONBLOCKING_REQUESTS']


Code review comments:
diff_hunk: @@ -73,8 +73,13 @@
 # Constants based on profiling the peak memory usage of
 # various sky commands. See `tests/load_test/` for details.
 # Max memory consumption for each request.
-_PER_BLOCKING_REQUEST_MEM_GB = 0.25
-_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.15
+# NOTE(dev): update these constants for each release according to the load
+# test results.
+# TODO(aylei): maintaining these constants is error-prone, we may need to
+# automatically tune parallelism at runtime according to system usage stats
+# in the future.
+_PER_BLOCKING_REQUEST_MEM_GB = 0.4
path: sky/server/requests/executor.py
text: @user1:
We have recently renamed the BLOCKING/NON_BLOCKING to `LONG` and `SHORT`. We should update these variable names.

@author:
sounds good

diff_hunk: @@ -73,8 +73,13 @@
 # Constants based on profiling the peak memory usage of
 # various sky commands. See `tests/load_test/` for details.
 # Max memory consumption for each request.
-_PER_BLOCKING_REQUEST_MEM_GB = 0.25
-_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.15
+# NOTE(dev): update these constants for each release according to the load
+# test results.
path: sky/server/requests/executor.py
text: @user1:
We previously used an average number for the resource limitation, so that we don't limit the total number of parallelism because a too conservative number. Setting these to peak number could avoid overscribe the resources but may affect the normal use cases. Would it be possible we have some kind of solution to ensure a large parallelism but avoid causing OOM

@author:
```
Peak Non-blocking Executor Memory: 0.40GB
Peak Non-blocking Executor Memory Average: 0.23GB
Peak Blocking Executor Memory: 0.44GB
Peak Blocking Executor Memory Average: 0.38GB
```
The average of peak time is used (with ceiling) since reaching `memory.max` causes a server hung or OOM kill (typically server hung, see https://github.com/skypilot-org/skypilot/issues/4732).

> Would it be possible we have some kind of solution to ensure a large parallelism but avoid causing OOM

Still investigating this, current directions:
- reduce the memory usage of executor process, then we can tune the constants to get more parallelism
- a mixed parallelism of threads and processes for executors, which is non-trivial and need more investigation and design I think

@author:
Besides, making memory throttling elegant (so that we can use average over lifetime instead of average of peak time) would also help increase overall utilization.

@author:
Need some benchmark

diff_hunk: @@ -301,34 +306,33 @@ def schedule_request(request_id: str,
     _get_queue(schedule_type).put(input_tuple)
 
 
+def executor_initializer(sub_title: str):
+    setproctitle.setproctitle(
+        f'{sub_title}:{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.
     """
-    logger.info(f'Starting {worker} with pid '
-                f'{multiprocessing.current_process().pid}')
     setproctitle.setproctitle(
         f'SkyPilot:worker:{worker.schedule_type.value}-{worker.id}')
+    sub_title = f'SkyPilot:executor:{worker.schedule_type.value}-{worker.id}'
path: sky/server/requests/executor.py
text: @user1:
reuse this variable in the line above as well?

diff_hunk: @@ -73,8 +73,13 @@
 # Constants based on profiling the peak memory usage of
 # various sky commands. See `tests/load_test/` for details.
 # Max memory consumption for each request.
-_PER_BLOCKING_REQUEST_MEM_GB = 0.25
-_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.15
+# NOTE(dev): update these constants for each release according to the load
+# test results.
+# TODO(aylei): maintaining these constants is error-prone, we may need to
+# automatically tune parallelism at runtime according to system usage stats
+# in the future.
+_PER_BLOCKING_REQUEST_MEM_GB = 0.4
+_PER_NON_BLOCKING_REQUEST_MEM_GB = 0.25
path: sky/server/requests/executor.py
text: @user1:
we could also add a note that these values can vary significantly depending on which clouds are enabled, which I observed during load testing on jobs controller
FWIW the "equivalent" value for jobs controller for `_PER_BLOCKING_REQUEST_MEM_GB` is set to 350MiB at the moment, which is fine for AWS but a bit sketchy on some other clouds
https://github.com/skypilot-org/skypilot/blob/a759f24649aa9faff5b88f21b8f08b367c55c170/sky/jobs/scheduler.py#L63-L64

diff_hunk: @@ -443,19 +462,48 @@ def start(deploy: bool) -> List[multiprocessing.Process]:
 
 
 @annotations.lru_cache(scope='global', maxsize=1)
-def _max_parallel_size_for_blocking(cpu_count: int, mem_size_gb: float) -> int:
+def _max_parallel_size_for_blocking(cpu_count: int,
+                                    mem_size_gb: float,
+                                    local=False) -> int:
     """Max parallelism for blocking requests."""
+    # since the estimation can be off by a lot, we allow users to override
+    # the default value by setting the environment variable.
+    env_parallel_size = os.getenv('SKYPILOT_MAX_PARALLEL_BLOCKING_REQUESTS')
path: sky/server/requests/executor.py
text: @user1:
We should have a centralized place to have all env var options, e.g. we can have a constant in `sky.skylet.constants`.