Skip to content

Commit ecbe023

Browse files
authored
Serialize the model instance registration in resource manager (#194)
* Serialize the model instance registration in resource manager * Fix the comment * Update copyright
1 parent e91834d commit ecbe023

File tree

2 files changed

+12
-2
lines changed

2 files changed

+12
-2
lines changed

src/rate_limiter.cc

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -72,8 +72,17 @@ RateLimiter::RegisterModelInstance(
7272
model_context.AddSpecificRequestQueue();
7373

7474
if (!ignore_resources_and_priority_) {
75+
// As there can be mulitple models being loaded concurrently, need
76+
// to hold a lock to protect the resource counts.
77+
// Without this serialization instances of other models might fail
78+
// to load because of the resource constraints in this instance.
79+
std::lock_guard<std::mutex> lk(resource_manager_mtx_);
7580
resource_manager_->AddModelInstance(model_instances.back().get());
76-
RETURN_IF_ERROR(resource_manager_->UpdateResourceLimits());
81+
const auto& status = resource_manager_->UpdateResourceLimits();
82+
if (!status.IsOk()) {
83+
resource_manager_->RemoveModelInstance(model_instances.back().get());
84+
return status;
85+
}
7786
}
7887
}
7988

src/rate_limiter.h

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
// Copyright 2020-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
1+
// Copyright 2020-2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
22
//
33
// Redistribution and use in source and binary forms, with or without
44
// modification, are permitted provided that the following conditions
@@ -314,6 +314,7 @@ class RateLimiter {
314314

315315
// Manager to keep track of the resource allocations
316316
std::unique_ptr<ResourceManager> resource_manager_;
317+
std::mutex resource_manager_mtx_;
317318

318319
// Mutex to serialize Payload [de]allocation
319320
std::mutex payload_mu_;

0 commit comments

Comments
 (0)