def update_azureml_config(olive_config):
    """Update the azureml config in the olive config."""
    subscription_id = os.environ.get("WORKSPACE_SUBSCRIPTION_ID")
    if subscription_id is None:
        raise Exception("Please set the environment variable WORKSPACE_SUBSCRIPTION_ID")

    resource_group = os.environ.get("WORKSPACE_RESOURCE_GROUP")
    if resource_group is None:
        raise Exception("Please set the environment variable WORKSPACE_RESOURCE_GROUP")

    workspace_name = os.environ.get("WORKSPACE_NAME")
    if workspace_name is None:
        raise Exception("Please set the environment variable WORKSPACE_NAME")

    exclude_managed_identity_credential = (
        {"exclude_managed_identity_credential": True} if "EXCLUDE_MANAGED_IDENTITY_CREDENTIAL" in os.environ else {}
    )

    client_id = os.environ.get("MANAGED_IDENTITY_CLIENT_ID")
    if client_id is None and not exclude_managed_identity_credential:
        raise Exception("Please set the environment variable MANAGED_IDENTITY_CLIENT_ID")

    keyvault_name = os.environ.get("KEYVAULT_NAME")

    olive_config["azureml_client"] = {
        "subscription_id": subscription_id,
        "resource_group": resource_group,
        "workspace_name": workspace_name,
        # pipeline agents have multiple managed identities, so we need to specify the client_id
        "default_auth_params": {"managed_identity_client_id": client_id, **exclude_managed_identity_credential},
        "keyvault_name": keyvault_name,
    }
Update the azureml config in the olive config.
examples/test/utils.py