feat: Allow self-ssh in single-node single-container compute sessions (…

…#2032)
lablup · Apr 18, 2024 · 5d69e72 · 5d69e72
1 parent 0176e13
commit 5d69e72
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 6 deletions.
diff --git a/changes/2032.feature.md b/changes/2032.feature.md
@@ -0,0 +1 @@
+Allow self-ssh in single-node single-container compute sessions.
diff --git a/src/ai/backend/common/types.py b/src/ai/backend/common/types.py
@@ -1055,7 +1055,7 @@ class ClusterInfo(TypedDict):
     size: int
     replicas: Mapping[str, int]  # per-role kernel counts
     network_name: Optional[str]
-    ssh_keypair: Optional[ClusterSSHKeyPair]
+    ssh_keypair: ClusterSSHKeyPair
     cluster_ssh_port_mapping: Optional[ClusterSSHPortMapping]
 
 

diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py
@@ -1444,11 +1444,7 @@ async def start_session(
             size=scheduled_session.cluster_size,
             replicas=replicas,
             network_name=network_name,
-            ssh_keypair=(
-                await self.create_cluster_ssh_keypair()
-                if scheduled_session.cluster_size > 1
-                else None
-            ),
+            ssh_keypair=await self.create_cluster_ssh_keypair(),
             cluster_ssh_port_mapping=cast(
                 Optional[ClusterSSHPortMapping], cluster_ssh_port_mapping
             ),