from mlflow.tracking import MlflowClient


class RegistryManager:

    def __init__(self):
        self.client = MlflowClient()

    def promote(self, model_name, version, stage):
        self.client.transition_model_version_stage(
            name=model_name,
            version=version,
            stage=stage
        )
