Skip to content

Commit 09cd6eb

Browse files
committed
limit automlx build time
1 parent 3986ceb commit 09cd6eb

File tree

1 file changed

+4
-0
lines changed

1 file changed

+4
-0
lines changed

tests/operators/forecast/test_datasets.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -179,6 +179,8 @@ def test_pandas_to_historical(model):
179179
yaml_i["spec"]["datetime_column"]["name"] = DATETIME_COL
180180
yaml_i["spec"]["horizon"] = PERIODS
181181
yaml_i["spec"]["output_directory"]["url"] = output_data_path
182+
if model == "automlx":
183+
yaml_i["spec"]["model_kwargs"] = {"time_budget": 2}
182184
operator_config = ForecastOperatorConfig.from_dict(yaml_i)
183185
forecast_operate(operator_config)
184186
check_output_for_errors(output_data_path)
@@ -201,6 +203,8 @@ def test_pandas_to_historical_test(model):
201203
yaml_i["spec"]["datetime_column"]["name"] = DATETIME_COL
202204
yaml_i["spec"]["horizon"] = PERIODS
203205
yaml_i["spec"]["output_directory"]["url"] = output_data_path
206+
if model == "automlx":
207+
yaml_i["spec"]["model_kwargs"] = {"time_budget": 2}
204208
operator_config = ForecastOperatorConfig.from_dict(yaml_i)
205209
forecast_operate(operator_config)
206210
check_output_for_errors(output_data_path)

0 commit comments

Comments
 (0)