codekingpro/portable-devtools

1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation.  All rights reserved.
3# Licensed under the MIT License.  See License.txt in the project root for
4# license information.
5# --------------------------------------------------------------------------
6
7import datetime
8import json
9
10import pandas as pd
11
12
13class BaseObject:
14    def __init__(self):
15        self.customized = {}
16
17    def to_dict(self):
18        default_values = self.__dict__.copy()
19        default_values.pop("customized", None)
20        default_values.update(self.customized)
21
22        for k, v in default_values.items():
23            if isinstance(v, BaseObject):
24                default_values[k] = v.to_dict()
25
26        return {k: v for k, v in default_values.items() if v}
27
28
29class ModelInfo(BaseObject):
30    def __init__(
31        self,
32        full_name: str | None = None,
33        is_huggingface: bool | None = False,
34        is_text_generation: bool | None = False,
35        short_name: str | None = None,
36    ):
37        super().__init__()
38        self.full_name = full_name
39        self.is_huggingface = is_huggingface
40        self.is_text_generation = is_text_generation
41        self.short_name = short_name
42        self.input_shape = []
43
44
45class BackendOptions(BaseObject):
46    def __init__(
47        self,
48        enable_profiling: bool | None = False,
49        execution_provider: str | None = None,
50        use_io_binding: bool | None = False,
51    ):
52        super().__init__()
53        self.enable_profiling = enable_profiling
54        self.execution_provider = execution_provider
55        self.use_io_binding = use_io_binding
56
57
58class Config(BaseObject):
59    def __init__(
60        self,
61        backend: str | None = "onnxruntime",
62        batch_size: int | None = 1,
63        seq_length: int | None = 0,
64        precision: str | None = "fp32",
65        warmup_runs: int | None = 1,
66        measured_runs: int | None = 10,
67    ):
68        super().__init__()
69        self.backend = backend
70        self.batch_size = batch_size
71        self.seq_length = seq_length
72        self.precision = precision
73        self.warmup_runs = warmup_runs
74        self.measured_runs = measured_runs
75        self.model_info = ModelInfo()
76        self.backend_options = BackendOptions()
77
78
79class Metadata(BaseObject):
80    def __init__(
81        self,
82        device: str | None = None,
83        package_name: str | None = None,
84        package_version: str | None = None,
85        platform: str | None = None,
86        python_version: str | None = None,
87    ):
88        super().__init__()
89        self.device = device
90        self.package_name = package_name
91        self.package_version = package_version
92        self.platform = platform
93        self.python_version = python_version
94
95
96class Metrics(BaseObject):
97    def __init__(
98        self,
99        latency_ms_mean: float | None = 0.0,
100        throughput_qps: float | None = 0.0,
101        max_memory_usage_GB: float | None = 0.0,
102    ):
103        super().__init__()
104        self.latency_ms_mean = latency_ms_mean
105        self.throughput_qps = throughput_qps
106        self.max_memory_usage_GB = max_memory_usage_GB
107
108
109class BenchmarkRecord:
110    def __init__(
111        self,
112        model_name: str,
113        precision: str,
114        backend: str,
115        device: str,
116        package_name: str,
117        package_version: str,
118        batch_size: int | None = 1,
119        warmup_runs: int | None = 1,
120        measured_runs: int | None = 10,
121        trigger_date: str | None = None,
122    ):
123        self.config = Config()
124        self.metrics = Metrics()
125        self.metadata = Metadata()
126        self.trigger_date = trigger_date or datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
127
128        self.config.model_info.full_name = model_name
129        self.config.precision = precision
130        self.config.backend = backend
131        self.config.batch_size = batch_size
132        self.config.warmup_runs = warmup_runs
133        self.config.measured_runs = measured_runs
134        self.metadata.device = device
135        self.metadata.package_name = package_name
136        self.metadata.package_version = package_version
137
138    def to_dict(self) -> dict:
139        return {
140            "config": self.config.to_dict(),
141            "metadata": self.metadata.to_dict(),
142            "metrics": self.metrics.to_dict(),
143            "trigger_date": self.trigger_date,
144        }
145
146    def to_json(self) -> str:
147        return json.dumps(self.to_dict(), default=str)
148
149    @classmethod
150    def save_as_csv(cls, file_name: str, records: list) -> None:
151        if records is None or len(records) == 0:
152            return
153        rds = [record.to_dict() for record in records]
154        df = pd.json_normalize(rds)
155        df.to_csv(file_name, index=False)
156
157    @classmethod
158    def save_as_json(cls, file_name: str, records: list) -> None:
159        if records is None or len(records) == 0:
160            return
161        rds = [record.to_dict() for record in records]
162        with open(file_name, "w") as f:
163            json.dump(rds, f, indent=4, default=str)
164