Source code
Revision control
Copy as Markdown
Other Tools
/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
// On macOS, `mach gtest` sets MOZ_XRE_DIR to <objdir>/dist/bin. The
// HWInference process is sandboxed and may only map executable code from the
// application bundle, so it cannot load libmozinference from there and these
// tests fail. Run them locally with MOZ_DISABLE_UTILITY_SANDBOX=1. CI passes
// --xre-path inside the bundle and runs them sandboxed.
#include "gtest/gtest.h"
#include "mozilla/Preferences.h"
#include "mozilla/Result.h"
#include "mozilla/SpinEventLoopUntil.h"
#include "mozilla/TimeStamp.h"
#include "mozilla/dom/TextGeneratorBinding.h"
#include "mozilla/gtest/WaitFor.h"
#include "mozilla/gtest/ipc/TestUtilityProcess.h"
#include "mozilla/hwinference/HWInferenceProcess.h"
#include "mozilla/hwinference/HWInferenceParent.h"
#include "mozilla/hwinference/TextGenerationParent.h"
#include "mozilla/ipc/FileDescriptor.h"
#include "mozilla/ipc/UtilityProcessManager.h"
#include "nsIProcessToolsService.h"
#include "nsServiceManagerUtils.h"
#include "nsThreadUtils.h"
#include "prio.h"
#include "private/pprio.h"
using namespace mozilla;
using namespace mozilla::hwinference;
namespace {
constexpr char kModelFile[] = "TinyStories-656K.Q8_0.gguf";
constexpr char kGracePref[] =
"browser.ml.hwinference.browser_idle_shutdown_grace_ms";
// Generous for a debug CI machine; a healthy run needs a fraction of it.
constexpr uint32_t kStepTimeoutSeconds = 120;
ipc::FileDescriptor OpenModelFd(const char* aPath) {
PRFileDesc* fd = PR_Open(aPath, PR_RDONLY, 0);
if (!fd) {
return ipc::FileDescriptor();
}
ipc::FileDescriptor result(
ipc::FileDescriptor::PlatformHandleType(PR_FileDesc2NativeHandle(fd)));
PR_Close(fd);
return result;
}
template <typename R, typename E, bool Excl>
Maybe<Result<R, E>> WaitForOrTimeout(RefPtr<MozPromise<R, E, Excl>> aPromise,
uint32_t aTimeoutSeconds) {
auto result = MakeRefPtr<media::Refcountable<Maybe<Result<R, E>>>>();
aPromise->Then(
GetCurrentSerialEventTarget(), __func__,
[result](R aResolve) {
*result = Some(Result<R, E>(std::move(aResolve)));
},
[result](E aReject) {
*result = Some(Result<R, E>(Err(std::move(aReject))));
});
TimeStamp deadline =
TimeStamp::Now() + TimeDuration::FromSeconds(aTimeoutSeconds);
SpinEventLoopUntil<ProcessFailureBehavior::IgnoreAndContinue>(
"TestTextGeneration WaitForOrTimeout"_ns,
[&] { return result->isSome() || TimeStamp::Now() > deadline; });
return std::move(*result);
}
template <typename Condition>
bool SpinUntil(const char* aName, uint32_t aTimeoutSeconds,
Condition&& aCondition) {
TimeStamp deadline =
TimeStamp::Now() + TimeDuration::FromSeconds(aTimeoutSeconds);
bool held = false;
SpinEventLoopUntil<ProcessFailureBehavior::IgnoreAndContinue>(
nsDependentCString(aName), [&] {
held = aCondition();
return held || TimeStamp::Now() > deadline;
});
return held;
}
// Nothing while the process is still launching, as its actor is not bound
// yet, and once it is gone.
Maybe<base::ProcessId> ProcessPid() {
if (!HWInferenceProcess::Browser().IsUp()) {
return Nothing();
}
RefPtr<HWInferenceParent> actor = HWInferenceProcess::Browser().Actor();
return actor->CanSend() ? Some(actor->OtherPid()) : Nothing();
}
bool ProcessIsGone() { return !HWInferenceProcess::Browser().IsUp(); }
void ExpectProcessGone(const char* aWhy) {
EXPECT_TRUE(SpinUntil("Process exit", kStepTimeoutSeconds, ProcessIsGone))
<< aWhy;
}
// Kills the process the way a crash would, and waits for the manager to
// notice.
void Crash(base::ProcessId aPid) {
nsCOMPtr<nsIProcessToolsService> tools =
do_GetService("@mozilla.org/processtools-service;1");
ASSERT_TRUE(tools);
ASSERT_TRUE(NS_SUCCEEDED(tools->Kill(aPid)));
ASSERT_TRUE(SpinUntil("Crash noticed", kStepTimeoutSeconds, ProcessIsGone));
}
// top-k:1 is greedy, but llama.cpp needs the chain to end with a dist sampler.
CopyableTArray<Sampler> GreedySamplers() {
CopyableTArray<Sampler> samplers;
// WebIDL defaults for the fields each sampler ignores.
Sampler topK(dom::TextGenerationSamplerType::Top_k, /* topK */ 1,
/* topP */ 0.95f, /* temp */ 0.80f,
/* seed */ 0xFFFFFFFF, CopyableTArray<LogitBias>());
Sampler dist = topK;
dist.type() = dom::TextGenerationSamplerType::Dist;
samplers.AppendElement(topK);
samplers.AppendElement(dist);
return samplers;
}
GenerateRequest StoryRequest(uint32_t aMaxTokens = 16,
uint32_t aBufferLength = 4) {
GenerateRequest request;
request.messages().AppendElement(ChatMessage(
dom::TextGenerationRole::System, "You are a friendly storyteller."_ns));
request.messages().AppendElement(
ChatMessage(dom::TextGenerationRole::User,
"Once upon a time there was a small mouse who"_ns));
request.maxTokens() = aMaxTokens;
request.bufferLength() = aBufferLength;
request.samplers() = GreedySamplers();
request.stopOnEndOfGenerationTokens() = true;
return request;
}
struct Session {
RefPtr<TextGenerationParent> mGenerator;
nsCString mStreamed;
};
// One fd per generator: sending a FileDescriptor consumes its handle.
Session CreateGenerator(const ipc::FileDescriptor& aModel) {
Session session;
session.mGenerator = TextGenerationParent::Create(
aModel,
TextGenerationOptions(512, 2, 2, 2048, 512,
dom::TextGenerationKVCacheDtype::F16, false));
if (!session.mGenerator) {
printf_stderr("CreateGenerator: TextGenerationParent::Create() failed\n");
}
return session;
}
Session CreateGenerator() { return CreateGenerator(OpenModelFd(kModelFile)); }
// Some() carries the load failure; Nothing() means the model loaded.
Maybe<TextGenerationParent::LoadFailure> WaitForReady(Session& aSession) {
auto ready =
WaitForOrTimeout(aSession.mGenerator->WhenReady(), kStepTimeoutSeconds);
EXPECT_TRUE(ready.isSome()) << "Ready timed out";
if (ready.isNothing()) {
return Some(TextGenerationParent::LoadFailure{
TextGenerationParent::LoadFailure::Cause::ActorGone,
"Ready timed out"_ns});
}
if (ready->isOk()) {
return Nothing();
}
return Some(ready->unwrapErr());
}
auto StartGenerate(Session& aSession, const GenerateRequest& aRequest) {
aSession.mGenerator->SetDeltaHandler(
[session = &aSession](const nsCString& aText) {
session->mStreamed.Append(aText);
});
return aSession.mGenerator->SendGenerate(aRequest);
}
// aErrorExpected only suppresses the diagnostic printf.
GenerateResponse Generate(Session& aSession, const GenerateRequest& aRequest,
bool aErrorExpected = false) {
auto result =
WaitForOrTimeout(StartGenerate(aSession, aRequest), kStepTimeoutSeconds);
EXPECT_TRUE(result.isSome()) << "Generate timed out";
if (result.isNothing()) {
return GenerateResponse();
}
EXPECT_TRUE(result->isOk());
GenerateResponse response =
result->isOk() ? result->unwrap() : GenerateResponse();
if (!aErrorExpected && response.type() == GenerateResponse::TGenerateError) {
printf_stderr("GenerateError: %s\n",
response.get_GenerateError().message().get());
}
return response;
}
void Teardown(Session& aSession) {
if (aSession.mGenerator) {
aSession.mGenerator->Close();
}
ExpectProcessGone("last generator gone but the HWInference process survived");
}
} // namespace
class TextGenerationTest : public mozilla::gtest::ipc::TestUtilityProcess {
protected:
// Zero grace: Teardown() asserts the process exits immediately.
void SetUp() override { Preferences::SetUint(kGracePref, 0); }
};
TEST_F(TextGenerationTest, GenerateFromModel) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
EXPECT_FALSE(session.mGenerator->ProcessReused());
EXPECT_TRUE(WaitForReady(session).isNothing())
<< "construction should load the model";
GenerateResponse response = Generate(session, StoryRequest());
ASSERT_EQ(response.type(), GenerateResponse::TGenerateResult);
const GenerateResult& result = response.get_GenerateResult();
EXPECT_FALSE(result.content().IsEmpty());
EXPECT_EQ(session.mStreamed, result.content());
EXPECT_GT(result.usage().promptTokens(), 0u);
EXPECT_GT(result.usage().generatedTokens(), 0u);
EXPECT_GT(result.usage().promptCharacters(), result.usage().promptTokens());
EXPECT_TRUE(result.reason() == dom::TextGenerationFinishReason::Length ||
result.reason() == dom::TextGenerationFinishReason::Eos);
EXPECT_GE(result.resources().after().cpuTimeMs(),
result.resources().before().cpuTimeMs());
EXPECT_GT(result.resources().after().memoryBytes(), 0u);
Teardown(session);
}
TEST_F(TextGenerationTest, ClearResetsHistory) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
EXPECT_TRUE(WaitForReady(session).isNothing())
<< "construction should load the model";
GenerateResponse first = Generate(session, StoryRequest());
ASSERT_EQ(first.type(), GenerateResponse::TGenerateResult);
session.mStreamed.Truncate();
GenerateResponse appended = Generate(session, StoryRequest());
ASSERT_EQ(appended.type(), GenerateResponse::TGenerateResult);
EXPECT_GT(appended.get_GenerateResult().usage().promptTokens(),
first.get_GenerateResult().usage().promptTokens())
<< "second Generate should prefill the accumulated history";
// Fire-and-forget; in-order delivery lands it before the next Generate.
EXPECT_TRUE(session.mGenerator->SendClear());
session.mStreamed.Truncate();
GenerateResponse cleared = Generate(session, StoryRequest());
ASSERT_EQ(cleared.type(), GenerateResponse::TGenerateResult);
EXPECT_EQ(cleared.get_GenerateResult().usage().promptTokens(),
first.get_GenerateResult().usage().promptTokens());
EXPECT_EQ(cleared.get_GenerateResult().usage().promptCharacters(),
first.get_GenerateResult().usage().promptCharacters());
#if defined(XP_MACOSX) && defined(__x86_64__)
EXPECT_FALSE(cleared.get_GenerateResult().content().IsEmpty());
#else
EXPECT_EQ(cleared.get_GenerateResult().content(),
first.get_GenerateResult().content())
<< "greedy generation after Clear should reproduce the first run";
#endif
Teardown(session);
}
TEST_F(TextGenerationTest, ZeroTokenResultCommitsHistory) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
ASSERT_TRUE(WaitForReady(session).isNothing());
GenerateResponse first = Generate(session, StoryRequest(1));
ASSERT_EQ(first.type(), GenerateResponse::TGenerateResult);
ASSERT_TRUE(session.mGenerator->SendClear());
GenerateResponse zero = Generate(session, StoryRequest(0));
ASSERT_EQ(zero.type(), GenerateResponse::TGenerateResult);
EXPECT_EQ(zero.get_GenerateResult().usage().generatedTokens(), 0u);
GenerateRequest next = StoryRequest(1);
next.messages().Clear();
GenerateResponse continued = Generate(session, next);
ASSERT_EQ(continued.type(), GenerateResponse::TGenerateResult);
EXPECT_EQ(continued.get_GenerateResult().usage().promptTokens(),
first.get_GenerateResult().usage().promptTokens());
EXPECT_EQ(continued.get_GenerateResult().usage().promptCharacters(),
first.get_GenerateResult().usage().promptCharacters());
EXPECT_GT(continued.get_GenerateResult().usage().generatedTokens(), 0u);
Teardown(session);
}
TEST_F(TextGenerationTest, CancelResolvesWithPartialResult) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
EXPECT_TRUE(WaitForReady(session).isNothing());
auto pending = StartGenerate(
session, StoryRequest(/* aMaxTokens */ 512, /* aBufferLength */ 1));
ASSERT_TRUE(SpinUntil("first delta", kStepTimeoutSeconds,
[&] { return !session.mStreamed.IsEmpty(); }));
EXPECT_TRUE(session.mGenerator->SendCancel());
auto result = WaitForOrTimeout(std::move(pending), kStepTimeoutSeconds);
ASSERT_TRUE(result.isSome())
<< "cancelled Generate never resolved";
ASSERT_TRUE(result->isOk());
GenerateResponse response = result->unwrap();
ASSERT_EQ(response.type(), GenerateResponse::TGenerateResult);
const GenerateResult& generated = response.get_GenerateResult();
EXPECT_EQ(generated.reason(), dom::TextGenerationFinishReason::Cancelled);
EXPECT_EQ(session.mStreamed, generated.content());
EXPECT_LT(generated.usage().generatedTokens(), 512u)
<< "cancel should have stopped generation early";
GenerateRequest next = StoryRequest(1);
next.messages().Clear();
GenerateResponse continued = Generate(session, next);
ASSERT_EQ(continued.type(), GenerateResponse::TGenerateResult);
EXPECT_EQ(continued.get_GenerateResult().usage().promptTokens(),
generated.usage().promptTokens());
EXPECT_EQ(continued.get_GenerateResult().usage().promptCharacters(),
generated.usage().promptCharacters());
Teardown(session);
}
TEST_F(TextGenerationTest, WhenReadyAnswersAfterTheLoadReported) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
ASSERT_TRUE(WaitForReady(session).isNothing())
<< "construction should load the model";
ASSERT_TRUE(WaitForReady(session).isNothing())
<< "a late WhenReady() must still see the successful load";
Teardown(session);
}
TEST_F(TextGenerationTest, InvalidFdFailsAtConstruction) {
Session session = CreateGenerator(ipc::FileDescriptor());
ASSERT_TRUE(session.mGenerator);
Maybe<TextGenerationParent::LoadFailure> loadError = WaitForReady(session);
ASSERT_TRUE(loadError.isSome());
EXPECT_EQ(loadError->cause,
TextGenerationParent::LoadFailure::Cause::Backend);
EXPECT_FALSE(loadError->message.IsEmpty());
GenerateResponse response =
Generate(session, StoryRequest(), /* aErrorExpected */ true);
ASSERT_EQ(response.type(), GenerateResponse::TGenerateError);
EXPECT_FALSE(response.get_GenerateError().message().IsEmpty());
Teardown(session);
}
TEST_F(TextGenerationTest, CloseBeforeReadyRetiresTheProcess) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
session.mGenerator->Close();
Maybe<TextGenerationParent::LoadFailure> loadError = WaitForReady(session);
ASSERT_TRUE(loadError.isSome());
EXPECT_EQ(loadError->cause,
TextGenerationParent::LoadFailure::Cause::ActorGone);
ExpectProcessGone("a generator closed before ready pinned the process");
}
TEST_F(TextGenerationTest, DISABLED_LaunchFailureEndsTheGenerator) {
auto manager = ipc::UtilityProcessManager::GetSingleton();
ASSERT_TRUE(manager);
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
// Nothing spins the event loop in between, so the launch is still pending
// and this fails it.
manager->CleanShutdownAllProcesses();
Maybe<TextGenerationParent::LoadFailure> loadError = WaitForReady(session);
ASSERT_TRUE(loadError.isSome())
<< "a generator on a failed launch must report it";
EXPECT_EQ(loadError->cause,
TextGenerationParent::LoadFailure::Cause::ActorGone);
EXPECT_FALSE(session.mGenerator->CanSend());
session.mGenerator = nullptr;
Session next = CreateGenerator();
ASSERT_TRUE(next.mGenerator);
EXPECT_FALSE(next.mGenerator->ProcessReused());
EXPECT_TRUE(WaitForReady(next).isNothing())
<< "a fresh launch must follow a failed one";
Teardown(next);
NS_ProcessPendingEvents(nullptr);
}
TEST_F(TextGenerationTest, ProcessDeathDuringGenerateRejects) {
Session session = CreateGenerator();
ASSERT_TRUE(session.mGenerator);
EXPECT_TRUE(WaitForReady(session).isNothing());
Maybe<base::ProcessId> pid = ProcessPid();
ASSERT_TRUE(pid.isSome());
auto pending = StartGenerate(
session, StoryRequest(/* aMaxTokens */ 4096, /* aBufferLength */ 1));
ASSERT_TRUE(SpinUntil("first delta", kStepTimeoutSeconds,
[&] { return !session.mStreamed.IsEmpty(); }));
Crash(*pid);
auto result = WaitForOrTimeout(std::move(pending), kStepTimeoutSeconds);
ASSERT_TRUE(result.isSome())
<< "Generate never settled after the crash";
EXPECT_TRUE(result->isErr()) << "Generate must reject with its process";
EXPECT_FALSE(session.mGenerator->CanSend());
Session next = CreateGenerator();
ASSERT_TRUE(next.mGenerator);
EXPECT_FALSE(next.mGenerator->ProcessReused());
EXPECT_TRUE(WaitForReady(next).isNothing());
Maybe<base::ProcessId> newPid = ProcessPid();
ASSERT_TRUE(newPid.isSome());
EXPECT_NE(*newPid, *pid);
session.mGenerator = nullptr;
EXPECT_EQ(ProcessPid(), newPid)
<< "a generator of the dead process must not touch its replacement";
Teardown(next);
}
TEST_F(TextGenerationTest, IdleGraceKeepsProcessForReuse) {
constexpr uint32_t kGraceMs = 2000;
Preferences::SetUint(kGracePref, kGraceMs);
Session first = CreateGenerator();
ASSERT_TRUE(first.mGenerator);
ASSERT_TRUE(WaitForReady(first).isNothing());
const auto pid = ProcessPid();
ASSERT_TRUE(pid.isSome());
first.mGenerator->Close();
Session second = CreateGenerator();
ASSERT_TRUE(second.mGenerator);
ASSERT_TRUE(WaitForReady(second).isNothing());
EXPECT_TRUE(second.mGenerator->ProcessReused());
EXPECT_EQ(ProcessPid(), pid);
const TimeStamp released = TimeStamp::Now();
second.mGenerator->Close();
ExpectProcessGone("process survived the last generator's idle grace");
EXPECT_GE((TimeStamp::Now() - released).ToMilliseconds(), kGraceMs * 0.8);
}