diff --git a/modal/io_streams.py b/modal/io_streams.py index 867707599d..46edac36b8 100644 --- a/modal/io_streams.py +++ b/modal/io_streams.py @@ -1,6 +1,7 @@ # Copyright Modal Labs 2022 import asyncio import time +import uuid from collections.abc import AsyncGenerator, AsyncIterator from typing import ( TYPE_CHECKING, @@ -370,6 +371,7 @@ def __init__(self, object_id: str, object_type: Literal["sandbox", "container_pr self._client = client self._is_closed = False self._buffer = bytearray() + self.writer_id = str(uuid.uuid4()) def _get_next_index(self) -> int: index = self._index @@ -449,7 +451,11 @@ async def drain(self) -> None: await retry_transient_errors( self._client.stub.SandboxStdinWrite, api_pb2.SandboxStdinWriteRequest( - sandbox_id=self._object_id, index=index, eof=self._is_closed, input=data + sandbox_id=self._object_id, + index=index, + eof=self._is_closed, + input=data, + writer_id=self.writer_id, ), ) else: diff --git a/modal_proto/api.proto b/modal_proto/api.proto index 5c571b886c..3dad58ff84 100644 --- a/modal_proto/api.proto +++ b/modal_proto/api.proto @@ -2606,6 +2606,7 @@ message RuntimeInputMessage { bytes message = 1; uint64 message_index = 2; bool eof = 3; + string writer_id = 4; } message RuntimeOutputBatch { @@ -2876,6 +2877,7 @@ message SandboxStdinWriteRequest { bytes input = 2; uint32 index = 3; bool eof = 4; + string writer_id = 5; } message SandboxStdinWriteResponse {