pyrox.dev
1{
2 lib,
3 stdenv,
4 fetchFromGitHub,
5
6 ## wandb-core
7 buildGoModule,
8 git,
9 versionCheckHook,
10
11 ## gpu-stats
12 rustPlatform,
13
14 ## wandb
15 buildPythonPackage,
16 replaceVars,
17
18 # build-system
19 hatchling,
20
21 # dependencies
22 click,
23 docker-pycreds,
24 gitpython,
25 platformdirs,
26 protobuf,
27 psutil,
28 pydantic,
29 pyyaml,
30 requests,
31 sentry-sdk,
32 setproctitle,
33 setuptools,
34 pythonOlder,
35 eval-type-backport,
36 typing-extensions,
37
38 # tests
39 pytestCheckHook,
40 azure-core,
41 azure-containerregistry,
42 azure-identity,
43 azure-storage-blob,
44 bokeh,
45 boto3,
46 coverage,
47 flask,
48 google-cloud-artifact-registry,
49 google-cloud-compute,
50 google-cloud-storage,
51 hypothesis,
52 jsonschema,
53 kubernetes,
54 kubernetes-asyncio,
55 matplotlib,
56 moviepy,
57 pandas,
58 parameterized,
59 pillow,
60 plotly,
61 pyfakefs,
62 pyte,
63 pytest-asyncio,
64 pytest-cov-stub,
65 pytest-mock,
66 pytest-timeout,
67 pytest-xdist,
68 rdkit,
69 responses,
70 scikit-learn,
71 soundfile,
72 tenacity,
73 torch,
74 torchvision,
75 tqdm,
76 writableTmpDirAsHomeHook,
77}:
78
79let
80 version = "0.19.10";
81 src = fetchFromGitHub {
82 owner = "wandb";
83 repo = "wandb";
84 tag = "v${version}";
85 hash = "sha256-SptjImAK0hDAr7UZjUMWnAVB1bxyzoATlyrPek5Tm48=";
86 };
87
88 gpu-stats = rustPlatform.buildRustPackage {
89 pname = "gpu-stats";
90 version = "0.3.0";
91 inherit src;
92
93 sourceRoot = "${src.name}/gpu_stats";
94
95 useFetchCargoVendor = true;
96 cargoHash = "sha256-kkWvTLduxFVIEvi4e65QQ7S0kHTRJ8XW028o430L91s=";
97
98 checkFlags = [
99 # fails in sandbox
100 "--skip=gpu_amd::tests::test_gpu_amd_new"
101 ];
102
103 nativeInstallCheckInputs = [
104 versionCheckHook
105 ];
106 versionCheckProgram = "${placeholder "out"}/bin/gpu_stats";
107 versionCheckProgramArg = "--version";
108 doInstallCheck = true;
109
110 meta = {
111 mainProgram = "gpu_stats";
112 };
113 };
114
115 wandb-core = buildGoModule rec {
116 pname = "wandb-core";
117 inherit src version;
118
119 sourceRoot = "${src.name}/core";
120
121 # hardcode the `gpu_stats` binary path.
122 postPatch = ''
123 substituteInPlace pkg/monitor/gpuresourcemanager.go \
124 --replace-fail \
125 'cmdPath, err := getGPUCollectorCmdPath()' \
126 'cmdPath, err := "${lib.getExe gpu-stats}", error(nil)'
127 '';
128
129 vendorHash = null;
130
131 nativeBuildInputs = [
132 git
133 ];
134
135 nativeInstallCheckInputs = [
136 versionCheckHook
137 ];
138 versionCheckProgramArg = "--version";
139 doInstallCheck = true;
140
141 checkFlags =
142 let
143 skippedTests = [
144 # gpu sampling crashes in the sandbox
145 "TestSystemMonitor_BasicStateTransitions"
146 "TestSystemMonitor_RepeatedCalls"
147 "TestSystemMonitor_UnexpectedTransitions"
148 "TestSystemMonitor_FullCycle"
149 ];
150 in
151 [ "-skip=^${lib.concatStringsSep "$|^" skippedTests}$" ];
152
153 __darwinAllowLocalNetworking = true;
154
155 meta.mainProgram = "wandb-core";
156 };
157in
158
159buildPythonPackage rec {
160 pname = "wandb";
161 pyproject = true;
162
163 inherit src version;
164
165 patches = [
166 # Replace git paths
167 (replaceVars ./hardcode-git-path.patch {
168 git = lib.getExe git;
169 })
170 ];
171
172 # Hard-code the path to the `wandb-core` binary in the code.
173 postPatch = ''
174 substituteInPlace wandb/util.py \
175 --replace-fail \
176 'bin_path = pathlib.Path(__file__).parent / "bin" / "wandb-core"' \
177 'bin_path = pathlib.Path("${lib.getExe wandb-core}")'
178 '';
179
180 env = {
181 # Prevent the install script to try building and embedding the `gpu_stats` and `wandb-core`
182 # binaries in the wheel.
183 # Their path have been patched accordingly in the `wandb-core` and `wanbd` source codes.
184 # https://github.com/wandb/wandb/blob/v0.18.5/hatch_build.py#L37-L47
185 WANDB_BUILD_SKIP_GPU_STATS = true;
186 WANDB_BUILD_UNIVERSAL = true;
187 };
188
189 build-system = [
190 hatchling
191 ];
192
193 dependencies =
194 [
195 click
196 docker-pycreds
197 gitpython
198 platformdirs
199 protobuf
200 psutil
201 pydantic
202 pyyaml
203 requests
204 sentry-sdk
205 setproctitle
206 # setuptools is necessary since pkg_resources is required at runtime.
207 setuptools
208 ]
209 ++ lib.optionals (pythonOlder "3.10") [
210 eval-type-backport
211 ]
212 ++ lib.optionals (pythonOlder "3.12") [
213 typing-extensions
214 ];
215
216 __darwinAllowLocalNetworking = true;
217
218 nativeCheckInputs = [
219 pytestCheckHook
220 azure-core
221 azure-containerregistry
222 azure-identity
223 azure-storage-blob
224 boto3
225 bokeh
226 coverage
227 flask
228 google-cloud-artifact-registry
229 google-cloud-compute
230 google-cloud-storage
231 hypothesis
232 jsonschema
233 kubernetes
234 kubernetes-asyncio
235 matplotlib
236 moviepy
237 pandas
238 parameterized
239 pillow
240 plotly
241 pyfakefs
242 pyte
243 pytest-asyncio
244 pytest-cov-stub
245 pytest-mock
246 pytest-timeout
247 pytest-xdist
248 rdkit
249 responses
250 scikit-learn
251 soundfile
252 tenacity
253 torch
254 torchvision
255 tqdm
256 writableTmpDirAsHomeHook
257 ];
258
259 # test_matplotlib_image_with_multiple_axes may take >60s
260 pytestFlagsArray = [
261 "--timeout=1024"
262 ];
263
264 disabledTestPaths = [
265 # Require docker access
266 "tests/system_tests"
267
268 # broke somewhere between sentry-sdk 2.15.0 and 2.22.0
269 "tests/unit_tests/test_analytics/test_sentry.py"
270 ];
271
272 disabledTests =
273 [
274 # Probably failing because of lack of internet access
275 # AttributeError: module 'wandb.sdk.launch.registry' has no attribute 'azure_container_registry'. Did you mean: 'elastic_container_registry'?
276 "test_registry_from_uri"
277
278 # Require docker
279 "test_get_requirements_section_pyproject"
280 "test_local_custom_env"
281 "test_local_custom_port"
282 "test_local_default"
283
284 # Expects python binary to be named `python3` but nix provides `python3.12`
285 # AssertionError: assert ['python3.12', 'main.py'] == ['python3', 'main.py']
286 "test_get_entrypoint"
287
288 # Require internet access
289 "test_audio_refs"
290 "test_bind_image"
291 "test_check_cors_configuration"
292 "test_check_wandb_version"
293 "test_from_path_project_type"
294 "test_image_accepts_bounding_boxes"
295 "test_image_accepts_bounding_boxes_optional_args"
296 "test_image_accepts_masks"
297 "test_image_accepts_masks_without_class_labels"
298 "test_image_seq_to_json"
299 "test_max_images"
300 "test_media_keys_escaped_as_glob_for_publish"
301 "test_parse_path"
302 "test_parse_project_path"
303 "test_translates_azure_err_to_normal_err"
304
305 # tests assertion if filesystem is compressed
306 "test_artifact_file_cache_cleanup"
307
308 # Tries to access a storage disk but there are none in the sandbox
309 # psutil.test_disk_out() returns None
310 "test_disk_in"
311 "test_disk_out"
312
313 # AssertionError: assert is_available('http://localhost:9400/metrics')
314 "test_dcgm"
315
316 # Error in the moviepy package:
317 # TypeError: must be real number, not NoneType
318 "test_video_numpy_mp4"
319
320 # AssertionError: assert not _IS_INTERNAL_PROCESS
321 "test_disabled_can_pickle"
322 "test_disabled_context_manager"
323 "test_mode_disabled"
324
325 # AssertionError: "one of name or plugin needs to be specified"
326 "test_opener_works_across_filesystem_boundaries"
327 "test_md5_file_hashes_on_mounted_filesystem"
328
329 # AttributeError: 'bytes' object has no attribute 'read'
330 "test_rewinds_on_failure"
331 "test_smoke"
332 "test_handles_multiple_calls"
333
334 # wandb.sdk.launch.errors.LaunchError: Found invalid name for agent MagicMock
335 "test_monitor_preempted"
336 "test_monitor_failed"
337 "test_monitor_running"
338 "test_monitor_job_deleted"
339
340 # Timeout >1024.0s
341 "test_log_media_prefixed_with_multiple_slashes"
342 "test_log_media_saves_to_run_directory"
343 "test_log_media_with_path_traversal"
344 ]
345 ++ lib.optionals stdenv.hostPlatform.isDarwin [
346 # AssertionError: assert not copy2_mock.called
347 "test_copy_or_overwrite_changed_no_copy"
348
349 # Fatal Python error: Aborted
350 "test_convert_plots"
351 "test_gpu_apple"
352 "test_image_from_matplotlib_with_image"
353 "test_make_plot_media_from_matplotlib_with_image"
354 "test_make_plot_media_from_matplotlib_without_image"
355 "test_matplotlib_contains_images"
356 "test_matplotlib_image"
357 "test_matplotlib_plotly_with_multiple_axes"
358 "test_matplotlib_to_plotly"
359 "test_plotly_from_matplotlib_with_image"
360
361 # RuntimeError: *** -[__NSPlaceholderArray initWithObjects:count:]: attempt to insert nil object from objects[1]
362 "test_wandb_image_with_matplotlib_figure"
363 ];
364
365 pythonImportsCheck = [ "wandb" ];
366
367 meta = {
368 description = "CLI and library for interacting with the Weights and Biases API";
369 homepage = "https://github.com/wandb/wandb";
370 changelog = "https://github.com/wandb/wandb/raw/v${version}/CHANGELOG.md";
371 license = lib.licenses.mit;
372 maintainers = with lib.maintainers; [ samuela ];
373 broken = gpu-stats.meta.broken || wandb-core.meta.broken;
374 };
375}