Trouble connecting to GCP TPU VM - google-cloud-platform

Question

I followed along with the instructions to create a cloud TPU VM and run a custom neural network as directed by the Run Tensorflow on TPU pod slices to a T. It's important to note that I have been able to initialize the cloud TPUs when running this model on google colab, but require more resources than can be provided there even when explicitly managing the memory of the code used.
When I create the VM, I use the following command:
gcloud compute tpus tpu-vm create test-tpu-vm --zone=us-central1-b --accelerator-type=v2-8 --version=tpu-vm-tf-2.11.0
Next I log into the instance like so:
gcloud compute tpus tpu-vm ssh test-tpu-vm --zone us-central1-b --project <project_id>
where I clone down the code repo as follows:
git clone https://github.com/messerb5467/kaggle-competitions.git
and need to install the pandas library as required by the script:
pip install pandas
After doing all of this, I run the script and get the following issue:
messerb5467#t1v-n-3b61e142-w-0:~/kaggle-competitions/allstate-insurance-claims$ ./allstate_claims_data_nn.py test-tpu-vm
2022-12-29 23:06:08.675448: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations: SSE3 SSE4.1 SSE4.2 AVX AVX2 AVX512F FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2022-12-29 23:06:08.855095: I tensorflow/core/tpu/tpu_initializer_helper.cc:275] Libtpu path is: libtpu.so
D1229 23:06:09.026751179 12938 config.cc:113] gRPC EXPERIMENT tcp_frame_size_tuning OFF (default:OFF)
D1229 23:06:09.026775139 12938 config.cc:113] gRPC EXPERIMENT tcp_read_chunks OFF (default:OFF)
D1229 23:06:09.026782994 12938 config.cc:113] gRPC EXPERIMENT tcp_rcv_lowat OFF (default:OFF)
D1229 23:06:09.026790374 12938 config.cc:113] gRPC EXPERIMENT peer_state_based_framing OFF (default:OFF)
D1229 23:06:09.026797554 12938 config.cc:113] gRPC EXPERIMENT flow_control_fixes OFF (default:OFF)
D1229 23:06:09.026804675 12938 config.cc:113] gRPC EXPERIMENT memory_pressure_controller OFF (default:OFF)
D1229 23:06:09.026812324 12938 config.cc:113] gRPC EXPERIMENT periodic_resource_quota_reclamation ON (default:ON)
D1229 23:06:09.026819471 12938 config.cc:113] gRPC EXPERIMENT unconstrained_max_quota_buffer_size OFF (default:OFF)
D1229 23:06:09.026826517 12938 config.cc:113] gRPC EXPERIMENT new_hpack_huffman_decoder OFF (default:OFF)
D1229 23:06:09.026833747 12938 config.cc:113] gRPC EXPERIMENT event_engine_client OFF (default:OFF)
D1229 23:06:09.026840808 12938 config.cc:113] gRPC EXPERIMENT monitoring_experiment ON (default:ON)
D1229 23:06:09.026847921 12938 config.cc:113] gRPC EXPERIMENT promise_based_client_call OFF (default:OFF)
I1229 23:06:09.027065091 12938 ev_epoll1_linux.cc:121] grpc epoll fd: 7
D1229 23:06:09.027080773 12938 ev_posix.cc:141] Using polling engine: epoll1
D1229 23:06:09.027107304 12938 dns_resolver_ares.cc:824] Using ares dns resolver
D1229 23:06:09.027394086 12938 lb_policy_registry.cc:45] registering LB policy factory for "priority_experimental"
D1229 23:06:09.027405540 12938 lb_policy_registry.cc:45] registering LB policy factory for "outlier_detection_experimental"
D1229 23:06:09.027414241 12938 lb_policy_registry.cc:45] registering LB policy factory for "weighted_target_experimental"
D1229 23:06:09.027422457 12938 lb_policy_registry.cc:45] registering LB policy factory for "pick_first"
D1229 23:06:09.027430746 12938 lb_policy_registry.cc:45] registering LB policy factory for "round_robin"
D1229 23:06:09.027444142 12938 lb_policy_registry.cc:45] registering LB policy factory for "ring_hash_experimental"
D1229 23:06:09.027470472 12938 lb_policy_registry.cc:45] registering LB policy factory for "grpclb"
D1229 23:06:09.027508895 12938 lb_policy_registry.cc:45] registering LB policy factory for "rls_experimental"
D1229 23:06:09.027531154 12938 lb_policy_registry.cc:45] registering LB policy factory for "xds_cluster_manager_experimental"
D1229 23:06:09.027539743 12938 lb_policy_registry.cc:45] registering LB policy factory for "xds_cluster_impl_experimental"
D1229 23:06:09.027548580 12938 lb_policy_registry.cc:45] registering LB policy factory for "cds_experimental"
D1229 23:06:09.027556928 12938 lb_policy_registry.cc:45] registering LB policy factory for "xds_cluster_resolver_experimental"
D1229 23:06:09.027565714 12938 certificate_provider_registry.cc:35] registering certificate provider factory for "file_watcher"
I1229 23:06:09.051036639 12938 socket_utils_common_posix.cc:336] TCP_USER_TIMEOUT is available. TCP_USER_TIMEOUT will be used thereafter
2022-12-29 23:06:09.090117: I tensorflow/core/tpu/tpu_initializer_helper.cc:225] GetTpuPjrtApi not found
Traceback (most recent call last):
File "./allstate_claims_data_nn.py", line 145, in <module>
main()
File "./allstate_claims_data_nn.py", line 141, in main
allstate_nn_model = AllStateModelTrainer(args.tpu_name)
File "./allstate_claims_data_nn.py", line 34, in __init__
os.environ['TPU_LOAD_LIBRARY'] = 0
File "/usr/lib/python3.8/os.py", line 680, in __setitem__
value = self.encodevalue(value)
File "/usr/lib/python3.8/os.py", line 750, in encode
raise TypeError("str expected, not %s" % type(value).__name__)
TypeError: str expected, not int
D1229 23:06:12.209423135 12938 init.cc:190] grpc_shutdown starts clean-up now
messerb5467#t1v-n-3b61e142-w-0:~/kaggle-competitions/allstate-insurance-claims$
Even if I follow the message and make the 0 a string, it continues on to produce a core dump instead of running as one would expect. Any help would be mighty appreciated.
I've tried using a string for TPU_LOAD_LIBRARY instead of the documented integer:
export TPU_LOAD_LIBRARY=0
Use a TPU_NAME of local instead of test-tpu-vm since TPU vms run directly on a TPU host.
Unfortunately when following this, the errors start spinning out of control and I'm not able register with the TPU nodes at all despite the initialization working just fine in colab.
I imagine it has to be something simple and I'm just missing something somewhere.

Answer 1

This code uses a TPU pod to run. So, you would need to follow the instructions given here to create the pod. Note that you need to use the version tpu-vm-tf-2.11.0-pod and not tpu-vm-tf-2.11.0 when creating the pod.
For e.g.,
gcloud compute tpus tpu-vm create test-tpu-vm \
--zone=us-central1-a --accelerator-type=v2-32 \
--version=tpu-vm-tf-2.11.0-pod
For line 33 we should pass pod name, in your case test-tpu-vm. So, the call to the trainer would be ./allstate_claims_data_nn.py test-tpu-vm.
However, in line 34 it is trying to set the environment variable with an integer. This will not work because this needs to be string when setting environment variable from inside python code in Ubuntu. However, if you set it as a string TPU would throw errors because TPUs need this environment variable as an integer. So, I would recommend skipping the init function and following this and use
export TPU_NAME=test-tpu-vm
export TPU_LOAD_LIBRARY=0
./allstate_claims_data_nn.py test-tpu-vm
(or modify the code to skip taking tpu name as an argument)
This will get you past the TPU setup errors, there will be more code logic error from line 73 which you can continue to work on.