pythonpython-3.xmlflow

How to decorate all callable methods of an imported lib


I am using MLFlow lib to write data to a MLFlow instance running in the cloud.

Currently I have to deal with authentication before calling some methods of the lib mlflow:

# First I deal with auth explicitly
os.environ["MLFLOW_TRACKING_TOKEN"] = my_custom_get_token_function()
...

# Only then I use the mlflow methods
mlflow.log_metric(...)

The above code runs successfuly.

I would like that all callable mlflow methods first execute my custom auth logic, and only then execute the mlflow logic. For that I though of decorating all callable mlflow methods.

The idea is to do something like:

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

This way, when I run mlflow.log_metric(...), it would first execute my custom auth logic, since mlflow now is really the decorated mlflow, and only then it would make the mlflow API request behind the log_metric method.

To accomplish that I came up with the following code:

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() # this functions has my custom auth logic
        return func(*args, **kwargs)

    return wrapper

import mlflow
mlflow = authenticate_mlflow(mlflow)

mlflow.log_metric(...)

But when I run the above code with python 3.11, it yields the error:

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

This is a good way of accomplishing the task "run my custom auth logic automatically before every callable mlflow method"? If so, what I am getting wrong in the implementation?


Solution

  • You've made the code unnecessarily complicated.

    The logic is;

    Get all attributes of the module.
    Filter those attributes if that start with _ an underscore because they are private. Replace them if they're callable.

    def authenticate_mlflow(mod):
        for attr_name in dir(mod):
            if attr_name.startswith('_'):
                continue
            attr = getattr(mod, attr_name)
            if callable(attr):
                setattr(mod, attr_name, authenticate_mlflow_method(attr))
        
    

    Since the object passed to this function refers to a module object and we're setting it's attributes, there's no need to return the module from this function and reassign it to the module name.