目录

基本概念

这描述了TorchX背后的核心概念和项目结构。 有关如何创建和运行应用程序,请参阅快速入门指南

项目结构

TorchX 的顶层模块有:

  1. torchx.specs: 应用程序规范(作业定义)API

  2. torchx.components: 预定义(内置)应用规范

  3. torchx.runner: 给定一个应用规格,将其作为调度器上的作业提交

  4. torchx.schedulers: 后端作业调度器,该运行器支持

  5. torchx.pipelines: 转换器,将给定的应用程序规范转换为ML管道平台中的一个“阶段”

  6. torchx.runtime: 可用于编写应用程序(非应用程序特定)的工具和抽象库

  7. torchx.cli: 命令行工具

下面是UML图

_images/torchx_module_uml.jpg

概念

AppDefs

在TorchX中,一个AppDef只是一个具有实际应用程序定义的结构。在调度员术语中,这是一个JobDefinition,而在Kubernetes中类似的概念是spec.yaml。为了区分应用程序二进制文件(逻辑)和规范,我们通常将TorchXAppDef称为“应用规范”或specs.AppDefspecs.AppDef是被torchx.runnertorchx.pipelines理解的常见接口,允许您将应用程序作为独立任务运行或作为机器学习管道中的一个阶段。

下面是简单地回显“你好,世界”的 specs.AppDef 示例

import torchx.specs as specs

specs.AppDef(
   name="echo",
   roles=[
       specs.Role(
           name="echo",
           entrypoint="/bin/echo",
           image="/tmp",
           args=["hello world"],
           num_replicas=1
       )
   ]
)

正如你所见,specs.AppDef 是一个纯 Python 数据类, 简单地编码了主二进制文件(入口点)的名称、传递给它的参数以及其他一些运行时参数,如 num_replicas, 以及要在其中运行的容器的信息 (entrypoint=/bin/echo).

该应用程序规范是灵活的,可以为各种应用程序拓扑编码规范。 例如,num_replicas > 1 表示应用程序是分布式的。 指定多个 specs.Roles 可以表示一个非同构的分布式应用程序,例如那些需要一个“协调者”和许多“工作者”的应用程序。

请参阅 torchx.specs API文档 以了解更多。

是什么让应用规格灵活的,也使其包含许多字段。好消息是在大多数情况下,你不需要从头开始构建一个应用规格。相反,你会使用一个名为components的模板化应用规格。

组件

TorchX中的一个组件就是一个模板化的 spec.AppDef。你可以将其视为spec.AppDef的方便“工厂方法”。

注意

与应用程序不同,组件并不会映射到实际的Python数据类。 而是一个返回spec.AppDef的工厂函数被称为组件。

应用程序规格模板化的粒度有所不同。某些组件,例如上面的echo示例是可直接运行的,这意味着它们包含硬编码的应用程序二进制文件。而其他组件,例如ddp(分布式数据并行)规格,只指定应用程序的拓扑结构。下面是可能的一种ddp风格训练应用程序规格的模板化,它指定了一个同构节点拓扑:

import torchx.specs as specs

def ddp(jobname: str, nnodes: int, image: str, entrypoint: str, *script_args: str):
   single_gpu = specs.Resources(cpu=4, gpu=1, memMB=1024)
   return specs.AppDef(
           name=jobname,
           roles=[
               specs.Role(
                   name="trainer",
                   entrypoint=entrypoint,
                   image=image,
                   resource=single_gpu,
                   args=script_args,
                   num_replicas=nnodes
               )
           ]
   )

正如你所见,参数化的程度完全取决于组件作者。创建一个组件的努力不过是编写一个 Python 函数。不要试图通过参数化一切来过度泛化组件。组件易于创建且成本低廉,基于重复用例,你可以根据需要创建任意多个组件。

提示 1: 由于组件是 Python 函数,组件组合可以通过 Python 函数组合实现,而不是通过对象组合。 但是出于可维护性的考虑,我们不建议进行组件组合。

小贴士2: 要定义组件之间的依赖关系,请使用管道化DSL。 请参阅下面的管道适配器部分,了解TorchX组件 在管道上下文中是如何使用的。

在编写自己的组件之前,请浏览包含在TorchX中的 组件库,看看是否有符合您需求的。

运行器和调度器

一个 Runner 完全符合您的预期——给定应用程序规格后,它会通过作业调度程序在集群上启动该应用程序作为作业。

有两种方式可以访问 TorchX 中的运行器:

  1. CLI: torchx run ~/app_spec.py

  2. Programmatically: torchx.runner.get_runner().run(appspec)

请参阅 调度器 以获取运行程序可以启动应用程序的调度器列表。

管道适配器

虽然运行者将组件作为独立作业启动,但torchx.pipelines 使得将组件插入到ML管道/工作流中成为可能。对于特定的目标管道平台(例如kubeflow pipelines),TorchX定义了一个适配器,将TorchX应用规范转换为目标平台的“阶段”表示形式。例如,torchx.pipelines.kfp适配器将应用规范转换为kfp.ContainerOp(更准确地说,是一个kfp“组件规范”yaml)。

在大多数情况下,应用程序规范会映射到流水线中的一个“阶段”(或节点)。 然而,高级组件,尤其是那些具有自己微型控制流的组件(例如超参数优化HPO),可能会映射到一个“子流水线”或“内联流水线”。 这些高级组件如何映射到流水线的确切语义取决于目标流水线平台。例如,如果流水线DSL允许从上游阶段动态添加阶段到流水线中,则TorchX可能会利用此功能将子流水线“内联”到主流水线中。TorchX通常尽量将其应用程序规范适配为目标流水线平台的< strong>最规范表示形式。

请参阅 Pipelines 以获取支持的管道平台列表。

运行时

重要的

torchx.runtime 绝对不是使用TorchX的要求。 如果你的基础设施是固定的,并且你不需要你的应用程序在不同类型的调度器和管道之间移植, 你可以跳过这一部分。

您的应用程序(不是应用程序规范,而是实际的应用程序二进制文件)对TorchX没有任何依赖关系 (例如,/bin/echo 不使用TorchX,但可以为它创建一个echo_torchx.py组件)。

注意

torchx.runtime 是您在编写应用程序二进制文件时唯一应该使用的模块!

然而,由于TorchX本质上允许您的应用程序在任何地方运行,因此建议以一种与调度程序/基础架构无关的方式编写应用程序。

这通常意味着在与调度器/基础设施交互的接触点添加一个API层。 例如,以下应用程序不是无基础设施的

import boto3

def main(input_path: str):
   s3 = boto3.session.Session().client("s3")
   path = s3_input_path.split("/")
   bucket = path[0]
   key = "/".join(path[1:])
   s3.download_file(bucket, key, "/tmp/input")
   input = torch.load("/tmp/input")
   # ...<rest of code omitted for brevity>...

上述二进制文件隐含地假设input_path是一个AWS S3路径。要使该训练器与存储无关,一种方法是引入一个FileSystem抽象层。对于文件系统,像PyTorch Lightning这样的框架已经定义了io层(Lightning在后台使用fsspec)。上述二进制文件可以通过使用Lightning重写,使其与存储无关。

import pytorch_lightning.utilities.io as io

def main(input_url: str):
   fs = io.get_filesystem(input_url)
   with fs.open(input_url, "rb") as f:
       input = torch.load(f)
   # ...<rest of code omitted for brevity>...

现在 main 可以被称为 main("s3://foo/bar")main("file://foo/bar") 使其与存储在各种存储中的输入兼容。

FileSystem 中,已经存在定义文件系统抽象的库。 在 torchx.runtime 中,你会发现提供各种功能抽象的库或指向其他库的指针, 这些功能可能是你需要编写一个基础设施无关的应用程序时所需的。 理想情况下,torchx.runtime 中的功能应及时上游合并到像 lightning 这样的库中, 这些库旨在用于编写你的应用程序。但为这些抽象找到合适的永久归宿可能需要时间, 甚至可能需要创建一个新的开源项目。 在这一切发生之前,这些功能可以通过 torchx.runtime 模块成熟并供用户使用。

下一步

查看 快速入门指南 以了解如何创建和运行组件。

文档

访问 PyTorch 的全面开发人员文档

查看文档

教程

获取面向初学者和高级开发人员的深入教程

查看教程

资源

查找开发资源并解答您的问题

查看资源