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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   pythonOlder,
   5   fetchFromGitHub,
   6   setuptools,
   7   wheel,
   8   aiofiles,
   9   aiohttp,
  10   importlib-metadata,
  11   nest-asyncio,
  12   psutil,
  13   pyyaml,
  14   torch,
  15   typing-extensions,
  16   pytest-asyncio,
  17   pytestCheckHook,
  18   pythonAtLeast,
  19   stdenv,
  20 }:
  21
  22 buildPythonPackage rec {
  23   pname = "torchsnapshot";
  24   version = "0.1.0";
  25   pyproject = true;
  26
  27   disabled = pythonOlder "3.7";
  28
  29   src = fetchFromGitHub {
  30     owner = "pytorch";
  31     repo = "torchsnapshot";
  32     rev = "refs/tags/${version}";
  33     hash = "sha256-F8OaxLH8BL6MPNLFv1hBuVmeEdnEQ5w2Qny6by1wP6k=";
  34   };
  35
  36   build-system = [
  37     setuptools
  38     wheel
  39   ];
  40
  41   dependencies = [
  42     aiofiles
  43     aiohttp
  44     importlib-metadata
  45     nest-asyncio
  46     psutil
  47     pyyaml
  48     torch
  49     typing-extensions
  50   ];
  51
  52   pythonImportsCheck = [ "torchsnapshot" ];
  53
  54   nativeCheckInputs = [
  55     pytest-asyncio
  56     pytestCheckHook
  57   ];
  58
  59   meta = with lib; {
  60     description = "Performant, memory-efficient checkpointing library for PyTorch applications, designed with large, complex distributed workloads in mind";
  61     homepage = "https://github.com/pytorch/torchsnapshot/";
  62     changelog = "https://github.com/pytorch/torchsnapshot/releases/tag/${version}";
  63     license = licenses.bsd3;
  64     maintainers = with maintainers; [ GaetanLepage ];
  65     broken =
  66       # https://github.com/pytorch/torchsnapshot/issues/175
  67       pythonAtLeast "3.12"
  68       # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
  69       || stdenv.hostPlatform.isDarwin;
  70   };
  71 }