pyrox.dev
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 ]
79 ++ lib.optionals stdenv.hostPlatform.isLinux [ nvidia-ml-py ];
80
81 pythonRemoveDeps = [
82 "nvidia-ml-py3"
83 ]
84 ++ lib.optionals stdenv.hostPlatform.isDarwin [ "nvidia-ml-py" ];
85
86 __darwinAllowLocalNetworking = true;
87
88 nativeCheckInputs = [ pytestCheckHook ];
89
90 checkInputs = [
91 hypothesis
92 numpy
93 ];
94
95 disabledTests = [
96 # Flaky -- socket collision
97 "test_show_browser"
98 ];
99
100 # remove scalene directory to prevent pytest import confusion
101 preCheck = ''
102 rm -rf scalene
103 '';
104
105 pythonImportsCheck = [ "scalene" ];
106
107 meta = with lib; {
108 description = "High-resolution, low-overhead CPU, GPU, and memory profiler for Python with AI-powered optimization suggestions";
109 homepage = "https://github.com/plasma-umass/scalene";
110 changelog = "https://github.com/plasma-umass/scalene/releases/tag/v${version}";
111 mainProgram = "scalene";
112 license = licenses.asl20;
113 maintainers = with maintainers; [ sarahec ];
114 badPlatforms = [
115 # The scalene doesn't seem to account for arm64 linux
116 "aarch64-linux"
117
118 # On darwin, builds 1) assume aarch64 and 2) mistakenly compile one part as
119 # x86 and the other as arm64 then tries to link them into a single binary
120 # which fails.
121 "x86_64-darwin"
122 "aarch64-darwin"
123 ];
124 };
125}