pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 fetchFromGitHub,
6 fetchpatch,
7
8 # build-system
9 setuptools,
10
11 # dependencies
12 addict,
13 distutils,
14 matplotlib,
15 numpy,
16 opencv4,
17 pyyaml,
18 rich,
19 termcolor,
20 yapf,
21
22 # tests
23 bitsandbytes,
24 coverage,
25 dvclive,
26 lion-pytorch,
27 lmdb,
28 mlflow,
29 parameterized,
30 pytestCheckHook,
31 transformers,
32 writableTmpDirAsHomeHook,
33}:
34
35buildPythonPackage rec {
36 pname = "mmengine";
37 version = "0.10.7";
38 pyproject = true;
39
40 src = fetchFromGitHub {
41 owner = "open-mmlab";
42 repo = "mmengine";
43 tag = "v${version}";
44 hash = "sha256-hQnwenuxHQwl+DwQXbIfsKlJkmcRvcHV1roK7q2X1KA=";
45 };
46
47 patches = [
48 # Explicitly disable weights_only in torch.load calls
49 # https://github.com/open-mmlab/mmengine/pull/1650
50 (fetchpatch {
51 name = "torch-2.6.0-compat.patch";
52 url = "https://github.com/open-mmlab/mmengine/pull/1650/commits/c21b8431b2c625560a3866c65328cff0380ba1f8.patch";
53 hash = "sha256-SLr030IdYD9wM/jPJuZd+Dr1jjFx/5/YkJj/IwhnNQg=";
54 })
55 ];
56
57 postPatch =
58 # Fails in python >= 3.13
59 # exec(compile(f.read(), version_file, "exec")) does not populate the locals() namesp
60 # In python 3.13, the locals() dictionary in a function does not automatically update with
61 # changes made by exec().
62 # https://peps.python.org/pep-0558/
63 ''
64 substituteInPlace setup.py \
65 --replace-fail \
66 "return locals()['__version__']" \
67 "return '${version}'"
68 ''
69 + ''
70 substituteInPlace tests/test_config/test_lazy.py \
71 --replace-fail "import numpy.compat" ""
72 '';
73
74 build-system = [ setuptools ];
75
76 dependencies = [
77 addict
78 distutils
79 matplotlib
80 numpy
81 opencv4
82 pyyaml
83 rich
84 termcolor
85 yapf
86 ];
87
88 pythonImportsCheck = [ "mmengine" ];
89
90 nativeCheckInputs = [
91 bitsandbytes
92 coverage
93 dvclive
94 lion-pytorch
95 lmdb
96 mlflow
97 parameterized
98 pytestCheckHook
99 transformers
100 writableTmpDirAsHomeHook
101 ];
102
103 preCheck =
104 # Otherwise, the backprop hangs forever. More precisely, this exact line:
105 # https://github.com/open-mmlab/mmengine/blob/02f80e8bdd38f6713e04a872304861b02157905a/tests/test_runner/test_activation_checkpointing.py#L46
106 # Solution suggested in https://github.com/pytorch/pytorch/issues/91547#issuecomment-1370011188
107 ''
108 export MKL_NUM_THREADS=1
109 '';
110
111 disabledTestPaths = [
112 # Require unpackaged aim
113 "tests/test_visualizer/test_vis_backend.py::TestAimVisBackend"
114
115 # Cannot find SSL certificate
116 # _pygit2.GitError: OpenSSL error: failed to load certificates: error:00000000:lib(0)::reason(0)
117 "tests/test_visualizer/test_vis_backend.py::TestDVCLiveVisBackend"
118
119 # AttributeError: type object 'MagicMock' has no attribute ...
120 "tests/test_fileio/test_backends/test_petrel_backend.py::TestPetrelBackend"
121 ]
122 ++ lib.optionals stdenv.hostPlatform.isDarwin [
123 # RuntimeError: attempt to insert nil object from objects[1]
124 "tests/test_visualizer/test_visualizer.py::TestVisualizer::test_draw_featmap"
125 "tests/test_visualizer/test_visualizer.py::TestVisualizer::test_show"
126
127 # AssertionError: torch.bfloat16 != torch.float32
128 "tests/test_runner/test_amp.py::TestAmp::test_autocast"
129
130 # ValueError: User specified autocast device_type must be cuda or cpu, but got mps
131 "tests/test_runner/test_runner.py::TestRunner::test_test"
132 "tests/test_runner/test_runner.py::TestRunner::test_val"
133 ];
134
135 disabledTests = [
136 # Require network access
137 "test_fileclient"
138 "test_http_backend"
139 "test_misc"
140
141 # RuntimeError
142 "test_dump"
143 "test_deepcopy"
144 "test_copy"
145 "test_lazy_import"
146
147 # AssertionError: os is not <module 'os' (frozen)>
148 "test_lazy_module"
149 ]
150 ++ lib.optionals stdenv.hostPlatform.isDarwin [
151 # Fails when max-jobs is set to use fewer processes than cores
152 # for example `AssertionError: assert 14 == 4`
153 "test_setup_multi_processes"
154 ];
155
156 # torch.distributed.DistNetworkError: The server socket has failed to bind.
157 __darwinAllowLocalNetworking = true;
158
159 meta = {
160 description = "Library for training deep learning models based on PyTorch";
161 homepage = "https://github.com/open-mmlab/mmengine";
162 changelog = "https://github.com/open-mmlab/mmengine/releases/tag/v${version}";
163 license = with lib.licenses; [ asl20 ];
164 maintainers = with lib.maintainers; [ rxiao ];
165 };
166}