Merge pull request #44 from neptune-ai/fixes

Fixes

Author: kshitij12345

CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### Changes
 - Updated MLflow exporter to work with the current neptune.ai API ([#40](https://github.com/neptune-ai/neptune-mlflow/pull/40))
+- Misc Fixes: ([#44](https://github.com/neptune-ai/neptune-mlflow/pull/44))
+    - Stop capturing hardware metrics
+    - Map MLflow `run_name` to `sys/name` in Neptune `Run`
 
 ### Fixes
 - Fix to support Python 3.7 and 3.8 ([#43](https://github.com/neptune-ai/neptune-mlflow/pull/43))

src/neptune_mlflow_exporter/impl/components/exporter.py

@@ -53,6 +53,7 @@ def export_run_info(neptune_run: NeptuneRun, mlflow_run: MlflowRun) -> None:
         info["end_time"] = datetime.fromtimestamp(mlflow_run.info.end_time / 1e3)
 
         neptune_run["run_info"] = info
+        neptune_run["sys/name"] = info["run_name"]
 
     def export_run_data(self, neptune_run: NeptuneRun, mlflow_run: MlflowRun) -> None:
         data_dict = mlflow_run.data.to_dictionary()

src/neptune_mlflow_exporter/impl/components/fetcher.py

@@ -16,10 +16,10 @@
 
 __all__ = ["Fetcher"]
 
-from typing import MutableMapping
 from dataclasses import dataclass
 from typing import (
     List,
+    MutableMapping,
     Set,
 )
 

src/neptune_mlflow_exporter/impl/orchestrator.py

@@ -48,7 +48,10 @@ def run(self) -> None:
             click.echo(f"Loading mlflow_run '{mlflow_run.info.run_name}'")
 
             with NeptuneRun(
-                project=self.config.project_name, api_token=self.config.api_token, custom_run_id=mlflow_run.info.run_id
+                project=self.config.project_name,
+                api_token=self.config.api_token,
+                custom_run_id=mlflow_run.info.run_id,
+                capture_hardware_metrics=False,
             ) as neptune_run:
                 try:
                     experiment = fetched_data.mlflow_experiments[mlflow_run.info.experiment_id]

tests/e2e/test_neptune_mlflow.py

@@ -58,7 +58,9 @@ def test_e2e(dataset, model, neptune_exporter_e2e):
     neptune_exporter_e2e.run()
 
     # check logged project metadata
-    experiment = mlflow.get_experiment_by_name("E2E neptune experiment")
+    experiment = mlflow.get_experiment_by_name(MLFLOW_EXPERIMENT_NAME)
+    mlflow_runs = mlflow.search_runs(experiment_names=[MLFLOW_EXPERIMENT_NAME])
+    run_name = mlflow_runs["tags.mlflow.runName"][0]
 
     # check logged run metadata
     neptune_run = neptune.init_run(custom_run_id=run_id)
@@ -72,6 +74,10 @@ def test_e2e(dataset, model, neptune_exporter_e2e):
     # run info
     assert neptune_run["run_info/lifecycle_stage"].fetch() == "active"
     assert neptune_run.exists("run_info/status")
+    assert neptune_run["run_info/run_name"].fetch() == run_name
+    # Currently, client overwrites `sys/name` when run is re-opened
+    # with `custom_run_id`.
+    # assert neptune_run["sys/name"].fetch() == run_name
 
     # run data
     assert set(MLFLOW_RUN_TAGS.items()).issubset(set(neptune_run["run_data/tags"].fetch().items()))