diff --git a/docs/guides/datasets.md b/docs/guides/datasets.md index 541120fc8..ddea45b48 100644 --- a/docs/guides/datasets.md +++ b/docs/guides/datasets.md @@ -20,7 +20,7 @@ The following arguments configure datasets and their processing: - `synthetic_text` — generates synthetic prompts on the fly. Required field: `prompt_tokens`. Optional: `output_tokens`, `turns`, `prefix_tokens`, `prefix_count`, `prefix_buckets`, and distribution controls (`prompt_tokens_stdev`, `output_tokens_stdev`, etc.). - `huggingface` (alias `hf`) — loads from HuggingFace Hub or a local directory/file. Required field: `source` (dataset ID or path). Pass dataset loading arguments (for example `split`, `name`) via `load_kwargs`. - `json_file`, `csv_file`, `text_file`, `parquet_file`, `arrow_file`, `hdf5_file`, `db_file`, `tar_file` — loads from a local file. Required field: `path`. - - `trace_synthetic`, `mooncake` — loads a JSONL, JSON, CSV, or Parquet trace file for replay benchmarking. Required field: `path`. Optional: `timestamp_column` (default: `timestamp`), `prompt_tokens_column` (default: `input_length`), `output_tokens_column` (default: `output_length`). + - `trace_synthetic`, `mooncake`, `weka` — loads a JSONL, JSON, CSV, or Parquet trace file for replay benchmarking. Required field: `path`. Optional: `timestamp_column` (default: `timestamp`), `prompt_tokens_column` (default: `input_length`), `output_tokens_column` (default: `output_length`). In addition, you can specify additional arguments to the dataset loading with the data argument `load_kwargs`: diff --git a/docs/guides/trace_replay.md b/docs/guides/trace_replay.md index 53e1835ad..52b6a922d 100644 --- a/docs/guides/trace_replay.md +++ b/docs/guides/trace_replay.md @@ -9,7 +9,8 @@ Detailed use of the replay profile and file-based datasets as a whole is explain These are passed to the `--data` argument as `kind=format`: - `trace_synthetic`: A trace format that does the bare minimum needed to complete a fully functioning trace replay benchmark with synthetic prompt generation -- `mooncake`: The trace format used by the serving platform Mooncake, as defined in [https://doi.org/10.48550/arXiv.2407.00079](https://doi.org/10.48550/arXiv.2407.00079) +- `mooncake`: The trace format used by the serving platform *Mooncake*, as defined in [https://doi.org/10.48550/arXiv.2407.00079](https://doi.org/10.48550/arXiv.2407.00079) +- `weka`: The trace format used by WEKA's *Augmented Memory Grid*, as specified [in the original research repository](https://github.com/callanjfox/agentic-coding-analysis/blob/master/docs/TRACE_FORMAT.md) ## Format-Agnostic Data Arguments @@ -24,8 +25,8 @@ All trace formats can accept the following optional data arguments: These are passed through the `--data` argument like below: ```bash -guidellm benchmark \ - --target http://localhost:8000 \ +guidellm run \ + --backend kind=openai_http,target=http://localhost:8000 \ --profile kind=replay \ --data "kind=trace_synthetic,path=replay.jsonl,timestamp_column=ts,prompt_tokens_column=input_tokens,output_tokens_column=generated_tokens" ``` @@ -36,9 +37,33 @@ guidellm benchmark \ ### `mooncake` -The Mooncake format expects an additional column for hash IDs. During prompt generation, hash IDs sharing the same previous ID are required to represent distinct blocks of token ids. +The Mooncake format expects an additional column for prefix-based cache hash IDs. During prompt generation, hash IDs sharing the same previous ID are required to represent distinct blocks of token ids. | Argument | Default | Description | | -------------------- | ---------- | --------------------------------------------------- | | `hash_ids_column` | "hash_ids" | Column name for lists of hash IDs in the trace file | | `hash_id_block_size` | 512 | Amount of tokens represented by one hash ID | + +### `weka` + +**NOTE:** :construction: While the format is accepted, some features such as subagent conversations, tool call events and non-linear histories are still in active development. The results from datasets including these features will be unreliable. + +The WEKA format expects a column with conversation UUIDs that is not wrapped within another column. The timestamp, input token length, output token length and hash IDs columns must either be top level columns, or must all be wrapped inside the same JSON column (ex. "requests"). + +Similar to Mooncake, WEKA uses prefix-based cache hash IDs. The original [specification](https://github.com/callanjfox/agentic-coding-analysis/blob/master/docs/TRACE_FORMAT.md) for the trace requires hash IDs to be 1 or greater, and for trailing hash IDs to be dropped if there are not enough input tokens to fill the hash ID block size. To accommodate for datasets which may not follow the specification exactly (ex. [semianalysisai/cc-traces-weka-no-subagents-051226](https://huggingface.co/datasets/semianalysisai/cc-traces-weka-no-subagents-051226)), GuideLLM will accept any non-negative integer as a valid hash ID, and will drop partially filled hash IDs if they exist. + +GuideLLM will generate prompts starting from the first conversation. When the conversation ends, the next conversation will be used. Hash IDs and relative timestamps are local to the conversation. After a conversation ends, the hash ID tree is reset and the relative timestamp returns to 0.0. + +| Argument | Default | Description | +| ------------------------ | ---------- | ---------------------------------------------------- | +| `conversation_id_column` | "id" | Column name for conversation UUIDs in the trace file | +| `hash_ids_column` | "hash_ids" | Column name for lists of hash IDs in the trace file | +| `hash_id_block_size` | 64 | Amount of tokens represented by one hash ID | + +Modified defaults: + +| Argument | New Default | +| ---------------------- | ----------- | +| `timestamp_column` | "t" | +| `prompt_tokens_column` | "in" | +| `output_tokens_column` | "out" | diff --git a/src/guidellm/data/deserializers/__init__.py b/src/guidellm/data/deserializers/__init__.py index f82bb38f5..98e436788 100644 --- a/src/guidellm/data/deserializers/__init__.py +++ b/src/guidellm/data/deserializers/__init__.py @@ -43,11 +43,13 @@ TraceDatasetDeserializer, TraceFormatBase, TraceFormatRegistry, + create_prompt_from_hash_ids, decode_prompt, generate_token_ids, ) from .trace_minimal import MinimalTraceFormatArgs from .trace_mooncake import MooncakeTraceFormatArgs +from .trace_weka import WEKATraceFormatArgs __all__ = [ "ArrowFileDatasetDeserializer", @@ -85,6 +87,8 @@ "TraceDatasetDeserializer", "TraceFormatBase", "TraceFormatRegistry", + "WEKATraceFormatArgs", + "create_prompt_from_hash_ids", "decode_prompt", "generate_token_ids", ] diff --git a/src/guidellm/data/deserializers/trace_common.py b/src/guidellm/data/deserializers/trace_common.py index f665def3d..9c9c6c131 100644 --- a/src/guidellm/data/deserializers/trace_common.py +++ b/src/guidellm/data/deserializers/trace_common.py @@ -6,9 +6,11 @@ from __future__ import annotations -from collections.abc import Callable, Iterable +import dataclasses +import enum +from collections.abc import Callable, Iterable, Sequence from pathlib import Path -from typing import Any, Protocol +from typing import Any, Protocol, cast import numpy as np from datasets import ( @@ -31,6 +33,13 @@ ) from guidellm.data.schemas import DataArgs from guidellm.utils.hf_datasets import load_dataset_from_file +from guidellm.utils.json_unwrap import ( + VirtualColumnLocation, + construct_virtual_column_locations, + get_json_column_names, + try_json_load, + unzip_virtual_column_locations, +) from guidellm.utils.registry import RegistryMixin __all__ = [ @@ -38,6 +47,7 @@ "TraceDatasetDeserializer", "TraceFormatBase", "TraceFormatRegistry", + "create_prompt_from_hash_ids", "decode_prompt", "generate_token_ids", ] @@ -59,7 +69,7 @@ def generate_token_ids( processor: PreTrainedTokenizerBase, faker: Faker, margin_of_safety: int = 8, -) -> list[int]: +) -> tuple[int]: """Generate `token_count` synthetic token ids for trace prompt construction. Ideally, `margin_of_safety` should be set to slighty more than @@ -69,77 +79,32 @@ def generate_token_ids( attempt += 1 # The Faker.text() can only generate text of at least 5 characters. num_chars = max(token_count * margin_of_safety * attempt, 5) - text = faker.text(max_nb_chars=num_chars) + text = faker.text(num_chars) token_ids = processor.encode(text) if len(token_ids) >= token_count: - return token_ids[:token_count] - - -def validate_trace_path(path: Path | str) -> Path: - path = Path(path) - if path.stat().st_size == 0: - raise ValueError(f"Trace file is empty: {path}") - return path - - -def check_and_raise_missing_columns( - required_columns: list[str], actual_columns: list[str] -) -> None: - missing = [c for c in required_columns if c not in actual_columns] - if missing: - raise KeyError(f"Trace row missing required columns: {missing}") - - -def load_trace_rows( - path: Path | str, - timestamp_column_name: str, - required_columns: Features, - **data_kwargs: Any, -) -> Dataset: - """ - Load trace file rows as a HuggingFace Dataset. + return tuple(token_ids[:token_count]) - Every column in required_columns must exist in the dataset; - otherwise KeyError is raised with a descriptive message. - Rows are sorted by column timestamp_column_name. - - :param path: Path to the trace file. - :param timestamp_column_name: Name of the timestamp column used to sort trace rows. - :param required_columns: List of column/fields that each row must have. Must contain - the timestamp column. - :param data_kwargs: Additional keyword arguments forwarded to load_dataset. - :return: HuggingFace Dataset (iterable as dicts, column-accessible). - :raises DataNotSupportedError: For any of the following reasons: - - The dataset is empty or has no valid rows - - A required column contains a NoneType - - A required column failed during cast to feature type - - :raises KeyError: If a required column is missing in the dataset. - :raises ValueError: If the file format is not .jsonl, .json, .csv or .parquet. - """ - path = validate_trace_path(path) - trace_dataset = load_dataset_from_file(path, **data_kwargs) - if required_columns: - check_and_raise_missing_columns( - required_columns.keys(), trace_dataset.column_names - ) - if not trace_dataset: - raise DataNotSupportedError(f"Trace file has no valid rows: {path}") - for name, val in required_columns.items(): - if trace_dataset.data[name].null_count != 0: - raise DataNotSupportedError(f"Missing column values in {name}") - try: - trace_dataset.cast_column(name, val) - except ValueError as e: - raise DataNotSupportedError(str(e)) from e +def create_prompt_from_hash_ids( + hash_ids: list[int], + hash_id_table: dict[int, tuple[int]], + processor: PreTrainedTokenizerBase, +) -> str: + """Returns a synthetic prompt from `hash_ids` using pre-generated token blocks. - return trace_dataset.sort(timestamp_column_name) + Precondition: All ids in `hash_ids` appear in `hash_id_table`.""" + prompt_token_ids = [ + token for hash_id in hash_ids for token in hash_id_table[hash_id] + ] + return decode_prompt(processor, prompt_token_ids) class TraceFormatBase(Protocol): def __init__(self) -> None: ... + def reset(self) -> None: + pass + def required_columns(self, config) -> Features: ... def validate_row(self, config, row: dict) -> None: @@ -188,16 +153,13 @@ class TraceDataArgs(DataArgs): default="output_length", description="Column name for output token counts in the trace file.", ) - - -def validate_row(row: dict, config: TraceDataArgs) -> None: - n_in = row[config.prompt_tokens_column] - n_out = row[config.output_tokens_column] - if n_in < 0 or n_out < 0: - raise DataNotSupportedError( - f"Trace token counts must be non-negative, got " - f"input_length={n_in}, output_length={n_out}" - ) + conversation_id_column: str | None = Field( + default=None, + description=( + "Column name for conversation IDs. Required for formats " + "with conversation-scoped trace data such as hash IDs." + ), + ) class TraceExamplesIterable(_BaseExamplesIterable): @@ -207,6 +169,7 @@ class TraceExamplesIterable(_BaseExamplesIterable): def __init__( self, config: TraceDataArgs, + trace_rows: Dataset, processor: PreTrainedTokenizerBase, random_seed: int, ): @@ -216,42 +179,27 @@ def __init__( self.processor = processor self.faker = Faker() self.faker.seed_instance(random_seed) - try: - self.trace_rows = load_trace_rows( - config.path, - config.timestamp_column, - required_columns=Features( - { - config.timestamp_column: Value("float"), - config.prompt_tokens_column: Value("int32"), - config.output_tokens_column: Value("int32"), - **dict(self.format.required_columns(self.config)), - } - ), - **config.load_kwargs, - ) - except (DatasetGenerationError, KeyError, ValueError) as e: - raise DataNotSupportedError(str(e)) from e - - for row in self.trace_rows: - validate_row(row, self.config) - self.format.validate_row(self.config, row) + self.trace_rows = trace_rows self.iteration_count = 0 def __iter__(self) -> Iterable[tuple[int, dict[str, Any]]]: self.iteration_count += 1 - row_idx = 0 timestamps = self.trace_rows[self.config.timestamp_column] - while True: - try: - row = self.trace_rows[row_idx] - except IndexError: - break + conv_col = self.config.conversation_id_column + current_conv = None + conv_start_ts = timestamps[0] + for row_idx, row in enumerate(self.trace_rows): + if conv_col: + conv_id = row[conv_col] + if conv_id != current_conv: + current_conv = conv_id + conv_start_ts = row[self.config.timestamp_column] + self.format.reset() prompt = self.format.create_prompt( self.config, row, self.processor, self.faker ) - relative_timestamp = timestamps[row_idx] - timestamps[0] + relative_timestamp = timestamps[row_idx] - conv_start_ts yield ( row_idx, { @@ -261,7 +209,6 @@ def __iter__(self) -> Iterable[tuple[int, dict[str, Any]]]: "relative_timestamp": relative_timestamp, }, ) - row_idx += 1 @property def is_typed(self) -> bool: @@ -312,10 +259,11 @@ class TraceDataset(IterableDataset): def __init__( self, config: TraceDataArgs, + trace_rows: Dataset, processor: PreTrainedTokenizerBase, random_seed: int, ): - ex_iterable = TraceExamplesIterable(config, processor, random_seed) + ex_iterable = TraceExamplesIterable(config, trace_rows, processor, random_seed) super().__init__( ex_iterable=ex_iterable, info=DatasetInfo( @@ -330,6 +278,242 @@ def set_epoch(self, epoch: int): self._ex_iterable.iteration_count = epoch +def _get_missing_columns( + required_columns: list[str], actual_columns: list[str] +) -> list[str]: + return [c for c in required_columns if c not in actual_columns] + + +class Status(enum.Enum): + FAILURE = enum.auto() + SUCCESS = enum.auto() + + +@dataclasses.dataclass +class ColumnSearchResult: + status: Status + checked_json_columns: bool + relevant_column_names: Sequence[str | VirtualColumnLocation] + + +def _is_supported_json_data_type(data: Any) -> bool: + """Currently, only JSON data in the form of a list of JSON objects is supported.""" + return isinstance(data, list) and (len(data) == 0 or isinstance(data[0], dict)) + + +def _find_virtual_columns( + sample_row: dict[str, Any], + json_column_names: list[str], + target_columns: list[str], +) -> ColumnSearchResult: + """Required columns must all be stored within the same column.""" + completely_missing = set(target_columns) + for col in json_column_names: + sample = sample_row[col] + parsed = try_json_load(sample) if isinstance(sample, str) else sample + if _is_supported_json_data_type(parsed) and len(parsed) > 0: + virtual_columns = [] if parsed is None else list(parsed[0].keys()) + missing = _get_missing_columns(target_columns, virtual_columns) + if not missing: + locations = construct_virtual_column_locations(col, target_columns) + return ColumnSearchResult(Status.SUCCESS, True, locations) + completely_missing = completely_missing.difference(missing) + if json_column_names: + return ColumnSearchResult(Status.FAILURE, True, list(completely_missing)) + return ColumnSearchResult(Status.FAILURE, False, []) + + +def _find_required_columns( + columns: list[str], dataset: Dataset, conversation_id_col: str | None = None +) -> ColumnSearchResult: + """Returns a list of all missing columns on failure. Otherwise returns a list of + the locations of any required columns embedded inside a JSON dict.""" + missing = _get_missing_columns(columns, dataset.column_names) + if missing: + # The conversation IDs should always be top-level. + if conversation_id_col: + if conversation_id_col in missing: + return ColumnSearchResult(Status.FAILURE, False, conversation_id_col) + columns.remove(conversation_id_col) + json_column_names = get_json_column_names(dataset) + sample = dataset[0] + result = _find_virtual_columns(sample, json_column_names, columns) + if result.status is Status.FAILURE and not result.checked_json_columns: + result.relevant_column_names = missing + return result + return ColumnSearchResult(Status.SUCCESS, False, []) + + +def _get_json_dicts(data: Any) -> Any: + if isinstance(data, str): + return try_json_load(data) + if isinstance(data, list) and len(data) > 0 and isinstance(data[0], str): + return list(map(try_json_load, data)) + return data + + +def _make_columns_from_virtual( + batch: dict[str, list], + *args, + wrapper_col: str, + virtual_cols: list[str], + conversation_id_col: str | None = None, +) -> dict[str, list]: + """Intended to be used with `datasets.Dataset.map()`.""" + indices = args[0] if args else [] + json_dicts = [] + conv_ids = [] + for batch_idx, json_dicts_list in enumerate(batch[wrapper_col]): + json_dicts.extend(_get_json_dicts(json_dicts_list)) + if conversation_id_col: + conv_ids.extend([indices[batch_idx]] * len(json_dicts_list)) + result = {c: [row[c] for row in json_dicts] for c in virtual_cols} + if conversation_id_col: + result[conversation_id_col] = conv_ids + return result + + +def _make_dataset_from_virtual( + dataset: Dataset, + columns: list[VirtualColumnLocation], + conversation_id_col: str | None = None, +) -> Dataset: + """Assumes all virtual columns are stored inside the same column. + (Currently ensured by `_is_supported_json_data_type`).""" + wrapper_cols, virt_cols = unzip_virtual_column_locations(columns) + return dataset.map( + _make_columns_from_virtual, + batched=True, + with_indices=conversation_id_col is not None, + remove_columns=dataset.column_names, + fn_kwargs={ + "wrapper_col": wrapper_cols[0], + "virtual_cols": virt_cols, + "conversation_id_col": conversation_id_col, + }, + ) + + +def _handle_column_search_result( + result: ColumnSearchResult, + dataset: Dataset, + conversation_id_col: str | None = None, +) -> Dataset: + """Returns an updated dataset where any required columns found wrapped inside + JSON dicts are unwrapped and added as columns to the dataset. + + :raises KeyError: If a required column is missing in the dataset.""" + if result.status is Status.FAILURE: + additional_info = "" + if result.checked_json_columns: + additional_info = ( + "Note: GuideLLM searched columns with lists of JSON objects after " + "failing to find them at the top level. " + "Ensure that all required columns are wrapped in the same column if " + "this is where they are intended to be found." + ) + raise KeyError( + f"Trace row missing required columns: {result.relevant_column_names} " + f"{additional_info}" + ) + if not result.checked_json_columns: + return dataset + return _make_dataset_from_virtual( + dataset, + cast("list[VirtualColumnLocation]", result.relevant_column_names), + conversation_id_col, + ) + + +def _load_trace_rows( + dataset: Dataset, + timestamp_column_name: str, + required_columns: Features, + conversation_id_column_name: str | None = None, +) -> Dataset: + """ + Load trace file rows as a HuggingFace Dataset. + + Every column in required_columns must exist in the dataset; + otherwise KeyError is raised with a descriptive message. + Rows are sorted by column timestamp_column_name. + + :param dataset: The dataset to load. + :param timestamp_column_name: Name of the timestamp column used to sort trace rows. + :param required_columns: List of column/fields that each row must have. Must contain + the timestamp column. + :param conversation_id_column_name: The conversation id column used to sort rows, + if applicable. + :return: HuggingFace Dataset (iterable as dicts, column-accessible). + :raises DataNotSupportedError: For any of the following reasons: + - The dataset is empty or has no valid rows + - A required column contains a NoneType + - A required column failed during cast to feature type + """ + result = _find_required_columns( + list(required_columns.keys()), dataset, conversation_id_column_name + ) + dataset = _handle_column_search_result(result, dataset, conversation_id_column_name) + + for name, val in required_columns.items(): + if dataset.data[name].null_count != 0: + raise DataNotSupportedError(f"Missing column values in {name}") + try: + dataset.cast_column(name, val) + except ValueError as e: + raise DataNotSupportedError(str(e)) from e + + if conversation_id_column_name: + return dataset.sort([conversation_id_column_name, timestamp_column_name]) + return dataset.sort(timestamp_column_name) + + +def validate_path(path: Path) -> None: + if not path.exists(): + raise DataNotSupportedError(f"Trace file not found: {path}") + if not path.is_file(): + raise DataNotSupportedError(f"Trace path is not a file: {path}") + if path.stat().st_size == 0: + raise DataNotSupportedError(f"Trace file is empty: {path}") + + +def try_load_trace(config: TraceDataArgs, dataset: Dataset) -> Dataset: + trace_format = TraceFormatRegistry.dispatch(config) + try: + return _load_trace_rows( + dataset, + config.timestamp_column, + required_columns=Features( + { + config.timestamp_column: Value("float"), + config.prompt_tokens_column: Value("int32"), + config.output_tokens_column: Value("int32"), + **dict(trace_format.required_columns(config)), + } + ), + conversation_id_column_name=config.conversation_id_column, + ) + except (DatasetGenerationError, KeyError, ValueError) as e: + raise DataNotSupportedError(str(e)) from e + + +def _validate_row(row: dict, config: TraceDataArgs) -> None: + n_in = row[config.prompt_tokens_column] + n_out = row[config.output_tokens_column] + if n_in < 0 or n_out < 0: + raise DataNotSupportedError( + f"Trace token counts must be non-negative, got " + f"input_length={n_in}, output_length={n_out}" + ) + + +def validate_rows(config: TraceDataArgs, trace_rows: Dataset) -> None: + trace_format = TraceFormatRegistry.dispatch(config) + for row in trace_rows: + _validate_row(row, config) + trace_format.validate_row(config, row) + + @DatasetDeserializerFactory.register(["trace_synthetic"]) class TraceDatasetDeserializer(DatasetDeserializer): """Dataset deserializer for all trace formats.""" @@ -340,8 +524,13 @@ def __call__( processor_factory: Callable[[], PreTrainedTokenizerBase], random_seed: int = 42, ) -> IterableDataset: - if not config.path.exists(): - raise DataNotSupportedError(f"Trace file not found: {config.path}") - if not config.path.is_file(): - raise DataNotSupportedError(f"Trace path is not a file: {config.path}") - return TraceDataset(config, processor_factory(), random_seed) + validate_path(config.path) + try: + dataset = load_dataset_from_file(config.path, **config.load_kwargs) + except ValueError as e: + raise DataNotSupportedError(str(e)) from e + if not dataset: + raise DataNotSupportedError(f"Trace file has no valid rows: {config.path}") + trace_rows = try_load_trace(config, dataset) + validate_rows(config, trace_rows) + return TraceDataset(config, trace_rows, processor_factory(), random_seed) diff --git a/src/guidellm/data/deserializers/trace_minimal.py b/src/guidellm/data/deserializers/trace_minimal.py index 8b9b5df91..9a552ae0f 100644 --- a/src/guidellm/data/deserializers/trace_minimal.py +++ b/src/guidellm/data/deserializers/trace_minimal.py @@ -64,4 +64,4 @@ def create_prompt( token_ids = generate_token_ids( row[config.prompt_tokens_column], processor, faker ) - return decode_prompt(processor, token_ids) + return decode_prompt(processor, list(token_ids)) diff --git a/src/guidellm/data/deserializers/trace_mooncake.py b/src/guidellm/data/deserializers/trace_mooncake.py index e2ea2d5bb..22077633f 100644 --- a/src/guidellm/data/deserializers/trace_mooncake.py +++ b/src/guidellm/data/deserializers/trace_mooncake.py @@ -26,7 +26,7 @@ TraceDatasetDeserializer, TraceFormatBase, TraceFormatRegistry, - decode_prompt, + create_prompt_from_hash_ids, generate_token_ids, ) from guidellm.data.schemas import DataArgs @@ -34,19 +34,6 @@ __all__ = ["MooncakeTraceFormatArgs"] -def _is_in_table(hash_id_table: list[Any], hash_id: int) -> bool: - return ( - hash_id < len(hash_id_table) - and hash_id >= 0 - and hash_id_table[hash_id] is not None - ) - - -def _resize_to_hold_id(hash_id_table: list[Any], hash_id: int) -> None: - num_new_entries = hash_id - (len(hash_id_table) - 1) - hash_id_table.extend(None for _ in range(num_new_entries)) - - def _calculate_required_prompt_tokens( config: MooncakeTraceFormatArgs, row: dict, hash_id: int ) -> int: @@ -61,11 +48,11 @@ def _calculate_required_prompt_tokens( def _create_distinct_token_block( block_size: int, - sibling_token_blocks: list[list[int]], + sibling_token_blocks: set[tuple[int, ...]], processor: PreTrainedTokenizerBase, faker: Faker, max_attempts: int = 20, -) -> list[int]: +) -> tuple[int]: """Constructs a new token block of `block_size` that does not appear in `sibling_token_blocks`.""" attempt = 0 @@ -79,20 +66,6 @@ def _create_distinct_token_block( ) -def _create_prompt_from_hash_ids( - hash_ids: list[int], - hash_id_table: list[list[int]], - processor: PreTrainedTokenizerBase, -) -> str: - """Returns a synthetic prompt from `hash_ids` using pre-generated token blocks. - - Precondition: All ids in `hash_ids` appear in `hash_id_table`.""" - prompt_token_ids = [ - token for hash_id in hash_ids for token in hash_id_table[hash_id] - ] - return decode_prompt(processor, prompt_token_ids) - - DatasetDeserializerFactory.register_decorator(TraceDatasetDeserializer, "mooncake") @@ -100,7 +73,7 @@ def _create_prompt_from_hash_ids( class MooncakeTraceFormatArgs(TraceDataArgs): kind: Literal["mooncake"] = Field( default="mooncake", - description="Type identifier for the trace Mooncake dataset deserializer.", + description="Type identifier for the Mooncake trace format.", ) hash_ids_column: str = Field( default="hash_ids", @@ -123,14 +96,15 @@ class MooncakeTraceFormat(TraceFormatBase): blocks in a prompt. The relationships of IDs forms a tree, where every first ID in a prompt has a parent node of `None`. Parent nodes can have an unbounded number of children. Two hash IDs can represent identical blocks of tokens so long - as they do not share the same parent (previous ID). For more details, see section 4 - of https://arxiv.org/pdf/2407.00079. + as they do not share the same parent (previous ID). + + For more details, see section 4 of https://arxiv.org/pdf/2407.00079. Generated prompts match the prompt token count of the row.""" def __init__(self) -> None: - self.hash_id_table: list[Any] = [] - self.sibling_token_blocks: dict[Any, list[list[int]]] = {} + self.hash_id_table: dict[int, tuple[int]] = {} + self.sibling_token_blocks: dict[Any, set[tuple[int, ...]]] = {} def required_columns(self, config: MooncakeTraceFormatArgs) -> Features: return Features({config.hash_ids_column: List(Value("int32"))}) @@ -160,16 +134,15 @@ def create_prompt( each hash ID that has not already been seen.""" ids = row[config.hash_ids_column] for idx, hash_id in enumerate(ids): - if not _is_in_table(self.hash_id_table, hash_id): - _resize_to_hold_id(self.hash_id_table, hash_id) + if hash_id not in self.hash_id_table: prev_id = None if idx == 0 else ids[idx - 1] num_tokens = _calculate_required_prompt_tokens(config, row, hash_id) - self.sibling_token_blocks.setdefault(prev_id, []) + self.sibling_token_blocks.setdefault(prev_id, set()) self.hash_id_table[hash_id] = _create_distinct_token_block( num_tokens, self.sibling_token_blocks[prev_id], processor, faker, ) - self.sibling_token_blocks[prev_id].append(self.hash_id_table[hash_id]) - return _create_prompt_from_hash_ids(ids, self.hash_id_table, processor) + self.sibling_token_blocks[prev_id].add(self.hash_id_table[hash_id]) + return create_prompt_from_hash_ids(ids, self.hash_id_table, processor) diff --git a/src/guidellm/data/deserializers/trace_weka.py b/src/guidellm/data/deserializers/trace_weka.py new file mode 100644 index 000000000..e4030e92e --- /dev/null +++ b/src/guidellm/data/deserializers/trace_weka.py @@ -0,0 +1,205 @@ +""" +The WEKA trace format and data arguments. + +Reads a trace file and yields one row per line with a +synthetic prompt matching the requested input_length for replay +benchmarks. Checks for distinctness between hash IDs that share the +same previous hash ID. + +Generates prompts starting from the first conversation. +When the conversation ends, the next conversation will be used. + +Some features such as subagent conversations, +tool call events and non-linear histories are still missing. +The results from datasets including these features will be unreliable. +""" + +from __future__ import annotations + +import math +from typing import Any, Literal + +from datasets import Features, List, Value +from faker import Faker +from pydantic import Field +from transformers import PreTrainedTokenizerBase + +from guidellm.data.deserializers.deserializer import ( + DataNotSupportedError, + DatasetDeserializerFactory, +) +from guidellm.data.deserializers.trace_common import ( + TraceDataArgs, + TraceDatasetDeserializer, + TraceFormatBase, + TraceFormatRegistry, + create_prompt_from_hash_ids, + decode_prompt, + generate_token_ids, +) +from guidellm.data.schemas import DataArgs + +__all__ = ["WEKATraceFormatArgs"] + + +def _create_distinct_token_block( + block_size: int, + sibling_token_blocks: set[tuple[int, ...]], + processor: PreTrainedTokenizerBase, + faker: Faker, + max_attempts: int = 20, +) -> tuple[int]: + """Constructs a new token block of `block_size` that does not appear in + `sibling_token_blocks`.""" + attempt = 0 + while attempt < max_attempts: + token_ids = generate_token_ids(block_size, processor, faker) + if token_ids not in sibling_token_blocks: + return token_ids + attempt += 1 + raise ValueError( + f"Failed to generate distinct synthetic token block after {attempt} attempts" + ) + + +def _generate_remaining_prompt( + num_tokens: int, processor: PreTrainedTokenizerBase, faker: Faker +) -> str: + if num_tokens == 0: + return "" + token_ids = generate_token_ids(num_tokens, processor, faker) + return decode_prompt(processor, list(token_ids)) + + +DatasetDeserializerFactory.register_decorator(TraceDatasetDeserializer, "weka") + + +@DataArgs.register("weka") +class WEKATraceFormatArgs(TraceDataArgs): + kind: Literal["weka"] = Field( + default="weka", + description="Type identifier for the WEKA trace format.", + ) + timestamp_column: str = Field( + default="t", + description="Column name for timestamps in the trace file.", + ) + prompt_tokens_column: str = Field( + default="in", + description="Column name for prompt token counts in the trace file.", + ) + output_tokens_column: str = Field( + default="out", + description="Column name for output token counts in the trace file.", + ) + conversation_id_column: str = Field( + default="id", + description="Column name for conversation UUIDs in the trace file.", + ) + hash_ids_column: str = Field( + default="hash_ids", + description="Column name for lists of hash IDs in the trace file.", + ) + hash_id_block_size: int = Field( + gt=0, + # Recommended in original github repository callanjfox/agentic-coding-analysis + default=64, + description="Amount of tokens represented by one hash ID.", + ) + + +@TraceFormatRegistry.register("weka") +class WEKATraceFormat(TraceFormatBase): + """WEKA trace format requires a column for timestamps, prompt token counts, + ouput token counts and lists of hash IDs. + + Hash IDs are unique identifiers based on the current and previous token + blocks in a prompt. The relationships of IDs forms a tree, where every first ID + in a prompt has a parent node of `None`. Parent nodes can have an unbounded + number of children. Two hash IDs can represent identical blocks of tokens so long + as they do not share the same parent (previous ID). + + For more details, see [the WEKA trace format specification][trace-spec]. + + [trace-spec]: https://github.com/callanjfox/agentic-coding-analysis/blob/master/docs/TRACE_FORMAT.md + + Generated prompts match the prompt token count of the row.""" + + def __init__(self) -> None: + self.hash_id_table: dict[int, tuple[int]] = {} + self.sibling_token_blocks: dict[Any, set[tuple[int, ...]]] = {} + + def reset(self) -> None: + self.hash_id_table = {} + self.sibling_token_blocks = {} + + def required_columns(self, config: WEKATraceFormatArgs) -> Features: + return Features( + { + config.conversation_id_column: Value("string"), + config.hash_ids_column: List(Value("int32")), + } + ) + + def validate_row(self, config: WEKATraceFormatArgs, row: dict) -> None: + """WEKA format drops what would be the partially filled hash ID at the end of + the chain. Some popular datasets + (e.g. `semianalysisai/cc-traces-weka-no-subagents-051226`) still contain the + trailing hash ID. + In this case, `validate_row` tolerates the addition, and handles it in + `create_prompt`.""" + n_in = row[config.prompt_tokens_column] + n_blocks = len(row[config.hash_ids_column]) + for hash_id in row[config.hash_ids_column]: + if hash_id < 0: + raise DataNotSupportedError( + f"Hash ID must be non-negative, got {hash_id}" + ) + expected = n_in / config.hash_id_block_size + if math.floor(expected) != n_blocks and math.ceil(expected) != n_blocks: + raise DataNotSupportedError( + f"Input token count of {n_in} split into blocks of size " + f"{config.hash_id_block_size} full blocks and " + f"{config.hash_id_block_size} full blocks + partially filled " + f"trailing block does not match given {n_blocks} blocks" + ) + + def create_prompt( + self, + config: WEKATraceFormatArgs, + row: dict, + processor: PreTrainedTokenizerBase, + faker: Faker, + ) -> str: + """Before generating the prompt, this first generates a block of tokens for + each hash ID that has not already been seen. + + Hash IDs that are partially filled are discarded to match the specification. + Remainder of the prompt is created after the creation via hash IDs token + blocks.""" + ids = row[config.hash_ids_column] + expected = row[config.prompt_tokens_column] / config.hash_id_block_size + if math.floor(expected) != len(ids) and math.ceil(expected) == len(ids): + ids.pop() + for idx, hash_id in enumerate(ids): + if hash_id not in self.hash_id_table: + prev_id = None if idx == 0 else ids[idx - 1] + self.sibling_token_blocks.setdefault(prev_id, set()) + self.hash_id_table[hash_id] = _create_distinct_token_block( + config.hash_id_block_size, + self.sibling_token_blocks[prev_id], + processor, + faker, + ) + self.sibling_token_blocks[prev_id].add(self.hash_id_table[hash_id]) + prompt = create_prompt_from_hash_ids(ids, self.hash_id_table, processor) + remainder = _generate_remaining_prompt( + row[config.prompt_tokens_column] % config.hash_id_block_size, + processor, + faker, + ) + if not prompt: + return remainder + if not remainder: + return prompt + return f"{prompt} {remainder}" diff --git a/src/guidellm/utils/json_unwrap.py b/src/guidellm/utils/json_unwrap.py new file mode 100644 index 000000000..d8275a62d --- /dev/null +++ b/src/guidellm/utils/json_unwrap.py @@ -0,0 +1,61 @@ +import dataclasses +import json +from typing import Any + +from datasets import Dataset, IterableDataset + + +@dataclasses.dataclass +class VirtualColumnLocation: + wrapper_column: str + virtual_column: str + + +def construct_virtual_column_locations( + wrapper_column: str, virtual_columns: list[str] +) -> list[VirtualColumnLocation]: + return [VirtualColumnLocation(wrapper_column, c) for c in virtual_columns] + + +def unzip_virtual_column_locations( + column_locations: list[VirtualColumnLocation], +) -> tuple[tuple[str], tuple[str]]: + """Returns a tuple of wrapper columns and a tuple of virtual columns, + in that order.""" + wrapper_cols, virt_cols = tuple( + zip(*(dataclasses.astuple(c) for c in column_locations), strict=True) + ) or ((), ()) + return (wrapper_cols, virt_cols) + + +def try_json_load(json_string: str) -> Any: + try: + return json.loads(json_string) + except (TypeError, json.JSONDecodeError): + return None + + +def is_json_serializable(obj: Any) -> bool: + try: + json.dumps(obj) + return True + except (TypeError, OverflowError): + return False + + +def get_json_column_names(dataset: Dataset | IterableDataset) -> list[str]: + """Assumes dataset has at least one column and at least one row. + + Returns a list of all columns in the dataset containing valid JSON. This includes + columns containing lists of valid JSON, as well as Python dictionaries that can be + serialized to JSON with no issue. + """ + sample = next(iter(dataset)) + sample = {k: (v[0] if isinstance(v, list) and v else v) for k, v in sample.items()} + column_names = dataset.column_names or list(next(iter(dataset)).keys()) + return [ + col + for col in column_names + if (isinstance(sample[col], dict) and is_json_serializable(sample[col])) + or try_json_load(sample[col]) is not None + ] diff --git a/tests/unit/data/deserializers/test_trace_common.py b/tests/unit/data/deserializers/test_trace_common.py index babd2f389..fc4ca4f94 100644 --- a/tests/unit/data/deserializers/test_trace_common.py +++ b/tests/unit/data/deserializers/test_trace_common.py @@ -1,4 +1,5 @@ import dataclasses +import json from collections.abc import Callable from pathlib import Path from typing import Any @@ -20,7 +21,7 @@ from guidellm.data.deserializers.trace_minimal import MinimalTraceFormatArgs -def _mock_processor() -> Mock: +def mock_processor() -> Mock: """Tokenizer where each whitespace-delimited word is one token.""" proc = Mock() proc.encode.side_effect = lambda text: list(range(len(text.split()))) @@ -40,21 +41,21 @@ def _mock_processor() -> Mock: ], ) def test_decode_prompt(token_ids, expected): - proc = _mock_processor() + proc = mock_processor() assert decode_prompt(proc, token_ids) == expected @pytest.mark.parametrize( ("token_count", "expected"), [ - (0, []), - (1, [0]), - (10, [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]), - (1000, list(range(1000))), + (0, ()), + (1, (0,)), + (10, (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)), + (1000, tuple(range(1000))), ], ) def test_generate_token_ids(token_count, expected): - proc = _mock_processor() + proc = mock_processor() faker = Faker() res = generate_token_ids(token_count, proc, faker) assert len(res) == len(expected) @@ -75,22 +76,25 @@ class TraceColumnGenerator: data_generator: Callable[[int], Any] -def _write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: +def write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: path = tmp_path / f"trace{suffix}" path.write_text(content) return path -def _generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: +def generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: + """Returns valid JSON lines.""" return "\n".join( "{" - + ", ".join(f'"{col.name}": {col.data_generator(idx)}' for col in columns) + + ", ".join( + f'"{col.name}": {col.data_generator(idx)}' for col in columns + ).replace("'", '"') + "}" for idx in range(num_rows) ) -def _get_from_kwargs(keys, kwargs) -> dict: +def get_from_kwargs(keys, kwargs) -> dict: return {k: v for k, v in kwargs.items() if k in keys} @@ -99,8 +103,8 @@ class TestTraceDatasetDeserializer: def deserializer(self) -> TraceDatasetDeserializer: return TraceDatasetDeserializer() - def _deserialize(self, deserializer, data, **kwargs): - col_kwargs = _get_from_kwargs( + def deserialize(self, deserializer, data, **kwargs): + col_kwargs = get_from_kwargs( ( "timestamp_column", "prompt_tokens_column", @@ -111,7 +115,7 @@ def _deserialize(self, deserializer, data, **kwargs): config = MinimalTraceFormatArgs(path=data, **col_kwargs) return deserializer( config=config, - processor_factory=_mock_processor, + processor_factory=mock_processor, random_seed=42, ) @@ -121,13 +125,13 @@ def _deserialize(self, deserializer, data, **kwargs): [".json", ".jsonl"], ) def test_loads_json(self, tmp_path: Path, deserializer, suffix): - trace = _write_trace( + trace = write_trace( tmp_path, '{"timestamp": 1, "input_length": 10, "output_length": 1}\n' '{"timestamp": 2, "input_length": 20, "output_length": 2}\n', suffix=suffix, ) - ds = self._deserialize(deserializer, trace) + ds = self.deserialize(deserializer, trace) for i, row in enumerate(ds): assert row["relative_timestamp"] == i assert row["prompt_tokens_count"] == (i + 1) * 10 @@ -135,49 +139,101 @@ def test_loads_json(self, tmp_path: Path, deserializer, suffix): @pytest.mark.sanity def test_loads_csv(self, tmp_path: Path, deserializer): - trace = _write_trace( + trace = write_trace( tmp_path, "timestamp,input_length,output_length\n1,10,1\n2,20,2\n", suffix=".csv", ) - ds = self._deserialize(deserializer, trace) + ds = self.deserialize(deserializer, trace) for i, row in enumerate(ds): assert row["relative_timestamp"] == i assert row["prompt_tokens_count"] == (i + 1) * 10 assert row["output_tokens_count"] == i + 1 - @pytest.mark.smoke - def test_loads_sorted_rows_and_keeps_token_columns_aligned( - self, tmp_path: Path, deserializer - ): - n_rows = 10 - trace = _write_trace( - tmp_path, - _generate_trace( + @pytest.fixture + def example_trace(self): + def generator(tmp_path: Path, n_rows: int) -> str: + trace_rows = generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: n_rows - i), TraceColumnGenerator("input_length", lambda i: n_rows - i), TraceColumnGenerator("output_length", lambda i: (n_rows - i) * 10), ], - ), - ) - ds = self._deserialize(deserializer, trace) + ) + return write_trace(tmp_path, trace_rows) + + return generator + + @pytest.fixture + def example_list_json_trace(self): + def generator(tmp_path: Path, n_rows: int) -> str: + trace_rows = generate_trace( + n_rows, + [ + TraceColumnGenerator("timestamp", lambda i: n_rows - i), + TraceColumnGenerator("input_length", lambda i: n_rows - i), + TraceColumnGenerator("output_length", lambda i: (n_rows - i) * 10), + ], + ) + json_data = [json.loads(s) for s in trace_rows.split("\n")] + return write_trace( + tmp_path, + generate_trace( + 1, [TraceColumnGenerator("requests", lambda _: json_data)] + ), + ) + + return generator + + @pytest.mark.parametrize( + "example", + ["example_trace", "example_list_json_trace"], + ) + @pytest.mark.smoke + def test_loads_sorted_rows_and_keeps_token_columns_aligned( + self, tmp_path: Path, deserializer, example, request + ): + trace_factory = request.getfixturevalue(example) + trace = trace_factory(tmp_path, 10) + ds = self.deserialize(deserializer, trace) assert isinstance(ds, IterableDataset) - proc = _mock_processor() + proc = mock_processor() for i, row in enumerate(ds): assert row["prompt_tokens_count"] == i + 1 assert row["output_tokens_count"] == (i + 1) * 10 assert len(proc.encode(row["prompt"])) == row["prompt_tokens_count"] + @pytest.mark.sanity + def test_loads_requests_column_stored_as_json_string( + self, tmp_path: Path, deserializer + ): + """Unwrap when the wrapper column is a JSON string of a list, not a native list. + + ## WRITTEN BY AI ## + """ + trace = write_trace( + tmp_path, + '{"requests": "[{\\"timestamp\\": 0, \\"input_length\\": 10,' + ' \\"output_length\\": 5}, {\\"timestamp\\": 1, \\"input_length\\": 20,' + ' \\"output_length\\": 10}]"}\n', + ) + ds = self.deserialize(deserializer, trace) + rows = list(ds) + assert len(rows) == 2 + assert rows[0]["prompt_tokens_count"] == 10 + assert rows[0]["output_tokens_count"] == 5 + assert rows[1]["prompt_tokens_count"] == 20 + assert rows[1]["output_tokens_count"] == 10 + @pytest.mark.smoke def test_emits_relative_timestamp_column_sorted_from_trace( self, tmp_path: Path, deserializer ): n_rows = 5 - trace = _write_trace( + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: i + 3), @@ -186,18 +242,18 @@ def test_emits_relative_timestamp_column_sorted_from_trace( ], ), ) - ds = self._deserialize(deserializer, trace) + ds = self.deserialize(deserializer, trace) for i, row in enumerate(ds): assert row["relative_timestamp"] == i @pytest.mark.smoke def test_rejects_invalid_path(self, deserializer): with pytest.raises(ValidationError, match="not a valid path"): - self._deserialize(deserializer, 123) + self.deserialize(deserializer, 123) with pytest.raises(DataNotSupportedError, match="file not found"): - self._deserialize(deserializer, "bad_path.jsonl") + self.deserialize(deserializer, "bad_path.jsonl") with pytest.raises(DataNotSupportedError, match="not a file"): - self._deserialize(deserializer, Path.cwd()) + self.deserialize(deserializer, Path.cwd()) @pytest.mark.sanity @pytest.mark.parametrize( @@ -252,17 +308,27 @@ def test_rejects_invalid_path(self, deserializer): def test_trace_validation_raises( self, tmp_path: Path, deserializer, content, kwargs, match ): - trace = _write_trace(tmp_path, content) + trace = write_trace(tmp_path, content) with pytest.raises(DataNotSupportedError, match=match): - self._deserialize(deserializer, trace, **kwargs) + self.deserialize(deserializer, trace, **kwargs) + + @pytest.mark.sanity + def test_malformed_json_columns_raises( + self, tmp_path: Path, deserializer, example_list_json_trace + ): + trace = example_list_json_trace(tmp_path, 2) + data = trace.read_text().replace("timestamp", "ts") + trace.write_text(data) + with pytest.raises(DataNotSupportedError, match="lists of JSON objects"): + self.deserialize(deserializer, trace) @pytest.mark.sanity def test_unsupported_file_suffix_raises(self, tmp_path: Path, deserializer): - trace = _write_trace( + trace = write_trace( tmp_path, '{"timestamp": 0, "input_length": 10, "output_length": 5, ' '"hash_ids": [0]}\n', suffix=".txt", ) with pytest.raises(DataNotSupportedError, match=r"Unsupported.*\.txt"): - self._deserialize(deserializer, trace) + self.deserialize(deserializer, trace) diff --git a/tests/unit/data/deserializers/test_trace_minimal.py b/tests/unit/data/deserializers/test_trace_minimal.py index fddda4f38..ee9e0c8b2 100644 --- a/tests/unit/data/deserializers/test_trace_minimal.py +++ b/tests/unit/data/deserializers/test_trace_minimal.py @@ -14,7 +14,7 @@ from guidellm.data.deserializers.trace_minimal import MinimalTraceFormatArgs -def _mock_processor() -> Mock: +def mock_processor() -> Mock: """Tokenizer where each whitespace-delimited word is one token.""" proc = Mock() proc.encode.side_effect = lambda text: list(range(len(text.split()))) @@ -24,7 +24,7 @@ def _mock_processor() -> Mock: return proc -def _write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: +def write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: path = tmp_path / f"trace{suffix}" path.write_text(content) return path @@ -37,7 +37,7 @@ class TraceColumnGenerator: data_generator: Callable[[int], Any] -def _generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: +def generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: return "\n".join( "{" + ", ".join(f'"{col.name}": {col.data_generator(idx)}' for col in columns) @@ -46,20 +46,20 @@ def _generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: ) -def _get_from_kwargs(keys, kwargs) -> dict: +def get_from_kwargs(keys, kwargs) -> dict: return {k: v for k, v in kwargs.items() if k in keys} class TestMinimalTraceFormat: @pytest.mark.regression def test_format_registered_with_deserializer(self, tmp_path: Path): - trace = _write_trace( + trace = write_trace( tmp_path, '{"timestamp": 0.0, "input_length": 10, "output_length": 5}\n', ) DatasetDeserializerFactory.deserialize( config=MinimalTraceFormatArgs(path=trace), - processor_factory=_mock_processor, + processor_factory=mock_processor, random_seed=42, ) @@ -67,8 +67,8 @@ def test_format_registered_with_deserializer(self, tmp_path: Path): def deserializer(self) -> TraceDatasetDeserializer: return TraceDatasetDeserializer() - def _deserialize(self, deserializer, data, **kwargs): - col_kwargs = _get_from_kwargs( + def deserialize(self, deserializer, data, **kwargs): + col_kwargs = get_from_kwargs( ( "timestamp_column", "prompt_tokens_column", @@ -79,19 +79,19 @@ def _deserialize(self, deserializer, data, **kwargs): config = MinimalTraceFormatArgs(path=data, **col_kwargs) return deserializer( config=config, - processor_factory=_mock_processor, + processor_factory=mock_processor, random_seed=42, ) @pytest.mark.smoke def test_honors_custom_column_names(self, tmp_path: Path, deserializer): - trace = _write_trace( + trace = write_trace( tmp_path, '{"ts": 3.0, "input_tokens": 4, "generated_tokens": 40}\n' '{"ts": 1.0, "input_tokens": 2, "generated_tokens": 20}\n', ) - ds = self._deserialize( + ds = self.deserialize( deserializer, trace, timestamp_column="ts", @@ -115,9 +115,9 @@ def test_generates_large_trace_prompts_from_reusable_base( output_lengths = [random.randint(3, 800) for _ in range(n_rows)] times = [0.0, 0.5, 1.0, 2.0] timestamps = [times[int(i / n_rows * len(times))] for i in range(n_rows)] - trace = _write_trace( + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: timestamps[i]), @@ -126,7 +126,7 @@ def test_generates_large_trace_prompts_from_reusable_base( ], ), ) - processor = _mock_processor() + processor = mock_processor() config = MinimalTraceFormatArgs(path=trace) ds = deserializer( config=config, diff --git a/tests/unit/data/deserializers/test_trace_mooncake.py b/tests/unit/data/deserializers/test_trace_mooncake.py index 0875315c1..a9506d098 100644 --- a/tests/unit/data/deserializers/test_trace_mooncake.py +++ b/tests/unit/data/deserializers/test_trace_mooncake.py @@ -15,7 +15,7 @@ from guidellm.data.schemas import DataNotSupportedError -def _ascending_processor() -> Mock: +def ascending_processor() -> Mock: """Tokenizer where each whitespace-delimited word is assigned a token in ascending order starting from 0. This is incompatible with most mooncake traces as there is no way to generate distinct token blocks for sibling @@ -28,7 +28,7 @@ def _ascending_processor() -> Mock: return proc -def _compatible_processor() -> Mock: +def compatible_processor() -> Mock: """Tokenizer where each whitespace-delimited word is assigned a token selected from a range of random integers. This is compatible with most mooncake traces as there is a way to generate distinct token blocks for @@ -44,20 +44,19 @@ def _compatible_processor() -> Mock: return proc -def _write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: +def write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: path = tmp_path / f"trace{suffix}" path.write_text(content) return path -def _make_valid_hash_ids( - n_rows: int, prompt_lengths: list[int], block_size: int -) -> list[list[int]]: +def make_valid_hash_ids(prompt_lengths: list[int], block_size: int) -> list[list[int]]: """The final token block of every row may be less than the hash id block size due to the prompt length not being divisible by it. Use this when testing large trace prompts to avoid including token blocks with less than the block size in the middle of later rows.""" tail_hash_ids = [] + n_rows = len(prompt_lengths) original_prompt_positions = dict(zip(prompt_lengths, range(n_rows), strict=False)) sorted_lengths = copy.deepcopy(prompt_lengths) sorted_lengths.sort() @@ -73,11 +72,11 @@ def _make_valid_hash_ids( return hash_ids -def _all_equal(items: list): +def all_equal(items: list): return len(set(items)) == 1 -def _all_distinct(items: list): +def all_distinct(items: list): seen = set() return not any(i in seen or seen.add(i) for i in items) @@ -89,7 +88,7 @@ class TraceColumnGenerator: data_generator: Callable[[int], Any] -def _generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: +def generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: return "\n".join( "{" + ", ".join(f'"{col.name}": {col.data_generator(idx)}' for col in columns) @@ -98,30 +97,34 @@ def _generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: ) -def _get_from_kwargs(keys, kwargs) -> dict: +def get_from_kwargs(keys, kwargs) -> dict: return {k: v for k, v in kwargs.items() if k in keys} class TestMooncakeTraceFormat: @pytest.mark.regression def test_format_registered_with_deserializer(self, tmp_path: Path): - trace = _write_trace( + trace = write_trace( tmp_path, '{"timestamp": 0.0, "input_length": 10, "output_length": 5, ' '"hash_ids": [0]}\n', ) DatasetDeserializerFactory.deserialize( config=MooncakeTraceFormatArgs(path=trace), - processor_factory=_ascending_processor, + processor_factory=ascending_processor, random_seed=42, ) + @pytest.fixture + def default_block_size(self, tmp_path: Path) -> int: + return MooncakeTraceFormatArgs(path=tmp_path).hash_id_block_size + @pytest.fixture def deserializer(self) -> TraceDatasetDeserializer: return TraceDatasetDeserializer() - def _deserialize(self, deserializer, data, **kwargs): - col_kwargs = _get_from_kwargs( + def deserialize(self, deserializer, data, **kwargs): + col_kwargs = get_from_kwargs( ( "timestamp_column", "prompt_tokens_column", @@ -134,16 +137,16 @@ def _deserialize(self, deserializer, data, **kwargs): config = MooncakeTraceFormatArgs(path=data, **col_kwargs) return deserializer( config=config, - processor_factory=_ascending_processor, + processor_factory=ascending_processor, random_seed=42, ) @pytest.mark.smoke def test_honors_custom_column_names(self, tmp_path: Path, deserializer): n_rows = 3 - trace = _write_trace( + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("ts", lambda i: i), @@ -153,7 +156,7 @@ def test_honors_custom_column_names(self, tmp_path: Path, deserializer): ], ), ) - self._deserialize( + self.deserialize( deserializer, trace, timestamp_column="ts", @@ -166,9 +169,9 @@ def test_honors_custom_column_names(self, tmp_path: Path, deserializer): def test_custom_hash_id_block_size(self, tmp_path: Path, deserializer): n_rows = 1 n_in = 1000 - trace = _write_trace( + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: i), @@ -180,21 +183,22 @@ def test_custom_hash_id_block_size(self, tmp_path: Path, deserializer): ], ), ) - self._deserialize(deserializer, trace, hash_id_block_size=n_in / 5) + self.deserialize(deserializer, trace, hash_id_block_size=n_in / 5) @pytest.mark.smoke - def test_generates_large_trace_prompts(self, tmp_path: Path, deserializer): + def test_generates_large_trace_prompts( + self, tmp_path: Path, deserializer, default_block_size + ): random.seed(0) n_rows = 25 prompt_lengths = [random.randint(2000, 100000) for _ in range(n_rows)] output_lengths = [random.randint(3, 800) for _ in range(n_rows)] times = [0.0, 0.5, 1.0, 2.0] timestamps = [times[int(i / n_rows * len(times))] for i in range(n_rows)] - block_size = MooncakeTraceFormatArgs(path=tmp_path).hash_id_block_size - hash_ids = _make_valid_hash_ids(n_rows, prompt_lengths, block_size) - trace = _write_trace( + hash_ids = make_valid_hash_ids(prompt_lengths, default_block_size) + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: timestamps[i]), @@ -204,14 +208,8 @@ def test_generates_large_trace_prompts(self, tmp_path: Path, deserializer): ], ), ) - processor = _ascending_processor() - config = MooncakeTraceFormatArgs(path=trace) - ds = deserializer( - config=config, - processor_factory=lambda: processor, - random_seed=42, - ) - + processor = ascending_processor() + ds = self.deserialize(deserializer, trace) for i, row in enumerate(ds): in_cnt = row["prompt_tokens_count"] assert in_cnt == prompt_lengths[i] @@ -221,6 +219,36 @@ def test_generates_large_trace_prompts(self, tmp_path: Path, deserializer): if actual_prompt_length != in_cnt: pytest.fail(f"{actual_prompt_length} != {in_cnt}") + @pytest.mark.smoke + def test_prompt_matching_or_bordering_block_size( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 3 + n_in = list(range(default_block_size - 1, default_block_size + 2)) + hash_ids = make_valid_hash_ids(n_in, default_block_size) + trace = write_trace( + tmp_path, + generate_trace( + n_rows, + [ + TraceColumnGenerator("timestamp", lambda i: i), + TraceColumnGenerator("input_length", lambda i: n_in[i]), + TraceColumnGenerator("output_length", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda i: hash_ids[i]), + ], + ), + ) + processor = compatible_processor() + ds = deserializer( + config=MooncakeTraceFormatArgs(path=trace), + processor_factory=lambda: processor, + random_seed=42, + ) + for row in ds: + actual_prompt_length = len(processor.encode(row["prompt"])) + if actual_prompt_length != row["prompt_tokens_count"]: + pytest.fail(f"{actual_prompt_length} != {row['prompt_tokens_count']}") + @pytest.mark.sanity @pytest.mark.parametrize( ("content", "kwargs", "match"), @@ -242,20 +270,23 @@ def test_generates_large_trace_prompts(self, tmp_path: Path, deserializer): def test_trace_validation_raises( self, tmp_path: Path, deserializer, content, kwargs, match ): - trace = _write_trace(tmp_path, content) + trace = write_trace(tmp_path, content) with pytest.raises(DataNotSupportedError, match=match): - self._deserialize(deserializer, trace, **kwargs) + self.deserialize(deserializer, trace, **kwargs) @pytest.mark.sanity - def test_incompatible_encoding_raises(self, tmp_path: Path, deserializer): + def test_incompatible_encoding_raises( + self, tmp_path: Path, deserializer, default_block_size + ): n_rows = 2 - trace = _write_trace( + n_in = default_block_size * 2 + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: i), - TraceColumnGenerator("input_length", lambda _: 1024), + TraceColumnGenerator("input_length", lambda _: n_in), TraceColumnGenerator("output_length", lambda _: 5), TraceColumnGenerator("hash_ids", lambda i: [0, i + 1]), ], @@ -264,7 +295,7 @@ def test_incompatible_encoding_raises(self, tmp_path: Path, deserializer): config = MooncakeTraceFormatArgs(path=trace) ds = deserializer( config=config, - processor_factory=_ascending_processor, + processor_factory=ascending_processor, random_seed=42, ) with pytest.raises(ValueError, match="generate distinct"): @@ -275,9 +306,9 @@ def test_incompatible_encoding_raises(self, tmp_path: Path, deserializer): def test_token_block_distinctness(self, tmp_path: Path, deserializer): n_rows = 4 n_in = 1024 - trace = _write_trace( + trace = write_trace( tmp_path, - _generate_trace( + generate_trace( n_rows, [ TraceColumnGenerator("timestamp", lambda i: i), @@ -287,15 +318,14 @@ def test_token_block_distinctness(self, tmp_path: Path, deserializer): ], ), ) - config = MooncakeTraceFormatArgs(path=trace) ds = deserializer( - config=config, - processor_factory=_compatible_processor, + config=MooncakeTraceFormatArgs(path=trace), + processor_factory=compatible_processor, random_seed=42, ) root_blocks, sibling_blocks = zip( *[(row["prompt"][: n_in // 2], row["prompt"][n_in // 2 :]) for row in ds], strict=False, ) - assert _all_equal(root_blocks) - assert _all_distinct(sibling_blocks) + assert all_equal(root_blocks) + assert all_distinct(sibling_blocks) diff --git a/tests/unit/data/deserializers/test_trace_weka.py b/tests/unit/data/deserializers/test_trace_weka.py new file mode 100644 index 000000000..9996a0dfa --- /dev/null +++ b/tests/unit/data/deserializers/test_trace_weka.py @@ -0,0 +1,458 @@ +import dataclasses +import json +import math +import random +from collections.abc import Callable +from pathlib import Path +from typing import Any +from unittest.mock import Mock + +import pytest + +from guidellm.data.deserializers import DatasetDeserializerFactory +from guidellm.data.deserializers.trace_common import TraceDatasetDeserializer +from guidellm.data.deserializers.trace_weka import WEKATraceFormatArgs +from guidellm.data.schemas import DataNotSupportedError + + +def ascending_processor() -> Mock: + """Tokenizer where each whitespace-delimited word is assigned a token + in ascending order starting from 0. This is incompatible with most WEKA + traces as there is no way to generate distinct token blocks for sibling + nodes.""" + proc = Mock() + proc.encode.side_effect = lambda text: list(range(len(text.split()))) + proc.decode.side_effect = lambda tokens, skip_special_tokens=False: " ".join( + f"tok{i}" for i, _ in enumerate(tokens) + ) + return proc + + +def compatible_processor() -> Mock: + """Tokenizer where each whitespace-delimited word is assigned a token + selected from a range of random integers. This is compatible with most + WEKA traces as there is a way to generate distinct token blocks for + sibling nodes.""" + random.seed(0) + proc = Mock() + proc.encode.side_effect = lambda text: [ + random.randint(0, 1000) for _ in range(len(text.split())) + ] + proc.decode.side_effect = lambda tokens, skip_special_tokens=False: " ".join( + f"tok{t}" for t in tokens + ) + return proc + + +def write_trace(tmp_path: Path, content: str, suffix: str = ".jsonl") -> Path: + path = tmp_path / f"trace{suffix}" + path.write_text(content) + return path + + +def make_valid_hash_ids(prompt_lengths: list[int], block_size: int) -> list[list[int]]: + """Hash IDs in WEKA format start at 1, and only create IDs for full token + blocks.""" + hash_ids = [] + for length in prompt_lengths: + n_ids = math.floor(length / block_size) + hash_ids.append(list(range(1, n_ids + 1))) + return hash_ids + + +@dataclasses.dataclass +class TraceColumnGenerator: + name: str + # Function with row index as the one argument + data_generator: Callable[[int], Any] + + +def generate_trace(num_rows: int, columns: list[TraceColumnGenerator]) -> str: + """Returns valid JSON lines.""" + return "\n".join( + "{" + + ", ".join( + f'"{col.name}": {col.data_generator(idx)}' for col in columns + ).replace("'", '"') + + "}" + for idx in range(num_rows) + ) + + +def generate_weka_trace( + num_rows: int, + num_virtual_rows: int, + non_wrapper_cols: list[TraceColumnGenerator], + virtual_cols: list[TraceColumnGenerator], +) -> str: + trace_rows = generate_trace(num_virtual_rows, virtual_cols) + json_data = [json.loads(s) for s in trace_rows.split("\n")] + cols = non_wrapper_cols + [TraceColumnGenerator("requests", lambda _: json_data)] + return generate_trace(num_rows, cols) + + +def get_from_kwargs(keys, kwargs) -> dict: + return {k: v for k, v in kwargs.items() if k in keys} + + +def all_equal(items: list): + return len(set(items)) == 1 + + +def all_distinct(items: list): + seen = set() + return not any(i in seen or seen.add(i) for i in items) + + +class TestWEKATraceFormat: + @pytest.mark.regression + def test_format_registered_with_deserializer(self, tmp_path: Path): + trace = write_trace( + tmp_path, + '{"id": "conv0", "requests": [{"t": 0, "in": 10,' + '"out": 5, "hash_ids": []}]}\n', + ) + DatasetDeserializerFactory.deserialize( + config=WEKATraceFormatArgs(path=trace), + processor_factory=ascending_processor, + random_seed=42, + ) + + @pytest.fixture + def default_block_size(self, tmp_path: Path) -> int: + return WEKATraceFormatArgs(path=tmp_path).hash_id_block_size + + @pytest.fixture + def deserializer(self) -> TraceDatasetDeserializer: + return TraceDatasetDeserializer() + + def deserialize(self, deserializer, data, **kwargs): + col_kwargs = get_from_kwargs( + ( + "conversation_id_column", + "timestamp_column", + "prompt_tokens_column", + "output_tokens_column", + "hash_ids_column", + "hash_id_block_size", + ), + kwargs, + ) + config = WEKATraceFormatArgs(path=data, **col_kwargs) + return deserializer( + config=config, + processor_factory=ascending_processor, + random_seed=42, + ) + + @pytest.mark.smoke + def test_honors_custom_column_names(self, tmp_path: Path, deserializer): + n_rows = 1 + n_virtual_rows = 3 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("conv_id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("timestamp", lambda i: i), + TraceColumnGenerator("input_tokens", lambda i: i + 1), + TraceColumnGenerator("generated_tokens", lambda i: (i + 1) * 10), + TraceColumnGenerator("ids", lambda _: []), + ], + ), + ) + self.deserialize( + deserializer, + trace, + conversation_id_column="conv_id", + timestamp_column="timestamp", + prompt_tokens_column="input_tokens", + output_tokens_column="generated_tokens", + hash_ids_column="ids", + ) + + @pytest.mark.smoke + def test_custom_hash_id_block_size(self, tmp_path: Path, deserializer): + n_rows = 1 + n_virtual_rows = 1 + n_in = 1000 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: n_in), + TraceColumnGenerator("out", lambda i: i + 1), + # Would throw a DataNotSupportedError with default block size 64 + # See row validation in trace_weka.py + TraceColumnGenerator("hash_ids", lambda _: [1, 2, 3, 4, 5]), + ], + ), + ) + self.deserialize(deserializer, trace, hash_id_block_size=n_in / 5) + + @pytest.mark.smoke + def test_generates_large_trace_prompts( + self, tmp_path: Path, deserializer, default_block_size + ): + random.seed(0) + n_rows = 1 + n_virtual_rows = 25 + prompt_lengths = [random.randint(2000, 100000) for _ in range(n_virtual_rows)] + output_lengths = [random.randint(3, 800) for _ in range(n_virtual_rows)] + times = [0.0, 0.5, 1.0, 2.0] + timestamps = [ + times[int(i / n_virtual_rows * len(times))] for i in range(n_virtual_rows) + ] + hash_ids = make_valid_hash_ids(prompt_lengths, default_block_size) + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: timestamps[i]), + TraceColumnGenerator("in", lambda i: prompt_lengths[i]), + TraceColumnGenerator("out", lambda i: output_lengths[i]), + TraceColumnGenerator("hash_ids", lambda i: hash_ids[i]), + ], + ), + ) + processor = ascending_processor() + ds = self.deserialize(deserializer, trace) + for i, row in enumerate(ds): + in_cnt = row["prompt_tokens_count"] + assert in_cnt == prompt_lengths[i] + assert row["output_tokens_count"] == output_lengths[i] + + actual_prompt_length = len(processor.encode(row["prompt"])) + if actual_prompt_length != in_cnt: + pytest.fail(f"{actual_prompt_length} != {in_cnt}") + + @pytest.mark.smoke + def test_prompt_matching_or_bordering_block_size( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 1 + n_virtual_rows = 3 + n_in = list(range(default_block_size - 1, default_block_size + 2)) + hash_ids = make_valid_hash_ids(n_in, default_block_size) + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda i: n_in[i]), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda i: hash_ids[i]), + ], + ), + ) + processor = ascending_processor() + ds = self.deserialize(deserializer, trace) + for row in ds: + actual_prompt_length = len(processor.encode(row["prompt"])) + if actual_prompt_length != row["prompt_tokens_count"]: + pytest.fail(f"{actual_prompt_length} != {row['prompt_tokens_count']}") + + @pytest.mark.sanity + def test_removes_partially_filled_hash_ids( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 1 + n_virtual_rows = 2 + n_in = default_block_size + 1 + hash_ids = make_valid_hash_ids([n_in], default_block_size)[0] + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: n_in), + TraceColumnGenerator("out", lambda i: i), + TraceColumnGenerator("hash_ids", lambda i: hash_ids + [i + 2]), + ], + ), + ) + processor = ascending_processor() + ds = self.deserialize(deserializer, trace) + for row in ds: + actual_prompt_length = len(processor.encode(row["prompt"])) + if actual_prompt_length != row["prompt_tokens_count"]: + pytest.fail(f"{actual_prompt_length} != {row['prompt_tokens_count']}") + + @pytest.mark.sanity + @pytest.mark.parametrize( + ("content", "kwargs", "match"), + [ + ( + '{"id": "conv0", "requests": [{"t": 0, "in": 10,' + '"out": 5, "hash_ids": [-1]}]}\n', + {}, + "non-negative", + ), + ( + '{"id": "conv0", "requests": [{"t": 0, "in": 1024,' + '"out": 5, "hash_ids": [1]}]}\n', + {}, + "given 1 blocks", + ), + ], + ) + def test_trace_validation_raises( + self, tmp_path: Path, deserializer, content, kwargs, match + ): + trace = write_trace(tmp_path, content) + with pytest.raises(DataNotSupportedError, match=match): + self.deserialize(deserializer, trace, **kwargs) + + @pytest.mark.sanity + def test_incompatible_encoding_raises( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 1 + n_virtual_rows = 2 + n_in = default_block_size * 2 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: n_in), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda i: [1, i + 2]), + ], + ), + ) + ds = deserializer( + config=WEKATraceFormatArgs(path=trace), + processor_factory=ascending_processor, + random_seed=42, + ) + with pytest.raises(ValueError, match="generate distinct"): + for _ in ds: + ... + + @pytest.mark.smoke + def test_token_block_distinctness( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 1 + n_virtual_rows = 4 + n_in = default_block_size * 2 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: n_in), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda i: [1, i + 2]), + ], + ), + ) + ds = deserializer( + config=WEKATraceFormatArgs(path=trace), + processor_factory=compatible_processor, + random_seed=42, + ) + root_blocks, sibling_blocks = zip( + *[(row["prompt"][: n_in // 2], row["prompt"][n_in // 2 :]) for row in ds], + strict=False, + ) + assert all_equal(root_blocks) + assert all_distinct(sibling_blocks) + + @pytest.mark.smoke + def test_multi_conversation_resets_relative_timestamp( + self, tmp_path: Path, deserializer + ): + n_rows = 2 + n_virtual_rows = 3 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: 10), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda _: []), + ], + ), + ) + ds = self.deserialize(deserializer, trace) + timestamps = [row["relative_timestamp"] for row in ds] + assert timestamps[0] == 0.0 + assert timestamps[1] != 0.0 + assert timestamps[3] == 0.0 + + @pytest.mark.sanity + def test_multi_conversation_resets_hash_id_state( + self, tmp_path: Path, deserializer, default_block_size + ): + n_rows = 2 + n_virtual_rows = 2 + n_in = default_block_size * 2 + trace = write_trace( + tmp_path, + generate_weka_trace( + n_rows, + n_virtual_rows, + [TraceColumnGenerator("id", lambda i: f'"conv{i}"')], + [ + TraceColumnGenerator("t", lambda i: i), + TraceColumnGenerator("in", lambda _: n_in), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda i: [1, i + 2]), + ], + ), + ) + ds = deserializer( + config=WEKATraceFormatArgs(path=trace), + processor_factory=compatible_processor, + random_seed=42, + ) + prompts = [row["prompt"] for row in ds] + assert prompts[0] != prompts[2] + + @pytest.mark.sanity + def test_zero_prompt_tokens_empty_hash_ids(self, tmp_path: Path, deserializer): + trace = write_trace( + tmp_path, + generate_weka_trace( + 1, + 1, + [TraceColumnGenerator("id", lambda _: '"conv0"')], + [ + TraceColumnGenerator("t", lambda _: 0.0), + TraceColumnGenerator("in", lambda _: 0), + TraceColumnGenerator("out", lambda _: 5), + TraceColumnGenerator("hash_ids", lambda _: []), + ], + ), + ) + ds = self.deserialize(deserializer, trace) + rows = list(ds) + assert rows[0]["prompt_tokens_count"] == 0 + assert rows[0]["output_tokens_count"] == 5 diff --git a/tests/unit/utils/test_json_unwrap.py b/tests/unit/utils/test_json_unwrap.py new file mode 100644 index 000000000..2f7e6c6a3 --- /dev/null +++ b/tests/unit/utils/test_json_unwrap.py @@ -0,0 +1,128 @@ +import datetime + +import pytest +from datasets import Dataset, IterableDataset + +from guidellm.utils.json_unwrap import ( + VirtualColumnLocation, + construct_virtual_column_locations, + get_json_column_names, + try_json_load, + unzip_virtual_column_locations, +) + + +@pytest.mark.smoke +@pytest.mark.parametrize( + ("wrapper_col", "virtual_cols", "expected"), + [ + ("", [], []), + ( + "wrapper", + ["field_1", "field_2"], + [ + VirtualColumnLocation("wrapper", "field_1"), + VirtualColumnLocation("wrapper", "field_2"), + ], + ), + ("wrapper", [], []), + ], +) +def test_construct_virtual_column_locations(wrapper_col, virtual_cols, expected): + actual = construct_virtual_column_locations(wrapper_col, virtual_cols) + if not actual: + assert actual == expected + for location in actual: + assert isinstance(location, VirtualColumnLocation) + + +@pytest.mark.smoke +@pytest.mark.parametrize( + ("arg", "expected"), + [ + ("", None), + (r"{}", {}), + ("'string'", None), + ("'1'", None), + ( + r'{"field_1": 1, "field_2": "two", "field_3": [3, 4]}', + {"field_1": 1, "field_2": "two", "field_3": [3, 4]}, + ), + ( + r'[{"field_1": 1, "field_2": "two"}, {"field_3": [3, 4]}]', + [{"field_1": 1, "field_2": "two"}, {"field_3": [3, 4]}], + ), + ], +) +def test_try_json_load(arg, expected): + assert try_json_load(arg) == expected + + +def one_row_generator(row: dict): + yield {key: value[0] for key, value in row.items()} + + +@pytest.mark.smoke +@pytest.mark.parametrize( + ("data", "expected"), + [ + ({"field": [1]}, []), + ({"wrapper": [r'{"inner_field": 2}']}, ["wrapper"]), + ({"field": [1], "wrapper": [r'{"inner_field": 2}']}, ["wrapper"]), + ( + {"field": [1], "wrapper": [r'{"inner_field": 2}'], "field_2": ["two"]}, + ["wrapper"], + ), + ({"wrapper": [r"{}"]}, ["wrapper"]), + ( + { + "wrapper": [r'{"inner_field": 2}'], + "field": [1], + "wrapper_2": [[r'{"inner_field": 3}']], + "field_2": [4], + "wrapper_3": [r'{"inner_field": 5}'], + }, + ["wrapper", "wrapper_2", "wrapper_3"], + ), + ( + { + "valid_dict": [{"field_1": 1, "field_2": "two", "field_3": [3, 4]}], + "valid_dict_2": [[{"field_1": 1, "field_2": "two", "field_3": [3, 4]}]], + "invalid_dict": [{"time": datetime.datetime.now()}], + "valid_dict_3": [{}], + }, + ["valid_dict", "valid_dict_2", "valid_dict_3"], + ), + ], +) +def test_get_json_column_names(data, expected): + dataset = Dataset.from_dict(data) + iterable_dataset = IterableDataset.from_generator( + one_row_generator, gen_kwargs={"row": data} + ) + for ds in (dataset, iterable_dataset): + assert get_json_column_names(ds) == expected + + +@pytest.mark.smoke +@pytest.mark.parametrize( + ("locations", "expected"), + [ + ([], ((), ())), + ( + [ + VirtualColumnLocation("wrapper", "field_1"), + VirtualColumnLocation("wrapper_2", "field_1"), + VirtualColumnLocation("wrapper", "field_2"), + ], + (("wrapper", "wrapper_2", "wrapper"), ("field_1", "field_1", "field_2")), + ), + ([VirtualColumnLocation("wrapper", "field")], (("wrapper",), ("field",))), + ], +) +def test_unzip_virtual_column_locations(locations, expected): + actual = unzip_virtual_column_locations(locations) + assert isinstance(actual, tuple) + for idx in range(len(actual)): + assert isinstance(actual[idx], tuple) + assert actual[idx] == expected[idx]