tjh.dev
1{ buildPythonPackage
2, fetchFromGitHub
3, lib
4, pandas
5, pytestCheckHook
6, torch
7, tensorboard
8, torchvision
9}:
10
11let
12 version = "0.3.1";
13 repo = fetchFromGitHub {
14 owner = "pytorch";
15 repo = "kineto";
16 rev = "v${version}";
17 hash = "sha256-Yg001XzOPDmz9wEP2b7Ggz/uU6x5PFzaaBeUBwWKFS0=";
18 };
19in
20buildPythonPackage 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 nativeCheckInputs = [ 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}