const uint32_t num_threads,

const std::string& threadpool_name) {

std::unique_ptr<Executor> executor;

if (num_threads == 0) {

executor.reset(new InlineExecutor());

} else {

executor.reset(new ThreadPoolExecutor(env, threadpool_name, num_threads));

}

return executor;

bool operator()(const ServableRequest& lhs,

const ServableRequest& rhs) const {

const int strcmp_result = lhs.name.compare(rhs.name);

if (strcmp_result != 0) {

return strcmp_result < 0;

}

DCHECK(lhs.version);

DCHECK(rhs.version);

return lhs.version.value() < rhs.version.value();

}

std::unique_ptr<BasicManager>* manager) {

manager->reset(new BasicManager(

options.env, options.num_load_threads, options.num_unload_threads,

options.max_num_load_retries, std::move(options.should_retry_model_load),

options.load_retry_interval_micros, options.flush_filesystem_caches,

std::move(options.resource_tracker), options.servable_event_bus,

std::move(options.pre_load_hook)));

return absl::OkStatus();

Env* const env, const uint32_t num_load_threads,

const uint32_t num_unload_threads, uint32_t max_num_load_retries,

std::function<bool(absl::Status)> should_retry_model_load,

int64_t load_retry_interval_micros, bool flush_filesystem_caches,

std::unique_ptr<ResourceTracker> resource_tracker,

EventBus<ServableState>* servable_event_bus,

std::function<void(const ServableId&)> pre_load_hook)

: servable_event_bus_(servable_event_bus),

should_retry_model_load_(std::move(should_retry_model_load)),

env_(env),

num_load_threads_(num_load_threads),

flush_filesystem_caches_(flush_filesystem_caches),

pre_load_hook_(std::move(pre_load_hook)) {

harness_options_.max_num_load_retries = max_num_load_retries;

harness_options_.load_retry_interval_micros = load_retry_interval_micros;

harness_options_.error_callback = [this](const ServableId& id,

ServableModelType model_type,

const absl::Status& error) {

PublishOnEventBus(

{id, ServableState::ManagerState::kEnd, error, model_type});

};

{

mutex_lock l(load_executor_mu_);

load_executor_ =

CreateExecutor(env_, num_load_threads, "BasicManager_Load_ThreadPool");

}

unload_executor_ = CreateExecutor(env_, num_unload_threads,

"BasicManager_Unload_ThreadPool");

resource_tracker_ = std::move(resource_tracker);

LOG(INFO) << "Unload all remaining servables in the manager.";

absl::Status status = absl::OkStatus();

{

mutex_lock l(mu_);

for (auto it = managed_map_.begin(); it != managed_map_.end(); ++it) {

LoaderHarness* const harness = it->second.get();

if (harness->state() == LoaderHarness::State::kReady) {

status.Update(harness->UnloadRequested());

status.Update(harness->StartQuiescing());

status.Update(harness->DoneQuiescing());

status.Update(harness->Unload());

}

if (harness->state() == LoaderHarness::State::kQuiescing) {

status.Update(harness->DoneQuiescing());

status.Update(harness->Unload());

}

if (harness->state() == LoaderHarness::State::kQuiesced) {

status.Update(harness->Unload());

}

}

}

return status;

const ServableRequest& request,

std::unique_ptr<UntypedServableHandle>* const untyped_handle) {

return serving_map_.GetUntypedServableHandle(request, untyped_handle);

// This blocks until the last handle given out by the old serving map is

// freed.

serving_map_.Update(managed_map_);

const ServableId& id) {

const auto range = managed_map_.equal_range(id.name);

for (auto iter = range.first; iter != range.second; ++iter) {

if (iter->second->id().version == id.version) {

return iter;

}

}

return managed_map_.end();

ServableData<std::unique_ptr<Loader>> servable,

std::function<std::shared_ptr<LoaderHarness>(const ServableId&,

std::unique_ptr<Loader>)>

harness_creator) {

VLOG(1) << "Request to start managing servable " << servable.id();

mutex_lock l(mu_);

const auto iter = BasicManager::FindHarnessInMap(servable.id());

if (iter != managed_map_.end()) {

return absl::FailedPreconditionError(

absl::StrCat("This servable is already being managed: ",

servable.id().DebugString()));

}

std::unique_ptr<Loader> loader;

if (servable.status().ok()) {

loader = servable.ConsumeDataOrDie();

}

std::shared_ptr<LoaderHarness> harness =

harness_creator(servable.id(), std::move(loader));

if (should_retry_model_load_) {

harness->set_should_retry(should_retry_model_load_);

}

if (!servable.status().ok()) {

harness->Error(servable.status());

} else {

PublishOnEventBus({harness->id(), ServableState::ManagerState::kStart,

harness->status(), harness->model_type()});

}

managed_map_.emplace(servable.id().name, harness);

return absl::OkStatus();

ServableData<std::unique_ptr<Loader>> servable) {

return ManageServableInternal(

std::move(servable),

[this](const ServableId& id, std::unique_ptr<Loader> loader) {

return std::make_shared<LoaderHarness>(id, std::move(loader),

harness_options_);

});

VLOG(1) << "Request to stop managing servable " << id;

mutex_lock l(mu_);

const auto it = FindHarnessInMap(id);

if (it == managed_map_.end()) {

LOG(ERROR) << "Request to delete harness for " << id

<< ", but no such harness found in managed_map_";

return absl::FailedPreconditionError(

absl::StrCat("This servable is not being managed: ", id.DebugString()));

}

const auto state = it->second->state();

if (state != LoaderHarness::State::kNew &&

state != LoaderHarness::State::kError &&

state != LoaderHarness::State::kDisabled) {

LOG(ERROR) << "Request to delete harness for " << id

<< ", but it is not in a new or end state. State: " << state;

return absl::FailedPreconditionError(absl::StrCat(

"This servable is not in a new or end state and we cannot stop "

"managing it: ",

id.DebugString(), " ", LoaderHarness::StateDebugString(state)));

}

managed_map_.erase(it);

return absl::OkStatus();

LoaderHarness** harness) {

// Look up the request servable's harness.

auto iter = FindHarnessInMap(id);

if (iter == managed_map_.end()) {

return absl::NotFoundError(

absl::StrCat("This servable is not being managed by the manager: ",

id.DebugString()));

}

TF_RETURN_IF_ERROR(iter->second->status());

*harness = iter->second.get();

return absl::OkStatus();

const {

std::vector<const Loader*> loaders;

for (const auto& entry : managed_map_) {

const LoaderHarness& harness = *entry.second;

bool uses_resources;

switch (harness.state()) {

case LoaderHarness::State::kNew:

uses_resources = false;

break;

case LoaderHarness::State::kLoadRequested:

uses_resources = false;

break;

case LoaderHarness::State::kLoadApproved:

uses_resources = true;

break;

case LoaderHarness::State::kLoading:

uses_resources = true;

break;

case LoaderHarness::State::kReady:

uses_resources = true;

break;

case LoaderHarness::State::kQuiescing:

uses_resources = true;

break;

case LoaderHarness::State::kQuiesced:

uses_resources = true;

break;

case LoaderHarness::State::kUnloadRequested:

uses_resources = true;

break;

case LoaderHarness::State::kUnloading:

uses_resources = true;

break;

case LoaderHarness::State::kDisabled:

uses_resources = false;

break;

case LoaderHarness::State::kError:

uses_resources = false;

break;

}

if (uses_resources) {

loaders.push_back(harness.loader());

}

}

return loaders;

mutex_lock l(mu_);

std::vector<std::string> servable_names;

for (auto iter = managed_map_.begin(); iter != managed_map_.end();

iter = managed_map_.equal_range(iter->first).second) {

servable_names.push_back(iter->first);

}

return servable_names;

PublishOnEventBus({harness->id(), ServableState::ManagerState::kLoading,

harness->status(), harness->model_type()});

// We save the id of the harness so that we can publish it after Load(). (We

// can't query harness again after Load() as it may be deleted by another

// thread that called StopManagingServable().)

const ServableId id = harness->id();

ServableModelType model_type = harness->model_type();

if (pre_load_hook_) {

pre_load_hook_(id);

}

// We don't hold the lock while calling Load() as it may block.

const absl::Status status = harness->Load();

// Whether the load succeeded or failed, flush filesystem caches if there is

// only one load thread.

if (flush_filesystem_caches_ && num_load_threads() <= 1) {

const absl::Status flush_status = Env::Default()->FlushFileSystemCaches();

if (!flush_status.ok()) {

LOG(WARNING) << "flushing filesystem caches failed: " << flush_status;

}

}

TF_RETURN_IF_ERROR(status);

{

mutex_lock l(mu_);

UpdateServingMap();

const auto it = FindHarnessInMap(id);

if (it != managed_map_.end()) {

// Re-capture the model type: for any model that implements `Servable`, it

// is only known after the model is loaded.

model_type = it->second->model_type();

}

}

PublishOnEventBus({id, ServableState::ManagerState::kAvailable,

absl::OkStatus(), model_type});

return absl::OkStatus();

const DoneCallback done_callback) {

VLOG(1) << "Request to load servable " << id;

LoadOrUnloadRequest request;

request.kind = LoadOrUnloadRequest::Kind::kLoad;

request.servable_id = id;

LoadOrUnloadServable(request, done_callback);

mutex_lock l(mu_);

LoaderHarness* harness;

const absl::Status status = GetHealthyHarness(id, &harness);

if (!status.ok()) {

return;

}

harness->set_should_retry([](absl::Status status) { return false; });

// We save the id of the harness so that we can publish it after Unload(). (We

// can't query harness again after Unload() as it may be deleted by another

// thread that called StopManagingServable().)

const ServableId id = harness->id();

const ServableModelType model_type = harness->model_type();

{

// StartQuiescing() would have been already called.

mutex_lock l(mu_);

PublishOnEventBus({id, ServableState::ManagerState::kUnloading,

harness->status(), model_type});

UpdateServingMap();

TF_RETURN_IF_ERROR(harness->DoneQuiescing());

}

// We don't hold the lock while calling Unload() as it may block.

TF_RETURN_IF_ERROR(harness->Unload());

PublishOnEventBus(

{id, ServableState::ManagerState::kEnd, absl::OkStatus(), model_type});

return absl::OkStatus();

const DoneCallback done_callback) {

VLOG(1) << "Request to unload servable " << id;

LoadOrUnloadRequest request;

request.kind = LoadOrUnloadRequest::Kind::kUnload;

request.servable_id = id;

LoadOrUnloadServable(request, done_callback);

const LoadOrUnloadRequest& request, LoaderHarness* harness) {

absl::Status execution_status;

switch (request.kind) {

case LoadOrUnloadRequest::Kind::kLoad:

execution_status = ExecuteLoad(harness);

break;

case LoadOrUnloadRequest::Kind::kUnload:

execution_status = ExecuteUnload(harness);

break;

}

{

mutex_lock l(mu_);

--num_ongoing_load_unload_executions_;

DCHECK_GE(num_ongoing_load_unload_executions_, 0);

num_ongoing_load_unload_executions_cv_.notify_all();

}

return execution_status;

// ThreadPoolExecutor destructor, implicitly calling ThreadPool destructor,

// waits for all scheduled work to finish. Should we wait within

// `load_executor_mu_` or outside?

//

// When changing `num_load_threads_` from M to N, the effective number of

// threads changes like this:

// - From M to 0 then to N, if destruct within lock; or

// - From M to (up to) M+N then to N, if destruct outside lock.

// The former is more intuitive when M and N are small (e.g. client intention

// is inline or single threaded loading), while the latter makes more sense

// when they are large.

const uint32_t old_num_threads = num_load_threads_.load();

if (old_num_threads < 2 || num_load_threads < 2) { // destruct within lock

mutex_lock l(load_executor_mu_);

load_executor_.reset();

num_load_threads_.store(num_load_threads);

load_executor_ =

CreateExecutor(env_, num_load_threads, "BasicManager_Load_ThreadPool");

} else { // destruct outside lock

std::unique_ptr<Executor> old_executor;

{

mutex_lock l(load_executor_mu_);

old_executor = std::move(load_executor_);

num_load_threads_.store(num_load_threads);

load_executor_ = CreateExecutor(env_, num_load_threads,

"BasicManager_Load_ThreadPool");

}

}

DoneCallback done_callback) {

const absl::Status status = [&]() {

mutex_lock l(mu_);

LoaderHarness* harness;

TF_RETURN_IF_ERROR(GetHealthyHarness(request.servable_id, &harness));

// Calling {Load,Unload}Request() synchronously here prevents any other

// concurrent calls to Load/UnloadServable() from proceeding.

switch (request.kind) {

case LoadOrUnloadRequest::Kind::kLoad:

TF_RETURN_IF_ERROR(harness->LoadRequested());

break;

case LoadOrUnloadRequest::Kind::kUnload:

TF_RETURN_IF_ERROR(harness->UnloadRequested());

break;

}

return absl::OkStatus();

}();

if (!status.ok()) {

done_callback(status);

return;

}

switch (request.kind) {

case LoadOrUnloadRequest::Kind::kLoad: {

mutex_lock l(load_executor_mu_);

load_executor_->Schedule([this, request, done_callback]() {

HandleLoadOrUnloadRequest(request, done_callback);

});

break;

}

case LoadOrUnloadRequest::Kind::kUnload: {

unload_executor_->Schedule([this, request, done_callback]() {

HandleLoadOrUnloadRequest(request, done_callback);

});

break;

}

}

DoneCallback done_callback) {

// Decision phase.

absl::Status decision_status;

LoaderHarness* harness;

{

// We serialize the decision phases of the requests. We will make a decision

// about the present request before allowing other requests to enter their

// decision phase. See the .h file for more explanation and rationale.

mutex_lock l(load_unload_decision_phase_mu_);

decision_status = ApproveLoadOrUnload(request, &harness);

}

if (!decision_status.ok()) {

done_callback(decision_status);

return;

}

// Execution phase.

const absl::Status execution_status = ExecuteLoadOrUnload(request, harness);

done_callback(execution_status);

const LoadOrUnloadRequest& request, LoaderHarness** harness) {

mutex_lock l(mu_);

TF_RETURN_IF_ERROR(GetHealthyHarness(request.servable_id, harness));

switch (request.kind) {

case LoadOrUnloadRequest::Kind::kLoad: {

TF_RETURN_IF_ERROR(ApproveLoad(*harness, &l));

break;

}

case LoadOrUnloadRequest::Kind::kUnload: {

TF_RETURN_IF_ERROR(ApproveUnload(*harness));

break;

}

}

++num_ongoing_load_unload_executions_;

return absl::OkStatus();

mutex_lock* mu_lock) {

if (resource_tracker_ != nullptr) {

// Attempt to reserve resources for the load.

const absl::Status resource_reservation_status =

ReserveResources(harness, mu_lock);

if (!resource_reservation_status.ok()) {

LOG(WARNING) << resource_reservation_status;

harness->Error(resource_reservation_status);

PublishOnEventBus({harness->id(), ServableState::ManagerState::kEnd,

resource_reservation_status, harness->model_type()});

return resource_reservation_status;

}

}

// Transition to state kLoadApproved inside the decision phase. We rely

// on this state to know whether resources have been reserved in

// GetLoadersCurrentlyUsingResources().

TF_RETURN_IF_ERROR(harness->LoadApproved());

return absl::OkStatus();

// Transition to state kQuiescing inside the decision phase, to prevent any

// concurrent unload requests from executing.

TF_RETURN_IF_ERROR(harness->StartQuiescing());

return absl::OkStatus();

mutex_lock* mu_lock) {

while (true) {

TF_RETURN_IF_ERROR(resource_tracker_->RecomputeUsedResources(

GetLoadersCurrentlyUsingResources()));

bool resources_reserved;

// We retry reserving resources because it may involve transiently failing

// operations like file-reads.

const absl::Status reserve_resources_status = Retry(

absl::StrCat("Reserving resources for servable: ",

harness->id().DebugString()),

harness_options_.max_num_load_retries,

harness_options_.load_retry_interval_micros,

[&]() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {

return resource_tracker_->ReserveResources(*harness->loader(),

&resources_reserved);

},

[&](absl::Status status) { return harness->should_retry(status); });

if (!reserve_resources_status.ok()) {

return absl::Status(

reserve_resources_status.code(),

absl::StrCat(

"Error while attempting to reserve resources to load servable ",

harness->id().DebugString(), ": ",

reserve_resources_status.message()));

}

if (resources_reserved) {

// Woohoo! We got our resources.

LOG(INFO) << "Successfully reserved resources to load servable "

<< harness->id().DebugString();

return absl::OkStatus();

}

// We weren't able to reserve the resources. See if there are any

// ongoing load/unload executions that may be temporarily tying up

// resources.

if (num_ongoing_load_unload_executions_ == 0) {

// There are no ongoing load/unloads, so we really are out of

// resources for this servable.

return absl::ResourceExhaustedError(

absl::StrCat("Insufficient resources to load servable ",

harness->id().DebugString()));

} else {

// Wait until at least one load/unload request finishes, then retry.

VLOG(1) << "Waiting for another load/unload request to finish";

num_ongoing_load_unload_executions_cv_.wait(*mu_lock);

}

}

if (servable_event_bus_ != nullptr) {

servable_event_bus_->Publish(state);

}