fix: report PIS bake completion

e2e/scenario/provision.go

@@ -47,7 +47,11 @@ func scriptlessUnsupported(s *Scenario) bool {
 func runVHDCachingScenario(ctx context.Context, name string, original *Scenario) error {
 	bakeScenario := freshScenario(original)
 	bakeScenario.artifactName = vhdStageArtifactName(original, name, "vhd-bake")
+	scriptless := !config.Config.DisableScriptless && !original.IsWindows() &&
+		len(original.Config.CustomDataWriteFiles) == 0 && original.VHD != nil &&
+		original.VHD.SupportsScriptless() && config.Config.SIGVersionTagName == "buildId"
 
+	bakeScenario.Runtime = &ScenarioRuntime{EnableScriptlessNBCCSECmd: scriptless}
 	bakeScenario.Config.SkipDefaultValidation = true
 	bakeScenario.Config.Validator = func(ctx context.Context, scenario *Scenario) error {
 		var validationErr error
@@ -62,13 +66,18 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 				ValidateWindowsServiceIsRunning(ctx, scenario, "containerd"),
 			)
 		} else {
-			validationErr = errors.Join(
+			checks := []error{
 				ValidateFileExists(ctx, scenario, "/etc/containerd/config.toml"),
 				ValidateFileExists(ctx, scenario, "/opt/azure/containers/base_prep.complete"),
-				ValidateFileDoesNotExist(ctx, scenario, "/opt/azure/containers/provision.complete"),
 				ValidateSystemdUnitIsRunning(ctx, scenario, "containerd"),
 				ValidateSystemdUnitIsNotRunning(ctx, scenario, "kubelet"),
-			)
+			}
+			if scriptless {
+				checks = append(checks, ValidateFileExists(ctx, scenario, "/opt/azure/containers/provision.complete"))
+			} else {
+				checks = append(checks, ValidateFileDoesNotExist(ctx, scenario, "/opt/azure/containers/provision.complete"))
+			}
+			validationErr = errors.Join(checks...)
 		}
 		return validationErr
 	}
@@ -93,7 +102,7 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 				}
 			}
 			nbc.PreProvisionOnly = true
-			nbc.EnableScriptlessNBCCSECmd = false
+			nbc.EnableScriptlessNBCCSECmd = scriptless
 			nbc.EnableScriptlessAKSNodeConfig = false
 			// Bake-only mutation: lets a scenario deliberately diverge bake-time
 			// state from provision-time state (e.g. a stale sentinel bootstrap token).
@@ -117,6 +126,15 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 	}
 
 	logging.Log(ctx, "=== Creating VHD Image ===")
+	if !bakeScenario.IsWindows() {
+		// Match RP image generalization by removing per-run provisioning state before capture.
+		cleanup := "sudo rm -f /opt/azure/containers/provision.complete /var/log/azure/aks/provision.json" +
+			" /opt/azure/containers/aks-node-controller-config.json /opt/azure/containers/aks-node-controller-nbc-cmd.sh" +
+			" /opt/azure/containers/boothook.sh /opt/bin/boothook.sh"
+		if _, err = execScriptOnVMForScenarioValidateExitCode(ctx, bakeScenario, cleanup, 0, "generalize VM before image capture"); err != nil {
+			return err
+		}
+	}
 	bakeScenario.failed = true
 	customVHD, err := CreateImage(ctx, bakeScenario)
 	if err != nil {
@@ -131,12 +149,16 @@ func runVHDCachingScenario(ctx context.Context, name string, original *Scenario)
 	provisionScenario := freshScenario(original)
 	provisionScenario.artifactName = vhdStageArtifactName(original, name, "vhd-provision")
 	provisionScenario.Config.VHD = customVHD
+	provisionScenario.Runtime = &ScenarioRuntime{EnableScriptlessNBCCSECmd: scriptless}
 	provisionScenario.Config.Validator = func(ctx context.Context, s *Scenario) error {
 		var markerErr error
 		if s.IsWindows() {
 			markerErr = ValidateFileExists(ctx, s, "C:\\AzureData\\provision.complete")
 		} else {
 			markerErr = ValidateFileExists(ctx, s, "/opt/azure/containers/provision.complete")
+			if markerErr == nil && scriptless {
+				markerErr = ValidateFileHasContent(ctx, s, "/var/log/azure/cluster-provision.log", "Skipping basePrep - base_prep.complete file exists")
+			}
 		}
 		if markerErr != nil {
 			return markerErr

e2e/scenario/vmss.go

@@ -619,7 +619,6 @@ func usesScriptlessProvisioning(s *Scenario) bool {
 	}
 	nbc := s.Runtime.NBC
 	return (nbc.EnableScriptlessAKSNodeConfig || nbc.EnableScriptlessNBCCSECmd) &&
-		!nbc.PreProvisionOnly &&
 		s.VHD.SupportsScriptless()
 }
 

parts/linux/cloud-init/artifacts/cse_start.sh

@@ -117,15 +117,17 @@ upload_logs() {
     fi
 }
 # Create the marker for the completed provisioning stage.
+mkdir -p /opt/azure/containers
 if [ "${PRE_PROVISION_ONLY}" = "true" ]; then
     # Stage 1: Create marker indicating Stage 2 is needed
-    mkdir -p /opt/azure/containers && touch /opt/azure/containers/base_prep.complete
+    touch /opt/azure/containers/base_prep.complete
     echo "Stage 1 complete - kubelet configuration skipped, Stage 2 required" >> /var/log/azure/cluster-provision.log
     echo "Created base_prep.complete marker file" >> /var/log/azure/cluster-provision.log
-else
-    # provision.complete signals that a normal provisioning attempt finished.
-    mkdir -p /opt/azure/containers && touch /opt/azure/containers/provision.complete
 fi
+# provision.complete signals that this CSE run finished, so `aks-node-controller provision-wait`
+# reports the result of a pre-provision run as well as a normal one. Image generalization removes
+# it before capture, so a node created from the image never inherits it.
+touch /opt/azure/containers/provision.complete
 
 if [ "$EXIT_CODE" -ne 0 ]; then
     upload_logs

pkg/agent/baker.go

@@ -284,7 +284,10 @@ func (t *TemplateGenerator) getScriptlessConfiguration(config *datamodel.NodeBoo
 }
 
 func supportsScriptlessProvisioning(config *datamodel.NodeBootstrappingConfiguration) bool {
-	return (config.EnableScriptlessAKSNodeConfig || (config.EnableScriptlessNBCCSECmd) || config.IsCustomDataOnlyProvisioningEnabled()) && !config.PreProvisionOnly
+	// NBC scriptless phase 2 supports pre-provision runs because cse_start.sh writes
+	// provision.complete after the bake result. Other scriptless modes remain unsupported.
+	return config.EnableScriptlessNBCCSECmd ||
+		(config.EnableScriptlessAKSNodeConfig || config.IsCustomDataOnlyProvisioningEnabled()) && !config.PreProvisionOnly
 }
 
 func validateScriptlessAKSNodeConfig(config *datamodel.NodeBootstrappingConfiguration) error {

pkg/agent/baker_test.go

@@ -343,7 +343,12 @@ var _ = Describe("Assert generated customData and cseCmd", func() {
 		})
 
 		Describe(".supportsScriptlessProvisioning()", func() {
-			It("given EnableScriptlessNBCCSECmd, PreProvisionOnly is true and no CustomCATrustConfig, it returns false", func() {
+			It("given EnableScriptlessNBCCSECmd and PreProvisionOnly, it returns true", func() {
+				config.EnableScriptlessNBCCSECmd = true
+				config.PreProvisionOnly = true
+				Expect(supportsScriptlessProvisioning(config)).To(BeTrue())
+			})
+			It("given PreProvisionOnly without EnableScriptlessNBCCSECmd, it returns false", func() {
 				config.PreProvisionOnly = true
 				Expect(supportsScriptlessProvisioning(config)).To(BeFalse())
 			})
@@ -1537,27 +1542,20 @@ var _ = Describe("getLinuxNodeBootstrappingPayload", func() {
 		Expect(nodeCustomData).To(ContainSubstring("encoding: gzip"))
 	})
 
-	It("should fall back to regular custom data when pre-provisioning is enabled", func() {
+	It("should use scriptless custom data when pre-provisioning is enabled", func() {
 		templateGenerator := InitializeTemplateGenerator()
 		config := newConfig(true)
 
 		payload := templateGenerator.getLinuxNodeBootstrappingPayload(config)
 		decodedPayload, err := base64.StdEncoding.DecodeString(payload)
 		Expect(err).NotTo(HaveOccurred())
 
-		decompressedPayload, err := getGzipDecodedValue(decodedPayload)
-		Expect(err).NotTo(HaveOccurred())
-
-		expectedCustomData := getCustomDataFromJSON(templateGenerator.getLinuxNodeCustomDataJSONObject(config))
-
-		Expect(string(decompressedPayload)).To(Equal(expectedCustomData))
-		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNodeCustomDataFilepath))
-		Expect(string(decompressedPayload)).NotTo(ContainSubstring(aksNbcCmdFilepath))
+		Expect(string(decodedPayload)).To(ContainSubstring(aksNodeCustomDataFilepath))
 	})
 })
 
 var _ = Describe("getNodeBootstrappingCmd", func() {
-	It("should use the regular linux CSE command when pre-provisioning is enabled", func() {
+	It("should use the provision-wait CSE command when pre-provisioning is enabled", func() {
 		templateGenerator := InitializeTemplateGenerator()
 		agentPoolProfile := &datamodel.AgentPoolProfile{
 			Name:   "nodepool1",
@@ -1589,8 +1587,7 @@ var _ = Describe("getNodeBootstrappingCmd", func() {
 			PreProvisionOnly:          true,
 		}
 
-		Expect(templateGenerator.getNodeBootstrappingCmd(config)).To(Equal(templateGenerator.getLinuxNodeCSECommand(config)))
-		Expect(templateGenerator.getNodeBootstrappingCmd(config)).NotTo(Equal("/opt/azure/containers/aks-node-controller provision-wait"))
+		Expect(templateGenerator.getNodeBootstrappingCmd(config)).To(Equal("/opt/azure/containers/aks-node-controller provision-wait"))
 	})
 
 	newScriptlessCmdTestConfig := func() *datamodel.NodeBootstrappingConfiguration {

spec/parts/linux/cloud-init/artifacts/cse_start_spec.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# provision.complete is what `aks-node-controller provision-wait` waits for, so it must be written
+# for a pre-provision run as well as a normal one. base_prep.complete is durable phase state and is
+# written only by a bake, so a node created from the image skips basePrep.
+
+Describe 'cse_start.sh completion markers'
+    markers_for() {
+        root="$(mktemp -d)"
+        mkdir -p "${root}/var/log/azure"
+        PRE_PROVISION_ONLY="${1}" bash -c "$(
+            awk '/^# Create the marker for the completed provisioning stage\.$/,/^touch .*provision\.complete$/' \
+                ./parts/linux/cloud-init/artifacts/cse_start.sh |
+                sed "s|/opt/azure/containers|${root}/opt/azure/containers|g; s|/var/log/azure|${root}/var/log/azure|g"
+        )"
+        ls "${root}/opt/azure/containers" | sort
+        rm -rf "${root}"
+    }
+
+    It 'writes base_prep.complete and provision.complete for a pre-provision run'
+        When call markers_for "true"
+        The line 1 of output should equal "base_prep.complete"
+        The line 2 of output should equal "provision.complete"
+        The lines of output should equal 2
+    End
+
+    It 'writes only provision.complete for a normal provisioning run'
+        When call markers_for "false"
+        The output should equal "provision.complete"
+    End
+End