def testMakeOutputDirsAndRemoveOutputDirs(self):
output_artifacts = self._output_resolver.generate_output_artifacts(1)
outputs_utils.make_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
self.assertTrue(tf.io.gfile.exists(artifact.uri))
outputs_utils.remove_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
self.assertFalse(tf.io.gfile.exists(artifact.uri))
def _run_executor(
self, execution_info: data_types.ExecutionInfo
) -> execution_result_pb2.ExecutorOutput:
"""Executes underlying component implementation."""
logging.info('Going to run a new execution: %s', execution_info)
outputs_utils.make_output_dirs(execution_info.output_dict)
try:
return self._executor_operator.run_executor(execution_info)
except Exception: # pylint: disable=broad-except
outputs_utils.remove_output_dirs(execution_info.output_dict)
raise
def testMakeOutputDirsAndRemoveOutputDirs(self):
output_artifacts = self._output_resolver.generate_output_artifacts(1)
outputs_utils.make_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
if isinstance(artifact, ValueArtifact):
self.assertFalse(fileio.isdir(artifact.uri))
else:
self.assertTrue(fileio.isdir(artifact.uri))
self.assertTrue(fileio.exists(artifact.uri))
outputs_utils.remove_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
self.assertFalse(fileio.exists(artifact.uri))
def _run_executor(
self, execution_info: base_executor_operator.ExecutionInfo
) -> execution_result_pb2.ExecutorOutput:
"""Executes underlying component implementation."""
logging.info('Going to run a new execution: %s', execution_info)
outputs_utils.make_output_dirs(execution_info.output_dict)
try:
return self._executor_operator.run_executor(execution_info)
except Exception as e: # pylint: disable=broad-except
outputs_utils.remove_output_dirs(execution_info.output_dict)
logging.error(
'Execution failed with error %s '
'and this is the stack trace \n %s', e, traceback.format_exc())
raise
def _run_executor(
self, execution_info: data_types.ExecutionInfo
) -> execution_result_pb2.ExecutorOutput:
"""Executes underlying component implementation."""
logging.info('Going to run a new execution: %s', execution_info)
outputs_utils.make_output_dirs(execution_info.output_dict)
try:
executor_output = self._executor_operator.run_executor(execution_info)
code = executor_output.execution_result.code
if code != 0:
result_message = executor_output.execution_result.result_message
err = (f'Execution {execution_info.execution_id} '
f'failed with error code {code} and '
f'error message {result_message}')
logging.error(err)
raise _ExecutionFailedError(err, executor_output)
return executor_output
except Exception: # pylint: disable=broad-except
outputs_utils.remove_output_dirs(execution_info.output_dict)
raise
def testMakeClearAndRemoveOutputDirs(self):
output_artifacts = self._output_resolver().generate_output_artifacts(1)
outputs_utils.make_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
if isinstance(artifact, ValueArtifact):
self.assertFalse(fileio.isdir(artifact.uri))
else:
self.assertTrue(fileio.isdir(artifact.uri))
with fileio.open(os.path.join(artifact.uri, 'output'),
'w') as f:
f.write('')
self.assertTrue(fileio.exists(artifact.uri))
outputs_utils.clear_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
if not isinstance(artifact, ValueArtifact):
self.assertEqual(fileio.listdir(artifact.uri), [])
outputs_utils.remove_output_dirs(output_artifacts)
for _, artifact_list in output_artifacts.items():
for artifact in artifact_list:
self.assertFalse(fileio.exists(artifact.uri))
def _remove_output_dirs(task: task_lib.ExecNodeTask,
ts_result: ts.TaskSchedulerResult) -> None:
outputs_utils.remove_output_dirs(task.output_artifacts)
if isinstance(ts_result.output, ts.ImporterNodeOutput):
if ts_result.output.output_artifacts is not None:
outputs_utils.remove_output_dirs(ts_result.output.output_artifacts)
