message: [tf.data] Add vlogging for time spent restoring each stage of a dataset iterator.
diff: @@ -665,7 +665,11 @@ class IteratorBase {
  // Saves the state of this iterator.
  virtual Status Save(SerializationContext* ctx, IteratorStateWriter* writer) {
  - return SaveInternal(ctx, writer);
  + int64 start_us = EnvTime::NowMicros();
  + TF_RETURN_IF_ERROR(SaveInternal(ctx, writer));
  + VLOG(1) << "Saved " << prefix() << " in "
  + << (EnvTime::NowMicros() - start_us) << "us";
  + return Status::OK();
  }
  protected:
  @@ -675,21 +679,33 @@ class IteratorBase {
  // Restores the state of this iterator.
  Status Restore(IteratorContext* ctx, IteratorStateReader* reader) {
  - return RestoreInternal(ctx, reader);
  + int64 start_us = EnvTime::NowMicros();
  + TF_RETURN_IF_ERROR(RestoreInternal(ctx, reader));
  + VLOG(1) << "Restored " << prefix() << " in "
  + << (EnvTime::NowMicros() - start_us) << "us";
  + return Status::OK();
  }
  // This is needed so that sub-classes of IteratorBase can call
  // `SaveInternal` on their input iterators.
  Status SaveInput(SerializationContext* ctx, IteratorStateWriter* writer,
  const std::unique_ptr<IteratorBase>& input) {
  - return input->SaveInternal(ctx, writer);
  + int64 start_us = EnvTime::NowMicros();
  + TF_RETURN_IF_ERROR(input->SaveInternal(ctx, writer));
  + VLOG(2) << "Saved " << input->prefix() << " in "
  + << (EnvTime::NowMicros() - start_us) << "us";
  + return Status::OK();
  }
  // This is needed so that sub-classes of IteratorBase can call
  // `RestoreInternal` on their input iterators.
  Status RestoreInput(IteratorContext* ctx, IteratorStateReader* reader,
  const std::unique_ptr<IteratorBase>& input) {
  - return input->RestoreInternal(ctx, reader);
  + int64 start_us = EnvTime::NowMicros();
  + TF_RETURN_IF_ERROR(input->RestoreInternal(ctx, reader));
  + VLOG(2) << "Restored " << input->prefix() << " in "
  + << (EnvTime::NowMicros() - start_us) << "us";
  + return Status::OK();
  }
  // Saves the state of this iterator.
  