Add coderabbit suggestions

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

Author: fynnsu

scripts/data_generation_offline.py

@@ -171,21 +171,21 @@ def parse_args():
 def get_existing_hidden_state_indices(output_path: Path) -> list[int]:
     """Find existing `hs_i.safetensors` files (where i is the file index)"""
 
-    existing_file_indices = []
+    existing_file_indices_set: set[int] = set()
 
     if not output_path.exists():
-        return existing_file_indices
+        return []
 
     for file_path in output_path.iterdir():
         if file_path.name.startswith("hs_") and file_path.name.endswith(".safetensors"):
             index_str = file_path.stem[3:]  # Remove "hs_" prefix
             try:
                 file_index = int(index_str)
-                existing_file_indices.append(file_index)
+                existing_file_indices_set.add(file_index)
             except ValueError:
                 continue
 
-    return sorted(existing_file_indices)
+    return sorted(existing_file_indices_set)
 
 
 def get_indices_to_process(
@@ -248,7 +248,7 @@ def check_safetensors_file(path: Path, tokens: list[int]):
             )
 
 
-async def worker(
+async def worker(  # noqa: C901
     client,
     model: str,
     queue: "asyncio.Queue[dict[str, Any]]",

src/speculators/data_generation/vllm_client.py

@@ -137,9 +137,11 @@ async def wait_for_lock_async(lock_path, timeout=10.0, poll_interval=0.1):
     fd = os.open(lock_path, os.O_RDONLY)
     try:
         await asyncio.wait_for(_poll_lock_async(fd, poll_interval), timeout=timeout)
-    finally:
+    except BaseException:
         os.close(fd)
-        os.remove(lock_path)
+        raise
+    os.close(fd)
+    os.remove(lock_path)
 
 
 def wait_for_lock(lock_path, timeout=10.0, poll_interval=0.1):
@@ -149,16 +151,18 @@ def wait_for_lock(lock_path, timeout=10.0, poll_interval=0.1):
         while True:
             try:
                 fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
-                return
+                break
             except BlockingIOError:
                 if time.monotonic() >= deadline:
                     raise TimeoutError(
                         f"Timed out waiting for lock: {lock_path}"
                     ) from None
                 time.sleep(poll_interval)
-    finally:
+    except BaseException:
         os.close(fd)
-        os.remove(lock_path)
+        raise
+    os.close(fd)
+    os.remove(lock_path)
 
 
 @with_retries

src/speculators/train/data.py

@@ -285,19 +285,22 @@ def _maybe_generate_hs(self, index: int) -> dict[str, torch.Tensor] | None:
                 timeout=self.request_timeout,
                 max_retries=self.max_retries,
             )
-        except Exception as e:  # noqa: BLE001
-            warnings.warn(str(e), stacklevel=1)
-            return None
 
-        loaded_hs = load_file(hs_filepath)
+            loaded_hs = load_file(hs_filepath)
 
-        match self.on_generate:
-            case "cache":
-                file_idx = self._map_to_file_idx(index)
-                target_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
-                shutil.move(hs_filepath, target_path)
-            case "delete":
-                Path(hs_filepath).unlink()
+            match self.on_generate:
+                case "cache":
+                    file_idx = self._map_to_file_idx(index)
+                    target_path = self.hidden_states_path / f"hs_{file_idx}.safetensors"
+                    shutil.move(hs_filepath, target_path)
+                case "delete":
+                    Path(hs_filepath).unlink()
+        except Exception as e:  # noqa: BLE001
+            warnings.warn(
+                f"Failed to load/cache hidden states for sample {index}: {e}",
+                stacklevel=1,
+            )
+            return None
 
         return loaded_hs