generalized dict and error message

Author: khatwanimohit

tests/generate/vllm_sampler_test.py

@@ -395,11 +395,11 @@ def test_expert_parallel_size_plumbed_to_sharding(self):
     self.assertEqual(sampler.args["tensor_parallel_size"], 4)
     self.assertEqual(sampler.args["data_parallel_size"], 1)
 
-  def test_sharding_keys_in_engine_kwargs_raise_value_error(self):
-    # Sharding parallelism sizes (tp, dp, ep) are tunix-owned VllmConfig fields
-    # and must be set directly on VllmConfig, not smuggled through engine_kwargs.
-    # Passing them via engine_kwargs should raise a ValueError at config
-    # construction time before any vLLM engine args are assembled.
+  def test_reserved_keys_in_engine_kwargs_raise_value_error(self):
+    # Reserved VllmConfig fields (e.g. tp, dp, ep) must be set directly on
+    # VllmConfig, not smuggled through engine_kwargs. Passing them via
+    # engine_kwargs should raise a ValueError at config construction time
+    # before any vLLM engine args are assembled.
     mesh = self._make_mock_mesh(8)
     for key in ("expert_parallel_size", "tensor_parallel_size", "data_parallel_size"):
       with self.subTest(key=key):

tunix/generate/vllm_sampler.py

@@ -71,10 +71,10 @@ class VllmConfig:
       init=False, default_factory=dict
   )
 
-  # Parallelism sizes are tunix-owned VllmConfig fields that require special
-  # processing before being passed to vLLM. They must not be passed via
-  # engine_kwargs, which is a direct pass-through to vLLM EngineArgs.
-  _SHARDING_KEYS: frozenset[str] = dataclasses.field(
+  # VllmConfig fields that require special processing before being passed to
+  # vLLM and must not be passed via engine_kwargs, which is a raw pass-through
+  # to vLLM EngineArgs.
+  _RESERVED_KEYS: frozenset[str] = dataclasses.field(
       default=frozenset(
           {"tensor_parallel_size", "data_parallel_size", "expert_parallel_size"}
       ),
@@ -85,11 +85,11 @@ class VllmConfig:
 
   def __post_init__(self, engine_kwargs: Optional[Dict[str, Any]]):
     engine_kwargs = engine_kwargs or {}
-    illegal = self._SHARDING_KEYS & engine_kwargs.keys()
+    illegal = self._RESERVED_KEYS & engine_kwargs.keys()
     if illegal:
       raise ValueError(
-          f"Sharding parallelism sizes must be set directly on VllmConfig, not"
-          f" passed via engine_kwargs: {sorted(illegal)}"
+          f"VllmConfig fields must be set directly on VllmConfig, not passed"
+          f" via engine_kwargs: {sorted(illegal)}"
       )
     self._processed_engine_kwargs = engine_kwargs
     if engine_kwargs: