fixes related to cube.download refactoring

mlcommons · Feb 29, 2024 · b04ee0c · b04ee0c
1 parent 23301fd
commit b04ee0c
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 1 deletion.
diff --git a/cli/medperf/commands/aggregator/run.py b/cli/medperf/commands/aggregator/run.py
@@ -80,6 +80,7 @@ def prepare_cube(self):
     def __get_cube(self, uid: int, name: str) -> Cube:
         self.ui.text = f"Retrieving {name} cube"
         cube = Cube.get(uid)
+        cube.download_run_files()
         self.ui.print(f"> {name} cube download complete")
         return cube
 

diff --git a/cli/medperf/commands/training/run.py b/cli/medperf/commands/training/run.py
@@ -83,6 +83,7 @@ def prepare_cube(self):
     def __get_cube(self, uid: int, name: str) -> Cube:
         self.ui.text = f"Retrieving {name} cube"
         cube = Cube.get(uid)
+        cube.download_run_files()
         self.ui.print(f"> {name} cube download complete")
         return cube
 

diff --git a/cli/medperf/commands/training/submit.py b/cli/medperf/commands/training/submit.py
@@ -41,7 +41,9 @@ def __init__(self, training_exp_info: dict):
 
     def get_mlcube(self):
         mlcube_id = self.training_exp.fl_mlcube
-        Cube.get(mlcube_id)
+        cube = Cube.get(mlcube_id)
+        # TODO: do we want to download run files?
+        cube.download_run_files()
 
     def submit(self):
         updated_body = self.training_exp.upload()