Resolve bug with agent_restarts where it was equal to max_restarts

Summary:
Resolve bug with agent_restarts where it was equal to max_restarts

Make sure that different schedulers start retry attempt with the same initial value

Reviewed By: kiukchung

Differential Revision: D26431772

fbshipit-source-id: 51cd4d59a680c505d2bb0863dc0510df17e8a472

Author: aivanou

torchelastic/agent/server/api.py

@@ -745,7 +745,7 @@ def _construct_event(
             "rdzv_backend": spec.rdzv_handler.get_backend(),
             "raw_error": raw_error,
             "metadata": md_str,
-            "agent_restarts": spec.max_restarts,
+            "agent_restarts": spec.max_restarts - self._remaining_restarts,
         }
         return Event(
             f"torchelastic.worker.status.{state}", source=source, metadata=metadata

torchelastic/agent/server/test/api_test.py

@@ -541,8 +541,9 @@ def test_get_agent_status_event(self):
         self.assertEqual(spec.role, actual_event.metadata["role"])
 
     def test_get_worker_status_event(self):
-        spec = self._get_worker_spec(max_restarts=1)
+        spec = self._get_worker_spec(max_restarts=4)
         agent = TestAgent(spec)
+        agent._remaining_restarts = spec.max_restarts - 2
         actual_event = agent._construct_event(
             state=WorkerState.SUCCEEDED.value,
             source="WORKER",
@@ -552,3 +553,4 @@ def test_get_worker_status_event(self):
         self.assertEqual("etcd", actual_event.metadata["rdzv_backend"])
         self.assertEqual(WorkerState.SUCCEEDED.value, actual_event.metadata["state"])
         self.assertEqual(spec.role, actual_event.metadata["role"])
+        self.assertEqual(2, actual_event.metadata["agent_restarts"])