update webui

Author: hasan7n

cli/medperf/commands/asset/submit.py

@@ -1,6 +1,7 @@
 import os
 import logging
 
+from medperf.commands.model.submit_util import ContinueAssetSubmission
 from medperf.comms.entity_resources import resources
 import medperf.config as config
 from medperf.entities.asset import Asset
@@ -28,6 +29,9 @@ def run(
             asset_path (str): Local path to the asset file.
             asset_url (str): URL to download the asset from.
             operational (bool): Whether to submit as OPERATIONAL state.
+
+        Returns:
+            str: The ID of submitted asset.
         """
         ui = config.ui
         submission = cls(name, asset_path, asset_url, operational)
@@ -38,10 +42,15 @@ def run(
             submission.prepare_asset()
             submission.create_asset_object()
             submission.validate_asset_file()
-            ui.text = "Submitting asset to MedPerf"
-            updated_body = submission.upload()
-            submission.to_permanent_path(updated_body)
-            submission.write(updated_body)
+
+        if submission.asset.is_model():
+            return ContinueAssetSubmission.run(submission)
+
+        with ui.interactive():
+            ui.text = "Submitting Asset to MedPerf"
+            updated_asset_dict = submission.upload()
+            submission.to_permanent_path(updated_asset_dict)
+            submission.write(updated_asset_dict)
 
         ui.print(f"Asset registered with UID: {submission.asset.id}")
         return submission.asset.id
@@ -103,6 +112,7 @@ def validate_asset_file(self):
         self.asset.prepare_asset_files()
 
     def upload(self):
+        # Note: this is not used for now
         updated_body = self.asset.upload()
         return updated_body
 

cli/medperf/commands/mlcube/submit.py

@@ -4,6 +4,7 @@
 from medperf.entities.cube import Cube
 from medperf.exceptions import InvalidArgumentError
 from medperf.utils import remove_path, sanitize_path, store_decryption_key
+from medperf.commands.model.submit_util import ContinueCubeSubmission
 import logging
 
 import yaml
@@ -22,19 +23,26 @@ def run(
 
         Args:
             submit_info (dict): Dictionary containing the cube information.
+
+        Returns:
+            str: The ID of the submitted cube
         """
         ui = config.ui
         submission = cls(
             submit_info, container_config, parameters_config, decryption_key
         )
         submission.read_config_files()
         submission.create_cube_object()
-        submission.raise_if_model()
 
         with ui.interactive():
             ui.text = "Validating Container can be downloaded"
             submission.validate()
             submission.download_run_files()
+
+        if submission.cube.is_model():
+            return ContinueCubeSubmission.run(submission)
+
+        with ui.interactive():
             ui.text = "Submitting Container to MedPerf"
             updated_cube_dict = submission.upload()
             submission.to_permanent_path(updated_cube_dict)
@@ -104,13 +112,6 @@ def create_cube_object(self):
         config.tmp_paths.append(self.cube.path)
         os.makedirs(self.cube.path, exist_ok=True)
 
-    def raise_if_model(self):
-        if self.cube.is_model():
-            raise InvalidArgumentError(
-                "The provided container config file seems to be for a model, not a cube. "
-                "Please use the 'model submit' command for model submissions."
-            )
-
     def validate(self):
         if self.cube.is_encrypted() and not self.decryption_key:
             raise InvalidArgumentError(

cli/medperf/commands/model/submit.py

@@ -1,12 +1,5 @@
-import os
-
-from git import Optional
-
-import medperf.config as config
-from medperf.entities.model import Model
+from typing import Optional
 from medperf.exceptions import InvalidArgumentError
-from medperf.utils import remove_path
-from medperf.enums import ModelType
 from medperf.commands.mlcube.submit import SubmitCube
 from medperf.commands.asset.submit import SubmitAsset
 
@@ -42,7 +35,6 @@ def run(
             asset_url (Optional[str]): URL to download the asset from (for ASSET type).
             operational (bool): Whether to submit as OPERATIONAL state.
         """
-        ui = config.ui
         submission = cls(
             name,
             operational,
@@ -56,24 +48,9 @@ def run(
             asset_url,
         )
         submission.validate_inputs()
-
-        with ui.interactive():
-            ui.text = "Submitting model to MedPerf"
-            if submission.model_type == ModelType.CONTAINER.value:
-                submission.prepare_container_submission()
-                submission.create_model_object()
-                updated_model_body, updated_container_body = submission.upload()
-                submission.finalize_container_submission(updated_container_body)
-            else:
-                submission.prepare_asset_submission()
-                submission.create_model_object()
-                updated_model_body, updated_asset_body = submission.upload()
-                submission.finalize_asset_submission(updated_asset_body)
-            submission.to_permanent_path(updated_model_body)
-            submission.write(updated_model_body)
-
-        ui.print(f"Model registered with UID: {submission.model.id}")
-        return submission.model.id
+        if submission.is_container:
+            return submission.run_container_submission()
+        return submission.run_asset_submission()
 
     def __init__(
         self,
@@ -98,9 +75,7 @@ def __init__(
         self.asset_path = asset_path
         self.asset_url = asset_url
         self.operational = operational
-        self.model = None
-        self.model_type = None
-        self.subentity_submission = None
+        self.is_container: bool = None
 
     def validate_inputs(self):
         container_provided = (
@@ -124,90 +99,24 @@ def validate_inputs(self):
                 raise InvalidArgumentError(
                     "Container config file is required for container-backed model."
                 )
-        self.model_type = (
-            ModelType.CONTAINER.value if container_provided else ModelType.ASSET.value
-        )
+        self.is_container = container_provided
 
-    def prepare_container_submission(self):
+    def run_container_submission(self):
         submit_info = {
             "name": self.name,
             "image_hash": self.image_hash,
             "additional_files_tarball_url": self.additional_file,
             "additional_files_tarball_hash": self.additional_hash,
             "state": "OPERATION" if self.operational else "DEVELOPMENT",
         }
-        ui = config.ui
-        submission = SubmitCube(
+        return SubmitCube.run(
             submit_info,
             self.container_config_file,
             self.parameters_config_file,
             self.decryption_key,
         )
-        submission.read_config_files()
-        submission.create_cube_object()
 
-        ui.text = "Validating Container can be downloaded"
-        submission.validate()
-        submission.download_run_files()
-        self.subentity_submission = submission
-
-    def prepare_asset_submission(self):
-        ui = config.ui
-        submission = SubmitAsset(
+    def run_asset_submission(self):
+        return SubmitAsset.run(
             self.name, self.asset_path, self.asset_url, self.operational
         )
-        submission.validate_inputs()
-
-        ui.text = "Preparing asset"
-        submission.prepare_asset()
-        submission.create_asset_object()
-        submission.validate_asset_file()
-        self.subentity_submission = submission
-
-    def create_model_object(self):
-        state = "OPERATION" if self.operational else "DEVELOPMENT"
-        container = None
-        asset = None
-        if self.model_type == ModelType.CONTAINER.value:
-            container = self.subentity_submission.cube
-        else:
-            asset = self.subentity_submission.asset
-
-        self.model = Model(
-            name=self.name,
-            type=self.model_type,
-            container=container,
-            asset=asset,
-            state=state,
-        )
-        config.tmp_paths.append(self.model.path)
-        os.makedirs(self.model.path, exist_ok=True)
-
-    def upload(self):
-        updated_body = self.model.upload()
-        subentity_body = (
-            updated_body["container"]
-            if self.model_type == ModelType.CONTAINER.value
-            else updated_body["asset"]
-        )
-        return updated_body, subentity_body
-
-    def finalize_container_submission(self, updated_body: dict):
-        self.subentity_submission.to_permanent_path(updated_body)
-        self.subentity_submission.write(updated_body)
-        self.subentity_submission.store_decryption_key()
-
-    def finalize_asset_submission(self, updated_body: dict):
-        self.subentity_submission.to_permanent_path(updated_body)
-        self.subentity_submission.write(updated_body)
-
-    def to_permanent_path(self, model_dict):
-        old_path = self.model.path
-        updated_model = Model(**model_dict)
-        new_path = updated_model.path
-        remove_path(new_path)
-        os.rename(old_path, new_path)
-
-    def write(self, updated_body):
-        self.model = Model(**updated_body)
-        self.model.write()

cli/medperf/commands/model/submit_util.py

@@ -0,0 +1,112 @@
+import os
+from typing import TYPE_CHECKING
+
+import medperf.config as config
+from medperf.entities.model import Model
+from medperf.utils import remove_path
+from medperf.enums import ModelType
+
+if TYPE_CHECKING:
+    from medperf.commands.mlcube.submit import SubmitCube
+    from medperf.commands.asset.submit import SubmitAsset
+
+
+class ContinueCubeSubmission:
+    @classmethod
+    def run(cls, container_submission: SubmitCube):
+        """Finalizes the container submission.
+
+        Args:
+            container_submission (SubmitCube): The submission object for the cube.
+        """
+        submission = cls(container_submission)
+        submission.create_model_object()
+        updated_body, container_body = submission.upload()
+        submission.finalize_container_submission(container_body)
+        submission.to_permanent_path(updated_body)
+        submission.write(updated_body)
+        return submission.model.container.id
+
+    def __init__(self, container_submission: SubmitCube):
+        self.container_submission = container_submission
+        self.model = None
+
+    def create_model_object(self):
+        self.model = Model(
+            name=self.container_submission.cube.name,
+            type=ModelType.CONTAINER.value,
+            container=self.container_submission.cube,
+            state=self.container_submission.cube.state,
+        )
+        config.tmp_paths.append(self.model.path)
+        os.makedirs(self.model.path, exist_ok=True)
+
+    def upload(self):
+        updated_body = self.model.upload()
+        return updated_body, updated_body["container"]
+
+    def finalize_container_submission(self, updated_body: dict):
+        self.container_submission.to_permanent_path(updated_body)
+        self.container_submission.write(updated_body)
+        self.container_submission.store_decryption_key()
+
+    def to_permanent_path(self, model_dict):
+        old_path = self.model.path
+        updated_model = Model(**model_dict)
+        new_path = updated_model.path
+        remove_path(new_path)
+        os.rename(old_path, new_path)
+
+    def write(self, updated_body):
+        self.model = Model(**updated_body)
+        self.model.write()
+
+
+class ContinueAssetSubmission:
+    @classmethod
+    def run(cls, asset_submission: SubmitAsset):
+        """Finalizes the asset submission.
+
+        Args:
+            asset_submission (SubmitAsset): The submission object for the asset.
+        """
+        submission = cls(asset_submission)
+        submission.create_model_object()
+        updated_body, asset_body = submission.upload()
+        submission.finalize_asset_submission(asset_body)
+        submission.to_permanent_path(updated_body)
+        submission.write(updated_body)
+        return submission.model.asset.id
+
+    def __init__(self, asset_submission: SubmitAsset):
+        self.asset_submission = asset_submission
+        self.model = None
+
+    def create_model_object(self):
+        self.model = Model(
+            name=self.asset_submission.asset.name,
+            type=ModelType.ASSET.value,
+            asset=self.asset_submission.asset,
+            state=self.asset_submission.asset.state,
+        )
+        config.tmp_paths.append(self.model.path)
+        os.makedirs(self.model.path, exist_ok=True)
+
+    def upload(self):
+        updated_body = self.model.upload()
+        return updated_body, updated_body["asset"]
+
+    def finalize_asset_submission(self, updated_body: dict):
+        self.asset_submission.to_permanent_path(updated_body)
+        self.asset_submission.write(updated_body)
+
+    def to_permanent_path(self, model_dict):
+        old_path = self.model.path
+        updated_model = Model(**model_dict)
+        new_path = updated_model.path
+        remove_path(new_path)
+        os.rename(old_path, new_path)
+
+    def write(self, updated_body):
+        self.model = Model(**updated_body)
+        self.model.write()

cli/medperf/comms/rest.py

@@ -360,6 +360,19 @@ def get_container_model(self, container_id: int) -> dict:
         error_msg = f"Could not retrieve model for container {container_id}"
         return self.__get(url, error_msg)
 
+    def get_asset_model(self, asset_id: int) -> dict:
+        """Retrieves the model for the given asset ID
+
+        Args:
+            asset_id (int): Asset ID
+
+        Returns:
+            dict: Model metadata
+        """
+        url = f"{self.server_url}/assets/{asset_id}/model/"
+        error_msg = f"Could not retrieve model for asset {asset_id}"
+        return self.__get(url, error_msg)
+
     # get list
     def get_benchmarks(self, filters=dict()) -> List[dict]:
         """Retrieves all benchmarks in the platform.

cli/medperf/entities/asset.py

@@ -52,6 +52,12 @@ def __init__(self, *args, **kwargs):
     def local_id(self):
         return self.name
 
+    def is_local(self) -> bool:
+        return self.asset_url == "local"
+
+    def is_model(self) -> bool:
+        return True
+
     @staticmethod
     def remote_prefilter(filters: dict) -> callable:
         comms_fn = config.comms.get_assets