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