Introduction To Training Pipelines The typical machine learning training process consists of various tasks such as data preparation, model training, model validation, model deployment and others. Because you may want to iterate on your model even after you deploy it to production, it is paramount to automate these tasks, e.g. to run them with new data, hyper-parameters, changes in the code, etc. Training pipelines help automate these tasks. At the same time, it is crucial to keep the pipeline reproducible, in the sense that you as well as others can replicate it and obtain similar results. In order to ensure that pipelines are reproducible, it is crucial to track infrastructure, code, data, hyper-parameters, experiment metrics, etc. In this article, I’ll give an overview on how WandB and dstack together can help build reproducible training pipelines. Experiment tracking First, we’ll start with experiment tracking. Experiment tracking involves tracking hyper-parameters (e.g., step-size, batch size, etc.) experiment metrics (e.g., accuracy, training loss, validation loss, etc.) hardware system metrics (e.g., the utilization of GPUs, memory, etc.) WandB provides one of the easiest ways to perform experiment tracking, as it is easy to integrate into the ML code (addition of only a few lines of code) uses cloud to store the results so you can always find these metrics by your run name (don’t underestimate the importance of this, as storing metrics on you local machine may be rather an anti-pattern in regard to reproducibility) automatically provides the visualization of various metrics with nice charts WandB is very easy to integrate into your existing codebase. Below is an example of how to do it if you use PyTorch Lightning. We need to import the object. Then, instantiate the object and pass a project name as its argument, as shown below. WandbLogger from pytorch_lightning.loggers import WandbLogger

wandb_logger = WandbLogger(project="my-test-project") The created can be passed into the logger argument in the object as shown below. wandb_logger Trainer trainer = Trainer(logger=wandb_logger) Once the run is completed, the information of the run is assigned to . my-test-project You may check and click your project for the results (loss, accuracy, system metrics etc) pertaining to various runs of the project, as shown below. wandb.ai In order to store hyper-parameters or other configuration parameters, is used and its integration into the ML code is easy as shown below. wandb.config import wandb
wandb.init()
wandb.config.epochs = 10
wandb.config.batch_size = 32 The experiment configuration for each run is automatically stored in the cloud, and it can be found in the file in the tab of a run at . config.yaml Files wandb.ai Automating workflows Now, let’s talk about the automation of tasks and tracking the rest of our pipeline, which may include data, code, and infrastructure. For this purpose, we use dstack. Here’s a brief list of what dstack can do: version data, code, and infrastructure automate training tasks through declarative configuration files automatically provision infrastructure in a linked cloud account (it supports AWS, GCP, and Azure), Most importantly, it can be used from your favorite IDE (or terminal) To automate workflows with dstack, one needs to define them in the file. ./dstack/workflows.yaml Here’s a very simple example: workflows:
  - name: prepare
    provider: python
    script: "prepare.py"
    requirements: "requirements.txt"
    artifacts: ["data"]

 - name: train
   provider: python
   version: "3.9"
   requirements: "requirements.txt"
   script: "train.py"
   depends-on:
      - prepare
   artifacts: ["model"]
   resources:
     gpu: 4 Here, we can define multiple workflows (i.e. tasks) and configure dependencies between them. In this particular example, we have two workflows: and . prepare train The workflow depends on the workflow. train prepare As you see, each workflow defines how to run the code, including what folders to store as output artifacts, and what infrastructure is needed (e.g. the number of GPUs, amount of memory, etc). dstack CLI is used to run the workflow from your local terminal. To run the above-mentioned workflow, execute the following command in your local terminal. dstack run train Once the run is submitted, the user can access relevant logs, changes in code, the artifacts etc in the dstack UI or CLI. In order to monitor the run, login to your account, and the contents in the dstack UI should look like below. dstack.ai Artifacts can be browsed through the user interface. Moreover, it is possible to download the contents of artifacts using the dstack CLI with the following command. dstack artifacts download <run-name> Versioning data In order to ensure reproducibility of the training pipelines, it is crucial to track data too. With dstack, data artifacts can be versioned by assigning tags, which can later be referenced in other workflows. In the example above, the workflow depended on the workflow. Each time you ran the workflow, dstack also ran the workflow, and then passed the output artifacts of the workflow to the workflow. train prepare train download prepare train Now let’s imagine that we’d like to run the workflow independently and then reuse the output artifacts of that particular run with the workflows. prepare train In order to do that, you have to run the workflow, and then assign a tag to it (e.g. through the UI as shown below or the CLI). prepare Then, you can refer to this tag from the workflow: train workflows:
  - name: prepare
    provider: python
    script: "prepare.py"
    requirements: "requirements.txt"
    artifacts: ["data"]

 - name: train
   provider: python
   version: "3.9"
   requirements: "requirements.txt"
   script: "train.py"
   depends-on:
      - prepare:latest
   artifacts: ["model"]
   resources:
     gpu: 4 By decoupling the workflows of preparing data and training the model, it becomes easier to train the model iteratively and keep every run reproducible. dstack + WandB Configuration It is possible to seamlessly use dstack and WandB together. To do that, obtain the WandB API key from “Settings” in Wandb as shown below. And add it to dstack Secrets. You can do it  via  “Settings” in . Click the “Add secret” button and add the key secret with the copied  WandB API key value. dstack.ai WANDB_API_KEY The dstack settings should look like below. In order to use the same  run names across dstack and WandB, you can use the environment variable to get the dstack run name and pass it as a display name to as shown below: RUN_NAME WandbLogger run_name = os.environ['RUN_NAME']

wandb log results to a project

wandb_logger = WandbLogger(name=run_name, project="my-test-project") That’s it for now. The source code used in this blog post can be found . here You’re very welcome to share your feedback, ask questions, and of course give a spin to WandB and dstack yourself. I want to try it. Where do I start? dstack Quickstart WandB Quickstart

Reproducible ML Training Pipelines With dstack And WandB

About Author

Comments

TOPICS

THIS ARTICLE WAS FEATURED IN

Related Stories

Untitled Story

How to Build a Training Pipeline on Multiple GPUs

The Noonification: Use This 7-Step McKinsey Framework to Solve Any Problem (1/10/2023)

The Noonification: A Taxonomy of Inclusiveness (1/11/2024)

The Noonification: What is the InfiniteNature-Zero AI Model? (11/19/2022)

10 Ways AI Has Changed Our Lives

How to Build a Training Pipeline on Multiple GPUs

The Noonification: Use This 7-Step McKinsey Framework to Solve Any Problem (1/10/2023)

The Noonification: A Taxonomy of Inclusiveness (1/11/2024)

The Noonification: What is the InfiniteNature-Zero AI Model? (11/19/2022)

10 Ways AI Has Changed Our Lives

Light-Mode

Classic

Newspaper

Dark-Mode

Neon Noir

Minty

HN StartUps