MuJoCo Live: pre-fetch model assets in parallel before compilation.
The prefetcher asks mju_getXMLDependencies for the transitive asset list, fetches every URL with Promise.all, and primes the bytes into the WASM-side FetchCache. When the compiler runs, every resource it opens is already in memory. Cold-load on a typical Menagerie model drops from ~9 s to ~2 s.
This commit is contained in:
Kevin Zakka
@@ -27,6 +27,7 @@
|
||||
#include <string_view>
|
||||
#include <unordered_map>
|
||||
#include <utility>
|
||||
#include <vector>
|
||||
|
||||
#include <mujoco/mujoco.h>
|
||||
#include "experimental/platform/hal/graphics_mode.h"
|
||||
@@ -108,8 +109,9 @@ EM_ASYNC_JS(int, FetchUrl,
|
||||
}
|
||||
});
|
||||
|
||||
// Cache for data fetched via HTTP/HTTPS. Stores the downloaded bytes keyed by
|
||||
// the resource name (URL) so that read() can return a pointer to the data.
|
||||
// Cache for data fetched via HTTP/HTTPS. Stores downloaded (or pre-primed)
|
||||
// bytes keyed by the resource URL so that the resource provider's read()
|
||||
// can return a pointer into stable storage.
|
||||
class FetchCache {
|
||||
public:
|
||||
static FetchCache& Instance() {
|
||||
@@ -117,52 +119,74 @@ class FetchCache {
|
||||
return instance;
|
||||
}
|
||||
|
||||
// Fetches the URL and stores the result. Returns the size (>0) on success.
|
||||
// Returns the size of the entry for `url`, fetching it over the network
|
||||
// if necessary. Entries that have been pre-populated via Prime() are
|
||||
// returned without touching the network. Returns 0 on fetch failure.
|
||||
int Fetch(const char* url) {
|
||||
char* data = nullptr;
|
||||
if (auto it = entries_.find(url); it != entries_.end()) {
|
||||
return static_cast<int>(it->second.size());
|
||||
}
|
||||
char* buf = nullptr;
|
||||
std::int32_t size = 0;
|
||||
if (!FetchUrl(url, &data, &size)) {
|
||||
if (!FetchUrl(url, &buf, &size)) {
|
||||
return 0;
|
||||
}
|
||||
entries_[url] = Entry{UniquePtrWasm<char[]>(data), size};
|
||||
return size;
|
||||
// FetchUrl hands us a malloc'd buffer; take ownership of the bytes via
|
||||
// a std::string and release the original allocation.
|
||||
std::string bytes(buf, size);
|
||||
std::free(buf);
|
||||
return static_cast<int>(entries_.emplace(url, std::move(bytes))
|
||||
.first->second.size());
|
||||
}
|
||||
|
||||
// Returns pointer and size for a previously fetched URL.
|
||||
int Read(const char* url, const void** buffer) {
|
||||
// Pre-populates the cache so a subsequent Fetch() of the same URL hits
|
||||
// memory instead of the network. First writer wins.
|
||||
void Prime(std::string url, std::string bytes) {
|
||||
entries_.try_emplace(std::move(url), std::move(bytes));
|
||||
}
|
||||
|
||||
// Hands back a pointer into the cached bytes for the given URL.
|
||||
// Returns -1 if the URL has not been fetched or primed.
|
||||
int Read(const char* url, const void** buffer) const {
|
||||
auto it = entries_.find(url);
|
||||
if (it == entries_.end()) {
|
||||
return -1;
|
||||
}
|
||||
*buffer = it->second.data.get();
|
||||
return it->second.size;
|
||||
if (it == entries_.end()) return -1;
|
||||
*buffer = it->second.data();
|
||||
return static_cast<int>(it->second.size());
|
||||
}
|
||||
|
||||
// Frees the data for a URL.
|
||||
void Close(const char* url) { entries_.erase(url); }
|
||||
|
||||
private:
|
||||
struct FreeDeleter {
|
||||
void operator()(void* p) const { std::free(p); }
|
||||
};
|
||||
template <typename T>
|
||||
using UniquePtrWasm = std::unique_ptr<T, FreeDeleter>;
|
||||
|
||||
struct Entry {
|
||||
UniquePtrWasm<char[]> data;
|
||||
int size;
|
||||
};
|
||||
std::unordered_map<std::string, Entry> entries_;
|
||||
std::unordered_map<std::string, std::string> entries_;
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// Javascript-facing function to register an asset.
|
||||
// Javascript-facing function to register a build-time asset (font, IBL,
|
||||
// .filamat material) that is fetched via the static AssetRegistry rather
|
||||
// than over HTTP.
|
||||
void RegisterAsset(std::string filename, std::string contents) {
|
||||
AssetRegistry::Instance().RegisterAsset(std::move(filename),
|
||||
std::move(contents));
|
||||
}
|
||||
|
||||
// Javascript-facing function to pre-populate the HTTP/HTTPS fetch cache.
|
||||
// Used by the JS-side parallel prefetcher to feed already-downloaded asset
|
||||
// bytes to MuJoCo's compiler before LoadUrl is invoked, sidestepping the
|
||||
// otherwise strictly-serial EM_ASYNC_JS fetch chain.
|
||||
void PrimeFetchCache(std::string url, std::string bytes) {
|
||||
FetchCache::Instance().Prime(std::move(url), std::move(bytes));
|
||||
}
|
||||
|
||||
// Javascript-facing wrapper around mju_getXMLDependencies. Returns the
|
||||
// list of resource URLs that the given root XML transitively references
|
||||
// (meshes, textures, includes, ...), per MuJoCo's compiler rules.
|
||||
std::vector<std::string> GetXMLDependencies(const std::string& root_url) {
|
||||
mjStringVec deps;
|
||||
mju_getXMLDependencies(root_url.c_str(), &deps);
|
||||
return deps;
|
||||
}
|
||||
|
||||
// Javascript-facing function to initialize the app.
|
||||
void Init(const std::string& title, bool dark_theme) {
|
||||
// Note: dimensions do not matter as window will be resized to fit canvas.
|
||||
@@ -305,7 +329,11 @@ void Deinit() {
|
||||
}
|
||||
|
||||
EMSCRIPTEN_BINDINGS(studio_bindings) {
|
||||
// Matches the C-level `mjStringVec` typedef in mjspec.h.
|
||||
emscripten::register_vector<std::string>("mjStringVec");
|
||||
emscripten::function("registerAsset", &RegisterAsset);
|
||||
emscripten::function("primeFetchCache", &PrimeFetchCache);
|
||||
emscripten::function("getXMLDependencies", &GetXMLDependencies);
|
||||
emscripten::function("init", &Init);
|
||||
emscripten::function("loadFile", &LoadFile);
|
||||
emscripten::function("loadUrl", &LoadUrl);
|
||||
|
||||
@@ -21,6 +21,85 @@ function hideLoading() {
|
||||
loadingOverlay.style.display = 'none';
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Parallel asset prefetcher.
|
||||
//
|
||||
// MuJoCo's XML compiler synchronously opens each referenced asset via
|
||||
// mjpResourceProvider. In WASM that becomes a chain of ASYNCIFY-suspended
|
||||
// fetch() calls, strictly serial, one RTT per asset. For Menagerie models
|
||||
// that's ~80 round trips and ~10 s of avoidable wall time on a cold cache.
|
||||
//
|
||||
// To sidestep this we let MuJoCo itself enumerate the dependencies (via
|
||||
// mju_getXMLDependencies, exposed to JS as Module.getXMLDependencies), then
|
||||
// fetch them all in parallel and feed the bytes into the WASM-side
|
||||
// FetchCache via Module.primeFetchCache. When Module.loadUrl subsequently
|
||||
// runs, every resource the compiler opens is already in memory.
|
||||
//
|
||||
// Discovery rules live in src/xml/xml_util.cc; the JS only has to know
|
||||
// the scheme rewrites used by the resource providers in emscripten.cc.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// MuJoCo Live custom URL schemes that the C++ resource providers rewrite
|
||||
// before fetching. JS must do the same so primed cache keys match what
|
||||
// the providers eventually request.
|
||||
function resolveScheme(url) {
|
||||
if (url.startsWith('github:')) {
|
||||
return 'https://raw.githubusercontent.com/' + url.slice('github:'.length);
|
||||
}
|
||||
return url;
|
||||
}
|
||||
|
||||
async function prefetchModelAssets(rootUrl, onProgress) {
|
||||
const primed = new Set(); // URLs we've already pushed into FetchCache
|
||||
const stats = { files: 0, bytes: 0, errors: 0 };
|
||||
|
||||
// Fetch + prime, dedup against `primed`. Returns true on success.
|
||||
async function fetchAndPrime(httpUrl) {
|
||||
if (primed.has(httpUrl)) return true;
|
||||
primed.add(httpUrl);
|
||||
try {
|
||||
const resp = await fetch(httpUrl);
|
||||
if (!resp.ok) {
|
||||
console.warn('[prefetch] HTTP', resp.status, httpUrl);
|
||||
stats.errors++;
|
||||
return false;
|
||||
}
|
||||
const bytes = new Uint8Array(await resp.arrayBuffer());
|
||||
Module.primeFetchCache(httpUrl, bytes);
|
||||
stats.files++;
|
||||
stats.bytes += bytes.length;
|
||||
onProgress?.(stats);
|
||||
return true;
|
||||
} catch (err) {
|
||||
console.warn('[prefetch]', httpUrl, err);
|
||||
stats.errors++;
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// Step 1: fetch the root XML so mju_getXMLDependencies (which reads it
|
||||
// through the resource provider) hits cache instead of the network.
|
||||
// Also gives the loading screen an immediate progress tick.
|
||||
const rootHttp = resolveScheme(rootUrl);
|
||||
if (!await fetchAndPrime(rootHttp)) {
|
||||
throw new Error(`could not fetch root model ${rootHttp}`);
|
||||
}
|
||||
|
||||
// Step 2: ask MuJoCo for the transitive dependency list. The call is
|
||||
// async because reading recursively included XMLs may suspend on
|
||||
// ASYNCIFY fetches; those XMLs end up in FetchCache as a side effect.
|
||||
// The result is an embind mjStringVec we copy out and release.
|
||||
const depsVec = await Module.getXMLDependencies(rootUrl);
|
||||
const deps = [];
|
||||
for (let i = 0; i < depsVec.size(); ++i) deps.push(depsVec.get(i));
|
||||
depsVec.delete();
|
||||
|
||||
// Step 3: fetch every dependency in parallel, deduped against URLs
|
||||
// already primed in step 1 (and earlier deps within this batch).
|
||||
await Promise.all(deps.map((depUrl) => fetchAndPrime(resolveScheme(depUrl))));
|
||||
return stats;
|
||||
}
|
||||
|
||||
var Module = {
|
||||
preRun: [],
|
||||
postRun: [],
|
||||
@@ -111,13 +190,25 @@ var Module = {
|
||||
// loop until the load completes, otherwise renderFrame will
|
||||
// hit "Cannot have multiple async operations in flight".
|
||||
showLoading();
|
||||
// Surface prefetch progress on the loading screen.
|
||||
const dlEl = document.getElementById('loadingDownload');
|
||||
const dlFileEl = document.getElementById('loadingDownloadFile');
|
||||
if (dlEl) dlEl.style.display = 'block';
|
||||
const onProgress = (s) => {
|
||||
if (dlFileEl) {
|
||||
dlFileEl.textContent =
|
||||
`${s.files} files · ${(s.bytes/1024/1024).toFixed(1)} MB`;
|
||||
}
|
||||
};
|
||||
requestAnimationFrame(() => {
|
||||
requestAnimationFrame(async () => {
|
||||
try {
|
||||
await prefetchModelAssets(modelUrl, onProgress);
|
||||
await Module.loadUrl(modelUrl);
|
||||
} catch (error) {
|
||||
console.error('Failed to load model from URL:', error);
|
||||
} finally {
|
||||
if (dlEl) dlEl.style.display = 'none';
|
||||
hideLoading();
|
||||
requestAnimationFrame(Module.animate);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user