org.deeplearning4j.aws.ec2.provision.ClusterSetup Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of deeplearning4j-aws Show documentation
There is a newer version: 0.9.1
/*
 *
 *  * Copyright 2015 Skymind,Inc.
 *  *
 *  *    Licensed 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.deeplearning4j.aws.ec2.provision;

import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;


import com.amazonaws.regions.Regions;
import org.deeplearning4j.aws.ec2.Ec2BoxCreator;
import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Sets up a DL4J cluster
 * @author Adam Gibson
 *
 */
public class ClusterSetup {

	@Option(name = "-w",usage = "Number of workers")
	private int numWorkers = 1;
	@Option(name = "-ami",usage = "Amazon machine image: default, amazon linux (only works with RHEL right now")
	private String ami = "ami-fb8e9292";
	@Option(name = "-s",usage = "size of instance: default m1.medium")
	private String size = "m3.xlarge";
	@Option(name = "-sg",usage = "security group, this needs to be applyTransformToDestination")
	private String securityGroupName;
	@Option(name = "-kp",usage = "key pair name, also needs to be applyTransformToDestination.")
	private String keyPairName;
	@Option(name = "-kpath",usage = "path to private key - needs to be applyTransformToDestination, this is used to login to amazon.")
	private String pathToPrivateKey;
    @Option(name = "-wscript", usage = "path to worker script to run, this will allow customization of dependencies")
	private String workerSetupScriptPath;
	@Option(name = "-mscript", usage = "path to master script to run this will allow customization of the dependencies")
	private String masterSetupScriptPath;
    @Option(name = "-region",usage = "specify a region")
    private String region = Regions.US_EAST_1.getName();

    private ExecutorService as;

	private static final Logger log = LoggerFactory.getLogger(ClusterSetup.class);


	public ClusterSetup(String[] args) {
		CmdLineParser parser = new CmdLineParser(this);
		try {
			parser.parseArgument(args);
		} catch (CmdLineException e) {
			parser.printUsage(System.err);
			log.error("Unable to parse args",e);
		}


	}

	public void exec() {
		//master + workers
		Ec2BoxCreator boxCreator = new Ec2BoxCreator(ami,numWorkers,size,securityGroupName,keyPairName);
        boxCreator.setRegion(Regions.fromName(region));
		boxCreator.create();
		boxCreator.blockTillAllRunning();
		List hosts = boxCreator.getHosts();
		//provisionMaster(hosts.getFromOrigin(0));
		provisionWorkers(hosts);


	}






	private void provisionWorkers(List workers) {
		as = Executors.newFixedThreadPool(Runtime.getRuntime().availableProcessors());
		for(final String workerHost : workers) {
			try {
              as.submit(new Runnable() {
				  @Override
				  public void run() {
					  HostProvisioner uploader = new HostProvisioner(workerHost, "ec2-user");
					  try {
						  uploader.addKeyFile(pathToPrivateKey);
						  //uploader.runRemoteCommand("sudo hostname " + workerHost);
						  uploader.uploadAndRun(workerSetupScriptPath, "");
					  } catch (Exception e) {
						  e.printStackTrace();
					  }

				  }
			  });

			}catch(Exception e) {
				log.error("Error ",e);
			}
		}
	}


	/**
	 * @param args
	 */
	public static void main(String[] args) {
		new ClusterSetup(args).exec();
	}

}