Release candidate: v1.91.0 (GoogleCloudPlatform#5651)

Author: sudheer-quad

.github/workflows/cache-metadata.yml

@@ -439,6 +439,16 @@ it is not supported for Toolkit development due to GNU specific shell scripts.
 If developing on a mac, a workaround is to install GNU tooling by installing
 `coreutils` and `findutils` from a package manager such as homebrew or conda.
 
+## Privacy notice
+
+To help improve Cluster Toolkit, feature usage statistics are collected and sent to Google. You can opt-out at any time by executing the following command:
+
+```shell
+./gcluster telemetry off
+```
+
+Cluster Toolkit telemetry overall is handled in accordance with the [Google Privacy Policy](https://policies.google.com/privacy). When you use Cluster Toolkit to interact with or utilize GCP Services, your information is handled in accordance with the [Google Cloud Privacy Notice](https://cloud.google.com/terms/cloud-privacy-notice).
+
 ### Contributing
 
 Please refer to the [contributing file](CONTRIBUTING.md) in our GitHub

README.md

@@ -188,6 +188,42 @@ func checkK8sDependencies(state *PrereqState, missing *[]missingPrereq) {
 	}
 }
 
+// isDockerCredsConfigured checks if Docker is configured to use gcloud credentials for the required registries.
+func isDockerCredsConfigured(region string) bool {
+	configDir := os.Getenv("DOCKER_CONFIG")
+	if configDir == "" {
+		homeDir, err := os.UserHomeDir()
+		if err != nil {
+			return false
+		}
+		configDir = filepath.Join(homeDir, ".docker")
+	}
+	configPath := filepath.Join(configDir, "config.json")
+	data, err := os.ReadFile(configPath)
+	if err != nil {
+		return false
+	}
+
+	var config struct {
+		CredHelpers map[string]string `json:"credHelpers"`
+		CredsStore  string            `json:"credsStore"`
+	}
+	if err := json.Unmarshal(data, &config); err != nil {
+		logging.Error("Failed to unmarshal Docker config from %s: %v", configPath, err)
+		return false
+	}
+
+	if config.CredsStore == "gcloud" {
+		return true
+	}
+
+	if config.CredHelpers["gcr.io"] != "gcloud" {
+		return false
+	}
+	pkgDevReg := fmt.Sprintf("%s-docker.pkg.dev", region)
+	return config.CredHelpers[pkgDevReg] == "gcloud"
+}
+
 // EnsurePrerequisites checks all necessary gcloud and kubectl prerequisites.
 func ensurePrerequisites(cmd *cobra.Command, projectID *string, location string) error {
 	if dryRunManifest != "" {
@@ -228,15 +264,17 @@ func ensurePrerequisites(cmd *cobra.Command, projectID *string, location string)
 	checkK8sDependencies(&state, &missing)
 
 	// Check Docker creds
-	if !state.DockerCredsConfigured {
-		region := shell.ExtractRegion(location)
+	region := shell.ExtractRegion(location)
+	if !isDockerCredsConfigured(region) {
+		cmds := []string{"gcloud auth configure-docker gcr.io --quiet"}
+		if region != "" {
+			cmds = append(cmds, fmt.Sprintf("gcloud auth configure-docker %s-docker.pkg.dev --quiet", region))
+		}
 		missing = append(missing, missingPrereq{
-			name: "Docker Credentials",
-			commands: []string{
-				"gcloud auth configure-docker gcr.io --quiet",
-				fmt.Sprintf("gcloud auth configure-docker %s-docker.pkg.dev --quiet", region),
-			},
+			name:     "Docker Credentials",
+			commands: cmds,
 		})
+	} else {
 		state.DockerCredsConfigured = true
 	}
 

cmd/job/prereq.go

@@ -320,11 +320,105 @@ func TestEnsureApplicationDefaultCredentials_Failure(t *testing.T) {
 	}
 }
 
+func TestIsDockerCredsConfigured(t *testing.T) {
+	tempDir := t.TempDir()
+	t.Setenv("HOME", tempDir)
+
+	dockerDir := filepath.Join(tempDir, ".docker")
+	configPath := filepath.Join(dockerDir, "config.json")
+
+	tests := []struct {
+		name          string
+		configContent string
+		shouldExist   bool
+		region        string
+		expected      bool
+	}{
+		{
+			name:        "No config file",
+			shouldExist: false,
+			region:      "us-central1",
+			expected:    false,
+		},
+		{
+			name:          "Corrupted config file",
+			configContent: "invalid",
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      false,
+		},
+		{
+			name:          "CredsStore set to gcloud",
+			configContent: `{"credsStore": "gcloud"}`,
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      true,
+		},
+		{
+			name:          "CredHelpers configured correctly",
+			configContent: `{"credHelpers": {"gcr.io": "gcloud", "us-central1-docker.pkg.dev": "gcloud"}}`,
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      true,
+		},
+		{
+			name:          "CredHelpers missing gcr.io",
+			configContent: `{"credHelpers": {"us-central1-docker.pkg.dev": "gcloud"}}`,
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      false,
+		},
+		{
+			name:          "CredHelpers missing regional registry",
+			configContent: `{"credHelpers": {"gcr.io": "gcloud"}}`,
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      false,
+		},
+		{
+			name:          "CredHelpers wrong helper",
+			configContent: `{"credHelpers": {"gcr.io": "desktop", "us-central1-docker.pkg.dev": "gcloud"}}`,
+			shouldExist:   true,
+			region:        "us-central1",
+			expected:      false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if !tt.shouldExist {
+				os.Remove(configPath)
+			} else {
+				if err := os.MkdirAll(dockerDir, 0755); err != nil {
+					t.Fatal(err)
+				}
+				if err := os.WriteFile(configPath, []byte(tt.configContent), 0644); err != nil {
+					t.Fatal(err)
+				}
+			}
+
+			got := isDockerCredsConfigured(tt.region)
+			if got != tt.expected {
+				t.Errorf("isDockerCredsConfigured() = %v, want %v", got, tt.expected)
+			}
+		})
+	}
+}
+
 func TestEnsurePrerequisites_DockerCreds(t *testing.T) {
+	tempDir := t.TempDir()
+	t.Setenv("HOME", tempDir)
+
 	origExecuteCommand := shell.ExecuteCommand
 	defer func() { shell.ExecuteCommand = origExecuteCommand }()
 
 	shell.ExecuteCommand = func(name string, args ...string) shell.CommandResult {
+		if name == "gcloud" && len(args) > 1 && args[0] == "auth" && args[1] == "list" {
+			return shell.CommandResult{ExitCode: 0, Stdout: "user@example.com"}
+		}
+		if name == "gcloud" && len(args) > 1 && args[0] == "services" && args[1] == "list" {
+			return shell.CommandResult{ExitCode: 0, Stdout: "artifactregistry.googleapis.com"}
+		}
 		return shell.CommandResult{ExitCode: 0}
 	}
 

cmd/job/prereq_test.go

@@ -16,6 +16,7 @@ package job
 
 import (
 	"fmt"
+	"hpc-toolkit/pkg/config"
 	"os"
 	"slices"
 	"time"
@@ -28,21 +29,23 @@ import (
 )
 
 var (
-	imageName       string
-	baseImage       string
-	buildContext    string
-	commandToRun    string
-	acceleratorType string
-	dryRunManifest  string
+	imageName      string
+	baseImage      string
+	buildContext   string
+	commandToRun   string
+	computeType    string
+	dryRunManifest string
 
 	workloadName     string
 	kueueQueueName   string
-	numSlicesOrNodes int
-	vmsPerSlice      int
-	maxRestarts      int
+	numNodes         int
+	numSlices        int
+	restarts         int
 	ttlAfterFinished string
 	gracePeriodStr   string
 
+	gkeDisableParallelContainers bool
+
 	placementPolicy string
 	nodeConstraint  map[string]string
 
@@ -104,17 +107,18 @@ func init() {
 	SubmitCmd.Flags().StringVarP(&baseImage, "base-image", "B", "", "Name of the base image for Crane to build upon (e.g., python:3.9-slim). Requires --build-context.")
 	SubmitCmd.Flags().StringVarP(&buildContext, "build-context", "b", "", "Path to the build context directory for Crane (e.g., .). Required with --base-image.")
 	SubmitCmd.Flags().StringVarP(&commandToRun, "command", "e", "", "Command to execute in the container (e.g., 'python train.py'). Required.")
-	SubmitCmd.Flags().StringVarP(&acceleratorType, "accelerator", "a", "", "Type of accelerator to request (e.g., 'nvidia-tesla-a100').")
+	SubmitCmd.Flags().StringVar(&computeType, "compute-type", "", "Type of compute to request (e.g., 'n2-standard-32', 'nvidia-l4', 'v6e-8').")
 	SubmitCmd.Flags().StringVarP(&dryRunManifest, "dry-run-out", "o", "", "Path to output the generated Kubernetes manifest instead of applying it.")
 	SubmitCmd.Flags().StringVarP(&platform, "platform", "f", "linux/amd64", "Target platform for the image build (e.g., 'linux/amd64', 'linux/arm64'). Used with --base-image.")
 
 	SubmitCmd.Flags().StringVarP(&workloadName, "name", "n", "", "Name of the workload to create. Required.")
 	SubmitCmd.Flags().StringVarP(&kueueQueueName, "queue", "q", "", "Name of the Kueue LocalQueue to submit the workload to. If empty, it will be auto-discovered.")
-	SubmitCmd.Flags().IntVar(&numSlicesOrNodes, "nodes", 1, "Number of JobSet replicas (or Slices for TPUs).")
-	SubmitCmd.Flags().IntVar(&vmsPerSlice, "vms-per-slice", 0, "Number of VMs (pods) per slice. Defaults to auto-calculated value for TPUs.")
-	SubmitCmd.Flags().IntVar(&maxRestarts, "max-restarts", 1, "Maximum number of restarts for the JobSet before failing.")
+	SubmitCmd.Flags().IntVar(&numNodes, "num-nodes", 1, "The number of nodes to use per group/slice. Defaults to 1 for CPU/GPU, or auto-calculated for TPUs.")
+	SubmitCmd.Flags().IntVar(&numSlices, "num-slices", 1, "The number of independent groups/slices to use.")
+	SubmitCmd.Flags().IntVar(&restarts, "restarts", 1, "Maximum number of restarts for the JobSet before failing.")
 	SubmitCmd.Flags().StringVar(&ttlAfterFinished, "gke-ttl-after-finished", "1h", "Time to retain the JobSet after it finishes (e.g. 5m, 1h).")
 	SubmitCmd.Flags().StringVar(&gracePeriodStr, "grace-period", "30s", "Time to wait before forcefully terminating a pod (e.g. 30s, 2m). Gives the workload time to save checkpoints or clean up distributed state during cancellation or preemption events (like Spot VM evictions).")
+	SubmitCmd.Flags().BoolVar(&gkeDisableParallelContainers, "gke-disable-parallel-containers", false, "Disable parallel containers for TPU v7/v7x on GKE.")
 
 	SubmitCmd.Flags().StringVar(&placementPolicy, "placement-policy", "", "Name of the GKE placement policy to use.")
 	SubmitCmd.Flags().StringToStringVar(&nodeConstraint, "node-constraint", nil, "Key=value pairs for node labels to target specific nodes. Maps to nodeSelector in GKE, and to SLURM's --constraint.")
@@ -147,10 +151,13 @@ func init() {
 
 	_ = SubmitCmd.MarkFlagRequired("command")
 	_ = SubmitCmd.MarkFlagRequired("name")
-	_ = SubmitCmd.MarkFlagRequired("accelerator")
+	_ = SubmitCmd.MarkFlagRequired("compute-type")
 }
 
 func runSubmitCmd(cmd *cobra.Command, args []string) error {
+	if err := config.ValidateHardwareRequest(computeType, topology, placementPolicy); err != nil {
+		return err
+	}
 	ttlSeconds, err := parseDurationToSeconds(ttlAfterFinished, "--gke-ttl-after-finished")
 	if err != nil {
 		return err
@@ -175,22 +182,26 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		awaitJobCompletion = true
 	}
 
+	if config.IsTPU(computeType) && cmd.Flags().Changed("num-nodes") {
+		return fmt.Errorf("--num-nodes cannot be used with TPU jobs (it is calculated automatically from topology)")
+	}
+
 	jobDef := orchestrator.JobDefinition{
 		ImageName:                     imageName,
 		BaseImage:                     baseImage,
 		BuildContext:                  buildContext,
 		Platform:                      platform,
 		CommandToRun:                  commandToRun,
-		AcceleratorType:               acceleratorType,
+		ComputeType:                   computeType,
 		DryRunManifest:                dryRunManifest,
 		ProjectID:                     projectID,
 		ClusterName:                   clusterName,
 		ClusterLocation:               location,
 		WorkloadName:                  workloadName,
 		KueueQueueName:                kueueQueueName,
-		NumSlices:                     numSlicesOrNodes,
-		VmsPerSlice:                   vmsPerSlice,
-		MaxRestarts:                   maxRestarts,
+		NumSlices:                     numSlices,
+		NodesPerSlice:                 numNodes,
+		MaxRestarts:                   restarts,
 		TtlSecondsAfterFinished:       ttlSeconds,
 		TerminationGracePeriodSeconds: gracePeriodSeconds,
 		PlacementPolicy:               placementPolicy,
@@ -202,6 +213,7 @@ func runSubmitCmd(cmd *cobra.Command, args []string) error {
 		Topology:                      topology,
 		GKEScheduler:                  gkeScheduler,
 		AwaitJobCompletion:            awaitJobCompletion,
+		UseParallelContainers:         !gkeDisableParallelContainers,
 		Timeout:                       timeoutStr,
 		PriorityClassName:             priorityClassName,
 		IsPathwaysJob:                 isPathwaysJob,