bad wip

Signed-off-by: technillogue <technillogue@gmail.com>
replicate · Nov 28, 2023 · 34d05af · 34d05af
1 parent 40a9f48
commit 34d05af
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/python/cog/server/eventtypes.py b/python/cog/server/eventtypes.py
@@ -7,6 +7,7 @@
 #
 @define
 class PredictionInput:
+    id: str
     payload: Dict[str, Any]
 
 
@@ -25,16 +26,19 @@ class Log:
 
 @define
 class PredictionOutput:
+    id: str
     payload: Any
 
 
 @define
 class PredictionOutputType:
+    id: str
     multi: bool = False
 
 
 @define
 class Done:
+    id: str
     canceled: bool = False
     error: bool = False
     error_detail: str = ""

diff --git a/python/cog/server/worker.py b/python/cog/server/worker.py
@@ -58,12 +58,12 @@ def setup(self) -> Iterable[_PublicEventType]:
         return self._wait(raise_on_error="Predictor errored during setup")
 
     def predict(
-        self, payload: Dict[str, Any], poll: Optional[float] = None
+        self, id: str, payload: Dict[str, Any], poll: Optional[float] = None
     ) -> None: #Iterator[_PublicEventType]:
         self._assert_state(WorkerState.READY)
         self._state = WorkerState.PROCESSING
         self._allow_cancel = True
-        self._events.send(PredictionInput(payload=payload))
+        self._events.send(PredictionInput(id=id, payload=payload))
 
         # return self._wait(poll=poll)