Merge pull request #261 from frappe/mergify/bp/master/pr-260

feat(registry): Registry retry (backport #260)

Author: Aradhya-Tripathi

agent/builder.py

@@ -9,11 +9,11 @@
 from typing import TYPE_CHECKING
 
 import docker
-import requests
 
 from agent.base import Base
 from agent.exceptions import RegistryDownException
 from agent.job import Job, Step, job, step
+from agent.utils import is_registry_healthy
 
 if TYPE_CHECKING:
     from typing import Literal
@@ -138,21 +138,6 @@ def _publish_docker_build_output(self, result):
             self._publish_throttled_output(False)
         self._publish_throttled_output(True)
 
-    def is_registry_healthy(self) -> bool:
-        """Check if production registry (only) is healthy in the push cycle"""
-        headers = {"Accept": "application/vnd.docker.distribution.manifest.v2+json"}
-
-        if self.registry["url"] != "registry.frappe.cloud":
-            return True
-
-        response = requests.get(
-            f"https://{self.registry['url']}/v2",
-            auth=(self.registry["username"], self.registry["password"]),
-            headers=headers,
-        )
-
-        return response.ok
-
     def _wait_for_registry_recovery(self):
         """Wait for registry to recover after restart"""
         time.sleep(60)
@@ -165,12 +150,16 @@ def _push_docker_image(self):
 
         for attempt in range(max_retries):
             try:
-                if not self.is_registry_healthy():
+                if not is_registry_healthy(
+                    self.registry["url"], self.registry["username"], self.registry["password"]
+                ):
                     raise RegistryDownException("Registry is currently down")
 
                 self._push_image(client)
 
-                if not self.is_registry_healthy():
+                if not is_registry_healthy(
+                    self.registry["url"], self.registry["username"], self.registry["password"]
+                ):
                     raise RegistryDownException("Registry became unhealthy after push")
 
                 return self.output["push"]

agent/server.py

@@ -22,11 +22,11 @@
 )
 from agent.base import AgentException, Base
 from agent.bench import Bench
-from agent.exceptions import BenchNotExistsException
+from agent.exceptions import BenchNotExistsException, RegistryDownException
 from agent.job import Job, Step, job, step
 from agent.patch_handler import run_patches
 from agent.site import Site
-from agent.utils import get_supervisor_processes_status
+from agent.utils import get_supervisor_processes_status, is_registry_healthy
 
 
 class Server(Base):
@@ -55,8 +55,22 @@ def docker_login(self, registry):
         password = registry["password"]
         return self.execute(f"docker login -u {username} -p {password} {url}")
 
+    def establish_connection_with_registry(self, max_retries: int, registry: dict[str, str]):
+        """Given the attempt count try and establish connection with the registry else Raise"""
+        for attempt in range(max_retries):
+            try:
+                if not is_registry_healthy(registry["url"], registry["username"], registry["password"]):
+                    raise RegistryDownException("Registry is not available")
+                break
+            except RegistryDownException as e:
+                if attempt == max_retries - 1:
+                    raise Exception("Failed to pull image") from e
+
+                time.sleep(60)
+
     @step("Initialize Bench")
-    def bench_init(self, name, config):
+    def bench_init(self, name, config, registry: dict[str, str]):
+        self.establish_connection_with_registry(max_retries=3, registry=registry)
         bench_directory = os.path.join(self.benches_directory, name)
         os.mkdir(bench_directory)
         directories = ["logs", "sites", "config"]
@@ -98,7 +112,7 @@ def dump(self):
     @job("New Bench", priority="low")
     def new_bench(self, name, bench_config, common_site_config, registry, mounts=None):
         self.docker_login(registry)
-        self.bench_init(name, bench_config)
+        self.bench_init(name, bench_config, registry)
         bench = Bench(name, self, mounts=mounts)
         bench.update_config(common_site_config, bench_config)
         if bench.bench_config.get("single_container"):

agent/utils.py

@@ -72,6 +72,18 @@ def get_size(folder, ignore_dirs=None):
     return total_size
 
 
+def is_registry_healthy(url: str, username: str, password: str) -> bool:
+    """Check if production registry (only) is healthy in the push cycle"""
+    headers = {"Accept": "application/vnd.docker.distribution.manifest.v2+json"}
+
+    if url != "registry.frappe.cloud":
+        return True
+
+    response = requests.get(f"https://{url}/v2", auth=(username, password), headers=headers)
+
+    return response.ok
+
+
 def cint(x):
     """Convert to integer"""
     if x is None: