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.
To get started, install the prerequisites, XManager itself and follow the tutorial to create and run a launch script.
See CONTRIBUTING.md for guidance on contributions.
Prerequisites
The codebase assumes Python 3.7+.
Install Docker
If you use xmanager.xm.PythonDocker to run XManager experiments, you need to install Docker.
- Follow the steps to install Docker.
- And if you are a Linux user, follow the steps to enable sudoless Docker.
Install Bazel
If you use xmanager.xm_local.BazelContainer or xmanager.xm_local.BazelBinary to run XManager experiments, you need to install Bazel.
- Follow the steps to install Bazel.
Create a GCP project
If you use xm_local.Caip (Cloud AI Platform) to run XManager experiments, you need to have a GCP project in order to be able to access CAIP to run jobs.
Create a GCP project.
Install gcloud.
Associate your Google Account (Gmail account) with your GCP project by running:
export GCP_PROJECT=<GCP PROJECT ID>
gcloud auth login
gcloud auth application-default login
gcloud config set project $GCP_PROJECT
Set up gcloud to work with Docker by running:
gcloud auth configure-docker
Enable Google Cloud Platform APIs.
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=<GOOGLE_CLOUD_BUCKET_NAME>
gsutil mb -l us-central1 gs://$GOOGLE_CLOUD_BUCKET_NAME
Add GOOGLE_CLOUD_BUCKET_NAME to the environment variables or your .bashrc:
export GOOGLE_CLOUD_BUCKET_NAME=<GOOGLE_CLOUD_BUCKET_NAME>
Install XManager
pip install git+https://github.com/deepmind/xmanager.git
Or, alternatively, a PyPI project is also available.
pip install xmanager
Writing XManager launch scripts
A snippet for the impatient
The basic structure of an XManager launch script can be summarized by these steps:
Create an experiment and acquire its context.
from xmanager import xm
from xmanager import xm_local
with xm_local.create_experiment(experiment_title='cifar10') as experiment:
Define specifications of executables you want to run.
spec = xm.PythonContainer(
path='/path/to/python/folder',
entrypoint=xm.ModuleName('cifar10'),
)
Package your executables.
from xmanager import xm_local
[executable] = experiment.package([
xm.Packageable(
executable_spec=spec,
executor_spec=xm_local.Caip.Spec(),
),
])
Define your hyperparameters.
import itertools
batch_sizes = [64, 1024]
learning_rates = [0.1, 0.001]
trials = list(
dict([('batch_size', bs), ('learning_rate', lr)])
for (bs, lr) in itertools.product(batch_sizes, learning_rates)
)
Define resource requirements for each job.
requirements = xm.JobRequirements(T4=1)
For each trial, add a job / job groups to launch them.
for hyperparameters in trials:
experiment.add(xm.Job(
executable=executable,
executor=xm_local.Caip(requirements=requirements),
args=hyperparameters,
))
Now we should be ready to run the launch script.
To learn more about different executables and executors follow 'Components'.
Run XManager
xmanager launch ./xmanager/examples/cifar10_tensorflow/launcher.py
In order to run multi-job experiments, the --xm_wrap_late_bindings flag might be required:
xmanager launch ./xmanager/examples/cifar10_tensorflow/launcher.py -- --xm_wrap_late_bindings
Components
Executable specifications
XManager executable specifications define what should be packaged in the form of binaries, source files, and other input dependencies required for job execution. Executable specifications are reusable are generally platform-independent.
Container
Container defines a pre-built Docker image located at a URL (or locally).
xm.Container(path='gcr.io/project-name/image-name:latest')
xm.container is a shortener for packageable construction.
assert xm.container(
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
...
) == xm.Packageable(
executable_spec=xm.Container(...),
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
)
BazelBinary
BazelBinary defines a Bazel binary target identified by a label.
xm.Binary(path='//path/to/target:label')
xm.bazel_binary is a shortener for packageable construction.
assert xm.bazel_binary(
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
...
) == xm.Packageable(
executable_spec=xm.BazelBinary(...),
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
)
PythonContainer
PythonContainer defines a Python project that is packaged into a Docker container.
xm.PythonContainer(
entrypoint: xm.ModuleName('<module name>'),
# Optionals.
path: '/path/to/python/project/', # Defaults to the current directory of the launch script.
base_image: '<image>[:<tag>]',
docker_instructions: ['RUN ...', 'COPY ...', ...],
)
A simple form of PythonContainer is to just launch a Python module with default docker_intructions.
xm.PythonContainer(entrypoint=xm.ModuleName('cifar10'))
That specification produces a Docker image that runs the following command:
python3 -m cifar10 fixed_arg1 fixed_arg2
An advanced form of PythonContainer allows you to override the entrypoint command as well as the Docker instructions.
xm.PythonContainer(
entrypoint=xm.CommandList([
'./pre_process.sh',
'python3 -m cifar10 $@',
'./post_process.sh',
]),
docker_instructions=[
'COPY pre_process.sh pre_process.sh',
'RUN chmod +x ./pre_process.sh',
'COPY cifar10.py',
'COPY post_process.sh post_process.sh',
'RUN chmod +x ./post_process.sh',
],
)
That specification produces a Docker image that runs the following commands:
./pre_process.sh
python3 -m cifar10 fixed_arg1 fixed_arg2
./post_process.sh
IMPORTANT: Note the use of $@ which accepts command-line arguments. Otherwise, all command-line arguments are ignored by your entrypoint.
xm.python_container is a shortener for packageable construction.
assert xm.python_container(
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
...
) == xm.Packageable(
executable_spec=xm.PythonContainer(...),
executor_spec=xm_local.Local.Spec(),
args=args,
env_vars=env_vars,
)
Executors
XManager executors define a platform where the job runs and resource requirements for the job.
Each executor also has a specification which describes how an executable specification should be prepared and packaged.
Cloud AI Platform (CAIP)
The Caip executor declares that an executable will be run on the CAIP platform.
The Caip executor takes in a resource requirements object.
xm_local.Caip(
xm.JobRequirements(
cpu=1, # Measured in vCPUs.
ram=4 * xm.GiB,
T4=1, # NVIDIA Tesla T4.
),
)
xm_local.Caip(
xm.JobRequirements(
cpu=1, # Measured in vCPUs.
ram=4 * xm.GiB,
TPU_V2=8, # TPU v2.
),
)
As of June 2021, the currently supported accelerator types are:
P100V100P4T4A100TPU_V2TPU_V3
IMPORTANT: Note that for TPU_V2 and TPU_V3 the only currently supported count is 8.
Caip Specification
The CAIP executor allows you specify a remote image repository to push to.
xm_local.Caip.Spec(
push_image_tag='gcr.io/<project>/<image>:<tag>',
)
Local
The local executor declares that an executable will be run on the same machine from which the launch script is invoked.
Kubernetes (experimental)
The Kubernetes executor declares that an executable will be run on a Kubernetes cluster. As of October 2021, Kubernetes is not fully supported.
The Kubernetes executor pulls from your local kubeconfig. The XManager command-line has helpers to set up a Google Kubernetes Engine (GKE) cluster.
pip install caliban==0.4.1
xmanager cluster create
# cleanup
xmanager cluster delete
You can store the GKE credentials in your kubeconfig:
gcloud container clusters get-credentials <cluster-name>
Kubernetes Specification
The Kubernetes executor allows you specify a remote image repository to push to.
xm_local.Kubernetes.Spec(
push_image_tag='gcr.io/<project>/<image>:<tag>',
)
Job / JobGroup
A Job represents a single executable on a particular executor, while a JobGroup unites a group of Jobs providing a gang scheduling concept: Jobs inside them are scheduled / descheduled simultaneously. Same Job and JobGroup instances can be added multiple times.
Job
A Job accepts an executable and an executor along with hyperparameters which can either be command-line arguments or environment variables.
Command-line arguments can be passed in list form, [arg1, arg2, arg3]:
binary arg1 arg2 arg3
They can also be passed in dictionary form, {key1: value1, key2: value2}:
binary --key1=value1 --key2=value2
Environment variables are always passed in Dict[str, str] form:
export KEY=VALUE
Jobs are defined like this:
[executable] = xm.Package(...)
executor = xm_local.Caip(...)
xm.Job(
executable=executable,
executor=executor,
args={
'batch_size': 64,
},
env_vars={
'NCCL_DEBUG': 'INFO',
},
)
JobGroup
A JobGroup accepts jobs in a kwargs form. The keyword can be any valid Python identifier. For example, you can call your jobs 'agent' and 'observer'.
agent_job = xm.Job(...)
observer_job = xm.Job(...)
xm.JobGroup(agent=agent_job, observer=observer_job)
deepmind
