feat: add machine_spec to metadata

Signed-off-by: Sunyanan Choochotkaew <[email protected]>

Author: sunya-ch

Makefile

@@ -57,7 +57,7 @@ run-estimator-with-test-spec:
 	$(CTR_CMD) run --rm -d --platform linux/amd64 \
 		--name estimator \
 		$(TEST_IMAGE) \
-		/bin/bash -c "estimator --machine-spec tests/data/machine/spec.json"
+		/bin/bash -c "estimator --machine-spec tests/data/machine/spec.json --log-level debug"
 
 run-collector-client:
 	$(CTR_CMD) exec estimator /bin/bash -c \
@@ -129,7 +129,7 @@ run-model-server-with-db:
 		-p 8100:8100 \
 		--name model-server $(TEST_IMAGE) \
 		model-server
-	while ! docker logs model-server 2>&1 | grep -q 'Running on all'; do \
+	while ! $(CTR_CMD) logs model-server 2>&1 | grep -q 'Running on all'; do \
 		echo "... waiting for model-server to serve";  sleep 5; \
 	done
 

src/kepler_model/estimate/estimator.py

@@ -12,9 +12,9 @@
 from kepler_model.estimate.archived_model import get_achived_model
 from kepler_model.estimate.model.model import load_downloaded_model
 from kepler_model.estimate.model_server_connector import is_model_server_enabled, make_request
-from kepler_model.train.profiler.node_type_index import get_machine_spec
+from kepler_model.train.profiler.node_type_index import NodeTypeSpec, discover_spec_values, get_machine_spec
 from kepler_model.util.config import SERVE_SOCKET, download_path, set_env_from_model_config
-from kepler_model.util.loader import get_download_output_path
+from kepler_model.util.loader import get_download_output_path, load_metadata
 from kepler_model.util.train_types import ModelOutputType, convert_enery_source, is_output_type_supported
 
 ###############################################
@@ -42,7 +42,7 @@ def __init__(self, metrics, values, output_type, source, system_features, system
 
 loaded_model = dict()
 
-def handle_request(data: str, machine_spec=None) -> dict:
+def handle_request(data: str, machine_spec=None, discovered_core=None) -> dict:
     try:
         power_request = json.loads(data, object_hook=lambda d: PowerRequest(**d))
     except Exception as e:
@@ -94,22 +94,41 @@ def handle_request(data: str, machine_spec=None) -> dict:
         if loaded_item is not None and loaded_item.estimator is not None:
             loaded_model[output_type.name][power_request.energy_source] = loaded_item
             logger.info(f"set model {loaded_item.model_name} for {output_type.name} ({power_request.energy_source})")
+        else:
+            msg = f"load item for {power_request.energy_source} is none"
+            logger.error(msg)
+            return {"powers": dict(), "msg": msg}
 
     model = loaded_model[output_type.name][power_request.energy_source]
     powers, msg = model.get_power(power_request.datapoint)
     if msg != "":
         logger.info(f"{model.model_name} failed to predict; removed: {msg}")
         if output_path != "" and os.path.exists(output_path):
             shutil.rmtree(output_path)
-
-    return {"powers": powers, "msg": msg}
-
+    response = {"powers": powers, "msg": msg}
+    # add core_ratio if applicable
+    core_ratio = 1
+    if discovered_core is not None and discovered_core > 0:
+        metadata = load_metadata(output_path)
+        if metadata is not None and "machine_spec" in metadata:
+            model_spec = NodeTypeSpec(**metadata["machine_spec"])
+            model_cores = model_spec.get_cores()
+            if model_cores > 0:
+                core_ratio = discovered_core/model_cores
+            logger.debug(f"model cores: {model_cores}")
+        logger.debug(f"metadata: {metadata}")
+    response["core_ratio"] = core_ratio
+
+    return response
 
 class EstimatorServer:
     def __init__(self, socket_path, machine_spec):
         self.socket_path = socket_path
         self.machine_spec = machine_spec
-        logger.info(f"initialize EstimatorServer with spec={machine_spec}")
+        spec_values = discover_spec_values()
+        discovered_spec = NodeTypeSpec(**spec_values)
+        self.discovered_core = discovered_spec.get_cores()
+        logger.info(f"initialize EstimatorServer with spec={machine_spec}, discovered_core={self.discovered_core}")
 
     def start(self):
         s = self.socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
@@ -135,11 +154,10 @@ def accepted(self, connection):
             if shunk is None or shunk.decode()[-1] == "}":
                 break
         decoded_data = data.decode()
-        y = handle_request(decoded_data, self.machine_spec)
+        y = handle_request(decoded_data, self.machine_spec, self.discovered_core)
         response = json.dumps(y)
         connection.send(response.encode())
 
-
 def clean_socket():
     logger.info("clean socket")
     if os.path.exists(SERVE_SOCKET):

src/kepler_model/estimate/model/model.py

@@ -165,6 +165,10 @@ def load_model(model_path):
         return None
 
     metadata["model_path"] = model_path
+    logger.info(f"load metadata {metadata}")
+    # need to delete machine_spec before load model
+    if "machine_spec" in metadata:
+        del metadata["machine_spec"]
     metadata_str = json.dumps(metadata)
     try:
         model = json.loads(metadata_str, object_hook=lambda d: Model(**d))

src/kepler_model/server/model_server.py

@@ -32,10 +32,11 @@
     is_matched_type,
     is_valid_model,
     load_json,
+    load_metadata,
     load_weight,
     parse_filters,
 )
-from kepler_model.util.saver import WEIGHT_FILENAME
+from kepler_model.util.saver import WEIGHT_FILENAME, save_metadata
 from kepler_model.util.train_types import (
     FeatureGroup,
     FeatureGroups,
@@ -216,6 +217,13 @@ def get_model():
     logger.info(f"response: model {best_model['model_name']} by {best_model['features']} with {ERROR_KEY}={best_model[ERROR_KEY]} selected with uncertainty={best_uncertainty}, looseness={best_looseness}")
     if req.weight:
         try:
+            # add this condition to provide compatibility to old version
+            # the old version always set default node_type
+            if req.node_type == any_node_type:
+                best_response["model_name"] = best_model["model_name"]
+                if "machine_spec" in best_model:
+                    best_response["machine_spec"] = best_model["machine_spec"]
+                best_response[ERROR_KEY] = best_model[ERROR_KEY]
             response = app.response_class(response=json.dumps(best_response), status=200, mimetype="application/json")
             return response
         except ValueError as err:
@@ -344,7 +352,30 @@ def load_init_pipeline():
         # remove downloaded zip
         os.remove(tmp_filepath)
     set_pipelines()
-
+    fill_machine_spec()
+
+def fill_machine_spec():
+    for energy_source in PowerSourceMap.keys():
+        if energy_source in pipelineName:
+            pipeline_name = pipelineName[energy_source]
+            if pipeline_name in nodeCollection:
+                node_collection = nodeCollection[pipeline_name]
+                for output_type in ModelOutputType:
+                    for feature_group in FeatureGroup:
+                        valid_group_path = get_model_group_path(model_toppath, output_type, feature_group, energy_source, pipeline_name=pipeline_name)
+                        for f in  os.listdir(valid_group_path):
+                            path = os.path.join(valid_group_path, f)
+                            if not os.path.isfile(path):
+                                metadata = load_metadata(path)
+                                if metadata is not None:
+                                    if "machine_spec" not in metadata and "model_name" in metadata:
+                                        model_name = metadata["model_name"]
+                                        node_type = get_node_type_from_name(model_name)
+                                        if node_type in node_collection.node_type_index:
+                                            metadata["machine_spec"] = node_collection.node_type_index[node_type].get_json()["attrs"]
+                                            save_metadata(path, metadata)
+                                            save_path = os.path.join(valid_group_path, model_name)
+                                            shutil.make_archive(save_path, "zip", save_path)
 
 @click.command()
 @click.option(

src/kepler_model/train/pipeline.py

@@ -44,6 +44,8 @@ def __init__(self, name, trainers, extractor, isolator):
         self.metadata["abs_trainers"] = [trainer.__class__.__name__ for trainer in trainers if trainer.node_level]
         self.metadata["dyn_trainers"] = [trainer.__class__.__name__ for trainer in trainers if not trainer.node_level]
         self.metadata["init_time"] = time_to_str(datetime.datetime.utcnow())
+        for trainer in trainers:
+            trainer.set_node_type_index(self.node_collection.node_type_index)
 
     def get_abs_data(self, query_results, energy_components, feature_group, energy_source, aggr):
         extracted_data, power_labels, _, _ = self.extractor.extract(query_results, energy_components, feature_group, energy_source, node_level=True, aggr=aggr)

src/kepler_model/train/profiler/node_type_index.py

@@ -156,8 +156,10 @@ def add_member(self, machine_id):
     def get_size(self):
         return len(self.members)
 
-    def get_cores(self):
-        return self.attrs[NodeAttribute.CORES]
+    def get_cores(self) -> int:
+        if attr_has_value(self.attrs, NodeAttribute.CORES):
+            return int(self.attrs[NodeAttribute.CORES])
+        return 0
 
     # check the comparing node-type spec is covered by this node-type spec
     def cover(self, compare_spec):

src/kepler_model/train/trainer/__init__.py

@@ -61,6 +61,10 @@ def __init__(self, model_class, energy_components, feature_group, energy_source,
         self.node_models = dict()
         self.node_scalers = dict()
         self.scaler_type = scaler_type
+        self.node_type_index = dict()
+
+    def set_node_type_index(self, node_type_index):
+        self.node_type_index = node_type_index
 
     def _get_save_path(self, node_type):
         save_path = get_save_path(self.group_path, self.trainer_name, node_type=node_type)
@@ -214,13 +218,17 @@ def save_metadata(self, node_type, mae, mae_map, mape, mape_map, item):
         save_path = self._get_save_path(node_type)
         model_name, model_file = self._model_filename(node_type)
         item["model_name"] = model_name
+        item["trainer"] = self.trainer_name
         item["model_class"] = self.model_class
         item["model_file"] = model_file
         item["features"] = self.features
         item["fe_files"] = [] if not hasattr(self, "fe_files") else self.fe_files
         item["output_type"] = self.output_type.name
         item["mae"] = mae
         item["mape"] = mape
+        if node_type in self.node_type_index:
+            item["node_type"] = node_type
+            item["machine_spec"] = self.node_type_index[node_type].get_json()["attrs"]
         item.update(mae_map)
         item.update(mape_map)
         self.metadata = item

tests/estimator_model_request_test.py

@@ -7,6 +7,11 @@
 # - kepler-model-server is not connected, but some achived models can be download via URL.
 #   - set sample model and make a dummy valid PowerRequest and another invalid PowerRequest
 #
+# Requires
+# - run `model-server``
+# - run `pytest tests/pipeline_test.py` (run once to get models)
+# - run `MODEL_PATH=$(pwd)/src/kepler_model/models python tests/http_server.py`
+#
 #########################
 # import external modules
 import json
@@ -66,7 +71,7 @@ def test_model_request():
         if url != "":
             print("Download: ", url)
             response = requests.get(url)
-            assert response.status_code == 200, "init url must be set and valid"
+            assert response.status_code == 200, f"init url {url} must be set and valid"
             output_path = get_download_output_path(download_path, energy_source, output_type)
             if output_type_name in loaded_model and energy_source in loaded_model[output_type.name]:
                 del loaded_model[output_type_name][energy_source]

tests/model_server_test.py

@@ -7,8 +7,9 @@
 import requests
 
 from kepler_model.server.model_server import MODEL_SERVER_PORT
-from kepler_model.train.profiler.node_type_index import NodeAttribute, attr_has_value
-from kepler_model.util.config import download_path
+from kepler_model.train.profiler.node_type_index import NodeAttribute, NodeTypeSpec, attr_has_value
+from kepler_model.util.config import ERROR_KEY, download_path
+from kepler_model.util.loader import any_node_type
 from kepler_model.util.train_types import FeatureGroup, FeatureGroups, ModelOutputType
 
 TMP_FILE = "tmp.zip"
@@ -18,17 +19,28 @@ def get_model_request_json(metrics, output_type, node_type, weight, trainer_name
     return {"metrics": metrics, "output_type": output_type.name, "node_type": node_type, "weight": weight, "trainer_name": trainer_name, "source": energy_source}
 
 
-def make_request(metrics, output_type, node_type=-1, weight=False, trainer_name="", energy_source="rapl-sysfs"):
+def make_request(metrics, output_type, node_type=any_node_type, weight=False, trainer_name="", energy_source="rapl-sysfs"):
     model_request = get_model_request_json(metrics, output_type, node_type, weight, trainer_name, energy_source)
     response = requests.post(f"http://localhost:{MODEL_SERVER_PORT}/model", json=model_request)
     assert response.status_code == 200, response.text
     if weight:
         weight_dict = json.loads(response.text)
         assert len(weight_dict) > 0, "weight dict must contain one or more than one component"
-        for weight_values in weight_dict.values():
-            weight_length = len(weight_values["All_Weights"]["Numerical_Variables"])
-            expected_length = len(metrics)
-            assert weight_length <= expected_length, f"weight metrics should covered by the requested {weight_length} > {expected_length}"
+        if node_type == any_node_type:
+            assert "model_name" in weight_dict
+            assert "machine_spec" in weight_dict
+            assert ERROR_KEY in weight_dict
+            assert len(weight_dict["model_name"]) > 0
+            spec_values = weight_dict["machine_spec"]
+            spec = NodeTypeSpec(**spec_values)
+            assert spec.get_cores() > 0
+
+        for key, values in weight_dict.items():
+            if key not in ["model_name", "machine_spec", ERROR_KEY]:
+                if "All_Weights" in values:
+                    weight_length = len(values["All_Weights"]["Numerical_Variables"])
+                    expected_length = len(metrics)
+                    assert weight_length <= expected_length, f"weight metrics should covered by the requested {weight_length} > {expected_length}"
     else:
         output_path = os.path.join(download_path, output_type.name)
         if os.path.exists(output_path):