From bbb53a2d88ab79b5de461061e410b4c2e6b16188 Mon Sep 17 00:00:00 2001
From: Ilya Panfilov <ipanfilo@amd.com>
Date: Thu, 28 May 2026 22:42:55 -0400
Subject: [PATCH 1/2] Install: use setuptools bdist_wheel; CI: call pytest as
 module

---
 ci/_utils.sh                        | 3 ++-
 setup.py                            | 5 ++++-
 transformer_engine/pytorch/setup.py | 5 ++++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/ci/_utils.sh b/ci/_utils.sh
index b4aae9cc7..35cba3203 100644
--- a/ci/_utils.sh
+++ b/ci/_utils.sh
@@ -266,6 +266,7 @@ pytest_run() {
     check_test_filter $_test_name_tag || return
     _start_ts=`date +%s`
     echo "Run [$_test_variant_tag] $@ at `time_elapsed $TEST_START_TS`"
-    pytest -v -rfEs `get_pytest_junitxml $_test_name_tag` $TEST_PYTEST_ARGS "$TEST_DIR/$@" || test_run_error "[$_test_variant_tag] $1"
+    python3 -m pytest -v -rfEs `get_pytest_junitxml $_test_name_tag` $TEST_PYTEST_ARGS "$TEST_DIR/$@"
+    test $? -eq 0 || test_run_error "[$_test_variant_tag] $1"
     echo "Done [$_test_variant_tag] $1 in `time_elapsed $_start_ts`"
 }
diff --git a/setup.py b/setup.py
index c66af39df..20ec983d3 100644
--- a/setup.py
+++ b/setup.py
@@ -16,7 +16,10 @@
 
 import setuptools
 from setuptools.command.egg_info import egg_info
-from wheel.bdist_wheel import bdist_wheel
+try:
+    from setuptools.command.bdist_wheel import bdist_wheel
+except ImportError:
+    from wheel.bdist_wheel import bdist_wheel
 
 from build_tools.build_ext import CMakeExtension, get_build_ext
 from build_tools.te_version import te_version
diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py
index a80dcacb8..616a50f44 100644
--- a/transformer_engine/pytorch/setup.py
+++ b/transformer_engine/pytorch/setup.py
@@ -15,7 +15,10 @@
 import platform
 import urllib
 import setuptools
-from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
+try:
+    from setuptools.command.bdist_wheel import bdist_wheel as _bdist_wheel
+except ImportError:
+    from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
 from packaging.version import parse
 
 try:

From 39c6d4e3690b5e14870d38685a9aca20905555dc Mon Sep 17 00:00:00 2001
From: Ilya Panfilov <ipanfilo@amd.com>
Date: Fri, 29 May 2026 21:23:14 -0400
Subject: [PATCH 2/2] Address review comments

---
 ci/_utils.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ci/_utils.sh b/ci/_utils.sh
index 35cba3203..9c0f4a847 100644
--- a/ci/_utils.sh
+++ b/ci/_utils.sh
@@ -234,7 +234,7 @@ start_message() {
     fi
     _rocm_path=`$REALPATH "$_rocm_path"`
     test -d "$_rocm_path" && echo "ROCm: $_rocm_path" || echo "ROCm path not found"
-    python --version
+    python3 --version
 }
 
 configure_omp_threads() {