azureml-mlflow 1.61.0 incompatible with mlflow 2.22.4 due to ArtifactRepository signature change

[mikehelder]

• Package Name: azureml-mlflow
• Package Version: 1.61.0
• Operating System: Linux 6.8.0-1026-azure
• Python Version: 3.9.25

Describe the bug
Incompatibility between azureml-mlflow 1.61.0 and mlflow 2.22.4:
mlflow 2.22.4 changed the function signature of its base class. In azureml-mlflow 1.61.0, the AzureMLflowArtifactRepository is calling its parent class init with 3 arguments (artifact_uri, tracking_uri, registry_uri), but the parent class in mlflow 2.22.4 only accepts 1 argument (plus self).

TypeError: __init__() takes 2 positional arguments but 4 were given

MLflow 2.22.4 (Parent Class):

class ArtifactRepository:
    def __init__(self, artifact_uri):
        #         ^^^^^^^^^^^^^^^^^^^^
        #         Only accepts 1 explicit parameter (+ self implicitly)
        pass

azureml-mlflow 1.61.0 (Child Class):

class AzureMLflowArtifactRepository(ArtifactRepository):
    def __init__(self, artifact_uri, tracking_uri=None, registry_uri=None):
        # Tries to call parent with ALL parameters:
        super().__init__(artifact_uri, tracking_uri, registry_uri)
        #                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        #                Passing 3 explicit arguments (+ self implicitly)

To Reproduce
Steps to reproduce the behavior:

import mlflow
import azureml.mlflow

# Step 1: Show the MLflow base class signature
print("\n1. MLflow ArtifactRepository.__init__ signature:")
from mlflow.store.artifact.artifact_repo import ArtifactRepository
import inspect

sig = inspect.signature(ArtifactRepository.__init__)
print(f"   {sig}")
print(f"   Parameters: {list(sig.parameters.keys())}")

# Step 2: Show the azureml-mlflow subclass signature
print("\n2. AzureMLflowArtifactRepository.__init__ signature:")
from azureml.mlflow._store.artifact.artifact_repo import AzureMLflowArtifactRepository

sig_azure = inspect.signature(AzureMLflowArtifactRepository.__init__)
print(f"   {sig_azure}")
print(f"   Parameters: {list(sig_azure.parameters.keys())}")

# Step 3: Read the source code to show the problematic super().__init__() call
print("\n3. Source code of the problematic super().__init__() call:")
import azureml.mlflow._store.artifact.artifact_repo as azure_module

with open(azure_module.__file__, 'r') as f:
    lines = f.readlines()
    for i, line in enumerate(lines[40:45], start=41):
        print(f"   Line {i}: {line.rstrip()}")

# Step 4: Try to instantiate and trigger the error
print("\n4. Attempting to instantiate AzureMLflowArtifactRepository...")
print("   This will trigger the TypeError:\n")

repo = AzureMLflowArtifactRepository(
    artifact_uri="azureml://test",
    tracking_uri="http://test-tracking-uri",
    registry_uri="http://test-registry-uri"
)

Expected behavior

1. Initialization succeeds
2. Artifacts can be logged to Azure ML workspace
3. Training pipeline completes successfully

Additional context
Works when downgrading to azureml-mlflow==1.59.0 and mlflow==2.19.0.

[github-actions]

Thanks for the feedback! We are routing this to the appropriate team for follow-up. cc @Azure/azure-ml-sdk @azureml-github.