[Mounts on images] add attach_local_file/attach_local_dir

modal/image.py

@@ -606,12 +606,70 @@ def build_dockerfile(version: ImageBuilderVersion) -> DockerfileSpec:
             context_mount=mount,
         )
 
+    def attach_local_file(self, local_path: Union[str, Path], remote_path: str, *, copy: bool = False) -> "_Image":
+        """Attach a local file to the image
+
+        By default, the file is attached as a lightweight mount on top of any
+        container that runs the returned Image. This makes this operation quick
+        and will not invalidate + rebuild the image when contents of the mount
+        change.
+
+        In order to run use the file in *subsequent* build steps, you explicitly
+        need to set `copy=True` on this function to instead copy the data into
+        the image itself.
+
+        In copy=True mode, this command works similar way to [`COPY`](https://docs.docker.com/engine/reference/builder/#copy)
+        works in a `Dockerfile`.
+        """
+        if not os.path.isabs(remote_path):
+            # TODO(elias): implement relative to absolute resolution using image workdir metadata
+            #  + make default remote_path="./"
+            #  This requires deferring the Mount creation until after "self" (the base image) has been resolved
+            #  so we know the workdir of the operation.
+            raise InvalidError("image.attach_local_file() currently only supports absolute remote_path values")
+
+        if remote_path.endswith("/"):
+            remote_path = remote_path + Path(local_path).name
+
+        mount = _Mount.from_local_file(local_path, remote_path)
+        return self._add_mount_layer_or_copy(mount, copy=copy)
+
+    def attach_local_dir(
+        self, local_path: Union[str, Path], remote_path: Union[str, Path], *, copy: bool = False
+    ) -> "_Image":
+        """Attach a local dir to the image
+
+        By default, the dir is attached as a lightweight mount on top of any
+        container that runs the returned Image. This makes this operation quick
+        and will not invalidate + rebuild the image when contents of the mount
+        change.
+
+        In order to run use the dir in *subsequent* build steps, you explicitly
+        need to set `copy=True` on this function to instead copy the data into
+        the image itself.
+
+        In copy=True mode, this command works similar way to [`COPY`](https://docs.docker.com/engine/reference/builder/#copy)
+        works in a `Dockerfile`.
+        """
+        # TODO(elias): distinguish adding *into* remote_path and *as* remote_path, similar
+        #  to shell `cp -r source dest` vs `cp -r source/ dest`
+        #  NOTE: docker COPY always copies the *contents* of the source directory, and you have
+        #  to specify the destination dir name explicitly
+        remote_path = PurePosixPath(remote_path)
+        if not remote_path.is_absolute():
+            # TODO(elias): implement relative to absolute resolution using image workdir metadata
+            #  + make default remote_path="./"
+            raise InvalidError("image.attach_local_dir() currently only supports absolute remote_path values")
+        mount = _Mount.from_local_dir(local_path, remote_path=remote_path)
+        return self._add_mount_layer_or_copy(mount, copy=copy)
+
     def copy_local_file(self, local_path: Union[str, Path], remote_path: Union[str, Path] = "./") -> "_Image":
         """Copy a file into the image as a part of building it.
 
         This works in a similar way to [`COPY`](https://docs.docker.com/engine/reference/builder/#copy)
         works in a `Dockerfile`.
         """
+        # TODO(elias): add pending deprecation - use attach
         basename = str(Path(local_path).name)
         mount = _Mount.from_local_file(local_path, remote_path=f"/{basename}")
 
@@ -1634,7 +1692,7 @@ def build_dockerfile(version: ImageBuilderVersion) -> DockerfileSpec:
             dockerfile_function=build_dockerfile,
         )
 
-    def workdir(self, path: str) -> "_Image":
+    def workdir(self, path: Union[str, PurePosixPath]) -> "_Image":
         """Set the working directory for subsequent image build steps and function execution.
 
         **Example**

modal/mount.py

@@ -378,7 +378,9 @@ def from_local_dir(
         )
 
     def add_local_file(
-        self, local_path: Union[str, Path], remote_path: Union[str, PurePosixPath, None] = None
+        self,
+        local_path: Union[str, Path],
+        remote_path: Union[str, PurePosixPath, None] = None,
     ) -> "_Mount":
         """
         Add a local file to the `Mount` object.

test/image_test.py

@@ -624,6 +624,87 @@ def tmp_path_with_content(tmp_path):
     return tmp_path
 
 
+@pytest.mark.parametrize(["copy"], [(True,), (False,)])
+@pytest.mark.parametrize(
+    ["remote_path", "expected_dest"],
+    [
+        ("/place/nice.txt", "/place/nice.txt"),
+        # Not supported yet, but soon:
+        ("/place/", "/place/data.txt"),  # use original basename if destination has a trailing slash
+        # ("output.txt", "/proj/output.txt")  # workdir relative target
+        # (None, "/proj/data.txt")  # default target - basename in current directory
+    ],
+)
+def test_image_attach_local_file(servicer, client, tmp_path_with_content, copy, remote_path, expected_dest):
+    app = App()
+
+    if remote_path is None:
+        remote_path_kwargs = {}
+    else:
+        remote_path_kwargs = {"remote_path": remote_path}
+
+    img = (
+        Image.from_registry("unknown_image")
+        .workdir("/proj")
+        .attach_local_file(tmp_path_with_content / "data.txt", **remote_path_kwargs, copy=copy)
+    )
+    app.function(image=img)(dummy)
+
+    with app.run(client=client):
+        if copy:
+            # check that dockerfile commands include COPY . .
+            layers = get_image_layers(img.object_id, servicer)
+            assert layers[0].dockerfile_commands == ["FROM base", "COPY . /"]
+            mount_id = layers[0].context_mount_id
+            # and then get the relevant context mount to check
+        if not copy:
+            assert len(img._mount_layers) == 1
+            mount_id = img._mount_layers[0].object_id
+
+        assert set(servicer.mount_contents[mount_id].keys()) == {expected_dest}
+
+
+@pytest.mark.parametrize(["copy"], [(True,), (False,)])
+@pytest.mark.parametrize(
+    ["remote_path", "expected_dest"],
+    [
+        ("/place/", "/place/sub"),  # copy full dir
+        ("/place", "/place/sub"),  # removing trailing slash on source makes no difference, unlike shell cp
+        # TODO: add support for relative paths:
+        # Not supported yet, but soon:
+        # ("place", "/proj/place/sub")  # workdir relative target
+        # (None, "/proj/sub")  # default target - copy into current directory
+    ],
+)
+def test_image_attach_local_dir(servicer, client, tmp_path_with_content, copy, remote_path, expected_dest):
+    app = App()
+
+    if remote_path is None:
+        remote_path_kwargs = {}
+    else:
+        remote_path_kwargs = {"remote_path": remote_path}
+
+    img = (
+        Image.from_registry("unknown_image")
+        .workdir("/proj")
+        .attach_local_dir(tmp_path_with_content / "data", **remote_path_kwargs, copy=copy)
+    )
+    app.function(image=img)(dummy)
+
+    with app.run(client=client):
+        if copy:
+            # check that dockerfile commands include COPY . .
+            layers = get_image_layers(img.object_id, servicer)
+            assert layers[0].dockerfile_commands == ["FROM base", "COPY . /"]
+            mount_id = layers[0].context_mount_id
+            # and then get the relevant context mount to check
+        if not copy:
+            assert len(img._mount_layers) == 1
+            mount_id = img._mount_layers[0].object_id
+
+        assert set(servicer.mount_contents[mount_id].keys()) == {expected_dest}
+
+
 def test_image_copy_local_dir(builder_version, servicer, client, tmp_path_with_content):
     app = App()
     app.image = Image.debian_slim().copy_local_dir(tmp_path_with_content, remote_path="/dummy")