diff --git a/src/tracking/mlflow.py b/src/tracking/mlflow.py index b38e8af..9275d9a 100644 --- a/src/tracking/mlflow.py +++ b/src/tracking/mlflow.py @@ -122,7 +122,7 @@ class MlflowTracker: }, ) version_number = str(version.version) - self._set_alias(client, self.registered_model_name, "prerelease-latest", version_number) + client.set_registered_model_alias(self.registered_model_name, "prerelease-latest", version_number) mlflow.set_tag("qc_cli.registered_model_name", self.registered_model_name) mlflow.set_tag("qc_cli.registered_model_version", version_number) return version_number @@ -147,6 +147,3 @@ class MlflowTracker: client.get_registered_model(name) except Exception: client.create_registered_model(name) - - def _set_alias(self, client: MlflowClient, name: str, alias: str, version: str) -> None: - client.set_registered_model_alias(name, alias, version)