fix: log better

Author: hanchchch

lavender_data/client/api.py

@@ -377,8 +377,7 @@ def get_next_item(
                 no_cache=no_cache,
                 max_retry_count=max_retry_count,
             )
-        content = self._check_response(response).payload.read()
-        return deserialize_sample(content)
+        return self._check_response(response).payload.read()
 
     def submit_next_item(
         self,
@@ -406,8 +405,7 @@ def get_submitted_result(self, iteration_id: str, cache_key: str):
             )
         if response.status_code == 202:
             raise LavenderDataApiError(response.content.decode("utf-8"))
-        content = self._check_response(response).payload.read()
-        return deserialize_sample(content)
+        return self._check_response(response).payload.read()
 
     def complete_index(self, iteration_id: str, index: int):
         with self._get_client() as client:

lavender_data/client/cli/api_call.py

@@ -1,4 +1,5 @@
 from lavender_data.client import LavenderDataClient
+from lavender_data.serialize import deserialize_sample
 
 
 def _api(api_url: str, api_key: str):
@@ -77,11 +78,13 @@ def get_next_item(
     no_cache: bool,
     max_retry_count: int,
 ):
-    return _api(api_url=api_url, api_key=api_key).get_next_item(
-        iteration_id=iteration_id,
-        rank=rank,
-        no_cache=no_cache,
-        max_retry_count=max_retry_count,
+    return deserialize_sample(
+        _api(api_url=api_url, api_key=api_key).get_next_item(
+            iteration_id=iteration_id,
+            rank=rank,
+            no_cache=no_cache,
+            max_retry_count=max_retry_count,
+        )
     )
 
 
@@ -102,8 +105,10 @@ def submit_next_item(
 
 
 def get_submitted_result(api_url: str, api_key: str, iteration_id: str, cache_key: str):
-    return _api(api_url=api_url, api_key=api_key).get_submitted_result(
-        iteration_id=iteration_id, cache_key=cache_key
+    return deserialize_sample(
+        _api(api_url=api_url, api_key=api_key).get_submitted_result(
+            iteration_id=iteration_id, cache_key=cache_key
+        )
     )
 
 

lavender_data/client/iteration.py

@@ -3,6 +3,7 @@
 from typing import Optional, Union, Literal
 from concurrent.futures import ThreadPoolExecutor, Future, as_completed
 
+from lavender_data.serialize import deserialize_sample, DeserializeException
 from lavender_data.client.api import (
     get_client,
     LavenderDataClient,
@@ -140,6 +141,8 @@ def __init__(
         self._api_key = api_key
         self._api = _api(self._api_url, self._api_key)
 
+        self._bytes = 0
+
         self.id = self._iteration_id
 
     def torch(
@@ -215,19 +218,21 @@ def _set_last_indices(self, sample_or_batch):
 
     def _get_next_item(self):
         try:
-            sample_or_batch = self._api.get_next_item(
+            serialized = self._api.get_next_item(
                 iteration_id=self._iteration_id,
                 rank=self._rank,
                 no_cache=self._no_cache,
                 max_retry_count=self._max_retry_count,
             )
+            self._bytes += len(serialized)
+            return deserialize_sample(serialized)
         except LavenderDataApiError as e:
             if "No more indices to pop" in str(e):
                 raise StopIteration
             else:
                 raise e
-
-        return sample_or_batch
+        except DeserializeException as e:
+            raise ValueError(f"Failed to deserialize sample: {e}")
 
     def _submit_next_item(self) -> str:
         cache_key = self._api.submit_next_item(
@@ -240,10 +245,12 @@ def _submit_next_item(self) -> str:
 
     def _get_submitted_result(self, cache_key: str):
         try:
-            return self._api.get_submitted_result(
+            serialized = self._api.get_submitted_result(
                 iteration_id=self._iteration_id,
                 cache_key=cache_key,
             )
+            self._bytes += len(serialized)
+            return deserialize_sample(serialized)
         except LavenderDataApiError as e:
             if "Data is still being processed" in str(e):
                 return None
@@ -253,6 +260,8 @@ def _get_submitted_result(self, cache_key: str):
                 raise StopIteration
             else:
                 raise e
+        except DeserializeException as e:
+            raise ValueError(f"Failed to deserialize sample: {e}")
 
     def __next__(self):
         self._complete_last_indices()

lavender_data/serialize.py

@@ -128,6 +128,10 @@ def serialize_sample(sample: dict):
     return header + body
 
 
+class DeserializeException(Exception):
+    pass
+
+
 def deserialize_sample(content: bytes, strict: bool = True):
     header_length, current = detach_length(content)
     keys = json.loads(current[:header_length].decode("utf-8"))
@@ -152,7 +156,7 @@ def deserialize_sample(content: bytes, strict: bool = True):
                 warnings.warn(msg)
                 values.append(None)
             else:
-                raise ValueError(msg)
+                raise DeserializeException(msg)
         current = value[value_length:]
         i += 1
 

lavender_data/server/app.py

@@ -126,8 +126,13 @@ async def add_process_time_header(request: Request, call_next):
     process_time = time.perf_counter() - start_time
 
     if log_filter(request, response):
+        path = (
+            f"{request.url.path}?{request.url.query}"
+            if request.url.query
+            else request.url.path
+        )
         logger.info(
-            f"{request.client.host}:{request.client.port} - {request.method} {request.url.path}?{request.url.query} {response.status_code} {process_time:.2f}s"
+            f"{request.client.host}:{request.client.port} - {request.method} {path} {response.status_code} {process_time:.2f}s"
         )
 
     return response

lavender_data/server/iteration/process.py

@@ -165,6 +165,21 @@ def process_next_samples(
                 raise error
 
 
+def _format_number(number: int):
+    if number < 1000:
+        return f"{number} "
+    elif number < 1000**2:
+        return f"{number/1000:.2f} K"
+    elif number < 1000**3:
+        return f"{number/1000**2:.2f} M"
+    else:
+        return f"{number/1000**3:.2f} G"
+
+
+def _ms(seconds: float):
+    return f"{seconds*1000:.2f} ms"
+
+
 @pool_task()
 def process_next_samples_and_store(
     params: ProcessNextSamplesParams,
@@ -184,7 +199,11 @@ def process_next_samples_and_store(
         shared_memory.set(cache_key, content, ex=cache_ttl)
         _store_time = time.perf_counter()
         logger.debug(
-            f"Done processing {cache_key} in {_store_time-_start:.2f}s, process: {_process_time-_start:.2f}s, serialize: {_serialize_time - _process_time:.2f}s, store: {_store_time - _serialize_time:.2f}s, size: {len(content)} bytes"
+            f"Done processing {cache_key} in {_ms(_store_time - _start)}, "
+            f"process: {_ms(_process_time - _start)}, "
+            f"serialize: {_ms(_serialize_time - _process_time)}, "
+            f"store: {_ms(_store_time - _serialize_time)}, "
+            f"size: {_format_number(len(content))}B"
         )
     except ProcessNextSamplesException as e:
         shared_memory.set(cache_key, f"processing_error:{e.json()}", ex=cache_ttl)