In Python versions prior to 3.11, TimeoutError is not a built-in exception, and asyncio.wait_for raises asyncio.TimeoutError. Using the bare TimeoutError here will cause a NameError on those Python versions when a timeout occurs. It is safer to use asyncio.TimeoutError for compatibility.

To maintain consistency with handle_loop and avoid logic duplication, consider refactoring the ZMQ error handling into a shared helper function. The tokenizer loop should handle zmq.error.Again and include a traceback for ZMQError to ensure transient timeouts don't unexpectedly terminate the task and to provide better debugging information.

Fixed — added zmq.error.Again handler and logger.debug on graceful shutdown to match handle_loop's pattern exactly.

Using time.time() for a profile ID can lead to collisions if multiple profiles are started in rapid succession and results in a string representation that may be less ideal for filenames (containing dots). Using a UUID is more robust and consistent with how request_id is generated elsewhere in this class.

Catch asyncio timeout in FlushCache handler

send_communicator_req() is backed by asyncio.wait_for, which raises asyncio.TimeoutError on Python 3.10, but this handler catches built-in TimeoutError, so timeout cases bypass this branch and get reported by the generic Exception path as INTERNAL instead of DEADLINE_EXCEEDED. This misclassifies real scheduler timeouts and breaks client retry/diagnostic behavior for FlushCache requests under load or partial scheduler stalls.

Useful? React with 👍 / 👎.

Bound gRPC flush calls with per-worker timeout

The gRPC fan-out path awaits grpc_client.flush_cache(0.0) without any timeout wrapper, unlike the HTTP branch which enforces REQUEST_TIMEOUT; as a result, one slow or unreachable gRPC worker can stall join_all and keep /flush_cache hanging for an unbounded/very long time. This turns an admin endpoint into a potential blocker during partial outages, so each gRPC flush should be wrapped with a per-worker timeout and treated as a failed worker on expiry.

Useful? React with 👍 / 👎.

⚠️ Potential issue | 🟠 Major

Clean up launched resources if the startup hook fails.

Line 120 awaits the hook after the scheduler processes and GrpcRequestManager already exist, but before serve_grpc() reaches its shutdown path. If that callback raises, startup aborts with child processes and ZMQ sockets still alive.

Move this hook under the same startup/teardown guard as the rest of the server bootstrap, or explicitly shut down request_manager, bootstrap_server, and scheduler_procs before re-raising.

Verify each finding against the current code and only fix it if needed.

In `@grpc_servicer/smg_grpc_servicer/sglang/server.py` around lines 120 - 121, The
on_request_manager_ready callback is awaited after creating request_manager and
scheduler_procs but before the main shutdown guard, so any exception leaks child
processes/sockets; wrap the await on_request_manager_ready(request_manager,
server_args, scheduler_info) in a try/except (or move the await inside the
existing startup/teardown guard) and on exception perform clean shutdown steps:
call request_manager.shutdown() (or its stop/close method),
stop/bootstrap_server.shutdown() (or the server close routine), and
terminate/join each process in scheduler_procs before re-raising the exception;
ensure you reference the same symbols (on_request_manager_ready,
request_manager, bootstrap_server, scheduler_procs) to locate and run their
existing cleanup paths.

⚠️ Potential issue | 🔴 Critical

#!/bin/bash
set -euo pipefail
python - <<'PY'
from pathlib import Path

path = Path("model_gateway/src/routers/grpc/client.rs")
text = path.read_text()

enum_block = text.split("pub enum GrpcClient", 1)[1].split("impl GrpcClient", 1)[0]
method_block = text.split("pub async fn flush_cache", 1)[1].split("/// Fetch tokenizer bundle", 1)[0]

print("enum_has_mlx =", "Mlx(" in enum_block)
print("flush_cache_handles_mlx =", "Self::Mlx" in method_block)
PY

Handle the GrpcClient::Mlx variant in this match.

The match self on line 238 is non-exhaustive. The GrpcClient enum includes the Mlx variant, but this method does not handle it, preventing compilation.

     pub async fn flush_cache(&self, timeout_s: f32) -> Result<(bool, String), tonic::Status> {
         match self {
             Self::Sglang(client) => {
                 let resp = client.flush_cache(timeout_s).await?;
                 Ok((resp.success, resp.message))
             }
-            Self::Vllm(_) | Self::Trtllm(_) => Err(tonic::Status::unimplemented(
+            Self::Vllm(_) | Self::Trtllm(_) | Self::Mlx(_) => Err(tonic::Status::unimplemented(
                 "flush_cache not supported for this backend",
             )),
         }
     }

Verify each finding against the current code and only fix it if needed.

In `@model_gateway/src/routers/grpc/client.rs` around lines 237 - 246, The match
in GrpcClient::flush_cache omits the Self::Mlx variant causing a non-exhaustive
match; add an explicit branch for Self::Mlx in the flush_cache method (alongside
Self::Sglang, Self::Vllm, Self::Trtllm). If the Mlx client exposes a flush_cache
method, call it the same way as in Self::Sglang (await, unpack resp.success and
resp.message, return Ok); if it does not support flush_cache, return
Err(tonic::Status::unimplemented(...)) for Self::Mlx so the match is exhaustive.

⚠️ Potential issue | 🟠 Major

#!/bin/bash
set -euo pipefail
echo "HTTP fan-out path:"
sed -n '52,84p' model_gateway/src/worker/manager.rs
echo
echo "gRPC flush path:"
sed -n '733,759p' model_gateway/src/worker/manager.rs

#!/bin/bash
set -euo pipefail

# Get full flush_cache_all function
echo "=== Full flush_cache_all function ==="
sed -n '691,780p' model_gateway/src/worker/manager.rs

echo
echo "=== REQUEST_TIMEOUT and MAX_CONCURRENT definitions ==="
rg 'REQUEST_TIMEOUT|MAX_CONCURRENT' model_gateway/src/worker/manager.rs -B 2 -A 2

#!/bin/bash
set -euo pipefail

# Check if flush_cache_all has a caller-level timeout wrapper
echo "=== Calls to flush_cache_all ==="
rg 'flush_cache_all' model_gateway/src --type rust -B 3 -A 3

echo
echo "=== gRPC client builder for timeouts ==="
fd -t f 'grpc_client' crates/grpc_client/src | head -5 | xargs -I {} sh -c 'echo "=== {} ==="; head -100 {}'

#!/bin/bash
set -euo pipefail

echo "=== Searching for flush_cache method definition in gRPC clients ==="
fd -t f '\.rs$' crates/grpc_client/src | xargs grep -l 'flush_cache' | head -3

echo
echo "=== flush_cache method signatures ==="
rg 'fn flush_cache' crates/grpc_client/src -A 5 | head -40

echo
echo "=== gRPC timeout patterns in model_gateway ==="
rg 'timeout|Timeout|Duration' model_gateway/src/routers/grpc --type rust -c | sort | uniq -c | head -10

#!/bin/bash
set -euo pipefail

echo "=== Check for gRPC request timeout patterns in model_gateway ==="
rg 'tokio::time::timeout|timeout\(' model_gateway/src/routers/grpc --type rust -B 2 -A 2 | head -60

#!/bin/bash
set -euo pipefail

echo "=== All flush_cache calls in the codebase ==="
rg 'flush_cache\(' model_gateway/src --type rust -B 2 -A 2

echo
echo "=== Check flush_cache implementation for timeout_s usage ==="
sed -n '/pub async fn flush_cache/,/^    }/p' crates/grpc_client/src/sglang_scheduler.rs | head -40

#!/bin/bash
set -euo pipeftp
echo "=== Check if timeout behavior is documented or if 0.0 is intentional ==="
git log --all --oneline --grep="flush_cache" -- model_gateway/src/worker/manager.rs | head -5
git log --all -p -- crates/grpc_client/src/sglang_scheduler.rs | grep -A 5 -B 5 'flush_cache' | head -40

#!/bin/bash
set -euo pipefail

echo "=== Check if there's any wrapper timeout around flush_cache_all calls ==="
sed -n '1,/^async fn flush_cache/p' model_gateway/src/server.rs | tail -20

echo
echo "=== Check other RPC calls in grpc/client.rs for timeout patterns ==="
rg 'tokio::time::timeout' model_gateway/src/routers/grpc/client.rs -A 3 -B 1 | head -40

Apply REQUEST_TIMEOUT and concurrency bounding to the gRPC fan-out, matching the HTTP path.

The HTTP fan-out path enforces .timeout(REQUEST_TIMEOUT) (5 seconds) and .buffer_unordered(MAX_CONCURRENT) (32 workers), but the gRPC path calls flush_cache(0.0) (infinite timeout) and uses join_all() (unbounded concurrency). A single stalled gRPC worker can hang the entire cache-flush operation indefinitely; firing all gRPC flushes at once risks overwhelming the fleet on deployments with many workers.

.timeout(REQUEST_TIMEOUT)
.buffer_unordered(MAX_CONCURRENT)

grpc_client.flush_cache(0.0).await
future::join_all(grpc_futures).await

Wrap each flush_cache() call with tokio::time::timeout(REQUEST_TIMEOUT, ...) and replace join_all() with stream::iter(grpc_futures).buffer_unordered(MAX_CONCURRENT).collect() to match the HTTP implementation.

Verify each finding against the current code and only fix it if needed.

In `@model_gateway/src/worker/manager.rs` around lines 734 - 753, gRPC fan-out
currently awaits grpc_client.flush_cache(0.0) with join_all, causing unbounded
concurrency and infinite waits; modify the async closure built from grpc_workers
(the one that clones Arc and calls get_grpc_client()) to wrap the
flush_cache(...) call with tokio::time::timeout(REQUEST_TIMEOUT, ...), handle
Timeout errors as an Err result string, and then replace the
future::join_all(grpc_futures).await aggregation with
stream::iter(grpc_futures).buffer_unordered(MAX_CONCURRENT).collect().await so
concurrency is bounded and each flush respects REQUEST_TIMEOUT; keep handling of
Ok(Some(grpc_client)), Ok(None), and Err from get_grpc_client unchanged but
adapt error strings for timeouts and cancelled tasks.

⚠️ Potential issue | 🟠 Major

#!/bin/bash
# Description: Inspect flush_cache call sites and existing timeout/deadline patterns.

set -euo pipefail

echo "Flush cache call sites:"
rg -n -C3 --type rust '\bflush_cache\s*\('

echo
echo "Existing timeout/deadline patterns:"
rg -n -C3 --type rust 'tokio::time::timeout|Request::set_timeout|REQUEST_TIMEOUT|deadline_exceeded'

#!/bin/bash
# Check all public async methods in sglang_scheduler.rs that make RPC calls
ast-grep --pattern $'impl SglangSchedulerClient {
  $$$
  pub async fn $_($$$) -> $_ {
    $$$
  }
  $$$
}'

Enforce a client-side deadline for the flush RPC.

timeout_s is only forwarded to the scheduler; the gRPC await itself can block indefinitely if the scheduler or transport never responds. Since WorkerManager::flush_cache_all awaits grpc_client.flush_cache(0.0), this can stall the admin flush path. Validate the timeout and wrap the RPC with a local deadline.

     pub async fn flush_cache(
         &self,
         timeout_s: f32,
     ) -> Result<proto::FlushCacheResponse, tonic::Status> {
+        let scheduler_wait_timeout = Duration::try_from_secs_f32(timeout_s).map_err(|_| {
+            tonic::Status::invalid_argument("timeout_s must be finite and non-negative")
+        })?;
+        let rpc_timeout = scheduler_wait_timeout.saturating_add(Duration::from_secs(5));
+
         debug!("Requesting cache flush (timeout_s={timeout_s})");
         let request = Request::new(proto::FlushCacheRequest { timeout_s });
 
         let mut client = self.client.clone();
-        let response = client.flush_cache(request).await?;
+        let response = tokio::time::timeout(rpc_timeout, client.flush_cache(request))
+            .await
+            .map_err(|_| tonic::Status::deadline_exceeded("flush_cache RPC timed out"))??;
         debug!("Flush cache response received");
         Ok(response.into_inner())
     }

Verify each finding against the current code and only fix it if needed.

In `@crates/grpc_client/src/sglang_scheduler.rs` around lines 304 - 310, The
flush_cache method currently forwards timeout_s to the server but does not
enforce a local deadline, so the .await on client.flush_cache(request) can hang;
update the method in sglang_scheduler.rs (the function that builds
proto::FlushCacheRequest and calls client.flush_cache(request)) to validate
timeout_s (treat <= 0.0 as "no local timeout") and, when timeout_s > 0.0, call
the RPC via tokio::time::timeout(Duration::from_secs_f32(timeout_s),
client.flush_cache(request)).await, mapping a tokio timeout error to a
tonic::Status::deadline_exceeded; if timeout_s <= 0.0 just await the RPC
normally. Ensure you still clone self.client as before and propagate other RPC
errors unchanged.

⚠️ Potential issue | 🟠 Major

#!/bin/bash
# Locate where sglang scheduler sends GetLoadsReqOutput (tokenizer vs detokenizer IPC).
python -c "import sglang, os; print(os.path.dirname(sglang.__file__))" 2>/dev/null || true

# Search installed sglang (if present) and the repo for send sites.
rg -nP -C3 '\bGetLoadsReqOutput\s*\(' --glob '!**/request_manager.py'
rg -nP -C3 'send_to_tokenizer.*GetLoads|send_to_detokenizer.*GetLoads'

# First, locate the request_manager.py file
fd -n "request_manager.py" | head -20

# Check the structure and confirm file exists
git ls-files | grep -E "request_manager\.py|sglang.*request_manager"

# Read the specific sections mentioned in the review
if [ -f "grpc_servicer/smg_grpc_servicer/sglang/request_manager.py" ]; then
  echo "=== Lines 520-530 (handle_loop GetLoadsReqOutput) ===" 
  sed -n '520,530p' grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
  echo ""
  echo "=== Lines 560-570 (_handle_tokenizer_loop GetLoadsReqOutput) ===" 
  sed -n '560,570p' grpc_servicer/smg_grpc_servicer/sglang/request_manager.py
fi

# Search for all GetLoadsReqOutput occurrences and routing
rg -n "GetLoadsReqOutput" grpc_servicer/smg_grpc_servicer/sglang/request_manager.py -B2 -A2

# Check for any comments or documentation about tokenizer vs detokenizer routing
rg -n "tokenizer.*[Gg]et[Ll]oads|detokenizer.*[Gg]et[Ll]oads" grpc_servicer/smg_grpc_servicer/sglang/ -i

# Check where handle_loop and _handle_tokenizer_loop are started
rg -n "handle_loop|_handle_tokenizer_loop" grpc_servicer/smg_grpc_servicer/sglang/request_manager.py | grep -E "(asyncio\.create_task|gather|Task|start)"

# Check the full auto_create_handle_loop method to see both loops started
sed -n '990,1020p' grpc_servicer/smg_grpc_servicer/sglang/request_manager.py

# Search for any comments or history about GetLoadsReqOutput and which path it should use
rg -n "GetLoads" grpc_servicer/smg_grpc_servicer/sglang/request_manager.py -B3 -A3 | head -80

# Check what sockets are being used and their purposes
rg -n "recv_from_detokenizer|recv_from_tokenizer" grpc_servicer/smg_grpc_servicer/sglang/request_manager.py | head -20

Remove dead GetLoadsReqOutput handling from handle_loop — scheduler sends via tokenizer socket only.

Both handle_loop (line 524–525) and _handle_tokenizer_loop (line 562–563) handle GetLoadsReqOutput, but the comment at lines 195–196 explicitly states the scheduler sends communicator responses (including GetLoadsReqOutput) directly via send_to_tokenizer, which is the new path used by _handle_tokenizer_loop. The handle_loop branch is dead code and should be removed to avoid confusion and reduce maintenance burden. Keep only the _handle_tokenizer_loop branch.

Verify each finding against the current code and only fix it if needed.

In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 546 -
579, Remove the dead branch that handles GetLoadsReqOutput in the handle_loop
function so only _handle_tokenizer_loop processes scheduler-sent communicator
responses; locate the GetLoadsReqOutput conditional inside handle_loop and
delete that elif branch (and any now-unused local references) leaving
GetLoadsReqOutput handling only in _handle_tokenizer_loop, ensuring no other
logic depends on the removed branch.

🧹 Nitpick | 🔵 Trivial

Optional: validate communicator_name via an explicit registry.

getattr(self, communicator_name) accepts any attribute on self and silently resolves to non-communicator objects (e.g., "send_to_scheduler"), failing later with a cryptic TypeError when invoked. A small registry makes the allowed set explicit and yields a clear error on misuse.

     async def send_communicator_req(self, req, communicator_name: str, timeout: float = 30.0):
         ...
         self.auto_create_handle_loop()
-        communicator = getattr(self, communicator_name)
-        return await communicator(req, timeout=timeout)
+        registry = {
+            "profile_communicator": self.profile_communicator,
+            "flush_cache_communicator": self.flush_cache_communicator,
+            "get_internal_state_communicator": self.get_internal_state_communicator,
+            "get_loads_communicator": self.get_loads_communicator,
+        }
+        try:
+            communicator = registry[communicator_name]
+        except KeyError as e:
+            raise ValueError(f"Unknown communicator: {communicator_name!r}") from e
+        return await communicator(req, timeout=timeout)

Verify each finding against the current code and only fix it if needed.

In `@grpc_servicer/smg_grpc_servicer/sglang/request_manager.py` around lines 965 -
982, The send_communicator_req method should validate communicator_name before
using getattr to avoid cryptic TypeError: add an explicit registry (e.g., a set
or dict like communicator_registry or _allowed_communicators on the same class)
listing valid communicator attribute names, then inside send_communicator_req
call auto_create_handle_loop(), check that communicator_name is in that
registry, retrieve the attribute with getattr(self, communicator_name), verify
it is callable, and if not present or not callable raise a clear ValueError
mentioning the invalid communicator_name and listing allowed names; update any
callers if needed and keep the rest of the method (await communicator(req,
timeout=timeout)) unchanged.

Catch asyncio timeout from FlushCache communicator call

send_communicator_req() waits on asyncio.wait_for, which raises asyncio.TimeoutError in this Python 3.10 environment, but this handler only catches built-in TimeoutError. When a scheduler flush response times out, it falls into the generic Exception branch and is reported as INTERNAL instead of DEADLINE_EXCEEDED, which misclassifies real timeout conditions and breaks client retry/diagnostic behavior during scheduler stalls.

Useful? React with 👍 / 👎.

Bound each gRPC flush fan-out call with a timeout

The gRPC branch fans out grpc_client.flush_cache(0.0) and awaits join_all without any per-worker timeout, unlike the HTTP branch that enforces REQUEST_TIMEOUT. If a single gRPC worker becomes slow or unreachable, /flush_cache can block indefinitely on that future and stall the whole admin endpoint during partial outages; wrap each gRPC flush call with tokio::time::timeout and record timeout as a failed worker.

Useful? React with 👍 / 👎.