Examples with WorkItemStatus com.google.api.services.dataflow.model.WorkItemStatus used on opensource projects

Example 21 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class StreamingDataflowWorker method sendWorkerUpdatesToDataflowService.

/**
 * Sends counter updates to Dataflow backend.
 */
private void sendWorkerUpdatesToDataflowService(CounterSet deltaCounters, CounterSet cumulativeCounters) throws IOException {
    // Throttle time is tracked by the windmillServer but is reported to DFE here.
    windmillQuotaThrottling.addValue(windmillServer.getAndResetThrottleTime());
    if (memoryMonitor.isThrashing()) {
        memoryThrashing.addValue(1);
    }
    List<CounterUpdate> counterUpdates = new ArrayList<>(128);
    if (publishCounters) {
        stageInfoMap.values().forEach(s -> counterUpdates.addAll(s.extractCounterUpdates()));
        counterUpdates.addAll(cumulativeCounters.extractUpdates(false, DataflowCounterUpdateExtractor.INSTANCE));
        counterUpdates.addAll(deltaCounters.extractModifiedDeltaUpdates(DataflowCounterUpdateExtractor.INSTANCE));
        if (hasExperiment(options, "beam_fn_api")) {
            Map<Object, List<CounterUpdate>> fnApiCounters = new HashMap<>();
            while (!this.pendingMonitoringInfos.isEmpty()) {
                final CounterUpdate item = this.pendingMonitoringInfos.poll();
                // WorkItem.
                if (item.getCumulative()) {
                    item.setCumulative(false);
                    // Group counterUpdates by counterUpdateKey so they can be aggregated before sending to
                    // dataflow service.
                    fnApiCounters.computeIfAbsent(getCounterUpdateKey(item), k -> new ArrayList<>()).add(item);
                } else {
                    // This is a safety check in case new counter type appears in FnAPI.
                    throw new UnsupportedOperationException("FnApi counters are expected to provide cumulative values." + " Please, update conversion to delta logic" + " if non-cumulative counter type is required.");
                }
            }
            // so we can avoid excessive I/Os for reporting to dataflow service.
            for (List<CounterUpdate> counterUpdateList : fnApiCounters.values()) {
                if (counterUpdateList.isEmpty()) {
                    continue;
                }
                List<CounterUpdate> aggregatedCounterUpdateList = CounterUpdateAggregators.aggregate(counterUpdateList);
                // updates.
                if (aggregatedCounterUpdateList.size() > 10) {
                    CounterUpdate head = aggregatedCounterUpdateList.get(0);
                    this.counterAggregationErrorCount.getAndIncrement();
                    // log warning message only when error count is the power of 2 to avoid spamming.
                    if (this.counterAggregationErrorCount.get() > 10 && Long.bitCount(this.counterAggregationErrorCount.get()) == 1) {
                        LOG.warn("Found non-aggregated counter updates of size {} with kind {}, this will likely " + "cause performance degradation and excessive GC if size is large.", counterUpdateList.size(), MoreObjects.firstNonNull(head.getNameAndKind(), head.getStructuredNameAndMetadata()));
                    }
                }
                counterUpdates.addAll(aggregatedCounterUpdateList);
            }
        }
    }
    // Handle duplicate counters from different stages. Store all the counters in a multi-map and
    // send the counters that appear multiple times in separate RPCs. Same logical counter could
    // appear in multiple stages if a step runs in multiple stages (as with flatten-unzipped stages)
    // especially if the counter definition does not set execution_step_name.
    ListMultimap<Object, CounterUpdate> counterMultimap = MultimapBuilder.hashKeys(counterUpdates.size()).linkedListValues().build();
    boolean hasDuplicates = false;
    for (CounterUpdate c : counterUpdates) {
        Object key = getCounterUpdateKey(c);
        if (counterMultimap.containsKey(key)) {
            hasDuplicates = true;
        }
        counterMultimap.put(key, c);
    }
    // Clears counterUpdates and enqueues unique counters from counterMultimap. If a counter
    // appears more than once, one of them is extracted leaving the remaining in the map.
    Runnable extractUniqueCounters = () -> {
        counterUpdates.clear();
        for (Iterator<Object> iter = counterMultimap.keySet().iterator(); iter.hasNext(); ) {
            List<CounterUpdate> counters = counterMultimap.get(iter.next());
            counterUpdates.add(counters.get(0));
            if (counters.size() == 1) {
                // There is single value. Remove the entry through the iterator.
                iter.remove();
            } else {
                // Otherwise remove the first value.
                counters.remove(0);
            }
        }
    };
    if (hasDuplicates) {
        extractUniqueCounters.run();
    } else {
        // Common case: no duplicates. We can just send counterUpdates, empty the multimap.
        counterMultimap.clear();
    }
    List<Status> errors;
    synchronized (pendingFailuresToReport) {
        errors = new ArrayList<>(pendingFailuresToReport.size());
        for (String stackTrace : pendingFailuresToReport) {
            errors.add(new Status().setCode(// rpc.Code.UNKNOWN
            2).setMessage(stackTrace));
        }
        // Best effort only, no need to wait till successfully sent.
        pendingFailuresToReport.clear();
    }
    WorkItemStatus workItemStatus = new WorkItemStatus().setWorkItemId(WINDMILL_COUNTER_UPDATE_WORK_ID).setErrors(errors).setCounterUpdates(counterUpdates);
    workUnitClient.reportWorkItemStatus(workItemStatus);
    // Send any counters appearing more than once in subsequent RPCs:
    while (!counterMultimap.isEmpty()) {
        extractUniqueCounters.run();
        workUnitClient.reportWorkItemStatus(new WorkItemStatus().setWorkItemId(WINDMILL_COUNTER_UPDATE_WORK_ID).setCounterUpdates(counterUpdates));
    }
}

Example 22 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class CounterShortIdCacheTest method testValidateShortIdsButNoUpdate.

@Test
public void testValidateShortIdsButNoUpdate() {
    CounterShortIdCache cache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    request.setWorkItemStatuses(Arrays.asList(new WorkItemStatus()));
    reply.setWorkItemServiceStates(createWorkServiceState(new Long[] { 1000L }));
    thrown.expect(IllegalArgumentException.class);
    thrown.expectMessage("Response has shortids but no corresponding CounterUpdate");
    cache.storeNewShortIds(request, reply);
}

Example 23 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class CounterShortIdCacheTest method testCacheStructuredName.

@Test
public void testCacheStructuredName() {
    CounterShortIdCache shortIdCache = new CounterShortIdCache();
    ReportWorkItemStatusRequest request = new ReportWorkItemStatusRequest();
    ReportWorkItemStatusResponse reply = new ReportWorkItemStatusResponse();
    // setup mock counters, three work statuses, one with two counters, one with one, one with none
    request.setWorkItemStatuses(createWorkStatusStructuredName(new String[] { "counter", "counter1" }, new String[] {}, new String[] { "counter2" }));
    reply.setWorkItemServiceStates(createWorkServiceState(new Long[] { 1000L, 1001L }, new Long[] {}, new Long[] { 1002L }));
    // Verify the empty case
    WorkItemStatus status1 = request.getWorkItemStatuses().get(0);
    WorkItemStatus status2 = request.getWorkItemStatuses().get(1);
    WorkItemStatus status3 = request.getWorkItemStatuses().get(2);
    shortIdCache.shortenIdsIfAvailable(status1.getCounterUpdates());
    for (CounterUpdate update : status1.getCounterUpdates()) {
        assertNull(update.getShortId());
    }
    // Add the shortIds
    shortIdCache.storeNewShortIds(request, reply);
    shortIdCache.shortenIdsIfAvailable(status1.getCounterUpdates());
    shortIdCache.shortenIdsIfAvailable(status2.getCounterUpdates());
    shortIdCache.shortenIdsIfAvailable(status3.getCounterUpdates());
    checkStatusAndShortIds(status1, 1000L, 1001L);
    checkStatusAndShortIds(status2);
    checkStatusAndShortIds(status3, 1002L);
}

Example 24 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class StreamingDataflowWorkerTest method testKeyCommitTooLargeException.

@Test
public void testKeyCommitTooLargeException() throws Exception {
    KvCoder<String, String> kvCoder = KvCoder.of(StringUtf8Coder.of(), StringUtf8Coder.of());
    List<ParallelInstruction> instructions = Arrays.asList(makeSourceInstruction(kvCoder), makeDoFnInstruction(new LargeCommitFn(), 0, kvCoder), makeSinkInstruction(kvCoder, 1));
    FakeWindmillServer server = new FakeWindmillServer(errorCollector);
    server.setExpectedExceptionCount(1);
    StreamingDataflowWorker worker = makeWorker(instructions, createTestingPipelineOptions(server), true);
    worker.setMaxWorkItemCommitBytes(1000);
    worker.start();
    server.addWorkToOffer(makeInput(1, 0, "large_key", DEFAULT_SHARDING_KEY));
    server.addWorkToOffer(makeInput(2, 0, "key", DEFAULT_SHARDING_KEY));
    server.waitForEmptyWorkQueue();
    Map<Long, Windmill.WorkItemCommitRequest> result = server.waitForAndGetCommits(1);
    assertEquals(2, result.size());
    assertEquals(makeExpectedOutput(2, 0, "key", DEFAULT_SHARDING_KEY, "key").build(), result.get(2L));
    assertTrue(result.containsKey(1L));
    WorkItemCommitRequest largeCommit = result.get(1L);
    assertEquals("large_key", largeCommit.getKey().toStringUtf8());
    assertEquals(makeExpectedTruncationRequestOutput(1, "large_key", DEFAULT_SHARDING_KEY, largeCommit.getEstimatedWorkItemCommitBytes()).build(), largeCommit);
    // Check this explicitly since the estimated commit bytes weren't actually
    // checked against an expected value in the previous step
    assertTrue(largeCommit.getEstimatedWorkItemCommitBytes() > 1000);
    // Spam worker updates a few times.
    int maxTries = 10;
    while (--maxTries > 0) {
        worker.reportPeriodicWorkerUpdates();
        Uninterruptibles.sleepUninterruptibly(1000, TimeUnit.MILLISECONDS);
    }
    // We should see an exception reported for the large commit but not the small one.
    ArgumentCaptor<WorkItemStatus> workItemStatusCaptor = ArgumentCaptor.forClass(WorkItemStatus.class);
    verify(mockWorkUnitClient, atLeast(2)).reportWorkItemStatus(workItemStatusCaptor.capture());
    List<WorkItemStatus> capturedStatuses = workItemStatusCaptor.getAllValues();
    boolean foundErrors = false;
    for (WorkItemStatus status : capturedStatuses) {
        if (!status.getErrors().isEmpty()) {
            assertFalse(foundErrors);
            foundErrors = true;
            String errorMessage = status.getErrors().get(0).getMessage();
            assertThat(errorMessage, Matchers.containsString("KeyCommitTooLargeException"));
        }
    }
    assertTrue(foundErrors);
}

Example 25 with WorkItemStatus

use of com.google.api.services.dataflow.model.WorkItemStatus in project beam by apache.

the class WorkItemStatusClientTest method populateSplitResultNativeReader.

@Test
public void populateSplitResultNativeReader() throws Exception {
    WorkItemStatus status = new WorkItemStatus();
    statusClient.setWorker(worker, executionContext);
    Position position = ReaderTestUtils.positionAtIndex(42L);
    DynamicSplitResult result = new NativeReader.DynamicSplitResultWithPosition(new DataflowReaderPosition(position));
    statusClient.populateSplitResult(status, result);
    assertThat(status.getStopPosition(), equalTo(position));
    assertThat(status.getDynamicSourceSplit(), nullValue());
}