def test_log_feature_importance_with_limit(limit: int) -> None:
with mlflow.start_run() as run:
default_path = _get_default_args(lg.log_feature_importance)["path"]
lg.log_feature_importance(["a", "b", "c"], [1, 2, 3],
"gain",
limit=limit)
assert_file_exists_in_artifacts(run, default_path)
def test_log_plt_figure() -> None:
fig, ax = plt.subplots()
ax.plot([0, 1], [0, 1])
with mlflow.start_run() as run:
path = "test.png"
lg.log_figure(fig, path)
assert_file_exists_in_artifacts(run, path)
def test_log_dict(path: str) -> None:
data = {"a": 0}
with mlflow.start_run() as run:
lg.log_dict(data, path)
assert_file_exists_in_artifacts(run, path)
artifacts_dir = run.info.artifact_uri.replace("file://", "")
loaded_data = _read_data(os.path.join(artifacts_dir, path))
assert loaded_data == data
def test_log_confusion_matrix() -> None:
with mlflow.start_run() as run:
default_path = _get_default_args(lg.log_confusion_matrix)["path"]
lg.log_confusion_matrix([[1, 2], [3, 4]])
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
path = "cm.png"
lg.log_confusion_matrix([[1, 2], [3, 4]], path)
assert_file_exists_in_artifacts(run, path)
def test_log_text(text: str) -> None:
path = "test.txt"
with mlflow.start_run() as run:
lg.log_text(text, path)
assert_file_exists_in_artifacts(run, path)
artifacts_dir = run.info.artifact_uri.replace("file://", "")
with open(os.path.join(artifacts_dir, path), "r") as f:
assert text == f.read()
def test_log_plotly_figure() -> None:
fig = go.Figure(data=[go.Bar(x=[1, 2, 3], y=[1, 3, 2])])
with mlflow.start_run() as run:
path = "test.html"
lg.log_figure(fig, path)
assert_file_exists_in_artifacts(run, path)
path = "test.png"
msg = '"{}" is not an HTML file.'.format(path)
with pytest.raises(ValueError, match=msg):
lg.log_figure(fig, path)
def test_log_numpy() -> None:
array = np.array([0])
path = "test.npy"
with mlflow.start_run() as run:
lg.log_numpy(array, path)
assert_file_exists_in_artifacts(run, path)
artifacts_dir = run.info.artifact_uri.replace("file://", "")
loaded_array = np.load(os.path.join(artifacts_dir, path))
np.testing.assert_array_equal(loaded_array, array)
def test_log_dict_with_fmt(fmt: str) -> None:
data = {"a": 0}
path = "test.{}".format(fmt.lstrip("."))
with mlflow.start_run() as run:
lg.log_dict(data, path, fmt)
assert_file_exists_in_artifacts(run, path)
artifacts_dir = run.info.artifact_uri.replace("file://", "")
loaded_data = _read_data(os.path.join(artifacts_dir, path))
assert loaded_data == data
def test_log_df(fmt: str) -> None:
df = pd.DataFrame({"a": [0]})
path = "test.{}".format(fmt)
with mlflow.start_run() as run:
lg.log_df(df, path, fmt)
assert_file_exists_in_artifacts(run, path)
artifacts_dir = run.info.artifact_uri.replace("file://", "")
readers = {"csv": pd.read_csv, "feather": pd.read_feather}
loaded_df = readers[fmt](os.path.join(artifacts_dir, path))
pd.testing.assert_frame_equal(loaded_df, df)
def test_log_pr_curve() -> None:
default_path = _get_default_args(lg.log_pr_curve)["path"]
with mlflow.start_run() as run:
lg.log_pr_curve([1, 0], [1, 0])
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
lg.log_pr_curve([1, 0], [1, 0], 0.5)
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
path = "pr.png"
lg.log_pr_curve([1, 0], [1, 0], path=path)
assert_file_exists_in_artifacts(run, path)
def test_log_roc_curve():
default_path = _get_default_args(lg.log_roc_curve)["path"]
with mlflow.start_run() as run:
lg.log_roc_curve([0, 1], [0, 1])
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
lg.log_roc_curve([0, 1], [0, 1], 0.5)
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
path = "roc.png"
lg.log_roc_curve([0, 1], [0, 1], path=path)
assert_file_exists_in_artifacts(run, path)
def test_log_feature_importance():
default_path = _get_default_args(lg.log_feature_importance)["path"]
with mlflow.start_run() as run:
lg.log_feature_importance(["a", "b", "c"], [1, 2, 3], "gain")
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
lg.log_feature_importance(["a", "b", "c"], [1, 2, 3], "gain", normalize=True)
assert_file_exists_in_artifacts(run, default_path)
with mlflow.start_run() as run:
path = "feat_imp.png"
lg.log_feature_importance(["a", "b", "c"], [1, 2, 3], "gain", path=path)
assert_file_exists_in_artifacts(run, path)
def test_artifact_context(path: str) -> None:
with mlflow.start_run() as run:
with lg._artifact_context(path) as tmp_path:
open(tmp_path, "w").close()
assert_file_exists_in_artifacts(run, path)
def test_log_pickle() -> None:
with mlflow.start_run() as run:
path = "test.pkl"
lg.log_pickle({"a": 0}, path)
assert_file_exists_in_artifacts(run, path)
