[Setup] Make platform-agnostic wheels

ghstack-source-id: 4889c03adeb9950a162d12963eda48e9b6d88513
Pull Request resolved: #2847

Author: vmoens

.github/scripts/td_script_agnostic.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+export TORCHRL_BUILD_VERSION=0.8.0
+export NO_CPP_BINARIES=1
+
+pip install git+https://github.com/pytorch/tensordict.git -U

.github/workflows/build-wheels-agnostic.yml

@@ -0,0 +1,48 @@
+name: Build Agnostic Wheels
+
+on:
+  pull_request:
+  push:
+    branches:
+      - nightly
+      - main
+      - release/*
+    tags:
+        # NOTE: Binary build pipelines should only get triggered on release candidate builds
+        # Release candidate tags look like: v1.11.0-rc1
+        - v[0-9]+.[0-9]+.[0-9]+-rc[0-9]+
+  workflow_dispatch:
+
+permissions:
+  id-token: write
+  contents: read
+
+jobs:
+  generate-matrix:
+    uses: pytorch/test-infra/.github/workflows/generate_binary_build_matrix.yml@main
+    with:
+      package-type: wheel
+      os: linux
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+  build:
+    needs: generate-matrix
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - repository: pytorch/rl
+            smoke-test-script: test/smoke_test.py
+            package-name: torchrl
+    name: pytorch/rl
+    uses: pytorch/test-infra/.github/workflows/build_wheels_linux.yml@main
+    with:
+      repository: ${{ matrix.repository }}
+      ref: ""
+      test-infra-repository: pytorch/test-infra
+      test-infra-ref: main
+      build-matrix: ${{ needs.generate-matrix.outputs.matrix }}
+      package-name: ${{ matrix.package-name }}
+      smoke-test-script: ${{ matrix.smoke-test-script }}
+      trigger-event: ${{ github.event_name }}
+      env-var-script: .github/scripts/td_script_agnostic.sh

setup.py

@@ -118,7 +118,26 @@ def run(self):
 #         return None
 
 
+def strtobool(val):
+    """Convert a string representation of truth to true (1) or false (0).
+
+    True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values
+    are 'n', 'no', 'f', 'false', 'off', and '0'.  Raises ValueError if
+    'val' is anything else.
+    """
+    val = val.lower()
+    if val in (1, True, "y", "yes", "t", "true", "on", "1"):
+        return 1
+    elif val in (0, False, "n", "no", "f", "false", "off", "0"):
+        return 0
+    else:
+        raise ValueError(f"invalid truth value {val!r}")
+
+
 def get_extensions():
+    no_cpp = strtobool(os.getenv("NO_CPP_BINARIES", "0"))
+    if no_cpp:
+        return []
     extension = CppExtension
 
     extra_link_args = []

torchrl/_extension.py

@@ -52,8 +52,10 @@ def _is_nightly(version):
 
 else:
     EXTENSION_WARNING = (
-        "Failed to import torchrl C++ binaries. Some modules (eg, prioritized replay buffers) may not work with your installation. "
-        "This is likely due to a discrepancy between your package version and the PyTorch version. Make sure both are compatible. "
-        "Usually, torchrl majors follow the pytorch majors within a few days around the release. "
+        "Failed to import torchrl C++ binaries. Some modules (e.g., prioritized replay buffers) may not work with your installation. "
+        "This could be because you are using a platform-agnostic version of TorchRL, which does not include C++ binaries. "
+        "If a more specific version is available for your platform, consider installing it for full functionality. "
+        "Additionally, ensure that your TorchRL version is compatible with your PyTorch version. "
+        "TorchRL major versions typically align with PyTorch major versions shortly after their release. "
         "For instance, TorchRL 0.5 requires PyTorch 2.4.0, and TorchRL 0.6 requires PyTorch 2.5.0."
     )