Fix waiting for lock acquisition during online training (#566)

<!-- markdownlint-disable -->

PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT
THE BOTTOM) HAVE BEEN CONSIDERED.

## Purpose

#424 accidentally added lock waiting logic but only for the first load
attempt of the hidden states file. When that first load (intended for
pre-existing data) fails, the generate hidden states is called and new
hidden states are generated. Unfortunately the new hidden states were
loaded without waiting for lock.

<!--- Why your changes are needed -->

## Description

Make `_maybe_load_hs_file` a helper function so that it can be used by
both loads.

<!--- High-level concise summary of changes -->

## Related Issue

<!--- Link related issue if applicable -->

## Tests

Tested locally that this fixes these errors (which subsequently led to
training crashing):
```
speculators/src/speculators/train/data.py:300: UserWarning: Failed to load/cache hidden states for sample 29: No such file or directory:
/tmp/pytest-of-fynnsu/pytest-55/test_online_smoke_Qwen_Qwen3_V0/hidden_states/chatcmpl-b549e54beee9b9c2-93723e8e.safetensors
```

<!--- Please describe in detail how you tested your changes. -->

I have filled in:

- [x] The purpose of the PR, such as "Fix some issue (link existing
issues this PR will resolve)".
- [x] The test plan/results, such as providing test command and pasting
the results.
- [ ] (Optional) The necessary documentation update.
- [x] I (a human) have written or reviewed the code in this pr to the
best of my ability.

Signed-off-by: Fynn Schmitt-Ulms <fschmitt@redhat.com>

Author: fynnsu

src/speculators/train/data.py

@@ -178,6 +178,17 @@ def __getitem__(self, index) -> BatchType | None:
         return data
 
 
+def _maybe_load_hs_file(file_path: Path) -> dict[str, torch.Tensor] | None:
+    lock_path = str(file_path) + ".lock"
+    if Path(lock_path).exists():
+        wait_for_lock(lock_path)
+
+    if file_path.exists():
+        return load_file(file_path)
+
+    return None
+
+
 class ArrowDataset(BaseDataset):
     def __init__(
         self,
@@ -258,19 +269,6 @@ def _compute_approx_lengths(self) -> list[int]:
         """Get lengths of the dataset samples."""
         return list(self.data.with_format(None)["seq_len"])
 
-    def _maybe_load_hs_file(self, index: int) -> dict[str, torch.Tensor] | None:
-        file_idx = self._map_to_file_idx(index)
-        candidate_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
-
-        lock_path = str(candidate_path) + ".lock"
-        if Path(lock_path).exists():
-            wait_for_lock(lock_path)
-
-        if candidate_path.exists():
-            return load_file(candidate_path)
-
-        return None
-
     def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
         if not self.client:
             self._setup_client()
@@ -287,7 +285,7 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
                 max_retries=self.max_retries,
             )
 
-            loaded_hs = load_file(hs_filepath)
+            loaded_hs = _maybe_load_hs_file(Path(hs_filepath))
 
             match self.on_generate:
                 case "cache":
@@ -306,7 +304,9 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
         return loaded_hs
 
     def _get_raw_data(self, index):
-        loaded_hs = self._maybe_load_hs_file(index)
+        file_idx = self._map_to_file_idx(index)
+        candidate_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
+        loaded_hs = _maybe_load_hs_file(candidate_path)
 
         if loaded_hs is None:
             match self.on_missing: