From 3e5e39ce919e1b03830bbcf701426025fe209f76 Mon Sep 17 00:00:00 2001 From: MinhThieu145 Date: Wed, 19 Jun 2024 21:19:07 -0600 Subject: [PATCH 1/3] Add admin actions for ec2 start and stop workers --- apps/challenges/admin.py | 54 +++++++++++++++++++ apps/challenges/aws_utils.py | 100 +++++++++++++++++++++++++++++++++-- 2 files changed, 149 insertions(+), 5 deletions(-) diff --git a/apps/challenges/admin.py b/apps/challenges/admin.py index f056d2c679..99cb0df1cd 100644 --- a/apps/challenges/admin.py +++ b/apps/challenges/admin.py @@ -11,6 +11,9 @@ scale_workers, start_workers, stop_workers, + start_ec2_workers_list, + stop_ec2_workers_list + ) from .admin_filters import ChallengeFilter @@ -85,8 +88,59 @@ class ChallengeAdmin(ImportExportTimeStampedAdmin): "scale_selected_workers", "restart_selected_workers", "delete_selected_workers", + "start_selected_ec2_instance_workers", + "stop_selected_ec2_instance_workers" ] action_form = UpdateNumOfWorkersForm + + + def start_selected_ec2_instance_workers(self, request, queryset): + response = start_ec2_workers_list(queryset) + count, failures = response["count"], response["failures"] + + if count == queryset.count(): + message = "All selected challenge workers successfully started." + messages.success(request, message) + else: + messages.success( + request, + "{} challenge workers were succesfully started.".format(count), + ) + for fail in failures: + challenge_pk, message = fail["challenge_pk"], fail["message"] + display_message = "Challenge {}: {}".format( + challenge_pk, message + ) + messages.error(request, display_message) + + start_selected_ec2_instance_workers.short_description = ( + "Start all selected ec2 instance challenge workers." + ) + + + def stop_selected_ec2_instance_workers(self, request, queryset): + response = stop_ec2_workers_list(queryset) + count, failures = response["count"], response["failures"] + + if count == queryset.count(): + message = "All selected challenge workers successfully stopped." + messages.success(request, message) + else: + messages.success( + request, + "{} challenge workers were succesfully stopped.".format(count), + ) + for fail in failures: + challenge_pk, message = fail["challenge_pk"], fail["message"] + display_message = "Challenge {}: {}".format( + challenge_pk, message + ) + messages.error(request, display_message) + + stop_selected_ec2_instance_workers.short_description = ( + "Stop all selected ec2 instance challenge workers." + ) + def start_selected_workers(self, request, queryset): response = start_workers(queryset) diff --git a/apps/challenges/aws_utils.py b/apps/challenges/aws_utils.py index 9278cb4487..99c397f272 100644 --- a/apps/challenges/aws_utils.py +++ b/apps/challenges/aws_utils.py @@ -483,9 +483,15 @@ def stop_ec2_instance(challenge): Returns: dict: A dictionary containing the status and message of the stop operation. """ + from .utils import get_aws_credentials_for_challenge + + for obj in serializers.deserialize("json", challenge): + challenge_obj = obj.object + challenge_aws_keys = get_aws_credentials_for_challenge(challenge_obj.pk) + target_instance_id = challenge.ec2_instance_id - ec2 = get_boto3_client("ec2", aws_keys) + ec2 = get_boto3_client("ec2", challenge_aws_keys) status_response = ec2.describe_instance_status(InstanceIds=[target_instance_id]) if status_response["InstanceStatuses"]: @@ -555,7 +561,7 @@ def describe_ec2_instance(challenge): } -def start_ec2_instance(challenge): +def start_ec2_instance(challenge, aws_keys): """ Start the EC2 instance associated with a challenge. @@ -675,7 +681,7 @@ def terminate_ec2_instance(challenge): } -def create_ec2_instance(challenge, ec2_storage=None, worker_instance_type=None, worker_image_url=None): +def create_ec2_instance(challenge, aws_keys, ec2_storage=None, worker_instance_type=None, worker_image_url=None): """ Create the EC2 instance associated with a challenge. @@ -799,6 +805,47 @@ def update_sqs_retention_period(challenge): } +def start_ec2_workers_list(queryset): + """ + The function called by the admin action method to start all the selected ec2 workers. + Calls the setup ec2 method. Before calling, checks if it uses ec2 worker. + Parameters: + queryset (): The queryset of selected challenges in the django admin page. + Returns: + dict: keys-> 'count': the number of workers successfully started. + 'failures': a dict of all the failures with their error messages and the challenge pk + """ + + if settings.DEBUG: + failures = [] + for challenge in queryset: + failures.append( + { + "message": "Workers cannot be started on AWS in development environment", + "challenge_pk": challenge.pk, + } + ) + return {"count": 0, "failures": failures} + + count = 0 + failures = [] + for challenge in queryset: + if challenge.uses_ec2_worker: + response = setup_ec2(challenge) + if "error" in response: + failures.append( + {"message": response["error"], "challenge_pk": challenge.pk} + ) + else: + count += 1 + else: + response = "Please select challenge with inactive workers only." + failures.append( + {"message": response, "challenge_pk": challenge.pk} + ) + return {"count": count, "failures": failures} + + def start_workers(queryset): """ The function called by the admin action method to start all the selected workers. @@ -848,6 +895,46 @@ def start_workers(queryset): return {"count": count, "failures": failures} +def stop_ec2_workers_list(queryset): + """ + The function called by the admin action method to start all the selected ec2 workers. + Calls the stop ec2 instance method. Before calling, checks if it uses ec2 worker. + Parameters: + queryset (): The queryset of selected challenges in the django admin page. + Returns: + dict: keys-> 'count': the number of workers successfully started. + 'failures': a dict of all the failures with their error messages and the challenge pk + """ + if settings.DEBUG: + failures = [] + for challenge in queryset: + failures.append( + { + "message": "Workers cannot be started on AWS in development environment", + "challenge_pk": challenge.pk, + } + ) + return {"count": 0, "failures": failures} + + count = 0 + failures = [] + for challenge in queryset: + if challenge.uses_ec2_worker: + response = stop_ec2_instance(challenge) + if "error" in response: + failures.append( + {"message": response["error"], "challenge_pk": challenge.pk} + ) + else: + count += 1 + else: + response = "Please select challenge with inactive workers only." + failures.append( + {"message": response, "challenge_pk": challenge.pk} + ) + return {"count": count, "failures": failures} + + def stop_workers(queryset): """ The function called by the admin action method to stop all the selected workers. @@ -1817,11 +1904,14 @@ def setup_ec2(challenge): Arguments: challenge {} -- instance of the model calling the post hook """ + from .utils import get_aws_credentials_for_challenge + for obj in serializers.deserialize("json", challenge): challenge_obj = obj.object + challenge_aws_keys = get_aws_credentials_for_challenge(challenge_obj.pk) if challenge_obj.ec2_instance_id: - return start_ec2_instance(challenge_obj) - return create_ec2_instance(challenge_obj) + return start_ec2_instance(challenge_obj, challenge_aws_keys) + return create_ec2_instance(challenge_obj, challenge_aws_keys) @app.task From f87237323ada4abded8de5b6af8f19f93b7d7694 Mon Sep 17 00:00:00 2001 From: MinhThieu145 Date: Wed, 19 Jun 2024 21:30:59 -0600 Subject: [PATCH 2/3] Update format --- apps/challenges/admin.py | 5 +---- apps/challenges/aws_utils.py | 30 +++++++++++++++++++++--------- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/apps/challenges/admin.py b/apps/challenges/admin.py index 99cb0df1cd..989f9dbe3d 100644 --- a/apps/challenges/admin.py +++ b/apps/challenges/admin.py @@ -89,10 +89,9 @@ class ChallengeAdmin(ImportExportTimeStampedAdmin): "restart_selected_workers", "delete_selected_workers", "start_selected_ec2_instance_workers", - "stop_selected_ec2_instance_workers" + "stop_selected_ec2_instance_workers", ] action_form = UpdateNumOfWorkersForm - def start_selected_ec2_instance_workers(self, request, queryset): response = start_ec2_workers_list(queryset) @@ -117,7 +116,6 @@ def start_selected_ec2_instance_workers(self, request, queryset): "Start all selected ec2 instance challenge workers." ) - def stop_selected_ec2_instance_workers(self, request, queryset): response = stop_ec2_workers_list(queryset) count, failures = response["count"], response["failures"] @@ -141,7 +139,6 @@ def stop_selected_ec2_instance_workers(self, request, queryset): "Stop all selected ec2 instance challenge workers." ) - def start_selected_workers(self, request, queryset): response = start_workers(queryset) count, failures = response["count"], response["failures"] diff --git a/apps/challenges/aws_utils.py b/apps/challenges/aws_utils.py index 99c397f272..b735c5f074 100644 --- a/apps/challenges/aws_utils.py +++ b/apps/challenges/aws_utils.py @@ -484,15 +484,17 @@ def stop_ec2_instance(challenge): dict: A dictionary containing the status and message of the stop operation. """ from .utils import get_aws_credentials_for_challenge - + for obj in serializers.deserialize("json", challenge): challenge_obj = obj.object challenge_aws_keys = get_aws_credentials_for_challenge(challenge_obj.pk) - + target_instance_id = challenge.ec2_instance_id ec2 = get_boto3_client("ec2", challenge_aws_keys) - status_response = ec2.describe_instance_status(InstanceIds=[target_instance_id]) + status_response = ec2.describe_instance_status( + InstanceIds=[target_instance_id] + ) if status_response["InstanceStatuses"]: instance_status = status_response["InstanceStatuses"][0] @@ -504,8 +506,12 @@ def stop_ec2_instance(challenge): if instance_state == "running": try: - response = ec2.stop_instances(InstanceIds=[target_instance_id]) - message = "Instance for challenge {} successfully stopped.".format(challenge.pk) + response = ec2.stop_instances( + InstanceIds=[target_instance_id] + ) + message = "Instance for challenge {} successfully stopped.".format( + challenge.pk + ) return { "response": response, "message": message, @@ -516,17 +522,23 @@ def stop_ec2_instance(challenge): "error": e.response, } else: - message = "Instance for challenge {} is not running. Please ensure the instance is running.".format(challenge.pk) + message = "Instance for challenge {} is not running. Please ensure the instance is running.".format( + challenge.pk + ) return { "error": message, } else: - message = "Instance status checks are not ready for challenge {}. Please wait for the status checks to pass.".format(challenge.pk) + message = "Instance status checks are not ready for challenge {}. Please wait for the status checks to pass.".format( + challenge.pk + ) return { "error": message, } else: - message = "Instance for challenge {} not found. Please ensure the instance exists.".format(challenge.pk) + message = "Instance for challenge {} not found. Please ensure the instance exists.".format( + challenge.pk + ) return { "error": message, } @@ -815,7 +827,7 @@ def start_ec2_workers_list(queryset): dict: keys-> 'count': the number of workers successfully started. 'failures': a dict of all the failures with their error messages and the challenge pk """ - + if settings.DEBUG: failures = [] for challenge in queryset: From 4dc05f360cd9afd59b03317e47b757b8efa113d0 Mon Sep 17 00:00:00 2001 From: MinhThieu145 Date: Wed, 19 Jun 2024 22:02:04 -0600 Subject: [PATCH 3/3] update format --- apps/challenges/aws_utils.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/apps/challenges/aws_utils.py b/apps/challenges/aws_utils.py index b735c5f074..ecc506c118 100644 --- a/apps/challenges/aws_utils.py +++ b/apps/challenges/aws_utils.py @@ -826,7 +826,7 @@ def start_ec2_workers_list(queryset): Returns: dict: keys-> 'count': the number of workers successfully started. 'failures': a dict of all the failures with their error messages and the challenge pk - """ + """ if settings.DEBUG: failures = [] @@ -846,7 +846,10 @@ def start_ec2_workers_list(queryset): response = setup_ec2(challenge) if "error" in response: failures.append( - {"message": response["error"], "challenge_pk": challenge.pk} + { + "message": response["error"], + "challenge_pk": challenge.pk, + } ) else: count += 1