Execute Jupyter notebook as a pipeline step | by Bart Grasza | Towards Data Science

This article introduces the python kf-notebook-component project which allows the execution of Jupyter Notebook as a separate step of a Kubeflow pipeline.
Sidenote: yes, I’m aware that Airflow has Papermill operator, but please bear with me to see why I think my solution is preferable.
I think everyone agrees that Jupyter notebook is a great tool for data wrangling. Kubeflow has an inbuilt “Jupyter Notebook Servers” module that allows you to do your research directly in the notebooks. When your experimentation phase is done and you want to take it to production, you can wrap your code in the pipeline steps and execute it directly from notebook in “Kubeflow Pipelines” module.
While this approach allows you to move pretty fast, I found that often wrapping code (sometimes hundreds of lines) into a single function that is later transformed into a single Kubeflow component makes the code messy. Also, during “wrapping” you need to remove e.g. plots that you used to experiment with data. Later when you come back and want to make some changes, you’d often need to partially “unwrap” your function and write this “experimentation” code again.
One day I wished I could just execute my entire notebook as a separate step of the pipeline. That’s how the project kf-notebook-component was born.
The idea is simple. Kubeflow notebook as a pipeline step should:
Have a quick and easy way of injecting data (simple variables and/or entire data objects) into the notebook.
Allow you to easily pass data (variables/data objects) produced/transformed inside notebook to the next step of the pipeline, no matter if it’s another Jupyter notebook or any other Kubeflow pipeline component.
Allow the display of notebook outputs for quality control and reporting when executed in Kubeflow Pipeline (e.g. production).
After your code is wrapped in the pipeline, very often you will want to make some changes to it. The notebook as a pipeline step should allow you to make changes (e.g. with sample data in your development environment) and after you’re done, submit the modified version of the pipeline without any other changes to the pipeline itself.
It should avoid having notebook ( ipynb) files being put into version control (if you ever tried to git merge changes made to your notebook then you know what I’m talking about) so that you can easily see changes made to each notebook.
Allow using other Kubeflow “goodies” like e.g. metrics.
What’s the result? This is how you define a pipeline using kf-notebook-component (only a single step to keep it short):
Of course, you can define any number of inputs and outputs.
When executed successfully (depending on your notebook code) you should see something like:
At that point you probably have the following questions:
Why does the notebook in the code snippet above have .pyextension, not .ipynb?How is INPUT data being “injected” into notebook and when to use add_input_param vs add_input_artifact?How to get OUTPUTs from the notebook and when to use add_output_param vs add_output_artifact?
Why does the notebook in the code snippet above have .pyextension, not .ipynb?
How is INPUT data being “injected” into notebook and when to use add_input_param vs add_input_artifact?
How to get OUTPUTs from the notebook and when to use add_output_param vs add_output_artifact?
Before answering them, let’s have a quick peek inside the kf-notebook-component project. The two most important external modules used in it are Papermill and Jupytext:
Papermill is responsible for notebook execution and injecting variables into the notebook (but not artifacts, I’ll explain why soon).
Jupytext can transform a notebook’s .ipynb files into other formats (and vice versa).
Papermill injects (overwrites) variables by appending them after the cell with tag “parameters” tag:
The kf-notebook-component requires notebooks to be in a specific Jupytext format called “Notebook paired with percent script”. You can learn more about it here but simply put, to generate a “percent script” follow these steps:
Install Jupytext as a Jupyter notebook extension
Create a new notebook and pair it with a percent script. This will generate a new .py file.
Put your .ipynb file in gitignore so that only the .py is going to be checked in your version control
Now you can even delete the .ipynb file and open .py in your notebook. It will look the same as a regular .ipynb file.
1.Why does the notebook in code snippet above have .pyextension, not .ipynb?You should already know the answer. It’s a “percent script” file type! When the pipeline step is run inside Kubeflow, your .py file will be transformed back into .ipynbso that Papermil could execute your notebook.The percent script is used to avoid pushing .ipynbfiles into your source control (it’s responsible for other things as well but I’ll skip the details).
2.How is INPUT data being “injected” into notebook and when to use add_input_param vs add_input_artifact?Papermill allows to “inject” only simple variables like int, stror float. Anything more complex (even list or dictionary) should be passed with artifacts.Luckily, injecting artifacts inside notebook works similarly to Papermill (and it’s handled by the kf-notebook-component). Specifically, it adds a cell (after the one tagged as “parameters”) with special code that loads your artifact data into a variable:
Input artifacts can handle both pickled binaries and regular text files.
3.How to get OUTPUTs from the notebook and when to use add_output_param vs add_output_artifact?
Output variables/artifacts are also handled by the kf-notebook-component. Similarly to inputs, use int, str, etc. with add_output_param and anything else with add_output_artifact .All you have to do is to map your notebook variables to pipeline outputs:
Special 3 dictionaries are outputs, artifactsand metrics.
What are the metrics? They are an awesome feature in Kubeflow. They allow you to output any numeric value as a metric. Later you can compare your pipeline runs:
Remember: once your code gets executed in a Kubeflow pipeline, don’t be surprised to see additional code cells that were injected by Papermill/kf-notebook-component to map your inputs/outputs to the “outside” of the notebook.
Define your component inputs/outputs and point at the notebook file it should execute (use percent format!).
Make sure your docker image can handle your notebook code (details below)
Make sure your notebook has a cell with tag parameters containing variables that will be overwritten when run with Kubeflow.
Whatever data you want to pass to the next step in the pipeline, should be put into outputs and/or artifacts dictionaries.
Base docker imagesTo keep things simple, in the code snippet above, I cheated a bit when I wrote:
# Define Docker image used to run your componentbase_image = 'tensorflow/tensorflow-1.14.0-py3:latest'
This would not work because this image doesn’t have papermill , jupytext and nbconvert (converts notebook to HTML) installed. Any Docker image that you want to use to run your notebook as a pipeline step should have them pip installed.
You might find it useful to display your notebook in a “report style” where all input cells are hidden. Simply, set remove_nb_inputsto True :
nb = NbComponentBuilder(    op_name='Task 1',     inject_notebook_path='notebook_code.py',     remove_nb_inputs=True)
Thanks for reading! Head out to the github repo to see the full pipeline code example.