diff --git a/src/aimanager/azext_aimanager/_client_factory.py b/src/aimanager/azext_aimanager/_client_factory.py index c954a2f1690..63067e9881d 100644 --- a/src/aimanager/azext_aimanager/_client_factory.py +++ b/src/aimanager/azext_aimanager/_client_factory.py @@ -29,3 +29,7 @@ def cf_model_deployments(cli_ctx, *_): def cf_ai_models(cli_ctx, *_): return get_aimanager_client(cli_ctx).ai_models + + +def cf_model_sources(cli_ctx, *_): + return get_aimanager_client(cli_ctx).model_sources diff --git a/src/aimanager/azext_aimanager/_help.py b/src/aimanager/azext_aimanager/_help.py index 211cb0fc147..60e7fb43d79 100644 --- a/src/aimanager/azext_aimanager/_help.py +++ b/src/aimanager/azext_aimanager/_help.py @@ -76,6 +76,67 @@ text: az aimanager get-credentials --name my-ai-manager -g myrg -f - """ +helps['aimanager modelsource'] = """ + type: group + short-summary: Manage model sources within an AI Manager. + long-summary: |- + A model source tells the platform where to pull model artifacts from and, for gated or + private sources, which credential to authenticate with. Model sources are referenced by + 'az aimanager namespace modeldeployment add --model-source-resource-id'. +""" + +helps['aimanager modelsource add'] = """ + type: command + short-summary: Add a model source to an AI Manager. + examples: + - name: Add a public Hugging Face model source + text: az aimanager modelsource add -g myrg --aimanager-name my-ai-manager -n hf --source-type HuggingFace + - name: Add a Hugging Face model source with an access token for gated models + text: az aimanager modelsource add -g myrg --aimanager-name my-ai-manager -n hf -s HuggingFace --token hf_xxx --description "Gated models" +""" + +helps['aimanager modelsource update'] = """ + type: command + short-summary: Update a model source within an AI Manager. + long-summary: |- + The source type is immutable after creation and is always preserved. Omitted properties + keep their current values. + examples: + - name: Rotate the access token of a model source + text: az aimanager modelsource update -g myrg --aimanager-name my-ai-manager -n hf --token hf_yyy + - name: Update the description of a model source + text: az aimanager modelsource update -g myrg --aimanager-name my-ai-manager -n hf --description "Internal mirror" +""" + +helps['aimanager modelsource show'] = """ + type: command + short-summary: Show the details of a model source within an AI Manager. + examples: + - name: Show a model source + text: az aimanager modelsource show -g myrg --aimanager-name my-ai-manager -n hf +""" + +helps['aimanager modelsource list'] = """ + type: command + short-summary: List the model sources within an AI Manager. + examples: + - name: List model sources + text: az aimanager modelsource list -g myrg --aimanager-name my-ai-manager +""" + +helps['aimanager modelsource delete'] = """ + type: command + short-summary: Delete a model source from an AI Manager. + examples: + - name: Delete a model source + text: az aimanager modelsource delete -g myrg --aimanager-name my-ai-manager -n hf +""" + +helps['aimanager modelsource wait'] = """ + type: command + short-summary: Wait for an AI Manager model source to reach a desired state. +""" + helps['aimanager namespace'] = """ type: group short-summary: Manage namespaces within an AI Manager. diff --git a/src/aimanager/azext_aimanager/_params.py b/src/aimanager/azext_aimanager/_params.py index 16800dde306..cf429b5f9cc 100644 --- a/src/aimanager/azext_aimanager/_params.py +++ b/src/aimanager/azext_aimanager/_params.py @@ -10,12 +10,17 @@ get_resource_name_completion_list, ) from azure.cli.core.commands.validators import get_default_location_from_resource_group -from azext_aimanager.constants import DELETE_POLICIES, MODEL_DEPLOYMENT_PERFORMANCE_MODES +from azext_aimanager.constants import ( + DELETE_POLICIES, + MODEL_DEPLOYMENT_PERFORMANCE_MODES, + MODEL_SOURCE_TYPES, +) from azext_aimanager._validators import ( validate_ai_manager_name, validate_namespace_name, validate_model_deployment_name, validate_ai_model_name, + validate_model_source_name, validate_labels, validate_annotations, validate_overrides, @@ -54,6 +59,32 @@ def load_arguments(self, _): c.argument('aks_custom_headers', options_list=['--aks-custom-headers'], help='Comma-separated key=value pairs to specify custom headers.') + with self.argument_context('aimanager modelsource') as c: + c.argument('ai_manager_name', options_list=['--aimanager-name'], + validator=validate_ai_manager_name, + help='The name of the AI Manager resource.') + c.argument('model_source_name', options_list=['--name', '-n'], + validator=validate_model_source_name, + help='The name of the model source.') + + with self.argument_context('aimanager modelsource list') as c: + c.ignore('model_source_name') + + with self.argument_context('aimanager modelsource add') as c: + c.argument('source_type', options_list=['--source-type', '-s'], required=True, + arg_type=get_enum_type(MODEL_SOURCE_TYPES), + help='The type of the model source. Immutable after creation.') + + for scope in ['aimanager modelsource add', 'aimanager modelsource update']: + with self.argument_context(scope) as c: + c.argument('description', + help='An optional, free-form description of the model source.') + c.argument('token', + help='Access token used by the platform to authenticate to the source. ' + 'Optional for public sources such as ungated Hugging Face models.') + c.argument('aks_custom_headers', options_list=['--aks-custom-headers'], + help='Comma-separated key=value pairs to specify custom headers.') + with self.argument_context('aimanager namespace') as c: c.argument('ai_manager_name', options_list=['--manager', '-m'], validator=validate_ai_manager_name, diff --git a/src/aimanager/azext_aimanager/_validators.py b/src/aimanager/azext_aimanager/_validators.py index ea536c3f406..f4725bf1bca 100644 --- a/src/aimanager/azext_aimanager/_validators.py +++ b/src/aimanager/azext_aimanager/_validators.py @@ -28,6 +28,12 @@ def validate_ai_model_name(namespace): raise InvalidArgumentValueError("--name/-n is not a valid AI model name.") +def validate_model_source_name(namespace): + name = getattr(namespace, "model_source_name", None) + if name is not None and not name.strip(): + raise InvalidArgumentValueError("--name/-n is not a valid model source name.") + + def _validate_key_value_pairs(values, option): if not values: return diff --git a/src/aimanager/azext_aimanager/commands.py b/src/aimanager/azext_aimanager/commands.py index 6d529534be8..aecfa53e7cc 100644 --- a/src/aimanager/azext_aimanager/commands.py +++ b/src/aimanager/azext_aimanager/commands.py @@ -14,6 +14,7 @@ cf_ai_manager_namespaces, cf_model_deployments, cf_ai_models, + cf_model_sources, ) @@ -63,6 +64,21 @@ def load_command_table(self, _): g.custom_command("get-credentials", "aimanager_namespace_get_credentials") g.wait_command("wait") + model_sources_sdk = CliCommandType( + operations_tmpl="azext_aimanager.vendored_sdks.v2026_05_02_preview.operations._operations#ModelSourcesOperations.{}", + operation_group="model_sources", + client_factory=cf_model_sources + ) + + # aimanager modelsource command group + with self.command_group("aimanager modelsource", model_sources_sdk, client_factory=cf_model_sources) as g: + g.custom_command("add", "add_modelsource", supports_no_wait=True) + g.custom_command("update", "update_modelsource", supports_no_wait=True) + g.custom_show_command("show", "show_modelsource") + g.custom_command("list", "list_modelsource") + g.custom_command("delete", "delete_modelsource", supports_no_wait=True, confirmation=True) + g.custom_wait_command("wait", "show_modelsource") + # aimanager namespace modeldeployment command group with self.command_group("aimanager namespace modeldeployment", model_deployments_sdk, client_factory=cf_model_deployments) as g: diff --git a/src/aimanager/azext_aimanager/constants.py b/src/aimanager/azext_aimanager/constants.py index fd9388a4874..224765d0cd8 100644 --- a/src/aimanager/azext_aimanager/constants.py +++ b/src/aimanager/azext_aimanager/constants.py @@ -20,3 +20,6 @@ "VmHourlyPrice:vmHourlyPrice, TotalHourlyPrice:totalHourlyPrice, " "MaxAvailableReplicas:maxAvailableReplicas, Quantization:quantization}" ) + +# Supported model source types for an AI Manager model source. +MODEL_SOURCE_TYPES = ["HuggingFace"] diff --git a/src/aimanager/azext_aimanager/custom.py b/src/aimanager/azext_aimanager/custom.py index 17a3b4de492..4591d1a108c 100644 --- a/src/aimanager/azext_aimanager/custom.py +++ b/src/aimanager/azext_aimanager/custom.py @@ -302,6 +302,108 @@ def aimanager_namespace_get_credentials(cmd, # endregion +# region Model source + +def _construct_modelsource(cmd, source_type, description=None, token=None): + properties_model = _get_model(cmd, "ModelSourceProperties", "model_sources") + source_model = _get_model(cmd, "ModelSource", "model_sources") + + credential = None + if token is not None: + credential_model = _get_model(cmd, "CredentialValue", "model_sources") + inline_model = _get_model(cmd, "InlineCredential", "model_sources") + credential = credential_model(inline=inline_model(value=token)) + + return source_model(properties=properties_model( + source_type=source_type, + description=description, + credential=credential, + )) + + +# pylint: disable=unused-argument +def add_modelsource(cmd, + client, + resource_group_name, + ai_manager_name, + model_source_name, + source_type, + description=None, + token=None, + aks_custom_headers=None, + no_wait=False): + try: + client.get(resource_group_name, ai_manager_name, model_source_name) + except ResourceNotFoundError: + pass + else: + raise ClientRequestError( + f"Model source '{model_source_name}' already exists. " + "Please use 'az aimanager modelsource update' to update it.") + + model_source = _construct_modelsource(cmd, source_type, description, token) + headers = get_aks_custom_headers(aks_custom_headers) + return sdk_no_wait( + no_wait, client.begin_create_or_update, resource_group_name, ai_manager_name, + model_source_name, model_source, headers=headers) + + +# pylint: disable=unused-argument +def update_modelsource(cmd, + client, + resource_group_name, + ai_manager_name, + model_source_name, + description=None, + token=None, + aks_custom_headers=None, + no_wait=False): + try: + existing = client.get(resource_group_name, ai_manager_name, model_source_name) + except ResourceNotFoundError: + raise ClientRequestError( + f"Model source '{model_source_name}' doesn't exist. " + "Please use 'az aimanager modelsource list' to get the current list of model sources.") + + existing_properties = existing.properties + if description is None and existing_properties is not None: + description = existing_properties.description + # sourceType is immutable after creation, so it is always carried over from the existing + # resource on this create-or-replace PUT. + source_type = existing_properties.source_type if existing_properties is not None else None + + model_source = _construct_modelsource(cmd, source_type, description, token) + headers = get_aks_custom_headers(aks_custom_headers) + return sdk_no_wait( + no_wait, client.begin_create_or_update, resource_group_name, ai_manager_name, + model_source_name, model_source, headers=headers) + + +def show_modelsource(cmd, client, resource_group_name, ai_manager_name, + model_source_name): # pylint: disable=unused-argument + return client.get(resource_group_name, ai_manager_name, model_source_name) + + +def list_modelsource(cmd, client, resource_group_name, + ai_manager_name): # pylint: disable=unused-argument + return client.list(resource_group_name, ai_manager_name) + + +def delete_modelsource(cmd, client, resource_group_name, ai_manager_name, model_source_name, + no_wait=False): # pylint: disable=unused-argument + try: + client.get(resource_group_name, ai_manager_name, model_source_name) + except ResourceNotFoundError: + raise ClientRequestError( + f"Model source '{model_source_name}' doesn't exist. " + "Please use 'az aimanager modelsource list' to get the current list of model sources.") + + return sdk_no_wait( + no_wait, client.begin_delete, resource_group_name, ai_manager_name, model_source_name) + +# endregion + + # region Model deployment def _construct_scaling_profile(cmd, replicas=None, min_replicas=None, max_replicas=None, diff --git a/src/aimanager/azext_aimanager/tests/latest/test_modelsource.py b/src/aimanager/azext_aimanager/tests/latest/test_modelsource.py new file mode 100644 index 00000000000..d948f02cee9 --- /dev/null +++ b/src/aimanager/azext_aimanager/tests/latest/test_modelsource.py @@ -0,0 +1,112 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +import unittest +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +from azure.cli.core.azclierror import ClientRequestError, InvalidArgumentValueError +from azure.core.exceptions import ResourceNotFoundError + +from azext_aimanager import custom +from azext_aimanager._validators import validate_model_source_name +from azext_aimanager.vendored_sdks.v2026_05_02_preview import models + + +class MockCmd: + def get_models(self, name, **_): + return getattr(models, name) + + +class TestModelSource(unittest.TestCase): + + def setUp(self): + self.cmd = MockCmd() + self.client = MagicMock() + + def test_construct_without_token_omits_credential(self): + source = custom._construct_modelsource(self.cmd, "HuggingFace", "desc") + + self.assertEqual(source.properties.source_type, "HuggingFace") + self.assertEqual(source.properties.description, "desc") + self.assertIsNone(source.properties.credential) + + def test_construct_with_token_sets_inline_credential(self): + source = custom._construct_modelsource(self.cmd, "HuggingFace", None, "hf_token") + + self.assertEqual(source.properties.credential.inline.value, "hf_token") + + def test_add_rejects_existing_source(self): + self.client.get.return_value = object() + + with self.assertRaises(ClientRequestError): + custom.add_modelsource( + self.cmd, self.client, "rg", "manager", "source", "HuggingFace") + + def test_update_rejects_missing_source(self): + self.client.get.side_effect = ResourceNotFoundError() + + with self.assertRaises(ClientRequestError): + custom.update_modelsource(self.cmd, self.client, "rg", "manager", "source") + + @patch.object(custom, "sdk_no_wait") + @patch.object(custom, "_construct_modelsource") + def test_update_preserves_source_type_and_description( + self, construct_modelsource, sdk_no_wait): + self.client.get.return_value = SimpleNamespace( + properties=SimpleNamespace( + source_type="HuggingFace", + description="existing description", + )) + construct_modelsource.return_value = "model-source" + sdk_no_wait.return_value = "result" + + result = custom.update_modelsource( + self.cmd, self.client, "rg", "manager", "source", token="hf_new") + + self.assertEqual(result, "result") + construct_modelsource.assert_called_once_with( + self.cmd, "HuggingFace", "existing description", "hf_new") + sdk_no_wait.assert_called_once_with( + False, + self.client.begin_create_or_update, + "rg", + "manager", + "source", + "model-source", + headers={}, + ) + + def test_delete_rejects_missing_source(self): + self.client.get.side_effect = ResourceNotFoundError() + + with self.assertRaises(ClientRequestError): + custom.delete_modelsource(self.cmd, self.client, "rg", "manager", "source") + + def test_list_uses_ai_manager_scope(self): + self.client.list.return_value = ["source"] + + result = custom.list_modelsource(self.cmd, self.client, "rg", "manager") + + self.assertEqual(result, ["source"]) + self.client.list.assert_called_once_with("rg", "manager") + + +class TestModelSourceValidators(unittest.TestCase): + + def test_valid_name(self): + validate_model_source_name(SimpleNamespace(model_source_name="hf")) + + def test_missing_name_is_allowed(self): + validate_model_source_name(SimpleNamespace(model_source_name=None)) + validate_model_source_name(SimpleNamespace()) + + def test_blank_name_is_rejected(self): + with self.assertRaises(InvalidArgumentValueError): + validate_model_source_name(SimpleNamespace(model_source_name=" ")) + + +if __name__ == '__main__': + unittest.main() diff --git a/src/aimanager/azext_aimanager/tests/latest/test_modelsource_scenario.py b/src/aimanager/azext_aimanager/tests/latest/test_modelsource_scenario.py new file mode 100644 index 00000000000..8cb2d0c8fdf --- /dev/null +++ b/src/aimanager/azext_aimanager/tests/latest/test_modelsource_scenario.py @@ -0,0 +1,75 @@ +# -------------------------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for license information. +# -------------------------------------------------------------------------------------------- + +from unittest.mock import MagicMock, patch + +from azure.cli.testsdk import ScenarioTest +from azure.core.exceptions import ResourceNotFoundError + +from azext_aimanager.vendored_sdks.v2026_05_02_preview import models + + +class ModelSourceScenarioTest(ScenarioTest): + + def test_modelsource_commands(self): + model_source = models.ModelSource({ + 'name': 'hf', + 'properties': { + 'sourceType': 'HuggingFace', + 'description': 'Hugging Face registry', + 'provisioningState': 'Succeeded', + }, + }) + + operations = MagicMock() + operations.get.side_effect = [ + # add -> not found, then show / update / delete each read the resource once + ResourceNotFoundError(), + model_source, + model_source, + model_source, + ] + operations.list.return_value = [model_source] + service_client = MagicMock() + service_client.model_sources = operations + + command_prefix = 'aimanager modelsource {} -g rg --aimanager-name manager' + + with patch('azext_aimanager._client_factory.get_aimanager_client', + return_value=service_client): + self.cmd( + command_prefix.format('add') + + ' -n hf -s HuggingFace --token hf_xxx --description "Hugging Face registry" --no-wait', + checks=[self.is_empty()]) + + self.cmd( + command_prefix.format('show') + ' -n hf', + checks=[ + self.check('name', 'hf'), + self.check('properties.sourceType', 'HuggingFace'), + self.check('properties.description', 'Hugging Face registry'), + ]) + + self.cmd( + command_prefix.format('list'), + checks=[self.check("length([?name=='hf'])", 1)]) + + self.cmd( + command_prefix.format('update') + ' -n hf --token hf_yyy --no-wait', + checks=[self.is_empty()]) + + self.cmd( + command_prefix.format('delete') + ' -n hf --yes --no-wait', + checks=[self.is_empty()]) + + self.assertEqual(operations.get.call_count, 4) + self.assertEqual(operations.begin_create_or_update.call_count, 2) + operations.list.assert_called_once_with('rg', 'manager') + operations.begin_delete.assert_called_once() + + # the source type is immutable and must be carried over on update + update_payload = operations.begin_create_or_update.call_args_list[1][0][3] + self.assertEqual(update_payload.properties.source_type, 'HuggingFace') + self.assertEqual(update_payload.properties.credential.inline.value, 'hf_yyy')