XManager：管理机器学习实验的框架 🧑‍🔬

XManager是一个用于打包、运行和跟踪机器学习实验的平台。它目前允许用户在本地或在Google Cloud Platform (GCP)上启动实验。通过XManager的API和Python 启动脚本 与实验进行交互。查看这些幻灯片以获取更详细的介绍。

要开始使用，请安装XManager、所需的前提条件（如果有的话），并按照教程或一个codelab（Colab笔记本 / Jupyter笔记本）创建并运行启动脚本。

查看CONTRIBUTING.md以获取有关贡献的指导。

安装XManager

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

或者，你也可以选择使用PyPI项目。

pip install xmanager

在基于Debian的系统上，可以通过克隆此存储库并运行以下命令来安装和设置XManager及其所有依赖项：

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

先决条件

代码库假定Python 3.9+。

安装Docker（可选）

如果你使用xmanager.xm.PythonDocker来运行XManager实验，你需要安装Docker。

1. 按照这些步骤安装Docker。

2. 如果你是Linux用户，按照这些步骤启用无需sudo的Docker。

安装Bazel（可选）

如果你使用xmanager.xm_local.BazelContainer或xmanager.xm_local.BazelBinary来运行XManager实验，你需要安装Bazel。

1. 按照这些步骤安装Bazel。

创建GCP项目（可选）

如果你使用xm_local.Vertex（Vertex AI）来运行XManager实验，你需要有一个GCP项目才能访问Vertex AI来运行作业。

1. 创建一个GCP项目。

2. 安装 gcloud。

3. 通过运行以下命令将你的Google帐户（Gmail帐户）与你的GCP项目关联起来：

   export GCP_PROJECT=<GCP PROJECT ID>
   gcloud auth login
   gcloud auth application-default login
   gcloud config set project $GCP_PROJECT
   

4. 通过运行以下命令将gcloud设置为与Docker一起使用：

   gcloud auth configure-docker
   

5. 启用Google Cloud Platform API。

   • 启用 IAM。

   • 启用 'Cloud AI Platform'。

   • 启用 'Container Registry'。

6. 如果你还没有一个，请在us-central1中创建一个临时存储桶。这个存储桶应该用于保存实验工件，如TensorFlow日志文件，这些文件可以由TensorBoard读取。如果你的Docker镜像是在远程构建的，这个存储桶也可以用于存储构建Docker镜像的文件。

   export GOOGLE_CLOUD_BUCKET_NAME=<GOOGLE_CLOUD_BUCKET_NAME>
   gsutil mb -l us-central1 gs://$GOOGLE_CLOUD_BUCKET_NAME
   

   将GOOGLE_CLOUD_BUCKET_NAME添加到环境变量或你的.bashrc中

   export GOOGLE_CLOUD_BUCKET_NAME=<GOOGLE_CLOUD_BUCKET_NAME>
   

编写XManager启动脚本

# 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 explicitly awaited via
  # `wait_for_completion`).

一个XManager启动脚本的基本结构可以概括为以下步骤：

1. 创建一个实验并获取其上下文。

   from xmanager import xm
   from xmanager import xm_local
   
   with xm_local.create_experiment(experiment_title='cifar10') as experiment:
   

2. 定义你想要运行的执行程序的规范。

   spec = xm.PythonContainer(
       path='/path/to/python/folder',
       entrypoint=xm.ModuleName('cifar10'),
   )
   

3. 打包你的执行程序。

   [executable] = experiment.package([
     xm.Packageable(
       executable_spec=spec,
       executor_spec=xm_local.Vertex.Spec(),
     ),
   ])
   

4. 定义你的超参数。

   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)
   )
   

5. 定义每个作业的资源需求。

   requirements = xm.JobRequirements(T4=1)
   

6. 对于每个试验，添加一个作业/作业组以启动它们。

   for hyperparameters in trials:
     experiment.add(xm.Job(
         executable=executable,
         executor=xm_local.Vertex(requirements=requirements),
         args=hyperparameters,
       ))
   

现在我们应该准备好运行启动脚本了。

要了解更多关于不同执行程序和执行器的信息，请参阅'组件'。

运行XManager

xmanager launch ./xmanager/examples/cifar10_tensorflow/launcher.py

为了运行多作业实验，可能需要--xm_wrap_late_bindings标志

xmanager launch ./xmanager/examples/cifar10_tensorflow/launcher.py -- --xm_wrap_late_bindings

组件

可执行规范

XManager可执行规范定义了作业执行所需打包的格式为二进制、源文件和其他输入依赖项。可执行规范是可重用的，并且通常与平台无关。

有关每个可执行规范的详细信息，请参阅executable_specs.md。

执行器

XManager执行器定义了作业运行的平台和作业的资源需求。

每个执行器也都有一个规范，该规范描述了可执行规范应该如何准备和打包。

有关每个执行器的详细信息，请参阅 executors.md。

作业 / 作业组

Job 代表特定执行器上的单个可执行程序，而 JobGroup 将一组 Job 合并在一起，提供了一种群体调度概念：它们内部的 Job 将会同时调度/取消调度。相同的 Job 和 JobGroup 实例可以被多次 添加。

作业

作业接受一个可执行程序和一个执行器，以及可以是以命令行参数或环境变量形式存在的超参数。

命令行参数可以以列表形式传递，例如 [arg1, arg2, arg3]

binary arg1 arg2 arg3

它们也可以以字典形式传递，例如 {key1: value1, key2: value2}

binary --key1=value1 --key2=value2

环境变量总是以 Dict[str, str] 形式传递

export KEY=VALUE

作业的定义如下

[executable] = xm.Package(...)

executor = xm_local.Vertex(...)

xm.Job(
    executable=executable,
    executor=executor,
    args={
        'batch_size': 64,
    },
    env_vars={
        'NCCL_DEBUG': 'INFO',
    },
)

作业组

作业组接受以关键字参数形式的工作。关键字可以是任何有效的 Python 标识符。例如，您可以称您的作业为 'agent' 和 'observer'。

agent_job = xm.Job(...)
observer_job = xm.Job(...)

xm.JobGroup(agent=agent_job, observer=observer_job)