automlx version upgrade

Author: prasankh

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

@@ -32,6 +32,7 @@ def _build_model(self) -> pd.DataFrame:
             logger.info("Ray already initialized")
         date_column = self.spec.datetime_column.name
         anomaly_output = AnomalyOutput(date_column=date_column)
+        time_budget = self.spec.model_kwargs.pop("time_budget", -1)
 
         # Iterate over the full_data_dict items
         for target, df in self.datasets.full_data_dict.items():
@@ -47,6 +48,7 @@ def _build_model(self) -> pd.DataFrame:
                 contamination=self.spec.contamination
                 if self.y_valid_dict is not None
                 else None,
+                time_budget=time_budget,
             )
             y_pred = est.predict(df)
             scores = est.predict_proba(df)

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

@@ -45,7 +45,7 @@ def set_kwargs(self):
             model_kwargs_cleaned.get("score_metric", AUTOMLX_DEFAULT_SCORE_METRIC),
         )
         model_kwargs_cleaned.pop("task", None)
-        time_budget = model_kwargs_cleaned.pop("time_budget", None)
+        time_budget = model_kwargs_cleaned.pop("time_budget", -1)
         model_kwargs_cleaned[
             "preprocessing"
         ] = self.spec.preprocessing or model_kwargs_cleaned.get("preprocessing", True)
@@ -116,7 +116,8 @@ def _build_model(self) -> pd.DataFrame:
                     )
                     model.fit(
                         X=data_i.drop(target, axis=1),
-                        y=data_i[[target]]
+                        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_}")

tests/operators/forecast/test_errors.py

@@ -210,7 +210,7 @@ def populate_yaml(
     if model == "autots":
         yaml_i["spec"]["model_kwargs"] = {"model_list": "superfast"}
     if model == "automlx":
-        yaml_i["spec"]["model_kwargs"] = {"time_budget": 1}
+        yaml_i["spec"]["model_kwargs"] = {"time_budget": 50}
 
     return yaml_i, output_data_path