diff --git a/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py b/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py index dd56960e0..9b300bdb4 100644 --- a/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py +++ b/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py @@ -11,6 +11,8 @@ import asyncio import logging +import math +import os from datetime import datetime from typing import Any @@ -31,6 +33,86 @@ logger = logging.getLogger(__name__) +# Default botocore timeouts. Every Rekognition/S3 call in this module runs via +# ``asyncio.to_thread``, i.e. on the *shared* default executor. Without an +# explicit read timeout a stalled AWS request pins one of the pool's limited +# worker threads indefinitely, which would starve every other ``to_thread`` +# user in the process. Bounding the request bounds the worker. +_DEFAULT_CONNECT_TIMEOUT = 10.0 +_DEFAULT_READ_TIMEOUT = 60.0 +_DEFAULT_MAX_ATTEMPTS = 3 + + +def _positive_finite_float_env(name: str, default: float) -> float: + """Read a strictly positive, finite float from the environment.""" + raw = os.environ.get(name) + if raw is None or not raw.strip(): + return default + try: + value = float(raw) + except (TypeError, ValueError) as exc: + raise ConfigurationError( + f"{name} must be a number, got {raw!r}", + provider=CloudAIProvider.AWS_REKOGNITION.value, + ) from exc + if not math.isfinite(value) or value <= 0: + raise ConfigurationError( + f"{name} must be a positive finite number, got {raw!r}", + provider=CloudAIProvider.AWS_REKOGNITION.value, + ) + return value + + +def _positive_int_env(name: str, default: int) -> int: + """Read a strictly positive integer from the environment.""" + raw = os.environ.get(name) + if raw is None or not raw.strip(): + return default + try: + value = int(raw) + except (TypeError, ValueError) as exc: + raise ConfigurationError( + f"{name} must be an integer, got {raw!r}", + provider=CloudAIProvider.AWS_REKOGNITION.value, + ) from exc + if value <= 0: + raise ConfigurationError( + f"{name} must be a positive integer, got {raw!r}", + provider=CloudAIProvider.AWS_REKOGNITION.value, + ) + return value + + +def _timeout_config_kwargs() -> dict: + """Parse botocore timeout settings from the environment. + + Kept separate from ``initialize`` so a bad value raises ConfigurationError + with its precise message instead of being re-wrapped as a generic + CloudAIError by that method's catch-all handler. + """ + return { + 'connect_timeout': _positive_finite_float_env( + 'AWS_REKOGNITION_CONNECT_TIMEOUT', _DEFAULT_CONNECT_TIMEOUT + ), + 'read_timeout': _positive_finite_float_env( + 'AWS_REKOGNITION_READ_TIMEOUT', _DEFAULT_READ_TIMEOUT + ), + 'retries': { + 'max_attempts': _positive_int_env( + 'AWS_REKOGNITION_MAX_ATTEMPTS', _DEFAULT_MAX_ATTEMPTS + ), + 'mode': 'standard', + }, + } + + +def _read_file_bytes(path: str) -> bytes: + """Read a file's bytes. Module-level so it can run in a worker thread.""" + with open(path, 'rb') as handle: + return handle.read() + + + class AWSRekognition(BaseCloudAI): """Amazon Rekognition video and image analysis integration.""" @@ -57,8 +139,12 @@ def _validate_config(self) -> None: async def initialize(self) -> None: """Initialize AWS Rekognition client.""" + # Parsed before the try so a misconfigured value surfaces as a + # ConfigurationError rather than the catch-all CloudAIError below. + timeout_kwargs = _timeout_config_kwargs() try: import boto3 + from botocore.config import Config as BotoConfig from botocore.exceptions import ClientError, NoCredentialsError # Create session with credentials @@ -68,9 +154,15 @@ async def initialize(self) -> None: region_name=self.config["region"] ) + # Bound every request so a stalled call cannot hold a shared + # executor thread forever (all SDK calls here run in to_thread). + client_config = BotoConfig(**timeout_kwargs) + # Initialize clients - self._rekognition_client = session.client('rekognition') - self._s3_client = session.client('s3') + self._rekognition_client = session.client( + 'rekognition', config=client_config + ) + self._s3_client = session.client('s3', config=client_config) # Test connection await self._test_connection() @@ -103,8 +195,9 @@ async def _test_connection(self) -> None: """Test AWS Rekognition connection.""" try: # Simple API call to test connectivity - self._rekognition_client.describe_collection( - CollectionId='non-existent-collection' + await asyncio.to_thread( + self._rekognition_client.describe_collection, + CollectionId='non-existent-collection', ) except Exception as e: if "ResourceNotFoundException" in str(e): @@ -174,27 +267,31 @@ async def analyze_image(self, image_url: str, results = {} if AnalysisType.LABEL_DETECTION in analysis_types: - results['labels'] = self._rekognition_client.detect_labels( + results['labels'] = await asyncio.to_thread( + self._rekognition_client.detect_labels, Image=image_data, MaxLabels=50, - MinConfidence=0.5 + MinConfidence=0.5, ) if AnalysisType.FACE_DETECTION in analysis_types: - results['faces'] = self._rekognition_client.detect_faces( + results['faces'] = await asyncio.to_thread( + self._rekognition_client.detect_faces, Image=image_data, - Attributes=['ALL'] + Attributes=['ALL'], ) if AnalysisType.TEXT_DETECTION in analysis_types: - results['text'] = self._rekognition_client.detect_text( - Image=image_data + results['text'] = await asyncio.to_thread( + self._rekognition_client.detect_text, + Image=image_data, ) if AnalysisType.CONTENT_MODERATION in analysis_types: - results['moderation'] = self._rekognition_client.detect_moderation_labels( + results['moderation'] = await asyncio.to_thread( + self._rekognition_client.detect_moderation_labels, Image=image_data, - MinConfidence=0.5 + MinConfidence=0.5, ) processing_time = (datetime.utcnow() - start_time).total_seconds() @@ -219,8 +316,9 @@ async def get_service_status(self) -> dict[str, Any]: # Test with describe_collection call try: - self._rekognition_client.describe_collection( - CollectionId='health-check-collection' + await asyncio.to_thread( + self._rekognition_client.describe_collection, + CollectionId='health-check-collection', ) except Exception as e: if "ResourceNotFoundException" in str(e): @@ -278,28 +376,32 @@ async def _start_video_analysis(self, s3_bucket: str, s3_key: str, # Start different analysis operations based on requested types if AnalysisType.LABEL_DETECTION in analysis_types: - response = self._rekognition_client.start_label_detection( + response = await asyncio.to_thread( + self._rekognition_client.start_label_detection, Video=video_input, - MinConfidence=0.5 + MinConfidence=0.5, ) job_ids['labels'] = response['JobId'] if AnalysisType.FACE_DETECTION in analysis_types: - response = self._rekognition_client.start_face_detection( - Video=video_input + response = await asyncio.to_thread( + self._rekognition_client.start_face_detection, + Video=video_input, ) job_ids['faces'] = response['JobId'] if AnalysisType.TEXT_DETECTION in analysis_types: - response = self._rekognition_client.start_text_detection( - Video=video_input + response = await asyncio.to_thread( + self._rekognition_client.start_text_detection, + Video=video_input, ) job_ids['text'] = response['JobId'] if AnalysisType.CONTENT_MODERATION in analysis_types: - response = self._rekognition_client.start_content_moderation( + response = await asyncio.to_thread( + self._rekognition_client.start_content_moderation, Video=video_input, - MinConfidence=0.5 + MinConfidence=0.5, ) job_ids['moderation'] = response['JobId'] @@ -326,13 +428,25 @@ async def _wait_for_job_completion(self, job_id: str, analysis_type: str) -> dic while elapsed_time < max_wait_time: try: if analysis_type == 'labels': - response = self._rekognition_client.get_label_detection(JobId=job_id) + response = await asyncio.to_thread( + self._rekognition_client.get_label_detection, + JobId=job_id, + ) elif analysis_type == 'faces': - response = self._rekognition_client.get_face_detection(JobId=job_id) + response = await asyncio.to_thread( + self._rekognition_client.get_face_detection, + JobId=job_id, + ) elif analysis_type == 'text': - response = self._rekognition_client.get_text_detection(JobId=job_id) + response = await asyncio.to_thread( + self._rekognition_client.get_text_detection, + JobId=job_id, + ) elif analysis_type == 'moderation': - response = self._rekognition_client.get_content_moderation(JobId=job_id) + response = await asyncio.to_thread( + self._rekognition_client.get_content_moderation, + JobId=job_id, + ) else: raise ValueError(f"Unknown analysis type: {analysis_type}") @@ -369,9 +483,8 @@ async def _prepare_image_input(self, image_url: str) -> dict[str, Any]: response = await client.get(image_url) return {'Bytes': response.content} else: - # Local file - with open(image_url, 'rb') as image_file: - return {'Bytes': image_file.read()} + # Local file - read off the event loop + return {'Bytes': await asyncio.to_thread(_read_file_bytes, image_url)} def _process_video_results(self, results: dict[str, Any], video_id: str, analysis_types: list[AnalysisType], diff --git a/tests/unit/test_aws_rekognition_provider.py b/tests/unit/test_aws_rekognition_provider.py index 2f0ad4000..ea88ae298 100644 --- a/tests/unit/test_aws_rekognition_provider.py +++ b/tests/unit/test_aws_rekognition_provider.py @@ -2,6 +2,8 @@ from __future__ import annotations +import math +import os import sys from pathlib import Path from unittest.mock import AsyncMock, MagicMock, patch @@ -131,6 +133,7 @@ async def test_initialize_sets_rekognition_client(self): mock_boto3.Session.return_value = mock_session with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore": MagicMock(), + "botocore.config": MagicMock(), "botocore.exceptions": MagicMock( ClientError=Exception, NoCredentialsError=Exception )}): @@ -147,6 +150,7 @@ async def test_initialize_sets_s3_client(self): mock_boto3.Session.return_value = mock_session with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore": MagicMock(), + "botocore.config": MagicMock(), "botocore.exceptions": MagicMock( ClientError=Exception, NoCredentialsError=Exception )}): @@ -169,6 +173,7 @@ async def test_initialize_uses_correct_region(self): mock_boto3.Session.return_value = mock_session with patch.dict("sys.modules", {"boto3": mock_boto3, "botocore": MagicMock(), + "botocore.config": MagicMock(), "botocore.exceptions": MagicMock( ClientError=Exception, NoCredentialsError=Exception )}): @@ -1004,3 +1009,298 @@ def test_cost_scales_with_duration(self): cost1 = provider.estimate_cost(60.0, [AnalysisType.LABEL_DETECTION]) cost2 = provider.estimate_cost(120.0, [AnalysisType.LABEL_DETECTION]) assert cost2 == pytest.approx(cost1 * 2) + + +# =========================================================================== +# Event-loop responsiveness: every boto3 call must run off the loop +# =========================================================================== + +async def _count_heartbeats(coro, tick: float = 0.005) -> tuple[object, int]: + """Run ``coro`` while a heartbeat task ticks; return (result, ticks). + + If the awaited work performs blocking I/O directly on the event loop the + heartbeat never gets scheduled and ``ticks`` stays at 0. + """ + import asyncio + import contextlib + + ticks = 0 + stop = False + + async def _beat(): + nonlocal ticks + while not stop: + await asyncio.sleep(tick) + ticks += 1 + + beat = asyncio.create_task(_beat()) + await asyncio.sleep(0) # let the heartbeat reach its first await first + try: + result = await coro + finally: + stop = True + beat.cancel() + with contextlib.suppress(asyncio.CancelledError): + await beat + return result, ticks + + +def _slow(return_value, delay: float = 0.12): + """A synchronous callable that blocks for ``delay`` seconds.""" + import time + + def _call(*_args, **_kwargs): + time.sleep(delay) + return return_value + + return _call + + +class TestRekognitionDoesNotBlockEventLoop: + @pytest.mark.parametrize( + ("analysis_type", "client_method", "response"), + [ + (AnalysisType.LABEL_DETECTION, "detect_labels", {"Labels": []}), + (AnalysisType.FACE_DETECTION, "detect_faces", {"FaceDetails": []}), + (AnalysisType.TEXT_DETECTION, "detect_text", {"TextDetections": []}), + ( + AnalysisType.CONTENT_MODERATION, + "detect_moderation_labels", + {"ModerationLabels": []}, + ), + ], + ) + async def test_each_image_detection_runs_off_the_loop( + self, analysis_type, client_method, response + ): + provider = _make_provider() + mock_client = _make_rekognition_client() + getattr(mock_client, client_method).side_effect = _slow(response) + provider._rekognition_client = mock_client + + with patch.object( + provider, + "_prepare_image_input", + new=AsyncMock(return_value={"Bytes": b"img"}), + ): + result, ticks = await _count_heartbeats( + provider.analyze_image("http://e.com/i.jpg", [analysis_type]) + ) + + assert isinstance(result, VideoAnalysisResult) + getattr(mock_client, client_method).assert_called_once() + assert ticks > 0, f"{client_method} blocked the event loop" + + @pytest.mark.parametrize( + ("analysis_type", "client_method", "result_key"), + [ + (AnalysisType.LABEL_DETECTION, "start_label_detection", "labels"), + (AnalysisType.FACE_DETECTION, "start_face_detection", "faces"), + (AnalysisType.TEXT_DETECTION, "start_text_detection", "text"), + ( + AnalysisType.CONTENT_MODERATION, + "start_content_moderation", + "moderation", + ), + ], + ) + async def test_each_video_start_runs_off_the_loop( + self, analysis_type, client_method, result_key + ): + provider = _make_provider() + mock_client = MagicMock() + getattr(mock_client, client_method).side_effect = _slow({"JobId": "j-1"}) + provider._rekognition_client = mock_client + + result, ticks = await _count_heartbeats( + provider._start_video_analysis( + "my-bucket", "video.mp4", [analysis_type] + ) + ) + + assert result == {result_key: "j-1"} + assert ticks > 0, f"{client_method} blocked the event loop" + + @pytest.mark.parametrize( + ("analysis_key", "client_method"), + [ + ("labels", "get_label_detection"), + ("faces", "get_face_detection"), + ("text", "get_text_detection"), + ("moderation", "get_content_moderation"), + ], + ) + async def test_each_video_poll_runs_off_the_loop( + self, analysis_key, client_method + ): + provider = _make_provider() + mock_client = MagicMock() + getattr(mock_client, client_method).side_effect = _slow( + {"JobStatus": "SUCCEEDED"} + ) + provider._rekognition_client = mock_client + + result, ticks = await _count_heartbeats( + provider._wait_for_job_completion("job-1", analysis_key) + ) + + assert result["JobStatus"] == "SUCCEEDED" + assert ticks > 0, f"{client_method} blocked the event loop" + + @pytest.mark.parametrize("provider_method", ["_test_connection", "get_service_status"]) + async def test_each_describe_collection_runs_off_the_loop(self, provider_method): + provider = _make_provider() + mock_client = MagicMock() + mock_client.describe_collection.side_effect = _slow({}) + provider._rekognition_client = mock_client + + _result, ticks = await _count_heartbeats( + getattr(provider, provider_method)() + ) + + mock_client.describe_collection.assert_called_once() + assert ticks > 0, ( + f"describe_collection in {provider_method} blocked the event loop" + ) + + async def test_local_image_read_runs_off_the_event_loop_thread(self, tmp_path): + """ + Deterministic counterpart to the heartbeat tests. + + The local read is far too fast to time reliably on a loaded runner, so + instead of measuring elapsed ticks this asserts the property directly: + the read must execute on a worker thread, not on the thread running the + event loop. This is immune to scheduler load and needs no sleeps. + """ + import threading + + image = tmp_path / "frame.jpg" + image.write_bytes(b"BINARY-IMAGE-PAYLOAD") + provider = _make_provider() + + # Patch the exact namespace that _prepare_image_input resolves from. + # Other test modules evict cloud_ai modules from sys.modules during + # collection, so re-importing the module here can patch a stale object. + method_globals = type(provider)._prepare_image_input.__globals__ + real_read = method_globals["_read_file_bytes"] + + loop_thread_id = threading.get_ident() + observed: dict[str, int] = {} + + def _recording_read(path): + observed['thread_id'] = threading.get_ident() + return real_read(path) + + with patch.dict(method_globals, {"_read_file_bytes": _recording_read}): + result = await provider._prepare_image_input(str(image)) + + assert result == {'Bytes': b"BINARY-IMAGE-PAYLOAD"} + assert 'thread_id' in observed, "_read_file_bytes was never called" + assert observed['thread_id'] != loop_thread_id, ( + "the local image read ran on the event loop thread instead of " + "being dispatched to a worker thread" + ) + + async def test_local_image_bytes_are_read_correctly(self, tmp_path): + """Guard: offloading must not change what is returned.""" + image = tmp_path / "frame.png" + payload = bytes(range(256)) * 8 + image.write_bytes(payload) + provider = _make_provider() + + result = await provider._prepare_image_input(str(image)) + + assert result == {'Bytes': payload} + + +# --------------------------------------------------------------------------- +# botocore client timeouts +# --------------------------------------------------------------------------- +class TestClientTimeoutConfiguration: + """Every SDK call here runs on the shared default executor via + ``asyncio.to_thread``. Without an explicit read timeout a stalled AWS + request would pin one of that pool's limited worker threads forever and + starve every other ``to_thread`` user in the process, so the clients must + always be built with a bounded botocore config. + """ + + @staticmethod + async def _initialize(env: dict | None = None): + provider = _make_provider() + mock_session = MagicMock() + mock_session.client.return_value = _make_rekognition_client() + mock_boto3 = MagicMock() + mock_boto3.Session.return_value = mock_session + config_mod = MagicMock() + + stubs = { + "boto3": mock_boto3, + "botocore": MagicMock(), + "botocore.config": config_mod, + "botocore.exceptions": MagicMock( + ClientError=Exception, NoCredentialsError=Exception + ), + } + with patch.dict("sys.modules", stubs), patch.dict(os.environ, env or {}): + await provider.initialize() + return mock_session, config_mod + + async def test_both_clients_are_built_with_a_bounded_config(self): + mock_session, config_mod = await self._initialize() + + built = config_mod.Config.return_value + services = {} + for call in mock_session.client.call_args_list: + services[call.args[0]] = call.kwargs.get("config") + + assert set(services) == {"rekognition", "s3"} + assert services["rekognition"] is built + assert services["s3"] is built + + async def test_default_timeouts_are_finite_and_positive(self): + _session, config_mod = await self._initialize() + + kwargs = config_mod.Config.call_args.kwargs + assert 0 < kwargs["connect_timeout"] < math.inf + assert 0 < kwargs["read_timeout"] < math.inf + assert kwargs["retries"]["max_attempts"] >= 1 + + async def test_timeouts_are_overridable_from_the_environment(self): + _session, config_mod = await self._initialize( + { + "AWS_REKOGNITION_CONNECT_TIMEOUT": "2.5", + "AWS_REKOGNITION_READ_TIMEOUT": "7", + "AWS_REKOGNITION_MAX_ATTEMPTS": "5", + } + ) + + kwargs = config_mod.Config.call_args.kwargs + assert kwargs["connect_timeout"] == 2.5 + assert kwargs["read_timeout"] == 7.0 + assert kwargs["retries"]["max_attempts"] == 5 + + @pytest.mark.parametrize( + "var", + ["AWS_REKOGNITION_CONNECT_TIMEOUT", "AWS_REKOGNITION_READ_TIMEOUT"], + ) + @pytest.mark.parametrize("bad", ["inf", "-inf", "nan", "0", "-1", "abc"]) + async def test_non_finite_or_non_positive_timeouts_are_rejected(self, var, bad): + with pytest.raises(ConfigurationError): + await self._initialize({var: bad}) + + @pytest.mark.parametrize("bad", ["0", "-3", "1.5", "abc"]) + async def test_invalid_max_attempts_is_rejected(self, bad): + with pytest.raises(ConfigurationError): + await self._initialize({"AWS_REKOGNITION_MAX_ATTEMPTS": bad}) + + @pytest.mark.parametrize( + "var", + [ + "AWS_REKOGNITION_CONNECT_TIMEOUT", + "AWS_REKOGNITION_READ_TIMEOUT", + "AWS_REKOGNITION_MAX_ATTEMPTS", + ], + ) + async def test_blank_values_fall_back_to_defaults(self, var): + _session, config_mod = await self._initialize({var: " "}) + assert config_mod.Config.call_args is not None