github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/portability/abstract_job_service.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.
#
# pytype: skip-file

import copy
import itertools
import json
import logging
import shutil
import tempfile
import uuid
import zipfile
from concurrent import futures
from typing import TYPE_CHECKING
from typing import Dict
from typing import Iterator
from typing import Optional
from typing import Tuple
from typing import Union

import grpc
from google.protobuf import json_format
from google.protobuf import timestamp_pb2

from apache_beam.portability.api import beam_artifact_api_pb2_grpc
from apache_beam.portability.api import beam_job_api_pb2
from apache_beam.portability.api import beam_job_api_pb2_grpc
from apache_beam.portability.api import endpoints_pb2
from apache_beam.runners.portability import artifact_service
from apache_beam.utils.timestamp import Timestamp

if TYPE_CHECKING:
  # pylint: disable=ungrouped-imports
  from typing import BinaryIO
  from google.protobuf import struct_pb2
  from apache_beam.portability.api import beam_runner_api_pb2

_LOGGER = logging.getLogger(__name__)

StateEvent = Tuple[int, Union[timestamp_pb2.Timestamp, Timestamp]]


def make_state_event(state, timestamp):
  if isinstance(timestamp, Timestamp):
    proto_timestamp = timestamp.to_proto()
  elif isinstance(timestamp, timestamp_pb2.Timestamp):
    proto_timestamp = timestamp
  else:
    raise ValueError(
        "Expected apache_beam.utils.timestamp.Timestamp, "
        "or google.protobuf.timestamp_pb2.Timestamp. "
        "Got %s" % type(timestamp))

  return beam_job_api_pb2.JobStateEvent(state=state, timestamp=proto_timestamp)


class AbstractJobServiceServicer(beam_job_api_pb2_grpc.JobServiceServicer):
  """Manages one or more pipelines, possibly concurrently.
  Experimental: No backward compatibility guaranteed.
  Servicer for the Beam Job API.
  """
  def __init__(self):
    self._jobs = {}  # type: Dict[str, AbstractBeamJob]

  def create_beam_job(self,
                      preparation_id,  # stype: str
                      job_name,  # type: str
                      pipeline,  # type: beam_runner_api_pb2.Pipeline
                      options  # type: struct_pb2.Struct
                     ):
    # type: (...) -> AbstractBeamJob

    """Returns an instance of AbstractBeamJob specific to this servicer."""
    raise NotImplementedError(type(self))

  def Prepare(self,
              request,  # type: beam_job_api_pb2.PrepareJobRequest
              context=None,
              timeout=None
             ):
    # type: (...) -> beam_job_api_pb2.PrepareJobResponse
    _LOGGER.debug('Got Prepare request.')
    preparation_id = '%s-%s' % (request.job_name, uuid.uuid4())
    self._jobs[preparation_id] = self.create_beam_job(
        preparation_id,
        request.job_name,
        request.pipeline,
        request.pipeline_options)
    self._jobs[preparation_id].prepare()
    _LOGGER.debug("Prepared job '%s' as '%s'", request.job_name, preparation_id)
    return beam_job_api_pb2.PrepareJobResponse(
        preparation_id=preparation_id,
        artifact_staging_endpoint=self._jobs[preparation_id].
        artifact_staging_endpoint(),
        staging_session_token=preparation_id)

  def Run(self,
          request,  # type: beam_job_api_pb2.RunJobRequest
          context=None,
          timeout=None
         ):
    # type: (...) -> beam_job_api_pb2.RunJobResponse
    # For now, just use the preparation id as the job id.
    job_id = request.preparation_id
    _LOGGER.info("Running job '%s'", job_id)
    self._jobs[job_id].run()
    return beam_job_api_pb2.RunJobResponse(job_id=job_id)

  def GetJobs(self,
              request,  # type: beam_job_api_pb2.GetJobsRequest
              context=None,
              timeout=None
             ):
    # type: (...) -> beam_job_api_pb2.GetJobsResponse
    return beam_job_api_pb2.GetJobsResponse(
        job_info=[job.to_runner_api() for job in self._jobs.values()])

  def GetState(
      self,
      request,  # type: beam_job_api_pb2.GetJobStateRequest
      context=None):
    # type: (...) -> beam_job_api_pb2.JobStateEvent
    return make_state_event(*self._jobs[request.job_id].get_state())

  def GetPipeline(self,
                  request,  # type: beam_job_api_pb2.GetJobPipelineRequest
                  context=None,
                  timeout=None
                 ):
    # type: (...) -> beam_job_api_pb2.GetJobPipelineResponse
    return beam_job_api_pb2.GetJobPipelineResponse(
        pipeline=self._jobs[request.job_id].get_pipeline())

  def Cancel(self,
             request,  # type: beam_job_api_pb2.CancelJobRequest
             context=None,
             timeout=None
            ):
    # type: (...) -> beam_job_api_pb2.CancelJobResponse
    self._jobs[request.job_id].cancel()
    return beam_job_api_pb2.CancelJobResponse(
        state=self._jobs[request.job_id].get_state()[0])

  def GetStateStream(self, request, context=None, timeout=None):
    # type: (...) -> Iterator[beam_job_api_pb2.JobStateEvent]

    """Yields state transitions since the stream started.
      """
    if request.job_id not in self._jobs:
      raise LookupError("Job {} does not exist".format(request.job_id))

    job = self._jobs[request.job_id]
    for state, timestamp in job.get_state_stream():
      yield make_state_event(state, timestamp)

  def GetMessageStream(self, request, context=None, timeout=None):
    # type: (...) -> Iterator[beam_job_api_pb2.JobMessagesResponse]

    """Yields messages since the stream started.
      """
    if request.job_id not in self._jobs:
      raise LookupError("Job {} does not exist".format(request.job_id))

    job = self._jobs[request.job_id]
    for msg in job.get_message_stream():
      if isinstance(msg, tuple):
        resp = beam_job_api_pb2.JobMessagesResponse(
            state_response=make_state_event(*msg))
      else:
        resp = beam_job_api_pb2.JobMessagesResponse(message_response=msg)
      yield resp

  def DescribePipelineOptions(self, request, context=None, timeout=None):
    # type: (...) -> beam_job_api_pb2.DescribePipelineOptionsResponse
    return beam_job_api_pb2.DescribePipelineOptionsResponse()


class AbstractBeamJob(object):
  """Abstract baseclass for managing a single Beam job."""

  def __init__(self,
               job_id,  # type: str
               job_name,  # type: str
               pipeline,  # type: beam_runner_api_pb2.Pipeline
               options  # type: struct_pb2.Struct
              ):
    self._job_id = job_id
    self._job_name = job_name
    self._pipeline_proto = pipeline
    self._pipeline_options = options
    self._state_history = [(beam_job_api_pb2.JobState.STOPPED, Timestamp.now())]

  def prepare(self):
    # type: () -> None

    """Called immediately after this class is instantiated"""
    raise NotImplementedError(self)

  def run(self):
    # type: () -> None
    raise NotImplementedError(self)

  def cancel(self):
    # type: () -> Optional[beam_job_api_pb2.JobState.Enum]
    raise NotImplementedError(self)

  def artifact_staging_endpoint(self):
    # type: () -> Optional[endpoints_pb2.ApiServiceDescriptor]
    raise NotImplementedError(self)

  def get_state_stream(self):
    # type: () -> Iterator[StateEvent]
    raise NotImplementedError(self)

  def get_message_stream(self):
    # type: () -> Iterator[Union[StateEvent, Optional[beam_job_api_pb2.JobMessage]]]
    raise NotImplementedError(self)

  @property
  def state(self):
    """Get the latest state enum."""
    return self.get_state()[0]

  def get_state(self):
    """Get a tuple of the latest state and its timestamp."""
    # this is safe: initial state is set in __init__
    return self._state_history[-1]

  def set_state(self, new_state):
    """Set the latest state as an int enum and update the state history.

    :param new_state: int
      latest state enum
    :return: Timestamp or None
      the new timestamp if the state has not changed, else None
    """
    if new_state != self._state_history[-1][0]:
      timestamp = Timestamp.now()
      self._state_history.append((new_state, timestamp))
      return timestamp
    else:
      return None

  def with_state_history(self, state_stream):
    """Utility to prepend recorded state history to an active state stream"""
    return itertools.chain(self._state_history[:], state_stream)

  def get_pipeline(self):
    # type: () -> beam_runner_api_pb2.Pipeline
    return self._pipeline_proto

  @staticmethod
  def is_terminal_state(state):
    from apache_beam.runners.portability import portable_runner
    return state in portable_runner.TERMINAL_STATES

  def to_runner_api(self):
    # type: () -> beam_job_api_pb2.JobInfo
    return beam_job_api_pb2.JobInfo(
        job_id=self._job_id,
        job_name=self._job_name,
        pipeline_options=self._pipeline_options,
        state=self.state)


class JarArtifactManager(object):
  def __init__(self, jar_path, root):
    self._root = root
    self._zipfile_handle = zipfile.ZipFile(jar_path, 'a')

  def close(self):
    self._zipfile_handle.close()

  def file_writer(self, path):
    # type: (str) -> Tuple[BinaryIO, str]

    """Given a relative path, returns an open handle that can be written to
    and an reference that can later be used to read this file."""
    full_path = '%s/%s' % (self._root, path)
    return self._zipfile_handle.open(
        full_path, 'w', force_zip64=True), 'classpath://%s' % full_path

  def zipfile_handle(self):
    return self._zipfile_handle


class UberJarBeamJob(AbstractBeamJob):
  """Abstract baseclass for creating a Beam job. The resulting job will be
  packaged and run in an executable uber jar."""

  # These must agree with those defined in PortablePipelineJarUtils.java.
  PIPELINE_FOLDER = 'BEAM-PIPELINE'
  PIPELINE_MANIFEST = PIPELINE_FOLDER + '/pipeline-manifest.json'

  # We only stage a single pipeline in the jar.
  PIPELINE_NAME = 'pipeline'
  PIPELINE_PATH = '/'.join([PIPELINE_FOLDER, PIPELINE_NAME, "pipeline.json"])
  PIPELINE_OPTIONS_PATH = '/'.join(
      [PIPELINE_FOLDER, PIPELINE_NAME, 'pipeline-options.json'])
  ARTIFACT_FOLDER = '/'.join([PIPELINE_FOLDER, PIPELINE_NAME, 'artifacts'])

  def __init__(
      self,
      executable_jar,
      job_id,
      job_name,
      pipeline,
      options,
      artifact_port=0):
    super().__init__(job_id, job_name, pipeline, options)
    self._executable_jar = executable_jar
    self._jar_uploaded = False
    self._artifact_port = artifact_port

  def prepare(self):
    # Copy the executable jar, injecting the pipeline and options as resources.
    with tempfile.NamedTemporaryFile(suffix='.jar') as tout:
      self._jar = tout.name
    shutil.copy(self._executable_jar, self._jar)
    self._start_artifact_service(self._jar, self._artifact_port)

  def _start_artifact_service(self, jar, requested_port):
    self._artifact_manager = JarArtifactManager(self._jar, self.ARTIFACT_FOLDER)
    self._artifact_staging_service = artifact_service.ArtifactStagingService(
        self._artifact_manager.file_writer)
    self._artifact_staging_service.register_job(
        self._job_id,
        {
            env_id: env.dependencies
            for (env_id,
                 env) in self._pipeline_proto.components.environments.items()
        })
    options = [("grpc.http2.max_pings_without_data", 0),
               ("grpc.http2.max_ping_strikes", 0)]
    self._artifact_staging_server = grpc.server(
        futures.ThreadPoolExecutor(), options=options)
    port = self._artifact_staging_server.add_insecure_port(
        '[::]:%s' % requested_port)
    beam_artifact_api_pb2_grpc.add_ArtifactStagingServiceServicer_to_server(
        self._artifact_staging_service, self._artifact_staging_server)
    self._artifact_staging_endpoint = endpoints_pb2.ApiServiceDescriptor(
        url='localhost:%d' % port)
    self._artifact_staging_server.start()
    _LOGGER.info('Artifact server started on port %s', port)
    return port

  def _stop_artifact_service(self):
    self._artifact_staging_server.stop(1)

    # Update dependencies to point to staged files.
    pipeline = copy.copy(self._pipeline_proto)
    if any(env.dependencies
           for env in pipeline.components.environments.values()):
      for env_id, deps in self._artifact_staging_service.resolved_deps(
          self._job_id).items():
        # Slice assignment not supported for repeated fields.
        env = self._pipeline_proto.components.environments[env_id]
        del env.dependencies[:]
        env.dependencies.extend(deps)

    # Copy the pipeline definition and metadata into the jar.
    z = self._artifact_manager.zipfile_handle()
    with z.open(self.PIPELINE_PATH, 'w') as fout:
      fout.write(
          json_format.MessageToJson(self._pipeline_proto).encode('utf-8'))
    with z.open(self.PIPELINE_OPTIONS_PATH, 'w') as fout:
      fout.write(
          json_format.MessageToJson(self._pipeline_options).encode('utf-8'))
    with z.open(self.PIPELINE_MANIFEST, 'w') as fout:
      fout.write(
          json.dumps({
              'defaultJobName': self.PIPELINE_NAME
          }).encode('utf-8'))

    # Closes the jar file.
    self._artifact_manager.close()

  def artifact_staging_endpoint(self):
    return self._artifact_staging_endpoint