Tag worker events with prediction (#2020)

* Formatting, moving `Envelope`, and using "tag"

* Include tag in Cancel event

and clean up TODO comments

* Add pydantic v2 support to child worker `_apredict`

* Keep a map of thread IDs to tags on _ChildWorker

* Use context var to pass tag instead of dict

since context vars are meant to be async-safe

* don't use the contextvar in sync StreamRedirector

* fix StreamPredictor working with tags

* fix record_event with tags

We were unintentionally dropping the tag from record_event() with sync
predictions.

This is fixed by introducing a self._current_tag property which Does The Right
Thing whether in sync or async context.

As part of this, we can also go back to a single _stream_write_hook()
implementation.

---------

Co-authored-by: Dominic Baggott <[email protected]>
Co-authored-by: Philip Potter <[email protected]>

Author: 3 people

python/cog/server/eventtypes.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Union
 
 from attrs import define, field, validators
 
@@ -7,7 +7,6 @@
 #
 @define
 class Cancel:
-    # TODO: identify which prediction!
     pass
 
 
@@ -50,3 +49,23 @@ class Done:
     canceled: bool = False
     error: bool = False
     error_detail: str = ""
+
+
+@define
+class Envelope:
+    """
+    Envelope contains an arbitrary event along with an optional tag used to
+    tangle/untangle concurrent work.
+    """
+
+    event: Union[
+        Cancel,
+        PredictionInput,
+        Shutdown,
+        Log,
+        PredictionMetric,
+        PredictionOutput,
+        PredictionOutputType,
+        Done,
+    ]
+    tag: Optional[str] = None