Add more exact types to request context interfaces and validate them (#465)

If we don't validate ourselfs then we'd be getting internal errors now
because we're using pydantic models to pass user data from request
context client to server.

Author: eabatalov

src/tensorlake/applications/interface/request_context.py

@@ -32,7 +32,7 @@ def get(self, key: str, default: Any | None = None) -> Any | None:
 class RequestMetrics:
     """Abstract interface for reporting application request metrics."""
 
-    def timer(self, name: str, value: float) -> None:
+    def timer(self, name: str, value: int | float) -> None:
         """Records a duration metric with the supplied name and value.
 
         Raises TensorlakeError on error.
@@ -54,8 +54,8 @@ class FunctionProgress:
 
     def update(
         self,
-        current: float,
-        total: float,
+        current: int | float,
+        total: int | float,
         message: str | None = None,
         attributes: dict[str, str] | None = None,
     ) -> None:

src/tensorlake/applications/request_context/http_client/metrics.py

@@ -22,7 +22,14 @@ def __init__(self, request_id: str, allocation_id: str, http_client: httpx.Clien
         self._allocation_id: str = allocation_id
         self._http_client: httpx.Client = http_client
 
-    def timer(self, name: str, value: float):
+    def timer(self, name: str, value: int | float):
+        # If we don't validate user supplied inputs here then there will be a Pydantic validation error
+        # below which will raise an InternalError instead of SDKUsageError.
+        if not isinstance(name, str):
+            raise SDKUsageError(f"Timer name must be a string, got: {name}")
+        if not isinstance(value, (int, float)):
+            raise SDKUsageError(f"Timer value must be a number, got: {value}")
+
         request_payload: AddMetricsRequest = AddMetricsRequest(
             request_id=self._request_id,
             allocation_id=self._allocation_id,
@@ -32,6 +39,13 @@ def timer(self, name: str, value: float):
         self._run_add_request(request_payload)
 
     def counter(self, name: str, value: int = 1):
+        # If we don't validate user supplied inputs here then there will be a Pydantic validation error
+        # below which will raise an InternalError instead of SDKUsageError.
+        if not isinstance(name, str):
+            raise SDKUsageError(f"Counter name must be a string, got: {name}")
+        if not isinstance(value, int):
+            raise SDKUsageError(f"Counter value must be an int, got: {value}")
+
         request_payload: AddMetricsRequest = AddMetricsRequest(
             request_id=self._request_id,
             allocation_id=self._allocation_id,

src/tensorlake/applications/request_context/http_client/progress.py

@@ -2,7 +2,7 @@
 
 import httpx
 
-from ...interface.exceptions import InternalError, SDKUsageError, SerializationError
+from ...interface.exceptions import InternalError, SDKUsageError
 from ...interface.request_context import FunctionProgress
 from ..http_server.handlers.progress_update import (
     PROGRESS_UPDATE_PATH,
@@ -24,12 +24,20 @@ def __init__(self, request_id: str, allocation_id: str, http_client: httpx.Clien
 
     def update(
         self,
-        current: float,
-        total: float,
+        current: int | float,
+        total: int | float,
         message: str | None = None,
         attributes: dict[str, str] | None = None,
     ) -> None:
-        # Instead of handling serialization errors on the Server, just validate attributes on client side.
+        # If we don't validate user supplied inputs here then there will be a Pydantic validation error
+        # below which will raise an InternalError instead of SDKUsageError.
+        if not isinstance(current, (int, float)):
+            raise SDKUsageError(f"'current' needs to be a number, got: {current}")
+        if not isinstance(total, (int, float)):
+            raise SDKUsageError(f"'total' needs to be a number, got: {total}")
+        if message is not None and not isinstance(message, str):
+            raise SDKUsageError(f"'message' needs to be a string, got: {message}")
+
         if attributes is not None:
             if not isinstance(attributes, dict):
                 raise SDKUsageError(

src/tensorlake/applications/request_context/http_client/state.py

@@ -53,6 +53,10 @@ def set(self, key: str, value: Any) -> None:
         # NB: This is called from user code, user code is blocked.
         # Any exception raised here goes directly to user code.
 
+        # If we don't validate user supplied inputs here then there will be a Pydantic validation error
+        # below which will raise an InternalError instead of SDKUsageError.
+        if not isinstance(key, str):
+            raise SDKUsageError(f"State key must be a string, got: {key}")
         # Raises SerializationError to customer code on failure.
         serialized_value: bytes = REQUEST_STATE_USER_DATA_SERIALIZER.serialize(value)
 
@@ -75,6 +79,12 @@ def set(self, key: str, value: Any) -> None:
     def get(self, key: str, default: Any | None = None) -> Any | None:
         # NB: This is called from user code, user code is blocked.
         # Any exception raised here goes directly to user code.
+
+        # If we don't validate user supplied inputs here then there will be a Pydantic validation error
+        # below which will raise an InternalError instead of SDKUsageError.
+        if not isinstance(key, str):
+            raise SDKUsageError(f"State key must be a string, got: {key}")
+
         try:
             blob: BLOB | None = self._get_read_only_blob(key=key)
             if blob is None:

src/tensorlake/applications/request_context/http_server/handlers/add_metrics.py

@@ -8,7 +8,7 @@
 
 class AddTimerRequest(BaseModel):
     name: str
-    value: float
+    value: int | float
 
 
 class AddCounterRequest(BaseModel):