Fix bug in runner where server ips are not gotten if stage is retried (#1330)

Summary:
Pull Request resolved: #1330

## What
* always get the server_ips when running a joint stage

## Why
* if a joint stage was retried and had STARTED status, no server ips would be gotten, causing an error

Reviewed By: jrodal98

Differential Revision: D38054016

fbshipit-source-id: 9c0ac7867e9c1d576d21e37f5fbc559983d0fcee

Author: Andrew Zhuang

fbpcs/bolt/bolt_runner.py

@@ -141,22 +141,22 @@ async def run_next_stage(
         publisher_status = (
             await self.publisher_client.update_instance(publisher_id)
         ).pc_instance_status
-        server_ips = None
         if publisher_status not in [stage.started_status, stage.completed_status]:
             # don't retry if started or completed status
             self.logger.info(f"Publisher {publisher_id} starting stage {stage.name}.")
             await self.publisher_client.run_stage(instance_id=publisher_id, stage=stage)
-            if stage.is_joint_stage:
-                server_ips = await self.get_server_ips_after_start(
-                    instance_id=publisher_id,
-                    stage=stage,
-                    timeout=stage.timeout,
-                    poll_interval=poll_interval,
+        server_ips = None
+        if stage.is_joint_stage:
+            server_ips = await self.get_server_ips_after_start(
+                instance_id=publisher_id,
+                stage=stage,
+                timeout=stage.timeout,
+                poll_interval=poll_interval,
+            )
+            if server_ips is None:
+                raise NoServerIpsException(
+                    f"{stage.name} requires server ips but got none."
                 )
-                if server_ips is None:
-                    raise NoServerIpsException(
-                        f"{stage.name} requires server ips but got none."
-                    )
         partner_status = (
             await self.partner_client.update_instance(partner_id)
         ).pc_instance_status

fbpcs/bolt/test/test_bolt_runner.py

@@ -121,6 +121,34 @@ async def test_non_joint_stage(
             server_ips=None,
         )
 
+    @mock.patch("fbpcs.bolt.bolt_runner.asyncio.sleep")
+    async def test_joint_stage_retry_gets_ips(self, mock_sleep) -> None:
+        # test that server ips are gotten when a joint stage is retried with STARTED status
+        # specifically, publisher status STARTED and partner status FAILED
+        server_ips = ["1.1.1.1"]
+        self.test_runner.get_server_ips_after_start = mock.AsyncMock(
+            return_value=server_ips
+        )
+        self.test_runner.publisher_client.update_instance = mock.AsyncMock(
+            return_value=BoltState(DummyJointStageFlow.JOINT_STAGE.started_status)
+        )
+        self.test_runner.partner_client.update_instance = mock.AsyncMock(
+            return_value=BoltState(DummyJointStageFlow.JOINT_STAGE.failed_status)
+        )
+        mock_partner_run_stage = mock.AsyncMock()
+        self.test_runner.partner_client.run_stage = mock_partner_run_stage
+        await self.test_runner.run_next_stage(
+            publisher_id="publisher_id",
+            partner_id="partner_id",
+            stage=DummyJointStageFlow.JOINT_STAGE,
+            poll_interval=5,
+        )
+        mock_partner_run_stage.assert_called_once_with(
+            instance_id="partner_id",
+            stage=DummyJointStageFlow.JOINT_STAGE,
+            server_ips=server_ips,
+        )
+
     @mock.patch("fbpcs.bolt.bolt_runner.asyncio.sleep")
     @mock.patch("fbpcs.bolt.bolt_job.BoltPlayerArgs")
     @mock.patch("fbpcs.bolt.bolt_job.BoltPlayerArgs")