This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 2c9af0f500 [Runtime] Allow aborting fetchNDArray through AbortSignal
(#17208)
2c9af0f500 is described below
commit 2c9af0f500c04383aa7220ab2c9220a608f75cbf
Author: Nestor Qin <[email protected]>
AuthorDate: Mon Jul 29 08:17:55 2024 -0400
[Runtime] Allow aborting fetchNDArray through AbortSignal (#17208)
[Runtime] Allow aborting fetchNDArray
---
web/src/artifact_cache.ts | 11 ++++++-----
web/src/runtime.ts | 13 +++++++++----
2 files changed, 15 insertions(+), 9 deletions(-)
diff --git a/web/src/artifact_cache.ts b/web/src/artifact_cache.ts
index f833df1be5..9690ed3320 100644
--- a/web/src/artifact_cache.ts
+++ b/web/src/artifact_cache.ts
@@ -58,13 +58,14 @@ export interface ArtifactCacheTemplate {
*
* @param url: The url to the data to be cached.
* @param storetype: Only applies to `ArtifactIndexedDBCache`. Since
`indexedDB` stores the actual
+ * @param signal: An optional AbortSignal to abort data retrival
* data rather than a request, we specify `storagetype`. There are two
options:
* 1. "json": IndexedDB stores `fetch(url).json()`
* 2. "arraybuffer": IndexedDB stores `fetch(url).arrayBuffer()`
*
* @note This is an async function.
*/
- addToCache(url: string, storetype?: string): Promise<void>;
+ addToCache(url: string, storetype?: string, signal?: AbortSignal):
Promise<void>;
/**
* check if cache has all keys in Cache
@@ -126,8 +127,8 @@ export class ArtifactCache implements ArtifactCacheTemplate
{
}
// eslint-disable-next-line @typescript-eslint/no-unused-vars
- async addToCache(url: string, storetype?: string) {
- const request = new Request(url);
+ async addToCache(url: string, storetype?: string, signal?: AbortSignal) {
+ const request = new Request(url, signal ? { signal } : undefined);
if (this.cache === undefined) {
this.cache = await caches.open(this.scope);
}
@@ -282,7 +283,7 @@ export class ArtifactIndexedDBCache implements
ArtifactCacheTemplate {
});
}
- async addToCache(url: string, storetype?: string): Promise<void> {
+ async addToCache(url: string, storetype?: string, signal?: AbortSignal):
Promise<void> {
await this.initDB(); // await the initDB process
// If already cached, nothing to do
const isInDB = await this.isUrlInDB(url);
@@ -290,7 +291,7 @@ export class ArtifactIndexedDBCache implements
ArtifactCacheTemplate {
return;
}
try {
- const response = await fetch(url);
+ const response = await fetch(url, signal ? { signal } : undefined);
if (!response.ok) {
throw new Error('Network response was not ok');
}
diff --git a/web/src/runtime.ts b/web/src/runtime.ts
index fd7bcc6ab2..d71c98e7d1 100644
--- a/web/src/runtime.ts
+++ b/web/src/runtime.ts
@@ -1444,13 +1444,15 @@ export class Instance implements Disposable {
* @param device The device to be fetched to.
* @param cacheScope The scope identifier of the cache
* @param cacheType The type of the cache: "cache" or "indexedDB"
+ * @param signal An optional AbortSignal to abort the fetch
* @returns The meta data
*/
async fetchNDArrayCache(
ndarrayCacheUrl: string,
device: DLDevice,
cacheScope = "tvmjs",
- cacheType = "cache"
+ cacheType = "cache",
+ signal?: AbortSignal,
): Promise<any> {
let artifactCache: ArtifactCacheTemplate;
if (cacheType === undefined || cacheType.toLowerCase() === "cache") {
@@ -1465,7 +1467,8 @@ export class Instance implements Disposable {
const list = await artifactCache.fetchWithCache(jsonUrl, "json");
await this.fetchNDArrayCacheInternal(
ndarrayCacheUrl,
- list["records"] as Array<NDArrayShardEntry>, device, artifactCache);
+ list["records"] as Array<NDArrayShardEntry>, device, artifactCache,
+ signal);
this.cacheMetadata = { ...this.cacheMetadata, ...(list["metadata"] as
Record<string, any>) };
}
@@ -1477,12 +1480,14 @@ export class Instance implements Disposable {
* @param list The list of array data.
* @param device The device to store the data to.
* @param artifactCache The artifact cache
+ * @param signal An optional AbortSignal to abort the fetch
*/
private async fetchNDArrayCacheInternal(
ndarrayCacheUrl: string,
list: Array<NDArrayShardEntry>,
device: DLDevice,
- artifactCache: ArtifactCacheTemplate
+ artifactCache: ArtifactCacheTemplate,
+ signal?: AbortSignal,
) {
const perf = compact.getPerformance();
const tstart = perf.now();
@@ -1537,7 +1542,7 @@ export class Instance implements Disposable {
const shard = list[i];
const dataUrl = new URL(shard.dataPath, ndarrayCacheUrl).href;
try {
- await artifactCache.addToCache(dataUrl, "arraybuffer");
+ await artifactCache.addToCache(dataUrl, "arraybuffer", signal);
} catch (err) {
this.env.logger("Error: Cannot fetch " + dataUrl + " err= " + err);
throw err;
