class_docstr: Google TPU accelerators.
docstr: Get additional resources required for TPU nodes.
  
  This will populate the TPU pod type and the TPU name which
  is used for TPU pod execution.
  
  When running workloads on a TPU pod, we need a way to run
  the same binary on every worker in the TPU pod.
  
  See https://jax.readthedocs.io/en/latest/multi_process.html
  for more information.
  
  To do this in ray, we take advantage of custom resources. We
  mark worker 0 of the TPU pod as a "coordinator" that identifies
  the other workers in the TPU pod. We therefore need:
  - worker 0 to be targetable.
  - all workers in the TPU pod to have a unique identifier consistent
  within a TPU pod.
  
  So assuming we want to run the following workload:
  
  @ray.remote
  def my_jax_fn():
      import jax
      return jax.device_count()
  
  We could broadcast this on a TPU pod (e.g. a v4-16) as follows:
  
  @ray.remote(resources={"TPU-v4-16-head"})
  def run_jax_fn(executable):
      # Note this will execute on worker 0
      tpu_name = ray.util.accelerators.tpu.get_tpu_pod_name()
      num_workers = ray.util.accelerators.tpu.get_tpu_num_workers()
      tpu_executable = executable.options(resources={"TPU": 4, tpu_name: 1})
      return [tpu_executable.remote() for _ in range(num_workers)]
  
  Returns:
      A dictionary representing additional resources that may be
      necessary for a particular accelerator type.
docstr_tok: ["Get","additional","resources","required","for","TPU","nodes",".","This","will","populate","the","TPU","pod","type","and","the","TPU","name","which","is","used","for","TPU","pod","execution",".","When","running","workloads","on","a","TPU","pod",",","we","need","a","way","to","run","the","same","binary","on","every","worker","in","the","TPU","pod",".","See","https",":","\/\/jax.readthedocs.io\/en\/latest\/multi_process.html","for","more","information",".","To","do","this","in","ray",",","we","take","advantage","of","custom","resources",".","We","mark","worker","0","of","the","TPU","pod","as","a","``","coordinator","''","that","identifies","the","other","workers","in","the","TPU","pod",".","We","therefore","need",":","-","worker","0","to","be","targetable",".","-","all","workers","in","the","TPU","pod","to","have","a","unique","identifier","consistent","within","a","TPU","pod",".","So","assuming","we","want","to","run","the","following","workload",":","@","ray.remote","def","my_jax_fn","(",")",":","import","jax","return","jax.device_count","(",")","We","could","broadcast","this","on","a","TPU","pod","(","e.g",".","a","v4-16",")","as","follows",":","@","ray.remote","(","resources=","{","``","TPU-v4-16-head","''","}",")","def","run_jax_fn","(","executable",")",":","#","Note","this","will","execute","on","worker","0","tpu_name","=","ray.util.accelerators.tpu.get_tpu_pod_name","(",")","num_workers","=","ray.util.accelerators.tpu.get_tpu_num_workers","(",")","tpu_executable","=","executable.options","(","resources=","{","``","TPU","''",":","4",",","tpu_name",":","1","}",")","return","[","tpu_executable.remote","(",")","for","_","in","range","(","num_workers",")","]","Returns",":","A","dictionary","representing","additional","resources","that","may","be","necessary","for","a","particular","accelerator","type","."]
code: def get_current_node_additional_resources() -> (
      Optional[Dict[str, float]]
  ):
      """Get additional resources required for TPU nodes.
  
      This will populate the TPU pod type and the TPU name which
      is used for TPU pod execution.
  
      When running workloads on a TPU pod, we need a way to run
      the same binary on every worker in the TPU pod.
  
      See https://jax.readthedocs.io/en/latest/multi_process.html
      for more information.
  
      To do this in ray, we take advantage of custom resources. We
      mark worker 0 of the TPU pod as a "coordinator" that identifies
      the other workers in the TPU pod. We therefore need:
      - worker 0 to be targetable.
      - all workers in the TPU pod to have a unique identifier consistent
      within a TPU pod.
  
      So assuming we want to run the following workload:
  
      @ray.remote
      def my_jax_fn():
          import jax
          return jax.device_count()
  
      We could broadcast this on a TPU pod (e.g. a v4-16) as follows:
  
      @ray.remote(resources={"TPU-v4-16-head"})
      def run_jax_fn(executable):
          # Note this will execute on worker 0
          tpu_name = ray.util.accelerators.tpu.get_tpu_pod_name()
          num_workers = ray.util.accelerators.tpu.get_tpu_num_workers()
          tpu_executable = executable.options(resources={"TPU": 4, tpu_name: 1})
          return [tpu_executable.remote() for _ in range(num_workers)]
  
      Returns:
          A dictionary representing additional resources that may be
          necessary for a particular accelerator type.
  
      """
      resources = {}
      tpu_name = TPUAcceleratorManager.get_current_node_tpu_name()
      worker_id = (
          TPUAcceleratorManager._get_current_node_tpu_worker_id()
      )
      tpu_pod_type = (
          TPUAcceleratorManager._get_current_node_tpu_pod_type()
      )
  
      if tpu_name and worker_id is not None and tpu_pod_type:
          pod_head_resource_name = f"TPU-{tpu_pod_type}-head"
          # Add the name of the TPU to the resource.
          resources[tpu_name] = 1
          # Only add in the TPU pod type resource to worker 0.
          if worker_id == 0:
              resources[pod_head_resource_name] = 1
      else:
          logging.info(
              "Failed to configure TPU pod. Got: "
              "tpu_name: %s, worker_id: %s, accelerator_type: %s",
              tpu_name,
              worker_id,
              tpu_pod_type,
          )
      if resources:
          return resources
      return None
  
imports: import os
  import re
  import glob
  import requests
  import logging
  from typing import Dict, Optional, List, Tuple
  from ray._private.accelerators.accelerator import AcceleratorManager
label_desc: Use image node_id 13 for calling the TPUAcceleratorManager obj's underlying member method code with example usage: obj.get_current_node_additional_resources() and returns: None, resources