add support for enabling tcpx/o in a3 and a3mega vm, provide script for injecting rxdm sidecar and other required components into user workload

[chengcongdu]

This change add support for enabling tcpx/o in a3 and a3mega vm, it include:

• Install the correct version of NCCL plugin and NRI plugin daemonsets for A3 and A3Mega nodepool
• provide script for injecting rxdm sidecar and other required components into user workload
• provide a sample workload in the format of Kubernetes Job manifest to show how GPUDirect work and what performance it brings

---

SAMPLE OUTPUT of the added enable_tcpxo_in_workload:
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0]: Creating...
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0]: Provisioning with 'local-exec'...
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0] (local-exec): Executing: ["/bin/sh" "-c" "python3 modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/scripts/enable-tcpxo-in-workload.py --file modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/sample-tcpxo-workload-job.yaml --rxdm v1.0.10"]
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0] (local-exec): A new manifest has been generated and updated to have TCPXO enabled based on the provided workload
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0] (local-exec): It can be found in /home/user/fork/hpc-toolkit/gke-a3-mega/primary/modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/sample-tcpxo-workload-job-tcpxo.yaml
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0] (local-exec): You can use the following commands to submit the sample job:
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0] (local-exec): kubectl create -f /home/user/fork/hpc-toolkit/gke-a3-mega/primary/modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/sample-tcpxo-workload-job-tcpxo.yaml
module.a3-megagpu_pool.null_resource.enable_tcpxo_in_workload[0]: Creation complete after 0s [id=1932468328196680693]
...
...
Apply complete! Resources: 19 added, 1 changed, 1 destroyed.

Outputs:

instructions_a3-megagpu_pool = <<EOT
Since you are using a3-megagpu-8g machine type that has GPUDirect support, your nodepool had been configured with the required plugins.
To use the GPUDirect you will have to add the some components into your workload manifest. Details below

A sample GKE job that had GPUDirect enabled and NCCL test included has been generated locally at:
/home/user/fork/hpc-toolkit/gke-a3-mega/primary/modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/sample-tcpxo-workload-job-tcpxo.yaml

You can use the following commands to submit the sample job:
kubectl create -f /home/user/fork/hpc-toolkit/gke-a3-mega/primary/modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/sample-tcpxo-workload-job-tcpxo.yaml

Follow below steps to enable GPUDirect for your own workload:
export WORKLOAD_PATH=<>
python3 /home/user/fork/hpc-toolkit/gke-a3-mega/primary/modules/embedded/modules/compute/gke-node-pool/gpu-direct-workload/scripts/enable-tcpxo-in-workload.py --file $WORKLOAD_PATH --rxdm v1.0.10
WARNING
The "--rxdm" version is tide to the nccl-tcpx/o-installer that had been deployed to your cluster, changing it to other value might have impact on performance
WARNING

Or you can also follow our GPUDirect user guide to update your workload
https://cloud.google.com/kubernetes-engine/docs/how-to/gpu-bandwidth-gpudirect-tcpx#add-gpudirect-manifests

After the command an updated manifest will be generated and you can deploy it to the clsuter using kubectl

---

SAMPLE OUTPUT of running python script directly:
A new manifest has been generated and updated to have TCPXO enabled based on the provided workload
It can be found in /home/user/chdu/sample-tcpxo-workload-job-tcpxo.yaml
You can use the following commands to submit the sample job:
kubectl create -f /home/user/chdu/sample-tcpxo-workload-job-tcpxo.yaml