From 8582148b344bc1ab004ff923971a0aebd26cc6a3 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Wed, 22 Jul 2026 11:51:09 +0000 Subject: [PATCH 1/8] feat(mtc): Enable MTC at job submission stage --- cmd/job/submit.go | 9 +++- cmd/job/submit_test.go | 12 ++--- docs/gcluster_job_guide.md | 10 ++++- pkg/orchestrator/gke/gke_job_orchestrator.go | 12 ++++- .../gke/gke_job_orchestrator_test.go | 3 +- pkg/orchestrator/gke/manifest_generator.go | 4 ++ pkg/orchestrator/gke/storage.go | 22 +++++++-- .../gke/templates/pathways_jobset.tmpl | 45 +++++++------------ pkg/orchestrator/gke/types.go | 17 ++++++- pkg/orchestrator/orchestrator.go | 7 +-- 10 files changed, 91 insertions(+), 50 deletions(-) diff --git a/cmd/job/submit.go b/cmd/job/submit.go index d90fd16903..053070d7b7 100644 --- a/cmd/job/submit.go +++ b/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, diff --git a/cmd/job/submit_test.go b/cmd/job/submit_test.go index 9f0c462f98..63e3ceaf2f 100644 --- a/cmd/job/submit_test.go +++ b/cmd/job/submit_test.go @@ -1000,20 +1000,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) } } diff --git a/docs/gcluster_job_guide.md b/docs/gcluster_job_guide.md index 026c300310..c1ad2b0819 100644 --- a/docs/gcluster_job_guide.md +++ b/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.* diff --git a/pkg/orchestrator/gke/gke_job_orchestrator.go b/pkg/orchestrator/gke/gke_job_orchestrator.go index c0f0ecb219..6b49a1fd1f 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator.go +++ b/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") @@ -501,6 +501,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 } @@ -1366,6 +1372,8 @@ func (g *GKEOrchestrator) prepareJobSetTemplateData(opts ManifestOptions, comman PathwaysWorkerEnv: sortedEnvVars(opts.Pathways.WorkerEnv), IsTPU: isTPU, IsGPU: isGPU, + MTCEnabled: opts.MTCEnabled, + RamdiskDirectory: opts.RamdiskDirectory, } } diff --git a/pkg/orchestrator/gke/gke_job_orchestrator_test.go b/pkg/orchestrator/gke/gke_job_orchestrator_test.go index b8d0996606..353fed6169 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator_test.go +++ b/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{ diff --git a/pkg/orchestrator/gke/manifest_generator.go b/pkg/orchestrator/gke/manifest_generator.go index aa6121f497..00e62998c1 100644 --- a/pkg/orchestrator/gke/manifest_generator.go +++ b/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 { diff --git a/pkg/orchestrator/gke/storage.go b/pkg/orchestrator/gke/storage.go index 4e78506ec0..6da5d95231 100644 --- a/pkg/orchestrator/gke/storage.go +++ b/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 { diff --git a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl index e6a138a573..0e856c04b7 100644 --- a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl +++ b/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 }} diff --git a/pkg/orchestrator/gke/types.go b/pkg/orchestrator/gke/types.go index 1cbd065994..147d79af10 100644 --- a/pkg/orchestrator/gke/types.go +++ b/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 @@ -280,6 +283,7 @@ type gkeCluster struct { NodePools []gkeJobNodePool `json:"nodePools"` Autoscaling gkeClusterAutoscaling `json:"autoscaling"` ControlPlaneEndpointsConfig *controlPlaneEndpointsConfig `json:"controlPlaneEndpointsConfig,omitempty"` + AddonsConfig *gkeAddonsConfig `json:"addonsConfig,omitempty"` } type controlPlaneEndpointsConfig struct { @@ -287,7 +291,16 @@ type controlPlaneEndpointsConfig struct { } type dnsEndpointConfig struct { - AllowExternalTraffic bool `json:"allowExternalTraffic,omitempty"` + AllowExternalTraffic bool `json:"allowExternalTraffic,omitempty"` + 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 @@ -362,6 +375,8 @@ type jobSetTemplateData struct { PathwaysWorkerEnv []EnvVar IsTPU bool IsGPU bool + MTCEnabled bool + RamdiskDirectory string } // Types for parsing kubectl get nodes -o json diff --git a/pkg/orchestrator/orchestrator.go b/pkg/orchestrator/orchestrator.go index db63a3e1b2..b4e31792a7 100644 --- a/pkg/orchestrator/orchestrator.go +++ b/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 From 5bd5575eb2e597e206fe5326b28dba0b381f4b89 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Thu, 23 Jul 2026 09:37:45 +0000 Subject: [PATCH 2/8] Move MTC validation to GKE orchestrator --- cmd/job/submit.go | 5 ++++ pkg/orchestrator/gke/gke_job_orchestrator.go | 27 ++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/cmd/job/submit.go b/cmd/job/submit.go index 053070d7b7..388550cbb1 100644 --- a/cmd/job/submit.go +++ b/cmd/job/submit.go @@ -380,6 +380,11 @@ func validateGKENAPFlags() error { if gkeNapProvisioning != "reservation" && gkeNapReservation != "" { 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 } diff --git a/pkg/orchestrator/gke/gke_job_orchestrator.go b/pkg/orchestrator/gke/gke_job_orchestrator.go index 6b49a1fd1f..0f8e436e91 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator.go +++ b/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" @@ -97,6 +99,12 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error { return err } + if job.Pathways.MTCEnabled { + if err := g.checkMTCAddonEnabled(job.ProjectID, job.ClusterLocation, job.ClusterName); err != nil { + return err + } + } + if err := g.fetchClusterState(&job); err != nil { return err } @@ -133,6 +141,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 --pathways-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) { From 2e8627470446e0614c1719d3dc22806add7dc7d3 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Thu, 23 Jul 2026 10:07:59 +0000 Subject: [PATCH 3/8] fix: address PR comments for mtc flags --- cmd/job/submit.go | 39 +- cmd/job/submit_test.go | 8 +- pathways_jobset.tmpl.orig | 382 ++++++++++++++++++ pkg/orchestrator/gke/gke_job_orchestrator.go | 24 +- .../gke/gke_job_orchestrator_test.go | 3 +- pkg/orchestrator/gke/manifest_generator.go | 2 +- pkg/orchestrator/gke/storage.go | 7 +- .../gke/templates/pathways_jobset.tmpl | 16 +- pkg/orchestrator/gke/types.go | 4 +- pkg/orchestrator/orchestrator.go | 6 +- 10 files changed, 441 insertions(+), 50 deletions(-) create mode 100644 pathways_jobset.tmpl.orig diff --git a/cmd/job/submit.go b/cmd/job/submit.go index 388550cbb1..0142de1617 100644 --- a/cmd/job/submit.go +++ b/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 } diff --git a/cmd/job/submit_test.go b/cmd/job/submit_test.go index 63e3ceaf2f..9f5598d3f4 100644 --- a/cmd/job/submit_test.go +++ b/cmd/job/submit_test.go @@ -279,7 +279,7 @@ func setupSubmitTestEnv(t *testing.T) { gkeScheduler = "" platform = "linux/amd64" awaitJobCompletion = false - priorityClassName = "medium" + priority = "medium" isPathwaysJob = false pathways = orchestrator.PathwaysJobDefinition{MaxSliceRestarts: 1} gkeNapProvisioning = "" @@ -296,6 +296,8 @@ func setupSubmitTestEnv(t *testing.T) { gkeOrchestratorFactory = func() orchestrator.JobOrchestrator { return &mockOrchestrator{} } + mtcEnabled = false + mtcRamdiskDirectory = "" } type mockOrchestrator struct { @@ -1012,8 +1014,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) } } diff --git a/pathways_jobset.tmpl.orig b/pathways_jobset.tmpl.orig new file mode 100644 index 0000000000..0e856c04b7 --- /dev/null +++ b/pathways_jobset.tmpl.orig @@ -0,0 +1,382 @@ +apiVersion: jobset.x-k8s.io/v1alpha2 +kind: JobSet +metadata: + name: {{.WorkloadName}} + labels: + gcluster.google.com/workload: {{.WorkloadName}} + kueue.x-k8s.io/queue-name: {{.KueueQueueName}} + annotations: + jobset.sigs.k8s.io/hack: "true" +spec: + suspend: false + ttlSecondsAfterFinished: {{.TtlSecondsAfterFinished}} + network: + enableDNSHostnames: true + publishNotReadyAddresses: true + coordinator: + replicatedJob: pathways-head + startupPolicy: + startupPolicyOrder: InOrder +{{- if not .Pathways.Headless}} + successPolicy: + operator: All + targetReplicatedJobs: + - pathways-head +{{- end }} + failurePolicy: + maxRestarts: {{.MaxRestarts}} + restartStrategy: BlockingRecreate + rules: + - action: FailJobSet + onJobFailureReasons: + - PodFailurePolicy + replicatedJobs: + - name: pathways-head + replicas: 1 + template: + metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname + spec: + completionMode: Indexed + parallelism: 1 + completions: 1 + backoffLimit: 0 +{{- if .PodFailurePolicy }} + podFailurePolicy: +{{(StructuralData .PodFailurePolicy)}} +{{- end }} + template: + metadata: + annotations: + kueue.x-k8s.io/safe-to-forcefully-delete: "true" +{{- if .GCSFuseEnabled }} + gke-gcsfuse/volumes: "true" +{{- end }} + spec: + nodeSelector: + cloud.google.com/gke-nodepool: {{.Pathways.HeadNodePool}} + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + restartPolicy: Never +{{- if .PriorityClassName }} + priorityClassName: {{.PriorityClassName}} +{{- end }} +{{- if .ServiceAccountName }} + serviceAccountName: {{.ServiceAccountName}} +{{- end }} +{{- if .ImagePullSecrets }} + imagePullSecrets: +{{(StructuralData .ImagePullSecrets)}} +{{- end }} +{{- if not .Pathways.Headless}} + initContainers: +{{- else }} + containers: +{{- end }} + - name: pathways-proxy + image: {{.Pathways.ProxyServerImage}} + imagePullPolicy: Always + ports: + - containerPort: 29000 + args: + - --server_port=29000 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - "--gcs_scratch_location={{.Pathways.GCSLocation}}" + {{- if gt .Pathways.ElasticSlices 0}} + - --num_elastic_slices={{.Pathways.ElasticSlices}} + {{- end}} + {{- range .ProxyArgsList }} + - {{.}} + {{- end }} + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: ABSL_FLAGS + value: "--pathways_pipe_unreachable_timeout=60s" + {{- range $.PathwaysProxyEnv }} + - name: {{ .Name }} + value: {{ printf "%q" .Value }} + {{- end }} + {{- if not .Pathways.Headless}} + restartPolicy: Always + {{- end }} + resources: + limits: + cpu: "16" + memory: "100Gi" + - name: pathways-rm + image: {{.Pathways.ServerImage}} + imagePullPolicy: Always + ports: + - containerPort: 29001 + - containerPort: 29002 + args: + - --server_port=29001 + - "--gcs_scratch_location={{.Pathways.GCSLocation}}" + - --node_type=resource_manager + - --instance_count={{.NumSlices}} + - "--instance_type={{.PathwaysInstanceType}}" + {{- range .ServerArgsList }} + - {{.}} + {{- end }} + env: + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: HOST_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: TPU_SKIP_MDS_QUERY + value: "true" + - name: ABSL_FLAGS + value: "--pathways_pipe_unreachable_timeout=60s" + {{- range $.PathwaysServerEnv }} + - name: {{ .Name }} + value: {{ printf "%q" .Value }} + {{- end }} + {{- if not .Pathways.Headless}} + restartPolicy: Always + {{- end }} + resources: + limits: + cpu: "8" + memory: "32Gi" + {{- if not .Pathways.Headless}} + containers: + - name: workload-container + image: {{.FullImageName}} + imagePullPolicy: Always + securityContext: + privileged: true + resources: + limits: + cpu: "24" + memory: "100Gi" + requests: + cpu: "2" + memory: "8Gi" + env: + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: JAX_PLATFORMS + value: proxy + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: JAX_BACKEND_TARGET + value: grpc://$(PATHWAYS_HEAD):29000 + - name: ABSL_FLAGS + value: "--pathways_pipe_unreachable_timeout=60s" + {{- range $.Env }} + - name: {{ .Name }} + value: {{ printf "%q" .Value }} + {{- end }} + command: + - "/bin/bash" + - "-c" + - | + echo "GCluster Start: $(date)" + _sigterm() { + if [ -n "$PID" ]; then + kill -SIGTERM $PID 2>/dev/null + wait $PID + fi + exit 143 + } + trap _sigterm SIGTERM + ( + {{.CommandToRun}} + ) & PID=$! + wait $PID + EXIT_CODE=$? + echo "GCluster End: $(date)" + echo "Exit code: $EXIT_CODE" + exit $EXIT_CODE + volumeMounts: + - name: shared-tmp + mountPath: /tmp +{{- if .VolumeMountsYAML }} +{{(StructuralData .VolumeMountsYAML)}} +{{- end }} + {{- end}} + volumes: + - name: shared-tmp + hostPath: + path: /tmp + type: DirectoryOrCreate +{{- if .VolumesYAML }} +{{(StructuralData .VolumesYAML)}} +{{- end }} + - name: worker + replicas: {{.NumSlices}} + template: + metadata: + annotations: + alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool + spec: + completionMode: Indexed + parallelism: {{.NodesPerSlice}} + completions: {{.NodesPerSlice}} + backoffLimit: {{.WorkerBackoffLimit}} + backoffLimitPerIndex: 4000 + podReplacementPolicy: Failed + maxFailedIndexes: 0 +{{- if .PodFailurePolicy }} + podFailurePolicy: +{{(StructuralData .PodFailurePolicy)}} +{{- end }} + template: + metadata: + labels: + gcluster.google.com/workload: {{.WorkloadName}} + annotations: + kueue.x-k8s.io/safe-to-forcefully-delete: "true" + cloud.google.com/skip-tpu-webhook-check: "true" +{{- if .TopologyAnnotation }} +{{(StructuralData .TopologyAnnotation)}} +{{- end }} +{{- if .GCSFuseEnabled }} + gke-gcsfuse/volumes: "true" +{{- end }} + spec: + hostNetwork: true + dnsPolicy: ClusterFirstWithHostNet + restartPolicy: OnFailure +{{- if .ServiceAccountName }} + serviceAccountName: {{.ServiceAccountName}} +{{- end }} +{{- if .ImagePullSecrets }} + imagePullSecrets: +{{(StructuralData .ImagePullSecrets)}} +{{- end }} +{{- if .PriorityClassName }} + priorityClassName: {{.PriorityClassName}} +{{- end }} + terminationGracePeriodSeconds: {{.TerminationGracePeriodSeconds}} + {{- if .Pathways.ColocatedPythonSidecarImage}} + initContainers: + - name: colocated-python-sidecar + image: {{.Pathways.ColocatedPythonSidecarImage}} + imagePullPolicy: Always + restartPolicy: Always + ports: + - containerPort: 50051 + protocol: TCP + resources: {} + env: + - name: TCMALLOC_RELEASE_RATE + value: "10" + - name: GRPC_SERVER_ADDRESS + value: '0.0.0.0:50051' + - name: CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY + value: /tmp/sidecar + - name: PYTHONUNBUFFERED + value: "1" + volumeMounts: + - name: shared-tmp + mountPath: /tmp + {{- if .MTCEnabled}} + - name: cache + mountPath: {{.RamdiskDirectory}} + - name: sidecar-shared-memory + mountPath: /tmp/sidecar + {{- end}} + {{- end}} + containers: + - name: pathways-worker + image: {{.Pathways.WorkerImage}} + imagePullPolicy: Always + ports: + - containerPort: 29005 + - containerPort: 29006 + - containerPort: 8471 + - containerPort: 8080 + args: + - --server_port=29005 + - --resource_manager_address=$(PATHWAYS_HEAD):29001 + - "--gcs_scratch_location={{.Pathways.GCSLocation}}" + {{- if .MTCEnabled}} + - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar + {{- end}} + {{- range .WorkerArgsList }} + - {{.}} + {{- end }} + env: + {{- if .Verbose }} + - name: TPU_MIN_LOG_LEVEL + value: "0" + - name: TF_CPP_MIN_LOG_LEVEL + value: "0" + {{- end }} + - name: XCLOUD_ENVIRONMENT + value: GCP + - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER + value: "false" + - name: MEGASCALE_NUM_SLICES + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] + - name: JOBSET_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] + - name: REPLICATED_JOB_NAME + valueFrom: + fieldRef: + fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] + - name: MEGASCALE_SLICE_ID + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] + - name: PATHWAYS_HEAD + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: MEGASCALE_COORDINATOR_ADDRESS + valueFrom: + fieldRef: + fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] + - name: ABSL_FLAGS + value: "--pathways_pipe_unreachable_timeout=60s" + {{- range $.PathwaysWorkerEnv }} + - name: {{ .Name }} + value: {{ printf "%q" .Value }} + {{- end }} +{{(StructuralData .ResourcesString)}} + volumeMounts: + - name: shared-tmp + mountPath: /tmp +{{- if .VolumeMountsYAML }} +{{(StructuralData .VolumeMountsYAML)}} +{{- end }} + volumes: + - name: shared-tmp + hostPath: + path: /tmp + type: DirectoryOrCreate + +{{- if .VolumesYAML }} +{{(StructuralData .VolumesYAML)}} +{{- end }} +{{- if .NodeSelector }} + nodeSelector: +{{(StructuralData .NodeSelector)}} +{{- end }} +{{- if .Affinity }} + affinity: +{{(StructuralData .Affinity)}} +{{- end }} +{{- if .Tolerations }} + tolerations: +{{(StructuralData .Tolerations)}} +{{- end }} diff --git a/pkg/orchestrator/gke/gke_job_orchestrator.go b/pkg/orchestrator/gke/gke_job_orchestrator.go index 0f8e436e91..821e3b771a 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator.go @@ -99,12 +99,6 @@ func (g *GKEOrchestrator) SubmitJob(job orchestrator.JobDefinition) error { return err } - if job.Pathways.MTCEnabled { - if err := g.checkMTCAddonEnabled(job.ProjectID, job.ClusterLocation, job.ClusterName); err != nil { - return err - } - } - if err := g.fetchClusterState(&job); err != nil { return err } @@ -396,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.MTCEnabled && job.RamdiskDirectory == "" { - job.RamdiskDirectory = "/tmp/mtc_checkpoints" + if job.MTCEnabled && job.MTCRamdiskDirectory == "" { + job.MTCRamdiskDirectory = "/tmp/mtc_checkpoints" } tmpl, err := yamltemplate.New("pathways_jobset.tmpl").ParseFS(templatesFS, "templates/pathways_jobset.tmpl") @@ -528,12 +522,6 @@ 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 } @@ -576,6 +564,12 @@ func (g *GKEOrchestrator) initializeJobSubmission(job *orchestrator.JobDefinitio return err } + if job.MTCEnabled { + if err := g.checkMTCAddonEnabled(job.ProjectID, job.ClusterLocation, job.ClusterName); err != nil { + return err + } + } + return nil } @@ -1400,7 +1394,7 @@ func (g *GKEOrchestrator) prepareJobSetTemplateData(opts ManifestOptions, comman IsTPU: isTPU, IsGPU: isGPU, MTCEnabled: opts.MTCEnabled, - RamdiskDirectory: opts.RamdiskDirectory, + MTCRamdiskDirectory: opts.MTCRamdiskDirectory, } } diff --git a/pkg/orchestrator/gke/gke_job_orchestrator_test.go b/pkg/orchestrator/gke/gke_job_orchestrator_test.go index 353fed6169..8c00e2b586 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator_test.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator_test.go @@ -536,8 +536,7 @@ func TestGeneratePathwaysManifest_MTC(t *testing.T) { GCSLocation: "gs://my-bucket", HeadNodePool: "pathways-np", }, - IsPathwaysJob: true, - MTCEnabled: true, + MTCEnabled: true, } mockResponses := map[string][]shell.CommandResult{ diff --git a/pkg/orchestrator/gke/manifest_generator.go b/pkg/orchestrator/gke/manifest_generator.go index 00e62998c1..e05d04da43 100644 --- a/pkg/orchestrator/gke/manifest_generator.go +++ b/pkg/orchestrator/gke/manifest_generator.go @@ -168,7 +168,7 @@ func (g *GKEOrchestrator) PrepareManifestOptions(job orchestrator.JobDefinition, IsPathwaysJob: job.IsPathwaysJob, Pathways: job.Pathways, MTCEnabled: job.MTCEnabled, - RamdiskDirectory: job.RamdiskDirectory, + MTCRamdiskDirectory: job.MTCRamdiskDirectory, } if err := g.fillManifestStrings(&opts, schedOpts, job, isDynamicSlicing, isStaticSlicing, profile.IsCPUMachine); err != nil { diff --git a/pkg/orchestrator/gke/storage.go b/pkg/orchestrator/gke/storage.go index 6da5d95231..2c8bab6977 100644 --- a/pkg/orchestrator/gke/storage.go +++ b/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 } @@ -302,8 +303,12 @@ func (sm *StorageManager) AddVolumeOptions(opts *ManifestOptions, vols []MountIn } if opts.MTCEnabled { + ramdiskDir := opts.MTCRamdiskDirectory + if ramdiskDir == "" { + ramdiskDir = "/tmp/mtc_checkpoints" + } mountSpecs = append(mountSpecs, - map[string]interface{}{"name": "cache", "mountPath": opts.RamdiskDirectory}, + 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"}}, diff --git a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl index 0e856c04b7..10cdeccbd1 100644 --- a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl +++ b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl @@ -287,7 +287,7 @@ spec: mountPath: /tmp {{- if .MTCEnabled}} - name: cache - mountPath: {{.RamdiskDirectory}} + mountPath: {{.MTCRamdiskDirectory}} - name: sidecar-shared-memory mountPath: /tmp/sidecar {{- end}} @@ -356,6 +356,12 @@ spec: volumeMounts: - name: shared-tmp mountPath: /tmp + {{- if .MTCEnabled}} + - name: cache + mountPath: {{.MTCRamdiskDirectory}} + - name: sidecar-shared-memory + mountPath: /tmp/sidecar + {{- end}} {{- if .VolumeMountsYAML }} {{(StructuralData .VolumeMountsYAML)}} {{- end }} @@ -364,6 +370,14 @@ spec: hostPath: path: /tmp type: DirectoryOrCreate + {{- if .MTCEnabled}} + - name: cache + csi: + driver: multitier-checkpoint.csi.storage.gke.io + - name: sidecar-shared-memory + emptyDir: + medium: Memory + {{- end}} {{- if .VolumesYAML }} {{(StructuralData .VolumesYAML)}} diff --git a/pkg/orchestrator/gke/types.go b/pkg/orchestrator/gke/types.go index 147d79af10..0ebc2cea17 100644 --- a/pkg/orchestrator/gke/types.go +++ b/pkg/orchestrator/gke/types.go @@ -181,7 +181,7 @@ type ManifestOptions struct { Pathways orchestrator.PathwaysJobDefinition IsPathwaysJob bool MTCEnabled bool - RamdiskDirectory string + MTCRamdiskDirectory string Verbose bool Env map[string]string AdditionalManifests []string @@ -376,7 +376,7 @@ type jobSetTemplateData struct { IsTPU bool IsGPU bool MTCEnabled bool - RamdiskDirectory string + MTCRamdiskDirectory string } // Types for parsing kubectl get nodes -o json diff --git a/pkg/orchestrator/orchestrator.go b/pkg/orchestrator/orchestrator.go index b4e31792a7..fb2c85b9e8 100644 --- a/pkg/orchestrator/orchestrator.go +++ b/pkg/orchestrator/orchestrator.go @@ -38,11 +38,9 @@ type PathwaysJobDefinition struct { ServerEnv map[string]string WorkerEnv map[string]string - // Pathways-specific sidecars ColocatedPythonSidecarImage string // Default: "" HeadNodePool string // Resolved node pool to use for the Pathways head job. - } type VolumeDefinition struct { @@ -96,8 +94,8 @@ type JobDefinition struct { Pathways PathwaysJobDefinition // Embedded struct for Pathways-specific args // Multi-Tier Checkpointing (MTC) - MTCEnabled bool - RamdiskDirectory string + MTCEnabled bool + MTCRamdiskDirectory string RawMounts []string Env map[string]string From eb06469232a4f90d2fcc23bed6d9c3520f64c80d Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Thu, 23 Jul 2026 10:21:11 +0000 Subject: [PATCH 4/8] chore: remove redundant MTC volumes from pathways template MTC volumes are dynamically generated by storage.go and injected into VolumeMountsYAML and VolumesYAML for both pathways and non-pathways jobs. The sidecar initContainer still uses hardcoded mounts as it intentionally bypasses the global VolumeMountsYAML block. --- cmd/job/submit_test.go | 2 + pathways_jobset.tmpl.orig | 382 ------------------ pkg/orchestrator/gke/gke_job_orchestrator.go | 2 +- .../gke/gke_job_orchestrator_test.go | 3 +- .../gke/templates/pathways_jobset.tmpl | 15 - pkg/orchestrator/gke/types.go | 12 +- pkg/orchestrator/orchestrator.go | 1 + 7 files changed, 7 insertions(+), 410 deletions(-) delete mode 100644 pathways_jobset.tmpl.orig diff --git a/cmd/job/submit_test.go b/cmd/job/submit_test.go index 9f5598d3f4..41d3d1e7da 100644 --- a/cmd/job/submit_test.go +++ b/cmd/job/submit_test.go @@ -278,6 +278,8 @@ func setupSubmitTestEnv(t *testing.T) { topology = "" gkeScheduler = "" platform = "linux/amd64" + mtcEnabled = false + mtcRamdiskDirectory = "" awaitJobCompletion = false priority = "medium" isPathwaysJob = false diff --git a/pathways_jobset.tmpl.orig b/pathways_jobset.tmpl.orig deleted file mode 100644 index 0e856c04b7..0000000000 --- a/pathways_jobset.tmpl.orig +++ /dev/null @@ -1,382 +0,0 @@ -apiVersion: jobset.x-k8s.io/v1alpha2 -kind: JobSet -metadata: - name: {{.WorkloadName}} - labels: - gcluster.google.com/workload: {{.WorkloadName}} - kueue.x-k8s.io/queue-name: {{.KueueQueueName}} - annotations: - jobset.sigs.k8s.io/hack: "true" -spec: - suspend: false - ttlSecondsAfterFinished: {{.TtlSecondsAfterFinished}} - network: - enableDNSHostnames: true - publishNotReadyAddresses: true - coordinator: - replicatedJob: pathways-head - startupPolicy: - startupPolicyOrder: InOrder -{{- if not .Pathways.Headless}} - successPolicy: - operator: All - targetReplicatedJobs: - - pathways-head -{{- end }} - failurePolicy: - maxRestarts: {{.MaxRestarts}} - restartStrategy: BlockingRecreate - rules: - - action: FailJobSet - onJobFailureReasons: - - PodFailurePolicy - replicatedJobs: - - name: pathways-head - replicas: 1 - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: kubernetes.io/hostname - spec: - completionMode: Indexed - parallelism: 1 - completions: 1 - backoffLimit: 0 -{{- if .PodFailurePolicy }} - podFailurePolicy: -{{(StructuralData .PodFailurePolicy)}} -{{- end }} - template: - metadata: - annotations: - kueue.x-k8s.io/safe-to-forcefully-delete: "true" -{{- if .GCSFuseEnabled }} - gke-gcsfuse/volumes: "true" -{{- end }} - spec: - nodeSelector: - cloud.google.com/gke-nodepool: {{.Pathways.HeadNodePool}} - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - restartPolicy: Never -{{- if .PriorityClassName }} - priorityClassName: {{.PriorityClassName}} -{{- end }} -{{- if .ServiceAccountName }} - serviceAccountName: {{.ServiceAccountName}} -{{- end }} -{{- if .ImagePullSecrets }} - imagePullSecrets: -{{(StructuralData .ImagePullSecrets)}} -{{- end }} -{{- if not .Pathways.Headless}} - initContainers: -{{- else }} - containers: -{{- end }} - - name: pathways-proxy - image: {{.Pathways.ProxyServerImage}} - imagePullPolicy: Always - ports: - - containerPort: 29000 - args: - - --server_port=29000 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - "--gcs_scratch_location={{.Pathways.GCSLocation}}" - {{- if gt .Pathways.ElasticSlices 0}} - - --num_elastic_slices={{.Pathways.ElasticSlices}} - {{- end}} - {{- range .ProxyArgsList }} - - {{.}} - {{- end }} - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: ABSL_FLAGS - value: "--pathways_pipe_unreachable_timeout=60s" - {{- range $.PathwaysProxyEnv }} - - name: {{ .Name }} - value: {{ printf "%q" .Value }} - {{- end }} - {{- if not .Pathways.Headless}} - restartPolicy: Always - {{- end }} - resources: - limits: - cpu: "16" - memory: "100Gi" - - name: pathways-rm - image: {{.Pathways.ServerImage}} - imagePullPolicy: Always - ports: - - containerPort: 29001 - - containerPort: 29002 - args: - - --server_port=29001 - - "--gcs_scratch_location={{.Pathways.GCSLocation}}" - - --node_type=resource_manager - - --instance_count={{.NumSlices}} - - "--instance_type={{.PathwaysInstanceType}}" - {{- range .ServerArgsList }} - - {{.}} - {{- end }} - env: - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: HOST_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: TPU_SKIP_MDS_QUERY - value: "true" - - name: ABSL_FLAGS - value: "--pathways_pipe_unreachable_timeout=60s" - {{- range $.PathwaysServerEnv }} - - name: {{ .Name }} - value: {{ printf "%q" .Value }} - {{- end }} - {{- if not .Pathways.Headless}} - restartPolicy: Always - {{- end }} - resources: - limits: - cpu: "8" - memory: "32Gi" - {{- if not .Pathways.Headless}} - containers: - - name: workload-container - image: {{.FullImageName}} - imagePullPolicy: Always - securityContext: - privileged: true - resources: - limits: - cpu: "24" - memory: "100Gi" - requests: - cpu: "2" - memory: "8Gi" - env: - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: JAX_PLATFORMS - value: proxy - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: JAX_BACKEND_TARGET - value: grpc://$(PATHWAYS_HEAD):29000 - - name: ABSL_FLAGS - value: "--pathways_pipe_unreachable_timeout=60s" - {{- range $.Env }} - - name: {{ .Name }} - value: {{ printf "%q" .Value }} - {{- end }} - command: - - "/bin/bash" - - "-c" - - | - echo "GCluster Start: $(date)" - _sigterm() { - if [ -n "$PID" ]; then - kill -SIGTERM $PID 2>/dev/null - wait $PID - fi - exit 143 - } - trap _sigterm SIGTERM - ( - {{.CommandToRun}} - ) & PID=$! - wait $PID - EXIT_CODE=$? - echo "GCluster End: $(date)" - echo "Exit code: $EXIT_CODE" - exit $EXIT_CODE - volumeMounts: - - name: shared-tmp - mountPath: /tmp -{{- if .VolumeMountsYAML }} -{{(StructuralData .VolumeMountsYAML)}} -{{- end }} - {{- end}} - volumes: - - name: shared-tmp - hostPath: - path: /tmp - type: DirectoryOrCreate -{{- if .VolumesYAML }} -{{(StructuralData .VolumesYAML)}} -{{- end }} - - name: worker - replicas: {{.NumSlices}} - template: - metadata: - annotations: - alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool - spec: - completionMode: Indexed - parallelism: {{.NodesPerSlice}} - completions: {{.NodesPerSlice}} - backoffLimit: {{.WorkerBackoffLimit}} - backoffLimitPerIndex: 4000 - podReplacementPolicy: Failed - maxFailedIndexes: 0 -{{- if .PodFailurePolicy }} - podFailurePolicy: -{{(StructuralData .PodFailurePolicy)}} -{{- end }} - template: - metadata: - labels: - gcluster.google.com/workload: {{.WorkloadName}} - annotations: - kueue.x-k8s.io/safe-to-forcefully-delete: "true" - cloud.google.com/skip-tpu-webhook-check: "true" -{{- if .TopologyAnnotation }} -{{(StructuralData .TopologyAnnotation)}} -{{- end }} -{{- if .GCSFuseEnabled }} - gke-gcsfuse/volumes: "true" -{{- end }} - spec: - hostNetwork: true - dnsPolicy: ClusterFirstWithHostNet - restartPolicy: OnFailure -{{- if .ServiceAccountName }} - serviceAccountName: {{.ServiceAccountName}} -{{- end }} -{{- if .ImagePullSecrets }} - imagePullSecrets: -{{(StructuralData .ImagePullSecrets)}} -{{- end }} -{{- if .PriorityClassName }} - priorityClassName: {{.PriorityClassName}} -{{- end }} - terminationGracePeriodSeconds: {{.TerminationGracePeriodSeconds}} - {{- if .Pathways.ColocatedPythonSidecarImage}} - initContainers: - - name: colocated-python-sidecar - image: {{.Pathways.ColocatedPythonSidecarImage}} - imagePullPolicy: Always - restartPolicy: Always - ports: - - containerPort: 50051 - protocol: TCP - resources: {} - env: - - name: TCMALLOC_RELEASE_RATE - value: "10" - - name: GRPC_SERVER_ADDRESS - value: '0.0.0.0:50051' - - name: CLOUD_PATHWAYS_SIDECAR_SHM_DIRECTORY - value: /tmp/sidecar - - name: PYTHONUNBUFFERED - value: "1" - volumeMounts: - - name: shared-tmp - mountPath: /tmp - {{- if .MTCEnabled}} - - name: cache - mountPath: {{.RamdiskDirectory}} - - name: sidecar-shared-memory - mountPath: /tmp/sidecar - {{- end}} - {{- end}} - containers: - - name: pathways-worker - image: {{.Pathways.WorkerImage}} - imagePullPolicy: Always - ports: - - containerPort: 29005 - - containerPort: 29006 - - containerPort: 8471 - - containerPort: 8080 - args: - - --server_port=29005 - - --resource_manager_address=$(PATHWAYS_HEAD):29001 - - "--gcs_scratch_location={{.Pathways.GCSLocation}}" - {{- if .MTCEnabled}} - - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar - {{- end}} - {{- range .WorkerArgsList }} - - {{.}} - {{- end }} - env: - {{- if .Verbose }} - - name: TPU_MIN_LOG_LEVEL - value: "0" - - name: TF_CPP_MIN_LOG_LEVEL - value: "0" - {{- end }} - - name: XCLOUD_ENVIRONMENT - value: GCP - - name: MEGASCALE_GRPC_ENABLE_XOR_TRACER - value: "false" - - name: MEGASCALE_NUM_SLICES - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/replicatedjob-replicas'] - - name: JOBSET_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/jobset-name'] - - name: REPLICATED_JOB_NAME - valueFrom: - fieldRef: - fieldPath: metadata.annotations['jobset.sigs.k8s.io/replicatedjob-name'] - - name: MEGASCALE_SLICE_ID - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/job-index'] - - name: PATHWAYS_HEAD - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: MEGASCALE_COORDINATOR_ADDRESS - valueFrom: - fieldRef: - fieldPath: metadata.labels['jobset.sigs.k8s.io/coordinator'] - - name: ABSL_FLAGS - value: "--pathways_pipe_unreachable_timeout=60s" - {{- range $.PathwaysWorkerEnv }} - - name: {{ .Name }} - value: {{ printf "%q" .Value }} - {{- end }} -{{(StructuralData .ResourcesString)}} - volumeMounts: - - name: shared-tmp - mountPath: /tmp -{{- if .VolumeMountsYAML }} -{{(StructuralData .VolumeMountsYAML)}} -{{- end }} - volumes: - - name: shared-tmp - hostPath: - path: /tmp - type: DirectoryOrCreate - -{{- if .VolumesYAML }} -{{(StructuralData .VolumesYAML)}} -{{- end }} -{{- if .NodeSelector }} - nodeSelector: -{{(StructuralData .NodeSelector)}} -{{- end }} -{{- if .Affinity }} - affinity: -{{(StructuralData .Affinity)}} -{{- end }} -{{- if .Tolerations }} - tolerations: -{{(StructuralData .Tolerations)}} -{{- end }} diff --git a/pkg/orchestrator/gke/gke_job_orchestrator.go b/pkg/orchestrator/gke/gke_job_orchestrator.go index 821e3b771a..c49b5d96cd 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator.go @@ -149,7 +149,7 @@ func (g *GKEOrchestrator) checkMTCAddonEnabled(projectID, location, clusterName } 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 --pathways-mtc-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 } diff --git a/pkg/orchestrator/gke/gke_job_orchestrator_test.go b/pkg/orchestrator/gke/gke_job_orchestrator_test.go index 8c00e2b586..353fed6169 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator_test.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator_test.go @@ -536,7 +536,8 @@ func TestGeneratePathwaysManifest_MTC(t *testing.T) { GCSLocation: "gs://my-bucket", HeadNodePool: "pathways-np", }, - MTCEnabled: true, + IsPathwaysJob: true, + MTCEnabled: true, } mockResponses := map[string][]shell.CommandResult{ diff --git a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl index 10cdeccbd1..eed13dc26f 100644 --- a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl +++ b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl @@ -356,12 +356,6 @@ spec: volumeMounts: - name: shared-tmp mountPath: /tmp - {{- if .MTCEnabled}} - - name: cache - mountPath: {{.MTCRamdiskDirectory}} - - name: sidecar-shared-memory - mountPath: /tmp/sidecar - {{- end}} {{- if .VolumeMountsYAML }} {{(StructuralData .VolumeMountsYAML)}} {{- end }} @@ -370,15 +364,6 @@ spec: hostPath: path: /tmp type: DirectoryOrCreate - {{- if .MTCEnabled}} - - name: cache - csi: - driver: multitier-checkpoint.csi.storage.gke.io - - name: sidecar-shared-memory - emptyDir: - medium: Memory - {{- end}} - {{- if .VolumesYAML }} {{(StructuralData .VolumesYAML)}} {{- end }} diff --git a/pkg/orchestrator/gke/types.go b/pkg/orchestrator/gke/types.go index 0ebc2cea17..884832993b 100644 --- a/pkg/orchestrator/gke/types.go +++ b/pkg/orchestrator/gke/types.go @@ -283,7 +283,6 @@ type gkeCluster struct { NodePools []gkeJobNodePool `json:"nodePools"` Autoscaling gkeClusterAutoscaling `json:"autoscaling"` ControlPlaneEndpointsConfig *controlPlaneEndpointsConfig `json:"controlPlaneEndpointsConfig,omitempty"` - AddonsConfig *gkeAddonsConfig `json:"addonsConfig,omitempty"` } type controlPlaneEndpointsConfig struct { @@ -291,16 +290,7 @@ type controlPlaneEndpointsConfig struct { } type dnsEndpointConfig struct { - AllowExternalTraffic bool `json:"allowExternalTraffic,omitempty"` - AddonsConfig *gkeAddonsConfig `json:"addonsConfig,omitempty"` -} - -type gkeAddonsConfig struct { - StatefulHaConfig *gkeStatefulHaConfig `json:"statefulHaConfig,omitempty"` -} - -type gkeStatefulHaConfig struct { - Enabled bool `json:"enabled"` + AllowExternalTraffic bool `json:"allowExternalTraffic,omitempty"` } // Types for JobSet status unmarshaling diff --git a/pkg/orchestrator/orchestrator.go b/pkg/orchestrator/orchestrator.go index fb2c85b9e8..ad4d1224c2 100644 --- a/pkg/orchestrator/orchestrator.go +++ b/pkg/orchestrator/orchestrator.go @@ -38,6 +38,7 @@ type PathwaysJobDefinition struct { ServerEnv map[string]string WorkerEnv map[string]string + // Pathways-specific sidecars ColocatedPythonSidecarImage string // Default: "" HeadNodePool string // Resolved node pool to use for the Pathways head job. From c1fa7d3aca880cad9db206ca4f531a4a79660eb0 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Sat, 25 Jul 2026 08:29:39 +0000 Subject: [PATCH 5/8] Update MTC flag names and struct fields to use gke prefix --- cmd/job/submit.go | 12 +++++------ cmd/job/submit_test.go | 20 +++++++++---------- docs/gcluster_job_guide.md | 8 ++++---- pkg/orchestrator/gke/gke_job_orchestrator.go | 10 +++++----- .../gke/gke_job_orchestrator_test.go | 2 +- pkg/orchestrator/gke/manifest_generator.go | 4 ++-- pkg/orchestrator/gke/storage.go | 4 ++-- .../gke/templates/pathways_jobset.tmpl | 6 +++--- pkg/orchestrator/gke/types.go | 8 ++++---- pkg/orchestrator/orchestrator.go | 4 ++-- 10 files changed, 39 insertions(+), 39 deletions(-) diff --git a/cmd/job/submit.go b/cmd/job/submit.go index 0142de1617..2b3fa3aa44 100644 --- a/cmd/job/submit.go +++ b/cmd/job/submit.go @@ -67,8 +67,8 @@ var ( verbose bool volumeStr []string - mtcEnabled bool - mtcRamdiskDirectory string + gkeMtcEnabled bool + gkeMtcRamdiskDirectory string isPathwaysJob bool pathways orchestrator.PathwaysJobDefinition @@ -184,8 +184,8 @@ func init() { 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(&mtcRamdiskDirectory, "mtc-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") @@ -258,8 +258,8 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error { Timeout: timeout, PriorityClassName: priority, Verbose: verbose, - MTCEnabled: mtcEnabled, - MTCRamdiskDirectory: mtcRamdiskDirectory, + GKEMTCEnabled: gkeMtcEnabled, + GKEMTCRamdiskDirectory: gkeMtcRamdiskDirectory, GKENAPProvisioning: gkeNapProvisioning, GKENAPReservation: gkeNapReservation, IsPathwaysJob: isPathwaysJob, diff --git a/cmd/job/submit_test.go b/cmd/job/submit_test.go index 41d3d1e7da..5718abf209 100644 --- a/cmd/job/submit_test.go +++ b/cmd/job/submit_test.go @@ -278,8 +278,8 @@ func setupSubmitTestEnv(t *testing.T) { topology = "" gkeScheduler = "" platform = "linux/amd64" - mtcEnabled = false - mtcRamdiskDirectory = "" + gkeMtcEnabled = false + gkeMtcRamdiskDirectory = "" awaitJobCompletion = false priority = "medium" isPathwaysJob = false @@ -298,8 +298,8 @@ func setupSubmitTestEnv(t *testing.T) { gkeOrchestratorFactory = func() orchestrator.JobOrchestrator { return &mockOrchestrator{} } - mtcEnabled = false - mtcRamdiskDirectory = "" + gkeMtcEnabled = false + gkeMtcRamdiskDirectory = "" } type mockOrchestrator struct { @@ -1004,20 +1004,20 @@ func TestSubmitCmd_PathwaysMTCFlags(t *testing.T) { "--project", "test-project", "--pathways-gcs-location", "gs://my-bucket", "--compute-type", "n2-standard-4", - "--mtc-enabled", - "--mtc-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 !mtcEnabled { - t.Errorf("expected mtcEnabled to be true") + if !gkeMtcEnabled { + t.Errorf("expected gkeMtcEnabled to be true") } - if mtcRamdiskDirectory != "/tmp/custom_mtc_dir" { - t.Errorf("expected mtcRamdiskDirectory to be /tmp/custom_mtc_dir, got %s", mtcRamdiskDirectory) + if gkeMtcRamdiskDirectory != "/tmp/custom_mtc_dir" { + t.Errorf("expected gkeMtcRamdiskDirectory to be /tmp/custom_mtc_dir, got %s", gkeMtcRamdiskDirectory) } } diff --git a/docs/gcluster_job_guide.md b/docs/gcluster_job_guide.md index c1ad2b0819..d928550b12 100644 --- a/docs/gcluster_job_guide.md +++ b/docs/gcluster_job_guide.md @@ -1099,13 +1099,12 @@ 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 `:[:]` format. Examples of ``: `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. | | `--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)* @@ -1146,9 +1145,10 @@ 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`). | diff --git a/pkg/orchestrator/gke/gke_job_orchestrator.go b/pkg/orchestrator/gke/gke_job_orchestrator.go index c49b5d96cd..cf759c9b78 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator.go @@ -390,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.MTCEnabled && job.MTCRamdiskDirectory == "" { - job.MTCRamdiskDirectory = "/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") @@ -564,7 +564,7 @@ func (g *GKEOrchestrator) initializeJobSubmission(job *orchestrator.JobDefinitio return err } - if job.MTCEnabled { + if job.GKEMTCEnabled { if err := g.checkMTCAddonEnabled(job.ProjectID, job.ClusterLocation, job.ClusterName); err != nil { return err } @@ -1393,8 +1393,8 @@ func (g *GKEOrchestrator) prepareJobSetTemplateData(opts ManifestOptions, comman PathwaysWorkerEnv: sortedEnvVars(opts.Pathways.WorkerEnv), IsTPU: isTPU, IsGPU: isGPU, - MTCEnabled: opts.MTCEnabled, - MTCRamdiskDirectory: opts.MTCRamdiskDirectory, + GKEMTCEnabled: opts.GKEMTCEnabled, + GKEMTCRamdiskDirectory: opts.GKEMTCRamdiskDirectory, } } diff --git a/pkg/orchestrator/gke/gke_job_orchestrator_test.go b/pkg/orchestrator/gke/gke_job_orchestrator_test.go index 353fed6169..f8b4c9da68 100644 --- a/pkg/orchestrator/gke/gke_job_orchestrator_test.go +++ b/pkg/orchestrator/gke/gke_job_orchestrator_test.go @@ -537,7 +537,7 @@ func TestGeneratePathwaysManifest_MTC(t *testing.T) { HeadNodePool: "pathways-np", }, IsPathwaysJob: true, - MTCEnabled: true, + GKEMTCEnabled: true, } mockResponses := map[string][]shell.CommandResult{ diff --git a/pkg/orchestrator/gke/manifest_generator.go b/pkg/orchestrator/gke/manifest_generator.go index e05d04da43..4b44761b30 100644 --- a/pkg/orchestrator/gke/manifest_generator.go +++ b/pkg/orchestrator/gke/manifest_generator.go @@ -167,8 +167,8 @@ func (g *GKEOrchestrator) PrepareManifestOptions(job orchestrator.JobDefinition, Env: job.Env, IsPathwaysJob: job.IsPathwaysJob, Pathways: job.Pathways, - MTCEnabled: job.MTCEnabled, - MTCRamdiskDirectory: job.MTCRamdiskDirectory, + GKEMTCEnabled: job.GKEMTCEnabled, + GKEMTCRamdiskDirectory: job.GKEMTCRamdiskDirectory, } if err := g.fillManifestStrings(&opts, schedOpts, job, isDynamicSlicing, isStaticSlicing, profile.IsCPUMachine); err != nil { diff --git a/pkg/orchestrator/gke/storage.go b/pkg/orchestrator/gke/storage.go index 2c8bab6977..812daf0449 100644 --- a/pkg/orchestrator/gke/storage.go +++ b/pkg/orchestrator/gke/storage.go @@ -302,8 +302,8 @@ func (sm *StorageManager) AddVolumeOptions(opts *ManifestOptions, vols []MountIn } } - if opts.MTCEnabled { - ramdiskDir := opts.MTCRamdiskDirectory + if opts.GKEMTCEnabled { + ramdiskDir := opts.GKEMTCRamdiskDirectory if ramdiskDir == "" { ramdiskDir = "/tmp/mtc_checkpoints" } diff --git a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl index eed13dc26f..b8674c3a02 100644 --- a/pkg/orchestrator/gke/templates/pathways_jobset.tmpl +++ b/pkg/orchestrator/gke/templates/pathways_jobset.tmpl @@ -285,9 +285,9 @@ spec: volumeMounts: - name: shared-tmp mountPath: /tmp - {{- if .MTCEnabled}} + {{- if .GKEMTCEnabled}} - name: cache - mountPath: {{.MTCRamdiskDirectory}} + mountPath: {{.GKEMTCRamdiskDirectory}} - 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 .MTCEnabled}} + {{- if .GKEMTCEnabled}} - --cloud_pathways_sidecar_shm_directory=/tmp/sidecar {{- end}} {{- range .WorkerArgsList }} diff --git a/pkg/orchestrator/gke/types.go b/pkg/orchestrator/gke/types.go index 884832993b..1da236a070 100644 --- a/pkg/orchestrator/gke/types.go +++ b/pkg/orchestrator/gke/types.go @@ -180,8 +180,8 @@ type ManifestOptions struct { IsCPUMachine bool Pathways orchestrator.PathwaysJobDefinition IsPathwaysJob bool - MTCEnabled bool - MTCRamdiskDirectory string + GKEMTCEnabled bool + GKEMTCRamdiskDirectory string Verbose bool Env map[string]string AdditionalManifests []string @@ -365,8 +365,8 @@ type jobSetTemplateData struct { PathwaysWorkerEnv []EnvVar IsTPU bool IsGPU bool - MTCEnabled bool - MTCRamdiskDirectory string + GKEMTCEnabled bool + GKEMTCRamdiskDirectory string } // Types for parsing kubectl get nodes -o json diff --git a/pkg/orchestrator/orchestrator.go b/pkg/orchestrator/orchestrator.go index ad4d1224c2..a3c2e9ae3c 100644 --- a/pkg/orchestrator/orchestrator.go +++ b/pkg/orchestrator/orchestrator.go @@ -95,8 +95,8 @@ type JobDefinition struct { Pathways PathwaysJobDefinition // Embedded struct for Pathways-specific args // Multi-Tier Checkpointing (MTC) - MTCEnabled bool - MTCRamdiskDirectory string + GKEMTCEnabled bool + GKEMTCRamdiskDirectory string RawMounts []string Env map[string]string From 246acf095a4fb24fdf04eca3a06e6dc715db474b Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Sat, 25 Jul 2026 08:33:54 +0000 Subject: [PATCH 6/8] Update MTC prerequisite docs to be generic --- docs/gcluster_job_guide.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/gcluster_job_guide.md b/docs/gcluster_job_guide.md index d928550b12..e40bd9ab56 100644 --- a/docs/gcluster_job_guide.md +++ b/docs/gcluster_job_guide.md @@ -31,9 +31,9 @@ If you use `--build-context` to build images on-the-fly, you must set: ### 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 }`). +If you plan to use Multi-Tier Checkpointing (`--gke-mtc-enabled` flag), ensure that the StateHA / Multi-Tier Checkpointing feature is enabled on your GKE cluster. -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. +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. ## 2. Prepare Sample Application Code From 12281e787563b1211599ca466b9b5cd8981c3b70 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Sat, 25 Jul 2026 08:36:22 +0000 Subject: [PATCH 7/8] Update docs to add gke-nap-provisioning and gke-nap-reservation flags to the GKE & Advanced Orchestration Flags table --- docs/gcluster_job_guide.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/gcluster_job_guide.md b/docs/gcluster_job_guide.md index e40bd9ab56..6b4942c7b6 100644 --- a/docs/gcluster_job_guide.md +++ b/docs/gcluster_job_guide.md @@ -31,7 +31,7 @@ If you use `--build-context` to build images on-the-fly, you must set: ### 1.1 Multi-Tier Checkpointing (MTC) Prerequisites -If you plan to use Multi-Tier Checkpointing (`--gke-mtc-enabled` flag), ensure that the StateHA / Multi-Tier Checkpointing feature is enabled on your GKE cluster. +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. @@ -1156,6 +1156,8 @@ The `gcluster job submit` command deploys a container image as a job (Kubernetes | `--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.* From e4f7537be1f3d5341842759f1b45a0455b41ddd1 Mon Sep 17 00:00:00 2001 From: Neelabh94 Date: Sat, 25 Jul 2026 08:38:59 +0000 Subject: [PATCH 8/8] Update docs to add link to official GKE MTC documentation --- docs/gcluster_job_guide.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/gcluster_job_guide.md b/docs/gcluster_job_guide.md index 6b4942c7b6..013338b92e 100644 --- a/docs/gcluster_job_guide.md +++ b/docs/gcluster_job_guide.md @@ -33,7 +33,7 @@ If you use `--build-context` to build images on-the-fly, you must set: 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. +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