Repository: sql-machine-learning/elasticdl

Issue:
[PS-1] create `ps/main.py` for PS main process and parse required arguments
Parse arguments required by PS, including to add a new argument `num_ps_pods`, with default value as 0.

num_ps_pods = 0 for ElasticlDL to use legacy PS (PS with master + Redis). We will change the default after we have implemented the new PS.


Discussion:

Resolving patch:
diff --git a/elasticdl/python/common/args.py b/elasticdl/python/common/args.py
--- a/elasticdl/python/common/args.py
+++ b/elasticdl/python/common/args.py
@@ -125,6 +125,9 @@ def add_common_params(parser):
     parser.add_argument(
         "--num_workers", type=int, help="Number of workers", default=0
     )
+    parser.add_argument(
+        "--num_ps_pods", type=int, help="Number of PS pods", default=0
+    )
     parser.add_argument(
         "--worker_resource_request",
         default="cpu=1,memory=4096Mi",
@@ -518,6 +521,19 @@ def parse_master_args(master_args=None):
     return args
 
 
+def parse_ps_args(ps_args=None):
+    parser = argparse.ArgumentParser(description="ElasticDL PS")
+    add_common_params(parser)
+    add_train_params(parser)
+    # TODO: add PS replica address for RPC stub creation
+
+    args, unknown_args = parser.parse_known_args(args=ps_args)
+    print_args(args, groups=ALL_ARGS_GROUPS)
+    if unknown_args:
+        logger.warning("Unknown arguments: %s", unknown_args)
+    return args
+
+
 def parse_worker_args(worker_args=None):
     parser = argparse.ArgumentParser(description="ElasticDL Worker")
     add_common_args_between_master_and_worker(parser)
diff --git a/elasticdl/python/ps/main.py b/elasticdl/python/ps/main.py
new file mode 100644
--- /dev/null
+++ b/elasticdl/python/ps/main.py
@@ -0,0 +1,19 @@
+from elasticdl.python.common.args import parse_ps_args
+from elasticdl.python.common.log_utils import get_logger
+from elasticdl.python.common.model_utils import (
+    get_module_file_path,
+    load_module,
+)
+
+
+def main():
+    args = parse_ps_args()
+    logger = get_logger("PS", level=args.log_level.upper())
+
+    model_module = load_module(
+        get_module_file_path(args.model_zoo, args.model_def)
+    ).__dict__
+    optimizer = model_module[args.optimizer]()
+
+    logger.info("Starting PS pod with optimizer as %s", optimizer._name)
+    # TODO: create `Parameters` class instance, and start PS RPC servicer.


Test patch:
diff --git a/elasticdl/python/tests/args_test.py b/elasticdl/python/tests/args_test.py
--- a/elasticdl/python/tests/args_test.py
+++ b/elasticdl/python/tests/args_test.py
@@ -4,6 +4,7 @@
 from elasticdl.python.common.args import (
     add_bool_param,
     build_arguments_from_parsed_result,
+    parse_ps_args,
 )
 
 
@@ -34,3 +35,33 @@ def test_build_arguments_from_parsed_result(self):
         value = "\t".join(sorted(original_arguments))
         target = "\t".join(sorted(["--bar", "False", "--foo", "4"]))
         self.assertEqual(value, target)
+
+    def test_parse_ps_args(self):
+        num_ps_pods = 2
+        minibatch_size = 16
+        num_minibatches_per_task = 8
+        model_zoo = "dummy_zoo"
+        model_def = "dummy_def"
+
+        original_args = [
+            "--num_ps_pods",
+            str(num_ps_pods),
+            "--minibatch_size",
+            str(minibatch_size),
+            "--model_zoo",
+            model_zoo,
+            "--model_def",
+            model_def,
+            "--job_name",
+            "test _args",
+            "--num_minibatches_per_task",
+            str(num_minibatches_per_task),
+        ]
+        parsed_args = parse_ps_args(original_args)
+        self.assertEqual(parsed_args.num_ps_pods, num_ps_pods)
+        self.assertEqual(parsed_args.minibatch_size, minibatch_size)
+        self.assertEqual(
+            parsed_args.num_minibatches_per_task, num_minibatches_per_task
+        )
+        self.assertEqual(parsed_args.model_zoo, model_zoo)
+        self.assertEqual(parsed_args.model_def, model_def)