Skip to content

Commit 209fef1

Browse files
yl231yl231
committed
Merge main into run-batch-evaluation
Accept main branch version for llm_inference/parallel_inference.py which uses FileLock for cross-platform compatibility.
2 parents de7e389 + a94b93f commit 209fef1

1 file changed

Lines changed: 68 additions & 92 deletions

File tree

‎llm_inference/parallel_inference.py‎

Lines changed: 68 additions & 92 deletions
Original file line numberDiff line numberDiff line change
@@ -14,14 +14,17 @@
1414
import json
1515
import os
1616
import logging
17-
from typing import List, Dict, Any, Optional, Set
17+
from typing import List, Dict, Any, Set, Optional
1818
from pathlib import Path
1919
import threading
2020
from concurrent.futures import ThreadPoolExecutor, as_completed
21-
import fcntl
21+
from filelock import FileLock
2222

2323
from model_inference import ModelInference
2424

25+
# Thread-local storage for ModelInference instances
26+
_thread_local = threading.local()
27+
2528
logger = logging.getLogger(__name__)
2629

2730

@@ -110,45 +113,35 @@ def load_cached_run_counts(self, model: str) -> Dict[str, Set[int]]:
110113
if cache_file.exists():
111114
logger.info(f"Loading cached run counts from {cache_file}")
112115

113-
# Use file lock to ensure consistent read (prevents reading while writing)
114-
with self.file_lock:
116+
# Use cross-platform file lock to ensure consistent read (prevents reading while writing)
117+
lock_file = FileLock(str(cache_file) + ".lock")
118+
with self.file_lock: # Thread-level lock for this process
115119
try:
116-
with open(cache_file, "r", encoding="utf-8") as f:
117-
# Acquire file lock for reading (Unix/Linux)
118-
try:
119-
fcntl.flock(
120-
f.fileno(), fcntl.LOCK_SH
121-
) # Shared lock for reading
122-
except (AttributeError, OSError):
123-
# Windows or fcntl not available, rely on threading lock
124-
pass
125-
126-
for line in f:
127-
line = line.strip()
128-
if not line:
129-
continue
130-
try:
131-
entry = json.loads(line)
132-
global_index = entry.get("global_index")
133-
if not global_index:
120+
with lock_file: # Process-level lock across all processes
121+
with open(cache_file, "r", encoding="utf-8") as f:
122+
for line in f:
123+
line = line.strip()
124+
if not line:
134125
continue
126+
try:
127+
entry = json.loads(line)
128+
global_index = entry.get("global_index")
129+
if not global_index:
130+
continue
135131

136-
# Get run_number (default to 1 for backward compatibility)
137-
run_number = entry.get("run_number", 1)
138-
139-
# Only count successful entries
140-
if entry.get("success", False):
141-
if global_index not in run_counts:
142-
run_counts[global_index] = set()
143-
run_counts[global_index].add(run_number)
144-
except json.JSONDecodeError:
145-
logger.warning(f"Failed to parse line in {cache_file}")
146-
continue
132+
# Get run_number (default to 1 for backward compatibility)
133+
run_number = entry.get("run_number", 1)
147134

148-
try:
149-
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
150-
except (AttributeError, OSError):
151-
pass
135+
# Only count successful entries
136+
if entry.get("success", False):
137+
if global_index not in run_counts:
138+
run_counts[global_index] = set()
139+
run_counts[global_index].add(run_number)
140+
except json.JSONDecodeError:
141+
logger.warning(
142+
f"Failed to parse line in {cache_file}"
143+
)
144+
continue
152145
except Exception as e:
153146
logger.error(f"Error loading cached run counts: {e}")
154147

@@ -246,7 +239,10 @@ def filter_uninferred_data(
246239
def clear_failed_entries(self, model: str) -> int:
247240
"""
248241
Remove all failed entries from the cache file, keeping only successful ones.
249-
Called once at the beginning to clean up the cache.
242+
243+
This allows failed entries to be retried in subsequent runs.
244+
Called automatically before processing each model to ensure failed entries
245+
are retried.
250246
251247
Args:
252248
model: Model name
@@ -282,12 +278,15 @@ def clear_failed_entries(self, model: str) -> int:
282278

283279
# Write back only successful entries if any were removed
284280
if failed_count > 0:
285-
with self.file_lock:
281+
# Use cross-platform file lock for process-safe file writing
282+
lock_file = FileLock(str(cache_file) + ".lock")
283+
with self.file_lock: # Thread-level lock for this process
286284
try:
287-
with open(cache_file, "w", encoding="utf-8") as f:
288-
for entry in successful_entries:
289-
json.dump(entry, f, ensure_ascii=False)
290-
f.write("\n")
285+
with lock_file: # Process-level lock across all processes
286+
with open(cache_file, "w", encoding="utf-8") as f:
287+
for entry in successful_entries:
288+
json.dump(entry, f, ensure_ascii=False)
289+
f.write("\n")
291290
logger.info(
292291
f"Cleared {failed_count} failed entries from {cache_file}"
293292
)
@@ -359,24 +358,14 @@ def save_single_result(self, result: Dict[str, Any], model: str):
359358
model_filename = model.replace("/", "_")
360359
cache_file = self.cache_dir / f"{model_filename}.jsonl"
361360

362-
# Use lock to ensure thread-safe file writing
363-
with self.file_lock:
361+
# Use cross-platform file lock to ensure process-safe file writing
362+
lock_file = FileLock(str(cache_file) + ".lock")
363+
with self.file_lock: # Thread-level lock for this process
364364
try:
365-
with open(cache_file, "a", encoding="utf-8") as f:
366-
# Acquire file lock (Unix/Linux)
367-
try:
368-
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
369-
except (AttributeError, OSError):
370-
# Windows or fcntl not available, rely on threading lock
371-
pass
372-
373-
json.dump(result, f, ensure_ascii=False)
374-
f.write("\n")
375-
376-
try:
377-
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
378-
except (AttributeError, OSError):
379-
pass
365+
with lock_file: # Process-level lock across all processes
366+
with open(cache_file, "a", encoding="utf-8") as f:
367+
json.dump(result, f, ensure_ascii=False)
368+
f.write("\n")
380369
except Exception as e:
381370
logger.error(f"Error saving single result: {e}")
382371

@@ -434,26 +423,16 @@ def save_all_results(self, results: Dict[str, Dict[str, Any]], model: str):
434423
# Merge: new results override existing ones
435424
merged_results = {**existing_results, **results}
436425

437-
# Write all results
438-
with open(cache_file, "w", encoding="utf-8") as f:
439-
# Acquire file lock (Unix/Linux) if available
440-
try:
441-
fcntl.flock(f.fileno(), fcntl.LOCK_EX)
442-
except (AttributeError, OSError):
443-
# Windows or fcntl not available, rely on threading lock
444-
pass
445-
446-
for result in merged_results.values():
447-
# Ensure run_number is present before saving
448-
if "run_number" not in result:
449-
result["run_number"] = 1
450-
json.dump(result, f, ensure_ascii=False)
451-
f.write("\n")
452-
453-
try:
454-
fcntl.flock(f.fileno(), fcntl.LOCK_UN)
455-
except (AttributeError, OSError):
456-
pass
426+
# Write all results with cross-platform file lock
427+
lock_file = FileLock(str(cache_file) + ".lock")
428+
with lock_file: # Process-level lock across all processes
429+
with open(cache_file, "w", encoding="utf-8") as f:
430+
for result in merged_results.values():
431+
# Ensure run_number is present before saving
432+
if "run_number" not in result:
433+
result["run_number"] = 1
434+
json.dump(result, f, ensure_ascii=False)
435+
f.write("\n")
457436

458437
logger.info(f"Saved {len(merged_results)} results to {cache_file}")
459438
except Exception as e:
@@ -493,8 +472,10 @@ def _process_single_entry(
493472
f"Processing {index + 1}/{total} | Global ID: {global_index} | Run: {run_number}"
494473
)
495474

496-
# Create model inferencer for this thread
497-
model_inferencer = ModelInference()
475+
# Get or create model inferencer for this thread (reuse per thread to avoid overhead)
476+
if not hasattr(_thread_local, "model_inferencer"):
477+
_thread_local.model_inferencer = ModelInference()
478+
model_inferencer = _thread_local.model_inferencer
498479

499480
try:
500481
inference_result = model_inferencer.infer(model.replace("_", "/"), prompt)
@@ -589,8 +570,11 @@ def process_single_model(
589570
logger.info(f"Target runs per query: {num_runs}")
590571
logger.info(f"{'=' * 80}")
591572

592-
# Note: We don't clear failed entries when using multiple runs,
593-
# as we want to track all attempts
573+
# Clear failed entries before processing to allow retries
574+
# This ensures that failed entries can be retried in this run
575+
failed_cleared = self.clear_failed_entries(model)
576+
if failed_cleared > 0:
577+
logger.info(f"Cleared {failed_cleared} failed entries to allow retries")
594578

595579
# Filter data to determine which queries need inference and which run number
596580
filtered_data = self.filter_uninferred_data(data, model, num_runs)
@@ -668,14 +652,6 @@ def process_single_model(
668652
f"Progress: {completed_count}/{len(filtered_data)} | "
669653
f"Success: {successful_count} | Failed: {failed_count}"
670654
)
671-
# Periodically consolidate results (reload and save all)
672-
with self.results_lock:
673-
# Reload cache to get any results saved by other threads
674-
current_cache = self.load_existing_cache(model)
675-
# Merge: existing cache + our in-memory updates
676-
merged_cache = {**current_cache, **cached_results}
677-
cached_results.update(merged_cache)
678-
self.save_all_results(cached_results, model)
679655

680656
except Exception as e:
681657
logger.error(f"Error processing future: {e}")

0 commit comments

Comments
 (0)