Adding debug vmodule flags

Author: lukebaumann

axlearn/cloud/gcp/pathways_utils.py

@@ -251,6 +251,8 @@ def _build_pathways_head_container(self) -> dict:
         # In Jax 0.6.2 and beyond this flag can be renamed to
         # IFRT_PROXY_USE_INSECURE_GRPC_CREDENTIALS as well.
         self._update_env_list(env_list, "TEST_UNDECLARED_OUTPUTS_DIR", "true")
+        self._update_env_list(env_list, "TF_CPP_MIN_LOG_LEVEL", "0")
+        self._update_env_list(env_list, "TF_CPP_VMODULE", "grpc_host_buffer=3,rpc_helper=3,host_buffer=3,ifrt_backend=3,grpc_service_impl=3")
 
         env_list.append(
             {
@@ -322,6 +324,7 @@ def _build_pathways_head_sidecar_containers(self) -> list[Nested[Any]]:
             f"--gcs_scratch_location={staging_location}",
             # This should be made configurable
             f"--num_elastic_slices={cfg.accelerator.num_replicas}",
+            "--vmodule=grpc_host_buffer=3,rpc_helper=3,host_buffer=3,ifrt_backend=3,grpc_service_impl=3",
         ]
         cmd_args.extend(xla_flags_from_options(self._xla_options).split())