def set_tag(self, run_id, tag):
"""
Set a tag on a run.
:param run_id: String ID of the run
:param tag: RunTag instance to log
"""
with self.ManagedSessionMaker() as session:
_validate_tag(tag.key, tag.value)
run = self._get_run(run_uuid=run_id, session=session)
self._check_run_is_active(run)
session.merge(SqlTag(run_uuid=run_id, key=tag.key, value=tag.value))
def record_logged_model(self, run_id, mlflow_model):
if not isinstance(mlflow_model, Model):
raise TypeError("Argument 'mlflow_model' should be mlflow.models.Model, got '{}'"
.format(type(mlflow_model)))
model_dict = mlflow_model.to_dict()
with self.ManagedSessionMaker() as session:
run = self._get_run(run_uuid=run_id, session=session)
self._check_run_is_active(run)
previous_tag = [t for t in run.tags if t.key == MLFLOW_LOGGED_MODELS]
if previous_tag:
value = json.dumps(json.loads(previous_tag[0].value) + [model_dict])
else:
value = json.dumps([model_dict])
_validate_tag(MLFLOW_LOGGED_MODELS, value)
session.merge(SqlTag(key=MLFLOW_LOGGED_MODELS, value=value, run_uuid=run_id))
def _set_tag(self, run: ElasticRun, tag: RunTag) -> None:
_validate_tag(tag.key, tag.value)
new_tag = ElasticTag(key=tag.key, value=tag.value)
run.tags.append(new_tag)
