Add programmatic retrieval of image logs (#2009)

* Add programmatic retrieval of image logs

* Set bool on proto that we want all logs

modal/image.py

@@ -9,7 +9,7 @@
 from dataclasses import dataclass
 from inspect import isfunction
 from pathlib import Path, PurePosixPath
-from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple, Union, get_args
+from typing import Any, AsyncGenerator, Callable, Dict, List, Literal, Optional, Sequence, Set, Tuple, Union, get_args
 
 from google.protobuf.message import Message
 from grpclib.exceptions import GRPCError, StreamTerminatedError
@@ -28,7 +28,7 @@
 from .gpu import GPU_T, parse_gpu_config
 from .mount import _Mount, python_standalone_mount_name
 from .network_file_system import _NetworkFileSystem
-from .object import _Object
+from .object import _Object, live_method_gen
 from .secret import _Secret
 from .volume import _Volume
 
@@ -1675,5 +1675,21 @@ def run_inside(self):
         """
         deprecation_error((2023, 12, 15), Image.run_inside.__doc__)
 
+    @live_method_gen
+    async def logs(self) -> AsyncGenerator[str, None]:
+        last_entry_id: Optional[str] = None
+
+        request = api_pb2.ImageJoinStreamingRequest(
+            image_id=self._object_id, timeout=55, last_entry_id=last_entry_id, include_logs_for_finished=True
+        )
+        async for response in unary_stream(self._client.stub.ImageJoinStreaming, request):
+            if response.result.status:
+                return
+            if response.entry_id:
+                last_entry_id = response.entry_id
+            for task_log in response.task_logs:
+                if task_log.data:
+                    yield task_log.data
+
 
 Image = synchronize_api(_Image)

modal_proto/api.proto

@@ -1357,6 +1357,7 @@ message ImageJoinStreamingRequest {
   string image_id = 1;
   float timeout = 2;
   string last_entry_id = 3;
+  bool include_logs_for_finished = 4;
 }
 
 message ImageJoinStreamingResponse {

test/conftest.py

@@ -853,13 +853,14 @@ async def ImageJoinStreaming(self, stream):
         if self.image_join_sleep_duration is not None:
             await asyncio.sleep(self.image_join_sleep_duration)
 
-        task_log_1 = api_pb2.TaskLogs(data="hello, world\n", file_descriptor=api_pb2.FILE_DESCRIPTOR_INFO)
+        task_log_1 = api_pb2.TaskLogs(data="build starting\n", file_descriptor=api_pb2.FILE_DESCRIPTOR_INFO)
         task_log_2 = api_pb2.TaskLogs(
             task_progress=api_pb2.TaskProgress(
                 len=1, pos=0, progress_type=api_pb2.IMAGE_SNAPSHOT_UPLOAD, description="xyz"
             )
         )
-        await stream.send_message(api_pb2.ImageJoinStreamingResponse(task_logs=[task_log_1, task_log_2]))
+        task_log_3 = api_pb2.TaskLogs(data="build finished\n", file_descriptor=api_pb2.FILE_DESCRIPTOR_INFO)
+        await stream.send_message(api_pb2.ImageJoinStreamingResponse(task_logs=[task_log_1, task_log_2, task_log_3]))
         await stream.send_message(
             api_pb2.ImageJoinStreamingResponse(
                 result=api_pb2.GenericResult(status=api_pb2.GenericResult.GENERIC_STATUS_SUCCESS)

test/image_test.py

@@ -1120,3 +1120,15 @@ async def MockImageJoinStreaming(self, stream):
         ctx.set_responder("ImageJoinStreaming", MockImageJoinStreaming)
         with parallel_app.run(client=client):
             pass
+
+
+@pytest.mark.asyncio
+async def test_logs(servicer, client):
+    app = App()
+    image = Image.debian_slim().pip_install("foobarbaz")
+    app.function(image=image)(dummy)
+    async with app.run.aio(client=client):
+        pass
+
+    logs = [data async for data in image.logs.aio()]
+    assert logs == ["build starting\n", "build finished\n"]