[performance] performance downgrade in Pathways when calling device_get

## Motivation Hi, community. I am using Pathways in GKE to test the inference of performance with SGLangJAX. According to my profile, I find out the device_get costs a lot of time compared with cost not in Pathways. ## Question How to tune the Pathways to improve the performance? Such like enable some configurations. Or device_get is not recommended in Pathways? ## Environment ```bash jax 0.8.1 jaxlib 0.8.1 ``` ### In Pathways tpu-v5p-8. I setup my Pathways GKE cluster according to the following script. I run `JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE='1' python3 test.py` on a e2-standard-32 CPU machine which locates in us-east5. ```bash CLUSTER=pathways-v5p PROJECT=xxx ZONE=us-east5-a REGION=us-east5 CLUSTER_VERSION="1.33.5-gke.1308000" PW_CPU_MACHINE_TYPE="n2-standard-64" CLUSTER_NODEPOOL_COUNT=1 TPU_MACHINE_TYPE="ct5p-hightpu-4t" WORKERS_PER_SLICE=1 TOPOLOGY="2x2x1" NUM_CPU_NODES=1 gcloud beta container clusters create ${CLUSTER} \ --project=${PROJECT} \ --zone=${ZONE} \ --scopes=storage-full,gke-default,cloud-platform \ --machine-type ${PW_CPU_MACHINE_TYPE} for i in $(seq 1 ${CLUSTER_NODEPOOL_COUNT}); do gcloud container node-pools create "tpu-np-${i}" \ --project=${PROJECT} \ --zone=${ZONE} \ --cluster=${CLUSTER} \ --machine-type=${TPU_MACHINE_TYPE} \ --num-nodes=${WORKERS_PER_SLICE} \ --placement-type=COMPACT \ --tpu-topology=${TOPOLOGY} \ --scopes=storage-full,gke-default,cloud-platform \ --workload-metadata=GCE_METADATA \ --spot \ --reservation-affinity=none done gcloud container node-pools create "cpu-pathways-np" \ --project ${PROJECT} \ --zone=${ZONE} \ --cluster ${CLUSTER} \ --machine-type ${PW_CPU_MACHINE_TYPE} \ --num-nodes ${NUM_CPU_NODES} \ --scopes=storage-full,gke-default,cloud-platform \ --workload-metadata=GCE_METADATA gcloud container clusters get-credentials ${CLUSTER} \ --zone=${ZONE} \ --project=${PROJECT} && kubectl config set-context --current --namespace=default kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.8.0/manifests.yaml kubectl apply --server-side -f https://github.com/google/pathways-job/releases/download/v0.1.2/install.yaml ``` ### Not in Pathways tpu-v6e-1. Note: I fail to launch another v5p vm due to quota problem. So I use v6e to replace it. ## Reproduce Example **_In Pathways_** Run `JAX_BACKEND_TARGET=grpc://127.0.0.1:29000 ENABLE_PATHWAYS_PERSISTENCE='1' python3 test.py`. ```bash In Pathways In Pathways y locates on device(9,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3) Use device_get: 8046.391ms In Pathways y locates on device(18,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3) Use device_get: 19442.270ms In Pathways y locates on device(27,TPU_DEVICE,coords=[0,0,0,0],vtask=0,slice=0,default_mem=device,mem_spaces=3) Use device_get: 18053.290ms ``` **_Not in Pathways_** Run `python3 test.py` ```bash Not in Pathways y locates on TPU_0(process=0,(0,0,0,0)) Use device_get: 3.373ms Not in Pathways y locates on TPU_0(process=0,(0,0,0,0)) Use device_get: 3.373ms Not in Pathways y locates on TPU_0(process=0,(0,0,0,0)) Use device_get: 3.371ms ``` ```python3 # test.py try: import pathwaysutils pathwaysutils.initialize() print("In Pathways") except Exception as e: print(f"Not in Pathways") import jax import time y = jax.numpy.ones((1000, 1000)) print(f"y locates on {y.device}") start_y = time.perf_counter() y_cpu=jax.device_get(y) print(f"Use device_get: {(time.perf_counter()-start_y)*1e3:.3f}ms") ```