diff --git a/terraform-gpu-devservers/lambda/reservation_processor/index.py b/terraform-gpu-devservers/lambda/reservation_processor/index.py index 804cb31..23eea1c 100644 --- a/terraform-gpu-devservers/lambda/reservation_processor/index.py +++ b/terraform-gpu-devservers/lambda/reservation_processor/index.py @@ -233,7 +233,7 @@ def scale_up_spot_asg(gpu_type: str, reservation_id: str = ""): "a100": {"instance_type": "p4d.24xlarge", "max_gpus": 8, "cpus": 96, "memory_gb": 1152, "efa_count": 4}, "h100": {"instance_type": "p5.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 32}, "h200": {"instance_type": "p5e.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 32}, - "b200": {"instance_type": "p6-b200.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 32}, + "b200": {"instance_type": "p6-b200.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 8}, # p6-b200.48xlarge advertises 8 EFA network cards (see main.tf efa_network_cards) "b300": {"instance_type": "p6-b300.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 8}, "cpu-arm": {"instance_type": "c7g.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64, "efa_count": 0}, "cpu-x86": {"instance_type": "c7i.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64, "efa_count": 0}, @@ -4790,22 +4790,15 @@ def get_pod_resource_limits(gpu_count: int, gpu_type: str, is_multinode: bool = "memory": f"{proportional_memory_limit}Gi" }) - # EFA optimization: Only use EFA for full-node multinode deployments (skip MIG slices) - use_efa = ( - gpu_type != "t4-small" and - not gpu_type.startswith("cpu-") and - "mig" not in gpu_type and - is_multinode and - gpu_count == max_gpus - ) - - if use_efa: - efa_count = config.get("efa_count", 1) + # EFA is attached to whole-host reservations (all GPUs on the node) for any + # EFA-capable SKU — single-node or multinode. See _pod_uses_efa. + if _pod_uses_efa(gpu_count, gpu_type): + efa_count = config["efa_count"] limits["vpc.amazonaws.com/efa"] = str(efa_count) limits["hugepages-2Mi"] = "5120Mi" - logger.info(f"Using EFA ({efa_count} interfaces) for multinode full-node deployment: {gpu_count}/{max_gpus} GPUs") + logger.info(f"Using EFA ({efa_count} interfaces) for whole-host deployment: {gpu_count}/{max_gpus} GPUs of {gpu_type}") else: - logger.info(f"Skipping EFA: multinode={is_multinode}, gpu_count={gpu_count}/{max_gpus}, gpu_type={gpu_type}") + logger.info(f"Skipping EFA: gpu_count={gpu_count}/{max_gpus}, gpu_type={gpu_type}") return limits @@ -4845,16 +4838,9 @@ def get_pod_resource_requests(gpu_count: int, gpu_type: str, is_multinode: bool "memory": f"{proportional_memory_request}Gi" }) - # EFA: Only for full-node multinode deployments (skip MIG slices) - use_efa = ( - gpu_type != "t4-small" and - not gpu_type.startswith("cpu-") and - "mig" not in gpu_type and - is_multinode and - gpu_count == max_gpus - ) - if use_efa: - efa_count = config.get("efa_count", 1) + # EFA: whole-host reservations on EFA-capable SKUs (see _pod_uses_efa). + if _pod_uses_efa(gpu_count, gpu_type): + efa_count = config["efa_count"] requests["vpc.amazonaws.com/efa"] = str(efa_count) requests["hugepages-2Mi"] = "5120Mi" @@ -4862,13 +4848,21 @@ def get_pod_resource_requests(gpu_count: int, gpu_type: str, is_multinode: bool def _pod_uses_efa(gpu_count: int, gpu_type: str, is_multinode: bool = False) -> bool: - """Check if pod will use EFA based on configuration""" + """Whether the pod should get EFA interfaces. + + EFA is attached to whole-host reservations — the pod takes every GPU on the + node (gpu_count == max_gpus) on an EFA-capable instance type (efa_count > 0). + This covers both single-node whole-host and multinode (each node is a whole + host) deployments. Partial-node ("subhost") pods, and MIG / CPU / t4 SKUs + (efa_count == 0), never get EFA. Because a whole-host pod owns the node it + can keep hostNetwork + sshd on 2222 with no co-tenant port conflict. + + is_multinode is accepted for call-site compatibility but no longer gates the + decision. + """ config = GPU_CONFIG.get(gpu_type, GPU_CONFIG_DEFAULT) - return ( - gpu_type != "t4-small" and - is_multinode and - gpu_count == config["max_gpus"] - ) + max_gpus = config["max_gpus"] + return config.get("efa_count", 0) > 0 and max_gpus > 0 and int(gpu_count) == max_gpus def get_cpu_thread_env_vars(gpu_count: int, gpu_type: str) -> list: diff --git a/tests/unit/lambda_fn/test_mig_gpu_config.py b/tests/unit/lambda_fn/test_mig_gpu_config.py index 00a822d..47494bc 100644 --- a/tests/unit/lambda_fn/test_mig_gpu_config.py +++ b/tests/unit/lambda_fn/test_mig_gpu_config.py @@ -237,9 +237,11 @@ def test_limits_full_gpu_multinode_gets_efa(lambda_index): assert lim["hugepages-2Mi"] == "5120Mi" -def test_limits_full_gpu_singlenode_no_efa(lambda_index): +def test_limits_full_gpu_singlenode_gets_efa(lambda_index): + # Whole-host single-node (8/8, not multinode) now gets EFA too. lim = lambda_index.get_pod_resource_limits(8, "h100", is_multinode=False) - assert "vpc.amazonaws.com/efa" not in lim + assert lim["vpc.amazonaws.com/efa"] == "32" + assert lim["hugepages-2Mi"] == "5120Mi" def test_requests_full_gpu_multinode_gets_efa(lambda_index): @@ -263,8 +265,9 @@ def test_pod_uses_efa_false_partial(lambda_index): assert lambda_index._pod_uses_efa(4, "h100", is_multinode=True) is False -def test_pod_uses_efa_false_singlenode(lambda_index): - assert lambda_index._pod_uses_efa(8, "h100", is_multinode=False) is False +def test_pod_uses_efa_true_singlenode_whole_host(lambda_index): + # Whole-host single-node reservation qualifies for EFA regardless of multinode. + assert lambda_index._pod_uses_efa(8, "h100", is_multinode=False) is True def test_pod_uses_efa_t4_small_excluded(lambda_index): diff --git a/tests/unit/lambda_fn/test_pod_resources.py b/tests/unit/lambda_fn/test_pod_resources.py index 018f7b9..617aba3 100644 --- a/tests/unit/lambda_fn/test_pod_resources.py +++ b/tests/unit/lambda_fn/test_pod_resources.py @@ -30,9 +30,13 @@ def test_h100_single_gpu_limits(self): assert lim == {"nvidia.com/gpu": "1", "cpu": "36", "memory": "256Gi"} def test_h100_full_node_limits_cpu_capped(self): - # ratio 1.0: fractional_cpu*1.5 = 288 but capped at config cpus (192) + # ratio 1.0: fractional_cpu*1.5 = 288 but capped at config cpus (192). + # Whole-host h100 also gets EFA (efa_count>0, gpu==max_gpus), even single-node. lim = index.get_pod_resource_limits(8, "h100") - assert lim == {"nvidia.com/gpu": "8", "cpu": "192", "memory": "2048Gi"} + assert lim == { + "nvidia.com/gpu": "8", "cpu": "192", "memory": "2048Gi", + "vpc.amazonaws.com/efa": "32", "hugepages-2Mi": "5120Mi", + } # The 1.5x boost must NOT exceed the node's physical cpu count. assert int(lim["cpu"]) == 192 @@ -59,9 +63,13 @@ def test_h100_single_gpu_requests(self): assert req == {"nvidia.com/gpu": "1", "cpu": "21", "memory": "230Gi"} def test_h100_full_node_requests(self): - # ratio 1.0: cpu=int(192*0.9)=172, mem=int(2048*0.9)=1843 + # ratio 1.0: cpu=int(192*0.9)=172, mem=int(2048*0.9)=1843. + # Whole-host h100 also carries the EFA + hugepages requests. req = index.get_pod_resource_requests(8, "h100") - assert req == {"nvidia.com/gpu": "8", "cpu": "172", "memory": "1843Gi"} + assert req == { + "nvidia.com/gpu": "8", "cpu": "172", "memory": "1843Gi", + "vpc.amazonaws.com/efa": "32", "hugepages-2Mi": "5120Mi", + } def test_requests_below_limits(self): # requests (0.9x) must be <= limits for the same allocation. @@ -184,7 +192,8 @@ def test_decimal_truncates_like_int(self): # --------------------------------------------------------------------------- # -# EFA gating: only multinode + full-node, never t4-small / MIG / cpu +# EFA gating: whole-host (all GPUs on the node) on EFA-capable SKUs, single-node +# or multinode; never partial-node ("subhost") / t4-small / MIG / cpu. # --------------------------------------------------------------------------- # class TestEfaGating: def test_efa_added_for_full_node_multinode_limits(self): @@ -197,41 +206,44 @@ def test_efa_added_for_full_node_multinode_requests(self): assert req["vpc.amazonaws.com/efa"] == "32" assert req["hugepages-2Mi"] == "5120Mi" - def test_no_efa_when_not_multinode(self): + def test_efa_added_for_whole_host_single_node(self): + # Whole-host single-node reservation (8/8, not multinode) now gets EFA: + # the pod owns the node, so hostNetwork + sshd:2222 stays conflict-free. lim = index.get_pod_resource_limits(8, "h100", is_multinode=False) - assert "vpc.amazonaws.com/efa" not in lim + req = index.get_pod_resource_requests(8, "h100", is_multinode=False) + assert lim["vpc.amazonaws.com/efa"] == "32" + assert lim["hugepages-2Mi"] == "5120Mi" + assert req["vpc.amazonaws.com/efa"] == "32" def test_no_efa_for_partial_node(self): - # multinode but not the whole node (4 of 8) => no EFA. - lim = index.get_pod_resource_limits(4, "h100", is_multinode=True) - assert "vpc.amazonaws.com/efa" not in lim - assert "hugepages-2Mi" not in lim + # Subhost (4 of 8) never gets EFA — even multinode. + assert "vpc.amazonaws.com/efa" not in index.get_pod_resource_limits(4, "h100", is_multinode=True) + assert "vpc.amazonaws.com/efa" not in index.get_pod_resource_limits(4, "h100", is_multinode=False) + assert "hugepages-2Mi" not in index.get_pod_resource_limits(4, "h100", is_multinode=False) - def test_t4_small_excluded_even_when_full_multinode(self): - # t4-small max_gpus==1, so gpu_count==max_gpus is satisfied, but the - # explicit t4-small exclusion still blocks EFA. + def test_t4_small_excluded_even_when_full(self): + # t4-small max_gpus==1 so gpu_count==max_gpus holds, but efa_count==0 blocks EFA. lim = index.get_pod_resource_limits(1, "t4-small", is_multinode=True) assert "vpc.amazonaws.com/efa" not in lim - def test_mig_full_node_multinode_no_efa(self): - # mig-2g max_gpus==8; full + multinode but MIG is excluded. + def test_mig_full_node_no_efa(self): + # mig-2g max_gpus==8; full node but efa_count==0 excludes MIG. lim = index.get_pod_resource_limits(8, "h100-mig-2g", is_multinode=True) assert "vpc.amazonaws.com/efa" not in lim def test_efa_count_per_gpu_type(self): - # a100 efa_count=4, b300 efa_count=8 — value pulled from GPU_CONFIG. - a100 = index.get_pod_resource_limits(8, "a100", is_multinode=True) - b300 = index.get_pod_resource_limits(8, "b300", is_multinode=True) - assert a100["vpc.amazonaws.com/efa"] == "4" - assert b300["vpc.amazonaws.com/efa"] == "8" + # value pulled from GPU_CONFIG: a100=4, b200=8, b300=8. + assert index.get_pod_resource_limits(8, "a100")["vpc.amazonaws.com/efa"] == "4" + assert index.get_pod_resource_limits(8, "b200")["vpc.amazonaws.com/efa"] == "8" + assert index.get_pod_resource_limits(8, "b300")["vpc.amazonaws.com/efa"] == "8" def test_pod_uses_efa_helper_agreement(self): - # _pod_uses_efa should agree with the presence of the efa key in limits, - # except it omits the MIG/cpu guards present in the limits function. + # _pod_uses_efa gates on efa_count>0 and whole-host, independent of multinode. assert index._pod_uses_efa(8, "h100", is_multinode=True) is True - assert index._pod_uses_efa(8, "h100", is_multinode=False) is False - assert index._pod_uses_efa(4, "h100", is_multinode=True) is False + assert index._pod_uses_efa(8, "h100", is_multinode=False) is True + assert index._pod_uses_efa(4, "h100", is_multinode=False) is False assert index._pod_uses_efa(1, "t4-small", is_multinode=True) is False + assert index._pod_uses_efa(0, "cpu-x86", is_multinode=True) is False # --------------------------------------------------------------------------- #