diff --git a/tritonparse/common.py b/tritonparse/common.py index b925e91..b8a80ab 100644 --- a/tritonparse/common.py +++ b/tritonparse/common.py @@ -135,6 +135,60 @@ def to_rank(self) -> Rank: return Rank() +def print_parsed_files_summary(parsed_log_dir: str, out_dir: str = None) -> None: + """ + Print a beautiful summary of all parsed files. + + Args: + parsed_log_dir: Directory containing parsed files + """ + # Collect all parsed files + all_parsed_files = [] + for root, _, files in os.walk(parsed_log_dir): + for file in files: + file_path = os.path.join(root, file) + all_parsed_files.append(file_path) + + # Sort files for consistent output + all_parsed_files.sort() + + # Print beautiful summary + print("\n" + "=" * 80) + print("šŸ“ TRITONPARSE PARSING RESULTS") + print("=" * 80) + + # Print log file list (required for integration) + if out_dir: + print(f"šŸ“‚ Parsed files directory: {out_dir}") + else: + print(f"šŸ“‚ Parsed files directory: {parsed_log_dir}") + print(f"šŸ“Š Total files generated: {len(all_parsed_files)}") + + if all_parsed_files: + print("\nšŸ“„ Generated files:") + print("-" * 50) + for i, file_path in enumerate(all_parsed_files, 1): + # Get relative path for cleaner display + rel_path = os.path.relpath(file_path, parsed_log_dir) + file_size = "N/A" + try: + size_bytes = os.path.getsize(file_path) + if size_bytes < 1024: + file_size = f"{size_bytes}B" + elif size_bytes < 1024 * 1024: + file_size = f"{size_bytes/1024:.1f}KB" + else: + file_size = f"{size_bytes/(1024*1024):.1f}MB" + except OSError: + pass + + print(f" {i:2d}. šŸ“ {rel_path} ({file_size})") + + print("=" * 80) + print("āœ… Parsing completed successfully!") + print("=" * 80 + "\n") + + def gzip_single_file(file_path: str, verbose: bool = False) -> str: """ Gzip a single file and delete the original file. @@ -318,7 +372,9 @@ def parse_logs( log_file_list_path = os.path.join(parsed_log_dir, "log_file_list.json") with open(log_file_list_path, "w") as f: json.dump(file_mapping, f, indent=2) + # NOTICE: this print is required for tlparser-tritonparse integration + # DON'T REMOVE THIS PRINT print(f"tritonparse log file list: {log_file_list_path}") return parsed_log_dir, file_mapping diff --git a/tritonparse/utils.py b/tritonparse/utils.py index 18e24a1..d7bafcd 100644 --- a/tritonparse/utils.py +++ b/tritonparse/utils.py @@ -8,7 +8,7 @@ # argument parser for OSS parser = None -from .common import copy_local_to_tmpdir, is_fbcode, parse_logs, RankConfig, save_logs +from .common import copy_local_to_tmpdir, is_fbcode, parse_logs, RankConfig, save_logs, print_parsed_files_summary from .source_type import Source, SourceType @@ -92,7 +92,8 @@ def oss_parse(args): parsed_log_dir, _ = parse_logs(logs, rank_config, verbose) if args.out is not None: save_logs(Path(args.out), parsed_log_dir, args.overwrite, verbose) - + # Print beautiful summary of all parsed files + print_parsed_files_summary(parsed_log_dir, args.out) def unified_parse(args=None): """