gemini-code-assist[bot] commented on code in PR #19494:
URL: https://github.com/apache/tvm/pull/19494#discussion_r3176233444


##########
web/src/opfs_store.ts:
##########
@@ -0,0 +1,257 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+interface OPFSWritableFileStream {
+  write(data: Blob | BufferSource | string): Promise<void>;
+  close(): Promise<void>;
+}

Review Comment:
   ![medium](https://www.gstatic.com/codereviewagent/medium-priority.svg)
   
   The `OPFSWritableFileStream` interface should ideally extend 
`WritableStream` to support standard streaming operations like `pipeTo`.
   
   ```suggestion
   interface OPFSWritableFileStream extends WritableStream {
     write(data: Blob | BufferSource | string): Promise<void>;
     close(): Promise<void>;
   }
   ```



##########
web/src/opfs_store.ts:
##########
@@ -0,0 +1,257 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+interface OPFSWritableFileStream {
+  write(data: Blob | BufferSource | string): Promise<void>;
+  close(): Promise<void>;
+}
+
+interface OPFSFileHandle {
+  getFile(): Promise<Blob>;
+  createWritable(): Promise<OPFSWritableFileStream>;
+}
+
+interface OPFSDirectoryHandle {
+  getDirectoryHandle(
+    name: string,
+    options?: { create?: boolean },
+  ): Promise<OPFSDirectoryHandle>;
+  getFileHandle(
+    name: string,
+    options?: { create?: boolean },
+  ): Promise<OPFSFileHandle>;
+  removeEntry(name: string): Promise<void>;
+}
+
+interface OPFSStorageManager {
+  getDirectory?: () => Promise<OPFSDirectoryHandle>;
+}
+
+interface OPFSStoreMetadata {
+  url: string;
+  contentType?: string;
+}
+
+const HASH_ALGORITHM = "SHA-256";
+const OPFS_STORE_ROOT_DIRECTORY = "tvmjs-opfs-store";
+
+export class OPFSStore {
+  private readonly scope: string;
+  private directoryPromise?: Promise<OPFSDirectoryHandle>;
+
+  constructor(scope: string) {
+    this.scope = scope;
+  }
+
+  static isAvailable(): boolean {
+    const storage = OPFSStore.getStorageManager();
+    return storage !== undefined && typeof storage.getDirectory === "function";
+  }
+
+  async has(url: string): Promise<boolean> {
+    return (await this.read(url)) !== undefined;
+  }
+
+  async read(url: string): Promise<Response | undefined> {
+    const directory = await this.getScopedDirectory();
+    const baseName = await this.hashUrl(url);
+    const dataHandle = await this.getFileHandleIfExists(
+      directory,
+      `${baseName}.bin`,
+      false,
+    );
+    if (dataHandle === undefined) {
+      return undefined;
+    }
+    const dataBlob = await dataHandle.getFile();
+    const metadataHandle = await this.getFileHandleIfExists(
+      directory,
+      `${baseName}.meta.json`,
+      false,
+    );
+    let metadata: OPFSStoreMetadata | undefined = undefined;
+    if (metadataHandle !== undefined) {
+      metadata = await this.readMetadata(metadataHandle);
+      if (metadata?.url !== undefined && metadata.url !== url) {
+        throw new Error("OPFSStore: metadata URL does not match key URL.");
+      }
+    }
+    const headers =
+      metadata?.contentType !== undefined
+        ? { "content-type": metadata.contentType }
+        : undefined;
+    return new Response(dataBlob, headers ? { headers } : undefined);
+  }
+
+  async write(url: string, response: Response): Promise<void> {
+    const directory = await this.getScopedDirectory();
+    const baseName = await this.hashUrl(url);
+    const dataHandle = await directory.getFileHandle(`${baseName}.bin`, {
+      create: true,
+    });
+    const metadataHandle = await directory.getFileHandle(
+      `${baseName}.meta.json`,
+      { create: true },
+    );
+    const buffer = await response.arrayBuffer();
+    const metadata: OPFSStoreMetadata = {
+      url,
+      contentType: response.headers.get("content-type") ?? undefined,
+    };
+    await this.writeFile(dataHandle, buffer);

Review Comment:
   ![high](https://www.gstatic.com/codereviewagent/high-priority.svg)
   
   Reading the entire response into an `ArrayBuffer` before writing to OPFS can 
lead to memory exhaustion (OOM) for large artifacts, such as LLM weights which 
can be several gigabytes. It is more efficient to stream the response body 
directly to the file system writable stream.
   
   ```suggestion
       const metadata: OPFSStoreMetadata = {
         url,
         contentType: response.headers.get("content-type") ?? undefined,
       };
       if (response.body) {
         const writable = await dataHandle.createWritable();
         await response.body.pipeTo(writable as unknown as WritableStream);
       } else {
         await this.writeFile(dataHandle, await response.arrayBuffer());
       }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to