diff --git a/v1/image.go b/v1/image.go
index 7ec4c49..57d21ee 100644
--- a/v1/image.go
+++ b/v1/image.go
@@ -40,6 +40,9 @@ func ValidateInstanceImage(ctx context.Context, instance Instance, privateKey st
 		}
 	}()
 
+	// If cloud-init is present, wait until it has finished (e.g. UFW/runcmd done).
+	_, _, _ = sshClient.RunCommand(ctx, "cloud-init status --wait 2>/dev/null || true")
+
 	arch, err := validateArchitecture(ctx, sshClient)
 	if err != nil {
 		return err
diff --git a/v1/providers/nebius/instance.go b/v1/providers/nebius/instance.go
index 6f43534..3460efd 100644
--- a/v1/providers/nebius/instance.go
+++ b/v1/providers/nebius/instance.go
@@ -1748,7 +1748,10 @@ func (c *NebiusClient) cleanupOrphanedBootDisks(ctx context.Context, testID stri
 // This is inspired by Shadeform's LaunchConfiguration approach but uses cloud-init instead of base64 scripts
 func generateCloudInitUserData(publicKey string, firewallRules v1.FirewallRules) string {
 	// Start with cloud-init header
-	script := "#cloud-config\n"
+	script := `#cloud-config
+packages:
+  - ufw
+`
 
 	// Add SSH key configuration if provided
 	if publicKey != "" {