feat(mtc): Enable MTC at job submission stage

Author: Neelabh94

cmd/job/submit.go

@@ -70,6 +70,9 @@ var (
 	volumeStr []string
 	pathways  orchestrator.PathwaysJobDefinition
 
+	mtcEnabled       bool
+	ramdiskDirectory string
+
 	gkeNapProvisioning string
 	gkeNapReservation  string
 
@@ -180,8 +183,8 @@ 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(&pathways.MTCEnabled, "pathways-mtc-enabled", false, "Enable Multi-Tier Checkpointing (MTC) for Pathways.")
-	SubmitCmd.Flags().StringVar(&pathways.RamdiskDirectory, "pathways-ramdisk-directory", "", "The ramdisk directory path for local checkpoints in MTC.")
+	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.MarkFlagRequired("name")
 	_ = SubmitCmd.MarkFlagRequired("compute-type")
@@ -257,6 +260,8 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		GKENAPReservation:             gkeNapReservation,
 		IsPathwaysJob:                 isPathwaysJob,
 		Pathways:                      pathways,
+		MTCEnabled:                    mtcEnabled,
+		RamdiskDirectory:              ramdiskDirectory,
 		RawMounts:                     volumeStr,
 		Env:                           parseEnvFlags(envVars),
 		Verbose:                       verbose,

cmd/job/submit_test.go

@@ -984,20 +984,20 @@ func TestSubmitCmd_PathwaysMTCFlags(t *testing.T) {
 		"--project", "test-project",
 		"--pathways-gcs-location", "gs://my-bucket",
 		"--compute-type", "n2-standard-4",
-		"--pathways-mtc-enabled",
-		"--pathways-ramdisk-directory", "/tmp/custom_mtc_dir",
+		"--mtc-enabled",
+		"--mtc-ramdisk-directory", "/tmp/custom_mtc_dir",
 	)
 
 	if err != nil {
 		t.Fatalf("command failed with error: %v", err)
 	}
 
-	if !pathways.MTCEnabled {
-		t.Errorf("expected pathways.MTCEnabled to be true")
+	if !mtcEnabled {
+		t.Errorf("expected mtcEnabled to be true")
 	}
 
-	if pathways.RamdiskDirectory != "/tmp/custom_mtc_dir" {
-		t.Errorf("expected pathways.RamdiskDirectory to be /tmp/custom_mtc_dir, got %s", pathways.RamdiskDirectory)
+	if ramdiskDirectory != "/tmp/custom_mtc_dir" {
+		t.Errorf("expected ramdiskDirectory to be /tmp/custom_mtc_dir, got %s", ramdiskDirectory)
 	}
 
 }

docs/gcluster_job_guide.md

@@ -29,6 +29,12 @@ If you use `--build-context` to build images on-the-fly, you must set:
 >
 > Successful checks are remembered in `~/.gcluster/job_prereq_state.json` to optimize subsequent runs. Checks are re-run if the state is older than 24 hours or if you switch projects.
 
+### 1.1 Multi-Tier Checkpointing (MTC) Prerequisites
+
+If you plan to use Multi-Tier Checkpointing (`--mtc-enabled` flag), ensure that your cluster was created with `enable_multi_tier_checkpointing = true` (or `stateful_ha_config { enabled = true }`).
+
+To use this feature, the cluster administrator must deploy the required `CheckpointConfiguration` Custom Resource globally at cluster creation. This is done by specifying the `mtc_target_bucket` variable in the Terraform blueprint. Job submitters then only need to pass `--mtc-enabled` to their jobs.
+
 ## 2. Prepare Sample Application Code
 
 Create a directory named `job_details` and place your application files inside it. This will serve as your build context for the job. The tool will package all files in this directory and add them to the image.
@@ -1098,6 +1104,8 @@ The `gcluster job submit` command deploys a container image as a job (Kubernetes
 | `--env` | `stringArray` | Custom environment variables to pass exclusively to the user's workload container in KEY=VALUE format (e.g. `--env KEY=VALUE`). Applies to both standard and Pathways workloads. Can be specified multiple times. |
 | `--await-job-completion` | `bool` | If true, the CLI waits for the job to complete before exiting. |
 | `--timeout` | `string` | Time to wait for job completion (e.g., `1h`, `10m`). Used with `--await-job-completion`. |
+| `--mtc-enabled` | `flag` | If present, enables Multi-Tier Checkpointing (MTC) for the workload. |
+| `--mtc-ramdisk-directory` | `string` | The ramdisk directory path for local checkpoints in MTC (defaults to `/tmp/mtc_checkpoints`). |
 | `--verbose` | `bool` | Enable verbose logging for the workload. |
 
 *(Note: `--cluster`, `--location`, and `--project` are also supported as common flags, see 9.1)*
@@ -1124,8 +1132,6 @@ The `gcluster job submit` command deploys a container image as a job (Kubernetes
 | `--pathways-worker-env` | `stringArray` | Custom environment variables injected specifically into the Pathways worker containers (KEY=VALUE). |
 | `--pathways-colocated-python-sidecar-image` | `string` | Image for an optional Python-based sidecar container running alongside workers. |
 | `--pathways-head-np` | `string` | The node pool name to target for the Pathways head job. |
-| `--pathways-mtc-enabled` | `flag` | If present, enables Multi-Tier Checkpointing (MTC) for the Pathways workload. |
-| `--pathways-ramdisk-directory` | `string` | The ramdisk directory path for local checkpoints in MTC (defaults to `/tmp/mtc_checkpoints`). |
 
 #### 9.3.3 GPU Related Flags
 *Use these flags to tune specialized multi-GPU topologies and related node parameters.*

pkg/orchestrator/gke/gke_job_orchestrator.go

@@ -369,8 +369,8 @@ func (g *GKEOrchestrator) GeneratePathwaysManifest(job orchestrator.JobDefinitio
 		// WorkerImage defaults to ServerImage if not explicitly set
 		job.Pathways.WorkerImage = job.Pathways.ServerImage
 	}
-	if job.Pathways.MTCEnabled && job.Pathways.RamdiskDirectory == "" {
-		job.Pathways.RamdiskDirectory = "/tmp/mtc_checkpoints"
+	if job.MTCEnabled && job.RamdiskDirectory == "" {
+		job.RamdiskDirectory = "/tmp/mtc_checkpoints"
 	}
 
 	tmpl, err := yamltemplate.New("pathways_jobset.tmpl").ParseFS(templatesFS, "templates/pathways_jobset.tmpl")
@@ -499,6 +499,12 @@ func (g *GKEOrchestrator) populateClusterMetadata(job *orchestrator.JobDefinitio
 	g.nodePoolSAs = nodePoolSAs
 	logging.Info("Calculated cluster capacity: %+v", g.capacity)
 
+	if job.MTCEnabled {
+		if clusterDesc.AddonsConfig == nil || clusterDesc.AddonsConfig.StatefulHaConfig == nil || !clusterDesc.AddonsConfig.StatefulHaConfig.Enabled {
+			return fmt.Errorf("MTC is not enabled on cluster '%s'. If you created your cluster using Cluster Toolkit, you can enable MTC by updating your cluster blueprint's gke-cluster module settings with:\n  enable_multi_tier_checkpointing: true\n  mtc_target_bucket: \"gs://YOUR_BUCKET\"\n  mtc_cache_size: \"50Gi\"\nThen run 'gcluster deploy'. For clusters provisioned otherwise or for underlying GKE concepts, see: https://docs.cloud.google.com/kubernetes-engine/docs/how-to/machine-learning/training/multi-tier-checkpointing", job.ClusterName)
+		}
+	}
+
 	return nil
 }
 
@@ -1373,6 +1379,8 @@ func (g *GKEOrchestrator) prepareJobSetTemplateData(opts ManifestOptions, comman
 		PathwaysWorkerEnv:             sortedEnvVars(opts.Pathways.WorkerEnv),
 		IsTPU:                         isTPU,
 		IsGPU:                         isGPU,
+		MTCEnabled:                    opts.MTCEnabled,
+		RamdiskDirectory:              opts.RamdiskDirectory,
 	}
 }
 

pkg/orchestrator/gke/gke_job_orchestrator_test.go

@@ -535,8 +535,9 @@ func TestGeneratePathwaysManifest_MTC(t *testing.T) {
 			ColocatedPythonSidecarImage: "sidecar:latest",
 			GCSLocation:                 "gs://my-bucket",
 			HeadNodePool:                "pathways-np",
-			MTCEnabled:                  true,
 		},
+		IsPathwaysJob: true,
+		MTCEnabled:    true,
 	}
 
 	mockResponses := map[string][]shell.CommandResult{

pkg/orchestrator/gke/manifest_generator.go

@@ -165,6 +165,10 @@ func (g *GKEOrchestrator) PrepareManifestOptions(job orchestrator.JobDefinition,
 		Topology:                      schedOpts.Topology,
 		Verbose:                       job.Verbose,
 		Env:                           job.Env,
+		IsPathwaysJob:                 job.IsPathwaysJob,
+		Pathways:                      job.Pathways,
+		MTCEnabled:                    job.MTCEnabled,
+		RamdiskDirectory:              job.RamdiskDirectory,
 	}
 
 	if err := g.fillManifestStrings(&opts, schedOpts, job, isDynamicSlicing, isStaticSlicing, profile.IsCPUMachine); err != nil {

pkg/orchestrator/gke/storage.go

@@ -289,10 +289,6 @@ func sanitizePVCName(name string) string {
 
 // AddVolumeOptions marshals and indents the volume and volume mount specifications into the manifest options.
 func (sm *StorageManager) AddVolumeOptions(opts *ManifestOptions, vols []MountInfo) {
-	if len(vols) == 0 {
-		return
-	}
-
 	var volSpecs []map[string]interface{}
 	var mountSpecs []map[string]interface{}
 	gcsFuseEnabled := false
@@ -305,6 +301,24 @@ func (sm *StorageManager) AddVolumeOptions(opts *ManifestOptions, vols []MountIn
 		}
 	}
 
+	if opts.MTCEnabled {
+		mountSpecs = append(mountSpecs,
+			map[string]interface{}{"name": "cache", "mountPath": opts.RamdiskDirectory},
+		)
+		volSpecs = append(volSpecs,
+			map[string]interface{}{"name": "cache", "csi": map[string]interface{}{"driver": "multitier-checkpoint.csi.storage.gke.io"}},
+		)
+
+		if opts.IsPathwaysJob {
+			mountSpecs = append(mountSpecs, map[string]interface{}{"name": "sidecar-shared-memory", "mountPath": "/tmp/sidecar"})
+			volSpecs = append(volSpecs, map[string]interface{}{"name": "sidecar-shared-memory", "emptyDir": map[string]interface{}{"medium": "Memory"}})
+		}
+	}
+
+	if len(volSpecs) == 0 {
+		return
+	}
+
 	opts.GCSFuseEnabled = gcsFuseEnabled
 
 	if b, err := yaml.Marshal(mountSpecs); err == nil {

pkg/orchestrator/gke/templates/pathways_jobset.tmpl

@@ -203,17 +203,17 @@ spec:
                 echo "Exit code: $EXIT_CODE"
                 exit $EXIT_CODE
               volumeMounts:
-              - name: shared-tmp
-                mountPath: /tmp
+                - name: shared-tmp
+                  mountPath: /tmp
 {{- if .VolumeMountsYAML }}
 {{(StructuralData .VolumeMountsYAML)}}
 {{- end }}
             {{- end}}
             volumes:
-            - name: shared-tmp
-              hostPath:
-                path: /tmp
-                type: DirectoryOrCreate
+              - name: shared-tmp
+                hostPath:
+                  path: /tmp
+                  type: DirectoryOrCreate
 {{- if .VolumesYAML }}
 {{(StructuralData .VolumesYAML)}}
 {{- end }}
@@ -285,9 +285,9 @@ spec:
               volumeMounts:
               - name: shared-tmp
                 mountPath: /tmp
-              {{- if .Pathways.MTCEnabled}}
+              {{- if .MTCEnabled}}
               - name: cache
-                mountPath: {{.Pathways.RamdiskDirectory}}
+                mountPath: {{.RamdiskDirectory}}
               - name: sidecar-shared-memory
                 mountPath: /tmp/sidecar
               {{- end}}
@@ -305,7 +305,7 @@ spec:
                 - --server_port=29005
                 - --resource_manager_address=$(PATHWAYS_HEAD):29001
                 - "--gcs_scratch_location={{.Pathways.GCSLocation}}"
-                {{- if .Pathways.MTCEnabled}}
+                {{- if .MTCEnabled}}
                 - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar
                 {{- end}}
                 {{- range .WorkerArgsList }}
@@ -354,30 +354,17 @@ spec:
               {{- end }}
 {{(StructuralData .ResourcesString)}}
               volumeMounts:
-              - name: shared-tmp
-                mountPath: /tmp
-              {{- if .Pathways.MTCEnabled}}
-              - name: cache
-                mountPath: {{.Pathways.RamdiskDirectory}}
-              - name: sidecar-shared-memory
-                mountPath: /tmp/sidecar
-              {{- end}}
+                - name: shared-tmp
+                  mountPath: /tmp
 {{- if .VolumeMountsYAML }}
 {{(StructuralData .VolumeMountsYAML)}}
 {{- end }}
             volumes:
-            - name: shared-tmp
-              hostPath:
-                path: /tmp
-                type: DirectoryOrCreate
-            {{- if .Pathways.MTCEnabled}}
-            - name: cache
-              csi:
-                driver: multitier-checkpoint.csi.storage.gke.io
-            - name: sidecar-shared-memory
-              emptyDir:
-                medium: Memory
-            {{- end}}
+              - name: shared-tmp
+                hostPath:
+                  path: /tmp
+                  type: DirectoryOrCreate
+
 {{- if .VolumesYAML }}
 {{(StructuralData .VolumesYAML)}}
 {{- end }}

pkg/orchestrator/gke/types.go

@@ -179,6 +179,9 @@ type ManifestOptions struct {
 	IsStaticSlicing               bool
 	IsCPUMachine                  bool
 	Pathways                      orchestrator.PathwaysJobDefinition
+	IsPathwaysJob                 bool
+	MTCEnabled                    bool
+	RamdiskDirectory              string
 	Verbose                       bool
 	Env                           map[string]string
 	AdditionalManifests           []string
@@ -276,9 +279,18 @@ type gkeClusterAutoscaling struct {
 }
 
 type gkeCluster struct {
-	Locations   []string              `json:"locations"`
-	NodePools   []gkeJobNodePool      `json:"nodePools"`
-	Autoscaling gkeClusterAutoscaling `json:"autoscaling"`
+	Locations    []string              `json:"locations"`
+	NodePools    []gkeJobNodePool      `json:"nodePools"`
+	Autoscaling  gkeClusterAutoscaling `json:"autoscaling"`
+	AddonsConfig *gkeAddonsConfig      `json:"addonsConfig,omitempty"`
+}
+
+type gkeAddonsConfig struct {
+	StatefulHaConfig *gkeStatefulHaConfig `json:"statefulHaConfig,omitempty"`
+}
+
+type gkeStatefulHaConfig struct {
+	Enabled bool `json:"enabled"`
 }
 
 // Types for JobSet status unmarshaling
@@ -353,6 +365,8 @@ type jobSetTemplateData struct {
 	PathwaysWorkerEnv             []EnvVar
 	IsTPU                         bool
 	IsGPU                         bool
+	MTCEnabled                    bool
+	RamdiskDirectory              string
 }
 
 // Types for parsing kubectl get nodes -o json

pkg/orchestrator/orchestrator.go

@@ -43,9 +43,6 @@ type PathwaysJobDefinition struct {
 
 	HeadNodePool string // Resolved node pool to use for the Pathways head job.
 
-	// Multi-Tier Checkpointing (MTC)
-	MTCEnabled       bool
-	RamdiskDirectory string
 }
 
 type VolumeDefinition struct {
@@ -98,6 +95,10 @@ type JobDefinition struct {
 	IsPathwaysJob bool
 	Pathways      PathwaysJobDefinition // Embedded struct for Pathways-specific args
 
+	// Multi-Tier Checkpointing (MTC)
+	MTCEnabled       bool
+	RamdiskDirectory string
+
 	RawMounts []string
 	Env       map[string]string