[release/air] Fix air_example_gptj_deepspeed_fine_tuning.gce failing...

doc/source/ray-air/examples/gptj_deepspeed_fine_tuning.ipynb

@@ -223,7 +223,7 @@
     "            \"aws\",\n",
     "            \"s3\",\n",
     "            \"sync\",\n",
-    "            \"--quiet\",\n",
+    "            \"--no-sign-request\",\n",
     "            \"s3://large-dl-models-mirror/models--EleutherAI--gpt-j-6B/main/\",\n",
     "            os.path.join(path, \"snapshots\", \"main\"),\n",
     "        ]\n",

python/ray/train/huggingface/transformers/_transformers_utils.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Iterator, Optional, Tuple, Type
 
 import datasets.iterable_dataset
 import transformers.trainer
@@ -62,19 +62,22 @@ def get_train_dataloader(self):
 
 
 # TODO(ml-team): Replace with a Datasets-HuggingFace integration when available.
-class RayDatasetHFIterable(datasets.iterable_dataset.ExamplesIterable):
-    """HF ExamplesIterable backed by a Dataset."""
+class RayDatasetHFIterable(datasets.iterable_dataset._BaseExamplesIterable):
+    """HF ``_BaseExamplesIterable`` backed by a ``ray.data.DataIterator``.
+
+    The other abstract methods of shuffling and sharding the data are not implemented,
+    since those operations should be done by Ray Data. For example, the dataset
+    is already sharded to each data parallel worker and is disabled
+    (see ``wrap_transformers_trainer`` above).
+    """
 
     def __init__(self, dataset: DataIterator) -> None:
+        super().__init__()
         self.dataset = dataset
-        self.generate_examples_fn = self.dataset.iter_rows
-
-        # Required for the superclass
-        self.kwargs = {}
 
-    def __iter__(self):
-        for row in self.generate_examples_fn(**self.kwargs):
-            yield (0, {k: v for k, v in row.items()})
+    def __iter__(self) -> Iterator[Tuple[int, dict]]:
+        for idx, row in enumerate(self.dataset.iter_rows()):
+            yield (idx, {k: v for k, v in row.items()})
 
 
 def process_dataset_for_hf(

release/release_tests.yaml

@@ -810,7 +810,7 @@
     cluster_compute: gptj_deepspeed_compute_aws.yaml
 
   run:
-    timeout: 3600
+    timeout: 4500
     script: python test_myst_doc.py --path gptj_deepspeed_fine_tuning.ipynb
 
   variations: