Package: xmanager
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
onGoogle Cloud Platform (GCP). Interaction with
experiments is done via XManager's APIs through Pythonlaunch scripts. Check
outthese slidesfor a more detailed introduction.To get started, installXManager, itsprerequisitesif needed and followthe
tutorialor a codelab
(Colab Notebook/Jupyter Notebook)
to create and run a launch script.SeeCONTRIBUTING.mdfor guidance on contributions.Install XManagerpipinstallgit+https://github.com/deepmind/xmanager.gitOr, alternatively,a PyPI projectis also
available.pipinstallxmanagerOn Debian-based systems, XManager and all its dependencies can be installed and
set up by cloning this repository and then runningcdxmanager/setup_scripts&&chmod+xsetup_all.sh&&../setup_all.shPrerequisitesThe codebase assumes Python 3.9+.Install Docker (optional)If you usexmanager.xm.PythonDockerto run XManager experiments,
you need to install Docker.Followthe stepsto install Docker.And if you are a Linux user, followthe stepsto enable sudoless Docker.Install Bazel (optional)If you usexmanager.xm_local.BazelContainerorxmanager.xm_local.BazelBinaryto run XManager experiments, you need to install Bazel.Followthe stepsto
install Bazel.Create a GCP project (optional)If you usexm_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.Createa GCP project.Installgcloud.Associate your Google Account (Gmail account) with your GCP project by
running:exportGCP_PROJECT=<GCPPROJECTID>
gcloudauthlogin
gcloudauthapplication-defaultlogin
gcloudconfigsetproject$GCP_PROJECTSet upgcloudto work with Docker by running:gcloudauthconfigure-dockerEnable Google Cloud Platform APIs.EnableIAM.Enablethe 'Cloud AI Platfrom'.Enablethe 'Container Registry'.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.exportGOOGLE_CLOUD_BUCKET_NAME=<GOOGLE_CLOUD_BUCKET_NAME>
gsutilmb-lus-central1gs://$GOOGLE_CLOUD_BUCKET_NAMEAddGOOGLE_CLOUD_BUCKET_NAMEto the environment variables or your .bashrc:exportGOOGLE_CLOUD_BUCKET_NAME=<GOOGLE_CLOUD_BUCKET_NAME>Writing XManager launch scriptsA snippet for the impatient 🙂# Contains core primitives and APIs.fromxmanagerimportxm# 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.fromxmanagerimportxm_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`.withxm_local.create_experiment(experiment_title='cifar10')asexperiment:# 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.forbatch_sizein[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 explicitly awaited via# `wait_for_completion`).The basic structure of an XManager launch script can be summarized by these
steps:Create an experiment and acquire its context.fromxmanagerimportxmfromxmanagerimportxm_localwithxm_local.create_experiment(experiment_title='cifar10')asexperiment:Define specifications of executables you want to run.spec=xm.PythonContainer(path='/path/to/python/folder',entrypoint=xm.ModuleName('cifar10'),)Package your executables.[executable]=experiment.package([xm.Packageable(executable_spec=spec,executor_spec=xm_local.Vertex.Spec(),),])Define your hyperparameters.importitertoolsbatch_sizes=[64,1024]learning_rates=[0.1,0.001]trials=list(dict([('batch_size',bs),('learning_rate',lr)])for(bs,lr)initertools.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.forhyperparametersintrials:experiment.add(xm.Job(executable=executable,executor=xm_local.Vertex(requirements=requirements),args=hyperparameters,))Now we should be readyto runthe launch script.To learn more about differentexecutablesandexecutorsfollow'Components'.Run XManagerxmanagerlaunch./xmanager/examples/cifar10_tensorflow/launcher.pyIn order to run multi-job experiments, the--xm_wrap_late_bindingsflag might
be required:xmanagerlaunch./xmanager/examples/cifar10_tensorflow/launcher.py----xm_wrap_late_bindingsComponentsExecutable specificationsXManager 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 and generally platform-independent.Seeexecutable_specs.mdfor details on each executable specification.NameDescriptionxmanager.xm.ContainerA pre-built.tarimage.xmanager.xm.BazelContainerABazeltarget producing a.tarimage.xmanager.xm.BinaryA pre-built binary.xmanager.xm.BazelBinaryABazeltarget producing a self-contained binary.xmanager.xm.PythonContainerA directory with Python modules to be packaged as a Docker container.ExecutorsXManager 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.Seeexecutors.mdfor details on each executor.NameDescriptionxmanager.xm_local.LocalRuns a binary or a container locally.xmanager.xm_local.VertexRuns a container onVertex AI.xmanager.xm_local.KubernetesRuns a container on Kubernetes.Job / JobGroupAJobrepresents a single executable on a particular executor, while aJobGroupunites a group ofJobs providing a gang scheduling concept:Jobs inside them are scheduled / descheduled simultaneously. SameJobandJobGroupinstances can beadded multiple times.JobA 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]:binaryarg1arg2arg3They can also be passed in dictionary form,{key1: value1, key2: value2}:binary--key1=value1--key2=value2Environment variables are always passed inDict[str, str]form:exportKEY=VALUEJobs are defined like this:[executable]=xm.Package(...)executor=xm_local.Vertex(...)xm.Job(executable=executable,executor=executor,args={'batch_size':64,},env_vars={'NCCL_DEBUG':'INFO',},)JobGroupA 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)