Add serialization and deserialization classes and methods with pydantic

Author: markurtz

src/guidellm/core/distribution.py

@@ -3,135 +3,99 @@
 import numpy as np
 from loguru import logger
 
+from guidellm.core.serializable import Serializable
+
 __all__ = ["Distribution"]
 
 
-class Distribution:
+class Distribution(Serializable):
     """
-    A class to represent a statistical distribution and perform various statistical
-    analyses.
-
-    :param data: List of numerical data points (int or float) to initialize the
-        distribution.
-    :type data: List[Union[int, float]], optional
+    A class to represent a statistical distribution and perform various
+    statistical analyses.
     """
 
-    def __init__(self, data: Optional[Union[List[int], List[float]]] = None):
-        """
-        Initialize the Distribution with optional data.
-
-        :param data: List of numerical data points to initialize the distribution,
-            defaults to None.
-        :type data: List[Union[int, float]], optional
-        """
-        self._data = list(data) if data else []
-        logger.debug(f"Initialized Distribution with data: {self._data}")
+    def __init__(self, **data):
+        super().__init__(**data)
+        logger.debug(f"Initialized Distribution with data: {self.data}")
 
     def __str__(self) -> str:
         """
         Return a string representation of the Distribution.
-
-        :return: String representation of the Distribution.
-        :rtype: str
         """
         return (
             f"Distribution(mean={self.mean:.2f}, median={self.median:.2f}, "
-            f"min={self.min}, max={self.max}, count={len(self._data)})"
+            f"min={self.min}, max={self.max}, count={len(self.data)})"
         )
 
     def __repr__(self) -> str:
         """
         Return an unambiguous string representation of the Distribution for debugging.
-
-        :return: Unambiguous string representation of the Distribution.
-        :rtype: str
-        """
-        return f"Distribution(data={self._data})"
-
-    @property
-    def data(self) -> List[Union[int, float]]:
-        """
-        Return the data points of the distribution.
-
-        :return: The data points of the distribution.
-        :rtype: List[Union[int, float]]
         """
-        return self._data
+        return f"Distribution(data={self.data})"
 
     @property
     def mean(self) -> float:
         """
         Calculate and return the mean of the distribution.
-
         :return: The mean of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate mean.")
             return 0.0
 
-        mean_value = np.mean(self._data).item()
+        mean_value = np.mean(self.data).item()
         logger.debug(f"Calculated mean: {mean_value}")
         return mean_value
 
     @property
     def median(self) -> float:
         """
         Calculate and return the median of the distribution.
-
         :return: The median of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate median.")
             return 0.0
 
-        median_value = np.median(self._data).item()
+        median_value = np.median(self.data).item()
         logger.debug(f"Calculated median: {median_value}")
         return median_value
 
     @property
     def variance(self) -> float:
         """
         Calculate and return the variance of the distribution.
-
         :return: The variance of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate variance.")
             return 0.0
 
-        variance_value = np.var(self._data).item()
+        variance_value = np.var(self.data).item()
         logger.debug(f"Calculated variance: {variance_value}")
         return variance_value
 
     @property
     def std_deviation(self) -> float:
         """
         Calculate and return the standard deviation of the distribution.
-
         :return: The standard deviation of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate standard deviation.")
             return 0.0
 
-        std_deviation_value = np.std(self._data).item()
+        std_deviation_value = np.std(self.data).item()
         logger.debug(f"Calculated standard deviation: {std_deviation_value}")
         return std_deviation_value
 
     def percentile(self, percentile: float) -> float:
         """
         Calculate and return the specified percentile of the distribution.
-
         :param percentile: The desired percentile to calculate (0-100).
-        :type percentile: float
         :return: The specified percentile of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate percentile.")
             return 0.0
 
@@ -142,61 +106,52 @@ def percentile(self, percentile: float) -> float:
     def percentiles(self, percentiles: List[float]) -> List[float]:
         """
         Calculate and return the specified percentiles of the distribution.
-
         :param percentiles: A list of desired percentiles to calculate (0-100).
-        :type percentiles: List[float]
         :return: A list of the specified percentiles of the distribution.
-        :rtype: List[float]
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate percentiles.")
             return [0.0] * len(percentiles)
 
-        percentiles_values = np.percentile(self._data, percentiles).tolist()
+        percentiles_values = np.percentile(self.data, percentiles).tolist()
         logger.debug(f"Calculated percentiles {percentiles}: {percentiles_values}")
         return percentiles_values
 
     @property
     def min(self) -> float:
         """
         Return the minimum value of the distribution.
-
         :return: The minimum value of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate minimum.")
             return 0.0
 
-        min_value = np.min(self._data)
+        min_value = np.min(self.data)
         logger.debug(f"Calculated min: {min_value}")
         return min_value
 
     @property
     def max(self) -> float:
         """
         Return the maximum value of the distribution.
-
         :return: The maximum value of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate maximum.")
             return 0.0
 
-        max_value = np.max(self._data)
+        max_value = np.max(self.data)
         logger.debug(f"Calculated max: {max_value}")
         return max_value
 
     @property
     def range(self) -> float:
         """
         Calculate and return the range of the distribution (max - min).
-
         :return: The range of the distribution.
-        :rtype: float
         """
-        if not self._data:
+        if not self.data:
             logger.warning("No data points available to calculate range.")
             return 0.0
 
@@ -207,9 +162,7 @@ def range(self) -> float:
     def describe(self) -> dict:
         """
         Return a dictionary describing various statistics of the distribution.
-
         :return: A dictionary with statistical summaries of the distribution.
-        :rtype: dict
         """
         description = {
             "mean": self.mean,
@@ -230,19 +183,15 @@ def describe(self) -> dict:
     def add_data(self, new_data: Union[List[int], List[float]]):
         """
         Add new data points to the distribution.
-
         :param new_data: A list of new numerical data points to add.
-        :type new_data: List[Union[int, float]]
         """
-        self._data.extend(new_data)
+        self.data.extend(new_data)
         logger.debug(f"Added new data: {new_data}")
 
     def remove_data(self, remove_data: Union[List[int], List[float]]):
         """
         Remove specified data points from the distribution.
-
         :param remove_data: A list of numerical data points to remove.
-        :type remove_data: List[Union[int, float]]
         """
-        self._data = [item for item in self._data if item not in remove_data]
+        self.data = [item for item in self.data if item not in remove_data]
         logger.debug(f"Removed data: {remove_data}")

src/guidellm/core/request.py

@@ -1,56 +1,32 @@
 import uuid
-from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Dict, Optional, Any
 
-from loguru import logger
+from guidellm.core.serializable import Serializable
 
 
-@dataclass(frozen=True)
-class TextGenerationRequest:
+class TextGenerationRequest(Serializable):
     """
     A class to represent a text generation request for generative AI workloads.
-
-    :param prompt: The input prompt for the text generation request.
-    :type prompt: str
-    :param prompt_token_count: The number of tokens in the prompt, defaults to None.
-    :type prompt_token_count: Optional[int]
-    :param generated_token_count: The number of tokens to generate, defaults to None.
-    :type generated_token_count: Optional[int]
-    :param params: Optional parameters for the text generation request,
-        defaults to None.
-    :type params: Optional[Dict[str, Any]]
     """
 
+    id: str
     prompt: str
-    id: uuid.UUID = field(default_factory=uuid.uuid4)
-    prompt_token_count: Optional[int] = None
-    generated_token_count: Optional[int] = None
-    params: Dict = field(default_factory=dict)
-
-    def __post_init__(self) -> None:
-        """
-        Log the initialization of the TextGenerationRequest instance.
-        """
-
-        logger.debug(
-            f"Initialized TextGenerationRequest with id={self.id}, "
-            f"prompt={self.prompt}, prompt_token_count={self.prompt_token_count}, "
-            f"generated_token_count={self.generated_token_count}, params={self.params}"
-        )
-
-    def __repr__(self) -> str:
-        """
-        Return a string representation of the TextGenerationRequest.
-
-        :return: String representation of the TextGenerationRequest.
-        :rtype: str
-        """
-
-        return (
-            f"TextGenerationRequest("
-            f"id={self.id}, "
-            f"prompt={self.prompt}, "
-            f"prompt_token_count={self.prompt_token_count}, "
-            f"generated_token_count={self.generated_token_count}, "
-            f"params={self.params})"
+    prompt_token_count: Optional[int]
+    generated_token_count: Optional[int]
+    params: Dict[str, Any]
+
+    def __init__(
+        self,
+        prompt: str,
+        prompt_token_count: Optional[int] = None,
+        generated_token_count: Optional[int] = None,
+        params: Optional[Dict[str, Any]] = None,
+        id: Optional[str] = None,
+    ):
+        super().__init__(
+            id=str(uuid.uuid4()) if id is None else id,
+            prompt=prompt,
+            prompt_token_count=prompt_token_count,
+            generated_token_count=generated_token_count,
+            params=params or {},
         )