Python PyTorch LocalElasticAgent用法及代碼示例

本文簡要介紹python語言中 torch.distributed.elastic.agent.server.local_elastic_agent.LocalElasticAgent 的用法。

用法:

class torch.distributed.elastic.agent.server.local_elastic_agent.LocalElasticAgent(spec, start_method='spawn', exit_barrier_timeout=300, log_dir=None)

處理 host-local 工人的 torchelastic.agent.server.ElasticAgent 的實現。此代理按主機部署，並配置為生成n 工作人員。使用 GPU 時，n 映射到主機上可用的 GPU 數量。

本地代理不與部署在其他主機上的其他本地代理進行通信，即使工作人員可以通信inter-host。 worker id 被解釋為本地進程。代理作為一個單元啟動和停止所有工作進程。

傳遞給工作函數的工作函數和參數必須與 python 多處理兼容。要將多處理數據結構傳遞給工作人員，您可以在與指定start_method 相同的多處理上下文中創建數據結構並將其作為函數參數傳遞。

exit_barrier_timeout 指定等待其他代理完成的時間量(以秒為單位)。這可以作為一個安全網來處理工人在不同時間完成的情況，以防止代理查看作為 scale-down 事件提前完成的工人。強烈建議用戶代碼處理確保以同步方式終止工作人員而不是依賴於exit_barrier_timeout。

示例啟動函數

def trainer(args) -> str:
    return "do train"

def main():
    start_method="spawn"
    shared_queue= multiprocessing.get_context(start_method).Queue()
    spec = WorkerSpec(
                role="trainer",
                local_world_size=nproc_per_process,
                entrypoint=trainer,
                args=("foobar",),
                ...<OTHER_PARAMS...>)
    agent = LocalElasticAgent(spec, start_method)
    results = agent.run()

    if results.is_failed():
        print("trainer failed")
    else:
        print(f"rank 0 return value: {results.return_values[0]}")
        # prints -> rank 0 return value: do train

示例啟動二進製文件

def main():
    spec = WorkerSpec(
                role="trainer",
                local_world_size=nproc_per_process,
                entrypoint="/usr/local/bin/trainer",
                args=("--trainer_args", "foobar"),
                ...<OTHER_PARAMS...>)
    agent = LocalElasticAgent(spec)
    results = agent.run()

    if not results.is_failed():
        print("binary launches do not have return values")