Merge branch 'develop' into wstcliyu/kueue-tpu-flavor

Author: pawloch00

.github/workflows/build_tests.yaml

@@ -20,10 +20,10 @@ on:
       tpu-type:
         description: 'TPU Type'
         required: true
-        default: 'v4-2x2x1'
+        default: 'v4-8'
         type: choice
         options:
-          - v4-2x2x1
+          - v4-8
   push:
     branches: ["main","develop"]
   pull_request: # By default this runs for types assigned, opened and synchronize.
@@ -36,6 +36,7 @@ jobs:
       cancel-in-progress: true
     outputs:
       cluster-name: ${{ steps.set-cluster-name.outputs.cluster-name }}
+      cluster-name-dws: ${{ steps.set-cluster-name-dws.outputs.cluster-name-dws }}
       group-name: ${{ steps.set-group-name.outputs.group-name }}
       zone: ${{ steps.set-zone.outputs.zone }}
       tpu-type: ${{ steps.set-tpu-type.outputs.tpu-type }}
@@ -59,6 +60,10 @@ jobs:
       id: set-cluster-name
       run: |
         echo cluster-name=build-xpk-2-nodepools-${{steps.set-run-id.outputs.run-id}} >> $GITHUB_OUTPUT
+    - name: set cluster-name-dws
+      id: set-cluster-name-dws
+      run: |
+        echo cluster-name-dws=build-xpk-2-nodepools-dws-${{steps.set-run-id.outputs.run-id}} >> $GITHUB_OUTPUT
     - name: set group-name
       id: set-group-name
       run: |    
@@ -70,7 +75,7 @@ jobs:
     - name: set tpu-type
       id: set-tpu-type
       run: |
-        echo tpu-type=v4-2x2x1 >> $GITHUB_OUTPUT
+        echo tpu-type=v4-8 >> $GITHUB_OUTPUT
     - name: set location
       id: set-location
       run: |    
@@ -158,6 +163,7 @@ jobs:
       cancel-in-progress: true
     uses: ./.github/workflows/reusable_cluster_create.yaml
     with:
+      cluster-name-dws: '${{needs.set-variables.outputs.cluster-name-dws}}'
       cluster-name: '${{needs.set-variables.outputs.cluster-name}}'
       tpu-type: '${{needs.set-variables.outputs.tpu-type || inputs.tpu-type}}'
       zone: '${{needs.set-variables.outputs.zone}}'
@@ -172,6 +178,7 @@ jobs:
       cancel-in-progress: true
     with:
       cluster-name: ${{needs.set-variables.outputs.cluster-name}}
+      cluster-name-dws: '${{needs.set-variables.outputs.cluster-name-dws}}'
       tpu-type: ${{needs.set-variables.outputs.tpu-type}}
       zone: ${{needs.set-variables.outputs.zone}}
       run-id: '${{needs.set-variables.outputs.run-id}}'
@@ -204,6 +211,7 @@ jobs:
     needs: [set-variables, storage-tests]
     uses: ./.github/workflows/reusable_cluster_delete.yaml
     with:
+      cluster-name-dws: ${{needs.set-variables.outputs.cluster-name-dws}}
       cluster-name: ${{needs.set-variables.outputs.cluster-name}}
       run-id: ${{needs.set-variables.outputs.run-id}}
       zone: ${{needs.set-variables.outputs.zone}}

.github/workflows/reusable_cluster_create.yaml

@@ -20,14 +20,26 @@ on:
       cluster-name:
         type: string
         required: true
+      cluster-name-dws:
+        type: string
+        required: true
       tpu-type:
         description: 'TPU Type'
         required: false
-        default: 'v4-2x2x1'
+        default: 'v4-8'
+        type: string
+      tpu-type-dws:
+        description: 'TPU Type for DWS flex nodepool'
         type: string
+        required: false
+        default: v5p-8
       zone:
         type: string
         required: true
+      zone-dws:
+        type: string
+        required: false
+        default: us-east5-a
       location:
         type: string
         required: true
@@ -39,6 +51,7 @@ on:
 env:
   # Names must be unique in parallel running tests.
   CLUSTER_ARGUMENTS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME}} --maintenance-window=23:50"
+  CLUSTER_ARGUMENTS_DWS: "--network=${{secrets.NETWORK_NAME}} --subnetwork=${{secrets.SUBNETWORK_NAME_DWS}} --maintenance-window=23:50"
 
 jobs:
   cluster-create:
@@ -76,6 +89,8 @@ jobs:
         gcloud config get compute/zone
     - name: Check xpk installation
       run: xpk --help
+    - name: Create a DWS flex queued xpk cluster 
+      run: python xpk.py cluster create --cluster ${{inputs.cluster-name-dws}} --tpu-type=${{inputs.tpu-type-dws}}  --num-slices=1 --zone=${{inputs.zone-dws}} --default-pool-cpu-num-nodes=2 --flex --enable-gcpfilestore-csi-driver --enable-gcsfuse-csi-driver --custom-cluster-arguments="${CLUSTER_ARGUMENTS_DWS}"
     - name: Create a Pathways-enabled XPK Cluster with 2x ${{inputs.tpu-type}} nodepools. Larger num-nodes to avoid master resizing.
       run: python xpk.py cluster create-pathways --cluster ${{inputs.cluster-name}} --tpu-type=${{inputs.tpu-type}}  --num-slices=1 --zone=${{inputs.zone}} --default-pool-cpu-machine-type=n1-standard-16 --default-pool-cpu-num-nodes=2 --reservation='${{ secrets.GCP_TPU_V4_RESERVATION }}' --enable-gcpfilestore-csi-driver --enable-gcsfuse-csi-driver --custom-cluster-arguments="${CLUSTER_ARGUMENTS}"
 

.github/workflows/reusable_cluster_delete.yaml

@@ -20,6 +20,13 @@ on:
       zone:
         required: true
         type: string
+      cluster-name-dws:
+        required: true
+        type: string
+      zone-dws:
+        required: false
+        type: string
+        default: us-east5-a
       run-id:
         required: true
         type: string
@@ -61,4 +68,7 @@ jobs:
       run: xpk --help
     - name: Delete the cluster created
       if: always()
-      run: python xpk.py cluster delete --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --force
+      run: python xpk.py cluster delete --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --force
+    - name: Delete the cluster created
+      if: always()
+      run: python xpk.py cluster delete --cluster ${{inputs.cluster-name-dws}} --zone=${{inputs.zone-dws}} --force

.github/workflows/reusable_cluster_private.yaml

@@ -23,7 +23,7 @@ on:
       tpu-type:
         description: 'TPU Type'
         required: false
-        default: 'v4-2x2x1'
+        default: 'v4-8'
         type: string
       zone:
         type: string

.github/workflows/reusable_storage_create.yaml

@@ -21,7 +21,7 @@ on:
       tpu-type:
         description: 'TPU Type'
         required: false
-        default: 'v4-2x2x1'
+        default: 'v4-8'
         type: string
       zone:
         type: string

.github/workflows/reusable_storage_tests.yaml

@@ -23,7 +23,7 @@ on:
       tpu-type:
         description: 'TPU Type'
         required: false
-        default: 'v4-2x2x1'
+        default: 'v4-8'
         type: string
       zone:
         type: string

.github/workflows/reusable_workload_tests.yaml

@@ -18,18 +18,30 @@ on:
       cluster-name:
         required: true
         type: string
+      cluster-name-dws:
+        required: true
+        type: string
       tpu-type:
         required: true
         type: string
+      tpu-type-dws:
+        required: false
+        type: string
+        default: v5p-8
       zone:
         required: true
         type: string
+      zone-dws:
+        required: false
+        type: string
+        default: us-east5-a
       run-id:
         required: true
         type: string
 
 env:
   WORKLOAD_NAME: xpktest-build-${{ github.run_attempt }}
+  WORKLOAD_NAME_DWS: xpktest-build-${{ github.run_attempt }}-dws
   PRIVATE_IMAGE_WORKLOAD_NAME: xpktest-build-priv-${{ github.run_attempt }}
   PATHWAYS_WORKLOAD_NAME: xpkpw-build-${{ github.run_attempt }}
 
@@ -78,6 +90,10 @@ jobs:
       run: python3 xpk.py inspector --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}}  --workload $WORKLOAD_NAME
     - name: Wait for workload completion and confirm it succeeded
       run: python3 xpk.py workload list --cluster ${{inputs.cluster-name}} --zone=${{inputs.zone}} --wait-for-job-completion $WORKLOAD_NAME --timeout 300
+    - name: Run dws flex queued TPU workload
+      run: python3 xpk.py workload create --workload $WORKLOAD_NAME_DWS --cluster ${{inputs.cluster-name-dws}} --zone=${{inputs.zone-dws}} --tpu-type=${{inputs.tpu-type-dws}}  --flex --command "echo foo" --num-slices=1
+    - name: Wait for workload completion and confirm it succeeded
+      run: python3 xpk.py workload list --cluster ${{inputs.cluster-name-dws}} --zone=${{inputs.zone-dws}} --wait-for-job-completion $WORKLOAD_NAME_DWS --timeout 500
     - name: Run a Pathways workload on Ubuntu base image
       run: python xpk.py workload create-pathways --cluster ${{inputs.cluster-name}} --workload $PATHWAYS_WORKLOAD_NAME --docker-image='marketplace.gcr.io/google/ubuntu2004' --tpu-type=${{inputs.tpu-type}} --num-slices=1 --zone=${{inputs.zone}} --command "echo \"Hello world from a test script! \""
     - name: Wait for Pathways workload completion and confirm it succeeded

README.md

@@ -227,6 +227,13 @@ all zones.
     --num-slices=4 --spot
     ```
 
+* Cluster Create (DWS flex queued capacity):
+    ```shell
+        python3 xpk.py cluster create \
+        --cluster xpk-test --tpu-type=v5litepod-16 \
+        --num-slices=4 --flex
+    ```
+
 * Cluster Create for Pathways:
     Pathways compatible cluster can be created using `cluster create-pathways`.
     ```shell
@@ -463,6 +470,7 @@ Currently, the below flags/arguments are supported for A3 Mega, A3 Ultra and A4
   * `--reservation`
   * `--spot`
   * `--on-demand` (A3 Mega only)
+  * `--flex`
 
 ## Running XPK on existing clusters
 
@@ -486,9 +494,10 @@ Currently XPK supports the below types of storages:
 - [Google Cloud Filestore](#filestore)
 - [Google Cloud Parallelstore](#parallelstore)
 - [Google Cloud Block storages (Persistent Disk, Hyperdisk)](#block-storage-persistent-disk-hyperdisk)
+- [Google Cloud Managed Lustre](#managed-lustre)
 
 ### FUSE
-A FUSE adapter lets you mount and access Cloud Storage buckets as local file systems, so applications can read and write objects in your bucket using standard file system semantics.
+A FUSE adapter lets you mount and access Cloud Storage buckets as local file systems, so workloads can read and write objects in your bucket using standard file system semantics.
 
 To use the GCS FUSE with XPK you need to create a [Storage Bucket](https://console.cloud.google.com/storage/).
 
@@ -515,7 +524,7 @@ Parameters:
 
 ### Filestore
 
-A Filestore adapter lets you mount and access [Filestore instances](https://cloud.google.com/filestore/) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
+A Filestore adapter lets you mount and access [Filestore instances](https://cloud.google.com/filestore/) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
 
 To create and attach a GCP Filestore instance to your cluster use `xpk storage create` command with `--type=gcpfilestore`:
 
@@ -551,7 +560,7 @@ Commands `xpk storage create` and `xpk storage attach` with `--type=gcpfilestore
 
 ### Parallelstore
 
-A Parallelstore adapter lets you mount and access [Parallelstore instances](https://cloud.google.com/parallelstore/) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
+A Parallelstore adapter lets you mount and access [Parallelstore instances](https://cloud.google.com/parallelstore/) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
 
 To use the GCS Parallelstore with XPK you need to create a [Parallelstore Instance](https://console.cloud.google.com/parallelstore/).
 
@@ -575,7 +584,7 @@ Parameters:
 
 ### Block storage (Persistent Disk, Hyperdisk)
 
-A PersistentDisk adapter lets you mount and access Google Cloud Block storage solutions ([Persistent Disk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#pd), [Hyperdisk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#hyperdisk)) as local file systems, so applications can read and write files in your volumes using standard file system semantics.
+A PersistentDisk adapter lets you mount and access Google Cloud Block storage solutions ([Persistent Disk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#pd), [Hyperdisk](https://cloud.google.com/kubernetes-engine/docs/concepts/storage-overview#hyperdisk)) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
 
 To use the GCE PersistentDisk with XPK you need to create a [disk in GCE](https://cloud.google.com/compute/docs/disks). Please consider that the disk type you are creating is [compatible with the VMs](https://cloud.google.com/compute/docs/machine-resource#machine_type_comparison) in the default and accelerator nodepools.
 
@@ -597,6 +606,30 @@ Parameters:
 - `--readonly` - if set to true, workload can only read from storage.
 - `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
 
+### Managed Lustre
+
+A Managed Lustre adaptor lets you mount and access [Google Cloud Managed Lustre instances](https://cloud.google.com/kubernetes-engine/docs/concepts/managed-lustre) as local file systems, so workloads can read and write files in your volumes using standard file system semantics.
+
+To use the GCP Managed Lustre with XPK you need to create [an instance](https://cloud.google.com/managed-lustre/docs/create-instance). Please make sure you enable GKE support when creating the instance (gcloud ex. `--gke-support-enabled`).
+
+Once it's ready you can use `xpk storage attach` with `--type=lustre` command to attach a Managed Lustre instance to your cluster. Currently, attaching a Managed Lustre instance is supported only by providing a manifest file.
+
+```shell
+python3 xpk.py storage attach test-lustre-storage --type=lustre \
+  --project=$PROJECT --cluster=$CLUSTER --zone=$ZONE \
+  --mount-point='/test-mount-point' --readonly=false \
+  --auto-mount=true \
+  --manifest='./examples/storage/lustre-manifest-attach.yaml'
+```
+
+Parameters:
+
+- `--type` - type of the storage `lustre`
+- `--auto-mount` - if set to true all workloads will have this storage mounted by default.
+- `--mount-point` - the path on which this storage should be mounted for a workload.
+- `--readonly` - if set to true, workload can only read from storage.
+- `--manifest` - path to the manifest file containing PersistentVolume and PresistentVolumeClaim definitions.
+
 ### List attached storages
 
 ```shell
@@ -638,8 +671,14 @@ python3 xpk.py storage delete test-fs-instance \
     python3 xpk.py workload create \
     --workload xpk-test-workload --command "echo goodbye" \
     --cluster xpk-test \
-    --tpu-type=v5litepod-16 --projet=$PROJECT
+    --tpu-type=v5litepod-16 --project=$PROJECT
     ```
+*   Workload create(DWS flex with queued provisioning):
+        ```shell
+    python3 xpk.py workload create \
+    --workload xpk-test-workload --command "echo goodbye" \
+    --cluster xpk-test  --flex \
+    --tpu-type=v5litepod-16 --project=$PROJECT
 
 *   Workload Create for Pathways:
     Pathways workload can be submitted using `workload create-pathways` on a Pathways enabled cluster (created with `cluster create-pathways`)

examples/storage/lustre-manifest-attach.yaml

@@ -0,0 +1,34 @@
+apiVersion: v1
+kind: PersistentVolume
+metadata:
+  name: xpk-lustre-pv
+spec:
+  storageClassName: ""
+  capacity:
+    storage: STORAGE_SIZE #ex. 36000Gi
+  accessModes:
+    - ReadWriteMany
+  persistentVolumeReclaimPolicy: Retain
+  volumeMode: Filesystem
+  claimRef:
+    namespace: default
+    name: xpk-lustre-pvc
+  csi:
+    driver: lustre.csi.storage.gke.io
+    volumeHandle: "PROJECT_ID/ZONE/INSTANCE_NAME"
+    volumeAttributes:
+      ip: IP_ADDRESS
+      filesystem: FILE_SYSTEM
+---
+  kind: PersistentVolumeClaim
+  apiVersion: v1
+  metadata:
+    name: xpk-lustre-pvc
+  spec:
+    accessModes:
+      - ReadWriteMany
+    storageClassName: ""
+    volumeName: xpk-lustre-pv
+    resources:
+      requests:
+        storage: STORAGE_SIZE

src/xpk/blueprints/a3mega/kueue-xpk-configuration.yaml.tftpl

@@ -15,8 +15,30 @@ metadata:
   name: "1xh100-mega-80gb-8"
 spec:
   nodeLabels:
-    cloud.google.com/gke-accelerator: nvidia-h100-mega-80gb
-  ${tas_name}
+    cloud.google.com/gke-accelerator: "nvidia-h100-mega-80gb"
+  %{~ if reservation==1 ~}
+  topologyName: "gke-default"
+  %{ endif }
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: ProvisioningRequestConfig
+metadata:
+  name: dws-config
+spec:
+  provisioningClassName: queued-provisioning.gke.io
+  managedResources:
+  - nvidia.com/gpu
+---
+apiVersion: kueue.x-k8s.io/v1beta1
+kind: AdmissionCheck
+metadata:
+  name: dws-prov
+spec:
+  controllerName: kueue.x-k8s.io/provisioning-request
+  parameters:
+    apiGroup: kueue.x-k8s.io
+    kind: ProvisioningRequestConfig
+    name: dws-config
 ---
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: ClusterQueue
@@ -35,6 +57,10 @@ spec:
         nominalQuota: 10000
       - name: "memory"
         nominalQuota: 10000Gi
+  %{~ if flex_start==1 ~}
+  admissionChecks:
+   - dws-prov
+  %{ endif }
 ---
 apiVersion: kueue.x-k8s.io/v1beta1
 kind: LocalQueue