[151] Runner executor

src/CommonUtilities/Models/NodeTask.cs

@@ -7,8 +7,11 @@ namespace Tes.Runner.Models
 {
     public class NodeTask
     {
+        public string? Id { get; set; }
+        public string? WorkflowId { get; set; }
         public string? ImageTag { get; set; }
         public string? ImageName { get; set; }
+        public string? ContainerWorkDir { get; set; }
         public List<string>? CommandsToExecute { get; set; }
         public List<FileInput>? Inputs { get; set; }
         public List<FileOutput>? Outputs { get; set; }
@@ -32,12 +35,20 @@ public class FileInput
         public string? Path { get; set; }
         public string? MountParentDirectory { get; set; }
         public string? SourceUrl { get; set; }
-        public TransformationStrategy? SasStrategy { get; set; }
+        public TransformationStrategy? TransformationStrategy { get; set; }
     }
 
     public class RuntimeOptions
     {
         public TerraRuntimeOptions? Terra { get; set; }
+
+        public string? NodeManagedIdentityResourceId { get; set; }
+    }
+
+    public class DockerCleanUpOptions
+    {
+        public bool ExecuteRmi { get; set; }
+        public bool ExecutePrune { get; set; }
     }
 
 

src/GenerateBatchVmSkus/GenerateBatchVmSkus.csproj

@@ -8,7 +8,7 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Azure.Identity" Version="1.8.2" />
+    <PackageReference Include="Azure.Identity" Version="1.10.1" />
     <PackageReference Include="Azure.ResourceManager.Batch" Version="1.1.1" />
     <PackageReference Include="Azure.ResourceManager.Compute" Version="1.1.0" />
     <PackageReference Include="Azure.Security.KeyVault.Secrets" Version="4.4.0" />

src/Tes.ApiClients/TerraWsmApiClient.cs

@@ -5,7 +5,6 @@
 using System.Text.Json;
 using System.Text.Json.Serialization;
 using Azure.Core;
-using Azure.Identity;
 using Microsoft.Extensions.Caching.Memory;
 using Microsoft.Extensions.Logging;
 using Tes.ApiClients.Models.Terra;
@@ -35,13 +34,13 @@ public class TerraWsmApiClient : TerraApiClient
 
         }
 
-        public static TerraWsmApiClient CreateTerraWsmApiClient(string apiUrl)
+        public static TerraWsmApiClient CreateTerraWsmApiClient(string apiUrl, TokenCredential tokenCredential)
         {
             var retryPolicyOptions = new RetryPolicyOptions();
             var cacheRetryHandler = new CachingRetryHandler(sharedMemoryCache,
                  Microsoft.Extensions.Options.Options.Create(retryPolicyOptions));
 
-            return new TerraWsmApiClient(apiUrl, new DefaultAzureCredential(), cacheRetryHandler, ApiClientsLoggerFactory.Create<TerraWsmApiClient>());
+            return new TerraWsmApiClient(apiUrl, tokenCredential, cacheRetryHandler, ApiClientsLoggerFactory.Create<TerraWsmApiClient>());
         }
 
         /// <summary>

src/Tes.ApiClients/Tes.ApiClients.csproj

@@ -7,8 +7,8 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="Azure.Core" Version="1.31.0" />
-    <PackageReference Include="Azure.Identity" Version="1.8.2" />
+    <PackageReference Include="Azure.Core" Version="1.35.0" />
+    <PackageReference Include="Azure.Identity" Version="1.10.1" />
     <PackageReference Include="Microsoft.Extensions.Caching.Abstractions" Version="7.0.0" />
     <PackageReference Include="Microsoft.Extensions.Caching.Memory" Version="7.0.0" />
     <PackageReference Include="Microsoft.Extensions.Logging" Version="7.0.0" />

src/Tes.Runner.Test/ResolutionPolicyHandlerTests.cs

@@ -72,9 +72,9 @@ public async Task ApplyResolutionPolicyAsync_WhenTestTaskInputsIsNotEmpty_Return
         {
             var testTaskInputs = new List<FileInput>
             {
-                new FileInput(){Path = "file", SourceUrl = "http://foo.bar", SasStrategy = TransformationStrategy.None},
-                new FileInput(){Path = "file1", SourceUrl = "http://foo1.bar", SasStrategy = TransformationStrategy.None},
-                new FileInput(){Path = "file2", SourceUrl = "http://foo2.bar", SasStrategy = TransformationStrategy.None}
+                new FileInput(){Path = "file", SourceUrl = "http://foo.bar", TransformationStrategy = TransformationStrategy.None},
+                new FileInput(){Path = "file1", SourceUrl = "http://foo1.bar", TransformationStrategy = TransformationStrategy.None},
+                new FileInput(){Path = "file2", SourceUrl = "http://foo2.bar", TransformationStrategy = TransformationStrategy.None}
             };
             var result = await resolutionPolicyHandler.ApplyResolutionPolicyAsync(testTaskInputs);
             Assert.IsNotNull(result);

src/Tes.Runner.Test/Storage/ArmUrlTransformationStrategyTests.cs

@@ -16,6 +16,7 @@ public class ArmUrlTransformationStrategyTests
         private ArmUrlTransformationStrategy armUrlTransformationStrategy = null!;
         private UserDelegationKey userDelegationKey = null!;
         const string StorageAccountName = "foo";
+        const string SasToken = "sv=2019-12-12&ss=bfqt&srt=sco&spr=https&st=2023-09-27T17%3A32%3A57Z&se=2023-09-28T17%3A32%3A57Z&sp=rwdlacupx&sig=SIGNATURE";
 
         [TestInitialize]
         public void SetUp()
@@ -56,6 +57,20 @@ public async Task TransformUrlWithStrategyAsync_InvalidBlobStorageUrl_UrlIsRetur
             Assert.AreEqual(blobUri.AbsoluteUri, transformedUrl.AbsoluteUri);
         }
 
+        [TestMethod]
+        [DataRow($"https://{StorageAccountName}.blob.core.windows.net?{SasToken}")]
+        [DataRow($"https://{StorageAccountName}.blob.core.windows.net/?{SasToken}")]
+        [DataRow($"https://{StorageAccountName}.blob.core.windows.net/cont?{SasToken}")]
+        [DataRow($"https://{StorageAccountName}.blob.core.windows.net/cont/blob?{SasToken}")]
+        public async Task TransformUrlWithStrategyAsync_BlobStorageUrlWithSasToken_UrlIsReturnAsIs(string sourceUrl)
+        {
+            var transformedUrl = await armUrlTransformationStrategy.TransformUrlWithStrategyAsync(sourceUrl, BlobSasPermissions.Read);
+
+            Assert.IsNotNull(transformedUrl);
+            var blobUri = new Uri(sourceUrl);
+            Assert.AreEqual(blobUri.AbsoluteUri, transformedUrl.AbsoluteUri);
+        }
+
         [TestMethod]
         public async Task TransformUrlWithStrategyAsync_CallTwiceForSameStorageAccount_CachesKey()
         {

src/Tes.Runner.Test/Storage/FileOperationResolverTests.cs

@@ -28,7 +28,7 @@ public void SetUp()
             {
                 Path = "/foo/bar",
                 SourceUrl = "https://foo.bar/cont?sig=sasToken",
-                SasStrategy = TransformationStrategy.None,
+                TransformationStrategy = TransformationStrategy.None,
             };
 
             singleFileOutput = new FileOutput

src/Tes.Runner/Authentication/CredentialsManager.cs

@@ -0,0 +1,64 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.Core;
+using Azure.Identity;
+using Microsoft.Extensions.Logging;
+using Polly;
+using Polly.Retry;
+using Tes.Runner.Models;
+using Tes.Runner.Transfer;
+
+namespace Tes.Runner.Authentication
+{
+    public class CredentialsManager
+    {
+        private readonly ILogger logger = PipelineLoggerFactory.Create<CredentialsManager>();
+
+        private readonly RetryPolicy retryPolicy;
+        private const int MaxRetryCount = 5;
+        private const int ExponentialBackOffExponent = 2;
+
+        public CredentialsManager()
+        {
+            retryPolicy = Policy
+                    .Handle<Exception>()
+                    .WaitAndRetry(MaxRetryCount,
+                    SleepDurationHandler);
+        }
+
+        private TimeSpan SleepDurationHandler(int attempt)
+        {
+            logger.LogInformation($"Attempt {attempt} to get token credential");
+            var duration = TimeSpan.FromSeconds(Math.Pow(ExponentialBackOffExponent, attempt));
+            logger.LogInformation($"Waiting {duration} before retrying");
+            return duration;
+        }
+
+        public TokenCredential GetTokenCredential(RuntimeOptions runtimeOptions)
+        {
+            return retryPolicy.Execute(() => GetTokenCredentialImpl(runtimeOptions));
+        }
+        private TokenCredential GetTokenCredentialImpl(RuntimeOptions runtimeOptions)
+        {
+            try
+            {
+                if (!string.IsNullOrWhiteSpace(runtimeOptions.NodeManagedIdentityResourceId))
+                {
+                    logger.LogInformation($"Token credentials with Managed Identity and resource ID: {runtimeOptions.NodeManagedIdentityResourceId}");
+
+                    return new ManagedIdentityCredential(new ResourceIdentifier(runtimeOptions.NodeManagedIdentityResourceId));
+                }
+
+                logger.LogInformation("Token credentials with DefaultAzureCredential");
+
+                return new DefaultAzureCredential();
+            }
+            catch (Exception e)
+            {
+                logger.LogError(e, "Failed to get token credential");
+                throw;
+            }
+        }
+    }
+}

src/Tes.Runner/Docker/ConsoleStreamLogReader.cs

@@ -0,0 +1,33 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+namespace Tes.Runner.Docker;
+
+public class ConsoleStreamLogReader : MultiplexedStreamLogReader
+{
+
+    public ConsoleStreamLogReader()
+    {
+
+    }
+
+    public override Task AppendStandardOutputAsync(string data)
+    {
+        Console.Write(data);
+        return Task.CompletedTask;
+    }
+
+    public override Task AppendStandardErrAsync(string data)
+    {
+        Console.Write(data);
+        return Task.CompletedTask;
+    }
+
+    public override void OnComplete(Exception? err)
+    {
+        if (err != null)
+        {
+            Console.Write(err.ToString());
+        }
+    }
+}

src/Tes.Runner/Docker/ContainerExecutionResult.cs

@@ -5,5 +5,5 @@
 
 namespace Tes.Runner.Docker
 {
-    public record ContainerExecutionResult(string Id, string? Error, long StatusCode, MultiplexedStream Logs);
+    public record ContainerExecutionResult(string Id, string? Error, long StatusCode);
 }

src/Tes.Runner/Docker/DockerExecutor.cs

@@ -4,6 +4,9 @@
 using Docker.DotNet;
 using Docker.DotNet.Models;
 using Microsoft.Extensions.Logging;
+using Microsoft.Extensions.Options;
+using Tes.ApiClients;
+using Tes.ApiClients.Options;
 using Tes.Runner.Transfer;
 
 namespace Tes.Runner.Docker
@@ -13,32 +16,43 @@ public class DockerExecutor
         private readonly IDockerClient dockerClient;
         private readonly ILogger logger = PipelineLoggerFactory.Create<DockerExecutor>();
         private readonly NetworkUtility networkUtility = new NetworkUtility();
+        private readonly RetryHandler retryHandler = new RetryHandler(Options.Create(new RetryPolicyOptions()));
+        private readonly MultiplexedStreamLogReader multiplexedStreamLogReader;
+
+        const int LogStreamingMaxWaitTimeInSeconds = 30;
 
         public DockerExecutor(Uri dockerHost)
         {
             dockerClient = new DockerClientConfiguration(dockerHost)
                 .CreateClient();
+            multiplexedStreamLogReader = new ConsoleStreamLogReader();
         }
 
-        public async Task<ContainerExecutionResult> RunOnContainerAsync(string? imageName, string? tag, List<string>? commandsToExecute, List<string>? volumeBindings)
+        public async Task<ContainerExecutionResult> RunOnContainerAsync(string? imageName, string? tag, List<string>? commandsToExecute, List<string>? volumeBindings, string? workingDir)
         {
             ArgumentException.ThrowIfNullOrEmpty(imageName);
-            ArgumentException.ThrowIfNullOrEmpty(tag);
             ArgumentNullException.ThrowIfNull(commandsToExecute);
 
-            await PullImageAsync(imageName, tag);
+            await PullImageWithRetriesAsync(imageName, tag);
 
             await ConfigureNetworkAsync();
 
-            var createResponse = await CreateContainerAsync(imageName, commandsToExecute, volumeBindings);
+            var createResponse = await CreateContainerAsync(imageName, tag, commandsToExecute, volumeBindings, workingDir);
 
             var logs = await StartContainerWithStreamingOutput(createResponse);
 
+            logger.LogInformation("Starting to read output from container");
+
+            multiplexedStreamLogReader.StartReadingFromLogStream(logs);
+
             var runResponse = await dockerClient.Containers.WaitContainerAsync(createResponse.ID);
 
-            return new ContainerExecutionResult(createResponse.ID, runResponse.Error?.Message, runResponse.StatusCode, logs);
+            await multiplexedStreamLogReader.WaitUntilAsync(TimeSpan.FromSeconds(LogStreamingMaxWaitTimeInSeconds));
+
+            return new ContainerExecutionResult(createResponse.ID, runResponse.Error?.Message, runResponse.StatusCode);
         }
 
+
         private async Task<MultiplexedStream> StartContainerWithStreamingOutput(CreateContainerResponse createResponse)
         {
             var logs = await StreamStdOutAndErrorAsync(createResponse.ID);
@@ -61,34 +75,50 @@ private async Task<MultiplexedStream> StreamStdOutAndErrorAsync(string container
                 });
         }
 
-        private async Task<CreateContainerResponse> CreateContainerAsync(string imageName, List<string> commandsToExecute, List<string>? volumeBindings)
+        private async Task<CreateContainerResponse> CreateContainerAsync(string imageName, string? imageTag,
+            List<string> commandsToExecute, List<string>? volumeBindings, string? workingDir)
         {
+            var imageWithTag = ToImageNameWithTag(imageName, imageTag);
+            logger.LogInformation($"Creating container with image name: {imageWithTag}");
+
             var createResponse = await dockerClient.Containers.CreateContainerAsync(
                 new CreateContainerParameters
                 {
-                    Image = imageName,
-                    Entrypoint = commandsToExecute,
+                    Image = imageWithTag,
+                    Cmd = commandsToExecute,
                     AttachStdout = true,
                     AttachStderr = true,
-                    WorkingDir = "/",
+                    WorkingDir = workingDir,
                     HostConfig = new HostConfig
                     {
+                        AutoRemove = true,
                         Binds = volumeBindings
                     }
                 });
             return createResponse;
         }
 
-        private async Task PullImageAsync(string imageName, string tag, AuthConfig? authConfig = null)
+        private static string ToImageNameWithTag(string imageName, string? imageTag)
         {
-            await dockerClient.Images.CreateImageAsync(
-                new ImagesCreateParameters()
-                {
-                    FromImage = imageName,
-                    Tag = tag
-                },
-                authConfig,
-                new Progress<JSONMessage>(message => logger.LogInformation(message.Status)));
+            if (string.IsNullOrWhiteSpace(imageTag))
+            {
+                return imageName;
+            }
+
+            return $"{imageName}:{imageTag}";
+        }
+
+        private async Task PullImageWithRetriesAsync(string imageName, string? tag, AuthConfig? authConfig = null)
+        {
+            logger.LogInformation($"Pulling image name: {imageName} image tag: {tag}");
+
+            await retryHandler.AsyncRetryPolicy.ExecuteAsync(async () =>
+            {
+                await dockerClient.Images.CreateImageAsync(
+                    new ImagesCreateParameters() { FromImage = imageName, Tag = tag },
+                    authConfig,
+                    new Progress<JSONMessage>(message => logger.LogDebug(message.Status)));
+            });
         }
 
         /// <summary>