sdks/python/apache_beam/transforms/environments.py

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

"""Environments concepts.

For internal use only. No backwards compatibility guarantees."""

# pytype: skip-file

import json
import logging
import sys
import tempfile
from collections.abc import Callable
from collections.abc import Iterable
from collections.abc import Iterator
from collections.abc import Mapping
from types import MappingProxyType
from typing import TYPE_CHECKING
from typing import Any
from typing import Optional
from typing import TypeVar
from typing import Union
from typing import overload

from google.protobuf import message

from apache_beam import coders
from apache_beam.options.pipeline_options import PortableOptions
from apache_beam.options.pipeline_options import SetupOptions
from apache_beam.options.pipeline_options import StandardOptions
from apache_beam.portability import common_urns
from apache_beam.portability import python_urns
from apache_beam.portability.api import beam_runner_api_pb2
from apache_beam.portability.api import endpoints_pb2
from apache_beam.runners.portability import stager
from apache_beam.transforms.resources import resource_hints_from_options
from apache_beam.utils import proto_utils

if TYPE_CHECKING:
  from apache_beam.options.pipeline_options import PipelineOptions
  from apache_beam.runners.pipeline_context import PipelineContext

__all__ = [
    'Environment',
    'AnyOfEnvironment',
    'DefaultEnvironment',
    'DockerEnvironment',
    'ProcessEnvironment',
    'ExternalEnvironment',
    'EmbeddedPythonEnvironment',
    'EmbeddedPythonGrpcEnvironment',
    'SubprocessSDKEnvironment',
    'PyPIArtifactRegistry'
]

T = TypeVar('T')
EnvironmentT = TypeVar('EnvironmentT', bound='Environment')
ConstructorFn = Callable[[
    Optional[Any],
    Iterable[str],
    Iterable[beam_runner_api_pb2.ArtifactInformation],
    Mapping[str, bytes],
    'PipelineContext'
],
                         Any]


def looks_like_json(s):
  import re
  return re.match(r'\s*\{.*\}\s*$', s)


APACHE_BEAM_DOCKER_IMAGE_PREFIX = 'apache/beam'
APACHE_BEAM_JAVA_CONTAINER_NAME_PREFIX = 'beam_java'
SDK_VERSION_CAPABILITY_PREFIX = 'beam:version:sdk_base:'


def is_apache_beam_container(container_image):
  return container_image and container_image.startswith(
      APACHE_BEAM_DOCKER_IMAGE_PREFIX)


class Environment(object):
  """Abstract base class for environments.

  Represents a type and configuration of environment.
  Each type of Environment should have a unique urn.

  For internal use only. No backwards compatibility guarantees.
  """

  _known_urns = {}  # type: dict[str, tuple[Optional[type], ConstructorFn]]
  _urn_to_env_cls = {}  # type: dict[str, type]

  def __init__(
      self,
      capabilities=(),  # type: Iterable[str]
      artifacts=(),  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints=None,  # type: Optional[Mapping[str, bytes]]
  ):
    # type: (...) -> None
    self._capabilities = capabilities
    self._artifacts = sorted(artifacts, key=lambda x: x.SerializeToString())
    # Hints on created environments should be immutable since pipeline context
    # stores environments in hash maps and we use hints to compute the hash.
    self._resource_hints = MappingProxyType(
        dict(resource_hints) if resource_hints else {})

  def __eq__(self, other):
    return (
        self.__class__ == other.__class__ and
        self._artifacts == other._artifacts
        # Assuming that we don't have instances of the same Environment subclass
        # with different set of capabilities.
        and self._resource_hints == other._resource_hints)

  def __hash__(self):
    # type: () -> int
    return hash((self.__class__, frozenset(self._resource_hints.items())))

  def artifacts(self):
    # type: () -> Iterable[beam_runner_api_pb2.ArtifactInformation]
    return self._artifacts

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, Optional[Union[message.Message, bytes, str]]]
    raise NotImplementedError

  def capabilities(self):
    # type: () -> Iterable[str]
    return self._capabilities

  def resource_hints(self):
    # type: () -> Mapping[str, bytes]
    return self._resource_hints

  @classmethod
  @overload
  def register_urn(
      cls,
      urn,  # type: str
      parameter_type,  # type: type[T]
  ):
    # type: (...) -> Callable[[Union[type, Callable[[T, Iterable[str], PipelineContext], Any]]], Callable[[T, Iterable[str], PipelineContext], Any]]
    pass

  @classmethod
  @overload
  def register_urn(
      cls,
      urn,  # type: str
      parameter_type,  # type: None
  ):
    # type: (...) -> Callable[[Union[type, Callable[[bytes, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any]]], Callable[[bytes, Iterable[str], PipelineContext], Any]]
    pass

  @classmethod
  @overload
  def register_urn(
      cls,
      urn,  # type: str
      parameter_type,  # type: type[T]
      constructor  # type: Callable[[T, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any]
  ):
    # type: (...) -> None
    pass

  @classmethod
  @overload
  def register_urn(
      cls,
      urn,  # type: str
      parameter_type,  # type: None
      constructor  # type: Callable[[bytes, Iterable[str], Iterable[beam_runner_api_pb2.ArtifactInformation], PipelineContext], Any]
  ):
    # type: (...) -> None
    pass

  @classmethod
  def register_urn(cls, urn, parameter_type, constructor=None):
    def register(constructor):
      if isinstance(constructor, type):
        constructor.from_runner_api_parameter = register(
            constructor.from_runner_api_parameter)
        # register environment urn to environment class
        cls._urn_to_env_cls[urn] = constructor
        return constructor

      else:
        cls._known_urns[urn] = parameter_type, constructor
        return staticmethod(constructor)

    if constructor:
      # Used as a statement.
      register(constructor)
    else:
      # Used as a decorator.
      return register

  @classmethod
  def get_env_cls_from_urn(cls, urn):
    # type: (str) -> type[Environment]
    return cls._urn_to_env_cls[urn]

  def to_runner_api(self, context):
    # type: (PipelineContext) -> beam_runner_api_pb2.Environment
    urn, typed_param = self.to_runner_api_parameter(context)
    return beam_runner_api_pb2.Environment(
        urn=urn,
        payload=typed_param.SerializeToString() if isinstance(
            typed_param, message.Message) else typed_param if
        (isinstance(typed_param, bytes) or
         typed_param is None) else typed_param.encode('utf-8'),
        capabilities=self.capabilities(),
        dependencies=self.artifacts(),
        resource_hints=self.resource_hints())

  @classmethod
  def from_runner_api(
      cls,
      proto,  # type: Optional[beam_runner_api_pb2.Environment]
      context  # type: PipelineContext
  ):
    # type: (...) -> Optional[Environment]
    if proto is None or not proto.urn:
      return None
    parameter_type, constructor = cls._known_urns[proto.urn]

    return constructor(
        proto_utils.parse_Bytes(proto.payload, parameter_type),
        proto.capabilities,
        proto.dependencies,
        proto.resource_hints,
        context)

  @classmethod
  def from_options(cls, options):
    # type: (type[EnvironmentT], PortableOptions) -> EnvironmentT

    """Creates an Environment object from PortableOptions.

    Args:
      options: The PortableOptions object.
    """
    if cls != Environment:
      raise NotImplementedError

    portable_options = options.view_as(PortableOptions)
    environment_type = portable_options.environment_type
    if not environment_type:
      environment_urn = common_urns.environments.DOCKER.urn
    elif environment_type.startswith('beam:env:'):
      environment_urn = environment_type
    elif environment_type == 'LOOPBACK':
      environment_urn = python_urns.EMBEDDED_PYTHON_LOOPBACK
    else:
      try:
        environment_urn = getattr(
            common_urns.environments, environment_type).urn
      except AttributeError:
        raise ValueError('Unknown environment type: %s' % environment_type)

    env_class = Environment.get_env_cls_from_urn(environment_urn)
    return env_class.from_options(portable_options)  # type: ignore


@Environment.register_urn(common_urns.environments.DEFAULT.urn, None)
class DefaultEnvironment(Environment):
  """Used as a stub when context is missing a default environment."""
  def to_runner_api_parameter(self, context):
    return common_urns.environments.DEFAULT.urn, None

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: beam_runner_api_pb2.DockerPayload
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> DefaultEnvironment
    return DefaultEnvironment(
        capabilities=capabilities,
        artifacts=artifacts,
        resource_hints=resource_hints)


@Environment.register_urn(
    common_urns.environments.DOCKER.urn, beam_runner_api_pb2.DockerPayload)
class DockerEnvironment(Environment):
  def __init__(
      self,
      container_image=None,  # type: Optional[str]
      capabilities=(),  # type: Iterable[str]
      artifacts=(),  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints=None,  # type: Optional[Mapping[str, bytes]]
  ):
    super().__init__(capabilities, artifacts, resource_hints)
    if container_image:
      logging.debug(
          'Using provided Python SDK container image: %s' % (container_image))
      self.container_image = container_image
    else:
      logging.debug('No image given, using default Python SDK image')
      self.container_image = self.default_docker_image()

    logging.debug(
        'Python SDK container image set to "%s" for Docker environment' %
        (self.container_image))

  def __eq__(self, other):
    return (
        super().__eq__(other) and self.container_image == other.container_image)

  def __hash__(self):
    return hash((super().__hash__(), self.container_image))

  def __repr__(self):
    return 'DockerEnvironment(container_image=%s)' % self.container_image

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.DockerPayload]
    return (
        common_urns.environments.DOCKER.urn,
        beam_runner_api_pb2.DockerPayload(container_image=self.container_image))

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: beam_runner_api_pb2.DockerPayload
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> DockerEnvironment
    return DockerEnvironment(
        container_image=payload.container_image,
        capabilities=capabilities,
        artifacts=artifacts,
        resource_hints=resource_hints)

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> DockerEnvironment
    if options.view_as(SetupOptions).prebuild_sdk_container_engine:
      # Imported here to avoid circular dependencies.
      # pylint: disable=wrong-import-order, wrong-import-position
      from apache_beam.runners.portability.sdk_container_builder import SdkContainerImageBuilder
      prebuilt_container_image = SdkContainerImageBuilder.build_container_image(
          options)
      return cls.from_container_image(
          container_image=prebuilt_container_image,
          artifacts=python_sdk_dependencies(options),
          resource_hints=resource_hints_from_options(options),
      )
    return cls.from_container_image(
        container_image=options.lookup_environment_option(
            'docker_container_image') or options.environment_config,
        artifacts=python_sdk_dependencies(options),
        resource_hints=resource_hints_from_options(options),
    )

  @classmethod
  def from_container_image(
      cls, container_image, artifacts=(), resource_hints=None):
    # type: (str, Iterable[beam_runner_api_pb2.ArtifactInformation], Optional[Mapping[str, bytes]]) -> DockerEnvironment
    return cls(
        container_image=container_image,
        capabilities=python_sdk_docker_capabilities(),
        artifacts=artifacts,
        resource_hints=resource_hints)

  @staticmethod
  def default_docker_image():
    # type: () -> str
    from apache_beam import version as beam_version

    sdk_version = beam_version.__version__
    version_suffix = '.'.join([str(i) for i in sys.version_info[0:2]])

    image = (
        APACHE_BEAM_DOCKER_IMAGE_PREFIX +
        '_python{version_suffix}_sdk:{tag}'.format(
            version_suffix=version_suffix, tag=sdk_version))
    return image


@Environment.register_urn(
    common_urns.environments.PROCESS.urn, beam_runner_api_pb2.ProcessPayload)
class ProcessEnvironment(Environment):
  def __init__(
      self,
      command,  # type: str
      os='',  # type: str
      arch='',  # type: str
      env=None,  # type: Optional[Mapping[str, str]]
      capabilities=(),  # type: Iterable[str]
      artifacts=(),  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints=None,  # type: Optional[Mapping[str, bytes]]
  ):
    # type: (...) -> None
    super().__init__(capabilities, artifacts, resource_hints)
    self.command = command
    self.os = os
    self.arch = arch
    self.env = env or {}

  def __eq__(self, other):
    return (
        super().__eq__(other) and self.command == other.command and
        self.os == other.os and self.arch == other.arch and
        self.env == other.env)

  def __hash__(self):
    # type: () -> int
    return hash((
        super().__hash__(),
        self.command,
        self.os,
        self.arch,
        frozenset(self.env.items())))

  def __repr__(self):
    # type: () -> str
    repr_parts = ['command=%s' % self.command]
    if self.os:
      repr_parts.append('os=%s' % self.os)
    if self.arch:
      repr_parts.append('arch=%s' % self.arch)
    repr_parts.append('env=%s' % self.env)
    return 'ProcessEnvironment(%s)' % ','.join(repr_parts)

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.ProcessPayload]
    return (
        common_urns.environments.PROCESS.urn,
        beam_runner_api_pb2.ProcessPayload(
            os=self.os, arch=self.arch, command=self.command, env=self.env))

  @staticmethod
  def from_runner_api_parameter(
      payload,
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> ProcessEnvironment
    return ProcessEnvironment(
        command=payload.command,
        os=payload.os,
        arch=payload.arch,
        env=payload.env,
        capabilities=capabilities,
        artifacts=artifacts,
        resource_hints=resource_hints,
    )

  @staticmethod
  def parse_environment_variables(variables):
    env = {}
    for var in variables:
      try:
        name, value = var.split('=', 1)
        env[name] = value
      except ValueError:
        raise ValueError(
            'Invalid process_variables "%s" (expected assignment in the '
            'form "FOO=bar").' % var)
    return env

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> ProcessEnvironment
    if options.environment_config:
      config = json.loads(options.environment_config)
      return cls(
          config.get('command'),
          os=config.get('os', ''),
          arch=config.get('arch', ''),
          env=config.get('env', ''),
          capabilities=python_sdk_capabilities(),
          artifacts=python_sdk_dependencies(options),
          resource_hints=resource_hints_from_options(options),
      )
    env = cls.parse_environment_variables(
        options.lookup_environment_option('process_variables').split(',')
        if options.lookup_environment_option('process_variables') else [])
    return cls(
        options.lookup_environment_option('process_command'),
        env=env,
        capabilities=python_sdk_capabilities(),
        artifacts=python_sdk_dependencies(options),
        resource_hints=resource_hints_from_options(options),
    )


@Environment.register_urn(
    common_urns.environments.EXTERNAL.urn, beam_runner_api_pb2.ExternalPayload)
class ExternalEnvironment(Environment):
  def __init__(
      self,
      url,  # type: str
      params=None,  # type: Optional[Mapping[str, str]]
      capabilities=(),  # type: Iterable[str]
      artifacts=(),  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints=None,  # type: Optional[Mapping[str, bytes]]
  ):
    super().__init__(capabilities, artifacts, resource_hints)
    self.url = url
    self.params = params

  def __eq__(self, other):
    return (
        super().__eq__(other) and self.url == other.url and
        self.params == other.params)

  def __hash__(self):
    # type: () -> int
    return hash((
        super().__hash__(),
        self.url,
        frozenset(self.params.items()) if self.params is not None else None))

  def __repr__(self):
    # type: () -> str
    return 'ExternalEnvironment(url=%s,params=%s)' % (self.url, self.params)

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.ExternalPayload]
    return (
        common_urns.environments.EXTERNAL.urn,
        beam_runner_api_pb2.ExternalPayload(
            endpoint=endpoints_pb2.ApiServiceDescriptor(url=self.url),
            params=self.params))

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: beam_runner_api_pb2.ExternalPayload
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> ExternalEnvironment
    return ExternalEnvironment(
        payload.endpoint.url,
        params=payload.params or None,
        capabilities=capabilities,
        artifacts=artifacts,
        resource_hints=resource_hints)

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> ExternalEnvironment
    if looks_like_json(options.environment_config):
      config = json.loads(options.environment_config)
      url = config.get('url')
      if not url:
        raise ValueError('External environment endpoint must be set.')
      params = config.get('params')
    elif options.environment_config:
      url = options.environment_config
      params = None
    else:
      url = options.lookup_environment_option('external_service_address')
      params = None

    return cls(
        url,
        params=params,
        capabilities=python_sdk_capabilities(),
        artifacts=python_sdk_dependencies(options),
        resource_hints=resource_hints_from_options(options))


def expand_anyof_environments(env_proto):
  if env_proto.urn == common_urns.environments.ANYOF.urn:
    for alt in beam_runner_api_pb2.AnyOfEnvironmentPayload.FromString(
        env_proto.payload).environments:
      yield from expand_anyof_environments(alt)
  else:
    yield env_proto


def resolve_anyof_environment(env_proto, *preferred_types):
  all_environments = list(expand_anyof_environments(env_proto))
  for preferred_type in preferred_types:
    for env in all_environments:
      if env.urn == preferred_type:
        return env
  return all_environments[0]


@Environment.register_urn(python_urns.EMBEDDED_PYTHON, None)
class EmbeddedPythonEnvironment(Environment):
  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, None]
    return python_urns.EMBEDDED_PYTHON, None

  @staticmethod
  def from_runner_api_parameter(
      unused_payload,  # type: None
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> EmbeddedPythonEnvironment
    return EmbeddedPythonEnvironment(capabilities, artifacts, resource_hints)

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> EmbeddedPythonEnvironment
    return cls(
        capabilities=python_sdk_capabilities(),
        artifacts=python_sdk_dependencies(options),
        resource_hints=resource_hints_from_options(options),
    )

  @classmethod
  def default(cls):
    # type: () -> EmbeddedPythonEnvironment
    return cls(capabilities=python_sdk_capabilities(), artifacts=())


@Environment.register_urn(python_urns.EMBEDDED_PYTHON_GRPC, bytes)
class EmbeddedPythonGrpcEnvironment(Environment):
  def __init__(
      self,
      state_cache_size=None,
      data_buffer_time_limit_ms=None,
      capabilities=(),
      artifacts=(),
      resource_hints=None,
  ):
    super().__init__(capabilities, artifacts, resource_hints)
    self.state_cache_size = state_cache_size
    self.data_buffer_time_limit_ms = data_buffer_time_limit_ms

  def __eq__(self, other):
    return (
        super().__eq__(other) and
        self.state_cache_size == other.state_cache_size and
        self.data_buffer_time_limit_ms == other.data_buffer_time_limit_ms)

  def __hash__(self):
    # type: () -> int
    return hash((
        super().__hash__(),
        self.state_cache_size,
        self.data_buffer_time_limit_ms))

  def __repr__(self):
    # type: () -> str
    repr_parts = []
    if not self.state_cache_size is None:
      repr_parts.append('state_cache_size=%d' % self.state_cache_size)
    if not self.data_buffer_time_limit_ms is None:
      repr_parts.append(
          'data_buffer_time_limit_ms=%d' % self.data_buffer_time_limit_ms)
    return 'EmbeddedPythonGrpcEnvironment(%s)' % ','.join(repr_parts)

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, bytes]
    params = {}
    if self.state_cache_size is not None:
      params['state_cache_size'] = self.state_cache_size
    if self.data_buffer_time_limit_ms is not None:
      params['data_buffer_time_limit_ms'] = self.data_buffer_time_limit_ms
    payload = json.dumps(params).encode('utf-8')
    return python_urns.EMBEDDED_PYTHON_GRPC, payload

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: bytes
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> EmbeddedPythonGrpcEnvironment
    if payload:
      config = EmbeddedPythonGrpcEnvironment.parse_config(
          payload.decode('utf-8'))
      return EmbeddedPythonGrpcEnvironment(
          state_cache_size=config.get('state_cache_size'),
          data_buffer_time_limit_ms=config.get('data_buffer_time_limit_ms'),
          capabilities=capabilities,
          artifacts=artifacts,
          resource_hints=resource_hints)
    else:
      return EmbeddedPythonGrpcEnvironment()

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> EmbeddedPythonGrpcEnvironment
    if options.environment_config:
      config = EmbeddedPythonGrpcEnvironment.parse_config(
          options.environment_config)
      return cls(
          state_cache_size=config.get('state_cache_size'),
          data_buffer_time_limit_ms=config.get('data_buffer_time_limit_ms'),
          capabilities=python_sdk_capabilities(),
          artifacts=python_sdk_dependencies(options))
    else:
      return cls(
          capabilities=python_sdk_capabilities(),
          artifacts=python_sdk_dependencies(options),
          resource_hints=resource_hints_from_options(options))

  @staticmethod
  def parse_config(s):
    # type: (str) -> dict[str, Any]
    if looks_like_json(s):
      config_dict = json.loads(s)
      if 'state_cache_size' in config_dict:
        config_dict['state_cache_size'] = int(config_dict['state_cache_size'])

      if 'data_buffer_time_limit_ms' in config_dict:
        config_dict['data_buffer_time_limit_ms'] = \
          int(config_dict['data_buffer_time_limit_ms'])
      return config_dict
    else:
      return {'state_cache_size': int(s)}

  @classmethod
  def default(cls):
    # type: () -> EmbeddedPythonGrpcEnvironment
    return cls(capabilities=python_sdk_capabilities(), artifacts=())


@Environment.register_urn(python_urns.EMBEDDED_PYTHON_LOOPBACK, None)
class PythonLoopbackEnvironment(EmbeddedPythonEnvironment):
  """Used as a stub when the loopback worker has not yet been started."""
  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, None]
    return python_urns.EMBEDDED_PYTHON_LOOPBACK, None

  @staticmethod
  def from_runner_api_parameter(
      unused_payload,  # type: None
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> PythonLoopbackEnvironment
    return PythonLoopbackEnvironment(
        capabilities=capabilities,
        artifacts=artifacts,
        resource_hints=resource_hints)


@Environment.register_urn(python_urns.SUBPROCESS_SDK, bytes)
class SubprocessSDKEnvironment(Environment):
  def __init__(
      self,
      command_string,  # type: str
      capabilities=(),  # type: Iterable[str]
      artifacts=(),  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints=None,  # type: Optional[Mapping[str, bytes]]
  ):
    super().__init__(capabilities, artifacts, resource_hints)
    self.command_string = command_string

  def __eq__(self, other):
    return (
        super().__eq__(other) and self.command_string == other.command_string)

  def __hash__(self):
    # type: () -> int
    return hash((super().__hash__(), self.command_string))

  def __repr__(self):
    # type: () -> str
    return 'SubprocessSDKEnvironment(command_string=%s)' % self.command_string

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, bytes]
    return python_urns.SUBPROCESS_SDK, self.command_string.encode('utf-8')

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: bytes
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> SubprocessSDKEnvironment
    return SubprocessSDKEnvironment(
        payload.decode('utf-8'), capabilities, artifacts, resource_hints)

  @classmethod
  def from_options(cls, options):
    # type: (PortableOptions) -> SubprocessSDKEnvironment
    return cls(
        options.environment_config,
        capabilities=python_sdk_capabilities(),
        artifacts=python_sdk_dependencies(options),
        resource_hints=resource_hints_from_options(options))

  @classmethod
  def from_command_string(cls, command_string):
    # type: (str) -> SubprocessSDKEnvironment
    return cls(
        command_string, capabilities=python_sdk_capabilities(), artifacts=())


@Environment.register_urn(
    common_urns.environments.ANYOF.urn,
    beam_runner_api_pb2.AnyOfEnvironmentPayload)
class AnyOfEnvironment(Environment):
  def __init__(self, environments):
    self._environments = environments

  def to_runner_api_parameter(self, context):
    # type: (PipelineContext) -> tuple[str, beam_runner_api_pb2.AnyOfEnvironmentPayload]
    return (
        common_urns.environments.ANYOF.urn,
        beam_runner_api_pb2.AnyOfEnvironmentPayload(
            environments=[
                env.to_runner_api(context) for env in self._environments
            ]))

  @staticmethod
  def from_runner_api_parameter(
      payload,  # type: beam_runner_api_pb2.AnyOfEnvironmentPayload
      capabilities,  # type: Iterable[str]
      artifacts,  # type: Iterable[beam_runner_api_pb2.ArtifactInformation]
      resource_hints,  # type: Mapping[str, bytes]
      context  # type: PipelineContext
  ):
    # type: (...) -> AnyOfEnvironment
    return AnyOfEnvironment([
        Environment.from_runner_api(env, context)
        for env in payload.environments
    ])

  @staticmethod
  def create_proto(
      environments: Iterable[beam_runner_api_pb2.Environment]
  ) -> beam_runner_api_pb2.Environment:
    return beam_runner_api_pb2.Environment(
        urn=common_urns.environments.ANYOF.urn,
        payload=beam_runner_api_pb2.AnyOfEnvironmentPayload(
            environments=environments).SerializeToString())

  def __hash__(self):
    return sum(hash(env) for env in self._environments)

  def __eq__(self, other):
    return (
        isinstance(other, AnyOfEnvironment) and
        set(self._environments) == set(other._environments))

  def to_runner_api(self, context):
    return self.create_proto(
        [env.to_runner_api(context) for env in self._environments])


class PyPIArtifactRegistry(object):
  _registered_artifacts = set()  # type: set[tuple[str, str]]

  @classmethod
  def register_artifact(cls, name, version):
    cls._registered_artifacts.add((name, version))

  @classmethod
  def get_artifacts(cls):
    for artifact in cls._registered_artifacts:
      yield artifact


def python_sdk_capabilities():
  # type: () -> list[str]
  return list(_python_sdk_capabilities_iter())


def python_sdk_docker_capabilities():
  return python_sdk_capabilities() + [common_urns.protocols.SIBLING_WORKERS.urn]


def sdk_base_version_capability():
  return (
      SDK_VERSION_CAPABILITY_PREFIX + DockerEnvironment.default_docker_image())


def _python_sdk_capabilities_iter():
  # type: () -> Iterator[str]
  for urn_spec in common_urns.coders.__dict__.values():
    if getattr(urn_spec, 'urn', None) in coders.Coder._known_urns:
      yield urn_spec.urn
  yield common_urns.protocols.LEGACY_PROGRESS_REPORTING.urn
  yield common_urns.protocols.HARNESS_MONITORING_INFOS.urn
  yield common_urns.protocols.WORKER_STATUS.urn
  yield python_urns.PACKED_COMBINE_FN
  yield sdk_base_version_capability()
  yield common_urns.sdf_components.TRUNCATE_SIZED_RESTRICTION.urn
  yield common_urns.primitives.TO_STRING.urn
  yield common_urns.protocols.DATA_SAMPLING.urn
  yield common_urns.protocols.SDK_CONSUMING_RECEIVED_DATA.urn
  yield common_urns.protocols.ORDERED_LIST_STATE.urn
  yield common_urns.protocols.NAMED_DATA_STREAMS.urn


def python_sdk_dependencies(options, tmp_dir=None):
  if tmp_dir is None:
    tmp_dir = tempfile.mkdtemp()
  skip_prestaged_dependencies = options.view_as(
      SetupOptions).prebuild_sdk_container_engine is not None
  runner = options.view_as(
      StandardOptions).runner or StandardOptions.DEFAULT_RUNNER
  log_submission_env_dependencies = runner.split(
      '.')[-1] not in StandardOptions.LOCAL_RUNNERS
  return stager.Stager.create_job_resources(
      options,
      tmp_dir,
      pypi_requirements=[
          artifact[0] + artifact[1]
          for artifact in PyPIArtifactRegistry.get_artifacts()
      ],
      skip_prestaged_dependencies=skip_prestaged_dependencies,
      log_submission_env_dependencies=log_submission_env_dependencies)