Simplify Kubernetes backend config

src/dstack/_internal/core/backends/kubernetes/compute.py

@@ -15,9 +15,12 @@
     get_instance_name,
 )
 from dstack._internal.core.backends.base.offers import match_requirements
-from dstack._internal.core.backends.kubernetes.client import get_api_from_config_data
 from dstack._internal.core.backends.kubernetes.config import KubernetesConfig
-from dstack._internal.core.errors import GatewayError
+from dstack._internal.core.backends.kubernetes.utils import (
+    get_api_from_config_data,
+    get_cluster_public_ip,
+)
+from dstack._internal.core.errors import BackendError, GatewayError
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.instances import (
     Disk,
@@ -97,16 +100,29 @@ def run_job(
         # as an ssh proxy jump to connect to all other services in Kubernetes.
         # Setup jump pod in a separate thread to avoid long-running run_job.
         # In case the thread fails, the job will be failed and resubmitted.
+        jump_pod_hostname = self.config.networking.ssh_host
+        if jump_pod_hostname is None:
+            jump_pod_hostname = get_cluster_public_ip(self.api)
+            if jump_pod_hostname is None:
+                raise BackendError(
+                    "Failed to acquire an IP for jump pod automatically. "
+                    "Specify ssh_host for Kubernetes backend."
+                )
+        jump_pod_port = _create_jump_pod_service_if_not_exists(
+            api=self.api,
+            project_name=run.project_name,
+            project_ssh_public_key=project_ssh_public_key.strip(),
+            jump_pod_port=self.config.networking.ssh_port,
+        )
         threading.Thread(
-            target=_setup_jump_pod,
+            target=_continue_setup_jump_pod,
             kwargs={
                 "api": self.api,
                 "project_name": run.project_name,
-                "project_ssh_public_key": project_ssh_public_key.strip(),
                 "project_ssh_private_key": project_ssh_private_key.strip(),
                 "user_ssh_public_key": run.run_spec.ssh_key_pub.strip(),
-                "jump_pod_host": self.config.networking.ssh_host,
-                "jump_pod_port": self.config.networking.ssh_port,
+                "jump_pod_host": jump_pod_hostname,
+                "jump_pod_port": jump_pod_port,
             },
         ).start()
         self.api.create_namespaced_pod(
@@ -157,9 +173,9 @@ def run_job(
             ssh_port=RUNNER_SSH_PORT,
             dockerized=False,
             ssh_proxy=SSHConnectionParams(
-                hostname=self.config.networking.ssh_host,
+                hostname=jump_pod_hostname,
                 username="root",
-                port=self.config.networking.ssh_port,
+                port=jump_pod_port,
             ),
             backend_data=None,
         )
@@ -301,21 +317,14 @@ def _get_gpus_from_node_labels(labels: Dict) -> List[Gpu]:
     return [Gpu(name=gpu_name, memory_mib=gpu_memory) for _ in range(gpu_count)]
 
 
-def _setup_jump_pod(
+def _continue_setup_jump_pod(
     api: client.CoreV1Api,
     project_name: str,
-    project_ssh_public_key: str,
     project_ssh_private_key: str,
     user_ssh_public_key: str,
     jump_pod_host: str,
     jump_pod_port: int,
 ):
-    _create_jump_pod_service_if_not_exists(
-        api=api,
-        project_name=project_name,
-        project_ssh_public_key=project_ssh_public_key,
-        jump_pod_port=jump_pod_port,
-    )
     _wait_for_pod_ready(
         api=api,
         pod_name=_get_jump_pod_name(project_name),
@@ -332,31 +341,32 @@ def _create_jump_pod_service_if_not_exists(
     api: client.CoreV1Api,
     project_name: str,
     project_ssh_public_key: str,
-    jump_pod_port: int,
-):
+    jump_pod_port: Optional[int],
+) -> int:
     try:
-        api.read_namespaced_service(
+        service = api.read_namespaced_service(
             name=_get_jump_pod_service_name(project_name),
             namespace=DEFAULT_NAMESPACE,
         )
     except client.ApiException as e:
         if e.status == 404:
-            _create_jump_pod_service(
+            service = _create_jump_pod_service(
                 api=api,
                 project_name=project_name,
                 project_ssh_public_key=project_ssh_public_key,
                 jump_pod_port=jump_pod_port,
             )
         else:
             raise
+    return service.spec.ports[0].node_port
 
 
 def _create_jump_pod_service(
     api: client.CoreV1Api,
     project_name: str,
     project_ssh_public_key: str,
-    jump_pod_port: int,
-):
+    jump_pod_port: Optional[int],
+) -> client.V1Service:
     # TODO use restricted ssh-forwarding-only user for jump pod instead of root.
     commands = _get_jump_pod_commands(authorized_keys=[project_ssh_public_key])
     pod_name = _get_jump_pod_name(project_name)
@@ -385,7 +395,7 @@ def _create_jump_pod_service(
             ),
         ),
     )
-    api.create_namespaced_service(
+    service_response = api.create_namespaced_service(
         namespace=DEFAULT_NAMESPACE,
         body=client.V1Service(
             metadata=client.V1ObjectMeta(name=_get_jump_pod_service_name(project_name)),
@@ -402,6 +412,7 @@ def _create_jump_pod_service(
             ),
         ),
     )
+    return service_response
 
 
 def _get_jump_pod_commands(authorized_keys: List[str]) -> List[str]:

src/dstack/_internal/core/backends/kubernetes/client.py → src/dstack/_internal/core/backends/kubernetes/utils.py

@@ -1,4 +1,4 @@
-from typing import Dict, List
+from typing import Dict, List, Optional
 
 import kubernetes
 import yaml
@@ -14,7 +14,20 @@ def get_api_from_config_dict(kubeconfig: Dict) -> kubernetes.client.CoreV1Api:
     return kubernetes.client.CoreV1Api(api_client=api_client)
 
 
+def get_cluster_public_ip(api_client: kubernetes.client.CoreV1Api) -> Optional[str]:
+    """
+    Returns public IP of any cluster node.
+    """
+    public_ips = get_cluster_public_ips(api_client)
+    if len(public_ips) == 0:
+        return None
+    return public_ips[0]
+
+
 def get_cluster_public_ips(api_client: kubernetes.client.CoreV1Api) -> List[str]:
+    """
+    Returns public IPs of all cluster nodes.
+    """
     public_ips = []
     for node in api_client.list_node().items:
         addresses = node.status.addresses

src/dstack/_internal/core/models/backends/kubernetes.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 
 from pydantic import BaseModel
 from typing_extensions import Literal
@@ -7,8 +7,8 @@
 
 
 class KubernetesNetworkingConfig(ForbidExtra):
-    ssh_host: str
-    ssh_port: int
+    ssh_host: Optional[str]
+    ssh_port: Optional[int]
 
 
 class KubernetesConfigInfo(BaseModel):

src/dstack/_internal/server/services/backends/configurators/kubernetes.py

@@ -1,8 +1,8 @@
 from typing import List
 
 from dstack._internal.core.backends.kubernetes import KubernetesBackend
-from dstack._internal.core.backends.kubernetes.client import get_api_from_config_data
 from dstack._internal.core.backends.kubernetes.config import KubernetesConfig
+from dstack._internal.core.backends.kubernetes.utils import get_api_from_config_data
 from dstack._internal.core.models.backends.base import BackendType
 from dstack._internal.core.models.backends.kubernetes import (
     AnyKubernetesConfigInfo,

src/dstack/_internal/server/services/config.py

@@ -82,8 +82,8 @@ def fill_data(cls, values):
 
 class KubernetesConfig(ForbidExtra):
     type: Literal["kubernetes"] = "kubernetes"
-    networking: KubernetesNetworkingConfig
     kubeconfig: KubeconfigConfig
+    networking: Optional[KubernetesNetworkingConfig]
 
 
 class LambdaConfig(ForbidExtra):
@@ -293,7 +293,12 @@ class _ConfigInfoWithCreds(BaseModel):
 
 
 def _config_to_internal_config(backend_config: BackendConfig) -> AnyConfigInfoWithCreds:
-    config_info = _ConfigInfoWithCreds.parse_obj(backend_config.dict())
+    backend_config_dict = backend_config.dict()
+    # Allow to not specify networking
+    if backend_config.type == "kubernetes":
+        if backend_config.networking is None:
+            backend_config_dict["networking"] = {}
+    config_info = _ConfigInfoWithCreds.parse_obj(backend_config_dict)
     if backend_config.type == "azure":
         config_info.__root__.locations = backend_config.regions
     return config_info.__root__