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