slalom/qai-cli
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Mlflow implementation #2

Merged
slalom merged 10 commits from ml-flow into main 2026-06-02 19:04:23 +00:00
Conversation 0 Commits 10 Files Changed 13 +1 -4
1 changed files with 1 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit cf6a561e2f - Show all commits

5
src/tracking/mlflow.py
View File

@@ -122,7 +122,7 @@ class MlflowTracker:
}, },
) )
version_number = str(version.version) 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_name", self.registered_model_name)
mlflow.set_tag("qc_cli.registered_model_version", version_number) mlflow.set_tag("qc_cli.registered_model_version", version_number)
return version_number return version_number
@@ -147,6 +147,3 @@ class MlflowTracker:
client.get_registered_model(name) client.get_registered_model(name)
except Exception: except Exception:
client.create_registered_model(name) 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)