Add a Python architecture input

README.md

@@ -15,6 +15,7 @@ Set up your GitHub Actions workflow with a specific version of [uv](https://docs
   - [Inputs](#inputs)
   - [Outputs](#outputs)
   - [Python version](#python-version)
+  - [Python architecture](#python-architecture)
   - [Working directory](#working-directory)
 - [Advanced Configuration](#advanced-configuration)
 - [How it works](#how-it-works)
@@ -56,6 +57,9 @@ Have a look under [Advanced Configuration](#advanced-configuration) for detailed
     # The version of Python to set UV_PYTHON to (overrides the Python version from .tool-versions)
     python-version: ""
 
+    # The Python architecture to set UV_PYTHON_ARCH to, e.g., x86_64 or aarch64
+    python-arch: ""
+
     # Use uv venv to activate a venv ready to be used by later steps
     activate-environment: "false"
 
@@ -193,6 +197,25 @@ jobs:
         run: uv run --frozen pytest
 ```
 
+### Python architecture
+
+Use `python-arch` to set `UV_PYTHON_ARCH` for the rest of the job. This selects the Python
+architecture independently of the version, including versions requested by `.python-version`.
+An explicit architecture or interpreter path in a Python request takes precedence.
+
+```yaml
+- uses: astral-sh/setup-uv@main
+  with:
+    python-version: "3.14"
+    python-arch: x86_64
+- run: uv run --frozen pytest
+```
+
+The input overrides an existing `UV_PYTHON_ARCH` value. When it is omitted, the action respects
+`UV_PYTHON_ARCH` from the environment. Both forms require a uv version that supports the variable;
+the action reports an error if the installed version does not support it. The selected architecture
+is included in the [cache key](docs/caching.md#cache-key).
+
 ### Working directory
 
 You can set the working directory with the `working-directory` input.

__tests__/cache/restore-cache.test.ts

@@ -1,9 +1,20 @@
-import { beforeEach, describe, expect, it, jest } from "@jest/globals";
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  jest,
+} from "@jest/globals";
 import { createSetupInputs } from "../helpers/setup-inputs";
 
 const mockRestoreCache = jest.fn();
 const mockSaveState = jest.fn();
 const mockSetOutput = jest.fn();
+const mockGetArch = jest.fn(() => "x86_64");
+const mockGetOSNameVersion = jest.fn(() => "ubuntu-24.04");
+const mockGetPlatform = jest.fn(async () => "unknown-linux-gnu");
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
 
 jest.unstable_mockModule("@actions/cache", () => ({
   restoreCache: mockRestoreCache,
@@ -24,9 +35,9 @@ jest.unstable_mockModule("../../src/utils/logging", () => ({
 }));
 
 jest.unstable_mockModule("../../src/utils/platforms", () => ({
-  getArch: jest.fn(() => "x86_64"),
-  getOSNameVersion: jest.fn(() => "ubuntu-24.04"),
-  getPlatform: jest.fn(async () => "unknown-linux-gnu"),
+  getArch: mockGetArch,
+  getOSNameVersion: mockGetOSNameVersion,
+  getPlatform: mockGetPlatform,
 }));
 
 const { restoreCache } = await import("../../src/cache/restore-cache");
@@ -38,7 +49,19 @@ function cacheKeyOutput(): string {
 }
 
 beforeEach(() => {
+  delete process.env.UV_PYTHON_ARCH;
   jest.clearAllMocks();
+  mockGetArch.mockReturnValue("x86_64");
+  mockGetOSNameVersion.mockReturnValue("ubuntu-24.04");
+  mockGetPlatform.mockResolvedValue("unknown-linux-gnu");
+});
+
+afterEach(() => {
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
 });
 
 describe("restoreCache", () => {
@@ -71,4 +94,39 @@ describe("restoreCache", () => {
       "setup-uv-2-x86_64-unknown-linux-gnu-ubuntu-24.04-3.11-dependencyhash-tests-3.11",
     );
   });
+
+  it("includes an inherited Python architecture in the cache key", async () => {
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    await restoreCache(createSetupInputs(), "3.14");
+
+    expect(cacheKeyOutput()).toContain("-3.14-python-aarch64-");
+  });
+
+  it.each(["aarch64", "x86_64"])(
+    "separates %s Python caches on a Windows ARM64 runner",
+    async (pythonArch) => {
+      mockGetArch.mockReturnValue("aarch64");
+      mockGetOSNameVersion.mockReturnValue("windows-11");
+      mockGetPlatform.mockResolvedValue("pc-windows-msvc");
+      const inputs = createSetupInputs({
+        cachePython: true,
+        pythonArch,
+        restoreCache: true,
+      });
+
+      await restoreCache(inputs, "3.14");
+
+      const cacheKey = `setup-uv-2-aarch64-pc-windows-msvc-windows-11-3.14-python-${pythonArch}-py-dependencyhash`;
+      expect(cacheKeyOutput()).toBe(cacheKey);
+      expect(mockRestoreCache).toHaveBeenCalledWith(
+        [inputs.cacheLocalPath?.path],
+        cacheKey,
+      );
+      expect(mockRestoreCache).toHaveBeenCalledWith(
+        [inputs.pythonDir],
+        `${cacheKey}-python`,
+      );
+    },
+  );
 });

__tests__/helpers/setup-inputs.ts

@@ -21,6 +21,7 @@ export function createSetupInputs(
     ignoreNothingToCache: false,
     noProject: false,
     pruneCache: false,
+    pythonArch: "",
     pythonDir: "/tmp/uv-python-dir",
     pythonVersion: "",
     quiet: false,

__tests__/utils/inputs.test.ts

@@ -19,6 +19,7 @@ const ORIGINAL_RUNNER_ENVIRONMENT = process.env.RUNNER_ENVIRONMENT;
 const ORIGINAL_RUNNER_TEMP = process.env.RUNNER_TEMP;
 const ORIGINAL_UV_CACHE_DIR = process.env.UV_CACHE_DIR;
 const ORIGINAL_UV_PYTHON = process.env.UV_PYTHON;
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
 const ORIGINAL_UV_PYTHON_INSTALL_DIR = process.env.UV_PYTHON_INSTALL_DIR;
 
 const mockDebug = jest.fn();
@@ -37,7 +38,9 @@ jest.unstable_mockModule("@actions/core", () => ({
   warning: mockWarning,
 }));
 
-const { CacheLocalSource, loadInputs } = await import("../../src/utils/inputs");
+const { CacheLocalSource, loadInputs, resolvePythonArch } = await import(
+  "../../src/utils/inputs"
+);
 
 function createTempProject(files: Record<string, string> = {}): string {
   const dir = fs.mkdtempSync(path.join(os.tmpdir(), "setup-uv-inputs-test-"));
@@ -62,6 +65,7 @@ function resetEnvironment(): void {
   delete process.env.RUNNER_TEMP;
   delete process.env.UV_CACHE_DIR;
   delete process.env.UV_PYTHON;
+  delete process.env.UV_PYTHON_ARCH;
   delete process.env.UV_PYTHON_INSTALL_DIR;
 }
 
@@ -77,6 +81,11 @@ function restoreEnvironment(): void {
   process.env.RUNNER_TEMP = ORIGINAL_RUNNER_TEMP;
   process.env.UV_CACHE_DIR = ORIGINAL_UV_CACHE_DIR;
   process.env.UV_PYTHON = ORIGINAL_UV_PYTHON;
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
   process.env.UV_PYTHON_INSTALL_DIR = ORIGINAL_UV_PYTHON_INSTALL_DIR;
 }
 
@@ -98,6 +107,7 @@ describe("loadInputs", () => {
       source: CacheLocalSource.Default,
     });
     expect(inputs.pythonDir).toBe("/runner-temp/uv-python-dir");
+    expect(inputs.pythonArch).toBe("");
     expect(inputs.venvPath).toBe("/workspace/.venv");
     expect(inputs.manifestFile).toBeUndefined();
     expect(inputs.resolutionStrategy).toBe("highest");
@@ -149,6 +159,28 @@ describe("loadInputs", () => {
     expect(inputs.pythonVersion).toBe("");
   });
 
+  it("prefers the python-arch input over UV_PYTHON_ARCH", () => {
+    mockInputs["working-directory"] = "/workspace";
+    mockInputs["python-arch"] = "x86_64";
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    const inputs = loadInputs();
+    expect(inputs.pythonArch).toBe("x86_64");
+    expect(resolvePythonArch(inputs.pythonArch)).toBe("x86_64");
+  });
+
+  it.each(["aarch64", "x86_64_v3", ""])(
+    "uses UV_PYTHON_ARCH when python-arch is omitted: %s",
+    (pythonArch) => {
+      mockInputs["working-directory"] = "/workspace";
+      process.env.UV_PYTHON_ARCH = pythonArch;
+
+      const inputs = loadInputs();
+      expect(inputs.pythonArch).toBe("");
+      expect(resolvePythonArch(inputs.pythonArch)).toBe(pythonArch);
+    },
+  );
+
   it.each(["pull_request_target", "workflow_run", "release"])(
     "disables automatic caching for the %s event",
     (eventName) => {

__tests__/utils/python-arch.test.ts

@@ -0,0 +1,152 @@
+import { promisify } from "node:util";
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  jest,
+} from "@jest/globals";
+
+const ORIGINAL_UV_PYTHON_ARCH = process.env.UV_PYTHON_ARCH;
+
+const mockExecFile =
+  jest.fn<
+    (...args: unknown[]) => Promise<{
+      stdout: string;
+      stderr: string;
+    }>
+  >();
+const mockExportVariable = jest.fn();
+
+jest.unstable_mockModule("node:child_process", () => ({
+  execFile: Object.assign(mockExecFile, { [promisify.custom]: mockExecFile }),
+}));
+jest.unstable_mockModule("@actions/core", () => ({
+  exportVariable: mockExportVariable,
+}));
+jest.unstable_mockModule("../../src/utils/logging", () => ({
+  info: jest.fn(),
+}));
+
+const { setupPythonArch } = await import("../../src/utils/python-arch");
+
+const supportedProbe = Object.assign(new Error("uv exited with code 2"), {
+  code: 2,
+  stderr:
+    "error: Failed to parse environment variable `UV_PYTHON_ARCH` with invalid value `setup-uv-probe`: Unknown architecture: setup-uv-probe\n",
+});
+const success = { stderr: "", stdout: "/cache\n" };
+
+beforeEach(() => {
+  delete process.env.UV_PYTHON_ARCH;
+  jest.clearAllMocks();
+  mockExecFile
+    .mockReset()
+    .mockRejectedValueOnce(supportedProbe)
+    .mockResolvedValueOnce(success);
+});
+
+afterEach(() => {
+  if (ORIGINAL_UV_PYTHON_ARCH === undefined) {
+    delete process.env.UV_PYTHON_ARCH;
+  } else {
+    process.env.UV_PYTHON_ARCH = ORIGINAL_UV_PYTHON_ARCH;
+  }
+});
+
+describe("setupPythonArch", () => {
+  it("leaves Python selection alone without an architecture", async () => {
+    await setupPythonArch("/tools/uv", "");
+
+    expect(mockExecFile).not.toHaveBeenCalled();
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("validates an inherited architecture without exporting it", async () => {
+    process.env.UV_PYTHON_ARCH = "aarch64";
+
+    await setupPythonArch("/tools/uv", "");
+
+    expect(mockExecFile).toHaveBeenNthCalledWith(
+      2,
+      "/tools/uv",
+      ["--no-config", "cache", "dir"],
+      expect.objectContaining({
+        env: expect.objectContaining({ UV_PYTHON_ARCH: "aarch64" }),
+      }),
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+    expect(process.env.UV_PYTHON_ARCH).toBe("aarch64");
+  });
+
+  it.each(["/runner temp/uv", "C:\\runner temp\\uv.exe"])(
+    "validates and exports the architecture using the installed uv: %s",
+    async (uvPath) => {
+      await setupPythonArch(uvPath, "x86_64");
+
+      expect(mockExecFile).toHaveBeenNthCalledWith(
+        1,
+        uvPath,
+        ["--no-config", "cache", "dir"],
+        expect.objectContaining({
+          encoding: "utf8",
+          env: expect.objectContaining({ UV_PYTHON_ARCH: "setup-uv-probe" }),
+        }),
+      );
+      expect(mockExecFile).toHaveBeenNthCalledWith(
+        2,
+        uvPath,
+        ["--no-config", "cache", "dir"],
+        expect.objectContaining({
+          env: expect.objectContaining({ UV_PYTHON_ARCH: "x86_64" }),
+        }),
+      );
+      expect(mockExportVariable).toHaveBeenCalledWith(
+        "UV_PYTHON_ARCH",
+        "x86_64",
+      );
+    },
+  );
+
+  it("rejects uv versions that ignore UV_PYTHON_ARCH", async () => {
+    mockExecFile.mockReset().mockResolvedValue(success);
+
+    await expect(setupPythonArch("/tools/uv", "x86_64")).rejects.toThrow(
+      "The installed version of uv does not support UV_PYTHON_ARCH",
+    );
+    expect(mockExecFile).toHaveBeenCalledTimes(1);
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("reports unexpected probe failures", async () => {
+    mockExecFile.mockReset().mockRejectedValue(
+      Object.assign(new Error("uv exited with code 2"), {
+        code: 2,
+        stderr: "error: unrelated configuration error\n",
+      }),
+    );
+
+    await expect(setupPythonArch("/tools/uv", "x86_64")).rejects.toThrow(
+      "Failed to check uv's support for UV_PYTHON_ARCH: error: unrelated configuration error",
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+
+  it("reports invalid architectures before exporting them", async () => {
+    mockExecFile
+      .mockReset()
+      .mockRejectedValueOnce(supportedProbe)
+      .mockRejectedValueOnce(
+        Object.assign(new Error("uv exited with code 2"), {
+          code: 2,
+          stderr: "error: Unknown architecture: invalid\n",
+        }),
+      );
+
+    await expect(setupPythonArch("/tools/uv", "invalid")).rejects.toThrow(
+      "Failed to set Python architecture to invalid: error: Unknown architecture: invalid",
+    );
+    expect(mockExportVariable).not.toHaveBeenCalled();
+  });
+});

action-types.yml

@@ -7,6 +7,8 @@ inputs:
     type: string
   python-version:
     type: string
+  python-arch:
+    type: string
   activate-environment:
     type: boolean
   venv-path:

action.yml

@@ -12,6 +12,9 @@ inputs:
   python-version:
     description: "The version of Python to set UV_PYTHON to. Overrides the Python version from .tool-versions."
     required: false
+  python-arch:
+    description: "The Python architecture to set UV_PYTHON_ARCH to, e.g., x86_64 or aarch64. Overrides UV_PYTHON_ARCH."
+    required: false
   activate-environment:
     description: "Use uv venv to activate a venv ready to be used by later steps. "
     default: "false"

dist/save-cache/index.cjs

@@ -63714,15 +63714,6 @@ function info2(msg) {
 }
 var warning2 = warning;
 
-// src/cache/restore-cache.ts
-var STATE_CACHE_KEY = "cache-key";
-var STATE_CACHE_MATCHED_KEY = "cache-matched-key";
-var STATE_PYTHON_CACHE_MATCHED_KEY = "python-cache-matched-key";
-
-// src/utils/constants.ts
-var STATE_UV_PATH = "uv-path";
-var STATE_UV_VERSION = "uv-version";
-
 // src/utils/inputs.ts
 var import_node_fs4 = __toESM(require("node:fs"), 1);
 var import_node_path = __toESM(require("node:path"), 1);
@@ -64447,6 +64438,7 @@ function loadInputs() {
   const version3 = getInput("version");
   const versionFile = getVersionFile(workingDirectory);
   const pythonVersion = getPythonVersion(versionFile);
+  const pythonArch = getInput("python-arch");
   const activateEnvironment = getBooleanInput("activate-environment");
   const noProject = getBooleanInput("no-project");
   const venvPath = getVenvPath(workingDirectory, activateEnvironment);
@@ -64490,6 +64482,7 @@ function loadInputs() {
     manifestFile,
     noProject,
     pruneCache: pruneCache2,
+    pythonArch,
     pythonDir,
     pythonVersion,
     quiet: quiet2,
@@ -64738,6 +64731,15 @@ function getResolutionStrategy() {
   );
 }
 
+// src/cache/restore-cache.ts
+var STATE_CACHE_KEY = "cache-key";
+var STATE_CACHE_MATCHED_KEY = "cache-matched-key";
+var STATE_PYTHON_CACHE_MATCHED_KEY = "python-cache-matched-key";
+
+// src/utils/constants.ts
+var STATE_UV_PATH = "uv-path";
+var STATE_UV_VERSION = "uv-version";
+
 // src/save-cache.ts
 function formatUnexpectedFailure(error2) {
   if (error2 instanceof Error) {

docs/caching.md

@@ -11,6 +11,7 @@ The cache key is automatically generated based on:
   `pc-windows-msvc`)
 - **OS version**: OS name and version (e.g., `ubuntu-22.04`, `macos-14`, `windows-2022`)
 - **Python version**: The Python version in use
+- **Python architecture**: The architecture selected by `python-arch` or `UV_PYTHON_ARCH`, when set
 - **Cache options**: Whether pruning and Python caching are enabled
 - **Dependency hash**: Hash of files matching `cache-dependency-glob`
 - **Suffix**: Optional `cache-suffix` if provided

docs/environment-and-tools.md

@@ -144,6 +144,7 @@ If you want to ignore this, set the `ignore-empty-workdir` input to `true`.
 This action sets several environment variables that influence uv's behavior and can be used by subsequent steps:
 
 - `UV_PYTHON`: Set when `python-version` is specified or the selected `.tool-versions` file contains a supported `python` entry. Controls which Python version uv uses.
+- `UV_PYTHON_ARCH`: Set when `python-arch` is specified. Controls the architecture of Python requests that do not specify one.
 - `UV_CACHE_DIR`: Set when caching is enabled (unless already configured in uv config files). Controls where uv stores its cache.
 - `UV_TOOL_DIR`: Set when `tool-dir` input is specified. Controls where uv installs tool environments.
 - `UV_TOOL_BIN_DIR`: Set when `tool-bin-dir` input is specified. Controls where uv installs tool binaries.
@@ -155,6 +156,7 @@ This action sets several environment variables that influence uv's behavior and
 - `UV_NO_MODIFY_PATH`: If set, prevents the action from modifying PATH. Cannot be used with `activate-environment`.
 - `UV_CACHE_DIR`: If already set, the action will respect it instead of setting its own cache directory.
 - `UV_PYTHON`: If already set and `python-version` is not specified, the action will respect it instead of using the `python` entry from `.tool-versions`.
+- `UV_PYTHON_ARCH`: If already set and `python-arch` is not specified, the action will use it and include it in the cache key. Requires a uv version that supports `UV_PYTHON_ARCH`.
 
 ```yaml
 - name: Example using environment variables

src/cache/restore-cache.ts

@@ -1,7 +1,7 @@
 import * as cache from "@actions/cache";
 import * as core from "@actions/core";
 import { hashFiles } from "../hash/hash-files";
-import type { SetupInputs } from "../utils/inputs";
+import { resolvePythonArch, type SetupInputs } from "../utils/inputs";
 import * as log from "../utils/logging";
 import { getArch, getOSNameVersion, getPlatform } from "../utils/platforms";
 
@@ -98,11 +98,15 @@ async function computeKeys(
     ? `-${encodeURIComponent(inputs.cacheSuffix)}`
     : "";
   const version = encodeURIComponent(pythonVersion ?? "unknown");
+  const pythonArch = resolvePythonArch(inputs.pythonArch);
+  const pythonArchKey = pythonArch
+    ? `-python-${encodeURIComponent(pythonArch)}`
+    : "";
   const platform = await getPlatform();
   const osNameVersion = getOSNameVersion();
   const pruned = inputs.pruneCache ? "-pruned" : "";
   const python = inputs.cachePython ? "-py" : "";
-  return `setup-uv-${CACHE_VERSION}-${getArch()}-${platform}-${osNameVersion}-${version}${pruned}${python}${cacheDependencyPathHash}${suffix}`;
+  return `setup-uv-${CACHE_VERSION}-${getArch()}-${platform}-${osNameVersion}-${version}${pythonArchKey}${pruned}${python}${cacheDependencyPathHash}${suffix}`;
 }
 
 function handleMatchResult(

src/setup-uv.ts

@@ -16,6 +16,7 @@ import {
   getPlatform,
   type Platform,
 } from "./utils/platforms";
+import { setupPythonArch } from "./utils/python-arch";
 import { getPythonRuntimeId } from "./utils/python-runtime";
 import { resolveUvVersion } from "./version/resolve";
 
@@ -86,6 +87,11 @@ async function run(): Promise<void> {
       throw new Error(`Unsupported architecture: ${process.arch}`);
     }
     const setupResult = await setupUv(inputs, platform, arch);
+    const uvPath = path.join(
+      setupResult.uvDir,
+      process.platform === "win32" ? "uv.exe" : "uv",
+    );
+    await setupPythonArch(uvPath, inputs.pythonArch);
 
     addToolBinToPath(inputs);
     addUvToPathAndOutput(setupResult.uvDir);

src/utils/inputs.ts

@@ -24,6 +24,7 @@ export interface SetupInputs {
   version: string;
   versionFile: string;
   pythonVersion: string;
+  pythonArch: string;
   activateEnvironment: boolean;
   noProject: boolean;
   venvPath: string;
@@ -54,6 +55,7 @@ export function loadInputs(): SetupInputs {
   const version = core.getInput("version");
   const versionFile = getVersionFile(workingDirectory);
   const pythonVersion = getPythonVersion(versionFile);
+  const pythonArch = core.getInput("python-arch");
   const activateEnvironment = core.getBooleanInput("activate-environment");
   const noProject = core.getBooleanInput("no-project");
   const venvPath = getVenvPath(workingDirectory, activateEnvironment);
@@ -100,6 +102,7 @@ export function loadInputs(): SetupInputs {
     manifestFile,
     noProject,
     pruneCache,
+    pythonArch,
     pythonDir,
     pythonVersion,
     quiet,
@@ -146,6 +149,10 @@ function getPythonVersion(versionFile: string): string {
   }
 }
 
+export function resolvePythonArch(pythonArch: string): string {
+  return pythonArch || process.env.UV_PYTHON_ARCH || "";
+}
+
 function getVenvPath(
   workingDirectory: string,
   activateEnvironment: boolean,

src/utils/python-arch.ts

@@ -0,0 +1,76 @@
+import { execFile } from "node:child_process";
+import { promisify } from "node:util";
+import * as core from "@actions/core";
+import { resolvePythonArch } from "./inputs";
+import * as log from "./logging";
+
+const execFileAsync = promisify(execFile);
+const PROBE_ARCH = "setup-uv-probe";
+
+export async function setupPythonArch(
+  uvPath: string,
+  pythonArchInput: string,
+): Promise<void> {
+  const pythonArch = resolvePythonArch(pythonArchInput);
+  if (pythonArch === "") {
+    return;
+  }
+
+  // Older uv releases ignore unknown environment variables. An invalid architecture
+  // confirms that uv recognizes UV_PYTHON_ARCH without finding or downloading Python.
+  const probe = await queryPythonArch(uvPath, PROBE_ARCH);
+  if (probe.exitCode === 0) {
+    throw new Error(
+      "The installed version of uv does not support UV_PYTHON_ARCH. Select a newer uv version or use an architecture-qualified python-version.",
+    );
+  }
+  if (
+    !probe.stderr.includes(
+      `environment variable \`UV_PYTHON_ARCH\` with invalid value \`${PROBE_ARCH}\``,
+    )
+  ) {
+    throw new Error(
+      `Failed to check uv's support for UV_PYTHON_ARCH: ${probe.stderr.trim() || `uv exited with code ${probe.exitCode}`}`,
+    );
+  }
+
+  const selected = await queryPythonArch(uvPath, pythonArch);
+  if (selected.exitCode !== 0) {
+    throw new Error(
+      `Failed to set Python architecture to ${pythonArch}: ${selected.stderr.trim() || `uv exited with code ${selected.exitCode}`}`,
+    );
+  }
+
+  if (pythonArchInput !== "") {
+    core.exportVariable("UV_PYTHON_ARCH", pythonArch);
+    log.info(`Set UV_PYTHON_ARCH to ${pythonArch}`);
+  }
+}
+
+async function queryPythonArch(
+  uvPath: string,
+  pythonArch: string,
+): Promise<{ exitCode: number; stderr: string }> {
+  try {
+    const { stderr } = await execFileAsync(
+      uvPath,
+      ["--no-config", "cache", "dir"],
+      {
+        encoding: "utf8",
+        env: { ...process.env, NO_COLOR: "1", UV_PYTHON_ARCH: pythonArch },
+      },
+    );
+    return { exitCode: 0, stderr };
+  } catch (error) {
+    if (
+      error instanceof Error &&
+      "code" in error &&
+      typeof error.code === "number" &&
+      "stderr" in error &&
+      typeof error.stderr === "string"
+    ) {
+      return { exitCode: error.code, stderr: error.stderr };
+    }
+    throw error;
+  }
+}