class_docstr: A context manager to enable CPU affinity for data loader workers
  (only used when running on CPU devices).
  
  Affinitization places data loader workers threads on specific CPU cores.
  In effect, it allows for more efficient local memory allocation and reduces
  remote memory calls.
  Every time a process or thread moves from one core to another, registers
  and caches need to be flushed and reloaded.
  This can become very costly if it happens often, and our threads may also
  no longer be close to their data, or be able to share data in a cache.
  
  See `here <https://pytorch-geometric.readthedocs.io/en/latest/advanced/
  cpu_affinity.html>`__ for the accompanying tutorial.
  
  .. warning::
  
      To correctly affinitize compute threads (*i.e.* with
      :obj:`KMP_AFFINITY`), please make sure that you exclude
      :obj:`loader_cores` from the list of cores available for the main
      process.
      This will cause core oversubsription and exacerbate performance.
  
  .. code-block:: python
  
      loader = NeigborLoader(data, num_workers=3)
      with loader.enable_cpu_affinity(loader_cores=[0, 1, 2]):
          for batch in loader:
              pass
docstr: Enables CPU affinity.
  
  Args:
      loader_cores ([int], optional): List of CPU cores to which data
          loader workers should affinitize to.
          By default, it will affinitize to :obj:`numa0` cores.
          If used with :obj:`"spawn"` multiprocessing context, it will
          automatically enable multithreading and use multiple cores
          per each worker.
docstr_tok: ["Enables","CPU","affinity",".","Args",":","loader_cores","(","[","int","]",",","optional",")",":","List","of","CPU","cores","to","which","data","loader","workers","should","affinitize","to",".","By","default",",","it","will","affinitize","to",":","obj",":","`","numa0","`","cores",".","If","used","with",":","obj",":","`","``","spawn","''","`","multiprocessing","context",",","it","will","automatically","enable","multithreading","and","use","multiple","cores","per","each","worker","."]
code: def enable_cpu_affinity(
      self,
      loader_cores: Optional[Union[List[List[int]], List[int]]] = None,
  ) -> None:
      r"""Enables CPU affinity.
  
      Args:
          loader_cores ([int], optional): List of CPU cores to which data
              loader workers should affinitize to.
              By default, it will affinitize to :obj:`numa0` cores.
              If used with :obj:`"spawn"` multiprocessing context, it will
              automatically enable multithreading and use multiple cores
              per each worker.
      """
      if not self.num_workers > 0:
          raise ValueError(
              f"'enable_cpu_affinity' should be used with at least one "
              f"worker (got {self.num_workers})"
          )
      if loader_cores and len(loader_cores) != self.num_workers:
          raise ValueError(
              f"The number of loader cores (got {len(loader_cores)}) "
              f"in 'enable_cpu_affinity' should match with the number "
              f"of workers (got {self.num_workers})"
          )
      if isinstance(self.data, HeteroData):
          warnings.warn(
              "Due to conflicting parallelization methods it is not advised "
              "to use affinitization with 'HeteroData' datasets. "
              "Use `enable_multithreading` for better performance."
          )
  
      self.loader_cores = loader_cores[:] if loader_cores else None
      if self.loader_cores is None:
          numa_info = get_numa_nodes_cores()
  
          if numa_info and len(numa_info[0]) > self.num_workers:
              # Take one thread per each node 0 core:
              node0_cores = [cpus[0] for core_id, cpus in numa_info[0]]
              node0_cores.sort()
          else:
              node0_cores = list(range(psutil.cpu_count(logical=False)))
  
          if len(node0_cores) < self.num_workers:
              raise ValueError(
                  f"More workers (got {self.num_workers}) than available "
                  f"cores (got {len(node0_cores)})"
              )
  
          # Set default loader core IDs:
          if torch.multiprocessing.get_context()._name == "spawn":
              work_thread_pool = int(
                  len(node0_cores) / self.num_workers
              )
              self.loader_cores = [
                  list(
                      range(
                          work_thread_pool * i,
                          work_thread_pool * (i + 1),
                      )
                  )
                  for i in range(self.num_workers)
              ]
          else:
              self.loader_cores = node0_cores[: self.num_workers]
  
      self._old_worker_init_fn = WorkerInitWrapper(self.worker_init_fn)
      try:
          self.worker_init_fn = self._aff_init_fn
          logging.debug(
              f"{self.num_workers} data loader workers are "
              f"assigned to CPUs {self.loader_cores}"
          )
          yield
      finally:
          self.worker_init_fn = self._old_worker_init_fn
  
imports: import glob
  import logging
  import os
  import os.path
  import warnings
  from contextlib import contextmanager
  from typing import Any, Callable, Dict, List, Optional, Union
  import psutil
  import torch
  from torch_geometric.data import HeteroData
label_desc: Use image node_id 2 for calling the AffinityMixin obj's underlying member method code with example usage: obj.enable_cpu_affinity(loader_cores) without return types