feat: Update the image alias cmd/func API to follow manager API updat…

…es (#211)

changes/211.feature.md

@@ -0,0 +1 @@
+Allow specifying the architecture when setting image aliases

src/ai/backend/client/cli/admin/image.py

@@ -89,11 +89,12 @@ async def rescan_images_impl(registry: str) -> None:
 @image.command()
 @click.argument('alias', type=str)
 @click.argument('target', type=str)
-def alias(alias, target):
+@click.option('--arch', type=str, default=None, help='Set an explicit architecture.')
+def alias(alias, target, arch):
     """Add an image alias."""
     with Session() as session:
         try:
-            result = session.Image.alias_image(alias, target)
+            result = session.Image.alias_image(alias, target, arch)
         except Exception as e:
             print_error(e)
             sys.exit(1)

src/ai/backend/client/func/image.py

@@ -1,4 +1,4 @@
-from typing import Sequence
+from typing import Optional, Sequence
 
 from ai.backend.client.output.fields import image_fields
 from ai.backend.client.output.types import FieldSpec
@@ -66,7 +66,12 @@ async def rescan_images(cls, registry: str):
 
     @api_function
     @classmethod
-    async def alias_image(cls, alias: str, target: str) -> dict:
+    async def alias_image(
+        cls,
+        alias: str,
+        target: str,
+        arch: Optional[str] = None,
+    ) -> dict:
         q = 'mutation($alias: String!, $target: String!) {' \
             '  alias_image(alias: $alias, target: $target) {' \
             '   ok msg' \
@@ -76,6 +81,8 @@ async def alias_image(cls, alias: str, target: str) -> dict:
             'alias': alias,
             'target': target,
         }
+        if arch:
+            variables = {'architecture': arch, **variables}
         data = await api_session.get().Admin._query(q, variables)
         return data['alias_image']