Method Compute

tensorflow/core/kernels/work_queue_ops.cc:340–369

  }

  void Compute(OpKernelContext* ctx) override {
    const Tensor* input_;
    OP_REQUIRES_OK(ctx, ctx->input("work", &input_));
    const string& work = input_->scalar<string>()();

    Env* const env = Env::Default();
    std::vector<string> work_parts = str_util::Split(work, "?");

    OP_REQUIRES(ctx, work_parts.size() == 2,
                errors::InvalidArgument("Invalid table path format: ", work));
    OP_REQUIRES_OK(ctx, env->IsDirectory(restore_works_dir_));

    uint64 start_time = env->NowMicros();
    string range = str_util::StringReplace(work_parts[1], "=", "_", true);
    range = str_util::StringReplace(range, "&", "_", true);
    string worker_id = strings::StrCat(job_name_, "_", task_index_);
    string worker_dir = strings::StrCat(restore_works_dir_, "/", worker_id);
    if (!env->IsDirectory(worker_dir).ok()) {
      OP_REQUIRES_OK(ctx, env->CreateDir(worker_dir));
    }
    string slice_file = strings::StrCat(worker_dir, "/", range);
    std::unique_ptr<WritableFile> wfile;
    OP_REQUIRES_OK(ctx, env->NewWritableFile(slice_file, &wfile));
    OP_REQUIRES_OK(ctx, wfile->Append(work));
    OP_REQUIRES_OK(ctx, wfile->Close());
    uint64 end_time = env->NowMicros();
    float time_use = (end_time - start_time) / 1000000.0;
    LOG(INFO) << "Job_name:" << job_name_ << ", time use:" << time_use
              << ", task_index:" << task_index_ << ", work_range:" << range;
  }

 private:
  string job_name_;