{aimanager} Show ProvisioningState in create/update -o table; add modelsource and model show table output

src/aimanager/HISTORY.rst

@@ -3,6 +3,18 @@
 Release History
 ===============
 
+2.0.1b2
++++++++
+* ``az aimanager create/update``, ``az aimanager namespace create/update`` and ``az aimanager
+  namespace modeldeployment create/update``: ``-o table`` output now uses the same columns as
+  ``show``/``list``, including ``ProvisioningState`` (and ``ModelId`` for model deployments).
+* ``az aimanager namespace modeldeployment create/update``: Without ``--no-wait``, the returned
+  resource now includes the top-level ``modelId`` field, matching ``show``.
+* ``az aimanager modelsource create/update/show/list``: Add ``-o table`` output with ``Name``,
+  ``ProvisioningState``, ``SourceType`` and ``Description`` columns.
+* ``az aimanager model show``: Add ``-o table`` output with the same ``Name``, ``ModelId`` and
+  ``Description`` columns as ``az aimanager model list``.
+
 2.0.1b1
 +++++++
 * ``az aimanager create``, ``az aimanager namespace create``, ``az aimanager modelsource

src/aimanager/azext_aimanager/_format.py

@@ -86,6 +86,37 @@ def namespace_list_table_format(results):
     return [namespace_table_format(r) for r in results]
 
 
+def modelsource_table_format(result):
+    """Format a single model source resource for display with "-o table"."""
+    properties = result.get('properties') or {}
+    return OrderedDict([
+        ('Name', result.get('name', '')),
+        ('ProvisioningState', properties.get('provisioningState', '')),
+        ('SourceType', properties.get('sourceType', '')),
+        ('Description', properties.get('description') or ''),
+    ])
+
+
+def modelsource_list_table_format(results):
+    """Format a list of model source resources for display with "-o table"."""
+    return [modelsource_table_format(r) for r in results]
+
+
+def aimodel_table_format(result):
+    """Format a single AI model resource for display with "-o table"."""
+    properties = result.get('properties') or {}
+    return OrderedDict([
+        ('Name', result.get('name', '')),
+        ('ModelId', properties.get('modelId', '')),
+        ('Description', properties.get('description') or ''),
+    ])
+
+
+def aimodel_list_table_format(results):
+    """Format a list of AI model resources for display with "-o table"."""
+    return [aimodel_table_format(r) for r in results]
+
+
 def _replica_display(value):
     """Render a replica count, using '-' when the count is not yet reported."""
     return str(value) if value is not None else '-'

src/aimanager/azext_aimanager/commands.py

@@ -5,16 +5,17 @@
 
 # pylint: disable=line-too-long
 from azure.cli.core.commands import CliCommandType
-from azext_aimanager.constants import (
-    AI_MODEL_TABLE_TRANSFORMER,
-)
 from azext_aimanager._format import (
     aimanager_table_format,
     aimanager_list_table_format,
     namespace_table_format,
     namespace_list_table_format,
+    modelsource_table_format,
+    modelsource_list_table_format,
     modeldeployment_table_format,
     modeldeployment_list_table_format,
+    aimodel_table_format,
+    aimodel_list_table_format,
     calculate_cost_table_format,
 )
 from azext_aimanager._client_factory import (
@@ -54,8 +55,10 @@ def load_command_table(self, _):
 
     # aimanager command group
     with self.command_group("aimanager", ai_managers_sdk, client_factory=cf_ai_managers, is_preview=True) as g:
-        g.custom_command("create", "create_aimanager", supports_no_wait=True)
-        g.custom_command("update", "update_aimanager", supports_no_wait=True)
+        g.custom_command("create", "create_aimanager", supports_no_wait=True,
+                         table_transformer=aimanager_table_format)
+        g.custom_command("update", "update_aimanager", supports_no_wait=True,
+                         table_transformer=aimanager_table_format)
         g.custom_show_command("show", "show_aimanager", table_transformer=aimanager_table_format)
         g.custom_command("list", "list_aimanager", table_transformer=aimanager_list_table_format)
         g.custom_command("delete", "delete_aimanager", supports_no_wait=True, confirmation=True)
@@ -64,8 +67,10 @@ def load_command_table(self, _):
 
     # aimanager namespace command group
     with self.command_group("aimanager namespace", ai_manager_namespaces_sdk, client_factory=cf_ai_manager_namespaces, is_preview=True) as g:
-        g.custom_command("create", "create_aimanager_namespace", supports_no_wait=True)
-        g.custom_command("update", "update_aimanager_namespace", supports_no_wait=True)
+        g.custom_command("create", "create_aimanager_namespace", supports_no_wait=True,
+                         table_transformer=namespace_table_format)
+        g.custom_command("update", "update_aimanager_namespace", supports_no_wait=True,
+                         table_transformer=namespace_table_format)
         g.custom_show_command("show", "show_aimanager_namespace", table_transformer=namespace_table_format)
         g.custom_command("list", "list_aimanager_namespace", table_transformer=namespace_list_table_format)
         g.custom_command("delete", "delete_aimanager_namespace", supports_no_wait=True, confirmation=True)
@@ -82,27 +87,30 @@ def load_command_table(self, _):
 
     # aimanager modelsource command group
     with self.command_group("aimanager modelsource", model_sources_sdk, client_factory=cf_model_sources) as g:
-        g.custom_command("create", "create_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("create", "create_modelsource", supports_no_wait=True,
+                         table_transformer=modelsource_table_format)
+        g.custom_command("update", "update_modelsource", supports_no_wait=True,
+                         table_transformer=modelsource_table_format)
+        g.custom_show_command("show", "show_modelsource", table_transformer=modelsource_table_format)
+        g.custom_command("list", "list_modelsource", table_transformer=modelsource_list_table_format)
         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, is_preview=True) as g:
-        g.custom_command("create", "create_modeldeployment", supports_no_wait=True)
-        g.custom_command("update", "update_modeldeployment", supports_no_wait=True)
+        g.custom_command("create", "create_modeldeployment", supports_no_wait=True,
+                         table_transformer=modeldeployment_table_format)
+        g.custom_command("update", "update_modeldeployment", supports_no_wait=True,
+                         table_transformer=modeldeployment_table_format)
         g.custom_show_command("show", "show_modeldeployment", table_transformer=modeldeployment_table_format)
         g.custom_command("list", "list_modeldeployment", table_transformer=modeldeployment_list_table_format)
         g.custom_command("delete", "delete_modeldeployment", supports_no_wait=True, confirmation=True)
         g.custom_wait_command("wait", "show_modeldeployment")
 
     # aimanager model command group
     with self.command_group("aimanager model", ai_models_sdk, client_factory=cf_ai_models) as g:
-        g.custom_show_command("show", "show_aimodel")
-        g.custom_command("list", "list_aimodel",
-                         table_transformer=AI_MODEL_TABLE_TRANSFORMER)
+        g.custom_show_command("show", "show_aimodel", table_transformer=aimodel_table_format)
+        g.custom_command("list", "list_aimodel", table_transformer=aimodel_list_table_format)
         g.custom_command("calculate-cost", "calculate_aimodel_cost",
                          table_transformer=calculate_cost_table_format)

src/aimanager/azext_aimanager/constants.py

@@ -20,10 +20,5 @@
 
 MODEL_DEPLOYMENT_PERFORMANCE_MODES = ["Balanced", "Latency", "Throughput"]
 
-# Table output projections for the 'az aimanager model' commands.
-AI_MODEL_TABLE_TRANSFORMER = (
-    "[].{Name:name, ModelId:properties.modelId, Description:properties.description}"
-)
-
 # Supported model source types for an AI Manager model source.
 MODEL_SOURCE_TYPES = ["HuggingFace"]

src/aimanager/azext_aimanager/custom.py

@@ -523,9 +523,10 @@ def create_modeldeployment(cmd, client, resource_group_name, ai_manager_name, na
         cmd, model_resource_id, vm_size, model_source_resource_id, performance_mode,
         scale, parse_key_value_list(overrides) if overrides is not None else None)
     headers = get_custom_headers(custom_headers)
-    return sdk_no_wait(
+    poller = sdk_no_wait(
         no_wait, client.begin_create_or_update, resource_group_name, ai_manager_name,
         namespace_name, model_deployment_name, deployment, headers=headers)
+    return _annotate_model_id_on_completion(cmd, poller, no_wait)
 
 
 # pylint: disable=unused-argument
@@ -558,10 +559,21 @@ def update_modeldeployment(cmd, client, resource_group_name, ai_manager_name, na
     headers = get_custom_headers(custom_headers)
     etag = existing.e_tag
     match_condition = MatchConditions.IfNotModified if etag is not None else None
-    return sdk_no_wait(
+    poller = sdk_no_wait(
         no_wait, client.begin_create_or_update, resource_group_name, ai_manager_name,
         namespace_name, model_deployment_name, deployment, headers=headers,
         etag=etag, match_condition=match_condition)
+    return _annotate_model_id_on_completion(cmd, poller, no_wait)
+
+
+def _annotate_model_id_on_completion(cmd, poller, no_wait):
+    """Wait for a create/update to finish and annotate the result with ``modelId``, so the
+    ``-o table`` ModelId column is populated as it is for ``show``. With --no-wait the poller
+    is returned unchanged."""
+    if no_wait:
+        return poller
+    result = LongRunningOperation(cmd.cli_ctx)(poller)  # blocks until Succeeded; raises on failure
+    return _annotate_model_ids(cmd, [result])[0]
 
 
 def show_modeldeployment(cmd, client, resource_group_name, ai_manager_name, namespace_name,

src/aimanager/azext_aimanager/tests/latest/test_aimanager_format.py

@@ -10,8 +10,12 @@
     aimanager_list_table_format,
     namespace_table_format,
     namespace_list_table_format,
+    modelsource_table_format,
+    modelsource_list_table_format,
     modeldeployment_table_format,
     modeldeployment_list_table_format,
+    aimodel_table_format,
+    aimodel_list_table_format,
     calculate_cost_table_format,
 )
 
@@ -117,6 +121,59 @@ def test_list_table_format_namespace(self):
         self.assertEqual(results[0]["Name"], "ns1")
 
 
+class TestModelSourceTableFormat(unittest.TestCase):
+    """Test cases for model source table output formatting."""
+
+    def _sample(self):
+        return {
+            "id": (
+                "/subscriptions/26fe00f8-0000-0000-0000-bb1d2e00343a"
+                "/resourceGroups/yiralirg"
+                "/providers/Microsoft.ContainerService/aiManagers/aimbyo"
+                "/modelSources/hf1"
+            ),
+            "name": "hf1",
+            "eTag": "ff459bfb-b983-436f-b55b-afe701d1c896",
+            "properties": {
+                "provisioningState": "Succeeded",
+                "sourceType": "HuggingFace",
+                "description": "gated models",
+            },
+        }
+
+    def test_table_format_columns_modelsource(self):
+        result = modelsource_table_format(self._sample())
+        self.assertEqual(
+            list(result.keys()),
+            ["Name", "ProvisioningState", "SourceType", "Description"],
+        )
+
+    def test_table_format_values_modelsource(self):
+        result = modelsource_table_format(self._sample())
+        self.assertEqual(result["Name"], "hf1")
+        self.assertEqual(result["ProvisioningState"], "Succeeded")
+        self.assertEqual(result["SourceType"], "HuggingFace")
+        self.assertEqual(result["Description"], "gated models")
+
+    def test_table_format_null_properties_modelsource(self):
+        result = modelsource_table_format({"name": "hf1", "properties": None})
+        self.assertEqual(result["Name"], "hf1")
+        self.assertEqual(result["ProvisioningState"], "")
+        self.assertEqual(result["SourceType"], "")
+        self.assertEqual(result["Description"], "")
+
+    def test_table_format_null_description_modelsource(self):
+        # A null description must still render as a (blank) column: knack drops None cells.
+        sample = self._sample()
+        sample["properties"]["description"] = None
+        self.assertEqual(modelsource_table_format(sample)["Description"], "")
+
+    def test_list_table_format_modelsource(self):
+        results = modelsource_list_table_format([self._sample(), self._sample()])
+        self.assertEqual(len(results), 2)
+        self.assertEqual(results[0]["Name"], "hf1")
+
+
 class TestModelDeploymentTableFormat(unittest.TestCase):
     """Test cases for model deployment table output formatting."""
 
@@ -185,6 +242,117 @@ def test_list_table_format_modeldeployment(self):
         self.assertEqual(results[0]["ModelId"], "meta-llama/Llama-3-8B")
 
 
+class TestAIModelTableFormat(unittest.TestCase):
+    """Test cases for AI model table output formatting."""
+
+    def _sample(self):
+        return {
+            "id": (
+                "/subscriptions/26fe00f8-0000-0000-0000-bb1d2e00343a"
+                "/providers/Microsoft.ContainerService/locations/westus2"
+                "/aiModels/9806f0c862fdd920"
+            ),
+            "name": "9806f0c862fdd920",
+            "properties": {
+                "modelId": "microsoft/Phi-4-mini-instruct",
+                "description": "Phi-4 mini",
+                "spec": {"parameterCount": "3.8B"},
+            },
+        }
+
+    def test_table_format_columns_aimodel(self):
+        result = aimodel_table_format(self._sample())
+        self.assertEqual(list(result.keys()), ["Name", "ModelId", "Description"])
+
+    def test_table_format_values_aimodel(self):
+        result = aimodel_table_format(self._sample())
+        self.assertEqual(result["Name"], "9806f0c862fdd920")
+        self.assertEqual(result["ModelId"], "microsoft/Phi-4-mini-instruct")
+        self.assertEqual(result["Description"], "Phi-4 mini")
+
+    def test_table_format_null_properties_aimodel(self):
+        result = aimodel_table_format({"name": "m1", "properties": None})
+        self.assertEqual(result["Name"], "m1")
+        self.assertEqual(result["ModelId"], "")
+        self.assertEqual(result["Description"], "")
+
+    def test_table_format_null_description_aimodel(self):
+        # A null description must still render as a (blank) column: knack drops None cells.
+        sample = self._sample()
+        sample["properties"]["description"] = None
+        self.assertEqual(aimodel_table_format(sample)["Description"], "")
+
+    def test_list_table_format_aimodel(self):
+        results = aimodel_list_table_format([self._sample(), self._sample()])
+        self.assertEqual(len(results), 2)
+        self.assertEqual(results[0]["ModelId"], "microsoft/Phi-4-mini-instruct")
+
+
+class TestCommandTableTransformers(unittest.TestCase):
+    """Every create/update/show/list command must render "-o table" with a formatter, since
+    azure-cli's default table drops nested fields such as properties.provisioningState."""
+
+    def _load(self):
+        from azext_aimanager.commands import load_command_table
+
+        transformers = {}
+
+        class Group:
+            def __init__(self, name):
+                self.name = name
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *_):
+                return False
+
+            def _record(self, verb, *_, **kwargs):
+                transformers["{} {}".format(self.name, verb)] = kwargs.get("table_transformer")
+
+            custom_command = custom_show_command = _record
+
+            def wait_command(self, *_, **__):
+                pass
+
+            custom_wait_command = wait_command
+
+        class Loader:
+            def command_group(self, name, *_, **__):
+                return Group(name)
+
+        load_command_table(Loader(), None)
+        return transformers
+
+    def test_create_update_show_list_have_table_transformers(self):
+        transformers = self._load()
+        expected = {
+            "aimanager": (aimanager_table_format, aimanager_list_table_format),
+            "aimanager namespace": (namespace_table_format, namespace_list_table_format),
+            "aimanager modelsource": (modelsource_table_format, modelsource_list_table_format),
+            "aimanager namespace modeldeployment": (
+                modeldeployment_table_format, modeldeployment_list_table_format),
+        }
+        for group, (single, many) in expected.items():
+            for verb in ("create", "update", "show"):
+                self.assertIs(transformers["{} {}".format(group, verb)], single, (group, verb))
+            self.assertIs(transformers["{} list".format(group)], many, group)
+
+    def test_aimodel_commands_have_table_transformers(self):
+        transformers = self._load()
+        self.assertIs(transformers["aimanager model show"], aimodel_table_format)
+        self.assertIs(transformers["aimanager model list"], aimodel_list_table_format)
+        self.assertIs(transformers["aimanager model calculate-cost"], calculate_cost_table_format)
+
+    def test_every_resource_command_has_callable_table_transformer(self):
+        # Guards against a newly added create/update/show/list command shipping without a
+        # formatter (or with a JMESPath string instead of a _format.py callable).
+        transformers = self._load()
+        for command, transformer in transformers.items():
+            if command.rsplit(" ", 1)[1] in ("create", "update", "show", "list"):
+                self.assertTrue(callable(transformer), command)
+
+
 class TestCalculateCostTableFormat(unittest.TestCase):
     """Test cases for 'az aimanager model calculate-cost' table output formatting."""