Neptune update (#22)

deepflare · Aug 8, 2023 · b2ceb50 · b2ceb50
1 parent 711bacc
commit b2ceb50
Show file tree

Hide file tree

Showing 3 changed files with 1,661 additions and 664 deletions.
diff --git a/exphub/download/neptune_downloader.py b/exphub/download/neptune_downloader.py
@@ -3,9 +3,9 @@
 import pandas as pd
 from typing import Optional, Union, List
 import os
-import neptune.new as neptune
+from neptune import Project, Run
+
 from exphub.download.experiment import Experiment
-from exphub.utils.noise import Suppressor
 from exphub.utils.paths import shorten_paths
 
 
@@ -44,7 +44,8 @@ def __init__(self, project_name: str, api_token: Optional[str] = None):
         else:
             self.api_token = api_token
             os.environ[NeptuneDownloader.NEPTUNE_API_TOKEN] = api_token
-        self.project = neptune.init_project(project=self.project_name, mode="read-only", api_token=self.api_token)
+
+        self.project = Project(project=self.project_name, api_token=self.api_token, mode="read-only")
 
     def download(self,
                  id: Optional[Union[str, List[str]]] = None,
@@ -100,13 +101,11 @@ def _download_series(self,
         if all([id is None, state is None, owner is None, tag is None]):
             raise ValueError('At least one of id, state, owner, or tag must be provided.')
 
-        ids = self.project.fetch_runs_table(
-            owner=owner, id=id, state=state, tag=tag, columns='sys/id').to_pandas()['sys/id'].values
+        ids = self.project.fetch_runs_table(owner=owner, id=id, state=state, tag=tag, columns='sys/id').to_pandas()['sys/id'].values
 
         # Run initialization
         runs = [
-            Suppressor.exec_no_stdout(
-                neptune.init_run, project=self.project_name, with_id=run_id, mode="read-only", api_token=self.api_token)
+            Run(run_id, project=self.project_name, api_token=self.api_token, mode="read-only")
             for run_id in ids
         ]
 
@@ -120,8 +119,8 @@ def _fetch_values(col_label):
             missing = 0
             for id, run in zip(ids, runs):
                 try:
-                    id2value[id] = Suppressor.exec_no_stdout(run[col_label].fetch_values, include_timestamp=False)
-                except neptune.exceptions.NeptuneException:
+                    id2value[id] = run[col_label].fetch_values(include_timestamp=False)
+                except KeyError:
                     print(f'[WARNING] Run {id} does not have a column named {col_label}')
                     missing += 1
             if missing == len(ids):