高级用法

TorchX 定义了插件点，您可以配置 TorchX 以最佳方式支持您的基础设施设置。大多数配置通过 Python 的 入口点 完成。

注意

入口点需要安装包含它们的 Python 包。如果您没有 Python 包，建议您创建一个，以便您可以在团队和组织之间共享资源定义、调度器和组件。

以下描述的入口点可以在您的项目 setup.py 文件中指定为

from setuptools import setup

setup(
    name="project foobar",
    entry_points={
        "torchx.schedulers": [
            "my_scheduler = my.custom.scheduler:create_scheduler",
        ],
        "torchx.named_resources": [
            "gpu_x2 = my_module.resources:gpu_x2",
        ],
    }
)

注册自定义调度器

您可以通过实现 .. py::class torchx.schedulers.Scheduler 接口来实现自定义调度器。

该 create_scheduler 函数应具有以下函数签名

from torchx.schedulers import Scheduler

def create_scheduler(session_name: str, **kwargs: object) -> Scheduler:
    return MyScheduler(session_name, **kwargs)

然后，您可以通过将入口点定义添加到您的 Python 项目中来注册此自定义调度器。

# setup.py
...
entry_points={
    "torchx.schedulers": [
        "my_scheduler = my.custom.scheduler:create_schedule",
    ],
}

注册命名资源

命名资源是一组预定义的资源规范，这些规范被赋予一个字符串名称。当您的集群具有一组固定的实例类型时，这尤其有用。例如，如果您的 AWS 上的深度学习训练 Kubernetes 集群仅包含 p3.16xlarge（64 个 vcpu、8 个 gpu、488GB），那么您可能希望为容器列出 T 恤尺寸的资源规范，如下所示

from torchx.specs import Resource

def gpu_x1() -> Resource:
    return Resource(cpu=8,  gpu=1, memMB=61_000)

def gpu_x2() -> Resource:
    return Resource(cpu=16, gpu=2, memMB=122_000)

def gpu_x3() -> Resource:
    return Resource(cpu=32, gpu=4, memMB=244_000)

def gpu_x4() -> Resource:
    return Resource(cpu=64, gpu=8, memMB=488_000)

要使这些资源定义可用，您需要通过入口点注册它们

# setup.py
...
entry_points={
    "torchx.named_resources": [
        "gpu_x2 = my_module.resources:gpu_x2",
    ],
}

安装包含入口点定义的包后，就可以以以下方式使用命名资源

>>> from torchx.specs import get_named_resources
>>> get_named_resources("gpu_x2")
Resource(cpu=16, gpu=2, memMB=122000, ...)

# my_module.component
from torchx.specs import AppDef, Role, get_named_resources

def test_app(resource: str) -> AppDef:
    return AppDef(name="test_app", roles=[
        Role(
            name="...",
            image="...",
            resource=get_named_resources(resource),
        )
    ])

test_app("gpu_x2")

注册自定义组件

您可以使用 torchx CLI 编写和注册一组自定义组件，作为 CLI 的内置组件。这使得您可以自定义一组与您的团队或组织最相关的组件，并将其作为 CLI builtin 支持。这样，用户在运行时会看到您的自定义组件

$ torchx builtins

自定义组件可以通过 [torchx.components] 入口点注册。如果 my_project.bar 具有以下目录结构

$PROJECT_ROOT/my_project/bar/
    |- baz.py

并且 baz.py 只有一个名为 trainer 的组件（函数）

# baz.py
import torchx.specs as specs

def trainer(...) -> specs.AppDef: ...

并且入口点被添加为

# setup.py
...
entry_points={
    "torchx.components": [
        "foo = my_project.bar",
    ],
}

TorchX 将搜索模块 my_project.bar 以查找所有定义的组件，并将找到的组件分组到 foo.* 前缀下。在这种情况下，组件 my_project.bar.baz.trainer 将以 foo.baz.trainer 的名称注册。

注意

TorchX 只搜索 Python 包（包含 __init__.py 文件的目录），并且不会尝试递归进入命名空间包（没有 __init__.py 文件的目录）。但是，您可以注册顶级命名空间包。

torchx CLI 将通过以下方式显示注册的组件

$ torchx builtins
Found 1 builtin components:
1. foo.baz.trainer

自定义组件然后可以按以下方式使用

$ torchx run foo.baz.trainer -- --name "test app"

当您注册自己的组件时，TorchX 将不包括它自己的内置组件。要添加 TorchX 的内置组件，您必须指定另一个入口点，如下所示

# setup.py
...
entry_points={
    "torchx.components": [
        "foo = my_project.bar",
        "torchx = torchx.components",
    ],
}

这将添加回 TorchX 内置组件，但使用 torchx.* 组件名称前缀（例如 torchx.dist.ddp 与默认的 dist.ddp）。

如果有两个注册入口指向同一个组件，例如

# setup.py
...
entry_points={
    "torchx.components": [
        "foo = my_project.bar",
        "test = my_project",
    ],
}

那么，对于 my_project.bar 中的那些组件，将有两组重叠的组件，它们具有不同的前缀别名：foo.* 和 test.bar.*。具体来说，

$ torchx builtins
Found 2 builtin components:
1. foo.baz.trainer
2. test.bar.baz.trainer

要省略分组并使组件名称更短，请使用下划线（例如 _ 或 _0、_1 等）。例如

# setup.py
...
entry_points={
    "torchx.components": [
        "_0 = my_project.bar",
        "_1 = torchx.components",
    ],
}

这将使训练器组件以 baz.trainer 的形式公开（而不是 foo.baz.trainer），并将添加回 TorchX 的 vanilla 安装中的内置组件，但没有 torchx.* 前缀。

$ torchx builtins
Found 11 builtin components:
1. baz.trainer
2. dist.ddp
3. utils.python
4. ... <more builtins from torchx.components.* ...>