Repository: sql-machine-learning/elasticdl

Issue:
Master creates `ps_command` and `ps_args` for Go PS pod.



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
@@ -249,6 +249,13 @@ def add_common_params(parser):
         "the job automatically, and users need to launch the job through "
         "command `kubectl create -f path_to_yaml_file`.",
     )
+    # delete this argument after finishing Go-based PS implementation
+    add_bool_param(
+        parser=parser,
+        name="--use_go_ps",
+        default=False,
+        help="True for Go-based PS, False for Python-based PS",
+    )
 
 
 def add_train_params(parser):
diff --git a/elasticdl/python/master/master.py b/elasticdl/python/master/master.py
--- a/elasticdl/python/master/master.py
+++ b/elasticdl/python/master/master.py
@@ -339,49 +339,75 @@ def _create_instance_manager(self, args):
             ]
             worker_args.extend(build_arguments_from_parsed_result(args))
 
-            ps_command = ["python"]
-            ps_args = [
-                "-m",
-                "elasticdl.python.ps.main",
-                "--grads_to_wait",
-                str(args.grads_to_wait),
-                "--lr_staleness_modulation",
-                str(args.lr_staleness_modulation),
-                "--use_async",
-                str(args.use_async),
-                "--minibatch_size",
-                str(args.minibatch_size),
-                "--model_zoo",
-                args.model_zoo,
-                "--model_def",
-                args.model_def,
-                "--job_name",
-                args.job_name,
-                "--num_minibatches_per_task",
-                str(args.num_minibatches_per_task),
-                "--port",
-                "2222",
-                "--master_addr",
-                self.master_addr,
-                "--namespace",
-                args.namespace,
-                "--evaluation_steps",
-                str(args.evaluation_steps),
-                "--checkpoint_dir",
-                str(args.checkpoint_dir),
-                "--checkpoint_steps",
-                str(args.checkpoint_steps),
-                "--keep_checkpoint_max",
-                str(args.keep_checkpoint_max),
-                "--num_ps_pods",
-                str(args.num_ps_pods),
-                "--checkpoint_dir_for_init",
-                str(args.checkpoint_dir_for_init),
-                "--num_workers",
-                str(args.num_workers),
-                "--log_level",
-                str(args.log_level),
-            ]
+            if args.use_go_ps:
+                # TODO: rename the Go PS executable using a meaningful filename
+                ps_command = ["main"]
+                ps_args = [
+                    "-job_name=" + args.job_name,
+                    "-namespace=" + args.namespace,
+                    "-master_addr=" + self.master_addr,
+                    "-port=2222",
+                    "-use_async=" + ("true" if args.use_async else "false"),
+                    "-grads_to_wait=" + str(args.grads_to_wait),
+                    "-lr_staleness_modulation="
+                    + ("true" if args.lr_staleness_modulation else "false"),
+                    "-sync_version_tolerance="
+                    + str(args.sync_version_tolerance),
+                    "-evaluation_steps=" + str(args.evaluation_steps),
+                    "-num_ps_pods=" + str(args.num_ps_pods),
+                    "-num_workers=" + str(args.num_workers),
+                    "-checkpoint_dir=" + str(args.checkpoint_dir),
+                    "-checkpoint_steps=" + str(args.checkpoint_steps),
+                    "-keep_checkpoint_max=" + str(args.keep_checkpoint_max),
+                    "-checkpoint_dir_for_init="
+                    + str(args.checkpoint_dir_for_init),
+                ]
+            else:
+                ps_command = ["python"]
+                ps_args = [
+                    "-m",
+                    "elasticdl.python.ps.main",
+                    "--grads_to_wait",
+                    str(args.grads_to_wait),
+                    "--lr_staleness_modulation",
+                    str(args.lr_staleness_modulation),
+                    "--sync_version_tolerance",
+                    str(args.sync_version_tolerance),
+                    "--use_async",
+                    str(args.use_async),
+                    "--model_zoo",
+                    args.model_zoo,
+                    "--model_def",
+                    args.model_def,
+                    "--job_name",
+                    args.job_name,
+                    "--port",
+                    "2222",
+                    "--master_addr",
+                    self.master_addr,
+                    "--namespace",
+                    args.namespace,
+                    "--evaluation_steps",
+                    str(args.evaluation_steps),
+                    "--checkpoint_dir",
+                    str(args.checkpoint_dir),
+                    "--checkpoint_steps",
+                    str(args.checkpoint_steps),
+                    "--keep_checkpoint_max",
+                    str(args.keep_checkpoint_max),
+                    "--num_ps_pods",
+                    str(args.num_ps_pods),
+                    "--checkpoint_dir_for_init",
+                    str(args.checkpoint_dir_for_init),
+                    "--num_workers",
+                    str(args.num_workers),
+                    "--log_level",
+                    str(args.log_level),
+                    "--minibatch_size",
+                    str(args.minibatch_size),
+                    "--num_minibatches_per_task",
+                    str(args.num_minibatches_per_task),
+                ]
 
             env_dict = parse_envs(args.envs)
             env = []


Test patch: