chore: add type hint for select_best_model,is_valid_model

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

Author: sunya-ch

src/kepler_model/estimate/model_server_connector.py

@@ -75,15 +75,15 @@ def list_all_models(energy_source=None, output_type=None, feature_group=None, no
         endpoint = get_model_server_list_endpoint()
         params= {}
         if energy_source:
-            params[ModelListParam.EnergySource] = energy_source
+            params[ModelListParam.EnergySource.value] = energy_source
         if output_type:
-            params[ModelListParam.OutputType] = output_type
+            params[ModelListParam.OutputType.value] = output_type
         if feature_group:
-            params[ModelListParam.FeatureGroup] = feature_group
+            params[ModelListParam.FeatureGroup.value] = feature_group
         if node_type:
-           params[ModelListParam.NodeType] = node_type
+           params[ModelListParam.NodeType.value] = node_type
         if filter:
-            params[ModelListParam.Filter] = filter
+            params[ModelListParam.Filter.value] = filter
 
         response = requests.get(endpoint, params=params)
     except Exception as err:

src/kepler_model/server/model_server.py

@@ -101,7 +101,7 @@ class ModelListParam(enum.Enum):
 """
 
 
-def select_best_model(spec, valid_group_path, filters, energy_source, pipeline_name="", trainer_name="", node_type=any_node_type, weight=False):
+def select_best_model(spec, valid_group_path: str, filters: dict, energy_source: str, pipeline_name: str="", trainer_name: str="", node_type: int=any_node_type, weight: bool=False):
     model_names = [f for f in os.listdir(valid_group_path) if f != CHECKPOINT_FOLDERNAME and not os.path.isfile(os.path.join(valid_group_path, f)) and (trainer_name == "" or trainer_name in f)]
     if weight:
         model_names = [name for name in model_names if name.split("_")[0] in weight_support_trainers]
@@ -194,11 +194,11 @@ def get_model():
 # get_available_models: return name list of best-candidate pipelines
 @app.route(MODEL_SERVER_MODEL_LIST_PATH, methods=["GET"])
 def get_available_models():
-    fg = request.args.get(ModelListParam.FeatureGroup)
-    ot = request.args.get(ModelListParam.OutputType)
-    energy_source = request.args.get(ModelListParam.EnergySource)
-    node_type = request.args.get(ModelListParam.NodeType)
-    filter = request.args.get(ModelListParam.Filter)
+    fg = request.args.get(ModelListParam.FeatureGroup.value)
+    ot = request.args.get(ModelListParam.OutputType.value)
+    energy_source = request.args.get(ModelListParam.EnergySource.value)
+    node_type = request.args.get(ModelListParam.NodeType.value)
+    filter = request.args.get(ModelListParam.Filter.value)
 
     try:
         if fg is None:

src/kepler_model/util/loader.py

@@ -173,7 +173,7 @@ def parse_filters(filter):
     return filters
 
 
-def is_valid_model(metadata, filters):
+def is_valid_model(metadata, filters: dict):
     for attrb, val in filters.items():
         if not hasattr(metadata, attrb) or getattr(metadata, attrb) is None:
             print("{} has no {}".format(metadata["model_name"], attrb))