Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[tf.data] Set limit on number of threads used in threadpool_dataset. #53895

Merged
merged 1 commit into from Jan 24, 2022
Merged

[tf.data] Set limit on number of threads used in threadpool_dataset. #53895

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
26 changes: 19 additions & 7 deletions tensorflow/core/kernels/data/experimental/threadpool_dataset_op.cc
View file
Open in desktop
Expand Up @@ -39,6 +39,22 @@ namespace experimental {
PrivateThreadPoolDatasetOp::kDatasetType;
/* static */ constexpr const char* const PrivateThreadPoolDatasetOp::kDatasetOp;

namespace {
// To prevent integer overflow issues when allocating threadpool memory for an
// unreasonable number of threads.
constexpr int kThreadLimit = 65536;

Status ValidateNumThreads(int32_t num_threads) {
if (num_threads < 0) {
return errors::InvalidArgument("`num_threads` must be >= 0");
}
if (num_threads >= kThreadLimit) {
return errors::InvalidArgument("`num_threads` must be < ", kThreadLimit);
}
return Status::OK();
}
} // namespace

class ThreadPoolResource : public ResourceBase {
public:
ThreadPoolResource(Env* env, const ThreadOptions& thread_options,
Expand Down Expand Up @@ -83,9 +99,7 @@ class ThreadPoolHandleOp : public OpKernel {
OP_REQUIRES_OK(ctx, ctx->GetAttr("num_threads", &num_threads_));
OP_REQUIRES_OK(ctx, ctx->GetAttr("max_intra_op_parallelism",
&max_intra_op_parallelism_));
OP_REQUIRES(
ctx, num_threads_ > 0,
errors::InvalidArgument("`num_threads` must be greater than zero."));
OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads_));
}

// The resource is deleted from the resource manager only when it is private
Expand Down Expand Up @@ -529,8 +543,7 @@ void PrivateThreadPoolDatasetOp::MakeDatasetFromOptions(OpKernelContext* ctx,
DatasetBase* input,
int32_t num_threads,
DatasetBase** output) {
OP_REQUIRES(ctx, num_threads >= 0,
errors::InvalidArgument("`num_threads` must be >= 0"));
OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));
*output = new Dataset(ctx,
DatasetContext(DatasetContext::Params(
{PrivateThreadPoolDatasetOp::kDatasetType,
Expand All @@ -544,8 +557,7 @@ void PrivateThreadPoolDatasetOp::MakeDataset(OpKernelContext* ctx,
int64_t num_threads = 0;
OP_REQUIRES_OK(
ctx, ParseScalarArgument<int64_t>(ctx, "num_threads", &num_threads));
OP_REQUIRES(ctx, num_threads >= 0,
errors::InvalidArgument("`num_threads` must be >= 0"));
OP_REQUIRES_OK(ctx, ValidateNumThreads(num_threads));
*output = new Dataset(ctx, input, num_threads);
}

Expand Down