InfinStor Logo

MLflow Kernel

Capture the history of JupyterLab cell execution in MLflow.

Why use MLflow Kernel?

InfinStor's MLflow Kernel is an essential tool for MLOps.

Essential for MLOps

Captures history of the cell in the notebook, including visualizations

A Better Alternative

Vastly superior way to record development activity compared to Git.

Next Generation Technology

Current methods of saving notebooks are inadequate in this day and age.

How does MLflow Kernel work?

MLflow Kernel connects to an MLflow service and records all the data science activities in the notebook.

  • It captures the cell code, the cell output and the visualizations, and records them as MLflow artifacts.

  • Each cell execution is captured as a separate run.

    • This happens transparently, and the user doesn't have to manage the lifecycle of a run.

  • All the runs of these cells are grouped inside a parent run that stays alive until the life of the kernel.

    • A new parent run is created when the kernel is restarted.

  • The kernel works seamlessly with the autologging feature which is supported by many ML libraries.

    • The user just needs to enable autologging and everything is recorded as part of cell execution run.

Deploy MLflow Kernel

MLflow Kernel can be deployed in own JupyterLab or Sagemaker Studio. For any questions, we provide free email support at

MLflow Kernel for JupyterLab

Following these simple steps will get the MLflow Kernel up and connected to the Infinstor’s MLflow service.

  1. Install all required packages and setup Jupyter kernelspecs.

    • pip install mlflow_kernel

  2. Install MLflow Kernel in the Jupyter kernelspecs.

    • python -m mlflow_kernel.install

  3. Create a config file in the directory where the notebook is started.

    • ~/.jupyter/mlflow_kernel_config.json

    • Copy the content of the config as shown in the right.

    • The "mlflow_tracking_uri" points to the MLflow server.

      • For example, if you are using InfinStor’s MLflow, it would be infinstor://

    • If "default_experiment_name" is not specified, the runs are created in 'default' experiment with id 0. The user should be authorized for the specified experiment or for the 'default' experiment.

"mlflow_tracking_uri": [mlflow tracking url],
"debug_enabled" : "false",
"default_experiment_name": [experiment name]

MLflow Kernel for Sagemaker Studio

MLFlow Kernel can be used with Sagemaker Studio using Infinstor’s pre-built image, which is preconfigured to connect to Infinstor’s MLFlow service.

  1. Set up Sagemaker Studio with the kernel image URI on the right.

  2. The role attached to the kernel must have access to the artifact bucket.

    • The default artifact bucket can be found in the user's profile on the InfinStor's service dashboard.

  3. First time notebook access on the notebook will prompt MLflow login.

    • Require MLflow login and password. Login is required for the first time only, and afterwards tokens are automatically refreshed.

    • The kernel must be restarted after the login.