From ef7bc117cda2419f8370f5bc47199403ce2a3c8c Mon Sep 17 00:00:00 2001 From: Andrey Balakshiy Date: Fri, 15 Mar 2024 18:40:19 +0300 Subject: [PATCH 1/4] chech docker image platform when pulling docker image --- lzy/execution-env/pom.xml | 12 ++ .../ai/lzy/env/base/DockerEnvDescription.java | 19 +- .../ai/lzy/env/base/DockerEnvironment.java | 80 ++++++++- .../lzy/env/base/DockerEnvironmentTest.java | 165 ++++++++++++++++++ 4 files changed, 264 insertions(+), 12 deletions(-) create mode 100644 lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java diff --git a/lzy/execution-env/pom.xml b/lzy/execution-env/pom.xml index 4685b6b0d..063f3fa26 100644 --- a/lzy/execution-env/pom.xml +++ b/lzy/execution-env/pom.xml @@ -59,6 +59,18 @@ junit test + + org.mockito + mockito-core + 5.11.0 + test + + + org.mockito + mockito-inline + 5.2.0 + test + diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java index ba612cd3a..3f9468930 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java +++ b/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 envVars, // In format = @Nullable String networkMode, - DockerClientConfig dockerClientConfig + DockerClientConfig dockerClientConfig, + @Nonnull + Set allowedPlatforms // In format os/arch like "linux/amd64". Empty means all 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 envVars = new ArrayList<>(); String networkMode = null; DockerClientConfig dockerClientConfig; + Set 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 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( diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java index 84f13c452..6b3c6ae18 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java @@ -3,16 +3,27 @@ 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.api.model.BindOptions; +import com.github.dockerjava.api.model.BindPropagation; +import com.github.dockerjava.api.model.DeviceRequest; +import com.github.dockerjava.api.model.Frame; +import com.github.dockerjava.api.model.HostConfig; +import com.github.dockerjava.api.model.Mount; +import com.github.dockerjava.api.model.MountType; +import com.github.dockerjava.api.model.PruneType; +import com.github.dockerjava.api.model.PullResponseItem; 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 +39,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 +49,7 @@ 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"; @Nullable public String containerId = null; @@ -275,12 +288,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 +306,63 @@ private void prepareImage(String image, LogStream out) throws Exception { var msg = "Pulling image %s ...".formatted(image); LOG.info(msg); out.log(msg); + Set 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) { + if (e.getMessage().contains(NO_MATCHING_MANIFEST_ERROR)) { + LOG.info("Cannot find image = {} for platform = {}", image, platform); + } 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 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 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 platform = %s is not in allowed platforms = %s".formatted( + config.image(), platform, allowedPlatformsStr); + LOG.info(msg); + out.log(msg); + + throw new RuntimeException("Cached image platform = %s is not in allowed platforms = %s".formatted( + platform, allowedPlatformsStr)); + } + } } diff --git a/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java new file mode 100644 index 000000000..a7d061dbf --- /dev/null +++ b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java @@ -0,0 +1,165 @@ +package ai.lzy.env.base; + +import ai.lzy.env.logs.LogStream; +import com.github.dockerjava.api.DockerClient; +import com.github.dockerjava.api.async.ResultCallback; +import com.github.dockerjava.api.command.InspectImageCmd; +import com.github.dockerjava.api.command.InspectImageResponse; +import com.github.dockerjava.api.command.PullImageCmd; +import com.github.dockerjava.api.command.PullImageResultCallback; +import com.github.dockerjava.api.exception.DockerClientException; +import com.github.dockerjava.api.exception.NotFoundException; +import com.github.dockerjava.api.model.PullResponseItem; +import com.github.dockerjava.core.DefaultDockerClientConfig; +import com.github.dockerjava.core.DockerClientConfig; +import com.github.dockerjava.core.DockerClientImpl; +import org.apache.commons.lang3.RandomStringUtils; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import java.util.List; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class DockerEnvironmentTest { + private final DockerClient dockerClient = mock(DockerClientImpl.class); + + private final InspectImageCmd inspectImageCmd = mock(InspectImageCmd.class); + + private final InspectImageResponse inspectImageResponse = mock(InspectImageResponse.class); + + private final PullImageCmd pullImageCmd = mock(PullImageCmd.class); + + private final PullImageCmd pullImageCmdForRightPlatform = mock(PullImageCmd.class); + + private final ResultCallback.Adapter callbackAdapter = mock(ResultCallback.Adapter.class); + + private static final String IMAGE = RandomStringUtils.randomAlphanumeric(20).toLowerCase(); + + private final LogStream logStream = mock(LogStream.class); + + @Before + public void setUp() throws Exception { + when(dockerClient.inspectImageCmd(IMAGE)).thenReturn(inspectImageCmd); + when(inspectImageCmd.exec()).thenReturn(inspectImageResponse); + when(dockerClient.pullImageCmd(IMAGE)).thenReturn(pullImageCmd); + when(pullImageCmd.withPlatform("linux/amd64")).thenReturn(pullImageCmdForRightPlatform); + when(pullImageCmd.withPlatform(anyString())).thenAnswer((arg) -> { + if ("linux/amd64".equals(arg.getArguments()[0])) { + return pullImageCmdForRightPlatform; + } else { + return pullImageCmd; + } + }); + PullImageResultCallback pullImageResultCallback = mock(PullImageResultCallback.class); + when(pullImageCmdForRightPlatform.exec(any())).thenReturn(pullImageResultCallback); + when(pullImageCmd.exec(any())).thenThrow(new DockerClientException( + "Could not pull image: no matching manifest for %s in the manifest list entries")); + + when(pullImageResultCallback.awaitCompletion()).thenReturn(callbackAdapter); + } + + @Test + public void testPrepareImageCachedImage() throws Exception { + executeTest(this::doTestPrepareImageCachedImage); + } + + @Test + public void testPrepareImageCachedImageWithNotAllowedPlatform() throws Exception { + executeTest(this::doTestPrepareImageCachedImageWithNotAllowedPlatform); + } + + @Test + public void testPrepareImageNodCachedImageWithRightPlatform() throws Exception { + executeTest(this::doTestPrepareImageNodCachedImageWithRightPlatform); + } + + @Test + public void testPrepareImageNodCachedImageWithoutRightPlatform() throws Exception { + executeTest(this::doTestPrepareImageNodCachedImageWithoutRightPlatform); + } + + private void doTestPrepareImageCachedImage() throws Exception { + when(inspectImageResponse.getArch()).thenReturn("amd64"); + when(inspectImageResponse.getOs()).thenReturn("linux"); + + DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( + List.of("darwin/arm64", "linux/amd64"))); + + environment.prepareImage(IMAGE, logStream); + + verify(dockerClient, never()).pullImageCmd(any()); + } + + private void doTestPrepareImageCachedImageWithNotAllowedPlatform() { + when(inspectImageResponse.getOs()).thenReturn("not_existed_os"); + when(inspectImageResponse.getArch()).thenReturn("not_existed_arch"); + + DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( + List.of("darwin/arm64", "linux/amd64"))); + + RuntimeException exception = Assert.assertThrows(RuntimeException.class, + () -> environment.prepareImage(IMAGE, logStream)); + assertNotNull(exception); + assertEquals("Cached image platform = not_existed_os/not_existed_arch is not in allowed platforms = darwin/arm64, linux/amd64", + exception.getMessage()); + + verify(dockerClient, never()).pullImageCmd(any()); + } + + private void doTestPrepareImageNodCachedImageWithRightPlatform() throws Exception { + when(inspectImageCmd.exec()).thenThrow(new NotFoundException("com.github.dockerjava.api.exception.NotFoundException:" + + " Status 404: {\"message\":\"No such image: %s\"}\n".formatted(IMAGE))); + DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( + List.of("darwin/arm64", "linux/amd64"))); + + environment.prepareImage(IMAGE, logStream); + verify(dockerClient, times(2)).pullImageCmd(IMAGE); + } + + private void doTestPrepareImageNodCachedImageWithoutRightPlatform() throws Exception { + when(inspectImageCmd.exec()).thenThrow(new NotFoundException("com.github.dockerjava.api.exception.NotFoundException:" + + " Status 404: {\"message\":\"No such image: %s\"}\n".formatted(IMAGE))); + DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( + List.of("darwin/arm64", "linux/win32"))); + + RuntimeException exception = Assert.assertThrows(RuntimeException.class, + () -> environment.prepareImage(IMAGE, logStream)); + assertNotNull(exception); + assertEquals("Cannot pull image for allowed platforms = linux/win32, darwin/arm64", exception.getMessage()); + + verify(dockerClient, times(2)).pullImageCmd(IMAGE); + } + + private DockerEnvDescription createDockerEnvDescription(List allowedPlatforms) { + DockerClientConfig dockerClientConfig = DefaultDockerClientConfig.createDefaultConfigBuilder().build(); + return DockerEnvDescription.newBuilder() + .withDockerClientConfig(dockerClientConfig) + .withAllowedPlatforms(allowedPlatforms) + .build(); + } + + + private void executeTest(Executable test) throws Exception { + try (var mockedDockerClient = mockStatic(DockerClientImpl.class)) { + mockedDockerClient.when(() -> DockerClientImpl.getInstance(any(), any())).thenReturn(dockerClient); + test.execute(); + } + } + + + @FunctionalInterface + private interface Executable { + void execute() throws Exception; + } +} \ No newline at end of file From 81f65f1522a2b84c1cdb05fb5a6155143cd4b098 Mon Sep 17 00:00:00 2001 From: Balakshiy Andrey Date: Sat, 16 Mar 2024 23:53:38 +0300 Subject: [PATCH 2/4] chech docker image platform when pulling docker image * some fixes --- .../ai/lzy/env/base/DockerEnvDescription.java | 2 +- .../ai/lzy/env/base/DockerEnvironment.java | 18 +++++------------- .../ai/lzy/env/base/DockerEnvironmentTest.java | 14 +++++--------- 3 files changed, 11 insertions(+), 23 deletions(-) diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java index 3f9468930..0a831f9e3 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java @@ -24,7 +24,7 @@ public record DockerEnvDescription( String networkMode, DockerClientConfig dockerClientConfig, @Nonnull - Set allowedPlatforms // In format os/arch like "linux/amd64". Empty means all allowed + Set allowedPlatforms // In format os/arch like "linux/amd64". Empty means all are allowed ) { public static Builder newBuilder() { diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java index 6b3c6ae18..deaca0d4a 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java @@ -12,15 +12,7 @@ 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.BindOptions; -import com.github.dockerjava.api.model.BindPropagation; -import com.github.dockerjava.api.model.DeviceRequest; -import com.github.dockerjava.api.model.Frame; -import com.github.dockerjava.api.model.HostConfig; -import com.github.dockerjava.api.model.Mount; -import com.github.dockerjava.api.model.MountType; -import com.github.dockerjava.api.model.PruneType; -import com.github.dockerjava.api.model.PullResponseItem; +import com.github.dockerjava.api.model.*; import com.github.dockerjava.core.DockerClientImpl; import com.github.dockerjava.httpclient5.ApacheDockerHttpClient; import com.google.common.annotations.VisibleForTesting; @@ -329,7 +321,8 @@ void prepareImage(String image, LogStream out) throws Exception { }) ) { if (pullResponseItem == null) { - throw new RuntimeException("Cannot pull image for allowed platforms = %s".formatted(String.join(", ", allowedPlatforms))); + throw new RuntimeException("Cannot pull image for allowed platforms = %s".formatted( + String.join(", ", allowedPlatforms))); } } @@ -356,13 +349,12 @@ private void checkPlatform(InspectImageResponse inspectImageResponse, LogStream String platform = inspectImageResponse.getOs() + "/" + inspectImageResponse.getArch(); if (!allowedPlatforms.contains(platform)) { var allowedPlatformsStr = String.join(", ", allowedPlatforms); - var msg = "Image %s platform = %s is not in allowed platforms = %s".formatted( + 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("Cached image platform = %s is not in allowed platforms = %s".formatted( - platform, allowedPlatformsStr)); + throw new RuntimeException(msg); } } } diff --git a/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java index a7d061dbf..21b4eb386 100644 --- a/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java +++ b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java @@ -24,12 +24,7 @@ import static org.junit.Assert.assertNotNull; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.mockStatic; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.times; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; +import static org.mockito.Mockito.*; public class DockerEnvironmentTest { private final DockerClient dockerClient = mock(DockerClientImpl.class); @@ -111,8 +106,8 @@ private void doTestPrepareImageCachedImageWithNotAllowedPlatform() { RuntimeException exception = Assert.assertThrows(RuntimeException.class, () -> environment.prepareImage(IMAGE, logStream)); assertNotNull(exception); - assertEquals("Cached image platform = not_existed_os/not_existed_arch is not in allowed platforms = darwin/arm64, linux/amd64", - exception.getMessage()); + assertEquals(("Image %s with platform = not_existed_os/not_existed_arch " + + "is not in allowed platforms = darwin/arm64, linux/amd64").formatted(IMAGE), exception.getMessage()); verify(dockerClient, never()).pullImageCmd(any()); } @@ -144,6 +139,7 @@ private void doTestPrepareImageNodCachedImageWithoutRightPlatform() throws Excep private DockerEnvDescription createDockerEnvDescription(List allowedPlatforms) { DockerClientConfig dockerClientConfig = DefaultDockerClientConfig.createDefaultConfigBuilder().build(); return DockerEnvDescription.newBuilder() + .withImage(IMAGE) .withDockerClientConfig(dockerClientConfig) .withAllowedPlatforms(allowedPlatforms) .build(); @@ -162,4 +158,4 @@ private void executeTest(Executable test) throws Exception { private interface Executable { void execute() throws Exception; } -} \ No newline at end of file +} From 80b064861241bd33c5a7b049a840eb7712d1795d Mon Sep 17 00:00:00 2001 From: Andrey Balakshiy Date: Sun, 17 Mar 2024 23:56:14 +0300 Subject: [PATCH 3/4] check docker image platform when pulling docker image * added case for docker client exception --- .../ai/lzy/env/base/DockerEnvironment.java | 8 +- .../lzy/env/base/DockerEnvironmentTest.java | 77 +++++++++++++++---- 2 files changed, 68 insertions(+), 17 deletions(-) diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java index deaca0d4a..2c5ea5c87 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java @@ -42,6 +42,7 @@ public class DockerEnvironment extends BaseEnvironment { 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; @@ -309,8 +310,11 @@ void prepareImage(String image, LogStream out) throws Exception { try { return pullWithPlatform(image, platform); } catch (DockerClientException e) { - if (e.getMessage().contains(NO_MATCHING_MANIFEST_ERROR)) { - LOG.info("Cannot find image = {} for platform = {}", image, platform); + 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; } diff --git a/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java index 21b4eb386..5936aefd0 100644 --- a/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java +++ b/lzy/execution-env/src/test/java/ai/lzy/env/base/DockerEnvironmentTest.java @@ -58,8 +58,6 @@ public void setUp() throws Exception { }); PullImageResultCallback pullImageResultCallback = mock(PullImageResultCallback.class); when(pullImageCmdForRightPlatform.exec(any())).thenReturn(pullImageResultCallback); - when(pullImageCmd.exec(any())).thenThrow(new DockerClientException( - "Could not pull image: no matching manifest for %s in the manifest list entries")); when(pullImageResultCallback.awaitCompletion()).thenReturn(callbackAdapter); } @@ -84,12 +82,21 @@ public void testPrepareImageNodCachedImageWithoutRightPlatform() throws Exceptio executeTest(this::doTestPrepareImageNodCachedImageWithoutRightPlatform); } + @Test + public void testPrepareImageNodCachedImageWithoutRightPlatform2() throws Exception { + executeTest(this::doTestPrepareImageNodCachedImageWithoutRightPlatform2); + } + + @Test + public void testPrepareImageNodCachedImageDockerClientException() throws Exception { + executeTest(this::doTestPrepareImageNodCachedImageDockerClientException); + } + private void doTestPrepareImageCachedImage() throws Exception { when(inspectImageResponse.getArch()).thenReturn("amd64"); when(inspectImageResponse.getOs()).thenReturn("linux"); - DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( - List.of("darwin/arm64", "linux/amd64"))); + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/amd64"); environment.prepareImage(IMAGE, logStream); @@ -100,8 +107,7 @@ private void doTestPrepareImageCachedImageWithNotAllowedPlatform() { when(inspectImageResponse.getOs()).thenReturn("not_existed_os"); when(inspectImageResponse.getArch()).thenReturn("not_existed_arch"); - DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( - List.of("darwin/arm64", "linux/amd64"))); + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/amd64"); RuntimeException exception = Assert.assertThrows(RuntimeException.class, () -> environment.prepareImage(IMAGE, logStream)); @@ -113,20 +119,35 @@ private void doTestPrepareImageCachedImageWithNotAllowedPlatform() { } private void doTestPrepareImageNodCachedImageWithRightPlatform() throws Exception { - when(inspectImageCmd.exec()).thenThrow(new NotFoundException("com.github.dockerjava.api.exception.NotFoundException:" + - " Status 404: {\"message\":\"No such image: %s\"}\n".formatted(IMAGE))); - DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( - List.of("darwin/arm64", "linux/amd64"))); + mockDockerClientException("Could not pull image: no matching manifest for %s in the manifest list entries".formatted(IMAGE)); + mockNotFound(); + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/amd64"); environment.prepareImage(IMAGE, logStream); verify(dockerClient, times(2)).pullImageCmd(IMAGE); } - private void doTestPrepareImageNodCachedImageWithoutRightPlatform() throws Exception { - when(inspectImageCmd.exec()).thenThrow(new NotFoundException("com.github.dockerjava.api.exception.NotFoundException:" + - " Status 404: {\"message\":\"No such image: %s\"}\n".formatted(IMAGE))); - DockerEnvironment environment = new DockerEnvironment(createDockerEnvDescription( - List.of("darwin/arm64", "linux/win32"))); + private void doTestPrepareImageNodCachedImageWithoutRightPlatform() { + mockDockerClientException("Could not pull image: no matching manifest for %s in the manifest list entries".formatted(IMAGE)); + mockNotFound(); + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/win32"); + + RuntimeException exception = Assert.assertThrows(RuntimeException.class, + () -> environment.prepareImage(IMAGE, logStream)); + assertNotNull(exception); + assertEquals("Cannot pull image for allowed platforms = linux/win32, darwin/arm64", exception.getMessage()); + + verify(dockerClient, times(2)).pullImageCmd(IMAGE); + } + + private void doTestPrepareImageNodCachedImageWithoutRightPlatform2() { + mockDockerClientException(""" + com.github.dockerjava.api.exception.DockerClientException: Could not pull image: image with reference %s + was found but does not match the specified platform: wanted darwin/arm64, + actual: linux/amd64""".formatted(IMAGE)); + mockNotFound(); + + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/win32"); RuntimeException exception = Assert.assertThrows(RuntimeException.class, () -> environment.prepareImage(IMAGE, logStream)); @@ -136,6 +157,32 @@ private void doTestPrepareImageNodCachedImageWithoutRightPlatform() throws Excep verify(dockerClient, times(2)).pullImageCmd(IMAGE); } + private void doTestPrepareImageNodCachedImageDockerClientException() { + mockDockerClientException("another docker client exception"); + mockNotFound(); + DockerEnvironment environment = createEnvironment("darwin/arm64", "linux/win32"); + + RuntimeException exception = Assert.assertThrows(DockerClientException.class, + () -> environment.prepareImage(IMAGE, logStream)); + assertNotNull(exception); + assertEquals("another docker client exception", exception.getMessage()); + + verify(dockerClient, times(3)).pullImageCmd(IMAGE); + } + + private void mockDockerClientException(String message) { + when(pullImageCmd.exec(any())).thenThrow(new DockerClientException(message)); + } + + private void mockNotFound() { + when(inspectImageCmd.exec()).thenThrow(new NotFoundException("com.github.dockerjava.api.exception.NotFoundException:" + + " Status 404: {\"message\":\"No such image: %s\"}\n".formatted(IMAGE))); + } + + private DockerEnvironment createEnvironment(String... platforms) { + return new DockerEnvironment(createDockerEnvDescription(List.of(platforms))); + } + private DockerEnvDescription createDockerEnvDescription(List allowedPlatforms) { DockerClientConfig dockerClientConfig = DefaultDockerClientConfig.createDefaultConfigBuilder().build(); return DockerEnvDescription.newBuilder() From 61bb32953998f12c9906fb69917f0167f2126cf2 Mon Sep 17 00:00:00 2001 From: Andrey Balakshiy Date: Mon, 18 Mar 2024 11:36:58 +0300 Subject: [PATCH 4/4] check docker image platform when pulling docker image * code review --- lzy/execution-env/pom.xml | 2 -- .../main/java/ai/lzy/env/base/DockerEnvDescription.java | 9 +-------- .../src/main/java/ai/lzy/env/base/DockerEnvironment.java | 7 ++++--- parent/pom.xml | 4 ++-- 4 files changed, 7 insertions(+), 15 deletions(-) diff --git a/lzy/execution-env/pom.xml b/lzy/execution-env/pom.xml index 063f3fa26..57b9964f7 100644 --- a/lzy/execution-env/pom.xml +++ b/lzy/execution-env/pom.xml @@ -62,13 +62,11 @@ org.mockito mockito-core - 5.11.0 test org.mockito mockito-inline - 5.2.0 test diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java index 0a831f9e3..20b40e8bd 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvDescription.java @@ -1,17 +1,11 @@ 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.*; import java.util.stream.Collectors; public record DockerEnvDescription( @@ -23,7 +17,6 @@ public record DockerEnvDescription( @Nullable String networkMode, DockerClientConfig dockerClientConfig, - @Nonnull Set allowedPlatforms // In format os/arch like "linux/amd64". Empty means all are allowed ) { diff --git a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java index 2c5ea5c87..2c67734d4 100644 --- a/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java +++ b/lzy/execution-env/src/main/java/ai/lzy/env/base/DockerEnvironment.java @@ -322,8 +322,8 @@ void prepareImage(String image, LogStream out) throws Exception { } } return null; - }) - ) { + })) + { if (pullResponseItem == null) { throw new RuntimeException("Cannot pull image for allowed platforms = %s".formatted( String.join(", ", allowedPlatforms))); @@ -336,7 +336,8 @@ void prepareImage(String image, LogStream out) throws Exception { } private ResultCallback.Adapter pullWithPlatform(String image, @Nullable String platform) - throws InterruptedException { + throws InterruptedException + { var pullingImage = client.pullImageCmd(image); if (platform != null) { pullingImage = pullingImage.withPlatform(platform); diff --git a/parent/pom.xml b/parent/pom.xml index 312d5173c..70ea15ccf 100644 --- a/parent/pom.xml +++ b/parent/pom.xml @@ -122,8 +122,8 @@ 4.13.2 2.2 6.0.0 - 3.5.15 - 3.6.28 + 5.11.0 + 5.2.0 4.10.0 3.6.2 0.2.6