myco.env

Routines for querying compute environment resources, like gpu count.

assert_gpu_available()

Ensure a GPU is available for training or application.

Source code in myco/env.py

def assert_gpu_available() -> None:
    """Ensure a GPU is available for training or application."""
    assert tf.test.is_built_with_gpu_support()
    assert count_available_gpus() > 0

count_available_gpus()

Determine number of GPUs available.

Source code in myco/env.py

def count_available_gpus() -> int:
    """Determine number of GPUs available."""
    return len(tf.config.list_physical_devices("GPU"))

count_available_xla()

Checks for XLA devices, which are typically TPUs. See www.tensorflow.org/xla

Source code in myco/env.py

def count_available_xla() -> int:
    """Checks for XLA devices, which are typically TPUs. See https://www.tensorflow.org/xla"""
    return len(tf.config.list_physical_devices("XLA_CPU")) + len(
        tf.config.list_physical_devices("XLA_GPU")
    )

get_gce_info(parameters=['name', 'machine-type'])

Get GCE metadata, like project, instane name, and machine type

Source code in myco/env.py

def get_gce_info(parameters: List[str] = ["name", "machine-type"]):
    """Get GCE metadata, like project, instane name, and machine type"""
    server = "http://metadata/computeMetadata/v1/instance"
    header = {"Metadata-Flavor": "Google"}
    if on_gce():
        info = {
            f"gce-{param}": requests.get(f"{server}/{param}", headers=header).text
            for param in parameters
        }
    else:
        info = {f"gce-{param}": "local" for param in parameters}
    return info

list_available_gpus()

Get all GPU device names.

Source code in myco/env.py

def list_available_gpus() -> list:
    """Get all GPU device names."""
    devices = [device.name for device in tf.config.list_physical_devices("GPU")]
    return [device[device.index("GPU") :] for device in devices]

on_gce()

Check if running a GCE instance via DNS lookup to metadata server.

Source code in myco/env.py

def on_gce():
    """Check if running a GCE instance via DNS lookup to metadata server."""
    try:
        socket.getaddrinfo("metadata.google.internal", 80)
        return True
    except socket.gaierror:
        return False