-
Notifications
You must be signed in to change notification settings - Fork 3.3k
/
mlflow.py
294 lines (238 loc) · 10.1 KB
/
mlflow.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
MLflow Logger
-------------
"""
import logging
import os
import re
from argparse import Namespace
from time import time
from typing import Any, Dict, Mapping, Optional, Union
from lightning_utilities.core.imports import module_available
from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment
from pytorch_lightning.utilities.logger import _add_prefix, _convert_params, _flatten_dict
from pytorch_lightning.utilities.rank_zero import rank_zero_only, rank_zero_warn
log = logging.getLogger(__name__)
LOCAL_FILE_URI_PREFIX = "file:"
_MLFLOW_AVAILABLE = module_available("mlflow")
try:
import mlflow
from mlflow.tracking import context, MlflowClient
from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME
# todo: there seems to be still some remaining import error with Conda env
except ModuleNotFoundError:
_MLFLOW_AVAILABLE = False
mlflow, MlflowClient, context = None, None, None
MLFLOW_RUN_NAME = "mlflow.runName"
# before v1.1.0
if hasattr(context, "resolve_tags"):
from mlflow.tracking.context import resolve_tags
# since v1.1.0
elif hasattr(context, "registry"):
from mlflow.tracking.context.registry import resolve_tags
else:
def resolve_tags(tags: Optional[Dict] = None) -> Optional[Dict]:
"""
Args:
tags: A dictionary of tags to override. If specified, tags passed in this argument will
override those inferred from the context.
Returns: A dictionary of resolved tags.
Note:
See ``mlflow.tracking.context.registry`` for more details.
"""
return tags
class MLFlowLogger(Logger):
"""Log using `MLflow <https://mlflow.org>`_.
Install it with pip:
.. code-block:: bash
pip install mlflow
.. code-block:: python
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import MLFlowLogger
mlf_logger = MLFlowLogger(experiment_name="lightning_logs", tracking_uri="file:./ml-runs")
trainer = Trainer(logger=mlf_logger)
Use the logger anywhere in your :class:`~pytorch_lightning.core.module.LightningModule` as follows:
.. code-block:: python
from pytorch_lightning import LightningModule
class LitModel(LightningModule):
def training_step(self, batch, batch_idx):
# example
self.logger.experiment.whatever_ml_flow_supports(...)
def any_lightning_module_function_or_hook(self):
self.logger.experiment.whatever_ml_flow_supports(...)
Args:
experiment_name: The name of the experiment.
run_name: Name of the new run. The `run_name` is internally stored as a ``mlflow.runName`` tag.
If the ``mlflow.runName`` tag has already been set in `tags`, the value is overridden by the `run_name`.
tracking_uri: Address of local or remote tracking server.
If not provided, defaults to `MLFLOW_TRACKING_URI` environment variable if set, otherwise it falls
back to `file:<save_dir>`.
tags: A dictionary tags for the experiment.
save_dir: A path to a local directory where the MLflow runs get saved.
Defaults to `./mlflow` if `tracking_uri` is not provided.
Has no effect if `tracking_uri` is provided.
prefix: A string to put at the beginning of metric keys.
artifact_location: The location to store run artifacts. If not provided, the server picks an appropriate
default.
run_id: The run identifier of the experiment. If not provided, a new run is started.
Raises:
ModuleNotFoundError:
If required MLFlow package is not installed on the device.
"""
LOGGER_JOIN_CHAR = "-"
def __init__(
self,
experiment_name: str = "lightning_logs",
run_name: Optional[str] = None,
tracking_uri: Optional[str] = os.getenv("MLFLOW_TRACKING_URI"),
tags: Optional[Dict[str, Any]] = None,
save_dir: Optional[str] = "./mlruns",
prefix: str = "",
artifact_location: Optional[str] = None,
run_id: Optional[str] = None,
):
if mlflow is None:
raise ModuleNotFoundError(
"You want to use `mlflow` logger which is not installed yet, install it with `pip install mlflow`."
)
super().__init__()
if not tracking_uri:
tracking_uri = f"{LOCAL_FILE_URI_PREFIX}{save_dir}"
self._experiment_name = experiment_name
self._experiment_id: Optional[str] = None
self._tracking_uri = tracking_uri
self._run_name = run_name
self._run_id = run_id
self.tags = tags
self._prefix = prefix
self._artifact_location = artifact_location
self._initialized = False
self._mlflow_client = MlflowClient(tracking_uri)
@property # type: ignore[misc]
@rank_zero_experiment
def experiment(self) -> MlflowClient:
r"""
Actual MLflow object. To use MLflow features in your
:class:`~pytorch_lightning.core.module.LightningModule` do the following.
Example::
self.logger.experiment.some_mlflow_function()
"""
if self._initialized:
return self._mlflow_client
if self._run_id is not None:
run = self._mlflow_client.get_run(self._run_id)
self._experiment_id = run.info.experiment_id
self._initialized = True
return self._mlflow_client
if self._experiment_id is None:
expt = self._mlflow_client.get_experiment_by_name(self._experiment_name)
if expt is not None:
self._experiment_id = expt.experiment_id
else:
log.warning(f"Experiment with name {self._experiment_name} not found. Creating it.")
self._experiment_id = self._mlflow_client.create_experiment(
name=self._experiment_name, artifact_location=self._artifact_location
)
if self._run_id is None:
if self._run_name is not None:
self.tags = self.tags or {}
if MLFLOW_RUN_NAME in self.tags:
log.warning(
f"The tag {MLFLOW_RUN_NAME} is found in tags. The value will be overridden by {self._run_name}."
)
self.tags[MLFLOW_RUN_NAME] = self._run_name
run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=resolve_tags(self.tags))
self._run_id = run.info.run_id
self._initialized = True
return self._mlflow_client
@property
def run_id(self) -> Optional[str]:
"""Create the experiment if it does not exist to get the run id.
Returns:
The run id.
"""
_ = self.experiment
return self._run_id
@property
def experiment_id(self) -> Optional[str]:
"""Create the experiment if it does not exist to get the experiment id.
Returns:
The experiment id.
"""
_ = self.experiment
return self._experiment_id
@rank_zero_only
def log_hyperparams(self, params: Union[Dict[str, Any], Namespace]) -> None:
params = _convert_params(params)
params = _flatten_dict(params)
for k, v in params.items():
if len(str(v)) > 250:
rank_zero_warn(
f"Mlflow only allows parameters with up to 250 characters. Discard {k}={v}", category=RuntimeWarning
)
continue
self.experiment.log_param(self.run_id, k, v)
@rank_zero_only
def log_metrics(self, metrics: Mapping[str, float], step: Optional[int] = None) -> None:
assert rank_zero_only.rank == 0, "experiment tried to log from global_rank != 0"
metrics = _add_prefix(metrics, self._prefix, self.LOGGER_JOIN_CHAR)
timestamp_ms = int(time() * 1000)
for k, v in metrics.items():
if isinstance(v, str):
log.warning(f"Discarding metric with string value {k}={v}.")
continue
new_k = re.sub("[^a-zA-Z0-9_/. -]+", "", k)
if k != new_k:
rank_zero_warn(
"MLFlow only allows '_', '/', '.' and ' ' special characters in metric name."
f" Replacing {k} with {new_k}.",
category=RuntimeWarning,
)
k = new_k
self.experiment.log_metric(self.run_id, k, v, timestamp_ms, step)
@rank_zero_only
def finalize(self, status: str = "success") -> None:
if not self._initialized:
return
if status == "success":
status = "FINISHED"
elif status == "failed":
status = "FAILED"
if self.experiment.get_run(self.run_id):
self.experiment.set_terminated(self.run_id, status)
@property
def save_dir(self) -> Optional[str]:
"""The root file directory in which MLflow experiments are saved.
Return:
Local path to the root experiment directory if the tracking uri is local.
Otherwise returns `None`.
"""
if self._tracking_uri.startswith(LOCAL_FILE_URI_PREFIX):
return self._tracking_uri.lstrip(LOCAL_FILE_URI_PREFIX)
@property
def name(self) -> Optional[str]:
"""Get the experiment id.
Returns:
The experiment id.
"""
return self.experiment_id
@property
def version(self) -> Optional[str]:
"""Get the run id.
Returns:
The run id.
"""
return self.run_id