lit-mlflow 0.1.0__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.
- lit_mlflow-0.1.0/LICENSE +21 -0
- lit_mlflow-0.1.0/PKG-INFO +107 -0
- lit_mlflow-0.1.0/README.md +85 -0
- lit_mlflow-0.1.0/pyproject.toml +226 -0
- lit_mlflow-0.1.0/src/lit_mlflow/__about__.py +4 -0
- lit_mlflow-0.1.0/src/lit_mlflow/__init__.py +5 -0
- lit_mlflow-0.1.0/src/lit_mlflow/callback.py +498 -0
- lit_mlflow-0.1.0/src/lit_mlflow/cli.py +5 -0
- lit_mlflow-0.1.0/src/lit_mlflow/logger.py +76 -0
- lit_mlflow-0.1.0/src/lit_mlflow/py.typed +0 -0
- lit_mlflow-0.1.0/src/lit_mlflow/utils/dbx.py +74 -0
lit_mlflow-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 twsl
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: lit-mlflow
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: An improved Lightning mlflow logger
|
|
5
|
+
Home-page: https://twsl.github.io/lit-mlflow/
|
|
6
|
+
License: MIT
|
|
7
|
+
Author: twsl
|
|
8
|
+
Author-email: 45483159+twsl@users.noreply.github.com
|
|
9
|
+
Requires-Python: >=3.10,<4.0
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Requires-Dist: databricks-cli (>=0.18.0,<0.19.0)
|
|
17
|
+
Requires-Dist: lightning (>=2.4.0,<3.0.0)
|
|
18
|
+
Requires-Dist: mlflow (>=2.18.0,<3.0.0)
|
|
19
|
+
Project-URL: Repository, https://github.com/twsl/lit-mlflow
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# lit-mlflow
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
[](https://github.com/twsl/lit-mlflow/actions/workflows/build.yaml)
|
|
26
|
+
[](https://github.com/twsl/lit-mlflow/actions/workflows/docs.yaml)
|
|
27
|
+
|
|
28
|
+
[](https://pypi.org/project/lit-mlflow/)
|
|
29
|
+
[](https://pypi.org/project/lit-mlflow/)
|
|
30
|
+
[](https://github.com/twsl/lit-mlflow/pulls?utf8=%E2%9C%93&q=is:pr%20author:app/dependabot)
|
|
31
|
+
[](https://squidfunk.github.io/mkdocs-material/)
|
|
32
|
+
[](https://python-poetry.org/)
|
|
33
|
+
[](https://github.com/astral-sh/ruff)
|
|
34
|
+
[](.pre-commit-config.yaml)
|
|
35
|
+
[](https://microsoft.github.io/pyright/)
|
|
36
|
+
[](https://github.com/PyCQA/bandit)
|
|
37
|
+
[](https://github.com/twsl/lit-mlflow/releases)
|
|
38
|
+
[](https://github.com/copier-org/copier)
|
|
39
|
+
[](LICENSE)
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
An improved Lightning mlflow logger. Works seamlessly with PyTorch Lightning on Databricks and offers more control compared to the `mlflow.pytorch.autolog` function.
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
## Features
|
|
46
|
+
|
|
47
|
+
- Makes `MLflow` logging work with `lightning` and [Databricks](https://www.databricks.com/)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
## Installation
|
|
51
|
+
|
|
52
|
+
With `pip`:
|
|
53
|
+
```bash
|
|
54
|
+
python -m pip install lit-mlflow
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
With [`poetry`](https://python-poetry.org/):
|
|
58
|
+
```bash
|
|
59
|
+
poetry add lit-mlflow
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
## How to use it
|
|
64
|
+
Replace `mlflow.autolog()` with the `MlFlowAutoCallback`:
|
|
65
|
+
|
|
66
|
+
```python
|
|
67
|
+
from lit_mlflow import MlFlowAutoCallback
|
|
68
|
+
import lightning.pytorch as pl
|
|
69
|
+
|
|
70
|
+
trainer = pl.Trainer(
|
|
71
|
+
callbacks=[
|
|
72
|
+
MlFlowAutoCallback()
|
|
73
|
+
]
|
|
74
|
+
)
|
|
75
|
+
```
|
|
76
|
+
|
|
77
|
+
To support Databricks mlflow, use the `DbxMLFlowLogger` instead of the `MlFlowLogger`:
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
from lit_mlflow import DbxMLFlowLogger
|
|
81
|
+
import lightning.pytorch as pl
|
|
82
|
+
|
|
83
|
+
trainer = pl.Trainer(
|
|
84
|
+
logger=[
|
|
85
|
+
DbxMLFlowLogger()
|
|
86
|
+
]
|
|
87
|
+
)
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
## Docs
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
poetry run mkdocs build -f ./docs/mkdocs.yml -d ./_build/
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
## Update template
|
|
98
|
+
|
|
99
|
+
```bash
|
|
100
|
+
copier update --trust -A --vcs-ref=HEAD
|
|
101
|
+
```
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
## Credits
|
|
105
|
+
|
|
106
|
+
This project was generated with [](https://github.com/twsl/python-project-template)
|
|
107
|
+
|
|
@@ -0,0 +1,85 @@
|
|
|
1
|
+
# lit-mlflow
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
[](https://github.com/twsl/lit-mlflow/actions/workflows/build.yaml)
|
|
5
|
+
[](https://github.com/twsl/lit-mlflow/actions/workflows/docs.yaml)
|
|
6
|
+
|
|
7
|
+
[](https://pypi.org/project/lit-mlflow/)
|
|
8
|
+
[](https://pypi.org/project/lit-mlflow/)
|
|
9
|
+
[](https://github.com/twsl/lit-mlflow/pulls?utf8=%E2%9C%93&q=is:pr%20author:app/dependabot)
|
|
10
|
+
[](https://squidfunk.github.io/mkdocs-material/)
|
|
11
|
+
[](https://python-poetry.org/)
|
|
12
|
+
[](https://github.com/astral-sh/ruff)
|
|
13
|
+
[](.pre-commit-config.yaml)
|
|
14
|
+
[](https://microsoft.github.io/pyright/)
|
|
15
|
+
[](https://github.com/PyCQA/bandit)
|
|
16
|
+
[](https://github.com/twsl/lit-mlflow/releases)
|
|
17
|
+
[](https://github.com/copier-org/copier)
|
|
18
|
+
[](LICENSE)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
An improved Lightning mlflow logger. Works seamlessly with PyTorch Lightning on Databricks and offers more control compared to the `mlflow.pytorch.autolog` function.
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
## Features
|
|
25
|
+
|
|
26
|
+
- Makes `MLflow` logging work with `lightning` and [Databricks](https://www.databricks.com/)
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
## Installation
|
|
30
|
+
|
|
31
|
+
With `pip`:
|
|
32
|
+
```bash
|
|
33
|
+
python -m pip install lit-mlflow
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
With [`poetry`](https://python-poetry.org/):
|
|
37
|
+
```bash
|
|
38
|
+
poetry add lit-mlflow
|
|
39
|
+
```
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
## How to use it
|
|
43
|
+
Replace `mlflow.autolog()` with the `MlFlowAutoCallback`:
|
|
44
|
+
|
|
45
|
+
```python
|
|
46
|
+
from lit_mlflow import MlFlowAutoCallback
|
|
47
|
+
import lightning.pytorch as pl
|
|
48
|
+
|
|
49
|
+
trainer = pl.Trainer(
|
|
50
|
+
callbacks=[
|
|
51
|
+
MlFlowAutoCallback()
|
|
52
|
+
]
|
|
53
|
+
)
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
To support Databricks mlflow, use the `DbxMLFlowLogger` instead of the `MlFlowLogger`:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
from lit_mlflow import DbxMLFlowLogger
|
|
60
|
+
import lightning.pytorch as pl
|
|
61
|
+
|
|
62
|
+
trainer = pl.Trainer(
|
|
63
|
+
logger=[
|
|
64
|
+
DbxMLFlowLogger()
|
|
65
|
+
]
|
|
66
|
+
)
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
## Docs
|
|
70
|
+
|
|
71
|
+
```bash
|
|
72
|
+
poetry run mkdocs build -f ./docs/mkdocs.yml -d ./_build/
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
## Update template
|
|
77
|
+
|
|
78
|
+
```bash
|
|
79
|
+
copier update --trust -A --vcs-ref=HEAD
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
## Credits
|
|
84
|
+
|
|
85
|
+
This project was generated with [](https://github.com/twsl/python-project-template)
|
|
@@ -0,0 +1,226 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["poetry-core>=1.0.0", "poetry-dynamic-versioning>=1.0.0,<2.0.0"]
|
|
3
|
+
build-backend = "poetry_dynamic_versioning.backend"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "lit-mlflow"
|
|
7
|
+
requires-python = ">=3.10"
|
|
8
|
+
|
|
9
|
+
# TODO: https://github.com/python-poetry/poetry/issues/3332
|
|
10
|
+
[tool.poetry]
|
|
11
|
+
name = "lit-mlflow"
|
|
12
|
+
version = "0.1.0"
|
|
13
|
+
description = "An improved Lightning mlflow logger"
|
|
14
|
+
authors = ["twsl <45483159+twsl@users.noreply.github.com>"]
|
|
15
|
+
license = "MIT"
|
|
16
|
+
readme = "README.md"
|
|
17
|
+
repository = "https://github.com/twsl/lit-mlflow"
|
|
18
|
+
homepage = "https://twsl.github.io/lit-mlflow/"
|
|
19
|
+
packages = [ { include = "lit_mlflow", from = "src" } ]
|
|
20
|
+
keywords = []
|
|
21
|
+
|
|
22
|
+
[tool.poetry.dependencies]
|
|
23
|
+
python = "^3.10"
|
|
24
|
+
mlflow = "^2.18.0"
|
|
25
|
+
lightning = "^2.4.0"
|
|
26
|
+
databricks-cli = "^0.18.0"
|
|
27
|
+
|
|
28
|
+
[tool.poetry.group.dev.dependencies]
|
|
29
|
+
ruff = "^0.8.2"
|
|
30
|
+
bandit = "^1.8.0"
|
|
31
|
+
pre-commit = "^4.0.1"
|
|
32
|
+
pyright = "^1.1.390"
|
|
33
|
+
ipykernel = "^6.29.5"
|
|
34
|
+
ipywidgets = "^8.1.5"
|
|
35
|
+
ipython = "^8.30.0"
|
|
36
|
+
jupytext = " ^1.16.4"
|
|
37
|
+
nbconvert = "^7.16.3"
|
|
38
|
+
|
|
39
|
+
[tool.poetry.group.test.dependencies]
|
|
40
|
+
pytest = "^8.3.4"
|
|
41
|
+
pytest-cov = "^6.0.0"
|
|
42
|
+
pytest-mock = "^3.12.0"
|
|
43
|
+
pytest-cases = "^3.8.6"
|
|
44
|
+
pytest-benchmark = "^5.1.0"
|
|
45
|
+
|
|
46
|
+
[tool.poetry.group.docs.dependencies]
|
|
47
|
+
mkdocs = "^1.6.1"
|
|
48
|
+
mkdocs-material = "^9.5.48"
|
|
49
|
+
mkdocs-git-revision-date-localized-plugin = "^1.2.9"
|
|
50
|
+
mkdocs-awesome-pages-plugin = "^2.9.3"
|
|
51
|
+
mkdocs-git-authors-plugin = "^0.9.2"
|
|
52
|
+
mkapi = "^3.0.22"
|
|
53
|
+
markdown-callouts = "^0.4.0"
|
|
54
|
+
pygments = "^2.18.0"
|
|
55
|
+
mkdocs-jupyter = "^0.25.1"
|
|
56
|
+
|
|
57
|
+
[tool.poetry.group.debug]
|
|
58
|
+
optional = true
|
|
59
|
+
|
|
60
|
+
[tool.poetry.group.debug.dependencies]
|
|
61
|
+
ipdb = "^0.13.9"
|
|
62
|
+
line_profiler = "^4.2.0"
|
|
63
|
+
|
|
64
|
+
[tool.poetry.scripts]
|
|
65
|
+
lit-mlflow = "lit_mlflow.cli:main"
|
|
66
|
+
|
|
67
|
+
[tool.poetry-dynamic-versioning]
|
|
68
|
+
enable = false
|
|
69
|
+
vcs = "git"
|
|
70
|
+
style = "semver"
|
|
71
|
+
pattern = "default-unprefixed"
|
|
72
|
+
|
|
73
|
+
[tool.poetry-dynamic-versioning.substitution]
|
|
74
|
+
folders = [
|
|
75
|
+
{ path = "src/lit_mlflow" }
|
|
76
|
+
]
|
|
77
|
+
|
|
78
|
+
[tool.poetry-dynamic-versioning.files."src/lit_mlflow/__about__.py"]
|
|
79
|
+
persistent-substitution = false
|
|
80
|
+
initial-content = """
|
|
81
|
+
# file generated by poetry dynamic versioning during poetry build/install
|
|
82
|
+
# don't change, don't track in version control
|
|
83
|
+
__version__ = "0.0.0"
|
|
84
|
+
__version_tuple__ = (0, 0, 0)
|
|
85
|
+
"""
|
|
86
|
+
|
|
87
|
+
[tool.pytest.ini_options]
|
|
88
|
+
minversion = "8.0"
|
|
89
|
+
# https://code.visualstudio.com/docs/python/testing#_pytest-configuration-settings
|
|
90
|
+
addopts = "--doctest-glob='test_*.md' --cov=src/lit_mlflow --cov-report xml"
|
|
91
|
+
testpaths = [
|
|
92
|
+
"tests",
|
|
93
|
+
]
|
|
94
|
+
|
|
95
|
+
[tool.coverage.run]
|
|
96
|
+
branch = true
|
|
97
|
+
omit = [
|
|
98
|
+
"tests/*",
|
|
99
|
+
]
|
|
100
|
+
|
|
101
|
+
[tool.coverage.report]
|
|
102
|
+
exclude_lines = [
|
|
103
|
+
# Have to re-enable the standard pragma
|
|
104
|
+
"pragma: no cover",
|
|
105
|
+
|
|
106
|
+
# Don't complain about missing debug-only code:
|
|
107
|
+
"def __repr__",
|
|
108
|
+
"if self.debug:",
|
|
109
|
+
"if debug:",
|
|
110
|
+
"if DEBUG:",
|
|
111
|
+
|
|
112
|
+
# Don't complain if tests don't hit defensive assertion code:
|
|
113
|
+
"raise AssertionError",
|
|
114
|
+
"raise NotImplementedError",
|
|
115
|
+
|
|
116
|
+
# Don't complain if non-runnable code isn't run:
|
|
117
|
+
"if 0:",
|
|
118
|
+
"if False:",
|
|
119
|
+
"if __name__ == .__main__.:",
|
|
120
|
+
]
|
|
121
|
+
|
|
122
|
+
[tool.ruff]
|
|
123
|
+
line-length = 120 # Must agree with Black
|
|
124
|
+
extend-include = ["*.ipynb"]
|
|
125
|
+
exclude = [
|
|
126
|
+
"migrations",
|
|
127
|
+
"__pycache__",
|
|
128
|
+
"env",
|
|
129
|
+
".env",
|
|
130
|
+
"venv",
|
|
131
|
+
".venv",
|
|
132
|
+
]
|
|
133
|
+
|
|
134
|
+
[tool.ruff.lint]
|
|
135
|
+
select = [
|
|
136
|
+
"B", # flake8-bugbear
|
|
137
|
+
"C4", # flake8-comprehensions
|
|
138
|
+
"D", # pydocstyle
|
|
139
|
+
"E", # Error
|
|
140
|
+
"F", # pyflakes
|
|
141
|
+
"I", # isort
|
|
142
|
+
"ISC", # flake8-implicit-str-concat
|
|
143
|
+
"N", # pep8-naming
|
|
144
|
+
"PGH", # pygrep-hooks
|
|
145
|
+
"PTH", # flake8-use-pathlib
|
|
146
|
+
"Q", # flake8-quotes
|
|
147
|
+
"S", # bandit
|
|
148
|
+
"SIM", # flake8-simplify
|
|
149
|
+
"TRY", # tryceratops
|
|
150
|
+
"UP", # pyupgrade
|
|
151
|
+
"W", # Warning
|
|
152
|
+
"YTT", # flake8-2020
|
|
153
|
+
]
|
|
154
|
+
extend-select = [
|
|
155
|
+
"I",
|
|
156
|
+
]
|
|
157
|
+
ignore = [
|
|
158
|
+
"D100",
|
|
159
|
+
"D101",
|
|
160
|
+
"D102",
|
|
161
|
+
"D103",
|
|
162
|
+
"D104",
|
|
163
|
+
"D105",
|
|
164
|
+
"D106",
|
|
165
|
+
"D107",
|
|
166
|
+
"D200",
|
|
167
|
+
"D401",
|
|
168
|
+
"E402",
|
|
169
|
+
"E501",
|
|
170
|
+
"F401",
|
|
171
|
+
"TRY002", # Allow default exceptions
|
|
172
|
+
"TRY003", # Avoid specifying messages outside exception class; overly strict, especially for ValueError
|
|
173
|
+
]
|
|
174
|
+
|
|
175
|
+
[tool.ruff.lint.flake8-bugbear]
|
|
176
|
+
extend-immutable-calls = [
|
|
177
|
+
"chr",
|
|
178
|
+
"typer.Argument",
|
|
179
|
+
"typer.Option",
|
|
180
|
+
]
|
|
181
|
+
|
|
182
|
+
[tool.ruff.lint.pydocstyle]
|
|
183
|
+
convention = "google"
|
|
184
|
+
|
|
185
|
+
[tool.ruff.lint.pycodestyle]
|
|
186
|
+
max-doc-length = 120
|
|
187
|
+
max-line-length = 120
|
|
188
|
+
|
|
189
|
+
[tool.ruff.lint.per-file-ignores]
|
|
190
|
+
"test*/**/*.py" = [
|
|
191
|
+
"D100",
|
|
192
|
+
"D101",
|
|
193
|
+
"D102",
|
|
194
|
+
"D103",
|
|
195
|
+
"D104",
|
|
196
|
+
"D105",
|
|
197
|
+
"D106",
|
|
198
|
+
"D107",
|
|
199
|
+
"S101", # use of "assert"
|
|
200
|
+
"S102", # use of "exec"
|
|
201
|
+
"S106", # possible hardcoded password.
|
|
202
|
+
"PGH001", # use of "eval"
|
|
203
|
+
]
|
|
204
|
+
|
|
205
|
+
[tool.ruff.lint.pep8-naming]
|
|
206
|
+
staticmethod-decorators = [
|
|
207
|
+
"pydantic.validator",
|
|
208
|
+
"pydantic.root_validator",
|
|
209
|
+
]
|
|
210
|
+
|
|
211
|
+
[tool.ruff.lint.isort]
|
|
212
|
+
force-sort-within-sections = true
|
|
213
|
+
|
|
214
|
+
[tool.ruff.format]
|
|
215
|
+
docstring-code-format = true
|
|
216
|
+
quote-style = "double"
|
|
217
|
+
|
|
218
|
+
[tool.bandit.assert_used]
|
|
219
|
+
skips = ['*_test.py', '*/test_*.py']
|
|
220
|
+
|
|
221
|
+
[tool.pyright]
|
|
222
|
+
pythonVersion = "3.10"
|
|
223
|
+
pythonPlatform = "Linux"
|
|
224
|
+
include = ["src"]
|
|
225
|
+
reportMissingImports = true
|
|
226
|
+
reportMissingTypeStubs = false
|
|
@@ -0,0 +1,498 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
import tempfile
|
|
4
|
+
from typing import Any, cast
|
|
5
|
+
|
|
6
|
+
from lightning.fabric.utilities.rank_zero import (
|
|
7
|
+
rank_zero_info, # type: ignore # noqa: PGH003
|
|
8
|
+
rank_zero_only,
|
|
9
|
+
rank_zero_warn, # type: ignore # noqa: PGH003
|
|
10
|
+
)
|
|
11
|
+
import lightning.pytorch as pl
|
|
12
|
+
from lightning.pytorch.callbacks import Callback, DeviceStatsMonitor, EarlyStopping
|
|
13
|
+
from lightning.pytorch.core.optimizer import LightningOptimizer
|
|
14
|
+
from lightning.pytorch.loggers import Logger, MLFlowLogger
|
|
15
|
+
from lightning.pytorch.loggers.mlflow import _get_resolve_tags
|
|
16
|
+
from lightning.pytorch.trainer.states import TrainerFn
|
|
17
|
+
from lightning.pytorch.utilities.model_summary.model_summary import ModelSummary
|
|
18
|
+
from lightning.pytorch.utilities.types import STEP_OUTPUT
|
|
19
|
+
import mlflow
|
|
20
|
+
from mlflow import MlflowClient, MlflowException
|
|
21
|
+
from mlflow.entities.dataset import Dataset
|
|
22
|
+
from mlflow.entities.dataset_input import DatasetInput
|
|
23
|
+
from mlflow.entities.input_tag import InputTag
|
|
24
|
+
from mlflow.entities.run import Run
|
|
25
|
+
from mlflow.entities.run_status import RunStatus
|
|
26
|
+
from mlflow.models import Model
|
|
27
|
+
from mlflow.pytorch import pickle_module as mlflow_pytorch_pickle_module
|
|
28
|
+
from torch.optim import Optimizer # pyright: ignore[reportPrivateImportUsage]
|
|
29
|
+
|
|
30
|
+
from lit_mlflow.logger import DbxMLFlowLogger
|
|
31
|
+
from lit_mlflow.utils.dbx import get_databricks_tags
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class MlFlowAutoCallback(Callback):
|
|
35
|
+
def __init__(self, verbose: bool = True, patch_device_monitor: bool = True) -> None:
|
|
36
|
+
self.supported_loggers = (MLFlowLogger, DbxMLFlowLogger)
|
|
37
|
+
self.verbose = verbose
|
|
38
|
+
self.logger: MLFlowLogger | DbxMLFlowLogger | None = None
|
|
39
|
+
self.autologging_disabled = False
|
|
40
|
+
self.patch_device_monitor = patch_device_monitor
|
|
41
|
+
|
|
42
|
+
@property
|
|
43
|
+
def client(self) -> MlflowClient | None:
|
|
44
|
+
if self.logger:
|
|
45
|
+
return self.logger.experiment
|
|
46
|
+
return None
|
|
47
|
+
|
|
48
|
+
def _get_logger(self, loggers: list[Logger]) -> MLFlowLogger | DbxMLFlowLogger | None:
|
|
49
|
+
if isinstance(loggers, list):
|
|
50
|
+
if len(loggers) == 0:
|
|
51
|
+
rank_zero_warn("Cannot log artifacts because Trainer has no logger.")
|
|
52
|
+
return None
|
|
53
|
+
else:
|
|
54
|
+
rank_zero_info(f" Supported loggers are: {', '.join(str(x.__name__) for x in self.supported_loggers)}")
|
|
55
|
+
for logger in loggers:
|
|
56
|
+
if isinstance(logger, self.supported_loggers):
|
|
57
|
+
return logger
|
|
58
|
+
else:
|
|
59
|
+
rank_zero_warn(
|
|
60
|
+
f"{self.__class__.__name__} does not support logging with {logger.__class__.__name__}."
|
|
61
|
+
)
|
|
62
|
+
return None
|
|
63
|
+
|
|
64
|
+
def _prevent_entry(self, trainer: "pl.Trainer") -> bool:
|
|
65
|
+
return self.logger is None or not trainer.is_global_zero
|
|
66
|
+
|
|
67
|
+
def _get_optimizer(self, optimizer: LightningOptimizer | Optimizer) -> Optimizer:
|
|
68
|
+
return optimizer._optimizer if isinstance(optimizer, LightningOptimizer) else optimizer
|
|
69
|
+
|
|
70
|
+
def _get_optimizer_name(self, optimizer: LightningOptimizer | Optimizer) -> str:
|
|
71
|
+
opt = self._get_optimizer(optimizer)
|
|
72
|
+
return opt.__class__.__name__
|
|
73
|
+
|
|
74
|
+
def _get_optimizer_defaults(self, optimizer: LightningOptimizer | Optimizer) -> dict[str, Any]:
|
|
75
|
+
opt = self._get_optimizer(optimizer)
|
|
76
|
+
return opt.defaults if hasattr(opt, "defaults") else {}
|
|
77
|
+
|
|
78
|
+
def _log_early_stop_params(self, early_stop_callback: EarlyStopping) -> None:
|
|
79
|
+
"""Logs early stopping configuration parameters to MLflow."""
|
|
80
|
+
if self.logger is None:
|
|
81
|
+
return None
|
|
82
|
+
|
|
83
|
+
params = {
|
|
84
|
+
p: getattr(early_stop_callback, p)
|
|
85
|
+
for p in ["monitor", "mode", "patience", "min_delta", "stopped_epoch"]
|
|
86
|
+
if hasattr(early_stop_callback, p)
|
|
87
|
+
}
|
|
88
|
+
self.logger.log_hyperparams(params)
|
|
89
|
+
|
|
90
|
+
def _log_early_stop_metrics(self, early_stop_callback: EarlyStopping) -> None:
|
|
91
|
+
"""Logs early stopping behavior results (e.g. stopped epoch) as metrics to MLflow."""
|
|
92
|
+
if self.logger is None:
|
|
93
|
+
return None
|
|
94
|
+
|
|
95
|
+
if early_stop_callback is None or early_stop_callback.stopped_epoch == 0:
|
|
96
|
+
return None
|
|
97
|
+
|
|
98
|
+
metrics: dict[str, float] = {
|
|
99
|
+
"stopped_epoch": early_stop_callback.stopped_epoch,
|
|
100
|
+
"restored_epoch": early_stop_callback.stopped_epoch - max(1, early_stop_callback.patience),
|
|
101
|
+
}
|
|
102
|
+
|
|
103
|
+
if hasattr(early_stop_callback, "best_score"):
|
|
104
|
+
metrics["best_score"] = float(early_stop_callback.best_score)
|
|
105
|
+
|
|
106
|
+
if hasattr(early_stop_callback, "wait_count"):
|
|
107
|
+
metrics["wait_count"] = early_stop_callback.wait_count
|
|
108
|
+
|
|
109
|
+
self.logger.log_metrics(metrics)
|
|
110
|
+
|
|
111
|
+
def _resolve_early_stopping_callback(self, trainer: "pl.Trainer") -> EarlyStopping | None:
|
|
112
|
+
if hasattr(trainer, "callbacks"):
|
|
113
|
+
for callback in cast(list[Callback], trainer.callbacks): # pyright: ignore[reportAttributeAccessIssue]
|
|
114
|
+
if isinstance(callback, EarlyStopping):
|
|
115
|
+
return callback
|
|
116
|
+
return None
|
|
117
|
+
|
|
118
|
+
def _log_model_summary(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
119
|
+
summary = str(ModelSummary(pl_module, max_depth=-1))
|
|
120
|
+
artifact_path = "model_summary.txt"
|
|
121
|
+
if self.logger and self.logger._run_id and self.client:
|
|
122
|
+
with tempfile.TemporaryDirectory(prefix="test", suffix="test", dir=Path.cwd()) as tmp_dir:
|
|
123
|
+
with Path.open(Path(f"{tmp_dir}/{artifact_path}"), "w") as tmp_file_summary:
|
|
124
|
+
tmp_file_summary.write(summary)
|
|
125
|
+
run_id = str(self.logger.run_id)
|
|
126
|
+
self.client.log_artifacts(run_id, tmp_dir, artifact_path)
|
|
127
|
+
|
|
128
|
+
def _log_model(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
129
|
+
if self.logger and self.logger.run_id and self.client:
|
|
130
|
+
rank_zero_info("Saving the model and uploading to MLFlow!")
|
|
131
|
+
with tempfile.TemporaryDirectory(prefix="test", suffix="test", dir=Path.cwd()) as tmp_dir:
|
|
132
|
+
local_path = Path(tmp_dir) / "model"
|
|
133
|
+
artifact_path = "model"
|
|
134
|
+
mlflow_model = Model(artifact_path=artifact_path, run_id=self.logger.run_id)
|
|
135
|
+
mlflow.pytorch.save_model(
|
|
136
|
+
pytorch_model=pl_module,
|
|
137
|
+
path=local_path,
|
|
138
|
+
conda_env=None,
|
|
139
|
+
mlflow_model=mlflow_model,
|
|
140
|
+
code_paths=None,
|
|
141
|
+
pickle_module=mlflow_pytorch_pickle_module,
|
|
142
|
+
signature=None,
|
|
143
|
+
input_example=None,
|
|
144
|
+
requirements_file=None,
|
|
145
|
+
extra_files=None,
|
|
146
|
+
pip_requirements=None,
|
|
147
|
+
extra_pip_requirements=None,
|
|
148
|
+
)
|
|
149
|
+
self.client.log_artifacts(
|
|
150
|
+
run_id=self.logger.run_id,
|
|
151
|
+
local_dir=tmp_dir,
|
|
152
|
+
artifact_path=artifact_path,
|
|
153
|
+
)
|
|
154
|
+
try:
|
|
155
|
+
self.client._record_logged_model(run_id=self.logger.run_id, mlflow_model=mlflow_model)
|
|
156
|
+
except MlflowException:
|
|
157
|
+
rank_zero_warn(
|
|
158
|
+
f"Logging model metadata to the tracking server {self.logger._tracking_uri} has failed"
|
|
159
|
+
)
|
|
160
|
+
|
|
161
|
+
# info = mlflow_model.get_model_info()
|
|
162
|
+
return None
|
|
163
|
+
|
|
164
|
+
def _print_auto_logged_info(self) -> None:
|
|
165
|
+
if self.logger and self.logger.run_id and self.client:
|
|
166
|
+
run = mlflow.get_run(run_id=self.logger.run_id)
|
|
167
|
+
if run:
|
|
168
|
+
artifacts = [f.path for f in self.client.list_artifacts(run.info.run_id, "model")]
|
|
169
|
+
tags = {k: v for k, v in run.data.tags.items() if not k.startswith("mlflow.")}
|
|
170
|
+
rank_zero_info(f"run_id: {run.info.run_id}")
|
|
171
|
+
rank_zero_info(f"artifacts: {artifacts}")
|
|
172
|
+
rank_zero_info(f"params: {run.data.params}")
|
|
173
|
+
rank_zero_info(f"metrics: {run.data.metrics}")
|
|
174
|
+
rank_zero_info(f"tags: {tags}")
|
|
175
|
+
|
|
176
|
+
def _log_cluster_tags(self) -> None:
|
|
177
|
+
tags = get_databricks_tags()
|
|
178
|
+
|
|
179
|
+
if self.logger and self.logger.run_id and self.client:
|
|
180
|
+
for tag, value in tags.items():
|
|
181
|
+
self.client.set_tag(self.logger.run_id, key=tag, value=value)
|
|
182
|
+
|
|
183
|
+
def _log_dataset_info(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
184
|
+
if hasattr(trainer, "datamodule") and trainer.datamodule: # pyright: ignore[reportAttributeAccessIssue]
|
|
185
|
+
dm = trainer.datamodule.train_dataloader() # pyright: ignore[reportAttributeAccessIssue]
|
|
186
|
+
dataset = dm.dataset # pyright: ignore[reportAttributeAccessIssue]
|
|
187
|
+
if self.logger and self.logger.run_id and self.client:
|
|
188
|
+
meta_ds = Dataset(
|
|
189
|
+
name=dataset.__class__.__name__,
|
|
190
|
+
digest="",
|
|
191
|
+
source=dm.__class__.__name__,
|
|
192
|
+
source_type="",
|
|
193
|
+
)
|
|
194
|
+
ds_input = DatasetInput(
|
|
195
|
+
dataset=meta_ds,
|
|
196
|
+
tags=[
|
|
197
|
+
InputTag("class", dataset.__class__.__name__),
|
|
198
|
+
],
|
|
199
|
+
)
|
|
200
|
+
self.client.log_inputs(run_id=self.logger.run_id, datasets=[ds_input])
|
|
201
|
+
|
|
202
|
+
def _patch_device_stats_monitor(self, trainer: "pl.Trainer") -> None:
|
|
203
|
+
def _patched_prefix_metric_keys(
|
|
204
|
+
metrics_dict: dict[str, float], prefix: str, separator: str
|
|
205
|
+
) -> dict[str, float]:
|
|
206
|
+
return {prefix + separator + k: v for k, v in metrics_dict.items()}
|
|
207
|
+
|
|
208
|
+
def _patched_get_and_log_device_stats(self, trainer: "pl.Trainer", key: str) -> None:
|
|
209
|
+
if not trainer._logger_connector.should_update_logs:
|
|
210
|
+
return
|
|
211
|
+
|
|
212
|
+
device = trainer.strategy.root_device
|
|
213
|
+
if self._cpu_stats is False and device.type == "cpu":
|
|
214
|
+
# cpu stats are disabled
|
|
215
|
+
return
|
|
216
|
+
|
|
217
|
+
device_stats = trainer.accelerator.get_device_stats(device)
|
|
218
|
+
|
|
219
|
+
if self._cpu_stats and device.type != "cpu":
|
|
220
|
+
# Don't query CPU stats twice if CPU is accelerator
|
|
221
|
+
from lightning.pytorch.accelerators.cpu import get_cpu_stats
|
|
222
|
+
|
|
223
|
+
device_stats.update(get_cpu_stats())
|
|
224
|
+
|
|
225
|
+
for logger in trainer.loggers:
|
|
226
|
+
separator = logger.group_separator
|
|
227
|
+
prefixed_device_stats = _patched_prefix_metric_keys(device_stats, f"system/{key}", separator)
|
|
228
|
+
logger.log_metrics(prefixed_device_stats, step=trainer.fit_loop.epoch_loop._batches_that_stepped)
|
|
229
|
+
|
|
230
|
+
patched = False
|
|
231
|
+
if hasattr(trainer, "callbacks"):
|
|
232
|
+
for callback in cast(list[Callback], trainer.callbacks): # pyright: ignore[reportAttributeAccessIssue]
|
|
233
|
+
if isinstance(callback, DeviceStatsMonitor):
|
|
234
|
+
callback._get_and_log_device_stats = _patched_get_and_log_device_stats.__get__(
|
|
235
|
+
callback, DeviceStatsMonitor
|
|
236
|
+
)
|
|
237
|
+
patched = True
|
|
238
|
+
rank_zero_info("Lightning device stats monitoring enabled!")
|
|
239
|
+
|
|
240
|
+
if not patched:
|
|
241
|
+
rank_zero_info("Lightning device stats monitor has not been added to callbacks!")
|
|
242
|
+
|
|
243
|
+
@rank_zero_only
|
|
244
|
+
def setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
|
|
245
|
+
"""Called when fit, validate, test, predict, or tune begins."""
|
|
246
|
+
if not self.autologging_disabled:
|
|
247
|
+
rank_zero_info("Starting MLFlow Databricks logging!")
|
|
248
|
+
rank_zero_info("Default auto logging disabled!")
|
|
249
|
+
mlflow.autolog(disable=True)
|
|
250
|
+
self.autologging_disabled = True
|
|
251
|
+
|
|
252
|
+
if trainer.is_global_zero:
|
|
253
|
+
self.logger = self._get_logger(trainer.loggers)
|
|
254
|
+
self._log_cluster_tags()
|
|
255
|
+
self._log_dataset_info(trainer, pl_module)
|
|
256
|
+
|
|
257
|
+
if self.patch_device_monitor:
|
|
258
|
+
self._patch_device_stats_monitor(trainer)
|
|
259
|
+
|
|
260
|
+
if not self._prevent_entry(trainer) and self.logger and self.logger.run_id and self.client:
|
|
261
|
+
self.client.update_run(run_id=self.logger.run_id, status=RunStatus.to_string(RunStatus.RUNNING))
|
|
262
|
+
|
|
263
|
+
@rank_zero_only
|
|
264
|
+
def teardown(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", stage: str) -> None:
|
|
265
|
+
"""Called when fit, validate, test, predict, or tune ends."""
|
|
266
|
+
if self._prevent_entry(trainer):
|
|
267
|
+
return None
|
|
268
|
+
|
|
269
|
+
if self.logger and self.logger.run_id and self.client:
|
|
270
|
+
self.client.update_run(run_id=self.logger.run_id, status=RunStatus.to_string(RunStatus.FINISHED))
|
|
271
|
+
|
|
272
|
+
if stage == TrainerFn.FITTING:
|
|
273
|
+
self.client.set_terminated(run_id=self.logger.run_id, status=RunStatus.to_string(RunStatus.FINISHED))
|
|
274
|
+
|
|
275
|
+
def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
276
|
+
"""Called when fit begins."""
|
|
277
|
+
|
|
278
|
+
def on_fit_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
279
|
+
"""Called when fit ends."""
|
|
280
|
+
|
|
281
|
+
def on_sanity_check_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
282
|
+
"""Called when the validation sanity check starts."""
|
|
283
|
+
|
|
284
|
+
def on_sanity_check_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
285
|
+
"""Called when the validation sanity check ends."""
|
|
286
|
+
|
|
287
|
+
def on_train_batch_start(
|
|
288
|
+
self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", batch: Any, batch_idx: int
|
|
289
|
+
) -> None:
|
|
290
|
+
"""Called when the train batch begins."""
|
|
291
|
+
|
|
292
|
+
def on_train_batch_end(
|
|
293
|
+
self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", outputs: STEP_OUTPUT, batch: Any, batch_idx: int
|
|
294
|
+
) -> None:
|
|
295
|
+
"""Called when the train batch ends.
|
|
296
|
+
|
|
297
|
+
Note:
|
|
298
|
+
The value ``outputs["loss"]`` here will be the normalized value w.r.t ``accumulate_grad_batches`` of the
|
|
299
|
+
loss returned from ``training_step``.
|
|
300
|
+
|
|
301
|
+
"""
|
|
302
|
+
|
|
303
|
+
def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
304
|
+
"""Called when the train epoch begins."""
|
|
305
|
+
|
|
306
|
+
def on_train_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
307
|
+
"""Called when the train epoch ends.
|
|
308
|
+
|
|
309
|
+
To access all batch outputs at the end of the epoch, you can cache step outputs as an attribute of the
|
|
310
|
+
:class:`lightning.pytorch.core.LightningModule` and access them in this hook:
|
|
311
|
+
|
|
312
|
+
.. code-block:: python
|
|
313
|
+
|
|
314
|
+
class MyLightningModule(L.LightningModule):
|
|
315
|
+
def __init__(self):
|
|
316
|
+
super().__init__()
|
|
317
|
+
self.training_step_outputs = []
|
|
318
|
+
|
|
319
|
+
def training_step(self):
|
|
320
|
+
loss = ...
|
|
321
|
+
self.training_step_outputs.append(loss)
|
|
322
|
+
return loss
|
|
323
|
+
|
|
324
|
+
|
|
325
|
+
class MyCallback(L.Callback):
|
|
326
|
+
def on_train_epoch_end(self, trainer, pl_module):
|
|
327
|
+
# do something with all training_step outputs, for example:
|
|
328
|
+
epoch_mean = torch.stack(pl_module.training_step_outputs).mean()
|
|
329
|
+
pl_module.log("training_epoch_mean", epoch_mean)
|
|
330
|
+
# free up the memory
|
|
331
|
+
pl_module.training_step_outputs.clear()
|
|
332
|
+
|
|
333
|
+
"""
|
|
334
|
+
if self.logger:
|
|
335
|
+
metrics = {str(key): float(value) for key, value in trainer.callback_metrics.items()}
|
|
336
|
+
self.logger.log_metrics(metrics, pl_module.current_epoch)
|
|
337
|
+
|
|
338
|
+
def on_validation_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
339
|
+
"""Called when the val epoch begins."""
|
|
340
|
+
|
|
341
|
+
def on_validation_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
342
|
+
"""Called when the val epoch ends."""
|
|
343
|
+
|
|
344
|
+
def on_test_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
345
|
+
"""Called when the test epoch begins."""
|
|
346
|
+
|
|
347
|
+
def on_test_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
348
|
+
"""Called when the test epoch ends."""
|
|
349
|
+
|
|
350
|
+
def on_predict_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
351
|
+
"""Called when the predict epoch begins."""
|
|
352
|
+
|
|
353
|
+
def on_predict_epoch_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
354
|
+
"""Called when the predict epoch ends."""
|
|
355
|
+
|
|
356
|
+
def on_validation_batch_start(
|
|
357
|
+
self,
|
|
358
|
+
trainer: "pl.Trainer",
|
|
359
|
+
pl_module: "pl.LightningModule",
|
|
360
|
+
batch: Any,
|
|
361
|
+
batch_idx: int,
|
|
362
|
+
dataloader_idx: int = 0,
|
|
363
|
+
) -> None:
|
|
364
|
+
"""Called when the validation batch begins."""
|
|
365
|
+
|
|
366
|
+
def on_validation_batch_end(
|
|
367
|
+
self,
|
|
368
|
+
trainer: "pl.Trainer",
|
|
369
|
+
pl_module: "pl.LightningModule",
|
|
370
|
+
outputs: STEP_OUTPUT,
|
|
371
|
+
batch: Any,
|
|
372
|
+
batch_idx: int,
|
|
373
|
+
dataloader_idx: int = 0,
|
|
374
|
+
) -> None:
|
|
375
|
+
"""Called when the validation batch ends."""
|
|
376
|
+
|
|
377
|
+
def on_test_batch_start(
|
|
378
|
+
self,
|
|
379
|
+
trainer: "pl.Trainer",
|
|
380
|
+
pl_module: "pl.LightningModule",
|
|
381
|
+
batch: Any,
|
|
382
|
+
batch_idx: int,
|
|
383
|
+
dataloader_idx: int = 0,
|
|
384
|
+
) -> None:
|
|
385
|
+
"""Called when the test batch begins."""
|
|
386
|
+
|
|
387
|
+
def on_test_batch_end(
|
|
388
|
+
self,
|
|
389
|
+
trainer: "pl.Trainer",
|
|
390
|
+
pl_module: "pl.LightningModule",
|
|
391
|
+
outputs: STEP_OUTPUT,
|
|
392
|
+
batch: Any,
|
|
393
|
+
batch_idx: int,
|
|
394
|
+
dataloader_idx: int = 0,
|
|
395
|
+
) -> None:
|
|
396
|
+
"""Called when the test batch ends."""
|
|
397
|
+
|
|
398
|
+
def on_predict_batch_start(
|
|
399
|
+
self,
|
|
400
|
+
trainer: "pl.Trainer",
|
|
401
|
+
pl_module: "pl.LightningModule",
|
|
402
|
+
batch: Any,
|
|
403
|
+
batch_idx: int,
|
|
404
|
+
dataloader_idx: int = 0,
|
|
405
|
+
) -> None:
|
|
406
|
+
"""Called when the predict batch begins."""
|
|
407
|
+
|
|
408
|
+
def on_predict_batch_end(
|
|
409
|
+
self,
|
|
410
|
+
trainer: "pl.Trainer",
|
|
411
|
+
pl_module: "pl.LightningModule",
|
|
412
|
+
outputs: Any,
|
|
413
|
+
batch: Any,
|
|
414
|
+
batch_idx: int,
|
|
415
|
+
dataloader_idx: int = 0,
|
|
416
|
+
) -> None:
|
|
417
|
+
"""Called when the predict batch ends."""
|
|
418
|
+
|
|
419
|
+
@rank_zero_only
|
|
420
|
+
def on_train_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
421
|
+
"""Called when the train begins."""
|
|
422
|
+
if self._prevent_entry(trainer):
|
|
423
|
+
return None
|
|
424
|
+
|
|
425
|
+
if self.logger and self.logger.run_id and self.client:
|
|
426
|
+
run_id = str(self.logger.run_id)
|
|
427
|
+
|
|
428
|
+
self.client.set_tag(run_id=run_id, key="Mode", value="training")
|
|
429
|
+
|
|
430
|
+
self.client.log_param(run_id=run_id, key="epochs", value=trainer.max_epochs)
|
|
431
|
+
|
|
432
|
+
if hasattr(trainer, "optimizers"):
|
|
433
|
+
for i, optimizer in enumerate(trainer.optimizers):
|
|
434
|
+
self.client.log_param(
|
|
435
|
+
self.logger.run_id, key=f"optimizer{i}_name", value=self._get_optimizer_name(optimizer)
|
|
436
|
+
)
|
|
437
|
+
defaults = self._get_optimizer_defaults(optimizer)
|
|
438
|
+
for key, value in defaults.items():
|
|
439
|
+
self.client.log_param(self.logger.run_id, key=f"optimizer{i}_{key}", value=str(value))
|
|
440
|
+
# self.client.log_param(
|
|
441
|
+
# self.logger.run_id, key=f"optimizer{i}_defaults", value=str(optimizer.defaults)
|
|
442
|
+
# )
|
|
443
|
+
|
|
444
|
+
callback = self._resolve_early_stopping_callback(trainer)
|
|
445
|
+
if callback:
|
|
446
|
+
self._log_early_stop_params(callback)
|
|
447
|
+
|
|
448
|
+
@rank_zero_only
|
|
449
|
+
def on_train_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
450
|
+
"""Called when the train ends."""
|
|
451
|
+
if self._prevent_entry(trainer):
|
|
452
|
+
return None
|
|
453
|
+
|
|
454
|
+
callback = self._resolve_early_stopping_callback(trainer)
|
|
455
|
+
if callback:
|
|
456
|
+
self._log_early_stop_metrics(callback)
|
|
457
|
+
|
|
458
|
+
self._log_model_summary(trainer, pl_module)
|
|
459
|
+
|
|
460
|
+
self._log_model(trainer, pl_module)
|
|
461
|
+
|
|
462
|
+
@rank_zero_only
|
|
463
|
+
def on_validation_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
464
|
+
"""Called when the validation loop begins."""
|
|
465
|
+
if self.logger and self.logger.run_id and self.client:
|
|
466
|
+
run_id = str(self.logger.run_id)
|
|
467
|
+
|
|
468
|
+
self.client.set_tag(run_id=run_id, key="Mode", value="validating")
|
|
469
|
+
|
|
470
|
+
@rank_zero_only
|
|
471
|
+
def on_validation_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
472
|
+
"""Called when the validation loop ends."""
|
|
473
|
+
|
|
474
|
+
@rank_zero_only
|
|
475
|
+
def on_test_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
476
|
+
"""Called when the test begins."""
|
|
477
|
+
if self._prevent_entry(trainer):
|
|
478
|
+
return None
|
|
479
|
+
if self.logger and self.logger.run_id and self.client:
|
|
480
|
+
self.client.set_tag(self.logger.run_id, key="Mode", value="testing")
|
|
481
|
+
|
|
482
|
+
def on_test_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
483
|
+
"""Called when the test ends."""
|
|
484
|
+
# originally, mlflow.autolog changes the mode to testing here, but we do it in on_test_start
|
|
485
|
+
|
|
486
|
+
def on_predict_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
487
|
+
"""Called when the predict begins."""
|
|
488
|
+
|
|
489
|
+
def on_predict_end(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
|
|
490
|
+
"""Called when predict ends."""
|
|
491
|
+
|
|
492
|
+
@rank_zero_only
|
|
493
|
+
def on_exception(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", exception: BaseException) -> None:
|
|
494
|
+
"""Called when any trainer execution is interrupted by an exception."""
|
|
495
|
+
if self._prevent_entry(trainer):
|
|
496
|
+
return
|
|
497
|
+
if self.logger and self.logger.run_id and self.client:
|
|
498
|
+
self.client.set_terminated(run_id=self.logger.run_id, status=RunStatus.to_string(RunStatus.FAILED))
|
|
@@ -0,0 +1,76 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
from typing import Any, Literal
|
|
4
|
+
|
|
5
|
+
from lightning.fabric.loggers.logger import Logger, rank_zero_experiment
|
|
6
|
+
from lightning.fabric.utilities.rank_zero import (
|
|
7
|
+
rank_zero_info, # type: ignore # noqa: PGH003
|
|
8
|
+
rank_zero_only,
|
|
9
|
+
rank_zero_warn, # type: ignore # noqa: PGH003
|
|
10
|
+
)
|
|
11
|
+
from lightning.pytorch.loggers import MLFlowLogger
|
|
12
|
+
import mlflow
|
|
13
|
+
from mlflow.tracking import MlflowClient
|
|
14
|
+
from mlflow.utils.databricks_utils import get_databricks_run_url, is_in_databricks_notebook
|
|
15
|
+
|
|
16
|
+
from lit_mlflow.utils.dbx import get_experiment_id, is_in_databricks, patch_dbx_credentials
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class DbxMLFlowLogger(MLFlowLogger):
|
|
20
|
+
def __init__(
|
|
21
|
+
self,
|
|
22
|
+
experiment_name: str = "lightning_logs",
|
|
23
|
+
run_name: str | None = None,
|
|
24
|
+
tracking_uri: str | None = mlflow.get_tracking_uri(), # os.getenv("MLFLOW_TRACKING_URI"),
|
|
25
|
+
tags: dict[str, Any] | None = None,
|
|
26
|
+
save_dir: str | None = "./mlruns",
|
|
27
|
+
log_model: Literal[True, False, "all"] = False,
|
|
28
|
+
prefix: str = "",
|
|
29
|
+
artifact_location: str | None = None,
|
|
30
|
+
run_id: str | None = None,
|
|
31
|
+
):
|
|
32
|
+
super().__init__(
|
|
33
|
+
experiment_name, run_name, tracking_uri, tags, save_dir, log_model, prefix, artifact_location, run_id
|
|
34
|
+
)
|
|
35
|
+
if not is_in_databricks():
|
|
36
|
+
rank_zero_warn(f"You are running `{self.__class__.__name__}` outside of Databricks.")
|
|
37
|
+
else:
|
|
38
|
+
self._fix_logging()
|
|
39
|
+
if patch_dbx_credentials():
|
|
40
|
+
client = self.experiment
|
|
41
|
+
if client:
|
|
42
|
+
rank_zero_info("MLflow client created.")
|
|
43
|
+
else:
|
|
44
|
+
rank_zero_warn("Could not retrieve the MLflow client.")
|
|
45
|
+
else:
|
|
46
|
+
rank_zero_warn("Could not patch Databricks credentials.")
|
|
47
|
+
|
|
48
|
+
def _get_url(self) -> str | None:
|
|
49
|
+
run_url = get_databricks_run_url(self._tracking_uri or "databricks", self.run_id or "")
|
|
50
|
+
return run_url
|
|
51
|
+
|
|
52
|
+
@rank_zero_only
|
|
53
|
+
def _fix_logging(self) -> None:
|
|
54
|
+
"""Fix the logging level of the Py4J gateway to ERROR to prevent log spam."""
|
|
55
|
+
logging.getLogger("py4j.java_gateway").setLevel(logging.ERROR)
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
@rank_zero_experiment
|
|
59
|
+
def experiment(self) -> MlflowClient:
|
|
60
|
+
if is_in_databricks_notebook():
|
|
61
|
+
self._experiment_id = get_experiment_id(None)
|
|
62
|
+
return super().experiment
|
|
63
|
+
|
|
64
|
+
@property
|
|
65
|
+
def run_id(self) -> str | None:
|
|
66
|
+
_ = self.experiment
|
|
67
|
+
active_run = mlflow.active_run()
|
|
68
|
+
if active_run:
|
|
69
|
+
self._run_id = active_run.info.run_id or ""
|
|
70
|
+
return self._run_id
|
|
71
|
+
|
|
72
|
+
# @rank_zero_only
|
|
73
|
+
# def finalize(self, status: str = "success") -> None:
|
|
74
|
+
# final = super().finalize(status)
|
|
75
|
+
# rank_zero_info(f"MLflow run URL: {self._get_url()}")
|
|
76
|
+
# return final
|
|
File without changes
|
|
@@ -0,0 +1,74 @@
|
|
|
1
|
+
import json
|
|
2
|
+
import os
|
|
3
|
+
|
|
4
|
+
from mlflow.utils.databricks_utils import (
|
|
5
|
+
_get_property_from_spark_context,
|
|
6
|
+
get_notebook_id,
|
|
7
|
+
is_in_databricks_job,
|
|
8
|
+
is_in_databricks_model_serving_environment,
|
|
9
|
+
is_in_databricks_notebook,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def is_in_databricks() -> bool:
|
|
14
|
+
"""Check if the code is running in Databricks."""
|
|
15
|
+
return is_in_databricks_notebook() or is_in_databricks_job() or is_in_databricks_model_serving_environment()
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def patch_dbx_credentials() -> bool:
|
|
19
|
+
"""Patch/propagate the Databricks credentials to the environment variables, required for DDP.
|
|
20
|
+
|
|
21
|
+
Returns:
|
|
22
|
+
Returns True if the values have been patched.
|
|
23
|
+
"""
|
|
24
|
+
if is_in_databricks_notebook():
|
|
25
|
+
from databricks_cli.configure.provider import DatabricksConfig, get_config
|
|
26
|
+
|
|
27
|
+
config: DatabricksConfig = get_config()
|
|
28
|
+
if config:
|
|
29
|
+
for k, v in config.__dict__.items():
|
|
30
|
+
if v:
|
|
31
|
+
os.environ[f"DATABRICKS_{k.upper()}"] = str(v)
|
|
32
|
+
return True
|
|
33
|
+
return False
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def get_experiment_id(default_name: str | None = None) -> str | None:
|
|
37
|
+
"""Returns the experiment Id if running in Databricks.
|
|
38
|
+
|
|
39
|
+
Args:
|
|
40
|
+
default_name: The optional default name to return if not in Databricks.
|
|
41
|
+
|
|
42
|
+
Returns:
|
|
43
|
+
The experiment Id.
|
|
44
|
+
"""
|
|
45
|
+
if is_in_databricks_notebook():
|
|
46
|
+
notebook_id = get_notebook_id()
|
|
47
|
+
if notebook_id:
|
|
48
|
+
return notebook_id
|
|
49
|
+
return default_name
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def get_databricks_tags() -> dict[str, str]:
|
|
53
|
+
"""Return a selection of Databricks cluster tags.
|
|
54
|
+
|
|
55
|
+
Returns:
|
|
56
|
+
The tags dictionary.
|
|
57
|
+
"""
|
|
58
|
+
tags = {}
|
|
59
|
+
if is_in_databricks():
|
|
60
|
+
all_tags_str = _get_property_from_spark_context("spark.databricks.clusterUsageTags.clusterAllTags")
|
|
61
|
+
if all_tags_str:
|
|
62
|
+
tags_list = json.loads(all_tags_str)
|
|
63
|
+
for tag in tags_list:
|
|
64
|
+
tags[tag["key"]] = tags_list["value"]
|
|
65
|
+
node_type = _get_property_from_spark_context("spark.databricks.clusterUsageTags.clusterNodeType")
|
|
66
|
+
if node_type:
|
|
67
|
+
tags["clusterNodeType"] = str(node_type)
|
|
68
|
+
runtime_version = os.environ.get("DATABRICKS_RUNTIME_VERSION")
|
|
69
|
+
if runtime_version:
|
|
70
|
+
tags["runtimeVersion"] = str(runtime_version)
|
|
71
|
+
num_worker = _get_property_from_spark_context("spark.databricks.clusterUsageTags.clusterWorkers")
|
|
72
|
+
if num_worker:
|
|
73
|
+
tags["numWorkers"] = int(num_worker)
|
|
74
|
+
return tags
|