chore: simplify the multiarch -> single arch resolution on the registry side

Signed-off-by: Alex Gronskiy <[email protected]>

Author: agronskiy

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/common/container_metadata/loading.py

@@ -19,7 +19,6 @@
 import json
 import os
 import pathlib
-import platform
 import tarfile
 import tempfile
 from typing import Optional
@@ -479,65 +478,6 @@ def find_file_matching_pattern_in_image_layers(
         authenticator.authenticate(repository=repository)
         # Don't fail here - authentication may fail but public containers can still be accessed
 
-    def _normalize_arch(a: object) -> Optional[str]:
-        if not a:
-            return None
-        s = str(a).lower()
-        if s in {"amd64", "x86_64"}:
-            return "amd64"
-        if s in {"arm64", "aarch64"}:
-            return "arm64"
-        return None
-
-    def _preferred_arch() -> str:
-        return _normalize_arch(platform.machine()) or "amd64"
-
-    def _pick_platform_manifest_digest(index_manifest: dict) -> Optional[str]:
-        manifests = index_manifest.get("manifests")
-        if not isinstance(manifests, list) or not manifests:
-            return None
-
-        pref = _preferred_arch()
-        # Prefer linux + preferred arch, then linux/amd64, then linux/arm64, then first digest.
-        best_score = 10_000
-        best_digest: Optional[str] = None
-
-        for m in manifests:
-            if not isinstance(m, dict):
-                continue
-            plat = m.get("platform") or {}
-            if not isinstance(plat, dict):
-                continue
-            os_name = str(plat.get("os") or "").lower()
-            arch = _normalize_arch(plat.get("architecture"))
-            digest = m.get("digest")
-            if not (isinstance(digest, str) and digest.startswith("sha256:")):
-                continue
-            if os_name != "linux" or not arch:
-                continue
-
-            score = 100
-            if arch == pref:
-                score = 0
-            elif arch == "amd64":
-                score = 10
-            elif arch == "arm64":
-                score = 20
-
-            if score < best_score:
-                best_score = score
-                best_digest = digest
-
-        if best_digest:
-            return best_digest
-
-        for m in manifests:
-            if isinstance(m, dict):
-                d = m.get("digest")
-                if isinstance(d, str) and d.startswith("sha256:"):
-                    return d
-        return None
-
     # Get top-level manifest and digest (tag may resolve to multi-arch index).
     top_manifest, top_digest = authenticator.get_manifest_and_digest(
         repository, reference
@@ -554,13 +494,30 @@ def _pick_platform_manifest_digest(index_manifest: dict) -> Optional[str]:
     if isinstance(top_manifest, dict) and isinstance(
         top_manifest.get("manifests"), list
     ):
-        platform_digest = _pick_platform_manifest_digest(top_manifest)
-        if platform_digest:
-            resolved, _ = authenticator.get_manifest_and_digest(
-                repository, platform_digest
-            )
-            if resolved:
-                manifest = resolved
+        # Prefer registry's default platform resolver by requesting a single-image manifest.
+        single_accept = (
+            "application/vnd.oci.image.manifest.v1+json, "
+            "application/vnd.docker.distribution.manifest.v2+json"
+        )
+        resolved, _ = authenticator.get_manifest_and_digest(
+            repository, reference, accept=single_accept
+        )
+        if resolved:
+            manifest = resolved
+
+        # If a registry still returns an index (ignoring Accept), fall back to the
+        # first digest entry for layer inspection. This does NOT affect recorded digests.
+        if isinstance(manifest, dict) and isinstance(manifest.get("manifests"), list):
+            for m in manifest.get("manifests") or []:
+                if isinstance(m, dict):
+                    d = m.get("digest")
+                    if isinstance(d, str) and d.startswith("sha256:"):
+                        resolved2, _ = authenticator.get_manifest_and_digest(
+                            repository, d, accept=single_accept
+                        )
+                        if resolved2:
+                            manifest = resolved2
+                        break
 
     # Check cache with digest validation (always validates digest)
     # For pattern searches, use pattern-based cache key (not resolved path)

packages/nemo-evaluator-launcher/src/nemo_evaluator_launcher/common/container_metadata/registries.py

@@ -217,7 +217,9 @@ def _read_docker_credentials(registry_url: str) -> Optional[Tuple[str, str]]:
         return None
 
 
-def _retry_without_auth(url: str, stream: bool = False) -> Optional[requests.Response]:
+def _retry_without_auth(
+    url: str, stream: bool = False, accept: Optional[str] = None
+) -> Optional[requests.Response]:
     """Retry HTTP request without authentication headers.
 
     Used for accessing public containers that may return 401/403 even for
@@ -231,7 +233,7 @@ def _retry_without_auth(url: str, stream: bool = False) -> Optional[requests.Res
         Response object if successful, None otherwise
     """
     temp_session = requests.Session()
-    temp_session.headers.update({"Accept": _DOCKER_MANIFEST_MEDIA_TYPE})
+    temp_session.headers.update({"Accept": accept or _DOCKER_MANIFEST_MEDIA_TYPE})
     response = temp_session.get(url, stream=stream)
     if response.status_code == 200:
         return response
@@ -254,7 +256,7 @@ def authenticate(self, repository: Optional[str] = None) -> bool:
         pass
 
     def get_manifest_and_digest(
-        self, repository: str, reference: str
+        self, repository: str, reference: str, accept: Optional[str] = None
     ) -> Tuple[Optional[Dict], Optional[str]]:
         """Get the manifest and digest for a specific image reference.
 
@@ -276,7 +278,12 @@ def get_manifest_and_digest(
             url = f"https://{self.registry_url}/v2/{repository}/manifests/{reference}"
             logger.debug("Fetching manifest", url=url)
 
-            response = self.session.get(url)
+            accept_header = (
+                accept
+                or self.session.headers.get("Accept")
+                or _DOCKER_MANIFEST_MEDIA_TYPE
+            )
+            response = self.session.get(url, headers={"Accept": accept_header})
 
             if response.status_code == 200:
                 manifest = response.json()
@@ -322,7 +329,7 @@ def get_manifest_and_digest(
                     "Got 401/403, retrying without authentication",
                     status_code=response.status_code,
                 )
-                retry_response = _retry_without_auth(url)
+                retry_response = _retry_without_auth(url, accept=accept_header)
                 if retry_response:
                     manifest = retry_response.json()
                     headers_dict = dict(retry_response.headers)