Skip to content

Commit

Permalink
Check experiment for hotkey logging (#32285) (#32290)
Browse files Browse the repository at this point in the history
* Check experiment for hotkey logging

* Spotless
  • Loading branch information
damccorm authored Aug 22, 2024
1 parent 4173f54 commit b552e0f
Show file tree
Hide file tree
Showing 3 changed files with 7 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -423,7 +423,7 @@ public Job translate(List<DataflowPackage> packages) {
if (options.getDataflowKmsKey() != null) {
environment.setServiceKmsKeyName(options.getDataflowKmsKey());
}
if (options.isHotKeyLoggingEnabled()) {
if (options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) {
DebugOptions debugOptions = new DebugOptions();
debugOptions.setEnableHotKeyLogging(true);
environment.setDebugOptions(debugOptions);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@

import static org.apache.beam.runners.dataflow.util.TimeUtil.fromCloudDuration;
import static org.apache.beam.runners.dataflow.util.TimeUtil.fromCloudTime;
import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment;

import com.google.api.client.util.Clock;
import com.google.api.services.dataflow.model.ApproximateSplitRequest;
Expand Down Expand Up @@ -132,7 +133,7 @@ protected void reportProgressHelper() throws Exception {

// The key set the in BatchModeExecutionContext is only set in the GroupingShuffleReader
// which is the correct key. The key is also translated into a Java object in the reader.
if (options.isHotKeyLoggingEnabled()) {
if (options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging")) {
hotKeyLogger.logHotKeyDetection(
hotKeyDetection.getUserStepName(),
TimeUtil.fromCloudDuration(hotKeyDetection.getHotKeyAge()),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@
*/
package org.apache.beam.runners.dataflow.worker.windmill.work.processing;

import static org.apache.beam.sdk.options.ExperimentalOptions.hasExperiment;

import com.google.api.services.dataflow.model.MapTask;
import com.google.auto.value.AutoValue;
import java.util.Collection;
Expand Down Expand Up @@ -368,7 +370,8 @@ private ExecuteWorkResult executeWork(
Duration hotKeyAge = Duration.millis(hotKeyInfo.getHotKeyAgeUsec() / 1000);

String stepName = getShuffleTaskStepName(computationState.getMapTask());
if (options.isHotKeyLoggingEnabled() && keyCoder.isPresent()) {
if ((options.isHotKeyLoggingEnabled() || hasExperiment(options, "enable_hot_key_logging"))
&& keyCoder.isPresent()) {
hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge, executionKey);
} else {
hotKeyLogger.logHotKeyDetection(stepName, hotKeyAge);
Expand Down

0 comments on commit b552e0f

Please sign in to comment.