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

Added back persistent memory tracking in queue op. The new tracking logic has avoided the crash in previous implementation: the queue_ passed to CreateTypedQueue may be unreffed if the resource is already created by another resource op that shares the same resource name and type.

Browse Source 
PiperOrigin-RevId: 168284509
This commit is contained in:
Yuefeng Zhou 2017-09-11 14:21:15 -07:00 committed by TensorFlower Gardener
parent ed11359944
commit 08587d45b4
1 changed files with 8 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tensorflow/core/kernels/queue_op.h
View File

@ -41,6 +41,14 @@ class QueueOp : public ResourceOpKernel<QueueInterface> {
context->GetAttr("component_types", &component_types_)); context->GetAttr("component_types", &component_types_));
} }
void Compute(OpKernelContext* context) override {
ResourceOpKernel<QueueInterface>::Compute(context);
if (resource_ && context->track_allocations()) {
context->record_host_persistent_memory_allocation(
resource_->MemoryUsed());
}
}
protected: protected:
// Variables accessible by subclasses // Variables accessible by subclasses
int32 capacity_; int32 capacity_;
@ -63,12 +71,8 @@ class TypedQueueOp : public QueueOp {
return errors::ResourceExhausted("Failed to allocate queue."); return errors::ResourceExhausted("Failed to allocate queue.");
} }
*ret = queue; *ret = queue;
queue_ = queue;
return queue->Initialize(); return queue->Initialize();
} }
private:
QueueInterface* queue_ = nullptr;
}; };
} // namespace tensorflow } // namespace tensorflow