[k8s] Allow per-context identity configuration + in-cluster auth fixes

sky/authentication.py

@@ -380,8 +380,8 @@ def setup_kubernetes_authentication(config: Dict[str, Any]) -> Dict[str, Any]:
     secret_field_name = clouds.Kubernetes().ssh_key_secret_field_name
     context = config['provider'].get(
         'context', kubernetes_utils.get_current_kube_config_context_name())
-    if context == kubernetes_utils.IN_CLUSTER_REGION:
-        # If the context is set to IN_CLUSTER_REGION, we are running in a pod
+    if context == kubernetes_utils.get_in_cluster_context_name():
+        # If the context is a in-cluster context name, we are running in a pod
         # with in-cluster configuration. We need to set the context to None
         # to use the mounted service account.
         context = None

sky/clouds/kubernetes.py

@@ -182,8 +182,9 @@ def regions_with_offering(cls, instance_type: Optional[str],
             if context is None:
                 # If running in-cluster, we allow the region to be set to the
                 # singleton region since there is no context name available.
-                regions.append(clouds.Region(
-                    kubernetes_utils.IN_CLUSTER_REGION))
+                in_cluster_region = (
+                    kubernetes_utils.get_in_cluster_context_name())
+                regions.append(clouds.Region(in_cluster_region))
             else:
                 regions.append(clouds.Region(context))
 
@@ -376,20 +377,32 @@ def make_deploy_resources_variables(
         remote_identity = skypilot_config.get_nested(
             ('kubernetes', 'remote_identity'),
             schemas.get_default_remote_identity('kubernetes'))
-        if (remote_identity ==
+
+        if isinstance(remote_identity, dict):
+            # If remote_identity is a dict, use the service account for the
+            # current context
+            k8s_service_account_name = remote_identity.get(context, None)
+            if k8s_service_account_name is None:
+                err_msg = (f'Context {context!r} not found in '
+                           'remote identities from config.yaml')
+                raise ValueError(err_msg)
+        else:
+            # If remote_identity is not a dict, use
+            k8s_service_account_name = remote_identity
+
+        if (k8s_service_account_name ==
                 schemas.RemoteIdentityOptions.LOCAL_CREDENTIALS.value):
             # SA name doesn't matter since automounting credentials is disabled
             k8s_service_account_name = 'default'
             k8s_automount_sa_token = 'false'
-        elif (remote_identity ==
+        elif (k8s_service_account_name ==
               schemas.RemoteIdentityOptions.SERVICE_ACCOUNT.value):
             # Use the default service account
             k8s_service_account_name = (
                 kubernetes_utils.DEFAULT_SERVICE_ACCOUNT_NAME)
             k8s_automount_sa_token = 'true'
         else:
             # User specified a custom service account
-            k8s_service_account_name = remote_identity
             k8s_automount_sa_token = 'true'
 
         fuse_device_required = bool(resources.requires_fuse)
@@ -410,6 +423,14 @@ def make_deploy_resources_variables(
             ('kubernetes', 'provision_timeout'),
             10,
             override_configs=resources.cluster_config_overrides)
+
+        # Set environment variables for the pod. Note that SkyPilot env vars
+        # are set separately when the task is run. These env vars are
+        # independent of the SkyPilot task to be run.
+        k8s_env_vars = {
+            kubernetes_utils.IN_CLUSTER_CONTEXT_NAME_ENV_VAR: context
+        }
+
         deploy_vars = {
             'instance_type': resources.instance_type,
             'custom_resources': custom_resources,
@@ -431,6 +452,7 @@ def make_deploy_resources_variables(
             'k8s_skypilot_system_namespace': _SKYPILOT_SYSTEM_NAMESPACE,
             'k8s_spot_label_key': spot_label_key,
             'k8s_spot_label_value': spot_label_value,
+            'k8s_env_vars': k8s_env_vars,
             'image_id': image_id,
         }
 
@@ -566,7 +588,7 @@ def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
             # TODO: Remove this after 0.9.0.
             return region, zone
 
-        if region == kubernetes_utils.IN_CLUSTER_REGION:
+        if region == kubernetes_utils.get_in_cluster_context_name():
             # If running incluster, we set region to IN_CLUSTER_REGION
             # since there is no context name available.
             return region, zone
@@ -575,7 +597,7 @@ def validate_region_zone(self, region: Optional[str], zone: Optional[str]):
         if all_contexts == [None]:
             # If [None] context is returned, use the singleton region since we
             # are running in a pod with in-cluster auth.
-            all_contexts = [kubernetes_utils.IN_CLUSTER_REGION]
+            all_contexts = [kubernetes_utils.get_in_cluster_context_name()]
         if region not in all_contexts:
             raise ValueError(
                 f'Context {region} not found in kubeconfig. Kubernetes only '

sky/provision/kubernetes/utils.py

@@ -36,7 +36,13 @@
 
 # TODO(romilb): Move constants to constants.py
 DEFAULT_NAMESPACE = 'default'
-IN_CLUSTER_REGION = 'in-cluster'
+DEFAULT_IN_CLUSTER_REGION = 'in-cluster'
+
+# The name for the environment variable that stores the in-cluster context name
+# for Kubernetes clusters. This is used to associate a name with the current
+# context when running with in-cluster auth. If not set, the context name is
+# set to DEFAULT_IN_CLUSTER_REGION.
+IN_CLUSTER_CONTEXT_NAME_ENV_VAR = 'SKYPILOT_IN_CLUSTER_CONTEXT_NAME'
 
 DEFAULT_SERVICE_ACCOUNT_NAME = 'skypilot-service-account'
 
@@ -1996,9 +2002,9 @@ def set_autodown_annotations(handle: 'backends.CloudVmRayResourceHandle',
 def get_context_from_config(provider_config: Dict[str, Any]) -> Optional[str]:
     context = provider_config.get('context',
                                   get_current_kube_config_context_name())
-    if context == IN_CLUSTER_REGION:
-        # If the context (also used as the region) is set to IN_CLUSTER_REGION
-        # we need to use in-cluster auth.
+    if context == get_in_cluster_context_name():
+        # If the context (also used as the region) is in-cluster, we need to
+        # we need to use in-cluster auth by setting the context to None.
         context = None
     return context
 
@@ -2136,3 +2142,13 @@ def process_skypilot_pods(
         num_pods = len(cluster.pods)
         cluster.resources_str = f'{num_pods}x {cluster.resources}'
     return list(clusters.values()), jobs_controllers, serve_controllers
+
+
+def get_in_cluster_context_name() -> Optional[str]:
+    """Returns the name of the in-cluster context from the environment.
+
+    If the environment variable is not set, returns the default in-cluster
+    context name.
+    """
+    return (os.environ.get(IN_CLUSTER_CONTEXT_NAME_ENV_VAR) or
+            DEFAULT_IN_CLUSTER_REGION)

sky/templates/kubernetes-ray.yml.j2

@@ -370,6 +370,14 @@ available_node_types:
               }
               trap : TERM INT; log_tail || sleep infinity & wait
 
+          {% if k8s_env_vars is not none %}
+          env:
+          {% for key, value in k8s_env_vars.items() %}
+          - name: {{ key }}
+            value: {{ value }}
+          {% endfor %}
+          {% endif %}
+
           ports:
           - containerPort: 22  # Used for SSH
           - containerPort: {{ray_port}}  # Redis port

sky/utils/schemas.py

@@ -668,7 +668,14 @@ def get_default_remote_identity(cloud: str) -> str:
 
 _REMOTE_IDENTITY_SCHEMA_KUBERNETES = {
     'remote_identity': {
-        'type': 'string'
+        'anyOf': [{
+            'type': 'string'
+        }, {
+            'type': 'object',
+            'additionalProperties': {
+                'type': 'string'
+            }
+        }]
     },
 }