Support editable install with setuptools>=80.0.0 (#9428)

Co-authored-by: Bhavya Bahl <[email protected]>

Author: tengyifei

.devcontainer/tpu-contributor/devcontainer.json

@@ -23,4 +23,4 @@
       ]
     }
   }
-}
+}

.devcontainer/tpu-internal/devcontainer.json

@@ -8,7 +8,7 @@
   ],
   "containerEnv": {
     "BAZEL_REMOTE_CACHE": "1",
-    "SILO_NAME": "cache-silo-${localEnv:USER}-tpuvm"
+    "SILO_NAME": "cache-silo-${localEnv:USER}-tpuvm-312"
   },
   "initializeCommand": "docker pull us-central1-docker.pkg.dev/tpu-pytorch-releases/docker/development:tpu",
   "customizations": {

pyproject.toml

@@ -0,0 +1,48 @@
+[build-system]
+# These are the packages required to run `setup.py` in an isolated environment.
+# Pip will install these *before* executing the build backend.
+requires = [
+    "setuptools>=42",
+    "wheel",
+    "requests",
+    "numpy",
+    "pyyaml",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-xla"
+description = "XLA bridge for PyTorch"
+readme = "README.md"
+authors = [
+    { name = "PyTorch/XLA Dev Team", email = "[email protected]" },
+]
+license = { file = "LICENSE" }
+requires-python = ">=3.10"
+classifiers = [
+    "Development Status :: 5 - Production/Stable",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: BSD License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+keywords = ["pytorch", "xla", "tpu", "deep learning", "compiler"]
+
+# This tells build tools to get this info from setup.py instead of this file.
+dynamic = [
+    "version",
+    "dependencies",
+    "optional-dependencies",
+    "entry-points",
+    "scripts"
+]
+
+[project.urls]
+Homepage = "https://github.com/pytorch/xla"
+Repository = "https://github.com/pytorch/xla"
+"Bug Tracker" = "https://github.com/pytorch/xla/issues"

scripts/build_developer.sh

@@ -80,6 +80,7 @@ cd $_SCRIPT_DIR/..
 pip uninstall torch_xla torchax torch_xla2 -y
 
 # Build the wheel too, which is useful for other testing purposes.
+rm -f torch_xla.egg-info/SOURCES.txt
 python3 setup.py bdist_wheel
 
 # Link the source files for local development.

setup.py

@@ -52,6 +52,7 @@
 import contextlib
 import distutils.ccompiler
 import distutils.command.clean
+import importlib.util
 import os
 import re
 import requests
@@ -61,7 +62,13 @@
 import tempfile
 import zipfile
 
-import build_util
+# This gloop imports build_util.py such that it works in Python 3.12's isolated
+# build environment while also not contaminating sys.path which breaks bdist_wheel.
+_PROJECT_DIR = os.path.dirname(os.path.abspath(__file__))
+_build_util_path = os.path.join(_PROJECT_DIR, 'build_util.py')
+spec = importlib.util.spec_from_file_location('build_util', _build_util_path)
+build_util = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(build_util)
 
 import platform
 
@@ -270,15 +277,21 @@ def __init__(self, bazel_target):
 class BuildBazelExtension(build_ext.build_ext):
   """A command that runs Bazel to build a C/C++ extension."""
 
-  def run(self):
-    for ext in self.extensions:
-      self.bazel_build(ext)
-    command.build_ext.build_ext.run(self)  # type: ignore
+  def build_extension(self, ext: Extension) -> None:
+    """
+    This method is called by setuptools to build a single extension.
+    We override it to implement our custom Bazel build logic.
+    """
+    if not isinstance(ext, BazelExtension):
+      # If it's not our custom extension type, let setuptools handle it.
+      super().build_extension(ext)
+      return
 
-  def bazel_build(self, ext):
+    # 1. Ensure the temporary build directory exists
     if not os.path.exists(self.build_temp):
       os.makedirs(self.build_temp)
 
+    # 2. Prepare the Bazel command
     bazel_argv = [
         'bazel', 'build', ext.bazel_target,
         f"--symlink_prefix={os.path.join(self.build_temp, 'bazel-')}"
@@ -288,22 +301,31 @@ def bazel_build(self, ext):
     if build_cpp_tests:
       bazel_argv.append('//:cpp_tests')
 
-    import torch
-    cxx_abi = os.getenv('CXX_ABI') or getattr(torch._C,
-                                              '_GLIBCXX_USE_CXX11_ABI', None)
-    if cxx_abi is not None:
-      bazel_argv.append(f'--cxxopt=-D_GLIBCXX_USE_CXX11_ABI={int(cxx_abi)}')
+    cxx_abi = os.getenv('CXX_ABI')
+    if cxx_abi is None:
+      try:
+        import torch
+        cxx_abi = getattr(torch._C, '_GLIBCXX_USE_CXX11_ABI', None)
+      except:
+        pass
+    if cxx_abi is None:
+      # Default to building with C++11 ABI, which has been the case since PyTorch 2.7
+      cxx_abi = "1"
+    bazel_argv.append(f'--cxxopt=-D_GLIBCXX_USE_CXX11_ABI={int(cxx_abi)}')
 
     bazel_argv.extend(build_util.bazel_options_from_env())
 
+    # 3. Run the Bazel build
     self.spawn(bazel_argv)
 
+    # 4. Copy the output file to the location setuptools expects
     ext_bazel_bin_path = os.path.join(self.build_temp, 'bazel-bin', ext.relpath,
                                       ext.target_name)
     ext_dest_path = self.get_ext_fullpath(ext.name)
     ext_dest_dir = os.path.dirname(ext_dest_path)
     if not os.path.exists(ext_dest_dir):
       os.makedirs(ext_dest_dir)
+
     shutil.copyfile(ext_bazel_bin_path, ext_dest_path)
 
 
@@ -313,17 +335,28 @@ def bazel_build(self, ext):
   long_description = f.read()
 
 # Finds torch_xla and its subpackages
-packages_to_include = find_packages(include=['torch_xla*'])
-# Explicitly add torchax
-packages_to_include.extend(find_packages(where='torchax', include=['torchax*']))
+# 1. Find `torch_xla` and its subpackages automatically from the root.
+packages_to_include = find_packages(include=['torch_xla', 'torch_xla.*'])
+
+# 2. Explicitly find the contents of the nested `torchax` package.
+#    Find all sub-packages within the torchax directory (e.g., 'ops').
+torchax_source_dir = 'torchax/torchax'
+torchax_subpackages = find_packages(where=torchax_source_dir)
+#    Construct the full list of packages, starting with the top-level
+#    'torchax' and adding all the discovered sub-packages.
+packages_to_include.extend(['torchax'] +
+                           ['torchax.' + pkg for pkg in torchax_subpackages])
 
-# Map the top-level 'torchax' package name to its source location
-torchax_dir = os.path.join(cwd, 'torchax')
-package_dir_mapping = {'torch_xla': os.path.join(cwd, 'torch_xla')}
-package_dir_mapping['torchax'] = os.path.join(torchax_dir, 'torchax')
+# 3. The package_dir mapping explicitly tells setuptools where the 'torchax'
+#    package's source code begins. `torch_xla` source code is inferred.
+package_dir_mapping = {'torchax': torchax_source_dir}
 
 
 class Develop(develop.develop):
+  """
+  Custom develop command to build C++ extensions and create a .pth file
+  for a multi-package editable install.
+  """
 
   def run(self):
     # Build the C++ extension
@@ -348,44 +381,42 @@ def link_packages(self):
     (`python setup.py develop`). Nightly and release wheel builds work out of the box
     without egg-link/pth.
     """
+    import glob
+
     # Ensure paths like self.install_dir are set
     self.ensure_finalized()
 
-    # Get the site-packages directory
-    target_dir = self.install_dir
-
-    # Remove the standard .egg-link file
-    # It's usually named based on the distribution name
     dist_name = self.distribution.get_name()
-    egg_link_file = os.path.join(target_dir, dist_name + '.egg-link')
-    if os.path.exists(egg_link_file):
-      print(f"Removing default egg-link file: {egg_link_file}")
-      try:
-        os.remove(egg_link_file)
-      except OSError as e:
-        print(f"Warning: Could not remove {egg_link_file}: {e}")
-
-    # Create our custom .pth file with specific paths
-    cwd = os.path.dirname(__file__)
-    # Path containing 'torch_xla' package source: ROOT
-    path_for_torch_xla = os.path.abspath(cwd)
-    # Path containing 'torchax' package source: ROOT/torchax
-    path_for_torchax = os.path.abspath(os.path.join(cwd, 'torchax'))
-
-    paths_to_add = {path_for_torch_xla, path_for_torchax}
-
-    # Construct a suitable .pth filename (PEP 660 style is good practice)
-    version = self.distribution.get_version()
-    # Sanitize name and version for filename (replace runs of non-alphanumeric chars with '-')
-    sanitized_name = re.sub(r"[^a-zA-Z0-9.]+", "_", dist_name)
-    sanitized_version = re.sub(r"[^a-zA-Z0-9.]+", "_", version)
-    pth_filename = os.path.join(
-        target_dir, f"__editable_{sanitized_name}_{sanitized_version}.pth")
-
-    # Ensure site-packages exists
-    os.makedirs(target_dir, exist_ok=True)
-
-    # Write the paths to the .pth file, one per line
+    install_cmd = self.get_finalized_command('install')
+    target_dir = install_cmd.install_lib
+    assert target_dir is not None
+
+    # Use glob to robustly find and remove the conflicting files.
+    # This is safer than trying to guess the exact sanitized filename.
+    safe_name_part = re.sub(r"[^a-zA-Z0-9]+", "_", dist_name)
+
+    for pattern in [
+        # Remove `.pth` files generated in Python 3.12.
+        f"__editable__.*{safe_name_part}*.pth",
+        f"__editable___*{safe_name_part}*_finder.py",
+        # Also remove the legacy egg-link format.
+        f"{dist_name}.egg-link"
+    ]:
+      for filepath in glob.glob(os.path.join(target_dir, pattern)):
+        print(f"Cleaning up conflicting install file: {filepath}")
+        with contextlib.suppress(OSError):
+          os.remove(filepath)
+
+    # Finally, create our own simple, multi-path .pth file.
+    # We name it simply, e.g., "torch_xla.pth".
+    pth_filename = os.path.join(target_dir, f"{dist_name}.pth")
+
+    project_root = os.path.dirname(os.path.abspath(__file__))
+    paths_to_add = {
+        project_root,  # For `torch_xla`
+        os.path.abspath(os.path.join(project_root, 'torchax')),  # For `torchax`
+    }
+
     with open(pth_filename, "w", encoding='utf-8') as f:
       for path in sorted(paths_to_add):
         f.write(path + "\n")

torch_xla/experimental/gru.py

@@ -99,7 +99,7 @@ def __init__(self, *args, **kwargs):
     super().__init__(*args, **kwargs)
 
   def forward(self, input, hx=None):
-    """
+    r"""
     Args:
         input: Tensor of shape (seq_len, batch, input_size)
         hx: Optional initial hidden state of shape (num_layers, batch, hidden_size).