From 38ecf5cc0fb607cf05e4252be6bd64c405f745c9 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Fri, 15 Aug 2025 13:19:47 +0800 Subject: [PATCH 1/4] Dynamic platform detection for ImageSpec to support cross-platform development Signed-off-by: Barry Wu --- flytekit/image_spec/image_spec.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 8a8a599d03..4f348c2b34 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -4,6 +4,7 @@ import hashlib import os import pathlib +import platform import re import sys import typing @@ -136,6 +137,9 @@ def __post_init__(self): # Use the builder with the highest priority by default self.builder = max(builder_registry, key=lambda name: builder_registry[name][1]) + if platform.machine() == "arm64": + self.platform = "linux/arm64" + parameters_str_list = [ "packages", "conda_channels", From 052706bf00a658a7c56a6eb17bec5aafea3f4df3 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Sun, 17 Aug 2025 22:54:53 +0800 Subject: [PATCH 2/4] Change platform to ARM64 if platform is not set and pushing to local registry on ARM64 machine Signed-off-by: Barry Wu --- flytekit/image_spec/image_spec.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/flytekit/image_spec/image_spec.py b/flytekit/image_spec/image_spec.py index 4f348c2b34..482fa007a2 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -95,7 +95,7 @@ class ImageSpec: cuda: Optional[str] = None cudnn: Optional[str] = None base_image: Optional[Union[str, "ImageSpec"]] = None - platform: str = "linux/amd64" + platform: str = None pip_index: Optional[str] = None pip_extra_index_url: Optional[List[str]] = None pip_secret_mounts: Optional[List[Tuple[str, str]]] = None @@ -127,6 +127,17 @@ def __post_init__(self): f"- 'docker.io/username' (for docker hub)\n" ) + if self.platform is None: + if ( + self.registry + and self.registry.lower().startswith("localhost:") + and platform.machine().lower() in ("arm64", "aarch64") + ): + # Only change platform to ARM64 if platform is not set and pushing to local registry on ARM64 machine + self.platform = "linux/arm64" + else: + self.platform = "linux/amd64" + # If not set, help the user set this option as well, to support the older default behavior where existence # of the source root implied that copying of files was needed. if self.source_root is not None: @@ -137,9 +148,6 @@ def __post_init__(self): # Use the builder with the highest priority by default self.builder = max(builder_registry, key=lambda name: builder_registry[name][1]) - if platform.machine() == "arm64": - self.platform = "linux/arm64" - parameters_str_list = [ "packages", "conda_channels", From 81e8df7e24663f5f90d4f2095492c2087e90a07a Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Mon, 18 Aug 2025 23:26:04 +0800 Subject: [PATCH 3/4] Add Optional for type checking Signed-off-by: Barry Wu --- 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 482fa007a2..ace620aab4 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -95,7 +95,7 @@ class ImageSpec: cuda: Optional[str] = None cudnn: Optional[str] = None base_image: Optional[Union[str, "ImageSpec"]] = None - platform: str = None + platform: Optional[str] = None pip_index: Optional[str] = None pip_extra_index_url: Optional[List[str]] = None pip_secret_mounts: Optional[List[Tuple[str, str]]] = None From 068c872fca84707b7523084d218ef3bae08e9142 Mon Sep 17 00:00:00 2001 From: Barry Wu Date: Tue, 19 Aug 2025 12:47:35 +0800 Subject: [PATCH 4/4] Fix lint Signed-off-by: Barry Wu --- 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 ace620aab4..357f266d1e 100644 --- a/flytekit/image_spec/image_spec.py +++ b/flytekit/image_spec/image_spec.py @@ -49,7 +49,7 @@ class ImageSpec: cuda (Optional[str]): Version of cuda to install. cudnn (Optional[str]): Version of cudnn to install. base_image (Optional[Union[str, 'ImageSpec']]): Base image of the image. - platform (str): Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64). + platform (Optional[str]): Specify the target platforms for the build output (for example, windows/amd64 or linux/amd64,darwin/arm64). pip_index (Optional[str]): Specify the custom pip index url. pip_extra_index_url (Optional[List[str]]): Specify one or more pip index urls as a list. pip_secret_mounts (Optional[List[Tuple[str, str]]]): Specify a list of tuples to mount secret for pip install. Each tuple should contain the path to