We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Deepspeed elastic launcher is incompatible with torch 2.0
**Describe the bug**

If I run `deepspeed` with `--elastic_training`, I got the following error:

```
Traceback (most recent call last):
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 723, in run
    result = self._invoke_run(role)
  File "./lib/python3.9/site-packages/deepspeed/elasticity/elastic_agent.py", line 126, in _invoke_run
    self._initialize_workers(self._worker_group)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 692, in _initialize_workers
    self._rendezvous(worker_group)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py", line 129, in wrapper
    result = f(*args, **kwargs)
  File "./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py", line 556, in _rendezvous
    self._set_master_addr_port(
TypeError: _set_master_addr_port() takes 3 positional arguments but 4 were given
```

The torch repository's pull request [#88922](https://github.com/pytorch/pytorch/pull/88922) introduces an additional parameter called `local_addr`. You can find the specific code changes related to this parameter in [this section](https://github.com/pytorch/pytorch/commit/5d37890b8eb952560f22b1b8ff661a024bf4b6d9#diff-0332dff4db45fd26518da1eddd4a93eb965a006de743edddba2316b64e2b1f6cR507) of the commit. However, it appears that the deepspeed project is still using the outdated method and has not incorporated this new parameter yet.

**To Reproduce**

Run any script with `--elastic_training`

**System info (please complete the following information):**
 - Python 3.9
 - torch 2.0.1

**Launcher context**
`deepspeed` launcher



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `deepspeed/elasticity/elastic_agent.py`
Content:
```
1 # Copyright (c) Microsoft Corporation.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 # DeepSpeed Team
5 
6 from torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent
7 from typing import Any, Dict, Optional, Tuple
8 from datetime import datetime
9 from torch.distributed.elastic.agent.server.api import log, _get_socket_with_port
10 from torch.distributed.elastic.metrics import put_metric
11 from torch.distributed.elastic.agent.server.api import (
12     RunResult,
13     WorkerGroup,
14     WorkerSpec,
15     WorkerState,
16 )
17 from torch.distributed import Store
18 import time
19 import os
20 from torch.distributed.elastic.multiprocessing import start_processes
21 from torch.distributed.elastic.utils import macros
22 import shutil
23 import copy
24 from contextlib import closing
25 import subprocess
26 
27 
28 class DSElasticAgent(LocalElasticAgent):
29 
30     def __init__(
31         self,
32         spec: WorkerSpec,
33         env: Dict,
34         start_method="spawn",
35         exit_barrier_timeout: float = 300,
36         log_dir: Optional[str] = None,
37     ):
38         super().__init__(spec, start_method, exit_barrier_timeout, log_dir)
39         self.ds_env = env
40 
41     @staticmethod
42     def _set_master_addr_port(store: Store, master_addr: Optional[str], master_port: Optional[int]):
43         if master_port is None:
44             sock = _get_socket_with_port()
45             with closing(sock):
46                 master_port = sock.getsockname()[1]
47 
48         if master_addr is None:
49             # master_addr = _get_fq_hostname()
50             result = subprocess.check_output("hostname -I", shell=True)
51             master_addr = result.decode('utf-8').split()[0]
52 
53         store.set("MASTER_ADDR", master_addr.encode(encoding="UTF-8"))
54         store.set("MASTER_PORT", str(master_port).encode(encoding="UTF-8"))
55 
56     def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]:
57         spec = worker_group.spec
58         store = worker_group.store
59         assert store is not None
60         master_addr, master_port = super()._get_master_addr_port(store)
61         restart_count = spec.max_restarts - self._remaining_restarts
62 
63         use_agent_store = spec.rdzv_handler.get_backend() == "static"
64 
65         args: Dict[int, Tuple] = {}
66         envs: Dict[int, Dict[str, str]] = {}
67         for worker in worker_group.workers:
68             local_rank = worker.local_rank
69 
70             worker_env_ds = copy.deepcopy(self.ds_env)
71             worker_env_elastic = {
72                 "LOCAL_RANK": str(local_rank),
73                 "RANK": str(worker.global_rank),
74                 "GROUP_RANK": str(worker_group.group_rank),
75                 "ROLE_RANK": str(worker.role_rank),
76                 "ROLE_NAME": spec.role,
77                 "LOCAL_WORLD_SIZE": str(spec.local_world_size),
78                 "WORLD_SIZE": str(worker.world_size),
79                 "GROUP_WORLD_SIZE": str(worker_group.group_world_size),
80                 "ROLE_WORLD_SIZE": str(worker.role_world_size),
81                 "MASTER_ADDR": master_addr,
82                 "MASTER_PORT": str(master_port),
83                 "TORCHELASTIC_RESTART_COUNT": str(restart_count),
84                 "TORCHELASTIC_MAX_RESTARTS": str(spec.max_restarts),
85                 "TORCHELASTIC_RUN_ID": spec.rdzv_handler.get_run_id(),
86                 "TORCHELASTIC_USE_AGENT_STORE": str(use_agent_store),
87                 "NCCL_ASYNC_ERROR_HANDLING": os.getenv("NCCL_ASYNC_ERROR_HANDLING", str(1)),
88             }
89             worker_env_ds.update(worker_env_elastic)
90             if "OMP_NUM_THREADS" in os.environ:
91                 worker_env_ds["OMP_NUM_THREADS"] = os.environ["OMP_NUM_THREADS"]
92 
93             envs[local_rank] = worker_env_ds
94             worker_args = list(spec.args)
95             worker_args = macros.substitute(worker_args, str(local_rank))
96             args[local_rank] = tuple(worker_args)
97 
98         # scaling events do not count towards restarts (gets same attempt #)
99         # remove existing log dir if this restart is due to a scaling event
100         attempt_log_dir = os.path.join(self._log_dir, f"attempt_{restart_count}")
101         shutil.rmtree(attempt_log_dir, ignore_errors=True)
102         os.makedirs(attempt_log_dir)
103 
104         assert spec.entrypoint is not None
105         self._pcontext = start_processes(
106             name=spec.role,
107             entrypoint=spec.entrypoint,
108             args=args,
109             envs=envs,
110             log_dir=attempt_log_dir,
111             start_method=self._start_method,
112             redirects=spec.redirects,
113             tee=spec.tee,
114         )
115 
116         return self._pcontext.pids()
117 
118     def _invoke_run(self, role: str = "default") -> RunResult:
119         # NOTE: currently only works for a single role
120 
121         spec = self._worker_group.spec
122         role = spec.role
123 
124         log.info(f"[{role}] starting workers for entrypoint: {spec.get_entrypoint_name()}")
125 
126         self._initialize_workers(self._worker_group)
127         monitor_interval = spec.monitor_interval
128         rdzv_handler = spec.rdzv_handler
129 
130         participants = rdzv_handler._state_holder.state.participants
131 
132         while True:
133             assert self._worker_group.state != WorkerState.INIT
134             time.sleep(monitor_interval)
135             run_result = self._monitor_workers(self._worker_group)
136             state = run_result.state
137             self._worker_group.state = state
138 
139             expire_time = datetime.utcnow() - (rdzv_handler._settings.keep_alive_interval *
140                                                rdzv_handler._settings.keep_alive_max_attempt)
141             _dead_nodes = [
142                 node for node, last_heartbeat in rdzv_handler._state_holder.state.last_heartbeats.items()
143                 if last_heartbeat < expire_time
144             ]
145 
146             put_metric(f"workers.{role}.remaining_restarts", self._remaining_restarts)
147             put_metric(f"workers.{role}.{state.name.lower()}", 1)
148 
149             if state == WorkerState.SUCCEEDED:
150                 log.info(f"[{role}] worker group successfully finished."
151                          f" Waiting {self._exit_barrier_timeout} seconds for other agents to finish.")
152                 self._exit_barrier()
153                 return run_result
154             elif state in {WorkerState.UNHEALTHY, WorkerState.FAILED
155                            } or len(participants) > len(rdzv_handler._state_holder.state.participants):
156                 if self._remaining_restarts > 0:
157                     log.info(f"[{role}] Worker group {state.name}. "
158                              f"{self._remaining_restarts}/{spec.max_restarts} attempts left;"
159                              f" will restart worker group")
160                     self._remaining_restarts -= 1
161                     # rdzv_handler._state_holder.state.restart = False
162                     self._restart_workers(self._worker_group)
163                     participants = rdzv_handler._state_holder.state.participants
164 
165                 else:
166                     self._stop_workers(self._worker_group)
167                     self._worker_group.state = WorkerState.FAILED
168                     self._exit_barrier()
169                     return run_result
170             elif state == WorkerState.HEALTHY:
171                 # membership changes do not count as retries
172                 num_nodes_waiting = rdzv_handler.num_nodes_waiting()
173                 group_rank = self._worker_group.group_rank
174                 if num_nodes_waiting > 0:
175                     log.info(f"[{role}] Detected {num_nodes_waiting} "
176                              f"new nodes from group_rank={group_rank}; "
177                              f"will restart worker group")
178                     self._restart_workers(self._worker_group)
179                     participants = rdzv_handler._state_holder.state.participants
180             else:
181                 raise Exception(f"[{role}] Worker group in {state.name} state")
182
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/deepspeed/elasticity/elastic_agent.py b/deepspeed/elasticity/elastic_agent.py
--- a/deepspeed/elasticity/elastic_agent.py
+++ b/deepspeed/elasticity/elastic_agent.py
@@ -39,7 +39,10 @@
         self.ds_env = env
 
     @staticmethod
-    def _set_master_addr_port(store: Store, master_addr: Optional[str], master_port: Optional[int]):
+    def _set_master_addr_port(store: Store,
+                              master_addr: Optional[str],
+                              master_port: Optional[int],
+                              local_addr: Optional[str] = None):
         if master_port is None:
             sock = _get_socket_with_port()
             with closing(sock):

{"golden_diff": "diff --git a/deepspeed/elasticity/elastic_agent.py b/deepspeed/elasticity/elastic_agent.py\n--- a/deepspeed/elasticity/elastic_agent.py\n+++ b/deepspeed/elasticity/elastic_agent.py\n@@ -39,7 +39,10 @@\n         self.ds_env = env\n \n     @staticmethod\n-    def _set_master_addr_port(store: Store, master_addr: Optional[str], master_port: Optional[int]):\n+    def _set_master_addr_port(store: Store,\n+                              master_addr: Optional[str],\n+                              master_port: Optional[int],\n+                              local_addr: Optional[str] = None):\n         if master_port is None:\n             sock = _get_socket_with_port()\n             with closing(sock):\n", "issue": "[BUG] Deepspeed elastic launcher is incompatible with torch 2.0\n**Describe the bug**\r\n\r\nIf I run `deepspeed` with `--elastic_training`, I got the following error:\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py\", line 129, in wrapper\r\n    result = f(*args, **kwargs)\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py\", line 723, in run\r\n    result = self._invoke_run(role)\r\n  File \"./lib/python3.9/site-packages/deepspeed/elasticity/elastic_agent.py\", line 126, in _invoke_run\r\n    self._initialize_workers(self._worker_group)\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py\", line 129, in wrapper\r\n    result = f(*args, **kwargs)\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py\", line 692, in _initialize_workers\r\n    self._rendezvous(worker_group)\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/metrics/api.py\", line 129, in wrapper\r\n    result = f(*args, **kwargs)\r\n  File \"./lib/python3.9/site-packages/torch/distributed/elastic/agent/server/api.py\", line 556, in _rendezvous\r\n    self._set_master_addr_port(\r\nTypeError: _set_master_addr_port() takes 3 positional arguments but 4 were given\r\n```\r\n\r\nThe torch repository's pull request [#88922](https://github.com/pytorch/pytorch/pull/88922) introduces an additional parameter called `local_addr`. You can find the specific code changes related to this parameter in [this section](https://github.com/pytorch/pytorch/commit/5d37890b8eb952560f22b1b8ff661a024bf4b6d9#diff-0332dff4db45fd26518da1eddd4a93eb965a006de743edddba2316b64e2b1f6cR507) of the commit. However, it appears that the deepspeed project is still using the outdated method and has not incorporated this new parameter yet.\r\n\r\n**To Reproduce**\r\n\r\nRun any script with `--elastic_training`\r\n\r\n**System info (please complete the following information):**\r\n - Python 3.9\r\n - torch 2.0.1\r\n\r\n**Launcher context**\r\n`deepspeed` launcher\r\n\r\n\n", "before_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nfrom torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent\nfrom typing import Any, Dict, Optional, Tuple\nfrom datetime import datetime\nfrom torch.distributed.elastic.agent.server.api import log, _get_socket_with_port\nfrom torch.distributed.elastic.metrics import put_metric\nfrom torch.distributed.elastic.agent.server.api import (\n    RunResult,\n    WorkerGroup,\n    WorkerSpec,\n    WorkerState,\n)\nfrom torch.distributed import Store\nimport time\nimport os\nfrom torch.distributed.elastic.multiprocessing import start_processes\nfrom torch.distributed.elastic.utils import macros\nimport shutil\nimport copy\nfrom contextlib import closing\nimport subprocess\n\n\nclass DSElasticAgent(LocalElasticAgent):\n\n    def __init__(\n        self,\n        spec: WorkerSpec,\n        env: Dict,\n        start_method=\"spawn\",\n        exit_barrier_timeout: float = 300,\n        log_dir: Optional[str] = None,\n    ):\n        super().__init__(spec, start_method, exit_barrier_timeout, log_dir)\n        self.ds_env = env\n\n    @staticmethod\n    def _set_master_addr_port(store: Store, master_addr: Optional[str], master_port: Optional[int]):\n        if master_port is None:\n            sock = _get_socket_with_port()\n            with closing(sock):\n                master_port = sock.getsockname()[1]\n\n        if master_addr is None:\n            # master_addr = _get_fq_hostname()\n            result = subprocess.check_output(\"hostname -I\", shell=True)\n            master_addr = result.decode('utf-8').split()[0]\n\n        store.set(\"MASTER_ADDR\", master_addr.encode(encoding=\"UTF-8\"))\n        store.set(\"MASTER_PORT\", str(master_port).encode(encoding=\"UTF-8\"))\n\n    def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]:\n        spec = worker_group.spec\n        store = worker_group.store\n        assert store is not None\n        master_addr, master_port = super()._get_master_addr_port(store)\n        restart_count = spec.max_restarts - self._remaining_restarts\n\n        use_agent_store = spec.rdzv_handler.get_backend() == \"static\"\n\n        args: Dict[int, Tuple] = {}\n        envs: Dict[int, Dict[str, str]] = {}\n        for worker in worker_group.workers:\n            local_rank = worker.local_rank\n\n            worker_env_ds = copy.deepcopy(self.ds_env)\n            worker_env_elastic = {\n                \"LOCAL_RANK\": str(local_rank),\n                \"RANK\": str(worker.global_rank),\n                \"GROUP_RANK\": str(worker_group.group_rank),\n                \"ROLE_RANK\": str(worker.role_rank),\n                \"ROLE_NAME\": spec.role,\n                \"LOCAL_WORLD_SIZE\": str(spec.local_world_size),\n                \"WORLD_SIZE\": str(worker.world_size),\n                \"GROUP_WORLD_SIZE\": str(worker_group.group_world_size),\n                \"ROLE_WORLD_SIZE\": str(worker.role_world_size),\n                \"MASTER_ADDR\": master_addr,\n                \"MASTER_PORT\": str(master_port),\n                \"TORCHELASTIC_RESTART_COUNT\": str(restart_count),\n                \"TORCHELASTIC_MAX_RESTARTS\": str(spec.max_restarts),\n                \"TORCHELASTIC_RUN_ID\": spec.rdzv_handler.get_run_id(),\n                \"TORCHELASTIC_USE_AGENT_STORE\": str(use_agent_store),\n                \"NCCL_ASYNC_ERROR_HANDLING\": os.getenv(\"NCCL_ASYNC_ERROR_HANDLING\", str(1)),\n            }\n            worker_env_ds.update(worker_env_elastic)\n            if \"OMP_NUM_THREADS\" in os.environ:\n                worker_env_ds[\"OMP_NUM_THREADS\"] = os.environ[\"OMP_NUM_THREADS\"]\n\n            envs[local_rank] = worker_env_ds\n            worker_args = list(spec.args)\n            worker_args = macros.substitute(worker_args, str(local_rank))\n            args[local_rank] = tuple(worker_args)\n\n        # scaling events do not count towards restarts (gets same attempt #)\n        # remove existing log dir if this restart is due to a scaling event\n        attempt_log_dir = os.path.join(self._log_dir, f\"attempt_{restart_count}\")\n        shutil.rmtree(attempt_log_dir, ignore_errors=True)\n        os.makedirs(attempt_log_dir)\n\n        assert spec.entrypoint is not None\n        self._pcontext = start_processes(\n            name=spec.role,\n            entrypoint=spec.entrypoint,\n            args=args,\n            envs=envs,\n            log_dir=attempt_log_dir,\n            start_method=self._start_method,\n            redirects=spec.redirects,\n            tee=spec.tee,\n        )\n\n        return self._pcontext.pids()\n\n    def _invoke_run(self, role: str = \"default\") -> RunResult:\n        # NOTE: currently only works for a single role\n\n        spec = self._worker_group.spec\n        role = spec.role\n\n        log.info(f\"[{role}] starting workers for entrypoint: {spec.get_entrypoint_name()}\")\n\n        self._initialize_workers(self._worker_group)\n        monitor_interval = spec.monitor_interval\n        rdzv_handler = spec.rdzv_handler\n\n        participants = rdzv_handler._state_holder.state.participants\n\n        while True:\n            assert self._worker_group.state != WorkerState.INIT\n            time.sleep(monitor_interval)\n            run_result = self._monitor_workers(self._worker_group)\n            state = run_result.state\n            self._worker_group.state = state\n\n            expire_time = datetime.utcnow() - (rdzv_handler._settings.keep_alive_interval *\n                                               rdzv_handler._settings.keep_alive_max_attempt)\n            _dead_nodes = [\n                node for node, last_heartbeat in rdzv_handler._state_holder.state.last_heartbeats.items()\n                if last_heartbeat < expire_time\n            ]\n\n            put_metric(f\"workers.{role}.remaining_restarts\", self._remaining_restarts)\n            put_metric(f\"workers.{role}.{state.name.lower()}\", 1)\n\n            if state == WorkerState.SUCCEEDED:\n                log.info(f\"[{role}] worker group successfully finished.\"\n                         f\" Waiting {self._exit_barrier_timeout} seconds for other agents to finish.\")\n                self._exit_barrier()\n                return run_result\n            elif state in {WorkerState.UNHEALTHY, WorkerState.FAILED\n                           } or len(participants) > len(rdzv_handler._state_holder.state.participants):\n                if self._remaining_restarts > 0:\n                    log.info(f\"[{role}] Worker group {state.name}. \"\n                             f\"{self._remaining_restarts}/{spec.max_restarts} attempts left;\"\n                             f\" will restart worker group\")\n                    self._remaining_restarts -= 1\n                    # rdzv_handler._state_holder.state.restart = False\n                    self._restart_workers(self._worker_group)\n                    participants = rdzv_handler._state_holder.state.participants\n\n                else:\n                    self._stop_workers(self._worker_group)\n                    self._worker_group.state = WorkerState.FAILED\n                    self._exit_barrier()\n                    return run_result\n            elif state == WorkerState.HEALTHY:\n                # membership changes do not count as retries\n                num_nodes_waiting = rdzv_handler.num_nodes_waiting()\n                group_rank = self._worker_group.group_rank\n                if num_nodes_waiting > 0:\n                    log.info(f\"[{role}] Detected {num_nodes_waiting} \"\n                             f\"new nodes from group_rank={group_rank}; \"\n                             f\"will restart worker group\")\n                    self._restart_workers(self._worker_group)\n                    participants = rdzv_handler._state_holder.state.participants\n            else:\n                raise Exception(f\"[{role}] Worker group in {state.name} state\")\n", "path": "deepspeed/elasticity/elastic_agent.py"}], "after_files": [{"content": "# Copyright (c) Microsoft Corporation.\n# SPDX-License-Identifier: Apache-2.0\n\n# DeepSpeed Team\n\nfrom torch.distributed.elastic.agent.server.local_elastic_agent import LocalElasticAgent\nfrom typing import Any, Dict, Optional, Tuple\nfrom datetime import datetime\nfrom torch.distributed.elastic.agent.server.api import log, _get_socket_with_port\nfrom torch.distributed.elastic.metrics import put_metric\nfrom torch.distributed.elastic.agent.server.api import (\n    RunResult,\n    WorkerGroup,\n    WorkerSpec,\n    WorkerState,\n)\nfrom torch.distributed import Store\nimport time\nimport os\nfrom torch.distributed.elastic.multiprocessing import start_processes\nfrom torch.distributed.elastic.utils import macros\nimport shutil\nimport copy\nfrom contextlib import closing\nimport subprocess\n\n\nclass DSElasticAgent(LocalElasticAgent):\n\n    def __init__(\n        self,\n        spec: WorkerSpec,\n        env: Dict,\n        start_method=\"spawn\",\n        exit_barrier_timeout: float = 300,\n        log_dir: Optional[str] = None,\n    ):\n        super().__init__(spec, start_method, exit_barrier_timeout, log_dir)\n        self.ds_env = env\n\n    @staticmethod\n    def _set_master_addr_port(store: Store,\n                              master_addr: Optional[str],\n                              master_port: Optional[int],\n                              local_addr: Optional[str] = None):\n        if master_port is None:\n            sock = _get_socket_with_port()\n            with closing(sock):\n                master_port = sock.getsockname()[1]\n\n        if master_addr is None:\n            # master_addr = _get_fq_hostname()\n            result = subprocess.check_output(\"hostname -I\", shell=True)\n            master_addr = result.decode('utf-8').split()[0]\n\n        store.set(\"MASTER_ADDR\", master_addr.encode(encoding=\"UTF-8\"))\n        store.set(\"MASTER_PORT\", str(master_port).encode(encoding=\"UTF-8\"))\n\n    def _start_workers(self, worker_group: WorkerGroup) -> Dict[int, Any]:\n        spec = worker_group.spec\n        store = worker_group.store\n        assert store is not None\n        master_addr, master_port = super()._get_master_addr_port(store)\n        restart_count = spec.max_restarts - self._remaining_restarts\n\n        use_agent_store = spec.rdzv_handler.get_backend() == \"static\"\n\n        args: Dict[int, Tuple] = {}\n        envs: Dict[int, Dict[str, str]] = {}\n        for worker in worker_group.workers:\n            local_rank = worker.local_rank\n\n            worker_env_ds = copy.deepcopy(self.ds_env)\n            worker_env_elastic = {\n                \"LOCAL_RANK\": str(local_rank),\n                \"RANK\": str(worker.global_rank),\n                \"GROUP_RANK\": str(worker_group.group_rank),\n                \"ROLE_RANK\": str(worker.role_rank),\n                \"ROLE_NAME\": spec.role,\n                \"LOCAL_WORLD_SIZE\": str(spec.local_world_size),\n                \"WORLD_SIZE\": str(worker.world_size),\n                \"GROUP_WORLD_SIZE\": str(worker_group.group_world_size),\n                \"ROLE_WORLD_SIZE\": str(worker.role_world_size),\n                \"MASTER_ADDR\": master_addr,\n                \"MASTER_PORT\": str(master_port),\n                \"TORCHELASTIC_RESTART_COUNT\": str(restart_count),\n                \"TORCHELASTIC_MAX_RESTARTS\": str(spec.max_restarts),\n                \"TORCHELASTIC_RUN_ID\": spec.rdzv_handler.get_run_id(),\n                \"TORCHELASTIC_USE_AGENT_STORE\": str(use_agent_store),\n                \"NCCL_ASYNC_ERROR_HANDLING\": os.getenv(\"NCCL_ASYNC_ERROR_HANDLING\", str(1)),\n            }\n            worker_env_ds.update(worker_env_elastic)\n            if \"OMP_NUM_THREADS\" in os.environ:\n                worker_env_ds[\"OMP_NUM_THREADS\"] = os.environ[\"OMP_NUM_THREADS\"]\n\n            envs[local_rank] = worker_env_ds\n            worker_args = list(spec.args)\n            worker_args = macros.substitute(worker_args, str(local_rank))\n            args[local_rank] = tuple(worker_args)\n\n        # scaling events do not count towards restarts (gets same attempt #)\n        # remove existing log dir if this restart is due to a scaling event\n        attempt_log_dir = os.path.join(self._log_dir, f\"attempt_{restart_count}\")\n        shutil.rmtree(attempt_log_dir, ignore_errors=True)\n        os.makedirs(attempt_log_dir)\n\n        assert spec.entrypoint is not None\n        self._pcontext = start_processes(\n            name=spec.role,\n            entrypoint=spec.entrypoint,\n            args=args,\n            envs=envs,\n            log_dir=attempt_log_dir,\n            start_method=self._start_method,\n            redirects=spec.redirects,\n            tee=spec.tee,\n        )\n\n        return self._pcontext.pids()\n\n    def _invoke_run(self, role: str = \"default\") -> RunResult:\n        # NOTE: currently only works for a single role\n\n        spec = self._worker_group.spec\n        role = spec.role\n\n        log.info(f\"[{role}] starting workers for entrypoint: {spec.get_entrypoint_name()}\")\n\n        self._initialize_workers(self._worker_group)\n        monitor_interval = spec.monitor_interval\n        rdzv_handler = spec.rdzv_handler\n\n        participants = rdzv_handler._state_holder.state.participants\n\n        while True:\n            assert self._worker_group.state != WorkerState.INIT\n            time.sleep(monitor_interval)\n            run_result = self._monitor_workers(self._worker_group)\n            state = run_result.state\n            self._worker_group.state = state\n\n            expire_time = datetime.utcnow() - (rdzv_handler._settings.keep_alive_interval *\n                                               rdzv_handler._settings.keep_alive_max_attempt)\n            _dead_nodes = [\n                node for node, last_heartbeat in rdzv_handler._state_holder.state.last_heartbeats.items()\n                if last_heartbeat < expire_time\n            ]\n\n            put_metric(f\"workers.{role}.remaining_restarts\", self._remaining_restarts)\n            put_metric(f\"workers.{role}.{state.name.lower()}\", 1)\n\n            if state == WorkerState.SUCCEEDED:\n                log.info(f\"[{role}] worker group successfully finished.\"\n                         f\" Waiting {self._exit_barrier_timeout} seconds for other agents to finish.\")\n                self._exit_barrier()\n                return run_result\n            elif state in {WorkerState.UNHEALTHY, WorkerState.FAILED\n                           } or len(participants) > len(rdzv_handler._state_holder.state.participants):\n                if self._remaining_restarts > 0:\n                    log.info(f\"[{role}] Worker group {state.name}. \"\n                             f\"{self._remaining_restarts}/{spec.max_restarts} attempts left;\"\n                             f\" will restart worker group\")\n                    self._remaining_restarts -= 1\n                    # rdzv_handler._state_holder.state.restart = False\n                    self._restart_workers(self._worker_group)\n                    participants = rdzv_handler._state_holder.state.participants\n\n                else:\n                    self._stop_workers(self._worker_group)\n                    self._worker_group.state = WorkerState.FAILED\n                    self._exit_barrier()\n                    return run_result\n            elif state == WorkerState.HEALTHY:\n                # membership changes do not count as retries\n                num_nodes_waiting = rdzv_handler.num_nodes_waiting()\n                group_rank = self._worker_group.group_rank\n                if num_nodes_waiting > 0:\n                    log.info(f\"[{role}] Detected {num_nodes_waiting} \"\n                             f\"new nodes from group_rank={group_rank}; \"\n                             f\"will restart worker group\")\n                    self._restart_workers(self._worker_group)\n                    participants = rdzv_handler._state_holder.state.participants\n            else:\n                raise Exception(f\"[{role}] Worker group in {state.name} state\")\n", "path": "deepspeed/elasticity/elastic_agent.py"}]}