trainlib 0.3.1__tar.gz → 0.6.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (109) hide show
  1. trainlib-0.6.1/PKG-INFO +243 -0
  2. trainlib-0.6.1/README.md +206 -0
  3. {trainlib-0.3.1 → trainlib-0.6.1}/pyproject.toml +22 -11
  4. trainlib-0.6.1/test/test_cli.py +160 -0
  5. trainlib-0.6.1/test/test_config.py +87 -0
  6. trainlib-0.6.1/test/test_console_display.py +101 -0
  7. trainlib-0.6.1/test/test_construction.py +191 -0
  8. trainlib-0.6.1/test/test_experiment_runs.py +85 -0
  9. trainlib-0.6.1/test/test_objectives.py +56 -0
  10. trainlib-0.6.1/test/test_plot.py +53 -0
  11. trainlib-0.6.1/test/test_prepared_domain.py +86 -0
  12. trainlib-0.6.1/test/test_records.py +165 -0
  13. trainlib-0.6.1/test/test_report_outputs.py +52 -0
  14. trainlib-0.6.1/test/test_trainer.py +107 -0
  15. trainlib-0.6.1/trainlib/__main__.py +35 -0
  16. trainlib-0.6.1/trainlib/channel.py +78 -0
  17. trainlib-0.6.1/trainlib/channels/__init__.py +23 -0
  18. trainlib-0.6.1/trainlib/channels/duration.py +21 -0
  19. trainlib-0.6.1/trainlib/channels/loss.py +39 -0
  20. trainlib-0.6.1/trainlib/channels/regression.py +94 -0
  21. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/cli/__init__.py +2 -2
  22. trainlib-0.6.1/trainlib/cli/config.py +50 -0
  23. trainlib-0.6.1/trainlib/cli/experiment.py +196 -0
  24. trainlib-0.6.1/trainlib/cli/train.py +69 -0
  25. trainlib-0.6.1/trainlib/components/loss/__init__.py +26 -0
  26. trainlib-0.6.1/trainlib/components/optimizer/__init__.py +13 -0
  27. trainlib-0.6.1/trainlib/construction.py +452 -0
  28. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/dataloader.py +74 -26
  29. trainlib-0.6.1/trainlib/dataloaders/__init__.py +20 -0
  30. trainlib-0.6.1/trainlib/dataloaders/memory.py +48 -0
  31. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/dataset.py +95 -100
  32. trainlib-0.6.1/trainlib/datasets/__init__.py +13 -0
  33. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/datasets/disk.py +8 -15
  34. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/datasets/memory.py +29 -17
  35. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domain.py +11 -26
  36. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domains/disk.py +6 -7
  37. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/domains/functional.py +2 -6
  38. trainlib-0.6.1/trainlib/domains/prepared.py +212 -0
  39. trainlib-0.6.1/trainlib/estimator.py +148 -0
  40. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/__init__.py +9 -4
  41. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/mlp.py +16 -72
  42. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/rnn.py +53 -139
  43. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/estimators/tdnn.py +8 -12
  44. trainlib-0.6.1/trainlib/experiment.py +430 -0
  45. trainlib-0.6.1/trainlib/experiments/__init__.py +18 -0
  46. trainlib-0.6.1/trainlib/experiments/regression.py +16 -0
  47. trainlib-0.6.1/trainlib/experiments/train.py +107 -0
  48. trainlib-0.6.1/trainlib/loss.py +172 -0
  49. trainlib-0.6.1/trainlib/losses/__init__.py +17 -0
  50. trainlib-0.6.1/trainlib/losses/composite.py +53 -0
  51. trainlib-0.6.1/trainlib/losses/criteria.py +78 -0
  52. trainlib-0.6.1/trainlib/losses/simple.py +125 -0
  53. trainlib-0.6.1/trainlib/measure.py +95 -0
  54. trainlib-0.6.1/trainlib/measures/__init__.py +11 -0
  55. trainlib-0.6.1/trainlib/measures/regression.py +49 -0
  56. trainlib-0.6.1/trainlib/objective.py +39 -0
  57. trainlib-0.6.1/trainlib/optimizer.py +274 -0
  58. trainlib-0.6.1/trainlib/optimizers/__init__.py +15 -0
  59. trainlib-0.6.1/trainlib/optimizers/adversarial.py +40 -0
  60. trainlib-0.6.1/trainlib/optimizers/simple.py +49 -0
  61. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/plotter.py +187 -101
  62. trainlib-0.6.1/trainlib/plotters/__init__.py +16 -0
  63. trainlib-0.6.1/trainlib/plotters/regression.py +53 -0
  64. trainlib-0.6.1/trainlib/record.py +486 -0
  65. trainlib-0.6.1/trainlib/registry.py +103 -0
  66. trainlib-0.6.1/trainlib/reporter.py +251 -0
  67. trainlib-0.6.1/trainlib/reporters/__init__.py +20 -0
  68. trainlib-0.6.1/trainlib/run.py +590 -0
  69. trainlib-0.6.1/trainlib/tabulator.py +203 -0
  70. trainlib-0.6.1/trainlib/tabulators/__init__.py +8 -0
  71. trainlib-0.6.1/trainlib/trainer.py +820 -0
  72. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/transform.py +8 -9
  73. trainlib-0.6.1/trainlib/utils/__init__.py +0 -0
  74. trainlib-0.6.1/trainlib/utils/artifacts.py +108 -0
  75. trainlib-0.6.1/trainlib/utils/config.py +146 -0
  76. trainlib-0.6.1/trainlib/utils/console.py +527 -0
  77. trainlib-0.6.1/trainlib/utils/custom.mplstyle +50 -0
  78. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/job.py +4 -7
  79. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/module.py +14 -6
  80. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/op.py +2 -2
  81. trainlib-0.6.1/trainlib/utils/plot.py +102 -0
  82. trainlib-0.6.1/trainlib/utils/statistics.py +188 -0
  83. trainlib-0.6.1/trainlib/utils/tables.py +476 -0
  84. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/text.py +2 -0
  85. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/type.py +23 -1
  86. trainlib-0.6.1/trainlib.egg-info/PKG-INFO +243 -0
  87. trainlib-0.6.1/trainlib.egg-info/SOURCES.txt +95 -0
  88. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/entry_points.txt +1 -0
  89. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/requires.txt +2 -0
  90. trainlib-0.3.1/PKG-INFO +0 -136
  91. trainlib-0.3.1/README.md +0 -100
  92. trainlib-0.3.1/trainlib/__main__.py +0 -18
  93. trainlib-0.3.1/trainlib/cli/train.py +0 -164
  94. trainlib-0.3.1/trainlib/dataloaders/__init__.py +0 -12
  95. trainlib-0.3.1/trainlib/dataloaders/memory.py +0 -17
  96. trainlib-0.3.1/trainlib/datasets/__init__.py +0 -12
  97. trainlib-0.3.1/trainlib/estimator.py +0 -201
  98. trainlib-0.3.1/trainlib/trainer.py +0 -639
  99. trainlib-0.3.1/trainlib/utils/plot.py +0 -38
  100. trainlib-0.3.1/trainlib.egg-info/PKG-INFO +0 -136
  101. trainlib-0.3.1/trainlib.egg-info/SOURCES.txt +0 -40
  102. {trainlib-0.3.1 → trainlib-0.6.1}/setup.cfg +0 -0
  103. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/__init__.py +0 -0
  104. {trainlib-0.3.1/trainlib/domains → trainlib-0.6.1/trainlib/components}/__init__.py +0 -0
  105. {trainlib-0.3.1/trainlib/utils → trainlib-0.6.1/trainlib/domains}/__init__.py +0 -0
  106. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/map.py +0 -0
  107. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib/utils/session.py +1 -1
  108. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/dependency_links.txt +0 -0
  109. {trainlib-0.3.1 → trainlib-0.6.1}/trainlib.egg-info/top_level.txt +0 -0
@@ -0,0 +1,243 @@
1
+ Metadata-Version: 2.4
2
+ Name: trainlib
3
+ Version: 0.6.1
4
+ Summary: Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training.
5
+ License-Expression: MIT
6
+ Project-URL: Homepage, https://doc.olog.io/trainlib
7
+ Project-URL: Documentation, https://doc.olog.io/trainlib
8
+ Project-URL: Repository, https://git.olog.io/olog/trainlib
9
+ Project-URL: Issues, https://git.olog.io/olog/trainlib/issues
10
+ Keywords: machine-learning
11
+ Classifier: Programming Language :: Python
12
+ Classifier: Operating System :: OS Independent
13
+ Classifier: Development Status :: 3 - Alpha
14
+ Classifier: Intended Audience :: Developers
15
+ Classifier: Intended Audience :: End Users/Desktop
16
+ Requires-Python: >=3.12
17
+ Description-Content-Type: text/markdown
18
+ Requires-Dist: torch
19
+ Requires-Dist: pandas
20
+ Requires-Dist: colorama>=0.4.6
21
+ Requires-Dist: matplotlib>=3.10.8
22
+ Requires-Dist: numpy>=2.4.1
23
+ Requires-Dist: tensorboard>=2.20.0
24
+ Requires-Dist: tqdm>=4.67.1
25
+ Requires-Dist: rich>=13.7
26
+ Requires-Dist: setuptools<=81.0.0
27
+ Provides-Extra: dev
28
+ Requires-Dist: ipykernel; extra == "dev"
29
+ Provides-Extra: doc
30
+ Requires-Dist: furo; extra == "doc"
31
+ Requires-Dist: myst-parser; extra == "doc"
32
+ Requires-Dist: sphinx; extra == "doc"
33
+ Requires-Dist: sphinx-togglebutton; extra == "doc"
34
+ Requires-Dist: sphinx-autodoc-typehints; extra == "doc"
35
+ Provides-Extra: test
36
+ Requires-Dist: pytest; extra == "test"
37
+
38
+ # Overview
39
+ Minimal framework for ML modeling, supporting advanced dataset operations and
40
+ streamlined training workflows.
41
+
42
+ # Install
43
+ The `trainlib` package can be installed from PyPI:
44
+
45
+ ```sh
46
+ pip install trainlib
47
+ ```
48
+
49
+ # Development
50
+ - Initialize/synchronize the project with `uv sync`, creating a virtual
51
+ environment with base package dependencies.
52
+ - Depending on needs, install the development dependencies with `uv sync
53
+ --extra dev`.
54
+
55
+ ## Testing
56
+ - To run the unit tests, make sure to first have the test dependencies
57
+ installed with `uv sync --extra test`, then run `make test`.
58
+ - For notebook testing, run `make install-kernel` to make the environment
59
+ available as a Jupyter kernel (to be selected when running notebooks).
60
+
61
+ ## Documentation
62
+ - Install the documentation dependencies with `uv sync --extra doc`.
63
+ - Run `make docs-build` (optionally preceded by `make docs-clean`), and serve
64
+ locally with `make docs-serve`.
65
+
66
+ # Development remarks
67
+ - Across `Trainer` / `Estimator` / `Dataset`, I've considered a
68
+ `ParamSpec`-based typing scheme to better orchestrate alignment in the
69
+ `Trainer.train()` loop, e.g., so we can statically check whether a dataset
70
+ appears to be fulfilling the argument requirements for the estimator's
71
+ `forward()` / `metrics()` methods. Something like
72
+
73
+ ```py
74
+ class Estimator[**P](nn.Module):
75
+ def forward(
76
+ self,
77
+ input: Tensor,
78
+ *args: P.args,
79
+ **kwargs: P.kwargs,
80
+ ) -> Tensor:
81
+ ...
82
+
83
+ class Trainer[**P]:
84
+ def __init__(
85
+ self,
86
+ estimator: Estimator[P],
87
+ ...
88
+ ): ...
89
+ ```
90
+
91
+ might be how we begin threading signatures. But ensuring dataset items can
92
+ match `P` is challenging. You can consider a "packed" object where we
93
+ obfuscate passing data through `P`-signatures:
94
+
95
+ ```py
96
+ class PackedItem[**P]:
97
+ def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
98
+ self._args = args
99
+ self._kwargs = kwargs
100
+
101
+ def apply[R](self, func: Callable[P, R]) -> R:
102
+ return func(*self._args, **self._kwargs)
103
+
104
+
105
+ class BatchedDataset[U, R, I, **P](Dataset):
106
+ @abstractmethod
107
+ def _process_item_data(
108
+ self,
109
+ item_data: I,
110
+ item_index: int,
111
+ ) -> PackedItem[P]:
112
+ ...
113
+
114
+ def __iter__(self) -> Iterator[PackedItem[P]]:
115
+ ...
116
+ ```
117
+
118
+ Meaningfully shaping those signatures is what remains, but you can't really
119
+ do this, not with typical type expression flexibility. For instance, if I'm
120
+ trying to appropriately type my base `TupleDataset`:
121
+
122
+ ```py
123
+ class SequenceDataset[I, **P](HomogenousDataset[int, I, I, P]):
124
+ ...
125
+
126
+ class TupleDataset[I](SequenceDataset[tuple[I, ...], "?"]):
127
+ ...
128
+ ```
129
+
130
+ Here there's no way for me to shape a `ParamSpec` to indicate arbitrarily
131
+ many arguments of a fixed type (`I` in this case) to allow me to unpack my
132
+ item tuples into an appropriate `PackedItem` (since `ParamSpec` isn't exactly
133
+ controllable in that way; it's meant make a generic signature expressible as
134
+ a type variable, *not to express particular shapes* of the types in that
135
+ signature).
136
+
137
+ Until this (among other issues) becomes clearer, I'm setting up around a
138
+ simpler `TypedDict` type variable. We won't have particularly strong static
139
+ checks for item alignment inside `Trainer`, but this seems about as good as I
140
+ can get around the current infrastructure.
141
+
142
+ ## Typing planning
143
+ There are currently two package-wide typing strategies geared toward future
144
+ Python version bumps and type system changes:
145
+
146
+ - Using `Unpack` with a type variable with a `TypedDict` upper bound, rather
147
+ than with a `TypedDict` directly. This is commented on extensively in
148
+ `trainlib/estimator.py`, but the short version is we want our key objects to
149
+ remain generic in their method signatures when they are expected to operate
150
+ on session-specific structured batch data. For instance, you can imagine an
151
+ `Estimator`'s `forward` method accepting a typed batch like
152
+
153
+ ```py
154
+ def forward(self, batch: BatchKwargs): ...
155
+
156
+ # or
157
+
158
+ def forward(self, **kwargs: Unpack[Batchwargs]): ...
159
+ ```
160
+
161
+ These are both currently supported patterns. Problems arise, however, when we
162
+ want an inherited class to specifically operate on an inherited type of
163
+ `BatchKwargs`
164
+
165
+ ```py
166
+ class BatchKwargs(TypedDict):
167
+ inputs: Tensor
168
+
169
+ class ExtraBatchKwargs(BatchKwargs):
170
+ inputs: Tensor
171
+ labels: Tensor
172
+
173
+ class A(Estimator):
174
+ def forward(self, batch: BatchKwargs): ...
175
+
176
+ # want to require Estimator subclass handles both inputs and labels
177
+ class B(A):
178
+ def forward(self, batch: ExtraBatchKwargs): ...
179
+ ```
180
+
181
+ Note how `ExtraBatchKwargs` doesn't use `NotRequired` for the new `labels`
182
+ attribute; the intention here is to actually require that `B` should handle
183
+ *more attributes* than `A`. **That breaks LSP**, so this doesn't work. You'd
184
+ need to either break out the `Estimator` hierarchies, or decide that `B`
185
+ could have inputs typed as `BatchKwargs` and possibly implement some extra
186
+ handling to check if batches of type `ExtraBatchKwargs` are passed in (note
187
+ that passing in a batch of type `ExtraBatchKwargs` to `A.forward` is
188
+ perfectly fine) and supply labels. **Or...** you make the class generic in
189
+ the batch type, and during inheritance you can establish increasingly tight
190
+ upper bounds on the batch type to accomplish precisely the desired effect.
191
+ That is, you can formally require subclasses handle new batch attributes
192
+ without breaking LSP. This could look like:
193
+
194
+ ```py
195
+ class A[Kw: BatchKwargs](Estimator):
196
+ def forward(self, batch: Kw): ...
197
+
198
+ class B[Kw: ExtraBatchKwargs](A[Kw]):
199
+ def forward(self, batch: Kw): ...
200
+ ```
201
+
202
+ This formally allows `B` to inherit from `A` while depending on new,
203
+ non-optional attributes in its batch type; you can achieve that with refined
204
+ type upper bound in a generic subtype. So this pattern works *now*, but what
205
+ doesn't work is further *unpacking* `Kw` in this pattern into kwargs. That
206
+ is,
207
+
208
+ ```py
209
+ class A[Kw: BatchKwargs](Estimator):
210
+ def forward(self, **kwargs: Unpack[Kw]): ...
211
+
212
+ class B[Kw: ExtraBatchKwargs](A[Kw]):
213
+ def forward(self, **kwargs: Unpack[Kw]): ...
214
+ ```
215
+
216
+ As sound as this is (and it *is* sound, described in PEP 827), `Unpack` isn't
217
+ supported for type variables; it has to be used directly with `TypedDict`
218
+ references. I have the latter in place across the package anticipating
219
+ eventual support for this pattern, but things are subject to change. I may
220
+ simply toggle back to the currently supported pattern provided it's very
221
+ similar, and if/when PEP 827 is integrated, the change will be as simply as
222
+ adding back the `Unpack` statements and keyword expansions.
223
+ - At the time of writing, `trainlib` attempts to remain compatible with Python
224
+ 3.12. This means the convenient lazy evaluation of type annotations can't be
225
+ fully embraced across the package, but we use `__future__` where possible to
226
+ make the transition easy when eventually bumping to 3.14+ (since the imports
227
+ can just be removed at that time). For instance, class definitions that need
228
+ to reference the class type itself shouldn't use string annotations but
229
+ instead direct type references:
230
+
231
+ ```py
232
+ # not
233
+ class RecordIndex:
234
+ @overload
235
+ def __getitem__(self, key: slice) -> "RecordIndex": ...
236
+
237
+ # but
238
+ from __future__ import annotations
239
+
240
+ class RecordIndex:
241
+ @overload
242
+ def __getitem__(self, key: slice) -> RecordIndex: ...
243
+ ```
@@ -0,0 +1,206 @@
1
+ # Overview
2
+ Minimal framework for ML modeling, supporting advanced dataset operations and
3
+ streamlined training workflows.
4
+
5
+ # Install
6
+ The `trainlib` package can be installed from PyPI:
7
+
8
+ ```sh
9
+ pip install trainlib
10
+ ```
11
+
12
+ # Development
13
+ - Initialize/synchronize the project with `uv sync`, creating a virtual
14
+ environment with base package dependencies.
15
+ - Depending on needs, install the development dependencies with `uv sync
16
+ --extra dev`.
17
+
18
+ ## Testing
19
+ - To run the unit tests, make sure to first have the test dependencies
20
+ installed with `uv sync --extra test`, then run `make test`.
21
+ - For notebook testing, run `make install-kernel` to make the environment
22
+ available as a Jupyter kernel (to be selected when running notebooks).
23
+
24
+ ## Documentation
25
+ - Install the documentation dependencies with `uv sync --extra doc`.
26
+ - Run `make docs-build` (optionally preceded by `make docs-clean`), and serve
27
+ locally with `make docs-serve`.
28
+
29
+ # Development remarks
30
+ - Across `Trainer` / `Estimator` / `Dataset`, I've considered a
31
+ `ParamSpec`-based typing scheme to better orchestrate alignment in the
32
+ `Trainer.train()` loop, e.g., so we can statically check whether a dataset
33
+ appears to be fulfilling the argument requirements for the estimator's
34
+ `forward()` / `metrics()` methods. Something like
35
+
36
+ ```py
37
+ class Estimator[**P](nn.Module):
38
+ def forward(
39
+ self,
40
+ input: Tensor,
41
+ *args: P.args,
42
+ **kwargs: P.kwargs,
43
+ ) -> Tensor:
44
+ ...
45
+
46
+ class Trainer[**P]:
47
+ def __init__(
48
+ self,
49
+ estimator: Estimator[P],
50
+ ...
51
+ ): ...
52
+ ```
53
+
54
+ might be how we begin threading signatures. But ensuring dataset items can
55
+ match `P` is challenging. You can consider a "packed" object where we
56
+ obfuscate passing data through `P`-signatures:
57
+
58
+ ```py
59
+ class PackedItem[**P]:
60
+ def __init__(self, *args: P.args, **kwargs: P.kwargs) -> None:
61
+ self._args = args
62
+ self._kwargs = kwargs
63
+
64
+ def apply[R](self, func: Callable[P, R]) -> R:
65
+ return func(*self._args, **self._kwargs)
66
+
67
+
68
+ class BatchedDataset[U, R, I, **P](Dataset):
69
+ @abstractmethod
70
+ def _process_item_data(
71
+ self,
72
+ item_data: I,
73
+ item_index: int,
74
+ ) -> PackedItem[P]:
75
+ ...
76
+
77
+ def __iter__(self) -> Iterator[PackedItem[P]]:
78
+ ...
79
+ ```
80
+
81
+ Meaningfully shaping those signatures is what remains, but you can't really
82
+ do this, not with typical type expression flexibility. For instance, if I'm
83
+ trying to appropriately type my base `TupleDataset`:
84
+
85
+ ```py
86
+ class SequenceDataset[I, **P](HomogenousDataset[int, I, I, P]):
87
+ ...
88
+
89
+ class TupleDataset[I](SequenceDataset[tuple[I, ...], "?"]):
90
+ ...
91
+ ```
92
+
93
+ Here there's no way for me to shape a `ParamSpec` to indicate arbitrarily
94
+ many arguments of a fixed type (`I` in this case) to allow me to unpack my
95
+ item tuples into an appropriate `PackedItem` (since `ParamSpec` isn't exactly
96
+ controllable in that way; it's meant make a generic signature expressible as
97
+ a type variable, *not to express particular shapes* of the types in that
98
+ signature).
99
+
100
+ Until this (among other issues) becomes clearer, I'm setting up around a
101
+ simpler `TypedDict` type variable. We won't have particularly strong static
102
+ checks for item alignment inside `Trainer`, but this seems about as good as I
103
+ can get around the current infrastructure.
104
+
105
+ ## Typing planning
106
+ There are currently two package-wide typing strategies geared toward future
107
+ Python version bumps and type system changes:
108
+
109
+ - Using `Unpack` with a type variable with a `TypedDict` upper bound, rather
110
+ than with a `TypedDict` directly. This is commented on extensively in
111
+ `trainlib/estimator.py`, but the short version is we want our key objects to
112
+ remain generic in their method signatures when they are expected to operate
113
+ on session-specific structured batch data. For instance, you can imagine an
114
+ `Estimator`'s `forward` method accepting a typed batch like
115
+
116
+ ```py
117
+ def forward(self, batch: BatchKwargs): ...
118
+
119
+ # or
120
+
121
+ def forward(self, **kwargs: Unpack[Batchwargs]): ...
122
+ ```
123
+
124
+ These are both currently supported patterns. Problems arise, however, when we
125
+ want an inherited class to specifically operate on an inherited type of
126
+ `BatchKwargs`
127
+
128
+ ```py
129
+ class BatchKwargs(TypedDict):
130
+ inputs: Tensor
131
+
132
+ class ExtraBatchKwargs(BatchKwargs):
133
+ inputs: Tensor
134
+ labels: Tensor
135
+
136
+ class A(Estimator):
137
+ def forward(self, batch: BatchKwargs): ...
138
+
139
+ # want to require Estimator subclass handles both inputs and labels
140
+ class B(A):
141
+ def forward(self, batch: ExtraBatchKwargs): ...
142
+ ```
143
+
144
+ Note how `ExtraBatchKwargs` doesn't use `NotRequired` for the new `labels`
145
+ attribute; the intention here is to actually require that `B` should handle
146
+ *more attributes* than `A`. **That breaks LSP**, so this doesn't work. You'd
147
+ need to either break out the `Estimator` hierarchies, or decide that `B`
148
+ could have inputs typed as `BatchKwargs` and possibly implement some extra
149
+ handling to check if batches of type `ExtraBatchKwargs` are passed in (note
150
+ that passing in a batch of type `ExtraBatchKwargs` to `A.forward` is
151
+ perfectly fine) and supply labels. **Or...** you make the class generic in
152
+ the batch type, and during inheritance you can establish increasingly tight
153
+ upper bounds on the batch type to accomplish precisely the desired effect.
154
+ That is, you can formally require subclasses handle new batch attributes
155
+ without breaking LSP. This could look like:
156
+
157
+ ```py
158
+ class A[Kw: BatchKwargs](Estimator):
159
+ def forward(self, batch: Kw): ...
160
+
161
+ class B[Kw: ExtraBatchKwargs](A[Kw]):
162
+ def forward(self, batch: Kw): ...
163
+ ```
164
+
165
+ This formally allows `B` to inherit from `A` while depending on new,
166
+ non-optional attributes in its batch type; you can achieve that with refined
167
+ type upper bound in a generic subtype. So this pattern works *now*, but what
168
+ doesn't work is further *unpacking* `Kw` in this pattern into kwargs. That
169
+ is,
170
+
171
+ ```py
172
+ class A[Kw: BatchKwargs](Estimator):
173
+ def forward(self, **kwargs: Unpack[Kw]): ...
174
+
175
+ class B[Kw: ExtraBatchKwargs](A[Kw]):
176
+ def forward(self, **kwargs: Unpack[Kw]): ...
177
+ ```
178
+
179
+ As sound as this is (and it *is* sound, described in PEP 827), `Unpack` isn't
180
+ supported for type variables; it has to be used directly with `TypedDict`
181
+ references. I have the latter in place across the package anticipating
182
+ eventual support for this pattern, but things are subject to change. I may
183
+ simply toggle back to the currently supported pattern provided it's very
184
+ similar, and if/when PEP 827 is integrated, the change will be as simply as
185
+ adding back the `Unpack` statements and keyword expansions.
186
+ - At the time of writing, `trainlib` attempts to remain compatible with Python
187
+ 3.12. This means the convenient lazy evaluation of type annotations can't be
188
+ fully embraced across the package, but we use `__future__` where possible to
189
+ make the transition easy when eventually bumping to 3.14+ (since the imports
190
+ can just be removed at that time). For instance, class definitions that need
191
+ to reference the class type itself shouldn't use string annotations but
192
+ instead direct type references:
193
+
194
+ ```py
195
+ # not
196
+ class RecordIndex:
197
+ @overload
198
+ def __getitem__(self, key: slice) -> "RecordIndex": ...
199
+
200
+ # but
201
+ from __future__ import annotations
202
+
203
+ class RecordIndex:
204
+ @overload
205
+ def __getitem__(self, key: slice) -> RecordIndex: ...
206
+ ```
@@ -4,12 +4,10 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "trainlib"
7
- version = "0.3.1"
7
+ version = "0.6.1"
8
8
  description = "Minimal framework for ML modeling. Supports advanced dataset operations and streamlined training."
9
- requires-python = ">=3.13"
10
- authors = [
11
- { name="Sam Griesemer", email="git@olog.io" },
12
- ]
9
+ requires-python = ">=3.12"
10
+ authors = []
13
11
  readme = "README.md"
14
12
  license = "MIT"
15
13
  keywords = [
@@ -25,15 +23,18 @@ classifiers = [
25
23
  ]
26
24
  dependencies = [
27
25
  "torch",
26
+ "pandas",
28
27
  "colorama>=0.4.6",
29
28
  "matplotlib>=3.10.8",
30
29
  "numpy>=2.4.1",
31
30
  "tensorboard>=2.20.0",
32
31
  "tqdm>=4.67.1",
32
+ "rich>=13.7",
33
33
  "setuptools<=81.0.0", # here currently b/c tensorboard breaks @ v82.0.0
34
34
  ]
35
35
 
36
36
  [project.scripts]
37
+ trainlib = "trainlib.__main__:main"
37
38
  symconf = "trainlib.__main__:main"
38
39
 
39
40
  [project.optional-dependencies]
@@ -61,6 +62,10 @@ Issues = "https://git.olog.io/olog/trainlib/issues"
61
62
  [tool.setuptools.packages.find]
62
63
  include = ["trainlib*"]
63
64
 
65
+ [tool.pytest.ini_options]
66
+ testpaths = ["test"]
67
+ addopts = "--import-mode=importlib"
68
+
64
69
  # for static data files under package root
65
70
  # [tool.setuptools.package-data]
66
71
  # "<package>" = ["data/*.toml"]
@@ -68,8 +73,13 @@ include = ["trainlib*"]
68
73
  [tool.ruff]
69
74
  line-length = 79
70
75
 
76
+ [tool.ruff.format]
77
+ quote-style = "double"
78
+ indent-style = "space"
79
+ docstring-code-format = true
80
+
71
81
  [tool.ruff.lint]
72
- select = ["ANN", "E", "F", "UP", "B", "SIM", "I", "C4", "PERF"]
82
+ select = ["ANN", "D", "E", "F", "UP", "B", "SIM", "I", "C4", "PERF"]
73
83
 
74
84
  [tool.ruff.lint.isort]
75
85
  length-sort = true
@@ -77,13 +87,14 @@ order-by-type = false
77
87
  force-sort-within-sections = false
78
88
 
79
89
  [tool.ruff.lint.per-file-ignores]
80
- "tests/**" = ["S101"]
90
+ "test/**" = ["S101"]
81
91
  "**/__init__.py" = ["F401"]
82
92
 
83
- [tool.ruff.format]
84
- quote-style = "double"
85
- indent-style = "space"
86
- docstring-code-format = true
93
+ [tool.ruff.lint.pydocstyle]
94
+ convention = "google"
95
+
96
+ [tool.setuptools.package-data]
97
+ "trainlib.utils" = ["custom.mplstyle"]
87
98
 
88
99
  [tool.uv.sources]
89
100
  torch = { index = "pytorch" }
@@ -0,0 +1,160 @@
1
+ import json
2
+ from copy import deepcopy
3
+ from pathlib import Path
4
+
5
+ import pytest
6
+
7
+ from trainlib.cli import train, experiment, create_parser
8
+ from trainlib.run import expand_runs
9
+ from trainlib.trainer import Trainer
10
+ from trainlib.cli.experiment import ExperimentConfig
11
+ from trainlib.experiments.train import TrainExperiment
12
+
13
+ EXAMPLES = Path(__file__).resolve().parents[1] / "example" / "cli"
14
+
15
+
16
+ def test_cli_dispatches_non_training_experiment(tmp_path, monkeypatch):
17
+ from trainlib.run import Run
18
+ from trainlib.record import RecordIndex
19
+ from trainlib.experiment import Experiment
20
+ from trainlib.experiments import experiment_map
21
+
22
+ class Evaluation(Experiment[Run]):
23
+ def _execute(self, spec, directory, *, progress, console, output):
24
+ self.runstore.save(
25
+ directory, RecordIndex(), {"test": {"score": 1.0}}
26
+ )
27
+
28
+ monkeypatch.setitem(experiment_map, "test_evaluation", Evaluation)
29
+ result = experiment.run(
30
+ {
31
+ "experiment": "test_evaluation",
32
+ "output": str(tmp_path / "eval"),
33
+ "groups": {"raw": {}},
34
+ "seeds": [7],
35
+ "summary": {"plots": False, "resamples": 10},
36
+ },
37
+ quiet=True,
38
+ )
39
+ assert len(result["completed"]) == 1
40
+ assert "summary" in result
41
+
42
+
43
+ def test_single_run_cli_loads_and_trains_without_experiment(
44
+ tmp_path: Path,
45
+ ) -> None:
46
+ config = json.loads((EXAMPLES / "run.json").read_text())
47
+ config["data"]["dataset"]["kwargs"]["num_samples"] = 12
48
+ config["estimator"]["kwargs"]["verbose"] = False
49
+ config["reporter"]["kwargs"].update(
50
+ progress=False,
51
+ console=False,
52
+ tensorboard=False,
53
+ )
54
+ config.update(max_epochs=1, chkpt_every=1)
55
+ path = tmp_path / "train.json"
56
+ path.write_text(json.dumps(config))
57
+ output = tmp_path / "single"
58
+
59
+ trainer = train.run_from_json(
60
+ parameters_file=str(path),
61
+ load_only=True,
62
+ output=output,
63
+ )
64
+ assert isinstance(trainer, Trainer)
65
+ assert not len(trainer.record_index)
66
+ assert not (output / "checkpoints").exists()
67
+
68
+ args = create_parser().parse_args(
69
+ [
70
+ "train",
71
+ "--parameters-file",
72
+ str(path),
73
+ "--output",
74
+ str(output),
75
+ ]
76
+ )
77
+ args.func(args)
78
+
79
+ assert (output / "config.json").is_file()
80
+ assert list((output / "checkpoints").rglob("*.pth"))
81
+ assert not (output / "experiment.json").exists()
82
+ assert not (output / "runs").exists()
83
+
84
+
85
+ def test_group_seed_expansion_matches_explicit_runs_and_extends_stably(
86
+ tmp_path: Path,
87
+ ) -> None:
88
+ config = json.loads((EXAMPLES / "experiment.json").read_text())
89
+ original = deepcopy(config)
90
+ runs = expand_runs(config)
91
+ assert len(runs) == 9
92
+ assert list(runs)[:3] == ["mlp-small-1", "mlp-small-2", "mlp-small-3"]
93
+ assert runs["mlp-medium-2"]["group"] == "mlp-medium"
94
+ assert runs["mlp-medium-2"]["seed"] == 2
95
+ assert config == original
96
+
97
+ initial = TrainExperiment(tmp_path, runs, common=config["common"])
98
+ assert expand_runs({"runs": runs}) == runs
99
+ config["seeds"].append(4)
100
+ extended = TrainExperiment(
101
+ tmp_path,
102
+ expand_runs(config),
103
+ common=config["common"],
104
+ )
105
+ assert len(extended.runs) == 12
106
+ for name, spec in initial.runs.items():
107
+ assert extended.runs[name] == spec
108
+
109
+ runs["mlp-small-1"]["estimator"]["kwargs"]["hidden_dims"].append(99)
110
+ assert runs["mlp-small-2"]["estimator"]["kwargs"]["hidden_dims"] == [8]
111
+ assert original["groups"] == config["groups"]
112
+
113
+
114
+ @pytest.mark.parametrize(
115
+ "config",
116
+ [
117
+ {"runs": {"one": {}}, "groups": {"one": {}}, "seeds": [1]},
118
+ {"groups": {"one": {}}},
119
+ {"groups": {"one": {}}, "seeds": [1, 1]},
120
+ ],
121
+ )
122
+ def test_experiment_rejects_ambiguous_or_incomplete_expansion(
123
+ config: ExperimentConfig,
124
+ ) -> None:
125
+ with pytest.raises(ValueError):
126
+ expand_runs(config)
127
+
128
+
129
+ def test_compact_experiment_cli_writes_expanded_manifest_and_analysis(
130
+ tmp_path: Path,
131
+ ) -> None:
132
+ config = json.loads((EXAMPLES / "experiment.json").read_text())
133
+ config["groups"] = {"small": config["groups"]["mlp-small"]}
134
+ config["seeds"] = [1]
135
+ config["comparisons"] = {}
136
+ common = config["common"]
137
+ common["data"]["dataset"]["kwargs"]["num_samples"] = 12
138
+ common["reporter"]["kwargs"].update(tensorboard=False)
139
+ common.update(max_epochs=1, trace_graph=False)
140
+ config["summary"].update(plots=False, resamples=10)
141
+ path = tmp_path / "experiment.json"
142
+ path.write_text(json.dumps(config))
143
+ output = tmp_path / "experiment"
144
+
145
+ result = experiment.run_from_json(
146
+ parameters_file=str(path),
147
+ output=output,
148
+ quiet=True,
149
+ )
150
+
151
+ assert len(result["completed"]) == 1
152
+ assert result["failures"] == {}
153
+ assert result["summary"]["rows"]
154
+ manifest = json.loads((output / "experiment.json").read_text())
155
+ (spec,) = manifest["runs"].values()
156
+ assert spec["name"] == "small-1"
157
+ assert spec["group"] == "small"
158
+ assert spec["seed"] == 1
159
+ assert list(output.glob("runs/*/attempt-*/state.pt"))
160
+ assert (output / "analysis" / "summary.json").is_file()