FIX: add a request timeout to remote dataset downloads (#3154)

Co-authored-by: Roman Lutz <romanlutz13@gmail.com>

Author: praann07

pyrit/datasets/seed_datasets/remote/remote_dataset_loader.py

@@ -54,6 +54,9 @@ class _RemoteDatasetLoader(SeedDatasetProvider, ABC):
     and cached downloads.
     """
 
+    # Seconds to wait for a remote server before giving up, so a stalled download cannot hang forever.
+    REQUEST_TIMEOUT_SECONDS: ClassVar[int] = 60
+
     FILE_TYPE_HANDLERS: ClassVar[dict[str, dict[str, Callable[..., Any]]]] = {
         "json": {"read": read_json, "write": write_json},
         "jsonl": {"read": read_jsonl, "write": write_jsonl},
@@ -285,7 +288,7 @@ def _fetch_from_public_url(self, *, source: str, file_type: str) -> list[dict[st
             ValueError: If the file_type is invalid.
             Exception: If the request to fetch examples fails.
         """
-        response = requests.get(source)
+        response = requests.get(source, timeout=self.REQUEST_TIMEOUT_SECONDS)
         if response.status_code == 200:
             if file_type in self.FILE_TYPE_HANDLERS:
                 if file_type == "json":
@@ -528,7 +531,7 @@ def _download_and_parse() -> dict[str, list[dict[str, Any]]]:
                     temp_to_clean = zip_path
 
                     logger.info(f"Downloading zip archive from {source}")
-                    with requests.get(source, stream=True) as response:
+                    with requests.get(source, stream=True, timeout=self.REQUEST_TIMEOUT_SECONDS) as response:
                         response.raise_for_status()
                         with zip_path.open("wb") as fh:
                             for chunk in response.iter_content(chunk_size=1 << 16):

tests/unit/datasets/test_remote_dataset_loader.py

@@ -255,6 +255,15 @@ def test_fetch_from_public_url_non_json_file_type(self):
             result = loader._fetch_from_public_url(source="https://example.com/data.csv", file_type="csv")
         assert result == [{"header1": "value1", "header2": "value2"}]
 
+    def test_fetch_from_public_url_passes_request_timeout(self):
+        loader = ConcreteRemoteLoader()
+        mock_response = MagicMock()
+        mock_response.status_code = 200
+        mock_response.text = '[{"key": "value"}]'
+        with patch("requests.get", return_value=mock_response) as mock_get:
+            loader._fetch_from_public_url(source="https://example.com/data.json", file_type="json")
+        assert mock_get.call_args.kwargs["timeout"] == loader.REQUEST_TIMEOUT_SECONDS
+
     def test_fetch_from_public_url_invalid_file_type_raises(self):
         loader = ConcreteRemoteLoader()
         mock_response = MagicMock()
@@ -318,6 +327,23 @@ async def test_parses_multiple_inner_files(self, tmp_path, monkeypatch):
         assert result["folder/a.jsonl"] == [{"a": 1}, {"a": 2}]
         assert result["folder/b.jsonl"] == [{"b": 3}]
 
+    async def test_zip_download_passes_request_timeout(self, tmp_path, monkeypatch):
+        monkeypatch.setattr(
+            "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.DB_DATA_PATH",
+            tmp_path,
+        )
+        zip_bytes = self._make_zip_bytes({"x.json": '[{"k": "v"}]'})
+
+        mock_get = MagicMock(return_value=self._mock_streaming_response(zip_bytes))
+        with patch(
+            "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.requests.get",
+            mock_get,
+        ):
+            loader = ConcreteRemoteLoader()
+            await loader._fetch_zip_from_url_async(source=self.SOURCE, inner_files=["x.json"], cache=False)
+
+        assert mock_get.call_args.kwargs["timeout"] == loader.REQUEST_TIMEOUT_SECONDS
+
     async def test_caches_zip_on_disk(self, tmp_path, monkeypatch):
         monkeypatch.setattr(
             "pyrit.datasets.seed_datasets.remote.remote_dataset_loader.DB_DATA_PATH",