def get_tpu_head()

in ai-ml/gke-ray/rayserve/llm/model-composition/serve_tpu.py [0:0]

• 8 lines of code
• 4 McCabe index (conditional complexity)

def get_tpu_head() -> Optional[str]:
    if "TPU" not in ray.cluster_resources():
         # Pass in # TPU heads when the current Ray cluster resources can't be auto-detected.
        if os.environ.get('TPU_HEADS') is not None:
            return int(os.environ.get('TPU_HEADS'))
    # return the TPU-{accelerator}-head resource
    for key, _ in ray.cluster_resources().items():
        if key.endswith("head"):
            return key
    return None