[Setup] Making tensordict pytorch-agnostic

Vincent Moens · Vincent Moens · commit 0300a595488d · 2025-03-11T11:47:23.000Z
ghstack-source-id: b2a9c72 Pull Request resolved: #1256
diff --git a/.github/scripts/linux-pre-script.sh b/.github/scripts/linux-pre-script.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+
+yum update gcc
+yum update libstdc++
+
+${CONDA_RUN} conda install -c conda-forge pybind11 -y
diff --git a/.github/scripts/version_script.bat b/.github/scripts/version_script.bat
@@ -2,6 +2,14 @@
 set TENSORDICT_BUILD_VERSION=0.7.0
 echo TENSORDICT_BUILD_VERSION is set to %TENSORDICT_BUILD_VERSION%
 
+if "%CONDA_RUN%"=="" (
+    echo CONDA_RUN is not set. Please activate your conda environment or set CONDA_RUN.
+    exit /b 1
+)
+
+:: Run the pip install command
+%CONDA_RUN% conda install -c conda-forge pybind11 -y
+
 @echo on
 
 set VC_VERSION_LOWER=17
diff --git a/.github/scripts/version_script.sh b/.github/scripts/version_script.sh
@@ -1,3 +1,5 @@
 #!/bin/bash
 
 export TENSORDICT_BUILD_VERSION=0.7.0
+
+${CONDA_RUN} conda install -c conda-forge pybind11 -y
diff --git a/.github/scripts/win-pre-script.bat b/.github/scripts/win-pre-script.bat
@@ -0,0 +1,17 @@
+@echo off
+:: Check if CONDA_RUN is set, if not, set it to a default value
+if "%CONDA_RUN%"=="" (
+    echo CONDA_RUN is not set. Please activate your conda environment or set CONDA_RUN.
+    exit /b 1
+)
+
+:: Run the pip install command
+%CONDA_RUN% conda install -c conda-forge pybind11 -y
+
+:: Check if the installation was successful
+if errorlevel 1 (
+    echo Failed to install cmake and pybind11.
+    exit /b 1
+) else (
+    echo Successfully installed cmake and pybind11.
+)
diff --git a/.github/unittest/linux/scripts/environment.yml b/.github/unittest/linux/scripts/environment.yml
@@ -19,3 +19,6 @@ dependencies:
     - coverage
     - h5py
     - orjson
+    - ninja
+    - pybind11
+    - cmake
diff --git a/.github/unittest/linux_torchrec/scripts/environment.yml b/.github/unittest/linux_torchrec/scripts/environment.yml
@@ -18,3 +18,6 @@ dependencies:
     - coverage
     - h5py
     - orjson
+    - cmake
+    - ninja
+    - pybind11
diff --git a/.github/unittest/rl_linux_optdeps/scripts/environment.yml b/.github/unittest/rl_linux_optdeps/scripts/environment.yml
@@ -16,3 +16,6 @@ dependencies:
     - pyyaml
     - scipy
     - orjson
+    - cmake
+    - ninja
+    - pybind11
diff --git a/.github/workflows/build-wheels-linux.yml b/.github/workflows/build-wheels-linux.yml
@@ -33,6 +33,7 @@ jobs:
         include:
           - repository: pytorch/tensordict
             smoke-test-script: test/smoke_test.py
+            pre-script: .github/scripts/linux-pre-script.sh
             post-script: .github/scripts/linux-post-script.sh
             package-name: tensordict
     name: pytorch/tensordict
diff --git a/.github/workflows/build-wheels-windows.yml b/.github/workflows/build-wheels-windows.yml
@@ -32,7 +32,7 @@ jobs:
       matrix:
         include:
           - repository: pytorch/tensordict
-            pre-script: ""
+            pre-script: .github/scripts/win-pre-script.bat
             env-script: .github/scripts/version_script.bat
             post-script: "python packaging/wheel/relocate.py"
             smoke-test-script: test/smoke_test.py
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -6,6 +6,11 @@ possible.
 Install the library as suggested in the README. For advanced features,
 it is preferable to install the nightly built of pytorch.
 
+You will need the following packages to be installed:
+```bash
+pip install ninja cmake pybind11 -U
+```
+
 Make sure you install tensordict in develop mode by running
 ```
 python setup.py develop
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -3,3 +3,4 @@ include LICENSE
 
 recursive-exclude * __pycache__
 recursive-exclude * *.py[co]
+recursive-include tensordict *.so
diff --git a/setup.py b/setup.py
@@ -15,11 +15,23 @@
 from pathlib import Path
 from typing import List
 
-from setuptools import find_packages, setup
-from torch.utils.cpp_extension import BuildExtension, CppExtension
+from setuptools import Extension, find_packages, setup
+from setuptools.command.build_ext import build_ext
 
 ROOT_DIR = Path(__file__).parent.resolve()
 
+
+def get_python_executable():
+    # Check if we're running in a virtual environment
+    if "VIRTUAL_ENV" in os.environ:
+        # Get the virtual environment's Python executable
+        python_executable = os.path.join(os.environ["VIRTUAL_ENV"], "bin", "python")
+    else:
+        # Fall back to sys.executable
+        python_executable = sys.executable
+    return python_executable
+
+
 try:
     sha = (
         subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=ROOT_DIR)
@@ -69,7 +81,7 @@ def _get_pytorch_version(is_nightly, is_local):
         return "torch>=2.7.0.dev"
     if is_local:
         return "torch"
-    return "torch>=2.6.0"
+    return "torch>=2.5.0"
 
 
 def _get_packages():
@@ -99,51 +111,68 @@ def run(self):
                 shutil.rmtree(str(path), ignore_errors=True)
 
 
-def get_extensions():
-    extension = CppExtension
-
-    extra_link_args = []
-    extra_compile_args = {
-        "cxx": [
-            "-O3",
-            "-std=c++17",
-            "-fdiagnostics-color=always",
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=""):
+        super().__init__(name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+
+class CMakeBuild(build_ext):
+    def run(self):
+        for ext in self.extensions:
+            self.build_extension(ext)
+
+    def build_extension(self, ext):
+        # extdir = os.path.abspath(os.path.dirname(self.get_ext_fullpath(ext.name)))
+        is_develop = self.distribution.get_command_obj("develop").finalized
+        # Set the output directory based on the mode
+        if is_develop:
+            extdir = os.path.abspath(os.path.join(ROOT_DIR, "tensordict"))
+        else:
+            extdir = os.path.abspath(os.path.join(self.build_lib, "tensordict"))
+        cmake_args = [
+            f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}",
+            f"-DPYTHON_EXECUTABLE={get_python_executable()}",
+            f"-DPython3_EXECUTABLE={get_python_executable()}",
         ]
-    }
-    debug_mode = os.getenv("DEBUG", "0") == "1"
-    if debug_mode:
-        logging.info("Compiling in debug mode")
-        extra_compile_args = {
-            "cxx": [
-                "-O0",
-                "-fno-inline",
-                "-g",
-                "-std=c++17",
-                "-fdiagnostics-color=always",
-            ]
-        }
-        extra_link_args = ["-O0", "-g"]
-
-    this_dir = os.path.dirname(os.path.abspath(__file__))
-    extensions_dir = os.path.join(this_dir, "tensordict", "csrc")
-
-    extension_sources = {
-        os.path.join(extensions_dir, p)
-        for p in glob.glob(os.path.join(extensions_dir, "*.cpp"))
-    }
-    sources = list(extension_sources)
-
-    ext_modules = [
-        extension(
-            "tensordict._C",
-            sources,
-            include_dirs=[this_dir],
-            extra_compile_args=extra_compile_args,
-            extra_link_args=extra_link_args,
+        CONDA_PREFIX = os.environ.get("CONDA_PREFIX")
+        # if CONDA_PREFIX:
+        #     CMAKE_PREFIX_PATH = os.environ.get("CMAKE_PREFIX_PATH")
+        #     if CMAKE_PREFIX_PATH:
+        #         cmake_args.append(f"-DCMAKE_PREFIX_PATH={CONDA_PREFIX}:{CMAKE_PREFIX_PATH}")
+        #     else:
+        #         cmake_args.append(f"-DCMAKE_PREFIX_PATH={CONDA_PREFIX}")
+        if CONDA_PREFIX:
+            # Find pybind11
+            pybind11_dir = None
+            for config_file in ["pybind11Config.cmake", "pybind11-config.cmake"]:
+                config_path = glob.glob(
+                    os.path.join(CONDA_PREFIX, "**", config_file), recursive=True
+                )
+                if config_path:
+                    pybind11_dir = os.path.dirname(config_path[0])
+                    break
+            else:
+                raise RuntimeError(
+                    f"could not find any of 'pybind11Config.cmake', 'pybind11-config.cmake' in {CONDA_PREFIX}"
+                )
+            if pybind11_dir:
+                cmake_args.append(f"-DPYBIND11_DIR={pybind11_dir}")
+
+        build_args = []
+        if not os.path.exists(self.build_temp):
+            os.makedirs(self.build_temp)
+        subprocess.check_call(
+            ["cmake", ext.sourcedir] + cmake_args, cwd=self.build_temp
+        )
+        subprocess.check_call(
+            ["cmake", "--build", "."] + build_args, cwd=self.build_temp
         )
-    ]
 
-    return ext_modules
+
+def get_extensions():
+    extensions_dir = os.path.join(ROOT_DIR, "tensordict", "csrc")
+    return [CMakeExtension("tensordict._C", sourcedir=extensions_dir)]
 
 
 def _main(argv):
@@ -181,7 +210,7 @@ def _main(argv):
         ),
         ext_modules=get_extensions(),
         cmdclass={
-            "build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
+            "build_ext": CMakeBuild,
             "clean": clean,
         },
         install_requires=[
@@ -208,8 +237,13 @@ def _main(argv):
             "Programming Language :: Python :: 3.10",
             "Programming Language :: Python :: 3.11",
             "Programming Language :: Python :: 3.12",
+            "Programming Language :: Python :: 3.13",
             "Development Status :: 4 - Beta",
         ],
+        # include_package_data=True,
+        package_data={
+            "tensordict": ["*.so", "*.pyd"],
+        },
     )
 
 
diff --git a/tensordict/csrc/CMakeLists.txt b/tensordict/csrc/CMakeLists.txt
@@ -0,0 +1,26 @@
+cmake_minimum_required(VERSION 3.12)
+project(tensordict)
+
+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# Set the Python executable to the one from your virtual environment
+
+find_package(Python3 REQUIRED COMPONENTS Interpreter Development)
+find_package(pybind11 REQUIRED)
+
+file(GLOB SOURCES "*.cpp")
+
+add_library(_C MODULE ${SOURCES})
+
+set_target_properties(_C PROPERTIES
+    OUTPUT_NAME "_C"
+    PREFIX ""  # Remove 'lib' prefix
+    SUFFIX ".so"  # Ensure correct suffix for macOS/Linux
+)
+set_target_properties(_C PROPERTIES
+    LIBRARY_OUTPUT_DIRECTORY "${CMAKE_LIBRARY_OUTPUT_DIRECTORY}"
+)
+
+target_include_directories(_C PRIVATE ${PROJECT_SOURCE_DIR})
+target_link_libraries(_C PRIVATE Python3::Python pybind11::module)

-Original file line number
+Diff line change
     - pyyaml
     - scipy
     - orjson
 +    - cmake
 +    - ninja
 +    - pybind11