From 5581cad1cba4e41c3ad7bdf709fb338de477bf41 Mon Sep 17 00:00:00 2001
From: Sanghun Lee <sanghun@lablup.com>
Date: Mon, 30 Sep 2024 13:52:34 +0900
Subject: [PATCH] Handle KeyError to keep pulling if response data is in
 unexpected shape

---
 src/ai/backend/agent/docker/agent.py | 52 +++++++++++++++-------------
 1 file changed, 28 insertions(+), 24 deletions(-)

diff --git a/src/ai/backend/agent/docker/agent.py b/src/ai/backend/agent/docker/agent.py
index 920b76f5559..b3a47b0a471 100644
--- a/src/ai/backend/agent/docker/agent.py
+++ b/src/ai/backend/agent/docker/agent.py
@@ -1677,31 +1677,35 @@ async def handle_err_response(resp: PullErrorResponse) -> None:
             async for resp in docker.images.pull(
                 image_ref.canonical, auth=auth_config, stream=True
             ):
-                match resp:
-                    case dict() if resp.get("status"):
-                        _resp = PullResponse(status=resp["status"])
-                        if detail := resp.get("progressDetail"):
-                            _resp["progressDetail"] = detail
-                        if progress := resp.get("progress"):
-                            _resp["progress"] = progress
-                        if id := resp.get("id"):
-                            _resp["id"] = id
-                        if do_report_per_layer:
-                            await handle_response_for_each_layer(_resp, layer_to_reporter_id_map)
-                        else:
-                            await handle_response(_resp, layer_ids)
-                    case dict() if resp.get("error"):
-                        await handle_err_response(
-                            PullErrorResponse(
-                                error=resp["error"],
-                                errorDetail=resp["errorDetail"],
+                try:
+                    match resp:
+                        case dict() if resp.get("status"):
+                            _resp = PullResponse(status=resp["status"])
+                            if detail := resp.get("progressDetail"):
+                                _resp["progressDetail"] = detail
+                            if progress := resp.get("progress"):
+                                _resp["progress"] = progress
+                            if id := resp.get("id"):
+                                _resp["id"] = id
+                            if do_report_per_layer:
+                                await handle_response_for_each_layer(
+                                    _resp, layer_to_reporter_id_map
+                                )
+                            else:
+                                await handle_response(_resp, layer_ids)
+                        case dict() if resp.get("error"):
+                            err_msg = str(resp["error"])
+                            await handle_err_response(
+                                PullErrorResponse(
+                                    error=resp["error"],
+                                    errorDetail=resp["errorDetail"],
+                                )
                             )
-                        )
-                        raise ImagePullFailure(str(resp["error"]))
-                    case _:
-                        log.warning(
-                            f"Unable to deserialize pulling response. skip. (resp:{str(resp)})"
-                        )
+                        case _:
+                            raise KeyError
+                except KeyError:
+                    log.warning(f"Unable to deserialize pulling response. skip. (resp:{str(resp)})")
+        return err_msg
 
     async def check_image(
         self, image_ref: ImageRef, image_id: str, auto_pull: AutoPullBehavior