From f0dd3873a922610229ba3f1f3bb0eca237d57d2d Mon Sep 17 00:00:00 2001 From: xuyongliang_01 Date: Fri, 31 Oct 2025 07:07:26 +0000 Subject: [PATCH] update ray_adapter/_private/state.py. Signed-off-by: xuyongliang_01 --- ray_adapter/_private/state.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/ray_adapter/_private/state.py b/ray_adapter/_private/state.py index c4722c2..37987de 100644 --- a/ray_adapter/_private/state.py +++ b/ray_adapter/_private/state.py @@ -36,17 +36,19 @@ def build_pg_dict(rg_info: RgInfo): for index, bundle in enumerate(rg_info.bundles): bundle_dict = {} - for key, value in bundle.resources.resources.items(): + for value in bundle.resources.resources.items(): resource = value if resource.type == Resource.Type.SCALER: - if "GPU" in resource.name: - bundle_dict["GPU"] = resource.scalar.value - elif "NPU" in resource.name: - bundle_dict["NPU"] = resource.scalar.value - elif "Memory" in resource.name: - bundle_dict["memory"] = resource.scalar.value - else: - bundle_dict[resource.name] = resource.scalar.value + resource_mapping = { + "GPU": "GPU", + "NPU": "NPU", + "Memory": "memory" + } + target_key = next( + (v for k, v in resource_mapping.items() if k in resource.name), + resource.name + ) + bundle_dict[target_key] = resource.scalar.value rg_dict['bundles'][index] = bundle_dict rg_dict['bundles_to_node_id'][index] = bundle.functionProxyId -- Gitee