AzureMLOnBehalfOfCredential does not work with keyvault

[jonaslb]

• Package Name: azure-ai-ml
• Package Version: 1.25.0
• Operating System: Linux
• Python Version: 3.12

Describe the bug

Azure AI ML describes in its documentation that AzureMLOnBehalfOfCredential can be used to access for example keyvault. Due to what seems to be a programming error (TypeError: Session.request() got an unexpected keyword argument 'claims'), that is however not possible.

To Reproduce

Here's a minimal reproducer:

# hello.py
import sys
from os import environ

from azure.ai.ml import MLClient, command
from azure.ai.ml.identity import AzureMLOnBehalfOfCredential
from azure.ai.ml.entities import JobResourceConfiguration, UserIdentityConfiguration, Environment, BuildContext
from azure.identity import DefaultAzureCredential
from azure.keyvault.secrets import SecretClient


def do_submit() -> None:
    client = MLClient(
        credential=DefaultAzureCredential(),
        subscription_id=environ.get("AZURE_ML_SUBSCRIPTION_ID"),
        workspace_name=environ.get("AZURE_ML_WORKSPACE"),
        resource_group_name=environ.get("AZURE_ML_RESOURCE_GROUP"),
    )
    job = command(
        code=".",
        command="python hello.py kv --obo",
        environment=Environment(
            name="OboAuth-MRE-Env",
            version="1",
            build=BuildContext(
                path=".",
                dockerfile_path="Dockerfile",
            ),
        ),
        display_name="mre-run",
        experiment_name="az-auth-issue-mre",
        identity=UserIdentityConfiguration(),
        resources=JobResourceConfiguration(
            instance_type="Standard_E4_v3",
            instance_count=1,
        ),
        queue_settings={"job_tier": "Standard"},
    )
    res = client.jobs.create_or_update(job)
    print("Job submitted:", res)


def do_access_kv(use_obo: bool = False) -> None:
    if use_obo:
        credential = AzureMLOnBehalfOfCredential()
    else:
        credential = DefaultAzureCredential()
    client = SecretClient(
        vault_url="https://MY-KV-URL-HERE.vault.azure.net/",
        credential=credential,
    )
    value = client.get_secret("MY-SECRET-HERE").value
    assert value is not None, "Failed to access key vault"
    print("Successfully accessed key vault")


if __name__ == "__main__":
    if "submit" in sys.argv:
        do_submit()
    elif "kv" in sys.argv:
        do_access_kv(use_obo="--obo" in sys.argv)
    else:
        print("Usage: python hello.py [submit|kv] [--obo]")
        sys.exit(1)

And

FROM python:3.12
RUN pip install azure-ai-ml azure-identity azure-keyvault

• First, verify the output from running hello.py kv locally is "Successfully accessed key vault".
• Second, run hello.py submit to submit to Azure ML
• Third, wait and observe that the Azure ML job fails with the exception TypeError: Session.request() got an unexpected keyword argument 'claims'

Expected behavior

I expected the "on behalf of" credential to work essentially the same as the default credential, with my identity.

This is also how documentation in the examples repository thinks it should work.

[github-actions]

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

[achauhan-scc]

@jonaslb - Can you please confirm which compute you are using; from code snippet it looks like serverless?
I believe OBO is supported on AMLCompute only.

[github-actions]

Hi @jonaslb. Thank you for opening this issue and giving us the opportunity to assist. To help our team better understand your issue and the details of your scenario please provide a response to the question asked above or the information requested above. This will help us more accurately address your issue.

[jonaslb]

@achauhan-scc that's correct, the snippet starts a serverless job. Where is it documented that this doesn't work? In any case the error (TypeError from what i guess is the requests library) doesn't really communicate this

[rohit-ganguly]

Hi @jonaslb, not an expert on our ML SDKs but I do know that Azure Identity on its own has an OnBehalfOfCredential that might be useful.

cc @xiangyan99 who can speak more on Azure Identity q's

[jonaslb]

Thanks for looking for alternatives! That credential still looks like it requires a client secret that I don't have on azure ml. I think I need either a managed identity or an "identity forward" solution. I thought azureml-on-behalf-of was the easiest, since it was mentioned in the examples, but it might be not for serverless, as achauhan says. If that turns out to be the case, I'll need to dig into assigning a managed identity to the serverless job - or maybe if there was another intended way of using auth on serverless, then I'm certainly all ears :)

[achauhan-scc]

@jonaslb - Can you please explain your scenario, so that I look around for better suggestions.
I will get the documentation updated to OBO on ML to AMLCompute only.

[github-actions]

Hi @jonaslb. Thank you for opening this issue and giving us the opportunity to assist. To help our team better understand your issue and the details of your scenario please provide a response to the question asked above or the information requested above. This will help us more accurately address your issue.

[jonaslb]

The scenario is almost as simple as the reproducer snippet in the issue. We use the jobs to perform some ad-hoc "data engineering"/transforms/reduction tasks in preparation for other tasks. It's often submitted from a laptop, although it can also be from an app that uses a managed identity. But anyway, one of the things these jobs do, is access some external APIs. For this, secrets are needed, and they are in an Azure Vault.

We use environment variables currently, but these are not actually considered secret, and hence viewable on e.g. ml studio, so we want to stop doing that.

I should say that we might be in a situation where Azure ML serverless is "overkill" for some of the jobs. So we are considering other products, such as Azure Batch or other orchestration tools. But for now, things are on Azure ML.