Files
FLAML/test/spark/test_overtime.py
Li Jiang c2b25310fc Sync Fabric till 2cd1c3da (#1433)
* Sync Fabric till 2cd1c3da

* Remove synapseml from tag names

* Fix 'NoneType' object has no attribute 'DataFrame'

* Deprecated 3.8 support

* Fix 'NoneType' object has no attribute 'DataFrame'

* Still use python 3.8 for pydoc

* Don't run tests in parallel

* Remove autofe and lowcode
2025-05-23 10:19:31 +08:00

68 lines
2.0 KiB
Python

import os
import time
import numpy as np
import pytest
from sklearn.datasets import load_iris
from flaml import AutoML
try:
from test.spark.custom_mylearner import *
except ImportError:
from custom_mylearner import *
try:
import pyspark
from flaml.tune.spark.mylearner import lazy_metric
from flaml.tune.spark.utils import check_spark
os.environ["FLAML_MAX_CONCURRENT"] = "10"
spark = pyspark.sql.SparkSession.builder.appName("App4OvertimeTest").getOrCreate()
spark_available, _ = check_spark()
skip_spark = not spark_available
except ImportError:
skip_spark = True
pytestmark = [pytest.mark.skipif(skip_spark, reason="Spark is not installed. Skip all spark tests."), pytest.mark.spark]
def test_overtime():
time_budget = 15
df, y = load_iris(return_X_y=True, as_frame=True)
df["label"] = y
automl_experiment = AutoML()
automl_settings = {
"dataframe": df,
"label": "label",
"time_budget": time_budget,
"eval_method": "cv",
"metric": lazy_metric,
"task": "classification",
"log_file_name": "test/iris_custom.log",
"log_training_metric": True,
"log_type": "all",
"n_jobs": 1,
"model_history": True,
"sample_weight": np.ones(len(y)),
"pred_time_limit": 1e-5,
"estimator_list": ["lgbm"],
"n_concurrent_trials": 2,
"use_spark": True,
"force_cancel": True,
}
start_time = time.time()
automl_experiment.fit(**automl_settings)
elapsed_time = time.time() - start_time
print(f"time budget: {time_budget:.2f}s, actual elapsed time: {elapsed_time:.2f}s")
# assert abs(elapsed_time - time_budget) < 5 # cancel assertion because github VM sometimes is super slow, causing the test to fail
print(automl_experiment.predict(df))
print(automl_experiment.model)
print(automl_experiment.best_iteration)
print(automl_experiment.best_estimator)
if __name__ == "__main__":
test_overtime()