Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add admin actions for starting, stopping and creating EC2 workers #4384

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Add admin actions for starting, stopping and creating EC2 workers #4384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3e5e39c
Add admin actions for ec2 start and stop workers
MinhThieu145 Jun 20, 2024
f872373
Update format
MinhThieu145 Jun 20, 2024
4dc05f3
update format
MinhThieu145 Jun 20, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 51 additions & 0 deletions apps/challenges/admin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,9 @@
scale_workers,
start_workers,
stop_workers,
start_ec2_workers_list,
stop_ec2_workers_list

)

from .admin_filters import ChallengeFilter
Expand Down Expand Up @@ -85,9 +88,57 @@ 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)
count, failures = response["count"], response["failures"]
Expand Down
127 changes: 116 additions & 11 deletions apps/challenges/aws_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -483,10 +483,18 @@ 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)
status_response = ec2.describe_instance_status(InstanceIds=[target_instance_id])
ec2 = get_boto3_client("ec2", challenge_aws_keys)
status_response = ec2.describe_instance_status(
InstanceIds=[target_instance_id]
)

if status_response["InstanceStatuses"]:
instance_status = status_response["InstanceStatuses"][0]
Expand All @@ -498,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,
Expand All @@ -510,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,
}
Expand Down Expand Up @@ -555,7 +573,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.

Expand Down Expand Up @@ -675,7 +693,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.

Expand Down Expand Up @@ -799,6 +817,50 @@ 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 (<class 'django.db.models.query.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.
Expand Down Expand Up @@ -848,6 +910,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 (<class 'django.db.models.query.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.
Expand Down Expand Up @@ -1817,11 +1919,14 @@ def setup_ec2(challenge):
Arguments:
challenge {<class 'django.db.models.query.QuerySet'>} -- 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
Expand Down