[air] Record workload submission and source staging telemetry

cmd/air/runsubmit.go

@@ -10,6 +10,7 @@ import (
 	"path"
 	"strconv"
 	"strings"
+	"time"
 
 	"github.com/databricks/cli/libs/auth"
 	"github.com/databricks/cli/libs/cmdio"
@@ -321,17 +322,20 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl
 		})
 	}
 
-	if err := group.Wait(); err != nil {
-		return snapshotResult{}, err
-	}
-	return snap, nil
+	err := group.Wait()
+	// Preserve measurements from attempted snapshot phases even if staging fails.
+	return snap, err
 }
 
 // submitWorkload runs the submit happy path: ensure the experiment directory,
 // upload the launch artifacts, assemble the Jobs payload, and submit it. It
 // returns the new run_id and its dashboard URL. showProgress enables the stderr
 // staging spinner (text mode only).
-func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) {
+func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) {
+	start := time.Now()
+	var snap snapshotResult
+	defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }()
+
 	// Compute and validate the actual submission path before creating artifacts.
 	base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg)
 	if err != nil {
@@ -399,7 +403,6 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 		}
 	}
 
-	var snap snapshotResult
 	err = withSpinner(ctx, showProgress, "Staging run artifacts…", func() error {
 		var stageErr error
 		snap, stageErr = stageRunArtifacts(ctx, fc, items, stageSnapshot)
@@ -426,12 +429,12 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 	// Submit returns as soon as the run is created; we don't wait for it to finish.
 	// Permissions are granted by the caller, after the submit result is shown, so
 	// the best-effort grant never delays the success line.
-	runID, err := submitRun(ctx, w, payload, poolID, priorityClass, cfg.unityCatalogImagePath(), containers)
+	runID, err = submitRun(ctx, w, payload, poolID, priorityClass, cfg.unityCatalogImagePath(), containers)
 	if err != nil {
 		return 0, "", err
 	}
 
-	dashboardURL := strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10)
+	dashboardURL = strings.TrimRight(w.Config.Host, "/") + "/jobs/runs/" + strconv.FormatInt(runID, 10)
 	return runID, dashboardURL, nil
 }
 

cmd/air/runsubmit_test.go

@@ -781,6 +781,16 @@ code_source:
 	// zero new import-file calls (a real skip), not a re-upload to the same name.
 	assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath)
 	assert.Equal(t, afterFirst, snapshotUploads, "unchanged plain_tar should skip the second upload")
+	require.NotNil(t, first.SizeBytes)
+	assert.Positive(t, *first.SizeBytes)
+	assert.Equal(t, first.SizeBytes, second.SizeBytes)
+	assert.Equal(t, new(false), second.UsesGit)
+	assert.Equal(t, new(modePlainTar), first.PackagingMode)
+	assert.Equal(t, first.PackagingMode, second.PackagingMode)
+	assert.NotNil(t, first.PackagingDurationMs)
+	assert.NotNil(t, first.UploadDurationMs)
+	assert.Equal(t, new(int64(0)), second.PackagingDurationMs)
+	assert.Equal(t, new(int64(0)), second.UploadDurationMs)
 }
 
 // A git_archive snapshot is content-addressed by (commit, include_paths): submitting
@@ -837,6 +847,16 @@ code_source:
 	// (the second submit is a cache hit and moves no bytes).
 	assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath)
 	assert.Len(t, uploaded, 1, "git_archive cache hit should skip the second upload")
+	require.NotNil(t, first.SizeBytes)
+	assert.Positive(t, *first.SizeBytes)
+	assert.Equal(t, first.SizeBytes, second.SizeBytes)
+	assert.Equal(t, new(true), second.UsesGit)
+	assert.Equal(t, new(modeGitArchive), first.PackagingMode)
+	assert.Equal(t, first.PackagingMode, second.PackagingMode)
+	assert.NotNil(t, first.PackagingDurationMs)
+	assert.NotNil(t, first.UploadDurationMs)
+	assert.Equal(t, new(int64(0)), second.PackagingDurationMs)
+	assert.Equal(t, new(int64(0)), second.UploadDurationMs)
 }
 
 // When enabled, a code source uploads provenance sidecars (git_state.json and
@@ -877,6 +897,8 @@ code_source:
 
 	assert.Empty(t, snap.GitStatePath)
 	assert.Empty(t, snap.GitDiffPath)
+	assert.Equal(t, new(true), snap.UsesGit)
+	assert.Equal(t, new(modePlainTar), snap.PackagingMode)
 
 	_, err = sidecarStore.Read(ctx, gitStateName)
 	assert.ErrorIs(t, err, fs.ErrNotExist)

cmd/air/snapshot.go

@@ -10,13 +10,17 @@ import (
 	"github.com/databricks/cli/libs/env"
 )
 
-// snapshotResult holds the code_source_path wired into the submit payload (the
-// uploaded code archive's remote path) plus the remote paths of the best-effort git
-// provenance sidecars (empty when not a git repo or upload failed).
+// snapshotResult holds artifact paths and snapshot measurements. Measurements
+// survive staging failures; CodeSourcePath is set only after upload or a cache hit.
 type snapshotResult struct {
-	CodeSourcePath string
-	GitStatePath   string
-	GitDiffPath    string
+	CodeSourcePath      string
+	GitStatePath        string
+	GitDiffPath         string
+	SizeBytes           *int64
+	UsesGit             *bool
+	PackagingMode       *snapshotMode
+	PackagingDurationMs *int64
+	UploadDurationMs    *int64
 }
 
 // resolveRootPath resolves a code_source snapshot root_path: expand environment

cmd/air/snapshot_upload.go

@@ -49,7 +49,7 @@ func uploadSnapshot(ctx context.Context, w *databricks.WorkspaceClient, snap *sn
 	}
 	result, err := uploadSnapshotTarball(ctx, w, repoPath, plan, snapshotArtifactPath)
 	if err != nil {
-		return snapshotResult{}, err
+		return result, err
 	}
 
 	// Upload git provenance sidecars (git_state.json / git_diff.patch) next to the
@@ -158,49 +158,69 @@ const airArtifactInternalDir = ".internal"
 // for git_archive and by the working-tree metadata fingerprint for plain_tar — so if
 // the identical object already exists we skip packaging and upload entirely and reuse it.
 func uploadSnapshotTarball(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, artifactPath string) (snapshotResult, error) {
+	result := snapshotResult{UsesGit: &plan.isGitRepo, PackagingMode: &plan.mode}
 	f, uploadPath, err := snapshotUploadFiler(ctx, w, artifactPath)
 	if err != nil {
-		return snapshotResult{}, err
+		return result, err
 	}
 
 	tarName, files, err := snapshotTarName(ctx, repoPath, plan)
 	if err != nil {
-		return snapshotResult{}, err
+		return result, err
 	}
 	// code_source_path is content-addressed by tarName, so it is the same whether we
 	// upload the bytes now or reuse an object already in the store.
 	remote := path.Join(uploadPath, tarName)
 
-	exists, err := snapshotExists(ctx, f, tarName)
-	if err != nil {
-		return snapshotResult{}, err
-	}
-	if exists {
+	info, err := f.Stat(ctx, tarName)
+	if err == nil {
 		log.Debugf(ctx, "snapshot upload skipped; reusing %s", remote)
-		return snapshotResult{CodeSourcePath: remote}, nil
+		size := info.Size()
+		result.CodeSourcePath = remote
+		result.SizeBytes = &size
+		result.PackagingDurationMs = new(int64(0))
+		result.UploadDurationMs = new(int64(0))
+		return result, nil
+	}
+	if !errors.Is(err, fs.ErrNotExist) {
+		return result, fmt.Errorf("failed to check snapshot cache: %w", err)
 	}
 
 	tmp, err := os.MkdirTemp("", "air-snapshot-*")
 	if err != nil {
-		return snapshotResult{}, err
+		return result, err
 	}
 	defer os.RemoveAll(tmp)
 
 	tarball := filepath.Join(tmp, tarName)
-	if err := packageSnapshot(ctx, repoPath, plan, files, tarball); err != nil {
-		return snapshotResult{}, err
+	packagingStart := time.Now()
+	err = packageSnapshot(ctx, repoPath, plan, files, tarball)
+	result.PackagingDurationMs = new(time.Since(packagingStart).Milliseconds())
+	if err != nil {
+		return result, err
 	}
 
 	file, err := os.Open(tarball)
 	if err != nil {
-		return snapshotResult{}, err
+		return result, err
 	}
 	defer file.Close()
+	// Size collection is best-effort and must not prevent a valid upload.
+	if info, err := file.Stat(); err == nil {
+		size := info.Size()
+		result.SizeBytes = &size
+	} else {
+		log.Debugf(ctx, "failed to measure code snapshot: %v", err)
+	}
 	cmdio.LogProgress(ctx, fmt.Sprintf("Uploading %s...", tarName))
-	if err := f.Write(ctx, tarName, file, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil {
-		return snapshotResult{}, fmt.Errorf("failed to upload snapshot %s: %w", tarName, err)
+	uploadStart := time.Now()
+	err = f.Write(ctx, tarName, file, filer.OverwriteIfExists, filer.CreateParentDirectories)
+	result.UploadDurationMs = new(time.Since(uploadStart).Milliseconds())
+	if err != nil {
+		return result, fmt.Errorf("failed to upload snapshot %s: %w", tarName, err)
 	}
-	return snapshotResult{CodeSourcePath: remote}, nil
+	result.CodeSourcePath = remote
+	return result, nil
 }
 
 // snapshotUploadFiler returns a filer rooted at <artifactPath>/.internal plus that
@@ -223,17 +243,3 @@ func snapshotUploadFiler(ctx context.Context, w *databricks.WorkspaceClient, art
 	f, err := filer.NewWorkspaceFilesClient(w, uploadPath)
 	return f, uploadPath, err
 }
-
-// snapshotExists reports whether name already exists in the artifact store, used to
-// short-circuit a content-addressed upload (either mode). A not-found is a clean miss
-// (false, nil); any other error is surfaced.
-func snapshotExists(ctx context.Context, store filer.Filer, name string) (bool, error) {
-	_, err := store.Stat(ctx, name)
-	if err == nil {
-		return true, nil
-	}
-	if errors.Is(err, fs.ErrNotExist) {
-		return false, nil
-	}
-	return false, fmt.Errorf("failed to check snapshot cache: %w", err)
-}

cmd/air/telemetry.go

@@ -0,0 +1,46 @@
+package aircmd
+
+import (
+	"context"
+	"strconv"
+	"time"
+
+	"github.com/databricks/cli/libs/telemetry"
+	"github.com/databricks/cli/libs/telemetry/protos"
+)
+
+// logRunEvent records the submission outcome independently of a later --watch
+// outcome. cfg has already passed local validation before submission starts.
+func logRunEvent(ctx context.Context, cfg *runConfig, snap snapshotResult, runID int64, elapsed time.Duration, err error) {
+	perNode, _ := gpusPerNode(gpuType(cfg.Compute.AcceleratorType))
+	_, hasDependencies := cfg.inlineDependencies()
+	event := &protos.AirRunEvent{
+		GPUType:                       cfg.Compute.AcceleratorType,
+		NumGPUs:                       cfg.Compute.NumAccelerators,
+		NumNodes:                      cfg.Compute.NumAccelerators / perNode,
+		HasDockerImage:                cfg.unityCatalogImagePath() != "" || len(cfg.Containers) > 0, // every container sets its own image
+		HasCodeSnapshot:               cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil,
+		HasRequirements:               hasDependencies,
+		HasParameters:                 len(cfg.Parameters) > 0,
+		MaxRetries:                    cfg.maxRetries(),
+		HasTimeout:                    cfg.TimeoutMinutes != nil,
+		SubmittedSuccessfully:         err == nil,
+		SubmitLatencyMs:               elapsed.Milliseconds(),
+		CodeSourceUsesGit:             snap.UsesGit,
+		CodeSourceSizeBytes:           snap.SizeBytes,
+		CodeSourcePackagingDurationMs: snap.PackagingDurationMs,
+		CodeSourceUploadDurationMs:    snap.UploadDurationMs,
+	}
+	if snap.PackagingMode != nil {
+		switch *snap.PackagingMode {
+		case modeGitArchive:
+			event.CodeSourcePackagingMode = protos.AirPackagingModeGitArchive
+		case modePlainTar:
+			event.CodeSourcePackagingMode = protos.AirPackagingModePlainTar
+		}
+	}
+	if err == nil {
+		event.JobRunID = strconv.FormatInt(runID, 10)
+	}
+	telemetry.Log(ctx, protos.DatabricksCliLog{AirRunEvent: event})
+}