Support run tokens

python/cog/schema.py

@@ -77,6 +77,8 @@ class PredictionRequest(PredictionBaseModel):
         default=WebhookEvent.default_events(),
     )
 
+    run_token: Optional[str] = None
+
     @classmethod
     def with_types(cls, input_type: Type[Any]) -> Any:
         # [compat] Input is implicitly optional -- previous versions of the

python/cog/server/eventtypes.py

@@ -69,3 +69,4 @@ class Envelope:
         Done,
     ]
     tag: Optional[str] = None
+    run_token: Optional[str] = None

python/cog/server/runner.py

@@ -113,7 +113,9 @@ def predict(
             payload = prediction.input.copy()
 
         sid = self._worker.subscribe(task.handle_event, tag=tag)
-        task.track(self._worker.predict(payload, tag=tag))
+        task.track(
+            self._worker.predict(payload, tag=tag, run_token=prediction.run_token)
+        )
         task.add_done_callback(self._task_done_callback(tag, sid))
 
         return task

python/cog/server/worker.py

@@ -67,6 +67,9 @@
 _tag_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
     "tag", default=None
 )
+_run_token_var: contextvars.ContextVar[Optional[str]] = contextvars.ContextVar(
+    "run_token", default=None
+)
 
 _PublicEventType = Union[Done, Log, PredictionOutput, PredictionOutputType]
 
@@ -139,7 +142,10 @@ def setup(self) -> "Future[Done]":
         return self._setup_result
 
     def predict(
-        self, payload: Dict[str, Any], tag: Optional[str] = None
+        self,
+        payload: Dict[str, Any],
+        tag: Optional[str] = None,
+        run_token: Optional[str] = None,
     ) -> "Future[Done]":
         # TODO: tag is Optional, but it's required when in concurrent mode and
         # basically unnecessary in sequential mode. Should we have a separate
@@ -162,11 +168,13 @@ def predict(
             result = Future()
             self._predictions_in_flight[tag] = PredictionState(tag, payload, result)
 
-        self._prediction_start_pool.submit(self._start_prediction(tag, payload))
+        self._prediction_start_pool.submit(
+            self._start_prediction(tag, payload, run_token)
+        )
         return result
 
     def _start_prediction(
-        self, tag: Optional[str], payload: Dict[str, Any]
+        self, tag: Optional[str], payload: Dict[str, Any], run_token: Optional[str]
     ) -> Callable[[], None]:
         def start_prediction() -> None:
             try:
@@ -219,6 +227,7 @@ def start_prediction() -> None:
                     Envelope(
                         event=PredictionInput(payload=payload),
                         tag=tag,
+                        run_token=run_token,
                     )
                 )
             except Exception as e:
@@ -489,6 +498,10 @@ def _current_tag(self) -> Optional[str]:
             return tag
         return self._sync_tag
 
+    @property
+    def _current_run_token(self) -> Optional[str]:
+        return _run_token_var.get()
+
     def _load_predictor(self) -> Optional[BasePredictor]:
         done = Done()
         wait_for_env()
@@ -590,7 +603,7 @@ def _loop(
             elif isinstance(e.event, Shutdown):
                 break
             elif isinstance(e.event, PredictionInput):
-                self._predict(e.tag, e.event.payload, predict, redirector)
+                self._predict(e.tag, e.event.payload, predict, redirector, e.run_token)
             else:
                 print(f"Got unexpected event: {e.event}", file=sys.stderr)
 
@@ -624,7 +637,9 @@ async def _aloop(
                     break
                 elif isinstance(e.event, PredictionInput):
                     tasks[e.tag] = tg.create_task(
-                        self._apredict(e.tag, e.event.payload, predict, redirector)
+                        self._apredict(
+                            e.tag, e.event.payload, predict, redirector, e.run_token
+                        )
                     )
                 else:
                     print(f"Got unexpected event: {e.event}", file=sys.stderr)
@@ -635,7 +650,9 @@ def _predict(
         payload: Dict[str, Any],
         predict: Callable[..., Any],
         redirector: StreamRedirector,
+        run_token: Optional[str],
     ) -> None:
+        _run_token_var.set(run_token)
         with self._handle_predict_error(redirector, tag=tag):
             result = predict(**payload)
 
@@ -686,8 +703,10 @@ async def _apredict(
         payload: Dict[str, Any],
         predict: Callable[..., Any],
         redirector: SimpleStreamRedirector,
+        run_token: Optional[str],
     ) -> None:
         _tag_var.set(tag)
+        _run_token_var.set(run_token)
 
         with self._handle_predict_error(redirector, tag=tag):
             future_result = predict(**payload)