def testEntryPoint(self):
"""Test the entrypoint with toy inputs."""
# Test both current version metadata and legacy property/custom property
# metadata styles.
for serialized_metadata in [
self._serialized_metadata, self._serialized_metadata_legacy
]:
with _ArgsCapture() as args_capture:
args = [
"--executor_class_path",
"%s.%s" %
(_FakeExecutor.__module__, _FakeExecutor.__name__),
"--json_serialized_invocation_args", serialized_metadata
]
kubeflow_v2_run_executor.main(
kubeflow_v2_run_executor._parse_flags(args))
# TODO(b/131417512): Add equal comparison to types.Artifact class so we
# can use asserters.
self.assertEqual(set(args_capture.input_dict.keys()),
set(["input_1", "input_2"]))
self.assertEqual(
set(args_capture.output_dict.keys()),
set(["output", standard_component_specs.BLESSING_KEY]))
self.assertEqual(args_capture.exec_properties,
_EXEC_PROPERTIES)
# Test what's been output.
with open(_TEST_OUTPUT_METADATA_JSON) as output_meta_json:
actual_output = json.dumps(json.load(output_meta_json),
indent=2,
sort_keys=True)
self.assertEqual(actual_output, self._expected_output)
os.remove(_TEST_OUTPUT_METADATA_JSON)
def testEntryPointWithDriver(self):
"""Test the entrypoint with Driver's output metadata."""
# Mock the driver's output metadata.
output_metadata = pipeline_spec_pb2.ExecutorOutput()
output_metadata.parameters["key_1"].string_value = "driver"
output_metadata.parameters["key_3"].string_value = "driver3"
fileio.makedirs(os.path.dirname(_TEST_OUTPUT_METADATA_JSON))
with fileio.open(_TEST_OUTPUT_METADATA_JSON, "wb") as f:
f.write(json_format.MessageToJson(output_metadata, sort_keys=True))
with _ArgsCapture() as args_capture:
args = [
"--executor_class_path",
"%s.%s" % (_FakeExecutor.__module__, _FakeExecutor.__name__),
"--json_serialized_invocation_args", self._serialized_metadata
]
kubeflow_v2_run_executor.main(
kubeflow_v2_run_executor._parse_flags(args))
# TODO(b/131417512): Add equal comparison to types.Artifact class so we
# can use asserters.
self.assertEqual(set(args_capture.input_dict.keys()),
set(["input_1", "input_2"]))
self.assertEqual(
set(args_capture.output_dict.keys()),
set(["output", standard_component_specs.BLESSING_KEY]))
# Verify that exec_properties use driver's output metadata.
self.assertEqual(
args_capture.exec_properties,
{
"key_1": "driver", # Overwrite.
"key_2": 536870911,
"key_3": "driver3" # Append.
})
# Test what's been output.
with open(_TEST_OUTPUT_METADATA_JSON) as output_meta_json:
actual_output = json.dumps(json.load(output_meta_json),
indent=2,
sort_keys=True)
self.assertEqual(actual_output, self._expected_output)
os.remove(_TEST_OUTPUT_METADATA_JSON)
def testEntryPoint(self):
"""Test the entrypoint with toy inputs."""
with _ArgsCapture() as args_capture:
args = [
"--executor_class_path",
"%s.%s" % (_FakeExecutor.__module__, _FakeExecutor.__name__),
"--json_serialized_invocation_args", self._serialized_metadata
]
kubeflow_v2_run_executor.main(args)
# TODO(b/131417512): Add equal comparison to types.Artifact class so we
# can use asserters.
self.assertEqual(set(args_capture.input_dict.keys()),
set(["input_1", "input_2"]))
self.assertEqual(set(args_capture.output_dict.keys()),
set(["output", constants.BLESSING_KEY]))
self.assertEqual(args_capture.exec_properties, _EXEC_PROPERTIES)
# Test what's been output.
with open(_TEST_OUTPUT_METADATA_JSON) as output_meta_json:
actual_output = json.load(output_meta_json)
self.assertDictEqual(actual_output, self._expected_output)
