Google Colab Processor Unit / Runtime Type Detector.

colab_pu_detector.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "colab_pu_detector.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gwLFwcWubtA4"
      },
      "source": [
        "# Google Colab Processor Unit / Runtime Type\n",
        "\n",
        "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/gist/HeinrichAD/eea00d50c5702cbf63082983b6a2d763/colab_pu_detector.ipynb)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hoxAOOacmklJ"
      },
      "source": [
        "Import via `httpimport` directly from GitHub or copy the mini [`ProcessorUnit`](https://gist.github.com/HeinrichAD/eea00d50c5702cbf63082983b6a2d763#file-processorunit-py) class into your notebook."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "3sXDPH9tFaA2"
      },
      "source": [
        "import httpimport\n",
        "\n",
        "url = \"https://gist.githubusercontent.com/HeinrichAD/eea00d50c5702cbf63082983b6a2d763/raw/48a281f4d5738928e73f53f59364ea2d622f050e\"\n",
        "with httpimport.remote_repo([\"ProcessorUnit\"], url):\n",
        "  from ProcessorUnit import PU"
      ],
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "K4tBDgqPa7gh"
      },
      "source": [
        "## Examples\n",
        "\n",
        "Get current Processor Unit"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "WjvyK6AWhQGg",
        "outputId": "97dfd604-5498-41a8-9925-a625119dfb85"
      },
      "source": [
        "print(PU.getPU())"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "PU.TPU\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xEcBJUrOhU5Y"
      },
      "source": [
        "Check specific Processor Unit"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0jfC9vMNhXoE",
        "outputId": "3ffad5c9-800e-4e6f-86f8-da431f36f58d"
      },
      "source": [
        "print(PU.isTPU())"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "True\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lobIvny0hall"
      },
      "source": [
        "The usual Python Enum power"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZBrf6zlWhZ4T",
        "outputId": "166d92d2-12f0-44b5-94b3-73723a4a7163"
      },
      "source": [
        "for pu in PU:\n",
        "  print(pu)"
      ],
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "PU.CPU\n",
            "PU.GPU\n",
            "PU.TPU\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "RipMpUOpc8Y6"
      },
      "source": [
        "## Basis of the Processor Unit decision\n",
        "\n",
        "```python\n",
        "from os import environ\n",
        "print(environ[\"COLAB_GPU\"] == '1', \"TPU_NAME\" in environ)\n",
        "```\n",
        "```\n",
        "If ...\n",
        "   False False   CPU\n",
        "   True  False   GPU\n",
        "   False True    TPU\n",
        "```"
      ]
    }
  ]
}

ProcessorUnit.py

from enum import Enum
from os import environ

class PU(Enum):
  """
  Google Colab Processor Unit / Runtime Type Detector.
  """
  CPU = 1
  GPU = 2
  TPU = 3

  def isCPU():
    return not (PU.isGPU() or PU.isTPU())
  def isGPU():
    return environ["COLAB_GPU"] == '1'
  def isTPU():
    return "TPU_NAME" in environ

  def getPU():
    return PU.GPU if PU.isGPU() else PU.TPU if PU.isTPU() else PU.CPU