wehub-resource-sync
875 lines
31 KiB
Python
875 lines
31 KiB
Python
import inspect
|
|
import sys
|
|
import time
|
|
from typing import Any, NamedTuple
|
|
from unittest import mock
|
|
|
|
import pytest
|
|
|
|
import mlflow
|
|
from mlflow import MlflowClient
|
|
from mlflow.ml_package_versions import FLAVOR_TO_MODULE_NAME
|
|
from mlflow.utils import gorilla
|
|
from mlflow.utils.autologging_utils import (
|
|
AUTOLOGGING_INTEGRATIONS,
|
|
AutologgingEventLogger,
|
|
BatchMetricsLogger,
|
|
autologging_integration,
|
|
autologging_is_disabled,
|
|
batch_metrics_logger,
|
|
get_autologging_config,
|
|
get_instance_method_first_arg_value,
|
|
get_method_call_arg_value,
|
|
log_fn_args_as_params,
|
|
resolve_input_example_and_signature,
|
|
)
|
|
from mlflow.utils.autologging_utils.safety import AutologgingSession, _wrap_patch
|
|
from mlflow.utils.autologging_utils.versioning import (
|
|
_check_version_in_range,
|
|
_is_pre_or_dev_release,
|
|
_strip_dev_version_suffix,
|
|
_violates_pep_440,
|
|
is_flavor_supported_for_associated_package_versions,
|
|
)
|
|
|
|
from tests.autologging.fixtures import test_mode_off
|
|
|
|
# Example function signature we are testing on
|
|
# def fn(arg1, default1=1, default2=2):
|
|
# pass
|
|
|
|
|
|
two_default_test_args = [
|
|
(["arg1", "default1"], {"default2": 42}, ["arg1", "default1", "default2"], [1, 2], {}),
|
|
(["arg1", "default1", "default2"], {}, ["arg1", "default1", "default2"], [1, 2], {}),
|
|
(["arg1"], {"default1": 42, "default2": 42}, ["arg1", "default1", "default2"], [1, 2], {}),
|
|
(
|
|
[],
|
|
{"arg1": 42, "default1": 42, "default2": 42},
|
|
["arg1", "default1", "default2"],
|
|
[1, 2],
|
|
{},
|
|
),
|
|
(["user_arg"], {"default1": 42}, ["arg1", "default1", "default2"], [1, 2], {"default2": 2}),
|
|
(["user_arg"], {"default2": 42}, ["arg1", "default1", "default2"], [1, 2], {"default1": 1}),
|
|
([], {"arg1": 42, "default1": 42}, ["arg1", "default1", "default2"], [1, 2], {"default2": 2}),
|
|
(["arg1", "default1"], {}, ["arg1", "default1", "default2"], [1, 2], {"default2": 2}),
|
|
(["arg1"], {}, ["arg1", "default1", "default2"], [1, 2], {"default1": 1, "default2": 2}),
|
|
([], {"arg1": 42}, ["arg1", "default1", "default2"], [1, 2], {"default1": 1, "default2": 2}),
|
|
]
|
|
|
|
|
|
# Test function signature for the following tests
|
|
# def fn_default_default(default1=1, default2=2, default3=3):
|
|
# pass
|
|
|
|
|
|
three_default_test_args = [
|
|
(
|
|
[],
|
|
{},
|
|
["default1", "default2", "default3"],
|
|
[1, 2, 3],
|
|
{"default1": 1, "default2": 2, "default3": 3},
|
|
),
|
|
(
|
|
[],
|
|
{"default2": 42},
|
|
["default1", "default2", "default3"],
|
|
[1, 2, 3],
|
|
{"default1": 1, "default3": 3},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def start_run():
|
|
mlflow.start_run()
|
|
yield
|
|
mlflow.end_run()
|
|
|
|
|
|
def dummy_fn(arg1, arg2="value2", arg3="value3"):
|
|
pass
|
|
|
|
|
|
log_test_args = [
|
|
([], {"arg1": "value_x", "arg2": "value_y"}, ["value_x", "value_y", "value3"]),
|
|
(["value_x"], {"arg2": "value_y"}, ["value_x", "value_y", "value3"]),
|
|
(["value_x"], {"arg3": "value_z"}, ["value_x", "value2", "value_z"]),
|
|
(["value_x", "value_y"], {}, ["value_x", "value_y", "value3"]),
|
|
(["value_x", "value_y", "value_z"], {}, ["value_x", "value_y", "value_z"]),
|
|
(
|
|
[],
|
|
{"arg1": "value_x", "arg2": "value_y", "arg3": "value_z"},
|
|
["value_x", "value_y", "value_z"],
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.mark.parametrize(("args", "kwargs", "expected"), log_test_args)
|
|
def test_log_fn_args_as_params(args, kwargs, expected, start_run):
|
|
log_fn_args_as_params(dummy_fn, args, kwargs)
|
|
client = MlflowClient()
|
|
params = client.get_run(mlflow.active_run().info.run_id).data.params
|
|
for arg, value in zip(["arg1", "arg2", "arg3"], expected):
|
|
assert arg in params
|
|
assert params[arg] == value
|
|
|
|
|
|
def test_log_fn_args_as_params_ignores_unwanted_parameters(
|
|
start_run,
|
|
):
|
|
args = "arg1"
|
|
kwargs = {"arg2": "value"}
|
|
unlogged = ["arg1", "arg2", "arg3"]
|
|
log_fn_args_as_params(dummy_fn, args, kwargs, unlogged)
|
|
client = MlflowClient()
|
|
params = client.get_run(mlflow.active_run().info.run_id).data.params
|
|
assert len(params.keys()) == 0
|
|
|
|
|
|
def get_func_attrs(f):
|
|
assert callable(f)
|
|
|
|
return (f.__name__, f.__doc__, f.__module__, inspect.signature(f))
|
|
|
|
|
|
def test_wrap_patch_with_class():
|
|
class Math:
|
|
def add(self, a, b):
|
|
"""add"""
|
|
return a + b
|
|
|
|
def new_add(self, *args, **kwargs):
|
|
"""new add"""
|
|
orig = gorilla.get_original_attribute(self, "add")
|
|
return 2 * orig(*args, **kwargs)
|
|
|
|
_wrap_patch(Math, Math.add.__name__, new_add)
|
|
|
|
assert Math().add(1, 2) == 6
|
|
|
|
|
|
def sample_function_to_patch(a, b):
|
|
return a + b
|
|
|
|
|
|
def test_wrap_patch_with_module():
|
|
this_module = sys.modules[__name__]
|
|
|
|
def new_sample_function(a, b):
|
|
"""new mlflow.log_param"""
|
|
return a - b
|
|
|
|
assert sample_function_to_patch(10, 5) == 15
|
|
_wrap_patch(this_module, sample_function_to_patch.__name__, new_sample_function)
|
|
assert sample_function_to_patch(10, 5) == 5
|
|
|
|
|
|
@pytest.fixture
|
|
def logger():
|
|
return mock.Mock()
|
|
|
|
|
|
def get_input_example():
|
|
return "data"
|
|
|
|
|
|
def infer_model_signature(_):
|
|
return "signature"
|
|
|
|
|
|
def test_if_getting_input_example_fails(logger):
|
|
error_msg = "NoneType has no whatever"
|
|
|
|
def throws():
|
|
raise Exception(error_msg)
|
|
|
|
input_example, signature = resolve_input_example_and_signature(
|
|
throws, infer_model_signature, True, True, logger
|
|
)
|
|
|
|
assert input_example is None
|
|
assert signature is None
|
|
calls = [
|
|
mock.call("Failed to gather input example: " + error_msg),
|
|
mock.call(
|
|
"Failed to infer model signature: "
|
|
+ "could not sample data to infer model signature: "
|
|
+ error_msg
|
|
),
|
|
]
|
|
assert logger.warning.has_calls(calls)
|
|
|
|
|
|
def test_if_model_signature_inference_fails(logger):
|
|
error_msg = "stack overflow"
|
|
|
|
def throws(_):
|
|
raise Exception(error_msg)
|
|
|
|
input_example, signature = resolve_input_example_and_signature(
|
|
get_input_example, throws, True, True, logger
|
|
)
|
|
|
|
assert input_example == "data"
|
|
# When the signature inference fails but an input example is specified, `signature` is set
|
|
# to `False` to disable the automatic signature inference feature in `log_model` APIs.
|
|
assert signature is False
|
|
logger.warning.assert_called_with("Failed to infer model signature: " + error_msg)
|
|
|
|
|
|
def test_happy_path_works(logger):
|
|
input_example, signature = resolve_input_example_and_signature(
|
|
get_input_example, infer_model_signature, True, True, logger
|
|
)
|
|
|
|
assert input_example == "data"
|
|
assert signature == "signature"
|
|
logger.warning.assert_not_called()
|
|
|
|
|
|
def test_avoids_collecting_input_example_if_not_needed(logger):
|
|
# We create a get_input_example that modifies the value of x
|
|
# If get_input_example was not invoked, x should not have been modified.
|
|
|
|
x = {"data": 0}
|
|
|
|
def modifies():
|
|
x["data"] = 1
|
|
|
|
resolve_input_example_and_signature(modifies, infer_model_signature, False, False, logger)
|
|
|
|
assert x["data"] == 0
|
|
logger.warning.assert_not_called()
|
|
|
|
|
|
def test_avoids_inferring_signature_if_not_needed(logger):
|
|
# We create an infer_model_signature that modifies the value of x
|
|
# If infer_model_signature was not invoked, x should not have been modified.
|
|
|
|
x = {"data": 0}
|
|
|
|
def modifies(_):
|
|
x["data"] = 1
|
|
|
|
resolve_input_example_and_signature(get_input_example, modifies, True, False, logger)
|
|
|
|
assert x["data"] == 0
|
|
logger.warning.assert_not_called()
|
|
|
|
|
|
def test_batch_metrics_logger_logs_all_metrics(start_run):
|
|
run_id = mlflow.active_run().info.run_id
|
|
with batch_metrics_logger(run_id) as metrics_logger:
|
|
for i in range(100):
|
|
metrics_logger.record_metrics({hex(i): i}, i)
|
|
|
|
metrics_on_run = MlflowClient().get_run(run_id).data.metrics
|
|
|
|
for i in range(100):
|
|
assert hex(i) in metrics_on_run
|
|
assert metrics_on_run[hex(i)] == i
|
|
|
|
|
|
def test_batch_metrics_logger_flush_logs_to_mlflow(start_run):
|
|
run_id = mlflow.active_run().info.run_id
|
|
|
|
# Need to patch _should_flush() to return False, so that we can manually flush the logger
|
|
with mock.patch(
|
|
"mlflow.utils.autologging_utils.BatchMetricsLogger._should_flush", return_value=False
|
|
):
|
|
metrics_logger = BatchMetricsLogger(run_id)
|
|
metrics_logger.record_metrics({"my_metric": 10}, 5)
|
|
|
|
# Recorded metrics should not be logged to mlflow run before flushing BatchMetricsLogger
|
|
metrics_on_run = MlflowClient().get_run(run_id).data.metrics
|
|
assert "my_metric" not in metrics_on_run
|
|
|
|
metrics_logger.flush()
|
|
|
|
# Recorded metric should be logged to mlflow run after flushing BatchMetricsLogger
|
|
metrics_on_run = MlflowClient().get_run(run_id).data.metrics
|
|
assert "my_metric" in metrics_on_run
|
|
assert metrics_on_run["my_metric"] == 10
|
|
|
|
|
|
def test_batch_metrics_logger_runs_training_and_logging_in_correct_ratio(start_run):
|
|
with mock.patch.object(MlflowClient, "log_batch") as log_batch_mock:
|
|
run_id = mlflow.active_run().info.run_id
|
|
with batch_metrics_logger(run_id) as metrics_logger:
|
|
metrics_logger.record_metrics({"x": 1}, step=0) # data doesn't matter
|
|
|
|
# first metrics should be logged immediately to record a previous timestamp and
|
|
# batch log time
|
|
log_batch_mock.assert_called_once()
|
|
|
|
metrics_logger.total_log_batch_time = 1
|
|
metrics_logger.total_training_time = 1
|
|
|
|
log_batch_mock.reset_mock() # resets the 'calls' of this mock
|
|
|
|
# the above 'training' took 1 second. So with target training-to-logging time ratio of
|
|
# 10:1, 9 more 'training' should happen without sending the batch and then after the
|
|
# 10th training the batch should be sent.
|
|
for i in range(2, 11):
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
log_batch_mock.assert_not_called()
|
|
metrics_logger.total_training_time = i
|
|
|
|
# at this point, average log batch time is 1, and total training time is 9
|
|
# thus the next record_metrics call should send the batch.
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
log_batch_mock.assert_called_once()
|
|
|
|
# update log_batch time to reflect the 'mocked' training time
|
|
metrics_logger.total_log_batch_time = 2
|
|
|
|
log_batch_mock.reset_mock() # reset the recorded calls
|
|
|
|
for i in range(12, 21):
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
log_batch_mock.assert_not_called()
|
|
metrics_logger.total_training_time = i
|
|
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
log_batch_mock.assert_called_once()
|
|
|
|
|
|
def test_batch_metrics_logger_chunks_metrics_when_batch_logging(start_run):
|
|
with mock.patch.object(MlflowClient, "log_batch") as log_batch_mock:
|
|
run_id = mlflow.active_run().info.run_id
|
|
with batch_metrics_logger(run_id) as metrics_logger:
|
|
metrics_logger.record_metrics({hex(x): x for x in range(5000)}, step=0)
|
|
run_id = mlflow.active_run().info.run_id
|
|
|
|
for call_idx, call in enumerate(log_batch_mock.call_args_list):
|
|
_, kwargs = call
|
|
|
|
assert kwargs["run_id"] == run_id
|
|
assert len(kwargs["metrics"]) == 1000
|
|
for metric_idx, metric in enumerate(kwargs["metrics"]):
|
|
assert metric.key == hex(call_idx * 1000 + metric_idx)
|
|
assert metric.value == call_idx * 1000 + metric_idx
|
|
assert metric.step == 0
|
|
|
|
|
|
def test_batch_metrics_logger_records_time_correctly(start_run):
|
|
with mock.patch.object(MlflowClient, "log_batch", wraps=lambda *args, **kwargs: time.sleep(1)):
|
|
run_id = mlflow.active_run().info.run_id
|
|
with batch_metrics_logger(run_id) as metrics_logger:
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
|
|
assert metrics_logger.total_log_batch_time >= 1
|
|
|
|
time.sleep(2)
|
|
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
|
|
assert metrics_logger.total_training_time >= 2
|
|
|
|
|
|
def test_batch_metrics_logger_logs_timestamps_as_int_milliseconds(start_run):
|
|
with (
|
|
mock.patch.object(MlflowClient, "log_batch") as log_batch_mock,
|
|
mock.patch("time.time", return_value=123.45678901234567890),
|
|
):
|
|
run_id = mlflow.active_run().info.run_id
|
|
with batch_metrics_logger(run_id) as metrics_logger:
|
|
metrics_logger.record_metrics({"x": 1}, step=0)
|
|
|
|
_, kwargs = log_batch_mock.call_args
|
|
|
|
logged_metric = kwargs["metrics"][0]
|
|
|
|
assert logged_metric.timestamp == 123456
|
|
|
|
|
|
def test_autologging_integration_calls_underlying_function_correctly():
|
|
@autologging_integration("test_integration")
|
|
def autolog(foo=7, disable=False, silent=False):
|
|
return foo
|
|
|
|
assert autolog(foo=10) == 10
|
|
|
|
|
|
def test_autologging_integration_stores_and_updates_config():
|
|
@autologging_integration("test_integration")
|
|
def autolog(foo=7, bar=10, disable=False, silent=False):
|
|
return foo
|
|
|
|
autolog()
|
|
assert AUTOLOGGING_INTEGRATIONS["test_integration"] == {
|
|
"foo": 7,
|
|
"bar": 10,
|
|
"disable": False,
|
|
"silent": False,
|
|
}
|
|
autolog(bar=11)
|
|
assert AUTOLOGGING_INTEGRATIONS["test_integration"] == {
|
|
"foo": 7,
|
|
"bar": 11,
|
|
"disable": False,
|
|
"silent": False,
|
|
}
|
|
autolog(6, disable=True)
|
|
assert AUTOLOGGING_INTEGRATIONS["test_integration"] == {
|
|
"foo": 6,
|
|
"bar": 10,
|
|
"disable": True,
|
|
"silent": False,
|
|
}
|
|
autolog(1, 2, False, silent=True)
|
|
assert AUTOLOGGING_INTEGRATIONS["test_integration"] == {
|
|
"foo": 1,
|
|
"bar": 2,
|
|
"disable": False,
|
|
"silent": True,
|
|
}
|
|
|
|
|
|
def test_autologging_integration_forwards_positional_and_keyword_arguments_as_expected():
|
|
@autologging_integration("test_integration")
|
|
def autolog(foo=7, bar=10, disable=False, silent=False):
|
|
return foo, bar, disable
|
|
|
|
assert autolog(1, bar=2, disable=False) == (1, 2, False)
|
|
|
|
|
|
def test_autologging_integration_validates_structure_of_autolog_function():
|
|
def fn_missing_disable_conf():
|
|
pass
|
|
|
|
def fn_bad_disable_conf_1(disable=True):
|
|
pass
|
|
|
|
# Try to use a falsy value that isn't "false"
|
|
def fn_bad_disable_conf_2(disable=0):
|
|
pass
|
|
|
|
for fn in [fn_missing_disable_conf, fn_bad_disable_conf_1, fn_bad_disable_conf_2]:
|
|
with pytest.raises(Exception, match="must specify a 'disable' argument"):
|
|
autologging_integration("test")(fn)
|
|
|
|
# Failure to apply the @autologging_integration decorator should not create a
|
|
# placeholder for configuration state
|
|
assert "test" not in AUTOLOGGING_INTEGRATIONS
|
|
|
|
|
|
def test_autologging_integration_makes_expected_event_logging_calls():
|
|
@autologging_integration("test_success")
|
|
def autolog_success(foo, bar=7, disable=False, silent=False):
|
|
pass
|
|
|
|
@autologging_integration("test_failure")
|
|
def autolog_failure(biz, baz="val", disable=False, silent=False):
|
|
raise Exception("autolog failed")
|
|
|
|
class TestLogger(AutologgingEventLogger):
|
|
class LoggerCall(NamedTuple):
|
|
integration: Any
|
|
call_args: Any
|
|
call_kwargs: Any
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def reset(self):
|
|
self.calls = []
|
|
|
|
def log_autolog_called(self, integration, call_args, call_kwargs):
|
|
self.calls.append(TestLogger.LoggerCall(integration, call_args, call_kwargs))
|
|
|
|
logger = TestLogger()
|
|
AutologgingEventLogger.set_logger(logger)
|
|
|
|
autolog_success("a", bar=9, disable=True)
|
|
assert len(logger.calls) == 1
|
|
call = logger.calls[0]
|
|
assert call.integration == "test_success"
|
|
# NB: In MLflow > 1.13.1, the `call_args` argument to `log_autolog_called` is deprecated.
|
|
# Positional arguments passed to `autolog()` should be forwarded to `log_autolog_called`
|
|
# in keyword format
|
|
assert call.call_args == ()
|
|
assert call.call_kwargs == {"foo": "a", "bar": 9, "disable": True, "silent": False}
|
|
|
|
logger.reset()
|
|
|
|
with pytest.raises(Exception, match="autolog failed"):
|
|
autolog_failure(82, disable=False, silent=True)
|
|
assert len(logger.calls) == 1
|
|
call = logger.calls[0]
|
|
assert call.integration == "test_failure"
|
|
# NB: In MLflow > 1.13.1, the `call_args` argument to `log_autolog_called` is deprecated.
|
|
# Positional arguments passed to `autolog()` should be forwarded to `log_autolog_called`
|
|
# in keyword format
|
|
assert call.call_args == ()
|
|
assert call.call_kwargs == {"biz": 82, "baz": "val", "disable": False, "silent": True}
|
|
|
|
|
|
@pytest.mark.usefixtures(test_mode_off.__name__)
|
|
def test_autologging_integration_succeeds_when_event_logging_throws_in_standard_mode():
|
|
@autologging_integration("test")
|
|
def autolog(disable=False, silent=False):
|
|
return "result"
|
|
|
|
class ThrowingLogger(AutologgingEventLogger):
|
|
def __init__(self):
|
|
self.logged_event = False
|
|
|
|
def log_autolog_called(self, integration, call_args, call_kwargs):
|
|
self.logged_event = True
|
|
raise Exception("autolog failed")
|
|
|
|
logger = ThrowingLogger()
|
|
AutologgingEventLogger.set_logger(logger)
|
|
assert autolog() == "result"
|
|
assert logger.logged_event
|
|
|
|
|
|
def test_get_autologging_config_returns_configured_values_or_defaults_as_expected():
|
|
assert get_autologging_config("nonexistent_integration", "foo") is None
|
|
|
|
@autologging_integration("test_integration_for_config")
|
|
def autolog(foo="bar", t=7, disable=False, silent=False):
|
|
pass
|
|
|
|
# Before `autolog()` has been invoked, config values should not be available
|
|
assert get_autologging_config("test_integration_for_config", "foo") is None
|
|
assert get_autologging_config("test_integration_for_config", "disable") is None
|
|
assert get_autologging_config("test_integration_for_config", "silent") is None
|
|
assert get_autologging_config("test_integration_for_config", "t", 10) == 10
|
|
|
|
autolog()
|
|
|
|
assert get_autologging_config("test_integration_for_config", "foo") == "bar"
|
|
assert get_autologging_config("test_integration_for_config", "disable") is False
|
|
assert get_autologging_config("test_integration_for_config", "silent") is False
|
|
assert get_autologging_config("test_integration_for_config", "t", 10) == 7
|
|
assert get_autologging_config("test_integration_for_config", "nonexistent") is None
|
|
|
|
autolog(foo="baz", silent=True)
|
|
|
|
assert get_autologging_config("test_integration_for_config", "foo") == "baz"
|
|
assert get_autologging_config("test_integration_for_config", "silent") is True
|
|
|
|
|
|
def test_autologging_is_disabled_returns_expected_values():
|
|
assert autologging_is_disabled("nonexistent_integration") is True
|
|
|
|
@autologging_integration("test_integration_for_disable_check")
|
|
def autolog(disable=False, silent=False):
|
|
pass
|
|
|
|
# Before `autolog()` has been invoked, `autologging_is_disabled` should return False
|
|
assert autologging_is_disabled("test_integration_for_disable_check") is True
|
|
|
|
autolog(disable=True)
|
|
|
|
assert autologging_is_disabled("test_integration_for_disable_check") is True
|
|
|
|
autolog(disable=False)
|
|
|
|
assert autologging_is_disabled("test_integration_for_disable_check") is False
|
|
|
|
|
|
def test_autologging_disable_restores_behavior():
|
|
from sklearn.datasets import load_diabetes
|
|
from sklearn.linear_model import LinearRegression
|
|
|
|
X, y = load_diabetes(return_X_y=True, as_frame=True)
|
|
X = X.iloc[:50, :4]
|
|
y = y.iloc[:50]
|
|
|
|
# train a model
|
|
model = LinearRegression()
|
|
|
|
def is_autolog_on():
|
|
run = mlflow.start_run()
|
|
model.fit(X, y)
|
|
mlflow.end_run()
|
|
run = MlflowClient().get_run(run.info.run_id)
|
|
return run.data.metrics and run.data.params
|
|
|
|
# Turn on autologging
|
|
mlflow.sklearn.autolog()
|
|
assert is_autolog_on()
|
|
|
|
# Turn off autologging within a context manager
|
|
with mlflow.utils.autologging_utils.disable_autologging():
|
|
assert not is_autolog_on()
|
|
|
|
# Autologging should be turned back on
|
|
assert is_autolog_on()
|
|
|
|
# The context manager should exit correctly even if an exception is raised
|
|
with pytest.raises(Exception, match="test"): # noqa PT012
|
|
with mlflow.utils.autologging_utils.disable_autologging():
|
|
assert not is_autolog_on()
|
|
raise Exception("test")
|
|
|
|
# Autologging should be turned back on after the exception
|
|
assert is_autolog_on()
|
|
|
|
|
|
def test_autologging_event_logger_default_implementation_does_not_throw_for_valid_inputs():
|
|
AutologgingEventLogger.set_logger(AutologgingEventLogger())
|
|
|
|
class PatchObj:
|
|
def test_fn(self):
|
|
pass
|
|
|
|
# Test successful autologging workflow
|
|
AutologgingEventLogger.get_logger().log_autolog_called(
|
|
"test_integration", ("a"), {"b": 1, "c": "d"}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_patch_function_start(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_original_function_start(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_original_function_success(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_patch_function_success(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
|
|
# Test patch function failure autologging workflow
|
|
AutologgingEventLogger.get_logger().log_patch_function_start(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_patch_function_error(
|
|
AutologgingSession("test_integration", "123"),
|
|
PatchObj(),
|
|
"test_fn",
|
|
(1000),
|
|
{"a": 2},
|
|
Exception("patch error"),
|
|
)
|
|
|
|
# Test original function failure autologging workflow
|
|
AutologgingEventLogger.get_logger().log_patch_function_start(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_original_function_start(
|
|
AutologgingSession("test_integration", "123"), PatchObj(), "test_fn", (1000), {"a": 2}
|
|
)
|
|
AutologgingEventLogger.get_logger().log_patch_function_error(
|
|
AutologgingSession("test_integration", "123"),
|
|
PatchObj(),
|
|
"test_fn",
|
|
(1000),
|
|
{"a": 2},
|
|
Exception("patch error"),
|
|
)
|
|
|
|
|
|
def test_autologging_event_logger_default_impl_warns_for_log_autolog_called_with_deprecated_args():
|
|
AutologgingEventLogger.set_logger(AutologgingEventLogger())
|
|
|
|
with pytest.warns(DeprecationWarning, match="Received 1 positional arguments"):
|
|
AutologgingEventLogger.get_logger().log_autolog_called(
|
|
"test_integration",
|
|
# call_args is deprecated in MLflow > 1.13.1; specifying a non-empty
|
|
# value for this parameter should emit a warning
|
|
call_args=("a"),
|
|
call_kwargs={"b": "c"},
|
|
)
|
|
|
|
|
|
def test_check_version_in_range():
|
|
assert _check_version_in_range("1.0.2", "1.0.1", "1.0.3")
|
|
assert _check_version_in_range("1.0.1", "1.0.1", "1.0.3")
|
|
assert _check_version_in_range("1.0.3", "1.0.1", "1.0.3")
|
|
assert not _check_version_in_range("1.0.0", "1.0.1", "1.0.3")
|
|
assert not _check_version_in_range("1.0.4", "1.0.1", "1.0.3")
|
|
assert not _check_version_in_range("0.99.99", "1.0.1", "1.0.3")
|
|
assert not _check_version_in_range("1.1.0", "1.0.1", "1.0.3")
|
|
assert _check_version_in_range("1.0.3", "1.0.1", "1.0.3.post1")
|
|
|
|
|
|
def test_is_pre_or_dev_release():
|
|
assert _is_pre_or_dev_release("0.24.0rc1")
|
|
assert _is_pre_or_dev_release("0.24.0dev1")
|
|
assert not _is_pre_or_dev_release("0.24.0")
|
|
|
|
|
|
def test_strip_dev_version_suffix():
|
|
assert _strip_dev_version_suffix("1.0.dev0") == "1.0"
|
|
assert _strip_dev_version_suffix("1.0dev0") == "1.0"
|
|
assert _strip_dev_version_suffix("1.0.dev") == "1.0"
|
|
assert _strip_dev_version_suffix("1.0") == "1.0"
|
|
|
|
|
|
def test_violates_pep_440():
|
|
assert _violates_pep_440("0.24.0-SNAPSHOT")
|
|
assert not _violates_pep_440("0.24.0rc1")
|
|
assert not _violates_pep_440("0.24.0dev1")
|
|
assert not _violates_pep_440("0.24.0")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("flavor", "module_version", "expected_result"),
|
|
[
|
|
("sklearn", "1.5.1", True),
|
|
("sklearn", "0.20.2", False),
|
|
("sklearn", "0.23.0rc1", False),
|
|
("sklearn", "0.23.0dev0", False),
|
|
("sklearn", "0.23.0-SNAPSHOT", False),
|
|
("pytorch", "2.4.0", True),
|
|
("pytorch", "1.5.99", False),
|
|
("pyspark.ml", "3.5.1", True),
|
|
("pyspark.ml", "3.0.0", False),
|
|
("llama_index", "0.13.1", True),
|
|
("llama_index", "0.1.2", False),
|
|
],
|
|
)
|
|
def test_is_autologging_integration_supported(flavor, module_version, expected_result):
|
|
module_name = FLAVOR_TO_MODULE_NAME[flavor]
|
|
with mock.patch(module_name + ".__version__", module_version):
|
|
assert expected_result == is_flavor_supported_for_associated_package_versions(flavor)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("flavor", "module_version", "expected_result"),
|
|
[
|
|
("pyspark.ml", "99.0.0.dev0", False),
|
|
("pyspark.ml", "3.5.0.dev0", True),
|
|
("pyspark.ml", "3.3.0.dev0", True),
|
|
("pyspark.ml", "3.2.1.dev0", False),
|
|
("pyspark.ml", "3.1.2.dev0", False),
|
|
("pyspark.ml", "3.0.1.dev0", False),
|
|
("pyspark.ml", "3.0.0.dev0", False),
|
|
("pyspark.ml", "2.4.8.dev0", False),
|
|
],
|
|
)
|
|
def test_dev_version_pyspark_is_supported_in_databricks(flavor, module_version, expected_result):
|
|
module_name = FLAVOR_TO_MODULE_NAME[flavor]
|
|
with mock.patch(module_name + ".__version__", module_version):
|
|
# In Databricks
|
|
with mock.patch(
|
|
"mlflow.utils.autologging_utils.versioning.is_in_databricks_runtime",
|
|
return_value=True,
|
|
) as mock_runtime:
|
|
assert is_flavor_supported_for_associated_package_versions(flavor) == expected_result
|
|
mock_runtime.assert_called()
|
|
|
|
# Not in Databricks
|
|
assert is_flavor_supported_for_associated_package_versions(flavor) is False
|
|
|
|
|
|
def test_disable_for_unsupported_versions_warning_sklearn_integration():
|
|
log_warn_fn_name = "mlflow.utils.autologging_utils._logger.warning"
|
|
log_info_fn_name = "mlflow.tracking.fluent._logger.info"
|
|
|
|
def is_sklearn_warning_fired(log_warn_fn_args):
|
|
return "MLflow sklearn autologging is known to be compatible" in log_warn_fn_args[0][0]
|
|
|
|
def is_sklearn_autolog_enabled_info_fired(log_info_fn_args):
|
|
return (
|
|
"Autologging successfully enabled for " in log_info_fn_args[0][0]
|
|
and log_info_fn_args[0][1] == "sklearn"
|
|
)
|
|
|
|
with mock.patch("sklearn.__version__", "1.5.1"):
|
|
AUTOLOGGING_INTEGRATIONS.clear()
|
|
with (
|
|
mock.patch(log_warn_fn_name) as log_warn_fn,
|
|
mock.patch(log_info_fn_name) as log_info_fn,
|
|
):
|
|
mlflow.autolog(disable_for_unsupported_versions=True)
|
|
assert all(not is_sklearn_warning_fired(args) for args in log_warn_fn.call_args_list)
|
|
assert any(
|
|
is_sklearn_autolog_enabled_info_fired(args) for args in log_info_fn.call_args_list
|
|
)
|
|
with (
|
|
mock.patch(log_warn_fn_name) as log_warn_fn,
|
|
mock.patch(log_info_fn_name) as log_info_fn,
|
|
):
|
|
mlflow.autolog(disable_for_unsupported_versions=False)
|
|
assert all(not is_sklearn_warning_fired(args) for args in log_warn_fn.call_args_list)
|
|
assert any(
|
|
is_sklearn_autolog_enabled_info_fired(args) for args in log_info_fn.call_args_list
|
|
)
|
|
|
|
with mock.patch(log_warn_fn_name) as log_warn_fn:
|
|
mlflow.sklearn.autolog(disable_for_unsupported_versions=True)
|
|
log_warn_fn.assert_not_called()
|
|
with mock.patch(log_warn_fn_name) as log_warn_fn:
|
|
mlflow.sklearn.autolog(disable_for_unsupported_versions=False)
|
|
log_warn_fn.assert_not_called()
|
|
|
|
with mock.patch("sklearn.__version__", "0.20.2"):
|
|
AUTOLOGGING_INTEGRATIONS.clear()
|
|
with (
|
|
mock.patch(log_warn_fn_name) as log_warn_fn,
|
|
mock.patch(log_info_fn_name) as log_info_fn,
|
|
):
|
|
mlflow.autolog(disable_for_unsupported_versions=True)
|
|
assert all(not is_sklearn_warning_fired(args) for args in log_warn_fn.call_args_list)
|
|
assert all(
|
|
not is_sklearn_autolog_enabled_info_fired(args)
|
|
for args in log_info_fn.call_args_list
|
|
)
|
|
with (
|
|
mock.patch(log_warn_fn_name) as log_warn_fn,
|
|
mock.patch(log_info_fn_name) as log_info_fn,
|
|
):
|
|
mlflow.autolog(disable_for_unsupported_versions=False)
|
|
assert any(is_sklearn_warning_fired(args) for args in log_warn_fn.call_args_list)
|
|
assert any(
|
|
is_sklearn_autolog_enabled_info_fired(args) for args in log_info_fn.call_args_list
|
|
)
|
|
with mock.patch(log_warn_fn_name) as log_warn_fn:
|
|
mlflow.sklearn.autolog(disable_for_unsupported_versions=True)
|
|
log_warn_fn.assert_not_called()
|
|
with mock.patch(log_warn_fn_name) as log_warn_fn:
|
|
mlflow.sklearn.autolog(disable_for_unsupported_versions=False)
|
|
assert log_warn_fn.call_count == 1
|
|
assert is_sklearn_warning_fired(log_warn_fn.call_args)
|
|
|
|
|
|
def test_unsupported_versions_warning_should_not_shown_for_excluded_packages():
|
|
with mock.patch("langchain.__version__", "100.200.300"):
|
|
AUTOLOGGING_INTEGRATIONS.clear()
|
|
with mock.patch("mlflow.utils.autologging_utils._logger.warning") as log_warn_fn:
|
|
mlflow.langchain.autolog()
|
|
assert len(log_warn_fn.call_args_list) == 0 or (
|
|
"MLflow langchain autologging is known to be compatible"
|
|
not in log_warn_fn.call_args_list[0][0]
|
|
)
|
|
|
|
|
|
def test_get_instance_method_first_arg_value():
|
|
class Test:
|
|
def f1(self, ab1, cd2):
|
|
pass
|
|
|
|
def f2(self, *args):
|
|
pass
|
|
|
|
def f3(self, *kwargs):
|
|
pass
|
|
|
|
def f4(self, *args, **kwargs):
|
|
pass
|
|
|
|
assert get_instance_method_first_arg_value(Test.f1, [3, 4], {}) == 3
|
|
assert get_instance_method_first_arg_value(Test.f1, [3], {"cd2": 4}) == 3
|
|
assert get_instance_method_first_arg_value(Test.f1, [], {"ab1": 3, "cd2": 4}) == 3
|
|
assert get_instance_method_first_arg_value(Test.f2, [3, 4], {}) == 3
|
|
with pytest.raises(AssertionError, match=r".*"):
|
|
get_instance_method_first_arg_value(Test.f3, [], {"ab1": 3, "cd2": 4})
|
|
with pytest.raises(AssertionError, match=r".*"):
|
|
get_instance_method_first_arg_value(Test.f4, [], {"ab1": 3, "cd2": 4})
|
|
|
|
|
|
def test_get_method_call_arg_value():
|
|
# suppose we call on a method defined like: `def f1(a, b=3, *, c=4, e=5)`
|
|
assert get_method_call_arg_value(1, "b", 3, [1, 2], {}) == 2
|
|
assert get_method_call_arg_value(1, "b", 3, [1], {}) == 3
|
|
assert get_method_call_arg_value(1, "b", 3, [1], {"b": 2}) == 2
|