[Flytekit] Envd builder with extra copy commands (#2774)

Signed-off-by: mao3267 <[email protected]>

plugins/flytekit-envd/flytekitplugins/envd/image_builder.py

@@ -176,6 +176,44 @@ def build():
         else:
             envd_config += '    io.copy(source="./", target="/root")\n'
 
+    if image_spec.copy:
+
+        def add_envd_copy_command(src_path: pathlib.Path, envd_config: str) -> str:
+            envd_version = metadata.version("envd")
+            if Version(envd_version) <= Version("0.3.37"):
+                if src_path.is_dir():
+                    envd_config += (
+                        f'    io.copy(host_path="{src_path.as_posix()}", envd_path="/root/{src_path.as_posix()}/")\n'
+                    )
+                else:
+                    envd_config += f'    io.copy(host_path="{src_path.as_posix()}", envd_path="/root/{src_path.parent.as_posix()}/")\n'
+            else:
+                if src_path.is_dir():
+                    envd_config += (
+                        f'    io.copy(source="{src_path.as_posix()}", target="/root/{src_path.as_posix()}/")\n'
+                    )
+                else:
+                    envd_config += (
+                        f'    io.copy(source="{src_path.as_posix()}", target="/root/{src_path.parent.as_posix()}/")\n'
+                    )
+            return envd_config
+
+        for src in image_spec.copy:
+            src_path = pathlib.Path(src)
+
+            if src_path.is_absolute() or ".." in src_path.parts:
+                raise ValueError("Absolute paths or paths with '..' are not allowed in COPY command.")
+
+            dst_path = pathlib.Path(cfg_path).parent / src_path
+            dst_path.parent.mkdir(parents=True, exist_ok=True)
+
+            if src_path.is_dir():
+                shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
+            else:
+                shutil.copy(src_path, dst_path)
+
+            envd_config = add_envd_copy_command(src_path, envd_config)
+
     with open(cfg_path, "w+") as f:
         f.write(envd_config)
 

plugins/flytekit-envd/tests/test_image_spec.py

@@ -6,7 +6,7 @@
 from flytekitplugins.envd.image_builder import EnvdImageSpecBuilder, create_envd_config
 
 from flytekit.image_spec.image_spec import ImageBuildEngine, ImageSpec
-
+import tempfile
 
 @pytest.fixture(scope="module", autouse=True)
 def register_envd_higher_priority():
@@ -32,24 +32,29 @@ def test_image_spec():
     # so Envd can find the base image when building imageSpec below
     ImageBuildEngine._IMAGE_NAME_TO_REAL_NAME[base_image.image_name()] = "cr.flyte.org/flyteorg/flytekit:py3.8-latest"
 
-    image_spec = ImageSpec(
-        packages=["pandas"],
-        apt_packages=["git"],
-        python_version="3.8",
-        base_image=base_image,
-        pip_index="https://pypi.python.org/simple",
-        source_root=os.path.dirname(os.path.realpath(__file__)),
-    )
-
-    image_spec = image_spec.with_commands("echo hello")
-    ImageBuildEngine.build(image_spec)
-    image_spec.base_image = base_image.image_name()
-    config_path = create_envd_config(image_spec)
-    assert image_spec.platform == "linux/amd64"
-    contents = Path(config_path).read_text()
-    assert (
-        contents
-        == f"""# syntax=v1
+    with tempfile.TemporaryDirectory(dir=Path.cwd().as_posix()) as temp_dir:
+        copy_file = Path(temp_dir) / "copy_file.txt"
+        copy_file.write_text("copy_file_content")
+
+        image_spec = ImageSpec(
+            packages=["pandas"],
+            apt_packages=["git"],
+            python_version="3.8",
+            base_image=base_image,
+            pip_index="https://pypi.python.org/simple",
+            source_root=os.path.dirname(os.path.realpath(__file__)),
+            copy=[f"{copy_file.relative_to(Path.cwd()).as_posix()}"],
+        )
+
+        image_spec = image_spec.with_commands("echo hello")
+        ImageBuildEngine.build(image_spec)
+        image_spec.base_image = base_image.image_name()
+        config_path = create_envd_config(image_spec)
+        assert image_spec.platform == "linux/amd64"
+        contents = Path(config_path).read_text()
+        assert (
+            contents
+            == f"""# syntax=v1
 
 def build():
     base(image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", dev=False)
@@ -60,6 +65,7 @@ def build():
     config.pip_index(url="https://pypi.python.org/simple")
     install.python(version="3.8")
     io.copy(source="./", target="/root")
+    io.copy(source="{copy_file.relative_to(Path.cwd()).as_posix()}", target="/root/{copy_file.parent.relative_to(Path.cwd()).as_posix()}/")
 """
     )