Merge pull request #441 from vimalk78/remove-xgboost

fix(xgboost): Remove xgboost dependency to reduce image sizes

Author: Sunil Thaha

.github/workflows/integration-test.yml

@@ -35,7 +35,6 @@ env:
 
 jobs:
   run-integration:
-    if: ${{ (inputs.base_change != 'true') || ( (inputs.base_change == 'true') && (inputs.docker_secret == 'true') ) }}
     runs-on: ubuntu-20.04
     steps:
       - name: use Kepler action to deploy cluster

.github/workflows/tekton-test.yml

@@ -42,7 +42,6 @@ env:
 
 jobs:
   tekton-test:
-    if: ${{ (inputs.base_change != 'true') || ( (inputs.base_change == 'true') && (inputs.docker_secret == 'true') ) }}
     runs-on: ubuntu-latest
     steps:
       - name: checkout

pyproject.toml

@@ -40,7 +40,6 @@ dependencies = [
   "scipy==1.14.1",
   "seaborn==0.13.2",
   "Werkzeug==3.0.4",
-  "xgboost==2.1.1",
   "boto3==1.34.155",
   "pymarkdownlnt==0.9.22",
   "yamllint==1.35.1",

src/kepler_model/estimate/model/model.py

@@ -1,11 +1,15 @@
 import json
 import logging
+import importlib.util as import_util
 
 import pandas as pd
 
 from kepler_model.estimate.model.curvefit_model import CurveFitModelEstimator
 from kepler_model.estimate.model.scikit_model import ScikitModelEstimator
-from kepler_model.estimate.model.xgboost_model import XgboostModelEstimator
+
+if import_util.find_spec("xgboost"):
+    from kepler_model.estimate.model.xgboost_model import XgboostModelEstimator
+
 from kepler_model.util.config import download_path
 from kepler_model.util.loader import get_download_output_path, load_metadata
 from kepler_model.util.prom_types import valid_container_query
@@ -17,11 +21,13 @@
 # model wrapper
 MODELCLASS = {
     "scikit": ScikitModelEstimator,
-    "xgboost": XgboostModelEstimator,
     "curvefit": CurveFitModelEstimator,
     # 'keras': KerasModelEstimator,
 }
 
+if import_util.find_spec("xgboost"):
+    MODELCLASS["xgboost"] = XgboostModelEstimator
+
 
 def default_predicted_col_func(energy_component):
     return f"default_{energy_component}_power"

src/kepler_model/util/train_types.py

@@ -10,6 +10,7 @@
 
 import enum
 import random
+import importlib.util as import_util
 
 SYSTEM_FEATURES = ["node_info", "cpu_scaling_frequency_hertz"]
 
@@ -45,8 +46,11 @@
     "KNeighborsRegressorTrainer",
     "LinearRegressionTrainer",
     "SVRRegressorTrainer",
-    "XgboostFitTrainer",
 ]
+
+if import_util.find_spec("xgboost"):
+    no_weight_trainers.append("XgboostFitTrainer")
+
 weight_support_trainers = ["SGDRegressorTrainer", "LogarithmicRegressionTrainer", "LogisticRegressionTrainer", "ExponentialRegressionTrainer"]
 default_trainer_names = no_weight_trainers + weight_support_trainers
 default_trainers = ",".join(default_trainer_names)