Merge pull request #3 from comet-ml/update-mlflow-2.0.1

Update support to mlflow 2.0.1

.flake8

@@ -1,3 +1,3 @@
 [flake8]
-exclude = .git,__pycache__,data,tools
-ignore = E101, E111, E114, E115, E116, E117, E121, E122, E123, E124, E125, E126, E127, E128, E129, E131, E133, E2, E3, E5, E501, E701, E702, E703, E704, W1, W2, W3, W503, W504
+max-line-length = 100
+extend-ignore = E203

.github/workflows/package.yml

@@ -0,0 +1,31 @@
+name: package
+
+on:
+  push:
+  release:
+    types: [published]
+
+
+jobs:
+    build-package:
+      runs-on: ubuntu-latest
+      steps:
+        - uses: actions/checkout@v3
+        - name: Set up Python 3.9
+          uses: actions/setup-python@v3
+          with:
+              python-version: 3.9
+        - name: Build pip package
+          run: |
+              pip install -U pip build
+              python3 -m build --sdist --wheel --outdir dist/ .
+        - name: Archive Pypi artifacts
+          uses: actions/upload-artifact@v3
+          with:
+              name: pypi_dist
+              path: dist
+        - name: Publish package
+          if: startsWith(github.ref, 'refs/tags')
+          uses: pypa/gh-action-pypi-publish@release/v1
+          with:
+            password: ${{ secrets.TWINE_PROD_PASSWORD }}

.github/workflows/tests.yml

@@ -43,29 +43,3 @@ jobs:
           pip install pre-commit
     - name: Run pre-commit
       run: pre-commit run --all-files
-  package:
-    runs-on: ubuntu-18.04
-    needs: [test, lint]
-    steps:
-    - uses: actions/checkout@v1
-    - name: Set up Python 3.7
-      uses: actions/setup-python@v1
-      with:
-        python-version: 3.7
-    - name: Package and upload to Pypi Test
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.TWINE_TEST_PASSWORD }}
-      run: |
-        python -m pip install --upgrade pip
-        pip install -U wheel twine
-        python setup.py bdist_wheel sdist
-        ls dist/
-        twine upload --repository-url https://test.pypi.org/legacy/ --skip-existing dist/*
-    - name: Upload to Pypi production if it's a new tag
-      if: contains(github.ref, 'tags')
-      env:
-        TWINE_USERNAME: __token__
-        TWINE_PASSWORD: ${{ secrets.TWINE_PROD_PASSWORD }}
-      run: |
-        twine upload dist/*

.pre-commit-config.yaml

@@ -1,39 +1,34 @@
 repos:
--   repo: https://github.com/asottile/seed-isort-config
-    rev: v1.9.4
-    hooks:
-    -   id: seed-isort-config
--   repo: https://github.com/pre-commit/mirrors-isort
-    rev: v4.3.21
+-   repo: https://github.com/PyCQA/isort
+    rev: 5.10.1
     hooks:
     -   id: isort
 -   repo: https://github.com/ambv/black
-    rev: 19.10b0
+    rev: 22.10.0
     hooks:
     - id: black
--   repo: https://gitlab.com/pycqa/flake8
-    rev: 3.7.9
+-   repo: https://github.com/pycqa/flake8
+    rev: 5.0.4
     hooks:
     -   id: flake8
-        additional_dependencies: ['flake8-coding==1.3.2', 'flake8-copyright==0.2.2', 'flake8-debugger==3.2.1', 'flake8-mypy==17.8.0']
+        args: ['--config=.flake8']
+        additional_dependencies: ['flake8-coding==1.3.2', 'flake8-copyright==0.2.3', 'flake8-debugger==4.1.2', 'flake8-mypy==17.8.0']
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v2.5.0
+    rev: v4.3.0
     hooks:
-    -   id: trailing-whitespace
+    -   id: check-json
+    -   id: check-merge-conflict
     -   id: check-yaml
+    -   id: debug-statements
     -   id: mixed-line-ending
-    -   id: name-tests-test
-        args: ['--django']
     -   id: requirements-txt-fixer
+    -   id: trailing-whitespace
 -   repo: https://github.com/codespell-project/codespell
-    rev: v1.16.0
+    rev: v2.2.2
     hooks:
     -   id: codespell
         exclude_types: [json]
         args: ['--ignore-words-list=feld']
 -   repo: meta
     hooks:
-    -   id: check-hooks-apply
     -   id: check-useless-excludes
-default_language_version:
-    python: python3.7

comet_for_mlflow/cli.py

@@ -42,17 +42,22 @@ def main():
         "--no-upload",
         dest="upload",
         action="store_false",
-        help="Do not upload the prepared experiments to comet.ml; will not create comet.ml projects",
+        help="Do not upload the prepared experiments to comet.ml;"
+        " will not create comet.ml projects",
     )
     parser.set_defaults(upload=True)
 
     parser.add_argument(
         "--api-key",
-        help="Set the Comet API key; required with --upload (the default); can also be configured in the usual places",
+        help="Set the Comet API key; required with --upload (the default);"
+        " can also be configured in the usual places",
     )
     parser.add_argument(
         "--mlflow-store-uri",
-        help="Set the MLFlow store uri. The MLFlow store uri to used to retrieve MLFlow runs, given directly to MLFlow, and supports all MLFlow schemes (file:// or SQLAlchemy-compatible database connection strings). If not set, reads MLFLOW_TRACKING_URI environment variable",
+        help="Set the MLFlow store uri. The MLFlow store uri to used to retrieve"
+        " MLFlow runs, given directly to MLFlow, and supports all MLFlow schemes"
+        " (file:// or SQLAlchemy-compatible database connection strings)."
+        " If not set, reads MLFLOW_TRACKING_URI environment variable",
     )
 
     parser.add_argument(
@@ -83,7 +88,8 @@ def main():
         help="Answer all yes/no questions automatically with 'no'",
     )
     parser.add_argument(
-        "--email", help="Set email address if needed for creating a comet.ml account",
+        "--email",
+        help="Set email address if needed for creating a comet.ml account",
     )
 
     args = parser.parse_args()

comet_for_mlflow/comet_for_mlflow.py

@@ -41,13 +41,19 @@
 from comet_ml.exceptions import CometRestApiException
 from comet_ml.offline import upload_single_offline_experiment
 from mlflow.entities.run_tag import RunTag
-from mlflow.entities.view_type import ViewType
 from mlflow.tracking import _get_store
 from mlflow.tracking._model_registry.utils import _get_store as get_model_registry_store
 from mlflow.tracking.registry import UnsupportedModelRegistryStoreURIException
 from tabulate import tabulate
 from tqdm import tqdm
 
+from .compat import (
+    get_artifact_repository,
+    get_mlflow_model_name,
+    get_mlflow_run_id,
+    search_mlflow_store_experiments,
+    search_mlflow_store_runs,
+)
 from .file_writer import JsonLinesFile
 from .utils import (
     get_comet_project_name,
@@ -65,18 +71,10 @@
     pass
 
 
-try:
-    # MLFLOW version 1.4.0
-    from mlflow.store.artifact.artifact_repository_registry import (
-        get_artifact_repository,
-    )
-except ImportError:
-    # MLFLOW version < 1.4.0
-    from mlflow.store.artifact_repository_registry import get_artifact_repository
-
 logging.basicConfig(level=logging.INFO, format="%(message)s")
 LOGGER = logging.getLogger()
 
+
 # Install a global exception hook
 def except_hook(exc_type, exc_value, exc_traceback):
     Reporting.report(
@@ -137,8 +135,7 @@ def __init__(
         except UnsupportedModelRegistryStoreURIException:
             self.model_registry_store = None
 
-        # Most of list_experiments returns a list anyway
-        self.mlflow_experiments = list(self.store.list_experiments())
+        self.mlflow_experiments = search_mlflow_store_experiments(self.store)
         self.len_experiments = len(self.mlflow_experiments)  # We start counting at 0
 
         self.summary = {
@@ -239,22 +236,28 @@ def prepare(self):
 
         LOGGER.info("")
         LOGGER.info(
-            "If you need support, you can contact us at http://chat.comet.ml/ or https://comet.ml/docs/quick-start/#getting-support"
+            """If you need support, you can contact us at http://chat.comet.ml/"""
+            """ or https://comet.ml/docs/quick-start/#getting-support"""
         )
         LOGGER.info("")
 
     def prepare_mlflow_exp(
-        self, exp,
+        self,
+        exp,
     ):
-        runs_info = self.store.list_run_infos(exp.experiment_id, ViewType.ALL)
+        runs_info = search_mlflow_store_runs(self.store, exp.experiment_id)
         len_runs = len(runs_info)
 
         for run_number, run_info in enumerate(runs_info):
             try:
-                run_id = run_info.run_id
+                run_id = get_mlflow_run_id(run_info)
+
                 run = self.store.get_run(run_id)
                 LOGGER.info(
-                    "## Preparing run %d/%d [%s]", run_number + 1, len_runs, run_id,
+                    "## Preparing run %d/%d [%s]",
+                    run_number + 1,
+                    len_runs,
+                    run_id,
                 )
                 LOGGER.debug(
                     "## Preparing run %d/%d: %r", run_number + 1, len_runs, run
@@ -410,15 +413,25 @@ def prepare_single_mlflow_run(self, run, original_experiment_name):
                         break
 
                 if matching_model:
+                    model_name = get_mlflow_model_name(matching_model)
+
+                    prefix = "models/"
+                    if artifact_path.startswith(prefix):
+                        comet_artifact_path = artifact_path[len(prefix) :]
+                    else:
+                        comet_artifact_path = artifact_path
+
                     json_writer.log_artifact_as_model(
                         local_artifact_path,
-                        artifact_path,
+                        comet_artifact_path,
                         run_start_time,
-                        matching_model.registered_model.name,
+                        model_name,
                     )
                 else:
                     json_writer.log_artifact_as_asset(
-                        local_artifact_path, artifact_path, run_start_time,
+                        local_artifact_path,
+                        artifact_path,
+                        run_start_time,
                     )
 
         return self.compress_archive(run.info.run_id)
@@ -438,12 +451,15 @@ def upload(self, prepared_data):
                 project_note = experiment.tags.get("mlflow.note.content", None)
                 if project_note:
                     note_template = (
-                        u"/!\\ This project notes has been copied from MLFlow. It might be overwritten if you run comet_for_mlflow again/!\\ \n%s"
+                        "/!\\ This project notes has been copied from MLFlow."
+                        " It might be overwritten if you run comet_for_mlflow again/!\\ \n%s"
                         % project_note
                     )
                     # We don't support Unicode project notes yet
                     self.api_client.set_project_notes(
-                        self.workspace, project_name, note_template,
+                        self.workspace,
+                        project_name,
+                        note_template,
                     )
 
                 all_project_names.append(project_name)
@@ -487,7 +503,8 @@ def upload(self, prepared_data):
             LOGGER.info("\t- %s", url)
 
         LOGGER.info(
-            "Get deeper instrumentation by adding Comet SDK to your project: https://comet.ml/docs/python-sdk/mlflow/"
+            "Get deeper instrumentation by adding Comet SDK to your project:"
+            " https://comet.ml/docs/python-sdk/mlflow/"
         )
         LOGGER.info("")
 
@@ -598,19 +615,23 @@ def create_or_login(self):
             Reporting.report("mlflow_new_user", api_key=new_account["apiKey"])
 
             LOGGER.info(
-                "A Comet.ml account has been created for you and an email was sent to you to setup your password later."
+                "A Comet.ml account has been created for you and an email was sent to"
+                " you to setup your password later."
             )
             save_api_key(new_account["apiKey"])
             LOGGER.info(
-                "Your Comet API Key has been saved to ~/.comet.ini, it is also available on your Comet.ml dashboard."
+                "Your Comet API Key has been saved to ~/.comet.ini, it is also"
+                " available on your Comet.ml dashboard."
             )
             return (
                 new_account["apiKey"],
                 new_account["token"],
             )
         else:
             LOGGER.info(
-                "An account already exists for this account, please input your API Key below (you can find it in your Settings page, https://comet.ml/docs/quick-start/#getting-your-comet-api-key):"
+                "An account already exists for this account, please input your API Key"
+                " below (you can find it in your Settings page,"
+                " https://comet.ml/docs/quick-start/#getting-your-comet-api-key):"
             )
             api_key = input("API Key: ")