No public description

PiperOrigin-RevId: 541993402

google/colab/_serverextension/_resource_monitor.py

@@ -128,8 +128,17 @@ def get_ram_usage(kernel_manager):
     # is better performed in the frontend presentation layer. 2) was only a
     # requirement for the split (KMC/K) container, a feature that was dropped
     # (cl/470476143).
+    def get_pid(kernel):
+      # TODO(b/264409633): Eliminate this function after migration to
+      # jupyter-client 7.x is complete.
+      try:
+        pid = kernel.provisioner.pid
+      except:  # pylint: disable=bare-except
+        pid = kernel.kernel.pid
+      return str(pid)
+
     pids_to_kernel_ids = {
-        str(kernel_manager.get_kernel(kernel_id).kernel.pid): kernel_id
+        get_pid(kernel_manager.get_kernel(kernel_id)): kernel_id
         for kernel_id in kernel_manager.list_kernel_ids()
     }
 

google/colab/_shell.py

@@ -26,6 +26,7 @@
 from ipykernel import jsonutil
 from ipykernel import zmqshell
 from IPython.core import alias
+from IPython.core import compilerop
 from IPython.core import inputsplitter
 from IPython.core import interactiveshell
 from IPython.core import oinspect
@@ -52,6 +53,16 @@ def _show_pip_warning():
 class Shell(zmqshell.ZMQInteractiveShell):
   """Shell with additional Colab-specific features."""
 
+  def init_instance_attrs(self):
+    """Initializes instance attrs.
+
+    Since https://github.com/ipython/ipykernel/pull/597 introduces filename
+    changes for source files which we parse to get execution count, we use the
+    original CachingCompiler instead of the custom XCachingCompiler.
+    """
+    super().init_instance_attrs()
+    self.compile = compilerop.CachingCompiler()
+
   def init_events(self):
     self.events = EventManager(self, available_events)
     self.events.register('pre_execute', self._clear_warning_registry)

setup.py

@@ -20,7 +20,7 @@
     # Note: these dependency versions should be kept in-sync with the versions
     # specified in the docker container requirements files.
     'google-auth==2.17.3',
-    'ipykernel==5.5.6',
+    'ipykernel==6.21.3',
     'ipython==7.34.0',
     'notebook==6.5.5',
     'pandas==1.5.3',