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:
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 testMakeBeamLabelsArgs(self):
"""Test for make_beam_labels_args."""
beam_pipeline_args = telemetry_utils.make_beam_labels_args()
self.assertListEqual([
'--labels',
'tfx_py_version=%d-%d' %
(sys.version_info.major, sys.version_info.minor),
'--labels',
'tfx_version=%s' % version.__version__.replace('.', '-'),
], beam_pipeline_args)
def testMakeBeamLabelsArgs(self):
"""Test for make_beam_labels_args."""
beam_pipeline_args = telemetry_utils.make_beam_labels_args()
expected_beam_pipeline_args = [
'--labels',
'tfx_py_version=%d-%d' %
(sys.version_info.major, sys.version_info.minor),
'--labels',
'tfx_version=%s' % version.__version__.replace('.', '-'),
]
self.assertListEqual(expected_beam_pipeline_args, beam_pipeline_args)
with telemetry_utils.scoped_labels(
{telemetry_utils.LABEL_TFX_EXECUTOR: 'TestExecutor'}):
beam_pipeline_args = telemetry_utils.make_beam_labels_args()
expected_beam_pipeline_args = [
'--labels',
'tfx_executor=third_party_executor',
] + expected_beam_pipeline_args
self.assertListEqual(expected_beam_pipeline_args,
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 launch(self):
self.recorded_labels = telemetry_utils.make_beam_labels_args()
return mock.MagicMock()
