python3Packages.torchsnapshot: cleanup, add missing dependency (#537563)

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

@@ -1,5 +1,6 @@
 {
   lib,
+  stdenv,
   buildPythonPackage,
   fetchFromGitHub,
 
@@ -11,6 +12,7 @@
   aiohttp,
   importlib-metadata,
   nest-asyncio,
+  numpy,
   psutil,
   pyyaml,
   torch,
@@ -21,15 +23,16 @@
   pytestCheckHook,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "torchsnapshot";
   version = "0.1.0";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "pytorch";
     repo = "torchsnapshot";
-    tag = version;
+    tag = finalAttrs.version;
     hash = "sha256-F8OaxLH8BL6MPNLFv1hBuVmeEdnEQ5w2Qny6by1wP6k=";
   };
 
@@ -52,6 +55,7 @@ buildPythonPackage rec {
     aiohttp
     importlib-metadata
     nest-asyncio
+    numpy
     psutil
     pyyaml
     torch
@@ -69,12 +73,17 @@ buildPythonPackage rec {
     # torch.distributed.elastic.multiprocessing.errors.ChildFailedError:
     # AssertionError: "Socket Timeout" does not match "wait timeout after 5000ms
     "test_linear_barrier_timeout"
+  ]
+  ++ lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
+    # aarch64-linux fails cpuinfo test, because /sys/devices/system/cpu/ does not exist in the sandbox:
+    # RuntimeError: Failed to initialize cpuinfo!
+    "test_tensor_copy"
   ];
 
   meta = {
     description = "Performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind";
     homepage = "https://github.com/pytorch/torchsnapshot/";
-    changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${version}";
+    changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.bsd3;
     maintainers = with lib.maintainers; [ GaetanLepage ];
     badPlatforms = [
@@ -82,4 +91,4 @@ buildPythonPackage rec {
       lib.systems.inspect.patterns.isDarwin
     ];
   };
-}
+})