Wrap onedocker_svc.get_containers with Throttling Retry (#2219)

Summary:
Pull Request resolved: #2219

## Why
onedocker_svc.get_containers might cause Throttling, wrap the retry on it.

## What
- Wrap onedocker_svc.get_containers with Throttling Retry

Reviewed By: marksliva

Differential Revision: D44269329

fbshipit-source-id: 4cb6f2215d477e847abf831f625d78ebff3e0101

Author: joe1234wu

fbpcs/common/entity/stage_state_instance.py

@@ -12,10 +12,12 @@
 from typing import List, Optional, Union
 
 from fbpcp.entity.container_instance import ContainerInstance, ContainerInstanceStatus
+from fbpcp.error.pcp import ThrottlingError
 from fbpcp.service.onedocker import OneDockerService
 from fbpcp.util.typing import checked_cast
 from fbpcs.common.entity.instance_base import InstanceBase
 from fbpcs.common.entity.pcs_container_instance import PCSContainerInstance
+from fbpcs.common.service.retry_handler import BackoffType, RetryHandler
 
 
 class StageStateInstanceStatus(Enum):
@@ -88,9 +90,14 @@ def _get_updated_containers(
         self, onedocker_svc: OneDockerService
     ) -> List[ContainerInstance]:
         containers_to_update = self.get_containers_to_update(self.containers)
-        updated_containers = onedocker_svc.get_containers(
-            [self.containers[idx].instance_id for idx in containers_to_update]
-        )
+        with RetryHandler(
+            ThrottlingError, backoff_type=BackoffType.LINEAR
+        ) as retry_handler:
+            updated_containers = retry_handler.execute_sync(
+                onedocker_svc.get_containers,
+                [self.containers[idx].instance_id for idx in containers_to_update],
+            )
+
         new_updated_containers = self.containers.copy()
         for i, updated_container in zip(containers_to_update, updated_containers):
             if updated_container is not None:

fbpcs/private_computation/service/run_binary_base_service.py

@@ -16,7 +16,7 @@
 from fbpcp.entity.container_type import ContainerType
 from fbpcp.error.pcp import ThrottlingError
 from fbpcp.service.onedocker import OneDockerService
-from fbpcs.common.service.retry_handler import RetryHandler
+from fbpcs.common.service.retry_handler import BackoffType, RetryHandler
 from fbpcs.private_computation.service.constants import DEFAULT_CONTAINER_TIMEOUT_IN_SEC
 
 DEFAULT_WAIT_FOR_CONTAINER_POLL = 5
@@ -145,15 +145,28 @@ async def wait_for_containers_async(
         container_ids = [container.instance_id for container in containers]
         finished_containers = []
 
-        updated_containers = onedocker_svc.get_containers(container_ids)
+        with RetryHandler(
+            ThrottlingError, backoff_type=BackoffType.LINEAR
+        ) as retry_handler:
+            updated_containers = retry_handler.execute_sync(
+                onedocker_svc.get_containers,
+                container_ids,
+            )
         pending_container_ids = (
             RunBinaryBaseService._remove_finished_containers_from_container_ids(
                 onedocker_svc, updated_containers, container_ids, finished_containers
             )
         )
         while pending_container_ids:
             await asyncio.sleep(poll)
-            updated_containers = onedocker_svc.get_containers(pending_container_ids)
+            with RetryHandler(
+                ThrottlingError, backoff_type=BackoffType.LINEAR
+            ) as retry_handler:
+                updated_containers = retry_handler.execute_sync(
+                    onedocker_svc.get_containers,
+                    pending_container_ids,
+                )
+
             pending_container_ids = (
                 RunBinaryBaseService._remove_finished_containers_from_container_ids(
                     onedocker_svc,