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

rename and future steps

Browse Source
This commit is contained in:
slalom
2026-05-29 15:40:38 -04:00
parent a43c792cfd
commit 57a8a0a9c4
4 changed files with 52 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
src/tracking/mlflow.py
View File

@@ -78,7 +78,9 @@ class MlflowTracker:
self._log_params({f"hyperparameters.{key}": value for key, value in training_job.hyperparameters.items()})
mlflow.set_tags(
{
"qc_cli.stage": "prerelease",
"qc_cli.stage": "experiment",
"qc_cli.artifact_kind": "trained_source",
"qc_cli.source": "sagemaker",
"qc_cli.command": "train start",
"sagemaker.job_name": training_job.job_name,
}
@@ -117,12 +119,14 @@ class MlflowTracker:
source=training_job_status.model_artifacts,
run_id=run_id,
tags={
"qc_cli.stage": "prerelease",
"qc_cli.stage": "experiment",
"qc_cli.artifact_kind": "trained_source",
"qc_cli.source": "sagemaker",
"sagemaker.job_name": training_job_status.name,
},
)
version_number = str(version.version)
client.set_registered_model_alias(self.registered_model_name, "prerelease-latest", version_number)
client.set_registered_model_alias(self.registered_model_name, "experiment-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