From d053ff5c0133b8d0d2b9dd7bba8fb98eecc70e2d Mon Sep 17 00:00:00 2001 From: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com> Date: Tue, 9 Jun 2026 08:48:51 -0400 Subject: [PATCH] add modelcar image size & byte size Signed-off-by: Adysen Rothman <85646824+adysenrothman@users.noreply.github.com> --- .../modelcatalog/performance_metrics.go | 33 ++++++++-- .../modelcatalog/performance_metrics_test.go | 64 ++++++++++++++----- 2 files changed, 74 insertions(+), 23 deletions(-) diff --git a/catalog/internal/catalog/modelcatalog/performance_metrics.go b/catalog/internal/catalog/modelcatalog/performance_metrics.go index 05e7328b7c..451c7c3eef 100644 --- a/catalog/internal/catalog/modelcatalog/performance_metrics.go +++ b/catalog/internal/catalog/modelcatalog/performance_metrics.go @@ -22,13 +22,15 @@ import ( // metadataJSON represents the minimal structure needed from metadata.json files // Only the ID field is needed to look up existing models type metadataJSON struct { - ID string `json:"id"` // Maps to model name for lookup - OverallAccuracy *float64 `json:"overall_accuracy"` // Overall accuracy score for the model - Size *string `json:"size"` // Model parameter count (e.g., "8B params") - TensorType *string `json:"tensor_type"` // Data precision (e.g., "FP16", "INT4") - VariantGroupID *string `json:"variant_group_id"` // UUID linking model variants together - MinVRAMGB *string `json:"min_vram_gb"` // Minimum VRAM required (e.g., "265 GB") - ColdStartMatrix []coldStartEntry `json:"cold_start_matrix"` // Cold start times per GPU configuration + ID string `json:"id"` // Maps to model name for lookup + OverallAccuracy *float64 `json:"overall_accuracy"` // Overall accuracy score for the model + Size *string `json:"size"` // Model parameter count (e.g., "8B params") + TensorType *string `json:"tensor_type"` // Data precision (e.g., "FP16", "INT4") + VariantGroupID *string `json:"variant_group_id"` // UUID linking model variants together + MinVRAMGB *string `json:"min_vram_gb"` // Minimum VRAM required (e.g., "265 GB") + ModelcarImageSize *string `json:"modelcar_image_size"` // Human-readable modelcar image size (e.g., "230.17 GB") + ModelcarImageSizeBytes *int64 `json:"modelcar_image_size_bytes"` // Modelcar image size in bytes + ColdStartMatrix []coldStartEntry `json:"cold_start_matrix"` // Cold start times per GPU configuration } type coldStartEntry struct { @@ -765,6 +767,23 @@ func enrichCatalogModelFromMetadata(existingModel dbmodels.CatalogModel, metadat }) } + if metadata.ModelcarImageSize != nil && *metadata.ModelcarImageSize != "" { + customProperties = append(customProperties, models.Properties{ + Name: "modelcar_image_size", + StringValue: metadata.ModelcarImageSize, + IsCustomProperty: true, + }) + } + + if metadata.ModelcarImageSizeBytes != nil { + bytesStr := strconv.FormatInt(*metadata.ModelcarImageSizeBytes, 10) + customProperties = append(customProperties, models.Properties{ + Name: "modelcar_image_size_bytes", + StringValue: &bytesStr, + IsCustomProperty: true, + }) + } + if len(customProperties) == 0 { return nil // Nothing to update } diff --git a/catalog/internal/catalog/modelcatalog/performance_metrics_test.go b/catalog/internal/catalog/modelcatalog/performance_metrics_test.go index 07580d38f8..19307efbf7 100644 --- a/catalog/internal/catalog/modelcatalog/performance_metrics_test.go +++ b/catalog/internal/catalog/modelcatalog/performance_metrics_test.go @@ -956,15 +956,17 @@ func TestUnmarshalJSON_EdgeCases(t *testing.T) { func TestParseMetadataJSON_NewFields(t *testing.T) { tests := []struct { - name string - jsonData string - wantID string - wantSize *string - wantTensorType *string - wantVariantID *string - wantMinVRAMGB *string - wantColdStartMatrix []coldStartEntry - wantErr bool + name string + jsonData string + wantID string + wantSize *string + wantTensorType *string + wantVariantID *string + wantMinVRAMGB *string + wantModelcarImageSize *string + wantModelcarImageSizeBytes *int64 + wantColdStartMatrix []coldStartEntry + wantErr bool }{ { name: "complete metadata with all new fields", @@ -1124,6 +1126,8 @@ func TestParseMetadataJSON_NewFields(t *testing.T) { "tensor_type": "FP8", "variant_group_id": "vwx234mn-5678-901v-wx23-456789abcdef", "min_vram_gb": "265 GB", + "modelcar_image_size": "230.17 GB", + "modelcar_image_size_bytes": 230171650363, "cold_start_matrix": [ { "gpu_type": "A100-80", @@ -1137,11 +1141,13 @@ func TestParseMetadataJSON_NewFields(t *testing.T) { } ] }`, - wantID: "sample-model/test-405b-instruct", - wantSize: &[]string{"405B params"}[0], - wantTensorType: &[]string{"FP8"}[0], - wantVariantID: &[]string{"vwx234mn-5678-901v-wx23-456789abcdef"}[0], - wantMinVRAMGB: &[]string{"265 GB"}[0], + wantID: "sample-model/test-405b-instruct", + wantSize: &[]string{"405B params"}[0], + wantTensorType: &[]string{"FP8"}[0], + wantVariantID: &[]string{"vwx234mn-5678-901v-wx23-456789abcdef"}[0], + wantMinVRAMGB: &[]string{"265 GB"}[0], + wantModelcarImageSize: &[]string{"230.17 GB"}[0], + wantModelcarImageSizeBytes: &[]int64{230171650363}[0], wantColdStartMatrix: []coldStartEntry{ {GPUType: "A100-80", GPUCount: "4", ColdStartTimeToLoadSeconds: "587.3"}, {GPUType: "B200", GPUCount: "2", ColdStartTimeToLoadSeconds: "559.9"}, @@ -1220,6 +1226,16 @@ func TestParseMetadataJSON_NewFields(t *testing.T) { t.Errorf("parseMetadataJSON() MinVRAMGB = %v, want %v", got.MinVRAMGB, tt.wantMinVRAMGB) } + // Test ModelcarImageSize field + if (got.ModelcarImageSize == nil) != (tt.wantModelcarImageSize == nil) || (got.ModelcarImageSize != nil && tt.wantModelcarImageSize != nil && *got.ModelcarImageSize != *tt.wantModelcarImageSize) { + t.Errorf("parseMetadataJSON() ModelcarImageSize = %v, want %v", got.ModelcarImageSize, tt.wantModelcarImageSize) + } + + // Test ModelcarImageSizeBytes field + if (got.ModelcarImageSizeBytes == nil) != (tt.wantModelcarImageSizeBytes == nil) || (got.ModelcarImageSizeBytes != nil && tt.wantModelcarImageSizeBytes != nil && *got.ModelcarImageSizeBytes != *tt.wantModelcarImageSizeBytes) { + t.Errorf("parseMetadataJSON() ModelcarImageSizeBytes = %v, want %v", got.ModelcarImageSizeBytes, tt.wantModelcarImageSizeBytes) + } + // Test ColdStartMatrix field if len(got.ColdStartMatrix) != len(tt.wantColdStartMatrix) { t.Errorf("parseMetadataJSON() ColdStartMatrix length = %d, want %d", len(got.ColdStartMatrix), len(tt.wantColdStartMatrix)) @@ -1477,9 +1493,13 @@ func TestEnrichCatalogModelFromMetadata_NewFields(t *testing.T) { } minVRAM := "80GB" + modelcarSize := "230.17 GB" + modelcarSizeBytes := int64(230171650363) metadata := metadataJSON{ - ID: modelName, - MinVRAMGB: &minVRAM, + ID: modelName, + MinVRAMGB: &minVRAM, + ModelcarImageSize: &modelcarSize, + ModelcarImageSizeBytes: &modelcarSizeBytes, ColdStartMatrix: []coldStartEntry{ {GPUType: "A100", GPUCount: "2", ColdStartTimeToLoadSeconds: "127.3"}, {GPUType: "H100", GPUCount: "1", ColdStartTimeToLoadSeconds: "68.9"}, @@ -1510,6 +1530,18 @@ func TestEnrichCatalogModelFromMetadata_NewFields(t *testing.T) { } else if v != "80GB" { t.Errorf("min_vram_gb = %q, want %q", v, "80GB") } + + if v, ok := propMap["modelcar_image_size"]; !ok { + t.Error("expected custom property 'modelcar_image_size' to be set") + } else if v != "230.17 GB" { + t.Errorf("modelcar_image_size = %q, want %q", v, "230.17 GB") + } + + if v, ok := propMap["modelcar_image_size_bytes"]; !ok { + t.Error("expected custom property 'modelcar_image_size_bytes' to be set") + } else if v != "230171650363" { + t.Errorf("modelcar_image_size_bytes = %q, want %q", v, "230171650363") + } } func TestCreateColdStartArtifact(t *testing.T) {