Skip to content

Commit

Permalink
Merge pull request #2936 from cloudflare/dlapid/python-pools2
Browse files Browse the repository at this point in the history 
Workerd side of python isolate pool
  • Loading branch information
@danlapid
danlapid authored Dec 9, 2024
2 parents c876e52 + a8f550e commit 93a6037
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 103 deletions.
2 changes: 1 addition & 1 deletion src/workerd/api/pyodide/pyodide.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -490,7 +490,7 @@ jsg::JsValue SetupEmscripten::getModule(jsg::Lock& js) {
}

void SetupEmscripten::visitForGc(jsg::GcVisitor& visitor) {
visitor.visit(emscriptenRuntime.emscriptenRuntime);
visitor.visit(const_cast<EmscriptenRuntime&>(emscriptenRuntime).emscriptenRuntime);
}

bool hasPythonModules(capnp::List<server::config::Worker::Module>::Reader modules) {
Expand Down
6 changes: 3 additions & 3 deletions src/workerd/api/pyodide/pyodide.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -411,8 +411,8 @@ class SimplePythonLimiter: public jsg::Object {

class SetupEmscripten: public jsg::Object {
public:
SetupEmscripten(EmscriptenRuntime emscriptenRuntime)
: emscriptenRuntime(kj::mv(emscriptenRuntime)) {};
SetupEmscripten(const EmscriptenRuntime& emscriptenRuntime)
: emscriptenRuntime(emscriptenRuntime) {};

jsg::JsValue getModule(jsg::Lock& js);

Expand All @@ -421,7 +421,7 @@ class SetupEmscripten: public jsg::Object {
}

private:
EmscriptenRuntime emscriptenRuntime;
const EmscriptenRuntime& emscriptenRuntime;
void visitForGc(jsg::GcVisitor& visitor);
};

Expand Down
206 changes: 107 additions & 99 deletions src/workerd/server/workerd-api.c++
View file
Original file line number Diff line number Diff line change
Expand Up @@ -129,13 +129,100 @@ static const PythonConfig defaultConfig{
};
} // namespace

namespace {
kj::Path getPyodideBundleFileName(kj::StringPtr version) {
return kj::Path(kj::str("pyodide_", version, ".capnp.bin"));
}

kj::Maybe<kj::Own<const kj::ReadableFile>> getPyodideBundleFile(
const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir, kj::StringPtr version) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto file = dir->tryOpenFile(filename);

return file;
}

return kj::none;
}

void writePyodideBundleFileToDisk(const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir,
kj::StringPtr version,
kj::ArrayPtr<byte> bytes) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto replacer = dir->replaceFile(filename, kj::WriteMode::CREATE | kj::WriteMode::MODIFY);

replacer->get().writeAll(bytes);
replacer->commit();
}
}

kj::Maybe<jsg::Bundle::Reader> fetchPyodideBundle(
const api::pyodide::PythonConfig& pyConfig, kj::StringPtr version) {
if (pyConfig.pyodideBundleManager.getPyodideBundle(version) != kj::none) {
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

auto maybePyodideBundleFile = getPyodideBundleFile(pyConfig.pyodideDiskCacheRoot, version);
KJ_IF_SOME(pyodideBundleFile, maybePyodideBundleFile) {
auto body = pyodideBundleFile->readAllBytes();
pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

if (version == "dev") {
// the "dev" version is special and indicates we're using the tip-of-tree version built for testing
// so we shouldn't fetch it from the internet, only check for its existence in the disk cache
return kj::none;
}

{
KJ_LOG(INFO, "Loading Pyodide package from internet...");
kj::Thread([&]() {
kj::AsyncIoContext io = kj::setupAsyncIo();
kj::HttpHeaderTable table;

kj::TlsContext::Options options;
options.useSystemTrustStore = true;

kj::Own<kj::TlsContext> tls = kj::heap<kj::TlsContext>(kj::mv(options));
auto& network = io.provider->getNetwork();
auto tlsNetwork = tls->wrapNetwork(network);
auto& timer = io.provider->getTimer();

auto client = kj::newHttpClient(timer, table, network, *tlsNetwork);

kj::HttpHeaders headers(table);

kj::String url =
kj::str("https://pyodide.runtime-playground.workers.dev/pyodide-capnp-bin/pyodide_",
version, ".capnp.bin");

auto req = client->request(kj::HttpMethod::GET, url.asPtr(), headers);

auto res = req.response.wait(io.waitScope);
auto body = res.body->readAllBytes().wait(io.waitScope);

writePyodideBundleFileToDisk(pyConfig.pyodideDiskCacheRoot, version, body);

pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
});
}

KJ_LOG(INFO, "Loaded Pyodide package from internet");
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}
} // namespace

struct WorkerdApi::Impl final {
kj::Own<CompatibilityFlags::Reader> features;
kj::Maybe<kj::Own<jsg::modules::ModuleRegistry>> maybeOwnedModuleRegistry;
kj::Own<JsgIsolateObserver> observer;
JsgWorkerdIsolate jsgIsolate;
api::MemoryCacheProvider& memoryCacheProvider;
const PythonConfig& pythonConfig;
kj::Maybe<api::pyodide::EmscriptenRuntime> maybeEmscriptenRuntime;

class Configuration {
public:
Expand Down Expand Up @@ -169,7 +256,21 @@ struct WorkerdApi::Impl final {
observer(kj::atomicAddRef(*observerParam)),
jsgIsolate(v8System, Configuration(*this), kj::mv(observerParam), kj::mv(createParams)),
memoryCacheProvider(memoryCacheProvider),
pythonConfig(pythonConfig) {}
pythonConfig(pythonConfig) {
jsgIsolate.runInLockScope([&](JsgWorkerdIsolate::Lock& lock) {
if (features->getPythonWorkers()) {
auto pythonRelease = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(*features));
auto version = getPythonBundleName(pythonRelease);
auto bundle = KJ_ASSERT_NONNULL(
fetchPyodideBundle(pythonConfig, version), "Failed to get Pyodide bundle");
auto context = lock.newContext<api::ServiceWorkerGlobalScope>(lock.v8Isolate);
v8::Context::Scope scope(context.getHandle(lock));
// Init emscripten synchronously, the python script will import setup-emscripten and
// call setEmscriptenModele
maybeEmscriptenRuntime = api::pyodide::EmscriptenRuntime::initialize(lock, true, bundle);
}
});
}

static v8::Local<v8::String> compileTextGlobal(
JsgWorkerdIsolate::Lock& lock, capnp::Text::Reader reader) {
Expand Down Expand Up @@ -419,92 +520,6 @@ kj::Maybe<jsg::ModuleRegistry::ModuleInfo> WorkerdApi::tryCompileModule(jsg::Loc
KJ_UNREACHABLE;
}

namespace {
kj::Path getPyodideBundleFileName(kj::StringPtr version) {
return kj::Path(kj::str("pyodide_", version, ".capnp.bin"));
}

kj::Maybe<kj::Own<const kj::ReadableFile>> getPyodideBundleFile(
const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir, kj::StringPtr version) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto file = dir->tryOpenFile(filename);

return file;
}

return kj::none;
}

void writePyodideBundleFileToDisk(const kj::Maybe<kj::Own<const kj::Directory>>& maybeDir,
kj::StringPtr version,
kj::ArrayPtr<byte> bytes) {
KJ_IF_SOME(dir, maybeDir) {
kj::Path filename = getPyodideBundleFileName(version);
auto replacer = dir->replaceFile(filename, kj::WriteMode::CREATE | kj::WriteMode::MODIFY);

replacer->get().writeAll(bytes);
replacer->commit();
}
}

kj::Maybe<jsg::Bundle::Reader> fetchPyodideBundle(
const api::pyodide::PythonConfig& pyConfig, kj::StringPtr version) {
if (pyConfig.pyodideBundleManager.getPyodideBundle(version) != kj::none) {
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

auto maybePyodideBundleFile = getPyodideBundleFile(pyConfig.pyodideDiskCacheRoot, version);
KJ_IF_SOME(pyodideBundleFile, maybePyodideBundleFile) {
auto body = pyodideBundleFile->readAllBytes();
pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}

if (version == "dev") {
// the "dev" version is special and indicates we're using the tip-of-tree version built for testing
// so we shouldn't fetch it from the internet, only check for its existence in the disk cache
return kj::none;
}

{
KJ_LOG(INFO, "Loading Pyodide package from internet...");
kj::Thread([&]() {
kj::AsyncIoContext io = kj::setupAsyncIo();
kj::HttpHeaderTable table;

kj::TlsContext::Options options;
options.useSystemTrustStore = true;

kj::Own<kj::TlsContext> tls = kj::heap<kj::TlsContext>(kj::mv(options));
auto& network = io.provider->getNetwork();
auto tlsNetwork = tls->wrapNetwork(network);
auto& timer = io.provider->getTimer();

auto client = kj::newHttpClient(timer, table, network, *tlsNetwork);

kj::HttpHeaders headers(table);

kj::String url =
kj::str("https://pyodide.runtime-playground.workers.dev/pyodide-capnp-bin/pyodide_",
version, ".capnp.bin");

auto req = client->request(kj::HttpMethod::GET, url.asPtr(), headers);

auto res = req.response.wait(io.waitScope);
auto body = res.body->readAllBytes().wait(io.waitScope);

writePyodideBundleFileToDisk(pyConfig.pyodideDiskCacheRoot, version, body);

pyConfig.pyodideBundleManager.setPyodideBundleData(kj::str(version), kj::mv(body));
});
}

KJ_LOG(INFO, "Loaded Pyodide package from internet");
return pyConfig.pyodideBundleManager.getPyodideBundle(version);
}
} // namespace

void WorkerdApi::compileModules(jsg::Lock& lockParam,
config::Worker::Reader conf,
Worker::ValidationErrorReporter& errorReporter,
Expand All @@ -519,22 +534,15 @@ void WorkerdApi::compileModules(jsg::Lock& lockParam,
if (hasPythonModules(confModules)) {
KJ_REQUIRE(featureFlags.getPythonWorkers(),
"The python_workers compatibility flag is required to use Python.");
// Inject SetupEmscripten module
modules->addBuiltinModule("internal:setup-emscripten",
jsg::alloc<SetupEmscripten>(KJ_ASSERT_NONNULL(impl->maybeEmscriptenRuntime)),
workerd::jsg::ModuleRegistry::Type::INTERNAL);

auto pythonRelease = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
auto version = getPythonBundleName(pythonRelease);
auto bundle = KJ_ASSERT_NONNULL(
fetchPyodideBundle(impl->pythonConfig, version), "Failed to get Pyodide bundle");
// Inject SetupEmscripten module
{
auto& lock = kj::downcast<JsgWorkerdIsolate::Lock>(lockParam);
auto context = lock.newContext<api::ServiceWorkerGlobalScope>({}, lock.v8Isolate);
v8::Context::Scope scope(context.getHandle(lock));
// Init emscripten synchronously, the python script will import setup-emscripten and
// call setEmscriptenModele
auto emscriptenRuntime = api::pyodide::EmscriptenRuntime::initialize(lock, true, bundle);
modules->addBuiltinModule("internal:setup-emscripten",
jsg::alloc<SetupEmscripten>(kj::mv(emscriptenRuntime)),
workerd::jsg::ModuleRegistry::Type::INTERNAL);
}

// Inject Pyodide bundle
modules->addBuiltinBundle(bundle, kj::none);
Expand Down

0 comments on commit 93a6037

Please sign in to comment.