From 6f4b912f8d36ba1aa8276151681f93237b462c26 Mon Sep 17 00:00:00 2001
From: erman-gurses <99776114+erman-gurses@users.noreply.github.com>
Date: Fri, 26 Jun 2026 10:44:33 -0700
Subject: [PATCH] Pass ROCm package `rocm_package_find_links_url` through JAX
 release workflow (#65)

## Summary
Related PR: https://github.com/ROCm/TheRock/pull/6054

Add `rocm_package_find_links_url` to the rockrel JAX release workflow
and forward it to the reusable TheRock JAX build workflow.

This allows the JAX manylinux build path to consume ROCm packages from a
configurable package index instead of relying on a hardcoded URL.

## Changes

* Add `rocm_package_find_links_url` workflow input to
`multi_arch_release_linux_jax_wheels.yml`
* Forward `rocm_package_find_links_url` to the reusable
`multi_arch_build_linux_jax_wheels.yml` workflow
* Default the value to the current dev multi-arch package index:

  * `https://rocm.devreleases.amd.com/whl-multi-arch/`

## Motivation

The JAX manylinux build flow now installs ROCm packages from a package
index. Passing the package source as an input allows the workflow to
support different release channels (dev, nightly, prerelease) without
hardcoding a specific index URL in the build workflow.

Tests:
https://github.com/ROCm/rockrel/actions/runs/28086158341
---
 .../workflows/multi_arch_release_linux_jax_wheels.yml | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/multi_arch_release_linux_jax_wheels.yml b/.github/workflows/multi_arch_release_linux_jax_wheels.yml
index 62661f6b..2c7f0835 100644
--- a/.github/workflows/multi_arch_release_linux_jax_wheels.yml
+++ b/.github/workflows/multi_arch_release_linux_jax_wheels.yml
@@ -15,15 +15,19 @@ on:
         type: string
         required: true
       rocm_version:
-        description: ROCm package version to build against.
+        description: "ROCm package version to build against."
         type: string
         required: true
+      rocm_package_find_links_url:
+        description: "URL for pip --find-links to install ROCm packages for the JAX manylinux build."
+        type: string
+        default: ""
       tar_url:
-        description: URL to the shared TheRock tarball used for the build.
+        description: "URL to TheRock tarball used for the build."
         type: string
         required: true
       ref:
-        description: Branch, tag, or SHA to checkout.
+        description: "Branch, tag, or SHA to checkout."
         type: string
         default: ""
       repository:
@@ -45,6 +49,7 @@ jobs:
       test_amdgpu_family: ${{ inputs.test_amdgpu_family }}
       release_type: ${{ inputs.release_type }}
       rocm_version: ${{ inputs.rocm_version }}
+      rocm_package_find_links_url: ${{ inputs.rocm_package_find_links_url }}
       tar_url: ${{ inputs.tar_url }}
       repository: ${{ inputs.repository || 'ROCm/TheRock' }}
       ref: ${{ inputs.ref || '' }}