[js/webgpu] Manage model download with a specific unittest option (#22214)

Currently in debug mode, unit test will always download models to local
file system, which is a bit annoying. This PR fixes this by adding a
specific option to enable model download.
This commit is contained in:
Yang Gu 2024-10-01 09:27:43 +08:00 committed by GitHub
parent c75f4a09b7
commit 9e5153b688
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 19 additions and 3 deletions

View file

@ -48,6 +48,7 @@ Options:
node
bs (for BrowserStack tests)
-p, --profile Enable profiler.
-m, --download-model Enable model download.
Profiler will generate extra logs which include the information of events time consumption
-t, --trace Enable trace.
-P[=<...>], --perf[=<...>] Generate performance number. Cannot be used with flag --debug.
@ -171,6 +172,11 @@ export interface TestRunnerCliArgs {
*/
profile: boolean;
/**
* whether to enable model download
*/
downloadModel: boolean;
/**
* Whether to enable file cache
*/
@ -431,6 +437,9 @@ export function parseTestRunnerCliArgs(cmdlineArgs: string[]): TestRunnerCliArgs
logLevel = 'verbose';
}
// Option: -m, --download-model
const downloadModel = args['download-model'] || args.m ? true : false;
// Option: -t, --trace
const trace = parseBooleanArg(args.trace || args.t, false);
@ -517,6 +526,7 @@ export function parseTestRunnerCliArgs(cmdlineArgs: string[]): TestRunnerCliArgs
env: env as TestRunnerCliArgs['env'],
logConfig,
profile,
downloadModel,
times: perf ? times : undefined,
ioBindingMode: ioBindingMode as TestRunnerCliArgs['ioBindingMode'],
optimizedModelFilePath,

View file

@ -161,6 +161,7 @@ async function main() {
op: opTestGroups,
log: args.logConfig,
profile: args.profile,
downloadModel: args.downloadModel,
options: {
sessionOptions: {
graphOptimizationLevel: args.graphOptimizationLevel,

View file

@ -99,7 +99,11 @@ for (const group of ORT_WEB_TEST_CONFIG.op) {
before('Initialize Context', async () => {
context = useProtoOpTest
? new ProtoOpTestContext(test, ORT_WEB_TEST_CONFIG.options.sessionOptions)
? new ProtoOpTestContext(
test,
ORT_WEB_TEST_CONFIG.downloadModel,
ORT_WEB_TEST_CONFIG.options.sessionOptions,
)
: new OpTestContext(test);
await context.init();
if (ORT_WEB_TEST_CONFIG.profile) {

View file

@ -888,6 +888,7 @@ export class ProtoOpTestContext {
readonly ioBindingMode: Test.IOBindingMode;
constructor(
test: Test.OperatorTest,
private readonly downloadModel: boolean,
private readonly sessionOptions: ort.InferenceSession.SessionOptions = {},
) {
const opsetImport = onnx.OperatorSetIdProto.create(test.opset);
@ -1074,8 +1075,7 @@ export class ProtoOpTestContext {
this.ioBindingMode = test.ioBinding;
this.loadedData = onnx.ModelProto.encode(model).finish().slice();
// in debug mode, open a new tab in browser for the generated onnx model.
if (ort.env.debug) {
if (this.downloadModel) {
const modelFile = new File([this.loadedData], `op_test_generated_model_${test.name}.onnx`, {
type: 'application/octet-stream',
});

View file

@ -172,6 +172,7 @@ export declare namespace Test {
log: ReadonlyArray<{ category: string; config: Logger.Config }>;
profile: boolean;
downloadModel: boolean;
options: Options;
}
}