Update images, add annotation. Update sample config YAML to use Python image.

Author: RoshaniN

README.md

@@ -19,7 +19,6 @@ The user workload is typically on a Vertex AI notebook, so users can connect to
 - docker version 17.03+.
 - kubectl version v1.11.3+.
 - Access to a Kubernetes v1.11.3+ cluster.
-- JobSet //ToDo(roshanin) install JobSet
 
 ### To Deploy on the cluster
 **Build and push your image to the location specified by `IMG`:**

api/v1/pathwaysjob_types.go

@@ -61,12 +61,13 @@ type PathwaysJobSpec struct {
 	// Maximum number of times the JobSet is restarted.
 	MaxRestarts int32 `json:"maxRestarts,omitempty"`
 
-	// PathwaysDir is a persistent location like GCS at which temporary
+	// PathwaysDir is a persistent GCS location at which temporary
 	// Pathways artifacts can be stored like HBM state during interruptions.
 	// Currently, Pathways supports a precreated GCS directory only.
 	PathwaysDir string `json:"pathwaysDir,omitempty"`
 
-	// PathwaysVersion is the version of the Pathways client.
+	// PathwaysVersion is the version of the Pathways cluster.
+	// This indicates the version of the Pathways RM, Proxy and Workers.
 	PathwaysVersion string `json:"pathwaysVersion,omitempty"`
 
 	// The list of worker types created for the Pathways Job. Currently only

config/crd/bases/pathways-job.pathways.domain_pathwaysjobs.yaml

@@ -8058,12 +8058,14 @@ spec:
                 type: integer
               pathwaysDir:
                 description: |-
-                  PathwaysDir is a persistent location like GCS at which temporary
+                  PathwaysDir is a persistent GCS location at which temporary
                   Pathways artifacts can be stored like HBM state during interruptions.
                   Currently, Pathways supports a precreated GCS directory only.
                 type: string
               pathwaysVersion:
-                description: PathwaysVersion is the version of the Pathways client.
+                description: |-
+                  PathwaysVersion is the version of the Pathways cluster.
+                  This indicates the version of the Pathways RM, Proxy and Workers.
                 type: string
               workers:
                 description: |-

config/manager/kustomization.yaml

@@ -18,5 +18,5 @@ apiVersion: kustomize.config.k8s.io/v1beta1
 kind: Kustomization
 images:
 - name: controller
-  newName: us-docker.pkg.dev/cloud-tpu-multipod-dev/pathways/pathwaysjob
+  newName: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/pathwaysjob
   newTag: latest

config/samples/jobset_example.yaml

@@ -14,5 +14,5 @@
 
 ## Append samples of your project ##
 resources:
-- pathways-api_v1_pathwaysjob.yaml
+- pathways-job_v1_pathwaysjob.yaml
 # +kubebuilder:scaffold:manifestskustomizesamples

config/samples/kustomization.yaml

@@ -15,7 +15,7 @@
 apiVersion: pathways-job.pathways.domain/v1
 kind: PathwaysJob
 metadata:
-  name: pathways-trial62
+  name: pathways-trial
 spec:
   maxRestarts: 10
   workers:
@@ -37,85 +37,14 @@ spec:
           - name: JAX_PLATFORMS
             value: proxy
           - name: JAX_BACKEND_TARGET
-            value: grpc://pathways-trial62-proxy-0-0.pathways-trial62:29008
-          image: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
+            value: grpc://pathways-trial-proxy-0-0.pathways-trial:29008
+          image: python:3.13
           imagePullPolicy: Always
           command:
-            - bash
-            - -c
-            - |
-              (python3 MaxText/train.py MaxText/configs/base.yml base_output_directory=gs://cloud-pathways-staging dataset_path=gs://maxtext-dataset/ steps=10 run_name=roshanin-pathways1 enable_single_controller=true attention=dot_product monitor_goodput=False enable_tensorboard=True enable_checkpointing=False);
-          volumeMounts:
-          - mountPath: /tmp
-            name: shared-tmp
-          resources:
-            limits:
-              cpu: "20"
-              memory: 90G
-
-
-
-    # #Pod template for inference, colocate mode.
-
-
-    # deploymentMode: colocate
-    # template: # UserPodTemplate
-    #   spec:
-    #     containers:
-    #     - name: jetstream
-    #       image: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
-    #       imagePullPolicy: Always
-    #       ports:
-    #       - containerPort: 9000
-    #       env:
-    #       - name: XCLOUD_ENVIRONMENT
-    #         value: GCP
-    #       - name: JAX_PLATFORMS
-    #         value: proxy
-    #       - name: JAX_BACKEND_TARGET
-    #         value: grpc://pathways-trial61-leader-0-0.pathways-trial61:29008
-    #       command:
-    #         - bash
-    #         - -c
-    #         - 'echo Start: $(date);
-    #           _sigterm() ( kill -SIGTERM $! 2>/dev/null;);
-    #           trap _sigterm SIGTERM;
-    #           (JAX_TRACEBACK_FILTERING=off python3 MaxText/maxengine_server.py
-    #           MaxText/configs/inference_jetstream.yml tokenizer_path=assets/tokenizer.llama2
-    #           load_parameters_path=gs://runner-maxtext-logs/2024-05-07-23-34/unscanned_chkpt/checkpoints/0/items
-    #           max_prefill_predict_length=1024 max_target_length=2048 async_checkpointing=false
-    #           model_name=''llama2-70b'' steps=1 ici_fsdp_parallelism=1 ici_autoregressive_parallelism=-1
-    #           ici_tensor_parallelism=1 scan_layers=false weight_dtype=bfloat16
-    #           per_device_batch_size=2) & PID=$!;
-    #           while kill -0 $PID 2>/dev/null;
-    #           do sleep 5;
-    #           done;
-    #           wait $PID;
-    #           EXIT_CODE=$?
-    #           echo EXIT_CODE=$EXIT_CODE;
-    #           echo End sleep: $(date);
-    #           sleep infinity;'
-    #     - name: tester
-    #       image: us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/maxtext_jax_stable:latest
-    #       imagePullPolicy: Always
-    #       command:
-    #       - bash
-    #       - -c
-    #       - 'echo Start: $(date);
-    #         _sigterm() ( kill -SIGTERM $! 2>/dev/null;);
-    #         trap _sigterm SIGTERM;
-    #         for i in {1..5}; do
-    #           echo Sending request $i;
-    #           time python3 JetStream/jetstream/tools/requester.py --tokenizer assets/tokenizer.llama2 --max_tokens=16 --server=0.0.0.0 --text=\"why earth is round\";
-    #           EXIT_CODE=$?;
-    #           echo Completed request;
-    #           echo EXIT_CODE=$EXIT_CODE;
-    #           if [[ $EXIT_CODE -ne 0 ]]; then
-    #             break;
-    #           fi;
-    #         done;
-    #         echo Last EXIT_CODE=$EXIT_CODE;
-    #         echo End sleep: $(date);
-    #         sleep infinity;'
-    #       securityContext:
-    #         privileged: true
+          - /bin/sh
+          - -c
+          - |
+            pip install --upgrade pip
+            pip install -U --pre jax jaxlib -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html
+            pip install pathwaysutils
+            python -c "import jax; import pathwaysutils; print(\"Number of JAX devices is\", len(jax.devices()))"

config/samples/pathways-job_v1_pathwaysjob.yaml

@@ -58,7 +58,7 @@ var (
 var WorkerTypeToTPUVersionMap = map[string]string{
 	"tpu-v6e-slice":        "tpuv6e",
 	"tpu-v5p-slice":        "tpuv5",
-	"tpu-v5-lite-podslice": "tpuv5",
+	"tpu-v5-lite-podslice": "tpuv5e",
 	"tpu-v4-podslice":      "tpuv4",
 }
 
@@ -278,17 +278,15 @@ func validateTPUTopologyWithWorkerType(ctx context.Context, tpuGKEAcceleratorTyp
 // Calculate the number of VMs based on the Topology (- used in completions/parallelisms)
 func calculateVMsFromTopology(topology string) int32 {
 	parts := strings.Split(topology, "x") // Examples - 2x2x4 or 4x4
-	if len(parts) < 2 {
-		return 0
-	}
 	// Calculate the number of chips based on the Topology.
+	// The topology must have already been validated with the worker type.
 	chips := 1
 	for _, part := range parts {
 		num, _ := strconv.Atoi(part)
 		chips *= num
 	}
 	vms := 1
-	chipsperVM := 4
+	chipsperVM := 4 // ToDo (roshanin): Add support for VMs with 8 chips per host.
 	if chips >= chipsperVM {
 		vms = chips / chipsperVM
 	}
@@ -308,13 +306,24 @@ func calculateTPUInfo(ctx context.Context, pw *pathwaysjob.PathwaysJob) error {
 	return nil
 }
 
+// Construct image tag based on Pathways version
+func makeImageTagUsingPathwaysVersion(pw *pathwaysjob.PathwaysJob) string {
+	var tag string
+	if pw.Spec.PathwaysVersion != "" {
+		tag = string(pw.Spec.PathwaysVersion)
+	} else {
+		tag = "latest"
+	}
+	return tag
+}
+
 // Constructs the Pathways resource manager container spec for the underlying JobSet
 func MakeResourceManagerContainer(pw *pathwaysjob.PathwaysJob, rmJobName string) (*corev1.Container, error) {
 	truth := true
 
 	rmContainerSpec := corev1.Container{
 		Name:            "pathways-rm",
-		Image:           "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/sanitized_server:latest",
+		Image:           fmt.Sprintf("us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:%s", makeImageTagUsingPathwaysVersion(pw)),
 		ImagePullPolicy: "Always",
 		SecurityContext: &corev1.SecurityContext{Privileged: &truth},
 		Args: []string{
@@ -342,7 +351,7 @@ func MakeProxyContainer(pw *pathwaysjob.PathwaysJob, rmJobName string) (*corev1.
 
 	proxyContainerSpec := corev1.Container{
 		Name:            "pathways-proxy",
-		Image:           "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/sanitized_proxy_server:latest",
+		Image:           fmt.Sprintf("us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:%s", makeImageTagUsingPathwaysVersion(pw)),
 		ImagePullPolicy: "Always",
 		SecurityContext: &corev1.SecurityContext{Privileged: &truth},
 		Args: []string{
@@ -360,6 +369,15 @@ func MakeProxyContainer(pw *pathwaysjob.PathwaysJob, rmJobName string) (*corev1.
 func MakeWorkerJob(ctx context.Context, pw *pathwaysjob.PathwaysJob, rmJobName string) (jobsetv1alpha2.ReplicatedJob, error) {
 	truth := true
 	volumeSourceType := corev1.HostPathDirectoryOrCreate
+	objectMeta := metav1.ObjectMeta{}
+
+	if pw.Spec.Controller.DeploymentMode == pathwaysjob.Default {
+		objectMeta = metav1.ObjectMeta{
+			Annotations: map[string]string{
+				"alpha.jobset.sigs.k8s.io/exclusive-topology": "cloud.google.com/gke-nodepool",
+			},
+		}
+	}
 
 	workerJob := jobsetv1alpha2.ReplicatedJob{
 		Name:     "worker",
@@ -370,11 +388,12 @@ func MakeWorkerJob(ctx context.Context, pw *pathwaysjob.PathwaysJob, rmJobName s
 				Completions:  ptr.To(int32(NumVMs)), // number of workers remember to change
 				Parallelism:  ptr.To(int32(NumVMs)), // number of workers  remember to change
 				Template: corev1.PodTemplateSpec{
+					ObjectMeta: objectMeta,
 					Spec: corev1.PodSpec{
 						Containers: []corev1.Container{
 							{
 								Name:            "pathways-worker",
-								Image:           "us-docker.pkg.dev/cloud-tpu-v2-images-dev/pathways/sanitized_server:latest",
+								Image:           fmt.Sprintf("us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:%s", makeImageTagUsingPathwaysVersion(pw)),
 								ImagePullPolicy: "Always",
 								SecurityContext: &corev1.SecurityContext{Privileged: &truth},
 								Args: []string{
@@ -549,6 +568,10 @@ func MakeJobsForDefaultDeployment(ctx context.Context, pw *pathwaysjob.PathwaysJ
 				Parallelism:  ptr.To(int32(1)),
 				Template: corev1.PodTemplateSpec{
 					Spec: corev1.PodSpec{
+						NodeSelector: map[string]string{ // predictably place RM on CPUs
+							"cloud.google.com/machine-family":  "n2",
+							"node.kubernetes.io/instance-type": "n2-standard-64",
+						},
 						HostNetwork: true,                              // For performance == McJAX
 						DNSPolicy:   corev1.DNSClusterFirstWithHostNet, // For performance == McJAX
 						Tolerations: []corev1.Toleration{
@@ -586,6 +609,10 @@ func MakeJobsForDefaultDeployment(ctx context.Context, pw *pathwaysjob.PathwaysJ
 				Parallelism:  ptr.To(int32(1)),
 				Template: corev1.PodTemplateSpec{
 					Spec: corev1.PodSpec{
+						NodeSelector: map[string]string{ // predictably place RM on CPUs
+							"cloud.google.com/machine-family":  "n2",
+							"node.kubernetes.io/instance-type": "n2-standard-64",
+						},
 						HostNetwork: true,                              // For performance == McJAX
 						DNSPolicy:   corev1.DNSClusterFirstWithHostNet, // For performance == McJAX
 						Tolerations: []corev1.Toleration{
@@ -629,6 +656,10 @@ func MakeJobsForDefaultDeployment(ctx context.Context, pw *pathwaysjob.PathwaysJ
 					// },
 					Template: corev1.PodTemplateSpec{
 						Spec: corev1.PodSpec{
+							NodeSelector: map[string]string{ // predictably place RM on CPUs
+								"cloud.google.com/machine-family":  "n2",
+								"node.kubernetes.io/instance-type": "n2-standard-64",
+							},
 							HostNetwork: true,                              // For performance == McJAX
 							DNSPolicy:   corev1.DNSClusterFirstWithHostNet, // For performance == McJAX
 							Tolerations: []corev1.Toleration{ // tolerations are important here to not run this job on TPUs