cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enable defining and using internal executor within a subgraph

Browse Source 
PiperOrigin-RevId: 522449982
This commit is contained in:
MediaPipe Team 2023-04-06 15:16:50 -07:00 committed by Copybara-Service
parent 81a405af1b
commit 38f838513a
4 changed files with 112 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
mediapipe/framework/BUILD
View File

@ -1670,6 +1670,7 @@ cc_test(
name = "subgraph_test",
srcs = ["subgraph_test.cc"],
deps = [
":calculator_cc_proto",
":calculator_framework",
":graph_service_manager",
":subgraph",
@ -1681,6 +1682,7 @@ cc_test(
"//mediapipe/framework/port:parse_text_proto",
"//mediapipe/framework/tool:sink",
"//mediapipe/framework/tool/testdata:dub_quad_test_subgraph",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:str_format",
],
)

49
mediapipe/framework/subgraph_test.cc
View File

@ -16,8 +16,10 @@
#include <string>
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "mediapipe/calculators/core/constant_side_packet_calculator.pb.h"
#include "mediapipe/framework/calculator.pb.h"
#include "mediapipe/framework/calculator_framework.h"
#include "mediapipe/framework/graph_service_manager.h"
#include "mediapipe/framework/port/gmock.h"
@ -188,5 +190,52 @@ TEST_F(SubgraphTest, CheckSubgraphOptionsPassedIn) {
EXPECT_EQ(packet.value().Get<std::string>(), "test");
}
class SubgraphUsingInternalExecutor : public Subgraph {
public:
absl::StatusOr<CalculatorGraphConfig> GetConfig(
mediapipe::SubgraphContext* sc) override {
return mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(
R"pb(
output_side_packet: "string"
executor {
name: "xyz"
type: "ThreadPoolExecutor"
options {
[mediapipe.ThreadPoolExecutorOptions.ext] { num_threads: 1 }
}
}
node {
calculator: "ConstantSidePacketCalculator"
executor: "xyz"
output_side_packet: "PACKET:string"
options: {
[mediapipe.ConstantSidePacketCalculatorOptions.ext]: {
packet { string_value: "passed" }
}
}
}
)pb");
}
};
REGISTER_MEDIAPIPE_GRAPH(SubgraphUsingInternalExecutor);
TEST(SubgraphExecutorTest, SubgraphCanDefineAndUseInternalExecutor) {
CalculatorGraphConfig config =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
output_side_packet: "str"
node {
calculator: "SubgraphUsingInternalExecutor"
output_side_packet: "str"
}
)pb");
CalculatorGraph graph;
MP_ASSERT_OK(graph.Initialize(config));
MP_ASSERT_OK(graph.StartRun({}));
MP_ASSERT_OK(graph.WaitUntilDone());
auto packet = graph.GetOutputSidePacket("str");
MP_ASSERT_OK(packet);
EXPECT_EQ(packet->Get<std::string>(), "passed");
}
} // namespace
} // namespace mediapipe

3
mediapipe/framework/tool/subgraph_expansion.cc
View File

@ -319,6 +319,9 @@ absl::Status ExpandSubgraphs(CalculatorGraphConfig* config,
subgraph.status_handler().end(),
proto_ns::RepeatedPtrFieldBackInserter(
config->mutable_status_handler()));
std::copy(
subgraph.executor().begin(), subgraph.executor().end(),
proto_ns::RepeatedPtrFieldBackInserter(config->mutable_executor()));
}
}
return absl::OkStatus();

58
mediapipe/framework/tool/subgraph_expansion_test.cc
View File

@ -527,6 +527,64 @@ TEST(SubgraphExpansionTest, ExecutorFieldOfNodeInSubgraphPreserved) {
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
// A subgraph that defines and uses an internal executor.
class NodeWithInternalExecutorSubgraph : public Subgraph {
public:
absl::StatusOr<CalculatorGraphConfig> GetConfig(
const SubgraphOptions& options) override {
return mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: "IN:foo"
output_stream: "OUT:bar"
executor {
name: "xyz"
type: "ThreadPoolExecutor"
options {
[mediapipe.ThreadPoolExecutorOptions.ext] { num_threads: 1 }
}
}
node {
calculator: "PassThroughCalculator"
executor: "xyz"
input_stream: "foo"
output_stream: "bar"
}
)pb");
}
};
REGISTER_MEDIAPIPE_GRAPH(NodeWithInternalExecutorSubgraph);
TEST(SubgraphExpansionTest, ExecutorCanDefinedAndUsedWithinSubgraph) {
CalculatorGraphConfig supergraph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: "input"
node {
calculator: "NodeWithInternalExecutorSubgraph"
input_stream: "IN:input"
output_stream: "OUT:output"
}
)pb");
CalculatorGraphConfig expected_graph =
mediapipe::ParseTextProtoOrDie<CalculatorGraphConfig>(R"pb(
input_stream: "input"
executor {
name: "xyz"
type: "ThreadPoolExecutor"
options {
[mediapipe.ThreadPoolExecutorOptions.ext] { num_threads: 1 }
}
}
node {
calculator: "PassThroughCalculator"
name: "nodewithinternalexecutorsubgraph__PassThroughCalculator"
input_stream: "input"
output_stream: "output"
executor: "xyz"
}
)pb");
MP_EXPECT_OK(tool::ExpandSubgraphs(&supergraph));
EXPECT_THAT(supergraph, mediapipe::EqualsProto(expected_graph));
}
const mediapipe::GraphService<std::string> kStringTestService{
"mediapipe::StringTestService"};
class GraphServicesClientTestSubgraph : public Subgraph {