From 4b786cfecdbee60a0a5b9b023f136905ebfa4e56 Mon Sep 17 00:00:00 2001 From: Hood Chatham Date: Wed, 27 Mar 2024 12:16:37 +0100 Subject: [PATCH] Python: Reduce memory usage when restoring snapshot (#1875) Previously we paid for two copies of the snapshot memory: one copy in the wasm linear memory itself and a second copy in `BUNDLE_MEMORY_SNAPSHOT`. This ensures that we never have more memory than one copy of the linear memory heap by copying the memory directly from the snapshot to the linear memory. We also release the C++ memory when we are done with it. --- src/pyodide/internal/metadata.js | 7 ++- src/pyodide/internal/python.js | 61 ++++++++++++------------ src/pyodide/internal/setupPackages.js | 6 +-- src/pyodide/python-entrypoint-helper.js | 10 +++- src/workerd/api/pyodide/pyodide.c++ | 38 ++++++++++----- src/workerd/api/pyodide/pyodide.h | 63 ++++++++++++++++++------- 6 files changed, 119 insertions(+), 66 deletions(-) diff --git a/src/pyodide/internal/metadata.js b/src/pyodide/internal/metadata.js index 22e7116377f..854b28ff59b 100644 --- a/src/pyodide/internal/metadata.js +++ b/src/pyodide/internal/metadata.js @@ -1,10 +1,15 @@ import { default as MetadataReader } from "pyodide-internal:runtime-generated/metadata"; export { default as LOCKFILE } from "pyodide-internal:generated/pyodide-lock.json"; import { default as PYODIDE_BUCKET } from "pyodide-internal:generated/pyodide-bucket.json"; +import { default as ArtifactBundler } from "pyodide-internal:artifacts"; export const IS_WORKERD = MetadataReader.isWorkerd(); export const IS_TRACING = MetadataReader.isTracing(); export const WORKERD_INDEX_URL = PYODIDE_BUCKET.PYODIDE_PACKAGE_BUCKET_URL; export const REQUIREMENTS = MetadataReader.getRequirements(); export const MAIN_MODULE_NAME = MetadataReader.getMainModule(); -export const BUNDLE_MEMORY_SNAPSHOT = MetadataReader.getMemorySnapshot(); +export const MEMORY_SNAPSHOT_READER = MetadataReader.hasMemorySnapshot() + ? MetadataReader + : ArtifactBundler.hasMemorySnapshot() + ? ArtifactBundler + : undefined; diff --git a/src/pyodide/internal/python.js b/src/pyodide/internal/python.js index d3c115e1273..51c03cbaa6a 100644 --- a/src/pyodide/internal/python.js +++ b/src/pyodide/internal/python.js @@ -10,7 +10,7 @@ import { } from "pyodide-internal:setupPackages"; import { default as TarReader } from "pyodide-internal:packages_tar_reader"; import processScriptImports from "pyodide-internal:process_script_imports.py"; -import { BUNDLE_MEMORY_SNAPSHOT } from "pyodide-internal:metadata"; +import { MEMORY_SNAPSHOT_READER } from "pyodide-internal:metadata"; /** * This file is a simplified version of the Pyodide loader: @@ -38,7 +38,8 @@ import pyodideWasmModule from "pyodide-internal:generated/pyodide.asm.wasm"; */ import stdlib from "pyodide-internal:generated/python_stdlib.zip"; -const SHOULD_UPLOAD_SNAPSHOT = ArtifactBundler.isEnabled() || ArtifactBundler.isEwValidating(); +const SHOULD_UPLOAD_SNAPSHOT = + ArtifactBundler.isEnabled() || ArtifactBundler.isEwValidating(); const DEDICATED_SNAPSHOT = true; /** @@ -47,7 +48,9 @@ const DEDICATED_SNAPSHOT = true; * which is quite slow. Startup with snapshot is 3-5 times faster than without * it. */ -let MEMORY = undefined; +let READ_MEMORY = undefined; +let SNAPSHOT_SIZE = undefined; + /** * Record the dlopen handles that are needed by the MEMORY. */ @@ -234,7 +237,10 @@ function getEmscriptenSettings(lockfile, indexURL) { // important because the file system lives outside of linear memory. preRun: [prepareFileSystem, setEnv, preloadDynamicLibs], instantiateWasm, - noInitialRun: !!MEMORY, // skip running main() if we have a snapshot + // if SNAPSHOT_SIZE is defined, start with the linear memory big enough to + // fit the snapshot. If it's not defined, this falls back to the default. + INITIAL_MEMORY: SNAPSHOT_SIZE, + noInitialRun: !!READ_MEMORY, // skip running main() if we have a snapshot API, // Pyodide requires we pass this in. }; } @@ -278,16 +284,8 @@ async function instantiateEmscriptenModule(emscriptenSettings) { async function prepareWasmLinearMemory(Module) { // Note: if we are restoring from a snapshot, runtime is not initialized yet. mountLib(Module, SITE_PACKAGES_INFO); - if (MEMORY) { - if (!(MEMORY instanceof Uint8Array)) { - throw new TypeError("Expected MEMORY to be a Uint8Array"); - } - // resize linear memory to fit our snapshot. I think `growMemory` only - // exists if `-sALLOW_MEMORY_GROWTH` is passed to the linker but we'll - // probably always do that. - Module.growMemory(MEMORY.byteLength); - // restore memory from snapshot - Module.HEAP8.set(MEMORY); + if (READ_MEMORY) { + READ_MEMORY(Module); // Don't call adjustSysPath here: it was called in the other branch when we // were creating the snapshot so the outcome of that is already baked in. return; @@ -440,13 +438,20 @@ function encodeSnapshot(heap, dsoJSON) { /** * Decode heap and dsoJSON from the memory snapshot artifact we downloaded */ -function decodeSnapshot(memorySnapshot) { - const uint32View = new Uint32Array(memorySnapshot); - const snapshotOffset = uint32View[0]; - const jsonLength = uint32View[1]; - const jsonView = new Uint8Array(memorySnapshot, 8, jsonLength); - DSO_METADATA = JSON.parse(new TextDecoder().decode(jsonView)); - MEMORY = new Uint8Array(memorySnapshot, snapshotOffset); +function decodeSnapshot() { + const buf = new Uint32Array(2); + MEMORY_SNAPSHOT_READER.readMemorySnapshot(0, buf); + const snapshotOffset = buf[0]; + SNAPSHOT_SIZE = MEMORY_SNAPSHOT_READER.getMemorySnapshotSize() - snapshotOffset; + const jsonLength = buf[1]; + const jsonBuf = new Uint8Array(jsonLength); + MEMORY_SNAPSHOT_READER.readMemorySnapshot(8, jsonBuf); + DSO_METADATA = JSON.parse(new TextDecoder().decode(jsonBuf)); + READ_MEMORY = function(Module) { + // restore memory from snapshot + MEMORY_SNAPSHOT_READER.readMemorySnapshot(snapshotOffset, Module.HEAP8); + MEMORY_SNAPSHOT_READER.disposeMemorySnapshot(); + } } /** @@ -492,24 +497,22 @@ function simpleRunPython(emscriptenModule, code) { let TEST_SNAPSHOT = undefined; (function () { // Lookup memory snapshot from artifact store. - const memorySnapshot = BUNDLE_MEMORY_SNAPSHOT || ArtifactBundler.getMemorySnapshot(); - if (!memorySnapshot) { + if (!MEMORY_SNAPSHOT_READER) { // snapshots are disabled or there isn't one yet return; } - if (memorySnapshot.constructor.name !== "ArrayBuffer") { - throw new TypeError("Expected snapshot to be an ArrayBuffer"); - } // Simple sanity check to ensure this snapshot isn't corrupted. // // TODO(later): we need better detection when this is corrupted. Right now the isolate will // just die. - if (memorySnapshot.byteLength <= 100) { - TEST_SNAPSHOT = memorySnapshot; + const snapshotSize = MEMORY_SNAPSHOT_READER.getMemorySnapshotSize(); + if (snapshotSize <= 100) { + TEST_SNAPSHOT = new Uint8Array(snapshotSize); + MEMORY_SNAPSHOT_READER.readMemorySnapshot(0, TEST_SNAPSHOT); return; } - decodeSnapshot(memorySnapshot); + decodeSnapshot(); })(); export async function loadPyodide(lockfile, indexURL) { diff --git a/src/pyodide/internal/setupPackages.js b/src/pyodide/internal/setupPackages.js index e4eceecffbd..116a30d1346 100644 --- a/src/pyodide/internal/setupPackages.js +++ b/src/pyodide/internal/setupPackages.js @@ -21,7 +21,6 @@ const STDLIB_PACKAGES = Object.values(LOCKFILE.packages) .filter(({ install_dir }) => install_dir === "stdlib") .map(({ name }) => canonicalizePackageName(name)); - /** * This stitches together the view of the site packages directory. Each * requirement corresponds to a folder in the original tar file. For each @@ -186,6 +185,5 @@ function addPackageToLoad(lockfile, name, toLoad) { export { REQUIREMENTS }; export const TRANSITIVE_REQUIREMENTS = getTransitiveRequirements(); -export const [SITE_PACKAGES_INFO, SITE_PACKAGES_SO_FILES, USE_LOAD_PACKAGE] = buildSitePackages( - TRANSITIVE_REQUIREMENTS, -); +export const [SITE_PACKAGES_INFO, SITE_PACKAGES_SO_FILES, USE_LOAD_PACKAGE] = + buildSitePackages(TRANSITIVE_REQUIREMENTS); diff --git a/src/pyodide/python-entrypoint-helper.js b/src/pyodide/python-entrypoint-helper.js index a2cc386eaf0..0c350fc741b 100644 --- a/src/pyodide/python-entrypoint-helper.js +++ b/src/pyodide/python-entrypoint-helper.js @@ -1,7 +1,11 @@ // This file is a BUILTIN module that provides the actual implementation for the // python-entrypoint.js USER module. -import { loadPyodide, uploadArtifacts, getMemoryToUpload } from "pyodide-internal:python"; +import { + loadPyodide, + uploadArtifacts, + getMemoryToUpload, +} from "pyodide-internal:python"; import { enterJaegerSpan } from "pyodide-internal:jaeger"; import { REQUIREMENTS, @@ -109,7 +113,9 @@ function getMainModule() { mainModulePromise = (async function () { const pyodide = await getPyodide(); await setupPackages(pyodide); - return enterJaegerSpan("pyimport_main_module", () => pyimportMainModule(pyodide)); + return enterJaegerSpan("pyimport_main_module", () => + pyimportMainModule(pyodide), + ); })(); return mainModulePromise; }); diff --git a/src/workerd/api/pyodide/pyodide.c++ b/src/workerd/api/pyodide/pyodide.c++ index 0604a12fd53..61c45d2f7d6 100644 --- a/src/workerd/api/pyodide/pyodide.c++ +++ b/src/workerd/api/pyodide/pyodide.c++ @@ -1,20 +1,27 @@ #include "pyodide.h" +#include "kj/array.h" +#include "kj/common.h" +#include "kj/debug.h" namespace workerd::api::pyodide { -int PackagesTarReader::read(jsg::Lock& js, int offset, kj::Array buf) { - int tarSize = PYODIDE_PACKAGES_TAR->size(); - if (offset >= tarSize || offset < 0) { +static int readToTarget(kj::ArrayPtr source, int offset, kj::ArrayPtr buf) { + int size = source.size(); + if (offset >= size || offset < 0) { return 0; } int toCopy = buf.size(); - if (tarSize - offset < toCopy) { - toCopy = tarSize - offset; + if (size - offset < toCopy) { + toCopy = size - offset; } - memcpy(buf.begin(), &((*PYODIDE_PACKAGES_TAR)[0]) + offset, toCopy); + memcpy(buf.begin(), source.begin() + offset, toCopy); return toCopy; } +int PackagesTarReader::read(jsg::Lock& js, int offset, kj::Array buf) { + return readToTarget(PYODIDE_PACKAGES_TAR.get(), offset, buf); +} + kj::Array> PyodideMetadataReader::getNames(jsg::Lock& js) { auto builder = kj::heapArrayBuilder>(this->names.size()); for (auto i : kj::zeroTo(builder.capacity())) { @@ -44,16 +51,21 @@ int PyodideMetadataReader::read(jsg::Lock& js, int index, int offset, kj::Array< return 0; } auto& data = contents[index]; - int dataSize = data.size(); - if (offset >= dataSize || offset < 0) { + return readToTarget(data, offset, buf); +} + +int PyodideMetadataReader::readMemorySnapshot(int offset, kj::Array buf) { + if (memorySnapshot == kj::none) { return 0; } - int toCopy = buf.size(); - if (dataSize - offset < toCopy) { - toCopy = dataSize - offset; + return readToTarget(KJ_REQUIRE_NONNULL(memorySnapshot), offset, buf); +} + +int ArtifactBundler::readMemorySnapshot(int offset, kj::Array buf) { + if (existingSnapshot == kj::none) { + return 0; } - memcpy(buf.begin(), &data[0] + offset, toCopy); - return toCopy; + return readToTarget(KJ_REQUIRE_NONNULL(existingSnapshot), offset, buf); } jsg::Ref makePyodideMetadataReader(Worker::Reader conf) { diff --git a/src/workerd/api/pyodide/pyodide.h b/src/workerd/api/pyodide/pyodide.h index 427f9bd38b3..500722e76da 100644 --- a/src/workerd/api/pyodide/pyodide.h +++ b/src/workerd/api/pyodide/pyodide.h @@ -1,5 +1,7 @@ #pragma once +#include "kj/array.h" +#include "kj/debug.h" #include #include #include @@ -23,6 +25,7 @@ class PackagesTarReader : public jsg::Object { } }; + // A function to read a segment of the tar file into a buffer // Set up this way to avoid copying files that aren't accessed. class PyodideMetadataReader : public jsg::Object { @@ -56,10 +59,6 @@ class PyodideMetadataReader : public jsg::Object { return kj::str(this->mainModule); } - kj::Maybe> getMemorySnapshot() { - return kj::mv(memorySnapshot); - } - kj::Array> getNames(jsg::Lock& js); kj::Array> getRequirements(jsg::Lock& js); @@ -68,6 +67,21 @@ class PyodideMetadataReader : public jsg::Object { int read(jsg::Lock& js, int index, int offset, kj::Array buf); + bool hasMemorySnapshot() { + return memorySnapshot != kj::none; + } + int getMemorySnapshotSize() { + if (memorySnapshot == kj::none) { + return 0; + } + return KJ_REQUIRE_NONNULL(memorySnapshot).size(); + } + + void disposeMemorySnapshot() { + memorySnapshot = kj::none; + } + int readMemorySnapshot(int offset, kj::Array buf); + JSG_RESOURCE_TYPE(PyodideMetadataReader) { JSG_METHOD(isWorkerd); JSG_METHOD(isTracing); @@ -76,7 +90,10 @@ class PyodideMetadataReader : public jsg::Object { JSG_METHOD(getNames); JSG_METHOD(getSizes); JSG_METHOD(read); - JSG_METHOD(getMemorySnapshot); + JSG_METHOD(hasMemorySnapshot); + JSG_METHOD(getMemorySnapshotSize); + JSG_METHOD(readMemorySnapshot); + JSG_METHOD(disposeMemorySnapshot); } void visitForMemoryInfo(jsg::MemoryTracker& tracker) const { @@ -101,11 +118,13 @@ class ArtifactBundler : public jsg::Object { ArtifactBundler(kj::Maybe> existingSnapshot, kj::Function(kj::Array snapshot)> uploadMemorySnapshotCb) - : storedSnapshot(kj::none), + : + storedSnapshot(kj::none), existingSnapshot(kj::mv(existingSnapshot)), uploadMemorySnapshotCb(kj::mv(uploadMemorySnapshotCb)), hasUploaded(false), - isValidating(false) {}; + isValidating(false) + {}; ArtifactBundler(kj::Maybe> existingSnapshot) : storedSnapshot(kj::none), @@ -144,13 +163,6 @@ class ArtifactBundler : public jsg::Object { storedSnapshot = kj::mv(snapshot); } - jsg::Optional> getMemorySnapshot(jsg::Lock& js) { - KJ_IF_SOME(val, existingSnapshot) { - return kj::mv(val); - } - return kj::none; - } - bool isEnabled() { return uploadMemorySnapshotCb != kj::none; } @@ -159,6 +171,19 @@ class ArtifactBundler : public jsg::Object { return existingSnapshot != kj::none; } + int getMemorySnapshotSize() { + if (existingSnapshot == kj::none) { + return 0; + } + return KJ_REQUIRE_NONNULL(existingSnapshot).size(); + } + + int readMemorySnapshot(int offset, kj::Array buf); + void disposeMemorySnapshot() { + existingSnapshot = kj::none; + } + + // Determines whether this ArtifactBundler was created inside the validator. bool isEwValidating() { return isValidating; @@ -169,14 +194,18 @@ class ArtifactBundler : public jsg::Object { } void visitForMemoryInfo(jsg::MemoryTracker& tracker) const { - KJ_IF_SOME(snapshot, existingSnapshot) { - tracker.trackFieldWithSize("snapshot", snapshot.size()); + if (existingSnapshot == kj::none) { + return; } + tracker.trackFieldWithSize("snapshot", KJ_REQUIRE_NONNULL(existingSnapshot).size()); } JSG_RESOURCE_TYPE(ArtifactBundler) { JSG_METHOD(uploadMemorySnapshot); - JSG_METHOD(getMemorySnapshot); + JSG_METHOD(hasMemorySnapshot); + JSG_METHOD(getMemorySnapshotSize); + JSG_METHOD(readMemorySnapshot); + JSG_METHOD(disposeMemorySnapshot); JSG_METHOD(isEnabled); JSG_METHOD(isEwValidating); JSG_METHOD(storeMemorySnapshot);