From Notebook to Kubernetes: Scaling GenAI Pipelines with ZenML

Alex Strick van Linschoten, ZenML12:42 · Oct 2024 · 256 views
Thumbnail for From Notebook to Kubernetes: Scaling GenAI Pipelines with ZenML Watch on YouTube
TL;DR
  1. 1

    ZenML lets users move the same Python pipeline from a local stack to Kubernetes by changing the selected stack.

  2. 2

    A notebook workflow can handle annotation, augmentation, model training, inference, and artifact tracking while ZenML versions the code and outputs.

  3. 3

    ZenML caching avoids rerunning completed work, and the same pipeline can move to another backend such as Vertex AI without changing the pipeline code.

Summary

Alex Strick van Linschoten demonstrates a GenAI workflow that starts in a notebook and moves intensive work to Kubernetes through ZenML. He annotates cat images, augments them, copies them to the cloud, fine-tunes Flux and Stable Diffusion models, and runs inference. The notebook remains the place for experimentation, while ZenML tracks code, artifacts, configuration, packages, logs, models, prompts, and generated images. ZenML stacks let him switch between local execution and a Kubernetes machine with a GPU. Caching also lets him inspect earlier runs without repeating expensive work. Alex's main argument is practical: data scientists can begin with familiar notebook code, then extract code into Python modules and run it on production infrastructure when needed. Moving from Kubernetes to another backend such as Vertex AI requires changing the stack rather than rewriting the pipeline.

Key ideas
00:54

ZenML organizes reproducible workflows around stacks and pipelines

Alex describes ZenML as an open-source framework for getting machine learning workflows into production. A stack is a collection of tools and components, while a pipeline is composed of steps, which are ordinary Python functions. ZenML is designed to work with many tools and cloud environments. Its orchestration tracks runs so teams can reproduce work, inspect earlier results, and return to a previous run when someone needs to review it. The aim is to reduce the cost of changing technology stacks while keeping workflows organized across a team.

02:07

The same pipeline can move from a local machine to Kubernetes

A local stack runs the pipeline on the developer's machine. When a step needs a GPU, more memory, or more storage, Alex switches to a Kubernetes stack with a more powerful machine. He shows this change from inside the notebook and says the pipeline completed in 1 minute 23 seconds. The pipeline code stays the same. The execution environment changes through the selected stack, which gives the notebook a path to larger infrastructure without requiring a separate implementation.

03:25

Notebook-based annotation can feed a reproducible training workflow

Alex uses a small cat-image example to demonstrate the workflow. He annotates pictures of Bluepus and Arya in the notebook, keeps the Bluepus images, and saves the labels as JSON. He then moves the selected images into their own folder and augments them with brightness changes, noise, and sharpening. ZenML can integrate with different annotation tools, and the example shows how these early data steps can remain close to the experimentation process before the data is copied into the cloud.

05:19

Heavy GenAI training can run in Kubernetes while code stays familiar

The local machine is not powerful enough for fine-tuning the large Flux model, so Alex runs the training pipeline in Kubernetes. The training code comes from examples in the Diffusers library and can either remain in the notebook or be extracted into a separate train.py file. ZenML downloads and versions the code for the run, then records the stack, code, logs, configuration, Python packages, outputs, and trained model. This lets the training implementation stay ordinary Python while the expensive execution happens remotely.

06:51

Caching prevents completed model work from running again

When Alex revisits the training pipeline, ZenML recognizes that a cached model is available and avoids repeating the run. He explains that changing the code would invalidate the relevant cache and cause the work to run again. The same behavior appears in the inference example, where a previous result can be opened instead of launching another Kubernetes job. This is useful for a live notebook workflow because experimentation can continue without paying the time and compute cost for unchanged steps.

07:47

Model outputs and prompts can be inspected together

After training, the models are pushed to Hugging Face, where they can also be used for inference. Alex runs inference on Kubernetes and captures the prompts and generated images. The ZenML dashboard groups the outputs so he can inspect images of Bluepus in situations such as a data center, a server room, or wearing VR goggles. Stable Diffusion produces less convincing results than Flux in his comparison. A local Flux run then hits a CUDA out-of-memory error, which illustrates why the remote GPU stack matters for this workflow.

09:34

The notebook is an entry point, not the final boundary

Alex argues that users can keep the experimental feel of a notebook while gaining production features such as tracked runs, versioned artifacts, and dashboard inspection. He does not suggest that teams should keep everything in notebooks permanently. Once the design works, code can be moved into normal Python modules. The important part is that the same pipeline structure can continue into a production environment, with the stack controlling where it runs.

10:52

Changing infrastructure can mean changing the stack instead of the code

ZenML's abstraction allows the pipeline to move between execution environments. Alex says that switching from Kubernetes to Vertex AI would not require changing the pipeline code. The user would select a different stack. This keeps infrastructure choices separate from the Python steps that describe annotation, training, or inference, while still allowing teams to choose a backend that fits the workload.

"It's pretty easy like once you've designed your code and stuff within a notebook, it's pretty easy to pull stuff out then into like normal Python modules."11:55
Who should watch
  • You are building GenAI experiments in notebooks and need a straightforward path to GPU-backed execution.
  • Your team wants run history, code and artifact versioning, logs, configuration, and cached results without rewriting notebook code into a separate orchestration system.
  • You are comparing local execution with Kubernetes or another backend such as Vertex AI and want infrastructure changes to stay outside the pipeline logic.