[CI] Fix failing CI runs for release (#1206)

.github/scripts/version_script.bat

@@ -1,3 +1,53 @@
 @echo off
 set TENSORDICT_BUILD_VERSION=0.7.0
 echo TENSORDICT_BUILD_VERSION is set to %TENSORDICT_BUILD_VERSION%
+
+@echo on
+
+set VC_VERSION_LOWER=17
+set VC_VERSION_UPPER=18
+if "%VC_YEAR%" == "2019" (
+    set VC_VERSION_LOWER=16
+    set VC_VERSION_UPPER=17
+)
+if "%VC_YEAR%" == "2017" (
+    set VC_VERSION_LOWER=15
+    set VC_VERSION_UPPER=16
+)
+
+for /f "usebackq tokens=*" %%i in (`"%ProgramFiles(x86)%\Microsoft Visual Studio\Installer\vswhere.exe" -legacy -products * -version [%VC_VERSION_LOWER%^,%VC_VERSION_UPPER%^) -property installationPath`) do (
+    if exist "%%i" if exist "%%i\VC\Auxiliary\Build\vcvarsall.bat" (
+        set "VS15INSTALLDIR=%%i"
+        set "VS15VCVARSALL=%%i\VC\Auxiliary\Build\vcvarsall.bat"
+        goto vswhere
+    )
+)
+
+:vswhere
+if "%VSDEVCMD_ARGS%" == "" (
+    call "%VS15VCVARSALL%" x64 || exit /b 1
+) else (
+    call "%VS15VCVARSALL%" x64 %VSDEVCMD_ARGS% || exit /b 1
+)
+
+@echo on
+
+if "%CU_VERSION%" == "xpu" call "C:\Program Files (x86)\Intel\oneAPI\setvars.bat"
+
+set DISTUTILS_USE_SDK=1
+
+set args=%1
+shift
+:start
+if [%1] == [] goto done
+set args=%args% %1
+shift
+goto start
+
+:done
+if "%args%" == "" (
+    echo Usage: vc_env_helper.bat [command] [args]
+    echo e.g. vc_env_helper.bat cl /c test.cpp
+)
+
+%args% || exit /b 1

packaging/wheel/relocate.py

@@ -0,0 +1,138 @@
+# 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.
+
+"""Helper script to package wheels and relocate binaries."""
+
+import glob
+import hashlib
+
+# Standard library imports
+import os
+import os.path as osp
+import shutil
+import sys
+import zipfile
+from base64 import urlsafe_b64encode
+
+HERE = osp.dirname(osp.abspath(__file__))
+PACKAGE_ROOT = osp.dirname(osp.dirname(HERE))
+
+
+def rehash(path, blocksize=1 << 20):
+    """Return (hash, length) for path using hashlib.sha256()"""
+    h = hashlib.sha256()
+    length = 0
+    with open(path, "rb") as f:
+        while block := f.read(blocksize):
+            length += len(block)
+            h.update(block)
+    digest = "sha256=" + urlsafe_b64encode(h.digest()).decode("latin1").rstrip("=")
+    # unicode/str python2 issues
+    return (digest, str(length))  # type: ignore
+
+
+def unzip_file(file, dest):
+    """Decompress zip `file` into directory `dest`."""
+    with zipfile.ZipFile(file, "r") as zip_ref:
+        zip_ref.extractall(dest)
+
+
+def is_program_installed(basename):
+    """
+    Return program absolute path if installed in PATH.
+    Otherwise, return None
+    On macOS systems, a .app is considered installed if
+    it exists.
+    """
+    if sys.platform == "darwin" and basename.endswith(".app") and osp.exists(basename):
+        return basename
+
+    for path in os.environ["PATH"].split(os.pathsep):
+        abspath = osp.join(path, basename)
+        if osp.isfile(abspath):
+            return abspath
+
+
+def find_program(basename):
+    """
+    Find program in PATH and return absolute path
+    Try adding .exe or .bat to basename on Windows platforms
+    (return None if not found)
+    """
+    names = [basename]
+    if os.name == "nt":
+        # Windows platforms
+        extensions = (".exe", ".bat", ".cmd", ".dll")
+        if not basename.endswith(extensions):
+            names = [basename + ext for ext in extensions] + [basename]
+    for name in names:
+        path = is_program_installed(name)
+        if path:
+            return path
+
+
+def compress_wheel(output_dir, wheel, wheel_dir, wheel_name):
+    """Create RECORD file and compress wheel distribution."""
+    # ("Update RECORD file in wheel")
+    dist_info = glob.glob(osp.join(output_dir, "*.dist-info"))[0]
+    record_file = osp.join(dist_info, "RECORD")
+
+    with open(record_file, "w") as f:
+        for root, _, files in os.walk(output_dir):
+            for this_file in files:
+                full_file = osp.join(root, this_file)
+                rel_file = osp.relpath(full_file, output_dir)
+                if full_file == record_file:
+                    f.write(f"{rel_file},,\n")
+                else:
+                    digest, size = rehash(full_file)
+                    f.write(f"{rel_file},{digest},{size}\n")
+
+    # ("Compressing wheel")
+    base_wheel_name = osp.join(wheel_dir, wheel_name)
+    shutil.make_archive(base_wheel_name, "zip", output_dir)
+    os.remove(wheel)
+    shutil.move(f"{base_wheel_name}.zip", wheel)
+    shutil.rmtree(output_dir)
+
+
+def patch_win():
+    # Get dumpbin location
+    dumpbin = find_program("dumpbin")
+    if dumpbin is None:
+        raise FileNotFoundError(
+            "Dumpbin was not found in the system, please make sure that is available on the PATH."
+        )
+
+    # Find wheel
+    # ("Finding wheels...")
+    wheels = glob.glob(osp.join(PACKAGE_ROOT, "dist", "*.whl"))
+    output_dir = osp.join(PACKAGE_ROOT, "dist", ".wheel-process")
+
+    for wheel in wheels:
+        print(f"processing {wheel}")
+        if osp.exists(output_dir):
+            shutil.rmtree(output_dir)
+        print(f"creating output directory {output_dir}")
+        os.makedirs(output_dir)
+
+        # ("Unzipping wheel...")
+        wheel_file = osp.basename(wheel)
+        wheel_dir = osp.dirname(wheel)
+        # (f"{wheel_file}")
+        wheel_name, _ = osp.splitext(wheel_file)
+        print(f"unzipping {wheel} in {output_dir}")
+        unzip_file(wheel, output_dir)
+        print("compressing wheel")
+        compress_wheel(output_dir, wheel, wheel_dir, wheel_name)
+
+
+if __name__ == "__main__":
+    if sys.platform == "linux":
+        pass
+    elif sys.platform == "win32":
+        patch_win()
+    else:
+        raise NotImplementedError

setup.cfg

@@ -15,6 +15,7 @@ per-file-ignores =
     test/smoke_test_deps.py: F401
     test_*.py: E731, E266, TOR101
     tutorials/*/**.py: T201
+    packaging/*/**.py: T201
 exclude = venv
 extend-select = B901, C401, C408, C409
 

tensordict/nn/distributions/composite.py

@@ -304,9 +304,7 @@ def maybe_deterministic_sample(dist):
             if hasattr(dist, "deterministic_sample"):
                 return dist.deterministic_sample
             else:
-                from tensordict.nn.probabilistic import (
-                    DETERMINISTIC_REGISTER,
-                )
+                from tensordict.nn.probabilistic import DETERMINISTIC_REGISTER
 
                 # Fallbacks
                 tdist = type(dist)

test/test_compile.py

@@ -6,6 +6,7 @@
 import contextlib
 import importlib.util
 import inspect
+import platform
 from pathlib import Path
 from typing import Any, Callable
 
@@ -43,6 +44,8 @@
 
 _v2_5 = TORCH_VERSION >= version.parse("2.5.0")
 
+_IS_OSX = platform.system() == "Darwin"
+
 
 def test_vmap_compile():
     # Since we monkey patch vmap we need to make sure compile is happy with it
@@ -952,6 +955,10 @@ def to_numpy(tensor):
 @pytest.mark.skipif(
     TORCH_VERSION <= version.parse("2.4.1"), reason="requires torch>=2.5"
 )
+@pytest.mark.skipif(
+    (TORCH_VERSION <= version.parse("2.7.0")) and _IS_OSX,
+    reason="requires torch>=2.7 ons OSX",
+)
 @pytest.mark.parametrize("compiled", [False, True])
 class TestCudaGraphs:
     @pytest.fixture(scope="class", autouse=True)