tensorflow之ResourceMgr

ResourceMgr,就是资源管理器。核心功能

1. 创建资源
2. 查找资源
3. 删除资源

tensorflow的ResourceMgr内部如何组织资源？从其创建和查找接口中可以看出：ResourceMgr里有多个container,每个container里又有多个。可以用C++ std::map来表示。 Resource可以是任意类型。

template  
Status Create(const std::string& container, const std::string& name,
               T* resource) TF_MUST_USE_RESULT;
template 
Status Lookup(const std::string& container, const std::string& name,
               T** resource) const TF_MUST_USE_RESULT;

ResourceBase

• 有引用计数
• 返回内存用量

// Forward declaration to avoid introducing a dependency on headers in
// "tensorflow/core/graph/...".
class GraphDefBuilder;
class Node;

// This is the base class of all resource classes. Each resource must be
// represented as a sub-class of ResourceBase (which is reference counted) to be
// able to work with resource facilities such ResourceHandle and ResourceMgr.
class ResourceBase : public core::WeakRefCounted {
 public:
  // Returns a debug string for *this.
  virtual std::string DebugString() const = 0;

  // Returns memory used by this resource.
  virtual int64_t MemoryUsed() const { return 0; }

  // Writes a representation of this resource into `builder`, so that executing
  // `*out` will recreate this resource.
  virtual Status AsGraphDef(GraphDefBuilder* builder, Node** out) const {
    return errors::Unimplemented("AsGraphDef not implemented for resource ",
                                 DebugString());
  }
};

ResourceHandler

1.  device: 设备名：设备拥有资源(内存)
2. container: 资源所有容器
3. name: 资源唯一名字
4. has_code: 资源的hash_code, 同一个设备同一个execution里有效

// Protocol buffer representing a handle to a tensorflow resource. Handles are
// not valid across executions, but can be serialized back and forth from within
// a single run.
message ResourceHandleProto {
  // Unique name for the device containing the resource.
  string device = 1;

  // Container in which this resource is placed.
  string container = 2;

  // Unique name of this resource.
  string name = 3;

  // Hash code for the type of the resource. Is only valid in the same device
  // and in the same execution.
  uint64 hash_code = 4;

  // For debug-only, the name of the type pointed to by this handle, if
  // available.
  string maybe_type_name = 5;

  // Protocol buffer representing a pair of (data type, tensor shape).
  message DtypeAndShape {
    DataType dtype = 1;
    TensorShapeProto shape = 2;
  }

  // Data types and shapes for the underlying resource.
  repeated DtypeAndShape dtypes_and_shapes = 6;

  reserved 7;
}

tensorflow/core/framework/resource_handle.h里又定义了ResourceHandle类。目的是能让tensorflow不依赖protobuf

//handle里存放了资源的device, container, name,而且Own了资源
class ResourceHandle {
 public:
  ResourceHandle();
  ResourceHandle(const ResourceHandleProto& proto);
  ~ResourceHandle();

  // Use this factory method if the `proto` comes from user controlled input, to
  // prevent a denial of service.
  static Status BuildResourceHandle(const ResourceHandleProto& proto,
                                    ResourceHandle* out);

  // Unique name for the device containing the resource.
  const std::string& device() const { return device_; }

  void set_device(const std::string& device) { device_ = device; }

  // Container in which this resource is placed.
  const std::string& container() const { return container_; }
  void set_container(const std::string& container) { container_ = container; }

  // Unique name of this resource.
  const std::string& name() const { return name_; }
  void set_name(const std::string& name) { name_ = name; }

  // Hash code for the type of the resource. Is only valid in the same device
  // and in the same execution.
  uint64 hash_code() const { return hash_code_; }
  void set_hash_code(uint64 hash_code) { hash_code_ = hash_code; }

  // For debug-only, the name of the type pointed to by this handle, if
  // available.
  const std::string& maybe_type_name() const { return maybe_type_name_; }
  void set_maybe_type_name(const std::string& value) {
    maybe_type_name_ = value;
  }

  // Data types and shapes for the underlying resource.
  std::vector dtypes_and_shapes() const {
    return dtypes_and_shapes_;
  }
  void set_dtypes_and_shapes(
      const std::vector& dtypes_and_shapes) {
    dtypes_and_shapes_ = dtypes_and_shapes;
  }

  void set_definition_stack_trace(
      const absl::optional& definition_stack_trace) {
    definition_stack_trace_ = definition_stack_trace;
  }

  const absl::optional& definition_stack_trace() const {
    return definition_stack_trace_;
  }

  // Conversion to and from ResourceHandleProto
  void AsProto(ResourceHandleProto* proto) const;
  Status FromProto(const ResourceHandleProto& proto);

  // Serialization via ResourceHandleProto
  std::string SerializeAsString() const;
  bool ParseFromString(const std::string& s);

  std::string DebugString() const;

  std::string SummarizeValue() const;

  // GUID for anonymous resources. Resources with this shared_name will have
  // their shared_name replaced with a GUID at creation time
  static constexpr const char* ANONYMOUS_NAME =
      "cd2c89b7-88b7-44c8-ad83-06c2a9158347";
  //函数会占有资源的强引用。
  template 
  static ResourceHandle MakeRefCountingHandle(
      T* resource, const string& device_name,
      const std::vector& dtypes_and_shapes = {},
      const absl::optional& definition_stack_trace = {}) {
    return MakeRefCountingHandle(resource, device_name, TypeIndex::Make(),
                                 dtypes_and_shapes, definition_stack_trace);
  }

  static ResourceHandle MakeRefCountingHandle(
      ResourceBase* resource, const string& device_name,
      const TypeIndex& type_index,
      const std::vector& dtypes_and_shapes = {},
      const absl::optional& definition_stack_trace = {});

  //指向资源的指针
  const core::IntrusivePtr& resource() const { return resource_; }

  // Gets the resource pointer in `handle` as `T*`, or an error if the actual
  // resource type is not `T`.
  template 
  StatusOr GetResource() const {
    TF_RETURN_IF_ERROR(ValidateType());
    return down_cast(resource_.get());
  }

  // Returns True if the resource handle is ref-counting.
  // See MakeRefCountingHandle.
  bool IsRefCounting() const { return resource_.get() != nullptr; }

  // Validates that the resource type in `handle` is `T`.
  template 
  Status ValidateType() const {
    return ValidateType(TypeIndex::Make());
  }

  Status ValidateType(const TypeIndex& type_index) const;

  // Generates unique IDs (e.g. for names of anonymous variables)
  static int64_t GenerateUniqueId();

 private:
  std::string device_;
  std::string container_;
  std::string name_;
  uint64 hash_code_ = 0;
  std::string maybe_type_name_;
  std::vector dtypes_and_shapes_;
  absl::optional definition_stack_trace_;
  // A smart pointer to the actual resource. When this field is not empty, the
  // handle is in a "ref-counting" mode, owning the resource; otherwise it's in
  // a "weak-ref" mode, only containing the name of the resource (conceptually a
  // weak reference).
  core::IntrusivePtr resource_;
  static std::atomic current_id_;
};

 ResourceManager

单纯看.h的接口，只能有个逻辑抽象。具体还得看.cc中的实现。

template 
Status ResourceMgr::Create(const std::string& container,
                           const std::string& name, T* resource) {
  CheckDeriveFromResourceBase();
  CHECK(resource != nullptr);
  mutex_lock l(mu_);
  //加锁后创建，而且拿到了类型的唯一ID. 比如C++中RTTI type info的hash值作为ID
  return DoCreate(container, TypeIndex::Make(), name, resource,
                  /* owns_resource */ true);
}

DoCreate真正创建

//Container是定义在.h中的一个flat_hash_map
  typedef absl::flat_hash_map
      Container;
//containers_ 是.h中ResourceMgr的成员变量
absl::flat_hash_map containers_ TF_GUARDED_BY(mu_);

Status ResourceMgr::DoCreate(const string& container_name, TypeIndex type,
                             const string& name, ResourceBase* resource,
                             bool owns_resource) {
  //拿到container,没有就创建
  Container* container = [&]() TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
    Container** ptr = &containers_[container_name];
    if (*ptr == nullptr) {
      *ptr = new Container;
    }
    return *ptr;
  }();

  // NOTE: Separating out the construction of the map key and value so that the
  // key can contain a StringPiece that borrows from the string in the value.
  //插入
  ResourceAndName resource_and_name(name);

  StringPiece borrowed_name(*resource_and_name.name);
  //是否占有资源
  if (owns_resource) { 
    resource_and_name.resource = core::RefCountPtr(resource);
  } else {
    auto cleanup_fn = [this, container, type, borrowed_name]() {
      mutex_lock l(mu_);
      auto iter = container->find({type.hash_code(), borrowed_name});
      if (iter != container->end()) {
        container->erase(iter);
      }
    };
    resource_and_name.resource =
        core::WeakPtr(resource, cleanup_fn);
  }

  Container::value_type key_and_value(Key(type.hash_code(), borrowed_name),
                                      std::move(resource_and_name));

  auto st = container->insert(std::move(key_and_value));
  if (st.second) {
    TF_RETURN_IF_ERROR(InsertDebugTypeName(type.hash_code(), type.name()));
    return Status::OK();
  }
  return errors::AlreadyExists("Resource ", container_name, "/", name, "/",
                               type.name());
}

class ResourceMgr {
 public:
  ResourceMgr();
  explicit ResourceMgr(const std::string& default_container);
  ~ResourceMgr();

  // Returns the default container name for *this.
  const std::string& default_container() const { return default_container_; }

  // Creates a resource "name" in the "container".  The caller transfers
  // the ownership of one ref on "resource" to *this, regardless of whether this
  // operation succeeds or fails.
  //
  // REQUIRES: std::is_base_of
  // REQUIRES: resource != nullptr.
  template 
  Status Create(const std::string& container, const std::string& name,
                T* resource) TF_MUST_USE_RESULT;

  // Creates a unowned resource "name" in the "container".  The caller does NOT
  // transfer the ownership of any ref on "resource" to *this, regardless of
  // whether this operation succeeds or fails.
  //
  // After the resource is destroyed, lookups from the manager fail.
  // The caller must call this->Delete() on the name to free up the memory
  // entry of the name.
  //
  // REQUIRES: std::is_base_of
  // REQUIRES: resource != nullptr.
  template 
  Status CreateUnowned(const std::string& container, const std::string& name,
                       T* resource) TF_MUST_USE_RESULT;

  // If "container" has a resource "name", returns it in "*resource" and
  // the caller takes the ownership of one ref on "*resource".
  //
  // REQUIRES: std::is_base_of
  // REQUIRES: resource != nullptr
  template 
  Status Lookup(const std::string& container, const std::string& name,
                T** resource) const TF_MUST_USE_RESULT;

  // If the resource manager has a resource matching "handle", returns it in
  // "*resource" and the caller takes the ownership of one ref on "*resource".
  //
  // REQUIRES: resource != nullptr
  Status Lookup(const ResourceHandle& handle,
                ResourceBase** resource) const TF_MUST_USE_RESULT;

  // Similar to Lookup, but looks up multiple resources at once, with only a
  // single lock acquisition.  If containers_and_names[i] is uninitialized
  // then this function does not modify resources[i].
  template 
  Status LookupMany(absl::Span const>
                        containers_and_names,
                    std::vector>*
                        resources) const TF_MUST_USE_RESULT;

  // If "container" has a resource "name", returns it in
  // "*resource". Otherwise, invokes creator() to create the resource.
  // The caller takes the ownership of one ref on "*resource".
  //
  // WARNING: creator() must not call any methods on ResourceMgr during its
  // execution, because a non-reentrant lock is held during the creator() call
  // in order to guarantee atomicity of LookupOrCreate().
  //
  // REQUIRES: std::is_base_of
  // REQUIRES: resource != nullptr
  template 
  Status LookupOrCreate(const std::string& container, const std::string& name,
                        T** resource,
                        std::function creator) TF_MUST_USE_RESULT;

  // Deletes the resource "name" from the "container".
  //
  // REQUIRES: std::is_base_of
  template 
  Status Delete(const std::string& container,
                const std::string& name) TF_MUST_USE_RESULT;

  // Deletes the resource pointed by "handle".
  Status Delete(const ResourceHandle& handle) TF_MUST_USE_RESULT;

  // Deletes all resources from the "container" and removes the container.
  Status Cleanup(const std::string& container) TF_MUST_USE_RESULT;

  // Deletes all resources in all containers.
  void Clear();

  // Returns a text description for all resources.
  std::string DebugString() const;

 private:
  typedef std::pair Key;
  struct KeyHash {
    std::size_t operator()(const Key& k) const {
      return Hash64(k.second.data(), k.second.size(), k.first);
    }
  };
  struct KeyEqual {
    bool operator()(const Key& x, const Key& y) const {
      return (x.second == y.second) && (x.first == y.first);
    }
  };
  struct ResourceAndName {
    absl::variant, core::WeakPtr>
        resource;
    std::unique_ptr name;

    ResourceAndName();
    explicit ResourceAndName(const string& name);
    ResourceAndName(ResourceAndName&& other) noexcept;
    ~ResourceAndName();

    ResourceAndName& operator=(ResourceAndName&&) noexcept;

    // Returns a strong reference to resource, or nullptr if the resource is
    // no longer valid.
    core::RefCountPtr GetResource() const;

   private:
    TF_DISALLOW_COPY_AND_ASSIGN(ResourceAndName);
  };
  typedef absl::flat_hash_map
      Container;

  const std::string default_container_;
  mutable mutex mu_;
  absl::flat_hash_map containers_ TF_GUARDED_BY(mu_);

  template 
  Status LookupInternal(const std::string& container, const std::string& name,
                        T** resource) const
      TF_SHARED_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;
  Status LookupInternal(const std::string& container, uint64 type_hash_code,
                        const std::string& name, ResourceBase** resource) const
      TF_SHARED_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;

  Status DoCreate(const std::string& container, TypeIndex type,
                  const std::string& name, ResourceBase* resource,
                  bool owns_resource)
      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;

  Status DoLookup(const std::string& container, TypeIndex type,
                  const std::string& name, ResourceBase** resource) const
      TF_SHARED_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;
  Status DoLookup(const std::string& container, uint64 type_hash_code,
                  const std::string& type_name,
                  const std::string& resource_name,
                  ResourceBase** resource) const
      TF_SHARED_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;

  Status DoDelete(const std::string& container, uint64 type_hash_code,
                  const std::string& resource_name,
                  const std::string& type_name) TF_MUST_USE_RESULT;
  Status DoDelete(const std::string& container, TypeIndex type,
                  const std::string& resource_name) TF_MUST_USE_RESULT;

  // Pops the ResourceAndName entry. The entry is moved from the list to
  // the output argument `resource_and_name`.
  Status PopResourceAndName(
      const std::string& container, uint64 type_hash_code,
      const std::string& resource_name, const std::string& type_name,
      ResourceAndName& resource_and_name) TF_MUST_USE_RESULT;
  // Inserts the type name for 'hash_code' into the hash_code to type name map.
  Status InsertDebugTypeName(uint64 hash_code, const std::string& type_name)
      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_) TF_MUST_USE_RESULT;

  // Returns the type name for the 'hash_code'.
  // Returns "" if a resource with such a type was never inserted into
  // the container.
  const char* DebugTypeName(uint64 hash_code) const
      TF_EXCLUSIVE_LOCKS_REQUIRED(mu_);

  // Map from type hash_code to type name.
  std::unordered_map debug_type_names_ TF_GUARDED_BY(mu_);

  TF_DISALLOW_COPY_AND_ASSIGN(ResourceMgr);
};