Skip to content

Commit

Permalink
Add additional ts benchmarks which read/write multiple tensorstores c…
Browse files Browse the repository at this point in the history 
…oncurrently.

Decompose a few functions so that profile cost is more visible.

* multi_genspec generates arrays for the multi tensorstore test.
* multi_read_benchmark reads multiple tensorstores from a "base" tensorstore.
* multi_write_benchmark writes multiple tensorstores from a "base" tensorstore.

Step 1: Generate a spec/config file for shards.
```
  bazel run tensorstore/internal/benchmark/multi_genspec  \
      --base_spec='{"driver": "zarr3",
                    "kvstore": { "driver": "ocdbt",
                    "base": "file:///tmp/checkpoint" }
      }'  \
    --config='[
  { "name": "sub1", "shape": [1024, 1024], "chunks": [256, 256] }
]'  > config.json
```

Step 2a: Run the write benchmark.
```
  bazel run tensorstore/internal/benchmark/multi_write_benchmark  \
      --context_spec='{"file_io_concurrency": { "limit": 128 }}'  \
      --base_spec='{"driver": "zarr3",
                    "kvstore": { "driver": "ocdbt",
                    "base": "file:///tmp/checkpoint" }
      }'  \
      --write_config=config.json
```

Step 2b: Run the read benchmark.
```
  bazel run tensorstore/internal/benchmark/multi_read_benchmark  \
      --context_spec='{"file_io_concurrency": { "limit": 128 }}'  \
      --repeat_reads=1000 \
      --base_spec='{"driver": "zarr3",
                    "kvstore": { "driver": "ocdbt",
                    "base": "file:///tmp/checkpoint" }
      }'  \
      --read_config=config.json
```

PiperOrigin-RevId: 681493672
Change-Id: I4614686c673d9ac8a811e394cdf4abcc0246ced1
  • Loading branch information
@laramiel @copybara-github
laramiel authored and copybara-github committed Oct 2, 2024
1 parent d465a9c commit 4d0196c
Show file tree
Hide file tree
Showing 31 changed files with 2,418 additions and 250 deletions.
1 change: 1 addition & 0 deletions tensorstore/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -714,6 +714,7 @@ tensorstore_cc_library(
"//tensorstore/serialization",
"//tensorstore/util:option",
"//tensorstore/util:result",
"//tensorstore/util:status",
"//tensorstore/util/garbage_collection",
"@com_github_nlohmann_json//:json",
"@com_google_absl//absl/meta:type_traits",
Expand Down
1 change: 1 addition & 0 deletions tensorstore/driver/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,7 @@ tensorstore_cc_library(
"//tensorstore/util:status",
"//tensorstore/util:str_cat",
"//tensorstore/util/execution",
"//tensorstore/util/execution:any_receiver",
"//tensorstore/util/execution:sender",
"//tensorstore/util/execution:sender_util",
"//tensorstore/util/garbage_collection",
Expand Down
3 changes: 1 addition & 2 deletions tensorstore/driver/copy.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,13 +35,12 @@
#include "tensorstore/internal/nditerable_data_type_conversion.h"
#include "tensorstore/internal/nditerable_util.h"
#include "tensorstore/internal/tagged_ptr.h"
#include "tensorstore/internal/type_traits.h"
#include "tensorstore/open_mode.h"
#include "tensorstore/progress.h"
#include "tensorstore/read_write_options.h"
#include "tensorstore/resize_options.h"
#include "tensorstore/transaction.h"
#include "tensorstore/util/execution/sender.h"
#include "tensorstore/util/execution/any_receiver.h"
#include "tensorstore/util/executor.h"
#include "tensorstore/util/future.h"
#include "tensorstore/util/result.h"
Expand Down
4 changes: 4 additions & 0 deletions tensorstore/driver/downsample/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -286,7 +286,11 @@ tensorstore_cc_library(
deps = [
"//tensorstore:array",
"//tensorstore:box",
"//tensorstore:contiguous_layout",
"//tensorstore:data_type",
"//tensorstore:index",
"//tensorstore:index_interval",
"//tensorstore:rank",
"//tensorstore/util:iterate",
"//tensorstore/util:span",
"@com_google_absl//absl/container:flat_hash_map",
Expand Down
10 changes: 10 additions & 0 deletions tensorstore/driver/downsample/downsample_nditerable.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -507,6 +507,11 @@ struct DownsampleImpl {
/// ElementwiseFunction LoopTemplate implementation for accumulating the
/// total.
struct ProcessInput {
template <typename ArrayAccessor>
static constexpr auto GetLoopFn() {
return &Loop<ArrayAccessor>;
}

template <typename ArrayAccessor>
static bool Loop(void* accumulate_buffer,
internal::IterationBufferShape output_block_shape,
Expand Down Expand Up @@ -607,6 +612,11 @@ struct DownsampleImpl {
/// ElementwiseFunction LoopTemplate implementation for computing the output
/// from the accumulated values.
struct ComputeOutput {
template <typename ArrayAccessor>
static constexpr auto GetLoopFn() {
return &Loop<ArrayAccessor>;
}

template <typename ArrayAccessor>
static bool Loop(void* accumulate_buffer,
internal::IterationBufferShape output_block_shape,
Expand Down
7 changes: 7 additions & 0 deletions tensorstore/driver/downsample/grid_occupancy_map.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,19 @@

#include "tensorstore/driver/downsample/grid_occupancy_map.h"

#include <algorithm>
#include <cassert>
#include <cstddef>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "tensorstore/array.h"
#include "tensorstore/box.h"
#include "tensorstore/contiguous_layout.h"
#include "tensorstore/data_type.h"
#include "tensorstore/index.h"
#include "tensorstore/index_interval.h"
#include "tensorstore/rank.h"
#include "tensorstore/util/iterate.h"
#include "tensorstore/util/span.h"

Expand Down
8 changes: 3 additions & 5 deletions tensorstore/internal/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -333,7 +333,6 @@ tensorstore_cc_library(
"//tensorstore:index",
"//tensorstore/util:byte_strided_pointer",
"@com_google_absl//absl/meta:type_traits",
"@com_google_absl//absl/status",
],
)

Expand Down Expand Up @@ -362,10 +361,12 @@ tensorstore_cc_library(
"//tensorstore/internal/riegeli:json_input",
"//tensorstore/internal/riegeli:json_output",
"//tensorstore/util:endian",
"//tensorstore/util:status",
"//tensorstore/util:str_cat",
"//tensorstore/util:utf8_string",
"@com_github_nlohmann_json//:json",
"@com_google_absl//absl/status",
"@com_google_riegeli//riegeli/bytes:reader",
"@com_google_riegeli//riegeli/bytes:writer",
],
)

Expand Down Expand Up @@ -1126,7 +1127,6 @@ tensorstore_cc_library(
"//tensorstore:data_type",
"//tensorstore:index",
"//tensorstore/util:span",
"//tensorstore/util:status",
"@com_google_absl//absl/status",
],
)
Expand Down Expand Up @@ -1225,8 +1225,6 @@ tensorstore_cc_library(
"//tensorstore:rank",
"//tensorstore/util:iterate",
"//tensorstore/util:span",
"//tensorstore/util:status",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/status",
],
)
Expand Down
156 changes: 155 additions & 1 deletion tensorstore/internal/benchmark/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,23 @@ tensorstore_cc_library(
],
)

tensorstore_cc_library(
name = "vector_flag",
hdrs = ["vector_flag.h"],
deps = [
"@com_google_absl//absl/flags:marshalling",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:string_view",
],
)

tensorstore_cc_binary(
name = "ts_benchmark",
testonly = True,
srcs = ["ts_benchmark.cc"],
deps = [
":metric_utils",
":vector_flag",
"//tensorstore:all_drivers",
"//tensorstore:context",
"//tensorstore:index",
Expand All @@ -34,7 +45,6 @@ tensorstore_cc_binary(
"//tensorstore/util:result",
"//tensorstore/util:status",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/flags:marshalling",
"@com_google_absl//absl/flags:parse",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/random",
Expand Down Expand Up @@ -103,3 +113,147 @@ tensorstore_cc_binary(
"@com_google_absl//absl/time",
],
)

tensorstore_cc_library(
name = "multi_spec",
srcs = ["multi_spec.cc"],
hdrs = ["multi_spec.h"],
deps = [
"//tensorstore:box",
"//tensorstore:index",
"//tensorstore:json_serialization_options_base",
"//tensorstore/internal/json_binding",
"//tensorstore/internal/json_binding:bindable",
"@com_github_nlohmann_json//:json",
"@com_google_absl//absl/log:absl_log",
"@com_google_riegeli//riegeli/bytes:fd_reader",
"@com_google_riegeli//riegeli/bytes:read_all",
],
)

tensorstore_cc_binary(
name = "multi_read_benchmark",
testonly = True,
srcs = ["multi_read_benchmark.cc"],
deps = [
":metric_utils",
":multi_spec",
":vector_flag",
"//tensorstore",
"//tensorstore:all_drivers",
"//tensorstore:array",
"//tensorstore:box",
"//tensorstore:context",
"//tensorstore:open",
"//tensorstore:open_mode",
"//tensorstore:spec",
"//tensorstore/index_space:dim_expression",
"//tensorstore/internal/json_binding:bindable",
"//tensorstore/internal/metrics",
"//tensorstore/internal/metrics:metadata",
"//tensorstore/kvstore",
"//tensorstore/util:future",
"//tensorstore/util:json_absl_flag",
"//tensorstore/util:result",
"//tensorstore/util:span",
"//tensorstore/util:status",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/flags:parse",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:absl_check",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
],
)

tensorstore_cc_binary(
name = "multi_write_benchmark",
testonly = True,
srcs = ["multi_write_benchmark.cc"],
deps = [
":metric_utils",
":multi_spec",
"//tensorstore",
"//tensorstore:all_drivers",
"//tensorstore:array",
"//tensorstore:box",
"//tensorstore:chunk_layout",
"//tensorstore:context",
"//tensorstore:data_type",
"//tensorstore:index",
"//tensorstore:open",
"//tensorstore:open_mode",
"//tensorstore:schema",
"//tensorstore:spec",
"//tensorstore:staleness_bound",
"//tensorstore/index_space:dim_expression",
"//tensorstore/internal:data_type_random_generator",
"//tensorstore/internal/json_binding:bindable",
"//tensorstore/internal/metrics",
"//tensorstore/internal/metrics:metadata",
"//tensorstore/kvstore",
"//tensorstore/util:division",
"//tensorstore/util:future",
"//tensorstore/util:iterate",
"//tensorstore/util:json_absl_flag",
"//tensorstore/util:result",
"//tensorstore/util:span",
"//tensorstore/util:status",
"@com_github_nlohmann_json//:json",
"@com_google_absl//absl/container:flat_hash_map",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/flags:parse",
"@com_google_absl//absl/hash",
"@com_google_absl//absl/log:absl_check",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/random",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/time",
],
)

tensorstore_cc_binary(
name = "multi_genspec",
srcs = ["multi_genspec.cc"],
deps = [
":multi_spec",
":vector_flag",
"//tensorstore",
"//tensorstore:all_drivers",
"//tensorstore:array_storage_statistics",
"//tensorstore:box",
"//tensorstore:chunk_layout",
"//tensorstore:context",
"//tensorstore:data_type",
"//tensorstore:open",
"//tensorstore:open_mode",
"//tensorstore:schema",
"//tensorstore:spec",
"//tensorstore/index_space:dim_expression",
"//tensorstore/internal/json_binding",
"//tensorstore/internal/json_binding:bindable",
"//tensorstore/kvstore",
"//tensorstore/util:division",
"//tensorstore/util:future",
"//tensorstore/util:iterate",
"//tensorstore/util:json_absl_flag",
"//tensorstore/util:result",
"//tensorstore/util:span",
"//tensorstore/util:status",
"@com_github_nlohmann_json//:json",
"@com_google_absl//absl/flags:flag",
"@com_google_absl//absl/flags:parse",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:absl_check",
"@com_google_absl//absl/log:absl_log",
"@com_google_absl//absl/status",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/synchronization",
],
)
Loading

0 comments on commit 4d0196c

Please sign in to comment.