Move tailscale_auth_key to env and set it on controller.

skypilot-org · Sep 29, 2024 · c137e1f · c137e1f
1 parent 779eba7
commit c137e1f
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 5 deletions.
diff --git a/sky/serve/service_spec.py b/sky/serve/service_spec.py
@@ -183,9 +183,9 @@ def from_yaml_config(config: Dict[str, Any]) -> 'SkyServiceSpec':
         vpn_section: Dict[str, Any] = config.get('vpn', None)
         if vpn_section is not None:
             # Get the tailscale auth key from the environment.
-            tailscale_auth_key = vpn_section.get('tailscale_auth_key', None)
-            if tailscale_auth_key is not None:
-                service_config['tailscale_auth_key'] = tailscale_auth_key
+            if vpn_section.get('tailscale'):
+                service_config['tailscale_auth_key'] = os.getenv(
+                    'TAILSCALE_AUTH_KEY')
 
         return SkyServiceSpec(**service_config)
 

diff --git a/sky/skylet/constants.py b/sky/skylet/constants.py
@@ -287,3 +287,6 @@
 ROLE_ASSIGNMENT_FAILURE_ERROR_MSG = (
     'Failed to assign Storage Blob Data Owner role to the '
     'storage account {storage_account_name}.')
+
+# Constants for Tailscale VPN
+TAILSCALE_AUTH_KEY_ENV_VAR = 'TAILSCALE_AUTH_KEY'
diff --git a/sky/utils/controller_utils.py b/sky/utils/controller_utils.py
@@ -420,6 +420,12 @@ def shared_controller_vars_to_fill(
         # Only set the SKYPILOT_CONFIG env var if the user has a config file.
         env_vars[
             skypilot_config.ENV_VAR_SKYPILOT_CONFIG] = remote_user_config_path
+    # Check if VPN auth key should be set.
+    if task_config:
+        vpn_config = task_config.get('service', {}).get('vpn', {})
+        if vpn_config.get('tailscale'):
+            env_vars[constants.TAILSCALE_AUTH_KEY_ENV_VAR] = task_config[
+                'tailscale_auth_key']
     vars_to_fill['controller_envs'] = env_vars
     return vars_to_fill
 

diff --git a/sky/utils/schemas.py b/sky/utils/schemas.py
@@ -379,8 +379,8 @@ def get_service_schema():
                 'type': 'object',
                 'additionalProperties': False,
                 'properties': {
-                    'tailscale_auth_key': {
-                        'type': 'string',
+                    'tailscale': {
+                        'type': 'boolean',
                     },
                 }
             },