add dataset and model size to run info

src/guidellm/presentation/data_models.py

@@ -1,13 +1,18 @@
 import random
 from collections import defaultdict
 from math import ceil
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
+import httpx
 from pydantic import BaseModel, computed_field
 
 if TYPE_CHECKING:
     from guidellm.benchmark.benchmark import GenerativeBenchmark
 
+from guidellm.dataset.file import FileDatasetCreator
+from guidellm.dataset.hf_datasets import HFDatasetsCreator
+from guidellm.dataset.in_memory import InMemoryDatasetCreator
+from guidellm.dataset.synthetic import SyntheticDatasetConfig, SyntheticDatasetCreator
 from guidellm.objects.statistics import DistributionSummary
 
 
@@ -58,6 +63,39 @@ class Model(BaseModel):
 class Dataset(BaseModel):
     name: str
 
+    @classmethod
+    def from_data(cls, request_loader: Any):
+        creators = [
+            InMemoryDatasetCreator,
+            SyntheticDatasetCreator,
+            FileDatasetCreator,
+            HFDatasetsCreator,
+        ]
+        dataset_name = None
+        data = request_loader.data
+        data_args = request_loader.data_args
+        processor = request_loader.processor
+        processor_args = request_loader.processor_args
+
+        for creator in creators:
+            if not creator.is_supported(data, None):
+                continue
+            random_seed = 42
+            dataset = creator.handle_create(
+                data, data_args, processor, processor_args, random_seed
+            )
+            dataset_name = creator.extract_dataset_name(dataset)
+            if dataset_name is None or dataset_name == "":
+                if creator == SyntheticDatasetCreator:
+                    data_dict = SyntheticDatasetConfig.parse_str(data)
+                    dataset_name = data_dict.source
+                if creator in (FileDatasetCreator, HFDatasetsCreator):
+                    dataset_name = data
+                if creator == InMemoryDatasetCreator:
+                    dataset_name = "In-memory"
+                break
+        return cls(name=dataset_name or "")
+
 
 class RunInfo(BaseModel):
     model: Model
@@ -71,11 +109,14 @@ def from_benchmarks(cls, benchmarks: list["GenerativeBenchmark"]):
         timestamp = max(
             bm.run_stats.start_time for bm in benchmarks if bm.start_time is not None
         )
+        response = httpx.get(f"https://huggingface.co/api/models/{model}")
+        model_json = response.json()
+
         return cls(
-            model=Model(name=model, size=0),
+            model=Model(name=model, size=model_json.get("usedStorage", 0)),
             task="N/A",
             timestamp=timestamp,
-            dataset=Dataset(name="N/A"),
+            dataset=Dataset.from_data(benchmarks[0].request_loader),
         )
 
 

src/ui/lib/components/PageHeader/PageHeader.component.tsx

@@ -2,12 +2,14 @@
 import { Box, Typography } from '@mui/material';
 
 import { useGetRunInfoQuery } from '../../store/slices/runInfo';
-import { formateDate } from '../../utils/helpers';
+import { formateDate, getFileSize } from '../../utils/helpers';
 import { SpecBadge } from '../SpecBadge';
 import { HeaderCell, HeaderWrapper } from './PageHeader.styles';
 
 export const Component = () => {
   const { data } = useGetRunInfoQuery();
+  const modelSize = getFileSize(data?.model?.size || 0);
+
   return (
     <Box py={2}>
       <Typography variant="subtitle2" color="surface.onSurfaceAccent">
@@ -24,11 +26,24 @@ export const Component = () => {
             variant="metric2"
             withTooltip
           />
+          <SpecBadge
+            label="Model size"
+            value={data?.model?.size ? `${modelSize?.size} ${modelSize?.units}` : '0B'}
+            variant="body1"
+          />
+        </HeaderCell>
+        <HeaderCell item xs={5} withDivider>
+          <SpecBadge
+            label="Dataset"
+            value={data?.dataset?.name || 'N/A'}
+            variant="caption"
+            withTooltip
+          />
         </HeaderCell>
         <HeaderCell item xs={2} sx={{ paddingRight: 0 }}>
           <SpecBadge
             label="Time Stamp"
-            value={data?.timestamp ? formateDate(data?.timestamp) : 'n/a'}
+            value={data?.timestamp ? formateDate(data?.timestamp) : 'N/A'}
             variant="caption"
           />
         </HeaderCell>

tests/unit/presentation/test_data_models.py

@@ -1,6 +1,13 @@
+from unittest.mock import MagicMock, patch
+
 import pytest
 
-from guidellm.presentation.data_models import Bucket
+from guidellm.dataset.file import FileDatasetCreator
+from guidellm.dataset.hf_datasets import HFDatasetsCreator
+from guidellm.dataset.in_memory import InMemoryDatasetCreator
+from guidellm.dataset.synthetic import SyntheticDatasetCreator
+from guidellm.presentation.data_models import Bucket, Dataset
+from tests.unit.mock_benchmark import mock_generative_benchmark
 
 
 @pytest.mark.smoke
@@ -18,3 +25,114 @@ def test_bucket_from_data():
     assert buckets[1].value == 8.0
     assert buckets[1].count == 5
     assert bucket_width == 1
+
+
+def mock_processor(cls):
+    return mock_generative_benchmark().request_loader.processor
+
+
+def new_handle_create(cls, *args, **kwargs):
+    return MagicMock()
+
+
+def new_extract_dataset_name(cls, *args, **kwargs):
+    return "data:prideandprejudice.txt.gz"
+
+
+@pytest.mark.smoke
+def test_dataset_from_data_uses_extracted_dataset_name():
+    mock_benchmark = mock_generative_benchmark()
+    with (
+        patch.object(SyntheticDatasetCreator, "handle_create", new=new_handle_create),
+        patch.object(
+            SyntheticDatasetCreator,
+            "extract_dataset_name",
+            new=new_extract_dataset_name,
+        ),
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "data:prideandprejudice.txt.gz"
+
+
+def new_is_supported(cls, *args, **kwargs):
+    return True
+
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_in_memory_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    with patch.object(InMemoryDatasetCreator, "is_supported", new=new_is_supported):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "In-memory"
+
+
+def hardcoded_isnt_supported(cls, *args, **kwargs):
+    return False
+
+
+def new_extract_dataset_name_none(cls, *args, **kwargs):
+    return None
+
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_synthetic_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    with (
+        patch.object(SyntheticDatasetCreator, "handle_create", new=new_handle_create),
+        patch.object(
+            InMemoryDatasetCreator, "is_supported", new=hardcoded_isnt_supported
+        ),
+        patch.object(SyntheticDatasetCreator, "is_supported", new=new_is_supported),
+        patch.object(
+            SyntheticDatasetCreator,
+            "extract_dataset_name",
+            new=new_extract_dataset_name_none,
+        ),
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "data:prideandprejudice.txt.gz"
+
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_file_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    mock_benchmark.request_loader.data = "dataset.yaml"
+    with (
+        patch.object(FileDatasetCreator, "handle_create", new=new_handle_create),
+        patch.object(
+            InMemoryDatasetCreator, "is_supported", new=hardcoded_isnt_supported
+        ),
+        patch.object(
+            SyntheticDatasetCreator, "is_supported", new=hardcoded_isnt_supported
+        ),
+        patch.object(FileDatasetCreator, "is_supported", new=new_is_supported),
+        patch.object(
+            FileDatasetCreator,
+            "extract_dataset_name",
+            new=new_extract_dataset_name_none,
+        ),
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "dataset.yaml"
+
+
+@pytest.mark.smoke
+def test_dataset_from_data_with_hf_dataset():
+    mock_benchmark = mock_generative_benchmark()
+    mock_benchmark.request_loader.data = "openai/gsm8k"
+    with (
+        patch.object(HFDatasetsCreator, "handle_create", new=new_handle_create),
+        patch.object(
+            InMemoryDatasetCreator, "is_supported", new=hardcoded_isnt_supported
+        ),
+        patch.object(
+            SyntheticDatasetCreator, "is_supported", new=hardcoded_isnt_supported
+        ),
+        patch.object(FileDatasetCreator, "is_supported", new=hardcoded_isnt_supported),
+        patch.object(HFDatasetsCreator, "is_supported", new=new_is_supported),
+        patch.object(
+            HFDatasetsCreator, "extract_dataset_name", new=new_extract_dataset_name_none
+        ),
+    ):
+        dataset = Dataset.from_data(mock_benchmark.request_loader)
+        assert dataset.name == "openai/gsm8k"