timeout on waitfor

apache · Dec 5, 2023 · acd6875 · acd6875
1 parent 85124fb
commit acd6875
Showing 1 changed file with 10 additions and 2 deletions.
diff --git a/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationUtils.java b/src/main/java/org/apache/sysds/runtime/controlprogram/federated/FederationUtils.java
@@ -24,6 +24,7 @@
 import java.util.List;
 import java.util.Optional;
 import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
 
 import io.netty.channel.ChannelInboundHandlerAdapter;
 import io.netty.channel.ChannelOutboundHandlerAdapter;
@@ -450,8 +451,15 @@ else if (aop.aggOp.increOp.fn instanceof Builtin &&
 
 	public static void waitFor(List<Future<FederatedResponse>> responses) {
 		try {
-			for(Future<FederatedResponse> fr : responses)
-				fr.get();
+			final int timeout = ConfigurationManager.getFederatedTimeout();
+			if(timeout > 0){
+				for(Future<FederatedResponse> fr : responses)
+					fr.get(timeout, TimeUnit.SECONDS);
+			}
+			else {
+				for(Future<FederatedResponse> fr : responses)
+					fr.get();
+			}
 		}
 		catch(Exception ex) {
 			throw new DMLRuntimeException(ex);