Prep workflow

cli/medperf/commands/dataset/prepare.py

@@ -96,11 +96,10 @@ def run(
         if preparation.should_run_prepare():
             with preparation.ui.interactive():
                 preparation.run_prepare()
-
         with preparation.ui.interactive():
-            preparation.run_sanity_check()
             preparation.run_statistics()
 
+        preparation.check_statistics()
         preparation.mark_dataset_as_ready()
 
         return preparation.dataset.id
@@ -211,25 +210,26 @@ def run_prepare(self):
         self.ui.print("> Container execution complete")
         report_sender.stop("finished")
 
-    def run_sanity_check(self):
-        sanity_check_timeout = config.sanity_check_timeout
+    def run_statistics(self):
+        # The statistics task runs sanity_check then Statistics in the workflow.
+        statistics_timeout = config.statistics_timeout
         out_datapath = self.out_datapath
         out_labelspath = self.out_labelspath
 
-        # Specify parameters for the tasks
-        sanity_check_mounts = {
+        statistics_mounts = {
             "data_path": out_datapath,
             "labels_path": out_labelspath,
+            "output_path": self.out_statistics_path,
         }
         if self.metadata_specified:
-            sanity_check_mounts["metadata_path"] = self.metadata_path
+            statistics_mounts["metadata_path"] = self.metadata_path
 
-        self.ui.text = "Running sanity check..."
+        self.ui.text = "Running sanity checks and statistics..."
         try:
             self.cube.run(
-                task="sanity_check",
-                timeout=sanity_check_timeout,
-                mounts=sanity_check_mounts,
+                task="statistics",
+                timeout=statistics_timeout,
+                mounts=statistics_mounts,
             )
         except ExecutionError:
             self.dataset.unmark_as_ready()
@@ -244,35 +244,14 @@ def run_sanity_check(self):
                 self.ui.print_warning(msg)
                 raise CleanExit(medperf_status_code=1)
 
-            msg = "The sanity check process failed"
+            msg = "The sanity check and statistics process failed"
             raise ExecutionError(msg)
-        self.ui.print("> Sanity checks complete")
-
-    def run_statistics(self):
-        statistics_timeout = config.statistics_timeout
-        out_datapath = self.out_datapath
-        out_labelspath = self.out_labelspath
-
-        statistics_mounts = {
-            "data_path": out_datapath,
-            "labels_path": out_labelspath,
-            "output_path": self.out_statistics_path,
-        }
-
-        if self.metadata_specified:
-            statistics_mounts["metadata_path"] = self.metadata_path
-
-        self.ui.text = "Generating statistics..."
+        self.ui.print("> Sanity checks and statistics complete")
 
-        try:
-            self.cube.run(
-                task="statistics",
-                timeout=statistics_timeout,
-                mounts=statistics_mounts,
-            )
-        except ExecutionError as e:
+    def check_statistics(self):
+        if not os.path.exists(self.out_statistics_path):
             self.dataset.unmark_as_ready()
-            raise e
+            raise ExecutionError("Statistics file was not created.")
 
         with open(self.out_statistics_path) as f:
             stats = yaml.safe_load(f)
@@ -281,8 +260,6 @@ def run_statistics(self):
             self.dataset.unmark_as_ready()
             raise ExecutionError("Statistics file is empty.")
 
-        self.ui.print("> Statistics complete")
-
     def mark_dataset_as_ready(self):
         self.dataset.mark_as_ready()
 

cli/medperf/config.py

@@ -283,7 +283,6 @@
 shm_size = None
 platform = "docker"
 prepare_timeout = None
-sanity_check_timeout = None
 statistics_timeout = None
 infer_timeout = None
 evaluate_timeout = None
@@ -313,7 +312,6 @@
 inline_parameters = [
     "loglevel",
     "prepare_timeout",
-    "sanity_check_timeout",
     "statistics_timeout",
     "infer_timeout",
     "evaluate_timeout",
@@ -326,7 +324,6 @@
 configurable_parameters = [
     "loglevel",
     "prepare_timeout",
-    "sanity_check_timeout",
     "statistics_timeout",
     "infer_timeout",
     "evaluate_timeout",

cli/medperf/containers/parsers/simple_container.py

@@ -2,6 +2,12 @@
 from medperf.exceptions import InvalidContainerSpec
 from medperf.enums import ContainerTypes
 import logging
+import shlex
+
+# Combined validation+statistics MedPerf task. It must start the in-container
+# workflow at the sanity_check step (which then runs Statistics).
+STATISTICS_TASK = "statistics"
+SANITY_CHECK_START_ARG = "--start=sanity_check"
 
 
 class SimpleContainerParser(Parser):
@@ -34,6 +40,63 @@ def check_schema(self) -> str:
                         "Mount type should be either a file or a directory."
                     )
 
+        self._check_data_preparator_constraints()
+
+    def _check_data_preparator_constraints(self):
+        tasks = self.container_config.get("tasks", {})
+        if "prepare" not in tasks:
+            return
+
+        if STATISTICS_TASK not in tasks:
+            raise InvalidContainerSpec(
+                "Data preparator containers must define 'statistics' "
+                "alongside 'prepare'."
+            )
+
+        if "sanity_check" in tasks:
+            raise InvalidContainerSpec(
+                "Data preparator containers must run sanity checks as part of "
+                "'statistics' (via --start=sanity_check), not as a separate "
+                "'sanity_check' task."
+            )
+
+        if "check_no_prepare" in tasks:
+            raise InvalidContainerSpec(
+                "Data preparator containers must use the 'statistics' task "
+                "(not 'check_no_prepare') for sanity checks and statistics."
+            )
+
+        stats_task = tasks[STATISTICS_TASK]
+        command = stats_task.get("run_args", {}).get("command", "")
+        if SANITY_CHECK_START_ARG not in self._tokenize_command(command):
+            raise InvalidContainerSpec(
+                "Data preparator 'statistics' task command must include "
+                f"'{SANITY_CHECK_START_ARG}'."
+            )
+
+        stats_outputs = stats_task.get("output_volumes", {})
+        if "output_path" not in stats_outputs:
+            raise InvalidContainerSpec(
+                "Data preparator 'statistics' task must define an 'output_path' "
+                "output volume for generated statistics."
+            )
+
+    @staticmethod
+    def _tokenize_command(command) -> list:
+        """Normalize a task command into tokens.
+
+        Container commands may be given as a shell string or a pre-split list
+        (both are accepted by the runners), so mirror that here instead of
+        assuming a string.
+        """
+        if isinstance(command, str):
+            return shlex.split(command)
+        if isinstance(command, (list, tuple)):
+            return [str(token) for token in command]
+        raise InvalidContainerSpec(
+            "Data preparator 'statistics' task 'command' must be a string " "or a list."
+        )
+
     def check_task_schema(self, task: str) -> str:
         tasks = self.container_config["tasks"]
         logging.debug(f"Available tasks: {tasks}")

cli/medperf/decorators.py

@@ -100,11 +100,6 @@ def wrapper(
             "--prepare_timeout",
             help="Maximum time in seconds before interrupting prepare task",
         ),
-        sanity_check_timeout: int = typer.Option(
-            config.sanity_check_timeout,
-            "--sanity_check_timeout",
-            help="Maximum time in seconds before interrupting sanity_check task",
-        ),
         statistics_timeout: int = typer.Option(
             config.statistics_timeout,
             "--statistics_timeout",
@@ -186,11 +181,6 @@ def wrapper(
             "--prepare_timeout",
             help="Maximum time in seconds before interrupting prepare task",
         ),
-        sanity_check_timeout: int = typer.Option(
-            config.sanity_check_timeout,
-            "--sanity_check_timeout",
-            help="Maximum time in seconds before interrupting sanity_check task",
-        ),
         statistics_timeout: int = typer.Option(
             config.statistics_timeout,
             "--statistics_timeout",

cli/medperf/storage/__init__.py

@@ -117,6 +117,22 @@ def __apply_cc_migrations(config_p: ConfigManager):
             config_p.storage[folder] = config_p.storage["benchmarks_folder"]
 
 
+def __apply_preparation_migrations(config_p: ConfigManager):
+    # sanity_check was folded into the statistics task, so a single
+    # statistics_timeout now covers both.
+    for profile in config_p.profiles.values():
+        if "sanity_check_timeout" not in profile:
+            continue
+        statistics_timeout = profile["statistics_timeout"]
+        sanity_check_timeout = profile["sanity_check_timeout"]
+        if statistics_timeout is not None and sanity_check_timeout is not None:
+            profile["statistics_timeout"] += profile["sanity_check_timeout"]
+        elif statistics_timeout is None and sanity_check_timeout is not None:
+            profile["statistics_timeout"] = sanity_check_timeout
+
+        del profile["sanity_check_timeout"]
+
+
 def apply_configuration_migrations():
     if not os.path.exists(config.config_path):
         return
@@ -128,5 +144,6 @@ def apply_configuration_migrations():
     __apply_results_to_executions_migrations(config_p)
     __apply_trusted_ca_migrations(config_p)
     __apply_cc_migrations(config_p)
+    __apply_preparation_migrations(config_p)
 
     write_config(config_p)

cli/medperf/templates/data_preparator_container/{{cookiecutter.project_slug}}/container_config.yaml

@@ -3,7 +3,7 @@ image: {{ cookiecutter.image_name }}
 tasks:
   prepare:
     run_args:
-      command: python3 /project/prepare_data.py
+      command: python3 /project/main.py
     input_volumes:
       data_path:
         mount_path: /mlcommons/volumes/raw_data
@@ -21,22 +21,9 @@ tasks:
       output_labels_path:
         mount_path: /mlcommons/volumes/labels
         type: directory
-  sanity_check:
-    run_args:
-      command: python3 /project/check.py
-    input_volumes:
-      data_path:
-        mount_path: /mlcommons/volumes/data
-        type: directory
-      labels_path:
-        mount_path: /mlcommons/volumes/labels
-        type: directory
-      parameters_file:
-        mount_path: /mlcommons/volumes/parameters/parameters_file.yaml
-        type: file
   statistics:
     run_args:
-      command: python3 /project/statistics.py
+      command: python3 /project/main.py --start=sanity_check
     input_volumes:
       data_path:
         mount_path: /mlcommons/volumes/data

cli/medperf/tests/commands/dataset/test_prepare.py

@@ -187,7 +187,7 @@ def test_prepare_checks_report_and_metadata_path(
 @pytest.mark.parametrize(
     "report_specified,exception", [[False, ExecutionError], [True, CleanExit]]
 )
-def test_sanity_checks_unmarks_the_dataset_as_ready_on_failure(
+def test_statistics_unmarks_the_dataset_as_ready_on_failure(
     mocker, data_preparation, cube, report_specified, exception
 ):
     # Arrange
@@ -200,63 +200,56 @@ def _failure_run(*args, **kwargs):
 
     # Act & assert
     with pytest.raises(exception):
-        data_preparation.run_sanity_check()
-
-    # Assert
-    unmark_spy.assert_called_once()
-
-
-def test_statistics_unmarks_the_dataset_as_ready_on_failure(
-    mocker, data_preparation, cube
-):
-    # Arrange
-    def _failure_run(*args, **kwargs):
-        raise ExecutionError()
-
-    mocker.patch.object(cube, "run", side_effect=_failure_run)
-    unmark_spy = mocker.patch.object(data_preparation.dataset, "unmark_as_ready")
-
-    # Act & assert
-    with pytest.raises(ExecutionError):
         data_preparation.run_statistics()
 
     # Assert
     unmark_spy.assert_called_once()
 
 
 @pytest.mark.parametrize("metadata_specified", [False, True])
-def test_statistics_checks_metadata_path(
-    mocker, data_preparation, metadata_specified, cube, fs
+def test_statistics_generates_statistics_and_checks_paths(
+    mocker, data_preparation, metadata_specified, cube
 ):
     # Arrange
     spy = mocker.patch.object(cube, "run")
     data_preparation.metadata_specified = metadata_specified
     data_preparation.out_statistics_path = "test.yaml"
-    fs.create_file(data_preparation.out_statistics_path, contents="")
-    mocker.patch("yaml.safe_load", return_value={})
 
     # Act
     data_preparation.run_statistics()
 
     # Assert
+    assert spy.call_args.kwargs["task"] == "statistics"
+    assert spy.call_args.kwargs["mounts"]["output_path"] == "test.yaml"
     if metadata_specified:
         assert "metadata_path" in spy.call_args.kwargs["mounts"].keys()
     else:
         assert "metadata_path" not in spy.call_args.kwargs["mounts"].keys()
 
 
-def test_preparation_fails_if_statistics_is_none(mocker, data_preparation, cube, fs):
-
+def test_preparation_fails_if_statistics_is_none(mocker, data_preparation, fs):
     # Arrange
     unmark_spy = mocker.patch.object(data_preparation.dataset, "unmark_as_ready")
-    mocker.patch.object(cube, "run")
     data_preparation.out_statistics_path = "test.yaml"
     fs.create_file(data_preparation.out_statistics_path, contents="")
     mocker.patch("yaml.safe_load", return_value=None)
 
     # Act
     with pytest.raises(ExecutionError):
-        data_preparation.run_statistics()
+        data_preparation.check_statistics()
+
+    # Assert
+    unmark_spy.assert_called_once()
+
+
+def test_preparation_fails_if_statistics_were_not_created(mocker, data_preparation):
+    # Arrange
+    unmark_spy = mocker.patch.object(data_preparation.dataset, "unmark_as_ready")
+    data_preparation.out_statistics_path = "missing.yaml"
+
+    # Act
+    with pytest.raises(ExecutionError, match="not created"):
+        data_preparation.check_statistics()
 
     # Assert
     unmark_spy.assert_called_once()

cli/medperf/tests/containers/__init__.py

@@ -0,0 +1 @@
+