Merge pull request #335259 from GaetanLepage/torchrl

python312Packages.torchrl: 0.4.0 -> 0.5.0

pkgs/development/python-modules/tensordict/default.nix

@@ -3,21 +3,27 @@
   buildPythonPackage,
   pythonOlder,
   fetchFromGitHub,
+
+  # build-system
   setuptools,
   torch,
-  wheel,
   which,
+
+  # dependencies
   cloudpickle,
   numpy,
+  orjson,
+
+  # checks
   h5py,
   pytestCheckHook,
+
   stdenv,
-  pythonAtLeast,
 }:
 
 buildPythonPackage rec {
   pname = "tensordict";
-  version = "0.4.0";
+  version = "0.5.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -26,19 +32,19 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "tensordict";
     rev = "refs/tags/v${version}";
-    hash = "sha256-wKEzNaaazGEkoElzp93RIlq/r5uRUdM7UyDy/DygIEc=";
+    hash = "sha256-jnRlN9gefR77pioIXf0qM1CP6EtpeQkBvVIecGkb/pw=";
   };
 
   build-system = [
     setuptools
     torch
-    wheel
     which
   ];
 
   dependencies = [
     cloudpickle
     numpy
+    orjson
     torch
   ];
 
@@ -55,8 +61,22 @@ buildPythonPackage rec {
   ];
 
   disabledTests =
-    # Hangs forever
-    [ "test_copy_onto" ]
+    [
+      # Hangs forever
+      "test_copy_onto"
+
+      # EOFError (MPI related)
+      # AssertionError: assert tensor(False)
+      # +  where tensor(False) = <built-in method all of Tensor object at 0x7ffe49bf87d0>()
+      "test_mp"
+
+      # torch._dynamo.exc.BackendCompilerFailed
+      # Requires a more recent version of triton
+      # Re-enable when https://github.com/NixOS/nixpkgs/pull/328247 is merged
+      "test_linear"
+      "test_seq"
+      "test_seq_lmbda"
+    ]
     ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-linux") [
       # RuntimeError: internal error
       "test_add_scale_sequence"
@@ -65,10 +85,19 @@ buildPythonPackage rec {
 
       # _queue.Empty errors in multiprocessing tests
       "test_isend"
+
+      # hangs forever
+      "test_map_iter_interrupt_early"
     ];
 
-  # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
-  disabledTestPaths = lib.optionals stdenv.isDarwin [ "test/test_distributed.py" ];
+  disabledTestPaths = lib.optionals stdenv.isDarwin [
+    # torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
+    # OpenMP support not found.
+    "test/test_compile.py"
+
+    # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
+    "test/test_distributed.py"
+  ];
 
   meta = {
     description = "Pytorch dedicated tensor container";
@@ -76,7 +105,5 @@ buildPythonPackage rec {
     homepage = "https://github.com/pytorch/tensordict";
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [ GaetanLepage ];
-    # No python 3.12 support yet: https://github.com/pytorch/rl/issues/2035
-    broken = pythonAtLeast "3.12";
   };
 }

pkgs/development/python-modules/torchrl/default.nix

@@ -3,13 +3,20 @@
   buildPythonPackage,
   pythonOlder,
   fetchFromGitHub,
+
+  # build-system
   ninja,
   setuptools,
-  wheel,
   which,
+
+  # dependencies
   cloudpickle,
   numpy,
+  packaging,
+  tensordict,
   torch,
+
+  # optional-dependencies
   ale-py,
   gym,
   pygame,
@@ -30,19 +37,20 @@
   hydra-core,
   tensorboard,
   wandb,
-  packaging,
-  tensordict,
+
+  # checks
   imageio,
   pytest-rerunfailures,
   pytestCheckHook,
   pyyaml,
   scipy,
+
   stdenv,
 }:
 
 buildPythonPackage rec {
   pname = "torchrl";
-  version = "0.4.0";
+  version = "0.5.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -51,13 +59,12 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "rl";
     rev = "refs/tags/v${version}";
-    hash = "sha256-8wSyyErqveP9zZS/UGvWVBYyylu9BuA447GEjXIzBIk=";
+    hash = "sha256-uDpOdOuHTqKFKspHOpl84kD9adEKZjvO2GnYuL27H5c=";
   };
 
   build-system = [
     ninja
     setuptools
-    wheel
     which
   ];
 
@@ -69,7 +76,7 @@ buildPythonPackage rec {
     torch
   ];
 
-  passthru.optional-dependencies = {
+  optional-dependencies = {
     atari = [
       ale-py
       gym
@@ -117,6 +124,7 @@ buildPythonPackage rec {
 
   nativeCheckInputs =
     [
+      h5py
       gymnasium
       imageio
       pytest-rerunfailures
@@ -125,9 +133,9 @@ buildPythonPackage rec {
       scipy
       torchvision
     ]
-    ++ passthru.optional-dependencies.atari
-    ++ passthru.optional-dependencies.gym-continuous
-    ++ passthru.optional-dependencies.rendering;
+    ++ optional-dependencies.atari
+    ++ optional-dependencies.gym-continuous
+    ++ optional-dependencies.rendering;
 
   disabledTests = [
     # mujoco.FatalError: an OpenGL platform library has not been loaded into this process, this most likely means that a valid OpenGL context has not been created before mjr_makeContext was called
@@ -155,7 +163,7 @@ buildPythonPackage rec {
     "test_distributed_collector_updatepolicy"
     "test_timeit"
 
-    # On a 24
+    # On a 24 threads system
     # assert torch.get_num_threads() == max(1, init_threads - 3)
     # AssertionError: assert 23 == 21
     "test_auto_num_threads"
@@ -168,6 +176,5 @@ buildPythonPackage rec {
     license = lib.licenses.mit;
     maintainers = with lib.maintainers; [ GaetanLepage ];
     # ~3k tests fail with: RuntimeError: internal error
-    broken = stdenv.isLinux && stdenv.isAarch64;
   };
 }