進階用法

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",
    ],
}

註冊具名資源

具名資源是一組預先定義的資源規格，並賦予字串名稱。當您的叢集具有固定的執行個體類型集時，這特別有用。例如，如果您在 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 定義的套件後，就可以按以下方式使用具名資源

>>> 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 內建元件 支援成為可能。這樣一來，當使用者執行時，就會看到您的自訂元件

$ 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: ...

並且將 entrypoints 新增為

# 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",
    ],
}

這會將 trainer 元件公開為 baz.trainer（而不是 foo.baz.trainer），並以 Torchx 的原始安裝方式新增回內建元件，而不使用 torchx.* 前置詞。

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