Source code for apache_beam.transforms.resources

# 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.
#

"""A module for defining resource requirements for execution of transforms.

Pipeline authors can use resource hints to provide additional information to
runners about the desired aspects of the execution environment.

Resource hints can be specified on a transform level for parts of the pipeline,
or globally via --resource_hint pipeline option.

See also: PTransforms.with_resource_hints().
"""

import re
from typing import Any
from typing import Dict
from typing import Mapping
from typing import Optional

from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.portability.common_urns import resource_hints

__all__ = [
    'ResourceHint',
    'AcceleratorHint',
    'MinRamHint',
    'CpuCountHint',
    'merge_resource_hints',
    'parse_resource_hints',
    'resource_hints_from_options',
]


[docs] class ResourceHint: """A superclass to define resource hints.""" # A unique URN, one per Resource Hint class. urn: Optional[str] = None _urn_to_known_hints: Dict[str, type] = {} _name_to_known_hints: Dict[str, type] = {}
[docs] @classmethod def parse(cls, value: str) -> Dict[str, bytes]: """Describes how to parse the hint. Override to specify a custom parsing logic.""" assert cls.urn is not None # Override this method to have a custom parsing logic. return {cls.urn: ResourceHint._parse_str(value)}
[docs] @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: """Reconciles values of a hint when the hint specified on a transform is also defined in an outer context, for example on a composite transform, or specified in the transform's execution environment. Override to specify a custom merging logic. """ # Defaults to the inner value as it is the most specific one. return inner_value
[docs] @staticmethod def get_by_urn(urn): return ResourceHint._urn_to_known_hints[urn]
[docs] @staticmethod def get_by_name(name): return ResourceHint._name_to_known_hints[name]
[docs] @staticmethod def is_registered(name): return name in ResourceHint._name_to_known_hints
[docs] @staticmethod def register_resource_hint(hint_name: str, hint_class: type) -> None: assert issubclass(hint_class, ResourceHint) assert hint_class.urn is not None ResourceHint._name_to_known_hints[hint_name] = hint_class ResourceHint._urn_to_known_hints[hint_class.urn] = hint_class
@staticmethod def _parse_str(value): if not isinstance(value, str): raise ValueError("Input must be a string.") return value.encode('ascii') @staticmethod def _parse_int(value): if isinstance(value, str): value = int(value) if not isinstance(value, int): raise ValueError("Input must be an integer.") return str(value).encode('ascii') @staticmethod def _parse_storage_size_str(value): """Parses a human-friendly storage size string into a number of bytes. """ if isinstance(value, int): return ResourceHint._parse_int(value) if not isinstance(value, str): raise ValueError("Input must be a string or integer.") value = value.strip().replace(" ", "") units = { 'PiB': 2**50, 'TiB': 2**40, 'GiB': 2**30, 'MiB': 2**20, 'KiB': 2**10, 'PB': 10**15, 'TB': 10**12, 'GB': 10**9, 'MB': 10**6, 'KB': 10**3, 'B': 1, } match = re.match(r'.*?(\D+)$', value) if not match: raise ValueError("Unrecognized value pattern.") suffix = match.group(1) if suffix not in units: raise ValueError("Unrecognized unit.") multiplier = units[suffix] value = value[:-len(suffix)] return str(round(float(value) * multiplier)).encode('ascii') @staticmethod def _use_max(v1, v2): return str(max(int(v1), int(v2))).encode('ascii')
[docs] class AcceleratorHint(ResourceHint): """Describes desired hardware accelerators in execution environment.""" urn = resource_hints.ACCELERATOR.urn
ResourceHint.register_resource_hint('accelerator', AcceleratorHint)
[docs] class MinRamHint(ResourceHint): """Describes min RAM requirements for transform's execution environment.""" urn = resource_hints.MIN_RAM_BYTES.urn
[docs] @classmethod def parse(cls, value: str) -> Dict[str, bytes]: return {cls.urn: ResourceHint._parse_storage_size_str(value)}
[docs] @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value)
ResourceHint.register_resource_hint('min_ram', MinRamHint) # Alias for interoperability with SDKs preferring camelCase. ResourceHint.register_resource_hint('minRam', MinRamHint)
[docs] class CpuCountHint(ResourceHint): """Describes number of CPUs available in transform's execution environment.""" urn = resource_hints.CPU_COUNT.urn
[docs] @classmethod def get_merged_value(cls, outer_value: bytes, inner_value: bytes) -> bytes: return ResourceHint._use_max(outer_value, inner_value)
ResourceHint.register_resource_hint('cpu_count', CpuCountHint) # Alias for interoperability with SDKs preferring camelCase. ResourceHint.register_resource_hint('cpuCount', CpuCountHint)
[docs] def parse_resource_hints(hints: Dict[Any, Any]) -> Dict[str, bytes]: parsed_hints = {} for hint, value in hints.items(): try: hint_cls = ResourceHint.get_by_name(hint) try: parsed_hints.update(hint_cls.parse(value)) except ValueError: raise ValueError(f"Resource hint {hint} has invalid value {value}.") except KeyError: raise ValueError(f"Unknown resource hint: {hint}.") return parsed_hints
[docs] def resource_hints_from_options( options: Optional[PipelineOptions]) -> Dict[str, bytes]: if options is None: return {} hints = {} option_specified_hints = options.view_as(StandardOptions).resource_hints for hint in option_specified_hints: if '=' in hint: k, v = hint.split('=', maxsplit=1) hints[k] = v else: hints[hint] = None return parse_resource_hints(hints)
[docs] def merge_resource_hints( outer_hints: Mapping[str, bytes], inner_hints: Mapping[str, bytes]) -> Dict[str, bytes]: merged_hints = dict(inner_hints) for urn, outer_value in outer_hints.items(): if urn in inner_hints: merged_value = ResourceHint.get_by_urn(urn).get_merged_value( outer_value=outer_value, inner_value=inner_hints[urn]) else: merged_value = outer_value merged_hints[urn] = merged_value return merged_hints