Are you a data scientist or an AI engineer trying to bridge local development and scalable deployment of ML projects without the hassle? 🤔 As always, you are faced with the age-old dilemma — how can you move fast developing and testing new features while collaborating with your colleagues and sharing results easily? To add another complication, LLMs can add considerable costs and development time, lowering your project’s ROI if not managed correctly. How do we reconcile these challenges, you ask? ML pipelines, of course 🚀!

_{“To a man with a hammer, everything looks like a nail.” Well, in the MLOps world we don’t shy away from over-reliance — for a good reason — on pipelines 👷 Photo by Martti Salmi on Unsplash}

To increase collaboration, all ML experiments should run off the same underlying code and data sources, eliminating the common “works on my laptop” problem 🙈. A centralised ML platform, such as Vertex AI Pipelines on Google Cloud, provides a scalable environment to run ML jobs — with MLOps best-practice features like artifact versioning, data lineage, and experiment tracking out-of-the-box. These tools are vital for tweaking input parameters and recording resultant metrics effectively (e.g. “how does accuracy change with a range of temperature values?”).

A pipeline is an effective solution to an ML task, as you can break it into isolated steps like “get data,” “call LLM,” and “evaluate results”. Each step runs as its own isolated container, linked to other steps through inputs and outputs. This design enables checkpoints and caching — key features for rapid development cycles.

_{A simple pipeline which retrieves some data (e.g. prompts), makes a number of calls to an LLM and evaluates its outputs.}

What we have so far is a recipe for success. However, the final ask of your ML pipelines should be the ability to iterate, test, and build locally 💪. The entire experience should be available end-to-end on your laptop to shorten development cycles and get fast feedback when writing new code or creating pipelines. You want to avoid a situation where you have to wait minutes until your remote job fails due to a typo or an import error. Instead, local running can be used to ensure the whole process works locally, and then use a cloud environment to run multiple pipelines at scale and with confidence. This setup can be achieved with Docker Desktop and Kubeflow Pipelines (v2).

Local pipelines on your laptop

The end goal is to be able to trigger our pipeline with a single command:

and see an output like:

with our simple demo pipeline looking like this (the full pipeline code is available in this demo repository):

where our components (steps) are linked to each other via their inputs and outputs:

The “magic” that makes it work locally, is the kfp (Kubeflow Pipelines) local.DockerRunner() method:

We also need to have the kfp-component-local:latest container available on our system, which we prepare with this docker-compose-kfp-local.yml:

and this Dockerfile (the full Dockerfile is available here):

By using — mount=type=cache at each layer we are utilising our local Docker cache, such that subsequent builds are taking less than ~10 seconds (on Apple M3 Max) 🚀 We are also using uv to help us manage and lock our project’s dependencies effectively (think of uv as a modern and fast plug-and-play replacement for poetry).

We also have multi-stage builds, with the local stage only executing when calling this Dockerfile via our docker-compose (target: local field). We need this for injecting our user’s Application Default Credentials (ADC) JSON file into the container. When this file is available inside the container, running commands that need to authenticate with Google Cloud, will work just like running Python function locally. ⚠️ However, you should be vigilant to never push this container remotely (as it essentially has your user’s 24-hour authentication token). The way to pass ADC to Dockefile is via the — build-arg argument:

So, now you have a robust way of making pipeline changes (e.g. adding a new package to your pyproject.toml or creating a new component) and running the whole pipeline (including connecting to GCP):

1. Re-build the local container (if necessary) with docker compose -f ...
2. Trigger the pipeline with python main.py

And see the output locally within seconds 😎

_{Running the pipeline locally, and seeing the final output within seconds.}

Remote pipelines on Google Cloud

Now that we have a way to build, test and run our pipelines locally. When we run on Google Cloud’s Vertex AI platform, we can leverage remote cloud compute to run pipelines at scale and be able to share the results easily.

What we need now is a way to switch between local and remote pipeline execution at run time, which we can do by extending our original pipeline code to include:

Where we have a JobConstants class that holds destination details (project, zone) as we all as the base image value, allowing us to define which container is being used by our components in this decorator @component(base_image=job_constants.BASE_IMAGE) via config.cfg:

We can (automatically) update config.cfg before triggering the pipeline: for local execution, we set base_image to kfp-component-local:latest that has our ADC, while for remote we use kfp-component:USER-latest. We also re-use the previous Dockerfile without changes, but now using cloudbuild-vertex.yml to submit the job to Cloud Build to compile the container remotely (the full Cloud Build yaml file is available here):

We use BuildKit and cached Docker images on Cloud Build to speed up our build 🚀 There is no need to optimise Cloud Build resources further (i.e. using non-standard VM for builds) as any increase in build performance will be wiped out by cold starts, as the whole remote build is now under one minute due to caching 🔥

The full pipeline is available here, where we also have the following helper functions defined:

where get_date_time_nowand os.getenv(‘USER’) allow us to add a unique timestamp and user name for each pipeline. And typer helps us to control the execution with --local argument (think of this as a modern argparse for Python).

So now, all we gotta do is:

1. Re-build the remote container (if necessary) with

2. Run the pipeline

And see our pipeline executed remotely on Vertex AI 🤠:

Our demo pipeline on Vertex AI.

Tracking experiments

When working with LLMs you want to build flexibly, with evaluation tools in place, to be able to switch LLM providers at will — make your systems data-centric and model-agnostic. This gives you agency to control costs and improve your accuracy. To do that, you need to compare pipeline runs as part of a larger experiment. For example, an LLM becomes a parameter and accuracy becomes the target. In fact, you can populate metrics_output.metadata with any custom metrics relevant to your experiment

and then compare the results for multiple pipeline runs

_{Comparing accuracy for different models.}

This way, we track all input parameters, save all intermediate artefacts, and log all metrics, greatly empowering our experimentation capability 💪.

Afterword

I hope you found this end-to-end implementation of local and remote Kubeflow pipelines with Vertex AI helpful. To make it even more accessible, I have made the full solution available as a self-contained repository 🚀 . Along with it, I have included a variety of infrastructure tools: GitHub CI/CD & pre-commit config for linting & testing, local scripts with typer , project dependency management with uv, and static checking with mypy. Have you discovered any other tools that have been game-changers in your ML journey? Feel free to share them in the comments below — I’d love to hear your thoughts! 🤝

My renewed passion for pipelines was ignited during my work at Electric Twin, where we are developing intelligence-as-a-service platforms that simulate human behaviour in real-time. This helps organisations deliver the right message to the right person, at the right time 😎 Collaborating and experimenting at scale on Google Cloud’s Vertex AI platform is pivotal in enhancing the accuracy of our engine, allowing us to provide reliable, actionable insights to our customers. For the latest updates, check out our LinkedIn page! 👋

Originally posted at: