Mounts on images [CLI-4] - image.add_local_python_package

modal/functions.py

@@ -601,6 +601,7 @@ def from_args(
 
         if is_local():
             entrypoint_mounts = info.get_entrypoint_mount()
+
             all_mounts = [
                 _get_client_mount(),
                 *explicit_mounts,
@@ -633,6 +634,7 @@ def from_args(
         if proxy:
             # HACK: remove this once we stop using ssh tunnels for this.
             if image:
+                # TODO(elias): this will cause an error if users use prior `.add_local_*` commands without copy=True
                 image = image.apt_install("autossh")
 
         function_spec = _FunctionSpec(
@@ -849,7 +851,7 @@ async def _load(self: _Function, resolver: Resolver, existing_object_id: Optiona
                     )
                     for path, volume in validated_volumes
                 ]
-                loaded_mount_ids = {m.object_id for m in all_mounts}
+                loaded_mount_ids = {m.object_id for m in all_mounts} | {m.object_id for m in image._mount_layers}
 
                 # Get object dependencies
                 object_dependencies = []

modal/image.py

@@ -5,6 +5,7 @@
 import re
 import shlex
 import sys
+import textwrap
 import typing
 import warnings
 from dataclasses import dataclass
@@ -274,11 +275,21 @@ class _Image(_Object, type_prefix="im"):
     force_build: bool
     inside_exceptions: List[Exception]
     _used_local_mounts: typing.FrozenSet[_Mount]  # used for mounts watching
+    _mounts: Sequence[_Mount]  # added as mounts on any container referencing the Image, see `def _mount_layers`
     _metadata: Optional[api_pb2.ImageMetadata] = None  # set on hydration, private for now
 
     def _initialize_from_empty(self):
         self.inside_exceptions = []
         self._used_local_mounts = frozenset()
+        self._mounts = ()
+        self.force_build = False
+
+    def _initialize_from_other(self, other: "_Image"):
+        # used by .clone()
+        self.inside_exceptions = other.inside_exceptions
+        self.force_build = other.force_build
+        self._used_local_mounts = other._used_local_mounts
+        self._mounts = other._mounts
 
     def _hydrate_metadata(self, message: Optional[Message]):
         env_image_id = config.get("image_id")  # set as an env var in containers
@@ -292,6 +303,57 @@ def _hydrate_metadata(self, message: Optional[Message]):
             assert isinstance(message, api_pb2.ImageMetadata)
             self._metadata = message
 
+    def _add_mount_layer_or_copy(self, mount: _Mount, copy: bool = False):
+        if copy:
+            return self.copy_mount(mount, remote_path="/")
+
+        base_image = self
+
+        async def _load(self2: "_Image", resolver: Resolver, existing_object_id: Optional[str]):
+            self2._hydrate_from_other(base_image)  # same image id as base image as long as it's lazy
+            self2._mounts = tuple(base_image._mounts) + (mount,)
+            self2._used_local_mounts = base_image._used_local_mounts | ({mount} if mount.is_local() else set())
+
+        return _Image._from_loader(_load, "ImageWithMounts()", deps=lambda: [base_image, mount])
+
+    @property
+    def _mount_layers(self) -> typing.Tuple[_Mount]:
+        """Non-evaluated mount layers on the image
+
+        When the image is used by a Modal container, these mounts need to be attached as well to
+        represent the full image content, as they haven't yet been represented as a layer in the
+        image.
+
+        When the image is used as a base image for a new layer (that is not itself a mount layer)
+        these mounts need to first be inserted as a copy operation (.copy_mount) into the image.
+        """
+        return self._mounts
+
+    def _assert_no_mount_layers(self):
+        if self._mount_layers:
+            raise InvalidError(
+                textwrap.dedent(
+                    """
+            An image tried to run a build step after using `image.add_local_*` to include local files.
+
+            Run `image.add_local_*` commands last in your image build to avoid rebuilding images with every local
+            filechange. Modal will then mount these files as a thin layer when starting your container, saving build
+            time.
+
+            If you need these files earlier in the build, set `copy=True` to copy the files directly into the image,
+            though this will increase build time.
+
+            Example:
+
+            my_image = (
+                Image.debian_slim()
+                .add_local_python_packages("mypak", copy=True)
+                .run_commands("python -m mypak")  # this now works!
+            )
+            """
+                )
+            )
+
     @staticmethod
     def _from_args(
         *,
@@ -306,9 +368,11 @@ def _from_args(
         force_build: bool = False,
         # For internal use only.
         _namespace: int = api_pb2.DEPLOYMENT_NAMESPACE_WORKSPACE,
+        _do_assert_no_mount_layers: bool = True,
     ):
         if base_images is None:
             base_images = {}
+
         if secrets is None:
             secrets = []
         if gpu_config is None:
@@ -334,6 +398,11 @@ def _deps() -> List[_Object]:
             return deps
 
         async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[str]):
+            if _do_assert_no_mount_layers:
+                for image in base_images.values():
+                    # base images can't have
+                    image._assert_no_mount_layers()
+
             environment = await _get_environment_cached(resolver.environment_name or "", resolver.client)
             # A bit hacky,but assume that the environment provides a valid builder version
             image_builder_version = cast(ImageBuilderVersion, environment._settings.image_builder_version)
@@ -349,7 +418,9 @@ async def _load(self: _Image, resolver: Resolver, existing_object_id: Optional[s
                     "No commands were provided for the image — have you tried using modal.Image.debian_slim()?"
                 )
             if dockerfile.commands and build_function:
-                raise InvalidError("Cannot provide both a build function and Dockerfile commands!")
+                raise InvalidError(
+                    "Cannot provide both build function and Dockerfile commands in the same image layer!"
+                )
 
             base_images_pb2s = [
                 api_pb2.BaseImage(
@@ -487,7 +558,7 @@ async def join():
                 local_mounts.add(context_mount)
             self._used_local_mounts = frozenset(local_mounts)
 
-        rep = "Image()"
+        rep = f"Image({dockerfile_function})"
         obj = _Image._from_loader(_load, rep, deps=_deps)
         obj.force_build = force_build
         return obj
@@ -553,6 +624,23 @@ def build_dockerfile(version: ImageBuilderVersion) -> DockerfileSpec:
             context_mount=mount,
         )
 
+    def add_local_python_packages(self, *packages: Union[str, Path], copy: bool = False) -> "_Image":
+        """Attaches local Python packages to the container running the image
+
+        Packages are added to the /root directory which is on the PYTHONPATH of any
+        executed Modal functions.
+
+        By default (copy=False) the packages are layered on top of the image when
+        the container starts up and not built as an image layer.
+
+        Set copy=True to force the packages to be added as an image layer instead.
+        This can be slower since it requires a rebuild of the image whenever the
+        required package files change, but it allows you to run additional build
+        steps after this operation.
+        """
+        mount = _Mount.from_local_python_packages(*packages)
+        return self._add_mount_layer_or_copy(mount, copy=copy)
+
     def copy_local_dir(self, local_path: Union[str, Path], remote_path: Union[str, Path] = ".") -> "_Image":
         """Copy a directory into the image as a part of building the image.
 

modal/sandbox.py

@@ -148,7 +148,7 @@ async def _load(self: _Sandbox, resolver: Resolver, _existing_object_id: Optiona
             definition = api_pb2.Sandbox(
                 entrypoint_args=entrypoint_args,
                 image_id=image.object_id,
-                mount_ids=[mount.object_id for mount in mounts],
+                mount_ids=[mount.object_id for mount in mounts] + [mount.object_id for mount in image._mount_layers],
                 secret_ids=[secret.object_id for secret in secrets],
                 timeout_secs=timeout,
                 workdir=workdir,

test/conftest.py

@@ -2026,3 +2026,8 @@ def import_fail_for_rich(name: str, *args, **kwargs) -> ModuleType:
 def disable_auto_mount(monkeypatch):
     monkeypatch.setenv("MODAL_AUTOMOUNT", "0")
     yield
+
+
+@pytest.fixture()
+def supports_on_path(supports_dir, monkeypatch):
+    monkeypatch.syspath_prepend(str(supports_dir))

test/image_test.py

@@ -10,6 +10,7 @@
 from typing import List, Literal, get_args
 from unittest import mock
 
+import modal
 from modal import App, Image, Mount, Secret, build, environments, gpu, method
 from modal._serialization import serialize
 from modal._utils.async_utils import synchronizer
@@ -24,6 +25,7 @@
     _validate_python_version,
 )
 from modal.mount import PYTHON_STANDALONE_VERSIONS
+from modal.runner import deploy_app
 from modal_proto import api_pb2
 
 from .supports.skip import skip_windows
@@ -35,8 +37,16 @@
 ]
 
 
-def dummy():
-    ...
+@pytest.fixture(autouse=True)
+def no_automount(monkeypatch):
+    # no tests in here use automounting, but a lot of them implicitly create
+    # functions w/ lots of modules is sys.modules which will automount
+    # which takes a lot of time, so we disable it
+    monkeypatch.setenv("MODAL_AUTOMOUNT", "0")
+
+
+def dummy() -> None:
+    return None
 
 
 def test_supported_python_series():
@@ -284,8 +294,6 @@ def test_image_pip_install_pyproject(builder_version, servicer, client):
     app.function(image=image)(dummy)
     with app.run(client=client):
         layers = get_image_layers(image.object_id, servicer)
-
-        print(layers[0].dockerfile_commands)
         assert any("pip install 'banana >=1.2.0' 'potato >=0.1.0'" in cmd for cmd in layers[0].dockerfile_commands)
 
 
@@ -298,7 +306,6 @@ def test_image_pip_install_pyproject_with_optionals(builder_version, servicer, c
     with app.run(client=client):
         layers = get_image_layers(image.object_id, servicer)
 
-        print(layers[0].dockerfile_commands)
         assert any(
             "pip install 'banana >=1.2.0' 'linting-tool >=0.0.0' 'potato >=0.1.0' 'pytest >=1.2.0'" in cmd
             for cmd in layers[0].dockerfile_commands
@@ -1209,3 +1216,137 @@ async def test_logs(servicer, client):
 
     logs = [data async for data in image._logs.aio()]
     assert logs == ["build starting\n", "build finished\n"]
+
+
+def hydrate_image(img, client):
+    # there should be a more straight forward way to do this?
+    app = App()
+    app.function(serialized=True, image=img)(lambda: None)
+    with app.run(client=client):
+        pass
+
+
+def test_add_local_lazy_vs_copy(client, servicer, set_env_client, supports_on_path):
+    deb = Image.debian_slim()
+    image_with_mount = deb.add_local_python_packages("pkg_a")
+
+    hydrate_image(image_with_mount, client)
+    assert image_with_mount.object_id == deb.object_id
+    assert len(image_with_mount._mount_layers) == 1
+
+    image_additional_mount = image_with_mount.add_local_python_packages("pkg_b")
+    hydrate_image(image_additional_mount, client)
+    assert len(image_additional_mount._mount_layers) == 2  # another mount added to lazy layer
+    assert len(image_with_mount._mount_layers) == 1  # original image should not be affected
+
+    # running commands
+    image_non_mount = image_with_mount.run_commands("echo 'hello'")
+    with pytest.raises(InvalidError, match="copy=True"):
+        # error about using non-copy add commands before other build steps
+        hydrate_image(image_non_mount, client)
+
+    image_with_copy = deb.add_local_python_packages("pkg_a", copy=True)
+    hydrate_image(image_with_copy, client)
+    assert len(image_with_copy._mount_layers) == 0
+
+    # do the same exact image using copy=True
+    image_with_copy_and_commands = deb.add_local_python_packages("pkg_a", copy=True).run_commands("echo 'hello'")
+    hydrate_image(image_with_copy_and_commands, client)
+    assert len(image_with_copy_and_commands._mount_layers) == 0
+
+    layers = get_image_layers(image_with_copy_and_commands.object_id, servicer)
+
+    echo_layer = layers[0]
+    assert echo_layer.dockerfile_commands == ["FROM base", "RUN echo 'hello'"]
+
+    copy_layer = layers[1]
+    assert copy_layer.dockerfile_commands == ["FROM base", "COPY . /"]
+    copied_files = servicer.mount_contents[copy_layer.context_mount_id].keys()
+    assert len(copied_files) == 8
+    assert all(fn.startswith("/root/pkg_a/") for fn in copied_files)
+
+
+def test_add_locals_are_attached_to_functions(servicer, client, supports_on_path):
+    deb_slim = Image.debian_slim()
+    img = deb_slim.add_local_python_packages("pkg_a")
+    app = App("my-app")
+    control_fun: modal.Function = app.function(serialized=True, image=deb_slim, name="control")(
+        dummy
+    )  # no mounts on image
+    fun: modal.Function = app.function(serialized=True, image=img, name="fun")(dummy)  # mounts on image
+    deploy_app(app, client=client)
+
+    control_func_mounts = set(servicer.app_functions[control_fun.object_id].mount_ids)
+    fun_def = servicer.app_functions[fun.object_id]
+    added_mounts = set(fun_def.mount_ids) - control_func_mounts
+    assert len(added_mounts) == 1
+    assert added_mounts == {img._mount_layers[0].object_id}
+
+
+def test_add_locals_are_attached_to_classes(servicer, client, supports_on_path, set_env_client):
+    deb_slim = Image.debian_slim()
+    img = deb_slim.add_local_python_packages("pkg_a")
+    app = App("my-app")
+    control_fun: modal.Function = app.function(serialized=True, image=deb_slim, name="control")(
+        dummy
+    )  # no mounts on image
+
+    class A:
+        some_arg: str = modal.parameter()
+
+    ACls = app.cls(serialized=True, image=img)(A)  # mounts on image
+    deploy_app(app, client=client)
+
+    control_func_mounts = set(servicer.app_functions[control_fun.object_id].mount_ids)
+    fun_def = servicer.function_by_name("A.*")  # class service function
+    added_mounts = set(fun_def.mount_ids) - control_func_mounts
+    assert len(added_mounts) == 1
+    assert added_mounts == {img._mount_layers[0].object_id}
+
+    obj = ACls(some_arg="foo")  # type: ignore
+    # hacky way to force hydration of the *parameter bound* function (instance service function):
+    obj.keep_warm(0)  #  type: ignore
+
+    obj_fun_def = servicer.function_by_name("A.*", ((), {"some_arg": "foo"}))  # instance service function
+    added_mounts = set(obj_fun_def.mount_ids) - control_func_mounts
+    assert len(added_mounts) == 1
+    assert added_mounts == {img._mount_layers[0].object_id}
+
+
+@skip_windows("servicer sandbox implementation not working on windows")
+def test_add_locals_are_attached_to_sandboxes(servicer, client, supports_on_path):
+    deb_slim = Image.debian_slim()
+    img = deb_slim.add_local_python_packages("pkg_a")
+    app = App("my-app")
+    with app.run(client=client):
+        modal.Sandbox.create(image=img, app=app, client=client)
+        sandbox_def = servicer.sandbox_defs[0]
+
+    assert sandbox_def.image_id == deb_slim.object_id
+    assert sandbox_def.mount_ids == [img._mount_layers[0].object_id]
+    copied_files = servicer.mount_contents[sandbox_def.mount_ids[0]]
+    assert len(copied_files) == 8
+    assert all(fn.startswith("/root/pkg_a/") for fn in copied_files)
+
+
+def empty_fun():
+    pass
+
+
+def test_add_locals_build_function(servicer, client, supports_on_path):
+    deb_slim = Image.debian_slim()
+    img = deb_slim.add_local_python_packages("pkg_a")
+    img_with_build_function = img.run_function(empty_fun)
+    with pytest.raises(InvalidError):
+        # build functions could still potentially rewrite mount contents,
+        # so we still require them to use copy=True
+        # TODO(elias): what if someone wants do use an equivalent of `run_function(..., mounts=[...]) ?
+        hydrate_image(img_with_build_function, client)
+
+    img_with_copy = deb_slim.add_local_python_packages("pkg_a", copy=True)
+    hydrate_image(img_with_copy, client)  # this is fine
+
+
+# TODO: test modal shell w/ lazy mounts
+# this works since the image is passed on as is to a sandbox which will load it and
+# transfer any virtual mount layers from the image as mounts to the sandbox