Skip to content

Commit

Permalink
Update mlflow with using resolve_tags (#6746)
Browse files Browse the repository at this point in the history
* Update mlflow.py

#6745 adds additional info about the run, as in the native API

* Update mlflow.py

trying to fix some backward compatibility issues with `resolve_tags`

* wip on backward compatibility

added a default for `getattr` in case the `registry` object exists, but has no proper attribute (weird case but who knows...)

* fix pep

* impoert

* fix registry import

* try fix failing tests

removed the first if statement, so that `resolve_tags` would be defined either case

* fix formatting

Co-authored-by: Jirka Borovec <jirka.borovec@seznam.cz>
  • Loading branch information
stllfe and Borda authored Apr 8, 2021
1 parent eb15abc commit 3007872
Showing 1 changed file with 15 additions and 2 deletions.
17 changes: 15 additions & 2 deletions pytorch_lightning/loggers/mlflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,23 @@
try:
import mlflow
from mlflow.tracking import MlflowClient
from mlflow.tracking import context
# todo: there seems to be still some remaining import error with Conda env
except ImportError:
_MLFLOW_AVAILABLE = False
mlflow, MlflowClient = None, None
mlflow, MlflowClient, context = None, None, None


# 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=None):
return tags


class MLFlowLogger(LightningLoggerBase):
Expand Down Expand Up @@ -140,7 +153,7 @@ def experiment(self) -> MlflowClient:
)

if self._run_id is None:
run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=self.tags)
run = self._mlflow_client.create_run(experiment_id=self._experiment_id, tags=resolve_tags(self.tags))
self._run_id = run.info.run_id
return self._mlflow_client

Expand Down

0 comments on commit 3007872

Please sign in to comment.