Basic workload generation and saving to CSV

Author: lixueclaire

README.md

@@ -69,6 +69,7 @@ requests = generate_workload(pool, rate_fn, duration=3600)
 ```
 
 See `examples/` for more detailed examples:
+- `basic_usage.py`: Basic workload generation and saving to CSV
 - `generate_custom.py`: Custom workload patterns
 - `generate_realistic.py`: Realistic workload generation
 - `generate_advanced.py`: Multimodal and reasoning workloads

examples/basic_usage.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python3
+"""
+Example: Generate a workload and save it as a CSV file.
+
+This script demonstrates how to:
+1. Basic usage of load client data and generate a workload
+2. Save the generated workload to a CSV file
+"""
+
+from servegen import Category, ClientPool
+from servegen.construct import generate_workload
+from servegen.utils import save_requests_to_csv
+
+
+def main():
+    # Load client data
+    print("Loading client pool...")
+    pool = ClientPool(Category.LANGUAGE, "m-large")
+    print(f"Loaded {len(pool.clients)} clients.")
+
+    # Generate workload
+    print("Generating workload...")
+    rate_fn = {0: 100.0, 600: 150.0}  # requests per second
+    requests = generate_workload(pool, rate_fn, duration=1200)
+    print(f"Generated {len(requests)} requests.")
+
+    # Save to CSV
+    output_file = "workload.csv"
+    print(f"Saving to {output_file}...")
+    save_requests_to_csv(requests, output_file)
+    print(f"Saved to {output_file}.")
+
+    # Print first few rows
+    print("\nFirst 10 rows of the generated workload:")
+    with open(output_file, "r") as f:
+        lines = f.readlines()
+        for line in lines[:11]:  # header + 10 rows
+            print(line.strip())
+
+
+if __name__ == "__main__":
+    main()

servegen/__init__.py

@@ -2,5 +2,13 @@
 from servegen.workload_types import Category, ArrivalPat
 from servegen.clientpool import ClientPool, Client
 from servegen.construct import generate_workload
+from servegen.utils import save_requests_to_csv
 
-__all__ = ['Category', 'ArrivalPat', 'ClientPool', 'Client', 'generate_workload']
+__all__ = [
+    "Category",
+    "ArrivalPat",
+    "ClientPool",
+    "Client",
+    "generate_workload",
+    "save_requests_to_csv",
+]

servegen/utils.py

@@ -1,6 +1,8 @@
 from typing import Dict, List, Optional, Union, Tuple, Any
 import numpy as np
+import csv
 from servegen.clientpool import ClientPool, ClientPoolView, ClientWindow
+from servegen.construct import Request
 
 def get_constant_rate_fn(
     pool: Union[ClientPool, ClientPoolView],
@@ -132,3 +134,32 @@ def sample_from_cdf(
     indices = np.clip(indices, 0, len(values) - 1)
 
     return values[indices]
+
+def save_requests_to_csv(
+    requests: List[Request],
+    filename: str
+) -> None:
+    """Save list of Request objects to a CSV file.
+
+    Args:
+        requests: List of Request objects to save.
+        filename: Path to the output CSV file.
+
+    Raises:
+        ValueError: If no requests are provided.
+    """
+    if not requests:
+        raise ValueError("No requests to save.")
+
+    with open(filename, "w", newline="") as csvfile:
+        fieldnames = ["request_id", "timestamp"] + list(requests[0].data.keys())
+        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
+
+        writer.writeheader()
+        for req in requests:
+            row = {
+                "request_id": req.request_id,
+                "timestamp": req.timestamp,
+                **req.data  # Unpack the data dictionary
+            }
+            writer.writerow(row)