src/main/java/org/apache/sysds/runtime/instructions/cp/ParamservBuiltinCPInstruction.java - systemds - Git at Google

 /*
  * 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.sysds.runtime.instructions.cp;

 import static org.apache.sysds.parser.Statement.PS_AGGREGATION_FUN;
 import static org.apache.sysds.parser.Statement.PS_BATCH_SIZE;
 import static org.apache.sysds.parser.Statement.PS_EPOCHS;
 import static org.apache.sysds.parser.Statement.PS_FEATURES;
 import static org.apache.sysds.parser.Statement.PS_FREQUENCY;
 import static org.apache.sysds.parser.Statement.PS_HYPER_PARAMS;
 import static org.apache.sysds.parser.Statement.PS_LABELS;
 import static org.apache.sysds.parser.Statement.PS_MODE;
 import static org.apache.sysds.parser.Statement.PS_MODEL;
 import static org.apache.sysds.parser.Statement.PS_PARALLELISM;
 import static org.apache.sysds.parser.Statement.PS_SCHEME;
 import static org.apache.sysds.parser.Statement.PS_UPDATE_FUN;
 import static org.apache.sysds.parser.Statement.PS_UPDATE_TYPE;

 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;

 import org.apache.commons.lang3.concurrent.BasicThreadFactory;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.spark.network.server.TransportServer;
 import org.apache.spark.util.LongAccumulator;
 import org.apache.sysds.api.DMLScript;
 import org.apache.sysds.hops.recompile.Recompiler;
 import org.apache.sysds.lops.LopProperties;
 import org.apache.sysds.parser.Statement.PSFrequency;
 import org.apache.sysds.parser.Statement.PSModeType;
 import org.apache.sysds.parser.Statement.PSScheme;
 import org.apache.sysds.parser.Statement.PSUpdateType;
 import org.apache.sysds.runtime.DMLRuntimeException;
 import org.apache.sysds.runtime.controlprogram.LocalVariableMap;
 import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
 import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
 import org.apache.sysds.runtime.controlprogram.context.SparkExecutionContext;
 import org.apache.sysds.runtime.controlprogram.paramserv.LocalPSWorker;
 import org.apache.sysds.runtime.controlprogram.paramserv.LocalParamServer;
 import org.apache.sysds.runtime.controlprogram.paramserv.ParamServer;
 import org.apache.sysds.runtime.controlprogram.paramserv.ParamservUtils;
 import org.apache.sysds.runtime.controlprogram.paramserv.SparkPSBody;
 import org.apache.sysds.runtime.controlprogram.paramserv.SparkPSWorker;
 import org.apache.sysds.runtime.controlprogram.paramserv.dp.DataPartitionLocalScheme;
 import org.apache.sysds.runtime.controlprogram.paramserv.dp.LocalDataPartitioner;
 import org.apache.sysds.runtime.controlprogram.paramserv.rpc.PSRpcFactory;
 import org.apache.sysds.runtime.controlprogram.parfor.stat.InfrastructureAnalyzer;
 import org.apache.sysds.runtime.controlprogram.parfor.stat.Timing;
 import org.apache.sysds.runtime.matrix.operators.Operator;
 import org.apache.sysds.runtime.util.ProgramConverter;
 import org.apache.sysds.utils.Statistics;

 public class ParamservBuiltinCPInstruction extends ParameterizedBuiltinCPInstruction {
 	private static final Log LOG = LogFactory.getLog(ParamservBuiltinCPInstruction.class.getName());

 	private static final int DEFAULT_BATCH_SIZE = 64;
 	private static final PSFrequency DEFAULT_UPDATE_FREQUENCY = PSFrequency.EPOCH;
 	private static final PSScheme DEFAULT_SCHEME = PSScheme.DISJOINT_CONTIGUOUS;
 	private static final PSModeType DEFAULT_MODE = PSModeType.LOCAL;
 	private static final PSUpdateType DEFAULT_TYPE = PSUpdateType.ASP;

 	public ParamservBuiltinCPInstruction(Operator op, LinkedHashMap<String, String> paramsMap, CPOperand out, String opcode, String istr) {
 		super(op, paramsMap, out, opcode, istr);
 	}

 	@Override
 	public void processInstruction(ExecutionContext ec) {
 		PSModeType mode = getPSMode();
 		switch (mode) {
 			case LOCAL:
 				runLocally(ec, mode);
 				break;
 			case REMOTE_SPARK:
 				runOnSpark((SparkExecutionContext) ec, mode);
 				break;
 			default:
 				throw new DMLRuntimeException(String.format("Paramserv func: not support mode %s", mode));
 		}
 	}

 	@SuppressWarnings("resource")
 	private void runOnSpark(SparkExecutionContext sec, PSModeType mode) {
 		Timing tSetup = DMLScript.STATISTICS ? new Timing(true) : null;

 		int workerNum = getWorkerNum(mode);
 		String updFunc = getParam(PS_UPDATE_FUN);
 		String aggFunc = getParam(PS_AGGREGATION_FUN);

 		// Get the compiled execution context
 		LocalVariableMap newVarsMap = createVarsMap(sec);
 		// Level of par is 1 in spark backend because one worker will be launched per task
 		ExecutionContext newEC = ParamservUtils.createExecutionContext(sec, newVarsMap, updFunc, aggFunc, 1);

 		// Create the agg service's execution context
 		ExecutionContext aggServiceEC = ParamservUtils.copyExecutionContext(newEC, 1).get(0);

 		// Create the parameter server
 		ListObject model = sec.getListObject(getParam(PS_MODEL));
 		ParamServer ps = createPS(mode, aggFunc, getUpdateType(), workerNum, model, aggServiceEC);

 		// Get driver host
 		String host = sec.getSparkContext().getConf().get("spark.driver.host");

 		// Create the netty server for ps
 		TransportServer server = PSRpcFactory.createServer(sec.getSparkContext().getConf(),(LocalParamServer) ps, host); // Start the server

 		// Force all the instructions to CP type
 		Recompiler.recompileProgramBlockHierarchy2Forced(
 			newEC.getProgram().getProgramBlocks(), 0, new HashSet<>(), LopProperties.ExecType.CP);

 		// Serialize all the needed params for remote workers
 		SparkPSBody body = new SparkPSBody(newEC);
 		HashMap<String, byte[]> clsMap = new HashMap<>();
 		String program = ProgramConverter.serializeSparkPSBody(body, clsMap);

 		// Add the accumulators for statistics
 		LongAccumulator aSetup = sec.getSparkContext().sc().longAccumulator("setup");
 		LongAccumulator aWorker = sec.getSparkContext().sc().longAccumulator("workersNum");
 		LongAccumulator aUpdate = sec.getSparkContext().sc().longAccumulator("modelUpdate");
 		LongAccumulator aIndex = sec.getSparkContext().sc().longAccumulator("batchIndex");
 		LongAccumulator aGrad = sec.getSparkContext().sc().longAccumulator("gradCompute");
 		LongAccumulator aRPC = sec.getSparkContext().sc().longAccumulator("rpcRequest");
 		LongAccumulator aBatch = sec.getSparkContext().sc().longAccumulator("numBatches");
 		LongAccumulator aEpoch = sec.getSparkContext().sc().longAccumulator("numEpochs");

 		// Create remote workers
 		SparkPSWorker worker = new SparkPSWorker(getParam(PS_UPDATE_FUN), getParam(PS_AGGREGATION_FUN),
 			getFrequency(), getEpochs(), getBatchSize(), program, clsMap, sec.getSparkContext().getConf(),
 			server.getPort(), aSetup, aWorker, aUpdate, aIndex, aGrad, aRPC, aBatch, aEpoch);

 		if (DMLScript.STATISTICS)
 			Statistics.accPSSetupTime((long) tSetup.stop());

 		MatrixObject features = sec.getMatrixObject(getParam(PS_FEATURES));
 		MatrixObject labels = sec.getMatrixObject(getParam(PS_LABELS));
 		try {
 			ParamservUtils.doPartitionOnSpark(sec, features, labels, getScheme(), workerNum) // Do data partitioning
 				.foreach(worker); // Run remote workers
 		} catch (Exception e) {
 			throw new DMLRuntimeException("Paramserv function failed: ", e);
 		} finally {
 			server.close(); // Stop the netty server
 		}

 		// Accumulate the statistics for remote workers
 		if (DMLScript.STATISTICS) {
 			Statistics.accPSSetupTime(aSetup.value());
 			Statistics.incWorkerNumber(aWorker.value());
 			Statistics.accPSLocalModelUpdateTime(aUpdate.value());
 			Statistics.accPSBatchIndexingTime(aIndex.value());
 			Statistics.accPSGradientComputeTime(aGrad.value());
 			Statistics.accPSRpcRequestTime(aRPC.value());
 		}

 		// Fetch the final model from ps
 		sec.setVariable(output.getName(), ps.getResult());
 	}

 	private void runLocally(ExecutionContext ec, PSModeType mode) {
 		Timing tSetup = DMLScript.STATISTICS ? new Timing(true) : null;
 		int workerNum = getWorkerNum(mode);
 		BasicThreadFactory factory = new BasicThreadFactory.Builder()
 			.namingPattern("workers-pool-thread-%d").build();
 		ExecutorService es = Executors.newFixedThreadPool(workerNum, factory);
 		String updFunc = getParam(PS_UPDATE_FUN);
 		String aggFunc = getParam(PS_AGGREGATION_FUN);

 		// Get the compiled execution context
 		LocalVariableMap newVarsMap = createVarsMap(ec);
 		ExecutionContext newEC = ParamservUtils.createExecutionContext(ec, newVarsMap, updFunc, aggFunc, getParLevel(workerNum));

 		// Create workers' execution context
 		List<ExecutionContext> workerECs = ParamservUtils.copyExecutionContext(newEC, workerNum);

 		// Create the agg service's execution context
 		ExecutionContext aggServiceEC = ParamservUtils.copyExecutionContext(newEC, 1).get(0);

 		PSFrequency freq = getFrequency();
 		PSUpdateType updateType = getUpdateType();

 		// Create the parameter server
 		ListObject model = ec.getListObject(getParam(PS_MODEL));
 		ParamServer ps = createPS(mode, aggFunc, updateType, workerNum, model, aggServiceEC);

 		// Create the local workers
 		List<LocalPSWorker> workers = IntStream.range(0, workerNum)
 			.mapToObj(i -> new LocalPSWorker(i, updFunc, freq, getEpochs(), getBatchSize(), workerECs.get(i), ps))
 			.collect(Collectors.toList());

 		// Do data partition
 		PSScheme scheme = getScheme();
 		partitionLocally(scheme, ec, workers);

 		if (DMLScript.STATISTICS)
 			Statistics.accPSSetupTime((long) tSetup.stop());

 		if (LOG.isDebugEnabled()) {
 			LOG.debug(String.format("\nConfiguration of paramserv func: "
 				+ "\nmode: %s \nworkerNum: %d \nupdate frequency: %s "
 				+ "\nstrategy: %s \ndata partitioner: %s",
 				mode, workerNum, freq, updateType, scheme));
 		}

 		try {
 			// Launch the worker threads and wait for completion
 			for (Future<Void> ret : es.invokeAll(workers))
 				ret.get(); //error handling
 			// Fetch the final model from ps
 			ec.setVariable(output.getName(), ps.getResult());
 		} catch (InterruptedException | ExecutionException e) {
 			throw new DMLRuntimeException("ParamservBuiltinCPInstruction: some error occurred: ", e);
 		} finally {
 			es.shutdownNow();
 		}
 	}

 	private LocalVariableMap createVarsMap(ExecutionContext ec) {
 		// Put the hyperparam into the variables table
 		LocalVariableMap varsMap = new LocalVariableMap();
 		ListObject hyperParams = getHyperParams(ec);
 		if (hyperParams != null) {
 			varsMap.put(PS_HYPER_PARAMS, hyperParams);
 		}
 		return varsMap;
 	}

 	private PSModeType getPSMode() {
 		if (!getParameterMap().containsKey(PS_MODE)) {
 			return DEFAULT_MODE;
 		}
 		PSModeType mode;
 		try {
 			mode = PSModeType.valueOf(getParam(PS_MODE));
 		} catch (IllegalArgumentException e) {
 			throw new DMLRuntimeException(String.format("Paramserv function: not support ps execution mode '%s'", getParam(PS_MODE)));
 		}
 		return mode;
 	}

 	private int getEpochs() {
 		int epochs = Integer.valueOf(getParam(PS_EPOCHS));
 		if (epochs <= 0) {
 			throw new DMLRuntimeException(String.format("Paramserv function: "
 				+ "The argument '%s' could not be less than or equal to 0.", PS_EPOCHS));
 		}
 		return epochs;
 	}

 	private static int getParLevel(int workerNum) {
 		return Math.max((int)Math.ceil((double)getRemainingCores()/workerNum), 1);
 	}

 	private PSUpdateType getUpdateType() {
 		if (!getParameterMap().containsKey(PS_UPDATE_TYPE)) {
 			return DEFAULT_TYPE;
 		}
 		PSUpdateType updType;
 		try {
 			updType = PSUpdateType.valueOf(getParam(PS_UPDATE_TYPE));
 		} catch (IllegalArgumentException e) {
 			throw new DMLRuntimeException(String.format("Paramserv function: not support update type '%s'.", getParam(PS_UPDATE_TYPE)));
 		}
 		if (updType == PSUpdateType.SSP)
 			throw new DMLRuntimeException("Paramserv function: Not support update type SSP.");
 		return updType;
 	}

 	private PSFrequency getFrequency() {
 		if (!getParameterMap().containsKey(PS_FREQUENCY)) {
 			return DEFAULT_UPDATE_FREQUENCY;
 		}
 		try {
 			return PSFrequency.valueOf(getParam(PS_FREQUENCY));
 		} catch (IllegalArgumentException e) {
 			throw new DMLRuntimeException(String.format("Paramserv function: "
 				+ "not support '%s' update frequency.", getParam(PS_FREQUENCY)));
 		}
 	}

 	private static int getRemainingCores() {
 		return InfrastructureAnalyzer.getLocalParallelism();
 	}

 	/**
 	 * Get the worker numbers according to the vcores
 	 *
 	 * @param mode execution mode
 	 * @return worker numbers
 	 */
 	private int getWorkerNum(PSModeType mode) {
 		switch (mode) {
 			case LOCAL:
 				return getParameterMap().containsKey(PS_PARALLELISM) ?
 					Integer.valueOf(getParam(PS_PARALLELISM)) : getRemainingCores();
 			case REMOTE_SPARK:
 				return getParameterMap().containsKey(PS_PARALLELISM) ?
 					Integer.valueOf(getParam(PS_PARALLELISM)) : SparkExecutionContext.getDefaultParallelism(true);
 			default:
 				throw new DMLRuntimeException("Unsupported parameter server: " + mode.name());
 		}
 	}

 	/**
 	 * Create a server which serves the local or remote workers
 	 *
 	 * @return parameter server
 	 */
 	private static ParamServer createPS(PSModeType mode, String aggFunc, PSUpdateType updateType, int workerNum, ListObject model, ExecutionContext ec) {
 		switch (mode) {
 			case LOCAL:
 			case REMOTE_SPARK:
 				return LocalParamServer.create(model, aggFunc, updateType, ec, workerNum);
 			default:
 				throw new DMLRuntimeException("Unsupported parameter server: "+mode.name());
 		}
 	}

 	private long getBatchSize() {
 		if (!getParameterMap().containsKey(PS_BATCH_SIZE)) {
 			return DEFAULT_BATCH_SIZE;
 		}
 		long batchSize = Integer.valueOf(getParam(PS_BATCH_SIZE));
 		if (batchSize <= 0) {
 			throw new DMLRuntimeException(String.format("Paramserv function: the number "
 				+ "of argument '%s' could not be less than or equal to 0.", PS_BATCH_SIZE));
 		}
 		return batchSize;
 	}

 	private ListObject getHyperParams(ExecutionContext ec) {
 		ListObject hyperparams = null;
 		if (getParameterMap().containsKey(PS_HYPER_PARAMS)) {
 			hyperparams = ec.getListObject(getParam(PS_HYPER_PARAMS));
 		}
 		return hyperparams;
 	}

 	private void partitionLocally(PSScheme scheme, ExecutionContext ec, List<LocalPSWorker> workers) {
 		MatrixObject features = ec.getMatrixObject(getParam(PS_FEATURES));
 		MatrixObject labels = ec.getMatrixObject(getParam(PS_LABELS));
 		DataPartitionLocalScheme.Result result = new LocalDataPartitioner(scheme).doPartitioning(workers.size(), features.acquireReadAndRelease(), labels.acquireReadAndRelease());
 		List<MatrixObject> pfs = result.pFeatures;
 		List<MatrixObject> pls = result.pLabels;
 		if (pfs.size() < workers.size()) {
 			if (LOG.isWarnEnabled()) {
 				LOG.warn(String.format("There is only %d batches of data but has %d workers. "
 					+ "Hence, reset the number of workers with %d.", pfs.size(), workers.size(), pfs.size()));
 			}
 			workers = workers.subList(0, pfs.size());
 		}
 		for (int i = 0; i < workers.size(); i++) {
 			workers.get(i).setFeatures(pfs.get(i));
 			workers.get(i).setLabels(pls.get(i));
 		}
 	}

 	private PSScheme getScheme() {
 		PSScheme scheme = DEFAULT_SCHEME;
 		if (getParameterMap().containsKey(PS_SCHEME)) {
 			try {
 				scheme = PSScheme.valueOf(getParam(PS_SCHEME));
 			} catch (IllegalArgumentException e) {
 				throw new DMLRuntimeException(String.format("Paramserv function: not support data partition scheme '%s'", getParam(PS_SCHEME)));
 			}
 		}
 		return scheme;
 	}

 }
	/*
	* 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.sysds.runtime.instructions.cp;

	import static org.apache.sysds.parser.Statement.PS_AGGREGATION_FUN;
	import static org.apache.sysds.parser.Statement.PS_BATCH_SIZE;
	import static org.apache.sysds.parser.Statement.PS_EPOCHS;
	import static org.apache.sysds.parser.Statement.PS_FEATURES;
	import static org.apache.sysds.parser.Statement.PS_FREQUENCY;
	import static org.apache.sysds.parser.Statement.PS_HYPER_PARAMS;
	import static org.apache.sysds.parser.Statement.PS_LABELS;
	import static org.apache.sysds.parser.Statement.PS_MODE;
	import static org.apache.sysds.parser.Statement.PS_MODEL;
	import static org.apache.sysds.parser.Statement.PS_PARALLELISM;
	import static org.apache.sysds.parser.Statement.PS_SCHEME;
	import static org.apache.sysds.parser.Statement.PS_UPDATE_FUN;
	import static org.apache.sysds.parser.Statement.PS_UPDATE_TYPE;

	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.LinkedHashMap;
	import java.util.List;
	import java.util.concurrent.ExecutionException;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.Executors;
	import java.util.concurrent.Future;
	import java.util.stream.Collectors;
	import java.util.stream.IntStream;

	import org.apache.commons.lang3.concurrent.BasicThreadFactory;
	import org.apache.commons.logging.Log;
	import org.apache.commons.logging.LogFactory;
	import org.apache.spark.network.server.TransportServer;
	import org.apache.spark.util.LongAccumulator;
	import org.apache.sysds.api.DMLScript;
	import org.apache.sysds.hops.recompile.Recompiler;
	import org.apache.sysds.lops.LopProperties;
	import org.apache.sysds.parser.Statement.PSFrequency;
	import org.apache.sysds.parser.Statement.PSModeType;
	import org.apache.sysds.parser.Statement.PSScheme;
	import org.apache.sysds.parser.Statement.PSUpdateType;
	import org.apache.sysds.runtime.DMLRuntimeException;
	import org.apache.sysds.runtime.controlprogram.LocalVariableMap;
	import org.apache.sysds.runtime.controlprogram.caching.MatrixObject;
	import org.apache.sysds.runtime.controlprogram.context.ExecutionContext;
	import org.apache.sysds.runtime.controlprogram.context.SparkExecutionContext;
	import org.apache.sysds.runtime.controlprogram.paramserv.LocalPSWorker;
	import org.apache.sysds.runtime.controlprogram.paramserv.LocalParamServer;
	import org.apache.sysds.runtime.controlprogram.paramserv.ParamServer;
	import org.apache.sysds.runtime.controlprogram.paramserv.ParamservUtils;
	import org.apache.sysds.runtime.controlprogram.paramserv.SparkPSBody;
	import org.apache.sysds.runtime.controlprogram.paramserv.SparkPSWorker;
	import org.apache.sysds.runtime.controlprogram.paramserv.dp.DataPartitionLocalScheme;
	import org.apache.sysds.runtime.controlprogram.paramserv.dp.LocalDataPartitioner;
	import org.apache.sysds.runtime.controlprogram.paramserv.rpc.PSRpcFactory;
	import org.apache.sysds.runtime.controlprogram.parfor.stat.InfrastructureAnalyzer;
	import org.apache.sysds.runtime.controlprogram.parfor.stat.Timing;
	import org.apache.sysds.runtime.matrix.operators.Operator;
	import org.apache.sysds.runtime.util.ProgramConverter;
	import org.apache.sysds.utils.Statistics;

	public class ParamservBuiltinCPInstruction extends ParameterizedBuiltinCPInstruction {
	private static final Log LOG = LogFactory.getLog(ParamservBuiltinCPInstruction.class.getName());

	private static final int DEFAULT_BATCH_SIZE = 64;
	private static final PSFrequency DEFAULT_UPDATE_FREQUENCY = PSFrequency.EPOCH;
	private static final PSScheme DEFAULT_SCHEME = PSScheme.DISJOINT_CONTIGUOUS;
	private static final PSModeType DEFAULT_MODE = PSModeType.LOCAL;
	private static final PSUpdateType DEFAULT_TYPE = PSUpdateType.ASP;

	public ParamservBuiltinCPInstruction(Operator op, LinkedHashMap<String, String> paramsMap, CPOperand out, String opcode, String istr) {
	super(op, paramsMap, out, opcode, istr);
	}

	@Override
	public void processInstruction(ExecutionContext ec) {
	PSModeType mode = getPSMode();
	switch (mode) {
	case LOCAL:
	runLocally(ec, mode);
	break;
	case REMOTE_SPARK:
	runOnSpark((SparkExecutionContext) ec, mode);
	break;
	default:
	throw new DMLRuntimeException(String.format("Paramserv func: not support mode %s", mode));
	}
	}

	@SuppressWarnings("resource")
	private void runOnSpark(SparkExecutionContext sec, PSModeType mode) {
	Timing tSetup = DMLScript.STATISTICS ? new Timing(true) : null;

	int workerNum = getWorkerNum(mode);
	String updFunc = getParam(PS_UPDATE_FUN);
	String aggFunc = getParam(PS_AGGREGATION_FUN);

	// Get the compiled execution context
	LocalVariableMap newVarsMap = createVarsMap(sec);
	// Level of par is 1 in spark backend because one worker will be launched per task
	ExecutionContext newEC = ParamservUtils.createExecutionContext(sec, newVarsMap, updFunc, aggFunc, 1);

	// Create the agg service's execution context
	ExecutionContext aggServiceEC = ParamservUtils.copyExecutionContext(newEC, 1).get(0);

	// Create the parameter server
	ListObject model = sec.getListObject(getParam(PS_MODEL));
	ParamServer ps = createPS(mode, aggFunc, getUpdateType(), workerNum, model, aggServiceEC);

	// Get driver host
	String host = sec.getSparkContext().getConf().get("spark.driver.host");

	// Create the netty server for ps
	TransportServer server = PSRpcFactory.createServer(sec.getSparkContext().getConf(),(LocalParamServer) ps, host); // Start the server

	// Force all the instructions to CP type
	Recompiler.recompileProgramBlockHierarchy2Forced(
	newEC.getProgram().getProgramBlocks(), 0, new HashSet<>(), LopProperties.ExecType.CP);

	// Serialize all the needed params for remote workers
	SparkPSBody body = new SparkPSBody(newEC);
	HashMap<String, byte[]> clsMap = new HashMap<>();
	String program = ProgramConverter.serializeSparkPSBody(body, clsMap);

	// Add the accumulators for statistics
	LongAccumulator aSetup = sec.getSparkContext().sc().longAccumulator("setup");
	LongAccumulator aWorker = sec.getSparkContext().sc().longAccumulator("workersNum");
	LongAccumulator aUpdate = sec.getSparkContext().sc().longAccumulator("modelUpdate");
	LongAccumulator aIndex = sec.getSparkContext().sc().longAccumulator("batchIndex");
	LongAccumulator aGrad = sec.getSparkContext().sc().longAccumulator("gradCompute");
	LongAccumulator aRPC = sec.getSparkContext().sc().longAccumulator("rpcRequest");
	LongAccumulator aBatch = sec.getSparkContext().sc().longAccumulator("numBatches");
	LongAccumulator aEpoch = sec.getSparkContext().sc().longAccumulator("numEpochs");

	// Create remote workers
	SparkPSWorker worker = new SparkPSWorker(getParam(PS_UPDATE_FUN), getParam(PS_AGGREGATION_FUN),
	getFrequency(), getEpochs(), getBatchSize(), program, clsMap, sec.getSparkContext().getConf(),
	server.getPort(), aSetup, aWorker, aUpdate, aIndex, aGrad, aRPC, aBatch, aEpoch);

	if (DMLScript.STATISTICS)
	Statistics.accPSSetupTime((long) tSetup.stop());

	MatrixObject features = sec.getMatrixObject(getParam(PS_FEATURES));
	MatrixObject labels = sec.getMatrixObject(getParam(PS_LABELS));
	try {
	ParamservUtils.doPartitionOnSpark(sec, features, labels, getScheme(), workerNum) // Do data partitioning
	.foreach(worker); // Run remote workers
	} catch (Exception e) {
	throw new DMLRuntimeException("Paramserv function failed: ", e);
	} finally {
	server.close(); // Stop the netty server
	}

	// Accumulate the statistics for remote workers
	if (DMLScript.STATISTICS) {
	Statistics.accPSSetupTime(aSetup.value());
	Statistics.incWorkerNumber(aWorker.value());
	Statistics.accPSLocalModelUpdateTime(aUpdate.value());
	Statistics.accPSBatchIndexingTime(aIndex.value());
	Statistics.accPSGradientComputeTime(aGrad.value());
	Statistics.accPSRpcRequestTime(aRPC.value());
	}

	// Fetch the final model from ps
	sec.setVariable(output.getName(), ps.getResult());
	}

	private void runLocally(ExecutionContext ec, PSModeType mode) {
	Timing tSetup = DMLScript.STATISTICS ? new Timing(true) : null;
	int workerNum = getWorkerNum(mode);
	BasicThreadFactory factory = new BasicThreadFactory.Builder()
	.namingPattern("workers-pool-thread-%d").build();
	ExecutorService es = Executors.newFixedThreadPool(workerNum, factory);
	String updFunc = getParam(PS_UPDATE_FUN);
	String aggFunc = getParam(PS_AGGREGATION_FUN);

	// Get the compiled execution context
	LocalVariableMap newVarsMap = createVarsMap(ec);
	ExecutionContext newEC = ParamservUtils.createExecutionContext(ec, newVarsMap, updFunc, aggFunc, getParLevel(workerNum));

	// Create workers' execution context
	List<ExecutionContext> workerECs = ParamservUtils.copyExecutionContext(newEC, workerNum);

	// Create the agg service's execution context
	ExecutionContext aggServiceEC = ParamservUtils.copyExecutionContext(newEC, 1).get(0);

	PSFrequency freq = getFrequency();
	PSUpdateType updateType = getUpdateType();

	// Create the parameter server
	ListObject model = ec.getListObject(getParam(PS_MODEL));
	ParamServer ps = createPS(mode, aggFunc, updateType, workerNum, model, aggServiceEC);

	// Create the local workers
	List<LocalPSWorker> workers = IntStream.range(0, workerNum)
	.mapToObj(i -> new LocalPSWorker(i, updFunc, freq, getEpochs(), getBatchSize(), workerECs.get(i), ps))
	.collect(Collectors.toList());

	// Do data partition
	PSScheme scheme = getScheme();
	partitionLocally(scheme, ec, workers);

	if (DMLScript.STATISTICS)
	Statistics.accPSSetupTime((long) tSetup.stop());

	if (LOG.isDebugEnabled()) {
	LOG.debug(String.format("\nConfiguration of paramserv func: "
	+ "\nmode: %s \nworkerNum: %d \nupdate frequency: %s "
	+ "\nstrategy: %s \ndata partitioner: %s",
	mode, workerNum, freq, updateType, scheme));
	}

	try {
	// Launch the worker threads and wait for completion
	for (Future<Void> ret : es.invokeAll(workers))
	ret.get(); //error handling
	// Fetch the final model from ps
	ec.setVariable(output.getName(), ps.getResult());
	} catch (InterruptedException \| ExecutionException e) {
	throw new DMLRuntimeException("ParamservBuiltinCPInstruction: some error occurred: ", e);
	} finally {
	es.shutdownNow();
	}
	}

	private LocalVariableMap createVarsMap(ExecutionContext ec) {
	// Put the hyperparam into the variables table
	LocalVariableMap varsMap = new LocalVariableMap();
	ListObject hyperParams = getHyperParams(ec);
	if (hyperParams != null) {
	varsMap.put(PS_HYPER_PARAMS, hyperParams);
	}
	return varsMap;
	}

	private PSModeType getPSMode() {
	if (!getParameterMap().containsKey(PS_MODE)) {
	return DEFAULT_MODE;
	}
	PSModeType mode;
	try {
	mode = PSModeType.valueOf(getParam(PS_MODE));
	} catch (IllegalArgumentException e) {
	throw new DMLRuntimeException(String.format("Paramserv function: not support ps execution mode '%s'", getParam(PS_MODE)));
	}
	return mode;
	}

	private int getEpochs() {
	int epochs = Integer.valueOf(getParam(PS_EPOCHS));
	if (epochs <= 0) {
	throw new DMLRuntimeException(String.format("Paramserv function: "
	+ "The argument '%s' could not be less than or equal to 0.", PS_EPOCHS));
	}
	return epochs;
	}

	private static int getParLevel(int workerNum) {
	return Math.max((int)Math.ceil((double)getRemainingCores()/workerNum), 1);
	}

	private PSUpdateType getUpdateType() {
	if (!getParameterMap().containsKey(PS_UPDATE_TYPE)) {
	return DEFAULT_TYPE;
	}
	PSUpdateType updType;
	try {
	updType = PSUpdateType.valueOf(getParam(PS_UPDATE_TYPE));
	} catch (IllegalArgumentException e) {
	throw new DMLRuntimeException(String.format("Paramserv function: not support update type '%s'.", getParam(PS_UPDATE_TYPE)));
	}
	if (updType == PSUpdateType.SSP)
	throw new DMLRuntimeException("Paramserv function: Not support update type SSP.");
	return updType;
	}

	private PSFrequency getFrequency() {
	if (!getParameterMap().containsKey(PS_FREQUENCY)) {
	return DEFAULT_UPDATE_FREQUENCY;
	}
	try {
	return PSFrequency.valueOf(getParam(PS_FREQUENCY));
	} catch (IllegalArgumentException e) {
	throw new DMLRuntimeException(String.format("Paramserv function: "
	+ "not support '%s' update frequency.", getParam(PS_FREQUENCY)));
	}
	}

	private static int getRemainingCores() {
	return InfrastructureAnalyzer.getLocalParallelism();
	}

	/**
	* Get the worker numbers according to the vcores
	*
	* @param mode execution mode
	* @return worker numbers
	*/
	private int getWorkerNum(PSModeType mode) {
	switch (mode) {
	case LOCAL:
	return getParameterMap().containsKey(PS_PARALLELISM) ?
	Integer.valueOf(getParam(PS_PARALLELISM)) : getRemainingCores();
	case REMOTE_SPARK:
	return getParameterMap().containsKey(PS_PARALLELISM) ?
	Integer.valueOf(getParam(PS_PARALLELISM)) : SparkExecutionContext.getDefaultParallelism(true);
	default:
	throw new DMLRuntimeException("Unsupported parameter server: " + mode.name());
	}
	}

	/**
	* Create a server which serves the local or remote workers
	*
	* @return parameter server
	*/
	private static ParamServer createPS(PSModeType mode, String aggFunc, PSUpdateType updateType, int workerNum, ListObject model, ExecutionContext ec) {
	switch (mode) {
	case LOCAL:
	case REMOTE_SPARK:
	return LocalParamServer.create(model, aggFunc, updateType, ec, workerNum);
	default:
	throw new DMLRuntimeException("Unsupported parameter server: "+mode.name());
	}
	}

	private long getBatchSize() {
	if (!getParameterMap().containsKey(PS_BATCH_SIZE)) {
	return DEFAULT_BATCH_SIZE;
	}
	long batchSize = Integer.valueOf(getParam(PS_BATCH_SIZE));
	if (batchSize <= 0) {
	throw new DMLRuntimeException(String.format("Paramserv function: the number "
	+ "of argument '%s' could not be less than or equal to 0.", PS_BATCH_SIZE));
	}
	return batchSize;
	}

	private ListObject getHyperParams(ExecutionContext ec) {
	ListObject hyperparams = null;
	if (getParameterMap().containsKey(PS_HYPER_PARAMS)) {
	hyperparams = ec.getListObject(getParam(PS_HYPER_PARAMS));
	}
	return hyperparams;
	}

	private void partitionLocally(PSScheme scheme, ExecutionContext ec, List<LocalPSWorker> workers) {
	MatrixObject features = ec.getMatrixObject(getParam(PS_FEATURES));
	MatrixObject labels = ec.getMatrixObject(getParam(PS_LABELS));
	DataPartitionLocalScheme.Result result = new LocalDataPartitioner(scheme).doPartitioning(workers.size(), features.acquireReadAndRelease(), labels.acquireReadAndRelease());
	List<MatrixObject> pfs = result.pFeatures;
	List<MatrixObject> pls = result.pLabels;
	if (pfs.size() < workers.size()) {
	if (LOG.isWarnEnabled()) {
	LOG.warn(String.format("There is only %d batches of data but has %d workers. "
	+ "Hence, reset the number of workers with %d.", pfs.size(), workers.size(), pfs.size()));
	}
	workers = workers.subList(0, pfs.size());
	}
	for (int i = 0; i < workers.size(); i++) {
	workers.get(i).setFeatures(pfs.get(i));
	workers.get(i).setLabels(pls.get(i));
	}
	}

	private PSScheme getScheme() {
	PSScheme scheme = DEFAULT_SCHEME;
	if (getParameterMap().containsKey(PS_SCHEME)) {
	try {
	scheme = PSScheme.valueOf(getParam(PS_SCHEME));
	} catch (IllegalArgumentException e) {
	throw new DMLRuntimeException(String.format("Paramserv function: not support data partition scheme '%s'", getParam(PS_SCHEME)));
	}
	}
	return scheme;
	}

	}