google-deepmind/xmanager

google-deepmind/xmanager

Description: A platform for managing machine learning experiments

Language: Python

License: Apache-2.0

Stars: 912

Forks: 64

Open issues: 32

Created: 2021-04-26T17:03:51Z

Pushed: 2026-05-18T13:58:00Z

Default branch: main

Fork: no

Archived: no

README:

XManager: A framework for managing machine learning experiments 🧑‍🔬

XManager is a platform for packaging, running and keeping track of machine learning experiments. It currently enables one to launch experiments locally or on Google Cloud Platform (GCP). Interaction with experiments is done via XManager's APIs through Python *launch scripts*. Check out these slides for a more detailed introduction.

To get started, install [XManager](#install-xmanager), its [prerequisites](#prerequisites) if needed and follow [the tutorial](#writing-xmanager-launch-scripts) or a codelab (Colab Notebook / Jupyter Notebook) to create and run a launch script.

See CONTRIBUTING.md for guidance on contributions.

Install XManager

pip install git+https://github.com/deepmind/xmanager.git

Or, alternatively, a PyPI project is also available.

pip install xmanager

On Debian-based systems, XManager and all its dependencies can be installed and set up by cloning this repository and then running

cd xmanager/setup_scripts && chmod +x setup_all.sh && . ./setup_all.sh

Prerequisites

The codebase assumes Python 3.9+.

Install Docker (optional)

If you use xmanager.xm.PythonDocker to run XManager experiments, you need to install Docker.

1. Follow the steps to install Docker.

2. And if you are a Linux user, follow the steps to enable sudoless Docker.

Install Bazel (optional)

If you use xmanager.xm_local.BazelContainer or xmanager.xm_local.BazelBinary to run XManager experiments, you need to install Bazel.

1. Follow the steps to install Bazel.

Create a GCP project (optional)

If you use xm_local.Vertex (Vertex AI) to run XManager experiments, you need to have a GCP project in order to be able to access Vertex AI to run jobs.

1. Create a GCP project.

2. Install gcloud.

3. Associate your Google Account (Gmail account) with your GCP project by running:

export GCP_PROJECT=
gcloud auth login
gcloud auth application-default login
gcloud config set project $GCP_PROJECT

4. Set up gcloud to work with Docker by running:

gcloud auth configure-docker

5. Enable Google Cloud Platform APIs.

• Enable

IAM.

• Enable

the 'Cloud AI Platfrom'.

• Enable

the 'Container Registry'.

6. Create a staging bucket in us-central1 if you do not already have one. This bucket should be used to save experiment artifacts like TensorFlow log files, which can be read by TensorBoard. This bucket may also be used to stage files to build your Docker image if you build your images remotely.

export GOOGLE_CLOUD_BUCKET_NAME=
gcloud storage buckets create --location=us-central1 gs://$GOOGLE_CLOUD_BUCKET_NAME

Add GOOGLE_CLOUD_BUCKET_NAME to the environment variables or your .bashrc:

export GOOGLE_CLOUD_BUCKET_NAME=

Writing XManager launch scripts

A snippet for the impatient 🙂

# Contains core primitives and APIs.
from xmanager import xm
# Implementation of those core concepts for what we call 'the local backend',
# which means all executables are sent for execution from this machine,
# independently of whether they are actually executed on our machine or on GCP.
from xmanager import xm_local
#
# Creates an experiment context and saves its metadata to the database, which we
# can reuse later via `xm_local.list_experiments`, for example. Note that
# `experiment` has tracking properties such as `id`.
with xm_local.create_experiment(experiment_title='cifar10') as experiment:
# Packaging prepares a given *executable spec* for running with a concrete
# *executor spec*: depending on the combination, that may involve building
# steps and / or copying the results somewhere. For example, a
# `xm.python_container` designed to run on `Kubernetes` will be built via
#`docker build`, and the new image will be uploaded to the container registry.
# But for our simple case where we have a prebuilt Linux binary designed to
# run locally only some validations are performed -- for example, that the
# file exists.
#
# `executable` contains all the necessary information needed to launch the
# packaged blob via `.add`, see below.
[executable] = experiment.package([
xm.binary(
# What we are going to run.
path='/home/user/project/a.out',
# Where we are going to run it.
executor_spec=xm_local.Local.Spec(),
)
])
#
# Let's find out which `batch_size` is best -- presumably our jobs write the
# results somewhere.
for batch_size in [64, 1024]:
# `add` creates a new *experiment unit*, which is usually a collection of
# semantically united jobs, and sends them for execution. To pass an actual
# collection one may want to use `JobGroup`s (more about it later in the
# documentation), but for our purposes we are going to pass just one job.
experiment.add(xm.Job(
# The `a.out` we packaged earlier.
executable=executable,
# We are using the default settings here, but executors have plenty of
# arguments available to control execution.
executor=xm_local.Local(),
# Time to pass the batch size as a command-line argument!
args={'batch_size': batch_size},
# We can also pass environment variables.
env_vars={'HEAPPROFILE': '/tmp/a_out.hprof'},
))
#
# The context will wait for locally run things (but not for remote things such
# as jobs sent to GCP, although they can be…