Optionally automatically cancel predictions with fatal errors

If our application relies on the all-in-one predict method, there's a
good chance we don't want to be stuck with a running prediction we're
paying for but won't ever use.

Author: erbridge

lib/Model.js

@@ -48,10 +48,16 @@ export default class Model extends ReplicateObject {
 
   async predict(
     input,
-    { onUpdate = noop, onTemporaryError = noop } = {},
+    {
+      onUpdate = noop,
+      onTemporaryError = noop,
+      onCancel = noop,
+      onCancelError = noop,
+    } = {},
     {
       defaultPollingInterval = 500,
       backoffFn = (errorCount) => Math.pow(2, errorCount) * 100,
+      cancelOnFatalError = false,
     } = {}
   ) {
     if (!input) {
@@ -60,39 +66,57 @@ export default class Model extends ReplicateObject {
 
     let prediction = await this.createPrediction(input);
 
-    onUpdate(prediction);
+    try {
+      onUpdate(prediction);
 
-    let pollingInterval = defaultPollingInterval;
-    let errorCount = 0;
+      let pollingInterval = defaultPollingInterval;
+      let errorCount = 0;
 
-    while (!prediction.hasTerminalStatus()) {
-      await sleep(pollingInterval);
-      pollingInterval = defaultPollingInterval; // Reset to default each time.
+      while (!prediction.hasTerminalStatus()) {
+        await sleep(pollingInterval);
+        pollingInterval = defaultPollingInterval; // Reset to default each time.
 
-      try {
-        prediction = await this.client.prediction(prediction.id).load();
+        try {
+          prediction = await this.client.prediction(prediction.id).load();
 
-        onUpdate(prediction);
+          onUpdate(prediction);
 
-        errorCount = 0; // Reset because we've had a non-error response.
-      } catch (err) {
-        if (!err instanceof ReplicateResponseError) {
-          throw err;
-        }
+          errorCount = 0; // Reset because we've had a non-error response.
+        } catch (err) {
+          if (!err instanceof ReplicateResponseError) {
+            throw err;
+          }
 
-        if (
-          !err.status ||
-          (Math.floor(err.status / 100) !== 5 && err.status !== 429)
-        ) {
-          throw err;
-        }
+          if (
+            !err.status ||
+            (Math.floor(err.status / 100) !== 5 && err.status !== 429)
+          ) {
+            throw err;
+          }
 
-        errorCount += 1;
+          errorCount += 1;
 
-        onTemporaryError(err);
+          onTemporaryError(err);
 
-        pollingInterval = backoffFn(errorCount);
+          pollingInterval = backoffFn(errorCount);
+        }
+      }
+    } catch (err) {
+      if (cancelOnFatalError) {
+        // We intentionally don't await this, so we don't block.
+        prediction
+          .cancel()
+          .catch((e) => {
+            onCancelError(e);
+
+            throw e;
+          })
+          .then(() => {
+            onCancel();
+          });
       }
+
+      throw err;
     }
 
     return prediction;