如何自动为MLFlow库可调用方法添加前置认证？排查实现报错

当前手动认证方式 #

我正在使用MLFlow向云端实例写入数据，目前需要在调用MLFlow方法前手动处理认证：

# 首先显式处理认证
os.environ["MLFLOW_TRACKING_TOKEN"] = my_custom_get_token_function()
...

# 之后才能使用mlflow方法
mlflow.log_metric(...)

这段代码可以正常运行，但每次调用MLFlow方法前都要手动执行认证逻辑，比较繁琐。

期望的自动认证效果 #

我希望所有MLFlow可调用方法都能先执行自定义认证逻辑，再执行原方法，期望用法如下：

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

调用mlflow.log_metric(...)时，自动先执行认证逻辑，再发起API请求。

自行实现的代码 #

我尝试用类继承+装饰器的方式实现，代码如下：

def authenticate_mlflow(cls):
    class AuthenticatedMlflow(cls):
        @staticmethod
        def __getattribute__(cls, name):
            attr = super().__getattribute__(name)
            if callable(attr):
                return authenticate_mlflow_method(attr)
            return attr

    return AuthenticatedMlflow

def authenticate_mlflow_method(func):
    def wrapper(*args, **kwargs):
        handle_auth() # 该函数包含自定义认证逻辑
        return func(*args, **kwargs)

    return wrapper

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

报错信息 #

在Python 3.11中运行上述代码时，出现如下错误：

mlflow = authenticate_mlflow(mlflow)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^
File ..., in authenticate_mlflow
    class AuthenticatedMlflow(cls):
...
TypeError: module() takes at most 2 arguments (3 given)

问题 #

请问这种自动认证的方案是否可行？如果可行，我的代码存在什么问题？

这个自动认证的方案思路是可行的，但你的代码存在核心错误：试图让Python类继承模块对象。

代码错误原因 #

authenticate_mlflow函数中，你传入的cls参数是mlflow模块（不是类），但Python的类只能继承其他类，不能继承模块。当你写class AuthenticatedMlflow(cls):时，Python会尝试把模块当作父类来继承，这违反了语法规则，因此抛出TypeError。

修正后的实现方案 #

改用代理对象的方式来包装mlflow模块，通过重写__getattr__方法拦截属性访问，对可调用方法自动添加认证装饰器：

import types

def authenticate_mlflow(module):
    class MlflowProxy:
        def __getattr__(self, name):
            # 获取原模块的属性
            attr = getattr(module, name)
            # 如果是子模块，递归包装
            if isinstance(attr, types.ModuleType):
                return authenticate_mlflow(attr)
            # 如果是可调用对象，添加认证装饰器
            if callable(attr):
                return authenticate_mlflow_method(attr)
            # 非可调用属性直接返回
            return attr

    return MlflowProxy()

def authenticate_mlflow_method(func):
    def wrapper(*args, **kwargs):
        handle_auth()  # 执行自定义认证逻辑
        return func(*args, **kwargs)
    return wrapper

# 使用方式不变
import mlflow
mlflow = authenticate_mlflow(mlflow)

# 调用方法时自动执行认证
mlflow.log_metric("accuracy", 0.95)

方案说明 #

1. 代理对象MlflowProxy通过__getattr__拦截所有属性访问请求
2. 若访问的是mlflow的子模块（如mlflow.sklearn），则递归创建代理，确保嵌套模块的方法也能自动认证
3. 若访问的是可调用方法（如log_metric），则用authenticate_mlflow_method装饰器包装，先执行认证逻辑再调用原方法
4. 普通属性（如常量、配置值）直接返回原模块的对应值

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