#
# 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 logging
import re
import time
from datetime import datetime
from typing import Any
from typing import Optional
from google.cloud import monitoring_v3
from google.protobuf.duration_pb2 import Duration
import apache_beam.testing.load_tests.dataflow_cost_consts as costs
from apache_beam.metrics.execution import MetricResult
from apache_beam.runners.dataflow.dataflow_runner import DataflowPipelineResult
from apache_beam.runners.dataflow.internal.apiclient import DataflowApplicationClient
from apache_beam.runners.runner import PipelineState
from apache_beam.testing.load_tests.load_test import LoadTest
[docs]
class DataflowCostBenchmark(LoadTest):
"""Base class for Dataflow performance tests which export metrics to
external databases: BigQuery or/and InfluxDB. Calculates the expected cost
for running the job on Dataflow in region us-central1.
Refer to :class:`~apache_beam.testing.load_tests.LoadTestOptions` for more
information on the required pipeline options.
If using InfluxDB with Basic HTTP authentication enabled, provide the
following environment options: `INFLUXDB_USER` and `INFLUXDB_USER_PASSWORD`.
If the hardware configuration for the job includes use of a GPU, please
specify the version in use with the Accelerator enumeration. This is used to
calculate the cost of the job later, as different accelerators have different
billing rates per hour of use.
"""
WORKER_START_PATTERN = re.compile(
r'^All workers have finished the startup processes and '
r'began to receive work requests.*$')
WORKER_STOP_PATTERN = re.compile(r'^Stopping worker pool.*$')
def __init__(
self,
metrics_namespace: Optional[str] = None,
is_streaming: bool = False,
gpu: Optional[costs.Accelerator] = None,
pcollection: str = 'ProcessOutput.out0'):
"""
Initializes DataflowCostBenchmark.
Args:
metrics_namespace (Optional[str]): Namespace for metrics.
is_streaming (bool): Whether the pipeline is streaming or batch.
gpu (Optional[costs.Accelerator]): Optional GPU type.
pcollection (str): PCollection name to monitor throughput.
"""
self.is_streaming = is_streaming
self.gpu = gpu
self.pcollection = pcollection
super().__init__(metrics_namespace=metrics_namespace)
self.dataflow_client = DataflowApplicationClient(
self.pipeline.get_pipeline_options())
self.monitoring_client = monitoring_v3.MetricServiceClient()
[docs]
def run(self) -> None:
try:
self.test()
if not hasattr(self, 'result'):
self.result = self.pipeline.run()
state = self.result.wait_until_finish(duration=self.timeout_ms)
assert state != PipelineState.FAILED
logging.info(
'Pipeline complete, sleeping for 4 minutes to allow resource '
'metrics to populate.')
time.sleep(240)
self.extra_metrics = self._retrieve_cost_metrics(self.result)
additional_metrics = self._get_additional_metrics(self.result)
self.extra_metrics.update(additional_metrics)
logging.info(self.extra_metrics)
self._metrics_monitor.publish_metrics(self.result, self.extra_metrics)
finally:
self.cleanup()
def _retrieve_cost_metrics(self,
result: DataflowPipelineResult) -> dict[str, Any]:
"""Calculates estimated cost based on pipeline resource usage."""
job_id = result.job_id()
metrics = result.metrics().all_metrics(job_id)
metrics_dict = self._process_metrics_list(metrics)
cost = 0.0
if self.is_streaming:
cost += metrics_dict.get(
"TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_STREAMING
cost += metrics_dict.get(
"TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_STREAMING
cost += metrics_dict.get(
"TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_STREAMING
else:
cost += metrics_dict.get(
"TotalVcpuTime", 0.0) / 3600 * costs.VCPU_PER_HR_BATCH
cost += metrics_dict.get(
"TotalMemoryUsage", 0.0) / 1000 / 3600 * costs.MEM_PER_GB_HR_BATCH
cost += metrics_dict.get(
"TotalStreamingDataProcessed", 0.0) * costs.SHUFFLE_PER_GB_BATCH
if self.gpu:
rate = costs.ACCELERATOR_TO_COST[self.gpu]
cost += metrics_dict.get("TotalGpuTime", 0.0) / 3600 * rate
cost += metrics_dict.get("TotalPdUsage", 0.0) / 3600 * costs.PD_PER_GB_HR
cost += metrics_dict.get(
"TotalSsdUsage", 0.0) / 3600 * costs.PD_SSD_PER_GB_HR
metrics_dict["EstimatedCost"] = cost
return metrics_dict
def _process_metrics_list(self,
metrics: list[MetricResult]) -> dict[str, Any]:
system_metrics = {}
for entry in metrics:
metric_key = entry.key
metric = metric_key.metric
if metric_key.step == '' and metric.namespace == 'dataflow/v1b3':
system_metrics[metric.name] = entry.committed or 0.0
return system_metrics
def _get_worker_time_interval(
self, job_id: str) -> tuple[Optional[str], Optional[str]]:
"""Extracts worker start and stop times from job messages."""
messages, _ = self.dataflow_client.list_messages(
job_id=job_id,
start_time=None,
end_time=None,
minimum_importance='JOB_MESSAGE_DETAILED')
start_time, end_time = None, None
for message in messages:
text = message.messageText
if text:
if self.WORKER_START_PATTERN.match(text):
start_time = message.time
if self.WORKER_STOP_PATTERN.match(text):
end_time = message.time
return start_time, end_time
def _get_throughput_metrics(
self, project: str, job_id: str, start_time: str,
end_time: str) -> dict[str, float]:
interval = monitoring_v3.TimeInterval(
start_time=start_time, end_time=end_time)
aggregation = monitoring_v3.Aggregation(
alignment_period=Duration(seconds=60),
per_series_aligner=monitoring_v3.Aggregation.Aligner.ALIGN_MEAN)
requests = {
"Bytes": monitoring_v3.ListTimeSeriesRequest(
name=f"projects/{project}",
filter=f'metric.type='
f'"dataflow.googleapis.com/job/estimated_bytes_produced_count" '
f'AND metric.labels.job_id='
f'"{job_id}" AND metric.labels.pcollection="{self.pcollection}"',
interval=interval,
aggregation=aggregation),
"Elements": monitoring_v3.ListTimeSeriesRequest(
name=f"projects/{project}",
filter=f'metric.type="dataflow.googleapis.com/job/element_count" '
f'AND metric.labels.job_id="{job_id}" '
f'AND metric.labels.pcollection="{self.pcollection}"',
interval=interval,
aggregation=aggregation)
}
metrics = {}
for key, req in requests.items():
time_series = self.monitoring_client.list_time_series(request=req)
values = [
point.value.double_value for series in time_series
for point in series.points
]
metrics[f"AvgThroughput{key}"] = sum(values) / len(
values) if values else 0.0
return metrics
def _get_job_runtime(self, start_time: str, end_time: str) -> float:
"""Calculates the job runtime duration in seconds."""
start_dt = datetime.fromisoformat(start_time[:-1])
end_dt = datetime.fromisoformat(end_time[:-1])
return (end_dt - start_dt).total_seconds()
def _get_additional_metrics(self,
result: DataflowPipelineResult) -> dict[str, Any]:
job_id = result.job_id()
job = self.dataflow_client.get_job(job_id)
project = job.projectId
start_time, end_time = self._get_worker_time_interval(job_id)
if not start_time or not end_time:
logging.warning('Could not find valid worker start/end times.')
return {}
throughput_metrics = self._get_throughput_metrics(
project, job_id, start_time, end_time)
return {
**throughput_metrics,
"JobRuntimeSeconds": self._get_job_runtime(start_time, end_time),
}