Merge pull request #428 from simvue-io/hotfix/pre-send-checks

Check metrics before sending them

Author: kzscisoft

simvue/models.py

@@ -1,12 +1,13 @@
 from typing import Annotated, Dict, List, Optional, Union
-
 from pydantic import BaseModel, Field, StringConstraints, PositiveInt
 
 FOLDER_REGEX: str = r"^/.*"
 NAME_REGEX: str = r"^[a-zA-Z0-9\-\_\s\/\.:]+$"
+METRIC_KEY_REGEX: str = r"^[a-zA-Z0-9\-\_\s\/\.:=><]+$"
 
 MetadataKeyString = Annotated[str, StringConstraints(pattern=r"^[\w\-\s\.]+$")]
 TagString = Annotated[str, StringConstraints(pattern=r"^[\w\-\s\.]+$")]
+MetricKeyString = Annotated[str, StringConstraints(pattern=METRIC_KEY_REGEX)]
 
 
 # Pydantic class to validate run.init()

simvue/run.py

@@ -38,7 +38,7 @@
 from .executor import Executor
 from .factory.proxy import Simvue
 from .metrics import get_gpu_metrics, get_process_cpu, get_process_memory
-from .models import RunInput, FOLDER_REGEX, NAME_REGEX
+from .models import RunInput, FOLDER_REGEX, NAME_REGEX, MetricKeyString
 from .serialization import serialize_object
 from .system import get_system
 from .metadata import git_info
@@ -1114,7 +1114,7 @@ def _add_metrics_to_dispatch(
     @pydantic.validate_call
     def log_metrics(
         self,
-        metrics: dict[str, typing.Union[int, float]],
+        metrics: dict[MetricKeyString, typing.Union[int, float]],
         step: typing.Optional[int] = None,
         time: typing.Optional[float] = None,
         timestamp: typing.Optional[str] = None,

simvue/utilities.py

@@ -18,6 +18,9 @@
 
 logger = logging.getLogger(__name__)
 
+if typing.TYPE_CHECKING:
+    from simvue.run import Run
+
 
 def parse_validation_response(
     response: dict[str, list[dict[str, str]]],
@@ -147,7 +150,7 @@ def skip_if_failed(
 
     def decorator(class_func: typing.Callable) -> typing.Callable:
         @functools.wraps(class_func)
-        def wrapper(self, *args, **kwargs) -> typing.Any:
+        def wrapper(self: "Run", *args, **kwargs) -> typing.Any:
             if getattr(self, failure_attr, None) and getattr(
                 self, ignore_exc_attr, None
             ):
@@ -166,7 +169,7 @@ def wrapper(self, *args, **kwargs) -> typing.Any:
                     setattr(self, failure_attr, True)
                     logger.error(error_str)
                     return on_failure_return
-                raise RuntimeError(error_str)
+                self._error(error_str)
 
         setattr(wrapper, "__fail_safe", True)
         return wrapper