Use lower level invocation model

Author: rculbertson

modal-js/src/function.ts

@@ -5,23 +5,24 @@ import { createHash } from "node:crypto";
 import {
   DataFormat,
   DeploymentNamespace,
-  FunctionPutInputsItem,
+  FunctionCallInvocationType,
+  FunctionInput,
 } from "../proto/modal_proto/api";
 import type { LookupOptions } from "./app";
 import { client } from "./client";
 import { FunctionCall } from "./function_call";
 import { environmentName } from "./config";
-import { NotFoundError } from "./errors";
+import { InternalFailure, NotFoundError } from "./errors";
 import { dumps } from "./pickle";
 import { ClientError, Status } from "nice-grpc";
-import {
-  ControlPlaneStrategy,
-  pollControlPlaneForOutput,
-} from "./invocation_strategy";
+import { ControlPlaneInvocation } from "./invocation";
 
 // From: modal/_utils/blob_utils.py
 const maxObjectSizeBytes = 2 * 1024 * 1024; // 2 MiB
 
+// From: client/modal/_functions.py
+const maxSystemRetries = 8;
+
 /** Represents a deployed Modal Function, which can be invoked remotely. */
 export class Function_ {
   readonly functionId: string;
@@ -59,8 +60,25 @@ export class Function_ {
     kwargs: Record<string, any> = {},
   ): Promise<any> {
     const input = await this.#createInput(args, kwargs);
-    const invocationStrategy = new ControlPlaneStrategy(this.functionId, input);
-    return await invocationStrategy.remote();
+    const invocation = await ControlPlaneInvocation.create(
+      this.functionId,
+      input,
+      FunctionCallInvocationType.FUNCTION_CALL_INVOCATION_TYPE_SYNC,
+    );
+    // TODO(ryan): Add tests for retries.
+    let retryCount = 0;
+    while (true) {
+      try {
+        return await invocation.await();
+      } catch (err) {
+        if (err instanceof InternalFailure && retryCount <= maxSystemRetries) {
+          await invocation.retry(retryCount);
+          retryCount++;
+        } else {
+          throw err;
+        }
+      }
+    }
   }
 
   // Spawn a single input into a remote function.
@@ -69,15 +87,18 @@ export class Function_ {
     kwargs: Record<string, any> = {},
   ): Promise<FunctionCall> {
     const input = await this.#createInput(args, kwargs);
-    const invocationStrategy = new ControlPlaneStrategy(this.functionId, input);
-    const functionCallId = await invocationStrategy.spawn();
-    return new FunctionCall(functionCallId);
+    const invocation = await ControlPlaneInvocation.create(
+      this.functionId,
+      input,
+      FunctionCallInvocationType.FUNCTION_CALL_INVOCATION_TYPE_ASYNC,
+    );
+    return new FunctionCall(invocation.functionCallId);
   }
 
   async #createInput(
     args: any[] = [],
     kwargs: Record<string, any> = {},
-  ): Promise<FunctionPutInputsItem> {
+  ): Promise<FunctionInput> {
     const payload = dumps([args, kwargs]);
 
     let argsBlobId: string | undefined = undefined;
@@ -87,25 +108,15 @@ export class Function_ {
 
     // Single input sync invocation
     return {
-      idx: 0,
-      input: {
-        args: argsBlobId ? undefined : payload,
-        argsBlobId,
-        dataFormat: DataFormat.DATA_FORMAT_PICKLE,
-        methodName: this.methodName,
-        finalInput: false, // This field isn't specified in the Python client, so it defaults to false.
-      },
+      args: argsBlobId ? undefined : payload,
+      argsBlobId,
+      dataFormat: DataFormat.DATA_FORMAT_PICKLE,
+      methodName: this.methodName,
+      finalInput: false, // This field isn't specified in the Python client, so it defaults to false.
     };
   }
 }
 
-export async function pollFunctionOutput(
-  functionCallId: string,
-  timeout?: number, // in milliseconds
-): Promise<any> {
-  return pollControlPlaneForOutput(functionCallId, timeout);
-}
-
 async function blobUpload(data: Uint8Array): Promise<string> {
   const contentMd5 = createHash("md5").update(data).digest("base64");
   const contentSha256 = createHash("sha256").update(data).digest("base64");

modal-js/src/function_call.ts

@@ -1,7 +1,7 @@
 // Manage existing Function Calls (look-ups, polling for output, cancellation).
 
 import { client } from "./client";
-import { pollFunctionOutput } from "./function";
+import { ControlPlaneInvocation } from "./invocation";
 
 /** Options for `FunctionCall.get()`. */
 export type FunctionCallGetOptions = {
@@ -34,7 +34,10 @@ export class FunctionCall {
   /** Get the result of a function call, optionally waiting with a timeout. */
   async get(options: FunctionCallGetOptions = {}): Promise<any> {
     const timeout = options.timeout;
-    return await pollFunctionOutput(this.functionCallId, timeout);
+    const invocation = ControlPlaneInvocation.fromFunctionCallId(
+      this.functionCallId,
+    );
+    return invocation.await(timeout);
   }
 
   /** Cancel a running function call. */

modal-js/src/invocation.ts

@@ -30,18 +30,17 @@ export interface Invocation {
    */
   await(): Promise<any>;
 
-  retry(): Promise<void>;
+  retry(retryCount: number): Promise<void>;
 }
 
 /**
- * Implementation of InvocationStrategy which sends inputs to the control plane.
+ * Implementation of Invocation which sends inputs to the control plane.
  */
 export class ControlPlaneInvocation implements Invocation {
   readonly functionCallId: string;
   private readonly input?: FunctionInput;
   private readonly functionCallJwt?: string;
   private inputJwt?: string;
-  private retryCount: number = 0;
 
   private constructor(
     functionCallId: string,
@@ -81,7 +80,7 @@ export class ControlPlaneInvocation implements Invocation {
     return await pollControlPlaneForOutput(this.functionCallId, timeout);
   }
 
-  async retry(): Promise<void> {
+  async retry(retryCount: number): Promise<void> {
     // we do not expect this to happen
     if (!this.input) {
       throw new Error("Cannot retry function invocation - input missing");
@@ -90,15 +89,14 @@ export class ControlPlaneInvocation implements Invocation {
     const retryItem: FunctionRetryInputsItem = {
       inputJwt: this.inputJwt!,
       input: this.input,
-      retryCount: this.retryCount,
+      retryCount: retryCount,
     };
 
     const functionRetryResponse = await client.functionRetryInputs({
       functionCallJwt: this.functionCallJwt,
       inputs: [retryItem],
     });
     this.inputJwt = functionRetryResponse.inputJwts[0];
-    this.retryCount += 1;
   }
 
   private static async execFunctionCall(