Fix Big Query feature

Author: ycchenzheng

benchmarks/benchmark_db_utils.py

@@ -139,10 +139,10 @@ def write_run(
 
   # 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.run_summary_writer import run_summary_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
@@ -168,9 +168,9 @@ 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)
+      run_summary_writer.validate_model_id(options.model_id, options.is_test)
+      and run_summary_writer.validate_hardware_id(options.hardware_id, options.is_test)
+      and run_summary_writer.validate_software_id(options.software_id, options.is_test)
   ):
     summary = workload_benchmark_v2_schema.WorkloadBenchmarkV2Schema(
         run_id=f"run-{uuid.uuid4()}",
@@ -179,6 +179,7 @@ def get_db_client(
         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,

benchmarks/benchmark_db_writer/__init__.py

@@ -0,0 +1,3 @@
+import datetime
+
+__version__ = "1.0.0.dev" + datetime.datetime.now().strftime("%Y%m%d")

benchmarks/benchmark_db_writer/bigquery_types.py

@@ -0,0 +1,53 @@
+import datetime
+import decimal
+import enum
+from typing import Dict, NewType, Type
+
+
+class BigQueryFieldModes(str, enum.Enum):
+  NULLABLE = "NULLABLE"
+  REQUIRED = "REQUIRED"
+  REPEATED = "REPEATED"
+
+
+class BigQueryTypes(str, enum.Enum):
+  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,
+}

benchmarks/benchmark_db_writer/bq_info_writer/__init__.py

@@ -0,0 +1,111 @@
+"""TODO: Update hardware info in the main function & run the script."""
+
+import logging
+import os
+
+from benchmarks.benchmark_db_writer import bq_writer_utils
+from benchmarks.benchmark_db_writer.schema.workload_benchmark_v2 import (
+    hardware_info_schema,
+)
+
+logging.basicConfig(
+    format="%(asctime)s %(levelname)-8s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+def write_hardware_config(
+    project,
+    dataset,
+    table,
+    dataclass_type,
+    hardware_id,
+    gcp_accelerator_name,
+    chip_name,
+    bf_16_tflops,
+    memory,
+    hardware_type,
+    provider_name,
+    chips_per_node=None,
+    update_person_ldap=os.getenv("USER", "mrv2"),
+    description="",
+    other="",
+    host_memory=None,
+    host_vcpus=None,
+):
+
+  writer = bq_writer_utils.create_bq_writer_object(
+      project=project,
+      dataset=dataset,
+      table=table,
+      dataclass_type=dataclass_type,
+  )
+
+  hardware_info = writer.query(where={"hardware_id": hardware_id})
+  if hardware_info:
+    raise ValueError("Hardware id %s is already present in the %s table" % (hardware_id, table))
+
+  hardware_data = hardware_info_schema.HardwareInfo(
+      hardware_id=hardware_id,
+      gcp_accelerator_name=gcp_accelerator_name,
+      chip_name=chip_name,
+      bf_16_tflops=bf_16_tflops,
+      memory=memory,
+      chips_per_node=chips_per_node,
+      hardware_type=hardware_type,
+      provider_name=provider_name,
+      update_person_ldap=update_person_ldap,
+      description=description,
+      other=other,
+      host_memory=host_memory,
+      host_vcpus=host_vcpus,
+  )
+
+  logging.info("Writing Data %s to %s table.", hardware_data, table)
+  writer.write([hardware_data])
+
+
+if __name__ == "__main__":
+
+  table_configs = [
+      {
+          "project": "ml-workload-benchmarks",
+          "dataset": "benchmark_dataset_v2",
+          "table": "hardware_info",
+      },
+      {
+          "project": "supercomputer-testing",
+          "dataset": "mantaray_v2",
+          "table": "hardware_info",
+      },
+  ]
+
+  # Update it on every run
+  hardware_id = "a4"
+  gcp_accelerator_name = "A4"
+  chip_name = "B200"
+  bf_16_tflops = 2237
+  memory = 180
+  chips_per_node = 8
+  hardware_type = "GPU"
+  provider_name = "Nvidia"
+  description = ""
+
+  for table_config in table_configs:
+    write_hardware_config(
+        project=table_config["project"],
+        dataset=table_config["dataset"],
+        table=table_config["table"],
+        dataclass_type=hardware_info_schema.HardwareInfo,
+        hardware_id=hardware_id,
+        gcp_accelerator_name=gcp_accelerator_name,
+        chip_name=chip_name,
+        bf_16_tflops=bf_16_tflops,
+        memory=memory,
+        chips_per_node=chips_per_node,
+        description=description,
+        hardware_type=hardware_type,
+        provider_name=provider_name,
+    )

benchmarks/benchmark_db_writer/bq_info_writer/hardware_info_writer.py

@@ -0,0 +1,82 @@
+"""TODO: Update hardware info in the main function & run the script."""
+
+import logging
+import os
+
+from benchmarks.benchmark_db_writer import bq_writer_utils
+from benchmarks.benchmark_db_writer.schema.workload_benchmark_v2 import (
+    microbenchmark_workload_info_schema,
+)
+
+logging.basicConfig(
+    format="%(asctime)s %(levelname)-8s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+def write_microbenchmark_workload_config(
+    project,
+    dataset,
+    table,
+    dataclass_type,
+    workload_id,
+    update_person_ldap=os.getenv("USER", "imo-eng"),
+    description="",
+):
+
+  writer = bq_writer_utils.create_bq_writer_object(
+      project=project,
+      dataset=dataset,
+      table=table,
+      dataclass_type=dataclass_type,
+  )
+
+  microbenchmark_workload_info = writer.query(where={"workload_id": workload_id})
+  if microbenchmark_workload_info:
+    raise ValueError("Workload id %s is already present in the %s table" % (microbenchmark_workload_info, table))
+
+  workload_data = microbenchmark_workload_info_schema.MicrobenchmarkWorkloadInfo(
+      workload_id=workload_id,
+      update_person_ldap=update_person_ldap,
+      description=description,
+  )
+
+  logging.info("Writing Data %s to %s table.", workload_data, table)
+  writer.write([workload_data])
+
+
+def insert(workload_id, description=""):
+  table_configs = [
+      {
+          "project": "ml-workload-benchmarks",
+          "dataset": "benchmark_dataset_v2",
+          "table": "microbenchmark_workload_info",
+      },
+      {
+          "project": "supercomputer-testing",
+          "dataset": "mantaray_v2",
+          "table": "microbenchmark_workload_info",
+      },
+  ]
+
+  assert workload_id is not None
+
+  for table_config in table_configs:
+    write_microbenchmark_workload_config(
+        project=table_config["project"],
+        dataset=table_config["dataset"],
+        table=table_config["table"],
+        dataclass_type=microbenchmark_workload_info_schema.MicrobenchmarkWorkloadInfo,
+        workload_id=workload_id,
+        description=description,
+    )
+
+
+if __name__ == "__main__":
+
+  # workloads = ["all_gather", "ppermute", "psum", "psum_scatter"]
+  workloads = []
+  for workload in workloads:
+    insert(workload, "")

benchmarks/benchmark_db_writer/bq_info_writer/microbenchmark_workload_info_writer.py

@@ -0,0 +1,105 @@
+"""TODO: Update model info in the main function & run the script."""
+
+import logging
+import os
+
+from benchmarks.benchmark_db_writer import bq_writer_utils
+from benchmarks.benchmark_db_writer.schema.workload_benchmark_v2 import model_info_schema
+
+logging.basicConfig(
+    format="%(asctime)s %(levelname)-8s %(message)s",
+    level=logging.INFO,
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+
+def write_model_config(
+    project,
+    dataset,
+    table,
+    dataclass_type,
+    model_id,
+    name,
+    variant,
+    parameter_size_in_billions,
+    update_person_ldap=os.getenv("USER", "mrv2"),
+    description="",
+    details="",
+):
+
+  writer = bq_writer_utils.create_bq_writer_object(
+      project=project,
+      dataset=dataset,
+      table=table,
+      dataclass_type=dataclass_type,
+  )
+
+  model_info = writer.query(where={"model_id": model_id})
+  if model_info:
+    raise ValueError("Model id %s is already present in the %s table" % (model_id, table))
+
+  # Check if there is already a model info based on name,
+  # variant and parameter size
+  model_info = writer.query(
+      where={
+          "name": name,
+          "variant": variant,
+          "parameter_size_in_billions": parameter_size_in_billions,
+      }
+  )
+  if model_info:
+    raise ValueError(
+        "Model with name %s, variant %s and "
+        "parameter size %s is already present in the %s "
+        "table" % (name, variant, parameter_size_in_billions, table)
+    )
+
+  model_data = model_info_schema.ModelInfo(
+      model_id=model_id,
+      name=name,
+      variant=variant,
+      parameter_size_in_billions=parameter_size_in_billions,
+      update_person_ldap=update_person_ldap,
+      description=description,
+      details=details,
+  )
+
+  logging.info("Writing Data %s to %s table.", model_data, table)
+  writer.write([model_data])
+
+
+if __name__ == "__main__":
+
+  table_configs = [
+      {
+          "project": "ml-workload-benchmarks",
+          "dataset": "benchmark_dataset_v2",
+          "table": "model_info",
+      },
+      {
+          "project": "supercomputer-testing",
+          "dataset": "mantaray_v2",
+          "table": "model_info",
+      },
+  ]
+
+  # Update it on every run
+  model_id = "mistral-7b"
+  name = "Mistral"
+  variant = "7B"
+  parameter_size_in_billions = 7
+  description = "https://huggingface.co/mistralai/Mistral-7B-v0.3"
+
+  for table_config in table_configs:
+    write_model_config(
+        project=table_config["project"],
+        dataset=table_config["dataset"],
+        table=table_config["table"],
+        model_id=model_id,
+        dataclass_type=model_info_schema.ModelInfo,
+        name=name,
+        variant=variant,
+        parameter_size_in_billions=parameter_size_in_billions,
+        description=description,
+    )