Examples with DataflowStepContext org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext used on opensource projects

Example 11 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class ExecutionTimeMonitoringInfoToCounterUpdateTransformerTest method testTransformReturnsNullIfSpecValidationFails.

@Test
public void testTransformReturnsNullIfSpecValidationFails() {
    Map<String, DataflowStepContext> stepContextMapping = new HashMap<>();
    ExecutionTimeMonitoringInfoToCounterUpdateTransformer testObject = new ExecutionTimeMonitoringInfoToCounterUpdateTransformer(mockSpecValidator, stepContextMapping);
    Optional<String> error = Optional.of("Error text");
    when(mockSpecValidator.validate(any())).thenReturn(error);
    MonitoringInfo monitoringInfo = MonitoringInfo.newBuilder().setUrn(Urns.USER_SUM_INT64).setType(TypeUrns.SUM_INT64_TYPE).build();
    assertNull(testObject.transform(monitoringInfo));
}

Example 12 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserMonitoringInfoToCounterUpdateTransformer method validate.

private Optional<String> validate(MonitoringInfo monitoringInfo) {
    Optional<String> validatorResult = specValidator.validate(monitoringInfo);
    if (validatorResult.isPresent()) {
        return validatorResult;
    }
    String urn = monitoringInfo.getUrn();
    if (!urn.equals(Urns.USER_SUM_INT64)) {
        throw new RuntimeException(String.format("Received unexpected counter urn. Expected urn: %s, received: %s", Urns.USER_SUM_INT64, urn));
    }
    String type = monitoringInfo.getType();
    if (!type.equals(TypeUrns.SUM_INT64_TYPE)) {
        throw new RuntimeException(String.format("Received unexpected counter type. Expected type: %s, received: %s", TypeUrns.SUM_INT64_TYPE, type));
    }
    final String ptransform = monitoringInfo.getLabelsMap().get(MonitoringInfoConstants.Labels.PTRANSFORM);
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    if (stepContext == null) {
        return Optional.of("Encountered user-counter MonitoringInfo with unknown ptransformId: " + monitoringInfo.toString());
    }
    return Optional.empty();
}

Example 13 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class UserMonitoringInfoToCounterUpdateTransformer method transform.

/**
 * Transforms user counter MonitoringInfo to relevant CounterUpdate.
 *
 * @return Relevant CounterUpdate or null if transformation failed.
 */
@Override
@Nullable
public CounterUpdate transform(MonitoringInfo monitoringInfo) {
    Optional<String> validationResult = validate(monitoringInfo);
    if (validationResult.isPresent()) {
        LOG.debug(validationResult.get());
        return null;
    }
    long value = decodeInt64Counter(monitoringInfo.getPayload());
    Map<String, String> miLabels = monitoringInfo.getLabelsMap();
    final String ptransform = miLabels.get(MonitoringInfoConstants.Labels.PTRANSFORM);
    final String counterName = miLabels.get(MonitoringInfoConstants.Labels.NAME);
    final String counterNamespace = miLabels.get(MonitoringInfoConstants.Labels.NAMESPACE);
    CounterStructuredNameAndMetadata name = new CounterStructuredNameAndMetadata();
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    name.setName(new CounterStructuredName().setOrigin(Origin.USER.toString()).setName(counterName).setOriginalStepName(stepContext.getNameContext().originalName()).setOriginNamespace(counterNamespace)).setMetadata(new CounterMetadata().setKind(Kind.SUM.toString()));
    return new CounterUpdate().setStructuredNameAndMetadata(name).setCumulative(true).setInteger(DataflowCounterUpdateExtractor.longToSplitInt(value));
}

Example 14 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class ExecutionTimeMonitoringInfoToCounterUpdateTransformer method transform.

@Override
@Nullable
public CounterUpdate transform(MonitoringInfo monitoringInfo) {
    Optional<String> validationResult = validate(monitoringInfo);
    if (validationResult.isPresent()) {
        LOG.debug(validationResult.get());
        return null;
    }
    long value = decodeInt64Counter(monitoringInfo.getPayload());
    String urn = monitoringInfo.getUrn();
    final String ptransform = monitoringInfo.getLabelsMap().get(MonitoringInfoConstants.Labels.PTRANSFORM);
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    String counterName = URN_TO_COUNTER_NAME_MAPPING.get(urn);
    CounterStructuredNameAndMetadata name = new CounterStructuredNameAndMetadata();
    name.setName(new CounterStructuredName().setOrigin("SYSTEM").setName(counterName).setOriginalStepName(stepContext.getNameContext().originalName()).setExecutionStepName(stepContext.getNameContext().stageName())).setMetadata(new CounterMetadata().setKind(Kind.SUM.toString()));
    return new CounterUpdate().setStructuredNameAndMetadata(name).setCumulative(true).setInteger(DataflowCounterUpdateExtractor.longToSplitInt(value));
}

Example 15 with DataflowStepContext

use of org.apache.beam.runners.dataflow.worker.DataflowExecutionContext.DataflowStepContext in project beam by apache.

the class ExecutionTimeMonitoringInfoToCounterUpdateTransformer method validate.

/**
 * Validates provided monitoring info against specs and common safety checks.
 *
 * @param monitoringInfo to validate.
 * @return Optional.empty() all validation checks are passed. Optional with error text otherwise.
 * @throws RuntimeException if received unexpected urn.
 */
protected Optional<String> validate(MonitoringInfo monitoringInfo) {
    Optional<String> validatorResult = specValidator.validate(monitoringInfo);
    if (validatorResult.isPresent()) {
        return validatorResult;
    }
    String urn = monitoringInfo.getUrn();
    if (!URN_TO_COUNTER_NAME_MAPPING.keySet().contains(urn)) {
        throw new RuntimeException(String.format("Received unexpected counter urn: %s", urn));
    }
    String type = monitoringInfo.getType();
    if (!type.equals(TypeUrns.SUM_INT64_TYPE)) {
        throw new RuntimeException(String.format("Received unexpected counter type. Expected type: %s, received: %s", TypeUrns.SUM_INT64_TYPE, type));
    }
    final String ptransform = monitoringInfo.getLabelsMap().get(MonitoringInfoConstants.Labels.PTRANSFORM);
    DataflowStepContext stepContext = transformIdMapping.get(ptransform);
    if (stepContext == null) {
        return Optional.of("Encountered MSec MonitoringInfo with unknown ptransformId: " + monitoringInfo.toString());
    }
    return Optional.empty();
}