Adding rpath for /opt/rocm/libs since librccl.so etc is not found (#135)

* Adding rpath for /opt/rocm/libs since librccl.so etc is not found

* fix linting

* remove --force-rpath to enable user to set LD_LIBRARY_PATH

* black linting fixes

* edit pylint config to disable duplicate-code warning, will refactor the whole repo in the future.

Author: gulsumgudukbay

.pylintrc

@@ -1,3 +1,3 @@
 [MESSAGES CONTROL]
 
-disable = consider-using-f-string
+disable = consider-using-f-string,duplicate-code

jax_rocm_plugin/jaxlib_ext/tools/build_gpu_kernels_wheel.py

@@ -16,6 +16,11 @@
 # run via bazel run as part of the jax cuda plugin build process.
 
 # Most users should not run this script directly; use build.py instead.
+# pylint: disable=duplicate-code
+
+"""
+Script to build a JAX ROCm kernel plugin wheel. Intended for use via Bazel.
+"""
 
 import argparse
 import functools
@@ -25,6 +30,7 @@
 import subprocess
 import tempfile
 
+# pylint: disable=import-error
 from bazel_tools.tools.python.runfiles import runfiles
 from jaxlib_ext.tools import build_utils
 
@@ -58,120 +64,125 @@
 parser.add_argument(
     "--enable-cuda",
     default=False,
-    help="Should we build with CUDA enabled? Requires CUDA and CuDNN.")
+    help="Should we build with CUDA enabled? Requires CUDA and CuDNN.",
+)
 parser.add_argument(
-    "--enable-rocm",
-    default=False,
-    help="Should we build with ROCM enabled?")
+    "--enable-rocm", default=False, help="Should we build with ROCM enabled?"
+)
 args = parser.parse_args()
 
 r = runfiles.Create()
-pyext = "pyd" if build_utils.is_windows() else "so"
+PYEXT = "pyd" if build_utils.is_windows() else "so"
 
 
-def write_setup_cfg(sources_path, cpu):
-  tag = build_utils.platform_tag(cpu)
-  with open(sources_path / "setup.cfg", "w") as f:
-    f.write(f"""[metadata]
+def write_setup_cfg(setup_cfg_path, cpu):
+    """Write setup.cfg with platform tag."""
+    tag = build_utils.platform_tag(cpu)
+    with open(setup_cfg_path / "setup.cfg", "w", encoding="utf-8") as cfg_file:
+        cfg_file.write(
+            f"""[metadata]
 license_files = LICENSE.txt
 
 [bdist_wheel]
 plat_name={tag}
-""")
-
-
-def prepare_wheel_rocm(
-    sources_path: pathlib.Path, *, cpu, rocm_version
-):
-  """Assembles a source tree for the rocm kernel wheel in `sources_path`."""
-  copy_runfiles = functools.partial(build_utils.copy_file, runfiles=r)
-
-  copy_runfiles(
-      "__main__/jax_plugins/rocm/plugin_pyproject.toml",
-      dst_dir=sources_path,
-      dst_filename="pyproject.toml",
-  )
-  copy_runfiles(
-      "__main__/jax_plugins/rocm/plugin_setup.py",
-      dst_dir=sources_path,
-      dst_filename="setup.py",
-  )
-  build_utils.update_setup_with_rocm_version(sources_path, rocm_version)
-  write_setup_cfg(sources_path, cpu)
-
-  plugin_dir = sources_path / f"jax_rocm{rocm_version}_plugin"
-  copy_runfiles(
-      dst_dir=plugin_dir,
-      src_files=[
-          f"jax/jaxlib/rocm/_linalg.{pyext}",
-          f"jax/jaxlib/rocm/_prng.{pyext}",
-          f"jax/jaxlib/rocm/_solver.{pyext}",
-          f"jax/jaxlib/rocm/_sparse.{pyext}",
-          f"jax/jaxlib/rocm/_hybrid.{pyext}",
-          f"jax/jaxlib/rocm/_rnn.{pyext}",
-          f"jax/jaxlib/rocm/_triton.{pyext}",
-          f"jax/jaxlib/rocm/rocm_plugin_extension.{pyext}",
-          "jax/jaxlib/version.py",
-      ],
-  )
-
-  # NOTE(mrodden): this is a hack to change/set rpath values
-  # in the shared objects that are produced by the bazel build
-  # before they get pulled into the wheel build process.
-  # we have to do this change here because setting rpath
-  # using bazel requires the rpath to be valid during the build
-  # which won't be correct until we make changes to
-  # the xla/tsl/jax plugin build
-
-  try:
-    subprocess.check_output(["which", "patchelf"])
-  except subprocess.CalledProcessError as ex:
-    mesg = (
-        "rocm plugin and kernel wheel builds require patchelf. "
-        "please install 'patchelf' and run again"
+"""
+        )
+
+
+def prepare_wheel_rocm(rocm_sources_path: pathlib.Path, *, cpu, rocm_version):
+    """Assembles a source tree for the rocm kernel wheel in `rocm_sources_path`."""
+    copy_runfiles = functools.partial(build_utils.copy_file, runfiles=r)
+
+    copy_runfiles(
+        "__main__/jax_plugins/rocm/plugin_pyproject.toml",
+        dst_dir=rocm_sources_path,
+        dst_filename="pyproject.toml",
     )
-    raise Exception(mesg) from ex
-
-  files = [
-      f"_linalg.{pyext}",
-      f"_prng.{pyext}",
-      f"_solver.{pyext}",
-      f"_sparse.{pyext}",
-      f"_hybrid.{pyext}",
-      f"_rnn.{pyext}",
-      f"_triton.{pyext}",
-      f"rocm_plugin_extension.{pyext}",
-  ]
-  runpath = '$ORIGIN/../rocm/lib:$ORIGIN/../../rocm/lib'
-  # patchelf --force-rpath --set-rpath $RUNPATH $so
-  for f in files:
-    so_path = os.path.join(plugin_dir, f)
-    fix_perms = False
-    perms = os.stat(so_path).st_mode
-    if not perms & stat.S_IWUSR:
-        fix_perms = True
-        os.chmod(so_path, perms | stat.S_IWUSR)
-    subprocess.check_call(["patchelf", "--force-rpath", "--set-rpath", runpath, so_path])
-    if fix_perms:
-        os.chmod(so_path, perms)
-
-tmpdir = tempfile.TemporaryDirectory(prefix="jax_rocm_plugin")
-sources_path = tmpdir.name
-try:
-  os.makedirs(args.output_path, exist_ok=True)
-  prepare_wheel_rocm(
-      pathlib.Path(sources_path), cpu=args.cpu, rocm_version=args.platform_version
-  )
-  package_name = f"jax rocm{args.platform_version} plugin"
-  if args.editable:
-    build_utils.build_editable(sources_path, args.output_path, package_name)
-  else:
-    git_hash = build_utils.get_githash(args.jaxlib_git_hash)
-    build_utils.build_wheel(
-        sources_path,
-        args.output_path,
-        package_name,
-        git_hash=git_hash,
+    copy_runfiles(
+        "__main__/jax_plugins/rocm/plugin_setup.py",
+        dst_dir=rocm_sources_path,
+        dst_filename="setup.py",
     )
-finally:
-  tmpdir.cleanup()
+    build_utils.update_setup_with_rocm_version(rocm_sources_path, rocm_version)
+    write_setup_cfg(rocm_sources_path, cpu)
+
+    plugin_dir = rocm_sources_path / f"jax_rocm{rocm_version}_plugin"
+    copy_runfiles(
+        dst_dir=plugin_dir,
+        src_files=[
+            f"jax/jaxlib/rocm/_linalg.{PYEXT}",
+            f"jax/jaxlib/rocm/_prng.{PYEXT}",
+            f"jax/jaxlib/rocm/_solver.{PYEXT}",
+            f"jax/jaxlib/rocm/_sparse.{PYEXT}",
+            f"jax/jaxlib/rocm/_hybrid.{PYEXT}",
+            f"jax/jaxlib/rocm/_rnn.{PYEXT}",
+            f"jax/jaxlib/rocm/_triton.{PYEXT}",
+            f"jax/jaxlib/rocm/rocm_plugin_extension.{PYEXT}",
+            "jax/jaxlib/version.py",
+        ],
+    )
+
+    # NOTE(mrodden): this is a hack to change/set rpath values
+    # in the shared objects that are produced by the bazel build
+    # before they get pulled into the wheel build process.
+    # we have to do this change here because setting rpath
+    # using bazel requires the rpath to be valid during the build
+    # which won't be correct until we make changes to
+    # the xla/tsl/jax plugin build
+
+    try:
+        subprocess.check_output(["which", "patchelf"])
+    except subprocess.CalledProcessError as ex:
+        mesg = (
+            "rocm plugin and kernel wheel builds require patchelf. "
+            "please install 'patchelf' and run again"
+        )
+        raise RuntimeError(mesg) from ex
+
+    files = [
+        f"_linalg.{PYEXT}",
+        f"_prng.{PYEXT}",
+        f"_solver.{PYEXT}",
+        f"_sparse.{PYEXT}",
+        f"_hybrid.{PYEXT}",
+        f"_rnn.{PYEXT}",
+        f"_triton.{PYEXT}",
+        f"rocm_plugin_extension.{PYEXT}",
+    ]
+    runpath = "$ORIGIN/../rocm/lib:$ORIGIN/../../rocm/lib:/opt/rocm/lib"
+    # patchelf --set-rpath $RUNPATH $so
+    for fname in files:
+        so_path = os.path.join(plugin_dir, fname)
+        fix_perms = False
+        perms = os.stat(so_path).st_mode
+        if not perms & stat.S_IWUSR:
+            fix_perms = True
+            os.chmod(so_path, perms | stat.S_IWUSR)
+        subprocess.check_call(["patchelf", "--set-rpath", runpath, so_path])
+        if fix_perms:
+            os.chmod(so_path, perms)
+
+
+def main():
+    """Main entry point for building the ROCm kernel plugin wheel."""
+    with tempfile.TemporaryDirectory(prefix="jax_rocm_plugin") as tmpdir:
+        sources_path = tmpdir
+        os.makedirs(args.output_path, exist_ok=True)
+        prepare_wheel_rocm(
+            pathlib.Path(sources_path), cpu=args.cpu, rocm_version=args.platform_version
+        )
+        package_name = f"jax rocm{args.platform_version} plugin"
+        if args.editable:
+            build_utils.build_editable(sources_path, args.output_path, package_name)
+        else:
+            git_hash = build_utils.get_githash(args.jaxlib_git_hash)
+            build_utils.build_wheel(
+                sources_path,
+                args.output_path,
+                package_name,
+                git_hash=git_hash,
+            )
+
+
+if __name__ == "__main__":
+    main()