-
Notifications
You must be signed in to change notification settings - Fork 79
feat(inference-logging-client): .log reader + multi-sink outputs (Spark / pandas / CSV / JSONL / text / analyze) #383
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: develop
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,10 +8,44 @@ | |
| import sys | ||
| import tempfile | ||
|
|
||
| from . import decode_mplog, format_dataframe_floats, get_format_name, get_mplog_metadata | ||
| from . import ( | ||
| analyze_log_file, | ||
| decode_log_file, | ||
| decode_log_file_to_csv, | ||
| decode_log_file_to_jsonl, | ||
| decode_log_file_to_pandas, | ||
| decode_logs, | ||
| decode_logs_to_pandas, | ||
| decode_mplog, | ||
| format_dataframe_floats, | ||
| get_format_name, | ||
| get_mplog_metadata, | ||
| list_log_sources, | ||
| print_analysis, | ||
| write_parsed_log, | ||
| ) | ||
| from .types import Format | ||
|
|
||
|
|
||
| def _is_log_input(path: str) -> bool: | ||
| """Return True when ``path`` should be routed through the .log reader | ||
| (file, directory, or gs:// URI/prefix).""" | ||
| if path.startswith("gs://"): | ||
| return True | ||
| if path.lower().endswith(".log"): | ||
| return True | ||
| if os.path.isdir(path): | ||
| return True | ||
| return False | ||
|
|
||
|
|
||
| def _is_multi_source(path: str) -> bool: | ||
| """True when the CLI input targets multiple .log files (directory or GCS prefix).""" | ||
| if path.startswith("gs://"): | ||
| return not path.lower().endswith(".log") | ||
| return os.path.isdir(path) | ||
|
|
||
|
|
||
| def main(): | ||
| """Main entry point with CLI interface.""" | ||
| parser = argparse.ArgumentParser( | ||
|
|
@@ -33,9 +67,43 @@ def main(): | |
| """, | ||
| ) | ||
|
|
||
| parser.add_argument("input", help="Input file containing MPLog bytes (or - for stdin)") | ||
| parser.add_argument("--model-proxy-id", "-m", required=True, help="Model proxy config ID") | ||
| parser.add_argument("--version", "-v", type=int, required=True, help="Schema version") | ||
| parser.add_argument( | ||
| "input", | ||
| help=( | ||
| "Input source: path to MPLog bytes file, path to an asyncloguploader " | ||
| ".log file, a gs://bucket/key URI, or '-' for stdin (MPLog bytes only)" | ||
| ), | ||
| ) | ||
| parser.add_argument( | ||
| "--model-proxy-id", | ||
| "-m", | ||
| default=None, | ||
| help="Model proxy config ID (required for raw MPLog input; ignored for .log/GCS)", | ||
| ) | ||
| parser.add_argument( | ||
| "--version", | ||
| "-v", | ||
| type=int, | ||
| default=None, | ||
| help="Schema version (required for raw MPLog input; ignored for .log/GCS)", | ||
| ) | ||
| parser.add_argument( | ||
| "--strict", | ||
| action="store_true", | ||
| help="For .log inputs: raise on any malformed frame instead of skipping", | ||
| ) | ||
| parser.add_argument( | ||
| "--output-format", | ||
| choices=["spark", "pandas", "csv", "jsonl", "text", "analyze"], | ||
| default=None, | ||
| help=( | ||
| "For .log / gs:// inputs, choose the output sink: 'spark' (default, " | ||
| "prints a Spark DataFrame), 'pandas' (CSV-style print of a pandas " | ||
| "DataFrame), 'csv' / 'jsonl' (write to --output), 'text' " | ||
| "(asynclogparser .parsed.log format, requires --output), or " | ||
| "'analyze' (structural report, no decoding)" | ||
| ), | ||
| ) | ||
| parser.add_argument( | ||
| "--format", | ||
| "-f", | ||
|
|
@@ -63,33 +131,143 @@ def main(): | |
|
|
||
| args = parser.parse_args() | ||
|
|
||
| # Read input | ||
| if args.input == "-": | ||
| data = sys.stdin.buffer.read() | ||
| else: | ||
| with open(args.input, "rb") as f: | ||
| data = f.read() | ||
| log_mode = _is_log_input(args.input) | ||
| out_fmt = args.output_format or ("spark" if log_mode else None) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
For |
||
| if out_fmt and not log_mode and out_fmt != "spark": | ||
| print( | ||
| f"Error: --output-format={out_fmt} only applies to .log / gs:// inputs", | ||
| file=sys.stderr, | ||
| ) | ||
| sys.exit(2) | ||
|
|
||
| inference_host = args.inference_host or os.getenv("INFERENCE_HOST", "http://localhost:8082") | ||
|
|
||
| # Decode input format | ||
| if args.hex: | ||
| # Non-Spark .log sinks short-circuit before Spark startup. | ||
| multi = _is_multi_source(args.input) | ||
| if log_mode and out_fmt in {"pandas", "csv", "jsonl", "text", "analyze"}: | ||
| try: | ||
| data = bytes.fromhex(data.decode("utf-8").strip()) | ||
| except ValueError as e: | ||
| print(f"Error: Invalid hex input: {e}", file=sys.stderr) | ||
| sys.exit(1) | ||
| elif args.base64: | ||
| try: | ||
| data = base64.b64decode(data) | ||
| if out_fmt == "analyze": | ||
| if multi: | ||
| print( | ||
| "Note: --output-format analyze inspects one file at a time; " | ||
| "listing sources under the given directory/prefix and picking " | ||
| "the first.", | ||
| file=sys.stderr, | ||
| ) | ||
| sources = list_log_sources(args.input) | ||
| if not sources: | ||
| print("Error: no .log files found under the given source", file=sys.stderr) | ||
| sys.exit(1) | ||
| print_analysis(analyze_log_file(sources[0])) | ||
| else: | ||
| print_analysis(analyze_log_file(args.input)) | ||
| return | ||
| if out_fmt == "pandas": | ||
| pdf = decode_logs_to_pandas( | ||
| args.input, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| strict=args.strict, | ||
| ) | ||
| if args.output: | ||
| pdf.to_csv(args.output, index=False) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Decoded log values are attacker-controlled input, and writing them directly to CSV allows CSV formula injection when the file is opened in spreadsheet software, potentially causing data exfiltration or command execution via external links/macros. Values beginning with Same issue also flagged at:
|
||
| print(f"Output written to {args.output} ({len(pdf)} rows)") | ||
| else: | ||
| print(pdf.to_csv(index=False)) | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 💡 SUGGESTION — Stream pandas CSV output to stdout instead of building one giant CSV string with When no output path is supplied, |
||
| return | ||
| if out_fmt in {"csv", "jsonl", "text"}: | ||
| if not args.output: | ||
| print( | ||
| f"Error: --output-format={out_fmt} requires --output PATH", | ||
| file=sys.stderr, | ||
| ) | ||
| sys.exit(2) | ||
| if multi: | ||
| # Fall through per-file to avoid ambiguous single-file writers. | ||
| sources = list_log_sources(args.input) | ||
| if not sources: | ||
| print("Error: no .log files found under the given source", file=sys.stderr) | ||
| sys.exit(1) | ||
| print( | ||
| f"Multi-source mode: writing one {out_fmt.upper()} per file into {args.output}/", | ||
| file=sys.stderr, | ||
| ) | ||
| os.makedirs(args.output, exist_ok=True) | ||
| sink = { | ||
| "csv": decode_log_file_to_csv, | ||
| "jsonl": decode_log_file_to_jsonl, | ||
| "text": write_parsed_log, | ||
| }[out_fmt] | ||
| total = 0 | ||
| ext = {"csv": ".csv", "jsonl": ".jsonl", "text": ".parsed.log"}[out_fmt] | ||
| for s in sources: | ||
| base = os.path.basename(s.rstrip("/")).rsplit(".log", 1)[0] + ext | ||
|
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Using only os.path.basename causes files with the same name in different directories or GCS prefixes to overwrite each other in the output directory. Also flagged on this line:
|
||
| dest = os.path.join(args.output, base) | ||
| n = sink( | ||
| s, | ||
| dest, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| strict=args.strict, | ||
| ) | ||
| total += n | ||
| print( | ||
| f"Wrote {len(sources)} files to {args.output}/ " | ||
| f"({total} {'records' if out_fmt=='text' else 'rows'} total)" | ||
| ) | ||
| return | ||
| sink = { | ||
| "csv": decode_log_file_to_csv, | ||
| "jsonl": decode_log_file_to_jsonl, | ||
| "text": write_parsed_log, | ||
| }[out_fmt] | ||
| n = sink( | ||
| args.input, | ||
| args.output, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| strict=args.strict, | ||
| ) | ||
| print(f"Output written to {args.output} ({n} {'records' if out_fmt=='text' else 'rows'})") | ||
| return | ||
| except Exception as e: | ||
| print(f"Error: Invalid base64 input: {e}", file=sys.stderr) | ||
| print(f"Error: {e}", file=sys.stderr) | ||
| sys.exit(1) | ||
|
|
||
| data: bytes = b"" | ||
|
|
||
| if not log_mode: | ||
| # Raw MPLog bytes path — requires schema identifiers. | ||
| if args.model_proxy_id is None or args.version is None: | ||
| print( | ||
| "Error: --model-proxy-id and --version are required when input is " | ||
| "MPLog bytes (anything other than a .log file or gs:// URI)", | ||
| file=sys.stderr, | ||
| ) | ||
| sys.exit(2) | ||
|
|
||
| if args.input == "-": | ||
| data = sys.stdin.buffer.read() | ||
| else: | ||
| with open(args.input, "rb") as f: | ||
| data = f.read() | ||
|
|
||
| if args.hex: | ||
| try: | ||
| data = bytes.fromhex(data.decode("utf-8").strip()) | ||
| except ValueError as e: | ||
| print(f"Error: Invalid hex input: {e}", file=sys.stderr) | ||
| sys.exit(1) | ||
| elif args.base64: | ||
| try: | ||
| data = base64.b64decode(data) | ||
| except Exception as e: | ||
| print(f"Error: Invalid base64 input: {e}", file=sys.stderr) | ||
| sys.exit(1) | ||
|
|
||
| # Parse format (None for auto-detection) | ||
| format_type = None if args.format == "auto" else Format(args.format) | ||
|
|
||
| # Get inference host from argument or environment variable | ||
| inference_host = args.inference_host or os.getenv("INFERENCE_HOST", "http://localhost:8082") | ||
|
|
||
| # Create SparkSession for CLI | ||
| from pyspark.sql import SparkSession | ||
|
|
||
|
|
@@ -104,16 +282,25 @@ def main(): | |
| spark.sparkContext.setLogLevel("ERROR") | ||
|
|
||
| try: | ||
| # Decode MPLog | ||
| df = decode_mplog( | ||
| log_data=data, | ||
| model_proxy_id=args.model_proxy_id, | ||
| version=args.version, | ||
| spark=spark, | ||
| format_type=format_type, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| ) | ||
| if log_mode: | ||
| # decode_logs auto-handles single file, local dir, or gs:// prefix. | ||
| df = decode_logs( | ||
| source=args.input, | ||
| spark=spark, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| strict=args.strict, | ||
| ) | ||
| else: | ||
| df = decode_mplog( | ||
| log_data=data, | ||
| model_proxy_id=args.model_proxy_id, | ||
| version=args.version, | ||
| spark=spark, | ||
| format_type=format_type, | ||
| inference_host=inference_host, | ||
| decompress=not args.no_decompress, | ||
| ) | ||
|
|
||
| # Format floats before output | ||
| df = format_dataframe_floats(df) | ||
|
|
@@ -152,22 +339,25 @@ def main(): | |
| # Show table (only fetches 20 rows, no full collect) | ||
| df.show(truncate=False) | ||
|
|
||
| # Get metadata for summary | ||
| metadata = get_mplog_metadata(data, decompress=not args.no_decompress) | ||
| detected_format_name = get_format_name(metadata.format_type) | ||
|
|
||
| # Print summary | ||
| print("\n--- Summary ---", file=sys.stderr) | ||
| print( | ||
| f"Format: {detected_format_name} (from metadata)" | ||
| if args.format == "auto" | ||
| else f"Format: {args.format}", | ||
| file=sys.stderr, | ||
| ) | ||
| print(f"Version: {metadata.version}", file=sys.stderr) | ||
| print( | ||
| f"Compression: {'enabled' if metadata.compression_enabled else 'disabled'}", file=sys.stderr | ||
| ) | ||
| if log_mode: | ||
| print(f"Source: {args.input} (.log mode)", file=sys.stderr) | ||
| else: | ||
| # Get metadata for summary (raw-MPLog mode only) | ||
| metadata = get_mplog_metadata(data, decompress=not args.no_decompress) | ||
| detected_format_name = get_format_name(metadata.format_type) | ||
| print( | ||
| f"Format: {detected_format_name} (from metadata)" | ||
| if args.format == "auto" | ||
| else f"Format: {args.format}", | ||
| file=sys.stderr, | ||
| ) | ||
| print(f"Version: {metadata.version}", file=sys.stderr) | ||
| print( | ||
| f"Compression: {'enabled' if metadata.compression_enabled else 'disabled'}", | ||
| file=sys.stderr, | ||
| ) | ||
| # Avoid full count() for huge DataFrames: use limit(1).count() for empty check only | ||
| try: | ||
| row_count = df.count() | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
_is_log_input treats any GCS URI as a framed .log file, so raw MPLog bytes stored in GCS cannot be decoded by the CLI and will be interpreted as a frame header. This is a behavior regression for users whose input path is remote but not a .log container.