hotfix: better ui for mlcube retrieval while preparation (#568)

mlcommons · Mar 28, 2024 · c944f89 · c944f89
1 parent a9fb50d
commit c944f89
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/cli/medperf/commands/dataset/prepare.py b/cli/medperf/commands/dataset/prepare.py
@@ -80,7 +80,8 @@ def run(cls, dataset_id: int, approve_sending_reports: bool = False):
         preparation = cls(dataset_id, approve_sending_reports)
         preparation.get_dataset()
         preparation.validate()
-        preparation.get_prep_cube()
+        with preparation.ui.interactive():
+            preparation.get_prep_cube()
         preparation.setup_parameters()
 
         # TODO: make these more readable
@@ -135,7 +136,10 @@ def validate(self):
             raise InvalidArgumentError("This dataset is in operation mode")
 
     def get_prep_cube(self):
-        self.ui.text = "Retrieving data preparation cube"
+        self.ui.text = (
+            "Retrieving and setting up data preparation MLCube. "
+            "This may take some time."
+        )
         self.cube = Cube.get(self.dataset.data_preparation_mlcube)
         self.cube.download_run_files()
         self.ui.print("> Preparation cube download complete")