Eliminate the need to map nested kernel IDs to outer kernel IDs

Author: ojarjur

jupyter_server/services/kernels/routing.py

@@ -216,8 +216,6 @@ async def prepare(self):
 
 
 class RoutingKernelManager(ServerKernelManager):
-    kernel_id_map: dict[str, str] = {}
-
     @property
     def is_remote(self):
         if not self.kernel_name or not self.kernel_id:
@@ -234,8 +232,7 @@ def wrapped_multi_kernel_manager(self):
     def wrapped_kernel_manager(self):
         if not self.kernel_id:
             return None
-        wrapped_kernel_id = RoutingKernelManager.kernel_id_map.get(self.kernel_id, self.kernel_id)
-        return self.wrapped_multi_kernel_manager.get_kernel(wrapped_kernel_id)
+        return self.wrapped_multi_kernel_manager.get_kernel(self.kernel_id)
 
     @default("websocket_connection_class")
     def _default_websocket_connection_class(self):
@@ -259,44 +256,29 @@ def client(self, *args, **kwargs):
 
     @in_pending_state
     async def start_kernel(self, *args, **kwargs):
-        kernel_id: t.Optional[str] = kwargs.pop("kernel_id", self.kernel_id)
-        if kernel_id:
-            self.kernel_id = kernel_id
-
         km = self.wrapped_multi_kernel_manager
+        if isinstance(km, GatewayMappingKernelManager):
+            kwargs.pop("kernel_id")
         wrapped_kernel_id: str = await ensure_async(
             km.start_kernel(kernel_name=self.kernel_name, **kwargs)
         )
-        self.kernel_id = self.kernel_id or wrapped_kernel_id
-        RoutingKernelManager.kernel_id_map[self.kernel_id] = wrapped_kernel_id
-        self.log.debug(
-            f"Created kernel {self.kernel_id} corresponding to {wrapped_kernel_id} in {km}"
-        )
-        self.log.debug(RoutingKernelManager.kernel_id_map)
+        self.kernel_id = wrapped_kernel_id
+        self.log.debug(f"Created kernel {self.kernel_id} in {km}")
 
     async def shutdown_kernel(self, now=False, restart=False):
-        wrapped_kernel_id = RoutingKernelManager.kernel_id_map.get(self.kernel_id, self.kernel_id)
         km = self.wrapped_multi_kernel_manager
-        await ensure_async(km.shutdown_kernel(wrapped_kernel_id, now=now, restart=restart))
-        RoutingKernelManager.kernel_id_map.pop(self.kernel_id, None)
+        await ensure_async(km.shutdown_kernel(self.kernel_id, now=now, restart=restart))
 
     async def restart_kernel(self, now=False):
-        wrapped_kernel_id = RoutingKernelManager.kernel_id_map.get(self.kernel_id, self.kernel_id)
         km = self.wrapped_multi_kernel_manager
-        return await ensure_async(km.restart_kernel(wrapped_kernel_id, now=now))
+        return await ensure_async(km.restart_kernel(self.kernel_id, now=now))
 
     async def interrupt_kernel(self):
         km = self.wrapped_kernel_manager
         return await ensure_async(km.interrupt_kernel())
 
     async def model(self):
-        wrapped_kernel_id = RoutingKernelManager.kernel_id_map.get(self.kernel_id, self.kernel_id)
-        wrapped_model = await ensure_async(
-            self.wrapped_multi_kernel_manager.kernel_model(wrapped_kernel_id)
-        )
-        model = copy.deepcopy(wrapped_model)
-        model["id"] = self.kernel_id
-        return model
+        return await ensure_async(self.wrapped_multi_kernel_manager.kernel_model(self.kernel_id))
 
 
 class RoutingMappingKernelManager(AsyncMappingKernelManager):
@@ -339,6 +321,12 @@ def has_remote_kernels(self):
                 return True
         return False
 
+    def pre_start_kernel(self, *args, **kwargs):
+        km, kernel_name, kernel_id = super().pre_start_kernel(*args, **kwargs)
+        if km and km.kernel_id and km.kernel_id != kernel_id:
+            kernel_id = km.kernel_id
+        return km, kernel_name, kernel_id
+
     async def list_kernels(self):
         if self.has_remote_kernels():
             # We have remote kernels, so we must call `list_kernels` on the