Multi arch imageSpec (#1630)

Multi arch imageSpec (#1630)

flytekit/image_spec/image_spec.py

@@ -32,6 +32,7 @@ class ImageSpec:
         packages: list of python packages to install.
         apt_packages: list of apt packages 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
     """
 
     name: str = "flytekit"
@@ -43,6 +44,7 @@ class ImageSpec:
     packages: Optional[List[str]] = None
     apt_packages: Optional[List[str]] = None
     base_image: Optional[str] = None
+    platform: str = "linux/amd64"
 
     def image_name(self) -> str:
         """
@@ -147,7 +149,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(image_spec.to_json(), "utf-8")
+    image_spec_bytes = bytes(spec.to_json(), "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("=", ".")

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

@@ -16,7 +16,7 @@ class EnvdImageSpecBuilder(ImageSpecBuilder):
 
     def build_image(self, image_spec: ImageSpec):
         cfg_path = create_envd_config(image_spec)
-        command = f"envd build --path {pathlib.Path(cfg_path).parent}"
+        command = f"envd build --path {pathlib.Path(cfg_path).parent}  --platform {image_spec.platform}"
         if image_spec.registry:
             command += f" --output type=image,name={image_spec.image_name()},push=true"
         click.secho(f"Run command: {command} ", fg="blue")

plugins/flytekit-envd/requirements.txt

@@ -6,22 +6,57 @@
 #
 -e file:.#egg=flytekitplugins-envd
     # via -r requirements.in
+adlfs==2023.4.0
+    # via flytekit
+aiobotocore==2.5.0
+    # via s3fs
+aiohttp==3.8.4
+    # via
+    #   adlfs
+    #   aiobotocore
+    #   gcsfs
+    #   s3fs
+aioitertools==0.11.0
+    # via aiobotocore
+aiosignal==1.3.1
+    # via aiohttp
 arrow==1.2.3
     # via jinja2-time
+async-timeout==4.0.2
+    # via aiohttp
+attrs==23.1.0
+    # via aiohttp
+azure-core==1.26.4
+    # via
+    #   adlfs
+    #   azure-identity
+    #   azure-storage-blob
+azure-datalake-store==0.0.53
+    # via adlfs
+azure-identity==1.13.0
+    # via adlfs
+azure-storage-blob==12.16.0
+    # via adlfs
 binaryornot==0.4.4
     # via cookiecutter
+botocore==1.29.76
+    # via aiobotocore
 cachetools==5.3.0
     # via google-auth
 certifi==2022.12.7
     # via
     #   kubernetes
     #   requests
 cffi==1.15.1
-    # via cryptography
+    # via
+    #   azure-datalake-store
+    #   cryptography
 chardet==5.1.0
     # via binaryornot
 charset-normalizer==3.1.0
-    # via requests
+    # via
+    #   aiohttp
+    #   requests
 click==8.1.3
     # via
     #   cookiecutter
@@ -33,11 +68,16 @@ cookiecutter==2.1.1
 croniter==1.3.8
     # via flytekit
 cryptography==40.0.1
-    # via pyopenssl
+    # via
+    #   azure-identity
+    #   azure-storage-blob
+    #   msal
+    #   pyjwt
+    #   pyopenssl
 dataclasses-json==0.5.7
     # via flytekit
 decorator==5.1.1
-    # via retry
+    # via gcsfs
 deprecated==1.2.13
     # via flytekit
 diskcache==5.4.0
@@ -48,22 +88,55 @@ docker-image-py==0.1.12
     # via flytekit
 docstring-parser==0.15
     # via flytekit
-envd==0.3.16
+envd==0.3.22
     # via flytekitplugins-envd
 flyteidl==1.3.15
     # via flytekit
 flytekit==1.5.0
     # via flytekitplugins-envd
+frozenlist==1.3.3
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2023.5.0
+    # via
+    #   adlfs
+    #   flytekit
+    #   gcsfs
+    #   s3fs
+gcsfs==2023.5.0
+    # via flytekit
 gitdb==4.0.10
     # via gitpython
 gitpython==3.1.31
     # via flytekit
+google-api-core==2.11.0
+    # via
+    #   google-cloud-core
+    #   google-cloud-storage
 google-auth==2.17.1
-    # via kubernetes
+    # via
+    #   gcsfs
+    #   google-api-core
+    #   google-auth-oauthlib
+    #   google-cloud-core
+    #   google-cloud-storage
+    #   kubernetes
+google-auth-oauthlib==1.0.0
+    # via gcsfs
+google-cloud-core==2.3.2
+    # via google-cloud-storage
+google-cloud-storage==2.9.0
+    # via gcsfs
+google-crc32c==1.5.0
+    # via google-resumable-media
+google-resumable-media==2.5.0
+    # via google-cloud-storage
 googleapis-common-protos==1.59.0
     # via
     #   flyteidl
     #   flytekit
+    #   google-api-core
     #   grpcio-status
 grpcio==1.53.0
     # via
@@ -72,11 +145,15 @@ grpcio==1.53.0
 grpcio-status==1.53.0
     # via flytekit
 idna==3.4
-    # via requests
+    # via
+    #   requests
+    #   yarl
 importlib-metadata==6.1.0
     # via
     #   flytekit
     #   keyring
+isodate==0.6.1
+    # via azure-storage-blob
 jaraco-classes==3.2.3
     # via keyring
 jinja2==3.1.2
@@ -85,6 +162,8 @@ jinja2==3.1.2
     #   jinja2-time
 jinja2-time==0.2.0
     # via cookiecutter
+jmespath==1.0.1
+    # via botocore
 joblib==1.2.0
     # via flytekit
 keyring==23.13.1
@@ -104,6 +183,17 @@ marshmallow-jsonschema==0.13.0
     # via flytekit
 more-itertools==9.1.0
     # via jaraco-classes
+msal==1.22.0
+    # via
+    #   azure-datalake-store
+    #   azure-identity
+    #   msal-extensions
+msal-extensions==1.0.0
+    # via azure-identity
+multidict==6.0.4
+    # via
+    #   aiohttp
+    #   yarl
 mypy-extensions==1.0.0
     # via typing-inspect
 natsort==8.3.1
@@ -121,16 +211,17 @@ packaging==23.0
     #   marshmallow
 pandas==1.5.3
     # via flytekit
+portalocker==2.7.0
+    # via msal-extensions
 protobuf==4.22.1
     # via
     #   flyteidl
+    #   google-api-core
     #   googleapis-common-protos
     #   grpcio-status
     #   protoc-gen-swagger
 protoc-gen-swagger==0.1.0
     # via flyteidl
-py==1.11.0
-    # via retry
 pyarrow==10.0.1
     # via flytekit
 pyasn1==0.4.8
@@ -141,11 +232,14 @@ pyasn1-modules==0.2.8
     # via google-auth
 pycparser==2.21
     # via cffi
+pyjwt[crypto]==2.7.0
+    # via msal
 pyopenssl==23.1.1
     # via flytekit
 python-dateutil==2.8.2
     # via
     #   arrow
+    #   botocore
     #   croniter
     #   flytekit
     #   kubernetes
@@ -170,23 +264,34 @@ regex==2023.3.23
     # via docker-image-py
 requests==2.28.2
     # via
+    #   azure-core
+    #   azure-datalake-store
     #   cookiecutter
     #   docker
     #   flytekit
+    #   gcsfs
+    #   google-api-core
+    #   google-cloud-storage
     #   kubernetes
+    #   msal
     #   requests-oauthlib
     #   responses
 requests-oauthlib==1.3.1
-    # via kubernetes
+    # via
+    #   google-auth-oauthlib
+    #   kubernetes
 responses==0.23.1
     # via flytekit
-retry==0.9.2
-    # via flytekit
 rsa==4.9
     # via google-auth
+s3fs==2023.5.0
+    # via flytekit
 six==1.16.0
     # via
+    #   azure-core
+    #   azure-identity
     #   google-auth
+    #   isodate
     #   kubernetes
     #   python-dateutil
 smmap==5.0.0
@@ -201,12 +306,16 @@ types-pyyaml==6.0.12.9
     # via responses
 typing-extensions==4.5.0
     # via
+    #   aioitertools
+    #   azure-core
+    #   azure-storage-blob
     #   flytekit
     #   typing-inspect
 typing-inspect==0.8.0
     # via dataclasses-json
 urllib3==1.26.15
     # via
+    #   botocore
     #   docker
     #   flytekit
     #   kubernetes
@@ -220,8 +329,11 @@ wheel==0.40.0
     # via flytekit
 wrapt==1.15.0
     # via
+    #   aiobotocore
     #   deprecated
     #   flytekit
+yarl==1.9.2
+    # via aiohttp
 zipp==3.15.0
     # via importlib-metadata
 

plugins/flytekit-envd/setup.py

@@ -4,7 +4,7 @@
 
 microlib_name = f"flytekitplugins-{PLUGIN_NAME}"
 
-plugin_requires = ["flytekit", "envd"]
+plugin_requires = ["flytekit", "envd>=0.3.22"]
 
 __version__ = "0.0.0+develop"
 

plugins/flytekit-envd/tests/test_image_spec.py

@@ -16,6 +16,7 @@ def test_image_spec():
 
     EnvdImageSpecBuilder().build_image(image_spec)
     config_path = create_envd_config(image_spec)
+    assert image_spec.platform == "linux/amd64"
     contents = Path(config_path).read_text()
     assert (
         contents
@@ -25,7 +26,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': 'flytekit:yZ8jICcDTLoDArmNHbWNwg..'})
+    runtime.environ(env={'PYTHONPATH': '/root', '_F_IMG_ID': 'flytekit:OLFSrRjcG5_uXuRqd0TSdQ..'})
     install.python(version="3.8")
 """
     )

tests/flytekit/unit/cli/pyflyte/test_run.py

@@ -261,9 +261,9 @@ def test_list_default_arguments(wf_path):
 )
 
 ic_result_4 = ImageConfig(
-    default_image=Image(name="default", fqn="flytekit", tag="4VC-c-UDrUvfySJ0aS3qCw.."),
+    default_image=Image(name="default", fqn="flytekit", tag="mMxGzKCqxVk8msz0yV22-g.."),
     images=[
-        Image(name="default", fqn="flytekit", tag="4VC-c-UDrUvfySJ0aS3qCw.."),
+        Image(name="default", fqn="flytekit", tag="mMxGzKCqxVk8msz0yV22-g.."),
         Image(name="xyz", fqn="docker.io/xyz", tag="latest"),
         Image(name="abc", fqn="docker.io/abc", tag=None),
     ],

tests/flytekit/unit/core/image_spec/test_image_spec.py

@@ -27,7 +27,7 @@ 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:yZ8jICcDTLoDArmNHbWNwg.."
+    assert image_spec.image_name() == "flytekit:OLFSrRjcG5_uXuRqd0TSdQ.."
     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))
@@ -42,7 +42,7 @@ 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) == "yZ8jICcDTLoDArmNHbWNwg.."
+    assert calculate_hash_from_image_spec(image_spec) == "OLFSrRjcG5_uXuRqd0TSdQ.."
     assert image_spec.exist() is False
 
     with pytest.raises(Exception):

tests/flytekit/unit/core/test_python_function_task.py

@@ -75,7 +75,7 @@ def build_image(self, img):
     ImageBuildEngine.register("test", TestImageSpecBuilder())
     assert (
         get_registerable_container_image(ImageSpec(builder="test", python_version="3.7"), cfg)
-        == "flytekit:0N8X-XowtpEkDYWDlb8Abg.."
+        == "flytekit:usEl-DNx_srVn7zp2vHlBw.."
     )