# pylint: disable=invalid-name """Command line options of job submission script.""" import os import argparse def get_cache_file_set(args): """Get the list of files to be cached. Parameters ---------- args: ArgumentParser.Argument The arguments returned by the parser. Returns ------- cache_file_set: set of str The set of files to be cached to local execution environment. command: list of str The commands that get rewritten after the file cache is used. """ fset = set() cmds = [] if args.auto_file_cache: for i in range(len(args.command)): fname = args.command[i] if os.path.exists(fname): fset.add(fname) cmds.append('./' + fname.split('/')[-1]) else: cmds.append(fname) for fname in args.files: if os.path.exists(fname): fset.add(fname) return fset, cmds def get_memory_mb(mem_str): """Get the memory in MB from memory string. mem_str: str String representation of memory requirement. Returns ------- mem_mb: int Memory requirement in MB. """ mem_str = mem_str.lower() if mem_str.endswith('g'): return int(float(mem_str[:-1]) * 1024) elif mem_str.endswith('m'): return int(float(mem_str[:-1])) else: msg = 'Invalid memory specification %s, need to be a number follows g or m' % mem_str raise RuntimeError(msg) def get_opts(args=None): """Get options to launch the job. Returns ------- args: ArgumentParser.Argument The arguments returned by the parser. cache_file_set: set of str The set of files to be cached to local execution environment. """ parser = argparse.ArgumentParser(description='DMLC job submission.') parser.add_argument('--cluster', type=str, choices=['yarn', 'slurm', 'mpi', 'sge', 'local', 'ssh', 'mesos', 'kubernetes'], help=('Cluster type of this submission,' + 'default to env variable ${DMLC_SUBMIT_CLUSTER}.')) parser.add_argument('--num-workers', required=True, type=int, help='Number of worker proccess to be launched.') parser.add_argument('--worker-cores', default=1, type=int, help='Number of cores to be allocated for each worker process.') parser.add_argument('--worker-memory', default='1g', type=str, help=('Memory need to be allocated for each worker,' + ' need to ends with g or m')) parser.add_argument('--num-servers', default=0, type=int, help='Number of server process to be launched. Only used in PS jobs.') parser.add_argument('--server-cores', default=1, type=int, help=('Number of cores to be allocated for each server process.' + 'Only used in PS jobs.')) parser.add_argument('--server-memory', default='1g', type=str, help=('Memory need to be allocated for each server, ' + 'need to ends with g or m.')) parser.add_argument('--jobname', default=None, type=str, help='Name of the job.') parser.add_argument('--queue', default='default', type=str, help='The submission queue the job should goes to.') parser.add_argument('--log-level', default='INFO', type=str, choices=['INFO', 'DEBUG'], help='Logging level of the logger.') parser.add_argument('--log-file', default=None, type=str, help=('Output log to the specific log file, ' + 'the log is still printed on stderr.')) parser.add_argument('--host-ip', default=None, type=str, help=('Host IP addressed, this is only needed ' + 'if the host IP cannot be automatically guessed.')) parser.add_argument('--hdfs-tempdir', default='/tmp', type=str, help=('Temporary directory in HDFS, ' + ' only needed in YARN mode.')) parser.add_argument('--host-file', default=None, type=str, help=('The file contains the list of hostnames, needed for MPI and ssh.')) parser.add_argument('--sge-log-dir', default=None, type=str, help=('Log directory of SGD jobs, only needed in SGE mode.')) parser.add_argument( '--auto-file-cache', default=True, type=bool, help=('Automatically cache files appeared in the command line' + 'to local executor folder.' + ' This will also cause rewritten of all the file names in the command,' + ' e.g. `../../kmeans ../kmeans.conf` will be rewritten to `./kmeans kmeans.conf`')) parser.add_argument('--files', default=[], action='append', help=('The cached file list which will be copied to local environment,' + ' You may need this option to cache additional files.' + ' You --auto-file-cache is off')) parser.add_argument('--archives', default=[], action='append', help=('Same as cached files,' + ' but corresponds to archieve files that will be unziped locally,' + ' You can use this option to ship python libraries.' + ' Only valid in yarn jobs.')) parser.add_argument('--env', action='append', default=[], help='Client and ApplicationMaster environment variables.') parser.add_argument('--yarn-app-classpath', type=str, help=('Explicit YARN ApplicationMaster classpath.' + 'Can be used to override defaults.')) parser.add_argument('--yarn-app-dir', type=str, default=os.path.join(os.path.dirname(__file__), os.pardir, 'yarn'), help=('Directory to YARN appmaster. Only used in YARN mode.')) parser.add_argument('--mesos-master', type=str, help=('Mesos master, default to ${MESOS_MASTER}')), parser.add_argument('--ship-libcxx', default=None, type=str, help=('The path to the customized gcc lib folder.' + 'You can use this option to ship customized libstdc++' + ' library to the workers.')) parser.add_argument('--sync-dst-dir', type=str, help = 'if specificed, it will sync the current \ directory into remote machines\'s SYNC_DST_DIR') parser.add_argument('command', nargs='+', help='Command to be launched') parser.add_argument('--slurm-worker-nodes', default=None, type=int, help=('Number of nodes on which workers are run. Used only in SLURM mode.' + 'If not explicitly set, it defaults to number of workers.')) parser.add_argument('--slurm-server-nodes', default=None, type=int, help=('Number of nodes on which parameter servers are run. Used only in SLURM mode.' + 'If not explicitly set, it defaults to number of parameter servers.')) parser.add_argument('--kube-namespace', default="default", type=str, help=('A namespace in whitch all tasks are run. Used only in Kubernetes mode.' + 'If not explicitly set, it defaults to default.')) parser.add_argument('--kube-worker-image', default="mxnet/python", type=str, help=('Container image of workers. Used only in Kubernetes mode.' + 'If not explicitly set, it defaults to mxnet/python.')) parser.add_argument('--kube-server-image', default="mxnet/python", type=str, help=('Container image of servers. Used only in Kubernetes mode.' + 'If not explicitly set, it defaults to mxnet/python.')) parser.add_argument('--kube-worker-template', default=None, type=str, help=('Manifest template for workers. Used only in Kubernetes mode.' + 'Can be used to override defaults.')) parser.add_argument('--kube-server-template', default=None, type=str, help=('Manifest template for servers. Used only in Kubernetes mode.' + 'Can be used to override defaults.')) (args, unknown) = parser.parse_known_args(args) args.command += unknown if args.cluster is None: args.cluster = os.getenv('DMLC_SUBMIT_CLUSTER', None) if args.cluster is None: raise RuntimeError('--cluster is not specified, ' + 'you can also specify the default behavior via ' + 'environment variable DMLC_SUBMIT_CLUSTER') args.worker_memory_mb = get_memory_mb(args.worker_memory) args.server_memory_mb = get_memory_mb(args.server_memory) return args