Support more values for "metrics"

python/cog/server/eventtypes.py

@@ -32,7 +32,7 @@ class Log:
 @define
 class PredictionMetric:
     name: str
-    value: Union[float, int]
+    value: Optional[Union[bool, float, int, str]]
 
 
 @define

python/cog/server/runner.py

@@ -390,10 +390,15 @@ def append_logs(self, logs: str) -> None:
         self._p.logs += logs
         self._send_webhook(schema.WebhookEvent.LOGS)
 
-    def set_metric(self, key: str, value: Union[float, int]) -> None:
+    def set_metric(self, key: str, value: Optional[Union[bool, float, int, str]]) -> None:
         if self._p.metrics is None:
             self._p.metrics = {}
-        self._p.metrics[key] = value
+
+        if value is None:
+            if key in self._p.metrics:
+                del self._p.metrics[key]
+        else:
+            self._p.metrics[key] = value
 
     def succeeded(self) -> None:
         self._log.info(("prediction" if not self._is_train else "train") + " succeeded")

python/cog/server/scope.py

@@ -10,7 +10,7 @@
 
 @frozen
 class Scope:
-    record_metric: Callable[[str, Union[float, int]], None]
+    record_metric: Callable[[str, Optional[Union[bool, float, int, str]]], None]
     context: Dict[str, str] = {}
     _tag: Optional[str] = None
 

python/cog/server/worker.py

@@ -489,7 +489,7 @@ def send_cancel_signal(self) -> None:
         if self.is_alive() and self.pid:
             os.kill(self.pid, signal.SIGUSR1)
 
-    def record_metric(self, name: str, value: Union[float, int]) -> None:
+    def record_metric(self, name: str, value: Optional[Union[bool, float, int, str]]) -> None:
         self._events.send(
             Envelope(PredictionMetric(name, value), tag=self._current_tag)
         )

python/tests/server/test_worker.py

@@ -194,7 +194,12 @@ def handle_event(self, event: _PublicEventType):
         elif isinstance(event, PredictionMetric):
             if self.metrics is None:
                 self.metrics = {}
-            self.metrics[event.name] = event.value
+
+            if event.value is None:
+                if event.name in self.metrics:
+                    del self.metrics[event.name]
+            else:
+                self.metrics[event.name] = event.value
         elif isinstance(event, PredictionOutput):
             assert self.output_type, "Should get output type before any output"
             if self.output_type.multi: