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

Merge pull request #3293 from kevinzakka:studio-live-parallel-prefetch

Browse Source 
PiperOrigin-RevId: 922417778
Change-Id: Ia839747c4af0e705962308dc84e9a7c3926fc910
This commit is contained in:
Copybara-Service
2026-05-27 16:47:18 -07:00
parent 4035c66428 249e9c08b8
commit f755d718d3
4 changed files with 244 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/experimental/studio/emscripten.cc
+55 -27
View File
@@ -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);
src/experimental/studio/live.js
+91
View File
@@ -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);
}
src/xml/xml_util.cc
+36 -24
View File
@@ -222,34 +222,46 @@ void mjCopyError(char* dst, const char* src, int maxlen) {
}
void mju_getXMLDependencies(const char* filename, mjStringVec* dependencies) {
// load XML file or parse string
tinyxml2::XMLDocument doc;
doc.LoadFile(filename);
// error checking
if (doc.Error()) {
mju_error("Problem reading XML file '%s': %s", filename, doc.ErrorStr());
// Open through the resource provider so the function works against any
// registered backend (OS file system, VFS, HTTP, "github:", ...) rather
// than only the OS file system.
mjResource* resource = mju_openResource("", filename, nullptr, nullptr, 0);
if (resource == nullptr) {
mju_error("Could not open '%s'", filename);
}
// get top-level element
tinyxml2::XMLElement *root = doc.RootElement();
if (!root) {
mju_error("XML root element not found");
}
std::unordered_set<std::string> files = {filename};
std::optional<FilePath> model_dir = std::nullopt;
mjResource *resource = mju_openResource("", filename, nullptr,
nullptr, 0);
if (resource != nullptr) {
const char* dir;
int ndir;
mju_getResourceDir(resource, &dir, &ndir);
model_dir = FilePath(std::string(dir, ndir));
// Read the XML bytes from the resource.
const void* buffer = nullptr;
int size = mju_readResource(resource, &buffer);
if (size < 0 || !size) {
mju_closeResource(resource);
mju_error("Could not read '%s'", filename);
}
// Get file references from include and model tags.
AccumulateFiles(files, root, model_dir.value());
// Capture the model directory while the resource is still open.
const char* dir = nullptr;
int ndir = 0;
mju_getResourceDir(resource, &dir, &ndir);
FilePath model_dir(std::string(dir, ndir));
// Parse from buffer and close (the parsed DOM is independent of the
// resource buffer once Parse returns).
tinyxml2::XMLDocument doc;
tinyxml2::XMLError err =
doc.Parse(static_cast<const char*>(buffer), static_cast<size_t>(size));
mju_closeResource(resource);
if (err != tinyxml2::XML_SUCCESS) {
mju_error("Problem reading XML file '%s': %s", filename,
doc.ErrorStr() ? doc.ErrorStr() : "");
}
tinyxml2::XMLElement* root = doc.RootElement();
if (!root) {
mju_error("XML root element not found in '%s'", filename);
}
std::unordered_set<std::string> files = {filename};
AccumulateFiles(files, root, model_dir);
*dependencies = {files.begin(), files.end()};
}
test/xml/xml_utils_test.cc
+62
View File
@@ -13,11 +13,14 @@
// limitations under the License.
#include <map>
#include <set>
#include <string>
#include <vector>
#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <mujoco/mjplugin.h>
#include <mujoco/mujoco.h>
#include "test/fixture.h"
@@ -42,5 +45,64 @@ TEST_F(MujocoTest, GetXMLDependenciesTest) {
kModelPaths.end()};
EXPECT_EQ(dependency_set, expected_dependency_set);
}
// Custom resource provider that serves XML strings from an in-memory map,
// used to exercise the non-filesystem code path.
namespace memxml {
static const std::map<std::string, std::string>* g_files = nullptr;
int Open(mjResource* resource) {
return g_files && g_files->count(resource->name) ? 1 : 0;
}
int Read(mjResource* resource, const void** buffer) {
auto it = g_files->find(resource->name);
if (it == g_files->end()) return -1;
*buffer = it->second.data();
return static_cast<int>(it->second.size());
}
void Close(mjResource* resource) {}
} // namespace memxml
// Verifies mju_getXMLDependencies works against a non-filesystem resource
// provider (the case the WASM/HTTP build relies on).
TEST_F(MujocoTest, GetXMLDependenciesViaResourceProvider) {
const std::map<std::string, std::string> files = {
{"memxml:/scene.xml",
"<mujoco>\n"
" <include file=\"child.xml\"/>\n"
" <compiler meshdir=\"meshes\"/>\n"
" <asset>\n"
" <mesh name=\"m\" file=\"m.obj\"/>\n"
" </asset>\n"
"</mujoco>"},
{"memxml:/child.xml",
"<mujoco>\n"
" <asset>\n"
" <texture name=\"t\" type=\"2d\" file=\"t.png\"/>\n"
" </asset>\n"
"</mujoco>"},
};
memxml::g_files = &files;
mjpResourceProvider provider = {
.prefix = "memxml",
.open = memxml::Open,
.read = memxml::Read,
.close = memxml::Close,
};
ASSERT_GT(mjp_registerResourceProvider(&provider), 0);
mjStringVec dependencies;
mju_getXMLDependencies("memxml:/scene.xml", &dependencies);
std::set<std::string> dep_set(dependencies.begin(), dependencies.end());
EXPECT_THAT(dep_set, testing::UnorderedElementsAre(
"memxml:/scene.xml",
"memxml:/child.xml",
"memxml:/meshes/m.obj",
"memxml:/t.png"));
memxml::g_files = nullptr;
}
} // namespace
} // namespace mujoco