diff --git a/pkg/workflow/compiler_jobs.go b/pkg/workflow/compiler_jobs.go index e99cb8f29d..5c0e019af9 100644 --- a/pkg/workflow/compiler_jobs.go +++ b/pkg/workflow/compiler_jobs.go @@ -808,6 +808,9 @@ func (c *Compiler) buildCustomJobs(data *WorkflowData, activationJobCreated bool // don't inherit GITHUB_ENV from the agent job, so the gh CLI won't // know which host to target without this step. job.Steps = append(job.Steps, generateGHESHostConfigurationStep()) + if shouldInjectNodeSetupForGPUCustomJob(configMap) { + job.Steps = append(job.Steps, generateNodeSetupStepForCustomJob(data)) + } job.Steps = append(job.Steps, preSteps...) job.Steps = append(job.Steps, regularSteps...) } @@ -824,6 +827,91 @@ func (c *Compiler) buildCustomJobs(data *WorkflowData, activationJobCreated bool return nil } +func shouldInjectNodeSetupForGPUCustomJob(configMap map[string]any) bool { + if configMap == nil { + return false + } + runsOn, hasRunsOn := configMap["runs-on"] + if !hasRunsOn || !containsRunnerLabel(runsOn, "aw-gpu-runner-T4") { + return false + } + return !jobStepsContainSetupNode(configMap) +} + +func containsRunnerLabel(value any, target string) bool { + switch v := value.(type) { + case string: + return strings.EqualFold(strings.TrimSpace(v), target) + case []any: + for _, item := range v { + if containsRunnerLabel(item, target) { + return true + } + } + case map[string]any: + // Support object-form runs-on (e.g. {group: "...", labels: ["aw-gpu-runner-T4"]}). + // GitHub Actions allows object form for larger/self-hosted runners. + for _, item := range v { + if containsRunnerLabel(item, target) { + return true + } + } + } + return false +} + +func jobStepsContainSetupNode(configMap map[string]any) bool { + for _, fieldName := range []string{"pre-steps", "steps"} { + fieldValue, hasField := configMap[fieldName] + if !hasField { + continue + } + steps, ok := fieldValue.([]any) + if !ok { + continue + } + for _, step := range steps { + stepMap, ok := step.(map[string]any) + if !ok { + continue + } + usesValue, hasUses := stepMap["uses"] + if !hasUses { + continue + } + usesStr, ok := usesValue.(string) + if !ok { + continue + } + if strings.Contains(strings.ToLower(usesStr), "setup-node") { + return true + } + } + } + return false +} + +func generateNodeSetupStepForCustomJob(data *WorkflowData) string { + requirements := map[string]*RuntimeRequirement{} + nodeRuntime := findRuntimeByID("node") + if nodeRuntime == nil { + compilerJobsLog.Print("Node runtime definition not found; skipping Node setup step injection for GPU custom job") + return "" + } + + updateRequiredRuntime(nodeRuntime, "", requirements) + if data != nil && data.Runtimes != nil { + applyRuntimeOverrides(data.Runtimes, requirements) + } + nodeRequirement, exists := requirements["node"] + if !exists { + compilerJobsLog.Print("Node runtime requirement missing after overrides; skipping Node setup step injection for GPU custom job") + return "" + } + + return strings.Join(generateSetupStep(nodeRequirement), "\n") + "\n" +} + func (c *Compiler) applyBuiltinJobPreSteps(data *WorkflowData) error { if data == nil || data.Jobs == nil { return nil diff --git a/pkg/workflow/compiler_jobs_test.go b/pkg/workflow/compiler_jobs_test.go index 18c3a7e2f0..512a10e2c3 100644 --- a/pkg/workflow/compiler_jobs_test.go +++ b/pkg/workflow/compiler_jobs_test.go @@ -2576,6 +2576,123 @@ func TestBuildCustomJobsRunsOnForms(t *testing.T) { } } +func TestBuildCustomJobsAddsNodeSetupForAWGPURunner(t *testing.T) { + compiler := NewCompiler() + compiler.jobManager = NewJobManager() + + data := &WorkflowData{ + Name: "Test Workflow", + AI: "copilot", + Jobs: map[string]any{ + "gpu_job": map[string]any{ + "runs-on": "aw-gpu-runner-T4", + "steps": []any{ + map[string]any{"name": "Work", "run": "echo hi"}, + }, + }, + }, + } + + err := compiler.buildCustomJobs(data, false) + if err != nil { + t.Fatalf("buildCustomJobs() returned unexpected error: %v", err) + } + + job, exists := compiler.jobManager.GetJob("gpu_job") + if !exists { + t.Fatal("Expected gpu_job to be added") + } + + stepsContent := strings.Join(job.Steps, "") + if !strings.Contains(stepsContent, "name: Setup Node.js") { + t.Fatalf("Expected custom GPU job to include Node setup step, got:\n%s", stepsContent) + } + if !strings.Contains(stepsContent, "node-version: '24'") { + t.Fatalf("Expected default Node 24 for custom GPU job, got:\n%s", stepsContent) + } +} + +func TestBuildCustomJobsAddsNodeSetupForAWGPURunnerWithRuntimeOverride(t *testing.T) { + compiler := NewCompiler() + compiler.jobManager = NewJobManager() + + data := &WorkflowData{ + Name: "Test Workflow", + AI: "copilot", + Runtimes: map[string]any{ + "node": map[string]any{ + "version": "20", + }, + }, + Jobs: map[string]any{ + "gpu_job": map[string]any{ + "runs-on": "aw-gpu-runner-T4", + "steps": []any{ + map[string]any{"name": "Work", "run": "echo hi"}, + }, + }, + }, + } + + err := compiler.buildCustomJobs(data, false) + if err != nil { + t.Fatalf("buildCustomJobs() returned unexpected error: %v", err) + } + + job, exists := compiler.jobManager.GetJob("gpu_job") + if !exists { + t.Fatal("Expected gpu_job to be added") + } + + stepsContent := strings.Join(job.Steps, "") + if !strings.Contains(stepsContent, "name: Setup Node.js") { + t.Fatalf("Expected custom GPU job to include Node setup step, got:\n%s", stepsContent) + } + if !strings.Contains(stepsContent, "node-version: '20'") { + t.Fatalf("Expected Node version to respect runtime override for custom GPU job, got:\n%s", stepsContent) + } +} + +func TestBuildCustomJobsSkipsNodeSetupForAWGPURunnerWhenAlreadyPresent(t *testing.T) { + compiler := NewCompiler() + compiler.jobManager = NewJobManager() + + data := &WorkflowData{ + Name: "Test Workflow", + AI: "copilot", + Jobs: map[string]any{ + "gpu_job": map[string]any{ + "runs-on": "aw-gpu-runner-T4", + "steps": []any{ + map[string]any{ + "name": "Setup Node.js (manual)", + "uses": "actions/setup-node@v6", + "with": map[string]any{ + "node-version": "24", + }, + }, + map[string]any{"name": "Work", "run": "echo hi"}, + }, + }, + }, + } + + err := compiler.buildCustomJobs(data, false) + if err != nil { + t.Fatalf("buildCustomJobs() returned unexpected error: %v", err) + } + + job, exists := compiler.jobManager.GetJob("gpu_job") + if !exists { + t.Fatal("Expected gpu_job to be added") + } + + stepsContent := strings.Join(job.Steps, "") + if strings.Count(stepsContent, "uses: actions/setup-node@") != 1 { + t.Fatalf("Expected exactly one setup-node action step when already present, got:\n%s", stepsContent) + } +} + // TestBuildCustomJobsNewSimpleFields tests extraction of simple job fields via CompileWorkflow func TestBuildCustomJobsNewSimpleFields(t *testing.T) { tmpDir := testutil.TempDir(t, "new-simple-fields-test")