Add training interfaces to js/common (#17333)

### Description
Following the design document:
* Added CreateTrainingSessionHandler to the Backend interface
* All existing Backend implementations throw an error for the new method
createTrainingSessionHandler
* Created TrainingSession namespace, interface, and
TrainingSessionFactory interface
* Created TrainingSessionImpl class implementation 

As methods are implemented, the TrainingSession interface will be added
to or modified.

### Motivation and Context
Adding the public-facing interfaces to the onnxruntime-common package is
one of the first steps to support ORT training for web bindings.

---------

Co-authored-by: Caroline Zhu <carolinezhu@microsoft.com>

js/common/lib/backend-impl.ts

@@ -26,7 +26,7 @@ const backendsSortedByPriority: string[] = [];
  * @ignore
  */
 export const registerBackend = (name: string, backend: Backend, priority: number): void => {
-  if (backend && typeof backend.init === 'function' && typeof backend.createSessionHandler === 'function') {
+  if (backend && typeof backend.init === 'function' && typeof backend.createInferenceSessionHandler === 'function') {
     const currentBackend = backends.get(name);
     if (currentBackend === undefined) {
       backends.set(name, {backend, priority});

js/common/lib/backend.ts

@@ -3,6 +3,7 @@
 
 import {InferenceSession} from './inference-session.js';
 import {OnnxValue} from './onnx-value.js';
+import {TrainingSession} from './training-session.js';
 
 /**
  * @ignore
@@ -14,16 +15,23 @@ export declare namespace SessionHandler {
 }
 
 /**
- * Represent a handler instance of an inference session.
+ * Represents shared SessionHandler functionality
  *
  * @ignore
  */
-export interface SessionHandler {
+interface SessionHandler {
   dispose(): Promise<void>;
 
   readonly inputNames: readonly string[];
   readonly outputNames: readonly string[];
+}
 
+/**
+ * Represent a handler instance of an inference session.
+ *
+ * @ignore
+ */
+export interface InferenceSessionHandler extends SessionHandler {
   startProfiling(): void;
   endProfiling(): void;
 
@@ -31,6 +39,20 @@ export interface SessionHandler {
       options: InferenceSession.RunOptions): Promise<SessionHandler.ReturnType>;
 }
 
+/**
+ * Represent a handler instance of a training inference session.
+ *
+ * @ignore
+ */
+export interface TrainingSessionHandler extends SessionHandler {
+  runTrainStep(
+      feeds: SessionHandler.FeedsType, fetches: SessionHandler.FetchesType,
+      options: InferenceSession.RunOptions): Promise<SessionHandler.ReturnType>;
+
+  loadParametersBuffer(array: Uint8Array, trainableOnly: boolean): Promise<void>;
+  getContiguousParameters(trainableOnly: boolean): Promise<Uint8Array>;
+}
+
 /**
  * Represent a backend that provides implementation of model inferencing.
  *
@@ -42,8 +64,13 @@ export interface Backend {
    */
   init(): Promise<void>;
 
-  createSessionHandler(uriOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
-      Promise<SessionHandler>;
+  createInferenceSessionHandler(uriOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler>;
+
+  createTrainingSessionHandler?
+      (checkpointStateUriOrBuffer: TrainingSession.URIorBuffer, trainModelUriOrBuffer: TrainingSession.URIorBuffer,
+       evalModelUriOrBuffer: TrainingSession.URIorBuffer, optimizerModelUriOrBuffer: TrainingSession.URIorBuffer,
+       options: InferenceSession.SessionOptions): Promise<TrainingSessionHandler>;
 }
 
 export {registerBackend} from './backend-impl.js';

js/common/lib/env.ts

@@ -9,6 +9,7 @@ export declare namespace Env {
     'ort-wasm.wasm'?: string;
     'ort-wasm-threaded.wasm'?: string;
     'ort-wasm-simd.wasm'?: string;
+    'ort-training-wasm-simd.wasm'?: string;
     'ort-wasm-simd-threaded.wasm'?: string;
     /* eslint-enable @typescript-eslint/naming-convention */
   };

js/common/lib/index.ts

@@ -22,3 +22,4 @@ export * from './env.js';
 export * from './inference-session.js';
 export * from './tensor.js';
 export * from './onnx-value.js';
+export * from './training-session.js';

js/common/lib/inference-session-impl.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import {resolveBackend} from './backend-impl.js';
-import {SessionHandler} from './backend.js';
+import {InferenceSessionHandler} from './backend.js';
 import {InferenceSession as InferenceSessionInterface} from './inference-session.js';
 import {OnnxValue} from './onnx-value.js';
 import {Tensor} from './tensor.js';
@@ -14,7 +14,7 @@ type FetchesType = InferenceSessionInterface.FetchesType;
 type ReturnType = InferenceSessionInterface.ReturnType;
 
 export class InferenceSession implements InferenceSessionInterface {
-  private constructor(handler: SessionHandler) {
+  private constructor(handler: InferenceSessionHandler) {
     this.handler = handler;
   }
   run(feeds: FeedsType, options?: RunOptions): Promise<ReturnType>;
@@ -195,7 +195,7 @@ export class InferenceSession implements InferenceSessionInterface {
     const eps = options.executionProviders || [];
     const backendHints = eps.map(i => typeof i === 'string' ? i : i.name);
     const backend = await resolveBackend(backendHints);
-    const handler = await backend.createSessionHandler(filePathOrUint8Array, options);
+    const handler = await backend.createInferenceSessionHandler(filePathOrUint8Array, options);
     return new InferenceSession(handler);
   }
 
@@ -213,5 +213,5 @@ export class InferenceSession implements InferenceSessionInterface {
     return this.handler.outputNames;
   }
 
-  private handler: SessionHandler;
+  private handler: InferenceSessionHandler;
 }

js/common/lib/training-session-impl.ts

@@ -0,0 +1,49 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import {TrainingSessionHandler} from './backend.js';
+import {InferenceSession as InferenceSession} from './inference-session.js';
+import {TrainingSession as TrainingSessionInterface, TrainingSessionCreateOptions} from './training-session.js';
+
+type SessionOptions = InferenceSession.SessionOptions;
+
+export class TrainingSession implements TrainingSessionInterface {
+  private constructor(handler: TrainingSessionHandler) {
+    this.handler = handler;
+  }
+  private handler: TrainingSessionHandler;
+
+  get inputNames(): readonly string[] {
+    return this.handler.inputNames;
+  }
+  get outputNames(): readonly string[] {
+    return this.handler.outputNames;
+  }
+
+  static async create(_trainingOptions: TrainingSessionCreateOptions, _sessionOptions?: SessionOptions):
+      Promise<TrainingSession> {
+    throw new Error('Method not implemented');
+  }
+
+  async loadParametersBuffer(_array: Uint8Array, _trainableOnly: boolean): Promise<void> {
+    throw new Error('Method not implemented.');
+  }
+
+  async getContiguousParameters(_trainableOnly: boolean): Promise<Uint8Array> {
+    throw new Error('Method not implemented.');
+  }
+
+  runTrainStep(feeds: InferenceSession.OnnxValueMapType, options?: InferenceSession.RunOptions|undefined):
+      Promise<InferenceSession.OnnxValueMapType>;
+  runTrainStep(
+      feeds: InferenceSession.OnnxValueMapType, fetches: InferenceSession.FetchesType,
+      options?: InferenceSession.RunOptions|undefined): Promise<InferenceSession.OnnxValueMapType>;
+  async runTrainStep(_feeds: unknown, _fetches?: unknown, _options?: unknown):
+      Promise<InferenceSession.OnnxValueMapType> {
+    throw new Error('Method not implemented.');
+  }
+
+  async release(): Promise<void> {
+    return this.handler.dispose();
+  }
+}

js/common/lib/training-session.ts

@@ -0,0 +1,134 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+import {InferenceSession} from './inference-session.js';
+import {TrainingSession as TrainingSessionImpl} from './training-session-impl.js';
+
+/* eslint-disable @typescript-eslint/no-redeclare */
+
+export declare namespace TrainingSession {
+  /**
+   * Either URI file path (string) or Uint8Array containing model or checkpoint information.
+   */
+  type URIorBuffer = string|Uint8Array;
+}
+
+/**
+ * Represent a runtime instance of an ONNX training session,
+ * which contains a model that can be trained, and, optionally,
+ * an eval and optimizer model.
+ */
+export interface TrainingSession {
+  // #region run()
+
+  /**
+   * Run TrainStep asynchronously with the given feeds and options.
+   *
+   * @param feeds - Representation of the model input. See type description of `InferenceSession.InputType` for
+   detail.
+   * @param options - Optional. A set of options that controls the behavior of model training.
+   * @returns A promise that resolves to a map, which uses output names as keys and OnnxValue as corresponding values.
+   */
+  runTrainStep(feeds: InferenceSession.FeedsType, options?: InferenceSession.RunOptions):
+      Promise<InferenceSession.ReturnType>;
+
+  /**
+   * Run a single train step with the given inputs and options.
+   *
+   * @param feeds - Representation of the model input.
+   * @param fetches - Representation of the model output.
+   * detail.
+   * @param options - Optional. A set of options that controls the behavior of model inference.
+   * @returns A promise that resolves to a map, which uses output names as keys and OnnxValue as corresponding
+   values.
+   */
+  runTrainStep(
+      feeds: InferenceSession.FeedsType, fetches: InferenceSession.FetchesType,
+      options?: InferenceSession.RunOptions): Promise<InferenceSession.ReturnType>;
+
+  // #endregion
+
+  // #region copy parameters
+  /**
+   * Copies from a buffer containing parameters to the TrainingSession parameters.
+   *
+   * @param buffer - buffer containing parameters
+   * @param trainableOnly - True if trainable parameters only to be modified, false otherwise.
+   */
+  loadParametersBuffer(array: Uint8Array, trainableOnly: boolean): Promise<void>;
+
+  /**
+   * Copies from the TrainingSession parameters to a buffer.
+   *
+   * @param trainableOnly - True if trainable parameters only to be copied, false othrwise.
+   * @returns A promise that resolves to a buffer of the requested parameters.
+   */
+  getContiguousParameters(trainableOnly: boolean): Promise<Uint8Array>;
+  // #endregion
+
+  // #region release()
+
+  /**
+   * Release the inference session and the underlying resources.
+   */
+  release(): Promise<void>;
+  // #endregion
+
+  // #region metadata
+
+  /**
+   * Get input names of the loaded model.
+   */
+  readonly inputNames: readonly string[];
+
+  /**
+   * Get output names of the loaded model.
+   */
+  readonly outputNames: readonly string[];
+  // #endregion
+}
+
+/**
+ * Represents the optional parameters that can be passed into the TrainingSessionFactory.
+ */
+export interface TrainingSessionCreateOptions {
+  /**
+   * URI or buffer for a .ckpt file that contains the checkpoint for the training model.
+   */
+  checkpointState: TrainingSession.URIorBuffer;
+  /**
+   * URI or buffer for the .onnx training file.
+   */
+  trainModel: TrainingSession.URIorBuffer;
+  /**
+   * Optional. URI or buffer for the .onnx optimizer model file.
+   */
+  optimizerModel?: TrainingSession.URIorBuffer;
+  /**
+   * Optional. URI or buffer for the .onnx eval model file.
+   */
+  evalModel?: TrainingSession.URIorBuffer;
+}
+
+/**
+ * Defines method overload possibilities for creating a TrainingSession.
+ */
+export interface TrainingSessionFactory {
+  // #region create()
+
+  /**
+   * Creates a new TrainingSession and asynchronously loads any models passed in through trainingOptions
+   *
+   * @param trainingOptions specify models and checkpoints to load into the Training Session
+   * @param sessionOptions specify configuration for training session behavior
+   *
+   * @returns Promise that resolves to a TrainingSession object
+   */
+  create(trainingOptions: TrainingSessionCreateOptions, sessionOptions?: InferenceSession.SessionOptions):
+      Promise<TrainingSession>;
+
+  // #endregion
+}
+
+// eslint-disable-next-line @typescript-eslint/naming-convention
+export const TrainingSession: TrainingSessionFactory = TrainingSessionImpl;

js/node/lib/backend.ts

@@ -1,11 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import {Backend, InferenceSession, SessionHandler} from 'onnxruntime-common';
+import {Backend, InferenceSession, InferenceSessionHandler, SessionHandler} from 'onnxruntime-common';
 
 import {Binding, binding} from './binding';
 
-class OnnxruntimeSessionHandler implements SessionHandler {
+class OnnxruntimeSessionHandler implements InferenceSessionHandler {
   #inferenceSession: Binding.InferenceSession;
 
   constructor(pathOrBuffer: string|Uint8Array, options: InferenceSession.SessionOptions) {
@@ -53,8 +53,8 @@ class OnnxruntimeBackend implements Backend {
     return Promise.resolve();
   }
 
-  async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
-      Promise<SessionHandler> {
+  async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler> {
     return new Promise((resolve, reject) => {
       process.nextTick(() => {
         try {

js/react_native/lib/backend.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import {Backend, InferenceSession, SessionHandler, Tensor,} from 'onnxruntime-common';
+import {Backend, InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common';
 import {Platform} from 'react-native';
 
 import {binding, Binding, JSIBlob, jsiHelper} from './binding';
@@ -43,7 +43,7 @@ const normalizePath = (path: string): string => {
   return path;
 };
 
-class OnnxruntimeSessionHandler implements SessionHandler {
+class OnnxruntimeSessionHandler implements InferenceSessionHandler {
   #inferenceSession: Binding.InferenceSession;
   #key: string;
 
@@ -165,8 +165,8 @@ class OnnxruntimeBackend implements Backend {
     return Promise.resolve();
   }
 
-  async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
-      Promise<SessionHandler> {
+  async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler> {
     const handler = new OnnxruntimeSessionHandler(pathOrBuffer);
     await handler.loadModel(options || {});
     return handler;

js/web/lib/backend-onnxjs.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 /* eslint-disable import/no-internal-modules */
-import {Backend, InferenceSession, SessionHandler} from 'onnxruntime-common';
+import {Backend, InferenceSession, InferenceSessionHandler} from 'onnxruntime-common';
 
 import {Session} from './onnxjs/session';
 import {OnnxjsSessionHandler} from './onnxjs/session-handler';
@@ -11,8 +11,8 @@ class OnnxjsBackend implements Backend {
   // eslint-disable-next-line @typescript-eslint/no-empty-function
   async init(): Promise<void> {}
 
-  async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
-      Promise<SessionHandler> {
+  async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler> {
     // NOTE: Session.Config(from onnx.js) is not compatible with InferenceSession.SessionOptions(from
     // onnxruntime-common).
     //       In future we should remove Session.Config and use InferenceSession.SessionOptions.

js/web/lib/backend-wasm.ts

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import {Backend, env, InferenceSession, SessionHandler} from 'onnxruntime-common';
+import {Backend, env, InferenceSession, InferenceSessionHandler} from 'onnxruntime-common';
 import {cpus} from 'os';
 
 import {initializeWebAssemblyInstance} from './wasm/proxy-wrapper';
@@ -40,10 +40,12 @@ class OnnxruntimeWebAssemblyBackend implements Backend {
     // init wasm
     await initializeWebAssemblyInstance();
   }
-  createSessionHandler(path: string, options?: InferenceSession.SessionOptions): Promise<SessionHandler>;
-  createSessionHandler(buffer: Uint8Array, options?: InferenceSession.SessionOptions): Promise<SessionHandler>;
-  async createSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
-      Promise<SessionHandler> {
+  createInferenceSessionHandler(path: string, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler>;
+  createInferenceSessionHandler(buffer: Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler>;
+  async createInferenceSessionHandler(pathOrBuffer: string|Uint8Array, options?: InferenceSession.SessionOptions):
+      Promise<InferenceSessionHandler> {
     const handler = new OnnxruntimeWebAssemblySessionHandler();
     await handler.loadModel(pathOrBuffer, options);
     return Promise.resolve(handler);

js/web/lib/onnxjs/session-handler.ts

@@ -1,12 +1,12 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-import {InferenceSession, SessionHandler, Tensor} from 'onnxruntime-common';
+import {InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common';
 
 import {Session} from './session';
 import {Tensor as OnnxjsTensor} from './tensor';
 
-export class OnnxjsSessionHandler implements SessionHandler {
+export class OnnxjsSessionHandler implements InferenceSessionHandler {
   constructor(private session: Session) {
     this.inputNames = this.session.inputNames;
     this.outputNames = this.session.outputNames;

js/web/lib/wasm/session-handler.ts

@@ -2,7 +2,7 @@
 // Licensed under the MIT License.
 
 import {readFile} from 'fs';
-import {env, InferenceSession, SessionHandler, Tensor} from 'onnxruntime-common';
+import {env, InferenceSession, InferenceSessionHandler, SessionHandler, Tensor} from 'onnxruntime-common';
 import {promisify} from 'util';
 
 import {SerializableModeldata, TensorMetadata} from './proxy-messages';
@@ -40,7 +40,7 @@ const decodeTensorMetadata = (tensor: TensorMetadata): Tensor => {
   }
 };
 
-export class OnnxruntimeWebAssemblySessionHandler implements SessionHandler {
+export class OnnxruntimeWebAssemblySessionHandler implements InferenceSessionHandler {
   private sessionId: number;
 
   inputNames: string[];