chech docker image platform when pulling docker image

lzy/execution-env/pom.xml

@@ -59,6 +59,18 @@
             <artifactId>junit</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-core</artifactId>
+            <version>5.11.0</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.mockito</groupId>
+            <artifactId>mockito-inline</artifactId>
+            <version>5.2.0</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
 </project>

lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java

@@ -1,13 +1,17 @@
 package ai.lzy.env.base;
 
 import com.github.dockerjava.core.DockerClientConfig;
+import jakarta.annotation.Nonnull;
 import jakarta.annotation.Nullable;
 import org.apache.commons.lang3.RandomStringUtils;
 import org.apache.commons.lang3.StringUtils;
 
 import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.stream.Collectors;
 
 public record DockerEnvDescription(
@@ -18,7 +22,9 @@ public record DockerEnvDescription(
     List<String> envVars,  // In format <NAME>=<value>
     @Nullable
     String networkMode,
-    DockerClientConfig dockerClientConfig
+    DockerClientConfig dockerClientConfig,
+    @Nonnull
+    Set<String> allowedPlatforms // In format os/arch like "linux/amd64". Empty means all are allowed
 ) {
 
     public static Builder newBuilder() {
@@ -32,6 +38,7 @@ public String toString() {
             ", image='" + image + '\'' +
             ", needGpu=" + needGpu +
             ", networkMode=" + networkMode +
+            ", allowedPlatforms=" + String.join(", ", allowedPlatforms) +
             ", mounts=[" + mounts.stream()
                 .map(it -> it.source() + " -> " + it.target() + (it.isRshared() ? " (R_SHARED)" : ""))
                 .collect(Collectors.joining(", ")) + "]" +
@@ -52,6 +59,7 @@ public static class Builder {
         List<String> envVars = new ArrayList<>();
         String networkMode = null;
         DockerClientConfig dockerClientConfig;
+        Set<String> allowedPlatforms = new HashSet<>();
 
         public Builder withName(String name) {
             this.name = name;
@@ -93,13 +101,18 @@ public Builder withDockerClientConfig(DockerClientConfig dockerClientConfig) {
             return this;
         }
 
+        public Builder withAllowedPlatforms(Collection<String> allowedPlatforms) {
+            this.allowedPlatforms.addAll(allowedPlatforms);
+            return this;
+        }
+
         public DockerEnvDescription build() {
             if (StringUtils.isBlank(name)) {
                 name = "job-" + RandomStringUtils.randomAlphanumeric(5);
             }
-            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig);
+            return new DockerEnvDescription(name, image, mounts, gpu, envVars, networkMode, dockerClientConfig,
+                    allowedPlatforms);
         }
-
     }
 
     public record ContainerRegistryCredentials(

lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java

@@ -3,16 +3,19 @@
 import ai.lzy.env.EnvironmentInstallationException;
 import ai.lzy.env.logs.LogStream;
 import com.github.dockerjava.api.DockerClient;
+import com.github.dockerjava.api.async.ResultCallback;
 import com.github.dockerjava.api.async.ResultCallbackTemplate;
 import com.github.dockerjava.api.command.ExecCreateCmd;
 import com.github.dockerjava.api.command.ExecCreateCmdResponse;
+import com.github.dockerjava.api.command.InspectImageResponse;
 import com.github.dockerjava.api.command.PullImageResultCallback;
 import com.github.dockerjava.api.exception.DockerClientException;
 import com.github.dockerjava.api.exception.DockerException;
 import com.github.dockerjava.api.exception.NotFoundException;
 import com.github.dockerjava.api.model.*;
 import com.github.dockerjava.core.DockerClientImpl;
 import com.github.dockerjava.httpclient5.ApacheDockerHttpClient;
+import com.google.common.annotations.VisibleForTesting;
 import io.github.resilience4j.core.IntervalFunction;
 import io.github.resilience4j.retry.Retry;
 import io.github.resilience4j.retry.RetryConfig;
@@ -28,6 +31,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -37,6 +41,8 @@ public class DockerEnvironment extends BaseEnvironment {
     private static final Logger LOG = LogManager.getLogger(DockerEnvironment.class);
     private static final long GB_AS_BYTES = 1073741824;
     private static final String ROOT_USER_UID = "0";
+    private static final String NO_MATCHING_MANIFEST_ERROR = "no matching manifest";
+    private static final String NOT_MATCH_PLATFORM_ERROR = "was found but does not match the specified platform";
 
     @Nullable
     public String containerId = null;
@@ -275,12 +281,14 @@ public void close() throws Exception {
         }
     }
 
-    private void prepareImage(String image, LogStream out) throws Exception {
+    @VisibleForTesting
+    void prepareImage(String image, LogStream out) throws Exception {
         try {
-            client.inspectImageCmd(image).exec();
+            var inspectImageResponse = client.inspectImageCmd(image).exec();
             var msg = "Image %s exists".formatted(image);
             LOG.info(msg);
             out.log(msg);
+            checkPlatform(inspectImageResponse, out);
             return;
         } catch (NotFoundException ignored) {
             var msg = "Image %s not found in cached images".formatted(image);
@@ -291,16 +299,66 @@ private void prepareImage(String image, LogStream out) throws Exception {
         var msg = "Pulling image %s ...".formatted(image);
         LOG.info(msg);
         out.log(msg);
+        Set<String> allowedPlatforms = config.allowedPlatforms();
         AtomicInteger pullingAttempt = new AtomicInteger(0);
-        retry.executeCallable(() -> {
+        try (var pullResponseItem = retry.executeCallable(() -> {
             LOG.info("Pulling image {}, attempt {}", image, pullingAttempt.incrementAndGet());
-            final var pullingImage = client
-                .pullImageCmd(config.image())
-                .exec(new PullImageResultCallback());
-            return pullingImage.awaitCompletion();
-        });
+            if (allowedPlatforms.isEmpty()) {
+                return pullWithPlatform(image, null);
+            } else {
+                for (String platform : config.allowedPlatforms()) {
+                    try {
+                        return pullWithPlatform(image, platform);
+                    } catch (DockerClientException e) {
+                        String exceptionMessage = e.getMessage();
+                        if (exceptionMessage.contains(NO_MATCHING_MANIFEST_ERROR) ||
+                                exceptionMessage.contains(NOT_MATCH_PLATFORM_ERROR)) {
+                            LOG.info("Cannot find image = {} for platform = {}: message = {}",
+                                    image, platform, exceptionMessage);
+                        } else {
+                            throw e;
+                        }
+                    }
+                }
+            }
+            return null;
+        })
+        ) {
+            if (pullResponseItem == null) {
+                throw new RuntimeException("Cannot pull image for allowed platforms = %s".formatted(
+                        String.join(", ", allowedPlatforms)));
+            }
+        }
+
         msg = "Pulling image %s done".formatted(image);
         LOG.info(msg);
         out.log(msg);
     }
+
+    private ResultCallback.Adapter<PullResponseItem> pullWithPlatform(String image, @Nullable String platform)
+            throws InterruptedException {
+        var pullingImage = client.pullImageCmd(image);
+        if (platform != null) {
+            pullingImage = pullingImage.withPlatform(platform);
+        }
+        return pullingImage.exec(new PullImageResultCallback()).awaitCompletion();
+    }
+
+    private void checkPlatform(InspectImageResponse inspectImageResponse, LogStream out) {
+        Set<String> allowedPlatforms = config.allowedPlatforms();
+        if (allowedPlatforms.isEmpty()) {
+            return;
+        }
+
+        String platform = inspectImageResponse.getOs() + "/" + inspectImageResponse.getArch();
+        if (!allowedPlatforms.contains(platform)) {
+            var allowedPlatformsStr = String.join(", ", allowedPlatforms);
+            var msg = "Image %s with platform = %s is not in allowed platforms = %s".formatted(
+                    config.image(), platform, allowedPlatformsStr);
+            LOG.info(msg);
+            out.log(msg);
+
+            throw new RuntimeException(msg);
+        }
+    }
 }