From b13426b24762d7dcdb97a93cf8ac18f5cdb1d6c1 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 8 Jan 2025 00:37:36 +0100
Subject: [PATCH] python312Packages.torchrl: skip flaky test on aarch64-linux

---
 .../python-modules/torchrl/default.nix        | 74 ++++++++++---------
 1 file changed, 40 insertions(+), 34 deletions(-)

diff --git a/pkgs/development/python-modules/torchrl/default.nix b/pkgs/development/python-modules/torchrl/default.nix
index d84e66034801..759ce4b4b5dc 100644
--- a/pkgs/development/python-modules/torchrl/default.nix
+++ b/pkgs/development/python-modules/torchrl/default.nix
@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
 
@@ -132,46 +133,51 @@ buildPythonPackage rec {
     ++ optional-dependencies.gym-continuous
     ++ optional-dependencies.rendering;
 
-  disabledTests = [
-    # torchrl is incompatible with gymnasium>=1.0
-    # https://github.com/pytorch/rl/discussions/2483
-    "test_resetting_strategies"
-    "test_torchrl_to_gym"
+  disabledTests =
+    [
+      # torchrl is incompatible with gymnasium>=1.0
+      # https://github.com/pytorch/rl/discussions/2483
+      "test_resetting_strategies"
+      "test_torchrl_to_gym"
 
-    # 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
-    "test_vecenvs_env"
+      # 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
+      "test_vecenvs_env"
 
-    # ValueError: Can't write images with one color channel.
-    "test_log_video"
+      # ValueError: Can't write images with one color channel.
+      "test_log_video"
 
-    # Those tests require the ALE environments (provided by unpackaged shimmy)
-    "test_collector_env_reset"
-    "test_gym"
-    "test_gym_fake_td"
-    "test_recorder"
-    "test_recorder_load"
-    "test_rollout"
-    "test_parallel_trans_env_check"
-    "test_serial_trans_env_check"
-    "test_single_trans_env_check"
-    "test_td_creation_from_spec"
-    "test_trans_parallel_env_check"
-    "test_trans_serial_env_check"
-    "test_transform_env"
+      # Those tests require the ALE environments (provided by unpackaged shimmy)
+      "test_collector_env_reset"
+      "test_gym"
+      "test_gym_fake_td"
+      "test_recorder"
+      "test_recorder_load"
+      "test_rollout"
+      "test_parallel_trans_env_check"
+      "test_serial_trans_env_check"
+      "test_single_trans_env_check"
+      "test_td_creation_from_spec"
+      "test_trans_parallel_env_check"
+      "test_trans_serial_env_check"
+      "test_transform_env"
 
-    # undeterministic
-    "test_distributed_collector_updatepolicy"
-    "test_timeit"
+      # undeterministic
+      "test_distributed_collector_updatepolicy"
+      "test_timeit"
 
-    # On a 24 threads system
-    # assert torch.get_num_threads() == max(1, init_threads - 3)
-    # AssertionError: assert 23 == 21
-    "test_auto_num_threads"
+      # On a 24 threads system
+      # assert torch.get_num_threads() == max(1, init_threads - 3)
+      # AssertionError: assert 23 == 21
+      "test_auto_num_threads"
 
-    # Flaky (hangs indefinitely on some CPUs)
-    "test_gae_multidim"
-    "test_gae_param_as_tensor"
-  ];
+      # Flaky (hangs indefinitely on some CPUs)
+      "test_gae_multidim"
+      "test_gae_param_as_tensor"
+    ]
+    ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
+      # Flaky (AssertionError)
+      "test_vecnorm_parallel_auto"
+    ];
 
   meta = {
     description = "Modular, primitive-first, python-first PyTorch library for Reinforcement Learning";