[SC-11676] Improve API timeout support to prevent timeouts with large datasets

pyproject.toml

@@ -10,7 +10,7 @@ description = "ValidMind Library"
 license = "Commercial License"
 name = "validmind"
 readme = "README.pypi.md"
-version = "2.8.30"
+version = "2.8.31"
 
 [tool.poetry.dependencies]
 aiohttp = {extras = ["speedups"], version = "*"}

validmind/__version__.py

@@ -1 +1 @@
-__version__ = "2.8.30"
+__version__ = "2.8.31"

validmind/api_client.py

@@ -22,7 +22,7 @@
 
 from .client_config import client_config
 from .errors import MissingAPICredentialsError, MissingModelIdError, raise_api_error
-from .logging import get_logger, init_sentry, send_single_error
+from .logging import get_logger, init_sentry, log_api_operation, send_single_error
 from .utils import NumpyEncoder, is_html, md_to_html, run_async
 from .vm_models import Figure
 
@@ -85,7 +85,7 @@ def _get_session() -> aiohttp.ClientSession:
     if not __api_session or __api_session.closed:
         __api_session = aiohttp.ClientSession(
             headers=_get_api_headers(),
-            timeout=aiohttp.ClientTimeout(total=30),
+            timeout=aiohttp.ClientTimeout(total=int(os.getenv("VM_API_TIMEOUT", 30))),
         )
 
     return __api_session
@@ -304,6 +304,10 @@ async def alog_metadata(
         raise e
 
 
+@log_api_operation(
+    operation_name="Sending figure to ValidMind API",
+    extract_key=lambda figure: figure.key,
+)
 async def alog_figure(figure: Figure) -> Dict[str, Any]:
     """Logs a figure.
 

validmind/logging.py

@@ -170,6 +170,54 @@ async def wrap(*args: Any, **kwargs: Any) -> Any:
     return wrap
 
 
+def log_api_operation(
+    operation_name: Optional[str] = None,
+    logger: Optional[logging.Logger] = None,
+    extract_key: Optional[Callable] = None,
+    force: bool = False,
+) -> Callable[[F], F]:
+    """Decorator to log API operations like figure uploads.
+
+    Args:
+        operation_name (str, optional): The name of the operation. Defaults to function name.
+        logger (logging.Logger, optional): The logger to use. Defaults to None.
+        extract_key (Callable, optional): Function to extract a key from args for logging.
+        force (bool, optional): Whether to force logging even if env var is off.
+
+    Returns:
+        Callable: The decorated function.
+    """
+
+    def decorator(func: F) -> F:
+        # check if log level is set to debug
+        if _get_log_level() != logging.DEBUG and not force:
+            return func
+
+        nonlocal logger
+        if logger is None:
+            logger = get_logger()
+
+        nonlocal operation_name
+        if operation_name is None:
+            operation_name = func.__name__
+
+        async def wrapped(*args: Any, **kwargs: Any) -> Any:
+            # Try to extract a meaningful identifier from the arguments
+            identifier = ""
+            if extract_key and args:
+                try:
+                    identifier = f": {extract_key(args[0])}"
+                except (AttributeError, IndexError):
+                    pass
+
+            logger.debug(f"{operation_name}{identifier}")
+            return await func(*args, **kwargs)
+
+        return wrapped
+
+    return decorator
+
+
 def send_single_error(error: Exception) -> None:
     """Send a single error to Sentry.
 

validmind/vm_models/result/result.py

@@ -7,6 +7,7 @@
 """
 import asyncio
 import json
+import os
 from abc import abstractmethod
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union
@@ -20,7 +21,7 @@
 from ... import api_client
 from ...ai.utils import DescriptionFuture
 from ...errors import InvalidParameterError
-from ...logging import get_logger
+from ...logging import get_logger, log_api_operation
 from ...utils import (
     HumanReadableEncoder,
     NumpyEncoder,
@@ -476,9 +477,30 @@ async def log_async(
             )
 
         if self.figures:
-            tasks.extend(
-                [api_client.alog_figure(figure) for figure in (self.figures or [])]
+            batch_size = min(
+                len(self.figures), int(os.getenv("VM_FIGURE_MAX_BATCH_SIZE", 20))
             )
+            figure_batches = [
+                self.figures[i : i + batch_size]
+                for i in range(0, len(self.figures), batch_size)
+            ]
+
+            async def upload_figures_in_batches():
+                for batch in figure_batches:
+
+                    @log_api_operation(
+                        operation_name=f"Uploading batch of {len(batch)} figures"
+                    )
+                    async def process_batch():
+                        batch_tasks = [
+                            api_client.alog_figure(figure) for figure in batch
+                        ]
+                        return await asyncio.gather(*batch_tasks)
+
+                    await process_batch()
+
+            tasks.append(upload_figures_in_batches())
+
         if self.description:
             revision_name = (
                 AI_REVISION_NAME