diff --git a/python/ray/llm/_internal/serve/core/configs/accelerators.py b/python/ray/llm/_internal/serve/core/configs/accelerators.py index 3ec11d809dd4..db70160698e7 100644 --- a/python/ray/llm/_internal/serve/core/configs/accelerators.py +++ b/python/ray/llm/_internal/serve/core/configs/accelerators.py @@ -8,7 +8,11 @@ import ray.util.accelerators.accelerators as accelerators from ray.llm._internal.serve.observability.logging import get_logger from ray.util.placement_group import PlacementGroup, placement_group -from ray.util.tpu import get_tpu_version_from_type, slice_placement_group +from ray.util.tpu import ( + get_chips_per_host, + get_tpu_version_from_type, + slice_placement_group, +) logger = get_logger(__name__) @@ -91,14 +95,40 @@ def create_placement_group( ) -> PlacementGroup: pass - @property - def requires_deferred_placement_group(self) -> bool: - """ - If True, Ray Serve will not provision a placement group for the deployment. - Instead, creation is deferred to the replica at runtime. - Defaults to False. - """ - return False + def get_placement_group_bundle_label_selector( + self, accelerator_type_str: Optional[str] = None + ) -> Optional[Dict[str, str]]: + """Returns label selectors to apply to the placement group bundles.""" + return None + + def apply_placement_group_bundle_label_selector( + self, + deployment_options: Dict[str, Any], + accelerator_type_str: Optional[str], + num_bundles: int, + ) -> None: + """Safely applies hardware-specific label selectors to the deployment options.""" + accel_selectors = self.get_placement_group_bundle_label_selector( + accelerator_type_str + ) + if not accel_selectors: + return + + existing_selectors = ( + deployment_options.get("placement_group_bundle_label_selector") or [] + ) + merged_selectors = [] + + for i in range(num_bundles): + selector = ( + existing_selectors[i].copy() + if i < len(existing_selectors) and existing_selectors[i] + else {} + ) + selector.update(accel_selectors) + merged_selectors.append(selector) + + deployment_options["placement_group_bundle_label_selector"] = merged_selectors @property @abstractmethod @@ -180,10 +210,21 @@ def __init__(self, config: TPUConfig): def default_bundles( self, *, num_devices: int, accelerator_type_str: Optional[str] = None ): - bundle = {"TPU": 1} + if self._config.topology and accelerator_type_str: + version = get_tpu_version_from_type(accelerator_type_str) + chips_per_host = get_chips_per_host(self._config.topology, version) + + num_bundles = max(1, num_devices // chips_per_host) + bundle = {"TPU": chips_per_host} + else: + # Fallback to single-chip/single-host scheduling + num_bundles = num_devices + bundle = {"TPU": 1} + if accelerator_type_str: bundle[format_ray_accelerator_resource(accelerator_type_str)] = 0.001 - return [bundle.copy() for _ in range(num_devices)] + + return [bundle.copy() for _ in range(num_bundles)] def create_placement_group( self, @@ -211,7 +252,8 @@ def create_placement_group( if not tpu_bundles: worker_bundle = {"TPU": 1} else: - worker_bundle = tpu_bundles[0] + # Use the last bundle to avoid picking up merged CPU driver resources + worker_bundle = tpu_bundles[-1] # Ensure all TPU bundles are homogeneous if any(b != worker_bundle for b in tpu_bundles): @@ -239,38 +281,39 @@ def create_placement_group( ) return self._slice_pg_wrapper.placement_group - @property - def requires_deferred_placement_group(self) -> bool: - """ - If a TPU topology is specified, we defer PG creation so the replica can - provision a `SlicePlacementGroup` at runtime. This ensures multi-host - TPU slices are gang-scheduled atomically according to their physical - topology rather than fragmented across the cluster. - """ - return bool(self._config.topology) + def get_placement_group_bundle_label_selector( + self, accelerator_type_str: Optional[str] = None + ) -> Optional[Dict[str, str]]: + if self._config.topology and accelerator_type_str: + return { + "ray.io/tpu-topology": self._config.topology, + "ray.io/accelerator-type": accelerator_type_str, + } + return None @property def requires_remote_initialization(self) -> bool: return True def get_remote_options(self, accelerator_type_str: str = None): - # TPUs use custom resource strings rather than a native kwarg - options: Dict[str, Any] = {"resources": {"TPU": 0.001}} - + # The PlacementGroupSchedulingStrategy natively handles routing the task to + # the correct hardware. We omit TPU resource requests to avoid consuming + # chips that the model engine workers must use. + options: Dict[str, Any] = {"resources": {}} if accelerator_type_str: - options["accelerator_type"] = accelerator_type_str + # Pin the task to the TPU accelerator to avoid scheduling on a CPU bundle. + options["label_selector"] = { + "ray.io/accelerator-type": accelerator_type_str + } return options def shutdown(self): - if self._slice_pg_wrapper is not None: + slice_pg_wrapper = getattr(self, "_slice_pg_wrapper", None) + if slice_pg_wrapper is not None: try: logger.info("Shutting down TPU slice PG for server replica.") - self._slice_pg_wrapper.shutdown() + slice_pg_wrapper.shutdown() except Exception as e: logger.warning(f"Failed to shut down TPU slice PG: {e}") finally: self._slice_pg_wrapper = None - - def __del__(self): - """Ensure placement groups are cleaned up when this backend is garbage collected.""" - self.shutdown() diff --git a/python/ray/llm/_internal/serve/core/server/llm_server.py b/python/ray/llm/_internal/serve/core/server/llm_server.py index 79933bdd59a6..05d07f11961f 100644 --- a/python/ray/llm/_internal/serve/core/server/llm_server.py +++ b/python/ray/llm/_internal/serve/core/server/llm_server.py @@ -270,7 +270,14 @@ async def _maybe_add_request_id_to_request( """Add the request id to the request.""" request_id = get_serve_request_id() if request_id: - request.request_id = request_id + if hasattr(request, "request_id"): + return + try: + request.request_id = request_id + except ValueError: + # Pydantic v2 strict schemas reject unknown fields. + # Safely ignore as request_id is only used for internal Ray Serve logging. + pass async def _maybe_resolve_lora_from_multiplex(self) -> None: """Handle the lora model for the request.""" @@ -699,27 +706,40 @@ def get_deployment_options(cls, llm_config: "LLMConfig"): # deployment_options ray_actor_options = deployment_options.get("ray_actor_options", {}) - if not engine_config.accelerator.requires_deferred_placement_group: - replica_actor_resources = { - "CPU": ray_actor_options.get("num_cpus", 1), - "GPU": ray_actor_options.get("num_gpus", 0), - **ray_actor_options.get("resources", {}), - } - if "memory" in ray_actor_options: - replica_actor_resources["memory"] = ray_actor_options["memory"] + replica_actor_resources = { + "CPU": ray_actor_options.get("num_cpus", 1), + "GPU": ray_actor_options.get("num_gpus", 0), + **ray_actor_options.get("resources", {}), + } + if "memory" in ray_actor_options: + replica_actor_resources["memory"] = ray_actor_options["memory"] - # TODO: Move this _merge_replica_actor_and_child_actor_bundles to a - # more generic place. - pg_bundles = _merge_replica_actor_and_child_actor_bundles( - engine_config.placement_bundles, replica_actor_resources - ) + # TODO: Move this _merge_replica_actor_and_child_actor_bundles to a + # more generic place. + pg_bundles = _merge_replica_actor_and_child_actor_bundles( + engine_config.placement_bundles, replica_actor_resources + ) + + deployment_options.update( + { + "placement_group_bundles": pg_bundles, + "placement_group_strategy": engine_config.placement_strategy, + } + ) - deployment_options.update( - { - "placement_group_bundles": pg_bundles, - "placement_group_strategy": engine_config.placement_strategy, - } + # Append hardware-specific `bundle_label_selectors` to the deployment options if needed + accelerator_type_str = ( + getattr( + llm_config.accelerator_type, "value", str(llm_config.accelerator_type) ) + if llm_config.accelerator_type + else None + ) + engine_config.accelerator.apply_placement_group_bundle_label_selector( + deployment_options=deployment_options, + accelerator_type_str=accelerator_type_str, + num_bundles=len(pg_bundles), + ) # Handle env vars from runtime_env default_runtime_env = ray.get_runtime_context().runtime_env @@ -728,7 +748,6 @@ def get_deployment_options(cls, llm_config: "LLMConfig"): "worker_process_setup_hook" ] = "ray.llm._internal.serve._worker_process_setup_hook" - ray_actor_options = deployment_options.get("ray_actor_options", {}) ray_actor_options["runtime_env"] = { **default_runtime_env, # Existing runtime_env should take precedence over the default. diff --git a/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py b/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py index fa4097244e37..2e0ac95d02cb 100644 --- a/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py +++ b/python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py @@ -279,6 +279,10 @@ async def start(self) -> None: from vllm.entrypoints.openai.api_server import init_app_state callback = self.llm_config.get_or_create_callback() + if callback.ctx.placement_group: + # Ensure existing PG for the Serve deployment is scheduled + # before attempting to initialize the engine. + await callback.ctx.placement_group.ready() await callback.run_callback("on_before_node_init") if callback.ctx.run_init_node: await initialize_node(self.llm_config) diff --git a/python/ray/llm/tests/serve/cpu/configs/test_models.py b/python/ray/llm/tests/serve/cpu/configs/test_models.py index ff35dca1fe9e..4faaddf10101 100644 --- a/python/ray/llm/tests/serve/cpu/configs/test_models.py +++ b/python/ray/llm/tests/serve/cpu/configs/test_models.py @@ -6,8 +6,6 @@ from ray.llm._internal.common.utils.download_utils import NodeModelDownloadable from ray.llm._internal.serve.core.configs.accelerators import ( - CPUAccelerator, - GPUAccelerator, TPUAccelerator, TPUConfig, ) @@ -403,19 +401,34 @@ def test_engine_config_infers_tpu_from_accelerator_type_string(self): assert isinstance(engine_config.accelerator, TPUAccelerator) assert engine_config.accelerator_type == "TPU-V6E" - def test_requires_deferred_placement_group(self): - """Test that requires_deferred_placement_group correctly identifies deferred PG requirements.""" - cpu_accel = CPUAccelerator() - assert cpu_accel.requires_deferred_placement_group is False + def test_tpu_accelerator_get_placement_group_bundle_label_selector(self): + """Test that TPUAccelerator correctly generates topology labels for Serve.""" + tpu_accel_no_topology = TPUAccelerator(TPUConfig(kind="tpu")) + assert ( + tpu_accel_no_topology.get_placement_group_bundle_label_selector("TPU-V6E") + is None + ) - gpu_accel = GPUAccelerator() - assert gpu_accel.requires_deferred_placement_group is False + tpu_accel_with_topology = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4")) + assert tpu_accel_with_topology.get_placement_group_bundle_label_selector( + "TPU-V6E" + ) == { + "ray.io/tpu-topology": "4x4", + "ray.io/accelerator-type": "TPU-V6E", + } - tpu_accel_no_topo = TPUAccelerator(TPUConfig(kind="tpu")) - assert tpu_accel_no_topo.requires_deferred_placement_group is False + def test_tpu_accelerator_get_remote_options(self): + """Test that TPUAccelerator get_remote_options returns an empty resources dict and label selector.""" + tpu_accel = TPUAccelerator(TPUConfig(kind="tpu")) - tpu_accel_with_topo = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4")) - assert tpu_accel_with_topo.requires_deferred_placement_group is True + options_no_type = tpu_accel.get_remote_options() + assert options_no_type == {"resources": {}} + + options_with_type = tpu_accel.get_remote_options("TPU-V6E") + assert options_with_type == { + "resources": {}, + "label_selector": {"ray.io/accelerator-type": "TPU-V6E"}, + } if __name__ == "__main__": diff --git a/python/ray/llm/tests/serve/cpu/deployments/conftest.py b/python/ray/llm/tests/serve/cpu/deployments/conftest.py index c0ffb22d2bad..e124358729cf 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/conftest.py +++ b/python/ray/llm/tests/serve/cpu/deployments/conftest.py @@ -36,6 +36,7 @@ def ray_tpu_cluster(): "ray.io/tpu-slice-name": "test-slice", "ray.io/tpu-worker-id": str(i), "ray.io/tpu-pod-type": pod_type, + "ray.io/accelerator-type": "TPU-V6E", } resources = {"TPU": 4, "accelerator_type:TPU-V6E": 4} @@ -50,6 +51,6 @@ def ray_tpu_cluster(): env_vars=env_vars, ) - ray.init(address=cluster.address) + ray.init(address=cluster.address, ignore_reinit_error=True) yield cluster ray.shutdown() diff --git a/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py b/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py index 10fe2bfdf14d..43bbbba738ba 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py +++ b/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_engine_tpu.py @@ -14,134 +14,89 @@ ) from ray.llm._internal.serve.core.server.llm_server import LLMServer from ray.llm.tests.serve.mocks.mock_vllm_engine import PGCreationMockEngine +from ray.serve._private.utils import resolve_tpu_slice_kwargs from ray.serve.llm import LLMConfig, ModelLoadingConfig -from ray.util.placement_group import PlacementGroup, placement_group_table -def test_tpu_slice_placement_group_creation_default_resources(ray_tpu_cluster): +def test_tpu_serve_deployment_single_tpu_fallback(ray_tpu_cluster): """ - Verifies that requesting a multi-host TPU topology correctly intercepts - standard PG creation and returns a PACK SlicePlacementGroup. - """ - llm_config = LLMConfig( - model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), - accelerator_type="TPU-V6E", - accelerator_config={"kind": "tpu", "topology": "4x4"}, - ) - - engine_config = llm_config.get_engine_config() - pg = engine_config.get_or_create_pg() - - assert isinstance(pg, PlacementGroup) - - pg_table = placement_group_table(pg) - assert pg_table["strategy"] == "PACK" - - # 4x4 v6e = 16 chips. We default to 1 TPU chip per bundle. - assert len(pg_table["bundles"]) == 16 - for bundle in pg_table["bundles"].values(): - assert "TPU" in bundle - assert bundle["TPU"] == 1 - - # Let the backend tear down its own resources if it has any - engine_config.accelerator.shutdown() - try: - ray.util.remove_placement_group(pg) - except Exception: - pass # Already cleaned up by the wrapper - - -def test_tpu_slice_placement_group_creation_host_resources(ray_tpu_cluster): - """ - Verifies that explicitly providing host-level bundles via - placement_group_config correctly overrides the 1-chip default. + Verifies that requesting a TPU without a topology gracefully + falls back to standard single-host bundle packing without + triggering the slice placement group interception. """ llm_config = LLMConfig( model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), accelerator_type="TPU-V6E", - accelerator_config={"kind": "tpu", "topology": "4x4"}, - placement_group_config={ - "strategy": "STRICT_SPREAD", - "bundles": [{"TPU": 4}], - }, + # Explicitly omit topology + accelerator_config={"kind": "tpu"}, + engine_kwargs={"tensor_parallel_size": 1}, ) - engine_config = llm_config.get_engine_config() - pg = engine_config.get_or_create_pg() - - assert isinstance(pg, PlacementGroup) - - pg_table = placement_group_table(pg) - assert pg_table["strategy"] == "STRICT_SPREAD" - # We should provision 4 host-level bundles instead of the default 16 chip-level bundles. - assert len(pg_table["bundles"]) == 4 - for bundle in pg_table["bundles"].values(): - assert "TPU" in bundle - assert bundle["TPU"] == 4 + app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine) - # Let the backend tear down its own resources if it has any - engine_config.accelerator.shutdown() + serve.start(http_options={"port": 0}) try: - ray.util.remove_placement_group(pg) - except Exception: - pass # Already cleaned up by the wrapper - + serve.run(app, name="single_tpu_app", route_prefix="/single_tpu") + pg_table = ray.util.placement_group_table() + active_pgs = list( + {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values() + ) -def test_single_tpu_fallback(ray_tpu_cluster): + # Ensure we only have standard PGs, no TPU Head PGs + tpu_head_resource = "TPU-v6e-16-head" + head_pgs = [ + pg + for pg in active_pgs + if len(pg["bundles"]) == 1 + and tpu_head_resource in list(pg["bundles"].values())[0] + ] + assert len(head_pgs) == 0 + + # Verify the deployment PG has the default PACK strategy and 1 TPU + deployment_pgs = [ + pg + for pg in active_pgs + if any("TPU" in bundle for bundle in pg["bundles"].values()) + ] + assert len(deployment_pgs) >= 1 + + target_pg = deployment_pgs[0] + assert target_pg["strategy"] == "PACK" + + # Verify it allocated 1 TPU per bundle + for bundle in target_pg["bundles"].values(): + if "TPU" in bundle: + assert bundle["TPU"] == 1 + + finally: + serve.shutdown() + + +def test_tpu_accelerator_remote_options_scheduling(ray_tpu_cluster): """ - Verifies that requesting a TPU without a topology gracefully - falls back to standard single-host bundle packing. + Verifies that TPUAccelerator.get_remote_options returns a label_selector, + and successfully schedules a task without causing Ray Core validation errors. """ llm_config = LLMConfig( model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), accelerator_type="TPU-V6E", + accelerator_config={"kind": "tpu"}, ) - engine_config = llm_config.get_engine_config() - pg = engine_config.get_or_create_pg() - pg_table = placement_group_table(pg) + options = engine_config.accelerator.get_remote_options("TPU-V6E") - # Verify it falls back to the default PACK strategy for 1 GPU/TPU - assert len(pg_table["bundles"]) == 1 - assert pg_table["strategy"] == "PACK" + # Ensure it returns the expected options (i.e. label_selector only) + assert options == { + "resources": {}, + "label_selector": {"ray.io/accelerator-type": "TPU-V6E"}, + } - # Let the backend tear down its own resources if it has any - engine_config.accelerator.shutdown() - try: - ray.util.remove_placement_group(pg) - except Exception: - pass # Already cleaned up by the wrapper + @ray.remote(**options) + def probe_metadata(): + return True - -def test_tpu_slice_placement_group_creation_bundle_per_worker(ray_tpu_cluster): - """ - Verifies that specifying bundle_per_worker correctly expands to bundles, - includes the accelerator hint for TPU, and correctly identifies TPU usage. - """ - llm_config = LLMConfig( - model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), - accelerator_type="TPU-V6E", - accelerator_config={"kind": "tpu", "topology": "4x4"}, - placement_group_config={ - "bundle_per_worker": {"TPU": 1}, - }, - engine_kwargs={ - "tensor_parallel_size": 2, - }, - ) - - engine_config = llm_config.get_engine_config() - - # Validate the accelerator backend was correctly inferred - assert isinstance(engine_config.accelerator, TPUAccelerator) - - bundles = engine_config.placement_bundles - assert len(bundles) == 2 - for bundle in bundles: - assert bundle["TPU"] == 1 - assert "accelerator_type:TPU-V6E" in bundle - assert bundle["accelerator_type:TPU-V6E"] == 0.001 + assert ray.get(probe_metadata.remote()) is True def test_accelerator_inference_logic(): @@ -150,7 +105,6 @@ def test_accelerator_inference_logic(): when no explicit accelerator_config is provided, and passes it correctly to the engine. """ - # TPU string correctly infers TPUConfig and TPUAccelerator cfg1 = LLMConfig( model_loading_config={"model_id": "test"}, accelerator_type="TPU-V6E", @@ -180,88 +134,120 @@ def test_accelerator_inference_logic(): assert isinstance(cfg4.get_engine_config().accelerator, CPUAccelerator) -def test_tpu_slice_placement_group_creation_heterogeneous_tpu_bundles_fail(): +def test_tpu_deployment_options_bundle_selector_injection(): """ - Verifies that a ValueError is raised when heterogeneous TPU bundles are provided. + Verifies that LLMServer.get_deployment_options correctly injects TPU topology + labels into the placement group bundle selectors and creates the expected bundles. """ - accelerator = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4")) + llm_config = LLMConfig( + model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), + accelerator_type="TPU-V6E", + accelerator_config={"kind": "tpu", "topology": "4x4"}, + engine_kwargs={"tensor_parallel_size": 16}, + ) - with pytest.raises(ValueError, match="Heterogeneous TPU bundles are not supported"): - accelerator.create_placement_group( - bundles=[{"TPU": 4}, {"TPU": 2}], - strategy="PACK", - name="test-pg", - accelerator_type_str="TPU-V6E", - ) + options = LLMServer.get_deployment_options(llm_config) + + # Ensure PG creation is no longer deferred + assert "placement_group_bundles" in options + assert "placement_group_bundle_label_selector" in options + + pg_bundles = options["placement_group_bundles"] + selectors = options["placement_group_bundle_label_selector"] + + # 4x4 topology = 16 chips + # Default is 4 bundles of 4 TPUs + assert len(pg_bundles) == 4 + assert len(selectors) == 4 + + # The first bundle should contain the TPU allocation and the replica actor's CPU + assert pg_bundles[0].get("TPU") == 4 + assert pg_bundles[0].get("CPU", 0) >= 1 + + # The remaining worker bundles should strictly be for the TPU host + for i in range(1, 4): + assert pg_bundles[i].get("TPU") == 4 + assert pg_bundles[i].get("CPU", 0) == 0 + + for selector in selectors: + assert selector["ray.io/tpu-topology"] == "4x4" + assert selector["ray.io/accelerator-type"] == "TPU-V6E" -def test_tpu_slice_placement_group_creation_cpu_driver_homogeneous_tpu_bundles_pass( - ray_tpu_cluster, -): +def test_tpu_slice_kwargs_ignores_cpu_driver_bundle(): """ - Verifies that CPU-only driver bundles are ignored and do not trigger an error - if subsequent TPU bundles are homogeneous. + Verifies that resolve_tpu_slice_kwargs correctly ignores the merged CPU + resources from the replica actor in the first bundle, and extracts the + TPU requirement from the remaining bundles. """ - accelerator = TPUAccelerator(TPUConfig(kind="tpu", topology="4x4")) + labels = [{"ray.io/tpu-topology": "4x4", "ray.io/accelerator-type": "TPU-V6E"}] + bundles = [{"TPU": 4, "CPU": 1}, {"TPU": 4}, {"TPU": 4}, {"TPU": 4}] - pg = accelerator.create_placement_group( - bundles=[{"CPU": 2}, {"TPU": 4}, {"TPU": 4}], - strategy="PACK", - name="test-pg", - accelerator_type_str="TPU-V6E", - ) + topology, version, worker_bundle = resolve_tpu_slice_kwargs(labels, bundles) - # Verify valid PG creation - assert isinstance(pg, PlacementGroup) + assert topology == "4x4" + assert version == "v6e" + assert worker_bundle == {"TPU": 4} + assert "CPU" not in worker_bundle - accelerator.shutdown() - try: - ray.util.remove_placement_group(pg) - except Exception: - pass +def test_tpu_slice_kwargs_rejects_heterogeneous_bundles(): + """ + Verifies that a ValueError is raised when heterogeneous TPU bundles are provided + to the Serve gang-scheduler helper. + """ + labels = [{"ray.io/tpu-topology": "4x4", "ray.io/accelerator-type": "TPU-V6E"}] + bundles = [{"TPU": 4}, {"TPU": 2}] -def test_tpu_serve_deployment_default_chip_level_bundles(ray_tpu_cluster): + with pytest.raises(ValueError, match="Heterogeneous TPU bundles are not supported"): + resolve_tpu_slice_kwargs(labels, bundles) + + +def test_tpu_serve_deployment_default_host_level_bundles(ray_tpu_cluster): """ - Verifies that a Serve deployment created for a multi-host TPU slice defaults - to chip-level bundles when no placement_group_config is specified. + Verifies that a Serve deployment created for a multi-host TPU slice intercepts + the default creation and provisions the SlicePlacementGroup correctly. """ llm_config = LLMConfig( model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), accelerator_type="TPU-V6E", accelerator_config={"kind": "tpu", "topology": "4x4"}, + engine_kwargs={"tensor_parallel_size": 16}, ) app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine) - serve.run(app) - pg_table = ray.util.placement_group_table() - active_pgs = list( - {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values() - ) + serve.start(http_options={"port": 0}) + try: + serve.run(app, name="default_host_app", route_prefix="/default_host") + pg_table = ray.util.placement_group_table() + active_pgs = list( + {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values() + ) - assert ( - len(active_pgs) == 2 - ), "Expected 2 PGs - one for TPU Head, one for worker bundles" + assert ( + len(active_pgs) == 2 + ), "Expected 2 PGs - one for TPU Head, one for worker bundles" - tpu_head_resource = "TPU-v6e-16-head" - head_pgs = [ - pg - for pg in active_pgs - if len(pg["bundles"]) == 1 - and tpu_head_resource in list(pg["bundles"].values())[0] - ] - assert len(head_pgs) == 1 + tpu_head_resource = "TPU-v6e-16-head" + head_pgs = [ + pg + for pg in active_pgs + if len(pg["bundles"]) == 1 + and tpu_head_resource in list(pg["bundles"].values())[0] + ] + assert len(head_pgs) == 1 - worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0] + worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0] - assert worker_pg["strategy"] == "PACK" - # 4x4 topology = 16 chips. Default is 16 bundles of 1 TPU. - assert len(worker_pg["bundles"]) == 16 - for bundle in worker_pg["bundles"].values(): - assert bundle.get("TPU", 0) == 1 + assert worker_pg["strategy"] == "PACK" - serve.shutdown() + # 4x4 topology = 16 chips. Default is 4 bundles of 4 TPUs. + assert len(worker_pg["bundles"]) == 4 + for bundle in worker_pg["bundles"].values(): + assert bundle.get("TPU", 0) == 4 + finally: + serve.shutdown() def test_tpu_serve_deployment_explicit_host_level_bundles(ray_tpu_cluster): @@ -274,38 +260,56 @@ def test_tpu_serve_deployment_explicit_host_level_bundles(ray_tpu_cluster): accelerator_type="TPU-V6E", accelerator_config={"kind": "tpu", "topology": "4x4"}, placement_group_config={"bundle_per_worker": {"TPU": 4}}, + engine_kwargs={"tensor_parallel_size": 4}, ) app = serve.deployment(LLMServer).bind(llm_config, engine_cls=PGCreationMockEngine) - serve.run(app) - pg_table = ray.util.placement_group_table() - active_pgs = list( - {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values() - ) + serve.start(http_options={"port": 0}) + try: + serve.run(app, name="explicit_host_app", route_prefix="/explicit_host") + pg_table = ray.util.placement_group_table() + active_pgs = list( + {k: v for k, v in pg_table.items() if v["state"] == "CREATED"}.values() + ) - assert ( - len(active_pgs) == 2 - ), "Expected 2 PGs - one for TPU Head, one for worker bundles" + assert ( + len(active_pgs) == 2 + ), "Expected 2 PGs - one for TPU Head, one for worker bundles" - tpu_head_resource = "TPU-v6e-16-head" - head_pgs = [ - pg - for pg in active_pgs - if len(pg["bundles"]) == 1 - and tpu_head_resource in list(pg["bundles"].values())[0] - ] - assert len(head_pgs) == 1 + tpu_head_resource = "TPU-v6e-16-head" + head_pgs = [ + pg + for pg in active_pgs + if len(pg["bundles"]) == 1 + and tpu_head_resource in list(pg["bundles"].values())[0] + ] + assert len(head_pgs) == 1 - worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0] + worker_pg = [pg for pg in active_pgs if pg not in head_pgs][0] - assert worker_pg["strategy"] == "PACK" - # 4x4 topology = 16 chips. With 4 TPUs per bundle, expect exactly 4 bundles. - assert len(worker_pg["bundles"]) == 4 - for bundle in worker_pg["bundles"].values(): - assert bundle.get("TPU", 0) == 4 + assert worker_pg["strategy"] == "PACK" - serve.shutdown() + assert len(worker_pg["bundles"]) == 4 + for bundle in worker_pg["bundles"].values(): + assert bundle.get("TPU", 0) == 4 + finally: + serve.shutdown() + + +def test_tpu_slice_kwargs_rejects_missing_accelerator_type(): + """ + Verifies that a ValueError is raised when a TPU topology is requested + but the accelerator type label is missing from the bundle labels. + """ + labels = [{"ray.io/tpu-topology": "4x4"}] + bundles = [{"TPU": 4}] + + with pytest.raises( + ValueError, + match="A TPU topology was requested, but 'ray.io/accelerator-type' was not found", + ): + resolve_tpu_slice_kwargs(labels, bundles) if __name__ == "__main__": diff --git a/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_server.py b/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_server.py index 1711f2f94317..cb8a5bcb2276 100644 --- a/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_server.py +++ b/python/ray/llm/tests/serve/cpu/deployments/llm/test_llm_server.py @@ -28,6 +28,7 @@ def serve_handle(mock_llm_config, stream_batching_interval_ms=0): } app = serve.deployment(LLMServer).bind(mock_llm_config, engine_cls=MockVLLMEngine) + serve.start(http_options={"port": 0}) handle = serve.run(app) # We set stream=True because the interfaces are async generators regardless # of the stream flag on request. @@ -53,6 +54,7 @@ def multiplexed_serve_handle(mock_llm_config, stream_batching_interval_ms=0): engine_cls=MockVLLMEngine, model_downloader=FakeLoraModelLoader, ) + serve.start(http_options={"port": 0}) handle = serve.run(app) handle = handle.options(stream=True, multiplexed_model_id="test_model_id") yield handle @@ -430,7 +432,11 @@ async def test_request_id_handling( chunks.append(chunk) assert len(chunks) == 1 - assert chunks[0].id == "test_request_id" + # Ray Serve intentionally drops internal request IDs to comply with strict + # OpenAI schemas (Pydantic v2). Verify the engine successfully processed + # the request and generated its own valid OpenAI-style ID instead. + assert isinstance(chunks[0].id, str) + assert chunks[0].id.startswith("chatcmpl-") @pytest.mark.parametrize("api_type", ["chat", "completion"]) @pytest.mark.parametrize("stream", [False, True]) @@ -678,19 +684,44 @@ def test_get_serve_options_without_accelerator_type(self): in serve_options["ray_actor_options"]["runtime_env"] ) - def test_deferred_placement_group_for_tpu_topology(self): - """Test that Serve skips PG creation when deferred placement group is required.""" + def test_tpu_placement_group_label_selector_injection(self): + """Test that Serve requests a placement group with TPU topology labels.""" llm_config = LLMConfig( model_loading_config=ModelLoadingConfig(model_id="test-tpu-model"), accelerator_type="TPU-V6E", accelerator_config={"kind": "tpu", "topology": "4x4"}, + engine_kwargs={"tensor_parallel_size": 16}, llm_engine="vLLM", ) serve_options = LLMServer.get_deployment_options(llm_config) - assert "placement_group_bundles" not in serve_options - assert "placement_group_strategy" not in serve_options + assert "placement_group_bundles" in serve_options + assert serve_options["placement_group_strategy"] == "PACK" + + pg_bundles = serve_options["placement_group_bundles"] + + # 4x4 topology = 16 chips. Default is 4 host-level bundles of 4 TPUs. + assert len(pg_bundles) == 4 + + # Verify the first bundle successfully merged the replica actor's CPU + assert pg_bundles[0].get("TPU") == 4 + assert pg_bundles[0].get("CPU", 0) >= 1 + + # Verify the remaining worker bundles are strictly TPU hosts + for i in range(1, 4): + assert pg_bundles[i].get("TPU") == 4 + assert pg_bundles[i].get("CPU", 0) == 0 + + # Verify the PG bundle label selector is injected for TPU topology + assert "placement_group_bundle_label_selector" in serve_options + + selectors = serve_options["placement_group_bundle_label_selector"] + assert len(selectors) == 4 + + for selector in selectors: + assert selector.get("ray.io/tpu-topology") == "4x4" + assert selector.get("ray.io/accelerator-type") == "TPU-V6E" if __name__ == "__main__": diff --git a/python/ray/llm/tests/serve/mocks/mock_vllm_engine.py b/python/ray/llm/tests/serve/mocks/mock_vllm_engine.py index 87b2d4751115..e5a99a8fb517 100644 --- a/python/ray/llm/tests/serve/mocks/mock_vllm_engine.py +++ b/python/ray/llm/tests/serve/mocks/mock_vllm_engine.py @@ -162,7 +162,10 @@ async def chat( # Generate streaming response async for response in self._generate_chat_response( - request=request, prompt_text=prompt_text.strip(), max_tokens=max_tokens + request=request, + prompt_text=prompt_text.strip(), + max_tokens=max_tokens, + raw_request_info=raw_request_info, ): yield response @@ -324,11 +327,19 @@ async def detokenize( yield response async def _generate_chat_response( - self, request: ChatCompletionRequest, prompt_text: str, max_tokens: int + self, + request: ChatCompletionRequest, + prompt_text: str, + max_tokens: int, + raw_request_info: Optional[RawRequestInfo] = None, ) -> AsyncGenerator[Union[str, ChatCompletionResponse], None]: """Generate mock chat completion response.""" - request_id = request.request_id or f"chatcmpl-{random.randint(1000, 9999)}" + request_id = ( + (raw_request_info.request_id if raw_request_info else None) + or getattr(request, "request_id", None) + or f"chatcmpl-{random.randint(1000, 9999)}" + ) # # Use request.model if provided, otherwise fall back to llm_config.model_id model_name = request.model or self.llm_config.model_id lora_prefix = ( diff --git a/python/ray/serve/_private/default_impl.py b/python/ray/serve/_private/default_impl.py index a68efdffd038..3f22ac4ec0cf 100644 --- a/python/ray/serve/_private/default_impl.py +++ b/python/ray/serve/_private/default_impl.py @@ -41,8 +41,10 @@ get_head_node_id, inside_ray_client_context, resolve_deployment_response, + resolve_tpu_slice_kwargs, ) from ray.util.placement_group import PlacementGroup +from ray.util.tpu import slice_placement_group # NOTE: Please read carefully before changing! # @@ -61,6 +63,24 @@ def create_cluster_node_info_cache(gcs_client: GcsClient) -> ClusterNodeInfoCach def _default_create_placement_group( request: CreatePlacementGroupRequest, ) -> PlacementGroup: + """Creates a placement group for the given request.""" + tpu_kwargs = resolve_tpu_slice_kwargs( + request.bundle_label_selector, request.bundles + ) + if tpu_kwargs is not None: + # If TPU-specific bundle label selectors are present, we utilize + # Slice PG utility to ensure atomic scheduling of co-located slices. + tpu_topology, tpu_version, worker_bundle = tpu_kwargs + slice_pg = slice_placement_group( + topology=tpu_topology, + accelerator_version=tpu_version, + resources_per_bundle=worker_bundle, + strategy=request.strategy, + name=request.name, + lifetime="detached", + ) + return slice_pg.placement_group + return ray.util.placement_group( request.bundles, request.strategy, diff --git a/python/ray/serve/_private/utils.py b/python/ray/serve/_private/utils.py index c4901fa4d9e9..9b5c80fd0d0d 100644 --- a/python/ray/serve/_private/utils.py +++ b/python/ray/serve/_private/utils.py @@ -13,7 +13,7 @@ from decimal import ROUND_HALF_UP, Decimal from enum import Enum from functools import wraps -from typing import Any, Callable, Dict, List, Optional, Set, TypeVar, Union +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar, Union import requests @@ -835,3 +835,50 @@ def _wake_up_next(self): self._value += 1 fut.set_result(True) return + + +def resolve_tpu_slice_kwargs( + labels: Optional[List[Dict[str, str]]], + bundles: Optional[List[Dict[str, float]]] = None, +) -> Optional[Tuple[str, str, Dict[str, float]]]: + """Parses TPU topology, version, and validates bundles for a TPU slice. + + Args: + labels: A list of label selector dictionaries applied to the placement group bundles. + bundles: An optional list of resource dictionaries defining the placement group bundles. + + Returns: + A tuple of (topology, accelerator_version, worker_bundle) if this is + a TPU slice request, otherwise None. + """ + if not labels or not isinstance(labels, list): + return None + + first_bundle_labels = labels[0] + if "ray.io/tpu-topology" not in first_bundle_labels: + return None + + tpu_topology = first_bundle_labels["ray.io/tpu-topology"] + + if "ray.io/accelerator-type" not in first_bundle_labels: + raise ValueError( + "A TPU topology was requested, but 'ray.io/accelerator-type' " + "was not found in the placement group bundle labels. The accelerator type " + "(e.g. 'TPU-V6E') is required to provision a multi-host slice." + ) + + raw_version = first_bundle_labels["ray.io/accelerator-type"] + tpu_version = raw_version.lower().replace("tpu-", "") + + worker_bundle = {"TPU": 1} + if bundles: + tpu_bundles = [b for b in bundles if b.get("TPU", 0) > 0] + if tpu_bundles: + worker_bundle = tpu_bundles[-1].copy() + if any(b.get("TPU") != worker_bundle.get("TPU") for b in tpu_bundles): + raise ValueError( + "Heterogeneous TPU bundles are not supported when a TPU topology is requested. " + "A multi-host TPU slice requires homogeneous resource bundles across all workers." + ) + + return tpu_topology, tpu_version, worker_bundle