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

src/aimanager/HISTORY.rst

@@ -3,6 +3,16 @@
 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.
+
 2.0.1b1
 +++++++
 * ``az aimanager create``, ``az aimanager namespace create``, ``az aimanager modelsource

src/aimanager/azext_aimanager/_format.py

@@ -86,6 +86,22 @@ 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 _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

@@ -13,6 +13,8 @@
     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,
     calculate_cost_table_format,
@@ -54,8 +56,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 +68,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,18 +88,22 @@ 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)

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,6 +10,8 @@
     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,
     calculate_cost_table_format,
@@ -117,6 +119,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 +240,57 @@ def test_list_table_format_modeldeployment(self):
         self.assertEqual(results[0]["ModelId"], "meta-llama/Llama-3-8B")
 
 
+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)
+
+
 class TestCalculateCostTableFormat(unittest.TestCase):
     """Test cases for 'az aimanager model calculate-cost' table output formatting."""
 

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

@@ -219,10 +219,11 @@ def test_list_without_namespace_all_forbidden_raises(
         self.assertNotIn("--namespace", recs)
         self.assertEqual(2, client.list_by_ai_manager_namespace.call_count)
 
+    @patch.object(custom, "_annotate_model_id_on_completion", side_effect=lambda _cmd, p, _nw: p)
     @patch.object(custom, "sdk_no_wait")
     @patch.object(custom, "_construct_modeldeployment")
     def test_update_preserves_omitted_properties_and_uses_etag(
-            self, construct_modeldeployment, sdk_no_wait):
+            self, construct_modeldeployment, sdk_no_wait, _annotate):
         existing_scale = models.ScalingProfile(
             manual=models.ManualScalingProfile(replicas=2))
         existing_overrides = models.ModelDeploymentOverrides(
@@ -248,6 +249,7 @@ def test_update_preserves_omitted_properties_and_uses_etag(
             self.cmd, client, "rg", "manager", "namespace", "deployment")
 
         self.assertEqual(result, "result")
+        _annotate.assert_called_once_with(self.cmd, "result", False)
         construct_modeldeployment.assert_called_once_with(
             self.cmd,
             "/models/model-a",
@@ -270,6 +272,40 @@ def test_update_preserves_omitted_properties_and_uses_etag(
             match_condition=MatchConditions.IfNotModified,
         )
 
+    @patch.object(custom, "_annotate_model_id_on_completion", return_value="annotated")
+    @patch.object(custom, "sdk_no_wait", return_value="poller")
+    @patch.object(custom, "_construct_modeldeployment")
+    def test_create_routes_result_through_model_id_annotation(
+            self, _construct, _sdk_no_wait, annotate):
+        result = custom.create_modeldeployment(
+            self.cmd, MagicMock(), "rg", "manager", "namespace", "deployment",
+            "/models/model-a", "Standard_NC24ads_A100_v4", replicas=1)
+
+        self.assertEqual(result, "annotated")
+        annotate.assert_called_once_with(self.cmd, "poller", False)
+
+    @patch.object(custom, "_annotate_model_ids")
+    @patch.object(custom, "LongRunningOperation")
+    def test_annotate_model_id_on_completion_waits_and_annotates(self, lro, annotate):
+        lro.return_value.return_value = "deployment"
+        annotate.return_value = [{"name": "md1", "modelId": "meta-llama/Llama-3-8B"}]
+        cmd = SimpleNamespace(cli_ctx=object())
+
+        result = custom._annotate_model_id_on_completion(cmd, "poller", False)
+
+        lro.return_value.assert_called_once_with("poller")
+        annotate.assert_called_once_with(cmd, ["deployment"])
+        self.assertEqual(result["modelId"], "meta-llama/Llama-3-8B")
+
+    @patch.object(custom, "_annotate_model_ids")
+    @patch.object(custom, "LongRunningOperation")
+    def test_annotate_model_id_on_completion_no_wait_returns_poller(self, lro, annotate):
+        result = custom._annotate_model_id_on_completion(SimpleNamespace(), "poller", True)
+
+        self.assertEqual(result, "poller")
+        lro.assert_not_called()
+        annotate.assert_not_called()
+
     @patch("azext_aimanager._client_factory.cf_ai_models")
     def test_annotate_model_ids_returns_plain_dicts_with_model_id(self, cf_ai_models):
         # The AIModel GET resolves the human-readable modelId.

src/aimanager/setup.py

@@ -14,7 +14,7 @@
     from distutils import log as logger
     logger.warn("Wheel is not available, disabling bdist_wheel hook")
 
-VERSION = '2.0.1b1'
+VERSION = '2.0.1b2'
 
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers