From 625a8743a26ba832893c42403c562fefeb30a98e Mon Sep 17 00:00:00 2001 From: Juan Date: Thu, 7 Aug 2025 23:45:33 +0200 Subject: [PATCH 1/2] Make timeout configurable and send figures in batches --- validmind/api_client.py | 8 +++-- validmind/logging.py | 48 ++++++++++++++++++++++++++++ validmind/vm_models/result/result.py | 28 ++++++++++++++-- 3 files changed, 79 insertions(+), 5 deletions(-) diff --git a/validmind/api_client.py b/validmind/api_client.py index 7bee2290b..1eb1cd5c4 100644 --- a/validmind/api_client.py +++ b/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. diff --git a/validmind/logging.py b/validmind/logging.py index 1cb81ec73..3b0322b39 100644 --- a/validmind/logging.py +++ b/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. diff --git a/validmind/vm_models/result/result.py b/validmind/vm_models/result/result.py index ecc763af4..6edee7bbe 100644 --- a/validmind/vm_models/result/result.py +++ b/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 From b348d7d4d2d0fc7515caeb9520eb376b0bb0c0ef Mon Sep 17 00:00:00 2001 From: Juan Date: Fri, 8 Aug 2025 09:19:39 +0200 Subject: [PATCH 2/2] 2.8.31 --- pyproject.toml | 2 +- validmind/__version__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 59d04eff8..12ceea155 100644 --- a/pyproject.toml +++ b/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 = "*"} diff --git a/validmind/__version__.py b/validmind/__version__.py index 839fa31f2..3a5995d32 100644 --- a/validmind/__version__.py +++ b/validmind/__version__.py @@ -1 +1 @@ -__version__ = "2.8.30" +__version__ = "2.8.31"