parameters: (pod, config)
return_statement: return pod
docstring: Add node selector and cpu, memory requirement info for CPU pod.
  
      Args:
          pod: Pod configuration directory to for starting a Kubernetes pod.
          config: Configuration containing cluster-wide info.
  
      Returns:
          Potentially modified pod.
docstring_summary: Add node selector and cpu, memory requirement info for CPU pod.
docstring_tokens: Add node selector and cpu memory requirement info for CPU pod .
function: def enable_cpu_config(pod, config):
      """Add node selector and cpu, memory requirement info for CPU pod.
  
      Args:
          pod: Pod configuration directory to for starting a Kubernetes pod.
          config: Configuration containing cluster-wide info.
  
      Returns:
          Potentially modified pod.
      """
      # Ignore if cpuworker is not enabled
      enable_cpuworker = config.get("enable_cpuworker", False)
      if enable_cpuworker is False:
          return pod
  
      # Only works for 0-GPU job
      if "resourcegpu" not in pod or int(pod["resourcegpu"]) != 0:
          return pod
  
      # When cpuworker is enabled, CPU job should have gpuType=None
      if "nodeSelector" not in pod:
          pod["nodeSelector"] = {}
      pod["nodeSelector"]["gpuType"] = "None"
  
      job_training_type = pod.get("jobtrainingtype", None)
      dist_role = pod.get("distRole", None)
  
      # No special config for ps pod. It is always co-located with a worker
      if dist_role == "ps":
          return pod
  
      # Add node selector cpuworker=active
      pod["nodeSelector"]["cpuworker"] = "active"
  
      # Add node selector sku=<sku_value>
      default_cpu_sku = config.get("default_cpu_sku", None)
      if "sku" in pod:
          pod["nodeSelector"]["sku"] = pod["sku"]
      elif default_cpu_sku is not None and job_training_type == "PSDistJob":
          pod["nodeSelector"]["sku"] = default_cpu_sku
  
      # Assign resource requirement based on job type and default configuration.
      # Pod requiring a full node requires occupying cpu and memory as much as
      # possible on a node. We attempt to achieve this by using cluster-wide
      # cpu usable ratio multiplied into cpu count, and memory usable ratio
      # multiplied into memory size.
      default_cpu_request = None
      default_cpu_limit = None
      default_mem_request = None
      default_mem_limit = None
  
      if job_training_type == "PSDistJob" and dist_role == "worker":
          full_node = True
      else:
          full_node = False
  
      if full_node is True:
          sku = pod["nodeSelector"].get("sku", None)
          sku_info = get_sku_info(sku=sku, config=config)
          if sku_info is not None:
              # Do not restrict the limit for full node worker
              default_cpu_request = cpu_format(sku_info["cpu"],
                                               sku_info["cpu_ratio"])
              default_mem_request = mem_format(sku_info["memory"],
                                               sku_info["memory_ratio"])
      else:
          default_cpu_request = cpu_format(config.get("default_cpurequest"))
          default_cpu_limit = cpu_format(config.get("default_cpulimit"))
          default_mem_request = mem_format(config.get("default_memoryrequest"))
          default_mem_limit = mem_format(config.get("default_memorylimit"))
  
      if "cpurequest" not in pod and default_cpu_request is not None:
          pod["cpurequest"] = default_cpu_request
      if "cpulimit" not in pod and default_cpu_limit is not None:
          pod["cpulimit"] = default_cpu_limit
      if "memoryrequest" not in pod and default_mem_request is not None:
          pod["memoryrequest"] = default_mem_request
      if "memorylimit" not in pod and default_mem_limit is not None:
          pod["memorylimit"] = default_mem_limit
  
      return pod
function_tokens: def enable_cpu_config ( pod , config ) : # Ignore if cpuworker is not enabled enable_cpuworker = config . get ( "enable_cpuworker" , False ) if enable_cpuworker is False : return pod # Only works for 0-GPU job if "resourcegpu" not in pod or int ( pod [ "resourcegpu" ] ) != 0 : return pod # When cpuworker is enabled, CPU job should have gpuType=None if "nodeSelector" not in pod : pod [ "nodeSelector" ] = { } pod [ "nodeSelector" ] [ "gpuType" ] = "None" job_training_type = pod . get ( "jobtrainingtype" , None ) dist_role = pod . get ( "distRole" , None ) # No special config for ps pod. It is always co-located with a worker if dist_role == "ps" : return pod # Add node selector cpuworker=active pod [ "nodeSelector" ] [ "cpuworker" ] = "active" # Add node selector sku=<sku_value> default_cpu_sku = config . get ( "default_cpu_sku" , None ) if "sku" in pod : pod [ "nodeSelector" ] [ "sku" ] = pod [ "sku" ] elif default_cpu_sku is not None and job_training_type == "PSDistJob" : pod [ "nodeSelector" ] [ "sku" ] = default_cpu_sku # Assign resource requirement based on job type and default configuration. # Pod requiring a full node requires occupying cpu and memory as much as # possible on a node. We attempt to achieve this by using cluster-wide # cpu usable ratio multiplied into cpu count, and memory usable ratio # multiplied into memory size. default_cpu_request = None default_cpu_limit = None default_mem_request = None default_mem_limit = None if job_training_type == "PSDistJob" and dist_role == "worker" : full_node = True else : full_node = False if full_node is True : sku = pod [ "nodeSelector" ] . get ( "sku" , None ) sku_info = get_sku_info ( sku = sku , config = config ) if sku_info is not None : # Do not restrict the limit for full node worker default_cpu_request = cpu_format ( sku_info [ "cpu" ] , sku_info [ "cpu_ratio" ] ) default_mem_request = mem_format ( sku_info [ "memory" ] , sku_info [ "memory_ratio" ] ) else : default_cpu_request = cpu_format ( config . get ( "default_cpurequest" ) ) default_cpu_limit = cpu_format ( config . get ( "default_cpulimit" ) ) default_mem_request = mem_format ( config . get ( "default_memoryrequest" ) ) default_mem_limit = mem_format ( config . get ( "default_memorylimit" ) ) if "cpurequest" not in pod and default_cpu_request is not None : pod [ "cpurequest" ] = default_cpu_request if "cpulimit" not in pod and default_cpu_limit is not None : pod [ "cpulimit" ] = default_cpu_limit if "memoryrequest" not in pod and default_mem_request is not None : pod [ "memoryrequest" ] = default_mem_request if "memorylimit" not in pod and default_mem_limit is not None : pod [ "memorylimit" ] = default_mem_limit return pod