modal 0.62.115__py3-none-any.whl → 0.72.13__py3-none-any.whl

modal/cls.pyi

@@ -1,5 +1,6 @@
+import collections.abc
 import google.protobuf.message
-import modal._output
+import inspect
 import modal.app
 import modal.client
 import modal.functions
@@ -16,167 +17,210 @@ import typing_extensions
 
 T = typing.TypeVar("T")
 
-class _Obj:
-    _functions: typing.Dict[str, modal.functions._Function]
-    _inited: bool
-    _entered: bool
-    _local_obj: typing.Any
-    _local_obj_constr: typing.Union[typing.Callable[[], typing.Any], None]
-
-    def __init__(self, user_cls: type, output_mgr: typing.Union[modal._output.OutputManager, None], base_functions: typing.Dict[str, modal.functions._Function], from_other_workspace: bool, options: typing.Union[modal_proto.api_pb2.FunctionOptions, None], args, kwargs):
-        ...
-
-    def get_obj(self):
-        ...
-
-    def get_local_obj(self):
-        ...
-
-    def enter(self):
-        ...
+def _use_annotation_parameters(user_cls: type) -> bool: ...
+def _get_class_constructor_signature(user_cls: type) -> inspect.Signature: ...
+def _bind_instance_method(
+    service_function: modal.functions._Function, class_bound_method: modal.functions._Function
+): ...
 
+class _Obj:
+    _cls: _Cls
+    _functions: dict[str, modal.functions._Function]
+    _has_entered: bool
+    _user_cls_instance: typing.Optional[typing.Any]
+    _args: tuple[typing.Any, ...]
+    _kwargs: dict[str, typing.Any]
+    _instance_service_function: typing.Optional[modal.functions._Function]
+
+    def _uses_common_service_function(self): ...
+    def __init__(
+        self,
+        cls: _Cls,
+        user_cls: typing.Optional[type],
+        options: typing.Optional[modal_proto.api_pb2.FunctionOptions],
+        args,
+        kwargs,
+    ): ...
+    def _cached_service_function(self) -> modal.functions._Function: ...
+    def _get_parameter_values(self) -> dict[str, typing.Any]: ...
+    def _new_user_cls_instance(self): ...
+    async def keep_warm(self, warm_pool_size: int) -> None: ...
+    def _cached_user_cls_instance(self): ...
+    def _enter(self): ...
     @property
-    def entered(self):
-        ...
-
-    @entered.setter
-    def entered(self, val):
-        ...
-
-    async def aenter(self):
-        ...
-
-    def __getattr__(self, k):
-        ...
-
+    def _entered(self) -> bool: ...
+    @_entered.setter
+    def _entered(self, val: bool): ...
+    async def _aenter(self): ...
+    def __getattr__(self, k): ...
 
 class Obj:
-    _functions: typing.Dict[str, modal.functions.Function]
-    _inited: bool
-    _entered: bool
-    _local_obj: typing.Any
-    _local_obj_constr: typing.Union[typing.Callable[[], typing.Any], None]
-
-    def __init__(self, user_cls: type, output_mgr: typing.Union[modal._output.OutputManager, None], base_functions: typing.Dict[str, modal.functions.Function], from_other_workspace: bool, options: typing.Union[modal_proto.api_pb2.FunctionOptions, None], args, kwargs):
-        ...
-
-    def get_obj(self):
-        ...
-
-    def get_local_obj(self):
-        ...
-
-    def enter(self):
-        ...
-
+    _cls: Cls
+    _functions: dict[str, modal.functions.Function]
+    _has_entered: bool
+    _user_cls_instance: typing.Optional[typing.Any]
+    _args: tuple[typing.Any, ...]
+    _kwargs: dict[str, typing.Any]
+    _instance_service_function: typing.Optional[modal.functions.Function]
+
+    def __init__(
+        self,
+        cls: Cls,
+        user_cls: typing.Optional[type],
+        options: typing.Optional[modal_proto.api_pb2.FunctionOptions],
+        args,
+        kwargs,
+    ): ...
+    def _uses_common_service_function(self): ...
+    def _cached_service_function(self) -> modal.functions.Function: ...
+    def _get_parameter_values(self) -> dict[str, typing.Any]: ...
+    def _new_user_cls_instance(self): ...
+
+    class __keep_warm_spec(typing_extensions.Protocol):
+        def __call__(self, warm_pool_size: int) -> None: ...
+        async def aio(self, warm_pool_size: int) -> None: ...
+
+    keep_warm: __keep_warm_spec
+
+    def _cached_user_cls_instance(self): ...
+    def _enter(self): ...
     @property
-    def entered(self):
-        ...
-
-    @entered.setter
-    def entered(self, val):
-        ...
-
-    async def aenter(self):
-        ...
-
-    def __getattr__(self, k):
-        ...
-
+    def _entered(self) -> bool: ...
+    @_entered.setter
+    def _entered(self, val: bool): ...
+    async def _aenter(self): ...
+    def __getattr__(self, k): ...
 
 class _Cls(modal.object._Object):
-    _user_cls: typing.Union[type, None]
-    _functions: typing.Dict[str, modal.functions._Function]
-    _options: typing.Union[modal_proto.api_pb2.FunctionOptions, None]
-    _callables: typing.Dict[str, typing.Callable]
-    _from_other_workspace: typing.Union[bool, None]
-    _app: typing.Union[modal.app._App, None]
-
-    def _initialize_from_empty(self):
-        ...
-
-    def _initialize_from_other(self, other: _Cls):
-        ...
-
-    def _set_output_mgr(self, output_mgr: modal._output.OutputManager):
-        ...
-
-    def _hydrate_metadata(self, metadata: google.protobuf.message.Message):
-        ...
-
-    def _get_metadata(self) -> modal_proto.api_pb2.ClassHandleMetadata:
-        ...
-
+    _user_cls: typing.Optional[type]
+    _class_service_function: typing.Optional[modal.functions._Function]
+    _method_functions: typing.Optional[dict[str, modal.functions._Function]]
+    _options: typing.Optional[modal_proto.api_pb2.FunctionOptions]
+    _callables: dict[str, typing.Callable[..., typing.Any]]
+    _app: typing.Optional[modal.app._App]
+    _name: typing.Optional[str]
+
+    def _initialize_from_empty(self): ...
+    def _initialize_from_other(self, other: _Cls): ...
+    def _get_partial_functions(self) -> dict[str, modal.partial_function._PartialFunction]: ...
+    def _hydrate_metadata(self, metadata: google.protobuf.message.Message): ...
     @staticmethod
-    def from_local(user_cls, app, decorator: typing.Callable[[modal.partial_function.PartialFunction, type], modal.functions._Function]) -> _Cls:
-        ...
-
+    def validate_construction_mechanism(user_cls): ...
+    @staticmethod
+    def from_local(user_cls, app: modal.app._App, class_service_function: modal.functions._Function) -> _Cls: ...
+    def _uses_common_service_function(self): ...
     @classmethod
-    def from_name(cls: typing.Type[_Cls], app_name: str, tag: typing.Union[str, None] = None, namespace=1, environment_name: typing.Union[str, None] = None, workspace: typing.Union[str, None] = None) -> _Cls:
-        ...
-
-    def with_options(self: _Cls, cpu: typing.Union[float, None] = None, memory: typing.Union[int, typing.Tuple[int, int], None] = None, gpu: typing.Union[None, bool, str, modal.gpu._GPUConfig] = None, secrets: typing.Collection[modal.secret._Secret] = (), volumes: typing.Dict[typing.Union[str, os.PathLike], modal.volume._Volume] = {}, retries: typing.Union[int, modal.retries.Retries, None] = None, timeout: typing.Union[int, None] = None, concurrency_limit: typing.Union[int, None] = None, allow_concurrent_inputs: typing.Union[int, None] = None, container_idle_timeout: typing.Union[int, None] = None, allow_background_volume_commits: bool = False) -> _Cls:
-        ...
-
+    def from_name(
+        cls: type[_Cls],
+        app_name: str,
+        name: str,
+        namespace=1,
+        environment_name: typing.Optional[str] = None,
+        workspace: typing.Optional[str] = None,
+    ) -> _Cls: ...
+    def with_options(
+        self: _Cls,
+        cpu: typing.Union[float, tuple[float, float], None] = None,
+        memory: typing.Union[int, tuple[int, int], None] = None,
+        gpu: typing.Union[None, bool, str, modal.gpu._GPUConfig] = None,
+        secrets: collections.abc.Collection[modal.secret._Secret] = (),
+        volumes: dict[typing.Union[str, os.PathLike], modal.volume._Volume] = {},
+        retries: typing.Union[int, modal.retries.Retries, None] = None,
+        timeout: typing.Optional[int] = None,
+        concurrency_limit: typing.Optional[int] = None,
+        allow_concurrent_inputs: typing.Optional[int] = None,
+        container_idle_timeout: typing.Optional[int] = None,
+    ) -> _Cls: ...
     @staticmethod
-    async def lookup(app_name: str, tag: typing.Union[str, None] = None, namespace=1, client: typing.Union[modal.client._Client, None] = None, environment_name: typing.Union[str, None] = None, workspace: typing.Union[str, None] = None) -> _Cls:
-        ...
-
-    def __call__(self, *args, **kwargs) -> _Obj:
-        ...
-
-    def __getattr__(self, k):
-        ...
-
+    async def lookup(
+        app_name: str,
+        name: str,
+        namespace=1,
+        client: typing.Optional[modal.client._Client] = None,
+        environment_name: typing.Optional[str] = None,
+        workspace: typing.Optional[str] = None,
+    ) -> _Cls: ...
+    def __call__(self, *args, **kwargs) -> _Obj: ...
+    def __getattr__(self, k): ...
 
 class Cls(modal.object.Object):
-    _user_cls: typing.Union[type, None]
-    _functions: typing.Dict[str, modal.functions.Function]
-    _options: typing.Union[modal_proto.api_pb2.FunctionOptions, None]
-    _callables: typing.Dict[str, typing.Callable]
-    _from_other_workspace: typing.Union[bool, None]
-    _app: typing.Union[modal.app.App, None]
-
-    def __init__(self, *args, **kwargs):
-        ...
-
-    def _initialize_from_empty(self):
-        ...
-
-    def _initialize_from_other(self, other: Cls):
-        ...
-
-    def _set_output_mgr(self, output_mgr: modal._output.OutputManager):
-        ...
-
-    def _hydrate_metadata(self, metadata: google.protobuf.message.Message):
-        ...
-
-    def _get_metadata(self) -> modal_proto.api_pb2.ClassHandleMetadata:
-        ...
-
+    _user_cls: typing.Optional[type]
+    _class_service_function: typing.Optional[modal.functions.Function]
+    _method_functions: typing.Optional[dict[str, modal.functions.Function]]
+    _options: typing.Optional[modal_proto.api_pb2.FunctionOptions]
+    _callables: dict[str, typing.Callable[..., typing.Any]]
+    _app: typing.Optional[modal.app.App]
+    _name: typing.Optional[str]
+
+    def __init__(self, *args, **kwargs): ...
+    def _initialize_from_empty(self): ...
+    def _initialize_from_other(self, other: Cls): ...
+    def _get_partial_functions(self) -> dict[str, modal.partial_function.PartialFunction]: ...
+    def _hydrate_metadata(self, metadata: google.protobuf.message.Message): ...
     @staticmethod
-    def from_local(user_cls, app, decorator: typing.Callable[[modal.partial_function.PartialFunction, type], modal.functions.Function]) -> Cls:
-        ...
-
+    def validate_construction_mechanism(user_cls): ...
+    @staticmethod
+    def from_local(user_cls, app: modal.app.App, class_service_function: modal.functions.Function) -> Cls: ...
+    def _uses_common_service_function(self): ...
     @classmethod
-    def from_name(cls: typing.Type[Cls], app_name: str, tag: typing.Union[str, None] = None, namespace=1, environment_name: typing.Union[str, None] = None, workspace: typing.Union[str, None] = None) -> Cls:
-        ...
-
-    def with_options(self: Cls, cpu: typing.Union[float, None] = None, memory: typing.Union[int, typing.Tuple[int, int], None] = None, gpu: typing.Union[None, bool, str, modal.gpu._GPUConfig] = None, secrets: typing.Collection[modal.secret.Secret] = (), volumes: typing.Dict[typing.Union[str, os.PathLike], modal.volume.Volume] = {}, retries: typing.Union[int, modal.retries.Retries, None] = None, timeout: typing.Union[int, None] = None, concurrency_limit: typing.Union[int, None] = None, allow_concurrent_inputs: typing.Union[int, None] = None, container_idle_timeout: typing.Union[int, None] = None, allow_background_volume_commits: bool = False) -> Cls:
-        ...
+    def from_name(
+        cls: type[Cls],
+        app_name: str,
+        name: str,
+        namespace=1,
+        environment_name: typing.Optional[str] = None,
+        workspace: typing.Optional[str] = None,
+    ) -> Cls: ...
+    def with_options(
+        self: Cls,
+        cpu: typing.Union[float, tuple[float, float], None] = None,
+        memory: typing.Union[int, tuple[int, int], None] = None,
+        gpu: typing.Union[None, bool, str, modal.gpu._GPUConfig] = None,
+        secrets: collections.abc.Collection[modal.secret.Secret] = (),
+        volumes: dict[typing.Union[str, os.PathLike], modal.volume.Volume] = {},
+        retries: typing.Union[int, modal.retries.Retries, None] = None,
+        timeout: typing.Optional[int] = None,
+        concurrency_limit: typing.Optional[int] = None,
+        allow_concurrent_inputs: typing.Optional[int] = None,
+        container_idle_timeout: typing.Optional[int] = None,
+    ) -> Cls: ...
 
     class __lookup_spec(typing_extensions.Protocol):
-        def __call__(self, app_name: str, tag: typing.Union[str, None] = None, namespace=1, client: typing.Union[modal.client.Client, None] = None, environment_name: typing.Union[str, None] = None, workspace: typing.Union[str, None] = None) -> Cls:
-            ...
-
-        async def aio(self, *args, **kwargs) -> Cls:
-            ...
+        def __call__(
+            self,
+            app_name: str,
+            name: str,
+            namespace=1,
+            client: typing.Optional[modal.client.Client] = None,
+            environment_name: typing.Optional[str] = None,
+            workspace: typing.Optional[str] = None,
+        ) -> Cls: ...
+        async def aio(
+            self,
+            app_name: str,
+            name: str,
+            namespace=1,
+            client: typing.Optional[modal.client.Client] = None,
+            environment_name: typing.Optional[str] = None,
+            workspace: typing.Optional[str] = None,
+        ) -> Cls: ...
 
     lookup: __lookup_spec
 
-    def __call__(self, *args, **kwargs) -> Obj:
-        ...
+    def __call__(self, *args, **kwargs) -> Obj: ...
+    def __getattr__(self, k): ...
+
+class _NO_DEFAULT:
+    def __repr__(self): ...
+
+_no_default: _NO_DEFAULT
+
+class _Parameter:
+    default: typing.Any
+    init: bool
+
+    def __init__(self, default: typing.Any, init: bool): ...
+    def __get__(self, obj, obj_type=None) -> typing.Any: ...
 
-    def __getattr__(self, k):
-        ...
+def is_parameter(p: typing.Any) -> bool: ...
+def parameter(*, default: typing.Any = modal.cls._NO_DEFAULT(), init: bool = True) -> typing.Any: ...

modal/config.py

@@ -29,7 +29,7 @@ Setting tokens using the CLI
 
 You can set a token by running the command::
 
-```bash
+```
 modal token set \
   --token-id <token id> \
   --token-secret <token secret>
@@ -55,6 +55,11 @@ Other possible configuration options are:
   Defaults to True.
   By default, Modal automatically mounts modules imported in the current scope, that
   are deemed to be "local". This can be turned off by setting this to False.
+* `force_build` (in the .toml file) / `MODAL_FORCE_BUILD` (as an env var).
+  Defaults to False.
+  When set, ignores the Image cache and builds all Image layers. Note that this
+  will break the cache for all images based on the rebuilt layers, so other images
+  may rebuild on subsequent runs / deploys even if the config is reverted.
 * `traceback` (in the .toml file) / `MODAL_TRACEBACK` (as an env var).
   Defaults to False. Enables printing full tracebacks on unexpected CLI
   errors, which can be useful for debugging client issues.
@@ -75,14 +80,15 @@ import os
 import typing
 import warnings
 from textwrap import dedent
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from google.protobuf.empty_pb2 import Empty
 
 from modal_proto import api_pb2
 
+from ._utils.deprecation import deprecation_error
 from ._utils.logger import configure_logger
-from .exception import InvalidError, deprecation_warning
+from .exception import InvalidError
 
 # Locate config file and read it
 
@@ -98,14 +104,25 @@ def _is_remote() -> bool:
 
 
 def _read_user_config():
+    config_data = {}
     if not _is_remote() and os.path.exists(user_config_path):
         # Defer toml import so we don't need it in the container runtime environment
         import toml
 
-        with open(user_config_path) as f:
-            return toml.load(f)
-    else:
-        return {}
+        try:
+            with open(user_config_path) as f:
+                config_data = toml.load(f)
+        except Exception as exc:
+            config_problem = str(exc)
+        else:
+            if not all(isinstance(e, dict) for e in config_data.values()):
+                config_problem = "TOML file must contain table sections for each profile."
+            else:
+                config_problem = ""
+        if config_problem:
+            message = f"\nError when reading the modal configuration from `{user_config_path}`.\n\n{config_problem}"
+            raise InvalidError(message)
+    return config_data
 
 
 _user_config = _read_user_config()
@@ -161,11 +178,9 @@ def _check_config() -> None:
             Support for using an implicit 'default' profile is deprecated.
             Please use `modal profile activate` to activate one of your profiles.
             (Use `modal profile list` to see the options.)
-
-            This will become an error in a future update.
             """
         )
-        deprecation_warning((2024, 2, 6), message, show_source=False)
+        deprecation_error((2024, 2, 6), message)
 
 
 _profile = os.environ.get("MODAL_PROFILE") or _config_active_profile()
@@ -189,16 +204,15 @@ _SETTINGS = {
     "token_id": _Setting(),
     "token_secret": _Setting(),
     "task_id": _Setting(),
-    "task_secret": _Setting(),
     "serve_timeout": _Setting(transform=float),
     "sync_entrypoint": _Setting(),
     "logs_timeout": _Setting(10, float),
     "image_id": _Setting(),
     "automount": _Setting(True, transform=_to_boolean),
-    "profiling_enabled": _Setting(False, transform=_to_boolean),
     "heartbeat_interval": _Setting(15, float),
     "function_runtime": _Setting(),
     "function_runtime_debug": _Setting(False, transform=_to_boolean),  # For internal debugging use.
+    "runtime_perf_record": _Setting(False, transform=_to_boolean),  # For internal debugging use.
     "environment": _Setting(),
     "default_cloud": _Setting(None, transform=lambda x: x if x else None),
     "worker_id": _Setting(),  # For internal debugging use.
@@ -206,6 +220,9 @@ _SETTINGS = {
     "force_build": _Setting(False, transform=_to_boolean),
     "traceback": _Setting(False, transform=_to_boolean),
     "image_builder_version": _Setting(),
+    "strict_parameters": _Setting(False, transform=_to_boolean),  # For internal/experimental use
+    "snapshot_debug": _Setting(False, transform=_to_boolean),
+    "client_retries": _Setting(False, transform=_to_boolean),  # For internal testing.
 }
 
 
@@ -243,7 +260,7 @@ class Config:
             os.environ["MODAL_" + key.upper()] = value
         except KeyError:
             # Override env vars not available in config, e.g. NVIDIA_VISIBLE_DEVICES.
-            # This is used for restoring env vars from a checkpoint.
+            # This is used for restoring env vars from a memory snapshot.
             os.environ[key.upper()] = value
 
     def __getitem__(self, key):
@@ -253,7 +270,7 @@ class Config:
         return repr(self.to_dict())
 
     def to_dict(self):
-        return {key: self.get(key) for key in _SETTINGS.keys()}
+        return {key: self.get(key) for key in sorted(_SETTINGS)}
 
 
 config = Config()
@@ -267,7 +284,7 @@ configure_logger(logger, config["loglevel"], config["log_format"])
 
 
 def _store_user_config(
-    new_settings: Dict[str, Any], profile: Optional[str] = None, active_profile: Optional[str] = None
+    new_settings: dict[str, Any], profile: Optional[str] = None, active_profile: Optional[str] = None
 ):
     """Internal method, used by the CLI to set tokens."""
     if profile is None:

modal/container_process.py

@@ -0,0 +1,177 @@
+# Copyright Modal Labs 2024
+import asyncio
+import platform
+from typing import Generic, Optional, TypeVar
+
+from modal_proto import api_pb2
+
+from ._utils.async_utils import TaskContext, synchronize_api
+from ._utils.deprecation import deprecation_error
+from ._utils.grpc_utils import retry_transient_errors
+from ._utils.shell_utils import stream_from_stdin, write_to_fd
+from .client import _Client
+from .exception import InteractiveTimeoutError, InvalidError
+from .io_streams import _StreamReader, _StreamWriter
+from .stream_type import StreamType
+
+T = TypeVar("T", str, bytes)
+
+
+class _ContainerProcess(Generic[T]):
+    _process_id: Optional[str] = None
+    _stdout: _StreamReader[T]
+    _stderr: _StreamReader[T]
+    _stdin: _StreamWriter
+    _text: bool
+    _by_line: bool
+    _returncode: Optional[int] = None
+
+    def __init__(
+        self,
+        process_id: str,
+        client: _Client,
+        stdout: StreamType = StreamType.PIPE,
+        stderr: StreamType = StreamType.PIPE,
+        text: bool = True,
+        by_line: bool = False,
+    ) -> None:
+        self._process_id = process_id
+        self._client = client
+        self._text = text
+        self._by_line = by_line
+        self._stdout = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDOUT,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stdout,
+            text=text,
+            by_line=by_line,
+        )
+        self._stderr = _StreamReader[T](
+            api_pb2.FILE_DESCRIPTOR_STDERR,
+            process_id,
+            "container_process",
+            self._client,
+            stream_type=stderr,
+            text=text,
+            by_line=by_line,
+        )
+        self._stdin = _StreamWriter(process_id, "container_process", self._client)
+
+    @property
+    def stdout(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stdout stream."""
+        return self._stdout
+
+    @property
+    def stderr(self) -> _StreamReader[T]:
+        """StreamReader for the container process's stderr stream."""
+        return self._stderr
+
+    @property
+    def stdin(self) -> _StreamWriter:
+        """StreamWriter for the container process's stdin stream."""
+        return self._stdin
+
+    @property
+    def returncode(self) -> int:
+        if self._returncode is None:
+            raise InvalidError(
+                "You must call wait() before accessing the returncode. "
+                "To poll for the status of a running process, use poll() instead."
+            )
+        return self._returncode
+
+    async def poll(self) -> Optional[int]:
+        """Check if the container process has finished running.
+
+        Returns `None` if the process is still running, else returns the exit code.
+        """
+        if self._returncode is not None:
+            return self._returncode
+
+        req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=0)
+        resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(self._client.stub.ContainerExecWait, req)
+
+        if resp.completed:
+            self._returncode = resp.exit_code
+            return self._returncode
+
+        return None
+
+    async def wait(self) -> int:
+        """Wait for the container process to finish running. Returns the exit code."""
+
+        if self._returncode is not None:
+            return self._returncode
+
+        while True:
+            req = api_pb2.ContainerExecWaitRequest(exec_id=self._process_id, timeout=50)
+            resp: api_pb2.ContainerExecWaitResponse = await retry_transient_errors(
+                self._client.stub.ContainerExecWait, req
+            )
+            if resp.completed:
+                self._returncode = resp.exit_code
+                return self._returncode
+
+    async def attach(self, *, pty: Optional[bool] = None):
+        if platform.system() == "Windows":
+            print("interactive exec is not currently supported on Windows.")
+            return
+
+        if pty is not None:
+            deprecation_error(
+                (2024, 12, 9),
+                "The `pty` argument to `modal.container_process.attach(pty=...)` is deprecated, "
+                "as only PTY mode is supported. Please remove the argument.",
+            )
+
+        from rich.console import Console
+
+        console = Console()
+
+        connecting_status = console.status("Connecting...")
+        connecting_status.start()
+        on_connect = asyncio.Event()
+
+        async def _write_to_fd_loop(stream: _StreamReader):
+            # Don't skip empty messages so we can detect when the process has booted.
+            async for chunk in stream._get_logs(skip_empty_messages=False):
+                if chunk is None:
+                    break
+
+                if not on_connect.is_set():
+                    connecting_status.stop()
+                    on_connect.set()
+
+                await write_to_fd(stream.file_descriptor, chunk)
+
+        async def _handle_input(data: bytes, message_index: int):
+            self.stdin.write(data)
+            await self.stdin.drain()
+
+        async with TaskContext() as tc:
+            stdout_task = tc.create_task(_write_to_fd_loop(self.stdout))
+            stderr_task = tc.create_task(_write_to_fd_loop(self.stderr))
+
+            try:
+                # time out if we can't connect to the server fast enough
+                await asyncio.wait_for(on_connect.wait(), timeout=60)
+
+                async with stream_from_stdin(_handle_input, use_raw_terminal=True):
+                    await stdout_task
+                    await stderr_task
+
+                # TODO: this doesn't work right now.
+                # if exit_status != 0:
+                #     raise ExecutionError(f"Process exited with status code {exit_status}")
+
+            except (asyncio.TimeoutError, TimeoutError):
+                connecting_status.stop()
+                stdout_task.cancel()
+                stderr_task.cancel()
+                raise InteractiveTimeoutError("Failed to establish connection to container. Please try again.")
+
+
+ContainerProcess = synchronize_api(_ContainerProcess)