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