diff --git a/ais_bench/benchmark/utils/file/file.py b/ais_bench/benchmark/utils/file/file.py index f15a26c4..02062551 100644 --- a/ais_bench/benchmark/utils/file/file.py +++ b/ais_bench/benchmark/utils/file/file.py @@ -20,9 +20,29 @@ logger = AISLogger() + +import fcntl + + +def _lock_file(f): + """Blocking exclusive lock on an open file.""" + fcntl.lockf(f.fileno(), fcntl.LOCK_EX) + + +def _unlock_file(f): + """Release lock on an open file.""" + try: + fcntl.lockf(f.fileno(), fcntl.LOCK_UN) + except IOError as e: + logger.warning(f"Failed to release file lock: {e}") + + def write_status(file_path, status): """Write status to a JSON file, appending to existing content. + Uses file locking to safely handle concurrent writes from multiple + processes to the same file (read-modify-write is atomic). + Args: file_path: Path to the status file status: Status data to append @@ -30,32 +50,42 @@ def write_status(file_path, status): Returns: bool: True if successful, False otherwise """ - # read existing content - existing_data = [] - if os.path.exists(file_path): + # Ensure the file exists before opening in r+ mode + if not os.path.exists(file_path): try: - with open(file_path, 'r', encoding='utf-8') as f: - existing_data = json.load(f) - except json.JSONDecodeError as e: - logger.warning( - f"Failed to parse JSON from status file '{file_path}': {e}. " - f"Starting with empty status list." - ) - existing_data = [] - except IOError as e: - logger.warning( - f"Failed to read status file '{file_path}': {e}. " - f"Starting with empty status list." - ) - existing_data = [] - - # add new status - existing_data.append(status) + with open(file_path, 'w', encoding='utf-8') as f: + json.dump([], f) + except IOError: + pass - # write to file try: - with open(file_path, 'w', encoding='utf-8') as f: - json.dump(existing_data, f) + with open(file_path, 'r+', encoding='utf-8') as f: + _lock_file(f) + try: + # read existing content + content = f.read() + if content.strip(): + try: + existing_data = json.loads(content) + except json.JSONDecodeError as e: + logger.warning( + f"Failed to parse JSON from status file '{file_path}': {e}. " + f"Starting with empty status list." + ) + existing_data = [] + else: + existing_data = [] + + # add new status + existing_data.append(status) + + # write back + f.seek(0) + f.truncate() + json.dump(existing_data, f) + f.flush() + finally: + _unlock_file(f) return True except IOError as e: logger.warning(f"Failed to write status to '{file_path}': {e}") @@ -82,17 +112,29 @@ def read_and_clear_statuses(tmp_file_dir, tmp_file_name_list): logger.debug(f"Reading and clearing {len(abs_path_list)} status files from '{tmp_file_dir}'") for tmp_file in abs_path_list: - try: - # read existing content - with open(tmp_file, 'r', encoding='utf-8') as f: - data = json.load(f) - - status_count = len(data) - all_status.extend(data) + if not os.path.exists(tmp_file): + continue - # clear file content - with open(tmp_file, 'w', encoding='utf-8') as f: - json.dump([], f) + try: + with open(tmp_file, 'r+', encoding='utf-8') as f: + _lock_file(f) + try: + content = f.read() + if content.strip(): + data = json.loads(content) + else: + data = [] + + status_count = len(data) + all_status.extend(data) + + # clear file content + f.seek(0) + f.truncate() + json.dump([], f) + f.flush() + finally: + _unlock_file(f) logger.debug(f"Read {status_count} statuses from '{tmp_file}' and cleared file") diff --git a/tests/UT/utils/file/test_file.py b/tests/UT/utils/file/test_file.py index 3b33dce3..b5d5d967 100644 --- a/tests/UT/utils/file/test_file.py +++ b/tests/UT/utils/file/test_file.py @@ -120,11 +120,17 @@ def test_ioerror_on_clear_logs_warning(self): with open(p, "w", encoding="utf-8") as f: json.dump([{"k": 1}], f) - # Mock open to raise IOError when writing to the file + # Mock the file object's flush to raise IOError during the clear phase, + # since the new code uses a single r+ handle with truncate+flush. original_open = open - def mock_open(file_path, mode='r', *args, **kwargs): - if mode == 'w' and file_path == p: - raise IOError("Permission denied") + def mock_open(file_path, mode='r+', *args, **kwargs): + if file_path == p: + real_file = original_open(file_path, mode, *args, **kwargs) + original_flush = real_file.flush + def failing_flush(*a, **kw): + raise IOError("Permission denied") + real_file.flush = failing_flush + return real_file return original_open(file_path, mode, *args, **kwargs) with patch('builtins.open', side_effect=mock_open):