From 85ccfc4734a9fa471bf3f5e5771590a6ed866a33 Mon Sep 17 00:00:00 2001
From: rbrugaro <rita.brugarolas.brufau@intel.com>
Date: Wed, 21 Aug 2024 15:31:48 -0700
Subject: [PATCH] reorg imports, moved checks to import_utils, remove prints
 for logger

---
 optimum/intel/utils/__init__.py       |  1 +
 optimum/intel/utils/import_utils.py   | 13 ++++++++++++
 optimum/intel/utils/modeling_utils.py | 29 ++++++++++++++-------------
 3 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/optimum/intel/utils/__init__.py b/optimum/intel/utils/__init__.py
index d77588f896..50cdfa143e 100644
--- a/optimum/intel/utils/__init__.py
+++ b/optimum/intel/utils/__init__.py
@@ -22,6 +22,7 @@
     is_neural_compressor_available,
     is_neural_compressor_version,
     is_nncf_available,
+    is_numa_available,
     is_openvino_available,
     is_torch_version,
     is_transformers_available,
diff --git a/optimum/intel/utils/import_utils.py b/optimum/intel/utils/import_utils.py
index 6be0aac47a..5d0e3e2471 100644
--- a/optimum/intel/utils/import_utils.py
+++ b/optimum/intel/utils/import_utils.py
@@ -150,6 +150,15 @@
     except importlib_metadata.PackageNotFoundError:
         _accelerate_available = False
 
+_numa_available = importlib.util.find_spec("numa") is not None
+_numa_version = "N/A"
+
+if _numa_available:
+    try:
+        _numa_version = importlib_metadata.version("numa")
+    except importlib_metadata.PackageNotFoundError:
+        _numa_available = False
+
 
 def is_transformers_available():
     return _transformers_available
@@ -272,6 +281,10 @@ def is_accelerate_available():
     return _accelerate_available
 
 
+def is_numa_available():
+    return _numa_available
+
+
 # This function was copied from: https://github.com/huggingface/accelerate/blob/874c4967d94badd24f893064cc3bef45f57cadf7/src/accelerate/utils/versions.py#L319
 def compare_versions(library_or_version: Union[str, Version], operation: str, requirement_version: str):
     """
diff --git a/optimum/intel/utils/modeling_utils.py b/optimum/intel/utils/modeling_utils.py
index 95e4795de7..0eb92ae052 100644
--- a/optimum/intel/utils/modeling_utils.py
+++ b/optimum/intel/utils/modeling_utils.py
@@ -12,17 +12,25 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import logging
+import math
 import os
+import platform
 import re
 from pathlib import Path
 from typing import List, Optional, Union
 
+import psutil
 import torch
 from huggingface_hub import HfApi, HfFolder
 
+from .import_utils import is_numa_available
+
 
 MULTI_QUERY_ATTN_MODELS = {"falcon", "gpt_bigcode"}
 
+logger = logging.getLogger(__name__)
+
 
 def get_model_device(model: torch.nn.Module) -> torch.device:
     """
@@ -145,17 +153,10 @@ def bind_cores_for_best_perf():
         None
 
     """
-
-    import importlib.util
-    import platform
-
     system = platform.system()
     if system == "Linux":
-        if importlib.util.find_spec("numa") is not None:
-            import math
-
+        if is_numa_available():
             import numa
-            import psutil
 
             local_size = get_int_from_env(
                 ["MPI_LOCALNRANKS", "OMPI_COMM_WORLD_LOCAL_SIZE", "MV2_COMM_WORLD_LOCAL_SIZE"], 1
@@ -169,12 +170,11 @@ def bind_cores_for_best_perf():
             node_id = int(rank_id / rank_per_node)
             rank_offset_per_node = rank_id % rank_per_node
             if os.getenv("OMP_NUM_THREADS") is None:
-                # set OMP_NUM_THREADS to num of physical cores per socket
                 num_cpus_per_rank = max(int(num_cpus_per_nodes / rank_per_node), 1)
-                print("setting OMP_NUM_THREADS to", num_cpus_per_rank)
+                logger.info(f"Setting OMP_NUM_THREADS to {num_cpus_per_rank} for better performance")
             else:
                 num_cpus_per_rank = int(os.getenv("OMP_NUM_THREADS"))
-                print("OMP_NUM_THREADS already set to ", num_cpus_per_rank)
+                logger.info(f"OMP_NUM_THREADS already set to  {num_cpus_per_rank}")
             if len(numa.get_membind()) == nodes:
                 # if numa memory binding is not set, set it to the node where the rank is running
                 numa.set_membind([node_id])
@@ -188,8 +188,9 @@ def bind_cores_for_best_perf():
                     0,
                     list(numa.node_to_cpus(node_id))[cpu_start : cpu_start + num_cpus_per_rank],
                 )
-            print(f"affinity={numa.get_affinity(0)}, membind = {numa.get_membind()}")
+            logger.info(f"affinity={numa.get_affinity(0)}, membind = {numa.get_membind()}")
         else:
-            print("numa module not found, skipping binding cores")
+            logger.warning("numa module not found, skipping binding cores")
+
     else:
-        print("OS not supported, skipping binding cores")
+        logger.error("bind_cores_for_best_perf: OS not supported, skipping binding cores")