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)