Optimizing CI/CD model management and evaluation workflows with Weights & Biases

Understanding the challenges of CI/CD in machine learning

Data is also a source: ML uses data (not only code) as a source. This stands in contrast to software engineering, where code is typically the only source of truth required to build an application.
Training output is non-deterministic: When training a model multiple times from code and data that have not changed, the resulting weights and biases (i.e., the numbers that define the model's operation) aren’t always consistent. In contrast, no matter how many times a traditional software engineering application is built, if the code doesn't change, the result should remain exactly the same.
Scale and variability: ML models often require a staggering amount of training runs —hundreds, thousands, and often far more than that—before they're ready for prime time. Moreover, training different models doesn’t always involve completing the same number of iterations. In contrast, building a software application typically requires a single sequence of pre-determined steps.

While traditional software engineering CI/CD involves building source code into an application by means of a single sequence of steps, ML engineering also uses data as a source and often produces non-deterministic results after lengthy and costly training jobs.

Dataset tracking: In addition to code repositories, we need a way to track the progression and differences in our source datasets.
Model tracking: We also need a way to track the changes and progression of our trained models, because, due to the scale of the work involved in training, developing a model from scratch may be prohibitively costly. For example, think about a model that takes three or four days to train: when you identify an issue with this model in production, will you be able to afford waiting for it to train from scratch again?
Governance, reproducibility, and observability: Finally, we need to tie together all the sources, configurations and performance metrics associated with trained models, in a way that helps us understand and trace how each of our models was produced. That way, we can always explain what inputs were used to create those models. In other words, we can understand our model’s lineage. This kind of lineage observability requirement has been a standard in heavily regulated industries for some time, but has recently become a widespread concern due to the proliferation of AI solutions in the market, as well as recent regulations calling for more diligent governance and risk management practices from those involved in using or developing AI.

Data science teams who focus on collecting, cleaning, balancing, aggregating and enriching data from multiple sources, in order to create ‘golden’ datasets that other teams can use to train ML models.
ML engineering teams who consume the golden datasets produced by data scientists, and use them to experiment with different ways to train a model for a particular task.
ML operations (MLOps) teams who are tasked with making sure trained models are properly integrated, evaluated, staged and deployed, as these models progress through a given CI/CD workflow.

Using Weights & Biases to optimize CI/CD workflows

Tracking datasets

pip install wandb
wandb login

import wandb

# Start a run
with wandb.init(project="cicd-2024", job_type="data-logging") as run:

    # Establish the artifact
    dataset_artifact = wandb.Artifact(name="source-data", type="datasets")

    # Add items to the artifact
    dataset_artifact.add_file("source-data.csv")

    # Log the artifact
    run.log_artifact(dataset_artifact)

The artifacts page in Weights & Biases organizes your datasets in helpful collections and keeps track of different versions automatically.

Tracking models and extending the lineage

Lineage as an ML workflow

Enter: use_artifact()

with wandb.init(
    job_type="training",
    config={
        "epochs": 50,
        "dataset": "team-jdoc/cicd-2024/source-data:v0",
        "params": {
          "lr": 0.0003,
          "batch_size": 32
        }
    }
) as run:

    # extend Artifact lineage
    run.use_artifact(run.config["dataset"], type="datasets")

    for epoch in range(run.config["epochs"]):
        # TODO: Insert additional dataloader/batching loop
        checkpoint_file_path = mock_backprop_pass()
        accuracy, loss = mock_forward_pass()
        print(f'Epoch {epoch}, accuracy: {accuracy}, loss: {loss}')
        # Log metrics as a dictionary of {key: value} pairs
        run.log({"accuracy":accuracy, "loss":loss})

    # Log the model
    model_artifact = wandb.Artifact(name="dummy_model", type="models")
    model_artifact.add_file(checkpoint_file_path)
    run.log_artifact(model_artifact)

Let's get automating

Artifact changes as event triggers

A new version of an Artifact has been logged, or
A new alias has been added to an artifact (more on this later)

An outgoing webhook is sent to a third-party tool in your CI/CD stack, or
A job is added to a Launch queue

Automatically retraining a model in response to data drift, and
Automatically delivering newly improved models through a continuous evaluation cycle

Automatically retraining a model in response to data drift



Automatically delivering models through continuous evaluation



Every training or retraining job, uses the log_artifact() method to log a new model version to W&B.
Each logged model then triggers a job automation that executes an evaluation with the purpose of comparing the newly trained model to the best model in the 'production' Model registry. The evaluation also generates a report highlighting the results of this comparison, permanently associating it to the model's lineage.
If the newly trained model outperforms the latest version of the model in the Model registry, this newly trained model is linked to that registry, thus triggering a subsequent automation.
In our example, the second automation is again a webhook sent to the team's CI/CD infrastructure, containing all the information necessary to retrieve, integrate and deploy the newly logged model.

Conclusion

Understand the key concepts, challenges, and benefits of an effective CI/CD pipeline in the context of ML,
Use W&B Registry as the central model and dataset repository and single source of truth for production artifacts, and
Construct workflows to evaluate and compare datasets and models, and generate reports to share results with your teamates