pkgs/development/python-modules/torch-tb-profiler/default.nix

   1 { buildPythonPackage
   2 , fetchFromGitHub
   3 , lib
   4 , pandas
   5 , pytestCheckHook
   6 , torch
   7 , tensorboard
   8 , torchvision
   9 }:
  10
  11 let
  12   version = "0.3.1";
  13   repo = fetchFromGitHub {
  14     owner = "pytorch";
  15     repo = "kineto";
  16     rev = "v${version}";
  17     hash = "sha256-Yg001XzOPDmz9wEP2b7Ggz/uU6x5PFzaaBeUBwWKFS0=";
  18   };
  19 in
  20 buildPythonPackage rec {
  21   pname = "torch_tb_profiler";
  22   inherit version;
  23   format = "setuptools";
  24
  25   # See https://discourse.nixos.org/t/extracting-sub-directory-from-fetchgit-or-fetchurl-or-any-derivation/8830.
  26   src = "${repo}/tb_plugin";
  27
  28   propagatedBuildInputs = [ pandas tensorboard ];
  29
  30   checkInputs = [ pytestCheckHook torch torchvision ];
  31
  32   disabledTests = [
  33     # Tests that attempt to access the filesystem in naughty ways.
  34     "test_profiler_api_without_gpu"
  35     "test_tensorboard_end2end"
  36     "test_tensorboard_with_path_prefix"
  37     "test_tensorboard_with_symlinks"
  38   ];
  39
  40   pythonImportsCheck = [ "torch_tb_profiler" ];
  41
  42   meta = with lib; {
  43     description = "PyTorch Profiler TensorBoard Plugin";
  44     homepage = "https://github.com/pytorch/kineto";
  45     license = licenses.bsd3;
  46     maintainers = with maintainers; [ samuela ];
  47   };
  48 }