submitit_delayed_launcher

class torchrl.collectors.distributed.submitit_delayed_launcher(num_jobs, framework='distributed', backend='gloo', tcpport='10003', submitit_main_conf: dict = {'slurm_cpus_per_task': 32, 'slurm_gpus_per_node': 1, 'slurm_partition': 'train', 'timeout_min': 10}, submitit_collection_conf: dict = {'slurm_cpus_per_task': 32, 'slurm_gpus_per_node': 0, 'slurm_partition': 'train', 'timeout_min': 10})

submitit 的延迟启动器。

在某些情况下，启动的作业无法自行生成其他作业，而这只能在跳转主机级别完成。

在这些情况下，可以使用 submitit_delayed_launcher() 预先启动收集器节点，这些节点将等待主工作程序提供启动指令。

参数:

• num_jobs (int) – 要启动的收集作业数。

• framework (str, optional) – 要使用的框架。可以是 "distributed" 或 "rpc"。"distributed" 需要 DistributedDataCollector 收集器，而 "rpc" 需要 RPCDataCollector。默认为 "distributed"。

• backend (str, optional) – 如果 framework 指向 "distributed"，则为 torch.distributed 后端。此值必须与传递给收集器的值匹配，否则主节点和卫星节点将无法进行 rendezvous 并永远挂起（即不会引发异常！）。默认为 'gloo'。

• tcpport (int 或 str, optional) – 要使用的 TCP 端口。默认为 torchrl.collectors.distributed.default_configs.TCP_PORT

• submitit_main_conf (dict, optional) – 要传递给 submitit 的主节点配置。默认为 torchrl.collectors.distributed.default_configs.DEFAULT_SLURM_CONF_MAIN

• submitit_collection_conf (dict, optional) – 要传递给 submitit 的配置。默认为 torchrl.collectors.distributed.default_configs.DEFAULT_SLURM_CONF

示例

>>> num_jobs=2
>>> @submitit_delayed_launcher(num_jobs=num_jobs)
... def main():
...     from torchrl.envs.utils import RandomPolicy
        from torchrl.envs.libs.gym import GymEnv
...     from torchrl.data import BoundedContinuous
...     collector = DistributedDataCollector(
...         [EnvCreator(lambda: GymEnv("Pendulum-v1"))] * num_jobs,
...         policy=RandomPolicy(BoundedContinuous(-1, 1, shape=(1,))),
...         launcher="submitit_delayed",
...     )
...     for data in collector:
...         print(data)
...
>>> if __name__ == "__main__":
...     main()
...