PTFE-1402 enable or disable external IP GCP runner (#523)

Author: tcarmet

runner_manager/backend/gcloud.py

@@ -91,16 +91,16 @@ def disks(self) -> List[AttachedDisk]:
 
     @property
     def network_interfaces(self) -> List[NetworkInterface]:
-        return [
-            NetworkInterface(
-                network=self.instance_config.network,
-                access_configs=[
-                    AccessConfig(
-                        name="External NAT",
-                    )
-                ],
-            )
-        ]
+        network_interface: NetworkInterface = NetworkInterface(
+            network=self.instance_config.network,
+        )
+        if self.instance_config.enable_external_ip:
+            network_interface.access_configs = [
+                AccessConfig(
+                    name="External NAT",
+                )
+            ]
+        return [network_interface]
 
     @property
     def scheduling(self) -> Scheduling:

runner_manager/models/backend.py

@@ -104,6 +104,7 @@ class GCPInstanceConfig(InstanceConfig):
     machine_type: str = "e2-small"
     network: str = "global/networks/default"
     enable_nested_virtualization: bool = True
+    enable_external_ip: bool = True
     spot: bool = False
     disk_size_gb: int = 20
     disk_type: Literal["pd-ssd", "pd-standard"] = "pd-ssd"

tests/unit/backend/test_gcp.py

@@ -1,7 +1,7 @@
 import os
 from typing import List
 
-from google.cloud.compute import Image
+from google.cloud.compute import Image, NetworkInterface
 from pytest import fixture, mark, raises
 from redis_om import NotFoundError
 
@@ -57,6 +57,12 @@ def gcp_runner(runner: Runner, gcp_group: RunnerGroup) -> Runner:
 def test_gcp_network_interfaces(gcp_group: RunnerGroup):
     interfaces = gcp_group.backend.network_interfaces
     assert len(interfaces) == 1
+    assert interfaces[0].access_configs[0].name == "External NAT"
+    # Test disabling external IP
+    gcp_group.backend.instance_config.enable_external_ip = False
+    interfaces: List[NetworkInterface] = gcp_group.backend.network_interfaces
+    assert len(interfaces) == 1
+    assert len(interfaces[0].access_configs) == 0
 
 
 def test_gcp_group(gcp_group: RunnerGroup):