fix: address PR comments for mtc flags

Author: Neelabh94

cmd/job/submit.go

@@ -62,16 +62,16 @@ var (
 	platform           string
 
 	awaitJobCompletion bool
-	timeoutStr         string
-	priorityClassName  string
-	isPathwaysJob      bool
+	timeout            string
+	priority           string
 	verbose            bool
+	volumeStr          []string
 
-	volumeStr []string
-	pathways  orchestrator.PathwaysJobDefinition
+	mtcEnabled          bool
+	mtcRamdiskDirectory string
 
-	mtcEnabled       bool
-	ramdiskDirectory string
+	isPathwaysJob bool
+	pathways      orchestrator.PathwaysJobDefinition
 
 	gkeNapProvisioning string
 	gkeNapReservation  string
@@ -124,7 +124,7 @@ and JobSet/Kueue specific configurations like workload name, queue, nodes, and r
 			}
 		}
 
-		priorityClassName = strings.ToLower(priorityClassName)
+		priority = strings.ToLower(priority)
 
 		return nil
 	},
@@ -161,8 +161,8 @@ func init() {
 	SubmitCmd.Flags().StringVar(&topology, "topology", "", "TPU slice topology (e.g., 2x2x1).")
 	SubmitCmd.Flags().StringVar(&gkeScheduler, "gke-scheduler", "", "Kubernetes Scheduler name (e.g., gke.io/topology-aware-auto).")
 	SubmitCmd.Flags().BoolVar(&awaitJobCompletion, "await-job-completion", false, "If true, gcluster will wait for the submitted job to complete.")
-	SubmitCmd.Flags().StringVar(&timeoutStr, "timeout", "-1s", "Time to wait for job in seconds or string format (e.g. 1h, 10m). Default is max timeout (-1s).")
-	SubmitCmd.Flags().StringVar(&priorityClassName, "priority", "", "A priority class name (e.g., low, medium, high, or any custom PriorityClass defined in the cluster). If empty, the cluster's default priority class will be used.")
+	SubmitCmd.Flags().StringVar(&timeout, "timeout", "-1s", "Time to wait for job in seconds or string format (e.g. 1h, 10m). Default is max timeout (-1s).")
+	SubmitCmd.Flags().StringVar(&priority, "priority", "", "A priority class name (e.g., low, medium, high, or any custom PriorityClass defined in the cluster). If empty, the cluster's default priority class will be used.")
 	SubmitCmd.Flags().BoolVar(&verbose, "verbose", false, "Enable verbose logging for the workload (TPUs and GPUs).")
 	SubmitCmd.Flags().StringVar(&gkeNapProvisioning, "gke-nap-provisioning", "", "Compute provisioning model for GKE NAP. Allowed values: on-demand, spot, reservation.")
 	SubmitCmd.Flags().StringVar(&gkeNapReservation, "gke-nap-reservation", "", "Name of the Google Cloud Reservation for GKE NAP (required if --gke-nap-provisioning=reservation).")
@@ -183,8 +183,9 @@ func init() {
 	SubmitCmd.Flags().StringArrayVar(&pathwaysWorkerEnv, "pathways-worker-env", []string{}, "Custom environment variables for the Pathways worker container in KEY=VALUE format. Can be specified multiple times.")
 	SubmitCmd.Flags().StringVar(&pathways.ColocatedPythonSidecarImage, "pathways-colocated-python-sidecar-image", "", "Image for an optional Python-based sidecar container to run alongside the Pathways head components.")
 	SubmitCmd.Flags().StringVar(&pathways.HeadNodePool, "pathways-head-np", "", "The node pool to use for the Pathways head job. If empty, it will be auto-detected (looking for 'cpu-np' or 'pathways-np').")
+
 	SubmitCmd.Flags().BoolVar(&mtcEnabled, "mtc-enabled", false, "Enable Multi-Tier Checkpointing (MTC).")
-	SubmitCmd.Flags().StringVar(&ramdiskDirectory, "mtc-ramdisk-directory", "", "The ramdisk directory path for local checkpoints in MTC.")
+	SubmitCmd.Flags().StringVar(&mtcRamdiskDirectory, "mtc-ramdisk-directory", "", "The ramdisk directory path for local checkpoints in MTC.")
 
 	_ = SubmitCmd.MarkFlagRequired("name")
 	_ = SubmitCmd.MarkFlagRequired("compute-type")
@@ -212,7 +213,7 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		affinity["cpu-affinity"] = cpuAffinityStr
 	}
 
-	if timeoutStr != "-1s" {
+	if timeout != "-1s" {
 		awaitJobCompletion = true
 	}
 
@@ -254,17 +255,17 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		GKEScheduler:                  gkeScheduler,
 		AwaitJobCompletion:            awaitJobCompletion,
 		UseParallelContainers:         !gkeDisableParallelContainers,
-		Timeout:                       timeoutStr,
-		PriorityClassName:             priorityClassName,
+		Timeout:                       timeout,
+		PriorityClassName:             priority,
+		Verbose:                       verbose,
+		MTCEnabled:                    mtcEnabled,
+		MTCRamdiskDirectory:           mtcRamdiskDirectory,
 		GKENAPProvisioning:            gkeNapProvisioning,
 		GKENAPReservation:             gkeNapReservation,
 		IsPathwaysJob:                 isPathwaysJob,
 		Pathways:                      pathways,
-		MTCEnabled:                    mtcEnabled,
-		RamdiskDirectory:              ramdiskDirectory,
 		RawMounts:                     volumeStr,
 		Env:                           parseEnvFlags(envVars),
-		Verbose:                       verbose,
 	}
 
 	return orc.SubmitJob(jobDef)
@@ -381,10 +382,6 @@ func validateGKENAPFlags() error {
 		return fmt.Errorf("--gke-nap-reservation should only be provided when --gke-nap-provisioning=reservation")
 	}
 
-	if pathways.MTCEnabled {
-		// MTC Addon validation is now handled by the GKE Orchestrator preflight checks.
-	}
-
 	return nil
 }
 

cmd/job/submit_test.go

@@ -279,7 +279,7 @@ func resetSubmitCmdFlags() {
 	gkeScheduler = ""
 	platform = "linux/amd64"
 	awaitJobCompletion = false
-	priorityClassName = "medium"
+	priority = "medium"
 	isPathwaysJob = false
 	pathways = orchestrator.PathwaysJobDefinition{MaxSliceRestarts: 1}
 	gkeNapProvisioning = ""
@@ -288,6 +288,8 @@ func resetSubmitCmdFlags() {
 	pathwaysProxyEnv = nil
 	pathwaysServerEnv = nil
 	pathwaysWorkerEnv = nil
+	mtcEnabled = false
+	mtcRamdiskDirectory = ""
 }
 
 type mockOrchestrator struct {
@@ -996,8 +998,8 @@ func TestSubmitCmd_PathwaysMTCFlags(t *testing.T) {
 		t.Errorf("expected mtcEnabled to be true")
 	}
 
-	if ramdiskDirectory != "/tmp/custom_mtc_dir" {
-		t.Errorf("expected ramdiskDirectory to be /tmp/custom_mtc_dir, got %s", ramdiskDirectory)
+	if mtcRamdiskDirectory != "/tmp/custom_mtc_dir" {
+		t.Errorf("expected mtcRamdiskDirectory to be /tmp/custom_mtc_dir, got %s", mtcRamdiskDirectory)
 	}
 
 }