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!

You can't do that, because the decorator is called when the *class* is defined, long before any instance of the class exists. — chepner, Jun 01 '21 at 19:34
You left out the part of the decorator that actually *uses* the `experiment_name` argument, so it's hard to recommend a work-around. — chepner, Jun 01 '21 at 19:40
@chepner, sorry about that. I realized I did leave it out. I'll edit it back in. Also found this: https://stackoverflow.com/questions/11731136/class-method-decorator-with-self-arguments — dyao, Jun 01 '21 at 19:45

chepner · Accepted Answer · 2021-06-01T20:01:23.080

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.

Thanks, is it correct to say that the `self` object is not in scope of `mlflow_experiment` so when I try and pass it in, it doesn't know what `self` is? — dyao, Jun 01 '21 at 20:06
Right. `self` is *just* another parameter name in the method itself; when you *call* the method, the descriptor protocol ensures that the name is bound to the invoking object, but it has no special status in the language. It helps to desugar the decorator syntax; it's equivalent to defining `train_mlflow` like a normal method, then writing `train_mlflow = mlflow_experiment(autolog=True)(train_mlflow)`. — chepner, Jun 01 '21 at 20:08

Passing instance variables to a decorator

1 Answers1

Linked