[Serve] Add on_request_completed hook and supports_rejection_protocol to custom request router API#60845
[Serve] Add on_request_completed hook and supports_rejection_protocol to custom request router API#60845abrarsheikh wants to merge 2 commits intomasterfrom
on_request_completed hook and supports_rejection_protocol to custom request router API#60845Conversation
… to custom request router API Signed-off-by: abrar <abrar@anyscale.com>
There was a problem hiding this comment.
Code Review
This pull request introduces several valuable enhancements to the custom request router API. The addition of the on_request_completed lifecycle hook, along with the robust implementation that ensures it's called even in failure scenarios, is a great improvement for custom routers needing cleanup logic. The new supports_rejection_protocol property provides useful flexibility for routers that manage capacity client-side. Extracting the backoff logic into _compute_backoff_s and _backoff helpers is a good design choice that improves extensibility.
The changes are well-implemented and accompanied by thorough tests. My review includes a few minor suggestions to add return type hints for better code clarity and type safety. Overall, this is a solid contribution.
Signed-off-by: abrar <abrar@anyscale.com>
on_request_completed(replica_id, internal_request_id)lifecycle hook toRequestRouter, called when a request finishes (success or failure). This enables custom routers to perform cleanup such as releasing capacity tokens.ReplicaResult(not just when autoscaling metrics are enabled) so the router is notified of every completion. Add afinallyguard for the error path where the callback couldn't be registered.supports_rejection_protocolproperty so custom routers that guarantee capacity client-side can opt out of the rejection round-trip._compute_backoff_s/_backoffhelpers.