Is there a way to register a message handler callback on a google.colab.kernel.comms object?

Question

The example in advanced_outputs.ipynb sending data from the kernel to the client is almost exactly what I need. And that example references the Jupyter comms documentation. But am I right that currently, although I can open a new connection to the client from the kernel with data attached, I cannot yet send multiple messages over the same connection?

The workflow that I think I want is e.g., in javascript:

  (async () => {
    google.colab.kernel.comms.registerTarget('example_comms', (comm, message) => {
      document.body.appendChild(document.createTextNode('comm opened.'))
      comm.send('comm opened');
      console.log(comm)
      comm.on_msg(function(msg) {
        var p = document.createElement("p")
        var t = document.createTextNode(message.data.foo)
        p.appendChild(t)
        document.body.appendChild(p)
      })
    });
  })()

and in python

channel = comm.Comm(target_name='example_comms')
def handle_message(msg):
  print(f"python received message: {msg['content']['data']}")
channel.on_msg(handle_message)
for i in range(10):
  channel.send(data={'foo': i})

But it seems that google.colab.kernel.comms.Comm does not export on_msg, nor any other way that I can see to send multiple messages over the same channel? The ipykernel.comm.Comm object has a send method, but I don't think it can be used at present?

Here is a notebook with minimal modifications from the advanced_outputs example. Running it throws an error in the console that on_msg is not defined.

I suspect (but admittedly have not yet measured the performance to confirm) that making a new connection for every message is overhead that I want to avoid. Any help is appreciated!

Answer 1

I think this is what you are looking for:

from IPython.display import Javascript

comm_ = None
def target_func(comm, msg):
  global comm_
  comm_ = comm  # The connection you are looking for !

  ## Register handler for later messages
  @comm.on_msg
  def _recv(msg):
    # Use msg['content']['data'] for the data in the message
    comm.send({'echo': msg['content']['data']})
        
  @comm.on_close
  def _close(msg):
    global comm_
    comm_ = None
get_ipython().kernel.comm_manager.register_target('comm_target', target_func)

Javascript('''
(async () => {
  const channel = await google.colab.kernel.comms.open('comm_target');

  (async function() {
    for await (const message of channel.messages) {
      console.log(message.data)
    }
  })();
})()
''')

Contrary to what you are doing, it is establishing a comm channel from client to kernel, which is way more natural in my opinion. Then at this point you get a persistent connection comm_ that you can use to send messages to the javascript frontend:

comm_.send(data="test")

It is also possible to get the list of every open connection doing:

get_ipython().kernel.comm_manager.comms

So it is possible to get rid of the ugly global var and monitor instead changed of this attribute. It could also be used to wait for a connection to be established for instance.

Note that in the case of Jupyter Notebook the Python code is the same, but not the Javascript one. Here is the corresponding implementation:

const channel = await Jupyter.notebook.kernel.comm_manager.new_comm('comm_target')

channel.on_msg(function(message) {
    console.log(message.content.data);
});