如何修改mlflow装饰器以支持传入类实例变量作为参数

装饰器的参数解析执行时机是类定义阶段，此时类还未完成实例化，self根本不存在，所以直接传self.arg作为装饰器参数会触发未定义错误。

核心思路是修改装饰器逻辑，把参数解析延迟到方法实际调用阶段，此时实例已经生成，self可以正常访问，下面提供两种常用实现方案：

方案1：传入实例属性名称（简单易用） #

修改装饰器，支持传入要读取的实例属性名字符串，方法调用时自动从第一个参数（即self）中读取对应属性值：

修改后的装饰器代码 #

import mlflow
from functools import wraps

def mlflow_experiment(experiment_name=None, exp_from_attr: str = None):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            # 类方法的第一个参数就是实例本身self
            self = args[0]
            # 优先从指定的实例属性中读取实验名
            current_exp_name = getattr(self, exp_from_attr) if exp_from_attr else experiment_name
            mlflow.set_experiment(current_exp_name)
            with mlflow.start_run():
                return func(*args, **kwargs)
        return wrapper
    return decorator

用法示例 #

class MyTrainer:
    def __init__(self, experiment_name):
        self.experiment_name = experiment_name
    
    # 传入你要读取的实例属性的名称字符串即可
    @mlflow_experiment(exp_from_attr="experiment_name")
    def train(self):
        # 你的训练逻辑
        pass

方案2：传入可调用对象（灵活度更高） #

如果需要支持更复杂的动态参数逻辑（比如多属性拼接、条件判断等），可以让装饰器支持传入可调用对象，方法调用时将self传入可调用对象获取实际参数值：

修改后的装饰器代码 #

import mlflow
from functools import wraps

def mlflow_experiment(experiment_name=None):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            self = args[0]
            # 如果传入的是可调用对象，执行后获取实际实验名
            current_exp_name = experiment_name(self) if callable(experiment_name) else experiment_name
            mlflow.set_experiment(current_exp_name)
            with mlflow.start_run():
                return func(*args, **kwargs)
        return wrapper
    return decorator

用法示例 #

class MyTrainer:
    def __init__(self, exp_prefix, exp_name):
        self.exp_prefix = exp_prefix
        self.exp_name = exp_name
    
    # 传入lambda表达式实现灵活的参数拼接
    @mlflow_experiment(experiment_name=lambda self: f"{self.exp_prefix}_{self.exp_name}")
    def train(self):
        # 你的训练逻辑
        pass

• 方案1写法更简洁，适合仅需要读取单个实例属性的场景，可读性更高
• 方案2支持任意复杂的动态参数逻辑，可满足个性化的参数构造需求

内容的提问来源于stack exchange，提问作者dyao