trying to simplify

vmoens · vmoens · commit c5a2e8913410 · 2025-07-29T11:12:16.000+01:00
diff --git a/.github/workflows/nightly_build.yml b/.github/workflows/nightly_build.yml
@@ -60,20 +60,14 @@ jobs:
         run: |
           export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
           python3 -mpip install --pre torch --index-url https://download.pytorch.org/whl/nightly/${{ matrix.cuda_support[1] }}
-      - name: Prepare nightly build
-        env:
-          TORCHRL_NIGHTLY: 1
-        run: |
-          export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
-          python3 packaging/prepare_nightly_build.py
       - name: Build TorchRL Nightly
         env:
           TORCHRL_NIGHTLY: 1
         run: |
           rm -r dist || true
           export PATH="/opt/python/${{ matrix.python_version[1] }}/bin:$PATH"
-          python3 -mpip install wheel
-          python3 setup.py bdist_wheel
+          python3 -mpip install build wheel
+          ./build_nightly.sh
           find dist -name '*whl' -exec bash -c ' mv $0 ${0/linux/manylinux1}' {} \;
       # pytorch/pytorch binaries are also manylinux_2_17 compliant but they
       # pretend that they're manylinux1 compliant so we do the same.
@@ -233,20 +227,14 @@ jobs:
         shell: bash
         run: |
           python3 -mpip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cpu -U
-      - name: Prepare nightly build
-        env:
-          TORCHRL_NIGHTLY: 1
-        shell: bash
-        run: |
-          python3 packaging/prepare_nightly_build.py
       - name: Build TorchRL nightly
         env:
           TORCHRL_NIGHTLY: 1
         shell: bash
         run: |
           rm -r dist || true
-          python3 -mpip install wheel
-          python3 setup.py bdist_wheel
+          python3 -mpip install build wheel
+          ./build_nightly.sh
       - name: Upload wheel for the test-wheel job
         uses: actions/upload-artifact@v4
         with:
diff --git a/build_nightly.sh b/build_nightly.sh
@@ -0,0 +1,45 @@
+#!/bin/bash
+set -e
+
+# Check if we're in a nightly build
+if [ "$TORCHRL_NIGHTLY" != "1" ]; then
+    echo "Not a nightly build, exiting"
+    exit 0
+fi
+
+echo "Starting nightly build process..."
+
+# Create backups of original files
+cp pyproject.toml pyproject.toml.backup
+cp version.txt version.txt.backup
+
+# Function to restore original files
+restore_files() {
+    echo "Restoring original files..."
+    mv pyproject.toml.backup pyproject.toml
+    mv version.txt.backup version.txt
+}
+
+# Set up trap to restore files on exit (success or failure)
+trap restore_files EXIT
+
+# Modify pyproject.toml for nightly build
+echo "Modifying pyproject.toml for nightly build..."
+sed -i.bak 's/name = "torchrl"/name = "torchrl-nightly"/' pyproject.toml
+
+# Replace tensordict dependency with tensordict-nightly
+echo "Replacing tensordict with tensordict-nightly..."
+sed -i.bak 's/"tensordict[^"]*"/"tensordict-nightly"/g' pyproject.toml
+
+# Clean up sed backup files
+rm -f pyproject.toml.bak
+
+# Set nightly version (YYYY.MM.DD format)
+echo "Setting nightly version..."
+echo "$(date +%Y.%m.%d)" > version.txt
+
+# Build the package
+echo "Building nightly package..."
+python -m build
+
+echo "Nightly build completed successfully!" 
diff --git a/packaging/prepare_nightly_build.py b/packaging/prepare_nightly_build.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import os
+import re
 from pathlib import Path
 
 
@@ -20,9 +21,11 @@ def prepare_nightly_build():
     with open(pyproject_path) as f:
         content = f.read()
 
-    # Replace tensordict dependency with tensordict-nightly
-    if "tensordict>=0.9.0,<0.10.0" in content:
-        content = content.replace("tensordict>=0.9.0,<0.10.0", "tensordict-nightly")
+    # Replace tensordict dependency with tensordict-nightly using regex
+    # This pattern matches "tensordict" followed by any version constraints
+    tensordict_pattern = r'tensordict[^,\]]*'
+    if re.search(tensordict_pattern, content):
+        content = re.sub(tensordict_pattern, "tensordict-nightly", content)
         print("Replaced tensordict with tensordict-nightly in pyproject.toml")
     else:
         print("tensordict dependency not found in pyproject.toml")
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,6 +20,7 @@ classifiers = [
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
     "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
     "Operating System :: OS Independent",
     "Development Status :: 4 - Beta",
     "Intended Audience :: Developers",
diff --git a/setup.py b/setup.py
@@ -1,78 +1,11 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-#
-# This source code is licensed under the MIT license found in the
-# LICENSE file in the root directory of this source tree.
-import glob
-import logging
 import os
-import shutil
 import sys
+import glob
+import logging
 from pathlib import Path
-
-from setuptools import Command, setup
+from setuptools import setup
 from torch.utils.cpp_extension import BuildExtension, CppExtension
 
-cwd = os.path.dirname(os.path.abspath(__file__))
-ROOT_DIR = Path(__file__).parent.resolve()
-
-
-def write_version_file(version):
-    version_path = os.path.join(cwd, "torchrl", "version.py")
-    logging.info(f"Writing version file to: {version_path}")
-    logging.info(f"Version to write: {version}")
-
-    # Get PyTorch version during build
-    try:
-        import torch
-
-        pytorch_version = torch.__version__
-    except ImportError:
-        pytorch_version = "unknown"
-
-    # Get git sha
-    try:
-        import subprocess
-
-        sha = (
-            subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd)
-            .decode("ascii")
-            .strip()
-        )
-    except Exception:
-        sha = "Unknown"
-
-    with open(version_path, "w") as f:
-        f.write(f"__version__ = '{version}'\n")
-        f.write(f"git_version = {repr(sha)}\n")
-        f.write(f"pytorch_version = '{pytorch_version}'\n")
-
-    logging.info("Version file written successfully")
-
-
-class clean(Command):
-    user_options = []
-
-    def initialize_options(self):
-        pass
-
-    def finalize_options(self):
-        pass
-
-    def run(self):
-        # Remove torchrl extension
-        for path in (ROOT_DIR / "torchrl").glob("**/*.so"):
-            logging.info(f"removing '{path}'")
-            path.unlink()
-        # Remove build directory
-        build_dirs = [
-            ROOT_DIR / "build",
-        ]
-        for path in build_dirs:
-            if path.exists():
-                logging.info(f"removing '{path}' (and everything under it)")
-                shutil.rmtree(str(path), ignore_errors=True)
-
-
 def get_extensions():
     extension = CppExtension
 
@@ -104,7 +37,7 @@ def get_extensions():
     cpp_files = glob.glob(os.path.join(extensions_dir, "*.cpp"))
     sources = [os.path.relpath(f) for f in cpp_files]
 
-    include_dirs = ["."]
+    include_dirs = [".", "torchrl/csrc"]
     python_include_dir = os.getenv("PYTHON_INCLUDE_DIR")
     if python_include_dir is not None:
         include_dirs.append(python_include_dir)
@@ -120,127 +53,13 @@ def get_extensions():
 
     return ext_modules
 
-
-def _main():
-    # Always use "torchrl" as the project name for GitHub discovery
-    # The version will be read from pyproject.toml
-
-    # Handle nightly builds
-    is_nightly = (
-        any("nightly" in arg for arg in sys.argv) or os.getenv("TORCHRL_NIGHTLY") == "1"
-    )
-    logging.info(f"is_nightly: {is_nightly}")
-
-    # Read version from version.txt
-    version_txt = os.path.join(cwd, "version.txt")
-    with open(version_txt) as f:
-        base_version = f.readline().strip()
-
-    if os.getenv("TORCHRL_BUILD_VERSION"):
-        version = os.getenv("TORCHRL_BUILD_VERSION")
-    elif is_nightly:
-        from datetime import date
-
-        today = date.today()
-        version = f"{today.year}.{today.month}.{today.day}"
-        logging.info(f"Using nightly version: {version}")
-        # Update version.txt for nightly builds
-        with open(version_txt, "w") as f:
-            f.write(f"{version}\n")
-    else:
-        # For regular builds, append git hash for development versions
-        try:
-            import subprocess
-
-            git_sha = (
-                subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=cwd)
-                .decode("ascii")
-                .strip()[:7]
-            )
-            version = f"{base_version}+{git_sha}"
-            logging.info(f"Using development version: {version}")
-        except Exception:
-            version = base_version
-            logging.info(f"Using base version: {version}")
-
-    # Always write the version file to ensure it's up to date
-    write_version_file(version)
-    logging.info(f"Building torchrl-{version}")
-
-    # Verify the version file was written correctly
-    try:
-        with open(os.path.join(cwd, "torchrl", "version.py")) as f:
-            content = f.read()
-            if f"__version__ = '{version}'" in content:
-                logging.info(f"Version file correctly contains: {version}")
-            else:
-                logging.error(
-                    f"Version file does not contain expected version: {version}"
-                )
-    except Exception as e:
-        logging.error(f"Failed to verify version file: {e}")
-
-    # Handle package name for nightly builds
-    if is_nightly:
-        package_name = "torchrl-nightly"  # Use torchrl-nightly for PyPI uploads
-    else:
-        package_name = "torchrl"  # Use torchrl for regular builds and GitHub discovery
-
+def main():
     setup_kwargs = {
-        "name": package_name,
-        # Only C++ extension configuration
         "ext_modules": get_extensions(),
-        "cmdclass": {
-            "build_ext": BuildExtension.with_options(),
-            "clean": clean,
-        },
-        "zip_safe": False,
-        "package_data": {
-            "torchrl": ["version.py"],
-        },
-        "include_package_data": True,
-        "packages": ["torchrl"],
+        "cmdclass": {"build_ext": BuildExtension.with_options()},
     }
-
-    # Handle nightly tensordict dependency override
-    if is_nightly:
-        setup_kwargs["install_requires"] = [
-            "torch>=2.1.0",
-            "numpy",
-            "packaging",
-            "cloudpickle",
-            "tensordict-nightly",
-        ]
-
-    # Override pyproject.toml settings for nightly builds
-    if is_nightly:
-        # Add all the metadata from pyproject.toml but override the name
-        setup_kwargs.update(
-            {
-                "description": "A modular, primitive-first, python-first PyTorch library for Reinforcement Learning",
-                "long_description": (Path(__file__).parent / "README.md").read_text(
-                    encoding="utf8"
-                ),
-                "long_description_content_type": "text/markdown",
-                "author": "torchrl contributors",
-                "author_email": "vmoens@fb.com",
-                "url": "https://github.com/pytorch/rl",
-                "classifiers": [
-                    "Programming Language :: Python :: 3.9",
-                    "Programming Language :: Python :: 3.10",
-                    "Programming Language :: Python :: 3.11",
-                    "Programming Language :: Python :: 3.12",
-                    "Operating System :: OS Independent",
-                    "Development Status :: 4 - Beta",
-                    "Intended Audience :: Developers",
-                    "Intended Audience :: Science/Research",
-                    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-                ],
-            }
-        )
-
+    
     setup(**setup_kwargs)
 
-
 if __name__ == "__main__":
-    _main()
+    main()
diff --git a/torchrl/__init__.py b/torchrl/__init__.py
@@ -25,29 +25,7 @@
 try:
     from .version import __version__
 except ImportError:
-    # Fallback: try to read version from version.txt and append git hash
-    try:
-        import os
-        import subprocess
-
-        version_file = os.path.join(os.path.dirname(__file__), "..", "version.txt")
-        with open(version_file) as f:
-            base_version = f.read().strip()
-
-        # Try to get git hash
-        try:
-            git_sha = (
-                subprocess.check_output(
-                    ["git", "rev-parse", "HEAD"], cwd=os.path.dirname(version_file)
-                )
-                .decode("ascii")
-                .strip()[:7]
-            )
-            __version__ = f"{base_version}+{git_sha}"
-        except Exception:
-            __version__ = base_version
-    except Exception:
-        __version__ = "unknown"
+    __version__ = "0.0.0+unknown"
 
 try:
     from torch.compiler import is_dynamo_compiling
