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

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

要开始使用，请安装 XManager，如有需要，先安装其 先决条件，然后按照 教程 或代码实验室 (Colab Notebook / Jupyter Notebook) 来创建并运行一个启动脚本。

有关贡献的指导，请参阅 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 镜像，该存储桶也可用于暂存文件以构建镜像。

   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 启动脚本

# 包含核心原语和 API。
from xmanager import xm
# 这些核心概念在我们称为“本地后端”的实现中应用，
# 即所有可执行文件都从本机发送执行，
# 不管它们实际上是在本机还是在 GCP 上执行。
from xmanager import xm_local
#
# 创建一个实验上下文，并将其元数据保存到数据库中，
# 我们以后可以通过 `xm_local.list_experiments` 等方式重复使用。
注意，

# `experiment` 具有诸如 `id` 之类的跟踪属性。
with xm_local.create_experiment(experiment_title='cifar10') as experiment:
  # 打包会将给定的 *可执行文件规范* 准备好，以便使用具体的
  # *执行器规范* 来运行：根据组合的不同，这可能涉及构建步骤和/或
  # 将结果复制到某个地方。例如，一个设计在 `Kubernetes` 上运行的
  # `xm.python_container` 将通过 `docker build` 构建，并且新镜像会被
  # 上传到容器注册表。
  # 但在我们这个简单的例子中，我们有一个预先构建好的仅能在本地
  # 运行的 Linux 二进制文件，因此只进行一些验证——比如检查文件是否
  # 存在。
  #
  # `executable` 包含了通过 `.add` 启动打包后的文件所需的所有必要
  # 信息，详见下文。
  [executable] = experiment.package([
      xm.binary(
          # 我们要运行什么。
          path='/home/user/project/a.out',
          # 我们要在哪儿运行它。
          executor_spec=xm_local.Local.Spec(),
      )
  ])
  #
  # 让我们找出哪个 `batch_size` 最优——假设我们的作业会将结果写入
  # 某个位置。
  for batch_size in [64, 1024]:
    # `add` 会创建一个新的 *实验单元*，通常是一组语义上相关的作业，
  # 并将其发送去执行。如果要传递实际的一组作业，可以使用 `JobGroup`
  # （文档稍后会有更多介绍），但出于我们的目的，这里我们将只传递一
  # 个作业。
    experiment.add(xm.Job(
        # 我们之前打包的 `a.out`。
        executable=executable,
        # 这里我们使用默认设置，但执行器有许多参数可用于控制执行。
        executor=xm_local.Local(),
        # 现在该把批量大小作为命令行参数传入了！
        args={'batch_size': batch_size},
        # 我们还可以传递环境变量。
        env_vars={'HEAPPROFILE': '/tmp/a_out.hprof'},
    ))
  #
  # 上下文会等待本地运行的任务完成（但不会等待远程任务，比如发送到
  # GCP 的作业，尽管可以通过 `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 联合起来， 提供一种“gang scheduling”概念：其中的 Job 会同时被调度或取消调度。相同的 Job 和 JobGroup 实例可以被多次 add。

作业

作业接受一个可执行文件和一个执行器，以及超参数，这些超参数可以是命令行参数 或者环境变量。

命令行参数可以以列表形式传递，如 [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)

XManager 快速上手指南

XManager 是一个用于打包、运行和追踪机器学习实验的框架。它支持在本地或 Google Cloud Platform (GCP) 上启动实验，并通过 Python 启动脚本进行交互。

环境准备

系统要求

• Python: 3.9 或更高版本。

可选前置依赖

根据你选择的运行后端，可能需要安装以下工具：

1. Docker (可选)

   • 如果你使用 xmanager.xm.PythonDocker 运行实验。
   • Linux 用户注意：安装后需配置免 sudo 运行 Docker。
   • 安装参考：Docker 官方文档

2. Bazel (可选)

   • 如果你使用 xmanager.xm_local.BazelContainer 或 xmanager.xm_local.BazelBinary。
   • 安装参考：Bazel 官方文档

3. Google Cloud Platform (GCP) (可选)

   • 如果你使用 xm_local.Vertex (Vertex AI) 运行实验。
   • 需要创建 GCP 项目、安装 gcloud CLI、配置认证并启用相关 API (IAM, Cloud AI Platform, Container Registry)。
   • 需要创建一个 GCS bucket 用于存储实验产物。

安装步骤

方法一：通过 PyPI 安装（推荐）

pip install xmanager

方法二：从源码安装

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

方法三：Debian/Ubuntu 一键脚本

如果你使用的是基于 Debian 的系统，可以克隆仓库并运行设置脚本来自动安装所有依赖：

git clone https://github.com/deepmind/xmanager.git
cd xmanager/setup_scripts && chmod +x setup_all.sh && . ./setup_all.sh

基本使用

XManager 的核心工作流程分为：创建实验上下文 -> 定义可执行文件规范 -> 打包 -> 定义超参数 -> 提交任务。

以下是一个最简单的本地运行示例脚本 (launcher.py)：

# 导入核心模块
from xmanager import xm
from xmanager import xm_local

# 1. 创建实验上下文
with xm_local.create_experiment(experiment_title='cifar10_demo') as experiment:
  
  # 2. 定义可执行文件规范 (这里以预编译的二进制文件为例)
  # 如果是 Python 项目，通常使用 xm.PythonContainer
  spec = xm.Binary(
      path='/home/user/project/a.out', 
      executor_spec=xm_local.Local.Spec(),
  )

  # 3. 打包可执行文件
  # package 方法返回一个 executable 列表，用于后续任务定义
  [executable] = experiment.package([spec])

  # 4. 定义超参数组合
  batch_sizes = [64, 1024]

  # 5. 提交任务 (Job)
  for batch_size in batch_sizes:
    experiment.add(xm.Job(
        executable=executable,
        # 指定执行器为本地运行
        executor=xm_local.Local(),
        # 传递命令行参数 (字典格式会自动转换为 --key=value)
        args={'batch_size': batch_size},
        # 传递环境变量
        env_vars={'HEAPPROFILE': '/tmp/a_out.hprof'},
    ))
    
  # 上下文退出时，本地任务会自动等待完成

运行启动脚本

使用 xmanager launch 命令运行上述脚本：

xmanager launch ./launcher.py

如果运行包含多任务绑定的复杂实验，可能需要添加标志：

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

核心概念简述

• Executable Specs: 定义要打包的内容（如 Docker 镜像、二进制文件、Python 目录）。
• Executors: 定义运行平台（Local 本地, Vertex GCP Vertex AI, Kubernetes）。
• Job: 单个执行单元，包含可执行文件、执行器和参数。
• JobGroup: 一组需要同时调度/取消调度的 Job 集合。