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..747c3206c 100644 --- a/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py +++ b/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py @@ -31,6 +31,13 @@ logger = logging.getLogger(__name__) +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.""" @@ -103,8 +110,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 +182,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 +231,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 +291,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 +343,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 +398,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..0bbd35e26 100644 --- a/tests/unit/test_aws_rekognition_provider.py +++ b/tests/unit/test_aws_rekognition_provider.py @@ -1004,3 +1004,179 @@ 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: + async def test_analyze_image_does_not_stall_the_event_loop(self): + provider = _make_provider() + mock_client = _make_rekognition_client() + mock_client.detect_labels.side_effect = _slow({"Labels": []}) + 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", [AnalysisType.LABEL_DETECTION]) + ) + + assert isinstance(result, VideoAnalysisResult) + assert ticks > 0, "detect_labels blocked the event loop" + + async def test_every_detection_type_runs_off_the_loop(self): + provider = _make_provider() + mock_client = _make_rekognition_client() + # four blocking calls back to back + mock_client.detect_labels.side_effect = _slow({"Labels": []}, 0.06) + mock_client.detect_faces.side_effect = _slow({"FaceDetails": []}, 0.06) + mock_client.detect_text.side_effect = _slow({"TextDetections": []}, 0.06) + mock_client.detect_moderation_labels.side_effect = _slow({"ModerationLabels": []}, 0.06) + 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", + [AnalysisType.LABEL_DETECTION, AnalysisType.FACE_DETECTION, + AnalysisType.TEXT_DETECTION, AnalysisType.CONTENT_MODERATION], + ) + ) + + mock_client.detect_labels.assert_called_once() + mock_client.detect_moderation_labels.assert_called_once() + assert ticks > 0, "the detect_* chain blocked the event loop" + + async def test_job_polling_does_not_stall_the_event_loop(self): + provider = _make_provider() + mock_client = MagicMock() + mock_client.get_label_detection.side_effect = _slow( + {'JobStatus': 'SUCCEEDED', 'Labels': []} + ) + provider._rekognition_client = mock_client + + result, ticks = await _count_heartbeats( + provider._wait_for_job_completion("job-1", "labels") + ) + + assert result['JobStatus'] == 'SUCCEEDED' + assert ticks > 0, "get_label_detection blocked the event loop" + + async def test_start_video_analysis_does_not_stall_the_event_loop(self): + provider = _make_provider() + mock_client = MagicMock() + mock_client.start_label_detection.side_effect = _slow({'JobId': 'j-1'}) + provider._rekognition_client = mock_client + + result, ticks = await _count_heartbeats( + provider._start_video_analysis( + "my-bucket", "video.mp4", [AnalysisType.LABEL_DETECTION] + ) + ) + + assert result == {'labels': 'j-1'} + assert ticks > 0, "start_label_detection 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() + + loop_thread_id = threading.get_ident() + observed: dict[str, int] = {} + + # Patch the *exact* globals dict the running coroutine resolves names + # in — derived from the provider we actually call — rather than a + # separately re-imported module alias. A function's ``__globals__`` is + # frozen to its defining module at ``def`` time. Under EventRelay's + # dual-import hazard (``src.youtube_extension`` vs ``youtube_extension`` + # are two distinct module objects with independent dicts), a bare + # ``from ... import aws_rekognition`` inside the test can resolve to a + # different object than the one backing ``_prepare_image_input``'s + # ``_read_file_bytes`` lookup — which silently no-ops the patch and was + # the cause of the intermittent full-suite CI failure. Sourcing the + # dict from ``type(provider)`` makes the two identical by construction. + globs = type(provider)._prepare_image_input.__globals__ + real_read = globs['_read_file_bytes'] + + def _recording_read(path): + observed['thread_id'] = threading.get_ident() + return real_read(path) + + globs['_read_file_bytes'] = _recording_read + try: + result = await provider._prepare_image_input(str(image)) + finally: + globs['_read_file_bytes'] = real_read + + 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}