training: Add support for capacityUnitCount and update command job schema

[saanikaguptamicrosoft]

Aligns the training extension with Foundry's new job submission shape:

• YAML: add top-level capacity_unit_count and priority; drop gpu_count and trim resources block to instance_count only. The service now infers the SKU from the compute cluster; users express partial-SKU allocation via capacity_unit_count.

• Wire: CommandJob gains top-level capacityUnitCount and priority; ResourceConfig trimmed to instanceCount (instance_type, shm_size, docker_args, properties/AISuperComputer removed from the write path).

• Read/show: RunHistoryCompute keeps GPUCount, InstanceType and legacy Priority so 'job show' still displays useful info for older jobs; new jobs prefer the top-level CapacityUnitCount and Priority from Get Job, falling back to run history otherwise.

• Template: 'azd ai training init' scaffolds the new shape with capacity_unit_count and priority as commented-out optional fields.

[azure-pipelines]

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]

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)

[Copilot]

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_count → capacity_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]

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.

[jongio]

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.

[jongio]

Neither capacity_unit_count nor priority is checked in ValidateJobOffline, but the template this PR scaffolds documents hard constraints for both: capacity units must be 1, 2, 4, 8 or a multiple of 8, and priority is Regular, Low or High.

That's inconsistent with how the validator treats every other constrained field. services.type and distribution.type are enum-checked, and the Ray port range check exists specifically so "obvious typos (e.g. 99999) don't require a network round-trip to discover". capacity_unit_count: 3 or a misspelled priority has exactly that shape of failure, and right now it costs a full submit round-trip to find out.

[jongio]

The PR description says run history keeps InstanceType so job show still displays useful info for older jobs, but this field isn't read anywhere. printComputeSection renders VMSize, InstanceCount, CapacityUnitCount, GPUCount and Priority from run history, and the Instance Type line in printJobDetails comes from props.Resources.InstanceType on the Get Job response, not from here.

Is the run-history value meant to be a fallback for when Get Job doesn't echo the resources block, or is the field dead?