diff --git a/apps/challenges/aws_utils.py b/apps/challenges/aws_utils.py index a38aeeda2a..9914cc9109 100644 --- a/apps/challenges/aws_utils.py +++ b/apps/challenges/aws_utils.py @@ -962,6 +962,93 @@ def delete_challenge_cleanup_schedule(challenge): pass # Schedule already deleted or never created +def schedule_challenge_cleanup_soon(challenge, delay_minutes=1): + """ + Schedule the pending-aware cleanup Lambda to run shortly after now. + + Used when a challenge's end_date is moved into the past. EventBridge + Scheduler rejects ``at()`` expressions in the past, so we cannot reuse + ``schedule_challenge_cleanup`` / ``update_challenge_cleanup_schedule`` + with the new end_date. The Lambda still checks pending submissions + before deleting ECS resources, so queued/running work can drain. + + Parameters: + challenge (): Challenge to clean up. + delay_minutes (int): Minutes from now to fire the schedule (default 1). + """ + from datetime import timedelta + + from django.utils import timezone + + if settings.DEBUG: + logger.info( + "Skipping soon-cleanup schedule for challenge %s in development " + "environment.", + challenge.pk, + ) + return + + if not CHALLENGE_CLEANUP_LAMBDA_ARN or not EVENTBRIDGE_SCHEDULER_ROLE_ARN: + logger.warning( + "CHALLENGE_CLEANUP_LAMBDA_ARN or EVENTBRIDGE_SCHEDULER_ROLE_ARN " + "not set. Skipping soon-cleanup schedule for challenge %s.", + challenge.pk, + ) + return + + run_at = timezone.now() + timedelta(minutes=delay_minutes) + schedule_name = ( + f"evalai-cleanup-challenge-{settings.ENVIRONMENT}-{challenge.pk}" + ) + schedule_expression = "at({})".format(run_at.strftime("%Y-%m-%dT%H:%M:%S")) + target = { + "Arn": CHALLENGE_CLEANUP_LAMBDA_ARN, + "RoleArn": EVENTBRIDGE_SCHEDULER_ROLE_ARN, + "Input": json.dumps( + { + "challenge_pk": challenge.pk, + "queue_name": challenge.queue, + } + ), + } + + try: + scheduler_client = get_boto3_client("scheduler", aws_keys) + try: + scheduler_client.update_schedule( + Name=schedule_name, + ScheduleExpression=schedule_expression, + ScheduleExpressionTimezone="UTC", + FlexibleTimeWindow={"Mode": "OFF"}, + Target=target, + ActionAfterCompletion="DELETE", + ) + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code") + if error_code != "ResourceNotFoundException": + raise + # Schedule already fired/auto-deleted; create a fresh one. + scheduler_client.create_schedule( + Name=schedule_name, + ScheduleExpression=schedule_expression, + ScheduleExpressionTimezone="UTC", + FlexibleTimeWindow={"Mode": "OFF"}, + Target=target, + ActionAfterCompletion="DELETE", + ) + logger.info( + "Scheduled soon-cleanup for challenge %s at %s", + challenge.pk, + run_at, + ) + except ClientError as e: + logger.exception( + "Failed to schedule soon-cleanup for challenge %s: %s", + challenge.pk, + e, + ) + + def ensure_workers_for_submission(challenge): """ Ensures the worker stack (ECS service, auto-scaling, EventBridge cleanup) diff --git a/apps/challenges/models.py b/apps/challenges/models.py index 7c24c97293..40021b0736 100644 --- a/apps/challenges/models.py +++ b/apps/challenges/models.py @@ -910,9 +910,12 @@ def handle_end_date_change_for_challenge(sender, instance, created, **kwargs): # Resources still exist; just reschedule the cleanup. aws.update_challenge_cleanup_schedule(challenge) else: - # New end_date is in the past; trigger cleanup if resources exist. + # New end_date is in the past. Do not force-delete workers — + # that kills in-flight/queued evaluations. Route through the + # pending-aware cleanup Lambda (#5179) so the queue can drain + # before ECS resources are removed. if challenge.workers is not None: - aws.delete_workers([challenge]) + aws.schedule_challenge_cleanup_soon(challenge) class DatasetSplit(TimeStampedModel): diff --git a/tests/unit/challenges/test_aws_utils.py b/tests/unit/challenges/test_aws_utils.py index 17919a5f76..3fee1c1500 100644 --- a/tests/unit/challenges/test_aws_utils.py +++ b/tests/unit/challenges/test_aws_utils.py @@ -46,6 +46,7 @@ scale_resources, scale_workers, schedule_challenge_cleanup, + schedule_challenge_cleanup_soon, service_manager, setup_auto_scaling_for_service, setup_ec2, @@ -5560,6 +5561,93 @@ def test_update_schedule_other_client_error_logs_exception( mock_scheduler.update_schedule.assert_called_once() +class TestScheduleChallengeCleanupSoon(unittest.TestCase): + @patch( + "challenges.aws_utils.EVENTBRIDGE_SCHEDULER_ROLE_ARN", + "arn:aws:iam::123:role/scheduler-role", + ) + @patch( + "challenges.aws_utils.CHALLENGE_CLEANUP_LAMBDA_ARN", + "arn:aws:lambda:us-east-1:123:function:cleanup", + ) + @patch("challenges.aws_utils.settings.ENVIRONMENT", "staging") + @patch("challenges.aws_utils.get_boto3_client") + def test_schedule_soon_updates_existing_schedule( + self, mock_get_boto3_client + ): + mock_scheduler = MagicMock() + mock_get_boto3_client.return_value = mock_scheduler + + challenge = MagicMock() + challenge.pk = 42 + challenge.queue = "test_queue" + + schedule_challenge_cleanup_soon(challenge, delay_minutes=1) + + mock_scheduler.update_schedule.assert_called_once() + mock_scheduler.create_schedule.assert_not_called() + call_kwargs = mock_scheduler.update_schedule.call_args[1] + assert call_kwargs["Name"] == "evalai-cleanup-challenge-staging-42" + assert call_kwargs["ScheduleExpression"].startswith("at(") + assert call_kwargs["ScheduleExpressionTimezone"] == "UTC" + assert call_kwargs["ActionAfterCompletion"] == "DELETE" + assert "challenge_pk" in call_kwargs["Target"]["Input"] + + @patch( + "challenges.aws_utils.EVENTBRIDGE_SCHEDULER_ROLE_ARN", + "arn:aws:iam::123:role/scheduler-role", + ) + @patch( + "challenges.aws_utils.CHALLENGE_CLEANUP_LAMBDA_ARN", + "arn:aws:lambda:us-east-1:123:function:cleanup", + ) + @patch("challenges.aws_utils.settings.ENVIRONMENT", "staging") + @patch("challenges.aws_utils.get_boto3_client") + def test_schedule_soon_creates_when_missing(self, mock_get_boto3_client): + mock_scheduler = MagicMock() + mock_scheduler.update_schedule.side_effect = ClientError( + error_response={ + "Error": {"Code": "ResourceNotFoundException"}, + "ResponseMetadata": {"HTTPStatusCode": 404}, + }, + operation_name="UpdateSchedule", + ) + mock_get_boto3_client.return_value = mock_scheduler + + challenge = MagicMock() + challenge.pk = 42 + challenge.queue = "test_queue" + + schedule_challenge_cleanup_soon(challenge) + + mock_scheduler.create_schedule.assert_called_once() + call_kwargs = mock_scheduler.create_schedule.call_args[1] + assert call_kwargs["Name"] == "evalai-cleanup-challenge-staging-42" + assert call_kwargs["ScheduleExpression"].startswith("at(") + + @patch( + "challenges.aws_utils.CHALLENGE_CLEANUP_LAMBDA_ARN", + "", + ) + @patch("challenges.aws_utils.get_boto3_client") + def test_schedule_soon_skipped_without_lambda_arn( + self, mock_get_boto3_client + ): + challenge = MagicMock() + challenge.pk = 42 + + schedule_challenge_cleanup_soon(challenge) + + mock_get_boto3_client.assert_not_called() + + @patch("challenges.aws_utils.settings", DEBUG=True) + def test_schedule_soon_skipped_in_debug(self, mock_settings): + challenge = MagicMock() + challenge.pk = 42 + + schedule_challenge_cleanup_soon(challenge) + + class TestDeleteChallengeCleanupSchedule(unittest.TestCase): @patch("challenges.aws_utils.settings.ENVIRONMENT", "staging") @patch("challenges.aws_utils.get_boto3_client") @@ -5659,8 +5747,11 @@ def test_end_date_extended_with_active_workers_reschedules( mock_update_schedule.assert_called_once_with(challenge) @patch("challenges.aws_utils.delete_workers") - def test_end_date_set_to_past_triggers_cleanup(self, mock_delete_workers): - """When end_date is changed to the past, trigger cleanup.""" + @patch("challenges.aws_utils.schedule_challenge_cleanup_soon") + def test_end_date_set_to_past_triggers_cleanup( + self, mock_cleanup_soon, mock_delete_workers + ): + """Past end_date uses pending-aware cleanup, not force-delete.""" from datetime import timedelta from challenges.models import handle_end_date_change_for_challenge @@ -5675,13 +5766,14 @@ def test_end_date_set_to_past_triggers_cleanup(self, mock_delete_workers): challenge._original_end_date = timezone.now() + timedelta(days=30) challenge.end_date = timezone.now() - timedelta(days=1) - mock_delete_workers.return_value = {"count": 1, "failures": []} - handle_end_date_change_for_challenge( sender=None, instance=challenge, created=False ) - mock_delete_workers.assert_called_once_with([challenge]) + # Force-deleting here would kill queued/running submissions that the + # cleanup Lambda (#5179) is designed to drain past end_date. + mock_cleanup_soon.assert_called_once_with(challenge) + mock_delete_workers.assert_not_called() def test_end_date_change_skipped_for_docker_based(self): """Docker-based challenges should be skipped."""