Add support for input plane invocation

rculbertson · rculbertson · commit a9984937d66e · 2025-06-18T14:26:02.000-04:00
diff --git a/modal-js/src/client.ts b/modal-js/src/client.ts
@@ -13,6 +13,8 @@ import {
 import { ClientType, ModalClientDefinition } from "../proto/modal_proto/api";
 import { type Profile, profile } from "./config";
 
+let modalAuthToken: string | undefined;
+
 /** gRPC client middleware to add auth token to request. */
 function authMiddleware(profile: Profile): ClientMiddleware {
   return async function* authMiddleware<Request, Response>(
@@ -27,6 +29,30 @@ function authMiddleware(profile: Profile): ClientMiddleware {
     options.metadata.set("x-modal-client-version", "1.0.0"); // CLIENT VERSION: Behaves like this Python SDK version
     options.metadata.set("x-modal-token-id", profile.tokenId);
     options.metadata.set("x-modal-token-secret", profile.tokenSecret);
+    if (modalAuthToken) {
+      options.metadata.set("x-modal-auth-token", modalAuthToken);
+    }
+
+    const prevOnHeader = options.onHeader;
+    options.onHeader = (header) => {
+      const token = header.get("x-modal-auth-token");
+      if (token) {
+        modalAuthToken = token;
+      }
+      if (prevOnHeader) {
+        prevOnHeader(header);
+      }
+    };
+    const prevOnTrailer = options.onTrailer;
+    options.onTrailer = (trailer) => {
+      const token = trailer.get("x-modal-auth-token");
+      if (token) {
+        modalAuthToken = token;
+      }
+      if (prevOnTrailer) {
+        prevOnTrailer(trailer);
+      }
+    };
     return yield* call.next(call.request, options);
   };
 }
@@ -190,15 +216,37 @@ const retryMiddleware: ClientMiddleware<RetryOptions> =
     }
   };
 
-// Ref: https://github.com/modal-labs/modal-client/blob/main/modal/_utils/grpc_utils.py
-const channel = createChannel(profile.serverUrl, undefined, {
-  "grpc.max_receive_message_length": 100 * 1024 * 1024,
-  "grpc.max_send_message_length": 100 * 1024 * 1024,
-  "grpc-node.flow_control_window": 64 * 1024 * 1024,
-});
-
-export const client = createClientFactory()
-  .use(authMiddleware(profile))
-  .use(retryMiddleware)
-  .use(timeoutMiddleware)
-  .create(ModalClientDefinition, channel);
+/**
+ * Map of server URL => client.
+ * The us-east client talks to the control plane; all other clients talk to input planes.
+ */
+const clients: Record<string, ReturnType<typeof createClient>> = {};
+
+/** Returns a client for the given server URL, creating it if it doesn't exist. */
+export const getOrCreateClient = (
+  serverURL: string,
+): ReturnType<typeof createClient> => {
+  if (serverURL in clients) {
+    return clients[serverURL];
+  }
+
+  clients[serverURL] = createClient(serverURL);
+  return clients[serverURL];
+};
+
+const createClient = (serverURL: string) => {
+  // Ref: https://github.com/modal-labs/modal-client/blob/main/modal/_utils/grpc_utils.py
+  const channel = createChannel(serverURL, undefined, {
+    "grpc.max_receive_message_length": 100 * 1024 * 1024,
+    "grpc.max_send_message_length": 100 * 1024 * 1024,
+    "grpc-node.flow_control_window": 64 * 1024 * 1024,
+  });
+
+  return createClientFactory()
+    .use(authMiddleware(profile))
+    .use(retryMiddleware)
+    .use(timeoutMiddleware)
+    .create(ModalClientDefinition, channel);
+};
+
+export const client = getOrCreateClient(profile.serverUrl);
diff --git a/modal-js/src/function.ts b/modal-js/src/function.ts
@@ -15,7 +15,11 @@ import { environmentName } from "./config";
 import { InternalFailure, NotFoundError } from "./errors";
 import { dumps } from "./pickle";
 import { ClientError, Status } from "nice-grpc";
-import { ControlPlaneInvocation } from "./invocation";
+import {
+  ControlPlaneInvocation,
+  InputPlaneInvocation,
+  Invocation,
+} from "./invocation";
 
 // From: modal/_utils/blob_utils.py
 const maxObjectSizeBytes = 2 * 1024 * 1024; // 2 MiB
@@ -27,11 +31,13 @@ const maxSystemRetries = 8;
 export class Function_ {
   readonly functionId: string;
   readonly methodName: string | undefined;
+  readonly inputPlaneUrl: string | undefined;
 
   /** @ignore */
-  constructor(functionId: string, methodName?: string) {
+  constructor(functionId: string, methodName?: string, inputPlaneUrl?: string) {
     this.functionId = functionId;
     this.methodName = methodName;
+    this.inputPlaneUrl = inputPlaneUrl;
   }
 
   static async lookup(
@@ -46,7 +52,11 @@ export class Function_ {
         namespace: DeploymentNamespace.DEPLOYMENT_NAMESPACE_WORKSPACE,
         environmentName: environmentName(options.environment),
       });
-      return new Function_(resp.functionId);
+      return new Function_(
+        resp.functionId,
+        undefined,
+        resp.handleMetadata?.inputPlaneUrl,
+      );
     } catch (err) {
       if (err instanceof ClientError && err.code === Status.NOT_FOUND)
         throw new NotFoundError(`Function '${appName}/${name}' not found`);
@@ -60,11 +70,7 @@ export class Function_ {
     kwargs: Record<string, any> = {},
   ): Promise<any> {
     const input = await this.#createInput(args, kwargs);
-    const invocation = await ControlPlaneInvocation.create(
-      this.functionId,
-      input,
-      FunctionCallInvocationType.FUNCTION_CALL_INVOCATION_TYPE_SYNC,
-    );
+    const invocation = await this.#createRemoteInvocation(input);
     // TODO(ryan): Add tests for retries.
     let retryCount = 0;
     while (true) {
@@ -81,6 +87,22 @@ export class Function_ {
     }
   }
 
+  async #createRemoteInvocation(input: FunctionInput): Promise<Invocation> {
+    if (this.inputPlaneUrl) {
+      return await InputPlaneInvocation.create(
+        this.inputPlaneUrl,
+        this.functionId,
+        input,
+      );
+    }
+
+    return await ControlPlaneInvocation.create(
+      this.functionId,
+      input,
+      FunctionCallInvocationType.FUNCTION_CALL_INVOCATION_TYPE_SYNC,
+    );
+  }
+
   // Spawn a single input into a remote function.
   async spawn(
     args: any[] = [],
diff --git a/modal-js/src/invocation.ts b/modal-js/src/invocation.ts
@@ -2,15 +2,17 @@ import {
   DataFormat,
   FunctionCallInvocationType,
   FunctionCallType,
-  FunctionGetOutputsResponse,
+  FunctionGetOutputsItem,
   FunctionInput,
   FunctionMapResponse,
+  FunctionPutInputsItem,
   FunctionRetryInputsItem,
   GeneratorDone,
   GenericResult,
   GenericResult_GenericStatus,
+  ModalClientClient,
 } from "../proto/modal_proto/api";
-import { client } from "./client";
+import { client, getOrCreateClient } from "./client";
 import { FunctionTimeoutError, InternalFailure, RemoteError } from "./errors";
 import { loads } from "./pickle";
 
@@ -33,6 +35,14 @@ export interface Invocation {
   retry(retryCount: number): Promise<void>;
 }
 
+/**
+ * Signature of a function that fetches a single output. Used by `pollForOutputs` to fetch from either
+ * the control plane or the input plane, depending on the implementation.
+ */
+type GetOutput = (
+  timeoutMillis: number,
+) => Promise<FunctionGetOutputsItem | undefined>;
+
 /**
  * Implementation of Invocation which sends inputs to the control plane.
  */
@@ -77,7 +87,28 @@ export class ControlPlaneInvocation implements Invocation {
   }
 
   async await(timeout?: number): Promise<any> {
-    return await pollControlPlaneForOutput(this.functionCallId, timeout);
+    return await pollControlPlaneForOutput(
+      (timeoutMillis: number) => this.#getOutput(timeoutMillis),
+      timeout,
+    );
+  }
+
+  async #getOutput(
+    timeoutMillis: number,
+  ): Promise<FunctionGetOutputsItem | undefined> {
+    try {
+      const response = await client.functionGetOutputs({
+        functionCallId: this.functionCallId,
+        maxValues: 1,
+        timeout: timeoutMillis / 1000, // Backend needs seconds
+        lastEntryId: "0-0",
+        clearOnSuccess: true,
+        requestedAt: timeNowSeconds(),
+      });
+      return response.outputs ? response.outputs[0] : undefined;
+    } catch (err) {
+      throw new Error(`FunctionGetOutputs failed: ${err}`);
+    }
   }
 
   async retry(retryCount: number): Promise<void> {
@@ -118,12 +149,88 @@ export class ControlPlaneInvocation implements Invocation {
   }
 }
 
+/**
+ * Implementation of Invocation which sends inputs to the input plane.
+ */
+export class InputPlaneInvocation implements Invocation {
+  private readonly client: ModalClientClient;
+  private readonly functionId: string;
+  private readonly input: FunctionPutInputsItem;
+  private attemptToken: string;
+
+  constructor(
+    client: ModalClientClient,
+    functionId: string,
+    input: FunctionPutInputsItem,
+    attemptToken: string,
+  ) {
+    this.client = client;
+    this.functionId = functionId;
+    this.input = input;
+    this.attemptToken = attemptToken;
+  }
+
+  static async create(
+    inputPlaneUrl: string,
+    functionId: string,
+    input: FunctionInput,
+  ) {
+    const functionPutInputsItem = {
+      idx: 0,
+      input: input,
+    };
+    const client = getOrCreateClient(inputPlaneUrl);
+    // Single input sync invocation
+    const attemptStartResponse = await client.attemptStart({
+      functionId: functionId,
+      input: functionPutInputsItem,
+    });
+    return new InputPlaneInvocation(
+      client,
+      functionId,
+      functionPutInputsItem,
+      attemptStartResponse.attemptToken,
+    );
+  }
+
+  async await(timeout?: number): Promise<any> {
+    return await pollControlPlaneForOutput(
+      (timeoutMillis: number) => this.#getOutput(timeoutMillis),
+      timeout,
+    );
+  }
+
+  async #getOutput(
+    timeoutMillis: number,
+  ): Promise<FunctionGetOutputsItem | undefined> {
+    try {
+      const response = await this.client.attemptAwait({
+        attemptToken: this.attemptToken,
+        requestedAt: timeNowSeconds(),
+        timeoutSecs: timeoutMillis / 1000,
+      });
+      return response.output;
+    } catch (err) {
+      throw new Error(`AttemptAwait failed: ${err}`);
+    }
+  }
+
+  async retry(_retryCount: number): Promise<void> {
+    const attemptRetryResponse = await this.client.attemptRetry({
+      functionId: this.functionId,
+      input: this.input,
+      attemptToken: this.attemptToken,
+    });
+    this.attemptToken = attemptRetryResponse.attemptToken;
+  }
+}
+
 function timeNowSeconds() {
   return Date.now() / 1e3;
 }
 
-export async function pollControlPlaneForOutput(
-  functionCallId: string,
+async function pollControlPlaneForOutput(
+  getOutput: GetOutput,
   timeout?: number, // in milliseconds
 ): Promise<any> {
   const startTime = Date.now();
@@ -133,23 +240,9 @@ export async function pollControlPlaneForOutput(
   }
 
   while (true) {
-    let response: FunctionGetOutputsResponse;
-    try {
-      response = await client.functionGetOutputs({
-        functionCallId: functionCallId,
-        maxValues: 1,
-        timeout: pollTimeout / 1000, // Backend needs seconds
-        lastEntryId: "0-0",
-        clearOnSuccess: true,
-        requestedAt: timeNowSeconds(),
-      });
-    } catch (err) {
-      throw new Error(`FunctionGetOutputs failed: ${err}`);
-    }
-
-    const outputs = response.outputs;
-    if (outputs.length > 0) {
-      return await processResult(outputs[0].result, outputs[0].dataFormat);
+    const output = await getOutput(pollTimeout);
+    if (output) {
+      return await processResult(output.result, output.dataFormat);
     }
 
     if (timeout !== undefined) {
diff --git a/modal-js/test/function.test.ts b/modal-js/test/function.test.ts
@@ -34,3 +34,12 @@ test("FunctionNotFound", async () => {
   );
   await expect(promise).rejects.toThrowError(NotFoundError);
 });
+
+test("FunctionCallInputPlane", async () => {
+  const function_ = await Function_.lookup(
+    "libmodal-test-support",
+    "input_plane",
+  );
+  const result = await function_.remote(["hello"]);
+  expect(result).toBe("output: hello");
+});
diff --git a/test-support/libmodal_test_support.py b/test-support/libmodal_test_support.py
@@ -19,6 +19,9 @@ def sleep(t: int) -> None:
 def bytelength(buf: bytes) -> int:
     return len(buf)
 
+@app.function(min_containers=1, experimental_options={"input_plane_region": "us-west"})
+def input_plane(s: str) -> str:
+    return "output: " + s
 
 @app.cls(min_containers=1)
 class EchoCls:
