[Fix] Added support to specify the platform on which the runtime image should be built. (#4402)

Signed-off-by: dependabot[bot] <support@github.com>
Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com>
Co-authored-by: Xingyao Wang <xingyao@all-hands.dev>
Co-authored-by: mamoodi <mamoodiha@gmail.com>
Co-authored-by: tofarr <tofarr@gmail.com>
Co-authored-by: Robert Brennan <contact@rbren.io>

evaluation/swe_bench/run_infer.py

@@ -132,6 +132,8 @@ def get_config(
             use_host_network=False,
             # large enough timeout, since some testcases take very long to run
             timeout=300,
+            # Add platform to the sandbox config to solve issue 4401
+            platform='linux/amd64',
             api_key=os.environ.get('ALLHANDS_API_KEY', None),
             remote_runtime_api_url=os.environ.get('SANDBOX_REMOTE_RUNTIME_API_URL'),
             keep_remote_runtime_alive=False,

openhands/core/config/sandbox_config.py

@@ -30,6 +30,7 @@ class SandboxConfig:
             For example, for specifying the base url of website for browsergym evaluation.
         browsergym_eval_env: The BrowserGym environment to use for evaluation.
             Default is None for general purpose browsing. Check evaluation/miniwob and evaluation/webarena for examples.
+        platform: The platform on which the image should be built. Default is None.
     """
 
     remote_runtime_api_url: str = 'http://localhost:8000'
@@ -49,6 +50,7 @@ class SandboxConfig:
     runtime_extra_deps: str | None = None
     runtime_startup_env_vars: dict[str, str] = field(default_factory=dict)
     browsergym_eval_env: str | None = None
+    platform: str | None = None
 
     def defaults_to_dict(self) -> dict:
         """Serialize fields to a dict for the frontend, including type hints, defaults, and whether it's optional."""

openhands/runtime/builder/base.py

@@ -7,6 +7,7 @@ class RuntimeBuilder(abc.ABC):
         self,
         path: str,
         tags: list[str],
+        platform: str | None = None,
     ) -> str:
         """
         Build the runtime image.
@@ -14,7 +15,7 @@ class RuntimeBuilder(abc.ABC):
         Args:
             path (str): The path to the runtime image's build directory.
             tags (list[str]): The tags to apply to the runtime image (e.g., ["repo:my-repo", "sha:my-sha"]).
-
+            platform (str, optional): The target platform for the build. Defaults to None.
         Returns:
             str: The name:tag of the runtime image after build (e.g., "repo:sha").
                 This can be different from the tags input if the builder chooses to mutate the tags (e.g., adding a

openhands/runtime/builder/docker.py

@@ -27,6 +27,7 @@ class DockerRuntimeBuilder(RuntimeBuilder):
         self,
         path: str,
         tags: list[str],
+        platform: str | None = None,
         use_local_cache: bool = False,
         extra_build_args: list[str] | None = None,
     ) -> str:
@@ -35,6 +36,7 @@ class DockerRuntimeBuilder(RuntimeBuilder):
         Args:
             path (str): The path to the Docker build context.
             tags (list[str]): A list of image tags to apply to the built image.
+            platform (str, optional): The target platform for the build. Defaults to None.
             use_local_cache (bool, optional): Whether to use and update the local build cache. Defaults to True.
             extra_build_args (list[str], optional): Additional arguments to pass to the Docker build command. Defaults to None.
 
@@ -70,6 +72,10 @@ class DockerRuntimeBuilder(RuntimeBuilder):
             '--load',
         ]
 
+        # Include the platform argument only if platform is specified
+        if platform:
+            buildx_cmd.append(f'--platform={platform}')
+
         cache_dir = '/tmp/.buildx-cache'
         if use_local_cache and self._is_cache_usable(cache_dir):
             buildx_cmd.extend(

openhands/runtime/builder/remote.py

@@ -23,7 +23,7 @@ class RemoteRuntimeBuilder(RuntimeBuilder):
         self.session = requests.Session()
         self.session.headers.update({'X-API-Key': self.api_key})
 
-    def build(self, path: str, tags: list[str]) -> str:
+    def build(self, path: str, tags: list[str], platform: str | None = None) -> str:
         """Builds a Docker image using the Runtime API's /build endpoint."""
         # Create a tar archive of the build context
         tar_buffer = io.BytesIO()

openhands/runtime/client/runtime.py

@@ -183,6 +183,7 @@ class EventStreamRuntime(Runtime):
             self.runtime_container_image = build_runtime_image(
                 self.base_container_image,
                 self.runtime_builder,
+                platform=self.config.sandbox.platform,
                 extra_deps=self.config.sandbox.runtime_extra_deps,
                 force_rebuild=self.config.sandbox.force_rebuild_runtime,
             )

openhands/runtime/remote/runtime.py

@@ -181,6 +181,7 @@ class RemoteRuntime(Runtime):
         self.container_image = build_runtime_image(
             self.config.sandbox.base_container_image,
             self.runtime_builder,
+            platform=self.config.sandbox.platform,
             extra_deps=self.config.sandbox.runtime_extra_deps,
             force_rebuild=self.config.sandbox.force_rebuild_runtime,
         )

openhands/runtime/utils/runtime_build.py

@@ -98,6 +98,7 @@ def get_runtime_image_repo_and_tag(base_image: str) -> tuple[str, str]:
 def build_runtime_image(
     base_image: str,
     runtime_builder: RuntimeBuilder,
+    platform: str | None = None,
     extra_deps: str | None = None,
     build_folder: str | None = None,
     dry_run: bool = False,
@@ -109,6 +110,7 @@ def build_runtime_image(
     Parameters:
     - base_image (str): The name of the base Docker image to use
     - runtime_builder (RuntimeBuilder): The runtime builder to use
+    - platform (str): The target platform for the build (e.g. linux/amd64, linux/arm64)
     - extra_deps (str):
     - build_folder (str): The directory to use for the build. If not provided a temporary directory will be used
     - dry_run (bool): if True, it will only ready the build folder. It will not actually build the Docker image
@@ -128,6 +130,7 @@ def build_runtime_image(
                 extra_deps=extra_deps,
                 dry_run=dry_run,
                 force_rebuild=force_rebuild,
+                platform=platform,
             )
             return result
 
@@ -138,6 +141,7 @@ def build_runtime_image(
         extra_deps=extra_deps,
         dry_run=dry_run,
         force_rebuild=force_rebuild,
+        platform=platform,
     )
     return result
 
@@ -149,6 +153,7 @@ def build_runtime_image_in_folder(
     extra_deps: str | None,
     dry_run: bool,
     force_rebuild: bool,
+    platform: str | None = None,
 ) -> str:
     runtime_image_repo, _ = get_runtime_image_repo_and_tag(base_image)
     lock_tag = f'oh_v{oh_version}_{get_hash_for_lock_files(base_image)}'
@@ -165,6 +170,7 @@ def build_runtime_image_in_folder(
                 runtime_image_repo,
                 hash_tag,
                 lock_tag,
+                platform,
             )
         return hash_image_name
 
@@ -194,6 +200,7 @@ def build_runtime_image_in_folder(
             runtime_image_repo,
             hash_tag,
             lock_tag,
+            platform,
         )
 
     return hash_image_name
@@ -293,6 +300,7 @@ def _build_sandbox_image(
     runtime_image_repo: str,
     hash_tag: str,
     lock_tag: str,
+    platform: str | None = None,
 ):
     """Build and tag the sandbox image. The image will be tagged with all tags that do not yet exist"""
 
@@ -305,7 +313,9 @@ def _build_sandbox_image(
         if not runtime_builder.image_exists(name, False)
     ]
 
-    image_name = runtime_builder.build(path=str(build_folder), tags=names)
+    image_name = runtime_builder.build(
+        path=str(build_folder), tags=names, platform=platform
+    )
     if not image_name:
         raise RuntimeError(f'Build failed for image {names}')
 
@@ -319,6 +329,7 @@ if __name__ == '__main__':
     )
     parser.add_argument('--build_folder', type=str, default=None)
     parser.add_argument('--force_rebuild', action='store_true', default=False)
+    parser.add_argument('--platform', type=str, default=None)
     args = parser.parse_args()
 
     if args.build_folder is not None:
@@ -349,6 +360,7 @@ if __name__ == '__main__':
                 build_folder=temp_dir,
                 dry_run=True,
                 force_rebuild=args.force_rebuild,
+                platform=args.platform,
             )
 
             _runtime_image_repo, runtime_image_hash_tag = runtime_image_hash_name.split(
@@ -382,5 +394,7 @@ if __name__ == '__main__':
         # Dockerfile, we actually build the Docker image
         logger.info('Building image in a temporary folder')
         docker_builder = DockerRuntimeBuilder(docker.from_env())
-        image_name = build_runtime_image(args.base_image, docker_builder)
+        image_name = build_runtime_image(
+            args.base_image, docker_builder, platform=args.platform
+        )
         print(f'\nBUILT Image: {image_name}\n')

tests/unit/test_runtime_build.py

@@ -215,6 +215,7 @@ def test_build_runtime_image_from_scratch():
                 f'{get_runtime_image_repo()}:{OH_VERSION}_mock-lock-hash_mock-source-hash',
                 f'{get_runtime_image_repo()}:{OH_VERSION}_mock-lock-hash',
             ],
+            platform=None,
         )
         assert (
             image_name