Shortcuts

概述

../_images/components_diagram.jpg

注意

上图仅用于说明目的。并非所有功能目前都开箱即用。

该模块包含了一系列内置的TorchX组件。目录结构按组件类别组织。组件只是模板化的应用规范。可以将它们视为不同类型作业定义的工厂方法。在本模块中返回specs.AppDef的函数就是我们所说的组件。

您可以浏览torchx.components模块中的组件库 或访问我们的文档页面

使用内置函数

一旦你找到一个内置组件,你可以:

  1. 将组件作为作业运行

  2. 在工作流(管道)的上下文中使用组件

在这两种情况下,组件都将作为作业运行,区别在于作业将直接在调度程序上作为独立作业运行,或者在具有上游和/或下游依赖关系的工作流中作为“阶段”运行。

注意

根据组件的语义,作业可能是单节点或分布式的。例如,如果组件具有单一角色,其中role.num_replicas == 1,那么作业是单节点作业。如果组件具有多个角色和/或任何角色的num_replicas > 1,那么作业是多节点分布式作业。

不确定是否应该将组件作为作业或管道阶段运行? 使用这个经验法则:

  1. 刚开始吗?通过将其作为作业运行来熟悉该组件

  2. 需要作业依赖关系?将组件作为流水线阶段运行

  3. 不需要作业依赖关系?将组件作为作业运行

创作

由于组件只是一个返回specs.AppDef的Python函数,编写自己的组件就像编写一个遵循以下规则的Python函数一样简单:

  1. 组件函数必须返回一个specs.AppDef,并且必须指定返回类型

  2. All arguments of the component must be PEP 484 type annotated and the type must be one of
    1. 基本类型: int, float, str, bool

    2. 可选的原语:Optional[int], Optional[float], Optional[str]

    3. 基本类型的映射:Dict[Primitive_key, Primitive_value]

    4. 基本类型列表:List[Primitive_values]

    5. 可选集合:Optional[List], Optional[Dict]

    6. VAR_ARG: *arg (在将参数传递给入口点脚本时有用)

  3. (可选)一个google格式的文档字符串 (特别是参见function_with_pep484_type_annotations)。这个文档字符串纯粹是信息性的, 因为torchx cli使用它来自动生成一个信息性的--help消息,这在与他人共享组件时非常有用。 如果组件没有文档字符串,--help选项仍然有效,但参数将有一个固定的描述(见下文)。 请注意,当通过torchx.runner以编程方式运行组件时,torchx根本不会拾取文档字符串。

下面是一个启动DDP脚本的示例组件,它是torchx.components.dist.ddp()内置组件的简化版本。

import os
import torchx.specs as specs

def ddp(
    *script_args: str,
    image: str,
    script: str,
    host: str = "aws_p3.2xlarge",
    nnodes: int = 1,
    nproc_per_node: int = 1,
) -> specs.AppDef:
   return specs.AppDef(
       name=os.path.basename(script),
       roles=[
           spec.Role(
               name="trainer",
               image=image,
               resource=specs.named_resources[host],
               num_replicas=nnodes,
               entrypoint="python",
               args=[
                   "-m",
                   "torch.distributed.run",
                   "--rdzv_backend=c10d",
                   "--rdzv_endpoint=localhost:5900",
                   f"--nnodes={nnodes}",
                   f"--nprocs_per_node={nprocs_per_node}",
                   "-m",
                   script,
                   *script_args,
               ],
           ),
       ]
   )

假设上述组件保存在 example.py 中,我们可以运行 --help 如下:

$ torchx ./example.py:ddp --help
usage: torchx run ...torchx_params... ddp  [-h] --image IMAGE --script SCRIPT [--host HOST]
                                          [--nnodes NNODES] [--nproc_per_node NPROC_PER_NODE]
                                          ...

AppDef: ddp. TIP: improve this help string by adding a docstring ...<omitted for brevity>...

positional arguments:
  script_args           (required)

optional arguments:
  -h, --help            show this help message and exit
  --image IMAGE         (required)
  --script SCRIPT       (required)
  --host HOST           (default: aws_p3.2xlarge)
  --nnodes NNODES       (default: 1)
  --nproc_per_node NPROC_PER_NODE
                        (default: 1)

如果我们像这样包含一个文档字符串:

def ddp(...) -> specs.AppDef:
  """
  DDP Simplified.

  Args:
     image: name of the docker image containing the script + deps
     script: path of the script in the image
     script_args: arguments to the script
     host: machine type (one from named resources)
     nnodes: number of nodes to launch
     nproc_per_node: number of scripts to launch per node

  """

  # ... component body same as above ...
  pass

然后 --help 消息将反映文档字符串中的函数和参数描述,如下所示:

usage: torchx run ...torchx_params... ddp  [-h] --image IMAGE --script SCRIPT [--host HOST]
                                          [--nnodes NNODES] [--nproc_per_node NPROC_PER_NODE]
                                          ...

App spec: DDP simplified.

positional arguments:
  script_args           arguments to the script

optional arguments:
  -h, --help            show this help message and exit
  --image IMAGE         name of the docker image containing the script + deps
  --script SCRIPT       path of the script in the image
  --host HOST           machine type (one from named resources)
  --nnodes NNODES       number of nodes to launch
  --nproc_per_node NPROC_PER_NODE
                        number of scripts to launch per node

验证

要验证您是否正确定义了组件,您可以:

  1. (最简单)使用命令行界面(CLI)对你的组件的--help进行试运行:torchx run --dryrun ~/component.py:train --help

  2. 使用组件 linter (参见 dist_test.py 作为示例)

作为作业运行

你可以使用torchx 命令行界面将组件作为作业运行,或者使用torchx.runner以编程方式运行。两者是相同的,实际上cli在底层使用了runner,所以选择权在你手中。快速开始指南将带你了解入门的基础知识。

程序化运行

要以编程方式运行内置组件或您自己的组件,只需将组件作为常规的Python函数调用,并将其传递给torchx.runner。以下是一个调用utils.echo内置组件的示例:

from torchx.components.utils import echo
from torchx.runner import get_runner

get_runner().run(echo(msg="hello world"), scheduler="local_cwd")

CLI 运行(内置)

当从命令行运行组件时,您需要传递要调用的组件函数。 对于内置组件,其形式为{component_module}.{component_fn},其中 {component_module}是组件相对于torchx.components的模块路径, 而{component_fn}是该模块中的组件函数。因此,对于 torchx.components.utils.echo,我们会去掉torchx.components前缀并运行它。

$ torchx run utils.echo --msg "hello world"

查看 CLI 文档 获取更多信息。

CLI 运行(自定义)

要使用cli运行您的自定义组件,您必须使用稍微不同的语法形式{component_path}:{component_fn}。其中{component_path}是您的组件python文件的文件路径,{component_fn}是该文件中组件函数的名称。假设您的组件位于/home/bob/component.py,并且组件函数名为train(),您将这样运行它

# option 1. use absolute path
$ torchx run /home/bob/component.py:train --help

# option 2. let the shell do the expansion
$ torchx run ~/component.py:train --help

# option 3. same but after CWD to $HOME
$ cd ~/
$ torchx run ./component.py:train --help

# option 4. files can be relative to CWD
$ cd ~/
$ torchx run component.py:train --help

注意

如果你知道TorchX的安装目录,builtins也可以这样运行!

从CLI传递组件参数

由于组件只是简单的 Python 函数,以编程方式使用它们非常直接。 如上所述,当通过 CLI 的 run 子命令运行组件时,组件参数使用双破折号 + 参数名称语法作为程序参数传递(例如 --param1=1--param1 1)。 CLI 根据组件的文档字符串自动生成 argparse 解析器。以下是如何传递各种类型组件参数的摘要,假设组件定义如下:

# in comp.py
from typing import Dict, List
import torchx.specs as specs

def f(i: int, f: float, s: str, b: bool, l: List[str], d: Dict[str, str], *args) -> specs.AppDef:
   """
   Example component

   Args:
       i: int param
       f: float param
       s: string param
       b: bool param
       l: list param
       d: map param
       args: varargs param

   Returns: specs.AppDef
   """

   pass
  1. 帮助: torchx run comp.py:f --help

  2. 基本类型 (int, float, str): torchx run comp.py:f --i 1 --f 1.2 --s "bar"

  3. 布尔值: torchx run comp.py:f --b True (或 --b False)

  4. 映射:torchx run comp.py:f --d k1=v1,k2=v2,k3=v3

  5. 列表:torchx run comp.py:f --l a,b,c

  6. VAR_ARG: *args are passed as positionals rather than arguments, hence they are specified at the end of the command. The -- delimiter is used to start the VAR_ARGS section. This is useful when the component and application have the same arguments or when passing through the --help arg. Below are a few examples: * *args=["arg1", "arg2", "arg3"]: torchx run comp.py:f --i 1 arg1 arg2 arg3 * *args=["--flag", "arg1"]: torchx run comp.py:f --i 1 --flag arg1 `` * ``*args=["--help"]: torchx run comp.py:f -- --help * *args=["--i", "2"]: torchx run comp.py:f --i 1 -- --i 2

在管道中运行

torchx.pipelines 定义了适配器,这些适配器将 torchx 组件转换为目标管道平台中表示管道“阶段”的对象(有关支持的管道编排器列表,请参见 管道)。

其他资源

参见:

  1. 本模块中定义的组件作为说明性示例

  2. 定义你自己的组件 快速入门指南

  3. 组件最佳实践 指南

  4. 应用程序最佳实践 指南