diff --git a/edge0-infer/src/weights/runtime.cpp b/edge0-infer/src/weights/runtime.cpp index e672a341..0d799da3 100644 --- a/edge0-infer/src/weights/runtime.cpp +++ b/edge0-infer/src/weights/runtime.cpp @@ -309,9 +309,11 @@ exec::Status Weights::promote(const std::shared_ptr& write) { bool Weights::ready(const std::shared_ptr& write) const { if(!write) return true; std::lock_guard operation(write->operation_); - return std::all_of(write->futures_.begin(),write->futures_.end(),[](const auto& f) { - return !f.valid() || f.wait_for(std::chrono::seconds(0))==std::future_status::ready; - }); + const auto ready=[](const auto& future) { + return !future.valid() || future.wait_for(std::chrono::seconds(0))==std::future_status::ready; + }; + return std::all_of(write->futures_.begin(),write->futures_.end(),ready) && + std::all_of(write->writers_.begin(),write->writers_.end(),ready); } Result> Weights::complete(const std::shared_ptr& write) { if(auto error=promote(write)) return *error; diff --git a/edge0-infer/tests/unit/weights/lifecycle_test.cpp b/edge0-infer/tests/unit/weights/lifecycle_test.cpp index dfe37788..78733f59 100644 --- a/edge0-infer/tests/unit/weights/lifecycle_test.cpp +++ b/edge0-infer/tests/unit/weights/lifecycle_test.cpp @@ -50,8 +50,10 @@ int main(int argc,char** argv) { auto first=take(w.writable_regions(tx,0,0,8)); CHECK(!w.writable_regions(tx,0,4,8)); CHECK(!w.writable_regions(tx,0,16,1)); auto last=take(w.writable_regions(tx,0,8,8)); + CHECK(!w.ready(tx)); CHECK(first->write(7,data,2)); CHECK(!first->write(0,data,8)); CHECK(!first->finish(Error{ErrorCategory::backend_error,"injected"})); + CHECK(!w.ready(tx)); auto completion=std::async(std::launch::async,[&] { return w.complete(tx); }); CHECK(completion.wait_for(20ms)==std::future_status::timeout); CHECK(!last->write(0,data+2,8)); CHECK(!last->finish()); CHECK(!completion.get()); @@ -72,6 +74,7 @@ int main(int argc,char** argv) { auto target=take(w.create_storage(backend,{{{GGML_TYPE_F32,{4}},1}}, edge0::exec::Purpose::resident_weights)); tx=take(w.begin_write(target,{item(0,"custom")})); auto writer=take(w.writable_regions(tx,0,0,16)); CHECK(!writer->write(0,data,16)); CHECK(!writer->finish()); + CHECK(w.ready(tx)); auto binding=take(w.complete(tx)); CHECK(binding[0].generation==1); view=take(w.acquire_views(target,binding[0])); auto closing=std::async(std::launch::async,[&] { return w.close(target); });