多程序

一個庫，用於啟動和管理 n 個由函式或二進位制檔案指定的 worker 子程序。

對於函式，它使用 torch.multiprocessing（以及 Python 的 multiprocessing）來 spawn/fork worker 程序。對於二進位制檔案，它使用 Python 的 subprocessing.Popen 來建立 worker 程序。

用法 1：以函式形式啟動兩個訓練器

from torch.distributed.elastic.multiprocessing import Std, start_processes


def trainer(a, b, c):
    pass  # train


# runs two trainers
# LOCAL_RANK=0 trainer(1,2,3)
# LOCAL_RANK=1 trainer(4,5,6)
ctx = start_processes(
    name="trainer",
    entrypoint=trainer,
    args={0: (1, 2, 3), 1: (4, 5, 6)},
    envs={0: {"LOCAL_RANK": 0}, 1: {"LOCAL_RANK": 1}},
    log_dir="/tmp/foobar",
    redirects=Std.ALL,  # write all worker stdout/stderr to a log file
    tee={0: Std.ERR},  # tee only local rank 0's stderr to console
)

# waits for all copies of trainer to finish
ctx.wait()

用法 2：以二進位制形式啟動 2 個 echo worker

# same as invoking
# echo hello
# echo world > stdout.log
ctx = start_processes(
        name="echo"
        entrypoint="echo",
        log_dir="/tmp/foobar",
        args={0: "hello", 1: "world"},
        redirects={1: Std.OUT},
       )

與 torch.multiprocessing 類似，函式 start_processes() 的返回值是一個程序上下文（api.PContext）。如果啟動的是函式，則返回 api.MultiprocessContext；如果啟動的是二進位制檔案，則返回 api.SubprocessContext。兩者都是父類 api.PContext 的具體實現。

啟動多個 Worker

torch.distributed.elastic.multiprocessing.start_processes(name, entrypoint, args, envs, logs_specs, log_line_prefixes=None, start_method='spawn')

使用提供的選項啟動 n 個 entrypoint 程序的副本。

entrypoint 可以是一個 Callable（函式）或一個 str（二進位制檔案）。副本的數量由 args 和 envs 引數中的條目數量決定，它們需要具有相同的鍵集。

args 和 env 引數是傳遞給由副本索引（本地 rank）對映的 entrypoint 的引數和環境變數。所有本地 rank 都必須包含在內。也就是說，鍵集應該是 {0,1,...,(nprocs-1)}。

注意

當 entrypoint 是一個二進位制檔案（str）時，args 只能是字串。如果給定任何其他型別，它將被轉換為字串表示（例如 str(arg1)）。此外，只有當主函式使用 torch.distributed.elastic.multiprocessing.errors.record 進行註解時，二進位制檔案失敗才會寫入 error.json 錯誤檔案。對於函式啟動，這是預設行為，無需手動新增 @record 註解。

redirects 和 tee 是位掩碼，指定要重定向到 log_dir 中的日誌檔案的一個或多個標準流 (std stream)。有效的掩碼值在 Std 中定義。若要僅重定向/tee 特定本地 rank 的流，請將 redirects 作為以本地 rank 為鍵的對映傳遞，以指定其重定向行為。任何缺失的本地 rank 將預設為 Std.NONE。

tee 的作用類似於 Unix 命令 “tee”，它會重定向並列印到控制檯。要避免 worker 的 stdout/stderr 列印到控制檯，請使用 redirects 引數。

對於每個程序，log_dir 將包含

1. {local_rank}/error.json：如果程序失敗，則此檔案包含錯誤資訊

2. {local_rank}/stdout.json：如果 redirect & STDOUT == STDOUT

3. {local_rank}/stderr.json：如果 redirect & STDERR == STDERR

注意

期望 log_dir 存在、為空且是一個目錄。

示例

log_dir = "/tmp/test"

# ok; two copies of foo: foo("bar0"), foo("bar1")
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
)

# invalid; envs missing for local rank 1
start_processes(
   name="trainer",
   entrypoint=foo,
   args:{0:("bar0",), 1:("bar1",),
   envs:{0:{}},
   log_dir=log_dir
)

# ok; two copies of /usr/bin/touch: touch file1, touch file2
start_processes(
   name="trainer",
   entrypoint="/usr/bin/touch",
   args:{0:("file1",), 1:("file2",),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )

# caution; arguments casted to string, runs:
# echo "1" "2" "3" and echo "[1, 2, 3]"
start_processes(
   name="trainer",
   entrypoint="/usr/bin/echo",
   args:{0:(1,2,3), 1:([1,2,3],),
   envs:{0:{}, 1:{}},
   log_dir=log_dir
 )

引數

• name (str) – 一個人類可讀的短名稱，描述程序是什麼（在 teeing stdout/stderr 輸出時用作標題）

• entrypoint (Union[Callable, str]) – 可以是 Callable（函式）或 cmd（二進位制檔案）

• args (dict[int, tuple]) – 傳遞給每個副本的引數

• envs (dict[int, dict[str, str]]) – 傳遞給每個副本的環境變數

• log_dir – 用於寫入日誌檔案的目錄

• start_method (str) – 多程序啟動方法 (spawn, fork, forkserver)，對於二進位制檔案無效

• redirects – 要重定向到日誌檔案的標準流

• tee – 要重定向並列印到控制檯的標準流

• local_ranks_filter – 要列印到控制檯的哪些 rank 的日誌

返回型別

PContext

程序上下文

class torch.distributed.elastic.multiprocessing.api.PContext(name, entrypoint, args, envs, logs_specs, log_line_prefixes=None)

用於標準化透過不同機制啟動的一組程序上的操作的基類。

名稱 PContext 是有意為之，以區分 torch.multiprocessing.ProcessContext。

警告

stdouts 和 stderrs 應該始終是 tee_stdouts 和 tee_stderrs（分別）的超集，這是因為 tee 被實現為重定向 + tail -f <stdout/stderr.log>

class torch.distributed.elastic.multiprocessing.api.MultiprocessContext(name, entrypoint, args, envs, start_method, logs_specs, log_line_prefixes=None)

包含以函式形式呼叫的 worker 程序的 PContext。

class torch.distributed.elastic.multiprocessing.api.SubprocessContext(name, entrypoint, args, envs, logs_specs, log_line_prefixes=None)

包含以二進位制檔案形式呼叫的 worker 程序的 PContext。

class torch.distributed.elastic.multiprocessing.api.RunProcsResult(return_values=<factory>, failures=<factory>, stdouts=<factory>, stderrs=<factory>)

使用 start_processes() 啟動的程序完成執行的結果。由 PContext 返回。

請注意以下事項

1. 所有欄位均按本地 rank 對映

2. return_values - 僅為函式填充（不包括二進位制檔案）。

3. stdouts - stdout.log 的路徑（如果沒有重定向則為空字串）

4. stderrs - stderr.log 的路徑（如果沒有重定向則為空字串）

class torch.distributed.elastic.multiprocessing.api.DefaultLogsSpecs(log_dir=None, redirects=Std.NONE, tee=Std.NONE, local_ranks_filter=None)

Default LogsSpecs 實現

• 如果 log_dir 不存在，則會建立它

• 為每次嘗試和每個 rank 生成巢狀資料夾。

reify(envs)

使用以下方案構建日誌目標路徑

• <log_dir>/<rdzv_run_id>/attempt_<attempt>/<rank>/stdout.log

• <log_dir>/<rdzv_run_id>/attempt_<attempt>/<rank>/stderr.log

• <log_dir>/<rdzv_run_id>/attempt_<attempt>/<rank>/error.json

返回型別

LogsDest

class torch.distributed.elastic.multiprocessing.api.LogsDest(stdouts=<factory>, stderrs=<factory>, tee_stdouts=<factory>, tee_stderrs=<factory>, error_files=<factory>)

對於每種日誌型別，儲存著本地 rank ID 到檔案路徑的對映。

class torch.distributed.elastic.multiprocessing.api.LogsSpecs(log_dir=None, redirects=Std.NONE, tee=Std.NONE, local_ranks_filter=None)

定義每個 worker 程序的日誌處理和重定向。

引數

• log_dir (Optional[str]) – 將寫入日誌的基本目錄。

• redirects (Union[Std, dict[int, torch.distributed.elastic.multiprocessing.api.Std]]) – 要重定向到檔案的流。傳遞單個 Std 列舉以重定向所有 worker 的流，或傳遞以 local_rank 為鍵的對映以選擇性地重定向。

• tee (Union[Std, dict[int, torch.distributed.elastic.multiprocessing.api.Std]]) – 要複製到 stdout/stderr 的流。傳遞單個 Std 列舉以複製所有 worker 的流，或傳遞以 local_rank 為鍵的對映以選擇性地複製。

abstract reify(envs)

給定環境變數，為每個本地 rank 構建日誌檔案的目標路徑。

Envs 引數包含每個本地 rank 的環境變數字典，其中的條目在 _start_workers() 中定義。

返回型別

LogsDest