diff --git a/src/youtube_extension/integrations/cloud_ai/blocking_io.py b/src/youtube_extension/integrations/cloud_ai/blocking_io.py new file mode 100644 index 000000000..b73d71388 --- /dev/null +++ b/src/youtube_extension/integrations/cloud_ai/blocking_io.py @@ -0,0 +1,81 @@ +"""Dedicated executor for potentially blocking Cloud AI local file I/O.""" + +from __future__ import annotations + +import asyncio +import logging +import os +import threading +from concurrent.futures import ThreadPoolExecutor +from typing import Any, Callable, TypeVar + +logger = logging.getLogger(__name__) + +_T = TypeVar("_T") +_DEFAULT_IO_WORKERS = max(1, min(8, (os.cpu_count() or 1))) + + +def _io_max_workers_from_env() -> int: + raw = os.environ.get("CLOUD_AI_IO_MAX_WORKERS") + if raw is None or not raw.strip(): + return _DEFAULT_IO_WORKERS + try: + return max(1, int(raw)) + except ValueError: + logger.warning( + "Invalid CLOUD_AI_IO_MAX_WORKERS=%r; using default=%d", + raw, + _DEFAULT_IO_WORKERS, + ) + return _DEFAULT_IO_WORKERS + + +_IO_MAX_WORKERS = _io_max_workers_from_env() +_IO_EXECUTOR = ThreadPoolExecutor( + max_workers=_IO_MAX_WORKERS, + thread_name_prefix="cloud-ai-io", +) +_IO_ACTIVE = 0 +_IO_QUEUED = 0 +_IO_STATE_LOCK = threading.Lock() + + +def _log_pool_pressure(*, active: int, queued: int) -> None: + if active + queued < _IO_MAX_WORKERS: + return + logger.warning( + "blocking_io_pool_near_capacity active=%d queued=%d max_workers=%d", + active, + queued, + _IO_MAX_WORKERS, + ) + + +async def run_blocking( + func: Callable[..., _T], *args: Any, timeout: float | None = None, **kwargs: Any +) -> _T: + """Run blocking work on the Cloud AI local-I/O pool.""" + loop = asyncio.get_running_loop() + + with _IO_STATE_LOCK: + global _IO_QUEUED + _IO_QUEUED += 1 + active = _IO_ACTIVE + queued = _IO_QUEUED + _log_pool_pressure(active=active, queued=queued) + + def _invoke() -> _T: + global _IO_ACTIVE, _IO_QUEUED + with _IO_STATE_LOCK: + _IO_QUEUED -= 1 + _IO_ACTIVE += 1 + try: + return func(*args, **kwargs) + finally: + with _IO_STATE_LOCK: + _IO_ACTIVE -= 1 + + future = loop.run_in_executor(_IO_EXECUTOR, _invoke) + if timeout is None: + return await future + return await asyncio.wait_for(future, timeout=timeout) 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 a7586bee2..19480ae14 100644 --- a/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py +++ b/src/youtube_extension/integrations/cloud_ai/providers/aws_rekognition.py @@ -23,6 +23,7 @@ DetectionResult, VideoAnalysisResult, ) +from ..blocking_io import run_blocking from ..exceptions import ( AuthenticationError, CloudAIError, @@ -493,7 +494,7 @@ async def _prepare_image_input(self, image_url: str) -> dict[str, Any]: # traversal or symlink escape); the read then uses the resolved # path, off the event loop. safe_path = resolve_local_media_path(image_url, provider=self.provider.value) - return {'Bytes': await asyncio.to_thread(_read_file_bytes, str(safe_path))} + return {'Bytes': await run_blocking(_read_file_bytes, str(safe_path))} def _process_video_results(self, results: dict[str, Any], video_id: str, analysis_types: list[AnalysisType], diff --git a/src/youtube_extension/integrations/cloud_ai/providers/azure_vision.py b/src/youtube_extension/integrations/cloud_ai/providers/azure_vision.py index 8d4705198..a26309d27 100644 --- a/src/youtube_extension/integrations/cloud_ai/providers/azure_vision.py +++ b/src/youtube_extension/integrations/cloud_ai/providers/azure_vision.py @@ -20,6 +20,7 @@ DetectionResult, VideoAnalysisResult, ) +from ..blocking_io import run_blocking from ..exceptions import ( AuthenticationError, CloudAIError, @@ -268,7 +269,7 @@ async def _prepare_image_input(self, image_url: str) -> Optional[bytes]: # CLOUD_AI_MEDIA_ROOT before opening anything; the read then uses # the resolved path rather than the raw string, off the event loop. safe_path = resolve_local_media_path(image_url, provider=self.provider.value) - return await asyncio.to_thread(_read_file_bytes, str(safe_path)) + return await run_blocking(_read_file_bytes, str(safe_path)) async def _await_ocr_call(self, deadline: float, func: Any, *args: Any, **kwargs: Any) -> Any: """Run a blocking Azure SDK call in a worker thread, bounded by a diff --git a/src/youtube_extension/integrations/cloud_ai/providers/google_cloud.py b/src/youtube_extension/integrations/cloud_ai/providers/google_cloud.py index 24d8f576f..d00b28e55 100644 --- a/src/youtube_extension/integrations/cloud_ai/providers/google_cloud.py +++ b/src/youtube_extension/integrations/cloud_ai/providers/google_cloud.py @@ -21,6 +21,7 @@ DetectionResult, VideoAnalysisResult, ) +from ..blocking_io import run_blocking from ..exceptions import ( AuthenticationError, CloudAIError, @@ -194,7 +195,7 @@ async def analyze_image(self, image_url: str, safe_path = resolve_local_media_path( image_url, provider=self.provider.value ) - image.content = await asyncio.to_thread( + image.content = await run_blocking( _read_file_bytes, str(safe_path) ) diff --git a/tests/unit/test_aws_rekognition_provider.py b/tests/unit/test_aws_rekognition_provider.py index 4d2e2faf7..153a5c8ae 100644 --- a/tests/unit/test_aws_rekognition_provider.py +++ b/tests/unit/test_aws_rekognition_provider.py @@ -345,6 +345,23 @@ async def test_local_file_returns_bytes(self, tmp_path, monkeypatch): result = await provider._prepare_image_input(str(img_file)) assert result == {'Bytes': b"\xff\xd8\xff\xe0"} + async def test_local_file_read_routes_through_run_blocking(self, tmp_path, monkeypatch): + img_file = tmp_path / "test.jpg" + img_file.write_bytes(b"\xff\xd8\xff\xe0") + monkeypatch.setenv("CLOUD_AI_MEDIA_ROOT", str(tmp_path)) + provider = _make_provider() + method_globals = type(provider)._prepare_image_input.__globals__ + run_blocking = AsyncMock(return_value=b"worker-bytes") + + with patch.dict(method_globals, {"run_blocking": run_blocking}): + result = await provider._prepare_image_input(str(img_file)) + + assert result == {'Bytes': b"worker-bytes"} + run_blocking.assert_awaited_once() + read_func, read_path = run_blocking.await_args.args + assert read_func is method_globals["_read_file_bytes"] + assert read_path == str(img_file) + async def test_http_url_fetches_bytes(self): provider = _make_provider() mock_response = MagicMock() diff --git a/tests/unit/test_azure_vision_provider.py b/tests/unit/test_azure_vision_provider.py index 54b811d0e..bbe5e621a 100644 --- a/tests/unit/test_azure_vision_provider.py +++ b/tests/unit/test_azure_vision_provider.py @@ -309,6 +309,23 @@ async def test_local_file_reads_bytes_fixture(self, tmp_path, monkeypatch): result = await provider._prepare_image_input(str(img_file)) assert result == b"\xff\xd8\xff" + async def test_local_file_read_routes_through_run_blocking(self, tmp_path, monkeypatch): + provider = _make_provider() + img_file = tmp_path / "test.jpg" + img_file.write_bytes(b"\xff\xd8\xff") + monkeypatch.setenv("CLOUD_AI_MEDIA_ROOT", str(tmp_path)) + method_globals = type(provider)._prepare_image_input.__globals__ + run_blocking = AsyncMock(return_value=b"worker-bytes") + + with patch.dict(method_globals, {"run_blocking": run_blocking}): + result = await provider._prepare_image_input(str(img_file)) + + assert result == b"worker-bytes" + run_blocking.assert_awaited_once() + read_func, read_path = run_blocking.await_args.args + assert read_func is method_globals["_read_file_bytes"] + assert read_path == str(img_file) + # =========================================================================== # _convert_azure_bbox diff --git a/tests/unit/test_cloud_ai_blocking_io.py b/tests/unit/test_cloud_ai_blocking_io.py new file mode 100644 index 000000000..fc661a869 --- /dev/null +++ b/tests/unit/test_cloud_ai_blocking_io.py @@ -0,0 +1,78 @@ +from __future__ import annotations + +import asyncio +import logging +import sys +import threading +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path + +import pytest + +_SRC = Path(__file__).resolve().parents[2] / "src" +sys.path.insert(0, str(_SRC)) + +import youtube_extension.integrations.cloud_ai.blocking_io as blocking_io + + +@pytest.fixture +def isolated_pool(monkeypatch): + original_executor = blocking_io._IO_EXECUTOR + original_limit = blocking_io._IO_MAX_WORKERS + gate = threading.Event() + executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="test-cloud-ai-io") + monkeypatch.setattr(blocking_io, "_IO_EXECUTOR", executor) + monkeypatch.setattr(blocking_io, "_IO_MAX_WORKERS", 1) + try: + yield gate + finally: + gate.set() + executor.shutdown(wait=True) + monkeypatch.setattr(blocking_io, "_IO_EXECUTOR", original_executor) + monkeypatch.setattr(blocking_io, "_IO_MAX_WORKERS", original_limit) + + +@pytest.mark.asyncio +async def test_saturated_io_pool_does_not_starve_default_to_thread(isolated_pool): + entered = threading.Event() + + def blocker(): + entered.set() + isolated_pool.wait() + return "released" + + first = asyncio.create_task(blocking_io.run_blocking(blocker, timeout=0.05)) + await asyncio.wait_for(asyncio.to_thread(entered.wait, 1), timeout=1) + second = asyncio.create_task(blocking_io.run_blocking(blocker, timeout=0.05)) + + with pytest.raises(asyncio.TimeoutError): + await first + with pytest.raises(asyncio.TimeoutError): + await second + + result = await asyncio.wait_for(asyncio.to_thread(lambda: "default-ok"), timeout=0.5) + assert result == "default-ok" + + +@pytest.mark.asyncio +async def test_logs_warning_when_io_pool_is_at_capacity(isolated_pool, caplog): + entered = threading.Event() + + def blocker(): + entered.set() + isolated_pool.wait() + return "released" + + with caplog.at_level(logging.WARNING): + first = asyncio.create_task(blocking_io.run_blocking(blocker, timeout=0.05)) + await asyncio.wait_for(asyncio.to_thread(entered.wait, 1), timeout=1) + second = asyncio.create_task(blocking_io.run_blocking(blocker, timeout=0.05)) + + with pytest.raises(asyncio.TimeoutError): + await first + with pytest.raises(asyncio.TimeoutError): + await second + + assert any( + "blocking_io_pool_near_capacity" in record.getMessage() for record in caplog.records + ) diff --git a/tests/unit/test_google_cloud_provider.py b/tests/unit/test_google_cloud_provider.py index e25c5a17d..dd981b82a 100644 --- a/tests/unit/test_google_cloud_provider.py +++ b/tests/unit/test_google_cloud_provider.py @@ -1260,3 +1260,25 @@ async def test_missing_local_file_wrapped_in_cloud_ai_error(self, tmp_path, monk "offloading must preserve the underlying I/O error in the exception chain" ) assert "No such file or directory" in str(exc_info.value) + + async def test_local_file_read_routes_through_run_blocking(self, tmp_path, monkeypatch): + monkeypatch.setenv("CLOUD_AI_MEDIA_ROOT", str(tmp_path)) + provider = _make_provider() + provider._vision_client = self._client() + img_file = tmp_path / "frame.jpg" + img_file.write_bytes(b"disk-bytes") + mock_vision = self._vision_modules() + method_globals = type(provider).analyze_image.__globals__ + run_blocking = AsyncMock(return_value=b"worker-bytes") + + with ( + self._patched_modules(mock_vision), + patch.dict(method_globals, {"run_blocking": run_blocking}), + ): + await provider.analyze_image(str(img_file), [AnalysisType.LABEL_DETECTION]) + + assert mock_vision.Image.return_value.content == b"worker-bytes" + run_blocking.assert_awaited_once() + read_func, read_path = run_blocking.await_args.args + assert read_func is method_globals["_read_file_bytes"] + assert read_path == str(img_file)