Add arm64 platform support to dagster-cloud.pex, do the build inside a manylinux builder image

scripts/Dockerfile.dagster-cloud-pex-builder

@@ -0,0 +1,14 @@
+# Builds generated/gha/dagster-cloud.pex and writes it to /generated
+
+# Use an official manylinux builder (https://github.com/pypa/manylinux#docker-images)
+FROM --platform=linux/amd64 quay.io/pypa/manylinux_2_28_x86_64:latest
+
+ENV PATH="/opt/python/cp311-cp311/bin:$PATH"
+
+RUN python3.11 -m pip install typer rich pex
+
+COPY release.py release.py
+
+RUN mkdir -p /generated
+
+CMD ["python3.11", "release.py", "build-dagster-cloud-pex"]

scripts/release.py

@@ -81,8 +81,8 @@ def build_docker_action(version_tag: str, publish_docker_action: bool = True):
             print(output)
 
 
-@app.command(help="Build dagster-cloud.pex")
-def update_dagster_cloud_pex(
+@app.command(help="Build dagster-cloud.pex - invoked by the dagster-cloud-pex-builder image")
+def build_dagster_cloud_pex(
     dagster_internal_branch: Optional[str] = DAGSTER_INTERNAL_BRANCH_OPTION,
     dagster_oss_branch: Optional[str] = DAGSTER_OSS_BRANCH_OPTION,
 ):
@@ -123,13 +123,13 @@ def update_dagster_cloud_pex(
         "PyGithub",
         "-o=dagster-cloud.pex",
         *platform_args,
-        "--platform=macosx_12_0_x86_64-cp-311-cp311",
-        "--platform=macosx_12_0_arm64-cp-311-cp311",
+        # For arm64 github runners
+        "--platform=manylinux_2_17_aarch64-cp-312-cp312",
         "--pip-version=23.0",
         "--resolver-version=pip-2020-resolver",
         "--venv=prepend",
         "--python-shebang=/usr/bin/env python",
-        "-v",
+        "-vvvvv",
     ]
     print(f"Running {args}")
     output = subprocess.check_output(
@@ -142,6 +142,43 @@ def update_dagster_cloud_pex(
     info("Built generated/gha/dagster-cloud.pex")
 
 
+@app.command(help="Update dagster-cloud.pex")
+def update_dagster_cloud_pex(
+    dagster_internal_branch: Optional[str] = DAGSTER_INTERNAL_BRANCH_OPTION,
+    dagster_oss_branch: Optional[str] = DAGSTER_OSS_BRANCH_OPTION,
+):
+    # Map /generated on the docker image to our local generated folder
+    map_folders = {"/generated": os.path.join(os.path.dirname(__file__), "..", "generated")}
+
+    mount_args = []
+    for target_folder, source_folder in map_folders.items():
+        mount_args.extend(["--mount", f"type=bind,source={source_folder},target={target_folder}"])
+
+    cmd = [
+        "docker",
+        "build",
+        "--progress=plain",
+        "-t",
+        "dagster-cloud-pex-builder",
+        "--platform=linux/amd64",
+        "-f",
+        os.path.join(os.path.dirname(__file__), "Dockerfile.dagster-cloud-pex-builder"),
+        os.path.dirname(__file__),
+    ]
+
+    subprocess.run(cmd, check=True)
+
+    cmd = [
+        "docker",
+        "run",
+        "--platform=linux/amd64",
+        *mount_args,
+        "-t",
+        "dagster-cloud-pex-builder",
+    ]
+    subprocess.run(cmd, check=True)
+
+
 @app.command()
 def update_docker_action_references(
     previous_version_tag,

src/Dockerfile.test-dagster-cloud-pex

@@ -0,0 +1,6 @@
+# Used for testing the generated dagster-cloud pex
+FROM --platform=linux/amd64 python:3.10-slim
+
+COPY generated/gha/dagster-cloud.pex /dagster-cloud.pex
+
+ENTRYPOINT ["/dagster-cloud.pex"]

tests/test_pex_builder.py

@@ -1,62 +1,91 @@
 import os
 import subprocess
 import tempfile
-from contextlib import contextmanager
-from typing import List
+from typing import List, Mapping
 
+import pytest
+
+
+def run_dagster_cloud_serverless_cmd(args: List[str], map_folders: Mapping[str, str]):
+    mount_args = []
+    for target_folder, source_folder in map_folders.items():
+        mount_args.extend(["--mount", f"type=bind,source={source_folder},target={target_folder}"])
+
+    cmd = [
+        "docker",
+        "run",
+        "--platform=linux/amd64",
+        *mount_args,
+        "-t",
+        "test-dagster-cloud-pex",
+        "-m",
+        "dagster_cloud_cli.entrypoint",
+        "serverless",
+        *args,
+    ]
+
+    subprocess.run(cmd, encoding="utf-8", capture_output=False, check=True)
+
+
+@pytest.fixture
+def built_test_dagster_cloud_pex_image(repo_root: str):
+    src_dir = os.path.join(os.path.dirname(__file__), "..", "src")
+
+    cmd = [
+        "docker",
+        "build",
+        "--progress=plain",
+        "-t",
+        "test-dagster-cloud-pex",
+        "--platform=linux/amd64",
+        "-f",
+        os.path.join(src_dir, "Dockerfile.test-dagster-cloud-pex"),
+        repo_root,
+    ]
+
+    subprocess.run(cmd, check=True)
+
+
+def test_pex_build_only(repo_root, built_test_dagster_cloud_pex_image):
+    dagster_project1 = repo_root / "tests/test-repos/dagster_project1"
 
-@contextmanager
-def run_dagster_cloud_serverless_cmd(dagster_cloud_pex_path, args: List[str]):
     with tempfile.TemporaryDirectory() as build_output_dir:
-        proc = subprocess.run(
+        map_folders = {"/dagster_project1": dagster_project1, "/build_output_dir": build_output_dir}
+
+        run_dagster_cloud_serverless_cmd(
             [
-                dagster_cloud_pex_path,
-                "-m",
-                "dagster_cloud_cli.entrypoint",
-                "serverless",
-                *args,
-                build_output_dir,
+                "build-python-executable",
+                "/dagster_project1",
+                "--api-token=fake",
+                "--url=fake",
+                "--python-version=3.10",
+                "/build_output_dir",
             ],
-            capture_output=True,
-            check=False,
+            map_folders=map_folders,
         )
-        if proc.returncode:
-            raise ValueError(
-                "Failed to run dagster-cloud.pex:" + (proc.stdout + proc.stderr).decode("utf-8")
-            )
 
         all_files = os.listdir(build_output_dir)
         pex_files = {
             filename for filename in all_files if filename.endswith(".pex") and filename != ".pex"
         }
-        yield (build_output_dir, list(pex_files), list(set(all_files) - pex_files))
 
-
-def test_pex_build_only(repo_root, dagster_cloud_pex_path):
-    dagster_project1 = repo_root / "tests/test-repos/dagster_project1"
-    with run_dagster_cloud_serverless_cmd(
-        dagster_cloud_pex_path,
-        [
-            "build-python-executable",
-            str(dagster_project1),
-            "--api-token=fake",
-            "--url=fake",
-            "--python-version=3.11",
-        ],
-    ) as (
-        build_output_dir,
-        pex_files,
-        other_files,
-    ):
         # one source-HASH.pex and one deps-HASH.pex file are expected
         assert 2 == len(pex_files)
         pex_file_by_alias = {filename.split("-", 1)[0]: filename for filename in pex_files}
 
         assert {"source", "deps"} == set(pex_file_by_alias)
 
 
-def test_dagster_cloud_runnable(dagster_cloud_pex_path):
-    output = subprocess.check_output(
-        [dagster_cloud_pex_path, "-c", "print('hello')"], encoding="utf-8"
-    )
-    assert "hello" in output
+def test_dagster_cloud_runnable(built_test_dagster_cloud_pex_image):
+    cmd = [
+        "docker",
+        "run",
+        "--platform=linux/amd64",
+        "-t",
+        "test-dagster-cloud-pex",
+        "-c",
+        "print('hello')",
+    ]
+    output = subprocess.run(cmd, encoding="utf-8", capture_output=True, check=True)
+
+    assert "hello" in output.stdout