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?
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.