高级用法

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 项目中添加一个 entry_points 定义来注册这个自定义调度器。

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

注册命名资源

命名资源是一组预定义的资源规格，并被赋予一个字符串名称。这在你的集群拥有一组固定的实例类型时特别有用。例如，如果你在 AWS 上的深度学习训练 Kubernetes 集群仅由 p3.16xlarge（64 个虚拟 CPU、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)

要使这些资源定义可用，您需要通过 entry_points 进行注册：

# 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

自定义组件可以通过以下对entry_points的修改进行注册：

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

这条行注册了一个与模块 foo 关联的组。 TorchX 将递归地遍历与 my_project.bar 相关联的最低级别目录，并找到所有定义的组件。

注意

如果有两个注册表项，例如 foo = my_project.bar 和 test = my_project 将会有两组重叠的组件，具有不同的别名。

注册后，torchx cli 将通过以下方式显示已注册的组件：

$ torchx builtins

如果 my_project.bar 有以下目录结构：

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

并且 baz.py 定义了一个名为 trainer 的组件（函数）。然后可以以以下方式运行该组件：

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