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()