def __init__(self, context: Optional[Context] = None):
"""Constructs a beam based executor."""
self._context = context
self._beam_pipeline_args = context.beam_pipeline_args if context else None
if self._beam_pipeline_args:
if beam:
self._beam_pipeline_args = dependency_utils.make_beam_dependency_flags(
self._beam_pipeline_args)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
# TODO(zhitaoli): Rethink how we can add labels and only normalize them
# if the job is submitted against GCP.
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
self._beam_pipeline_args.extend(
telemetry_utils.make_beam_labels_args())
# TODO(b/174174381): Don't use beam_pipeline_args to set ABSL flags.
flags.FLAGS(sys.argv + self._beam_pipeline_args,
known_only=True)
else:
# TODO(b/156000550): We should not specialize `Context` to embed beam
# pipeline args. Instead, the `Context` should consists of generic
# purpose `extra_flags` which can be interpreted differently by
# different implementations of executors.
absl.logging.warning(
'Executor context\'s beam_pipeline_args is being ignored because '
'Apache Beam is not installed.')
def __init__(self, context: Optional[Context] = None):
"""Constructs a beam based executor."""
super().__init__(context)
self._beam_pipeline_args = None
if context:
if isinstance(context, BaseBeamExecutor.Context):
self._beam_pipeline_args = context.beam_pipeline_args
else:
raise ValueError(
'BaseBeamExecutor found initialized with '
'BaseExecutorSpec. Please use BeamEecutorSpec for '
'Beam Components instead.')
if self._beam_pipeline_args:
self._beam_pipeline_args = dependency_utils.make_beam_dependency_flags(
self._beam_pipeline_args)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
# TODO(zhitaoli): Rethink how we can add labels and only normalize them
# if the job is submitted against GCP.
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
self._beam_pipeline_args.extend(
telemetry_utils.make_beam_labels_args())
# TODO(b/174174381): Don't use beam_pipeline_args to set ABSL flags.
flags.FLAGS(sys.argv + self._beam_pipeline_args, known_only=True)
def __init__(self, context: Optional[Context] = None):
"""Constructs a beam based executor."""
self._context = context
self._beam_pipeline_args = context.beam_pipeline_args if context else None
if self._beam_pipeline_args:
self._beam_pipeline_args = dependency_utils.make_beam_dependency_flags(
self._beam_pipeline_args)
def __init__(self, context: Optional[Context] = None):
"""Constructs a beam based executor."""
self._context = context
self._beam_pipeline_args = context.beam_pipeline_args if context else None
if self._beam_pipeline_args:
self._beam_pipeline_args = dependency_utils.make_beam_dependency_flags(
self._beam_pipeline_args)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
# TODO(zhitaoli): Rethink how we can add labels and only normalize them
# if the job is submitted against GCP.
self._beam_pipeline_args.extend(
telemetry_utils.make_beam_labels_args(
tfx_executor=executor_class_path))
def __init__(self, context: Optional[Context] = None):
"""Constructs a beam based executor."""
self._context = context
self._beam_pipeline_args = context.beam_pipeline_args if context else None
if self._beam_pipeline_args:
self._beam_pipeline_args = dependency_utils.make_beam_dependency_flags(
self._beam_pipeline_args)
executor_class_path = '%s.%s' % (self.__class__.__module__,
self.__class__.__name__)
# TODO(zhitaoli): Rethink how we can add labels and only normalize them
# if the job is submitted against GCP.
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: executor_class_path}):
self._beam_pipeline_args.extend(telemetry_utils.make_beam_labels_args())
# TODO(b/174174381): Don't use beam_pipeline_args to set ABSL flags.
flags.FLAGS(sys.argv + self._beam_pipeline_args, known_only=True)
def testNoActionOnFlag(self, flag_value):
beam_pipeline_args = [flag_value]
self.assertListEqual(
[flag_value],
dependency_utils.make_beam_dependency_flags(beam_pipeline_args),
)
def testMakeBeamDependencyFlags(self, mock_build_ephemeral_package):
mock_build_ephemeral_package.return_value = 'mock_file'
beam_flags = dependency_utils.make_beam_dependency_flags(
beam_pipeline_args=[])
self.assertListEqual(['--extra_package=mock_file'], beam_flags)
mock_build_ephemeral_package.assert_called_with()
