Handle KeyError to keep pulling if response data is in unexpected shape

lablup · Nov 16, 2024 · 6ab232c · 6ab232c
1 parent d1d6ead
commit 6ab232c
Showing 1 changed file with 27 additions and 24 deletions.
diff --git a/src/ai/backend/agent/docker/agent.py b/src/ai/backend/agent/docker/agent.py
@@ -1673,31 +1673,34 @@ 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"):
-                        err_msg = str(resp["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"],
+                                )
                             )
-                        )
-                    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(