Sustainability-aware job scheduler for federated learning

Inspiration

Federated learning provides great opportunities for training complex machine learning models while honoring the privacy requirements of sensitive data sets. This is especially useful in the medical field, where patient-specific data sets can make use of and improve publically available models. The federated learning framework DISCO (Distributed Collaborative Machine Learning) enables the joint training of machine learning models while allowing each participating party only access to their own data sets. In addition, federated learning provides the great opportunity to incorporate a more complicated scheduling logic into the task distribution. Here, we make use of this idea by making the task scheduler act in a sustainable way, i.e., it distributes the model training steps such that sustainability factors and negative effects for the environment are minimized while training the models in an effective way.

What it does

We extend DISCO with code which records sustainability metrics for each client trying to train a model. When distributing updated weights and enabling model training, these metrics are used to prioritize sustainable and environmentally friendly clients. In this way, clients are incentivized to behave sustainable way. In particular, we use the following metrics:

Latency: prioritize models which are able to train their model in an efficient way to incentivize good resource allocations
Carbon intensity: prioritize models which run on machines powered by renewable and clean energy sources to protect our environment

How we built it

We used typescript to code both back-end and front-end. We retrieve energy usage data from https://www.electricitymaps.com/.

Challenges we ran into

Getting used to the already existing, big DISCO codebase and finding the best entry points for our extensions. At first it seemed intimidating, but after working through it for some time we found it to be well designed and very readable. Deciding on what are the best sustainability metrics.

Accomplishments that we're proud of

Adding a useful and working feature to a complicated production codebase.

What we learned

Writing reusable and efficient typescript. Which sustainability factors are especially relevant to the field of machine learning and what to watch out for when using machine learning methods in the future. Collaborating in a diverse and exciting team :) Understanding how to work with big codebases we did not write ourselves. Becoming more comfortable with contributing to open-source projects.