Add azure.ai.training extension

[saanikaguptamicrosoft]

Notes

• Relevant PRs (refer PR overview for details)
  • Add artifact resolution for code and inputs in job create #7153
  • Implement cancel job command for custom training #7272
  • Implement delete job command for custom training #7273
  • Add validate command for custom training #7407 , Share offline validation between job validate and job submit commands #8068
  • Implement stream command for training job #7939
  • Implement connect-ssh to job node #7964
  • Implement download command for training job #8029
  • Add support for gpuCount for partial SKU scenario #8067
  • Bugfix: Honor --subscription and --project-endpoint flags over stored env values #8093
  • Implement show services command #8121
  • Add validations for UAMI requirement #8122
  • Add template flag in init command #8123

Testing

• Help command

• Init command

• Job submit

• Job show

• Jobs list

Other operations like validate, download, stream, delete, cancel, show-services, connect-ssh are also supported.

[wbreza]

Re-review against the 16 new commits (3106077 → f1ab8f3)

Thanks for the rapid turnaround, @saanikaguptamicrosoft — a lot of prior feedback is cleanly resolved. Summary below covers (a) verification of prior findings and (b) new findings introduced by the delta.

✅ Prior findings — resolved

• azcopy silent parse errors — fixed in runner.go; first-occurrence warnings via warnedUnparseable / warnedProgressFormat 👍
• Zero-test critical paths (hash, upload_service) — hash_test.go and upload_service_test.go added with strong coverage (determinism, order-independence, zombie recovery, hash collision, missing SAS, sentinel tagging) 👍
• Empty PR description — now populated with related PRs and screenshots 👍
• job get vs job show naming — job_get.go renamed to job_show.go 👍
• proxyEndpointPattern allowing % — % → %% escape applied for OpenSSH ProxyCommand expansion 👍
• Design/ placement — renamed to lowercase design/, README references updated 👍

⚠️ Prior findings — partial

• ServiceEndpoint consolidation — new internal/utils/services.go ServiceEndpoint() helper exists and stream_service.go uses it, but internal/cmd/job_show_services.go:47 still calls its local extractServiceEndpointStr. One call site missed.
• Untested paths — coverage added for hash/upload_service, but pkg/client/download.go still has no tests.

❌ Prior findings — not addressed

1. --debug flag still prints request URLs and bodies unredacted
pkg/client/client.go:91–103. SAS URIs in query strings and credential fields in GetDatasetCredentials / GetModelCredentials request bodies remain printable.
Fix: strip query string from the URL print; redact sasUri, Authorization, sig, signature in the body print; or gate behind a developer-only env var and document the leak risk.

2. SSH proxy stdin goroutine still untracked
internal/cmd/job_ssh_proxy.go:167–182. No wg.Add, no os.Stdin close on cancel; relies on process exit. Acceptable for the short-lived ProxyCommand model but wg.Wait() at the end of runProxy cannot guarantee clean shutdown.
Fix: at minimum, close os.Stdin from the cancellation watcher so the upstream Read() returns.

🟥 New findings from the delta

Critical — DeleteJob long-running-operation polling is single-shot
pkg/client/jobs.go (pollDeleteOperationOnce / DeleteJob). Issues a single GET against the LRO Location URL; on 202 Accepted returns DeleteJobInProgress rather than continuing to poll. This violates Azure LRO semantics, gives users the "delete accepted but unverified" UX, and creates a race window for concurrent delete attempts.
Fix: loop with bounded retries + backoff respecting Retry-After, or adopt runtime.NewPoller from azcore.

High — Artifact contentinfo fan-out drops per-root errors
internal/cmd/job_download.go ~lines 688–720. Per-root goroutines record rootResults[idx].err, but after wg.Wait() the results are consumed without surfacing the per-root error; a permission-denied or 404 on a single root prefix is silently dropped → partial downloads reported as success.
Fix: after wg.Wait(), aggregate errors with errors.Join and either fail the command or print a loud per-root warning.

High — Debug URL logging extended to APIM data-plane path
pkg/client/client.go getAbsoluteDataPlane() adds a new fmt.Fprintf(os.Stderr, "[DEBUG] GET %s\n", parsed.String()) site. Same class of leak as the unresolved prior finding above but extended.
Fix: strip query string before printing.

Medium — redactSAS may pass URLs through unredacted
internal/azcopy/runner.go ~lines 388–397. If ? is absent the input URL is returned unchanged. Defense-in-depth: also check for known SAS markers (sv=, sig=, se=) and redact any token-like substring.

Medium — mapServiceType() silently passes unknown types
internal/cmd/job_submit.go ~lines 1188–1203. Returns input unchanged on unknown service type, relying on ValidateJobOffline() always running. If validation is bypassed or drifts, server returns a generic 400 instead of a clear CLI error.
Fix: return an error on unknown service type — keep offline validation as the primary check, this as a backstop.

Medium — SSH ProxyCommand shell quoting
internal/cmd/job_connect_ssh.go ~lines 434–440. %-escape handles OpenSSH expansion, but the rendered ProxyCommand isn't single-quoted, so shell metacharacters in job names / endpoints ($, backticks, !) could be interpreted by the parent shell.
Fix: wrap the constructed proxy command in single quotes when emitting it.

Medium — Distribution Ray fields lack client-side validation
pkg/models/common.go Distribution + buildDistribution() ~lines 1208–1227. *int Port/DashboardPort and *bool IncludeDashboard pointer copies are correct, but out-of-range/negative ports aren't validated → cryptic server error.

Medium — Root-folder fan-out semaphore leak on context cancel
internal/cmd/job_download.go. sem <- struct{}{} isn't context-aware; cancellation while blocked on acquire leaks the goroutine.
Fix: select { case sem <- struct{}{}: ... case <-ctx.Done(): return ctx.Err() }.

Medium — azcopy end-of-job diagnostics truncation
internal/azcopy/runner.go printEndOfJobDiagnostics() caps to first 5 failed transfers + 16 KB stderr; may hide root cause on batch failures. Consider writing full diagnostics to a side-file when truncation kicks in.

✅ Cleanly verified

• "Temp: Print API request response" (6d7f2d7) was fully reverted by 61c438f — no leftover prints.
• go fix modernization commit (16cff27) is mechanical only — no behavioral changes.
• Metrics path fix (305bdcc) consistent with APIM schema.
• New tests (hash_test.go, upload_service_test.go) exercise the failure paths called out previously — exactly the right targets.

Posted by Copilot CLI on behalf of the reviewer.

[JeffreyCA]

If the above comments can be addressed in a follow-up PR before the first release I'm good with that too.

Would like to at least get release-ext-azure-ai-training.yml merged soon so @vhvb1989 can help with creating the pipelines.

[saanikaguptamicrosoft]

@wbreza

⚠️ Prior findings — partial

ServiceEndpoint consolidation — new internal/utils/services.go ServiceEndpoint() helper exists and stream_service.go uses it, but internal/cmd/job_show_services.go:47 still calls its local extractServiceEndpointStr. One call site missed.

This comment is no longer applicable. We don't have any function extractServiceEndpointStr post we moved to APIM APIs.

Untested paths — coverage added for hash/upload_service, but pkg/client/download.go still has no tests.

Acknowledged — tracking as follow-up. These are thin HTTP wrappers around the shared doDataPlane transport and are exercised end-to-end by model download, dataset download, and job download. Will add focused httptest coverage in a follow-up so this PR stays scoped.

❌ Prior findings — not addressed

--debug flag still prints request URLs and bodies unredacted

Done

SSH proxy stdin goroutine still untracked

Intentional and consistent with existing AML CLI extension (_ssh_connector.py uses the same pattern: a thread-pool os.read(0, ...) that's not awaited and is reaped by process exit). Closing os.Stdin from a watcher would not actually unblock the in-flight read(2) syscall on either Linux pipes or Windows console handles, so the suggested "minimum fix" doesn't deliver the determinism it implies — SSH terminating the ProxyCommand process is the cleanup mechanism by design.

🟥 New findings from the delta

Critical — DeleteJob long-running-operation polling is single-shot

Addressed. Also added a --no-wait flag for single peek.

High — Artifact contentinfo fan-out drops per-root errors

Done

High — Debug URL logging extended to APIM data-plane path

Done

Medium — redactSAS may pass URLs through unredacted

Done

Medium — mapServiceType() silently passes unknown types

Tested empirically — backend silently drops unknown service types (no 400), and the case is already caught upstream by ValidateJobOffline() (job_validator.go:154-180) which rejects them with a clear field-scoped error before submit. Keeping the mapper as a pure translator and the validator as the single source of truth for the allow-list (duplicating it here would invert the drift risk, not reduce it).

Medium — SSH ProxyCommand shell quoting

The only non-constant inputs to the ProxyCommand string are (a) selfPath from os.Executable() and (b) the proxy endpoint URL. The endpoint is already character-class-validated by job_connect_ssh.go to exclude all shell metacharacters ($, backticks, ;, |, &, <, >, (, ), ', ", spaces, etc.), and the job name doesn't appear in the ProxyCommand string at all (destination is the fixed alias azureml-job). The remaining input selfPath would require an attacker-controlled install path of the extension binary, which already implies arbitrary code execution.

Also, single-quoting the whole ProxyCommand value would break it — when ssh spawns it via /bin/sh -c, the shell would try to execute the single-quoted string as a single binary name. The endpoint validation is the right layer for this defense and is already in place.

Medium — Distribution Ray fields lack client-side validation

Addressed

Medium — Root-folder fan-out semaphore leak on context cancel

Addressed

Medium — azcopy end-of-job diagnostics truncation

Done — bumped the stderr buffer cap to 64 KB and added maybeWriteFullDiagnostics that dumps the full EndOfJob payload (as JSON) plus captured stderr to a temp file (%TEMP%/azd-azcopy-diag-*.log) whenever either the failed-transfers list or the stderr buffer was truncated for terminal display; the path is printed to stderr so users can inspect everything for batch failures.

[vhvb1989]

/check-enforcer evaluate