<reponame>alexeykaplin/pipelines
# Copyright 2021 The Kubeflow Authors. All Rights Reserved.
#
# Licensed 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.

import logging
import time
import re
import googleapiclient.discovery as discovery
from google_cloud_pipeline_components.experimental.proto.gcp_resources_pb2 import GcpResources
from google.protobuf.json_format import Parse

_POLLING_INTERVAL_IN_SECONDS = 20
_CONNECTION_ERROR_RETRY_LIMIT = 5

_JOB_SUCCESSFUL_STATES = ['JOB_STATE_DONE']
_JOB_FAILED_STATES = [
    'JOB_STATE_STOPPED', 'JOB_STATE_FAILED', 'JOB_STATE_CANCELLED',
    'JOB_STATE_UPDATED', 'JOB_STATE_DRAINED'
]
_JOB_TERMINATED_STATES = _JOB_SUCCESSFUL_STATES + _JOB_FAILED_STATES
_DATAFLOW_URI_TEMPLATE = r'(https://dataflow.googleapis.com/v1b3/projects/(?P<project>.*)/locations/(?P<location>.*)/jobs/(?P<jobid>.*))'


def wait_gcp_resources(
    type,
    project,
    location,
    payload,
    gcp_resources,
):
    """
    Poll the gcp resources till it reaches a final state.
    """
    input_gcp_resources = Parse(payload, GcpResources())
    if len(input_gcp_resources.resources) != 1:
        raise ValueError(
            "Invalid payload: %s. Wait component support waiting on only one resource at this moment."
            % payload
        )

    if input_gcp_resources.resources[0].resource_type != 'DataflowJob':
        raise ValueError(
            "Invalid payload: %s. Wait component only support waiting on Dataflow job at this moment."
            % payload
        )

    dataflow_job_uri = input_gcp_resources.resources[0].resource_uri
    uri_pattern = re.compile(_DATAFLOW_URI_TEMPLATE)
    match = uri_pattern.match(dataflow_job_uri)
    # Get the project and location from the job URI instead from the parameter.
    try:
        project = match.group('project')
        location = match.group('location')
        job_id = match.group('jobid')
    except AttributeError as err:
        # TODO(ruifang) propagate the error.
        raise ValueError(
            'Invalid dataflow resource URI: {}. Expect: {}.'.format(
                dataflow_job_uri,
                'https://dataflow.googleapis.com/v1b3/projects/[project_id]/locations/[location]/jobs/[job_id]'
            )
        )

    # Propagate the GCP resources as the output of the wait component
    with open(gcp_resources, 'w') as f:
        f.write(payload)

    # Poll the job status
    retry_count = 0
    while True:
        try:
            df_client = discovery.build(
                'dataflow', 'v1b3', cache_discovery=False
            )
            job = df_client.projects().locations().jobs().get(
                projectId=project, jobId=job_id, location=location, view=None
            ).execute()
            retry_count = 0
        except ConnectionError as err:
            retry_count += 1
            if retry_count <= _CONNECTION_ERROR_RETRY_LIMIT:
                logging.warning(
                    'ConnectionError (%s) encountered when polling job: %s. Retrying.',
                    err, job_id
                )
            else:
                logging.error(
                    'Request failed after %s retries.',
                    _CONNECTION_ERROR_RETRY_LIMIT
                )
                # TODO(ruifang) propagate the error.
                raise

        job_state = job.get('currentState', None)
        # Write the job details as gcp_resources
        if job_state in _JOB_SUCCESSFUL_STATES:
            logging.info(
                'GetDataflowJob response state =%s. Job completed', job_state
            )
            return

        elif job_state in _JOB_TERMINATED_STATES:
            # TODO(ruifang) propagate the error.
            raise RuntimeError(
                'Job {} failed with error state: {}.'.format(job_id, job_state)
            )
        else:
            logging.info(
                'Job %s is in a non-final state %s. Waiting for %s seconds for next poll.',
                job_id, job_state, _POLLING_INTERVAL_IN_SECONDS
            )
            time.sleep(_POLLING_INTERVAL_IN_SECONDS)

# Copyright 2021 The Kubeflow Authors. All Rights Reserved.
#
# Licensed 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.

import logging
import time
import re
import googleapiclient.discovery as discovery
from google_cloud_pipeline_components.experimental.proto.gcp_resources_pb2 import GcpResources
from google.protobuf.json_format import Parse

_POLLING_INTERVAL_IN_SECONDS = 20
_CONNECTION_ERROR_RETRY_LIMIT = 5

_JOB_SUCCESSFUL_STATES = ['JOB_STATE_DONE']
_JOB_FAILED_STATES = [
    'JOB_STATE_STOPPED', 'JOB_STATE_FAILED', 'JOB_STATE_CANCELLED',
    'JOB_STATE_UPDATED', 'JOB_STATE_DRAINED'
]
_JOB_TERMINATED_STATES = _JOB_SUCCESSFUL_STATES + _JOB_FAILED_STATES
_DATAFLOW_URI_TEMPLATE = r'(https://dataflow.googleapis.com/v1b3/projects/(?P<project>.*)/locations/(?P<location>.*)/jobs/(?P<jobid>.*))'


def wait_gcp_resources(
    type,
    project,
    location,
    payload,
    gcp_resources,
):
    """
    Poll the gcp resources till it reaches a final state.
    """
    input_gcp_resources = Parse(payload, GcpResources())
    if len(input_gcp_resources.resources) != 1:
        raise ValueError(
            "Invalid payload: %s. Wait component support waiting on only one resource at this moment."
            % payload
        )

    if input_gcp_resources.resources[0].resource_type != 'DataflowJob':
        raise ValueError(
            "Invalid payload: %s. Wait component only support waiting on Dataflow job at this moment."
            % payload
        )

    dataflow_job_uri = input_gcp_resources.resources[0].resource_uri
    uri_pattern = re.compile(_DATAFLOW_URI_TEMPLATE)
    match = uri_pattern.match(dataflow_job_uri)
    # Get the project and location from the job URI instead from the parameter.
    try:
        project = match.group('project')
        location = match.group('location')
        job_id = match.group('jobid')
    except AttributeError as err:
        # TODO(ruifang) propagate the error.
        raise ValueError(
            'Invalid dataflow resource URI: {}. Expect: {}.'.format(
                dataflow_job_uri,
                'https://dataflow.googleapis.com/v1b3/projects/[project_id]/locations/[location]/jobs/[job_id]'
            )
        )

    # Propagate the GCP resources as the output of the wait component
    with open(gcp_resources, 'w') as f:
        f.write(payload)

    # Poll the job status
    retry_count = 0
    while True:
        try:
            df_client = discovery.build(
                'dataflow', 'v1b3', cache_discovery=False
            )
            job = df_client.projects().locations().jobs().get(
                projectId=project, jobId=job_id, location=location, view=None
            ).execute()
            retry_count = 0
        except ConnectionError as err:
            retry_count += 1
            if retry_count <= _CONNECTION_ERROR_RETRY_LIMIT:
                logging.warning(
                    'ConnectionError (%s) encountered when polling job: %s. Retrying.',
                    err, job_id
                )
            else:
                logging.error(
                    'Request failed after %s retries.',
                    _CONNECTION_ERROR_RETRY_LIMIT
                )
                # TODO(ruifang) propagate the error.
                raise

        job_state = job.get('currentState', None)
        # Write the job details as gcp_resources
        if job_state in _JOB_SUCCESSFUL_STATES:
            logging.info(
                'GetDataflowJob response state =%s. Job completed', job_state
            )
            return

        elif job_state in _JOB_TERMINATED_STATES:
            # TODO(ruifang) propagate the error.
            raise RuntimeError(
                'Job {} failed with error state: {}.'.format(job_id, job_state)
            )
        else:
            logging.info(
                'Job %s is in a non-final state %s. Waiting for %s seconds for next poll.',
                job_id, job_state, _POLLING_INTERVAL_IN_SECONDS
            )
            time.sleep(_POLLING_INTERVAL_IN_SECONDS)
# Copyright 2021 The Kubeflow Authors. All Rights Reserved. # # Licensed 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. Poll the gcp resources till it reaches a final state. # Get the project and location from the job URI instead from the parameter. # TODO(ruifang) propagate the error. # Propagate the GCP resources as the output of the wait component # Poll the job status # TODO(ruifang) propagate the error. # Write the job details as gcp_resources # TODO(ruifang) propagate the error.