Add modal shell im-* to shell into an image id

modal/cli/run.py

@@ -613,6 +613,12 @@ def shell(
     ```
     modal shell sb-abc123xyz
     ```
+
+    Launch a shell to explore an Image by ID:
+
+    ```
+    modal shell im-abc123xyz
+    ```
     """
     env = ensure_env(env)
 
@@ -624,7 +630,29 @@ def shell(
 
     app = App("modal shell")
 
-    if ref is not None:
+    if ref is None or ref.startswith("im-"):
+        if ref is None:
+            modal_image = Image.from_registry(image, add_python=add_python) if image else None
+        else:
+            if image is not None:
+                raise ClickException("image is not supported when running `modal shell im-*`")
+            modal_image = Image.from_id(ref)
+
+        volumes = {} if volume is None else {f"/mnt/{vol}": Volume.from_name(vol) for vol in volume}
+        secrets = [] if secret is None else [Secret.from_name(s) for s in secret]
+        start_shell = partial(
+            interactive_shell,
+            image=modal_image,
+            cpu=cpu,
+            memory=memory,
+            gpu=gpu,
+            cloud=cloud,
+            volumes=volumes,
+            secrets=secrets,
+            region=region.split(",") if region else [],
+            pty=pty,
+        )
+    else:
         # `modal shell` with a sandbox ID gets the task_id, that's then handled by the `ta-*` flow below.
         if ref.startswith("sb-") and len(ref[3:]) > 0 and ref[3:].isalnum():
             from ..sandbox import Sandbox
@@ -688,22 +716,6 @@ def shell(
             pty=pty,
             proxy=function_spec.proxy,
         )
-    else:
-        modal_image = Image.from_registry(image, add_python=add_python) if image else None
-        volumes = {} if volume is None else {f"/mnt/{vol}": Volume.from_name(vol) for vol in volume}
-        secrets = [] if secret is None else [Secret.from_name(s) for s in secret]
-        start_shell = partial(
-            interactive_shell,
-            image=modal_image,
-            cpu=cpu,
-            memory=memory,
-            gpu=gpu,
-            cloud=cloud,
-            volumes=volumes,
-            secrets=secrets,
-            region=region.split(",") if region else [],
-            pty=pty,
-        )
 
     # NB: invoking under bash makes --cmd a lot more flexible.
     cmds = shlex.split(f'/bin/bash -c "{cmd}"')

test/cli_test.py

@@ -631,6 +631,23 @@ def test_shell_preserve_token(servicer, set_env_client, mock_shell_pty, monkeypa
     assert captured_out == [(1, shell_prompt), (1, expected_output)]
 
 
+@skip_windows("modal shell is not supported on Windows.")
+def test_shell_modal_image(servicer, set_env_client, mock_shell_pty):
+    image_id = "im-123abcxyz"
+    _run(["shell", image_id])
+    sandbox_def = servicer.sandbox_defs[0]
+    assert sandbox_def.image_id == image_id
+
+
+@skip_windows("modal shell is not supported on Windows.")
+def test_shell_modal_image_error(servicer):
+    _run(
+        ["shell", "im-hello", "--image", "some-image"],
+        expected_exit_code=1,
+        expected_stderr="image is not supported",
+    )
+
+
 def test_shell_unsuported_cmds_fails_on_windows(servicer, set_env_client, mock_shell_pty):
     expected_exit_code = 1 if platform.system() == "Windows" else 0
     res = _run(["shell"], expected_exit_code=expected_exit_code)

test/conftest.py

@@ -1614,6 +1614,11 @@ async def ImageGetOrCreate(self, stream):
             )
         )
 
+    async def ImageFromId(self, stream):
+        request: api_pb2.ImageFromIdRequest = await stream.recv_message()
+        image_id = request.image_id
+        await stream.send_message(api_pb2.ImageFromIdResponse(image_id=image_id))
+
     async def ImageJoinStreaming(self, stream):
         req = await stream.recv_message()