From 52f61f32a3cb716d6b1ed79876e4a47fa9b3d758 Mon Sep 17 00:00:00 2001 From: dharma-shashank-meesho Date: Wed, 1 Jul 2026 12:17:37 +0530 Subject: [PATCH 1/2] feat(inference-logging-client): add .log reader + multi-sink outputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds first-class support for reading asyncloguploader .log containers directly (local disk, file-like, or gs://) into any of six output sinks — Spark, pandas, CSV, JSONL, human-readable text, or a structural analyze report — all backed by one shared decode path so row semantics are identical across sinks. - log_reader.py: new module with iter_log_records / iter_decoded_log_rows and six sink helpers (decode_log_file, decode_log_file_to_pandas, decode_log_file_to_csv/jsonl, write_parsed_log, analyze_log_file). - __init__.py: export the new public API. - cli.py: auto-route .log / gs:// inputs; add --output-format and --strict flags. - pyproject.toml: add [gcs] optional extra pulling google-cloud-storage. - readme.md: consolidated .log documentation into one Quick Start subsection + one API-reference block; added parity table with the asynclogparser standalone script. Verified end-to-end on a 128MB production .log (4,262 records, 9 frames): analyze / iter / pandas / CSV / JSONL / parsed-text sinks all decode correctly against schemas fetched from Horizon-v2. --- .../inference_logging_client/__init__.py | 22 + .../inference_logging_client/cli.py | 220 ++++- .../inference_logging_client/log_reader.py | 864 ++++++++++++++++++ .../inference_logging_client/pyproject.toml | 3 + py-sdk/inference_logging_client/readme.md | 241 ++++- 5 files changed, 1290 insertions(+), 60 deletions(-) create mode 100644 py-sdk/inference_logging_client/inference_logging_client/log_reader.py diff --git a/py-sdk/inference_logging_client/inference_logging_client/__init__.py b/py-sdk/inference_logging_client/inference_logging_client/__init__.py index 7efa8f2d..59100412 100644 --- a/py-sdk/inference_logging_client/inference_logging_client/__init__.py +++ b/py-sdk/inference_logging_client/inference_logging_client/__init__.py @@ -45,6 +45,18 @@ 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, + iter_decoded_log_rows, + iter_log_records, + 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 @@ -56,6 +68,16 @@ __all__ = [ "decode_mplog", "decode_mplog_dataframe", + "decode_log_file", + "decode_log_file_to_pandas", + "decode_log_file_to_csv", + "decode_log_file_to_jsonl", + "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", diff --git a/py-sdk/inference_logging_client/inference_logging_client/cli.py b/py-sdk/inference_logging_client/inference_logging_client/cli.py index a5a1ce93..a1c009c0 100644 --- a/py-sdk/inference_logging_client/inference_logging_client/cli.py +++ b/py-sdk/inference_logging_client/inference_logging_client/cli.py @@ -8,10 +8,29 @@ 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_mplog, + format_dataframe_floats, + get_format_name, + get_mplog_metadata, + 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.""" + if path.startswith("gs://"): + return True + return path.lower().endswith(".log") + + def main(): """Main entry point with CLI interface.""" parser = argparse.ArgumentParser( @@ -33,9 +52,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 +116,95 @@ 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) + 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) - # Decode input format - 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: + inference_host = args.inference_host or os.getenv("INFERENCE_HOST", "http://localhost:8082") + + # Non-Spark .log sinks short-circuit before Spark startup. + if log_mode and out_fmt in {"pandas", "csv", "jsonl", "text", "analyze"}: try: - data = base64.b64decode(data) + if out_fmt == "analyze": + print_analysis(analyze_log_file(args.input)) + return + if out_fmt == "pandas": + pdf = decode_log_file_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) + print(f"Output written to {args.output} ({len(pdf)} rows)") + else: + print(pdf.to_csv(index=False)) + 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) + 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 +219,24 @@ 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: + df = decode_log_file( + 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 +275,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() diff --git a/py-sdk/inference_logging_client/inference_logging_client/log_reader.py b/py-sdk/inference_logging_client/inference_logging_client/log_reader.py new file mode 100644 index 00000000..0d00311d --- /dev/null +++ b/py-sdk/inference_logging_client/inference_logging_client/log_reader.py @@ -0,0 +1,864 @@ +"""Read and decode .log files produced by asyncloguploader. + +Supports local file paths, ``gs://bucket/key`` GCS URIs, and open binary +file-like objects. Each .log file is a sequence of frames; each frame holds +length-prefixed records of ``[timestamp_ns][MPLog protobuf]``. + +Frame layout (all little-endian): + + ┌─────────────────────────────────────────────────────────────┐ + │ Frame header (8 bytes) │ + │ capacity (uint32) + valid_data_bytes (uint32) │ + ├─────────────────────────────────────────────────────────────┤ + │ Frame body (capacity - 8 bytes) │ + │ First valid_data_bytes contain records; rest is padding. │ + │ Record: [4B length][8B timestamp_ns][MPLog payload] │ + └─────────────────────────────────────────────────────────────┘ +""" + +from __future__ import annotations + +import os +import warnings +from contextlib import contextmanager +from pathlib import Path +from typing import IO, TYPE_CHECKING, Any, BinaryIO, Collection, Iterator, List, Optional, Tuple, Union + +from .exceptions import FormatError +from .formats import decode_arrow_features, decode_parquet_features, decode_proto_features +from .io import get_feature_schema, parse_mplog_protobuf +from .types import FORMAT_TYPE_MAP, FeatureInfo, Format + +if TYPE_CHECKING: + from pyspark.sql import DataFrame as SparkDataFrame + from pyspark.sql import SparkSession + + +_HEADER_SIZE = 8 +_LENGTH_PREFIX_SIZE = 4 +_TIMESTAMP_SIZE = 8 + +LogSource = Union[str, Path, BinaryIO, IO[bytes]] + + +def _is_gcs_uri(source: Any) -> bool: + return isinstance(source, str) and source.startswith("gs://") + + +def _parse_gcs_uri(uri: str) -> Tuple[str, str]: + """Split ``gs://bucket/key`` into ``(bucket, key)``.""" + if not uri.startswith("gs://"): + raise ValueError(f"Not a GCS URI: {uri!r}") + rest = uri[len("gs://") :] + if "/" not in rest: + raise ValueError(f"GCS URI missing object key: {uri!r}") + bucket, key = rest.split("/", 1) + if not bucket or not key: + raise ValueError(f"GCS URI must be gs:///: {uri!r}") + return bucket, key + + +@contextmanager +def _open_source(source: LogSource): + """Yield a binary, seekable-where-possible reader for ``source``. + + Local paths and GCS objects are opened here and closed on exit. File-like + objects are yielded as-is (caller retains ownership). + """ + if _is_gcs_uri(source): + try: + from google.cloud import storage # type: ignore + except ImportError as exc: + raise ImportError( + "Reading from gs:// requires the 'google-cloud-storage' package. " + "Install it with: pip install google-cloud-storage" + ) from exc + + bucket_name, blob_name = _parse_gcs_uri(source) # type: ignore[arg-type] + client = storage.Client() + blob = client.bucket(bucket_name).blob(blob_name) + handle = blob.open("rb") + try: + yield handle + finally: + handle.close() + return + + if isinstance(source, (str, Path)): + handle = open(source, "rb") + try: + yield handle + finally: + handle.close() + return + + # File-like + if hasattr(source, "read"): + yield source + return + + raise TypeError( + f"Unsupported source type: {type(source).__name__}. " + "Expected a path, gs:// URI, or binary file-like object." + ) + + +def _read_exact(stream: BinaryIO, n: int) -> bytes: + """Read exactly ``n`` bytes or return whatever was available (possibly empty).""" + chunks: List[bytes] = [] + remaining = n + while remaining > 0: + chunk = stream.read(remaining) + if not chunk: + break + chunks.append(chunk) + remaining -= len(chunk) + return b"".join(chunks) + + +def _iter_records_from_block(block: bytes, valid_data_bytes: int) -> Iterator[Tuple[int, bytes]]: + """Yield (timestamp_ns, mplog_payload) records from a single frame body.""" + offset = 0 + limit = min(valid_data_bytes, len(block)) + while offset + _LENGTH_PREFIX_SIZE <= limit: + record_length = int.from_bytes(block[offset : offset + _LENGTH_PREFIX_SIZE], "little") + offset += _LENGTH_PREFIX_SIZE + + if record_length == 0: + # Zero-length sentinel — keep scanning. + continue + if offset + record_length > limit: + # Record claims to extend past valid data; stop this frame. + break + if record_length < _TIMESTAMP_SIZE: + # Can't contain a timestamp; skip but keep moving. + offset += record_length + continue + + ts_ns = int.from_bytes(block[offset : offset + _TIMESTAMP_SIZE], "little") + payload = block[offset + _TIMESTAMP_SIZE : offset + record_length] + yield ts_ns, payload + offset += record_length + + +def iter_log_records( + source: LogSource, + strict: bool = False, +) -> Iterator[Tuple[int, bytes]]: + """Stream ``(timestamp_ns, mplog_payload_bytes)`` from a .log file. + + Args: + source: Local path, ``gs://bucket/key`` URI, or open binary file-like. + strict: If True, raise :class:`FormatError` on any malformed frame. + If False (default), warn and skip the offending frame, advancing by + ``capacity`` bytes when possible. Truncated trailing frames always + stop iteration. + + Yields: + Tuples of ``(timestamp_ns, mplog_payload_bytes)`` in file order. The + payload is the raw MPLog protobuf — feed it to + :func:`get_mplog_metadata` or the decode functions to extract features. + """ + with _open_source(source) as stream: + frame_index = 0 + while True: + header = _read_exact(stream, _HEADER_SIZE) + if len(header) == 0: + return + if len(header) < _HEADER_SIZE: + if strict: + raise FormatError( + f"Truncated frame header at frame {frame_index} " + f"(read {len(header)} of {_HEADER_SIZE} bytes)" + ) + warnings.warn( + f"Truncated frame header at frame {frame_index}; stopping", + UserWarning, + stacklevel=2, + ) + return + + capacity = int.from_bytes(header[0:4], "little") + valid_data_bytes = int.from_bytes(header[4:8], "little") + + if capacity < _HEADER_SIZE: + msg = ( + f"Frame {frame_index}: capacity {capacity} is smaller than the " + f"{_HEADER_SIZE}-byte header" + ) + if strict: + raise FormatError(msg) + warnings.warn(msg + "; stopping", UserWarning, stacklevel=2) + return + + data_size = capacity - _HEADER_SIZE + + if valid_data_bytes > capacity: + msg = ( + f"Frame {frame_index}: valid_data_bytes ({valid_data_bytes}) > " + f"capacity ({capacity})" + ) + if strict: + raise FormatError(msg) + warnings.warn(msg + "; skipping frame", UserWarning, stacklevel=2) + # Best-effort: advance by data_size and continue. + _read_exact(stream, data_size) + frame_index += 1 + continue + + if valid_data_bytes == 0: + # Empty frame — still consume the padding. + if data_size > 0: + _read_exact(stream, data_size) + frame_index += 1 + continue + + block = _read_exact(stream, data_size) + if len(block) < data_size: + # Truncated trailing frame. + msg = ( + f"Frame {frame_index}: expected {data_size} body bytes, " + f"got {len(block)} (EOF)" + ) + if strict: + raise FormatError(msg) + if len(block) < valid_data_bytes: + warnings.warn(msg + "; stopping", UserWarning, stacklevel=2) + return + # We have enough for valid_data_bytes — fall through. + + yield from _iter_records_from_block(block, valid_data_bytes) + frame_index += 1 + + +def read_log_file( + source: LogSource, + strict: bool = False, +) -> List[Tuple[int, bytes]]: + """Eagerly materialize all records from a .log file. + + Convenience wrapper around :func:`iter_log_records`. For large files + prefer the iterator form to avoid loading every record into memory. + """ + return list(iter_log_records(source, strict=strict)) + + +# Columns we attach per record (in addition to entity_id and feature cols). +_RECORD_METADATA_COLUMNS = ( + "timestamp_ns", + "mp_config_id", + "version", + "format_type", + "user_id", + "tracking_id", + "parent_entity", +) + + +def _decode_one_record( + payload: bytes, + timestamp_ns: int, + inference_host: Optional[str], + decompress: bool, + schema_override: Optional[List[FeatureInfo]], + needed_columns: Optional[Collection[str]], + schema_cache: dict, + stringify_values: bool, +) -> Iterator[dict]: + """Decode a single MPLog payload into per-entity dict rows.""" + # Optional zstd decompression of the outer payload. + working = payload + if decompress and len(working) >= 4 and working[:4] == b"\x28\xb5\x2f\xfd": + try: + import zstandard as zstd # type: ignore + except ImportError as exc: + raise ImportError( + "MPLog payload is zstd-compressed but the 'zstandard' package is not " + "installed. Install it with: pip install zstandard" + ) from exc + working = zstd.ZstdDecompressor().decompress(working) + + parsed = parse_mplog_protobuf(working) + encoded_per_entity: List[bytes] = list(getattr(parsed, "_encoded_features", []) or []) + if not encoded_per_entity: + return + + mp_config_id = parsed.model_proxy_config_id or "" + version = int(parsed.version or 0) + + # Resolve schema: explicit override > cache > fetch. + if schema_override is not None: + schema = schema_override + else: + cache_key = (mp_config_id, version) + schema = schema_cache.get(cache_key) + if schema is None: + if not mp_config_id: + return + try: + schema = get_feature_schema(mp_config_id, version, inference_host) + except Exception as exc: + warnings.warn( + f"Failed to fetch schema for {cache_key}: {exc}", + UserWarning, + stacklevel=2, + ) + return + schema_cache[cache_key] = schema + + fmt = FORMAT_TYPE_MAP.get(parsed.format_type, Format.PROTO) + if fmt == Format.ARROW: + decoder = decode_arrow_features + elif fmt == Format.PARQUET: + decoder = decode_parquet_features + else: + decoder = decode_proto_features + + entities = list(parsed.entities or []) + parents = list(parsed.parent_entity or []) + + base_row = { + "timestamp_ns": timestamp_ns, + "mp_config_id": mp_config_id, + "version": version, + "format_type": int(parsed.format_type or 0), + "user_id": parsed.user_id or "", + "tracking_id": parsed.tracking_id or "", + } + + for i, enc in enumerate(encoded_per_entity): + entity_id = str(entities[i]) if i < len(entities) else f"entity_{i}" + try: + decoded = decoder(enc, schema, needed_columns=needed_columns) + except Exception as exc: + warnings.warn( + f"Decode failed for {mp_config_id}@v{version} entity#{i}: {exc}", + UserWarning, + stacklevel=2, + ) + continue + + row = dict(base_row) + row["entity_id"] = entity_id + row["parent_entity"] = str(parents[i]) if i < len(parents) else "" + for name, value in decoded.items(): + if not stringify_values: + row[name] = value + continue + # Spark mode: stringify complex values so the inferred schema is stable. + if value is None: + row[name] = None + elif isinstance(value, (list, tuple)): + row[name] = str(list(value)) + elif isinstance(value, bytes): + row[name] = value.hex() + else: + row[name] = value + yield row + + +def iter_decoded_log_rows( + source: LogSource, + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, + stringify_values: bool = False, +) -> Iterator[dict]: + """Stream fully-decoded per-entity rows from a .log file. + + This is the single shared decode path used by every output sink + (Spark, pandas, CSV, JSONL, text). Use it directly when you want to + pipe rows somewhere custom without materializing the whole file. + + Args: + source: Local path, ``gs://bucket/key`` URI, or open binary file-like. + inference_host: Inference service URL. Defaults to ``INFERENCE_HOST`` env. + decompress: Attempt zstd decompression on each MPLog payload. + schema: Pre-fetched schema; bypasses per-record schema fetch. + needed_columns: Optional subset of feature names to keep. + strict: Propagated to :func:`iter_log_records`. + stringify_values: If True, coerce list/tuple/bytes feature values to + strings (used by the Spark sink to keep its inferred schema stable). + + Yields: + Dict rows with keys ``entity_id, timestamp_ns, mp_config_id, version, + format_type, user_id, tracking_id, parent_entity`` and one entry per + decoded feature. + """ + if inference_host is None: + inference_host = os.getenv("INFERENCE_HOST", "http://localhost:8082") + + schema_cache: dict = {} + for ts_ns, payload in iter_log_records(source, strict=strict): + yield from _decode_one_record( + payload, + ts_ns, + inference_host=inference_host, + decompress=decompress, + schema_override=schema, + needed_columns=needed_columns, + schema_cache=schema_cache, + stringify_values=stringify_values, + ) + + +def decode_log_file( + source: LogSource, + spark: "SparkSession", + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, +) -> "SparkDataFrame": + """Decode an asyncloguploader .log file into a Spark DataFrame. + + Args: + source: Local path, ``gs://bucket/key`` GCS URI, or open binary file-like. + spark: SparkSession used to materialize the result DataFrame. + inference_host: Inference service base URL. If ``None``, reads from the + ``INFERENCE_HOST`` env var (falls back to ``http://localhost:8082``). + decompress: Attempt zstd decompression of MPLog payloads. + schema: Optional pre-fetched schema. When set, schema lookups are + skipped (assumes a single schema across all records). + needed_columns: Optional subset of feature names to keep. + strict: If True, raise :class:`FormatError` on any malformed frame. + + Returns: + A DataFrame with one row per ``(record, entity)``. Columns are + ``entity_id`` followed by record metadata + (``timestamp_ns``, ``mp_config_id``, ``version``, ``format_type``, + ``user_id``, ``tracking_id``, ``parent_entity``) and then the decoded + feature columns. + """ + rows: List[dict] = list( + iter_decoded_log_rows( + source, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + stringify_values=True, + ) + ) + + if not rows: + from pyspark.sql.types import LongType, StringType, StructField, StructType + + fields = [StructField("entity_id", StringType(), True)] + for col in _RECORD_METADATA_COLUMNS: + dtype = LongType() if col in ("timestamp_ns", "version", "format_type") else StringType() + fields.append(StructField(col, dtype, True)) + return spark.createDataFrame([], StructType(fields)) + + # Build ordered columns: entity_id, record metadata (in fixed order), then features. + feature_names: List[str] = [] + seen = set(("entity_id",) + _RECORD_METADATA_COLUMNS) + for row in rows: + for k in row.keys(): + if k not in seen: + feature_names.append(k) + seen.add(k) + + ordered_cols = ["entity_id", *_RECORD_METADATA_COLUMNS, *feature_names] + # Backfill so every row has every column (Spark needs uniform dict keys). + for row in rows: + for col in ordered_cols: + row.setdefault(col, None) + + df = spark.createDataFrame(rows) + return df.select(ordered_cols) + + +def _ordered_columns(rows: List[dict]) -> List[str]: + """Return ``entity_id, , ...`` in stable order.""" + feature_names: List[str] = [] + seen = set(("entity_id",) + _RECORD_METADATA_COLUMNS) + for row in rows: + for k in row.keys(): + if k not in seen: + feature_names.append(k) + seen.add(k) + return ["entity_id", *_RECORD_METADATA_COLUMNS, *feature_names] + + +def decode_log_file_to_pandas( + source: LogSource, + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, +): + """Decode a .log file straight into a pandas DataFrame. + + Same row shape as :func:`decode_log_file` but skips Spark entirely. + Best for files that fit comfortably in driver memory. + + Requires ``pandas`` to be importable (transitively available via the + ``pyarrow`` dependency, but not declared explicitly — install it if your + environment lacks it). + """ + try: + import pandas as pd # type: ignore + except ImportError as exc: + raise ImportError( + "decode_log_file_to_pandas requires pandas. Install it with: pip install pandas" + ) from exc + + rows = list( + iter_decoded_log_rows( + source, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + stringify_values=False, + ) + ) + if not rows: + return pd.DataFrame(columns=["entity_id", *_RECORD_METADATA_COLUMNS]) + + cols = _ordered_columns(rows) + return pd.DataFrame(rows, columns=cols) + + +def decode_log_file_to_csv( + source: LogSource, + output_path: Union[str, Path], + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, +) -> int: + """Stream a .log file into a single CSV file. Returns rows written. + + Uses two passes: first collects column names by buffering rows, then writes + a single CSV with a stable header. Memory cost is bounded by the rows in + one file (typical asyncloguploader files yield well under 1M rows). + """ + import csv + + rows = list( + iter_decoded_log_rows( + source, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + stringify_values=True, + ) + ) + output_path = str(output_path) + if not rows: + with open(output_path, "w", newline="") as fh: + csv.writer(fh).writerow(["entity_id", *_RECORD_METADATA_COLUMNS]) + return 0 + + cols = _ordered_columns(rows) + with open(output_path, "w", newline="") as fh: + writer = csv.DictWriter(fh, fieldnames=cols) + writer.writeheader() + for row in rows: + writer.writerow({c: row.get(c) for c in cols}) + return len(rows) + + +def decode_log_file_to_jsonl( + source: LogSource, + output_path: Union[str, Path], + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, +) -> int: + """Stream a .log file into newline-delimited JSON. Returns rows written. + + Unlike CSV, rows are written as they are produced — memory cost is one row + at a time. + """ + import json + + output_path = str(output_path) + count = 0 + with open(output_path, "w") as fh: + for row in iter_decoded_log_rows( + source, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + stringify_values=False, + ): + # bytes values -> hex hex strings so json can serialize them + serializable = { + k: (v.hex() if isinstance(v, (bytes, bytearray)) else v) for k, v in row.items() + } + fh.write(json.dumps(serializable, default=str)) + fh.write("\n") + count += 1 + return count + + +def write_parsed_log( + source: LogSource, + output_path: Union[str, Path], + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, +) -> int: + """Write a human-readable parsed log (asynclogparser ``.parsed.log`` format). + + The output mirrors the layout produced by ``asynclogparse.py``:: + + === Record N === + Timestamp: (raw: ns) + User ID: ... + Tracking ID: ... + Model Config ID: ... + Schema Version: + Format Type: (proto|arrow|parquet) + Entities: + ... + --- Entity 1 --- + Entity ID: ... + Features: + : + + Returns the number of records written. + """ + import datetime as _dt + + output_path = str(output_path) + schema_cache: dict = {} + + if inference_host is None: + inference_host = os.getenv("INFERENCE_HOST", "http://localhost:8082") + + record_count = 0 + with open(output_path, "w") as out: + for ts_ns, payload in iter_log_records(source, strict=strict): + record_count += 1 + + # Re-run the decode for this single record (so the output groups + # entities by record, like asynclogparser does). + working = payload + if decompress and len(working) >= 4 and working[:4] == b"\x28\xb5\x2f\xfd": + try: + import zstandard as zstd # type: ignore + working = zstd.ZstdDecompressor().decompress(working) + except ImportError: + pass + + try: + parsed = parse_mplog_protobuf(working) + except Exception as e: + out.write(f"=== Record {record_count}: PARSE ERROR ===\nError: {e}\n\n") + continue + + out.write(f"=== Record {record_count} ===\n") + if ts_ns: + ts = _dt.datetime.fromtimestamp(ts_ns / 1e9, tz=_dt.timezone.utc) + out.write( + f"Timestamp: {ts.strftime('%Y-%m-%d %H:%M:%S.%f')} UTC (raw: {ts_ns} ns)\n" + ) + out.write(f"User ID: {parsed.user_id}\n") + out.write(f"Tracking ID: {parsed.tracking_id}\n") + out.write(f"Model Config ID: {parsed.model_proxy_config_id}\n") + out.write(f"Schema Version: {parsed.version}\n") + fmt_name = {0: "proto", 1: "arrow", 2: "parquet"}.get( + parsed.format_type, f"unknown({parsed.format_type})" + ) + out.write(f"Format Type: {parsed.format_type} ({fmt_name})\n") + entities = list(parsed.entities or []) + parents = list(parsed.parent_entity or []) + encoded = list(getattr(parsed, "_encoded_features", []) or []) + out.write(f"Entities: {len(entities)}\n") + if entities: + out.write(f"Entity IDs: {entities}\n") + out.write(f"Parent Entities: {len(parents)}\n") + if parents: + out.write(f"Parent Entity IDs: {parents}\n") + out.write(f"Encoded Features: {len(encoded)}\n\n") + + # Resolve schema once per record. + if schema is not None: + feat_schema = schema + else: + key = (parsed.model_proxy_config_id or "", int(parsed.version or 0)) + feat_schema = schema_cache.get(key) + if feat_schema is None and key[0]: + try: + feat_schema = get_feature_schema(key[0], key[1], inference_host) + schema_cache[key] = feat_schema + except Exception as e: + warnings.warn( + f"Schema fetch failed for {key}: {e}", UserWarning, stacklevel=2 + ) + feat_schema = None + + fmt = FORMAT_TYPE_MAP.get(parsed.format_type, Format.PROTO) + decoder = { + Format.ARROW: decode_arrow_features, + Format.PARQUET: decode_parquet_features, + }.get(fmt, decode_proto_features) + + for i, enc in enumerate(encoded): + out.write(f"--- Entity {i+1} ---\n") + eid = entities[i] if i < len(entities) else "" + if eid: + out.write(f"Entity ID: {eid}\n") + pe = parents[i] if i < len(parents) else "" + if pe: + out.write(f"Parent Entity: {pe}\n") + out.write("Features:\n") + if feat_schema is None: + out.write(f" (encoded_features length: {len(enc)} bytes — no schema)\n") + else: + try: + decoded = decoder(enc, feat_schema, needed_columns=needed_columns) + for k, v in decoded.items(): + out.write(f" {k}: {v}\n") + except Exception as e: + out.write(f" (decode error: {e})\n") + out.write("\n") + out.write("\n") + + return record_count + + +def analyze_log_file(source: LogSource, sample_records: int = 5) -> dict: + """Structural analysis of a .log file. No protobuf decoding, no schema fetch. + + Mirrors ``asynclogparser`` ``--analyze``: returns frame-by-frame structure, + record counts, data utilization, and timestamp-prefix detection. + + Returns a dict with keys: ``file_bytes``, ``frames`` (list), ``totals``, + ``timestamp_format`` (``True`` / ``False`` / ``None``). + """ + import datetime as _dt + + frames: List[dict] = [] + total_valid = 0 + total_records = 0 + sample: List[bytes] = [] + + with _open_source(source) as stream: + idx = 0 + while True: + header = _read_exact(stream, _HEADER_SIZE) + if len(header) == 0: + break + if len(header) < _HEADER_SIZE: + frames.append({"index": idx, "error": "truncated header"}) + break + + capacity = int.from_bytes(header[0:4], "little") + valid = int.from_bytes(header[4:8], "little") + info = {"index": idx, "capacity": capacity, "valid_data_bytes": valid, "records": 0} + + if capacity < _HEADER_SIZE or valid > capacity: + info["error"] = f"invalid (cap={capacity}, valid={valid})" + frames.append(info) + break + + data = _read_exact(stream, capacity - _HEADER_SIZE) + if len(data) < capacity - _HEADER_SIZE: + info["error"] = f"truncated body ({len(data)}/{capacity - _HEADER_SIZE})" + frames.append(info) + break + + off = 0 + limit = min(valid, len(data)) + while off + _LENGTH_PREFIX_SIZE <= limit: + rec_len = int.from_bytes(data[off : off + _LENGTH_PREFIX_SIZE], "little") + off += _LENGTH_PREFIX_SIZE + if rec_len == 0: + continue + if off + rec_len > limit: + break + info["records"] += 1 + if len(sample) < sample_records: + sample.append(data[off : off + rec_len]) + off += rec_len + + total_records += info["records"] + total_valid += valid + frames.append(info) + idx += 1 + + # Best-effort total file size (only works for seekable streams). + file_bytes: Optional[int] = None + try: + with _open_source(source) as s2: + s2.seek(0, 2) + file_bytes = s2.tell() + except Exception: + pass + + # Detect whether records carry an 8-byte UnixNano prefix (post-2024 format). + ts_format: Optional[bool] = None + if sample: + ts_format = True + for rec in sample: + if len(rec) < _TIMESTAMP_SIZE: + ts_format = False + break + ts_val = int.from_bytes(rec[:_TIMESTAMP_SIZE], "little") + # Reasonable UnixNano range (~2020-2033). + if not (1_600_000_000_000_000_000 < ts_val < 2_000_000_000_000_000_000): + ts_format = False + break + + return { + "file_bytes": file_bytes, + "frames": frames, + "totals": { + "frames": len(frames), + "valid_bytes": total_valid, + "records": total_records, + "utilization": (total_valid / file_bytes) if file_bytes else None, + }, + "timestamp_format": ts_format, + } + + +def print_analysis(result: dict) -> None: + """Pretty-print the output of :func:`analyze_log_file`.""" + fb = result.get("file_bytes") + print("=" * 60) + print(f"FILE ANALYSIS ({fb:,} bytes / {fb/(1024*1024):.2f} MB)" if fb else "FILE ANALYSIS") + print("=" * 60) + t = result["totals"] + print(f" frames: {t['frames']}, records: {t['records']:,}, valid: {t['valid_bytes']:,} bytes") + if t.get("utilization") is not None: + print(f" utilization: {t['utilization']*100:.1f}%") + print() + print("PER-FRAME:") + for fi in result["frames"]: + if "error" in fi: + print(f" frame {fi['index']:3d} ERROR: {fi['error']}") + continue + cap_mb = fi["capacity"] / (1024 * 1024) + valid_mb = fi["valid_data_bytes"] / (1024 * 1024) + pct = (fi["valid_data_bytes"] / fi["capacity"] * 100) if fi["capacity"] else 0 + print( + f" frame {fi['index']:3d} cap={cap_mb:7.2f}MB valid={valid_mb:7.2f}MB " + f"({pct:5.1f}%) records={fi['records']:,}" + ) + print() + tsf = result.get("timestamp_format") + if tsf is True: + print(" records carry an 8-byte UnixNano timestamp prefix.") + elif tsf is False: + print(" WARNING: records do NOT carry the expected timestamp prefix.") diff --git a/py-sdk/inference_logging_client/pyproject.toml b/py-sdk/inference_logging_client/pyproject.toml index 13366434..a5a6968f 100644 --- a/py-sdk/inference_logging_client/pyproject.toml +++ b/py-sdk/inference_logging_client/pyproject.toml @@ -34,6 +34,9 @@ dependencies = [ ] [project.optional-dependencies] +gcs = [ + "google-cloud-storage>=2.0.0", +] dev = [ "pytest>=7.0.0", "pytest-cov>=4.0.0", diff --git a/py-sdk/inference_logging_client/readme.md b/py-sdk/inference_logging_client/readme.md index b773af05..3764559d 100644 --- a/py-sdk/inference_logging_client/readme.md +++ b/py-sdk/inference_logging_client/readme.md @@ -2,6 +2,8 @@ A Python SDK for decoding MPLog feature logs from proto, arrow, or parquet format. This client enables you to decode binary-encoded feature data from machine learning inference logging pipelines into Spark DataFrames. +It also reads framed `.log` files produced by `asyncloguploader` — directly from a local path, an open file-like object, or a `gs://` GCS URI — and decodes every embedded MPLog record in one call. + --- ## Table of Contents @@ -13,6 +15,7 @@ A Python SDK for decoding MPLog feature logs from proto, arrow, or parquet forma - [Core API Reference](#core-api-reference) - [decode_mplog()](#decode_mplog) - [decode_mplog_dataframe()](#decode_mplog_dataframe) + - [Reading .log Files](#log-api-reference) - [get_mplog_metadata()](#get_mplog_metadata) - [get_feature_schema()](#get_feature_schema) - [clear_schema_cache()](#clear_schema_cache) @@ -41,6 +44,8 @@ The Inference Logging Client is designed to decode MPLog (Model Proxy Log) featu | **Arrow** | Arrow IPC format with binary columns | Columnar analytics | | **Parquet** | Parquet format with feature map | Long-term storage | +In addition, the SDK reads the **framed `.log` container** emitted by `asyncloguploader`. Each `.log` file holds many MPLog records — the SDK deframes the container and decodes every embedded record (one row per entity) in a single call. See [Reading .log files from GCS or local disk](#reading-log-files-from-gcs-or-local-disk). + ### Key Features - **Multi-format support**: Decode Proto, Arrow, and Parquet encoded logs @@ -48,7 +53,8 @@ The Inference Logging Client is designed to decode MPLog (Model Proxy Log) featu - **Zstd compression support**: Automatic decompression of zstd-compressed data - **Schema fetching**: Retrieves feature schemas from inference API with caching - **Spark integration**: Returns data as PySpark DataFrames -- **CLI tool**: Command-line interface for quick decoding +- **`.log` container reader**: Local path, file-like, or `gs://` URI — deframes and decodes in one call +- **CLI tool**: Command-line interface for quick decoding (auto-routes `.log` and `gs://` inputs) - **Thread-safe caching**: LRU cache for schemas with thread-safe access --- @@ -74,13 +80,22 @@ pip install -e . pip install -e ".[dev]" ``` +### With GCS Support (for `gs://` .log inputs) + +```bash +pip install "inference-logging-client[gcs]" +# or from source: +pip install -e ".[gcs]" +``` + ### Dependencies | Package | Version | Purpose | |---------|---------|---------| -| pyspark | >=3.3.0 | Spark DataFrame operations | +| pyspark | ==3.3.0 | Spark DataFrame operations (exact pin) | | pyarrow | >=5.0.0 | Arrow/Parquet format support | | zstandard | >=0.15.0 | Zstd decompression | +| google-cloud-storage | >=2.0.0 | **Optional** — only required to read `gs://` URIs | --- @@ -148,6 +163,61 @@ decoded_df.show() spark.stop() ``` +### Reading `.log` Files (local disk or GCS) + +`asyncloguploader` writes framed `.log` containers holding many MPLog records +per file, partitioned by model + date + hour: + +``` +gs://gcs-dsci-inferflow-async-logger-prd/ +└── async-logger-gcs-flush//// + └── ___.log +``` + +Decode one file — local path, file-like object, or `gs://` URI: + +```python +import inference_logging_client as ilc + +# Local +pdf = ilc.decode_log_file_to_pandas("./pod.log") + +# GCS (requires: pip install "inference-logging-client[gcs]") +uri = ("gs://gcs-dsci-inferflow-async-logger-prd/async-logger-gcs-flush/" + "search-ad-head-prepaid/2026-06-30/18/" + "search-ad-head-prepaid--prd-inferflow-search-ad-ssd-primary-" + "54c577fb4d-csbhq_2026-06-30_18-29-03_3.log") +pdf = ilc.decode_log_file_to_pandas(uri) +``` + +Fan out across a partition (files share the same schema — fetch it once): + +```python +from google.cloud import storage +client = storage.Client() +blobs = list(client.list_blobs( + "gcs-dsci-inferflow-async-logger-prd", + prefix="async-logger-gcs-flush/search-ad-head-prepaid/2026-06-30/18/", +)) +schema = ilc.get_feature_schema("search-ad-head-prepaid", version=4) +for b in blobs: + ilc.decode_log_file_to_csv( + f"gs://{b.bucket.name}/{b.name}", + f"./out/{b.name.rsplit('/', 1)[-1]}.csv", + schema=schema, + ) +``` + +Or use the CLI — same code path, no boilerplate: + +```bash +inference-logging-client ./pod.log --output-format csv -o pod.csv +inference-logging-client gs://.../pod.log --output-format analyze +``` + +Row schema, output-format table, and full API list live in +[Reading .log Files](#log-api-reference). + --- ## Configuration @@ -388,6 +458,111 @@ decoded = inference_logging_client.decode_mplog_dataframe( --- + +### Reading .log Files + +All `.log` reading routes through one shared iterator +(`iter_decoded_log_rows`), so every output sink emits the **same rows** — +only the container differs. Each function accepts a `source` that can be a +local path, an open binary file-like object, or a `gs://bucket/key` URI. + +#### Available sinks + +| Function | Returns / Writes | Best for | +|----------|------------------|----------| +| `decode_log_file(source, spark, ...)` | `pyspark.sql.DataFrame` | Big files, distributed downstream | +| `decode_log_file_to_pandas(source, ...)` | `pandas.DataFrame` | Notebooks, mid-sized files | +| `decode_log_file_to_csv(source, path, ...)` | rows written (int) | Handoff to BigQuery / Snowflake / Excel | +| `decode_log_file_to_jsonl(source, path, ...)` | rows written (int) | Streaming pipelines | +| `write_parsed_log(source, path, ...)` | records written (int) | Human-readable inspection (asynclogparser `.parsed.log` layout) | +| `analyze_log_file(source)` + `print_analysis(...)` | dict / prints | Structural diagnosis, **no schema fetch** | +| `iter_decoded_log_rows(source, ...)` | `Iterator[dict]` | Custom sinks (Kafka, ClickHouse, ...) | +| `iter_log_records(source, ...)` | `Iterator[(ts_ns, bytes)]` | Raw MPLog payloads for custom decoding | +| `read_log_file(source, ...)` | `list[(ts_ns, bytes)]` | Small files, tests | + +#### Common keyword arguments + +Every decoding function above accepts these — same defaults, same meaning: + +| Argument | Default | Meaning | +|----------|---------|---------| +| `inference_host` | `INFERENCE_HOST` env, else `http://localhost:8082` | Schema API base URL | +| `decompress` | `True` | Attempt zstd decompression per MPLog payload | +| `schema` | `None` | Pre-fetched `list[FeatureInfo]`; skips per-record schema fetch | +| `needed_columns` | `None` | Subset of feature names to keep | +| `strict` | `False` | `True` raises `FormatError` on any malformed frame | + +#### Row shape (every decoding sink) + +``` +entity_id, timestamp_ns, mp_config_id, version, format_type, +user_id, tracking_id, parent_entity, , , ... +``` + +One row per `(record, entity)`. The Spark sink stringifies `list`/`bytes` +values for schema stability; all other sinks preserve native Python types. + +#### Frame layout + +``` +Frame header (8B, little-endian): capacity (uint32) + valid_data_bytes (uint32) +Frame body (capacity − 8 bytes): first valid_data_bytes carry records, rest is padding +Record: [4B length][8B timestamp_ns][MPLog protobuf] +``` + +Records are packed back-to-back with no per-record alignment. The MPLog +protobuf is the same payload `get_mplog_metadata()` and `decode_mplog()` +consume — its metadata byte still drives per-record format selection +(proto / arrow / parquet). + +#### Examples + +```python +import inference_logging_client as ilc + +# Spark +from pyspark.sql import SparkSession +spark = SparkSession.builder.getOrCreate() +df = ilc.decode_log_file("gs://bucket/pod.log", spark) + +# pandas (no Spark) +pdf = ilc.decode_log_file_to_pandas("./pod.log") + +# Write to disk +ilc.decode_log_file_to_csv ("./pod.log", "./out.csv") +ilc.decode_log_file_to_jsonl("./pod.log", "./out.jsonl") +ilc.write_parsed_log ("./pod.log", "./out.parsed.log") # asynclogparser format + +# Structural diagnosis (no schema fetch, no network) +ilc.print_analysis(ilc.analyze_log_file("./pod.log")) + +# CI fail-fast on corrupt frames +ilc.decode_log_file_to_pandas("./golden.log", strict=True) + +# Custom sink via the shared iterator +for row in ilc.iter_decoded_log_rows("gs://bucket/pod.log"): + producer.send("decoded", row) +``` + +#### Parity with `asynclogparser` + +| asynclogparser | SDK equivalent | +|----------------|----------------| +| `python asynclogparse.py foo.log` | `write_parsed_log("foo.log", "foo.parsed.log")` | +| `python asynclogparse.py --analyze foo.log` | `analyze_log_file(...)` + `print_analysis(...)` | +| `deframe_log_file` | `iter_log_records(...)` — byte-level identical | +| Per-record protobuf parse | `parse_mplog_protobuf(...)` | +| Per-entity `decode_proto_features` | Auto-dispatched by `format_type` (proto/arrow/parquet) | + +Additions the SDK carries over the standalone script: +zstd decompression of MPLog payloads, cross-file schema cache keyed by +`(mp_config_id, version)`, and `strict` mode for CI. The ~250 lines of +byte-scan recovery heuristics in `asynclogparser` are intentionally not +carried — use `strict=True` for fail-fast, or the default `strict=False` to +warn-and-skip on a bad frame. + +--- + ### get_mplog_metadata() Extract metadata from MPLog bytes without full decoding. Useful for inspecting format and version. @@ -757,29 +932,42 @@ finally: ### Basic Usage +The CLI accepts three kinds of input and auto-routes them: + +| Input | Detected by | Path | +|-------|-------------|------| +| Raw MPLog bytes file or `-` (stdin) | default | `decode_mplog()` — needs `-m` / `-v` | +| `*.log` container (any local path) | `.log` extension | `decode_log_file()` | +| `gs://bucket/key` URI | `gs://` prefix | `decode_log_file()` | + ```bash -# Decode with auto-detection +# Raw MPLog bytes (the original mode) inference-logging-client --model-proxy-id my-model --version 1 input.bin -# Short form -inference-logging-client -m my-model -v 1 input.bin +# Asyncloguploader .log file — no -m/-v needed (per-record from metadata) +inference-logging-client ./pod-xyz_17-58-12.log -o decoded.csv + +# Same, but straight from GCS +inference-logging-client gs://my-bucket/asynclogs/pod-xyz_17-58-12.log --json ``` ### CLI Arguments | Argument | Short | Required | Default | Description | |----------|-------|----------|---------|-------------| -| `input` | - | Yes | - | Input file or `-` for stdin | -| `--model-proxy-id` | `-m` | Yes | - | Model proxy config ID | -| `--version` | `-v` | Yes | - | Schema version | -| `--format` | `-f` | No | `auto` | Format: `proto`, `arrow`, `parquet`, `auto` | +| `input` | - | Yes | - | MPLog bytes file, `*.log` path, `gs://...` URI, or `-` for stdin | +| `--model-proxy-id` | `-m` | Raw-bytes mode only | - | Model proxy config ID (ignored for `.log` / `gs://`) | +| `--version` | `-v` | Raw-bytes mode only | - | Schema version (ignored for `.log` / `gs://`) | +| `--format` | `-f` | No | `auto` | Format: `proto`, `arrow`, `parquet`, `auto` (raw-bytes mode only) | | `--inference-host` | - | No | env/localhost | Inference service URL | -| `--hex` | - | No | - | Input is hex-encoded | -| `--base64` | - | No | - | Input is base64-encoded | +| `--hex` | - | No | - | Input is hex-encoded (raw-bytes mode only) | +| `--base64` | - | No | - | Input is base64-encoded (raw-bytes mode only) | | `--no-decompress` | - | No | - | Skip zstd decompression | -| `--output` | `-o` | No | stdout | Output directory (CSV/JSON) | +| `--output` | `-o` | No | stdout | Output file path; CSV by default, JSON with `--json` | | `--json` | - | No | - | Output as JSON | | `--spark-master` | - | No | `local[*]` | Spark master URL | +| `--strict` | - | No | - | `.log` mode only: raise on malformed frames instead of skipping | +| `--output-format` | - | No | `spark` (in .log mode) | `.log`/`gs://` only: `spark`, `pandas`, `csv`, `jsonl`, `text`, or `analyze` | ### Examples @@ -811,6 +999,32 @@ inference-logging-client -m my-model -v 1 \ # Skip decompression (for pre-decompressed data) inference-logging-client -m my-model -v 1 --no-decompress input.bin + +# --- .log / GCS examples ----------------------------------------- + +# Decode a local asyncloguploader .log file to CSV +inference-logging-client ./pod-xyz_17-58-12.log -o decoded.csv + +# Decode straight from GCS (requires the [gcs] extra) +inference-logging-client gs://my-bucket/asynclogs/pod-xyz_17-58-12.log \ + --inference-host https://inference.prod.example.com \ + -o gs_decoded.csv + +# CI-style gate: fail on any corrupt frame +inference-logging-client gs://my-bucket/asynclogs/golden.log --strict --json + +# Pandas-style CSV print to stdout (no Spark) +inference-logging-client ./pod-xyz_17-58-12.log --output-format pandas + +# Write CSV / JSONL directly (no Spark) +inference-logging-client ./pod-xyz_17-58-12.log --output-format csv -o decoded.csv +inference-logging-client ./pod-xyz_17-58-12.log --output-format jsonl -o decoded.jsonl + +# Asynclogparser-compatible human-readable text +inference-logging-client ./pod-xyz_17-58-12.log --output-format text -o decoded.parsed.log + +# Structural diagnosis only (no schema fetch, no decoding) +inference-logging-client ./pod-xyz_17-58-12.log --output-format analyze ``` ### CLI Output Format @@ -985,11 +1199,12 @@ spark.stop() inference_logging_client/ ├── __init__.py # Public API exports, decode_mplog(), decode_mplog_dataframe() ├── __main__.py # Module execution entry point -├── cli.py # Command-line interface +├── cli.py # Command-line interface (auto-routes .log / gs://) ├── decoder.py # Core byte decoding, type conversion ├── exceptions.py # Exception classes ├── formats.py # Proto/Arrow/Parquet format decoders ├── io.py # Schema fetching, protobuf parsing +├── log_reader.py # .log frame deframer + GCS reader + decode_log_file() ├── types.py # Data type definitions (Format, FeatureInfo, DecodedMPLog) └── utils.py # Utility functions (type normalization, formatting) ``` From db800e21cd8e19729408b3ccbacfda7817a77d34 Mon Sep 17 00:00:00 2001 From: dharma-shashank-meesho Date: Wed, 1 Jul 2026 13:16:56 +0530 Subject: [PATCH 2/2] feat(inference-logging-client): directory / gs:// prefix aware decode_logs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds decode_logs, decode_logs_to_pandas, and list_log_sources so the same API accepts a single .log file, a local directory, a gs:// prefix, or an explicit list of sources — and returns one merged DataFrame. - log_reader.py: new list_log_sources that dispatches on path shape (single file / local dir / GCS URI / prefix / list); decode_logs and decode_logs_to_pandas union across every listed source with unionByName(allowMissingColumns=True) / pd.concat. - __init__.py: export the three new public functions. - cli.py: default log-mode Spark path now calls decode_logs so directories and prefixes work; --output-format csv/jsonl/text on a multi-source input writes one file per input into --output as a directory; --output-format analyze on a multi-source input inspects the first file (with a note). - readme.md: adds a "Single file vs whole hour partition" subsection and extends the sinks table with the three new functions. Verified on gs://gcs-dsci-inferflow-async-logger-prd/.../2026-06-30/23/ (hour-level prefix): list_log_sources found 130 .log files; decode_logs_to_pandas on a 3-file subset produced a 26,364 x 550 merged DataFrame; CLI --output-format analyze on the bare prefix worked. --- .../inference_logging_client/__init__.py | 6 + .../inference_logging_client/cli.py | 74 ++++++- .../inference_logging_client/log_reader.py | 192 ++++++++++++++++++ py-sdk/inference_logging_client/readme.md | 35 ++++ 4 files changed, 302 insertions(+), 5 deletions(-) diff --git a/py-sdk/inference_logging_client/inference_logging_client/__init__.py b/py-sdk/inference_logging_client/inference_logging_client/__init__.py index 59100412..e43f237c 100644 --- a/py-sdk/inference_logging_client/inference_logging_client/__init__.py +++ b/py-sdk/inference_logging_client/inference_logging_client/__init__.py @@ -51,8 +51,11 @@ 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, @@ -72,6 +75,9 @@ "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", diff --git a/py-sdk/inference_logging_client/inference_logging_client/cli.py b/py-sdk/inference_logging_client/inference_logging_client/cli.py index a1c009c0..9fdf089f 100644 --- a/py-sdk/inference_logging_client/inference_logging_client/cli.py +++ b/py-sdk/inference_logging_client/inference_logging_client/cli.py @@ -14,10 +14,13 @@ 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, ) @@ -25,10 +28,22 @@ def _is_log_input(path: str) -> bool: - """Return True when ``path`` should be routed through the .log reader.""" + """Return True when ``path`` should be routed through the .log reader + (file, directory, or gs:// URI/prefix).""" if path.startswith("gs://"): return True - return path.lower().endswith(".log") + 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(): @@ -128,13 +143,27 @@ def main(): inference_host = args.inference_host or os.getenv("INFERENCE_HOST", "http://localhost:8082") # 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: if out_fmt == "analyze": - print_analysis(analyze_log_file(args.input)) + 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_log_file_to_pandas( + pdf = decode_logs_to_pandas( args.input, inference_host=inference_host, decompress=not args.no_decompress, @@ -153,6 +182,40 @@ def main(): 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 + 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, @@ -220,7 +283,8 @@ def main(): try: if log_mode: - df = decode_log_file( + # decode_logs auto-handles single file, local dir, or gs:// prefix. + df = decode_logs( source=args.input, spark=spark, inference_host=inference_host, diff --git a/py-sdk/inference_logging_client/inference_logging_client/log_reader.py b/py-sdk/inference_logging_client/inference_logging_client/log_reader.py index 0d00311d..3bc01b16 100644 --- a/py-sdk/inference_logging_client/inference_logging_client/log_reader.py +++ b/py-sdk/inference_logging_client/inference_logging_client/log_reader.py @@ -862,3 +862,195 @@ def print_analysis(result: dict) -> None: print(" records carry an 8-byte UnixNano timestamp prefix.") elif tsf is False: print(" WARNING: records do NOT carry the expected timestamp prefix.") + + +# --------------------------------------------------------------------------- +# Directory / prefix handling +# --------------------------------------------------------------------------- + +def _looks_like_single_file(s: str, suffix: str) -> bool: + return s.endswith(suffix) + + +def list_log_sources( + source: Union[str, Path, List[Union[str, Path]]], + suffix: str = ".log", +) -> List[str]: + """Expand ``source`` into a concrete list of ``.log`` file paths / URIs. + + Rules: + - A ``list`` / ``tuple`` — returned as-is (stringified). + - A local path to a file ending with ``suffix`` — ``[source]``. + - A local directory — non-recursive listing of files ending with ``suffix``. + - A ``gs://bucket/key`` URI ending with ``suffix`` — ``[source]``. + - A ``gs://bucket/prefix/`` URI (or any URI without ``suffix``) — every + object under that prefix whose name ends with ``suffix``, returned as + ``gs://bucket/`` in lexicographic order. + + File-like objects can't be enumerated — pass them directly to the + single-file APIs. + """ + if isinstance(source, (list, tuple)): + return [str(s) for s in source] + + if isinstance(source, Path): + source = str(source) + + if not isinstance(source, str): + raise TypeError( + f"list_log_sources needs a path, gs:// URI, or list; got {type(source).__name__}" + ) + + # GCS + if source.startswith("gs://"): + if _looks_like_single_file(source, suffix): + return [source] + try: + from google.cloud import storage # type: ignore + except ImportError as exc: + raise ImportError( + "Listing gs:// prefixes requires 'google-cloud-storage'. " + "Install it with: pip install google-cloud-storage" + ) from exc + + rest = source[len("gs://") :] + if "/" in rest: + bucket_name, prefix = rest.split("/", 1) + else: + bucket_name, prefix = rest, "" + client = storage.Client() + blobs = client.list_blobs(bucket_name, prefix=prefix) + found = sorted( + f"gs://{bucket_name}/{b.name}" for b in blobs if b.name.endswith(suffix) + ) + return found + + # Local + if os.path.isdir(source): + entries = sorted( + os.path.join(source, name) + for name in os.listdir(source) + if name.endswith(suffix) and os.path.isfile(os.path.join(source, name)) + ) + return entries + + if os.path.isfile(source): + return [source] + + if _looks_like_single_file(source, suffix): + # Not an existing file, but user clearly asked for one — let the caller + # surface the FileNotFoundError. + return [source] + + raise FileNotFoundError( + f"list_log_sources: {source!r} is neither an existing file/dir nor a gs:// URI" + ) + + +def decode_logs( + source: Union[LogSource, List[Union[str, Path]]], + spark: "SparkSession", + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, + suffix: str = ".log", +) -> "SparkDataFrame": + """Decode a single ``.log`` file **or** every ``.log`` file under a directory / + GCS prefix / list, and return one unioned Spark DataFrame. + + Same row shape and semantics as :func:`decode_log_file`. Uses + :func:`list_log_sources` to expand the source, then decodes each file and + unions with ``unionByName(..., allowMissingColumns=True)`` so per-file + schema drift (e.g. different feature subsets) doesn't break the merge. + """ + # Fast path: caller passed a file-like directly. + if hasattr(source, "read") and not isinstance(source, (str, Path, list, tuple)): + return decode_log_file( + source, # type: ignore[arg-type] + spark, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + ) + + sources = list_log_sources(source, suffix=suffix) # type: ignore[arg-type] + if not sources: + # Return an empty DataFrame with the standard column set. + from pyspark.sql.types import LongType, StringType, StructField, StructType + + fields = [StructField("entity_id", StringType(), True)] + for col in _RECORD_METADATA_COLUMNS: + dtype = LongType() if col in ("timestamp_ns", "version", "format_type") else StringType() + fields.append(StructField(col, dtype, True)) + return spark.createDataFrame([], StructType(fields)) + + from functools import reduce + + dfs = [ + decode_log_file( + s, + spark, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + ) + for s in sources + ] + if len(dfs) == 1: + return dfs[0] + return reduce(lambda a, b: a.unionByName(b, allowMissingColumns=True), dfs) + + +def decode_logs_to_pandas( + source: Union[LogSource, List[Union[str, Path]]], + inference_host: Optional[str] = None, + decompress: bool = True, + schema: Optional[List[FeatureInfo]] = None, + needed_columns: Optional[Collection[str]] = None, + strict: bool = False, + suffix: str = ".log", +): + """Same as :func:`decode_logs` but returns a pandas DataFrame (no Spark). + + Files are decoded sequentially, then concatenated with + ``pd.concat(..., ignore_index=True)``. + """ + try: + import pandas as pd # type: ignore + except ImportError as exc: + raise ImportError( + "decode_logs_to_pandas requires pandas. Install it with: pip install pandas" + ) from exc + + if hasattr(source, "read") and not isinstance(source, (str, Path, list, tuple)): + return decode_log_file_to_pandas( + source, # type: ignore[arg-type] + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + ) + + sources = list_log_sources(source, suffix=suffix) # type: ignore[arg-type] + if not sources: + return pd.DataFrame(columns=["entity_id", *_RECORD_METADATA_COLUMNS]) + + frames = [ + decode_log_file_to_pandas( + s, + inference_host=inference_host, + decompress=decompress, + schema=schema, + needed_columns=needed_columns, + strict=strict, + ) + for s in sources + ] + return pd.concat(frames, ignore_index=True) diff --git a/py-sdk/inference_logging_client/readme.md b/py-sdk/inference_logging_client/readme.md index 3764559d..6c2f172f 100644 --- a/py-sdk/inference_logging_client/readme.md +++ b/py-sdk/inference_logging_client/readme.md @@ -472,6 +472,9 @@ local path, an open binary file-like object, or a `gs://bucket/key` URI. |----------|------------------|----------| | `decode_log_file(source, spark, ...)` | `pyspark.sql.DataFrame` | Big files, distributed downstream | | `decode_log_file_to_pandas(source, ...)` | `pandas.DataFrame` | Notebooks, mid-sized files | +| `decode_logs(source, spark, ...)` | `pyspark.sql.DataFrame` | **Single file OR a directory / GCS prefix** — auto-lists and unions | +| `decode_logs_to_pandas(source, ...)` | `pandas.DataFrame` | Same, pandas flavour | +| `list_log_sources(source)` | `list[str]` | Enumerate `.log` files under a dir / `gs://` prefix / list | | `decode_log_file_to_csv(source, path, ...)` | rows written (int) | Handoff to BigQuery / Snowflake / Excel | | `decode_log_file_to_jsonl(source, path, ...)` | rows written (int) | Streaming pipelines | | `write_parsed_log(source, path, ...)` | records written (int) | Human-readable inspection (asynclogparser `.parsed.log` layout) | @@ -480,6 +483,38 @@ local path, an open binary file-like object, or a `gs://bucket/key` URI. | `iter_log_records(source, ...)` | `Iterator[(ts_ns, bytes)]` | Raw MPLog payloads for custom decoding | | `read_log_file(source, ...)` | `list[(ts_ns, bytes)]` | Small files, tests | +#### Single file vs whole hour partition + +`decode_logs` / `decode_logs_to_pandas` accept any of these — they figure out +what you meant and produce one merged DataFrame: + +| `source` | Behaviour | +|----------|-----------| +| Single `.log` file (path or `gs://...log`) | Decodes that file. | +| Local directory | Non-recursive listing of `*.log` inside, unions results. | +| `gs://bucket/prefix/` (or any URI not ending in `.log`) | Lists every `.log` under the prefix (works at hour, day, or month level). | +| `list[str \| Path]` | Uses the list verbatim. | + +```python +import inference_logging_client as ilc + +# One hour's worth of logs → one Spark DataFrame +df = ilc.decode_logs( + "gs://gcs-dsci-inferflow-async-logger-prd/async-logger-gcs-flush/" + "search-organic-l2-ranker-prepaid-rtp-mall-hasp_scaleup/2026-06-30/23/", + spark, +) + +# Same, pandas +pdf = ilc.decode_logs_to_pandas("./hour_dump/") + +# See what would be enumerated without decoding +files = ilc.list_log_sources("gs://bucket/prefix/") # → list[str] +``` + +The CLI also picks up directories and GCS prefixes automatically — pass one +to `inference-logging-client` and it dispatches through `decode_logs`. + #### Common keyword arguments Every decoding function above accepts these — same defaults, same meaning: