def testLauncher_importer_mode_reimport_enabled(self):
handler = importer_node_handler.ImporterNodeHandler()
execution_metadata = handler.run(
mlmd_connection=self._mlmd_connection,
pipeline_node=self._importer,
pipeline_info=self._pipeline_info,
pipeline_runtime_spec=self._pipeline_runtime_spec)
with self._mlmd_connection as m:
[artifact] = m.store.get_artifacts_by_type('Schema')
self.assertProtoPartiallyEquals("""
id: 1
type_id: 5
uri: "my_url"
custom_properties {
key: "int_custom_property"
value {
int_value: 123
}
}
custom_properties {
key: "str_custom_property"
value {
string_value: "abc"
}
}
state: LIVE""",
artifact,
ignored_fields=[
'create_time_since_epoch',
'last_update_time_since_epoch'
])
[execution] = m.store.get_executions_by_id([execution_metadata.id])
self.assertProtoPartiallyEquals("""
id: 1
type_id: 4
last_known_state: COMPLETE
custom_properties {
key: "artifact_uri"
value {
string_value: "my_url"
}
}
custom_properties {
key: "reimport"
value {
int_value: 1
}
}
""",
execution,
ignored_fields=[
'create_time_since_epoch',
'last_update_time_since_epoch'
])
execution_metadata = handler.run(
mlmd_connection=self._mlmd_connection,
pipeline_node=self._importer,
pipeline_info=self._pipeline_info,
pipeline_runtime_spec=self._pipeline_runtime_spec)
with self._mlmd_connection as m:
new_artifact = m.store.get_artifacts_by_type('Schema')[1]
self.assertProtoPartiallyEquals("""
id: 2
type_id: 5
uri: "my_url"
custom_properties {
key: "int_custom_property"
value {
int_value: 123
}
}
custom_properties {
key: "str_custom_property"
value {
string_value: "abc"
}
}
state: LIVE""",
new_artifact,
ignored_fields=[
'create_time_since_epoch',
'last_update_time_since_epoch'
])
[execution] = m.store.get_executions_by_id([execution_metadata.id])
self.assertProtoPartiallyEquals("""
id: 2
type_id: 4
last_known_state: COMPLETE
custom_properties {
key: "artifact_uri"
value {
string_value: "my_url"
}
}
custom_properties {
key: "reimport"
value {
int_value: 1
}
}
""",
execution,
ignored_fields=[
'create_time_since_epoch',
'last_update_time_since_epoch'
])
def testLauncher_importer_mode_reimport_disabled(self):
self._importer.parameters.parameters[
'reimport'].field_value.int_value = 0
handler = importer_node_handler.ImporterNodeHandler()
execution_info = handler.run(
mlmd_connection=self._mlmd_connection,
pipeline_node=self._importer,
pipeline_info=self._pipeline_info,
pipeline_runtime_spec=self._pipeline_runtime_spec)
with self._mlmd_connection as m:
[artifact] = m.store.get_artifacts_by_type('Schema')
self.assertProtoPartiallyEquals("""
id: 1
uri: "my_url"
custom_properties {
key: "int_custom_property"
value {
int_value: 123
}
}
custom_properties {
key: "str_custom_property"
value {
string_value: "abc"
}
}
custom_properties {
key: "tfx_version"
value {
string_value: "0.123.4.dev"
}
}
state: LIVE""",
artifact,
ignored_fields=[
'type_id',
'create_time_since_epoch',
'last_update_time_since_epoch'
])
[execution
] = m.store.get_executions_by_id([execution_info.execution_id])
self.assertProtoPartiallyEquals("""
id: 1
last_known_state: COMPLETE
custom_properties {
key: "artifact_uri"
value {
string_value: "my_url"
}
}
custom_properties {
key: "reimport"
value {
int_value: 0
}
}
""",
execution,
ignored_fields=[
'type_id',
'create_time_since_epoch',
'last_update_time_since_epoch'
])
# Run the 2nd execution. Since the reimport is disabled, no new schema
# is imported and the corresponding execution is published as CACHED.
execution_info = handler.run(
mlmd_connection=self._mlmd_connection,
pipeline_node=self._importer,
pipeline_info=self._pipeline_info,
pipeline_runtime_spec=self._pipeline_runtime_spec)
with self._mlmd_connection as m:
# No new Schema is produced.
self.assertLen(m.store.get_artifacts_by_type('Schema'), 1)
[execution
] = m.store.get_executions_by_id([execution_info.execution_id])
self.assertProtoPartiallyEquals("""
id: 2
last_known_state: CACHED
custom_properties {
key: "artifact_uri"
value {
string_value: "my_url"
}
}
custom_properties {
key: "reimport"
value {
int_value: 0
}
}
""",
execution,
ignored_fields=[
'type_id',
'create_time_since_epoch',
'last_update_time_since_epoch'
])