[go: nahoru, domu]

Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update the client to allow for retrieving the latest alias #12342

Open
wants to merge 1 commit into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Update the client to allow for retrieving the latest alias
  • Loading branch information
kriscon-db committed Jun 13, 2024
commit 8e86dfaa38d7a79bfebececaad520985bd2ac646
2 changes: 1 addition & 1 deletion mlflow/store/model_registry/file_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -942,7 +942,7 @@ def delete_model_version_tag(self, name, version, key):

def _get_registered_model_alias_path(self, name, alias):
_validate_model_name(name)
_validate_model_alias_name(alias)
_validate_model_alias_name(alias, perform_latest_check=False)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

From offline discussion, we don't want to remove validation from the OSS storage backends (file/SQL store)

registered_model_path = self._get_registered_model_path(name)
if not exists(registered_model_path):
raise MlflowException(
Expand Down
2 changes: 1 addition & 1 deletion mlflow/store/model_registry/sqlalchemy_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -1173,7 +1173,7 @@ def get_model_version_by_alias(self, name, alias):
A single :py:class:`mlflow.entities.model_registry.ModelVersion` object.
"""
_validate_model_name(name)
_validate_model_alias_name(alias)
_validate_model_alias_name(alias, perform_latest_check=False)
with self.ManagedSessionMaker() as session:
existing_alias = self._get_registered_model_alias(session, name, alias)
if existing_alias is not None:
Expand Down
2 changes: 1 addition & 1 deletion mlflow/tracking/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4729,5 +4729,5 @@ def print_model_version_info(mv):
Aliases: ["test-alias"]
"""
_validate_model_name(name)
_validate_model_alias_name(alias)
_validate_model_alias_name(alias, perform_latest_check=False)
Copy link
Collaborator
@smurching smurching Jun 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It does make sense to remove validation from this client codepath (and in rest_store etc if we have validation there). I think on reads just removing validation entirely is fine, we can let the backend throw on an invalid name etc. For removing client-side write validation, we should make sure analogous logic is first present in all backends we maintain for at least one OSS release

return self._get_registry_client().get_model_version_by_alias(name, alias)
4 changes: 2 additions & 2 deletions mlflow/utils/validation.py
Original file line number Diff line number Diff line change
Expand Up @@ -397,7 +397,7 @@ def _validate_model_version(model_version):
)


def _validate_model_alias_name(model_alias_name):
def _validate_model_alias_name(model_alias_name, perform_latest_check=True):
Copy link
Collaborator
@smurching smurching Jun 14, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can probably get rid of this boolean / latest-specific logic if we switch to removing validation on reads + enforcing it on writes. And then this PR can be about removing client-side validation on fetching model versions by alias

if model_alias_name is None or model_alias_name == "":
raise MlflowException(
"Registered model alias name cannot be empty.", INVALID_PARAMETER_VALUE
Expand All @@ -410,7 +410,7 @@ def _validate_model_alias_name(model_alias_name):
_validate_length_limit(
"Registered model alias name", MAX_REGISTERED_MODEL_ALIAS_LENGTH, model_alias_name
)
if model_alias_name.lower() == "latest":
if perform_latest_check and model_alias_name.lower() == "latest":
raise MlflowException(
"'latest' alias name (case insensitive) is reserved.", INVALID_PARAMETER_VALUE
)
Expand Down
16 changes: 13 additions & 3 deletions tests/utils/test_validation.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
import copy

Check failure on line 1 in tests/utils/test_validation.py

View workflow job for this annotation

GitHub Actions / lint

Unformatted file. Run `ruff format .` or comment `@mlflow-automation autoformat` to format.

import pytest

Expand Down Expand Up @@ -59,7 +59,7 @@
"temp_V123",
]

BAD_ALIAS_NAMES = [
BAD_ALIAS_NAMES_WITHOUT_LATEST = [
"",
".",
"/",
Expand All @@ -76,12 +76,14 @@
"a" * 256,
None,
"$dgs",
"latest",
"Latest",
"v123",
"V1",
]

LATEST_TESTS = ["latest", "Latest"]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want all caps as well?

Suggested change
LATEST_TESTS = ["latest", "Latest"]
LATEST_TESTS = ["latest", "Latest", "LATEST"]

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sure. I can add.


BAD_ALIAS_NAMES = BAD_ALIAS_NAMES_WITHOUT_LATEST + LATEST_TESTS


@pytest.mark.parametrize(
("path", "expected"),
Expand Down Expand Up @@ -165,6 +167,14 @@
_validate_model_alias_name(alias_name)
assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)

@pytest.mark.parametrize("alias_name", BAD_ALIAS_NAMES_WITHOUT_LATEST)
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Update this to just be BAD_ALIAS_NAMES

def test_validate_model_alias_name_bad_but_allow_latest(alias_name):
if alias_name not in LATEST_TESTS:
with pytest.raises(MlflowException, match="alias name") as e:
_validate_model_alias_name(alias_name, perform_latest_check=False)
assert e.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)
else:
_validate_model_alias_name(alias_name)

@pytest.mark.parametrize(
"run_id",
Expand Down
Loading