diff --git a/components/src/dynamo/planner/monitoring/dgd_services.py b/components/src/dynamo/planner/monitoring/dgd_services.py index 7b599755e75d..80d66615545b 100644 --- a/components/src/dynamo/planner/monitoring/dgd_services.py +++ b/components/src/dynamo/planner/monitoring/dgd_services.py @@ -210,6 +210,15 @@ def get_component_from_type_or_name( component_names = [name for name, _ in matching_components] raise DuplicateSubComponentError(component_type.value, component_names) + if not matching_components and component_type == SubComponentType.DECODE: + generic_workers = [ + (name, component) + for name, component in components.items() + if get_component_type(component) == V1BETA1_GENERIC_WORKER_COMPONENT_TYPE + ] + if len(generic_workers) == 1: + matching_components = generic_workers + if not matching_components and component_name in components: component = components[component_name] if not _can_use_explicit_component_name(component, component_type): diff --git a/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py b/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py index 2f8bbe4611b1..5879f174dad5 100644 --- a/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py +++ b/components/src/dynamo/planner/tests/unit/test_kubernetes_connector.py @@ -378,6 +378,55 @@ def test_get_service_name_from_v1beta_worker_type_by_name(kubernetes_connector): assert service.number_replicas() == 2 +def test_get_service_name_from_unique_v1beta_worker_type_for_decode( + kubernetes_connector, +): + deployment = _deployment(_component("arbitrary-name", "worker", replicas=2)) + + service = get_component_from_type_or_name(deployment, SubComponentType.DECODE) + + assert service.name == "arbitrary-name" + assert service.number_replicas() == 2 + + +def test_get_service_name_from_multiple_v1beta_workers_by_name( + kubernetes_connector, +): + deployment = _deployment( + _component("prefill-name", "worker"), + _component("decode-name", "worker"), + ) + + with pytest.raises(SubComponentNotFoundError): + get_component_from_type_or_name(deployment, SubComponentType.DECODE) + + service = get_component_from_type_or_name( + deployment, SubComponentType.DECODE, "decode-name" + ) + assert service.name == "decode-name" + + +@pytest.mark.asyncio +async def test_validate_deployment_agg_worker_by_type( + kubernetes_connector, mock_kube_api +): + mock_kube_api.get_graph_deployment.return_value = _deployment( + _component("Frontend", "frontend", replicas=1), + _component( + "arbitrary-name", + "worker", + replicas=2, + args=["--model", "Qwen/Qwen3-8B"], + ), + ) + + await kubernetes_connector.validate_deployment( + decode_component_name="stale-default-name", + require_prefill=False, + require_decode=True, + ) + + def test_get_service_name_from_sub_component_type_not_found(kubernetes_connector): deployment = _deployment(_component("test-component-decode", "decode", replicas=3)) with pytest.raises(SubComponentNotFoundError) as exc_info: