From 60af948051cbeaf51d87397c40354c1496fcdd8f Mon Sep 17 00:00:00 2001 From: Lawrence Gripper Date: Thu, 16 Oct 2025 21:16:14 +0100 Subject: [PATCH] Custom Image: Preflight checks (#4081) Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- src/Runner.Common/Constants.cs | 2 + src/Runner.Worker/JobExtension.cs | 4 + .../SnapshotOperationProvider.cs | 27 +++ src/Test/L0/Worker/JobExtensionL0.cs | 188 ++++++++++++++++++ .../L0/Worker/SnapshotOperationProviderL0.cs | 1 + 5 files changed, 222 insertions(+) diff --git a/src/Runner.Common/Constants.cs b/src/Runner.Common/Constants.cs index f3550842b..45ce81e0e 100644 --- a/src/Runner.Common/Constants.cs +++ b/src/Runner.Common/Constants.cs @@ -170,6 +170,8 @@ namespace GitHub.Runner.Common public static readonly string AddCheckRunIdToJobContext = "actions_add_check_run_id_to_job_context"; public static readonly string DisplayHelpfulActionsDownloadErrors = "actions_display_helpful_actions_download_errors"; public static readonly string ContainerActionRunnerTemp = "actions_container_action_runner_temp"; + public static readonly string SnapshotPreflightHostedRunnerCheck = "actions_snapshot_preflight_hosted_runner_check"; + public static readonly string SnapshotPreflightImageGenPoolCheck = "actions_snapshot_preflight_image_gen_pool_check"; } // Node version migration related constants diff --git a/src/Runner.Worker/JobExtension.cs b/src/Runner.Worker/JobExtension.cs index 58e8929b4..14b798e6d 100644 --- a/src/Runner.Worker/JobExtension.cs +++ b/src/Runner.Worker/JobExtension.cs @@ -400,6 +400,10 @@ namespace GitHub.Runner.Worker if (snapshotRequest != null) { var snapshotOperationProvider = HostContext.GetService(); + // Check that that runner is capable of taking a snapshot + snapshotOperationProvider.RunSnapshotPreflightChecks(context); + + // Add postjob step to write snapshot file jobContext.RegisterPostJobStep(new JobExtensionRunner( runAsync: (executionContext, _) => snapshotOperationProvider.CreateSnapshotRequestAsync(executionContext, snapshotRequest), condition: snapshotRequest.Condition, diff --git a/src/Runner.Worker/SnapshotOperationProvider.cs b/src/Runner.Worker/SnapshotOperationProvider.cs index 73630d498..16024e067 100644 --- a/src/Runner.Worker/SnapshotOperationProvider.cs +++ b/src/Runner.Worker/SnapshotOperationProvider.cs @@ -1,15 +1,19 @@ #nullable enable +using System; using System.IO; using System.Threading.Tasks; using GitHub.DistributedTask.Pipelines; +using GitHub.DistributedTask.WebApi; using GitHub.Runner.Common; using GitHub.Runner.Sdk; +using GitHub.Runner.Worker.Handlers; namespace GitHub.Runner.Worker; [ServiceLocator(Default = typeof(SnapshotOperationProvider))] public interface ISnapshotOperationProvider : IRunnerService { Task CreateSnapshotRequestAsync(IExecutionContext executionContext, Snapshot snapshotRequest); + void RunSnapshotPreflightChecks(IExecutionContext jobContext); } public class SnapshotOperationProvider : RunnerService, ISnapshotOperationProvider @@ -24,9 +28,32 @@ public class SnapshotOperationProvider : RunnerService, ISnapshotOperationProvid } IOUtil.SaveObject(snapshotRequest, snapshotRequestFilePath); + executionContext.Output($"Image Name: {snapshotRequest.ImageName} Version: {snapshotRequest.Version}"); executionContext.Output($"Request written to: {snapshotRequestFilePath}"); executionContext.Output("This request will be processed after the job completes. You will not receive any feedback on the snapshot process within the workflow logs of this job."); executionContext.Output("If the snapshot process is successful, you should see a new image with the requested name in the list of available custom images when creating a new GitHub-hosted Runner."); return Task.CompletedTask; } + + public void RunSnapshotPreflightChecks(IExecutionContext context) + { + var shouldCheckRunnerEnvironment = context.Global.Variables.GetBoolean(Constants.Runner.Features.SnapshotPreflightHostedRunnerCheck) ?? false; + if (shouldCheckRunnerEnvironment && + context.Global.Variables.TryGetValue(WellKnownDistributedTaskVariables.RunnerEnvironment, out var runnerEnvironment) && + !string.IsNullOrEmpty(runnerEnvironment)) + { + context.Debug($"Snapshot: RUNNER_ENVIRONMENT={runnerEnvironment}"); + if (!string.Equals(runnerEnvironment, "github-hosted", StringComparison.OrdinalIgnoreCase)) + { + throw new ArgumentException("Snapshot workflows must be run on a GitHub Hosted Runner"); + } + } + var imageGenEnabled = StringUtil.ConvertToBoolean(Environment.GetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED")); + context.Debug($"Snapshot: GITHUB_ACTIONS_IMAGE_GEN_ENABLED={imageGenEnabled}"); + var shouldCheckImageGenPool = context.Global.Variables.GetBoolean(Constants.Runner.Features.SnapshotPreflightImageGenPoolCheck) ?? false; + if (shouldCheckImageGenPool && !imageGenEnabled) + { + throw new ArgumentException("Snapshot workflows must be run a hosted runner with Image Generation enabled"); + } + } } diff --git a/src/Test/L0/Worker/JobExtensionL0.cs b/src/Test/L0/Worker/JobExtensionL0.cs index 9ce99070b..60814998e 100644 --- a/src/Test/L0/Worker/JobExtensionL0.cs +++ b/src/Test/L0/Worker/JobExtensionL0.cs @@ -567,5 +567,193 @@ namespace GitHub.Runner.Common.Tests.Worker } } } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_HostedRunnerCheck_Enabled_GitHubHosted_Success() + { + using (TestHostContext hc = CreateTestContext()) + { + _jobEc.Global.Variables.Set(WellKnownDistributedTaskVariables.RunnerEnvironment, "github-hosted"); + + hc.SetSingleton(new SnapshotOperationProvider()); + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightHostedRunnerCheck, "true"); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + await jobExtension.InitializeJob(_jobEc, _message); + + var postJobSteps = _jobEc.PostJobSteps; + Assert.Equal(1, postJobSteps.Count); + } + + Environment.SetEnvironmentVariable("RUNNER_ENVIRONMENT", null); + } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_HostedRunnerCheck_Enabled_SelfHosted_ThrowsException() + { + using (TestHostContext hc = CreateTestContext()) + { + _jobEc.Global.Variables.Set(WellKnownDistributedTaskVariables.RunnerEnvironment, "self-hosted"); + hc.SetSingleton(new SnapshotOperationProvider()); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightHostedRunnerCheck, "true"); + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + var exception = await Assert.ThrowsAsync(() => jobExtension.InitializeJob(_jobEc, _message)); + Assert.Contains("Snapshot workflows must be run on a GitHub Hosted Runner", exception.Message); + } + + Environment.SetEnvironmentVariable("RUNNER_ENVIRONMENT", null); + } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_ImageGenPoolCheck_Enabled_ImageGenEnabled_Success() + { + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", "true"); + + using (TestHostContext hc = CreateTestContext()) + { + hc.SetSingleton(new SnapshotOperationProvider()); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightImageGenPoolCheck, "true"); + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + await jobExtension.InitializeJob(_jobEc, _message); + + var postJobSteps = _jobEc.PostJobSteps; + Assert.Equal(1, postJobSteps.Count); + } + + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", null); + } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_ImageGenPoolCheck_Enabled_ImageGen_False_ThrowsException() + { + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", "false"); + + using (TestHostContext hc = CreateTestContext()) + { + hc.SetSingleton(new SnapshotOperationProvider()); + _jobEc.SetRunnerContext("environment", "github-hosted"); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightImageGenPoolCheck, "true"); + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + var exception = await Assert.ThrowsAsync(() => jobExtension.InitializeJob(_jobEc, _message)); + Assert.Contains("Snapshot workflows must be run a hosted runner with Image Generation enabled", exception.Message); + } + + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", null); + } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_ImageGenPoolCheck_Enabled_ImageGen_Missing_ThrowsException() + { + using (TestHostContext hc = CreateTestContext()) + { + hc.SetSingleton(new SnapshotOperationProvider()); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightImageGenPoolCheck, "true"); + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + var exception = await Assert.ThrowsAsync(() => jobExtension.InitializeJob(_jobEc, _message)); + Assert.Contains("Snapshot workflows must be run a hosted runner with Image Generation enabled", exception.Message); + } + + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", null); + } + + [Fact] + [Trait("Level", "L0")] + [Trait("Category", "Worker")] + public async Task SnapshotPreflightChecks_BothChecks_Enabled_AllConditionsMet_Success() + { + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", "true"); + + using (TestHostContext hc = CreateTestContext()) + { + hc.SetSingleton(new SnapshotOperationProvider()); + + var jobExtension = new JobExtension(); + jobExtension.Initialize(hc); + + // Enable both preflight checks + _jobEc.Global.Variables.Set(WellKnownDistributedTaskVariables.RunnerEnvironment, "github-hosted"); + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightHostedRunnerCheck, "true"); + _jobEc.Global.Variables.Set(Constants.Runner.Features.SnapshotPreflightImageGenPoolCheck, "true"); + + var snapshot = new Pipelines.Snapshot("TestImageNameForPreflightCheck"); + var imageNameValueStringToken = new StringToken(null, null, null, snapshot.ImageName); + _message.Snapshot = imageNameValueStringToken; + + _actionManager.Setup(x => x.PrepareActionsAsync(It.IsAny(), It.IsAny>(), It.IsAny())) + .Returns(Task.FromResult(new PrepareResult(new List(), new Dictionary()))); + + await jobExtension.InitializeJob(_jobEc, _message); + + var postJobSteps = _jobEc.PostJobSteps; + Assert.Equal(1, postJobSteps.Count); + } + + Environment.SetEnvironmentVariable("RUNNER_ENVIRONMENT", null); + Environment.SetEnvironmentVariable("GITHUB_ACTIONS_IMAGE_GEN_ENABLED", null); + } } } diff --git a/src/Test/L0/Worker/SnapshotOperationProviderL0.cs b/src/Test/L0/Worker/SnapshotOperationProviderL0.cs index 4f747ae8e..d2d5260b6 100644 --- a/src/Test/L0/Worker/SnapshotOperationProviderL0.cs +++ b/src/Test/L0/Worker/SnapshotOperationProviderL0.cs @@ -38,6 +38,7 @@ public class SnapshotOperationProviderL0 Assert.NotNull(actualSnapshot); Assert.Equal(expectedSnapshot.ImageName, actualSnapshot!.ImageName); _ec.Verify(ec => ec.Write(null, $"Request written to: {_snapshotRequestFilePath}"), Times.Once); + _ec.Verify(ec => ec.Write(null, $"Image Name: {expectedSnapshot.ImageName} Version: {expectedSnapshot.Version}"), Times.Once); _ec.Verify(ec => ec.Write(null, "This request will be processed after the job completes. You will not receive any feedback on the snapshot process within the workflow logs of this job."), Times.Once); _ec.Verify(ec => ec.Write(null, "If the snapshot process is successful, you should see a new image with the requested name in the list of available custom images when creating a new GitHub-hosted Runner."), Times.Once); _ec.VerifyNoOtherCalls();