From b79c5b552aad8557e4aa580480290c70bfb563a8 Mon Sep 17 00:00:00 2001 From: Zheng Chen Date: Thu, 30 Oct 2025 18:24:53 +0000 Subject: [PATCH 1/2] Add Big Qurey necessary dependencies --- benchmarks/benchmark_db_utils.py | 109 +++--- .../benchmark_db_writer/bigquery_types.py | 86 +++++ .../benchmark_db_writer/bq_writer_utils.py | 57 ++++ .../dataclass_bigquery_writer.py | 323 ++++++++++++++++++ .../dataclass_converter_utils.py | 170 +++++++++ .../row_transformer_utils.py | 49 +++ .../workload_benchmark_v2_schema.py | 137 ++++++++ benchmarks/globals.py | 2 +- benchmarks/maxtext_xpk_runner.py | 13 +- benchmarks/recipes/runner_utils.py | 7 + benchmarks/recipes/user_configs.py | 5 + benchmarks/upload_metrics_to_bq.py | 8 +- .../tpu-requirements.txt | 1 + 13 files changed, 893 insertions(+), 74 deletions(-) create mode 100644 benchmarks/benchmark_db_writer/bigquery_types.py create mode 100644 benchmarks/benchmark_db_writer/bq_writer_utils.py create mode 100644 benchmarks/benchmark_db_writer/dataclass_bigquery_writer.py create mode 100644 benchmarks/benchmark_db_writer/dataclass_converter_utils.py create mode 100644 benchmarks/benchmark_db_writer/row_transformer_utils.py create mode 100644 benchmarks/benchmark_db_writer/schema/workload_benchmark_v2/workload_benchmark_v2_schema.py diff --git a/benchmarks/benchmark_db_utils.py b/benchmarks/benchmark_db_utils.py index edf54c18d6..605b8cc3ea 100644 --- a/benchmarks/benchmark_db_utils.py +++ b/benchmarks/benchmark_db_utils.py @@ -25,15 +25,12 @@ import dataclasses import getpass import os -import sys import uuid from argparse import Namespace -BQ_WRITER_PATH = "/benchmark-automation/benchmark_db_writer/src" temp_dir = gettempdir() DEFAULT_LOCAL_DIR = os.path.join(temp_dir, "") -# bq_writer_repo_root = get_bq_writer_path(DEFAULT_LOCAL_DIR) DEFAULT_TUNING_PARAMS_FILE = os.path.join(temp_dir, "tuning_params.json") @@ -114,7 +111,6 @@ def write_run( dataset: The dataset used in the run. num_of_superblock: The number of superblocks in the hardware. ( valid for GPUs) update_person_ldap: The LDAP ID of the person updating the record (default: current user). - is_test: Whether to use the testing project or the production project. metrics: Metrics object containing: median_step_time: The median step time of the run. e2e_step_time: The end-to-end time of the run. @@ -134,25 +130,20 @@ def write_run( Raises: ValueError: If any of the IDs are invalid. """ - bq_writer_repo_root = BQ_WRITER_PATH - sys.path.append(bq_writer_repo_root) - # pylint: disable=import-outside-toplevel - from benchmark_db_writer import bq_writer_utils - from benchmark_db_writer import dataclass_bigquery_writer - from benchmark_db_writer.run_summary_writer import sample_run_summary_writer - from benchmark_db_writer.schema.workload_benchmark_v2 import workload_benchmark_v2_schema + from benchmarks.benchmark_db_writer import bq_writer_utils + from benchmarks.benchmark_db_writer import dataclass_bigquery_writer + from benchmarks.benchmark_db_writer.schema.workload_benchmark_v2 import workload_benchmark_v2_schema def get_db_client( - project: str, dataset: str, table: str, dataclass_type: Type, is_test: bool = False + project: str, dataset: str, table: str, dataclass_type: Type ) -> dataclass_bigquery_writer.DataclassBigQueryWriter: """Creates a BigQuery client object. Args: table: The name of the BigQuery table. dataclass_type: The dataclass type corresponding to the table schema. - is_test: Whether to use the testing project or the production project. Returns: A BigQuery client object. @@ -167,53 +158,45 @@ def get_db_client( print(options.model_id) - if ( - sample_run_summary_writer.validate_model_id(options.model_id, options.is_test) - and sample_run_summary_writer.validate_hardware_id(options.hardware_id, options.is_test) - and sample_run_summary_writer.validate_software_id(options.software_id, options.is_test) - ): - summary = workload_benchmark_v2_schema.WorkloadBenchmarkV2Schema( - run_id=f"run-{uuid.uuid4()}", - model_id=options.model_id, - software_id=options.software_id, - hardware_id=options.hardware_id, - hardware_num_chips=number_of_chips, - hardware_num_nodes=number_of_nodes, - result_success=run_success, - configs_framework=framework_config_in_json, - configs_env=env_variables, - configs_container_version=options.container_image_name, - configs_xla_flags=options.xla_flags.replace(",", " "), - configs_dataset=options.dataset, - logs_artifact_directory="", - update_person_ldap=getpass.getuser(), - run_source="automation", - run_type=options.run_type, - run_release_status=run_release_status, - workload_precision=options.precision, - workload_gbs=int(options.global_batch_size), - workload_optimizer=options.optimizer, - workload_sequence_length=int(options.seq_length), - metrics_e2e_time=metrics.e2e_step_time, - metrics_mfu=mfu, - metrics_step_time=metrics.median_step_time, - metrics_tokens_per_second=metrics.avg_tokens_per_sec, - metrics_num_steps=number_of_steps, - metrics_other=other_metrics_in_json, - hardware_nccl_driver_nickname=nccl_driver_nickname, - hardware_topology=options.topology, - hardware_num_superblocks=0, - logs_comments=comment, - ) - - client = get_db_client( - options.db_project, - options.db_dataset, - "run_summary", - workload_benchmark_v2_schema.WorkloadBenchmarkV2Schema, - options.is_test, - ) - client.write([summary]) - - else: - raise ValueError("Could not upload data in run summary table") + summary = workload_benchmark_v2_schema.WorkloadBenchmarkV2Schema( + run_id=f"run-{uuid.uuid4()}", + model_id=options.model_id, + software_id=options.software_id, + hardware_id=options.hardware_id, + hardware_num_chips=number_of_chips, + hardware_num_nodes=number_of_nodes, + hardware_num_slices=options.hardware_num_slices, + result_success=run_success, + configs_framework=framework_config_in_json, + configs_env=env_variables, + configs_container_version=options.container_image_name, + configs_xla_flags=options.xla_flags.replace(",", " "), + configs_dataset=options.dataset, + logs_artifact_directory="", + update_person_ldap=getpass.getuser(), + run_source="automation", + run_type=options.run_type, + run_release_status=run_release_status, + workload_precision=options.precision, + workload_gbs=int(options.global_batch_size), + workload_optimizer=options.optimizer, + workload_sequence_length=int(options.seq_length), + metrics_e2e_time=metrics.e2e_step_time, + metrics_mfu=mfu, + metrics_step_time=metrics.median_step_time, + metrics_tokens_per_second=metrics.avg_tokens_per_sec, + metrics_num_steps=number_of_steps, + metrics_other=other_metrics_in_json, + hardware_nccl_driver_nickname=nccl_driver_nickname, + hardware_topology=options.topology, + hardware_num_superblocks=0, + logs_comments=comment, + ) + + client = get_db_client( + options.db_project, + options.db_dataset, + "run_summary", + workload_benchmark_v2_schema.WorkloadBenchmarkV2Schema, + ) + client.write([summary]) diff --git a/benchmarks/benchmark_db_writer/bigquery_types.py b/benchmarks/benchmark_db_writer/bigquery_types.py new file mode 100644 index 0000000000..bae72e3104 --- /dev/null +++ b/benchmarks/benchmark_db_writer/bigquery_types.py @@ -0,0 +1,86 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +This module defines enumerations for BigQuery data types (e.g., `STRING`, +`INT64`) and field modes (e.g., `NULLABLE`, `REQUIRED`). + +It also defines a primary mapping, `TypeMapping`, which translates these +BigQuery types into their corresponding standard Python types (like `str`, `int`, +`datetime.datetime`). Custom types (`TimeStamp`, `Geography`) are included +for specific BQ types not perfectly represented by Python built-ins. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/src/aotc/ +benchmark_db_writer/src/benchmark_db_writer/bigquery_types.py +""" +import datetime +import decimal +import enum +from typing import Dict, NewType, Type + + +class BigQueryFieldModes(str, enum.Enum): + """ + Enums for BigQueryFieldModes + """ + + NULLABLE = "NULLABLE" + REQUIRED = "REQUIRED" + REPEATED = "REPEATED" + + +class BigQueryTypes(str, enum.Enum): + """ + Enums for BigQueryTypes + """ + + STRING = "STRING" + BYTES = "BYTES" + INTEGER = "INT64" + INT64 = "INT64" + FLOAT64 = "FLOAT64" + FLOAT = "FLOAT64" + NUMERIC = "NUMERIC" + BOOL = "BOOL" + BOOLEAN = "BOOL" + STRUCT = "STRUCT" + RECORD = "STRUCT" + TIMESTAMP = "TIMESTAMP" + DATE = "DATE" + TIME = "TIME" + DATETIME = "DATETIME" + GEOGRAPHY = "GEOGRAPHY" + JSON = "JSON" + + +Geography = NewType("Geography", str) + + +class TimeStamp(datetime.datetime): + pass + + +TypeMapping: Dict[BigQueryTypes, Type] = { + BigQueryTypes.STRING: str, + BigQueryTypes.BYTES: bytes, + BigQueryTypes.INT64: int, + BigQueryTypes.FLOAT64: float, + BigQueryTypes.NUMERIC: decimal.Decimal, + BigQueryTypes.BOOL: bool, + BigQueryTypes.TIMESTAMP: TimeStamp, + BigQueryTypes.DATE: datetime.date, + BigQueryTypes.TIME: datetime.time, + BigQueryTypes.DATETIME: datetime.datetime, + BigQueryTypes.GEOGRAPHY: Geography, + BigQueryTypes.JSON: dict, +} diff --git a/benchmarks/benchmark_db_writer/bq_writer_utils.py b/benchmarks/benchmark_db_writer/bq_writer_utils.py new file mode 100644 index 0000000000..305e72e21d --- /dev/null +++ b/benchmarks/benchmark_db_writer/bq_writer_utils.py @@ -0,0 +1,57 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Utilities and factory functions for creating BigQuery writer clients. + +This module provides helper functions to simplify the instantiation of the +`DataclassBigQueryWriter`. It centralizes the configuration, such as +project and dataset IDs, making it easier to create database clients +for specific tables. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/ +src/aotc/benchmark_db_writer/src/benchmark_db_writer/bigquery_types.py +""" +from typing import Type +from benchmarks.benchmark_db_writer import dataclass_bigquery_writer + + +def create_bq_writer_object(project, dataset, table, dataclass_type): + """Creates a BQ writer config and uses it to create BQ writer object.""" + + config = dataclass_bigquery_writer.BigqueryWriterConfig(project, dataset, table) + + writer = dataclass_bigquery_writer.DataclassBigQueryWriter(dataclass_type, config) + + return writer + + +def get_db_client(table: str, dataclass_type: Type) -> create_bq_writer_object: + """Creates a BigQuery client object. + + Args: + table: The name of the BigQuery table. + dataclass_type: The dataclass type corresponding to the table schema. + + Returns: + A BigQuery client object. + """ + + project = "ml-workload-benchmarks" + dataset = "benchmark_dataset_v2" + return create_bq_writer_object( + project=project, + dataset=dataset, + table=table, + dataclass_type=dataclass_type, + ) diff --git a/benchmarks/benchmark_db_writer/dataclass_bigquery_writer.py b/benchmarks/benchmark_db_writer/dataclass_bigquery_writer.py new file mode 100644 index 0000000000..873bcab431 --- /dev/null +++ b/benchmarks/benchmark_db_writer/dataclass_bigquery_writer.py @@ -0,0 +1,323 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Core BigQuery writer class that maps Python dataclasses to BQ tables. + +This module provides the primary `DataclassBigQueryWriter`, a generic class +that uses a Python dataclass to define the schema of a BigQuery table. It +handles table creation, schema validation, schema updates, and provides +methods for writing dataclass instances to the table and reading BQ rows +back as dataclass instances. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/ +src/aotc/benchmark_db_writer/src/benchmark_db_writer/dataclass_bigquery_writer.py +""" +import copy +import dataclasses +import logging +import pprint +import time +from typing import Any, Generic, List, Optional, Sequence, Type, TypeVar + +from benchmarks.benchmark_db_writer import bigquery_types +from benchmarks.benchmark_db_writer import dataclass_converter_utils +from benchmarks.benchmark_db_writer import row_transformer_utils +import google.api_core.exceptions +from google.cloud import bigquery + +# The type of the generic dataclass +T = TypeVar("T") + +logger = logging.getLogger(__name__) + + +def _field_type_str(field_type: str): + """Normalizes the field type to a string. + + Args: + field_type: the field type to convert to a string. + + Returns: + The string representation of the field type. + """ + if isinstance(field_type, bigquery_types.BigQueryTypes): + return field_type.value + else: + return bigquery_types.BigQueryTypes[field_type].value + + +def _field_to_dict(field: bigquery.schema.SchemaField): + """A concise dict representation of a SchemaField. + + This is only to compare schemas to check if the schema fields have changed. + + Args: + field: the schema field to convert to a dict + + Returns: + A dict representation of the schema field. + """ + return { + "field_type": _field_type_str(field.field_type), + "mode": field.mode, + "fields": schema_to_dict(field.fields), + } + + +def schema_to_dict(schema: Sequence[bigquery.schema.SchemaField]): + """A concise dict representation of a bigquery schema. + + This is used to compare the current schema against the dataclass generated + schema. + + Args: + schema: the schema to convert to a dict. + + Returns: + A dict representation of the schema field. + """ + return {field.name: _field_to_dict(field) for field in schema} + + +def _recursive_struct_param( + name: str, schema: dict[str, Any], values: Optional[dict[str, Any]] = None +) -> bigquery.StructQueryParameter: + """Recursively builds a StructQueryParameter from schema and values. + + Args: + name: The name of the struct parameter. + schema: The concise schema dict for the struct (from `schema_to_dict`). + values: An optional dict of values for the struct. + + Returns: + A `bigquery.StructQueryParameter` object. + """ + params = [] + # match up schema to values + for field_name, field_schema in schema.items(): + value = values[field_name] if values else None + param = _query_param(field_name, field_schema, value) + assert param + params.append(param) + return bigquery.StructQueryParameter(name, *params) + + +def _query_param(name: str, schema_field: dict[str, Any], value: Any): # -> bigquery._AbstractQueryParameter: + """Builds a BigQuery query parameter from schema and a value. + + Handles nested STRUCTs by calling `_recursive_struct_param`. + + Args: + name: The name of the parameter (used as @name in the query). + schema_field: The concise schema dict for the field. + value: The value for the parameter. + + Returns: + A `bigquery.ScalarQueryParameter` or `bigquery.StructQueryParameter`. + """ + if schema_field["field_type"] == "STRUCT": + assert value is None or isinstance(value, dict) + # recurse the schema even for None/NULL values which we have to propagate + # all the way through the struct + return _recursive_struct_param(name, schema=schema_field["fields"], values=value) + else: + return bigquery.ScalarQueryParameter(name, schema_field["field_type"], value) + + +@dataclasses.dataclass +class BigqueryWriterConfig: + project: str + dataset: str + table: str + + +class DataclassBigQueryWriter(Generic[T]): + """Uses the `bq-schema` package to write a dataclass to a BigQuery table.""" + + def __init__(self, dataclass_type: Type[T], config: BigqueryWriterConfig): + """Initializes the writer. + + Args: + dataclass_type: the dataclass type to use as the schema + project: the GCP project to write to + dataset: the dataset to write to + table: the table to write to + """ + self.client = bigquery.Client(project=config.project) + self.row_transformer = None + self.table_id = f"{config.project}.{config.dataset}.{config.table}" + self.dataclass_type = dataclass_type + + self.input_data_schema = dataclass_converter_utils.dataclass_to_schema(self.dataclass_type) + # Get or create table + try: + self.table = self.client.get_table(self.table_id) + except google.api_core.exceptions.NotFound: + logger.warning("Table %s not found, creating it", self.table_id) + self.client.create_table(self.table_id) + self.table = self.client.get_table(self.table_id) + # When creating the table for the first time, always update schema. + self.update_schema() + + # Check schema of table and input dataclass + self.check_schema() + + def check_schema(self): + """Validates the dataclass schema against the live table schema. + + Raises: + ValueError: If the dataclass schema is incompatible with the + table schema. + - If a column exists in the dataclass but not the table. + - If a REQUIRED column exists in the table but not the dataclass. + """ + table_schema = schema_to_dict(self.table.schema) + data_schema = schema_to_dict(self.input_data_schema) + + # Check whether dataclass has any additional column + for dataclass_column in data_schema.keys(): + if dataclass_column not in table_schema: + raise ValueError( + f"Schema of table {self.table_id} is different than input data." + " Please check both schema and re-run.\n" + f"Column: {dataclass_column} is absent in table whereas it's " + "present in dataclass." + ) + + # Check whether big query table has any additional column which are not "nullable" + for table_column, column_attributes in table_schema.items(): + if table_column not in data_schema and column_attributes["mode"] != bigquery_types.BigQueryFieldModes.NULLABLE: + + raise ValueError( + f"Schema of table {self.table_id} is different than input data." + " Please check both schema and re-run.\n" + f"Column: {table_column} is absent in dataclass whereas it's " + "present in table & is of Required type." + ) + + def update_schema(self): + """When new table is created, this function gets called to update the schema.""" + logger.info( + "DataclassBigQueryWriter: updating schema to %s", + pprint.pformat(self.input_data_schema), + ) + old_schema = copy.deepcopy(self.table.schema) + try: + self.table.schema = self.input_data_schema + self.table = self.client.update_table(self.table, ["schema"]) + logger.info("BigQueryResultWriter: waiting for some time for the schema to" " propagate") + time.sleep(60) + except google.api_core.exceptions.GoogleAPICallError as e: + logger.exception("Failed to update bigquery schema with error %s", e) + self.table.schema = old_schema + + def transform(self, dataclass: T) -> dict: + return row_transformer_utils.RowTransformer.dataclass_instance_to_bq_row(dataclass) + + def read(self, where: Optional[str] = None) -> tuple[list[T], list[T]]: + """Reads the bigquery table using `where` as the WHERE clause. + + Args: + where: used as the `WHERE` expression when querying the database. + + Returns: + The list of bigquery entries as the dataclass T. + """ + row_transformer = row_transformer_utils.RowTransformer[T](self.dataclass_type) + query = "SELECT * FROM " + self.table_id + if where: + query += " WHERE " + where + raw_rows = [] + rows = [] + for bq_row in self.client.query(query=query): + raw_rows.append(bq_row) + dataclass = row_transformer.bq_row_to_dataclass_instance(bq_row) + assert isinstance(dataclass, self.dataclass_type) + rows.append(dataclass) + return rows, raw_rows + + def _get_field_schema_dict(self, field_name): + schema_dict = {"fields": schema_to_dict(self.input_data_schema)} + + field_dir = field_name.split(".") + for key in field_dir: + schema_dict = schema_dict["fields"][key] + return schema_dict + + def _get_query_for_value(self, field_name, value): # -> Tuple[str, bigquery._AbstractQueryParameter]: + if dataclasses.is_dataclass(value): + value = row_transformer_utils.RowTransformer.dataclass_instance_to_bq_row(value) + # # find schema for `field_name`: + field_schema = self._get_field_schema_dict(field_name) + at_name = "_".join(field_name.split(".")) + return f"{field_name} = @{at_name}", _query_param(at_name, field_schema, value) + + def query_column(self, column_name) -> List[Any]: + """Returns all values of the given column name.""" + + query_str = f"SELECT {column_name} FROM {self.table_id}" + query_result = self.client.query(query=query_str) + + return [row[0] for row in query_result] + + def query(self, where: Optional[dict[str, Any]] = None) -> list[T]: + """Reads the bigquery table using `where` dict as the WHERE clause. + + Args: + where: A dict with key value pair using which WHERE clause is constructed. + + Returns: + The list of bigquery entries as the dataclass T. + """ + if where is None: + where = {} + where_exprs = [] + params = [] + for field_name, value in where.items(): + where_expr, param = self._get_query_for_value(field_name, value) + params.append(param) + where_exprs.append(where_expr) + query_str = f"SELECT * FROM {self.table_id}" + if where_exprs: + where_stmt = " AND ".join(where_exprs) + query_str += f" WHERE {where_stmt}" + job_config = bigquery.QueryJobConfig(query_parameters=params) + + row_transformer = row_transformer_utils.RowTransformer[T](self.dataclass_type) + rows = [] + for bq_row in self.client.query(query=query_str, job_config=job_config): + dataclass = row_transformer.bq_row_to_dataclass_instance(bq_row) + assert isinstance(dataclass, self.dataclass_type) + rows.append(dataclass) + return rows + + def write(self, rows: List[T]): + """Bulk write to big query. + + Args: + rows: list of rows (dataclasses) to write to bigquery + """ + serialized_rows = [self.transform(row) for row in rows] + try: + logger.info("Writing to BigQuery: %d rows", len(serialized_rows)) + insert_errors = self.client.insert_rows(table=self.table, rows=serialized_rows) + if insert_errors: + logger.error( + "There were errors while writing to Bigquery:\n%s", + pprint.pformat(insert_errors), + ) + else: + logger.info("Successfully wrote to BigQuery") + except google.api_core.exceptions.GoogleAPICallError as e: + logger.exception("Failed to write to BigQuery with error %s", e) diff --git a/benchmarks/benchmark_db_writer/dataclass_converter_utils.py b/benchmarks/benchmark_db_writer/dataclass_converter_utils.py new file mode 100644 index 0000000000..708b6cd118 --- /dev/null +++ b/benchmarks/benchmark_db_writer/dataclass_converter_utils.py @@ -0,0 +1,170 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Convert a python dataclass into a BigQuery schema definition. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/ +src/aotc/benchmark_db_writer/src/benchmark_db_writer/dataclass_converter_utils.py +""" + +import dataclasses +import logging +from typing import Any, List, Optional, Type, Union, get_type_hints + +from benchmarks.benchmark_db_writer import bigquery_types +from google.cloud.bigquery import SchemaField +import typing_extensions + +logger = logging.getLogger(__name__) + +_BASIC_TYPES_TO_NAME = {primitive_type: bq_type for bq_type, primitive_type in bigquery_types.TypeMapping.items()} +_NoneType = type(None) + + +def parse_inner_type_of_list(list_type: Any) -> Type: + return typing_extensions.get_args(list_type)[0] + + +def parse_inner_type_of_optional(optional_type: Any) -> Type: + args = typing_extensions.get_args(optional_type) + if not (len(args) == 2 and any(arg is _NoneType for arg in args)): + raise TypeError(f"Unsupported type: {optional_type}.") + + return next(arg for arg in args if arg is not _NoneType) + + +def _parse_field_description(field: dataclasses.Field) -> Optional[str]: + if "description" in field.metadata: + return field.metadata["description"] + return None + + +def _parse_fields(field_type: Type) -> List[SchemaField]: + """Recursive call for nested dataclasses.""" + + if dataclasses.is_dataclass(field_type): + return dataclass_to_schema(field_type) + return [] + + +def _parse_list(field: dataclasses.Field) -> SchemaField: + field_type = parse_inner_type_of_list(field.type) + return SchemaField( + name=field.name, + field_type=_python_type_to_big_query_type(field_type), + mode=bigquery_types.BigQueryFieldModes.REPEATED, + description=_parse_field_description(field), + fields=_parse_fields(field_type), + ) + + +def _python_type_to_big_query_type( + field_type: Any, +) -> bigquery_types.BigQueryTypes: + """ + Args: + field_type: The Python type (e.g., `str`, `int`, a dataclass). + + Returns: + The corresponding `bigquery_types.BigQueryTypes` enum value. + + Raises: + TypeError: If the Python type is not supported or mapped. + """ + if dataclasses.is_dataclass(field_type): + return bigquery_types.BigQueryTypes.STRUCT + + bq_type = _BASIC_TYPES_TO_NAME.get(field_type) + if bq_type: + return bq_type + + raise TypeError(f"Unsupported type: {field_type}") + + +def _parse_optional(field: dataclasses.Field) -> SchemaField: + field_type = parse_inner_type_of_optional(field.type) + return SchemaField( + name=field.name, + field_type=_python_type_to_big_query_type(field_type), + mode=bigquery_types.BigQueryFieldModes.NULLABLE, + description=_parse_field_description(field), + fields=_parse_fields(field_type), + ) + + +def _field_to_schema(field: dataclasses.Field) -> SchemaField: + """ + Args: + field: The `dataclasses.Field` to convert. + + Returns: + A corresponding `SchemaField` object. + + Raises: + TypeError: If the field's type is complex and unsupported. + """ + field_type = _BASIC_TYPES_TO_NAME.get(field.type) + if field_type: + return SchemaField( + name=field.name, + field_type=field_type, + description=_parse_field_description(field), + mode=bigquery_types.BigQueryFieldModes.REQUIRED, + ) + + if dataclasses.is_dataclass(field.type): + return SchemaField( + name=field.name, + field_type=bigquery_types.BigQueryTypes.STRUCT, + mode=bigquery_types.BigQueryFieldModes.REQUIRED, + description=_parse_field_description(field), + fields=_parse_fields(field.type), + ) + + # typing.Optional is the same as typing.Union[SomeType, NoneType] + if typing_extensions.get_origin(field.type) is Union: + return _parse_optional(field) + + if typing_extensions.get_origin(field.type) is list: + return _parse_list(field) + + raise TypeError(f"Unsupported type: {field.type}.") + + +def dataclass_to_schema(dataclass: Type, localns: Optional[dict] = None) -> List[SchemaField]: + """Transfrom a dataclass into a list of SchemaField. + + If you want to transform a dataclass that is not defined in the + global scope you need to pass your locals. + + def my_func(): + @dataclass + class Example1: + a: int + + @dataclass + class Example2: + b: Example1 + + dataclass_to_schema(Example2, localns=locals()) + """ + if not dataclasses.is_dataclass(dataclass): + raise TypeError("Not a dataclass.") + + type_hints = get_type_hints(dataclass, localns=localns) + dataclass_fields = dataclasses.fields(dataclass) + + for field in dataclass_fields: + field.type = type_hints[field.name] + return [_field_to_schema(field) for field in dataclass_fields] diff --git a/benchmarks/benchmark_db_writer/row_transformer_utils.py b/benchmarks/benchmark_db_writer/row_transformer_utils.py new file mode 100644 index 0000000000..f8d6e9250b --- /dev/null +++ b/benchmarks/benchmark_db_writer/row_transformer_utils.py @@ -0,0 +1,49 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +""" +Provides a utility class for transforming data. + +This module contains the `RowTransformer`, a generic class that uses `dacite` +to convert `google.cloud.bigquery.table.Row` objects into specific +Python dataclass instances and vice-versa. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/src/ +aotc/benchmark_db_writer/src/benchmark_db_writer/row_transformer_utils.py +""" +import dataclasses +from typing import Generic, Type, TypeVar + +import dacite +from google.cloud.bigquery.table import Row + +T = TypeVar("T") # pylint: disable=invalid-name + + +class RowTransformer(Generic[T]): + """Serialized / deserialize rows.""" + + def __init__(self, schema: Type[T]): + self._schema: Type[T] = schema + + def bq_row_to_dataclass_instance(self, bq_row: Row) -> T: + """Create a dataclass instance from a row returned by the bq library.""" + + row_dict = dict(bq_row.items()) + + return dacite.from_dict(self._schema, row_dict, config=dacite.Config(check_types=False)) + + @staticmethod + def dataclass_instance_to_bq_row(instance: T) -> dict: + """Convert a dataclass instance into a dictionary, which can be inserted into bq.""" + return dataclasses.asdict(instance) diff --git a/benchmarks/benchmark_db_writer/schema/workload_benchmark_v2/workload_benchmark_v2_schema.py b/benchmarks/benchmark_db_writer/schema/workload_benchmark_v2/workload_benchmark_v2_schema.py new file mode 100644 index 0000000000..e582519822 --- /dev/null +++ b/benchmarks/benchmark_db_writer/schema/workload_benchmark_v2/workload_benchmark_v2_schema.py @@ -0,0 +1,137 @@ +# Copyright 2023–2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Defines the BigQuery schema for a V2 Workload Benchmark run. + +This module contains the `WorkloadBenchmarkV2Schema` dataclass, which defines +the structure of the "run_summary" table in BigQuery. Each instance of this +dataclass represents a single benchmark run. +Copied & Modified from https://github.com/AI-Hypercomputer/aotc/blob/main/src/ +aotc/benchmark_db_writer/src/benchmark_db_writer/schema/workload_benchmark_v2/ +workload_benchmark_v2_schema.py +""" +import dataclasses +import datetime +from typing import Optional +from benchmarks.benchmark_db_writer import bigquery_types + + +@dataclasses.dataclass +class WorkloadBenchmarkV2Schema: + """Dataclass representing the schema for the 'run_summary' BQ table. + + Attributes: + run_id: Primary key for the run. + model_id: Foreign key to the model info table. + software_id: Foreign key to the software info table. + hardware_id: Foreign key to the hardware info table. + hardware_num_chips: Number of chips used for this run. + result_success: Boolean indicating if the run was successful. + update_person_ldap: LDAP of the person who last updated this entry. + ... and other fields defining benchmark configuration, metrics, and logs. + """ + + run_id: str + + # Unique model id to map model info table + model_id: str + + # Foreign key to join with software info + software_id: str + # Foreign key to join with hardware info + hardware_id: str + hardware_num_chips: int + + result_success: bool + + update_person_ldap: str + configs_framework: Optional[str] = None + configs_container_version: Optional[str] = None + logs_artifact_directory: Optional[str] = None + configs_env: Optional[str] = None + hardware_num_nodes: Optional[int] = None + + # Foreign key to join with storage info + storage_id: Optional[str] = None + + run_source: str = "manual" + is_run_externally_visible: bool = False + run_type: str = "perf_optimization" + + run_release_status: Optional[str] = "local" + k8_jobset_yaml_file_path: Optional[str] = None + + benchmark_type: Optional[str] = None + + experiment_id: Optional[str] = None + + workload_gbs: Optional[int] = None + workload_mbs: Optional[int] = None + workload_precision: Optional[str] = None + workload_optimizer: Optional[str] = None + workload_others: Optional[str] = None + workload_manager: Optional[str] = None + workload_type: str = "training" + workload_sequence_length: Optional[int] = None + + metrics_step_time: Optional[float] = None + metrics_mfu: Optional[float] = None + metrics_tokens_per_second: Optional[float] = None + metrics_e2e_time: Optional[float] = None + metrics_num_steps: Optional[int] = None + metrics_other: Optional[str] = None + metrics_tflops_per_second: Optional[float] = None + + hardware_num_superblocks: Optional[str] = None + hardware_num_slices: Optional[int] = None + hardware_topology: Optional[str] = None + hardware_num_cores: Optional[int] = None + result_error: Optional[str] = None + hardware_nccl_driver_nickname: Optional[str] = None + + configs_xla_flags: Optional[str] = None + configs_dataset: Optional[str] = None + configs_reviewer: Optional[str] = None + configs_other: Optional[str] = None + + logs_profile: Optional[str] = None + logs_cloud_logs: Optional[str] = None + logs_comments: Optional[str] = None + logs_other: Optional[str] = None + + checkpointing_async: Optional[bool] = None + checkpointing_interval_every_n_steps: Optional[int] = None + checkpointing_size_in_gibs: Optional[float] = None + checkpointing_individual_file_size: Optional[int] = None + checkpointing_file_format: Optional[str] = None + + max_epochs: Optional[int] = None + max_steps: Optional[int] = None + training_dataset_samples: Optional[int] = None + data_loader_num_workers: Optional[int] = None + data_loader_prefetch_factor: Optional[int] = None + training_dataset_file_format: Optional[str] = None + + start_time: Optional[bigquery_types.TimeStamp] = None + end_time: Optional[bigquery_types.TimeStamp] = None + + gcs_metrics_bucket: Optional[str] = None + gcsfuse_csi_driver: Optional[str] = None + cloud_region: Optional[str] = None + source_bucket: Optional[str] = None + + cluster_name: Optional[str] = None + + reviewer_ldap: str = "" + update_timestamp: Optional[bigquery_types.TimeStamp] = datetime.datetime.now() diff --git a/benchmarks/globals.py b/benchmarks/globals.py index ba3a625b72..90528eadf4 100644 --- a/benchmarks/globals.py +++ b/benchmarks/globals.py @@ -17,7 +17,7 @@ import os.path # This is the MaxText root: with "max_utils.py"; &etc. TODO: Replace `os.path.basename` with `os.path.abspath` -MAXTEXT_PKG_DIR = os.environ.get("MAXTEXT_PKG_DIR", "MaxText") +MAXTEXT_PKG_DIR = os.environ.get("MAXTEXT_PKG_DIR", "src/MaxText") # This is the maxtext repo root: with ".git" folder; "README.md"; "pyproject.toml"; &etc. MAXTEXT_REPO_ROOT = os.environ.get( diff --git a/benchmarks/maxtext_xpk_runner.py b/benchmarks/maxtext_xpk_runner.py index d968d04908..a01894d1e8 100644 --- a/benchmarks/maxtext_xpk_runner.py +++ b/benchmarks/maxtext_xpk_runner.py @@ -112,7 +112,6 @@ class WorkloadConfig: num_devices_per_slice: int = dataclasses.field(init=False) db_project: str = "" db_dataset: str = "" - db_is_test: bool = True disruption_configs: DisruptionConfig = None xpk_storage: None | list[str] = None hlo_dump: None | bool = None @@ -126,12 +125,12 @@ def __post_init__(self): "device_type is None and generate_metrics_and_upload_to_big_query is enabled. " "Device_type is required for uploading run results to BigQuery" ) + size = int(self.device_type.split("-")[-1]) if ( self.device_type.startswith("v6e") or self.device_type.startswith("v5e") or self.device_type.startswith("v5litepod") ): - size = int(self.device_type.split("-")[-1]) if size == 256: self.num_devices_per_slice = 256 self.topology = "16x16" @@ -156,8 +155,11 @@ def __post_init__(self): else: raise ValueError(f"Unsupported v5e or v6e size: {size}") else: - self.num_devices_per_slice = int(self.device_type.split("-")[1]) / 2 + self.num_devices_per_slice = size / 2 self.topology = "" + self.hardware_id = self.device_type.split("-")[0] + if self.hardware_id == "v5litepod": + self.hardware_id = "v5e" def wait_for_xpk_workload_completion(cluster_config: XpkClusterConfig, workload_name, xpk_path) -> int: @@ -341,6 +343,7 @@ def _build_args_from_config(wl_config: WorkloadConfig) -> dict: "model_id": wl_config.model.model_type, "hardware_id": wl_config.hardware_id, "software_id": "jax_maxtext", + "hardware_num_slices": wl_config.num_slices, "number_of_chips": wl_config.num_devices_per_slice * wl_config.num_slices, "container_image_name": wl_config.base_docker_image, "global_batch_size": per_device_batch_size * wl_config.num_devices_per_slice * wl_config.num_slices, @@ -356,7 +359,6 @@ def _build_args_from_config(wl_config: WorkloadConfig) -> dict: "tuning_params": f"'{tuning_params_str}'", "db_project": wl_config.db_project, "db_dataset": wl_config.db_dataset, - "is_test": wl_config.db_is_test, } @@ -445,7 +447,8 @@ def build_user_command( f"base_output_directory={wl_config.base_output_directory}", f"{vertex_tensorboard}", f"{run_name_command}", - f"{enable_metrics_cmd}" f"{upload_hlo_dump}", + f"{enable_metrics_cmd}", + f"{upload_hlo_dump}", ] ) return command diff --git a/benchmarks/recipes/runner_utils.py b/benchmarks/recipes/runner_utils.py index 15102b3c36..4d65eb792a 100644 --- a/benchmarks/recipes/runner_utils.py +++ b/benchmarks/recipes/runner_utils.py @@ -44,6 +44,9 @@ def _create_workload_config( "xpk_path": user_config.xpk_path, "num_steps": num_steps, "priority": priority, + "generate_metrics_and_upload_to_big_query": user_config.bq_enable, + "db_project": user_config.bq_db_project, + "db_dataset": user_config.bq_db_dataset, } # Add any extra arguments, like disruption_configs, if they exist config_args.update(kwargs) @@ -81,6 +84,10 @@ def generate_and_run_workloads( """ Generates and executes XPK workloads, with or without disruptions. """ + if user_config.bq_enable and (not user_config.bq_db_project or not user_config.bq_db_dataset): + logging.error("Validation FAILED: BigQuery is enabled, but 'bq_db_project' or 'bq_db_dataset' is missing.") + return 1 + workload_configs = list( _generate_workloads( user_config, diff --git a/benchmarks/recipes/user_configs.py b/benchmarks/recipes/user_configs.py index 6c01c0dfc7..f50e912573 100644 --- a/benchmarks/recipes/user_configs.py +++ b/benchmarks/recipes/user_configs.py @@ -70,6 +70,11 @@ class UserConfig: selected_model_names: list[str] = dataclasses.field(default_factory=lambda: ["llama3_1_8b_8192"]) num_slices_list: list[int] = dataclasses.field(default_factory=lambda: [2]) + # BigQuery configuration + bq_enable: bool = False + bq_db_project: str = "" + bq_db_dataset: str = "" + # other configuration xpk_path: str = "~/xpk" max_restarts: int = 0 diff --git a/benchmarks/upload_metrics_to_bq.py b/benchmarks/upload_metrics_to_bq.py index 8576f3cc0e..9ffffe41df 100644 --- a/benchmarks/upload_metrics_to_bq.py +++ b/benchmarks/upload_metrics_to_bq.py @@ -43,7 +43,6 @@ from benchmarks.benchmark_db_utils import Metrics from benchmarks.benchmark_db_utils import recover_tuning_params from benchmarks.benchmark_db_utils import write_run -from benchmarks.benchmark_utils import str2bool from benchmarks.command_utils import run_command_with_updates @@ -180,11 +179,10 @@ def add_parser_arguments(parser: argparse.ArgumentParser): help="Dataset of the database", ) parser.add_argument( - "--is_test", - type=str2bool, + "--hardware_num_slices", + type=int, required=False, - default=True, - help="Whether to use the testing project or production project", + help="hardware slice number", ) diff --git a/dependencies/requirements/generated_requirements/tpu-requirements.txt b/dependencies/requirements/generated_requirements/tpu-requirements.txt index b9fa08ddf9..2fa1020bc0 100644 --- a/dependencies/requirements/generated_requirements/tpu-requirements.txt +++ b/dependencies/requirements/generated_requirements/tpu-requirements.txt @@ -80,6 +80,7 @@ grain>=0.2.14 grpc-google-iam-v1>=0.14.3 grpcio-status>=1.71.2 grpcio>=1.75.1 +gspread>=6.2.1 gviz-api>=1.10.0 h11>=0.16.0 h5py>=3.15.1 From 2c4ff3e1701a7d5b971962bc3af62b6e71768af2 Mon Sep 17 00:00:00 2001 From: Alfred Yu Date: Thu, 20 Nov 2025 17:04:48 +0800 Subject: [PATCH 2/2] test --- .pre-commit-config.yaml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 854fdb3f1b..aa7b1932e0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,6 +10,8 @@ repos: - '-w' - '--skip="*.txt,pylintrc,.*,src/MaxText/assets/*"' - '-L ND,nd,sems,TE,ROUGE,rouge,astroid,ags,dout' + - '-s' + - '--count' - '.' additional_dependencies: - tomli