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