#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""Pipeline options obtained from command line parsing."""

# pytype: skip-file

import argparse

import difflib

import json

import logging

import os

from typing import Any

from typing import Callable

from typing import Dict

from typing import List

from typing import Optional

from typing import Sequence

from typing import Type

from typing import TypeVar

import apache_beam as beam

from apache_beam.options.value_provider import RuntimeValueProvider

from apache_beam.options.value_provider import StaticValueProvider

from apache_beam.options.value_provider import ValueProvider

from apache_beam.transforms.display import HasDisplayData

from apache_beam.utils import logger

from apache_beam.utils import proto_utils

__all__ = [

'PipelineOptions',

'StandardOptions',

'TypeOptions',

'DirectOptions',

'GoogleCloudOptions',

'AzureOptions',

'HadoopFileSystemOptions',

'WorkerOptions',

'DebugOptions',

'ProfilingOptions',

'SetupOptions',

'TestOptions',

'S3Options'

]

PipelineOptionsT = TypeVar('PipelineOptionsT', bound='PipelineOptions')

_LOGGER = logging.getLogger(__name__)

# Map defined with option names to flag names for boolean options

# that have a destination(dest) in parser.add_argument() different

# from the flag name and whose default value is `None`.

_FLAG_THAT_SETS_FALSE_VALUE = {

'use_public_ips': 'no_use_public_ips',

'save_main_session': 'no_save_main_session'

}

# Set of options which should not be overriden when applying options from a

# different language. This is relevant when using x-lang transforms where the

# expansion service is started up with some pipeline options, and will

# impact which options are passed in to expanded transforms' expand functions.

_NON_OVERIDABLE_XLANG_OPTIONS = ['runner', 'experiments']

def _static_value_provider_of(value_type):

"""Helper function to plug a ValueProvider into argparse.

Args:

value_type: the type of the value. Since the type param of argparse's

add_argument will always be ValueProvider, we need to

preserve the type of the actual value.

Returns:

A partially constructed StaticValueProvider in the form of a function.

"""

def _f(value):

_f.__name__ = value_type.__name__

return StaticValueProvider(value_type, value)

return _f

class _BeamArgumentParser(argparse.ArgumentParser):

"""An ArgumentParser that supports ValueProvider options.

Example Usage::

class TemplateUserOptions(PipelineOptions):

@classmethod

def _add_argparse_args(cls, parser):

parser.add_value_provider_argument('--vp_arg1', default='start')

parser.add_value_provider_argument('--vp_arg2')

parser.add_argument('--non_vp_arg')

"""

def add_value_provider_argument(self, *args, **kwargs):

"""ValueProvider arguments can be either of type keyword or positional.

At runtime, even positional arguments will need to be supplied in the

key/value form.

"""

# Extract the option name from positional argument ['pos_arg']

assert args and len(args[0]) >= 1

if args[0][0] != '-':

option_name = args[0]

if kwargs.get('nargs') is None: # make them optionally templated

kwargs['nargs'] = '?'

else:

# or keyword arguments like [--kw_arg, -k, -w] or [--kw-arg]

option_name = [i.replace('--', '') for i in args if i[:2] == '--'][0]

# reassign the type to make room for using

# StaticValueProvider as the type for add_argument

value_type = kwargs.get('type') or str

kwargs['type'] = _static_value_provider_of(value_type)

# reassign default to default_value to make room for using

# RuntimeValueProvider as the default for add_argument

default_value = kwargs.get('default')

kwargs['default'] = RuntimeValueProvider(

option_name=option_name,

value_type=value_type,

default_value=default_value)

# have add_argument do most of the work

self.add_argument(*args, **kwargs)

# The argparse package by default tries to autocomplete option names. This

# results in an "ambiguous option" error from argparse when an unknown option

# matching multiple known ones are used. This suppresses that behavior.

def error(self, message):

if message.startswith('ambiguous option: '):

return

super().error(message)

class _DictUnionAction(argparse.Action):

"""

argparse Action take union of json loads values. If a key is specified in more

than one of the values, the last value takes precedence.

"""

def __call__(self, parser, namespace, values, option_string=None):

if not hasattr(namespace, self.dest) or getattr(namespace,

self.dest) is None:

setattr(namespace, self.dest, {})

getattr(namespace, self.dest).update(values)

class _GcsCustomAuditEntriesAction(argparse.Action):

"""

Argparse Action for GCS custom audit entries.

According to Google Cloud Storage audit logging documentation

(https://cloud.google.com/storage/docs/audit-logging#add-custom-metadata), the

following limitations apply:

- keys must be 64 characters or less,

- values 1,200 characters or less, and

- a maximum of four custom metadata entries are permitted per request.

"""

MAX_KEY_LENGTH = 64

MAX_VALUE_LENGTH = 1200

MAX_ENTRIES = 4

GCS_AUDIT_PREFIX = 'x-goog-custom-audit-'

def _exceed_entry_limit(self):

job_audit_entry = _GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX + 'job'

if job_audit_entry in self._custom_audit_entries:

return len(

self._custom_audit_entries) > _GcsCustomAuditEntriesAction.MAX_ENTRIES

else:

return len(self._custom_audit_entries) > (

_GcsCustomAuditEntriesAction.MAX_ENTRIES - 1)

def _add_entry(self, key, value):

if len(key) > _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH:

raise argparse.ArgumentError(

None,

"The key '%s' in GCS custom audit entries exceeds the %d-character limit." # pylint: disable=line-too-long

% (key, _GcsCustomAuditEntriesAction.MAX_KEY_LENGTH))

if len(value) > _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH:

raise argparse.ArgumentError(

None,

"The value '%s' in GCS custom audit entries exceeds the %d-character limit." # pylint: disable=line-too-long

% (value, _GcsCustomAuditEntriesAction.MAX_VALUE_LENGTH))

if key.startswith(_GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX):

self._custom_audit_entries[key] = value

else:

self._custom_audit_entries[_GcsCustomAuditEntriesAction.GCS_AUDIT_PREFIX +

key] = value

def __call__(self, parser, namespace, values, option_string=None):

if not hasattr(namespace, self.dest) or getattr(namespace,

self.dest) is None:

setattr(namespace, self.dest, {})

self._custom_audit_entries = getattr(namespace, self.dest)

if option_string == '--gcs_custom_audit_entries':

# in the format of {"key": "value"}

assert (isinstance(values, str))

sub_entries = json.loads(values)

for key, value in sub_entries.items():

self._add_entry(key, value)

else: # option_string == '--gcs_custom_audit_entry'

# in the format of 'key=value'

assert (isinstance(values, str))

parts = values.split('=', 1)

key = parts[0]

value = parts[1] if len(parts) > 1 else ''

self._add_entry(key, value)

if self._exceed_entry_limit():

raise argparse.ArgumentError(

None,

"The maximum allowed number of GCS custom audit entries (including the default x-goo-custom-audit-job) is %d." # pylint: disable=line-too-long

% _GcsCustomAuditEntriesAction.MAX_ENTRIES)

class _CommaSeparatedListAction(argparse.Action):

"""

Argparse Action that splits comma-separated values and appends them to

a list. This allows options like --experiments=abc,def to be treated

as separate experiments 'abc' and 'def', similar to how Java SDK handles

them.

If there are key=value experiments in a raw argument, the remaining part of

the argument are treated as values and won't split further. For example:

'abc,def,master_key=k1=v1,k2=v2' becomes

['abc', 'def', 'master_key=k1=v1,k2=v2'].

"""

def __call__(self, parser, namespace, values, option_string=None):

if not hasattr(namespace, self.dest) or getattr(namespace,

self.dest) is None:

setattr(namespace, self.dest, [])

# Split comma-separated values and extend the list

if isinstance(values, str):

# Smart splitting: only split at commas that are not part of

# key=value pairs

split_values = self._smart_split(values)

getattr(namespace, self.dest).extend(split_values)

else:

# If values is not a string, just append it

getattr(namespace, self.dest).append(values)

def _smart_split(self, values):

"""Split comma-separated values, but preserve commas within

key=value pairs."""

result = []

current = []

equals_depth = 0

i = 0

while i < len(values):

char = values[i]

if char == '=':

equals_depth += 1

current.append(char)

elif char == ',' and equals_depth <= 1:

# This comma is a top-level separator (not inside a complex value)

if current:

result.append(''.join(current).strip())

current = []

equals_depth = 0

elif char == ',' and equals_depth > 1:

# This comma is inside a complex value, keep it

current.append(char)

elif char == ' ' and not current:

# Skip leading spaces

pass

else:

current.append(char)

i += 1

# Add the last item

if current:

result.append(''.join(current).strip())

return [v for v in result if v] # Filter out empty values

class PipelineOptions(HasDisplayData):

# Set of options which should not be overriden when pipeline options are

# being merged (see from_runner_api). This primarily comes up when expanding

# the Python expansion service

"""This class and subclasses are used as containers for command line options.

These classes are wrappers over the standard argparse Python module

(see https://docs.python.org/3/library/argparse.html). To define one option

or a group of options, create a subclass from PipelineOptions.

Example Usage::

class XyzOptions(PipelineOptions):

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument('--abc', default='start')

parser.add_argument('--xyz', default='end')

The arguments for the add_argument() method are exactly the ones

described in the argparse public documentation.

Pipeline objects require an options object during initialization.

This is obtained simply by initializing an options class as defined above.

Example Usage::

p = Pipeline(options=XyzOptions())

if p.options.xyz == 'end':

raise ValueError('Option xyz has an invalid value.')

Instances of PipelineOptions or any of its subclass have access to values

defined by other PipelineOption subclasses (see get_all_options()), and

can be converted to an instance of another PipelineOptions subclass

(see view_as()). All views share the underlying data structure that stores

option key-value pairs.

By default the options classes will use command line arguments to initialize

the options.

"""

def __init__(self, flags: Optional[Sequence[str]] = None, **kwargs) -> None:

"""Initialize an options class.

The initializer will traverse all subclasses, add all their argparse

arguments and then parse the command line specified by flags or by default

the one obtained from sys.argv.

The subclasses of PipelineOptions do not need to redefine __init__.

Args:

flags: An iterable of command line arguments to be used. If not specified

then sys.argv will be used as input for parsing arguments.

**kwargs: Add overrides for arguments passed in flags. For overrides

of arguments, please pass the `option names` instead of

flag names.

Option names: These are defined as dest in the

parser.add_argument() for each flag. Passing flags

like {no_use_public_ips: True}, for which the dest is

defined to a different flag name in the parser,

would be discarded. Instead, pass the dest of

the flag (dest of no_use_public_ips is use_public_ips).

"""

# Initializing logging configuration in case the user did not set it up.

logging.basicConfig()

# self._flags stores a list of not yet parsed arguments, typically,

# command-line flags. This list is shared across different views.

# See: view_as().

if isinstance(flags, str):

# Unfortunately a single str passes the Iterable[str] test, as it is

# an iterable of single characters. This is almost certainly not the

# intent...

raise ValueError(

"Flags must be an iterable of of strings, not a single string.")

self._flags = flags

# Build parser that will parse options recognized by the [sub]class of

# PipelineOptions whose object is being instantiated.

parser = _BeamArgumentParser(allow_abbrev=False)

for cls in type(self).mro():

if cls == PipelineOptions:

break

elif '_add_argparse_args' in cls.__dict__:

cls._add_argparse_args(parser) # type: ignore

# The _visible_options attribute will contain options that were recognized

# by the parser.

self._visible_options, _ = parser.parse_known_args(flags)

# self._all_options is initialized with overrides to flag values,

# provided in kwargs, and will store key-value pairs for options recognized

# by current PipelineOptions [sub]class and its views that may be created.

# See: view_as().

# This dictionary is shared across different views, and is lazily updated

# as each new views are created.

# Users access this dictionary store via __getattr__ / __setattr__ methods.

self._all_options = kwargs

# Initialize values of keys defined by this class.

for option_name in self._visible_option_list():

# Note that options specified in kwargs will not be overwritten.

if option_name not in self._all_options:

self._all_options[option_name] = getattr(

self._visible_options, option_name)

def __getstate__(self):

# The impersonate_service_account option must be used only at submission of

# a Beam job. However, Beam IOs might store pipeline options

# within transform implementation that becomes serialized in RunnerAPI,

# causing this option to be inadvertently used at runtime.

# This serialization hook removes it.

if self.view_as(GoogleCloudOptions).impersonate_service_account:

dict_copy = dict(self.__dict__)

dict_copy['_all_options'] = dict(dict_copy['_all_options'])

dict_copy['_all_options']['impersonate_service_account'] = None

return dict_copy

else:

return self.__dict__

@classmethod

def _add_argparse_args(cls, parser: _BeamArgumentParser) -> None:

# Override this in subclasses to provide options.

pass

@classmethod

def from_dictionary(cls, options):

"""Returns a PipelineOptions from a dictionary of arguments.

Args:

options: Dictionary of argument value pairs.

Returns:

A PipelineOptions object representing the given arguments.

"""

flags = []

for k, v in options.items():

# Note: If a boolean flag is True in the dictionary,

# implicitly the method assumes the boolean flag is

# specified as a command line argument. If the

# boolean flag is False, this method simply discards them.

# Eg: {no_auth: True} is similar to python your_file.py --no_auth

# {no_auth: False} is similar to python your_file.py.

if isinstance(v, bool):

if v:

flags.append('--%s' % k)

elif k in _FLAG_THAT_SETS_FALSE_VALUE:

# Capture overriding flags, which have a different dest

# from the flag name defined in the parser.add_argument

# Eg: no_use_public_ips, which has the dest=use_public_ips

# different from flag name

flag_that_disables_the_option = (_FLAG_THAT_SETS_FALSE_VALUE[k])

flags.append('--%s' % flag_that_disables_the_option)

elif isinstance(v, list):

for i in v:

flags.append('--%s=%s' % (k, i))

elif isinstance(v, dict):

flags.append('--%s=%s' % (k, json.dumps(v)))

elif v is None:

# Don't process None type args here, they will be treated

# as strings when parsed by BeamArgumentParser..

logging.warning('Not setting flag with value None: %s', k)

else:

flags.append('--%s=%s' % (k, v))

return cls(flags)

@staticmethod

def _warn_on_unknown_options(unknown_args, parser):

if not unknown_args:

return

all_known_options = [

opt for action in parser._actions for opt in action.option_strings

]

for arg in unknown_args:

msg = f"Unparseable argument: {arg}"

if arg.startswith('--'):

arg_name = arg.split('=', 1)[0]

suggestions = difflib.get_close_matches(arg_name, all_known_options)

if suggestions:

msg += f". Did you mean '{suggestions[0]}'?'"

logger.log_first_n(logging.WARN, msg, key="message")

def get_all_options(

self,

drop_default=False,

add_extra_args_fn: Optional[Callable[[_BeamArgumentParser], None]] = None,

retain_unknown_options=False,

display_warnings=False,

current_only=False,

) -> Dict[str, Any]:

"""Returns a dictionary of all defined arguments.

Returns a dictionary of all defined arguments into a dictionary.

Args:

drop_default: If set to true, options that are equal to their default

values, are not returned as part of the result dictionary.

add_extra_args_fn: Callback to populate additional arguments, can be used

by runner to supply otherwise unknown args.

retain_unknown_options: If set to true, options not recognized by any

known pipeline options class will still be included in the result. If

set to false, they will be discarded.

current_only: If set to true, only returns options defined in this class.

Otherwise, arguments that are defined in any subclass of PipelineOptions

are returned (default).

Returns:

Dictionary of all args and values.

"""

# TODO(https://github.com/apache/beam/issues/18197): PipelineOption

# sub-classes in the main session might be repeated. Pick last unique

# instance of each subclass to avoid conflicts.

subset = {}

parser = _BeamArgumentParser(allow_abbrev=False)

if current_only:

subset.setdefault(str(type(self)), type(self))

else:

for cls in PipelineOptions.__subclasses__():

subset.setdefault(str(cls), cls)

for cls in subset.values():

cls._add_argparse_args(parser) # pylint: disable=protected-access

if add_extra_args_fn:

add_extra_args_fn(parser)

known_args, unknown_args = parser.parse_known_args(self._flags)

if display_warnings:

self._warn_on_unknown_options(unknown_args, parser)

if retain_unknown_options:

seen = set()

def add_new_arg(arg, **kwargs):

if arg not in seen:

parser.add_argument(arg, **kwargs)

seen.add(arg)

i = 0

while i < len(unknown_args):

# End of argument parsing.

if unknown_args[i] == '--':

break

# Treat all unary flags as booleans, and all binary argument values as

# strings.

if not unknown_args[i].startswith('-'):

i += 1

continue

if i + 1 >= len(unknown_args) or unknown_args[i + 1].startswith('-'):

split = unknown_args[i].split('=', 1)

if len(split) == 1:

add_new_arg(unknown_args[i], action='store_true')

else:

add_new_arg(split[0], type=str)

i += 1

elif unknown_args[i].startswith('--'):

add_new_arg(unknown_args[i], type=str)

i += 2

else:

# skip all binary flags used with '-' and not '--'.

# ex: using -f instead of --f (or --flexrs_goal) will prevent

# argument validation before job submission and can be incorrectly

# submitted to job.

_LOGGER.warning(

"Discarding flag %s, single dash flags are not allowed.",

unknown_args[i])

i += 2

continue

parsed_args, _ = parser.parse_known_args(self._flags)

else:

if unknown_args and not current_only:

_LOGGER.warning("Discarding unparseable args: %s", unknown_args)

parsed_args = known_args

result = vars(parsed_args)

overrides = self._all_options.copy()

# Apply the overrides if any

for k in list(result):

overrides.pop(k, None)

if k in self._all_options:

result[k] = self._all_options[k]

if (drop_default and parser.get_default(k) == result[k] and

not isinstance(parser.get_default(k), ValueProvider)):

del result[k]

if overrides:

if retain_unknown_options:

result.update(overrides)

elif not current_only:

_LOGGER.warning("Discarding invalid overrides: %s", overrides)

return result

def to_runner_api(self):

def to_struct_value(o):

if isinstance(o, (bool, int, str)):

return o

elif isinstance(o, (tuple, list)):

return [to_struct_value(e) for e in o]

elif isinstance(o, dict):

return {str(k): to_struct_value(v) for k, v in o.items()}

else:

return str(o) # Best effort.

return proto_utils.pack_Struct(

**{

f'beam:option:{k}:v1': to_struct_value(v)

for (k, v) in self.get_all_options(

drop_default=True, retain_unknown_options=True).items()

if v is not None

})

@classmethod

def from_runner_api(cls, proto_options, original_options=None):

def from_urn(key):

assert key.startswith('beam:option:')

assert key.endswith(':v1')

return key[12:-3]

parsed = {from_urn(key): value for (key, value) in proto_options.items()}

if original_options is None:

return cls(**parsed)

for (key, value) in parsed.items():

if value and key not in _NON_OVERIDABLE_XLANG_OPTIONS:

original_options._all_options[key] = value

return original_options

def display_data(self):

return self.get_all_options(drop_default=True, retain_unknown_options=True)

def view_as(self, cls: Type[PipelineOptionsT]) -> PipelineOptionsT:

"""Returns a view of current object as provided PipelineOption subclass.

Example Usage::

options = PipelineOptions(['--runner', 'Direct', '--streaming'])

standard_options = options.view_as(StandardOptions)

if standard_options.streaming:

# ... start a streaming job ...

Note that options objects may have multiple views, and modifications

of values in any view-object will apply to current object and other

view-objects.

Args:

cls: PipelineOptions class or any of its subclasses.

Returns:

An instance of cls that is initialized using options contained in current

object.

"""

view = cls(self._flags)

for option_name in view._visible_option_list():

# Initialize values of keys defined by a cls.

#

# Note that we do initialization only once per key to make sure that

# values in _all_options dict are not-recreated with each new view.

# This is important to make sure that values of multi-options keys are

# backed by the same list across multiple views, and that any overrides of

# pipeline options already stored in _all_options are preserved.

if option_name not in self._all_options:

self._all_options[option_name] = getattr(

view._visible_options, option_name)

# Note that views will still store _all_options of the source object.

view._all_options = self._all_options

return view

def is_compat_version_prior_to(self, breaking_change_version):

"""Check if update_compatibility_version is prior to a breaking change.

Returns True if the pipeline should use old behavior (i.e., the

update_compatibility_version is set and is earlier than the given version).

Returns False if update_compatibility_version is not set or is >= the

breaking change version.

Args:

breaking_change_version: Version string (e.g., "2.72.0") at which

the breaking change was introduced.

"""

v1 = self.view_as(StreamingOptions).update_compatibility_version

if v1 is None:

return False

v1_parts = (v1.split('.') + ['0', '0', '0'])[:3]

v2_parts = (breaking_change_version.split('.') + ['0', '0', '0'])[:3]

return tuple(map(int, v1_parts)) < tuple(map(int, v2_parts))

def _visible_option_list(self) -> List[str]:

return sorted(

option for option in dir(self._visible_options) if option[0] != '_')

def __dir__(self) -> List[str]:

return sorted(

dir(type(self)) + list(self.__dict__) + self._visible_option_list())

def __getattr__(self, name):

# Special methods which may be accessed before the object is

# fully constructed (e.g. in unpickling).

if name[:2] == name[-2:] == '__':

return object.__getattribute__(self, name)

elif name in self._visible_option_list():

return self._all_options[name]

else:

raise AttributeError(

"'%s' object has no attribute '%s'" % (type(self).__name__, name))

def __setattr__(self, name, value):

if name in ('_flags', '_all_options', '_visible_options'):

super().__setattr__(name, value)

elif name in self._visible_option_list():

self._all_options[name] = value

else:

raise AttributeError(

"'%s' object has no attribute '%s'" % (type(self).__name__, name))

def __str__(self):

return '%s(%s)' % (

type(self).__name__,

', '.join(

'%s=%s' % (option, getattr(self, option))

for option in self._visible_option_list()))

class StandardOptions(PipelineOptions):

DEFAULT_RUNNER = 'DirectRunner'

ALL_KNOWN_RUNNERS = (

'apache_beam.runners.dataflow.dataflow_runner.DataflowRunner',

'apache_beam.runners.direct.direct_runner.BundleBasedDirectRunner',

'apache_beam.runners.direct.direct_runner.DirectRunner',

'apache_beam.runners.direct.direct_runner.SwitchingDirectRunner',

'apache_beam.runners.interactive.interactive_runner.InteractiveRunner',

'apache_beam.runners.portability.flink_runner.FlinkRunner',

'apache_beam.runners.portability.fn_api_runner.FnApiRunner',

'apache_beam.runners.portability.portable_runner.PortableRunner',

'apache_beam.runners.portability.prism_runner.PrismRunner',

'apache_beam.runners.portability.spark_runner.SparkRunner',

'apache_beam.runners.test.TestDirectRunner',

'apache_beam.runners.test.TestDataflowRunner',

)

LOCAL_RUNNERS = (

'BundleBasedDirectRunner',

'DirectRunner',

'SwitchingDirectRunner',

'FnApiRunner',

'PrismRunner',

'TestDirectRunner',

)

KNOWN_RUNNER_NAMES = [path.split('.')[-1] for path in ALL_KNOWN_RUNNERS]

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument(

'--runner',

help=(

'Pipeline runner used to execute the workflow. Valid values are '

'one of %s, or the fully qualified name of a PipelineRunner '

'subclass. If unspecified, defaults to %s.' %

(', '.join(cls.KNOWN_RUNNER_NAMES), cls.DEFAULT_RUNNER)))

# Whether to enable streaming mode.

parser.add_argument(

'--streaming',

default=False,

action='store_true',

help='Whether to enable streaming mode.')

parser.add_argument(

'--resource_hint',

'--resource_hints',

dest='resource_hints',

action='append',

default=[],

help=(

'Resource hint to set in the pipeline execution environment.'

'Hints specified via this option override hints specified '

'at transform level. Interpretation of hints is defined by '

'Beam runners.'))

parser.add_argument(

'--auto_unique_labels',

default=False,

action='store_true',

help='Whether to automatically generate unique transform labels '

'for every transform. The default behavior is to raise an '

'exception if a transform is created with a non-unique label. '

'Using --auto_unique_labels could cause data loss when '

'updating a pipeline or reloading the job state. '

'This is not recommended for streaming jobs.')

parser.add_argument(

'--no_wait_until_finish',

default=False,

action='store_true',

help='By default, the "with" statement waits for the job to '

'complete. Set this flag to bypass this behavior and continue '

'execution immediately')

class StreamingOptions(PipelineOptions):

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument(

'--update_compatibility_version',

default=None,

nargs='?',

const=None,

help='Attempt to produce a pipeline compatible with the given prior '

'version of the Beam SDK. '

'See for example, https://cloud.google.com/dataflow/docs/guides/'

'updating-a-pipeline')

class CrossLanguageOptions(PipelineOptions):

@staticmethod

def _beam_services_from_enviroment():

return json.loads(os.environ.get('BEAM_SERVICE_OVERRIDES') or '{}')

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument(

'--beam_services',

type=lambda s: {

**cls._beam_services_from_enviroment(), **json.loads(s)

},

default=cls._beam_services_from_enviroment(),

help=(

'For convenience, Beam provides the ability to automatically '

'download and start various services (such as expansion services) '

'used at pipeline construction and execution. These services are '

'identified by gradle target. This option provides the ability to '

'use pre-started services or non-default pre-existing artifacts to '

'start the given service. '

'Should be a json mapping of gradle build targets to pre-built '

'artifacts (e.g. jar files) or expansion endpoints '

'(e.g. host:port). Defaults to the value of BEAM_SERVICE_OVERRIDES '

'from the environment.'))

parser.add_argument(

'--use_transform_service',

default=False,

action='store_true',

help='Use the Docker-composed-based transform service when expanding '

'cross-language transforms.')

def additional_option_ptransform_fn():

beam.transforms.ptransform.ptransform_fn_typehints_enabled = True

# Optional type checks that aren't enabled by default.

additional_type_checks: Dict[str, Callable[[], None]] = {

'ptransform_fn': additional_option_ptransform_fn,

}

def enable_all_additional_type_checks():

"""Same as passing --type_check_additional=all."""

for f in additional_type_checks.values():

f()

class TypeOptions(PipelineOptions):

@classmethod

def _add_argparse_args(cls, parser):

# TODO(laolu): Add a type inferencing option here once implemented.

parser.add_argument(

'--type_check_strictness',

default='DEFAULT_TO_ANY',

choices=['ALL_REQUIRED', 'DEFAULT_TO_ANY'],

help='The level of exhaustive manual type-hint '

'annotation required')

parser.add_argument(

'--type_check_additional',

default='',

help='Comma separated list of additional type checking features to '

'enable. Options: all, ptransform_fn. For details see:'

'https://beam.apache.org/documentation/sdks/python-type-safety/')

parser.add_argument(

'--no_pipeline_type_check',

dest='pipeline_type_check',

action='store_false',

help='Disable type checking at pipeline construction '

'time')

parser.add_argument(

'--runtime_type_check',

default=False,

action='store_true',

help='Enable type checking at pipeline execution '

'time. NOTE: only supported with the '

'DirectRunner')

parser.add_argument(

'--performance_runtime_type_check',

default=False,

action='store_true',

help='Enable faster type checking via sampling at pipeline execution '

'time. NOTE: only supported with portable runners '

'(including the DirectRunner)')

parser.add_argument(

'--allow_non_deterministic_key_coders',

default=False,

action='store_true',

help='Use non-deterministic coders (such as pickling) for key-grouping '

'operations such as GroupByKey. This is unsafe, as runners may group '

'keys based on their encoded bytes, but is available for backwards '

'compatibility. See BEAM-11719.')

parser.add_argument(

'--allow_unsafe_triggers',

default=False,

action='store_true',

help='Allow the use of unsafe triggers. Unsafe triggers have the '

'potential to cause data loss due to finishing and/or never having '

'their condition met. Some operations, such as GroupByKey, disallow '

'this. This exists for cases where such loss is acceptable and for '

'backwards compatibility. See BEAM-9487.')

parser.add_argument(

'--force_cloudpickle_deterministic_coders',

default=False,

action='store_true',

help=(

'Force the use of cloudpickle-based deterministic coders '

'instead of dill-based coders, even when '

'update_compatibility_version would normally trigger dill usage '

'for backward compatibility. This flag overrides automatic coder '

'selection to always use the modern cloudpickle serialization '

' path. Warning: May break pipeline update compatibility with '

' SDK versions prior to 2.68.0.'))

def validate(self, unused_validator):

errors = []

if beam.version.__version__ >= '3':

errors.append(

'Update --type_check_additional default to include all '

'available additional checks at Beam 3.0 release time.')

keys = self.type_check_additional.split(',')

for key in keys:

if not key:

continue

elif key == 'all':

enable_all_additional_type_checks()

elif key in additional_type_checks:

additional_type_checks[key]()

else:

errors.append('Unrecognized --type_check_additional feature: %s' % key)

return errors

class DirectOptions(PipelineOptions):

"""DirectRunner-specific execution options."""

@classmethod

def _add_argparse_args(cls, parser):

parser.add_argument(

'--no_direct_runner_use_stacked_bundle',

action='store_false',

dest='direct_runner_use_stacked_bundle',

help='DirectRunner uses stacked WindowedValues within a Bundle for '

'memory optimization. Set --no_direct_runner_use_stacked_bundle to '

'avoid it.')

parser.add_argument(

'--direct_runner_bundle_repeat',

type=int,

default=0,

help='replay every bundle this many extra times, for profiling'

'and debugging')

parser.add_argument(

'--direct_num_workers',

type=int,

default=1,

help='number of parallel running workers.')

parser.add_argument(

'--direct_running_mode',

default='in_memory',

choices=['in_memory', 'multi_threading', 'multi_processing'],

help='Workers running environment.')

parser.add_argument(

'--direct_embed_docker_python',

default=False,

action='store_true',

dest='direct_embed_docker_python',

help='DirectRunner uses the embedded Python environment when '

'the default Python docker environment is specified.')

parser.add_argument(

'--direct_test_splits',

default={},

type=json.loads,

help='Split test configuration of the json form '

'{"step_name": {"timings": [...], "fractions": [...]}, ...} '

'where step_name is the name of a step controlling the stage to which '

'splits will be sent, timings is a list of floating-point times '

'(in seconds) at which the split requests will be sent, and '

'fractions is a corresponding list of floating points to use in the '

'split requests themselves.')