code: def set_affinity(
      gpu_id,
      nproc_per_node=None,
      *,
      mode: Union[
          str, AffinityMode
      ] = AffinityMode.socket_unique_contiguous,
      cores="all_logical",
      balanced=True,
  ):
      """
      The process is assigned with a proper CPU affinity that matches CPU-GPU
      hardware architecture on a given platform. Usually, it improves and
      stabilizes the performance of deep learning training workloads.
  
      This function assumes that the workload runs in multi-process single-device
      mode (there are multiple training processes, and each process is running on
      a single GPU). This is typical for multi-GPU data-parallel training
      workloads (e.g., using `torch.nn.parallel.DistributedDataParallel`).
  
      Available affinity modes:
      * 'socket' - the process is assigned with all available physical CPU cores
      from the CPU socket connected to the GPU with a given id.
      * 'socket_single' - the process is assigned with the first available
      physical CPU core from the list of all CPU cores from the CPU socket
      connected to the GPU with a given id (multiple GPUs could be assigned with
      the same CPU core).
      * 'socket_single_unique' - the process is assigned with a single unique
      available physical CPU core from the list of all CPU cores from the CPU
      socket connected to the GPU with a given id.
      * 'socket_unique_interleaved' - the process is assigned with a unique
      subset of available physical CPU cores from the CPU socket connected to a
      GPU with a given id, cores are assigned with interleaved indexing pattern
      * 'socket_unique_contiguous' - (the default) the process is assigned with a
      unique subset of available physical CPU cores from the CPU socket connected
      to a GPU with a given id, cores are assigned with contiguous indexing
      pattern
  
      Available "cores" modes:
      * 'all_logical' - assigns the process with all logical cores associated with
      a given corresponding physical core (i.e., automatically includes all
      available hyperthreading siblings)
      * 'single_logical' - assigns the process with only one logical core
      associated with a given corresponding physical core (i.e., excludes
      hyperthreading siblings)
  
      'socket_unique_contiguous' is the recommended mode for deep learning
      training workloads on NVIDIA DGX machines.
  
      Args:
          gpu_id: integer index of a GPU, value from 0 to 'nproc_per_node' - 1
          nproc_per_node: number of processes per node
          mode: affinity mode
          balanced: assign an equal number of physical cores to each process,
              affects only 'socket_unique_interleaved' and
              'socket_unique_contiguous' affinity modes
          cores: 'all_logical' or 'single_logical'
  
      Returns a set of logical CPU cores on which the process is eligible to run.
  
      Example:
  
      import argparse
      import os
  
      import gpu_affinity
      import torch
  
  
      def main():
          parser = argparse.ArgumentParser()
          parser.add_argument(
              '--local_rank',
              type=int,
              default=os.getenv('LOCAL_RANK', 0),
          )
          args = parser.parse_args()
  
          nproc_per_node = torch.cuda.device_count()
  
          affinity = gpu_affinity.set_affinity(args.local_rank, nproc_per_node)
          print(f'{args.local_rank}: core affinity: {affinity}')
  
  
      if __name__ == "__main__":
          main()
  
      Launch the example with:
      python -m torch.distributed.launch --nproc_per_node <#GPUs> example.py
  
  
      WARNING: On DGX A100, only half of the CPU cores have direct access to GPUs.
      This function restricts execution only to the CPU cores directly connected
      to GPUs, so on DGX A100, it will limit the code to half of the CPU cores and
      half of CPU memory bandwidth (which may be fine for many DL models).
  
      WARNING: Intel's OpenMP implementation resets affinity on the first call to
      an OpenMP function after a fork. It's recommended to run with env variable:
      `KMP_AFFINITY=disabled` if the affinity set by gpu_affinity should be
      preserved after a fork (e.g. in PyTorch DataLoader workers).
      """
      if not isinstance(mode, AffinityMode):
          mode = AffinityMode[mode]
      pynvml.nvmlInit()
      if nproc_per_node is None:
          nproc_per_node = pynvml.nvmlDeviceGetCount()
  
      if mode == AffinityMode.none:
          pass
      elif mode == AffinityMode.socket:
          set_socket_affinity(gpu_id, nproc_per_node, cores)
      elif mode == AffinityMode.socket_single:
          set_socket_single_affinity(gpu_id, nproc_per_node, cores)
      elif mode == AffinityMode.socket_single_unique:
          set_socket_single_unique_affinity(
              gpu_id, nproc_per_node, cores
          )
      elif mode == AffinityMode.socket_unique_interleaved:
          set_socket_unique_affinity(
              gpu_id, nproc_per_node, cores, "interleaved", balanced
          )
      elif mode == AffinityMode.socket_unique_contiguous:
          set_socket_unique_affinity(
              gpu_id, nproc_per_node, cores, "contiguous", balanced
          )
      else:
          raise RuntimeError("Unknown affinity mode")
  
      affinity = os.sched_getaffinity(0)
      return affinity
  
imports: import collections
  import itertools
  import os
  import pathlib
  import re
  import pynvml
  from typing import Union
  from enum import Enum, auto
label_desc: Use image node_id 13 for calling a global function with example usage: set_affinity(gpu_id, nproc_per_node) and returns: affinity