From c35982cba5903d6156e844e63a09ff6a0dfd28c6 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Wed, 7 Jun 2023 10:40:38 -0700 Subject: [PATCH 01/10] wip Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 4 ++++ plugins/flytekit-envd/flytekitplugins/envd/image_builder.py | 1 + 2 files changed, 5 insertions(+) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 23e33cbd9c..1759a2e1f2 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -31,6 +31,8 @@ class ImageSpec: registry: registry of the image. packages: list of python packages to install. apt_packages: list of apt packages to install. + cuda: version of cuda to install. + cudnn: version of cudnn to install. base_image: base image of the image. platform: Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64 """ @@ -43,6 +45,8 @@ class ImageSpec: registry: Optional[str] = None packages: Optional[List[str]] = None apt_packages: Optional[List[str]] = None + cuda: Optional[str] = None + cudnn: Optional[str] = None base_image: Optional[str] = None platform: str = "linux/amd64" diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 2e107b7848..73381d95d3 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -48,6 +48,7 @@ def build(): base(image="{base_image}", dev=False) 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)))}]) + install.cuda(version="{image_spec.cuda}", cudnn="{image_spec.cudnn}") runtime.environ(env={env}) """ From ab3999aa6bf9efcdd1c5e027f4ec75b72ae9e196 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 9 Jun 2023 23:38:40 -0700 Subject: [PATCH 02/10] init Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 1759a2e1f2..21efc10aba 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -37,12 +37,12 @@ class ImageSpec: platform: Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64 """ + registry: str name: str = "flytekit" python_version: str = None # Use default python in the base image if None. builder: str = "envd" source_root: Optional[str] = None env: Optional[typing.Dict[str, str]] = None - registry: Optional[str] = None packages: Optional[List[str]] = None apt_packages: Optional[List[str]] = None cuda: Optional[str] = None From 7cacde840618ae0a32f9290e5ecb8f13635874af Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Fri, 9 Jun 2023 23:53:00 -0700 Subject: [PATCH 03/10] remove dataclass jsono Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 21efc10aba..abd7588e18 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -4,19 +4,17 @@ import typing from abc import abstractmethod from copy import copy -from dataclasses import dataclass +from dataclasses import asdict, dataclass from functools import lru_cache from typing import List, Optional import click import requests -from dataclasses_json import dataclass_json DOCKER_HUB = "docker.io" _F_IMG_ID = "_F_IMG_ID" -@dataclass_json @dataclass class ImageSpec: """ @@ -153,7 +151,7 @@ def calculate_hash_from_image_spec(image_spec: ImageSpec): # copy the image spec to avoid modifying the original image spec. otherwise, the hash will be different. spec = copy(image_spec) spec.source_root = hash_directory(image_spec.source_root) if image_spec.source_root else b"" - image_spec_bytes = bytes(spec.to_json(), "utf-8") + image_spec_bytes = asdict(spec).__str__().encode("utf-8") tag = base64.urlsafe_b64encode(hashlib.md5(image_spec_bytes).digest()).decode("ascii") # replace "=" with "." to make it a valid tag return tag.replace("=", ".") From 8c7c43ad310ec74b4e8c06c2b4604b6cb9692f86 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Sat, 10 Jun 2023 11:03:11 -0700 Subject: [PATCH 04/10] nit Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 2 +- .../flytekitplugins/envd/image_builder.py | 11 ++++++----- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index abd7588e18..94f6b57134 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -106,7 +106,7 @@ def exist(self) -> bool: return True def __hash__(self): - return hash(self.to_json()) + return hash(asdict(self).__str__()) class ImageSpecBuilder: diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 73381d95d3..d09cd036cd 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -35,7 +35,7 @@ def build_image(self, image_spec: ImageSpec): def create_envd_config(image_spec: ImageSpec) -> str: - base_image = DefaultImages.default_image() if image_spec.base_image is None else image_spec.base_image + base_image = "ubuntu20.04" if image_spec.base_image is None else image_spec.base_image packages = [] if image_spec.packages is None else image_spec.packages apt_packages = [] if image_spec.apt_packages is None else image_spec.apt_packages env = {"PYTHONPATH": "/root", _F_IMG_ID: image_spec.image_name()} @@ -48,17 +48,18 @@ def build(): base(image="{base_image}", dev=False) 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)))}]) - install.cuda(version="{image_spec.cuda}", cudnn="{image_spec.cudnn}") runtime.environ(env={env}) """ + ctx = context_manager.FlyteContextManager.current_context() + cfg_path = ctx.file_access.get_random_local_path("build.envd") + pathlib.Path(cfg_path).parent.mkdir(parents=True, exist_ok=True) if image_spec.python_version: # Indentation is required by envd envd_config += f' install.python(version="{image_spec.python_version}")\n' - ctx = context_manager.FlyteContextManager.current_context() - cfg_path = ctx.file_access.get_random_local_path("build.envd") - pathlib.Path(cfg_path).parent.mkdir(parents=True, exist_ok=True) + if image_spec.cuda and image_spec.cudnn: + envd_config += f' install.cuda(version="{image_spec.cuda}", cudnn="{image_spec.cudnn}")\n' if image_spec.source_root: shutil.copytree(image_spec.source_root, pathlib.Path(cfg_path).parent, dirs_exist_ok=True) From 05914e7b0fb24a23be63f328d855c9e88dd44f66 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 13:36:42 -0700 Subject: [PATCH 05/10] update Signed-off-by: Kevin Su --- .../flytekitplugins/envd/image_builder.py | 7 ++++++- .../unit/core/image_spec/test_image_spec.py | 13 +++++++++++-- .../flytekit/unit/core/test_python_function_task.py | 6 ++---- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index d09cd036cd..e840d88e7c 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -35,7 +35,12 @@ def build_image(self, image_spec: ImageSpec): def create_envd_config(image_spec: ImageSpec) -> str: - base_image = "ubuntu20.04" if image_spec.base_image is None else image_spec.base_image + base_image = DefaultImages.default_image() if image_spec.base_image is None else image_spec.base_image + if image_spec.cuda and image_spec.cudnn: + if image_spec.python_version is None: + raise Exception("python_version is required when cuda and cudnn are specified") + base_image = "ubuntu20.04" + packages = [] if image_spec.packages is None else image_spec.packages apt_packages = [] if image_spec.apt_packages is None else image_spec.apt_packages env = {"PYTHONPATH": "/root", _F_IMG_ID: image_spec.image_name()} diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index 46b5cb0244..dbd808a32b 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -1,4 +1,7 @@ +import base64 +import hashlib import os +from dataclasses import asdict import pytest @@ -27,7 +30,12 @@ def test_image_spec(): assert image_spec.source_root is None assert image_spec.env is None assert image_spec.is_container() is True - assert image_spec.image_name() == "flytekit:OLFSrRjcG5_uXuRqd0TSdQ.." + + image_spec.source_root = b"" + image_spec_bytes = asdict(image_spec).__str__().encode("utf-8") + tag = base64.urlsafe_b64encode(hashlib.md5(image_spec_bytes).digest()).decode("ascii") + tag = tag.replace("=", ".") + assert image_spec.image_name() == f"flytekit:{tag}" ctx = context_manager.FlyteContext.current_context() with context_manager.FlyteContextManager.with_context( ctx.with_execution_state(ctx.execution_state.with_params(mode=ExecutionState.Mode.TASK_EXECUTION)) @@ -41,8 +49,9 @@ def build_image(self, img): ImageBuildEngine.register("dummy", DummyImageSpecBuilder()) ImageBuildEngine._REGISTRY["dummy"].build_image(image_spec) + assert "dummy" in ImageBuildEngine._REGISTRY - assert calculate_hash_from_image_spec(image_spec) == "OLFSrRjcG5_uXuRqd0TSdQ.." + assert calculate_hash_from_image_spec(image_spec) == tag assert image_spec.exist() is False with pytest.raises(Exception): diff --git a/tests/flytekit/unit/core/test_python_function_task.py b/tests/flytekit/unit/core/test_python_function_task.py index 48d8e2487e..f1c6609e22 100644 --- a/tests/flytekit/unit/core/test_python_function_task.py +++ b/tests/flytekit/unit/core/test_python_function_task.py @@ -73,10 +73,8 @@ def build_image(self, img): ... ImageBuildEngine.register("test", TestImageSpecBuilder()) - assert ( - get_registerable_container_image(ImageSpec(builder="test", python_version="3.7"), cfg) - == "flytekit:usEl-DNx_srVn7zp2vHlBw.." - ) + image_spec = ImageSpec(builder="test", python_version="3.7", registry="") + assert get_registerable_container_image(image_spec, cfg) == image_spec.image_name() def test_get_registerable_container_image_no_images(): From 6313abc8ab11589acc135e4709a505252b140e7c Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 13:42:14 -0700 Subject: [PATCH 06/10] nit Signed-off-by: Kevin Su --- tests/flytekit/unit/core/image_spec/test_image_spec.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/flytekit/unit/core/image_spec/test_image_spec.py b/tests/flytekit/unit/core/image_spec/test_image_spec.py index dbd808a32b..2895427b44 100644 --- a/tests/flytekit/unit/core/image_spec/test_image_spec.py +++ b/tests/flytekit/unit/core/image_spec/test_image_spec.py @@ -18,6 +18,8 @@ def test_image_spec(): python_version="3.8", registry="", base_image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", + cuda="11.2.2", + cudnn="8", ) assert image_spec.python_version == "3.8" @@ -25,6 +27,8 @@ def test_image_spec(): assert image_spec.packages == ["pandas"] assert image_spec.apt_packages == ["git"] assert image_spec.registry == "" + assert image_spec.cuda == "11.2.2" + assert image_spec.cudnn == "8" assert image_spec.name == "flytekit" assert image_spec.builder == "envd" assert image_spec.source_root is None From 80ba0a917ca3683fde5ad361c1178f2f39b642ae Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 14:02:13 -0700 Subject: [PATCH 07/10] Fix tests Signed-off-by: Kevin Su --- plugins/flytekit-envd/tests/test_image_spec.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 603bc8979b..47ace4b2c0 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -18,16 +18,17 @@ def test_image_spec(): EnvdImageSpecBuilder().build_image(image_spec) config_path = create_envd_config(image_spec) assert image_spec.platform == "linux/amd64" + image_name = image_spec.image_name() contents = Path(config_path).read_text() assert ( contents - == """# syntax=v1 + == f"""# syntax=v1 def build(): base(image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", dev=False) install.python_packages(name = ["pandas"]) install.apt_packages(name = ["git"]) - runtime.environ(env={'PYTHONPATH': '/root', '_F_IMG_ID': 'flytekit:46qVNvYHJxppEvVIYrthdA..'}) + runtime.environ(env={'PYTHONPATH': '/root', '_F_IMG_ID': {image_name}}) config.pip_index(url = "https://private-pip-index/simple") install.python(version="3.8") """ From 5948744cf4baff7c50566995bb25f3c89b973684 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 14:45:16 -0700 Subject: [PATCH 08/10] Fix tests Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 2 +- tests/flytekit/unit/cli/pyflyte/test_run.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index a7ba45e06a..747f79b166 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -36,12 +36,12 @@ class ImageSpec: pip_index: Specify the custom pip index url """ - registry: str name: str = "flytekit" python_version: str = None # Use default python in the base image if None. builder: str = "envd" source_root: Optional[str] = None env: Optional[typing.Dict[str, str]] = None + registry: Optional[List[str]] = None packages: Optional[List[str]] = None apt_packages: Optional[List[str]] = None cuda: Optional[str] = None diff --git a/tests/flytekit/unit/cli/pyflyte/test_run.py b/tests/flytekit/unit/cli/pyflyte/test_run.py index 8ebd56c643..5585559c51 100644 --- a/tests/flytekit/unit/cli/pyflyte/test_run.py +++ b/tests/flytekit/unit/cli/pyflyte/test_run.py @@ -269,9 +269,9 @@ def test_list_default_arguments(wf_path): ) ic_result_4 = ImageConfig( - default_image=Image(name="default", fqn="flytekit", tag="6y6c8ofS_Pwa2FImlcm3Qg.."), + default_image=Image(name="default", fqn="flytekit", tag="eJgTB5QCJDOSksy6gE0lXA.."), images=[ - Image(name="default", fqn="flytekit", tag="6y6c8ofS_Pwa2FImlcm3Qg.."), + Image(name="default", fqn="flytekit", tag="eJgTB5QCJDOSksy6gE0lXA.."), Image(name="xyz", fqn="docker.io/xyz", tag="latest"), Image(name="abc", fqn="docker.io/abc", tag=None), ], From 67b3348653a659401b99363a2b0233ab0bbf6de8 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 15:46:07 -0700 Subject: [PATCH 09/10] update cudnn Signed-off-by: Kevin Su --- flytekit/image_spec/image_spec.py | 2 +- .../flytekit-envd/flytekitplugins/envd/image_builder.py | 7 ++++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 747f79b166..870d4a9df0 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -41,7 +41,7 @@ class ImageSpec: builder: str = "envd" source_root: Optional[str] = None env: Optional[typing.Dict[str, str]] = None - registry: Optional[List[str]] = None + registry: Optional[str] = None packages: Optional[List[str]] = None apt_packages: Optional[List[str]] = None cuda: Optional[str] = None diff --git a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py index 5b5a9c01dc..42ea6b6e5d 100644 --- a/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py +++ b/plugins/flytekit-envd/flytekitplugins/envd/image_builder.py @@ -36,7 +36,7 @@ def build_image(self, image_spec: ImageSpec): def create_envd_config(image_spec: ImageSpec) -> str: base_image = DefaultImages.default_image() if image_spec.base_image is None else image_spec.base_image - if image_spec.cuda and image_spec.cudnn: + if image_spec.cuda: if image_spec.python_version is None: raise Exception("python_version is required when cuda and cudnn are specified") base_image = "ubuntu20.04" @@ -65,8 +65,9 @@ def build(): # Indentation is required by envd envd_config += f' install.python(version="{image_spec.python_version}")\n' - if image_spec.cuda and image_spec.cudnn: - envd_config += f' install.cuda(version="{image_spec.cuda}", cudnn="{image_spec.cudnn}")\n' + if image_spec.cuda: + cudnn = image_spec.cudnn if image_spec.cudnn else "" + envd_config += f' install.cuda(version="{image_spec.cuda}", cudnn="{cudnn}")\n' if image_spec.source_root: shutil.copytree(image_spec.source_root, pathlib.Path(cfg_path).parent, dirs_exist_ok=True) From 93c8bbc18f4489fa189dae649cd6b6f48a5d3056 Mon Sep 17 00:00:00 2001 From: Kevin Su Date: Thu, 15 Jun 2023 23:31:11 -0700 Subject: [PATCH 10/10] nit Signed-off-by: Kevin Su --- plugins/flytekit-envd/tests/test_image_spec.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/plugins/flytekit-envd/tests/test_image_spec.py b/plugins/flytekit-envd/tests/test_image_spec.py index 47ace4b2c0..e62dfb2468 100644 --- a/plugins/flytekit-envd/tests/test_image_spec.py +++ b/plugins/flytekit-envd/tests/test_image_spec.py @@ -10,7 +10,6 @@ def test_image_spec(): packages=["pandas"], apt_packages=["git"], python_version="3.8", - registry="", base_image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", pip_index="https://private-pip-index/simple", ) @@ -28,7 +27,7 @@ def build(): base(image="cr.flyte.org/flyteorg/flytekit:py3.8-latest", dev=False) 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}'}}) config.pip_index(url = "https://private-pip-index/simple") install.python(version="3.8") """