Skip to content

Commit

Permalink
Force debug, better error handle
Browse files Browse the repository at this point in the history
  • Loading branch information
abuabraham-ttd committed Dec 19, 2024
1 parent 9ee0d14 commit a6650e0
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 9 deletions.
13 changes: 7 additions & 6 deletions scripts/aws/ec2.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import yaml

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, SecretNotFoundException
from confidential_compute import ConfidentialCompute, ConfidentialComputeConfig, SecretNotFoundException, ConfidentialComputeStartupException

class AWSConfidentialComputeConfig(ConfidentialComputeConfig):
enclave_memory_mb: int
Expand Down Expand Up @@ -100,8 +100,7 @@ def add_defaults(configs: Dict[str, any]) -> AWSConfidentialComputeConfig:
try:
client = boto3.client("secretsmanager", region_name=region)
except Exception as e:
# MissingInstanceProfile
raise RuntimeError("Please use IAM instance profile for your instance and make sure that has permission to access Secret Manager")
raise RuntimeError("Please use IAM instance profile for your instance and make sure that has permission to access Secret Manager", e)
try:
secret = add_defaults(json.loads(client.get_secret_value(SecretId=secret_identifier)["SecretString"]))
self.__validate_aws_specific_config(secret)
Expand Down Expand Up @@ -204,7 +203,7 @@ def __run_nitro_enclave(self):
"--enclave-cid", "42",
"--enclave-name", "uid2operator"
]
if self.configs["debug_mode"]:
if self.configs('debug_mode', True): #E2E override
print("Running in debug_mode")
command += ["--debug-mode", "--attach-console"]
self.run_command(command)
Expand Down Expand Up @@ -247,12 +246,14 @@ def __kill_auxiliaries(self) -> None:
parser = argparse.ArgumentParser(description="Manage EC2-based confidential compute workflows.")
parser.add_argument("-o", "--operation", choices=["stop", "start"], default="start", help="Operation to perform.")
args = parser.parse_args()
ec2 = EC2()
try:
ec2 = EC2()
if args.operation == "stop":
ec2.cleanup()
else:
ec2.run_compute()
except ConfidentialComputeStartupException as e:
print("Failed starting up Confidential Compute. Please find the error code and documentation", e)
except Exception as e:
print("Failed starting up Confidential Compute. Please contact uid2", e)
print("Unknown failure while starting up Confidential Compute. Please contact UID support team with this log", e)

10 changes: 7 additions & 3 deletions scripts/confidential_compute.py
Original file line number Diff line number Diff line change
Expand Up @@ -122,16 +122,20 @@ def run_command(command, seperate_process=False):
except Exception as e:
print(f"Failed to run command: {str(e)}")
raise RuntimeError (f"Failed to start {' '.join(command)} ")

class ConfidentialComputeStartupException(Exception):
def __init__(self, message):
super().__init__(message)

class MissingConfigError(Exception):
class MissingConfigError(ConfidentialComputeStartupException):
"""Custom exception to handle missing config keys."""
def __init__(self, missing_keys):
self.missing_keys = missing_keys
self.message = f"\n Missing configuration keys: {', '.join(missing_keys)} \n"
super().__init__(self.message)

class SecretNotFoundException(Exception):
class SecretNotFoundException(ConfidentialComputeStartupException):
"""Custom exception if secret manager is not found"""
def __init__(self, name):
self.message = f"Secret manager not found - {name}"
self.message = f"Secret manager not found - {name}. Please check if secret exist and the Instance Profile has permission to read it"
super().__init__(self.message)

0 comments on commit a6650e0

Please sign in to comment.