From 3e98afe8f96531ffb6332679083918f7ad67481e Mon Sep 17 00:00:00 2001
From: Romil Bhardwaj <romil.bhardwaj@berkeley.edu>
Date: Wed, 16 Oct 2024 20:32:38 -0700
Subject: [PATCH] [GCP] Add gVNIC support (#4095)

* add gvnic support through config.yaml

* lint

* docs
---
 docs/source/reference/config.rst | 9 +++++++++
 sky/clouds/gcp.py                | 5 +++++
 sky/provision/gcp/config.py      | 6 +++++-
 sky/templates/gcp-ray.yml.j2     | 3 +++
 sky/utils/schemas.py             | 3 +++
 5 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/docs/source/reference/config.rst b/docs/source/reference/config.rst
index 5c52e7487b9..b8255b46402 100644
--- a/docs/source/reference/config.rst
+++ b/docs/source/reference/config.rst
@@ -419,6 +419,15 @@ Available fields and semantics:
     # Default: 'LOCAL_CREDENTIALS'.
     remote_identity: LOCAL_CREDENTIALS
 
+    # Enable gVNIC (optional).
+    #
+    # Set to true to use gVNIC on GCP instances. gVNIC offers higher performance
+    # for multi-node clusters, but costs more.
+    # Reference: https://cloud.google.com/compute/docs/networking/using-gvnic
+    #
+    # Default: false.
+    enable_gvnic: false
+
   # Advanced Azure configurations (optional).
   # Apply to all new instances but not existing ones.
   azure:
diff --git a/sky/clouds/gcp.py b/sky/clouds/gcp.py
index bac297a92c3..0e02f9fd456 100644
--- a/sky/clouds/gcp.py
+++ b/sky/clouds/gcp.py
@@ -546,6 +546,11 @@ def _failover_disk_tier() -> Optional[resources_utils.DiskTier]:
         resources_vars[
             'force_enable_external_ips'] = skypilot_config.get_nested(
                 ('gcp', 'force_enable_external_ips'), False)
+
+        # Add gVNIC from config
+        resources_vars['enable_gvnic'] = skypilot_config.get_nested(
+            ('gcp', 'enable_gvnic'), False)
+
         return resources_vars
 
     def _get_feasible_launchable_resources(
diff --git a/sky/provision/gcp/config.py b/sky/provision/gcp/config.py
index 416f0c1a694..a8292669a7c 100644
--- a/sky/provision/gcp/config.py
+++ b/sky/provision/gcp/config.py
@@ -670,8 +670,12 @@ def _configure_subnet(region: str, cluster_name: str,
         'accessConfigs': [{
             'name': 'External NAT',
             'type': 'ONE_TO_ONE_NAT',
-        }],
+        }]
     }]
+    # Add gVNIC if specified in config
+    enable_gvnic = config.provider_config.get('enable_gvnic', False)
+    if enable_gvnic:
+        default_interfaces[0]['nicType'] = 'gVNIC'
     enable_external_ips = _enable_external_ips(config)
     if not enable_external_ips:
         # Removing this key means the VM will not be assigned an external IP.
diff --git a/sky/templates/gcp-ray.yml.j2 b/sky/templates/gcp-ray.yml.j2
index 5f06eef05c7..f3e6232d5d8 100644
--- a/sky/templates/gcp-ray.yml.j2
+++ b/sky/templates/gcp-ray.yml.j2
@@ -64,6 +64,9 @@ provider:
   # leakage.
   disable_launch_config_check: true
   use_managed_instance_group: {{ gcp_use_managed_instance_group }}
+{%- if enable_gvnic %}
+  enable_gvnic: {{ enable_gvnic }}
+{%- endif %}
 
 auth:
   ssh_user: gcpuser
diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
index 6e752f73ebc..94a6ed690e1 100644
--- a/sky/utils/schemas.py
+++ b/sky/utils/schemas.py
@@ -755,6 +755,9 @@ def get_config_schema():
                 'force_enable_external_ips': {
                     'type': 'boolean'
                 },
+                'enable_gvnic': {
+                    'type': 'boolean'
+                },
                 **_LABELS_SCHEMA,
                 **_NETWORK_CONFIG_SCHEMA,
             },