Cancel input downloads if one of them fails

python/cog/server/worker.py

@@ -186,7 +186,23 @@ def start_prediction() -> None:
                             self._input_download_pool.submit(item.convert) for item in v
                         ]
                         to_await += futs[k]
-                futures.wait(to_await, return_when=futures.FIRST_EXCEPTION)
+                done, not_done = futures.wait(to_await, return_when=futures.FIRST_EXCEPTION)
+
+                if len(not_done) > 0:
+                    # if any future isn't done, this is because one of the
+                    # futures raised an exception. first we cancel outstanding
+                    # work
+                    for fut in not_done:
+                        fut.cancel()
+                    # then we find an exception to raise
+                    for fut in done:
+                        fut.result()  # raises if the future finished with an exception
+                    # we should never get here
+                    raise Exception("Internal error: lost track of exception while downloading input files")
+
+                # all futures are done. some might still have raised an
+                # exception, but when we call fut.result() that will re-raise
+                # and do the right thing
                 for k, v in futs.items():
                     if isinstance(v, list):
                         payload[k] = []