Expose prediction IDs for model code (#2591)

Signed-off-by: Nikhil Sinha <nikhil@replicate.com>

Author: NikhilSinha1

python/cog/server/runner.py

@@ -113,6 +113,11 @@ def predict(
         else:
             payload = prediction.input.copy()
 
+        if prediction.context is None:
+            prediction.context = {}
+        if prediction.id is not None:
+            prediction.context["id"] = prediction.id
+
         sid = self._worker.subscribe(task.handle_event, tag=tag)
         task.track(self._worker.predict(payload, context=prediction.context, tag=tag))
         task.add_done_callback(self._task_done_callback(tag, sid))