diff --git a/cli/medperf/commands/mlcube/mlcube.py b/cli/medperf/commands/mlcube/mlcube.py index 9256f35f2..2d2ccb9cf 100644 --- a/cli/medperf/commands/mlcube/mlcube.py +++ b/cli/medperf/commands/mlcube/mlcube.py @@ -87,6 +87,13 @@ def submit( help="Identifier to download the image file. See the description above", ), image_hash: str = typer.Option("", "--image-hash", help="hash of image file"), + stages_file: str = typer.Option( + "", + "--stages-file", + "-s", + help="Identifier to download the stages file. See the description above" + ), + stages_hash: str = typer.Option("", "--stages-hash", help="Hash of the stages file"), operational: bool = typer.Option( False, "--operational", @@ -99,6 +106,7 @@ def submit( - parameters_file\n - additional_file\n - image_file\n + - stages_file\n are expected to be given in the following format: where `source_prefix` instructs the client how to download the resource, and `resource_identifier` is the identifier used to download the asset. The following are supported:\n @@ -117,6 +125,8 @@ def submit( "image_tarball_hash": image_hash, "additional_files_tarball_url": additional_file, "additional_files_tarball_hash": additional_hash, + "stages_url": stages_file, + "stages_hash": stages_hash, "state": "OPERATION" if operational else "DEVELOPMENT", } SubmitCube.run(mlcube_info) diff --git a/cli/medperf/entities/cube.py b/cli/medperf/entities/cube.py index 5fe2f3801..edfb9b1f2 100644 --- a/cli/medperf/entities/cube.py +++ b/cli/medperf/entities/cube.py @@ -230,6 +230,11 @@ def download_config_files(self): except InvalidEntityError as e: raise InvalidEntityError(f"MLCube {self.name} parameters file: {e}") + try: + self.download_stages() + except InvalidEntityError as e: + raise InvalidEntityError(f"MLCube {self.name} stages file: {e}") + def download_run_files(self): try: self.download_additional()