def test_return_a_couple_models(self):
cache = ModelCache()
def get_search_return_values(experiment_ids, *args, **kwargs):
df = pd.DataFrame(columns=list(cache.columns_of_interest.keys()),
data=[{
'run_id': "123",
'tags.mlflow.runName': 'my_run',
'tags.mlflow.BuildNumber': '4',
'end_time': datetime(2020, 8, 29, 8, 0, 0),
'params.MLPipelineParamsName': 'default',
'params.FeatureSetName': 'default',
'params.AlgorithmName': 'default',
'params.AlgorithmParamsName': 'default',
'tags.DidPassAcceptanceTest': 'no'
}, {
'run_id': "456",
'tags.mlflow.runName': 'my_second_run',
'tags.mlflow.BuildNumber': '5',
'end_time': datetime(2020, 8, 29, 9, 0, 0),
'params.MLPipelineParamsName': 'default',
'params.FeatureSetName': 'default',
'params.AlgorithmName': 'default',
'params.AlgorithmParamsName': 'default',
'tags.DidPassAcceptanceTest': 'yes'
}])
df["end_time"] = pd.to_datetime(df["end_time"])
return df
def get_experiment_id(scenario, *args, **kwargs):
if scenario == "houses":
return None
else:
return SimpleNamespace(**{'experiment_id': "id_" + scenario})
mlflow.get_experiment_by_name = MagicMock(
side_effect=get_experiment_id)
mlflow.search_runs = MagicMock(side_effect=get_search_return_values)
available_models = cache.list_available_models_from_ml_flow()
assert set(available_models.keys()) == {"groceries", "iris"}
assert len(available_models["groceries"]) == 2
assert [x["run_id"]
for x in available_models["groceries"]] == ['123', '456']
assert [
x["is_latest"] for x in available_models["groceries"]
if x["run_id"] == "456"
][0]
assert not [
x["is_latest"]
for x in available_models["groceries"] if x["run_id"] == "123"
][0]
def test_list_no_models_available(self):
cache = ModelCache()
def get_search_return_values(experiment_ids, *args, **kwargs):
df = pd.DataFrame(columns=list(cache.columns_of_interest.keys()))
df["end_time"] = pd.to_datetime(df["end_time"])
return df
def get_experiment_id(scenario, *args, **kwargs):
return SimpleNamespace(**{'experiment_id': "id_" + scenario})
mlflow.get_experiment_by_name = MagicMock(side_effect=get_experiment_id)
mlflow.search_runs = MagicMock(side_effect=get_search_return_values)
available_models = cache.list_available_models_from_ml_flow()
assert available_models == {"groceries": [], "houses": []}
