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, 可选) – 要使用的框架。可以是 "distributed" 或 "rpc"。 "distributed" 需要一个 DistributedDataCollector 收集器，而 "rpc" 需要一个 RPCDataCollector。默认值为 "distributed"。

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

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

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

• submitit_collection_conf (dict, 可选) – 要传递给 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 BoundedTensorSpec
...     collector = DistributedDataCollector(
...         [EnvCreator(lambda: GymEnv("Pendulum-v1"))] * num_jobs,
...         policy=RandomPolicy(BoundedTensorSpec(-1, 1, shape=(1,))),
...         launcher="submitit_delayed",
...     )
...     for data in collector:
...         print(data)
...
>>> if __name__ == "__main__":
...     main()
...