fix parallel push_to_hub in dataset_dict (#7613)

* fix parallel push_to_hub in dataset_dict

* style

* fix

* fix

* fix

* fix

* last one

Author: lhoestq

.github/workflows/ci.yml

@@ -10,7 +10,7 @@ on:
       - ci-*
 
 env:
-  HF_ALLOW_CODE_EVAL: 1
+  CI_HEADERS: ${{ secrets.CI_HEADERS }}
 
 jobs:
 

setup.py

@@ -152,6 +152,8 @@
 ]
 
 TESTS_REQUIRE = [
+    # fix pip install issues for windows
+    "numba>=0.56.4",  # to get recent versions of llvmlite for windows ci
     # test dependencies
     "absl-py",
     "decorator",

src/datasets/arrow_dataset.py

@@ -5517,7 +5517,7 @@ def _push_parquet_shards_to_hub(
             total=num_shards,
             desc=desc,
         )
-        with contextlib.nullcontext() if num_proc is None and num_proc > 1 else Pool(num_proc) as pool:
+        with contextlib.nullcontext() if num_proc is None or num_proc <= 1 else Pool(num_proc) as pool:
             update_stream = (
                 Dataset._push_parquet_shards_to_hub_single(**kwargs_iterable[0])
                 if pool is None

tests/fixtures/hub.py

@@ -8,6 +8,7 @@
 import requests
 from huggingface_hub.hf_api import HfApi, RepositoryNotFoundError
 from huggingface_hub.utils import hf_raise_for_status
+from huggingface_hub.utils._headers import _http_user_agent
 
 
 CI_HUB_USER = "__DUMMY_TRANSFORMERS_USER__"
@@ -19,17 +20,18 @@
 CI_HFH_HUGGINGFACE_CO_URL_TEMPLATE = CI_HUB_ENDPOINT + "/{repo_id}/resolve/{revision}/{filename}"
 
 
-@pytest.fixture
-def ci_hfh_hf_hub_url(monkeypatch):
-    monkeypatch.setattr(
-        "huggingface_hub.file_download.HUGGINGFACE_CO_URL_TEMPLATE", CI_HFH_HUGGINGFACE_CO_URL_TEMPLATE
-    )
-
-
 @pytest.fixture
 def ci_hub_config(monkeypatch):
     monkeypatch.setattr("datasets.config.HF_ENDPOINT", CI_HUB_ENDPOINT)
     monkeypatch.setattr("datasets.config.HUB_DATASETS_URL", CI_HUB_DATASETS_URL)
+    monkeypatch.setattr(
+        "huggingface_hub.file_download.HUGGINGFACE_CO_URL_TEMPLATE", CI_HFH_HUGGINGFACE_CO_URL_TEMPLATE
+    )
+    old_environ = dict(os.environ)
+    os.environ["HF_ENDPOINT"] = CI_HUB_ENDPOINT
+    yield
+    os.environ.clear()
+    os.environ.update(old_environ)
 
 
 @pytest.fixture
@@ -38,6 +40,22 @@ def set_ci_hub_access_token(ci_hub_config, monkeypatch):
     monkeypatch.setattr("huggingface_hub.constants.HF_HUB_DISABLE_IMPLICIT_TOKEN", False)
     old_environ = dict(os.environ)
     os.environ["HF_TOKEN"] = CI_HUB_USER_TOKEN
+    os.environ["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0"
+    yield
+    os.environ.clear()
+    os.environ.update(old_environ)
+
+
+def _http_ci_user_agent(*args, **kwargs):
+    ua = _http_user_agent(*args, **kwargs)
+    return ua + os.environ.get("CI_HEADERS", "")
+
+
+@pytest.fixture(autouse=True)
+def set_hf_ci_headers(monkeypatch):
+    old_environ = dict(os.environ)
+    os.environ["TRANSFORMERS_IS_CI"] = "1"
+    monkeypatch.setattr("huggingface_hub.utils._headers._http_user_agent", _http_ci_user_agent)
     yield
     os.environ.clear()
     os.environ.update(old_environ)
@@ -105,7 +123,7 @@ def _hf_gated_dataset_repo_txt_data(hf_api: HfApi, hf_token, text_file_content):
 
 
 @pytest.fixture()
-def hf_gated_dataset_repo_txt_data(_hf_gated_dataset_repo_txt_data, ci_hub_config, ci_hfh_hf_hub_url):
+def hf_gated_dataset_repo_txt_data(_hf_gated_dataset_repo_txt_data, ci_hub_config):
     return _hf_gated_dataset_repo_txt_data
 
 
@@ -129,7 +147,7 @@ def hf_private_dataset_repo_txt_data_(hf_api: HfApi, hf_token, text_file_content
 
 
 @pytest.fixture()
-def hf_private_dataset_repo_txt_data(hf_private_dataset_repo_txt_data_, ci_hub_config, ci_hfh_hf_hub_url):
+def hf_private_dataset_repo_txt_data(hf_private_dataset_repo_txt_data_, ci_hub_config):
     return hf_private_dataset_repo_txt_data_
 
 
@@ -153,9 +171,7 @@ def hf_private_dataset_repo_zipped_txt_data_(hf_api: HfApi, hf_token, zip_csv_wi
 
 
 @pytest.fixture()
-def hf_private_dataset_repo_zipped_txt_data(
-    hf_private_dataset_repo_zipped_txt_data_, ci_hub_config, ci_hfh_hf_hub_url
-):
+def hf_private_dataset_repo_zipped_txt_data(hf_private_dataset_repo_zipped_txt_data_, ci_hub_config):
     return hf_private_dataset_repo_zipped_txt_data_
 
 
@@ -179,7 +195,5 @@ def hf_private_dataset_repo_zipped_img_data_(hf_api: HfApi, hf_token, zip_image_
 
 
 @pytest.fixture()
-def hf_private_dataset_repo_zipped_img_data(
-    hf_private_dataset_repo_zipped_img_data_, ci_hub_config, ci_hfh_hf_hub_url
-):
+def hf_private_dataset_repo_zipped_img_data(hf_private_dataset_repo_zipped_img_data_, ci_hub_config):
     return hf_private_dataset_repo_zipped_img_data_

tests/test_hub.py

@@ -45,7 +45,7 @@ def test_dataset_url(repo_id, filename, revision):
     assert url == f"https://huggingface.co/datasets/{repo_id}/resolve/{revision or 'main'}/{quote(filename)}"
 
 
-def test_delete_from_hub(temporary_repo, hf_api, hf_token, csv_path, ci_hub_config, ci_hfh_hf_hub_url) -> None:
+def test_delete_from_hub(temporary_repo, hf_api, hf_token, csv_path, ci_hub_config) -> None:
     with temporary_repo() as repo_id:
         hf_api.create_repo(repo_id, token=hf_token, repo_type="dataset")
         hf_api.upload_file(

tests/test_upstream_hub.py

@@ -45,7 +45,7 @@
 
 
 @for_all_test_methods(xfail_if_500_502_http_error)
-@pytest.mark.usefixtures("ci_hub_config", "ci_hfh_hf_hub_url")
+@pytest.mark.usefixtures("ci_hub_config")
 class TestPushToHub:
     _api = HfApi(endpoint=CI_HUB_ENDPOINT)
     _token = CI_HUB_USER_TOKEN
@@ -969,7 +969,7 @@ def text_file_with_metadata(request, tmp_path, text_file):
 
 
 @for_all_test_methods(xfail_if_500_502_http_error)
-@pytest.mark.usefixtures("ci_hub_config", "ci_hfh_hf_hub_url")
+@pytest.mark.usefixtures("ci_hub_config")
 class TestLoadFromHub:
     _api = HfApi(endpoint=CI_HUB_ENDPOINT)
     _token = CI_HUB_USER_TOKEN