feat(processor): enable EFA for whole-host reservations (single-node, not just multinode)

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,30 +4838,31 @@ 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"
 
     return requests
 
 
 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:

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):

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
 
 
 # --------------------------------------------------------------------------- #