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