unpin envd and add extra_path to imagespec (flyteorg#2063)

Signed-off-by: Samhita Alla <[email protected]> Signed-off-by: Kevin Su <[email protected]> Co-authored-by: Kevin Su <[email protected]>
jasonlai1218 · Dec 21, 2023 · 1087356 · 1087356
1 parent 8af01f2
commit 1087356
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py
@@ -1,9 +1,11 @@
+import json
 import os
 import pathlib
 import shutil
 import subprocess
 
 import click
+from packaging.version import Version
 
 from flytekit.configuration import DefaultImages
 from flytekit.core import context_manager
@@ -19,16 +21,22 @@ class EnvdImageSpecBuilder(ImageSpecBuilder):
     def execute_command(self, command):
         click.secho(f"Run command: {command} ", fg="blue")
         p = subprocess.Popen(command.split(), stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+
+        result = []
         for line in iter(p.stdout.readline, ""):
             if p.poll() is not None:
                 break
             if line.decode().strip() != "":
-                click.secho(line.decode().strip(), fg="blue")
+                output = line.decode().strip()
+                click.secho(output, fg="blue")
+                result.append(output)
 
         if p.returncode != 0:
             _, stderr = p.communicate()
             raise Exception(f"failed to run command {command} with error {stderr}")
 
+        return result
+
     def build_image(self, image_spec: ImageSpec):
         cfg_path = create_envd_config(image_spec)
 
@@ -65,7 +73,7 @@ def build():
     run(commands=[{', '.join(map(str, map(lambda x: f'"{x}"', commands)))}])
     install.python_packages(name=[{', '.join(map(str, map(lambda x: f'"{x}"', packages)))}])
     install.apt_packages(name=[{', '.join(map(str, map(lambda x: f'"{x}"', apt_packages)))}])
-    runtime.environ(env={env})
+    runtime.environ(env={env}, extra_path=['/root'])
     config.pip_index(url="{pip_index}")
 """
     ctx = context_manager.FlyteContextManager.current_context()
@@ -87,8 +95,15 @@ def build():
 
     if image_spec.source_root:
         shutil.copytree(image_spec.source_root, pathlib.Path(cfg_path).parent, dirs_exist_ok=True)
+
+        version_command = "envd version -s -f json"
+        envd_version = json.loads(EnvdImageSpecBuilder().execute_command(version_command)[0])["envd"].replace("v", "")
+
         # Indentation is required by envd
-        envd_config += '    io.copy(host_path="./", envd_path="/root")'
+        if Version(envd_version) <= Version("0.3.37"):
+            envd_config += '    io.copy(host_path="./", envd_path="/root")'
+        else:
+            envd_config += '    io.copy(source="./", target="/root")'
 
     with open(cfg_path, "w+") as f:
         f.write(envd_config)

diff --git a/plugins/flytekit-envd/setup.py b/plugins/flytekit-envd/setup.py
@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit", "envd<=0.3.36"]
+plugin_requires = ["flytekit", "envd"]
 
 __version__ = "0.0.0+develop"
 

diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py
@@ -30,7 +30,7 @@ def build():
     run(commands=["echo hello"])
     install.python_packages(name=["pandas"])
     install.apt_packages(name=["git"])
-    runtime.environ(env={{'PYTHONPATH': '/root', '_F_IMG_ID': '{image_name}'}})
+    runtime.environ(env={{'PYTHONPATH': '/root', '_F_IMG_ID': '{image_name}'}}, extra_path=['/root'])
     config.pip_index(url="https://private-pip-index/simple")
     install.python(version="3.8")
 """