diff --git a/experimental/air/cmd/register_image.go b/experimental/air/cmd/register_image.go index d4195d0d8e2..c3537102a00 100644 --- a/experimental/air/cmd/register_image.go +++ b/experimental/air/cmd/register_image.go @@ -111,18 +111,15 @@ configuration (run ` + "`docker login`" + ` first); there are no credential flag timeout := time.Duration(timeoutMinutes) * time.Minute - // Discover credentials from the local Docker config and store them in a - // per-user secret for the registration call. If storage fails, registration - // still proceeds without credentials — a public image succeeds — and - // credErr is reported as the cause if the registry rejects anonymous access. - // credErr is not fatal on its own: it is reported by registrationError only - // if the registry then rejects anonymous access. - scope, key, credErr := discoverCredentials(ctx, w, c, dockerImageURL) + // credErr is not fatal on its own: registration proceeds without + // credentials (a public image still succeeds) and registrationError reports + // it only if the registry rejects anonymous access. + creds, credErr := discoverCredentials(ctx, w, c, dockerImageURL) if credErr != nil { log.Debugf(ctx, "could not store local Docker credentials: %v", credErr) } - updated, sha, err := registerWithCredentialFallback(ctx, c, dockerImageURL, scope, key, timeout) + updated, sha, err := registerWithCredentialFallback(ctx, c, dockerImageURL, creds, timeout) if err != nil { kind, retryable := classifyRegistrationError(err) return renderError(ctx, cmd, "REGISTRATION_FAILED", kind, retryable, @@ -142,16 +139,25 @@ configuration (run ` + "`docker login`" + ` first); there are no credential flag return cmd } +// imageCredentials is a secret reference passed to registration. discovered +// marks it as auto-discovered from the local Docker config rather than +// configured by the user, which decides whether a rejection may be retried +// anonymously. +type imageCredentials struct { + scope string + key string + discovered bool +} + // discoverCredentials resolves registry credentials from the local Docker config -// and stores them in a per-user secret, returning the (scope, key) reference for -// registration. It first probes whether the image is public: if so, no -// credentials are stored (avoiding a throwaway secret). Returns empty scope/key -// when the image is public or no local credentials exist, both with a nil error. -// A non-nil error means credentials were found but could not be stored (e.g. the -// user lacks permission to create a secret scope); it is advisory, so the caller -// can still attempt an anonymous registration and report this as the cause if -// that fails. -func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, dockerImageURL string) (scope, key string, err error) { +// and stores them in a per-user secret, returning the reference for registration. +// It first probes whether the image is public: if so, no credentials are stored +// (avoiding a throwaway secret). Returns empty credentials when the image is +// public or no local credentials exist, both with a nil error. A non-nil error +// means credentials were found but could not be stored (e.g. the user lacks +// permission to create a secret scope); it is advisory, so the caller can still +// attempt an anonymous registration and report this as the cause if that fails. +func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, dockerImageURL string) (creds imageCredentials, err error) { // readDockerCredentials keys off the registry host, so it needs the normalized // URL (e.g. bare "ubuntu" resolves to the Docker Hub host). normalized := normalizeDockerImageURL(dockerImageURL) @@ -161,20 +167,20 @@ func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c * // twice. username, password, ok := readDockerCredentials(ctx, normalized) if !ok { - return "", "", nil + return imageCredentials{}, nil } if public := c.checkImageAccess(ctx, dockerImageURL); public != nil && *public { log.Infof(ctx, "image is publicly accessible; skipping local Docker credentials") - return "", "", nil + return imageCredentials{}, nil } - scope, key, err = storeDockerCredentials(ctx, w, normalized, username, password) + scope, key, err := storeDockerCredentials(ctx, w, normalized, username, password) if err != nil { - return "", "", err + return imageCredentials{}, err } log.Infof(ctx, "using Docker credentials from local config (stored as %s/%s)", scope, key) - return scope, key, nil + return imageCredentials{scope: scope, key: key, discovered: true}, nil } // isAuthError reports whether err is an authentication or permission failure, @@ -218,14 +224,15 @@ func registrationError(dockerImageURL string, err, credErr error) error { return fmt.Errorf("image %q was not found or requires credentials: run `docker login` for its registry, then retry: %w", dockerImageURL, err) } -// registerWithCredentialFallback registers the image and, if the stored +// registerWithCredentialFallback registers the image and, if auto-discovered // credentials are rejected as an auth failure, retries once anonymously so a -// public image isn't blocked by stale local creds (e.g. a revoked PAT from an -// old `docker login`). The retry only fires when credentials were supplied. -func registerWithCredentialFallback(ctx context.Context, c *imageClient, dockerImageURL, scope, key string, timeout time.Duration) (updated bool, sha string, err error) { - updated, sha, err = resolveImage(ctx, c, dockerImageURL, scope, key, timeout) - if err != nil && scope != "" && isAuthError(err) { - log.Warnf(ctx, "stored Docker credentials were rejected (%v); retrying without credentials in case the image is public", err) +// public image isn't blocked by stale local creds (e.g. a revoked PAT from an old +// `docker login`). Credentials the user configured explicitly are never retried +// away: they asked for those specifically, so the rejection is the real answer. +func registerWithCredentialFallback(ctx context.Context, c *imageClient, dockerImageURL string, creds imageCredentials, timeout time.Duration) (updated bool, sha string, err error) { + updated, sha, err = resolveImage(ctx, c, dockerImageURL, creds.scope, creds.key, timeout) + if err != nil && creds.discovered && isAuthError(err) { + log.Warnf(ctx, "Docker credentials discovered from your local config were rejected (%v); retrying without credentials in case the image is public", err) return resolveImage(ctx, c, dockerImageURL, "", "", timeout) } return updated, sha, err diff --git a/experimental/air/cmd/register_image_test.go b/experimental/air/cmd/register_image_test.go index 8ad5383940b..5b7adb8779a 100644 --- a/experimental/air/cmd/register_image_test.go +++ b/experimental/air/cmd/register_image_test.go @@ -167,7 +167,7 @@ func credRejectingImageServer(t *testing.T, credentialedPOSTs *int) string { func TestRegisterWithCredentialFallbackRetriesAnonymously(t *testing.T) { var credentialedPOSTs int url := credRejectingImageServer(t, &credentialedPOSTs) - updated, sha, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "scope", "key", time.Second) + updated, sha, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{scope: "scope", key: "key", discovered: true}, time.Second) require.NoError(t, err) assert.True(t, updated) assert.Equal(t, "pubsha", sha) @@ -179,7 +179,17 @@ func TestRegisterWithCredentialFallbackNoRetryWithoutCreds(t *testing.T) { // failure surfaces directly. var credentialedPOSTs int url := credRejectingImageServer(t, &credentialedPOSTs) - _, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "", "", time.Second) + _, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{}, time.Second) require.NoError(t, err) // anonymous POST succeeds on this server assert.Equal(t, 0, credentialedPOSTs) } + +func TestRegisterWithCredentialFallbackNoRetryForExplicitCreds(t *testing.T) { + // The user named these credentials, so a rejection is the real answer: do not + // silently retry without them. + var credentialedPOSTs int + url := credRejectingImageServer(t, &credentialedPOSTs) + _, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{scope: "scope", key: "key"}, time.Second) + require.Error(t, err) + assert.Equal(t, 1, credentialedPOSTs, "should try once with creds and stop") +} diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go index 4cfbf3736d1..96dcf96d5a1 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -305,15 +305,54 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error { // dockerImageConfig is environment.docker_image. type dockerImageConfig struct { URL string `yaml:"url"` + // "latest" re-checks the registry for the tag's newest digest each run; + // "auto" (default) reuses the existing registration. + TagPolicy string `yaml:"tag_policy"` + // Credentials for a private image that "latest" re-resolves; discovered from + // the local Docker config when unset. + CredentialsScope string `yaml:"credentials_scope"` + CredentialsKey string `yaml:"credentials_key"` } +const ( + dockerTagPolicyAuto = "auto" + dockerTagPolicyLatest = "latest" +) + func (d *dockerImageConfig) validate() error { - if strings.TrimSpace(d.URL) == "" { + // Store the trimmed values: URL rides the submitted task, and the credential + // pairing check below must not treat blank-but-present as set. + d.URL = strings.TrimSpace(d.URL) + d.CredentialsScope = strings.TrimSpace(d.CredentialsScope) + d.CredentialsKey = strings.TrimSpace(d.CredentialsKey) + + if d.URL == "" { return errors.New("docker_image.url cannot be empty") } + + switch strings.ToLower(strings.TrimSpace(d.TagPolicy)) { + case "", dockerTagPolicyAuto, dockerTagPolicyLatest: + default: + return fmt.Errorf("invalid docker_image.tag_policy %q: must be %q or %q", d.TagPolicy, dockerTagPolicyAuto, dockerTagPolicyLatest) + } + + if (d.CredentialsScope != "") != (d.CredentialsKey != "") { + return errors.New("docker_image.credentials_scope and docker_image.credentials_key must be provided together") + } + + // Credentials are only consulted when re-resolving the tag, so accepting them + // under the default policy would silently ignore them. + if d.CredentialsScope != "" && !d.wantsLatest() { + return fmt.Errorf("docker_image.credentials_scope/credentials_key only apply with tag_policy %q; the image is otherwise used as already registered", dockerTagPolicyLatest) + } return nil } +// wantsLatest reports whether the image should be re-resolved before the run. +func (d *dockerImageConfig) wantsLatest() bool { + return strings.EqualFold(strings.TrimSpace(d.TagPolicy), dockerTagPolicyLatest) +} + // codeSourceConfig is the `code_source` block. Only the "snapshot" type exists. type codeSourceConfig struct { Type string `yaml:"type"` diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go index 1408b600736..60b9352bbba 100644 --- a/experimental/air/cmd/runconfig_launch.go +++ b/experimental/air/cmd/runconfig_launch.go @@ -26,16 +26,21 @@ func (c *runConfig) maxRetries() int { } // dockerImageURL returns the custom docker image URL, or "" when none is set. -// -// TODO: not wired into submission yet — the native ai_runtime_task carries no -// docker field, and full support needs image registration (pending the DCS work). func (c *runConfig) dockerImageURL() string { - if c.Environment != nil && c.Environment.DockerImage != nil { - return c.Environment.DockerImage.URL + if img := c.dockerImage(); img != nil { + return img.URL } return "" } +// dockerImage returns the environment.docker_image block, or nil when none is set. +func (c *runConfig) dockerImage() *dockerImageConfig { + if c.Environment == nil { + return nil + } + return c.Environment.DockerImage +} + // requirementsFile returns the path to a requirements file when // environment.dependencies is a string, and whether it was set. func (c *runConfig) requirementsFile() (string, bool) { diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index 45cae7cd0ae..cd7804cd4af 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -287,6 +287,46 @@ func TestEnvironmentConfigValidate(t *testing.T) { environmentConfig{DockerImage: &dockerImageConfig{URL: " "}}, "docker_image.url cannot be empty", }, + { + "tag policy latest ok", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "latest"}}, + "", + }, + { + "tag policy auto ok", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "AUTO"}}, + "", + }, + { + "invalid tag policy", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "newest"}}, + `invalid docker_image.tag_policy "newest"`, + }, + { + "credentials scope without key", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "s"}}, + "must be provided together", + }, + { + "credentials key without scope", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsKey: "k"}}, + "must be provided together", + }, + { + "credentials pair ok with latest", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "latest", CredentialsScope: "s", CredentialsKey: "k"}}, + "", + }, + { + "credentials without latest rejected", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "s", CredentialsKey: "k"}}, + `only apply with tag_policy "latest"`, + }, + { + "blank credentials scope is not treated as set", + environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: " "}}, + "", + }, { "version with file deps", environmentConfig{ diff --git a/experimental/air/cmd/rundockerimage.go b/experimental/air/cmd/rundockerimage.go new file mode 100644 index 00000000000..9645d367770 --- /dev/null +++ b/experimental/air/cmd/rundockerimage.go @@ -0,0 +1,113 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "time" + + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" +) + +// imageReadyTimeout is generous: replicating a large image takes minutes. +const imageReadyTimeout = time.Hour + +// digestDisplay abbreviates a manifest digest for logs. +func digestDisplay(sha string) string { + if sha == "" { + return "digest unknown" + } + return shortManifestSHA(sha) +} + +func notRegisteredError(dockerImageURL string) error { + return fmt.Errorf("docker image not registered: %s\nregister it first: databricks experimental air register-image %s", dockerImageURL, dockerImageURL) +} + +// prepareDockerImage verifies a run's custom image before submit, so a bad image +// fails here instead of deep in the launch stage where the cause is obscured. +func prepareDockerImage(ctx context.Context, w *databricks.WorkspaceClient, img *dockerImageConfig) error { + c, err := newImageClient(w) + if err != nil { + return err + } + + if img.wantsLatest() { + if err := resolveLatestDockerImage(ctx, w, c, img); err != nil { + return err + } + } + + return waitForRegisteredImage(ctx, c, img.URL) +} + +// resolveLatestDockerImage re-registers the image so the run picks up the tag's +// newest digest, using the config's credentials when set and otherwise the local +// Docker config. +func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, img *dockerImageConfig) error { + // Credentials the user configured are used as-is (discovered=false), so a + // rejection is not retried anonymously and reports the secret they named. + creds := imageCredentials{scope: img.CredentialsScope, key: img.CredentialsKey} + var credErr error + if creds.scope == "" { + creds, credErr = discoverCredentials(ctx, w, c, img.URL) + if credErr != nil { + log.Debugf(ctx, "could not store local Docker credentials: %v", credErr) + } + } + + // Visible at the default log level (WARN): this can block for minutes. + cmdio.LogString(ctx, fmt.Sprintf("Re-resolving %s against the source registry (tag_policy=latest)...", img.URL)) + if _, _, err := registerWithCredentialFallback(ctx, c, img.URL, creds, imageReadyTimeout); err != nil { + if !creds.discovered && creds.scope != "" && isAuthError(err) { + return fmt.Errorf("the credentials in secret %s/%s were rejected for image %q: %w", creds.scope, creds.key, img.URL, err) + } + return registrationError(img.URL, err, credErr) + } + return nil +} + +// waitForRegisteredImage requires an existing registration and blocks while it is +// still importing. Missing or FAILED is an error the user fixes by re-registering. +// +// An AVAILABLE registration whose stored credentials have since lost registry +// access is accepted here and only fails at pod launch. The Python CLI catches +// that with a :validateImageAccess probe (cli/sdk/_submit.py), which is not ported +// deliberately: this check belongs in the backend, so every client gets it and +// the CLI doesn't pay a round trip per submit. Port it there rather than here. +func waitForRegisteredImage(ctx context.Context, c *imageClient, dockerImageURL string) error { + reg, err := c.getImage(ctx, dockerImageURL) + if err != nil { + return err + } + if reg == nil { + return notRegisteredError(dockerImageURL) + } + + switch reg.Status { + case imageStatusAvailable: + log.Infof(ctx, "using image %s (%s)", dockerImageURL, digestDisplay(reg.ManifestSHA256)) + return nil + case imageStatusFailed: + msg := reg.StatusMessage + if msg == "" { + msg = "unknown error" + } + return fmt.Errorf("docker image registration failed: %s\nfix the issue and re-register: databricks experimental air register-image %s", msg, dockerImageURL) + case imageStatusPending, imageStatusImporting: + // This wait runs up to imageReadyTimeout, so it must be visible at the + // default log level (WARN); log.Infof would be silent. + cmdio.LogString(ctx, fmt.Sprintf("Docker image registration in progress (%s); waiting for it to become available...", reg.Status)) + final, err := c.waitForImageReady(ctx, dockerImageURL, imageReadyTimeout, imagePollInterval) + if err != nil { + return err + } + cmdio.LogString(ctx, "Image ready ("+digestDisplay(final.ManifestSHA256)+")") + return nil + } + + // Unreachable: normalizeStatus maps every unknown state to PENDING. + return errors.New("unexpected image status " + string(reg.Status)) +} diff --git a/experimental/air/cmd/rundockerimage_test.go b/experimental/air/cmd/rundockerimage_test.go new file mode 100644 index 00000000000..362e0030a3c --- /dev/null +++ b/experimental/air/cmd/rundockerimage_test.go @@ -0,0 +1,195 @@ +package aircmd + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/databricks/cli/libs/cmdio" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestDigestDisplay(t *testing.T) { + assert.Equal(t, "digest unknown", digestDisplay("")) + assert.Equal(t, "abc", digestDisplay("abc")) + assert.Equal(t, "0123456789abcdef...", digestDisplay("0123456789abcdefghij")) +} + +func TestWaitForRegisteredImageAvailable(t *testing.T) { + url := imageServer(t, `{}`, `{"state":"AVAILABLE","manifest_sha256":"abc"}`) + require.NoError(t, waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0")) +} + +func TestWaitForRegisteredImageNotRegistered(t *testing.T) { + // :get 404s, so the run must stop with registration guidance. + url := imageServer(t, `{}`, "") + err := waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0") + require.Error(t, err) + assert.Contains(t, err.Error(), "docker image not registered") + assert.Contains(t, err.Error(), "air register-image nvcr.io/org/img:1.0") +} + +func TestWaitForRegisteredImageFailed(t *testing.T) { + url := imageServer(t, `{}`, `{"state":"FAILED","status_message":"manifest not found"}`) + err := waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0") + require.Error(t, err) + assert.Contains(t, err.Error(), "registration failed: manifest not found") + assert.Contains(t, err.Error(), "re-register") +} + +func TestWaitForRegisteredImageWaitsWhileImporting(t *testing.T) { + // First poll is still importing; the next reports AVAILABLE. + url := imageServer(t, `{}`, + `{"state":"IMPORTING"}`, + `{"state":"AVAILABLE","manifest_sha256":"abc"}`) + require.NoError(t, waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0")) +} + +// latestImageServer serves a registry where the image is registered and +// AVAILABLE, counting POSTs so a test can assert whether a re-registration +// happened. +func latestImageServer(t *testing.T, posts *int) string { + t.Helper() + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + _, _ = w.Write([]byte(`{"state":"AVAILABLE","manifest_sha256":"abc"}`)) + case imagesAPIPath: + *posts++ + _, _ = w.Write([]byte(`{"image":{"state":"AVAILABLE","manifest_sha256":"abc"}}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + return srv.URL +} + +func TestPrepareDockerImageAutoDoesNotReregister(t *testing.T) { + var posts int + w := newTestWorkspaceClient(t, latestImageServer(t, &posts)) + img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0"} + require.NoError(t, prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img)) + assert.Zero(t, posts, "default tag policy must not re-register") +} + +// TestPrepareDockerImageLatestDiscoversCredentials covers the config-without- +// credentials path: creds come from the local Docker config and ride the POST. +func TestPrepareDockerImageLatestDiscoversCredentials(t *testing.T) { + var postBodies []string + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + _, _ = w.Write([]byte(`{"state":"AVAILABLE","manifest_sha256":"abc"}`)) + case imagesAPIPath + ":checkImageAccess": + _, _ = w.Write([]byte(`{"publicly_accessible": false}`)) + case imagesAPIPath: + body, _ := io.ReadAll(r.Body) + postBodies = append(postBodies, string(body)) + _, _ = w.Write([]byte(`{"image":{"state":"AVAILABLE","manifest_sha256":"abc"}}`)) + case "/api/2.0/secrets/scopes/list": + _, _ = w.Write([]byte(`{"scopes":[]}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + ctx := cmdio.MockDiscard(writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`)) + w := newTestWorkspaceClient(t, srv.URL) + img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0", TagPolicy: "latest"} + + require.NoError(t, prepareDockerImage(ctx, w, img)) + require.Len(t, postBodies, 1) + assert.Contains(t, postBodies[0], `"credentials_key":"nvcr.io-bob-local"`) + assert.Contains(t, postBodies[0], `"credentials_scope":"docker-credentials-`) +} + +// TestPrepareDockerImageLatestStorageDeniedReportsCause covers the journey where +// creds exist locally but can't be stored: the error must name that cause rather +// than tell the user to `docker login`. +func TestPrepareDockerImageLatestStorageDeniedReportsCause(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`)) + case imagesAPIPath + ":checkImageAccess": + _, _ = w.Write([]byte(`{"publicly_accessible": false}`)) + case "/api/2.0/secrets/scopes/create": + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"denied"}`)) + case imagesAPIPath: + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"cannot pull: unauthorized"}`)) + case "/api/2.0/secrets/scopes/list": + _, _ = w.Write([]byte(`{"scopes":[]}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + ctx := cmdio.MockDiscard(writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`)) + w := newTestWorkspaceClient(t, srv.URL) + img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0", TagPolicy: "latest"} + + err := prepareDockerImage(ctx, w, img) + require.Error(t, err) + assert.Contains(t, err.Error(), "could not be stored") + assert.NotContains(t, err.Error(), "run `docker login`") +} + +// TestPrepareDockerImageLatestExplicitCredsRejected asserts a rejected secret the +// user named reports that secret, and is not retried anonymously or blamed on a +// missing `docker login`. +func TestPrepareDockerImageLatestExplicitCredsRejected(t *testing.T) { + var credentialedPOSTs int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + switch r.URL.Path { + case imagesAPIPath + ":get": + w.WriteHeader(http.StatusNotFound) + _, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`)) + case imagesAPIPath: + body, _ := io.ReadAll(r.Body) + if strings.Contains(string(body), "credentials_scope") { + credentialedPOSTs++ + } + w.WriteHeader(http.StatusForbidden) + _, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"denied"}`)) + default: + _, _ = w.Write([]byte(`{}`)) + } + })) + t.Cleanup(srv.Close) + + w := newTestWorkspaceClient(t, srv.URL) + img := &dockerImageConfig{ + URL: "nvcr.io/org/img:1.0", + TagPolicy: "latest", + CredentialsScope: "myscope", + CredentialsKey: "mykey", + } + + err := prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img) + require.Error(t, err) + assert.Contains(t, err.Error(), "credentials in secret myscope/mykey were rejected") + assert.NotContains(t, err.Error(), "docker login") + assert.Equal(t, 1, credentialedPOSTs, "explicit credentials must not be retried anonymously") +} + +func TestPrepareDockerImageLatestReregisters(t *testing.T) { + var posts int + w := newTestWorkspaceClient(t, latestImageServer(t, &posts)) + img := &dockerImageConfig{ + URL: "nvcr.io/org/img:1.0", + TagPolicy: "latest", + CredentialsScope: "scope", + CredentialsKey: "key", + } + require.NoError(t, prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img)) + assert.Equal(t, 1, posts, "tag_policy=latest must re-register against the source registry") +} diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 3aa2436b827..54189fba956 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -80,6 +80,11 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri }, }}, CodeSourcePath: snap.CodeSourcePath, + // NOTE: docker_image_url is intentionally not set here yet. The field was + // added to jobs.AiRuntimeTask in databricks-sdk-go after v0.170.0, which the + // CLI has not bumped to. prepareDockerImage already verifies the image is + // registered; passing it on the task lands in the follow-up PR once the SDK + // bump + codegen is in. } if cfg.MLflowRunName != nil { task.MlflowRun = *cfg.MLflowRunName @@ -186,6 +191,15 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if err != nil { return 0, "", err } + + // After the cheap workspace checks (a tag_policy=latest refresh can block for + // minutes) but before any upload, so a bad image wastes no artifact work. + if img := cfg.dockerImage(); img != nil { + if err := prepareDockerImage(ctx, w, img); err != nil { + return 0, "", err + } + } + runName := "" if cfg.MLflowRunName != nil { runName = *cfg.MLflowRunName