import copy import socket from unittest.mock import patch import pytest from mlflow.entities import Metric, Param, RunTag from mlflow.environment_variables import MLFLOW_ARTIFACT_LOCATION_MAX_LENGTH from mlflow.exceptions import MlflowException from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, ErrorCode from mlflow.utils.os import is_windows from mlflow.utils.validation import ( MAX_TAG_VAL_LENGTH, _is_numeric, _parse_trace_archival_duration_config, _validate_batch_log_data, _validate_batch_log_limits, _validate_db_type_string, _validate_experiment_artifact_location, _validate_experiment_artifact_location_length, _validate_experiment_name, _validate_list_param, _validate_metric_name, _validate_model_alias_name, _validate_model_alias_name_reserved, _validate_model_name, _validate_model_renaming, _validate_param_name, _validate_run_id, _validate_tag_name, _validate_webhook_url, path_not_unique, ) GOOD_METRIC_OR_PARAM_NAMES = [ "a", "Ab-5_", "a/b/c", "a.b.c", ".a", "b.", "a..a/._./o_O/.e.", "a b/c d", ] BAD_METRIC_OR_PARAM_NAMES = [ "", ".", "/", "..", "//", "a//b", "a/./b", "/a", "a/", "\\", "./", "/./", ] GOOD_ALIAS_NAMES = [ "a", "Ab-5_", "test-alias", "1a2b5cDeFgH", "a" * 255, "lates", # spellchecker: disable-line "v123_temp", "123", "123v", "temp_V123", ] BAD_ALIAS_NAMES = [ "", ".", "/", "..", "//", "a b", "a/./b", "/a", "a/", ":", "\\", "./", "/./", "a" * 256, None, "$dgs", ] @pytest.mark.parametrize( ("path", "expected"), [ ("a", False), ("a/b/c", False), ("a.b/c", False), (".a", False), # Not unique paths ("./a", True), ("a/b/../c", True), (".", True), ("../a/b", True), ("/a/b/c", True), ], ) def test_path_not_unique(path, expected): assert path_not_unique(path) is expected def test_parse_trace_archival_duration_config_archive_now(): assert ( _parse_trace_archival_duration_config( '{"older_than": " 7d "}', duration_key="older_than", allow_missing_duration=True, ) == "7d" ) def test_parse_trace_archival_duration_config_experiment_retention(): assert ( _parse_trace_archival_duration_config( '{"type": "duration", "value": "12h"}', duration_key="value", expected_type="duration", ) == "12h" ) def test_parse_trace_archival_duration_config_rejects_non_object(): with pytest.raises(MlflowException, match="JSON object"): _parse_trace_archival_duration_config('["1d"]', duration_key="value") @pytest.mark.parametrize( ("value", "expected"), [ (0, True), (0.0, True), # Non-numeric cases (True, False), (False, False), ("0", False), (None, False), ], ) def test_is_numeric(value, expected): assert _is_numeric(value) is expected @pytest.mark.parametrize("metric_name", GOOD_METRIC_OR_PARAM_NAMES) def test_validate_metric_name_good(metric_name): _validate_metric_name(metric_name) def _bad_parameter_pattern(name): if name == "\\": return r"Invalid value \"\\\\\" for parameter" # Manually handle the backslash case elif name == "*****": return r"Invalid value \"\*\*\*\*\*\" for parameter" else: return f'Invalid value "{name}" for parameter' @pytest.mark.parametrize("metric_name", BAD_METRIC_OR_PARAM_NAMES) def test_validate_metric_name_bad(metric_name): with pytest.raises( MlflowException, match=_bad_parameter_pattern(metric_name), ) as e: _validate_metric_name(metric_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.parametrize("param_name", GOOD_METRIC_OR_PARAM_NAMES) def test_validate_param_name_good(param_name): _validate_param_name(param_name) @pytest.mark.parametrize("param_name", BAD_METRIC_OR_PARAM_NAMES) def test_validate_param_name_bad(param_name): with pytest.raises(MlflowException, match=_bad_parameter_pattern(param_name)) as e: _validate_param_name(param_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.skipif(not is_windows(), reason="Windows do not support colon in params and metrics") @pytest.mark.parametrize( "param_name", [ ":", "aa:bb:cc", ], ) def test_validate_colon_name_bad_windows(param_name): with pytest.raises(MlflowException, match=_bad_parameter_pattern(param_name)) as e: _validate_param_name(param_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.parametrize("tag_name", GOOD_METRIC_OR_PARAM_NAMES) def test_validate_tag_name_good(tag_name): _validate_tag_name(tag_name) @pytest.mark.parametrize("tag_name", BAD_METRIC_OR_PARAM_NAMES) def test_validate_tag_name_bad(tag_name): with pytest.raises(MlflowException, match=_bad_parameter_pattern(tag_name)) as e: _validate_tag_name(tag_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.parametrize("alias_name", GOOD_ALIAS_NAMES) def test_validate_model_alias_name_good(alias_name): _validate_model_alias_name(alias_name) @pytest.mark.parametrize("alias_name", BAD_ALIAS_NAMES) def test_validate_model_alias_name_bad(alias_name): with pytest.raises(MlflowException, match="alias name") as e: _validate_model_alias_name(alias_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.parametrize("alias_name", ["latest", "LATEST", "Latest", "v123", "V1"]) def test_validate_model_alias_name_reserved(alias_name): with pytest.raises(MlflowException, match="reserved") as e: _validate_model_alias_name_reserved(alias_name) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) @pytest.mark.parametrize( "run_id", [ "a" * 32, "f0" * 16, "abcdef0123456789" * 2, "a" * 33, "a" * 31, "a" * 256, "A" * 32, "g" * 32, "a_" * 32, "abcdefghijklmnopqrstuvqxyz", ], ) def test_validate_run_id_good(run_id): _validate_run_id(run_id) @pytest.mark.parametrize("run_id", ["a/bc" * 8, "", "a" * 400, "*" * 5]) def test_validate_run_id_bad(run_id): with pytest.raises(MlflowException, match=_bad_parameter_pattern(run_id)) as e: _validate_run_id(run_id) assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) def test_validate_batch_log_limits(): too_many_metrics = [Metric(f"metric-key-{i}", 1, 0, i * 2) for i in range(1001)] too_many_params = [Param(f"param-key-{i}", "b") for i in range(101)] too_many_tags = [RunTag(f"tag-key-{i}", "b") for i in range(101)] good_kwargs = {"metrics": [], "params": [], "tags": []} bad_kwargs = { "metrics": [too_many_metrics], "params": [too_many_params], "tags": [too_many_tags], } match = r"A batch logging request can contain at most \d+" for arg_name, arg_values in bad_kwargs.items(): for arg_value in arg_values: final_kwargs = copy.deepcopy(good_kwargs) final_kwargs[arg_name] = arg_value with pytest.raises(MlflowException, match=match): _validate_batch_log_limits(**final_kwargs) # Test the case where there are too many entities in aggregate with pytest.raises(MlflowException, match=match): _validate_batch_log_limits(too_many_metrics[:900], too_many_params[:51], too_many_tags[:50]) # Test that we don't reject entities within the limit _validate_batch_log_limits(too_many_metrics[:1000], [], []) _validate_batch_log_limits([], too_many_params[:100], []) _validate_batch_log_limits([], [], too_many_tags[:100]) def test_validate_batch_log_data(monkeypatch): metrics_with_bad_key = [ Metric("good-metric-key", 1.0, 0, 0), Metric("super-long-bad-key" * 1000, 4.0, 0, 0), ] metrics_with_bad_val = [Metric("good-metric-key", "not-a-double-val", 0, 0)] metrics_with_bool_val = [Metric("good-metric-key", True, 0, 0)] metrics_with_bad_ts = [Metric("good-metric-key", 1.0, "not-a-timestamp", 0)] metrics_with_neg_ts = [Metric("good-metric-key", 1.0, -123, 0)] metrics_with_bad_step = [Metric("good-metric-key", 1.0, 0, "not-a-step")] params_with_bad_key = [ Param("good-param-key", "hi"), Param("super-long-bad-key" * 1000, "but-good-val"), ] params_with_bad_val = [ Param("good-param-key", "hi"), Param("another-good-key", "but-bad-val" * 1000), ] tags_with_bad_key = [ RunTag("good-tag-key", "hi"), RunTag("super-long-bad-key" * 1000, "but-good-val"), ] tags_with_bad_val = [ RunTag("good-tag-key", "hi"), RunTag("another-good-key", "a" * (MAX_TAG_VAL_LENGTH + 1)), ] bad_kwargs = { "metrics": [ metrics_with_bad_key, metrics_with_bad_val, metrics_with_bool_val, metrics_with_bad_ts, metrics_with_neg_ts, metrics_with_bad_step, ], "params": [params_with_bad_key, params_with_bad_val], "tags": [tags_with_bad_key, tags_with_bad_val], } good_kwargs = {"metrics": [], "params": [], "tags": []} monkeypatch.setenv("MLFLOW_TRUNCATE_LONG_VALUES", "false") for arg_name, arg_values in bad_kwargs.items(): for arg_value in arg_values: final_kwargs = copy.deepcopy(good_kwargs) final_kwargs[arg_name] = arg_value with pytest.raises(MlflowException, match=r".+"): _validate_batch_log_data(**final_kwargs) # Test that we don't reject entities within the limit _validate_batch_log_data( metrics=[Metric("metric-key", 1.0, 0, 0)], params=[Param("param-key", "param-val")], tags=[RunTag("tag-key", "tag-val")], ) @pytest.mark.parametrize("location", ["abcde", None]) def test_validate_experiment_artifact_location_good(location): _validate_experiment_artifact_location(location) @pytest.mark.parametrize("location", ["runs:/blah/bleh/blergh"]) def test_validate_experiment_artifact_location_bad(location): with pytest.raises(MlflowException, match="Artifact location cannot be a runs:/ URI"): _validate_experiment_artifact_location(location) @pytest.mark.parametrize("experiment_name", ["validstring", b"test byte string".decode("utf-8")]) def test_validate_experiment_name_good(experiment_name): _validate_experiment_name(experiment_name) @pytest.mark.parametrize("experiment_name", ["", 12, 12.7, None, {}, []]) def test_validate_experiment_name_bad(experiment_name): with pytest.raises(MlflowException, match="Invalid experiment name"): _validate_experiment_name(experiment_name) @pytest.mark.parametrize("db_type", ["mysql", "mssql", "postgresql", "sqlite"]) def test_validate_db_type_string_good(db_type): _validate_db_type_string(db_type) @pytest.mark.parametrize("db_type", ["MySQL", "mongo", "cassandra", "sql", ""]) def test_validate_db_type_string_bad(db_type): with pytest.raises(MlflowException, match="Invalid database engine") as e: _validate_db_type_string(db_type) assert "Invalid database engine" in e.value.message @pytest.mark.parametrize( "artifact_location", [ "s3://test-bucket/", "file:///path/to/artifacts", "mlflow-artifacts:/path/to/artifacts", "dbfs:/databricks/mlflow-tracking/some-id", ], ) def test_validate_experiment_artifact_location_length_good(artifact_location): _validate_experiment_artifact_location_length(artifact_location) @pytest.mark.parametrize( "artifact_location", ["s3://test-bucket/" + "a" * 10000, "file:///path/to/" + "directory" * 1111], ids=["s3_long_path", "file_long_path"], ) def test_validate_experiment_artifact_location_length_bad(artifact_location): with pytest.raises(MlflowException, match="Invalid artifact path length"): _validate_experiment_artifact_location_length(artifact_location) def test_setting_experiment_artifact_location_env_var_works(monkeypatch): artifact_location = "file://aaaa" # length 11 # should not throw _validate_experiment_artifact_location_length(artifact_location) # reduce limit to 10 monkeypatch.setenv(MLFLOW_ARTIFACT_LOCATION_MAX_LENGTH.name, "10") with pytest.raises(MlflowException, match="Invalid artifact path length"): _validate_experiment_artifact_location_length(artifact_location) # increase limit to 11 monkeypatch.setenv(MLFLOW_ARTIFACT_LOCATION_MAX_LENGTH.name, "11") _validate_experiment_artifact_location_length(artifact_location) @pytest.mark.parametrize( "param_value", [ ["1", "2", "3"], [], [1, 2, 3], ], ) def test_validate_list_param_with_valid_list(param_value): _validate_list_param("experiment_ids", param_value) def test_validate_list_param_with_none_not_allowed(): with pytest.raises(MlflowException, match="experiment_ids must be a list"): _validate_list_param("experiment_ids", None, allow_none=False) def test_validate_list_param_with_none_allowed(): _validate_list_param("experiment_ids", None, allow_none=True) @pytest.mark.parametrize( ("param_name", "param_value", "expected_type"), [ ("experiment_ids", 4, "int"), ("param_name", "value", "str"), ("my_param", {"key": "value"}, "dict"), ], ) def test_validate_list_param_with_invalid_type(param_name, param_value, expected_type): with pytest.raises( MlflowException, match=rf"{param_name} must be a list, got {expected_type}" ) as exc_info: _validate_list_param(param_name, param_value) assert f"Did you mean to use {param_name}=[{param_value!r}]?" in str(exc_info.value) assert exc_info.value.error_code == "INVALID_PARAMETER_VALUE" # -- _validate_webhook_url tests -- def _mock_getaddrinfo(ip_str): return lambda host, port, *a, **kw: [(None, None, None, None, (ip_str, 0))] @pytest.mark.parametrize( ("url", "expected_match"), [ (123, "Webhook URL must be a string"), ("", "Webhook URL cannot be empty"), (" ", "Webhook URL cannot be empty"), ("ftp://example.com", "Invalid webhook URL scheme"), ("http://example.com", "Invalid webhook URL scheme"), ("https://", "must include a hostname"), ], ) def test_validate_webhook_url_rejects_invalid_input(url, expected_match): with pytest.raises(MlflowException, match=expected_match): _validate_webhook_url(url) @pytest.mark.parametrize( ("url", "resolved_ip"), [ ("https://127.0.0.1/callback", "127.0.0.1"), ("https://localhost/callback", "127.0.0.1"), ("https://internal.corp/hook", "10.0.0.1"), ("https://internal.corp/hook", "172.16.0.1"), ("https://internal.corp/hook", "192.168.1.1"), ("https://metadata.internal/hook", "169.254.169.254"), ("https://cgnat.internal/hook", "100.64.0.1"), ("https://ipv6-loopback.internal/hook", "::1"), ("https://ipv6-private.internal/hook", "fc00::1"), ], ) def test_validate_webhook_url_rejects_private_ips(url, resolved_ip): with patch( "mlflow.utils.validation.socket.getaddrinfo", side_effect=_mock_getaddrinfo(resolved_ip), ): with pytest.raises(MlflowException, match="must not resolve to a non-public"): _validate_webhook_url(url) def test_validate_webhook_url_rejects_unresolvable_hostname(): with patch( "mlflow.utils.validation.socket.getaddrinfo", side_effect=socket.gaierror("Name or service not known"), ): with pytest.raises(MlflowException, match="Cannot resolve webhook URL hostname"): _validate_webhook_url("https://does-not-exist.invalid/hook") def test_validate_webhook_url_rejects_if_any_resolved_address_is_private(): def multi_resolve(host, port, *a, **kw): return [ (None, None, None, None, ("8.8.8.8", 0)), (None, None, None, None, ("10.0.0.1", 0)), ] with patch("mlflow.utils.validation.socket.getaddrinfo", side_effect=multi_resolve): with pytest.raises(MlflowException, match="must not resolve to a non-public"): _validate_webhook_url("https://dual-homed.example.com/hook") def test_validate_webhook_url_accepts_public_ip(): with patch( "mlflow.utils.validation.socket.getaddrinfo", side_effect=_mock_getaddrinfo("8.8.8.8"), ): _validate_webhook_url("https://example.com/webhook") def test_validate_webhook_url_allow_private_ips_env_var(monkeypatch): monkeypatch.setenv("MLFLOW_WEBHOOK_ALLOW_PRIVATE_IPS", "true") with patch( "mlflow.utils.validation.socket.getaddrinfo", side_effect=_mock_getaddrinfo("127.0.0.1"), ): _validate_webhook_url("https://localhost/callback") @pytest.mark.parametrize("invalid_name", ["my/model", "model:v1", "name/with:both"]) def test_validate_model_name_invalid_chars(invalid_name): with pytest.raises( MlflowException, match="Names cannot contain '/' or ':'", check=lambda e: e.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE), ): _validate_model_name(invalid_name) @pytest.mark.parametrize("invalid_name", ["my/model", "model:v1", "name/with:both"]) def test_validate_model_renaming_invalid_chars(invalid_name): with pytest.raises( MlflowException, match="Names cannot contain '/' or ':'", check=lambda e: e.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE), ): _validate_model_renaming(invalid_name)