Skip to main content

Kernel API

The Kernel API allows you to create headless Kernels in the H2O AI Cloud. This is useful for running long-running tasks, such as training models, without needing to keep a notebook open. The API allows you to create Kernels with defined resources, queue code to run on the Kernel, and retrieve the results.

Retrieving Kernel Images​

The administrator has the ability to create Kernel Images which are the runtime environments for the Kernels. Kernel Images contain a specific version of Python and a set of Python packages.

You can retrieve the list of available Kernel Images using the list_all_kernel_images function. Note that only Python Kernel Images are supported in the Kernel API at this time.

import h2o_notebook

from h2o_notebook.clients.kernel_image.type import KernelImageType

clients = h2o_notebook.login()

all_images = clients.kernel_image_client.list_all_kernel_images()
images = [i for i in all_images if i.kernel_image_type == KernelImageType.TYPE_PYTHON]
print(images)

Defining Kernel Templates​

Next you will need to define the compute for the Kernel. This can be done through the Kernel Template resource. A Kernel Template can and should be reused for multiple Kernels.

Kernel Templates define the resources available to the Kernel, such as CPU, memory, and GPU. Additionally, you can specify environmental variables or override the Kubernetes Pod spec to add tolerations, mount volumes, etc.

import h2o_notebook

clients = h2o_notebook.login()

my_template = clients.kernel_template_client.create_kernel_template(
    kernel_template_id="my-template-3",
    milli_cpu_limit=4000, # 4 CPUs
    gpu=0,
    memory_bytes_limit="1Gi",
    max_idle_duration="1h",
    environmental_variables={"MY_ENV_VAR": "my-value"},
)

Creating Kernel​

With Kernel Images and Kernel Templates defined, you can now create a Kernel. It can take some time for Kernel to finish starting up, but you may start sending code to the Kernel immediately.

Terminating a Kernel will stop the execution of any code running on the Kernel and free up the resources but you may still access the output. Deleting a Kernel will also remove any output associated with the Kernel.

import h2o_notebook

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    display_name="My First Kernel", # Optional
    kernel_image_id="python",
    kernel_template_id="my-template",
    environmental_variables={"ANOTHER_ENV_VAR": "my-value"},
)

print(kernel)

clients.kernel_client.terminate_kernel(kernel_id=kernel.kernel_id) # TODO: Not implemented
clients.kernel_client.delete_kernel(kernel_id=kernel.kernel_id)

# Helper function to wait for the Kernel to free up resources
clients.kernel_client.wait_kernel_deleted(kernel_id=kernel.kernel_id)

Idle Kernels will be automatically terminated after the configured max_idle_duration in the Kernel Template.

Submitting Kernel Tasks​

You can submit code to run on the Kernel using the Kernel Task resource. The Task will be queued and run in the order it was submitted.

import h2o_notebook

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    kernel_image_id="python",
    kernel_template_id="my-template",
)

task = clients.kernel_task_client.create_kernel_task(
    kernel_id=kernel.kernel_id,
    code="""
        import time
        print('Start')
        time.sleep(5)
        print('End')
        """
)

print(task)

task = clients.kernel_task_client.wait_task_completed(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task.kernel_task_id,
)

print(task)

clients.kernel_client.delete_kernel(kernel_id=kernel.kernel_id)

Cancelling a Kernel Task​

You can cancel a Kernel Task that is currently running. The Task will be cancelled but the output will be preserved.

import h2o_notebook

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    kernel_image_id="python",
    kernel_template_id="my-template",
)

task1 = clients.kernel_task_client.create_kernel_task(
    kernel_id=kernel.kernel_id,
    code="""
        import time
        time.sleep(60)
        """
)

task2 = clients.kernel_task_client.create_kernel_task(
    kernel_id=kernel.kernel_id,
    code="""
        import time
        time.sleep(1)
        """
)

clients.kernel_task_client.cancel_kernel_task(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task1.kernel_task_id,
)

task2 = clients.kernel_task_client.wait_task_completed(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task2.kernel_task_id,
)

print(task2)

clients.kernel_client.delete_kernel(kernel_id=kernel.kernel_id)

Retrieving Kernel Task standard output​

Kernel Tasks produce output that can be retrieved using the Kernel Task Message resource. The output can be either standard output or standard error in std_out and std_err fields respectively.

The output is received over time in chunks (messages) and can be retrieved using the list_all_kernel_task_messages function. You can get the messages while the Kernel Task is running or after it has completed. You may need to aggregate the messages to get the full output.

import h2o_notebook

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    kernel_image_id="python",
    kernel_template_id="my-template",
)

task = clients.kernel_task_client.create_kernel_task(
    kernel_id=kernel.kernel_id,
    code="""
        import time
        for i in range(10):
            print('Hello for the {} time!'.format(i))
            time.sleep(1)
        """
)

task = clients.kernel_task_client.wait_task_completed(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task.kernel_task_id,
)

output = clients.kernel_task_client.list_all_kernel_task_messages(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task.kernel_task_id,
)

print(task)
print(output)

clients.kernel_client.delete_kernel(kernel_id=kernel.kernel_id)

Retrieving Kernel Task display output​

Kernel Tasks can also produce display output. The display output is available in the display_messages field of the Kernel Task Message. It contains the MIME type and the data of the display output. There can be multiple representations of the same data.

import os
import base64
import h2o_notebook

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    kernel_image_id="python",
    kernel_template_id="my-template",
)

task = clients.kernel_task_client.create_kernel_task(
    kernel_id=kernel.kernel_id,
    code="""
        import matplotlib.pyplot as plt
        import numpy as np
        
        xaxis = np.array([2, 8])
        yaxis = np.array([4, 9])
        plt.plot(xaxis, yaxis)
        plt.show()
        """
)

task = clients.kernel_task_client.wait_task_completed(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task.kernel_task_id,
)

output = clients.kernel_task_client.list_all_kernel_task_messages(
    kernel_id=kernel.kernel_id,
    kernel_task_id=task.kernel_task_id,
)

print(output[0].display_message.display_message_data["text/plain"])

png_data = output[0].display_message.display_message_data["image/png"]
with open("plot.png", "wb") as f:
    f.write(base64.b64decode(png_data))

os.system("open plot.png")

clients.kernel_client.delete_kernel(kernel_id=kernel.kernel_id)
Feedback