[Docker Plug-in] Don't pull image if already exists unless forced (#895)

* Don't pull image if already exists unless forced Signed-off-by: Victor Chang <[email protected]>
Project-MONAI · Oct 19, 2023 · 6f94803 · 6f94803
1 parent 5788146
commit 6f94803
Show file tree

Hide file tree

Showing 4 changed files with 122 additions and 6 deletions.
diff --git a/src/TaskManager/Plug-ins/Docker/DockerPlugin.cs b/src/TaskManager/Plug-ins/Docker/DockerPlugin.cs
@@ -15,6 +15,7 @@
  */
 
 using System.Globalization;
+using Amazon.Runtime.Internal.Transform;
 using Ardalis.GuardClauses;
 using Docker.DotNet;
 using Docker.DotNet.Models;
@@ -132,13 +133,17 @@ public override async Task<ExecutionStatus> ExecuteTask(CancellationToken cancel
 
             try
             {
-                var imageCreateParameters = new ImagesCreateParameters()
+                var alwaysPull = Event.TaskPluginArguments.ContainsKey(Keys.AlwaysPull) && Event.TaskPluginArguments[Keys.AlwaysPull].Equals("true", StringComparison.OrdinalIgnoreCase);
+                if (alwaysPull || !await ImageExistsAsync(cancellationToken).ConfigureAwait(false))
                 {
-                    FromImage = Event.TaskPluginArguments[Keys.ContainerImage],
-                };
-
-                // Pull image.
-                await _dockerClient.Images.CreateImageAsync(imageCreateParameters, new AuthConfig(), new Progress<JSONMessage>(), cancellationToken).ConfigureAwait(false);
+                    // Pull image.
+                    _logger.ImageDoesNotExist(Event.TaskPluginArguments[Keys.ContainerImage]);
+                    var imageCreateParameters = new ImagesCreateParameters()
+                    {
+                        FromImage = Event.TaskPluginArguments[Keys.ContainerImage],
+                    };
+                    await _dockerClient.Images.CreateImageAsync(imageCreateParameters, new AuthConfig(), new Progress<JSONMessage>(), cancellationToken).ConfigureAwait(false);
+                }
             }
             catch (Exception exception)
             {
@@ -199,6 +204,20 @@ public override async Task<ExecutionStatus> ExecuteTask(CancellationToken cancel
             };
         }
 
+        private async Task<bool> ImageExistsAsync(CancellationToken cancellationToken)
+        {
+            var imageListParameters = new ImagesListParameters
+            {
+                Filters = new Dictionary<string, IDictionary<string, bool>>
+                {
+                    { "reference", new Dictionary<string, bool> { { Event.TaskPluginArguments[Keys.ContainerImage], true } } }
+                }
+            };
+
+            var results = await _dockerClient.Images.ListImagesAsync(imageListParameters, cancellationToken);
+            return results?.Any() ?? false;
+        }
+
         public override async Task<ExecutionStatus> GetStatus(string identity, TaskCallbackEvent callbackEvent, CancellationToken cancellationToken = default)
         {
             Guard.Against.NullOrWhiteSpace(identity, nameof(identity));

diff --git a/src/TaskManager/Plug-ins/Docker/Keys.cs b/src/TaskManager/Plug-ins/Docker/Keys.cs
@@ -38,6 +38,11 @@ internal static class Keys
         /// </summary>
         public static readonly string Command = "command";
 
+        /// <summary>
+        /// Key to indicate whether to always pull the image.
+        /// </summary>
+        public static readonly string AlwaysPull = "always_pull";
+
         /// <summary>
         /// Key for task timeout value.
         /// </summary>

diff --git a/src/TaskManager/Plug-ins/Docker/Logging/Log.cs b/src/TaskManager/Plug-ins/Docker/Logging/Log.cs
@@ -100,5 +100,8 @@ public static partial class Log
 
         [LoggerMessage(EventId = 1026, Level = LogLevel.Error, Message = "Error terminating container '{container}'.")]
         public static partial void ErrorTerminatingContainer(this ILogger logger, string container, Exception ex);
+
+        [LoggerMessage(EventId = 1027, Level = LogLevel.Information, Message = "Image does not exist '{image}' locally, attempting to pull.")]
+        public static partial void ImageDoesNotExist(this ILogger logger, string image);
     }
 }
diff --git a/tests/UnitTests/TaskManager.Docker.Tests/DockerPluginTest.cs b/tests/UnitTests/TaskManager.Docker.Tests/DockerPluginTest.cs
@@ -228,6 +228,95 @@ public async Task ExecuteTask_WhenFailedToMonitorContainer_ExpectTaskToBeAccepte
             runner.Dispose();
         }
 
+        [Fact(DisplayName = "ExecuteTask - do not pull the image when the specified image exists")]
+        public async Task ExecuteTask_WhenImageExists_ExpectNotToPull()
+        {
+            var payloadFiles = new List<VirtualFileInfo>()
+            {
+                new VirtualFileInfo( "file.dcm",  "path/to/file.dcm", "etag", 1000)
+            };
+            var contianerId = Guid.NewGuid().ToString();
+
+            _dockerClient.Setup(p => p.Images.CreateImageAsync(
+                It.IsAny<ImagesCreateParameters>(),
+                It.IsAny<AuthConfig>(),
+                It.IsAny<IProgress<JSONMessage>>(),
+                It.IsAny<CancellationToken>()));
+            _dockerClient.Setup(p => p.Images.ListImagesAsync(
+                It.IsAny<ImagesListParameters>(),
+                It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new List<ImagesListResponse>() { new ImagesListResponse() });
+            _dockerClient.Setup(p => p.Containers.CreateContainerAsync(
+                It.IsAny<CreateContainerParameters>(),
+                It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new CreateContainerResponse { ID = contianerId, Warnings = new List<string>() { "warning" } });
+
+            _storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(payloadFiles);
+            _storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
+
+            var message = GenerateTaskDispatchEventWithValidArguments();
+
+            var runner = new DockerPlugin(_serviceScopeFactory.Object, _logger.Object, message);
+            var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(false);
+
+            _dockerClient.Verify(p => p.Images.CreateImageAsync(
+                It.IsAny<ImagesCreateParameters>(),
+                It.IsAny<AuthConfig>(),
+                It.IsAny<IProgress<JSONMessage>>(),
+                It.IsAny<CancellationToken>()), Times.Never());
+            _dockerClient.Verify(p => p.Images.ListImagesAsync(
+                It.IsAny<ImagesListParameters>(),
+                It.IsAny<CancellationToken>()), Times.Once());
+            runner.Dispose();
+        }
+
+        [Fact(DisplayName = "ExecuteTask - pull the image when force by the user even the specified image exists")]
+        public async Task ExecuteTask_WhenAlwaysPullIsSet_ExpectToPullEvenWhenImageExists()
+        {
+            var payloadFiles = new List<VirtualFileInfo>()
+            {
+                new VirtualFileInfo( "file.dcm",  "path/to/file.dcm", "etag", 1000)
+            };
+            var contianerId = Guid.NewGuid().ToString();
+
+            _dockerClient.Setup(p => p.Images.CreateImageAsync(
+                It.IsAny<ImagesCreateParameters>(),
+                It.IsAny<AuthConfig>(),
+                It.IsAny<IProgress<JSONMessage>>(),
+                It.IsAny<CancellationToken>()));
+            _dockerClient.Setup(p => p.Images.ListImagesAsync(
+                It.IsAny<ImagesListParameters>(),
+                It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new List<ImagesListResponse>() { new ImagesListResponse() });
+            _dockerClient.Setup(p => p.Containers.CreateContainerAsync(
+                It.IsAny<CreateContainerParameters>(),
+                It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new CreateContainerResponse { ID = contianerId, Warnings = new List<string>() { "warning" } });
+
+            _storageService.Setup(p => p.ListObjectsAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(payloadFiles);
+            _storageService.Setup(p => p.GetObjectAsync(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<CancellationToken>()))
+                .ReturnsAsync(new MemoryStream(Encoding.UTF8.GetBytes("hello")));
+
+            var message = GenerateTaskDispatchEventWithValidArguments();
+            message.TaskPluginArguments.Add(Keys.AlwaysPull, bool.TrueString);
+
+            var runner = new DockerPlugin(_serviceScopeFactory.Object, _logger.Object, message);
+            var result = await runner.ExecuteTask(CancellationToken.None).ConfigureAwait(false);
+
+            _dockerClient.Verify(p => p.Images.CreateImageAsync(
+                It.IsAny<ImagesCreateParameters>(),
+                It.IsAny<AuthConfig>(),
+                It.IsAny<IProgress<JSONMessage>>(),
+                It.IsAny<CancellationToken>()), Times.Once());
+            _dockerClient.Verify(p => p.Images.ListImagesAsync(
+                It.IsAny<ImagesListParameters>(),
+                It.IsAny<CancellationToken>()), Times.Never());
+            runner.Dispose();
+        }
+
         [Fact(DisplayName = "ExecuteTask - when called with a valid event expect task to be accepted and monitored in the background")]
         public async Task ExecuteTask_WhenCalledWithValidEvent_ExpectTaskToBeAcceptedAndMonitored()
         {