mirror of
https://github.com/microsoft/autogen.git
synced 2025-09-18 04:34:10 +00:00
65 lines
2.7 KiB
Python
65 lines
2.7 KiB
Python
|
import pytest
|
||
|
|
from pandas import DataFrame
|
||
|
|
from sklearn.datasets import load_iris
|
||
|
|
import mlflow
|
||
|
|
import mlflow.entities
|
||
|
|
from flaml import AutoML
|
||
|
|
|
||
|
|
|
||
|
|
class TestMLFlowLoggingParam:
|
||
|
|
def test_should_start_new_run_by_default(self, automl_settings):
|
||
|
|
with mlflow.start_run():
|
||
|
|
parent = mlflow.last_active_run()
|
||
|
|
automl = AutoML()
|
||
|
|
X_train, y_train = load_iris(return_X_y=True)
|
||
|
|
automl.fit(X_train=X_train, y_train=y_train, **automl_settings)
|
||
|
|
|
||
|
|
children = self._get_child_runs(parent)
|
||
|
|
assert len(children) >= 1, "Expected at least 1 child run, got {}".format(len(children))
|
||
|
|
|
||
|
|
def test_should_not_start_new_run_when_mlflow_logging_set_to_false_in_init(self, automl_settings):
|
||
|
|
with mlflow.start_run():
|
||
|
|
parent = mlflow.last_active_run()
|
||
|
|
automl = AutoML(mlflow_logging=False)
|
||
|
|
X_train, y_train = load_iris(return_X_y=True)
|
||
|
|
automl.fit(X_train=X_train, y_train=y_train, **automl_settings)
|
||
|
|
|
||
|
|
children = self._get_child_runs(parent)
|
||
|
|
assert len(children) == 0, "Expected 0 child runs, got {}".format(len(children))
|
||
|
|
|
||
|
|
def test_should_not_start_new_run_when_mlflow_logging_set_to_false_in_fit(self, automl_settings):
|
||
|
|
with mlflow.start_run():
|
||
|
|
parent = mlflow.last_active_run()
|
||
|
|
automl = AutoML()
|
||
|
|
X_train, y_train = load_iris(return_X_y=True)
|
||
|
|
automl.fit(X_train=X_train, y_train=y_train, mlflow_logging=False, **automl_settings)
|
||
|
|
|
||
|
|
children = self._get_child_runs(parent)
|
||
|
|
assert len(children) == 0, "Expected 0 child runs, got {}".format(len(children))
|
||
|
|
|
||
|
|
def test_should_start_new_run_when_mlflow_logging_set_to_true_in_fit(self, automl_settings):
|
||
|
|
with mlflow.start_run():
|
||
|
|
parent = mlflow.last_active_run()
|
||
|
|
automl = AutoML(mlflow_logging=False)
|
||
|
|
X_train, y_train = load_iris(return_X_y=True)
|
||
|
|
automl.fit(X_train=X_train, y_train=y_train, mlflow_logging=True, **automl_settings)
|
||
|
|
|
||
|
|
children = self._get_child_runs(parent)
|
||
|
|
assert len(children) >= 1, "Expected at least 1 child run, got {}".format(len(children))
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def _get_child_runs(parent_run: mlflow.entities.Run) -> DataFrame:
|
||
|
|
experiment_id = parent_run.info.experiment_id
|
||
|
|
return mlflow.search_runs(
|
||
|
|
[experiment_id], filter_string="tags.mlflow.parentRunId = '{}'".format(parent_run.info.run_id)
|
||
|
|
)
|
||
|
|
|
||
|
|
@pytest.fixture(scope="class")
|
||
|
|
def automl_settings(self):
|
||
|
|
return {
|
||
|
|
"time_budget": 2, # in seconds
|
||
|
|
"metric": "accuracy",
|
||
|
|
"task": "classification",
|
||
|
|
"log_file_name": "iris.log",
|
||
|
|
}
|