Avoid querying metadata query to check if it's GCE if `TPU_SKIP_MDS_Q…

…UERY` is set.

PiperOrigin-RevId: 687399234

jax/_src/clusters/cloud_tpu_cluster.py

@@ -180,6 +180,9 @@ def is_env_present(cls) -> bool:
     if not running_in_cloud_tpu_vm:
       logger.debug("Did not detect cloud TPU VM")
       return False
+    if os.environ.get("TPU_SKIP_MDS_QUERY") is not None:
+      logger.debug("TPU_SKIP_MDS_QUERY is set to True, so it's probably not a GCE TPU cluster.")
+      return False
     metadata_response, metadata_code = get_metadata('agent-worker-number')
     if metadata_code == metadata_response_code_success:
       logger.debug("Gce Tpu Cluster detected for Jax Distributed System")