Support priority request header override (#30811)

Co-authored-by: Abhishek Parmar <abhishek@orng.net>
This commit is contained in:
Lianmin Zheng
2026-07-10 14:36:43 -07:00
committed by GitHub
co-authored by Abhishek Parmar
parent bc82b06400
commit 4fcc994be1
2 changed files with 29 additions and 0 deletions
@@ -15,6 +15,7 @@ _HEADER_OVERRIDES = {
"x-override-conversation-id": ("conversation_id", str), "x-override-conversation-id": ("conversation_id", str),
"x-override-routed-dp-rank": ("routed_dp_rank", int), "x-override-routed-dp-rank": ("routed_dp_rank", int),
"x-override-disagg-prefill-dp-rank": ("disagg_prefill_dp_rank", int), "x-override-disagg-prefill-dp-rank": ("disagg_prefill_dp_rank", int),
"x-override-priority": ("priority", int),
} }
@@ -19,6 +19,7 @@ def _obj():
conversation_id=None, conversation_id=None,
routed_dp_rank=None, routed_dp_rank=None,
disagg_prefill_dp_rank=None, disagg_prefill_dp_rank=None,
priority=None,
) )
@@ -36,6 +37,7 @@ class TestApplyRoutingHeaders(unittest.TestCase):
"x-override-conversation-id": "c1", "x-override-conversation-id": "c1",
"x-override-routed-dp-rank": "3", "x-override-routed-dp-rank": "3",
"x-override-disagg-prefill-dp-rank": "5", "x-override-disagg-prefill-dp-rank": "5",
"x-override-priority": "7",
} }
), ),
) )
@@ -46,6 +48,7 @@ class TestApplyRoutingHeaders(unittest.TestCase):
self.assertEqual(obj.conversation_id, "c1") self.assertEqual(obj.conversation_id, "c1")
self.assertEqual(obj.routed_dp_rank, 3) self.assertEqual(obj.routed_dp_rank, 3)
self.assertEqual(obj.disagg_prefill_dp_rank, 5) self.assertEqual(obj.disagg_prefill_dp_rank, 5)
self.assertEqual(obj.priority, 7)
def test_absent_headers_leave_obj_unchanged(self): def test_absent_headers_leave_obj_unchanged(self):
obj = _obj() obj = _obj()
@@ -76,6 +79,31 @@ class TestApplyRoutingHeaders(unittest.TestCase):
obj, Headers({"x-override-bootstrap-port": "not-an-int"}) obj, Headers({"x-override-bootstrap-port": "not-an-int"})
) )
def test_priority_header_overrides_body_value(self):
# The scheduler reads obj.priority, so the header value must be the one
# that ends up on the object even when the body already set a priority.
obj = _obj()
obj.priority = 1
apply_header_overrides(obj, Headers({"x-override-priority": "5"}))
self.assertEqual(obj.priority, 5)
def test_negative_priority_header_is_applied(self):
obj = _obj()
obj.priority = 1
apply_header_overrides(obj, Headers({"x-override-priority": "-3"}))
self.assertEqual(obj.priority, -3)
def test_priority_body_value_preserved_when_header_absent(self):
obj = _obj()
obj.priority = 2
apply_header_overrides(obj, Headers({}))
self.assertEqual(obj.priority, 2)
def test_invalid_priority_fails_loud(self):
obj = _obj()
with self.assertRaises(HTTPException):
apply_header_overrides(obj, Headers({"x-override-priority": "high"}))
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()