From 8b8d8d6e6f459cd4613f066f3001778c41478c19 Mon Sep 17 00:00:00 2001
From: clearml <>
Date: Thu, 5 Dec 2024 22:33:52 +0200
Subject: [PATCH] Change model input_size field to string

---
 apiserver/apimodels/__init__.py        | 10 ++++++++++
 apiserver/apimodels/serving.py         |  4 ++--
 apiserver/schema/services/serving.conf |  8 ++++----
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/apiserver/apimodels/__init__.py b/apiserver/apimodels/__init__.py
index 9038b96..49123c7 100644
--- a/apiserver/apimodels/__init__.py
+++ b/apiserver/apimodels/__init__.py
@@ -1,4 +1,5 @@
 from enum import Enum
+from numbers import Number
 from typing import Union, Type, Iterable, Mapping
 
 import jsonmodels.errors
@@ -96,6 +97,15 @@ class ScalarField(fields.BaseField):
     types = (str, int, float, bool)
 
 
+class SafeStringField(fields.StringField):
+    """String field that can also accept numbers as input"""
+    def parse_value(self, value):
+        if isinstance(value, Number):
+            value = str(value)
+
+        return super().parse_value(value)
+
+
 class DictField(fields.BaseField):
     types = (dict,)
 
diff --git a/apiserver/apimodels/serving.py b/apiserver/apimodels/serving.py
index 95c2061..64e39d9 100644
--- a/apiserver/apimodels/serving.py
+++ b/apiserver/apimodels/serving.py
@@ -13,7 +13,7 @@ from jsonmodels.fields import (
 from jsonmodels import validators
 from jsonmodels.validators import Min
 
-from apiserver.apimodels import ListField, JsonSerializableMixin
+from apiserver.apimodels import ListField, JsonSerializableMixin, SafeStringField
 from apiserver.apimodels import ActualEnumField
 from apiserver.config_repo import config
 from .workers import MachineStats
@@ -36,7 +36,7 @@ class ServingModel(Base):
     model_version = StringField()
     preprocess_artifact = StringField()
     input_type = StringField()
-    input_size = IntField()
+    input_size = SafeStringField()
     tags = ListField(str)
     system_tags = ListField(str)
     reference: Sequence[ReferenceItem] = ListField(ReferenceItem)
diff --git a/apiserver/schema/services/serving.conf b/apiserver/schema/services/serving.conf
index 7a5106a..e280059 100644
--- a/apiserver/schema/services/serving.conf
+++ b/apiserver/schema/services/serving.conf
@@ -58,8 +58,8 @@ _definitions {
                 description: Input type
             }
             input_size {
-                type: integer
-                description: Input size in bytes
+                type: string
+                description: Input size
             }
             reference: ${_definitions.reference}
         }
@@ -178,8 +178,8 @@ _definitions {
                 description: Input type
             }
             input_size {
-                type: integer
-                description: Input size in bytes
+                type: string
+                description: Input size
             }
         }
     }