def create_registered_model(self, name, tags=None): """ Create a new registered model in backend store. :param name: Name of the new model. This is expected to be unique in the backend store. :param tags: A list of :py:class:`mlflow.entities.model_registry.RegisteredModelTag` instances associated with this registered model. :return: A single object of :py:class:`mlflow.entities.model_registry.RegisteredModel` created in the backend. """ _validate_model_name(name) for tag in tags or []: _validate_registered_model_tag(tag.key, tag.value) with self.ManagedSessionMaker() as session: try: creation_time = now() registered_model = SqlRegisteredModel( name=name, creation_time=creation_time, last_updated_time=creation_time) tags_dict = {} for tag in tags or []: tags_dict[tag.key] = tag.value registered_model.registered_model_tags = [ SqlRegisteredModelTag(key=key, value=value) for key, value in tags_dict.items() ] self._save_to_db(session, registered_model) session.flush() return registered_model.to_mlflow_entity() except sqlalchemy.exc.IntegrityError as e: raise MlflowException( 'Registered Model (name={}) already exists. ' 'Error: {}'.format(name, str(e)), RESOURCE_ALREADY_EXISTS)
def set_registered_model_tag(self, name, tag): """ Set a tag for the registered model. :param name: Registered model name. :param tag: :py:class:`mlflow.entities.model_registry.RegisteredModelTag` instance to log. :return: None """ _validate_model_name(name) _validate_registered_model_tag(tag.key, tag.value) with self.ManagedSessionMaker() as session: # check if registered model exists self._get_registered_model(session, name) session.merge(SqlRegisteredModelTag(name=name, key=tag.key, value=tag.value))