feat(job): Enable MTC at job submission stage

cmd/job/submit.go

@@ -62,13 +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
+	gkeMtcEnabled          bool
+	gkeMtcRamdiskDirectory string
+
+	isPathwaysJob bool
+	pathways      orchestrator.PathwaysJobDefinition
 
 	gkeNapProvisioning string
 	gkeNapReservation  string
@@ -121,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
 	},
@@ -158,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).")
@@ -180,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(&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(&gkeMtcEnabled, "gke-mtc-enabled", false, "Enable Multi-Tier Checkpointing (MTC).")
+	SubmitCmd.Flags().StringVar(&gkeMtcRamdiskDirectory, "gke-mtc-ramdisk-dir", "", "The ramdisk directory path for local checkpoints in MTC.")
 
 	_ = SubmitCmd.MarkFlagRequired("name")
 	_ = SubmitCmd.MarkFlagRequired("compute-type")
@@ -209,7 +213,7 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		affinity["cpu-affinity"] = cpuAffinityStr
 	}
 
-	if timeoutStr != "-1s" {
+	if timeout != "-1s" {
 		awaitJobCompletion = true
 	}
 
@@ -251,15 +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,
+		GKEMTCEnabled:                 gkeMtcEnabled,
+		GKEMTCRamdiskDirectory:        gkeMtcRamdiskDirectory,
 		GKENAPProvisioning:            gkeNapProvisioning,
 		GKENAPReservation:             gkeNapReservation,
 		IsPathwaysJob:                 isPathwaysJob,
 		Pathways:                      pathways,
 		RawMounts:                     volumeStr,
 		Env:                           parseEnvFlags(envVars),
-		Verbose:                       verbose,
 	}
 
 	return orc.SubmitJob(jobDef)
@@ -375,6 +381,7 @@ func validateGKENAPFlags() error {
 	if gkeNapProvisioning != "reservation" && gkeNapReservation != "" {
 		return fmt.Errorf("--gke-nap-reservation should only be provided when --gke-nap-provisioning=reservation")
 	}
+
 	return nil
 }
 

cmd/job/submit_test.go

@@ -278,8 +278,10 @@ func setupSubmitTestEnv(t *testing.T) {
 	topology = ""
 	gkeScheduler = ""
 	platform = "linux/amd64"
+	gkeMtcEnabled = false
+	gkeMtcRamdiskDirectory = ""
 	awaitJobCompletion = false
-	priorityClassName = "medium"
+	priority = "medium"
 	isPathwaysJob = false
 	pathways = orchestrator.PathwaysJobDefinition{MaxSliceRestarts: 1}
 	gkeNapProvisioning = ""
@@ -296,6 +298,8 @@ func setupSubmitTestEnv(t *testing.T) {
 	gkeOrchestratorFactory = func() orchestrator.JobOrchestrator {
 		return &mockOrchestrator{}
 	}
+	gkeMtcEnabled = false
+	gkeMtcRamdiskDirectory = ""
 }
 
 type mockOrchestrator struct {
@@ -1000,20 +1004,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",
+		"--gke-mtc-enabled",
+		"--gke-mtc-ramdisk-dir", "/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 !gkeMtcEnabled {
+		t.Errorf("expected gkeMtcEnabled 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 gkeMtcRamdiskDirectory != "/tmp/custom_mtc_dir" {
+		t.Errorf("expected gkeMtcRamdiskDirectory to be /tmp/custom_mtc_dir, got %s", gkeMtcRamdiskDirectory)
 	}
 
 }

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 (`--gke-mtc-enabled` flag), ensure that the Multi-Tier Checkpointing feature is enabled on your GKE cluster.
+
+To use this feature, the cluster administrator must also ensure that the required `CheckpointConfiguration` Custom Resource is deployed to the cluster, specifying the target cloud storage bucket for checkpoints. Once the cluster is configured, job submitters simply pass `--gke-mtc-enabled` to their jobs. For more details on cluster configuration, see the [GKE Multi-Tier Checkpointing documentation](https://cloud.google.com/kubernetes-engine/docs/how-to/multi-tier-checkpointing).
+
 ## 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.
@@ -1093,6 +1099,7 @@ The `gcluster job submit` command deploys a container image as a job (Kubernetes
 | `-o, --dry-run-out` | `string` | Local file path to save the generated Kubernetes manifest instead of applying it (must specify a file path, not a directory). |
 | `--num-slices` | `int` | Number of independent groups/slices to use (Default: `1`). |
 | `--num-nodes` | `int` | Number of nodes to use per group/slice (Default: `1`). Auto-calculated for TPUs based on topology. |
+| `--node-constraint` | `string` | Maps to Kubernetes node labels to target specific hardware instance types. Supports pipe separator (`|`) for multiple values. |
 | `--restarts` | `int` | Maximum number of restarts allowed for the JobSet before marked as failed (Default: `1`). |
 | `--mount` | `stringArray` | Mount storage volumes, buckets, filestore instances, or PVCs using the `<src>:<dest>[:<mode>]` format. Examples of `<src>`: `gs://my-bucket`, `filestore://my-instance/share`, `my-pvc` (for Lustre/etc), or `/host/path`. |
 | `--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. |
@@ -1124,8 +1131,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.*
@@ -1140,16 +1145,19 @@ The `gcluster job submit` command deploys a container image as a job (Kubernetes
 | :--- | :--- | :--- |
 | `-q, --queue` | `string` | Name of the Kueue `LocalQueue` to submit the job to (Auto-discovered by default). |
 | `--priority` | `string` | Priority class name assigned to the job queue (supports default classes like `low`, `medium`, `high`, or any custom PriorityClass defined in the cluster). If empty, the cluster's default priority class will be used. |
-| `--gke-ttl-after-finished` | `string` | Time duration to retain the JobSet resources after completion (Default: `1h`). |
 | `--grace-period` | `string` | Buffer period given to pods to save checkpoints before forced termination (Default: `30s`). |
-| `--node-constraint` | `string` | Maps to Kubernetes node labels to target specific hardware instance types. Supports pipe separator (`|`) for multiple values. |
+| `--gke-mtc-enabled` | `flag` | If present, enables Multi-Tier Checkpointing (MTC) for the workload. |
+| `--gke-mtc-ramdisk-dir` | `string` | The ramdisk directory path for local checkpoints in MTC (defaults to `/tmp/mtc_checkpoints`). |
+| `--gke-ttl-after-finished` | `string` | Time duration to retain the JobSet resources after completion (Default: `1h`). |
 | `--placement-policy` | `string` | Specifies a GCE Placement Policy name (e.g., `compact-placement`) to minimize latency. |
 | `--restart-on-exit-codes` | `string` | Comma-separated list of retriable exit codes that bypass the main restart budget. |
 | `--gke-scheduler` | `string` | Specific GKE scheduler selection (e.g., `gke.io/topology-aware-auto`). |
 | `--image-pull-secret` | `string` | Secret name required to authenticate and pull images from private container registries. |
 | `--service-account` | `string` | Kubernetes service account name used to provide fine-grained IAM roles to the job pods. |
 | `--cpu-affinity` | `string` | CPU affinity rules (e.g., `'numa'`). |
 | `--gke-disable-parallel-containers` | `bool` | Disable parallel containers for TPU v7/v7x on GKE. (Default: `false`) |
+| `--gke-nap-provisioning` | `string` | Compute provisioning model for GKE NAP. Allowed values: `on-demand`, `spot`, `reservation`. |
+| `--gke-nap-reservation` | `string` | Name of the Google Cloud Reservation for GKE NAP (required if `--gke-nap-provisioning=reservation`). |
 
 ### 9.4 `list` Flags
 *Use these flags to filter the list of jobs.*

pkg/orchestrator/gke/gke_job_orchestrator.go

@@ -32,6 +32,8 @@ import (
 	"strings"
 	"time"
 
+	container "google.golang.org/api/container/v1"
+
 	"github.com/google/safetext/yamltemplate"
 
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -133,6 +135,25 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error {
 	return nil
 }
 
+func (g *GKEOrchestrator) checkMTCAddonEnabled(projectID, location, clusterName string) error {
+	ctx := context.Background()
+	svc, err := container.NewService(ctx)
+	if err != nil {
+		return fmt.Errorf("failed to create cluster manager client: %v", err)
+	}
+
+	name := fmt.Sprintf("projects/%s/locations/%s/clusters/%s", projectID, location, clusterName)
+	resp, err := svc.Projects.Locations.Clusters.Get(name).Context(ctx).Do()
+	if err != nil {
+		return fmt.Errorf("failed to get cluster details for MTC validation: %v", err)
+	}
+
+	if resp.AddonsConfig == nil || resp.AddonsConfig.HighScaleCheckpointingConfig == nil || !resp.AddonsConfig.HighScaleCheckpointingConfig.Enabled {
+		return fmt.Errorf("Multi-Tier Checkpointing (MTC) requires the HighScaleCheckpointing addon to be enabled on the target GKE cluster. Please update your cluster blueprint to set 'enable_multi_tier_checkpointing: true' and deploy the cluster before submitting jobs with --mtc-enabled")
+	}
+	return nil
+}
+
 // ListJobs retrieves a list of jobs in the GKE cluster.
 // It filters jobs based on the provided ListOptions.
 func (g *GKEOrchestrator) ListJobs(opts orchestrator.ListOptions) ([]orchestrator.JobStatus, error) {
@@ -369,8 +390,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.GKEMTCEnabled && job.GKEMTCRamdiskDirectory == "" {
+		job.GKEMTCRamdiskDirectory = "/tmp/mtc_checkpoints"
 	}
 
 	tmpl, err := yamltemplate.New("pathways_jobset.tmpl").ParseFS(templatesFS, "templates/pathways_jobset.tmpl")
@@ -543,6 +564,12 @@ func (g *GKEOrchestrator) initializeJobSubmission(job *orchestrator.JobDefinitio
 		return err
 	}
 
+	if job.GKEMTCEnabled {
+		if err := g.checkMTCAddonEnabled(job.ProjectID, job.ClusterLocation, job.ClusterName); err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
 
@@ -1366,6 +1393,8 @@ func (g *GKEOrchestrator) prepareJobSetTemplateData(opts ManifestOptions, comman
 		PathwaysWorkerEnv:             sortedEnvVars(opts.Pathways.WorkerEnv),
 		IsTPU:                         isTPU,
 		IsGPU:                         isGPU,
+		GKEMTCEnabled:                 opts.GKEMTCEnabled,
+		GKEMTCRamdiskDirectory:        opts.GKEMTCRamdiskDirectory,
 	}
 }
 

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,
+		GKEMTCEnabled: 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,
+		GKEMTCEnabled:                 job.GKEMTCEnabled,
+		GKEMTCRamdiskDirectory:        job.GKEMTCRamdiskDirectory,
 	}
 
 	if err := g.fillManifestStrings(&opts, schedOpts, job, isDynamicSlicing, isStaticSlicing, profile.IsCPUMachine); err != nil {

pkg/orchestrator/gke/storage.go

@@ -106,6 +106,7 @@ func (sm *StorageManager) parseSingleVolume(vStr string) (src, dest string, read
 	if err != nil {
 		return "", "", false, err
 	}
+
 	if err := validateSrcScheme(src, vStr); err != nil {
 		return "", "", false, err
 	}
@@ -289,10 +290,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 +302,28 @@ func (sm *StorageManager) AddVolumeOptions(opts *ManifestOptions, vols []MountIn
 		}
 	}
 
+	if opts.GKEMTCEnabled {
+		ramdiskDir := opts.GKEMTCRamdiskDirectory
+		if ramdiskDir == "" {
+			ramdiskDir = "/tmp/mtc_checkpoints"
+		}
+		mountSpecs = append(mountSpecs,
+			map[string]interface{}{"name": "cache", "mountPath": ramdiskDir},
+		)
+		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 {