experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Fix iterator bug introduced in recent change to ASBS

Browse Source 
PiperOrigin-RevId: 331183329
Change-Id: Ie1aa494d32983b33b670712c7cdb47b2e3ff2a3f
This commit is contained in:
A. Unique TensorFlower 2020-09-11 11:07:23 -07:00 committed by TensorFlower Gardener
parent 721a91c7e9
commit 77aba03c2a
2 changed files with 16 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler.h
View File

@ -464,7 +464,7 @@ void AdaptiveSharedBatchScheduler<
it != batches_.end() && available_threads > 0;) {
if ((*it)->IsClosed()) {
const internal::ASBSBatch<TaskType>* batch = *it;
batches_.erase(it++);
it = batches_.erase(it);
batch->queue()->ReleaseBatch(batch);
batch_thread_pool_->Schedule(
std::bind(&AdaptiveSharedBatchScheduler<TaskType>::CallbackWrapper,

15
tensorflow/core/kernels/batching_util/adaptive_shared_batch_scheduler_test.cc
View File

@ -353,6 +353,21 @@ TEST(AdaptiveSharedBatchSchedulerTest, QueueCapacityInfo) {
EXPECT_EQ(queue->SchedulingCapacity(), 9 * 1000 + 300);
finish_processing.Notify();
}
TEST(AdaptiveSharedBatchSchedulerTest, FullBatches) {
std::shared_ptr<AdaptiveSharedBatchScheduler<FakeTask>> scheduler;
TF_ASSERT_OK(AdaptiveSharedBatchScheduler<FakeTask>::Create({}, &scheduler));
auto queue_callback = [](std::unique_ptr<Batch<FakeTask>> batch) {
ASSERT_TRUE(batch->IsClosed());
};
AdaptiveSharedBatchScheduler<FakeTask>::QueueOptions queue_options;
queue_options.max_batch_size = 100;
queue_options.batch_timeout_micros = 1000000000000;
std::unique_ptr<BatchScheduler<FakeTask>> queue;
TF_ASSERT_OK(scheduler->AddQueue(queue_options, queue_callback, &queue));
TF_ASSERT_OK(ScheduleTask(100, queue.get()));
// Full batches should not have to wait batch_timeout_micros.
}
} // namespace anonymous
} // namespace serving
} // namespace tensorflow