Refactor serving code - #113
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughThe PR reorganizes serving around canonical configuration and core types, model plugins, dedicated KV-cache components, replica engines, worker IPC services, shared HTTP response helpers, and architecture-boundary validation. Compatibility re-exports preserve legacy import paths while new tests cover routing, caching, lifecycle, registry, and worker behavior. ChangesServing architecture refactor
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (8)
tests/test_architecture.py (1)
28-36: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winIgnore comments and docstrings when detecting model branches.
This flags any incidental
qwen/deepseekmention, including documentation, as an architecture violation. Inspect executable AST branch expressions instead, so the test enforces its stated contract without blocking harmless documentation.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_architecture.py` around lines 28 - 36, Update the architecture-violation scan in the generic_roots loop to parse each Python file’s AST and inspect only executable branch expressions for “qwen” and “deepseek”, excluding comments and docstrings. Preserve the existing violation reporting format and relative-path behavior while avoiding raw source-text matching.pypto_serving/serving/composition.py (1)
60-67: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMigrate the engine lifecycle to a lifespan handler.
@app.on_event("startup")/@app.on_event("shutdown")is deprecated in FastAPI; use anasynccontextmanagerwithFastAPI(lifespan=lifespan)instead for the engine start/stop and merge profile shutdown flow.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/composition.py` around lines 60 - 67, Replace the startup and shutdown event handlers with an asynccontextmanager lifespan handler that starts the engine before yielding, then stops the engine and invokes shutdown_callback afterward. Pass this handler via FastAPI(lifespan=lifespan), preserving the existing engine lifecycle and profile shutdown behavior.pypto_serving/serving/worker/service.py (1)
363-375: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsolidate the duplicated
supports_device_embeddingbranches.
prev_token_tensoris assigned once in theelsebranch (host-embedding, line 371) and again via a separate re-check of the same condition (lines 374-375) for the device-embedding path. Both branches are mutually exclusive today, but splitting the assignment across two independentifchecks on the same condition is confusing and easy to break under future edits.♻️ Proposed consolidation
decode_token_tensor = torch.tensor(decode_tokens, dtype=torch.long, device=device) if self.executor.supports_device_embedding: # Device kernel embeds directly from token ids — do not build # host-side embedding tensors. decode_embeddings = None prev_embeddings = None + prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device) else: decode_embeddings = self.executor.lookup_embeddings(runtime_model, decode_token_tensor) prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device) prev_embeddings = self.executor.lookup_embeddings(runtime_model, prev_token_tensor) - - if self.executor.supports_device_embedding: - prev_token_tensor = torch.tensor(prev_tokens, dtype=torch.long, device=device)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/worker/service.py` around lines 363 - 375, Consolidate the duplicate supports_device_embedding checks in the decode-token setup so prev_token_tensor is assigned within one mutually exclusive branch structure. Preserve the existing behavior: create prev_token_tensor for device embedding without host embeddings, and create it before prev_embeddings in the host-embedding path.pypto_serving/serving/server/responses.py (2)
95-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
__all__is not sorted per Ruff (RUF022).♻️ Proposed fix
__all__ = [ - "FinalGeneration", - "SSE_DONE", - "collect_final_generation", - "cumulative_text_delta", - "map_finish_reason", - "sse_model", - "to_generate_config", - "usage_from_output", + "collect_final_generation", + "cumulative_text_delta", + "FinalGeneration", + "map_finish_reason", + "SSE_DONE", + "sse_model", + "to_generate_config", + "usage_from_output", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/server/responses.py` around lines 95 - 104, Sort the exported names in the module-level __all__ list alphabetically to satisfy Ruff rule RUF022, preserving all existing exports.Source: Linters/SAST tools
88-89: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winUse
model_dump_json()for SSE payload serialization.
sse_modelis on the streaming hot path, andmodel.model_dump_json()uses Pydantic’s direct JSON serializer, avoiding the unnecessarymodel_dump()→json.dumps()round trip. The output whitespace may differ from the currentjson.dumps()call, but it is still compact and matches Pydantic v2’s idiomatic API.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/server/responses.py` around lines 88 - 89, Update sse_model to serialize the BaseModel directly with model_dump_json() instead of model_dump() followed by json.dumps(), while preserving the existing SSE “data:” prefix and double-newline terminator.pypto_serving/model/registry.py (1)
46-63: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUnreachable ambiguity branch in
detect.The loop returns on the first matching non-fallback plugin, so control only reaches Line 57 when no plugin matched (and no fallback is registered). At that point
matchesis always empty, making thelen(matches) == 1and "matches multiple registered families" branches dead code — the intended multi-match ambiguity error can never fire. Consider simplifying to make the actual behavior explicit:♻️ Proposed simplification (behavior-preserving)
def detect(self, config_data: dict[str, object]) -> ModelPlugin: """Detect a model family, using the registered fallback last.""" fallback = None for plugin in self._plugins: if plugin.is_fallback: fallback = plugin continue if plugin.matches(config_data): return plugin if fallback is not None: return fallback - matches = [plugin for plugin in self._plugins if plugin.matches(config_data)] - if len(matches) == 1: - return matches[0] - if not matches: - raise ValueError("Could not detect a registered model family") - families = ", ".join(plugin.family for plugin in matches) - raise ValueError(f"Model metadata matches multiple registered families: {families}") + raise ValueError("Could not detect a registered model family")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/model/registry.py` around lines 46 - 63, Update ModelPluginRegistry.detect to remove the unreachable matches-count and ambiguity handling after the plugin loop. Preserve first-match behavior for non-fallback plugins, return the registered fallback when present, and otherwise raise the existing “Could not detect a registered model family” error.pypto_serving/serving/memory/kv_cache.py (2)
34-42: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
__all__not sorted (RUF022).Ruff flags this list as not case-insensitively sorted —
NONE_HASHshould come afterhash_block_tokens. Since coding guidelines require following theruff.tomlconfiguration, this will fail lint ifRUF022is enabled in CI.🔧 Proposed fix
__all__ = [ "FreeKVCacheBlockQueue", "KVCacheBlock", "KVCacheBlocks", "KVCacheCapacityError", "KvCacheManager", - "NONE_HASH", "hash_block_tokens", + "NONE_HASH", ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/memory/kv_cache.py` around lines 34 - 42, Sort the exported names in __all__ case-insensitively to satisfy RUF022, placing hash_block_tokens before NONE_HASH while preserving all existing exports.Source: Linters/SAST tools
299-307: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winManager reaches into private members of
GroupedKVCacheAllocator.
_candidate_group_partitions,_select_group_partition, and_group_poolcallself._group_allocator._candidate_partitions,._select_partition, and._pool— all underscore-prefixed (private) methods of another class. This crosses the encapsulation boundary the refactor is otherwise establishing betweenKvCacheManagerand the allocator.Consider promoting these to public methods on
GroupedKVCacheAllocator(or dropping the leading underscore) soKvCacheManagerdepends only on its public surface.Also applies to: 341-346
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/memory/kv_cache.py` around lines 299 - 307, Promote the allocator operations used by KvCacheManager to the public GroupedKVCacheAllocator API: expose candidate partition selection, partition selection, and pool access without underscore-prefixed names. Update _candidate_group_partitions, _select_group_partition, and _group_pool to call the new public symbols, and update all corresponding definitions and call sites consistently.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pypto_serving/serving/engine/async_engine.py`:
- Around line 95-97: Update AsyncEngine.stop to call asyncio.gather with
return_exceptions=True, ensuring every reversed self._cores replica completes
its stop coroutine even when one fails. Preserve propagation or handling of
shutdown errors according to the existing stop contract after all cores have
finished.
In `@pypto_serving/serving/engine/output.py`:
- Around line 110-141: Remove completed request contexts from
ReplicaEngineCore._request_contexts: in pypto_serving/serving/engine/output.py
lines 110-141, delete the entry after request_output.finished handling; in
pypto_serving/serving/engine/replica.py lines 509-520, change the request lookup
to pop(request_id, None) so worker-error completion also cleans up safely.
In `@pypto_serving/serving/server/responses.py`:
- Around line 42-49: Handle the backend failure reason emitted by
replica.py::_handle_step_error before it is normalized by map_finish_reason:
update collect_final_generation and the streaming handlers in server.py to
detect output.finish_reason == "error" and propagate a distinct client-visible
failure, using an HTTP 5xx for non-streaming responses and an SSE error event or
equivalent logged streaming error. Keep map_finish_reason’s existing default
behavior for genuinely unmapped reasons.
In `@pypto_serving/serving/worker/service.py`:
- Around line 151-157: Update the exception handler in busy_loop so it logs the
caught queue-read exception with the existing logger before breaking out of the
loop, matching the diagnostic behavior of the other exception handlers in the
file.
In `@tests/test_kv_cache.py`:
- Around line 55-67: Update test_grouped_cache_growth_is_atomic_across_groups to
use asymmetric capacities so ensure_group_blocks("second", ...) allocates from
an earlier group before failing in a later group. Assert the later failure
leaves group_request_partition("second") unset and the earlier group’s
allocation returned, then preserve the existing release and successful retry
checks with expectations adjusted to the asymmetric pool sizes.
---
Nitpick comments:
In `@pypto_serving/model/registry.py`:
- Around line 46-63: Update ModelPluginRegistry.detect to remove the unreachable
matches-count and ambiguity handling after the plugin loop. Preserve first-match
behavior for non-fallback plugins, return the registered fallback when present,
and otherwise raise the existing “Could not detect a registered model family”
error.
In `@pypto_serving/serving/composition.py`:
- Around line 60-67: Replace the startup and shutdown event handlers with an
asynccontextmanager lifespan handler that starts the engine before yielding,
then stops the engine and invokes shutdown_callback afterward. Pass this handler
via FastAPI(lifespan=lifespan), preserving the existing engine lifecycle and
profile shutdown behavior.
In `@pypto_serving/serving/memory/kv_cache.py`:
- Around line 34-42: Sort the exported names in __all__ case-insensitively to
satisfy RUF022, placing hash_block_tokens before NONE_HASH while preserving all
existing exports.
- Around line 299-307: Promote the allocator operations used by KvCacheManager
to the public GroupedKVCacheAllocator API: expose candidate partition selection,
partition selection, and pool access without underscore-prefixed names. Update
_candidate_group_partitions, _select_group_partition, and _group_pool to call
the new public symbols, and update all corresponding definitions and call sites
consistently.
In `@pypto_serving/serving/server/responses.py`:
- Around line 95-104: Sort the exported names in the module-level __all__ list
alphabetically to satisfy Ruff rule RUF022, preserving all existing exports.
- Around line 88-89: Update sse_model to serialize the BaseModel directly with
model_dump_json() instead of model_dump() followed by json.dumps(), while
preserving the existing SSE “data:” prefix and double-newline terminator.
In `@pypto_serving/serving/worker/service.py`:
- Around line 363-375: Consolidate the duplicate supports_device_embedding
checks in the decode-token setup so prev_token_tensor is assigned within one
mutually exclusive branch structure. Preserve the existing behavior: create
prev_token_tensor for device embedding without host embeddings, and create it
before prev_embeddings in the host-embedding path.
In `@tests/test_architecture.py`:
- Around line 28-36: Update the architecture-violation scan in the generic_roots
loop to parse each Python file’s AST and inspect only executable branch
expressions for “qwen” and “deepseek”, excluding comments and docstrings.
Preserve the existing violation reporting format and relative-path behavior
while avoiding raw source-text matching.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 87c160ba-4264-4a31-b8b3-bf0b9bc1ff68
📒 Files selected for processing (73)
.pre-commit-config.yamldocs/architecture/overview.mddocs/architecture/refactoring.mdexamples/model/qwen3_14b/npu_generate.pypypto_serving/__init__.pypypto_serving/cli/main.pypypto_serving/config/__init__.pypypto_serving/config/generation.pypypto_serving/config/runtime.pypypto_serving/config/serving.pypypto_serving/config/types.pypypto_serving/core/__init__.pypypto_serving/core/batch.pypypto_serving/core/model.pypypto_serving/core/request.pypypto_serving/core/tokenizer.pypypto_serving/model/common/executor/executor.pypypto_serving/model/common/executor/pypto_executor.pypypto_serving/model/common/executor/sampler.pypypto_serving/model/common/runner/model_runner.pypypto_serving/model/deepseek/cache.pypypto_serving/model/deepseek/contracts.pypypto_serving/model/deepseek/input_builder.pypypto_serving/model/deepseek/npu_executor.pypypto_serving/model/deepseek/npu_runner.pypypto_serving/model/deepseek/plugin.pypypto_serving/model/deepseek/runtime_types.pypypto_serving/model/model_loader.pypypto_serving/model/plugin.pypypto_serving/model/qwen/cache_policy.pypypto_serving/model/qwen/input_builder.pypypto_serving/model/qwen/l3_runtime.pypypto_serving/model/qwen/npu_executor.pypypto_serving/model/qwen/npu_runner.pypypto_serving/model/qwen/plugin.pypypto_serving/model/qwen/runtime_types.pypypto_serving/model/registry.pypypto_serving/model/tokenizer.pypypto_serving/serving/composition.pypypto_serving/serving/engine/async_engine.pypypto_serving/serving/engine/engine.pypypto_serving/serving/engine/output.pypypto_serving/serving/engine/replica.pypypto_serving/serving/engine/request_state.pypypto_serving/serving/engine/routing.pypypto_serving/serving/memory/block_pool.pypypto_serving/serving/memory/grouped_cache.pypypto_serving/serving/memory/host_cache.pypypto_serving/serving/memory/kv_cache.pypypto_serving/serving/sched/cache.pypypto_serving/serving/sched/cache_coordinator.pypypto_serving/serving/sched/scheduler.pypypto_serving/serving/sched/types.pypypto_serving/serving/server/ipc.pypypto_serving/serving/server/ports.pypypto_serving/serving/server/responses.pypypto_serving/serving/server/schemas.pypypto_serving/serving/server/server.pypypto_serving/serving/server/serving_worker.pypypto_serving/serving/worker/__init__.pypypto_serving/serving/worker/client.pypypto_serving/serving/worker/protocol.pypypto_serving/serving/worker/service.pytests/lint/check_architecture.pytests/test_architecture.pytests/test_batching.pytests/test_composition.pytests/test_kv_cache.pytests/test_model_registry.pytests/test_package.pytests/test_scheduler_policy.pytests/test_serving_worker.pytests/test_worker_client.py
| async def stop(self) -> None: | ||
| """Stop all DP engine cores.""" | ||
| """Stop all data-parallel replica cores.""" | ||
| await asyncio.gather(*(core.stop() for core in reversed(self._cores))) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
stop() can abandon other replicas' shutdown if one core fails.
Without return_exceptions=True, if one core's stop() raises, the exception propagates immediately while the other cores' stop() coroutines (already scheduled by gather) are left running unsupervised. If the caller (e.g. the FastAPI shutdown handler) doesn't keep the event loop alive long enough, other replicas' worker subprocesses may never be cleanly closed.
🛡️ Proposed fix
async def stop(self) -> None:
"""Stop all data-parallel replica cores."""
- await asyncio.gather(*(core.stop() for core in reversed(self._cores)))
+ results = await asyncio.gather(
+ *(core.stop() for core in reversed(self._cores)),
+ return_exceptions=True,
+ )
+ errors = [result for result in results if isinstance(result, BaseException)]
+ if errors:
+ raise errors[0]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async def stop(self) -> None: | |
| """Stop all DP engine cores.""" | |
| """Stop all data-parallel replica cores.""" | |
| await asyncio.gather(*(core.stop() for core in reversed(self._cores))) | |
| async def stop(self) -> None: | |
| """Stop all data-parallel replica cores.""" | |
| results = await asyncio.gather( | |
| *(core.stop() for core in reversed(self._cores)), | |
| return_exceptions=True, | |
| ) | |
| errors = [result for result in results if isinstance(result, BaseException)] | |
| if errors: | |
| raise errors[0] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/serving/engine/async_engine.py` around lines 95 - 97, Update
AsyncEngine.stop to call asyncio.gather with return_exceptions=True, ensuring
every reversed self._cores replica completes its stop coroutine even when one
fails. Preserve propagation or handling of shutdown errors according to the
existing stop contract after all cores have finished.
| def process( | ||
| self, | ||
| scheduler_output: SchedulerOutput, | ||
| new_tokens: dict[str, int | list[int]], | ||
| ) -> None: | ||
| request_outputs = self._scheduler.update_from_output(scheduler_output, new_tokens) | ||
|
|
||
| for request_output in request_outputs: | ||
| ctx = self._request_contexts.get(request_output.request_id) | ||
| if ctx is None: | ||
| continue | ||
|
|
||
| text = detokenize_incrementally(self._tokenizer, ctx) | ||
| self._detect_stop_string(request_output, ctx, text) | ||
|
|
||
| if request_output.finished: | ||
| text = finalize_detokenization(self._tokenizer, ctx) | ||
| self._schedule_worker_free(request_output.request_id) | ||
|
|
||
| if not ctx.stream and not request_output.finished: | ||
| continue | ||
|
|
||
| ctx.queue.put_nowait( | ||
| TokenOutput( | ||
| token_id=request_output.new_token_id, | ||
| text=text, | ||
| finished=request_output.finished, | ||
| finish_reason=request_output.finish_reason, | ||
| prompt_tokens=ctx.request.num_prompt_tokens, | ||
| completion_tokens=len(ctx.request.output_token_ids), | ||
| ) | ||
| ) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win
ReplicaEngineCore._request_contexts leaks on every normal/error completion. Both sites finish a request without removing its RequestContext from the shared _request_contexts dict; only abort_request's explicit .pop() and add_request's cancellation-only cleanup actually remove entries, so the two most common completion paths (success, worker error) leak for the life of the process.
pypto_serving/serving/engine/output.py#L110-L141: afterrequest_output.finishedis handled, adddel self._request_contexts[request_output.request_id].pypto_serving/serving/engine/replica.py#L509-L520: changeself._request_contexts.get(request_id)toself._request_contexts.pop(request_id, None).
📍 Affects 2 files
pypto_serving/serving/engine/output.py#L110-L141(this comment)pypto_serving/serving/engine/replica.py#L509-L520
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/serving/engine/output.py` around lines 110 - 141, Remove
completed request contexts from ReplicaEngineCore._request_contexts: in
pypto_serving/serving/engine/output.py lines 110-141, delete the entry after
request_output.finished handling; in pypto_serving/serving/engine/replica.py
lines 509-520, change the request lookup to pop(request_id, None) so
worker-error completion also cleans up safely.
| def map_finish_reason(reason: str) -> str: | ||
| mapping = { | ||
| "FINISHED_EOS": "stop", | ||
| "FINISHED_LENGTH": "length", | ||
| "FINISHED_STOP": "stop", | ||
| "FINISHED_ABORTED": "stop", | ||
| } | ||
| return mapping.get(reason, "stop") |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Unmapped "error" finish reason silently reported as "stop".
replica.py::_handle_step_error emits TokenOutput(finished=True, finish_reason="error") when the worker step fails, but "error" is absent from this mapping, so mapping.get(reason, "stop") returns "stop". Both collect_final_generation (non-streaming) and server.py's streaming paths use this mapping, so a backend failure mid-generation is silently reported to API clients as a normal, successful completion.
🐛 Proposed fix: surface backend errors distinctly
def map_finish_reason(reason: str) -> str:
mapping = {
"FINISHED_EOS": "stop",
"FINISHED_LENGTH": "length",
"FINISHED_STOP": "stop",
"FINISHED_ABORTED": "stop",
+ "error": "stop",
}
- return mapping.get(reason, "stop")
+ return mapping.get(reason, "stop")A generic default preserves current behavior for genuinely unmapped reasons, but callers (e.g. collect_final_generation, the streaming handlers in server.py) should additionally check output.finish_reason == "error" and propagate a distinct error signal to the client (e.g. an HTTP 5xx for the non-streaming path, or an SSE error event/log for streaming) instead of masking it as "stop".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/serving/server/responses.py` around lines 42 - 49, Handle the
backend failure reason emitted by replica.py::_handle_step_error before it is
normalized by map_finish_reason: update collect_final_generation and the
streaming handlers in server.py to detect output.finish_reason == "error" and
propagate a distinct client-visible failure, using an HTTP 5xx for non-streaming
responses and an SSE error event or equivalent logged streaming error. Keep
map_finish_reason’s existing default behavior for genuinely unmapped reasons.
| def busy_loop(self) -> None: | ||
| logger.info("Worker entering busy loop") | ||
| while True: | ||
| try: | ||
| raw: bytes = self.input_queue.get() | ||
| except Exception: | ||
| break |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
Log the exception before silently exiting the busy loop.
except Exception: break on the queue read discards the failure with no trace, unlike the other except blocks in this file. A worker that dies here leaves no diagnostic evidence, making a stalled replica hard to root-cause in production.
🪵 Proposed fix to log before breaking
try:
raw: bytes = self.input_queue.get()
- except Exception:
+ except Exception:
+ logger.error("Worker failed to read from input_queue", exc_info=True)
break📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def busy_loop(self) -> None: | |
| logger.info("Worker entering busy loop") | |
| while True: | |
| try: | |
| raw: bytes = self.input_queue.get() | |
| except Exception: | |
| break | |
| def busy_loop(self) -> None: | |
| logger.info("Worker entering busy loop") | |
| while True: | |
| try: | |
| raw: bytes = self.input_queue.get() | |
| except Exception: | |
| logger.error("Worker failed to read from input_queue", exc_info=True) | |
| break |
🧰 Tools
🪛 Ruff (0.15.21)
[warning] 156-156: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pypto_serving/serving/worker/service.py` around lines 151 - 157, Update the
exception handler in busy_loop so it logs the caught queue-read exception with
the existing logger before breaking out of the loop, matching the diagnostic
behavior of the other exception handlers in the file.
Source: Linters/SAST tools
| def test_grouped_cache_growth_is_atomic_across_groups() -> None: | ||
| manager = KvCacheManager(enable_prefix_cache=False) | ||
| manager.init_groups((_group("attention"), _group("state")), max_batch_size=1) | ||
| first = manager.ensure_group_blocks("first", token_count=4) | ||
|
|
||
| assert first == {"attention": [0, 1], "state": [0, 1]} | ||
| with pytest.raises(KVCacheCapacityError): | ||
| manager.ensure_group_blocks("second", token_count=2) | ||
| assert manager.group_request_partition("second") is None | ||
|
|
||
| manager.release_all_group_requests("first") | ||
| second = manager.ensure_group_blocks("second", token_count=2) | ||
| assert second == {"attention": [0], "state": [0]} |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Exercise a failure after one group has allocated.
first exhausts both identical two-block pools, so second fails in the first group and cannot verify rollback across groups. Give the groups asymmetric capacity and make a later group fail after an earlier group grows; then assert the earlier allocation was returned.
As per coding guidelines, maintain the repository’s unit tests.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/test_kv_cache.py` around lines 55 - 67, Update
test_grouped_cache_growth_is_atomic_across_groups to use asymmetric capacities
so ensure_group_blocks("second", ...) allocates from an earlier group before
failing in a later group. Assert the later failure leaves
group_request_partition("second") unset and the earlier group’s allocation
returned, then preserve the existing release and successful retry checks with
expectations adjusted to the asymmetric pool sizes.
Source: Coding guidelines
ecf8abf to
a9bb0f8
Compare
1 participant
Refactor serving code