code: def set_socket_unique_affinity(
      gpu_id, nproc_per_node, cores, mode, balanced=True
  ):
      """
      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.
      Assignment automatically includes hyperthreading siblings (if siblings are
      available).
  
      Args:
          gpu_id: index of a GPU
          nproc_per_node: number of processes per node
          cores: 'all_logical' or 'single_logical'
          mode: 'contiguous' or 'interleaved'
          balanced: assign an equal number of physical cores to each process,
      """
      grouped_socket_affinities = get_grouped_socket_affinities(
          nproc_per_node
      )
  
      grouped_socket_affinities_to_device_ids = collections.defaultdict(
          list
      )
  
      for idx, grouped_socket_affinity in enumerate(
          grouped_socket_affinities
      ):
          grouped_socket_affinities_to_device_ids[
              tuple(grouped_socket_affinity)
          ].append(idx)
  
      # compute minimal number of physical cores per GPU across all GPUs and
      # sockets, code assigns this number of cores per GPU if balanced == True
      min_physical_cores_per_gpu = min(
          [
              len(cores) // len(gpus)
              for cores, gpus in grouped_socket_affinities_to_device_ids.items()
          ]
      )
  
      grouped_unique_affinities = [None] * nproc_per_node
  
      for (
          grouped_socket_affinity,
          device_ids,
      ) in grouped_socket_affinities_to_device_ids.items():
          devices_per_group = len(device_ids)
          if balanced:
              cores_per_device = min_physical_cores_per_gpu
              grouped_socket_affinity = grouped_socket_affinity[
                  : devices_per_group * min_physical_cores_per_gpu
              ]
          else:
              cores_per_device = (
                  len(grouped_socket_affinity) // devices_per_group
              )
  
          for socket_subgroup_id, device_id in enumerate(device_ids):
              # In theory there should be no difference in performance between
              # 'interleaved' and 'contiguous' pattern on Intel-based DGX-1,
              # but 'contiguous' should be better for DGX A100 because on AMD
              # Rome 4 consecutive cores are sharing L3 cache.
              # TODO: code doesn't attempt to automatically detect layout of
              # L3 cache, also external environment may already exclude some
              # cores, this code makes no attempt to detect it and to align
              # mapping to multiples of 4.
  
              if mode == "interleaved":
                  unique_grouped_affinity = list(
                      grouped_socket_affinity[
                          socket_subgroup_id::devices_per_group
                      ]
                  )
              elif mode == "contiguous":
                  unique_grouped_affinity = list(
                      grouped_socket_affinity[
                          socket_subgroup_id
                          * cores_per_device : (socket_subgroup_id + 1)
                          * cores_per_device
                      ]
                  )
              else:
                  raise RuntimeError(
                      "Unknown set_socket_unique_affinity mode"
                  )
  
              grouped_unique_affinities[
                  device_id
              ] = unique_grouped_affinity
  
      ungrouped_affinities = ungroup_affinities(
          grouped_unique_affinities, cores
      )
      os.sched_setaffinity(0, ungrouped_affinities[gpu_id])
  
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 12 for calling a global function with example usage: set_socket_unique_affinity(gpu_id, nproc_per_node, cores, mode, balanced) without return types