Merge pull request #444 from sthaha/fix-estimator-improvements

fix(estimator): logging and machine-spec computation

Author: Sunil Thaha

src/kepler_model/estimate/estimator.py

@@ -136,7 +136,7 @@ def start(self):
         s = self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
         s.bind(self.socket_path)
         s.listen(1)
-        logger.info(f"started serving on {self.socket_path}")
+        logger.info(f"listening on {self.socket_path}")
         try:
             while True:
                 connection, _ = s.accept()
@@ -187,8 +187,13 @@ def sig_handler(signum, frame) -> None:
 )
 def run(log_level: str, machine_spec: str):
     level = getattr(logging, log_level.upper())
-    logging.basicConfig(level=level)
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s %(levelname)s %(filename)s:%(lineno)s: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
 
+    logger.info("starting estimator")
     set_env_from_model_config()
     clean_socket()
     signal.signal(signal.SIGTERM, sig_handler)

src/kepler_model/train/profiler/node_type_index.py

@@ -77,7 +77,7 @@ def discover_spec_values():
         vendor = format_vendor(cpu_info["vendor_id_raw"])
 
     cores = psutil.cpu_count(logical=True)
-    chips = max(1, int(subprocess.check_output('cat /proc/cpuinfo | grep "physical id" | sort -u | wc -l', shell=True)))
+    chips = max(1, int(subprocess.check_output('grep "physical id" /proc/cpuinfo | sort -u | wc -l', shell=True)))
     threads_per_core = max(1, cores // psutil.cpu_count(logical=False))
     memory = psutil.virtual_memory().total
     memory_gb = int(memory / GB)

src/kepler_model/util/config.py

@@ -13,12 +13,16 @@
 #################################################
 
 import os
+import logging
 
 import requests
 
 from .loader import base_model_url, default_init_model_name, default_pipelines, default_train_output_pipeline, get_pipeline_url, get_url
 from .train_types import FeatureGroup, ModelOutputType, is_output_type_supported
 
+logger = logging.getLogger(__name__)
+
+
 # must be writable (for shared volume mount)
 MNT_PATH = "/mnt"
 # can be read only (for configmap mount)
@@ -122,7 +126,7 @@ def set_env_from_model_config():
         splits = line.split("=")
         if len(splits) > 1:
             os.environ[splits[0].strip()] = splits[1].strip()
-            print(f"set {splits[0]} to {splits[1]}.")
+            logging.info(f"set {splits[0]} to {splits[1]}.")
 
 
 def is_estimator_enable(prefix):
@@ -152,7 +156,7 @@ def get_init_model_url(energy_source, output_type, model_topurl=model_topurl):
     for prefix in modelConfigPrefix:
         if get_energy_source(prefix) == energy_source:
             modelURL = get_init_url(prefix)
-            print("get init url", modelURL)
+            logger.info("get init url", modelURL)
             url = get_url(
                 feature_group=FeatureGroup.BPFOnly,
                 output_type=ModelOutputType[output_type],
@@ -175,7 +179,7 @@ def get_init_model_url(energy_source, output_type, model_topurl=model_topurl):
                         response = requests.get(url)
                         if response.status_code == 200:
                             modelURL = url
-                print(f"init URL is not set, use {modelURL}")
+                logger.info(f"init URL is not set, use {modelURL}")
             return modelURL
-    print(f"no match config for {output_type}, {energy_source}")
+    logger.info(f"no match config for {output_type}, {energy_source}")
     return ""

src/kepler_model/util/loader.py

@@ -76,7 +76,7 @@ def load_json(path: str, name: str = ""):
             res = json.load(f)
         return res
     except Exception as err:
-        logger.error(f"fail to load json {filepath}: {err}")
+        logger.error(f"failed to load json {filepath}: {err}")
         return None