diff --git a/pkgs/development/python-modules/tensordict/default.nix b/pkgs/development/python-modules/tensordict/default.nix index bd35eff65707..fff726f5beeb 100644 --- a/pkgs/development/python-modules/tensordict/default.nix +++ b/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) = () + "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"; }; } diff --git a/pkgs/development/python-modules/torchrl/default.nix b/pkgs/development/python-modules/torchrl/default.nix index 638f1844eaba..3023380920c7 100644 --- a/pkgs/development/python-modules/torchrl/default.nix +++ b/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; }; }