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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
