Fix duplicate download log messages in multi-process environment (#14299)

This commit is contained in:
alisonshao
2025-12-02 09:33:18 -08:00
committed by GitHub
parent df1f31241b
commit 25a6be4930
+22 -27
View File
@@ -322,12 +322,10 @@ def find_local_hf_snapshot_dir(
incomplete_files = glob.glob(os.path.join(blobs_dir, "*.incomplete")) incomplete_files = glob.glob(os.path.join(blobs_dir, "*.incomplete"))
if incomplete_files: if incomplete_files:
logger.info( log_info_on_rank0(
"Found %d .incomplete files in %s for %s. " logger,
"Will clean up and re-download.", f"Found {len(incomplete_files)} .incomplete files in {blobs_dir} for "
len(incomplete_files), f"{model_name_or_path}. Will clean up and re-download.",
blobs_dir,
model_name_or_path,
) )
_cleanup_corrupted_model_cache( _cleanup_corrupted_model_cache(
model_name_or_path, model_name_or_path,
@@ -367,22 +365,20 @@ def find_local_hf_snapshot_dir(
if not is_valid: if not is_valid:
if corrupted_files: if corrupted_files:
# Selective cleanup: only remove corrupted files # Selective cleanup: only remove corrupted files
logger.info( log_info_on_rank0(
"Found %d corrupted file(s) for %s: %s. " logger,
f"Found {len(corrupted_files)} corrupted file(s) for "
f"{model_name_or_path}: {error_msg}. "
"Will selectively clean and re-download only these files.", "Will selectively clean and re-download only these files.",
len(corrupted_files),
model_name_or_path,
error_msg,
) )
_cleanup_corrupted_files_selective(model_name_or_path, corrupted_files) _cleanup_corrupted_files_selective(model_name_or_path, corrupted_files)
return None return None
else: else:
# Cannot selectively clean (e.g., missing shards) - remove entire cache # Cannot selectively clean (e.g., missing shards) - remove entire cache
logger.info( log_info_on_rank0(
"Validation failed for %s: %s. " logger,
f"Validation failed for {model_name_or_path}: {error_msg}. "
"Will remove entire cache and re-download.", "Will remove entire cache and re-download.",
model_name_or_path,
error_msg,
) )
_cleanup_corrupted_model_cache( _cleanup_corrupted_model_cache(
model_name_or_path, found_local_snapshot_dir, error_msg model_name_or_path, found_local_snapshot_dir, error_msg
@@ -400,28 +396,27 @@ def find_local_hf_snapshot_dir(
"adapter_model.safetensors", "adapter_model.safetensors",
]: ]:
if not _validate_safetensors_file(f): if not _validate_safetensors_file(f):
logger.info( log_info_on_rank0(
"Corrupted model file %s for %s. " logger,
f"Corrupted model file {base_name} for {model_name_or_path}. "
"Will selectively clean and re-download this file.", "Will selectively clean and re-download this file.",
base_name,
model_name_or_path,
) )
# Selective cleanup for single file # Selective cleanup for single file
_cleanup_corrupted_files_selective(model_name_or_path, [f]) _cleanup_corrupted_files_selective(model_name_or_path, [f])
return None return None
if len(local_weight_files) > 0: if len(local_weight_files) > 0:
logger.info( log_info_on_rank0(
"Found local HF snapshot for %s at %s; skipping download.", logger,
model_name_or_path, f"Found local HF snapshot for {model_name_or_path} at "
found_local_snapshot_dir, f"{found_local_snapshot_dir}; skipping download.",
) )
return found_local_snapshot_dir return found_local_snapshot_dir
else: else:
logger.info( log_info_on_rank0(
"Local HF snapshot at %s has no files matching %s; will attempt download.", logger,
found_local_snapshot_dir, f"Local HF snapshot at {found_local_snapshot_dir} has no files matching "
allow_patterns, f"{allow_patterns}; will attempt download.",
) )
return None return None