Skip to content

Commit

Permalink
[CELEBORN-1700] Flink supports fallback to vanilla Flink built-in shu…
Browse files Browse the repository at this point in the history 
…ffle implementation
  • Loading branch information
@SteNicholas
SteNicholas committed Nov 22, 2024
1 parent 094fe28 commit dfc93da
Show file tree
Hide file tree
Showing 37 changed files with 1,572 additions and 261 deletions.
15 changes: 7 additions & 8 deletions ...mmon/src/main/java/org/apache/celeborn/plugin/flink/AbstractRemoteShuffleEnvironment.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
import static org.apache.flink.runtime.io.network.metrics.NettyShuffleMetricFactory.*;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
Expand Down Expand Up @@ -129,12 +128,8 @@ public ShuffleIOOwnerContext createShuffleIOOwnerContext(
nettyGroup.addGroup(METRIC_GROUP_INPUT));
}

public void releasePartitionsLocally(Collection<ResultPartitionID> partitionIds) {
throw new FlinkRuntimeException("Not implemented yet.");
}

public Collection<ResultPartitionID> getPartitionsOccupyingLocalResources() {
return new ArrayList<>();
return resultPartitionManager.getUnreleasedPartitions();
}

@VisibleForTesting
Expand Down Expand Up @@ -176,13 +171,17 @@ public List<IndexedInputGate> createInputGates(
IndexedInputGate[] inputGates = new IndexedInputGate[inputGateDescriptors.size()];
for (int gateIndex = 0; gateIndex < inputGates.length; gateIndex++) {
InputGateDeploymentDescriptor igdd = inputGateDescriptors.get(gateIndex);
IndexedInputGate inputGate = createInputGateInternal(ownerContext, gateIndex, igdd);
IndexedInputGate inputGate =
createInputGateInternal(ownerContext, producerStateProvider, gateIndex, igdd);
inputGates[gateIndex] = inputGate;
}
return Arrays.asList(inputGates);
}
}

abstract IndexedInputGate createInputGateInternal(
ShuffleIOOwnerContext ownerContext, int gateIndex, InputGateDeploymentDescriptor igdd);
ShuffleIOOwnerContext ownerContext,
PartitionProducerStateProvider producerStateProvider,
int gateIndex,
InputGateDeploymentDescriptor igdd);
}
181 changes: 121 additions & 60 deletions client-flink/common/src/main/java/org/apache/celeborn/plugin/flink/RemoteShuffleMaster.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,23 @@
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;

import javax.annotation.Nullable;

import org.apache.flink.api.common.BatchShuffleMode;
import org.apache.flink.api.common.JobID;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.ExecutionOptions;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.runtime.io.network.NettyShuffleServiceFactory;
import org.apache.flink.runtime.io.network.partition.ResultPartitionType;
import org.apache.flink.runtime.shuffle.JobShuffleContext;
import org.apache.flink.runtime.shuffle.NettyShuffleDescriptor;
import org.apache.flink.runtime.shuffle.NettyShuffleMaster;
import org.apache.flink.runtime.shuffle.PartitionDescriptor;
import org.apache.flink.runtime.shuffle.ProducerDescriptor;
import org.apache.flink.runtime.shuffle.ShuffleDescriptor;
Expand All @@ -43,16 +50,20 @@

import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.plugin.flink.fallback.ShuffleFallbackPolicyRunner;
import org.apache.celeborn.plugin.flink.utils.FlinkUtils;

public class RemoteShuffleMaster implements ShuffleMaster<RemoteShuffleDescriptor> {
public class RemoteShuffleMaster implements ShuffleMaster<ShuffleDescriptor> {
private static final Logger LOG = LoggerFactory.getLogger(RemoteShuffleMaster.class);
private final CelebornConf conf;
private final ShuffleMasterContext shuffleMasterContext;
// Flink JobId -> Celeborn register shuffleIds
private final Map<JobID, Set<Integer>> jobShuffleIds = JavaUtils.newConcurrentHashMap();
private final ConcurrentHashMap.KeySetView<JobID, Boolean> nettyJobIds =
ConcurrentHashMap.newKeySet();
private String celebornAppId;
private volatile LifecycleManager lifecycleManager;
private final ShuffleTaskInfo shuffleTaskInfo = new ShuffleTaskInfo();
Expand All @@ -62,20 +73,26 @@ public class RemoteShuffleMaster implements ShuffleMaster<RemoteShuffleDescripto
"celeborn-client-remote-shuffle-master-executor");
private final ResultPartitionAdapter resultPartitionDelegation;
private final long lifecycleManagerTimestamp;
private final NettyShuffleServiceFactory nettyShuffleServiceFactory;
private volatile NettyShuffleMaster nettyShuffleMaster;

public RemoteShuffleMaster(
ShuffleMasterContext shuffleMasterContext, ResultPartitionAdapter resultPartitionDelegation) {
ShuffleMasterContext shuffleMasterContext,
ResultPartitionAdapter resultPartitionDelegation,
@Nullable NettyShuffleServiceFactory nettyShuffleServiceFactory) {
Configuration configuration = shuffleMasterContext.getConfiguration();
checkShuffleConfig(configuration);
this.conf = FlinkUtils.toCelebornConf(configuration);
this.shuffleMasterContext = shuffleMasterContext;
this.resultPartitionDelegation = resultPartitionDelegation;
this.lifecycleManagerTimestamp = System.currentTimeMillis();
this.nettyShuffleServiceFactory = nettyShuffleServiceFactory;
}

@Override
public void registerJob(JobShuffleContext context) {
JobID jobID = context.getJobId();
LOG.info("Register job: {}.", jobID);
if (lifecycleManager == null) {
synchronized (RemoteShuffleMaster.class) {
if (lifecycleManager == null) {
Expand All @@ -87,17 +104,31 @@ public void registerJob(JobShuffleContext context) {
}
}

Set<Integer> previousShuffleIds = jobShuffleIds.putIfAbsent(jobID, new HashSet<>());
LOG.info("Register job: {}.", jobID);
if (previousShuffleIds != null) {
throw new RuntimeException("Duplicated registration job: " + jobID);
try {
if (nettyShuffleServiceFactory != null
&& ShuffleFallbackPolicyRunner.applyFallbackPolicies(context, conf, lifecycleManager)) {
LOG.warn("Fallback to vanilla Flink NettyShuffleMaster for job: {}.", jobID);
nettyJobIds.add(jobID);
nettyShuffleMaster().registerJob(context);
} else {
Set<Integer> previousShuffleIds = jobShuffleIds.putIfAbsent(jobID, new HashSet<>());
if (previousShuffleIds != null) {
throw new RuntimeException("Duplicated registration job: " + jobID);
}
shuffleResourceTracker.registerJob(context);
}
} catch (CelebornIOException e) {
throw new RuntimeException(e);
}
shuffleResourceTracker.registerJob(context);
}

@Override
public void unregisterJob(JobID jobID) {
LOG.info("Unregister job: {}.", jobID);
if (nettyJobIds.remove(jobID)) {
nettyShuffleMaster().unregisterJob(jobID);
return;
}
Set<Integer> shuffleIds = jobShuffleIds.remove(jobID);
if (shuffleIds != null) {
executor.execute(
Expand All @@ -116,46 +147,56 @@ public void unregisterJob(JobID jobID) {
}

@Override
public CompletableFuture<RemoteShuffleDescriptor> registerPartitionWithProducer(
public CompletableFuture<ShuffleDescriptor> registerPartitionWithProducer(
JobID jobID, PartitionDescriptor partitionDescriptor, ProducerDescriptor producerDescriptor) {
return CompletableFuture.supplyAsync(
() -> {
Set<Integer> shuffleIds = jobShuffleIds.get(jobID);
if (shuffleIds == null) {
throw new RuntimeException("Can not find job in lifecycleManager, job: " + jobID);
}
if (nettyJobIds.contains(jobID)) {
try {
return nettyShuffleMaster()
.registerPartitionWithProducer(jobID, partitionDescriptor, producerDescriptor)
.get();
} catch (InterruptedException | ExecutionException e) {
throw new RuntimeException(e);
}
} else {
Set<Integer> shuffleIds = jobShuffleIds.get(jobID);
if (shuffleIds == null) {
throw new RuntimeException("Can not find job in lifecycleManager, job: " + jobID);
}

FlinkResultPartitionInfo resultPartitionInfo =
new FlinkResultPartitionInfo(jobID, partitionDescriptor, producerDescriptor);
ShuffleResourceDescriptor shuffleResourceDescriptor =
shuffleTaskInfo.genShuffleResourceDescriptor(
resultPartitionInfo.getShuffleId(),
resultPartitionInfo.getTaskId(),
resultPartitionInfo.getAttemptId());
FlinkResultPartitionInfo resultPartitionInfo =
new FlinkResultPartitionInfo(jobID, partitionDescriptor, producerDescriptor);
ShuffleResourceDescriptor shuffleResourceDescriptor =
shuffleTaskInfo.genShuffleResourceDescriptor(
resultPartitionInfo.getShuffleId(),
resultPartitionInfo.getTaskId(),
resultPartitionInfo.getAttemptId());

synchronized (shuffleIds) {
shuffleIds.add(shuffleResourceDescriptor.getShuffleId());
}
synchronized (shuffleIds) {
shuffleIds.add(shuffleResourceDescriptor.getShuffleId());
}

RemoteShuffleResource remoteShuffleResource =
new RemoteShuffleResource(
lifecycleManager.getHost(),
lifecycleManager.getPort(),
lifecycleManagerTimestamp,
shuffleResourceDescriptor);
RemoteShuffleResource remoteShuffleResource =
new RemoteShuffleResource(
lifecycleManager.getHost(),
lifecycleManager.getPort(),
lifecycleManagerTimestamp,
shuffleResourceDescriptor);

shuffleResourceTracker.addPartitionResource(
jobID,
shuffleResourceDescriptor.getShuffleId(),
shuffleResourceDescriptor.getPartitionId(),
resultPartitionInfo.getResultPartitionId());
shuffleResourceTracker.addPartitionResource(
jobID,
shuffleResourceDescriptor.getShuffleId(),
shuffleResourceDescriptor.getPartitionId(),
resultPartitionInfo.getResultPartitionId());

return new RemoteShuffleDescriptor(
celebornAppId,
jobID,
resultPartitionInfo.getShuffleId(),
resultPartitionInfo.getResultPartitionId(),
remoteShuffleResource);
return new RemoteShuffleDescriptor(
celebornAppId,
jobID,
resultPartitionInfo.getShuffleId(),
resultPartitionInfo.getResultPartitionId(),
remoteShuffleResource);
}
},
executor);
}
Expand All @@ -164,30 +205,34 @@ public CompletableFuture<RemoteShuffleDescriptor> registerPartitionWithProducer(
public void releasePartitionExternally(ShuffleDescriptor shuffleDescriptor) {
executor.execute(
() -> {
if (!(shuffleDescriptor instanceof RemoteShuffleDescriptor)) {
if (shuffleDescriptor instanceof RemoteShuffleDescriptor) {
try {
RemoteShuffleDescriptor descriptor = (RemoteShuffleDescriptor) shuffleDescriptor;
RemoteShuffleResource shuffleResource = descriptor.getShuffleResource();
ShuffleResourceDescriptor resourceDescriptor =
shuffleResource.getMapPartitionShuffleDescriptor();
LOG.debug("release partition resource: {}.", resourceDescriptor);
lifecycleManager.releasePartition(
resourceDescriptor.getShuffleId(), resourceDescriptor.getPartitionId());
shuffleResourceTracker.removePartitionResource(
descriptor.getJobId(),
resourceDescriptor.getShuffleId(),
resourceDescriptor.getPartitionId());
} catch (Throwable throwable) {
// it is not a problem if we failed to release the target data partition
// because the session timeout mechanism will do the work for us latter
LOG.debug("Failed to release data partition {}.", shuffleDescriptor, throwable);
}
} else if (shuffleDescriptor instanceof NettyShuffleDescriptor) {
nettyShuffleMaster().releasePartitionExternally(shuffleDescriptor);
} else {
LOG.error(
"Only RemoteShuffleDescriptor is supported {}.",
shuffleDescriptor.getClass().getName());
"Unsupported shuffle descriptor {}. Only supports {} and {}",
shuffleDescriptor.getClass().getName(),
RemoteShuffleDescriptor.class.getName(),
NettyShuffleDescriptor.class.getName());
shuffleMasterContext.onFatalError(
new RuntimeException("Illegal shuffle descriptor type."));
return;
}
try {
RemoteShuffleDescriptor descriptor = (RemoteShuffleDescriptor) shuffleDescriptor;
RemoteShuffleResource shuffleResource = descriptor.getShuffleResource();
ShuffleResourceDescriptor resourceDescriptor =
shuffleResource.getMapPartitionShuffleDescriptor();
LOG.debug("release partition resource: {}.", resourceDescriptor);
lifecycleManager.releasePartition(
resourceDescriptor.getShuffleId(), resourceDescriptor.getPartitionId());
shuffleResourceTracker.removePartitionResource(
descriptor.getJobId(),
resourceDescriptor.getShuffleId(),
resourceDescriptor.getPartitionId());
} catch (Throwable throwable) {
// it is not a problem if we failed to release the target data partition
// because the session timeout mechanism will do the work for us latter
LOG.debug("Failed to release data partition {}.", shuffleDescriptor, throwable);
}
});
}
Expand Down Expand Up @@ -224,11 +269,16 @@ public MemorySize computeShuffleMemorySizeForTask(
@Override
public void close() throws Exception {
try {
nettyJobIds.clear();
jobShuffleIds.clear();
LifecycleManager manager = lifecycleManager;
if (null != manager) {
manager.stop();
}
if (nettyShuffleMaster != null) {
nettyShuffleMaster.close();
nettyShuffleMaster = null;
}
} catch (Exception e) {
LOG.warn("Encounter exception when shutdown: {}", e.getMessage(), e);
}
Expand All @@ -254,4 +304,15 @@ private void checkShuffleConfig(Configuration configuration) {
BatchShuffleMode.ALL_EXCHANGES_BLOCKING.name()));
}
}

private NettyShuffleMaster nettyShuffleMaster() {
if (nettyShuffleMaster == null) {
synchronized (this) {
if (nettyShuffleMaster == null) {
nettyShuffleMaster = nettyShuffleServiceFactory.createShuffleMaster(shuffleMasterContext);
}
}
}
return nettyShuffleMaster;
}
}
57 changes: 57 additions & 0 deletions ...k/common/src/main/java/org/apache/celeborn/plugin/flink/fallback/ForceFallbackPolicy.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
/*
* 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.
*/

package org.apache.celeborn.plugin.flink.fallback;

import org.apache.flink.runtime.shuffle.JobShuffleContext;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import org.apache.celeborn.client.LifecycleManager;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.protocol.FallbackPolicy;

public class ForceFallbackPolicy implements ShuffleFallbackPolicy {

private static final Logger LOG = LoggerFactory.getLogger(ForceFallbackPolicy.class);

public static final ForceFallbackPolicy INSTANCE = new ForceFallbackPolicy();
/**
* If celeborn.client.flink.shuffle.fallback.policy is ALWAYS, fallback to flink built-in shuffle
* implementation.
*
* @param shuffleContext The job shuffle context of Flink.
* @param celebornConf The configuration of Celeborn.
* @param lifecycleManager The {@link LifecycleManager} of Celeborn.
* @return Return true if celeborn.client.flink.shuffle.fallback.policy is ALWAYS, otherwise
* false.
*/
@Override
public boolean needFallback(
JobShuffleContext shuffleContext,
CelebornConf celebornConf,
LifecycleManager lifecycleManager) {
FallbackPolicy shuffleFallbackPolicy = celebornConf.flinkShuffleFallbackPolicy();
if (FallbackPolicy.ALWAYS.equals(shuffleFallbackPolicy)) {
LOG.warn(
"{} is {}, forcibly fallback to flink built-in shuffle implementation.",
CelebornConf.FLINK_SHUFFLE_FALLBACK_POLICY().key(),
FallbackPolicy.ALWAYS.name());
}
return FallbackPolicy.ALWAYS.equals(shuffleFallbackPolicy);
}
}
Loading

0 comments on commit dfc93da

Please sign in to comment.