Repository: sql-machine-learning/elasticdl

Issue:
More strict commandline argument check



Discussion:

Resolving patch:
diff --git a/elasticdl/python/elasticdl/client/client.py b/elasticdl/python/elasticdl/client/client.py
--- a/elasticdl/python/elasticdl/client/client.py
+++ b/elasticdl/python/elasticdl/client/client.py
@@ -14,9 +14,9 @@
 def _m_file_in_docker(model_file):
     return "/model/" + os.path.basename(model_file)
 
+
 def _build_docker_image(
-    m_file, image_name, image_base="elasticdl:dev",
-    repository=None
+    m_file, image_name, image_base="elasticdl:dev", repository=None
 ):
     DOCKER_TEMPLATE = """
 FROM {}
@@ -24,7 +24,11 @@ def _build_docker_image(
 """
 
     with tempfile.NamedTemporaryFile(mode="w+", delete=False) as df:
-        df.write(DOCKER_TEMPLATE.format(image_base, m_file, _m_file_in_docker(m_file)))
+        df.write(
+            DOCKER_TEMPLATE.format(
+                image_base, m_file, _m_file_in_docker(m_file)
+            )
+        )
 
     client = docker.APIClient(base_url="unix://var/run/docker.sock")
     print("===== Building Docker Image =====")
@@ -32,7 +36,9 @@ def _build_docker_image(
         dockerfile=df.name, path=".", rm=True, tag=image_name, decode=True
     ):
         if "error" in line:
-            raise RuntimeError("Docker image build failure: %s" % line["error"])
+            raise RuntimeError(
+                "Docker image build failure: %s" % line["error"]
+            )
         text = line.get("stream", None)
         if text:
             sys.stdout.write(text)
@@ -42,6 +48,7 @@ def _build_docker_image(
         for line in client.push(image_name, stream=True, decode=True):
             print(line)
 
+
 def _gen_master_def(image_name, model_file, job_name, args, argv):
     master_yaml = """
 apiVersion: v1
@@ -80,33 +87,56 @@ def _gen_master_def(image_name, model_file, job_name, args, argv):
         fieldRef:
           fieldPath: status.podIP
   restartPolicy: Never
-""" .format(m_file=_m_file_in_docker(model_file), image_name=image_name, job_name=job_name,
-        master_cpu_limit=args.master_cpu_limit, master_cpu_request=args.master_cpu_request, 
-        master_memory_limit=args.master_memory_limit, master_memory_request=args.master_memory_request,
-        worker_cpu_limit=args.worker_cpu_limit, worker_cpu_request=args.worker_cpu_request,
-        worker_memory_limit=args.worker_memory_limit, worker_memory_request=args.worker_memory_request,
-        image_pull_policy=args.image_pull_policy)
+""".format(
+        m_file=_m_file_in_docker(model_file),
+        image_name=image_name,
+        job_name=job_name,
+        master_cpu_limit=args.master_cpu_limit,
+        master_cpu_request=args.master_cpu_request,
+        master_memory_limit=args.master_memory_limit,
+        master_memory_request=args.master_memory_request,
+        worker_cpu_limit=args.worker_cpu_limit,
+        worker_cpu_request=args.worker_cpu_request,
+        worker_memory_limit=args.worker_memory_limit,
+        worker_memory_request=args.worker_memory_request,
+        image_pull_policy=args.image_pull_policy,
+    )
 
     master_def = yaml.safe_load(master_yaml)
 
     # Build master arguments
-    master_def['spec']['containers'][0]['args'].extend(argv)
+    master_def["spec"]["containers"][0]["args"].extend(argv)
 
     if args.master_pod_priority is not None:
-        master_def['spec']['priorityClassName'] = args.master_pod_priority
+        master_def["spec"]["priorityClassName"] = args.master_pod_priority
 
     if args.volume_name is not None and args.mount_path is not None:
-        persistent_volume_claim = {'claimName': 'fileserver-claim', 'readOnly': False}
-        volume = {'name': args.volume_name, 'persistentVolumeClaim': persistent_volume_claim}
-        master_def['spec']['volumes'] = [volume]
-        master_def['spec']['containers'][0]['volumeMounts'] = [
-            {'mountPath': args.mount_path, 'name': args.volume_name}]
-        master_def['spec']['containers'][0]['args'].extend(['--mount_path', 
-            args.mount_path, '--volume_name', args.volume_name,
-            '--image_pull_policy', args.image_pull_policy])
+        persistent_volume_claim = {
+            "claimName": "fileserver-claim",
+            "readOnly": False,
+        }
+        volume = {
+            "name": args.volume_name,
+            "persistentVolumeClaim": persistent_volume_claim,
+        }
+        master_def["spec"]["volumes"] = [volume]
+        master_def["spec"]["containers"][0]["volumeMounts"] = [
+            {"mountPath": args.mount_path, "name": args.volume_name}
+        ]
+        master_def["spec"]["containers"][0]["args"].extend(
+            [
+                "--mount_path",
+                args.mount_path,
+                "--volume_name",
+                args.volume_name,
+                "--image_pull_policy",
+                args.image_pull_policy,
+            ]
+        )
 
     return master_def
 
+
 def _submit(image_name, model_file, job_name, args, argv):
     master_def = _gen_master_def(image_name, model_file, job_name, args, argv)
     config.load_kube_config()
@@ -114,74 +144,107 @@ def _submit(image_name, model_file, job_name, args, argv):
     resp = api.create_namespaced_pod(body=master_def, namespace="default")
     print("Master launched. status='%s'" % str(resp.status))
 
-def _validate_params(args):
-    cpu_regexp = re.compile('([1-9]{1})([0-9]*)m$')
-    memory_regexp = re.compile('([1-9]{1})([0-9]*)(E|P|T|G|M|K|Ei|Pi|Ti|Gi|Mi|Ki)$') 
-    if cpu_regexp.match(args.master_cpu_request) is None:
-        raise ValueError("invalid master cpu request: " + args.master_cpu_request)
-    if cpu_regexp.match(args.master_cpu_limit) is None:
-        raise ValueError("invalid master cpu limit: " + args.master_cpu_limit)
-    if memory_regexp.match(args.master_memory_request) is None:
-        raise ValueError("invalid master memory request: " + args.master_memory_request)
-    if memory_regexp.match(args.master_memory_limit) is None:
-        raise ValueError("invalid master memory limit: " + args.master_memory_limit)
-
-    if cpu_regexp.match(args.worker_cpu_request) is None:
-        raise ValueError("invalid worker cpu request: " + args.worker_cpu_request)
-    if cpu_regexp.match(args.worker_cpu_limit) is None:
-        raise ValueError("invalid worker cpu limit: " + args.worker_cpu_limit)
-    if memory_regexp.match(args.worker_memory_request) is None:
-        raise ValueError("invalid worker memory request: " + args.worker_memory_request)
-    if memory_regexp.match(args.worker_memory_limit) is None:
-        raise ValueError("invalid worker memory limit: " + args.worker_memory_limit)
+
+def _valid_cpu_spec(arg):
+    regexp = re.compile("([1-9]{1})([0-9]*)m$")
+    if not regexp.match(arg):
+        raise ValueError("invalid cpu request spec: " + arg)
+
+
+def _valid_mem_spec(arg):
+    regexp = re.compile("([1-9]{1})([0-9]*)(E|P|T|G|M|K|Ei|Pi|Ti|Gi|Mi|Ki)$")
+    if not regexp.match(arg):
+        raise ValueError("invalid memory request spec: " + arg)
+
 
 def main():
     parser = argparse.ArgumentParser(description="ElasticDL Client")
     # Rewrite model_file argument and pass all other arguments to master.
-    parser.add_argument("--model_file", help="Path to Model file", required=True)
-    parser.add_argument("--image_base", help="Base image containing elasticdl runtime environment.", required=True)
-    parser.add_argument("--repository", help="The repository to push docker image to.")
+    parser.add_argument(
+        "--model_file", help="Path to Model file", required=True
+    )
+    parser.add_argument(
+        "--image_base",
+        help="Base image containing elasticdl runtime environment.",
+        required=True,
+    )
+    parser.add_argument(
+        "--repository", help="The repository to push docker image to."
+    )
     parser.add_argument("--job_name", help="ElasticDL job name", required=True)
-    parser.add_argument("--master_cpu_request", 
+    parser.add_argument(
+        "--master_cpu_request",
         default="100m",
-        help="the minimal cpu required by master in training")
-    parser.add_argument("--master_cpu_limit", 
+        type=_valid_cpu_spec,
+        help="the minimal cpu required by master in training",
+    )
+    parser.add_argument(
+        "--master_cpu_limit",
         default="100m",
-        help="the maximal cpu used by master in training")
-    parser.add_argument("--master_memory_request", 
+        type=_valid_cpu_spec,
+        help="the maximal cpu used by master in training",
+    )
+    parser.add_argument(
+        "--master_memory_request",
         default="1024Mi",
-        help="the minimal memory required by master in training")
-    parser.add_argument("--master_memory_limit", 
+        type=_valid_mem_spec,
+        help="the minimal memory required by master in training",
+    )
+    parser.add_argument(
+        "--master_memory_limit",
         default="1024Mi",
-        help="the maximal memory used by master in training")
-    parser.add_argument("--worker_cpu_request", 
+        type=_valid_mem_spec,
+        help="the maximal memory used by master in training",
+    )
+    parser.add_argument(
+        "--worker_cpu_request",
         default="1000m",
-        help="the minimal cpu required by worker in training")
-    parser.add_argument("--worker_cpu_limit", 
+        type=_valid_cpu_spec,
+        help="the minimal cpu required by worker in training",
+    )
+    parser.add_argument(
+        "--worker_cpu_limit",
         default="1000m",
-        help="the maximal cpu used by worker in training")
-    parser.add_argument("--worker_memory_request", 
+        type=_valid_cpu_spec,
+        help="the maximal cpu used by worker in training",
+    )
+    parser.add_argument(
+        "--worker_memory_request",
         default="4096Mi",
-        help="the minimal memory required by worker in training")
-    parser.add_argument("--worker_memory_limit", 
+        type=_valid_mem_spec,
+        help="the minimal memory required by worker in training",
+    )
+    parser.add_argument(
+        "--worker_memory_limit",
         default="4096Mi",
-        help="the maximal memory used by worker in training")
-    parser.add_argument("--master_pod_priority",
-        help="the requested priority of master pod")
-    parser.add_argument("--volume_name",
-        help="the volume name of network filesytem")
-    parser.add_argument("--mount_path",
-        help="the mount path in the docker container")
-    parser.add_argument("--image_pull_policy",
+        type=_valid_mem_spec,
+        help="the maximal memory used by worker in training",
+    )
+    parser.add_argument(
+        "--master_pod_priority", help="the requested priority of master pod"
+    )
+    parser.add_argument(
+        "--volume_name", help="the volume name of network filesytem"
+    )
+    parser.add_argument(
+        "--mount_path", help="the mount path in the docker container"
+    )
+    parser.add_argument(
+        "--image_pull_policy",
         default="Always",
-        help="the image pull policy of master and worker")
+        help="the image pull policy of master and worker",
+    )
     args, argv = parser.parse_known_args()
     _validate_params(args)
 
     job_name = args.job_name
-    image_name = args.image_base + '_' + job_name 
-    _build_docker_image(args.model_file, image_name, image_base=args.image_base,
-        repository=args.repository)
+    image_name = args.image_base + "_" + job_name
+    _build_docker_image(
+        args.model_file,
+        image_name,
+        image_base=args.image_base,
+        repository=args.repository,
+    )
     _submit(image_name, args.model_file, job_name, args, argv)
 
 
diff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py
--- a/elasticdl/python/elasticdl/master/main.py
+++ b/elasticdl/python/elasticdl/master/main.py
@@ -15,7 +15,10 @@
 from elasticdl.python.elasticdl.master.servicer import MasterServicer
 from elasticdl.python.elasticdl.master.task_queue import _TaskQueue
 from elasticdl.python.elasticdl.master.k8s_worker_manager import WorkerManager
-from elasticdl.python.elasticdl.common.model_helper import load_user_model, build_model
+from elasticdl.python.elasticdl.common.model_helper import (
+    load_user_model,
+    build_model,
+)
 
 
 def _make_task_queue(data_dir, record_per_task, num_epoch):
@@ -27,6 +30,22 @@ def _make_task_queue(data_dir, record_per_task, num_epoch):
     return _TaskQueue(f_records, record_per_task, num_epoch)
 
 
+def _pos_int(arg):
+    res = int(arg)
+    if res <= 0:
+        raise ValueError("Positive integer argument required. Got %s" % res)
+    return res
+
+
+def _non_neg_int(arg):
+    res = int(arg)
+    if res < 0:
+        raise ValueError(
+            "Non-negative integer argument required. Get %s" % res
+        )
+    return res
+
+
 def _parse_args():
     parser = argparse.ArgumentParser(description="ElasticDL Master")
     parser.add_argument(
@@ -39,23 +58,23 @@ def _parse_args():
         help="Training data directory. Files should be in RecordIO format",
         required=True,
     )
-    parser.add_argument("--record_per_task", type=int, required=True)
-    parser.add_argument("--num_epoch", type=int, required=True)
+    parser.add_argument("--record_per_task", type=_pos_int, required=True)
+    parser.add_argument("--num_epoch", type=_pos_int, required=True)
     parser.add_argument(
         "--grads_to_wait",
-        type=int,
+        type=_pos_int,
         help="Number of gradients to wait before updating model",
         required=True,
     )
     parser.add_argument(
         "--minibatch_size",
-        type=int,
+        type=_pos_int,
         help="Minibatch size used by workers to compute gradients",
         required=True,
     )
     parser.add_argument(
         "--num_worker",
-        type=int,
+        type=_pos_int,
         help="the number of workers used in training",
         default=0,
     )
@@ -66,13 +85,13 @@ def _parse_args():
     )
     parser.add_argument(
         "--checkpoint_steps",
-        type=int,
+        type=_non_neg_int,
         help="Save checkpoint every this many steps. If 0, no checkpoints to save.",
         default=0,
     )
     parser.add_argument(
         "--keep_checkpoint_max",
-        type=int,
+        type=_non_neg_int,
         help="The maximum number of recent checkpoint files to keep. If 0, keep all.",
         default=3,
     )
@@ -97,10 +116,13 @@ def _parse_args():
         default="4096Mi",
     )
     parser.add_argument(
-        "--worker_pod_priority",
-        help="the requested priority of worker pod")
+        "--worker_pod_priority", help="the requested priority of worker pod"
+    )
     parser.add_argument(
-        "--worker_image", help="docker image for worker", default=None
+        "--worker_image",
+        help="docker image for worker",
+        default=None,
+        required=True,
     )
     parser.add_argument("--job_name", help="job name", required=True)
     parser.add_argument(
@@ -109,10 +131,12 @@ def _parse_args():
         choices=["tf_example", "bytes"],
         help="Type of codec(tf_example or bytes)",
     )
-    parser.add_argument("--volume_name",
-        help="the volume name of network filesytem")
-    parser.add_argument("--mount_path",
-        help="the mount path in the docker container")
+    parser.add_argument(
+        "--volume_name", help="the volume name of network filesytem"
+    )
+    parser.add_argument(
+        "--mount_path", help="the mount path in the docker container"
+    )
     parser.add_argument(
         "--log_level",
         choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
@@ -120,8 +144,10 @@ def _parse_args():
         default="WARNING",
         help="the logging level",
     )
-    parser.add_argument("--image_pull_policy",
-        help="the image pull policy of master and worker")
+    parser.add_argument(
+        "--image_pull_policy",
+        help="the image pull policy of master and worker",
+    )
     return parser.parse_args()
 
 
@@ -133,8 +159,8 @@ def main():
 
     # Initialize logger
     logging.basicConfig(
-        format='%(asctime)s %(name)s %(levelname)-8s '
-        '[%(filename)s:%(lineno)d] %(message)s',
+        format="%(asctime)s %(name)s %(levelname)-8s "
+        "[%(filename)s:%(lineno)d] %(message)s"
     )
     # Set level for ROOT logger.
     logging.getLogger().setLevel(args.log_level)
@@ -179,7 +205,7 @@ def main():
             "--codec_type",
             args.codec_type,
             "--log_level",
-            args.log_level
+            args.log_level,
         ]
 
         worker_manager = WorkerManager(


Test patch: