[GitHub] [systemds] sebwrede commented on a change in pull request #1525: [SYSTEMDS-3280] WIP: Add homomorphic encryption to Federated Parameter Server

[GitBox <gi...@apache.org>]

sebwrede commented on a change in pull request #1525:
URL: https://github.com/apache/systemds/pull/1525#discussion_r793718603



##########
File path: src/main/java/org/apache/sysds/runtime/controlprogram/paramserv/FederatedPSControlThread.java
##########
@@ -160,21 +161,45 @@ public void setup(double weightingFactor) {
 			PROG_END);
 
 		// write program and meta data to worker
-		Future<FederatedResponse> udfResponse = _featuresData.executeFederatedOperation(
-			new FederatedRequest(RequestType.EXEC_UDF, _featuresData.getVarID(),
-				new SetupFederatedWorker(_batchSize, dataSize, _possibleBatchesPerLocalEpoch,
-					programSerialized, _inst.getNamespace(), _inst.getFunctionName(),
-					_ps.getAggInst().getFunctionName(), _ec.getListObject("hyperparams"),
-					_modelVarID, _nbatches, _modelAvg)));
+		Future<FederatedResponse> udfResponse;
+		//_use_homomorphic_encryption = false;
+		if (_use_homomorphic_encryption) {
+			// TODO: generate a here
+			byte[] a = ((HEParamServer)_ps).generateA();
+			// generate pk[i] on each client and return it
+			udfResponse = _featuresData.executeFederatedOperation(
+					new FederatedRequest(RequestType.EXEC_UDF, _featuresData.getVarID(),
+							new SetupHEFederatedWorker(_batchSize, dataSize, _possibleBatchesPerLocalEpoch,
+									programSerialized, _inst.getNamespace(), _inst.getFunctionName(),
+									_ps.getAggInst().getFunctionName(), _ec.getListObject("hyperparams"),
+									_modelVarID, _nbatches, a)));
+		} else {
+			udfResponse = _featuresData.executeFederatedOperation(
+					new FederatedRequest(RequestType.EXEC_UDF, _featuresData.getVarID(),
+							new SetupFederatedWorker(_batchSize, dataSize, _possibleBatchesPerLocalEpoch,
+									programSerialized, _inst.getNamespace(), _inst.getFunctionName(),
+									_ps.getAggInst().getFunctionName(), _ec.getListObject("hyperparams"),
+									_modelVarID, _nbatches, _modelAvg)));
+		}

Review comment:
       You could have less code duplication if the constructor call to SetupFederatedWorker and SetupHEFederatedWorker is in the if-else blocks and the method call `executeFederatedOperation` and `new FederatedRequest()` is after the if-else on whatever SetupFederatedWorker instance just created. 
   It could be even further reduced with new, shorter constructors for SetupFederatedWorker and SetupHEFederatedWorker followed by a method call to setParams which could set the general/shared params of the object. 
   This would make it easier to spot the difference between the two federated worker setups, but the functionality would be the same. 

##########
File path: src/main/java/org/apache/sysds/runtime/controlprogram/paramserv/FederatedPSControlThread.java
##########
@@ -324,6 +401,7 @@ protected ListObject pullModel() {
 	}
 
 	protected void weightAndPushGradients(ListObject gradients) {
+		assert (!(_weighting && _use_homomorphic_encryption)) : "weights and homomorphic encryption are not supported together";

Review comment:
       I think this could also be implemented as an exception. It is a relevant stopping condition if the user is trying to execute some parameter server setup which is not possible and the user needs to get this information. Usually we throw something like this as an exception in SystemDS.

##########
File path: src/main/java/org/apache/sysds/runtime/controlprogram/paramserv/FederatedPSControlThread.java
##########
@@ -354,11 +432,12 @@ protected void computeWithBatchUpdates() {
 				int localStartBatchNum = getNextLocalBatchNum(currentLocalBatchNumber++, _possibleBatchesPerLocalEpoch);
 				ListObject model = pullModel();
 				ListObject gradients = computeGradientsForNBatches(model, 1, localStartBatchNum);
-				if (_modelAvg)
+				if (_modelAvg && !_use_homomorphic_encryption)
+					// we can't call the agg fn if we use HE, because it is implemented homomorphically in SEALServer::aggregateCiphertexts

Review comment:
       We could consider throwing an exception here as well. 

##########
File path: src/main/java/org/apache/sysds/runtime/controlprogram/paramserv/FederatedPSControlThread.java
##########
@@ -258,9 +283,61 @@ public FederatedResponse execute(ExecutionContext ec, Data... data) {
 		}
 	}
 
-	/**
+	private static class SetupHEFederatedWorker extends SetupFederatedWorker {
+		private static final long serialVersionUID = 9128347291804980123L;
+
+		byte[] _partial_pubkey_a;
+
+        protected SetupHEFederatedWorker(long batchSize, long dataSize, int possibleBatchesPerLocalEpoch,
+                                         String programString, String namespace, String gradientsFunctionName, String aggregationFunctionName,
+                                         ListObject hyperParams, long modelVarID, int nbatches, byte[] partial_pubkey_a) {
+            // delegate everything to parent class. set modelAvg to true, as it is the only supported case
+            super(batchSize, dataSize, possibleBatchesPerLocalEpoch, programString, namespace, gradientsFunctionName,
+                    aggregationFunctionName, hyperParams, modelVarID, nbatches, true);
+
+            _partial_pubkey_a = partial_pubkey_a;
+        }
+
+        @Override
+        public FederatedResponse execute(ExecutionContext ec, Data... data) {
+			// TODO: set other CKKS parameters
+			// TODO generate partial public key
+			SEALClient sc = new SEALClient(_partial_pubkey_a);
+			ec.setSealClient(sc);
+			PublicKey partial_pubkey = sc.generatePartialPublicKey();
+
+			FederatedResponse res = super.execute(ec, data);
+			if (!res.isSuccessful()) {
+				return res;
+			}
+
+			return new FederatedResponse(FederatedResponse.ResponseType.SUCCESS, partial_pubkey);
+        }
+    }
+		/**
 	 * Teardown UDF executed on the federated worker
 	 */
+	private static class SetPublicKeyFederatedWorker extends FederatedUDF {
+		private static final long serialVersionUID = -1536502123123318969L;
+		private final PublicKey _public_key;
+
+		protected SetPublicKeyFederatedWorker(PublicKey public_key) {
+			super(new long[]{});
+			_public_key = public_key;
+		}
+
+		@Override
+		public FederatedResponse execute(ExecutionContext ec, Data... data) {
+			ec.getSealClient().setPublicKey(_public_key);
+			return new FederatedResponse(FederatedResponse.ResponseType.SUCCESS);
+		}

Review comment:
       Why does this need to be done in a separate UDF? 
   Isn't it possible to do it in `SetupHEFederatedWorker`?

##########
File path: src/main/java/org/apache/sysds/runtime/controlprogram/paramserv/HEParamServer.java
##########
@@ -0,0 +1,136 @@
+package org.apache.sysds.runtime.controlprogram.paramserv;
+
+import org.apache.sysds.parser.Statement;
+import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
+import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
+import org.apache.sysds.runtime.controlprogram.paramserv.homomorphicEncryption.PublicKey;
+import org.apache.sysds.runtime.controlprogram.paramserv.homomorphicEncryption.SEALServer;
+import org.apache.sysds.runtime.instructions.cp.CiphertextMatrix;
+import org.apache.sysds.runtime.instructions.cp.ListObject;
+import org.apache.sysds.runtime.instructions.cp.PlaintextMatrix;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/// This class implements Homomorphic Encryption (HE) for LocalParamServer. It only supports modelAvg=true.
+public class HEParamServer extends LocalParamServer {
+    private int _thread_counter = 0;
+    private final List<FederatedPSControlThread> _threads;
+    private final List<Object> _result_buffer; // one per thread
+    private Object _result;
+    private final SEALServer _seal_server;
+
+    public static HEParamServer create(ListObject model, String aggFunc, Statement.PSUpdateType updateType,
+                                          Statement.PSFrequency freq, ExecutionContext ec, int workerNum, String valFunc, int numBatchesPerEpoch,
+                                          MatrixObject valFeatures, MatrixObject valLabels, int nbatches)
+    {
+        return new HEParamServer(model, aggFunc, updateType, freq, ec,
+                workerNum, valFunc, numBatchesPerEpoch, valFeatures, valLabels, nbatches);
+    }
+
+    private HEParamServer(ListObject model, String aggFunc, Statement.PSUpdateType updateType,
+                             Statement.PSFrequency freq, ExecutionContext ec, int workerNum, String valFunc, int numBatchesPerEpoch,
+                             MatrixObject valFeatures, MatrixObject valLabels, int nbatches)
+    {
+        super(model, aggFunc, updateType, freq, ec, workerNum, valFunc, numBatchesPerEpoch, valFeatures, valLabels, nbatches, true);
+
+        _seal_server = new SEALServer();
+
+        _threads = Collections.synchronizedList(new ArrayList<>(workerNum));
+        for (int i = 0; i < getNumWorkers(); i++) {
+            _threads.add(null);
+        }
+
+        _result_buffer = new ArrayList<>(workerNum);
+        resetResultBuffer();
+    }
+
+    public void registerThread(int thread_id, FederatedPSControlThread thread) {
+        _threads.set(thread_id, thread);
+    }
+
+    private synchronized void resetResultBuffer() {
+        _result_buffer.clear();
+        for (int i = 0; i < getNumWorkers(); i++) {
+            _result_buffer.add(null);
+        }
+    }
+
+    public byte[] generateA() {
+        return _seal_server.generateA();
+    }
+
+    public PublicKey aggregatePartialPublicKeys(PublicKey[] partial_public_keys) {
+        return _seal_server.aggregatePartialPublicKeys(partial_public_keys);
+    }
+
+    // this method collects all T Objects from each worker into a list and then calls f once on this list to produce
+    // another T, which it returns.
+    private synchronized <T,U> U collectAndDo(int workerId, T obj, Function<List<T>, U> f) {
+        _result_buffer.set(workerId, obj);
+        _thread_counter++;
+
+        if (_thread_counter == getNumWorkers()) {
+            List<T> buf = _result_buffer.stream().map(x -> (T)x).collect(Collectors.toList());
+            _result = f.apply(buf);
+            resetResultBuffer();
+            _thread_counter = 0;
+            notifyAll();
+        } else {
+            try {
+                wait();
+            } catch (InterruptedException i) {
+                throw new RuntimeException("thread interrupted");
+            }
+        }
+
+        return (U) _result;
+    }
+
+    private CiphertextMatrix[] homomorphicAggregation(List<ListObject> encrypted_models) {
+        CiphertextMatrix[] result = new CiphertextMatrix[encrypted_models.get(0).getLength()];
+        for (int matrix_idx = 0; matrix_idx < encrypted_models.get(0).getLength(); matrix_idx++) {
+            CiphertextMatrix[] summands = new CiphertextMatrix[encrypted_models.size()];
+            for (int i = 0; i < encrypted_models.size(); i++) {
+                summands[i] = (CiphertextMatrix) encrypted_models.get(i).getData(matrix_idx);
+            }
+            result[matrix_idx] = _seal_server.accumulateCiphertexts(summands);;
+        }
+        return result;
+    }
+
+    private Void homomorphicAverage(CiphertextMatrix[] encrypted_sums, List<PlaintextMatrix[]> partial_decryptions) {

Review comment:
       Is it necessary to return `Void` instead of just making the method `void`. The use of `Void` return type is very rare in Java. 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: dev-unsubscribe@systemds.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org