ImageSpec: add arg files & envd_config io.copy()

Signed-off-by: Austin Liu <[email protected]> ImageSpec: add arg files Signed-off-by: Austin Liu <[email protected]> imageSpec: add files Signed-off-by: Austin Liu <[email protected]> fmt Signed-off-by: Austin Liu <[email protected]> fix path Signed-off-by: Austin Liu <[email protected]>
flyteorg · Mar 16, 2024 · 4d95e19 · 4d95e19
1 parent f5cd70d
commit 4d95e19
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 1 deletion.
diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py
@@ -51,6 +51,7 @@ class ImageSpec:
     builder: Optional[str] = None
     source_root: Optional[str] = None
     env: Optional[typing.Dict[str, str]] = None
+    files: Optional[List[str]] = None
     registry: Optional[str] = None
     packages: Optional[List[str]] = None
     conda_packages: Optional[List[str]] = None

diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py
@@ -73,7 +73,6 @@ def create_envd_config(image_spec: ImageSpec) -> str:
     if image_spec.env:
         env.update(image_spec.env)
     pip_index = "https://pypi.org/simple" if image_spec.pip_index is None else image_spec.pip_index
-
     envd_config = f"""# syntax=v1
 
 def build():
@@ -88,6 +87,13 @@ def build():
     cfg_path = ctx.file_access.get_random_local_path("build.envd")
     pathlib.Path(cfg_path).parent.mkdir(parents=True, exist_ok=True)
 
+    if len(image_spec.files) > 0:
+        for fd in image_spec.files:
+            fb = os.path.basename(fd)
+            file_path = f"{pathlib.Path(cfg_path).parent}{os.sep}{fb}"
+            shutil.copyfile(f"{fd}", file_path)
+        envd_config += '    io.copy(source="./", target="/root")'
+
     if conda_packages:
         envd_config += "    install.conda(use_mamba=True)\n"
         envd_config += f"    install.conda_packages(name=[{conda_packages}], channel=[{conda_channels}])\n"