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. KernelImages can exist only in workspace workspaces/global (global workspace).

import h2o_notebook

from h2o_notebook.clients.kernel_image.type import KernelImageType

clients = h2o_notebook.login()

# Print all python KernelImages from global workspace.
all_images = clients.kernel_image_client.list_all_kernel_images(parent="workspaces/global")
images = [i for i in all_images if i.kernel_image_type == KernelImageType.TYPE_PYTHON]
print(images)

# Get python KernelImage from global workspace.
python_kernel_image = clients.kernel_image_client.get_kernel_image(
    name="workspaces/global/kernelImages/python",
)
# python KernelImage name: workspaces/global/kernelImages/python
print(f"python KernelImage name: {python_kernel_image.name}")

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. KernelTemplates can exist only in workspace workspaces/default (default user workspace).

import h2o_notebook

from h2o_notebook.clients.kernel_template.kernel_template import KernelTemplate

clients = h2o_notebook.login()

# Create new KernelTemplate in global workspace.
my_template = clients.kernel_template_client.create_kernel_template(
    parent="workspaces/default",
    kernel_template=KernelTemplate(
        milli_cpu_limit=4000,  # 4 CPUs
        gpu=0,
        memory_bytes_limit="1Gi",
        max_idle_duration="1h",
        environmental_variables={"MY_ENV_VAR": "my-value"},
    ),
    kernel_template_id="my-template",
)
# my KernelTemplate name: workspaces/default/kernelTemplates/my-template
print(f"my KernelTemplate name: {my_template.name}")

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.

Kernels can exist only in workspace workspaces/default (default workspace).

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

from h2o_notebook.clients.kernel.kernel import Kernel

clients = h2o_notebook.login()

# Create new Kernel in default workspace.
kernel = clients.kernel_client.create_kernel(
    parent="workspaces/default",
    kernel=Kernel(
        kernel_image="workspaces/global/kernelImages/python",
        kernel_template="workspaces/global/kernelTemplates/my-template",
        display_name="My First Kernel",  # Optional
        environmental_variables={"ANOTHER_ENV_VAR": "my-value"},  # Optional
    ),
)
# New Kernel created with name "workspaces/default/kernels/{automatically-generated-kernel-id}".
print(kernel)

clients.kernel_client.terminate_kernel(name=kernel.name)
clients.kernel_client.delete_kernel(name=kernel.name)

# Helper function to wait for the Kernel to free up resources
print("waiting for kernel to be deleted...")
clients.kernel_client.wait_kernel_deleted(name=kernel.name)
print("kernel deleted")

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

from h2o_notebook.clients.kernel.kernel import Kernel
from h2o_notebook.clients.kernel_task.kernel_task import KernelTask

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    parent="workspaces/default",
    kernel=Kernel(
        kernel_image="workspaces/global/kernelImages/python",
        kernel_template="workspaces/global/kernelTemplates/my-template",
    ),
)

task = clients.kernel_task_client.create_kernel_task(
    parent=kernel.name,
    kernel_task=KernelTask(
        code="""
        import time
        print('Start')
        time.sleep(5)
        print('End')
        """
    ),
)
print(f"task created:\n{task}")

print("waiting for task to complete...")
task = clients.kernel_task_client.wait_task_completed(
    name=task.name,
)
print(f"task completed:\n{task}")

clients.kernel_client.delete_kernel(name=kernel.name)
print("waiting for kernel to be deleted...")
clients.kernel_client.wait_kernel_deleted(name=kernel.name)
print("kernel deleted")

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

from h2o_notebook.clients.kernel.kernel import Kernel
from h2o_notebook.clients.kernel_task.kernel_task import KernelTask

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    parent="workspaces/default",
    kernel=Kernel(
        kernel_image="workspaces/global/kernelImages/python",
        kernel_template="workspaces/global/kernelTemplates/my-template",
    ),
)

task1 = clients.kernel_task_client.create_kernel_task(
    parent=kernel.name,
    kernel_task=KernelTask(
        code="""
        import time
        time.sleep(60)
        """
    ),
)

task2 = clients.kernel_task_client.create_kernel_task(
    parent=kernel.name,
    kernel_task=KernelTask(
        code="""
        import time
        time.sleep(1)
        """,
    ),
)

print("cancel task1")
clients.kernel_task_client.cancel_kernel_task(name=task1.name)

print("waiting for task2 to complete...")
task2 = clients.kernel_task_client.wait_task_completed(name=task2.name)
print("task2 completed")

print(task2)

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

print("waiting for kernel to be deleted...")
clients.kernel_client.wait_kernel_deleted(name=kernel.name)
print("kernel deleted")

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.

If you want the messages to be aggregated in one output, you can use specialized function list

import h2o_notebook
from h2o_notebook.clients.kernel.kernel import Kernel
from h2o_notebook.clients.kernel_task.kernel_task import KernelTask

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    parent="workspaces/default",
    kernel=Kernel(
        kernel_image="workspaces/global/kernelImages/python",
        kernel_template="workspaces/global/kernelTemplates/my-template",
    ),
)

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

print("waiting for task to complete...")
task = clients.kernel_task_client.wait_task_completed(name=task.name)
print("task completed")

messages = clients.kernel_task_client.list_all_kernel_task_messages(kernel_task_name=task.name)
aggregated_output = clients.kernel_task_client.get_kernel_task_output(kernel_task_name=task.name)

print(f"task:\n{task}")
print(f"single messages:\n{messages}")
print(f"aggregated output:\n{aggregated_output}")

clients.kernel_client.delete_kernel(name=kernel.name)
print("waiting for kernel to be deleted...")
clients.kernel_client.wait_kernel_deleted(name=kernel.name)
print("kernel deleted")

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.

The same display output is also available in the aggregated KernelTaskOutput in its attachments.

import base64
import os

import h2o_notebook
from h2o_notebook.clients.kernel.kernel import Kernel
from h2o_notebook.clients.kernel_task.kernel_task import KernelTask

clients = h2o_notebook.login()

kernel = clients.kernel_client.create_kernel(
    parent="workspaces/default",
    kernel=Kernel(
        kernel_image="workspaces/global/kernelImages/python",
        kernel_template="workspaces/global/kernelTemplates/my-template",
    ),
)

task = clients.kernel_task_client.create_kernel_task(
    parent=kernel.name,
    kernel_task=KernelTask(
        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()
        """
    ),
)

print("waiting for task to complete...")
task = clients.kernel_task_client.wait_task_completed(name=task.name)
print("task completed")

messages = clients.kernel_task_client.list_all_kernel_task_messages(kernel_task_name=task.name)
text_data = messages[0].display_message.display_message_data["text/plain"]
print(f"message text data: {text_data}")
png_data = messages[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")

# Aggregated output contains display data in 'attachments'.
# Attachments data is not guaranteed to be ordered by mimeType.
output = clients.kernel_task_client.get_kernel_task_output(kernel_task_name=task.name)
for data in output.attachments[0].data:
    if data.mime_type == "text/plain":
        print(f"output text data: {data.data}")
    elif data.mime_type == "image/png":
        png_data = messages[0].display_message.display_message_data["image/png"]
        with open("plot2.png", "wb") as f:
            f.write(base64.b64decode(png_data))
        os.system("open plot2.png")

clients.kernel_client.delete_kernel(name=kernel.name)
print("waiting for kernel to be deleted...")
clients.kernel_client.wait_kernel_deleted(name=kernel.name)
print("kernel deleted")
Feedback