tpux 0.1.3

A set of Python scripts that makes your experience on TPU better

tpux: Enhance Your Google Cloud TPU Experience

Welcome to tpux, your essential toolkit designed to revolutionize the way you use Google Cloud TPUs. This suite of tools is tailored to simplify and streamline your TPU setup and operation processes, ensuring you maximize your productivity with minimal effort.

Pronunciation

To pronounce "tpux", first say "TPU" as you would in English, followed by "X" pronounced as /iks/ in French.

Why You Need tpux

Setting up Google Cloud TPU instances traditionally involves initializing empty VM instances, a process that can be tedious and repetitive. With tpux, this setup is greatly simplified, allowing you to focus on what truly matters—your work.

Features

• tpux: A user-friendly setup script that automates the configuration of your Google Cloud TPUs. This tool ensures that you are equipped with the latest practices and optimizations, keeping your operations cutting-edge.
• podrun: Seamlessly execute commands across all nodes in your TPU pods. Ideal for scaling applications and managing large-scale machine learning tasks, it enhances efficiency and effectiveness across your deployments.

Inspired by the comprehensive guide ayaka14732/tpu-starter, tpux incorporates best practices for TPU usage in open-source environments.

Setting Up Your TPU VM or TPU Pod with tpux

When Creating TPU VM or TPU Pod Instances

During the creation of a TPU VM instance, ensure to select the latest tpu-ubuntu2204-base software version to benefit from the most up-to-date system and software packages.

Besides using the web UI to create TPUs, you can also use the Google Cloud Shell. Here, your --version option should specify tpu-ubuntu2204-base. For example:

until gcloud alpha compute tpus tpu-vm create node-2 --zone us-central2-b --accelerator-type v4-32 --version tpu-ubuntu2204-base ; do : ; done

Using the tpux Command to Execute the Setup Script

After SSH into one of the hosts of your TPU VM or TPU Pod, you can perform the setup using the following method:

pip install tpux
export PATH="$HOME/.local/bin:$PATH"
tpux

Simply follow the on-screen prompts to complete the setup of your TPU VM or TPU Pod.

Executing Commands Across All Hosts with the podrun Command

After setting up with the tpux command, you can use the podrun command to execute specified commands across all TPU hosts.

podrun reads the command to be executed from stdin, for example:

echo echo meow | podrun -i

This command outputs "meow" on all hosts.

Using the -i parameter executes the command on all machines, while omitting -i executes on all hosts except the local one:

echo echo meow | podrun

This command outputs "meow" on all hosts except the local machine.

For more information on how to use the podrun command, type:

podrun -h

Verifying Successful Configuration of Your TPU Pod

Given the complexity of configuring a TPU Pod, after executing the tpux setup command, you may want to ensure it was successful. You can verify this by:

echo echo meow | podrun -i

If the TPU Pod is configured correctly, the above command should output multiple lines of "meow," where the number of lines corresponds to the number of TPU Pod hosts.

touch ~/nfs_share/meow
echo ls -l ~/nfs_share/meow | podrun -i

If configured correctly, the above commands should display the results of ls -l ~/nfs_share/meow on multiple lines, with the number of lines equaling the number of TPU Pod hosts.

Disclaimer

This is not an officially supported Google product.