Support priority request header override (#30811)
Co-authored-by: Abhishek Parmar <abhishek@orng.net>
This commit is contained in:
co-authored by
Abhishek Parmar
parent
bc82b06400
commit
4fcc994be1
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user