Skip to content

training: Align training job yaml schema with new command job schema - #9349

Open
saanikaguptamicrosoft wants to merge 2 commits into
Azure:mainfrom
saanikaguptamicrosoft:saanika/capacityUnitCount
Open

training: Align training job yaml schema with new command job schema#9349
saanikaguptamicrosoft wants to merge 2 commits into
Azure:mainfrom
saanikaguptamicrosoft:saanika/capacityUnitCount

Conversation

@saanikaguptamicrosoft

@saanikaguptamicrosoft saanikaguptamicrosoft commented Jul 29, 2026

Copy link
Copy Markdown
Member

Notes

Testing

  • Bare CPU — no priority, no gpu_count
image
  • CPU with top-level priority: high
image
  • GPU with gpu_count: 2, no resources block
image

Copilot AI balanced review requested due to automatic review settings July 29, 2026 09:28
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
Successfully started running 1 pipeline(s).
21 pipeline(s) were filtered out due to trigger conditions.
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Aligns training-job YAML, wire models, display logic, and scaffolding with Foundry’s capacity-unit submission schema.

Changes:

  • Replaces GPU sizing with top-level capacity units and priority.
  • Restricts resource configuration to instance count.
  • Updates job display fallbacks and generated YAML.

Reviewed changes

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
pkg/models/job.go Adds new command-job wire fields.
pkg/models/history.go Supports new and legacy compute metadata.
pkg/models/common.go Narrows resource configuration.
internal/utils/yaml_parser.go Updates the accepted YAML shape.
internal/cmd/job_submit.go Maps new fields into submissions.
internal/cmd/job_show.go Updates compute-detail rendering.
internal/cmd/init_template.go Scaffolds the new YAML format.
Comments suppressed due to low confidence (1)

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:483

  • [azd-code-reviewer] These top-level Get Job fields are still rendered only inside the History.Compute != nil branch. Run history is supplementary and may be unavailable, so a new job can show its compute name while omitting both Capacity Units and Priority. Render top-level values independently, then consult run history only for fallback values.
		// Prefer the user-submitted capacityUnitCount (from Get Job) when present — for
		// partial-SKU jobs it reflects the actual allocation. Fall back to the SKU GPU
		// count reported by run history for older jobs that pre-date capacityUnitCount.
		if d.Job.Properties.CapacityUnitCount > 0 {
			fmt.Fprintf(w, "Capacity Units:\t%d\n", d.Job.Properties.CapacityUnitCount)

Comment thread cli/azd/extensions/azure.ai.training/internal/cmd/init_template.go
Comment thread cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go Outdated
Comment thread cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go Outdated
@github-actions github-actions Bot added the ext-training azure.ai.training extension label Jul 29, 2026
Copilot AI review requested due to automatic review settings July 29, 2026 10:01

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go:85

  • [azd-code-reviewer] yaml.Unmarshal ignores unknown keys, so removing these fields makes existing files containing gpu_count, resources.instance_type, shm_size, docker_args, or properties validate successfully while silently discarding their allocation settings. Reject removed keys with migration guidance (for example, gpu_countcapacity_unit_count) in both submit and offline-validation parsing paths, or enable strict known-field decoding consistently.
// Only instance_count is honored; instance_type, slaTier, priority and the AISuperComputer
// properties block have been removed — the service now infers the SKU from the compute
// cluster, priority is a top-level job field, and users specify partial capacity via the
// top-level capacity_unit_count field.
type ResourceDefinition struct {

cli/azd/extensions/azure.ai.training/pkg/models/job.go:27

  • [azd-code-reviewer] Removing GPUCount here drops gpuCount while GetJob decodes older responses. JSON mode then re-marshals this typed model without fetching run history (job_show.go:83-92), so its documented backward-compatible “raw job response” silently loses the legacy field. Keep GPUCount on the read model; its zero value remains omitted from new submit payloads.
	CapacityUnitCount         int                  `json:"capacityUnitCount,omitempty"`
	Priority                  string               `json:"priority,omitempty"`

cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go:33

  • [azd-code-reviewer] The generated template documents constrained values for both new fields, but offline validation currently accepts any integer/string and also allows capacity_unit_count together with instance_count, despite describing them as alternatives. Invalid values therefore pass job validate and fail only at submission. Add validator rules and table-driven tests for the capacity-unit set, the priority enum, and mutually exclusive sizing options.
	CapacityUnitCount    int                          `yaml:"capacity_unit_count"`
	Priority             string                       `yaml:"priority"`

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:488

  • [azd-code-reviewer] These top-level Get Job fields are still inside the History.Compute != nil branch. If supplementary run-history retrieval fails or returns no compute object, CapacityUnitCount and Priority are available on d.Job but disappear from table output. Read history into an optional fallback, then render the top-level values independently.
		// Capacity Units: prefer the top-level CapacityUnitCount from Get Job,
		// then whatever run history reports, then the SKU-level GPU count.
		if d.Job.Properties.CapacityUnitCount > 0 {
			fmt.Fprintf(w, "Capacity Units:\t%d\n", d.Job.Properties.CapacityUnitCount)
		} else if c.CapacityUnitCount > 0 {

@jongio Jon Gallant (jongio) left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Three things worth addressing before this merges.

Medium

  • Silent migration gap: resources.properties.AISuperComputer.priority and resources.instance_type are dropped without warning for any job.yaml scaffolded by 0.0.2-preview. Inline on yaml_parser.go.
  • capacity_unit_count and priority skip offline validation even though the new template documents allowed values for both.

Low

  • RunHistoryCompute.InstanceType is retained per the PR description but never rendered.

The init_template.go sizing conflict and the gpu_count migration point already raised on this PR both still look open.

ShmSize string `yaml:"shm_size"`
DockerArgs string `yaml:"docker_args"`
Properties map[string]any `yaml:"properties"`
InstanceCount int `yaml:"instance_count"`

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The previous init template scaffolded this block with instance_type: <user to add> and properties.AISuperComputer.slaTier / priority: <user to add>, so every job.yaml generated by 0.0.2-preview has those keys filled in with values the user deliberately picked.

ParseJobFile calls yaml.Unmarshal rather than a decoder with KnownFields(true), so after this change those keys still parse and then get dropped on the floor. The user's priority selection silently reverts to the service default because the new top-level priority is empty, and nothing in the submit output tells them their setting was ignored.

Worth keeping a deprecated shadow of the old fields long enough to emit a migration error:

  • keep InstanceType and Properties on ResourceDefinition
  • in ValidateJobOffline, fail when either is set, with something like "resources.instance_type is no longer supported, the service infers the SKU from the compute cluster" and "resources.properties.AISuperComputer.priority moved to the top-level priority field"

That turns a silent behavior change into a one-time actionable error.

Comment thread cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go Outdated
Comment thread cli/azd/extensions/azure.ai.training/pkg/models/history.go Outdated
@microsoft-github-policy-service microsoft-github-policy-service Bot added the no-recent-activity identity issues with no activity label Aug 10, 2026
Copilot AI review requested due to automatic review settings August 14, 2026 08:58
@microsoft-github-policy-service microsoft-github-policy-service Bot removed the no-recent-activity identity issues with no activity label Aug 14, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (2)

cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go:33

  • The advertised YAML migration is not implemented here: JobDefinition still accepts gpu_count and has no capacity_unit_count. Since yaml.Unmarshal ignores unknown keys, a file using the new field silently submits zero; buildJobResource and CommandJob also continue emitting gpuCount, while the template/show paths still use GPU count. Replace the submit path with CapacityUnitCount (capacity_unit_count in YAML and capacityUnitCount on the wire), retaining GPUCount only in RunHistoryCompute for legacy reads.
	Priority             string                       `yaml:"priority"`

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:497

  • Top-level priority is only printed when run history contains a compute object. fetchJobDetails tolerates a missing/failed run-history lookup, so a valid Priority returned by Get Job disappears in that case instead of being preferred. Close the history-compute block before this fallback and consult run history only when top-level priority is empty.
		// Priority: prefer the top-level Priority from Get Job, otherwise use
		// whatever run history reports.
		if d.Job.Properties.Priority != "" {
			fmt.Fprintf(w, "Priority:\t%s\n", d.Job.Properties.Priority)
		} else if c.Priority != "" {
			fmt.Fprintf(w, "Priority:\t%s\n", c.Priority)

Copilot AI review requested due to automatic review settings August 14, 2026 10:26
@saanikaguptamicrosoft saanikaguptamicrosoft changed the title training: Add support for capacityUnitCount and update command job schema training: Align samples with new command job schema Aug 14, 2026
@saanikaguptamicrosoft saanikaguptamicrosoft changed the title training: Align samples with new command job schema training: Align training job yaml schema with new command job schema Aug 14, 2026

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (2)

cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go:85

  • [azd-code-reviewer] Removing these fields from the struct does not reject the old schema: both submit and validate use yaml.Unmarshal, which silently ignores unknown keys. Existing files with resources.instance_type, resources.properties.AISuperComputer.priority, and similar settings will still validate and submit while those settings are discarded. Decode with KnownFields(true) in both paths, or retain deprecated fields long enough to emit a targeted migration error, and update the command-job guide.
// Only instance_count is honored; instance_type, slaTier, priority and the AISuperComputer
// properties block have been removed — the service now infers the SKU from the compute
// cluster, priority is a top-level job field, and users specify partial GPU allocations
// via the top-level gpu_count field (GPU clusters only).
type ResourceDefinition struct {

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:493

  • [azd-code-reviewer] Top-level GPUCount and Priority are still printed only when run-history compute data exists. fetchJobDetails treats history errors and 404s as nonfatal, so job show omits these authoritative Get Job values whenever history is unavailable. Move the top-level checks outside this block and use history only as the fallback.
		// Priority: prefer the top-level Priority from Get Job, otherwise use
		// whatever run history reports.
		if d.Job.Properties.Priority != "" {
			fmt.Fprintf(w, "Priority:\t%s\n", d.Job.Properties.Priority)

…block

Replace the historical capacity_unit_count/gpu_count/nested-priority sizing knobs with the current MFE contract: top-level gpuCount for partial-GPU allocation on GPU clusters, top-level priority, and instance_count as the only field on the resources block. Update job show to render GPUs and inferred instance_type/priority from the run history.
Copilot AI review requested due to automatic review settings August 14, 2026 10:58

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 7 changed files in this pull request and generated no new comments.

Suppressed comments (3)

cli/azd/extensions/azure.ai.training/internal/utils/yaml_parser.go:87

  • [azd-code-reviewer] Removing these fields does not make old YAML invalid: ParseJobFile uses yaml.Unmarshal, which silently ignores unknown mapping keys. Existing files with resources.priority, instance_type, or properties will therefore submit successfully after those settings have been discarded; in particular, priority silently falls back to the service default. Retain the legacy keys long enough to return a migration error, or add strict/targeted decoding that rejects them with the new top-level replacements.
	InstanceCount int `yaml:"instance_count"`

cli/azd/extensions/azure.ai.training/internal/cmd/job_submit.go:184

  • [azd-code-reviewer] The new YAML-to-wire priority contract is not covered by the existing buildJobResource test, so a wrong nesting or JSON key could regress unnoticed. Add a table-driven case that marshals the payload and verifies priority is top-level, along with the CPU/GPU/resource combinations listed in the PR testing notes.
		Priority:                  def.Priority,

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:485

  • [azd-code-reviewer] These top-level values are still emitted only inside the History.Compute != nil branch. fetchJobDetails deliberately treats missing/failed run history as nonfatal, so job show can have valid GPUCount and Priority from Get Job but omit both from the display. Keep only the history fallbacks inside that condition and print the top-level values independently; cover the no-history case.
		// GPUs: prefer the top-level GPUCount from Get Job, otherwise use what run history reports.
		if d.Job.Properties.GPUCount > 0 {

Copilot AI review requested due to automatic review settings August 14, 2026 12:19

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

cli/azd/extensions/azure.ai.training/internal/cmd/job_submit.go:184

  • azd-code-reviewer: Add table-driven tests for the CPU, priority, and partial-GPU submission shapes described by this PR. The current buildJobResource test only checks outputs, so it does not verify that priority and gpuCount are serialized at the intended level or that only instanceCount remains under resources.
		GPUCount:                  def.GPUCount,
		Priority:                  def.Priority,

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:492

  • azd-code-reviewer: Both top-level values are still rendered only when run history contains a compute object. GetRunHistory legitimately returns nil on 404, so a queued or newly submitted job can have gpuCount/priority in Get Job but omit both from table output. Render the top-level values independently and use run history only as the fallback.
		// Priority: prefer the top-level Priority from Get Job, otherwise use
		// whatever run history reports.
		if d.Job.Properties.Priority != "" {

cli/azd/extensions/azure.ai.training/internal/cmd/job_show.go:399

  • azd-code-reviewer: The resource section no longer prints ShmSize, even though ResourceConfig still retains it specifically for Get Job responses. Existing jobs or service-inferred responses that include shmSize therefore lose information in table output. Keep the previous “Shared Memory” row on the read path.
			fmt.Fprintf(w, "Instance Type:\t%s\n", props.Resources.InstanceType)

ShmSize string `yaml:"shm_size"`
DockerArgs string `yaml:"docker_args"`
Properties map[string]any `yaml:"properties"`
InstanceCount int `yaml:"instance_count"`
@saanikaguptamicrosoft

Copy link
Copy Markdown
Member Author

/check-enforcer override

@saanikaguptamicrosoft
saanikaguptamicrosoft enabled auto-merge (squash) August 14, 2026 12:31
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ext-training azure.ai.training extension

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Align training job yaml schema with new command job schema

5 participants