chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

MuJoCo Web Viewer: Implement parallel chunked model downloading and in-place model reloading.

Browse Source 
This change introduces parallel chunked downloading for large model files (.mjb) directly into WASM linear memory, enables model loading without full page reloads and reduces memory overhead.

Key changes:
- Implement a chunked model endpoint in the Python web server to support range requests.
- Add parallel chunked fetching in the frontend with retry logic and a single-fetch fallback.
- Increase initial WASM memory to 3 GB to accommodate large models and prevent heap fragmentation.
- Support in-place model reloading in the C++ client, including texture cache invalidation when the Filament context is recreated.
- Display a model download progress bar and model parsing/loading banner to the UI.
- Fixes model drag and drop (caused by typo in sessionId, corrected to session_id).

PiperOrigin-RevId: 960249236
Change-Id: Icac89e6a4ca099882aaf9b111744c1b6ab0cc6c1
This commit is contained in:
Matija Kecman
2026-08-06 05:51:01 -07:00
committed by Copybara-Service
parent 7fd2061f5c
commit 84950fa371
11 changed files with 406 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/experimental/studio/web/index.html
+169 -9
View File
@@ -45,6 +45,160 @@ limitations under the License.
const WS_CLOSE_INACTIVE = 4003;
const WS_CLOSE_NOT_CONTROLLER = 4004;
// Parallel chunked model download.
//
// Fetches /model in up to PARALLEL concurrent 64 MiB requests
// and streams the chunks directly into a WASM linear heap buffer
// allocated via allocModelBuffer(totalSize). Each chunk is retried up
// to MAX_RETRIES times on transient errors (proxy resets, HTTP2
// protocol errors, etc.). Returns { ptr, size } on success, or null
// on failure.
async function fetchModelChunked() {
const CHUNK = 64 * 1024 * 1024; // 64 MiB per request
const PARALLEL = 6; // concurrent fetches
const MAX_RETRIES = 3; // per-chunk retry limit
const RETRY_DELAY = 500; // ms between retries
let ptr = 0;
try {
const totalSizeResp = await fetch("/model?total_bytes");
if (!totalSizeResp.ok) return null;
const totalSizeData = await totalSizeResp.json();
const totalSize = totalSizeData.total_bytes;
if (!totalSize || totalSize <= 0) return null;
ptr = Module.allocModelBuffer(totalSize);
if (!ptr) {
console.error("[model] failed to allocate WASM buffer of size:", totalSize);
return null;
}
const chunks = [];
for (let offset = 0; offset < totalSize; offset += CHUNK) {
chunks.push({ offset: offset, size: Math.min(CHUNK, totalSize - offset) });
}
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(0, totalSize, 0);
}
let bytesDownloaded = 0;
let maxRetry = 0;
// Bounded-parallelism worker pool with per-chunk retries.
let nextIndex = 0;
let errors = 0;
async function worker() {
while (nextIndex < chunks.length) {
const index = nextIndex++;
const chunk = chunks[index];
const url = "/model?offset_bytes=" + chunk.offset + "&size_bytes=" + chunk.size;
let ok = false;
for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
try {
if (attempt > 0) {
if (attempt > maxRetry) maxRetry = attempt;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
}
console.log("[model] chunk", index, "retry", attempt);
await new Promise((r) => setTimeout(r, RETRY_DELAY * attempt));
}
const resp = await fetch(url);
if (!resp.ok) continue;
const buffer = await resp.arrayBuffer();
HEAPU8.set(new Uint8Array(buffer), ptr + chunk.offset);
bytesDownloaded += chunk.size;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(bytesDownloaded, totalSize, maxRetry);
}
ok = true;
break;
} catch (e) {
console.warn("[model] chunk", index, "attempt", attempt, e);
}
}
if (!ok) {
console.error("[model] chunk", index, "failed after", MAX_RETRIES, "retries");
errors++;
}
}
}
const workers = [];
for (let i = 0; i < Math.min(PARALLEL, chunks.length); i++) {
workers.push(worker());
}
await Promise.all(workers);
if (errors > 0) return null; // finally frees ptr
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
}
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
// before C++ synchronous parsing blocks the thread.
await new Promise((resolve) => setTimeout(resolve, 60));
const result = { ptr: ptr, size: totalSize };
ptr = 0; // transfer ownership to caller
return result;
} catch (e) {
console.error("[model] chunked fetch failed:", e);
return null; // finally frees ptr
} finally {
if (ptr) Module.freeModelBuffer(ptr);
}
}
// Single-fetch fallback if chunked download using fetchModelChunked() fails.
async function fetchModelSingle() {
let ptr = 0;
try {
const resp = await fetch("/model");
if (!resp.ok) return null;
const buffer = await resp.arrayBuffer();
const totalSize = buffer.byteLength;
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(0, totalSize, 0);
}
ptr = Module.allocModelBuffer(totalSize);
if (!ptr) return null;
HEAPU8.set(new Uint8Array(buffer), ptr);
if (Module.updateModelDownloadProgress) {
Module.updateModelDownloadProgress(totalSize, totalSize, 0);
}
// Allow ~4 frames (60ms) to elapse so the 100% progress bar paints
// before C++ synchronous parsing blocks the thread.
await new Promise((resolve) => setTimeout(resolve, 60));
const result = { ptr: ptr, size: totalSize };
ptr = 0; // transfer ownership to caller
return result;
} catch (e) {
console.error("[model] single fetch failed:", e);
return null; // finally frees ptr
} finally {
if (ptr) Module.freeModelBuffer(ptr);
}
}
// Called from C++ (via EM_ASM) when the Python side swaps the model.
async function reloadModel() {
// Retry the chunked download up to 3 times with increasing delay.
// Falling back to a single-fetch for large models hits the same
// proxy errors, so retrying chunks is the only viable path.
for (let attempt = 0; attempt < 3; attempt++) {
if (attempt > 0) {
console.log("[model] retrying chunked download, attempt", attempt + 1);
await new Promise((r) => setTimeout(r, 1000 * attempt));
}
const modelRes = await fetchModelChunked();
if (modelRes) {
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
Module.freeModelBuffer(modelRes.ptr);
return;
}
}
console.error("[model] all chunked download attempts failed; trying single-fetch");
const modelRes = await fetchModelSingle();
if (modelRes) {
Module.parseModelBuffer(modelRes.ptr, modelRes.size);
Module.freeModelBuffer(modelRes.ptr);
}
}
var Module = {
canvas: (() => {
const canvas = document.getElementById("canvas");
@@ -92,20 +246,26 @@ limitations under the License.
];
const assetPromises = assetsToPrefetch.map(async (filename) => {
try {
const response = await fetch("assets/" + filename);
if (!response.ok) {
console.error(`Failed to fetch asset ${filename}: ${response.statusText}`);
for (let attempt = 0; attempt < 3; attempt++) {
try {
if (attempt > 0) {
await new Promise((r) => setTimeout(r, 500 * attempt));
}
const response = await fetch("assets/" + filename);
if (!response.ok) continue;
const buffer = await response.arrayBuffer();
Module.registerAsset(filename, new Uint8Array(buffer));
return;
} catch (error) {
console.warn(`Attempt ${attempt + 1} failed for ${filename}:`, error);
}
const buffer = await response.arrayBuffer();
Module.registerAsset(filename, new Uint8Array(buffer));
} catch (error) {
console.error(`Error prefetching asset ${filename}:`, error);
}
console.error(`Error prefetching asset ${filename} after 3 attempts`);
});
Promise.all(assetPromises).then(() => {
Module.startApp();
reloadModel();
});
},
};
@@ -186,7 +346,7 @@ limitations under the License.
const proto = location.protocol === "https:" ? "wss://" : "ws://";
// The session id tells the server which page is dropping; it only
// accepts drops from the controller.
const sid = (window.Module && Module.sessionId) || "";
const sid = (window.Module && Module.session_id) || "";
const dropUrl = proto + location.host + "/drop?sid=" + encodeURIComponent(sid);
const ws = new WebSocket(dropUrl);
ws.binaryType = "arraybuffer";