高级用法

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)

然后，您可以通过添加 entry_points 定义来注册此自定义调度程序 添加到您的 Python 项目中。

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

注册命名资源

命名资源是一组预定义的资源规范，它们被赋予了 string name 的 S S S T这特别有用 当您的集群具有一组固定的实例类型时。例如，如果您的 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)

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

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

安装具有 entry_points 定义的软件包后，命名的 然后，可以按以下方式使用 resource：

>>> 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")

注册自定义组件

可以使用 CLI 编写和注册一组自定义组件作为 CLI 的内置组件。这使得定制成为可能 一组与您的团队或组织和支持最相关的组件 它作为 CLI 。这样，用户将看到您的自定义组件 当他们运行时torchxbuiltin

$ torchx builtins

自定义组件可以通过以下修改来注册：entry_points

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

上面的行注册了一个与 module 关联的组。 TorchX 将递归遍历与 关联的最低级目录，并找到 所有定义的组件。foomy_project.barmy_project.bar

注意

如果有两个注册表项，例如 并且将有两组具有不同别名的重叠组件。foo = my_project.bartest = my_project

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

$ torchx builtins

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

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

并定义了一个名为 .然后，该组件可以按以下方式作为作业运行：baz.pytrainer

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