https://github.com/trisongz/tpu-vm-docker-containers

Docker File Templates to access TPUs within TPU VM in Containers
https://github.com/trisongz/tpu-vm-docker-containers

docker-compose docker-container google-cloud-platform machine-learning tpu

Last synced: 26 days ago
JSON representation

Docker File Templates to access TPUs within TPU VM in Containers

• Host: GitHub
• URL: https://github.com/trisongz/tpu-vm-docker-containers
• Owner: trisongz
• License: mit
• Created: 2021-06-11T06:11:38.000Z (over 3 years ago)
• Default Branch: main
• Last Pushed: 2021-06-11T06:11:42.000Z (over 3 years ago)
• Last Synced: 2024-12-23T17:43:39.667Z (about 1 month ago)
• Topics: docker-compose, docker-container, google-cloud-platform, machine-learning, tpu
• Language: Dockerfile
• Homepage:
• Size: 2.93 KB
• Stars: 6
• Watchers: 1
• Forks: 0
• Open Issues: 1
• Metadata Files:
  • Readme: README.md
  • License: LICENSE

Awesome Lists containing this project

README

        
# Accessing TPUs in Docker Containers with TPU VM

 Docker File Templates to access TPUs within TPU VM in Containers

This Github repo accompanies the [Medium Post](https://trisongz.medium.com/accessing-your-tpus-in-docker-containers-with-tpu-vm-e944f5909dd4) about how to get your containers to access the TPU directly in TPU VMs

## The Important Pieces

Below is the most important configs in the docker-compose.yaml file that you'll need to include in your specific container that is accessing the TPU. 

```yaml

cap_add:

    - ALL #unsure if necessary

environment:

    - TPU_NAME=tpu_name

    - TF_CPP_MIN_LOG_LEVEL=0

    - XRT_TPU_CONFIG="localservice;0;localhost:51011"

    - TF_XLA_FLAGS=--tf_xla_enable_xla_devices

volumes:

    - /var/run/docker.sock:/var/run/docker.sock #unsure if necessary

    - /usr/share/tpu/:/usr/share/tpu/

    - /lib/libtpu.so:/lib/libtpu.so

privileged: true

devices:

    - "/dev:/dev"

```

Included are a few Dockerfiles to help make sure google-cloud-sdk are installed with the service account that TPU VM is configured with.

After configuring your Dockerfile, docker-compose and tpu.env, you should be able to run

```bash

docker-compose -f docker-compose.yaml --env-file=tpu.env up # -d for daemon

```

If you this repository helpful, I'd appreciate a star, a fork, or even sharing it!