coderabbit comments

Author: djsaunde

src/axolotl/utils/data/rl.py

@@ -69,8 +69,10 @@ def _load_datasets():
 
     # Prepare datasets (with file locking logic for multiple ranks)
     loader = FileLockLoader(cfg)
-    train_dataset, eval_dataset = loader.load(_load_datasets)
-    loader.cleanup()
+    try:
+        train_dataset, eval_dataset = loader.load(_load_datasets)
+    finally:
+        loader.cleanup()
 
     # Apply deduplication if configured
     if cfg.dataset_exact_deduplication:
@@ -187,10 +189,10 @@ def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:
     Returns:
         Combined and processed dataset for the specified split.
     """
-    datasets = cfg.datasets if split == "train" else cfg.test_datasets
+    datasets_configs = cfg.datasets if split == "train" else cfg.test_datasets
     split_datasets: list[Dataset | DatasetDict] = []
 
-    for dataset_config in datasets_with_name_generator(datasets):
+    for dataset_config in datasets_with_name_generator(datasets_configs):
         dataset: Dataset | DatasetDict = load_dataset_with_config(
             dataset_config, cfg.hf_use_auth_token, streaming=False
         )
@@ -199,7 +201,7 @@ def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:
     tokenizer = load_tokenizer(cfg)
 
     for i, data_set in enumerate(split_datasets):
-        _type = datasets[i]["type"]
+        _type = datasets_configs[i]["type"]
         if _type:
             if isinstance(_type, DictDefault):
                 _type = "user_defined.default"
@@ -246,7 +248,7 @@ def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:
     if not cfg.skip_prepare_dataset:
         # Save preprocessed dataset
         dataset_hash = generate_dataset_hash_from_config(
-            cfg, cfg.datasets, tokenizer.name_or_path
+            cfg, datasets_configs, tokenizer.name_or_path
         )
         save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
 

src/axolotl/utils/data/sft.py

@@ -99,8 +99,10 @@ def _load_datasets():
 
     # Prepare datasets (with file locking logic for multiple ranks)
     loader = FileLockLoader(cfg)
-    train_dataset, eval_dataset, prompters = loader.load(_load_datasets)
-    loader.cleanup()
+    try:
+        train_dataset, eval_dataset, prompters = loader.load(_load_datasets)
+    finally:
+        loader.cleanup()
 
     # Validate sample packing configuration for evaluation
     if eval_dataset and cfg.sample_packing and cfg.eval_sample_packing is not False:
@@ -265,11 +267,11 @@ def _load_tokenized_prepared_datasets(
         Tuple of (dataset, prompters list).
     """
     # Select correct dataset configuration based on split
-    datasets_config = cfg.datasets if split == "train" else cfg.test_datasets
+    datasets_configs = cfg.datasets if split == "train" else cfg.test_datasets
 
     # Generate dataset hash for caching
     dataset_hash = generate_dataset_hash_from_config(
-        cfg, datasets_config, tokenizer.name_or_path
+        cfg, datasets_configs, tokenizer.name_or_path
     )
 
     # Try loading from hub if push_dataset_to_hub is configured
@@ -286,7 +288,7 @@ def _load_tokenized_prepared_datasets(
     if dataset is None:
         dataset, prompters = _load_raw_datasets(
             cfg,
-            datasets_config,
+            datasets_configs,
             tokenizer,
             split,
             processor,
@@ -298,7 +300,7 @@ def _load_tokenized_prepared_datasets(
 
 def _load_raw_datasets(
     cfg: DictDefault,
-    cfg_datasets: list,
+    datasets_configs: list,
     tokenizer: PreTrainedTokenizer,
     split: str,
     processor: ProcessorMixin | None = None,
@@ -315,7 +317,7 @@ def _load_raw_datasets(
     # Load and process individual datasets
     datasets = []
     prompters = []
-    for dataset_config in datasets_with_name_generator(cfg_datasets):
+    for dataset_config in datasets_with_name_generator(datasets_configs):
         dataset_wrapper, dataset_prompter = _load_and_process_single_dataset(
             dataset_config=dataset_config,
             cfg=cfg,
@@ -338,7 +340,7 @@ def _load_raw_datasets(
 
         # Save the prepared dataset
         dataset_hash = generate_dataset_hash_from_config(
-            cfg, cfg.datasets, tokenizer.name_or_path
+            cfg, datasets_configs, tokenizer.name_or_path
         )
         save_preprocessed_dataset(cfg, dataset, dataset_hash, split)
 

tests/e2e/multigpu/test_locking.py

@@ -90,11 +90,7 @@ def worker():
     def test_load_waiting_for_ready_flag(self, mock_sleep, loader):
         """Test that processes wait for the ready flag to appear."""
         mock_load_fn = Mock(return_value="waiting_data")
-
-        # Create a mock path object with controllable exists() behavior
-        mock_ready_flag_path = Mock()
-
-        # Track exists() calls
+        mock_ready_flag_path = Path(tempfile.mktemp())
         exists_call_count = 0
 
         def mock_exists():