Passing instance variables to a decorator

Question

I found this useful decorator that allows you to pass in some optional arguments

def mlflow_experiment(
    _func=None,
    *,
    experiment_name=None
  ):
      def experiment_decorator(func):
          @functools.wraps(func)
          def experiment_wrapper(self, *args, **kwargs):
              nonlocal experiment_name

              experiment_id = (
                  mlflow.set_experiment(experiment_name)
                  if experiment_name is not None
                  else None
              )
                ...

              value = func(self, *args, **kwargs)

              return value

          return experiment_wrapper

      if _func is None:
          return experiment_decorator
      else:
          return experiment_decorator(_func)

So in a use case like this where I just pass in a string to experiment_name, the code works flawlessly.

@mlflow_experiment(autolog=True, experiment_name = 'blarg')    
def train_mlflow(self, maxevals=50, model_id=0):
  ...

I've always had a hard time figuring out scope in decorators but I wasn't surprised that using passing an instance variable defined in __init__ does NOT work.

class LGBOptimizerMlfow:
    def __init__(self, arg):
        self.arg = arg

    @mlflow_experiment(autolog=True, experiment_name = self.arg)    
    def train_mlflow(self, maxevals=50, model_id=0):
        ...

>>> `NameError: name 'self' is not defined`

Just to see if scoping was an issue, I declared the variable outside the class and it worked.

And just for the heck of it I decided to declare a global variable inside the class which also works but its less than ideal, especially if I want to pass it into the class or a method as a optional argument.

class LGBOptimizerMlfow:
    global arg
    arg = 'hi'

    @mlflow_experiment(autolog=True, experiment_name = arg)    
    def train_mlflow(self, maxevals=50, model_id=0):
    ...
  

Any help to revise the code so that the decorator accepts an instance variable would be lovely.

Thank you!

Answer 1

Decorators are called while the class is being defined, and self is simply a parameter used for each instance method, not something the class itself provides. So self is not defined at the time you need it to be for use as an argument to your decorator.

You need to modify experiment_wrapper to take a name directly from its self argument, rather than from an argument to mflow_experiment. Something like

def mlflow_experiment(
    _func=None,
    *,
    experiment_name=None,
    tracking_uri=None,
    autolog=False,
    run_name=None,
    tags=None,
  ):
      def experiment_decorator(func):
          @functools.wraps(func)
          def experiment_wrapper(self, *args, **kwargs):
              nonlocal tracking_uri

              experiment_name = getattr(self, 'experiment_name', None)
              experiment_id = (
                  mlflow.set_experiment(experiment_name)
                  if experiment_name is not None
                  else None
              )
                ...

              with mlflow.start_run(experiment_id=experiment_id
                                  , run_name=run_name
                                  , tags=tags):
              value = func(self, *args, **kwargs)

              return value

          return experiment_wrapper

      if _func is None:
          return experiment_decorator
      else:
          return experiment_decorator(_func)

Then you need to make sure that each instance has an experiment name (or None) associated with it.

class LGBOptimizerMlfow:
    def __init__(self, arg, experiment_name=None):
        self.arg = arg
        self.experiment_name = experiment_name

    @mlflow_experiment(autolog=True, experiment_name = self.arg)    
    def train_mlflow(self, maxevals=50, model_id=0):
        ...

---

Another alternative is to make experiment_name an argument to train_mflow, making it easier to create different names with the same method. (This may be closer to what you were intending.)

class LGBOptimizerMlfow:
    def __init__(self, arg):
        self.arg = arg

    @mlflow_experiment(autolog=True)    
    def train_mlflow(self, maxevals=50, model_id=0, experiment_name=None):
        if experiment_name is None:
            self.experiment_name = self.arg
        ...

The definition of the decorator remains the same as shown above.