[Setup] statically link _C extension against the Python library

ghstack-source-id: 26f374e
Pull-Request-resolved: #1304
(cherry picked from commit af17524)

Author: Vincent Moens

.github/scripts/version_script.sh

@@ -7,7 +7,6 @@ export MACOSX_DEPLOYMENT_TARGET=15.0
 
 ${CONDA_RUN} pip install --upgrade pip
 
-${CONDA_RUN} conda install conda-forge::rust -y
 # for orjson
 export UNSAFE_PYO3_BUILD_FREE_THREADED=1
 

.github/scripts/win-pre-script.bat

@@ -7,7 +7,6 @@ if "%CONDA_RUN%"=="" (
 
 :: Run the pip install command
 %CONDA_RUN% conda install conda-forge::pybind11 -y
-%CONDA_RUN% conda install conda-forge::rust -y
 
 :: Check if the installation was successful
 if errorlevel 1 (

tensordict/csrc/CMakeLists.txt

@@ -36,13 +36,16 @@ else()
     )
 endif()
 
+find_package(Python COMPONENTS Development.Static)
+target_link_libraries(_C ${Python_STATIC_LIBRARIES})
+
 target_include_directories(_C PRIVATE ${PROJECT_SOURCE_DIR})
 
-if(APPLE OR WIN32) # Check if the target OS is OSX/macOS
-    target_link_libraries(_C PRIVATE pybind11::module)
-else()
-    target_link_libraries(_C PRIVATE Python3::Python pybind11::module)
-endif()
+#if(APPLE OR WIN32) # Check if the target OS is OSX/macOS
+target_link_libraries(_C PRIVATE pybind11::module)
+#else()
+#    target_link_libraries(_C PRIVATE Python3::Python pybind11::module)
+#endif()
 
 if(CMAKE_BUILD_TYPE STREQUAL "Debug")
     set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -O0 -fsanitize=address")

test/smoke_test.py

@@ -2,6 +2,11 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+import subprocess
+import sys
+from pathlib import Path
+
+_IS_LINUX = sys.platform.startswith("linux")
 
 
 def test_imports_deps():
@@ -23,6 +28,32 @@ def test_imports():
     print("version", tensordict.__version__)  # noqa
 
 
+def test_static_linking():
+    if not _IS_LINUX:
+        return
+    # Locate _C.so
+    try:
+        import tensordict._C
+    except ImportError as e:
+        raise RuntimeError(f"Failed to import tensordict._C: {e}")
+    # Get the path to _C.so
+    _C_path = Path(tensordict._C.__file__)
+    if not _C_path.exists():
+        raise RuntimeError(f"_C.so not found at {_C_path}")
+    # Run ldd on _C.so
+    try:
+        output = subprocess.check_output(["ldd", str(_C_path)]).decode("utf-8")
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"Failed to run ldd on {_C_path}: {e}")
+    # Check if libpython is dynamically linked
+    for line in output.splitlines():
+        if "libpython" in line and "=>" in line and "not found" not in line:
+            raise RuntimeError(
+                f"tensordict/_C.so is dynamically linked against {line.strip()}"
+            )
+    print("Test passed: tensordict/_C.so does not show dynamic linkage to libpython.")  # noqa
+
+
 if __name__ == "__main__":
     test_imports_deps()
     test_imports()