Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,21 @@
decode_proto_features,
)
from .io import clear_schema_cache, get_feature_schema, get_mplog_metadata, parse_mplog_protobuf
from .log_reader 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,
iter_decoded_log_rows,
iter_log_records,
list_log_sources,
print_analysis,
read_log_file,
write_parsed_log,
)
from .types import FORMAT_TYPE_MAP, DecodedMPLog, FeatureInfo, Format
from .utils import format_dataframe_floats, get_format_name, unpack_metadata_byte

Expand All @@ -56,6 +71,19 @@
__all__ = [
"decode_mplog",
"decode_mplog_dataframe",
"decode_log_file",
"decode_log_file_to_pandas",
"decode_log_file_to_csv",
"decode_log_file_to_jsonl",
"decode_logs",
"decode_logs_to_pandas",
"list_log_sources",
"write_parsed_log",
"analyze_log_file",
"print_analysis",
"iter_decoded_log_rows",
"iter_log_records",
"read_log_file",
"get_mplog_metadata",
"get_feature_schema",
"clear_schema_cache",
Expand Down
284 changes: 237 additions & 47 deletions py-sdk/inference_logging_client/inference_logging_client/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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",
Expand Down Expand Up @@ -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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ WARNING — Do not auto-route every gs:// URI through the .log reader regardless of file type.

_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.

out_fmt = args.output_format or ("spark" if log_mode else None)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ WARNING — Defaulting .log CLI inputs to the Spark sink makes simple -o exports pay Spark startup and driver materialization instead of using the streaming CSV/JSONL sinks.

For .log mode, out_fmt becomes spark unless the user explicitly passes --output-format, so inference-logging-client file.log -o out.csv routes through decode_log_file, materializes all rows on the driver, starts Spark, and then writes. Since this PR adds direct CSV/JSONL sinks, the CLI should route file-output .log exports to those streaming sinks by default or require an explicit --output-format spark.

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)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ WARNING — Escape or prefix spreadsheet formula metacharacters before writing decoded log fields with to_csv.

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 =, +, -, @, tabs, or carriage returns should be neutralized for CSV outputs.

Same issue also flagged at:

  • py-sdk/inference_logging_client/inference_logging_client/log_reader.py:569

print(f"Output written to {args.output} ({len(pdf)} rows)")
else:
print(pdf.to_csv(index=False))

Copy link
Copy Markdown

Choose a reason for hiding this comment

The 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 to_csv(index=False).

When no output path is supplied, print(pdf.to_csv(index=False)) allocates the entire CSV representation as a Python string in addition to the already-materialized DataFrame. For large .log files this can double memory and stall until formatting completes. Use pdf.to_csv(sys.stdout, index=False) or require --output above a safe row threshold.

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

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ WARNING — Preserve relative path components or otherwise disambiguate output names when writing multi-source sinks.

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:

  • ⚠️ WARNING — Multi-source CLI output can overwrite files when different source paths share the same basename.

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

Expand All @@ -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)
Expand Down Expand Up @@ -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()
Expand Down
Loading
Loading