Retraining automlx models

Author: prasankh

ads/opctl/operator/lowcode/common/utils.py

@@ -218,17 +218,12 @@ def _write_file(local_filename, remote_filename, storage_options, **kwargs):
 
 
 def load_pkl(filepath):
-    return _safe_write(fn=_load_pkl, filepath=filepath)
-
-
-def _load_pkl(filepath):
     storage_options = {}
     if ObjectStorageDetails.is_oci_path(filepath):
         storage_options = default_signer()
 
     with fsspec.open(filepath, "rb", **storage_options) as f:
         return cloudpickle.load(f)
-    return None
 
 
 def write_pkl(obj, filename, output_dir, storage_options):

ads/opctl/operator/lowcode/forecast/model/automlx.py

@@ -142,17 +142,20 @@ def _build_model(self) -> pd.DataFrame:
                 )
 
                 if self.loaded_models is not None and s_id in self.loaded_models:
-                    model = self.loaded_models[s_id]
-                else:
-                    model = Pipeline(
-                        task="forecasting",
-                        **model_kwargs,
-                    )
-                    model.fit(
-                        X=data_i.drop(target, axis=1),
-                        y=data_i[[target]],
-                        time_budget=time_budget,
-                    )
+                    model = self.loaded_models[s_id]["model"]
+                    model_kwargs["model_list"] = [model.selected_model_]
+                    model_kwargs["search_space"]={}
+                    model_kwargs["search_space"][model.selected_model_] = model.selected_model_params_
+
+                model = Pipeline(
+                    task="forecasting",
+                    **model_kwargs,
+                )
+                model.fit(
+                    X=data_i.drop(target, axis=1),
+                    y=data_i[[target]],
+                    time_budget=time_budget,
+                )
                 logger.debug(f"Selected model: {model.selected_model_}")
                 logger.debug(f"Selected model params: {model.selected_model_params_}")
                 summary_frame = model.forecast(