training: replace capacityUnitCount with gpuCount per service decision

Service decided against capacityUnitCount terminology. The MFE contract now exposes gpuCount at properties level for partial-GPU allocation on GPU clusters. CPU clusters continue to use full SKUs via instanceCount only.

Author: saanikaguptamicrosoft

cli/azd/extensions/azure.ai.training/internal/cmd/init_template.go

@@ -75,11 +75,11 @@ environment: <user to add>
 compute: <user to add>
 
 # Optional: choose one of the two ways to size the job.
-#   capacity_unit_count — total cores across the allocation (1, 2, 4, 8 or multiple of 8).
-#     The service picks the smallest SKU that fits and slices nodes as needed (partial SKU).
-#   resources.instance_count — whole nodes of the compute cluster's full SKU.
+#   gpu_count — number of GPUs to allocate (partial SKU on GPU clusters only).
+#     The service picks the smallest GPU SKU that fits and slices nodes as needed.
+#   resources.instance_count — whole nodes of the compute cluster's full SKU (CPU or GPU).
 #     Default when nothing is set: instance_count = 1.
-# capacity_unit_count: 4
+# gpu_count: 2
 resources:
   instance_count: 1
 

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

@@ -481,12 +481,9 @@ func printComputeSection(d *jobDetails) {
 		if c.InstanceCount > 0 {
 			fmt.Fprintf(w, "Nodes:\t%d\n", c.InstanceCount)
 		}
-		// 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 {
-			fmt.Fprintf(w, "Capacity Units:\t%d\n", c.CapacityUnitCount)
+		// GPUs: prefer the top-level GPUCount from Get Job, otherwise use what run history reports.
+		if d.Job.Properties.GPUCount > 0 {
+			fmt.Fprintf(w, "GPUs:\t%d\n", d.Job.Properties.GPUCount)
 		} else if c.GPUCount > 0 {
 			fmt.Fprintf(w, "GPUs:\t%d\n", c.GPUCount)
 		}

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

@@ -162,7 +162,7 @@ func buildJobResource(def *utils.JobDefinition) *models.JobResource {
 		Command:                   def.Command,
 		EnvironmentImageReference: def.Environment,
 		ComputeID:                 def.Compute,
-		CapacityUnitCount:         def.CapacityUnitCount,
+		GPUCount:                  def.GPUCount,
 		Priority:                  def.Priority,
 		UserAssignedIdentityID:    def.Identity,
 		CodeID:                    def.Code,
@@ -212,9 +212,9 @@ func buildJobResource(def *utils.JobDefinition) *models.JobResource {
 
 	// Resources — only instance_count survives on the wire. instance_type, slaTier and the
 	// AISuperComputer sub-block are now inferred by the service from the compute cluster;
-	// priority is a top-level CommandJob field; partial-SKU is expressed via top-level
-	// capacity_unit_count. Accept instance_count from either the structured resources block
-	// or the flat top-level field.
+	// priority is a top-level CommandJob field; partial GPU allocation is expressed via the
+	// top-level gpu_count field. Accept instance_count from either the structured resources
+	// block or the flat top-level field.
 	if def.Resources != nil && def.Resources.InstanceCount > 0 {
 		job.Resources = &models.ResourceConfig{
 			InstanceCount: def.Resources.InstanceCount,

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

@@ -29,7 +29,7 @@ type JobDefinition struct {
 	Distribution         *DistributionDefinition      `yaml:"distribution"`
 	Resources            *ResourceDefinition          `yaml:"resources"`
 	InstanceCount        int                          `yaml:"instance_count"`
-	CapacityUnitCount    int                          `yaml:"capacity_unit_count"`
+	GPUCount             int                          `yaml:"gpu_count"`
 	Priority             string                       `yaml:"priority"`
 	EnvironmentVariables map[string]string            `yaml:"environment_variables"`
 	Identity             string                       `yaml:"identity"`
@@ -80,8 +80,8 @@ type ServiceDefinition struct {
 // ResourceDefinition represents the compute resource configuration in a YAML job definition.
 // 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.
+// 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 {
 	InstanceCount int `yaml:"instance_count"`
 }

cli/azd/extensions/azure.ai.training/pkg/models/history.go

@@ -54,7 +54,6 @@ type RunHistoryCompute struct {
 	VMSize            string `json:"vmSize,omitempty"`
 	InstanceType      string `json:"instanceType,omitempty"`
 	InstanceCount     int    `json:"instanceCount,omitempty"`
-	CapacityUnitCount int    `json:"capacityUnitCount,omitempty"`
 	GPUCount          int    `json:"gpuCount,omitempty"`
 	Priority          string `json:"priority,omitempty"`
 }

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

@@ -23,7 +23,7 @@ type CommandJob struct {
 	EnvironmentImageReference string               `json:"environmentImageReference,omitempty"`
 	CodeID                    string               `json:"codeId,omitempty"`
 	ComputeID                 string               `json:"computeId,omitempty"`
-	CapacityUnitCount         int                  `json:"capacityUnitCount,omitempty"`
+	GPUCount                  int                  `json:"gpuCount,omitempty"`
 	Priority                  string               `json:"priority,omitempty"`
 	UserAssignedIdentityID    string               `json:"userAssignedIdentityId,omitempty"`
 	Inputs                    map[string]JobInput  `json:"inputs,omitempty"`