Refactor web UI with Alpine.js and TailwindCSS, enhance task management

cli/cli_tests_training.sh

@@ -163,16 +163,6 @@ echo "AGG_UID=$AGG_UID" >> "$LAST_ENV_FILE"
 
 echo "\n"
 
-##########################################################
-echo "====================================="
-echo "Running aggregator association step"
-echo "====================================="
-print_eval medperf aggregator associate -a $AGG_UID -t $TRAINING_UID -y
-checkFailed "aggregator association step failed"
-##########################################################
-
-echo "\n"
-
 ##########################################################
 echo "====================================="
 echo "Activate modelowner profile"
@@ -185,10 +175,10 @@ echo "\n"
 
 ##########################################################
 echo "====================================="
-echo "Approve aggregator association"
+echo "Running set aggregator step"
 echo "====================================="
-print_eval medperf association approve -t $TRAINING_UID -a $AGG_UID
-checkFailed "agg association approval failed"
+print_eval medperf training set_aggregator -t $TRAINING_UID -a $AGG_UID -y
+checkFailed "Setting aggregator failed"
 ##########################################################
 
 echo "\n"

cli/medperf/_version.py

@@ -1 +1 @@
-__version__ = "0.2.0"
+__version__ = "0.3.0"

cli/medperf/asset_management/asset_management.py

@@ -9,7 +9,7 @@
 from medperf.asset_management.asset_storage_manager import AssetStorageManager
 from medperf.asset_management.asset_policy_manager import AssetPolicyManager
 from medperf.asset_management.cc_operator import OperatorManager
-from medperf.utils import tar, generate_tmp_path
+from medperf.utils import tar, generate_tmp_path, remove_path
 import secrets
 from medperf.exceptions import MedperfException
 from medperf import config as medperf_config
@@ -42,55 +42,52 @@ def setup_dataset_for_cc(dataset: Dataset):
     cc_policy = dataset.get_cc_policy()
     __verify_cloud_environment(cc_config)
 
-    # create dataset asset
+    # policy setup
+    medperf_config.ui.text = "Generating encryption key"
+    encryption_key = generate_encryption_key()
+    asset_policy_manager = AssetPolicyManager(cc_config)
+    asset_policy_manager.setup_policy(cc_policy, encryption_key)
+
+    # storage
     medperf_config.ui.text = "Compressing dataset"
     asset_path = generate_tmp_path()
     tar(asset_path, [dataset.data_path, dataset.labels_path])
-
-    __setup_asset_for_cc(cc_config, cc_policy, asset_path)
+    asset_storage_manager = AssetStorageManager(cc_config, asset_path, encryption_key)
+    asset_storage_manager.store_asset()
+    del encryption_key
+    remove_path(asset_path)
 
 
 def setup_model_for_cc(model: Model):
     if not model.is_cc_configured():
         return
     cc_config = model.get_cc_config()
     cc_policy = model.get_cc_policy()
-    if model.type != "ASSET":
+    if not model.is_asset():
         raise MedperfException(
             f"Model {model.id} is not a file-based asset and cannot be set up for confidential computing."
         )
     asset = model.asset_obj
-    # create model asset
     asset_path = asset.get_archive_path()
 
     __verify_cloud_environment(cc_config)
-    __setup_asset_for_cc(cc_config, cc_policy, asset_path, for_model=True)
-
-
-def __verify_cloud_environment(cc_config: dict):
-    AssetStorageManager(cc_config, None, None).setup()
-
 
-def __setup_asset_for_cc(
-    cc_config: dict,
-    cc_policy: dict,
-    asset_path: str,
-    for_model: bool = False,
-):
-    # create encryption key
+    # policy setup
+    medperf_config.ui.text = "Generating encryption key"
     encryption_key = generate_encryption_key()
-
-    asset_storage_manager = AssetStorageManager(cc_config, asset_path, encryption_key)
-    asset_policy_manager = AssetPolicyManager(cc_config, for_model=for_model)
+    asset_policy_manager = AssetPolicyManager(cc_config, for_model=True)
+    asset_policy_manager.setup_policy(cc_policy, encryption_key)
 
     # storage
+    asset_storage_manager = AssetStorageManager(cc_config, asset_path, encryption_key)
     asset_storage_manager.store_asset()
-
-    # policy setup
-    asset_policy_manager.setup_policy(cc_policy, encryption_key)
     del encryption_key
 
 
+def __verify_cloud_environment(cc_config: dict):
+    AssetStorageManager(cc_config, None, None).setup()
+
+
 def update_dataset_cc_policy(dataset: Dataset, permitted_workloads: list[CCWorkloadID]):
     if not dataset.is_cc_configured():
         raise MedperfException(
@@ -108,7 +105,7 @@ def update_model_cc_policy(model: Model, permitted_workloads: list[CCWorkloadID]
             f"Model {model.id} does not have a configuration for confidential computing."
         )
     cc_config = model.get_cc_config()
-    if model.type != "ASSET":
+    if not model.is_asset():
         raise MedperfException(
             f"Model {model.id} is not a file-based asset and cannot be set up for confidential computing."
         )

cli/medperf/asset_management/asset_storage_manager.py

@@ -6,10 +6,15 @@
     remove_path,
 )
 from medperf.encryption import SymmetricEncryption
-from medperf.asset_management.gcp_utils import GCPAssetConfig, upload_file_to_gcs
+from medperf.asset_management.gcp_utils import (
+    GCPAssetConfig,
+    upload_from_file_object_to_gcs,
+)
 from medperf.asset_management.asset_check import verify_asset_owner_setup
+from medperf.asset_management.utils import CustomWriter, get_file_size
 from medperf.exceptions import MedperfException
 from medperf import config as medperf_config
+from tqdm import tqdm
 
 
 class AssetStorageManager:
@@ -30,12 +35,25 @@ def __encrypt_asset(self):
         asset_hash = get_file_hash(tmp_encrypted_asset_path)
         return tmp_encrypted_asset_path, asset_hash
 
-    def __upload_encrypted_asset(self, tmp_encrypted_asset_path):
-        upload_file_to_gcs(
-            self.config,
-            tmp_encrypted_asset_path,
-            self.config.encrypted_asset_bucket_file,
-        )
+    def __upload_encrypted_asset(self, tmp_encrypted_asset_path: str):
+        with open(tmp_encrypted_asset_path, "rb") as in_file:
+            with tqdm.wrapattr(
+                in_file,
+                "read",
+                total=get_file_size(in_file),
+                miniters=1,
+                desc="Uploading encrypted dataset to the bucket",
+                unit="B",
+                unit_scale=True,
+                unit_divisor=1024,
+                file=CustomWriter(),
+            ) as file_obj:
+                upload_from_file_object_to_gcs(
+                    self.config,
+                    file_obj,
+                    self.config.encrypted_asset_bucket_file,
+                )
+        remove_path(tmp_encrypted_asset_path)
 
     def setup(self):
         medperf_config.ui.text = "Verifying Cloud Environment"

cli/medperf/asset_management/gcp_utils/__init__.py

@@ -2,6 +2,7 @@
 from .kms import set_kms_iam_policy, encrypt_with_kms_key
 from .storage import (
     upload_file_to_gcs,
+    upload_from_file_object_to_gcs,
     upload_string_to_gcs,
     download_file_from_gcs,
     download_string_from_gcs,
@@ -19,6 +20,7 @@
     "set_kms_iam_policy",
     "encrypt_with_kms_key",
     "upload_file_to_gcs",
+    "upload_from_file_object_to_gcs",
     "upload_string_to_gcs",
     "download_file_from_gcs",
     "download_string_from_gcs",

cli/medperf/asset_management/gcp_utils/storage.py

@@ -13,6 +13,16 @@ def upload_file_to_gcs(
     blob.upload_from_filename(local_file)
 
 
+def upload_from_file_object_to_gcs(
+    config: Union[GCPAssetConfig, GCPOperatorConfig], file: object, gcs_path: str
+):
+    """Upload file to Google Cloud Storage."""
+    client = storage.Client()
+    bucket = client.bucket(config.bucket)
+    blob = bucket.blob(gcs_path)
+    blob.upload_from_file(file)
+
+
 def upload_string_to_gcs(
     config: Union[GCPAssetConfig, GCPOperatorConfig], content: bytes, gcs_path: str
 ):

cli/medperf/asset_management/utils.py

@@ -0,0 +1,21 @@
+from medperf import config
+import os
+
+
+class CustomWriter:
+    """class to use with tqdm to print progress using config.ui"""
+
+    def write(self, msg):
+        config.ui.print(msg)
+
+    def flush(self):
+        pass
+
+
+def get_file_size(file_object) -> int:
+    """Get the size of a file in bytes."""
+    try:
+        total_bytes = os.fstat(file_object.fileno()).st_size
+    except (AttributeError, OSError):
+        total_bytes = None
+    return total_bytes

cli/medperf/commands/aggregator/aggregator.py

@@ -5,7 +5,6 @@
 import medperf.config as config
 from medperf.decorators import clean_except
 from medperf.commands.aggregator.submit import SubmitAggregator
-from medperf.commands.aggregator.associate import AssociateAggregator
 from medperf.commands.aggregator.run import StartAggregator
 
 from medperf.commands.list import EntityList
@@ -33,22 +32,6 @@ def submit(
     config.ui.print("✅ Done!")
 
 
-@app.command("associate")
-@clean_except
-def associate(
-    aggregator_id: int = typer.Option(
-        ..., "--aggregator_id", "-a", help="UID of benchmark to associate with"
-    ),
-    training_exp_id: int = typer.Option(
-        ..., "--training_exp_id", "-t", help="UID of benchmark to associate with"
-    ),
-    approval: bool = typer.Option(False, "-y", help="Skip approval step"),
-):
-    """Associates an aggregator with a training experiment."""
-    AssociateAggregator.run(aggregator_id, training_exp_id, approved=approval)
-    config.ui.print("✅ Done!")
-
-
 @app.command("start")
 @clean_except
 def run(

cli/medperf/commands/aggregator/run.py

@@ -24,14 +24,14 @@ def run(
             training_exp_id (int): Training experiment UID.
         """
         execution = cls(training_exp_id, publish_on, overwrite)
-        execution.prepare()
-        execution.validate()
-        execution.check_existing_outputs()
-        execution.prepare_aggregator()
-        execution.prepare_participants_list()
-        execution.prepare_plan()
-        execution.prepare_pki_assets()
         with config.ui.interactive():
+            execution.prepare()
+            execution.validate()
+            execution.check_existing_outputs()
+            execution.prepare_aggregator()
+            execution.prepare_participants_list()
+            execution.prepare_plan()
+            execution.prepare_pki_assets()
             execution.run_experiment()
 
     def __init__(self, training_exp_id, publish_on, overwrite) -> None:

cli/medperf/commands/aggregator/submit.py

@@ -23,6 +23,7 @@ def run(cls, name: str, address: str, port: int, aggregation_mlcube: int):
             updated_benchmark_body = submission.submit()
         ui.print("Uploaded")
         submission.write(updated_benchmark_body)
+        return submission.aggregator.id
 
     def __init__(self, name: str, address: str, port: int, aggregation_mlcube: int):
         self.ui = config.ui
@@ -41,5 +42,5 @@ def submit(self):
 
     def write(self, updated_body):
         remove_path(self.aggregator.path)
-        aggregator = Aggregator(**updated_body)
-        aggregator.write()
+        self.aggregator = Aggregator(**updated_body)
+        self.aggregator.write()

cli/medperf/commands/association/approval.py

@@ -10,15 +10,14 @@ def run(
         training_exp_uid: int = None,
         dataset_uid: int = None,
         model_uid: int = None,
-        aggregator_uid: int = None,
     ):
-        """Sets approval status for an association between a benchmark and a dataset or mlcube
+        """Sets approval status for an association between a benchmark and a dataset or mlcube,
+        or between a training experiment and a dataset.
 
         Args:
             benchmark_uid (int): Benchmark UID.
             approval_status (str): Desired approval status to set for the association.
-            comms (Comms): Instance of Comms interface.
-            ui (UI): Instance of UI interface.
+            training_exp_uid (int, optional): Training experiment UID. Defaults to None.
             dataset_uid (int, optional): Dataset UID. Defaults to None.
             mlcube_uid (int, optional): MLCube UID. Defaults to None.
         """
@@ -28,7 +27,6 @@ def run(
             training_exp_uid,
             dataset_uid,
             model_uid,
-            aggregator_uid,
             approval_status.value,
         )
         update = {"approval_status": approval_status.value}
@@ -42,12 +40,7 @@ def run(
                 comms.update_benchmark_model_association(
                     benchmark_uid, model_uid, update
                 )
-        if training_exp_uid:
-            if dataset_uid:
-                comms.update_training_dataset_association(
-                    training_exp_uid, dataset_uid, update
-                )
-            if aggregator_uid:
-                comms.update_training_aggregator_association(
-                    training_exp_uid, aggregator_uid, update
-                )
+        if training_exp_uid and dataset_uid:
+            comms.update_training_dataset_association(
+                training_exp_uid, dataset_uid, update
+            )