Confidential Computing misc updates

cli/medperf/asset_management/asset_check.py

@@ -0,0 +1,42 @@
+from medperf.asset_management.gcp_utils import checks, get_user_credentials
+
+
+def verify_asset_owner_setup(bucket_name, kms_key_resource, wip_resource):
+    base_creds = get_user_credentials()
+    result = checks.check_user_role_on_bucket(
+        "user",
+        base_creds,
+        bucket_name,
+        "roles/storage.admin",
+    )
+    if result:
+        return False, result
+
+    result = checks.check_user_role_on_kms_key(
+        base_creds,
+        kms_key_resource,
+        "roles/cloudkms.cryptoKeyEncrypter",
+    )
+
+    if result:
+        return False, result
+
+    result = checks.check_user_role_on_kms_key(
+        base_creds,
+        kms_key_resource,
+        "roles/cloudkms.admin",
+    )
+
+    if result:
+        return False, result
+
+    result = checks.check_user_role_on_wip(
+        base_creds,
+        wip_resource,
+        "roles/iam.workloadIdentityPoolAdmin",
+    )
+
+    if result:
+        return False, result
+
+    return True, ""

cli/medperf/asset_management/asset_management.py

@@ -1,4 +1,8 @@
-from medperf.asset_management.gcp_utils import CCWorkloadID
+from medperf.asset_management.gcp_utils import (
+    CCWorkloadID,
+    GCPAssetConfig,
+    GCPOperatorConfig,
+)
 from medperf.entities.dataset import Dataset
 from medperf.entities.model import Model
 from medperf.entities.user import User
@@ -7,134 +11,117 @@
 from medperf.asset_management.cc_operator import OperatorManager
 from medperf.utils import tar, generate_tmp_path
 import secrets
-import os
-from medperf import config
+from medperf.exceptions import MedperfException
+from medperf import config as medperf_config
 
 
-def generate_encryption_key(encryption_key_file: str):
-    with open(encryption_key_file, "wb") as f:
-        pass
-    os.chmod(encryption_key_file, 0o700)
-    with open(encryption_key_file, "ab") as f:
-        f.write(secrets.token_bytes(32))
+def generate_encryption_key():
+    return secrets.token_bytes(32)
+
+
+def validate_cc_config(cc_config: dict, asset_name_prefix: str):
+    if cc_config == {}:
+        return
+
+    cc_config["encrypted_asset_bucket_file"] = asset_name_prefix + ".enc"
+    cc_config["encrypted_key_bucket_file"] = asset_name_prefix + "_key.enc"
+
+    GCPAssetConfig(**cc_config)
+
+
+def validate_cc_operator_config(cc_config: dict):
+    if cc_config == {}:
+        return
+    GCPOperatorConfig(**cc_config)
 
 
 def setup_dataset_for_cc(dataset: Dataset):
+    if not dataset.is_cc_configured():
+        return
     cc_config = dataset.get_cc_config()
     cc_policy = dataset.get_cc_policy()
-    if not cc_config:
-        raise ValueError(
-            f"Dataset {dataset.id} does not have a configuration for confidential computing."
-        )
-    if cc_policy is None:
-        raise ValueError(
-            f"Dataset {dataset.id} does not have a policy for confidential computing."
-        )
+    __verify_cloud_environment(cc_config)
+
     # create dataset asset
+    medperf_config.ui.text = "Compressing dataset"
     asset_path = generate_tmp_path()
     tar(asset_path, [dataset.data_path, dataset.labels_path])
 
-    # create encryption key
-    encryption_key_folder = os.path.join(
-        config.cc_artifacts_dir, "dataset" + str(dataset.id)
-    )
-    os.makedirs(encryption_key_folder, exist_ok=True)
-    encryption_key_file = os.path.join(encryption_key_folder, "encryption_key.bin")
-    generate_encryption_key(encryption_key_file)
-
-    __setup_asset_for_cc(cc_config, cc_policy, asset_path, encryption_key_file)
+    __setup_asset_for_cc(cc_config, cc_policy, 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 not cc_config:
-        raise ValueError(
-            f"Model {model.id} does not have a configuration for confidential computing."
-        )
-    if cc_policy is None:
-        raise ValueError(
-            f"Model {model.id} does not have a policy for confidential computing."
-        )
     if model.type != "ASSET":
-        raise ValueError(
+        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()
 
-    # create encryption key
-    encryption_key_folder = os.path.join(
-        config.cc_artifacts_dir, "model" + str(model.id)
-    )
-    os.makedirs(encryption_key_folder, exist_ok=True)
-    encryption_key_file = os.path.join(encryption_key_folder, "encryption_key.bin")
-    generate_encryption_key(encryption_key_file)
+    __verify_cloud_environment(cc_config)
+    __setup_asset_for_cc(cc_config, cc_policy, asset_path, for_model=True)
 
-    __setup_asset_for_cc(cc_config, cc_policy, asset_path, encryption_key_file)
+
+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, encryption_key_file: str
+    cc_config: dict,
+    cc_policy: dict,
+    asset_path: str,
+    for_model: bool = False,
 ):
-    # asset storage setup
-    asset_storage_manager = AssetStorageManager(
-        cc_config, asset_path, encryption_key_file
-    )
-    asset_storage_manager.setup()
+    # create 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)
+
+    # storage
     asset_storage_manager.store_asset()
 
     # policy setup
-    asset_policy_manager = AssetPolicyManager(cc_config, encryption_key_file)
-    asset_policy_manager.setup()
-    asset_policy_manager.setup_policy(cc_policy)
+    asset_policy_manager.setup_policy(cc_policy, encryption_key)
+    del encryption_key
 
 
 def update_dataset_cc_policy(dataset: Dataset, permitted_workloads: list[CCWorkloadID]):
-    cc_config = dataset.get_cc_config()
-    if not cc_config:
-        raise ValueError(
+    if not dataset.is_cc_configured():
+        raise MedperfException(
             f"Dataset {dataset.id} does not have a configuration for confidential computing."
         )
 
-    encryption_key_folder = os.path.join(
-        config.cc_artifacts_dir, "dataset" + str(dataset.id)
-    )
-    encryption_key_file = os.path.join(encryption_key_folder, "encryption_key.bin")
-
-    asset_policy_manager = AssetPolicyManager(cc_config, encryption_key_file)
+    cc_config = dataset.get_cc_config()
+    asset_policy_manager = AssetPolicyManager(cc_config)
     asset_policy_manager.configure_policy(permitted_workloads)
 
 
 def update_model_cc_policy(model: Model, permitted_workloads: list[CCWorkloadID]):
-    cc_config = model.get_cc_config()
-    if not cc_config:
-        raise ValueError(
+    if not model.is_cc_configured():
+        raise MedperfException(
             f"Model {model.id} does not have a configuration for confidential computing."
         )
+    cc_config = model.get_cc_config()
     if model.type != "ASSET":
-        raise ValueError(
+        raise MedperfException(
             f"Model {model.id} is not a file-based asset and cannot be set up for confidential computing."
         )
 
-    encryption_key_folder = os.path.join(
-        config.cc_artifacts_dir, "model" + str(model.id)
-    )
-    encryption_key_file = os.path.join(encryption_key_folder, "encryption_key.bin")
-
-    asset_policy_manager = AssetPolicyManager(cc_config, encryption_key_file)
+    asset_policy_manager = AssetPolicyManager(cc_config, for_model=True)
     asset_policy_manager.configure_policy(permitted_workloads)
 
 
 def setup_operator(user: User):
-    cc_config = user.get_cc_config()
-    if not cc_config:
-        raise ValueError(
-            "User does not have a configuration for confidential computing."
-        )
+    if not user.is_cc_configured():
+        return
 
+    cc_config = user.get_cc_config()
     operator_manager = OperatorManager(cc_config)
     operator_manager.setup()
 
@@ -156,6 +143,10 @@ def run_workload(
         model_cc_config,
         result_collector_public_key,
     )
+
+
+def wait_for_workload(workload: CCWorkloadID, operator_cc_config: dict):
+    operator_manager = OperatorManager(operator_cc_config)
     operator_manager.wait_for_workload_completion(workload)
 
 
@@ -168,3 +159,8 @@ def download_results(
     operator_manager = OperatorManager(operator_cc_config)
 
     operator_manager.download_results(workload, private_key_bytes, results_path)
+
+
+def workload_results_exists(operator_cc_config: dict, workload: CCWorkloadID) -> bool:
+    operator_manager = OperatorManager(operator_cc_config)
+    return operator_manager.results_exist(workload)