1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 cloudpickle,
7 cython,
8 hypothesis,
9 jinja2,
10 numpy,
11 nvidia-ml-py,
12 psutil,
13 pydantic,
14 pytestCheckHook,
15 pythonOlder,
16 rich,
17 setuptools-scm,
18 setuptools,
19}:
20
21let
22 heap-layers-src = fetchFromGitHub {
23 owner = "emeryberger";
24 repo = "heap-layers";
25 name = "Heap-Layers";
26 rev = "a2048eae91b531dc5d72be7a194e0b333c06bd4c";
27 sha256 = "sha256-vl3z30CBX7hav/DM/UE0EQ9lLxZF48tMJrYMXuSulyA=";
28 };
29
30 printf-src = fetchFromGitHub {
31 owner = "mpaland";
32 repo = "printf";
33 name = "printf";
34 tag = "v4.0.0";
35 sha256 = "sha256-tgLJNJw/dJGQMwCmfkWNBvHB76xZVyyfVVplq7aSJnI=";
36 };
37in
38
39buildPythonPackage rec {
40 pname = "scalene";
41 version = "1.5.52";
42 pyproject = true;
43 disabled = pythonOlder "3.9";
44
45 src = fetchFromGitHub {
46 owner = "plasma-umass";
47 repo = "scalene";
48 tag = "v${version}";
49 hash = "sha256-8WE/tR0tGwdNSPtieS90QAOFlS66h/JxaV2LvpZjx2E=";
50 };
51
52 patches = [
53 ./01-manifest-no-git.patch
54 ];
55
56 prePatch = ''
57 cp -r ${heap-layers-src} vendor/Heap-Layers
58 mkdir vendor/printf
59 cp ${printf-src}/printf.c vendor/printf/printf.cpp
60 cp -r ${printf-src}/* vendor/printf
61 sed -i 's/^#define printf printf_/\/\/&/' vendor/printf/printf.h
62 sed -i 's/^#define vsnprintf vsnprintf_/\/\/&/' vendor/printf/printf.h
63 '';
64
65 nativeBuildInputs = [
66 cython
67 setuptools
68 setuptools-scm
69 ];
70
71 propagatedBuildInputs = [
72 cloudpickle
73 jinja2
74 numpy
75 psutil
76 pydantic
77 rich
78 ] ++ lib.optionals stdenv.hostPlatform.isLinux [ nvidia-ml-py ];
79
80 pythonRemoveDeps = [
81 "nvidia-ml-py3"
82 ] ++ lib.optionals stdenv.hostPlatform.isDarwin [ "nvidia-ml-py" ];
83
84 __darwinAllowLocalNetworking = true;
85
86 nativeCheckInputs = [ pytestCheckHook ];
87
88 checkInputs = [
89 hypothesis
90 numpy
91 ];
92
93 disabledTests = [
94 # Flaky -- socket collision
95 "test_show_browser"
96 ];
97
98 # remove scalene directory to prevent pytest import confusion
99 preCheck = ''
100 rm -rf scalene
101 '';
102
103 pythonImportsCheck = [ "scalene" ];
104
105 meta = with lib; {
106 description = "High-resolution, low-overhead CPU, GPU, and memory profiler for Python with AI-powered optimization suggestions";
107 homepage = "https://github.com/plasma-umass/scalene";
108 changelog = "https://github.com/plasma-umass/scalene/releases/tag/v${version}";
109 mainProgram = "scalene";
110 license = licenses.asl20;
111 maintainers = with maintainers; [ sarahec ];
112 badPlatforms = [
113 # The scalene doesn't seem to account for arm64 linux
114 "aarch64-linux"
115
116 # On darwin, builds 1) assume aarch64 and 2) mistakenly compile one part as
117 # x86 and the other as arm64 then tries to link them into a single binary
118 # which fails.
119 "x86_64-darwin"
120 "aarch64-darwin"
121 ];
122 };
123}