Support more values for "metrics"

We now support boolean, null, and string values for some metrics
(technically, they're not metrics, because you can't sum over them, but
they're wired up the same and this is an experimental intenal function).

Author: erbridge

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 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: