Shortcuts

基本概念

这描述了TorchX背后的高级概念和项目结构。 有关如何创建和运行应用程序的信息,请查看快速入门指南

项目结构

TorchX中的顶级模块包括:

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

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

  3. torchx.workspace: 处理用于远程执行的镜像补丁

  4. torchx.cli: 命令行工具

  5. torchx.runner: 给定一个应用规范,将应用作为作业提交到调度器上

  6. torchx.schedulers: 运行器支持的后端作业调度器

  7. torchx.pipelines: 适配器,将给定的应用规范转换为ML管道平台中的“阶段”

  8. torchx.runtime: 您可以在编写应用程序时使用的实用程序和抽象库(非应用程序规范)

下面是一个UML图

_images/torchx_module_uml.jpg

概念

AppDefs

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

下面是一个简单的specs.AppDef示例,它回显“hello world”

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中有两种访问runners的方式:

  1. CLI: torchx run ~/app_spec.py

  2. 以编程方式:torchx.runner.get_runner().run(appspec)

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

管道适配器

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

在大多数情况下,应用程序规范会映射到管道中的一个“阶段”(或节点)。 然而,高级组件,特别是那些具有自己小型控制流的组件(例如HPO),可能会映射到一个“子管道”或“内联管道”。 这些高级组件如何映射到管道的具体语义取决于目标管道平台。例如,如果 管道DSL允许从上游阶段动态地向管道添加阶段,那么TorchX可能会利用此功能将 子管道“内联”到主管道中。TorchX通常会尽力将应用程序规范适应目标管道平台中的最规范表示。

查看管道以获取支持的管道平台列表。

运行时间

重要

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之类的库中, 这些库旨在用于编写您的应用程序。但为这些抽象找到一个合适的永久归宿可能需要时间, 甚至可能需要创建一个全新的OSS项目。在此之前,这些功能可以通过torchx.runtime模块成熟并为用户所用。

下一步

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