Repository: sql-machine-learning/elasticdl

Issue:
clean arguments
There are different roles in ElasticDL. And different roles require different arguments from the command-line. 
All arguments except those passed to master are constructed by master. So if we need to pass a new argument to any role, we first add the argument in the master's argument list and then pass it on to the specific role. There is a main function for each role currently. In the main function, it will parse the argument and pass it to the constructor of the role.

The problem is we have to add an argument twice before we can use it.
+ In master's argument list
+ In the specific role's constructor

So we suggest passing directly the parsed argument object instead of unpacking the arguments and passing each one of them as the argument of the constructor.


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
@@ -491,7 +491,6 @@ def add_common_args_between_master_and_worker(parser):
     parser.add_argument(
         "--distribution_strategy",
         type=str,
-        default="ParameterServerStrategy",
         help="Master will use a distribution policy on a list of devices "
         "according to the distributed strategy, "
         'e.g. "ParameterServerStrategy"',
@@ -591,7 +590,7 @@ def parse_worker_args(worker_args=None):
         "--worker_id", help="ID unique to the worker", type=int, required=True
     )
     parser.add_argument("--job_type", help="Job type", required=True)
-    parser.add_argument("--master_addr", help="Master ip:port", required=True)
+    parser.add_argument("--master_addr", help="Master ip:port")
     parser.add_argument(
         "--embedding_service_endpoint",
         type=str,
@@ -611,6 +610,8 @@ def parse_worker_args(worker_args=None):
         help="Addresses of parameter service pods, separated by comma",
     )
 
+    if worker_args:
+        worker_args = list(map(str, worker_args))
     args, unknown_args = parser.parse_known_args(args=worker_args)
     print_args(args, groups=ALL_ARGS_GROUPS)
     if unknown_args:
diff --git a/elasticdl/python/worker/main.py b/elasticdl/python/worker/main.py
--- a/elasticdl/python/worker/main.py
+++ b/elasticdl/python/worker/main.py
@@ -8,6 +8,8 @@
 
 def main():
     args = parse_worker_args()
+    if args.master_addr is None:
+        raise ValueError("master_addr is missing for worker")
     channel = grpc.insecure_channel(
         args.master_addr,
         options=[
@@ -30,25 +32,7 @@ def main():
     logger = log_utils.get_logger(__name__)
 
     logger.info("Starting worker %d", args.worker_id)
-    worker = Worker(
-        args.worker_id,
-        args.job_type,
-        args.minibatch_size,
-        args.model_zoo,
-        channel=channel,
-        embedding_service_endpoint=eval(args.embedding_service_endpoint),
-        dataset_fn=args.dataset_fn,
-        loss=args.loss,
-        optimizer=args.optimizer,
-        eval_metrics_fn=args.eval_metrics_fn,
-        model_def=args.model_def,
-        model_params=args.model_params,
-        data_reader_params=args.data_reader_params,
-        prediction_outputs_processor=args.prediction_outputs_processor,
-        get_model_steps=args.get_model_steps,
-        distribution_strategy=args.distribution_strategy,
-        ps_channels=ps_channels,
-    )
+    worker = Worker(args, channel=channel, ps_channels=ps_channels)
     worker.run()
 
 
diff --git a/elasticdl/python/worker/worker.py b/elasticdl/python/worker/worker.py
--- a/elasticdl/python/worker/worker.py
+++ b/elasticdl/python/worker/worker.py
@@ -45,56 +45,25 @@ class Worker(object):
 
     def __init__(
         self,
-        worker_id,
-        job_type,
-        minibatch_size,
-        model_zoo,
-        dataset_fn="dataset_fn",
-        loss="loss",
-        optimizer="optimizer",
-        eval_metrics_fn="eval_metrics_fn",
+        args,
         channel=None,
-        embedding_service_endpoint=None,
-        model_def=None,
-        model_params="",
-        data_reader_params="",
-        prediction_outputs_processor="PredictionOutputsProcessor",
-        max_minibatch_retry_num=DEFAULT_MAX_MINIBATCH_RETRY_NUM,
-        get_model_steps=1,
-        distribution_strategy=None,
         ps_channels=None,
+        max_minibatch_retry_num=DEFAULT_MAX_MINIBATCH_RETRY_NUM,
     ):
         """
         Arguments:
-            worker_id: The worker ID.
-            job_type: The job type.
-            minibatch_size: The size of the minibatch used for each iteration.
-            model_zoo: The directory that contains user-defined model files
-                or a specific model file.
-            dataset_fn: The name of the dataset function defined in the
-                model file.
-            loss: The name of the loss function defined in the model file.
-            optimizer: The name of the optimizer defined in the model file.
-            eval_metrics_fn: The name of the evaluation metrics function
-                defined in the model file.
             channel: The channel for the gRPC master service.
-            embedding_service_endpoint: The endpoint to the embedding service.
-            model_def: The import path to the model definition
-                function/class in the model zoo, e.g.
-                "cifar10_subclass.CustomModel".
-            model_params: The dictionary of model parameters in a string
-                separated by semi-colon used to instantiate the model,
-                e.g. "param1=1; param2=2".
-            data_reader_params: The data reader parameters in a string
-                separated by semi-colon used to instantiate the data reader,
-                e.g. "param1=1; param2=2".
-            prediction_outputs_processor: The name of the prediction output
-                processor class defined in the model file.
-            get_model_steps: Worker will perform `get_model` from the
-                parameter server every this many steps.
+            ps_channels: TODO
             max_minibatch_retry_num: The maximum number of a minibatch retry
                 as its results (e.g. gradients) are not accepted by master.
+
         """
+        self._args = args
+        if channel is None:
+            self._stub = None
+        else:
+            self._stub = elasticdl_pb2_grpc.MasterStub(channel)
+
         self._use_multi_ps = False
         if isinstance(ps_channels, list):
             if len(ps_channels) > 0:
@@ -103,10 +72,17 @@ def __init__(
                     elasticdl_pb2_grpc.PserverStub(c) for c in ps_channels
                 ]
                 self._var_to_ps = {}
+        self._max_minibatch_retry_num = max_minibatch_retry_num
+        self._init_from_args(args)
 
-        self._worker_id = worker_id
-        self._job_type = job_type
-        self._minibatch_size = minibatch_size
+    def _init_from_args(self, args):
+        """
+        Please refer to elastic/python/common/args.py for more
+        details about arguments of a worker.
+        """
+        self._worker_id = args.worker_id
+        self._job_type = args.job_type
+        self._minibatch_size = args.minibatch_size
         (
             model_inst,
             self._dataset_fn,
@@ -115,37 +91,35 @@ def __init__(
             self._eval_metrics_fn,
             self._prediction_outputs_processor,
         ) = get_model_spec(
-            model_zoo=model_zoo,
-            model_def=model_def,
-            dataset_fn=dataset_fn,
-            loss=loss,
-            optimizer=optimizer,
-            eval_metrics_fn=eval_metrics_fn,
-            model_params=model_params,
-            prediction_outputs_processor=prediction_outputs_processor,
+            model_zoo=args.model_zoo,
+            model_def=args.model_def,
+            dataset_fn=args.dataset_fn,
+            loss=args.loss,
+            optimizer=args.optimizer,
+            eval_metrics_fn=args.eval_metrics_fn,
+            model_params=args.model_params,
+            prediction_outputs_processor=args.prediction_outputs_processor,
         )
 
-        self._embedding_service_endpoint = embedding_service_endpoint
-
-        if channel is None:
-            self._stub = None
-        else:
-            self._stub = elasticdl_pb2_grpc.MasterStub(channel)
+        self._embedding_service_endpoint = eval(
+            args.embedding_service_endpoint
+        )
 
         self._model_handler = ModelHandler.get_model_handler(
-            distribution_strategy, stub=self._stub
+            args.distribution_strategy, stub=self._stub
         )
         model_inst = self._model_handler.get_model_to_train(model_inst)
         self.set_model(model_inst)
 
-        self._max_minibatch_retry_num = max_minibatch_retry_num
         self._model_version = -1
         self._task_data_service = TaskDataService(
             self,
             self._job_type == JobType.TRAINING_WITH_EVALUATION,
-            data_reader_params=get_dict_from_params_str(data_reader_params),
+            data_reader_params=get_dict_from_params_str(
+                args.data_reader_params
+            ),
         )
-        self._get_model_steps = get_model_steps
+        self._get_model_steps = args.get_model_steps
         if self._get_model_steps > 1:
             self._opt = self._opt_fn()
             self._non_embed_grads = None


Test patch:
diff --git a/elasticdl/python/tests/checkpoint_test.py b/elasticdl/python/tests/checkpoint_test.py
--- a/elasticdl/python/tests/checkpoint_test.py
+++ b/elasticdl/python/tests/checkpoint_test.py
@@ -3,6 +3,7 @@
 import unittest
 
 from elasticdl.proto import elasticdl_pb2
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import JobType
 from elasticdl.python.common.model_utils import (
     get_module_file_path,
@@ -73,14 +74,21 @@ def testMaxCheckpointVersions(self):
 
             batch_size = 2
             # Launch the training
-            worker = Worker(
+            arguments = [
+                "--worker_id",
                 1,
+                "--job_type",
                 JobType.TRAINING_ONLY,
+                "--minibatch_size",
                 batch_size,
+                "--model_zoo",
                 _model_zoo_path,
-                model_def="test_module.custom_model",
-                channel=None,
-            )
+                "--model_def",
+                "test_module.custom_model",
+            ]
+            args = parse_worker_args(arguments)
+            worker = Worker(args)
+
             filename = create_recordio_file(128, DatasetName.TEST_MODULE, 1)
             task_d = _TaskDispatcher(
                 {filename: (0, 128)}, {}, {}, records_per_task=64, num_epochs=1
diff --git a/elasticdl/python/tests/indices_slices_gradient_test.py b/elasticdl/python/tests/indices_slices_gradient_test.py
--- a/elasticdl/python/tests/indices_slices_gradient_test.py
+++ b/elasticdl/python/tests/indices_slices_gradient_test.py
@@ -4,6 +4,7 @@
 import numpy as np
 import tensorflow as tf
 
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import JobType
 from elasticdl.python.master.servicer import MasterServicer
 from elasticdl.python.tests.in_process_master import InProcessMaster
@@ -40,14 +41,21 @@ def _create_master_and_worker(self):
             checkpoint_service=None,
             evaluation_service=None,
         )
-        worker = Worker(
+        arguments = [
+            "--worker_id",
             1,
+            "--job_type",
             JobType.TRAINING_ONLY,
+            "--minibatch_size",
             2,
+            "--model_zoo",
             _model_zoo_path,
-            model_def="test_module.custom_model",
-            channel=None,
-        )
+            "--model_def",
+            "test_module.custom_model",
+        ]
+        args = parse_worker_args(arguments)
+
+        worker = Worker(args)
         worker.set_model(model_inst)
         worker._stub = InProcessMaster(master)
 
diff --git a/elasticdl/python/tests/report_gradients_of_bet_test.py b/elasticdl/python/tests/report_gradients_of_bet_test.py
--- a/elasticdl/python/tests/report_gradients_of_bet_test.py
+++ b/elasticdl/python/tests/report_gradients_of_bet_test.py
@@ -5,6 +5,7 @@
 import numpy as np
 import tensorflow as tf
 
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import JobType
 from elasticdl.python.elasticdl.layers.embedding import EmbeddingAndIds
 from elasticdl.python.master.embedding_service import EmbeddingService
@@ -87,14 +88,21 @@ def _create_master_and_worker(
             checkpoint_service=None,
             evaluation_service=None,
         )
-        worker = Worker(
+        arguments = [
+            "--worker_id",
             1,
+            "--job_type",
             JobType.TRAINING_ONLY,
+            "--minibatch_size",
             2,
+            "--model_zoo",
             _model_zoo_path,
-            model_def="test_module.custom_model",
-            channel=None,
-        )
+            "--model_def",
+            "test_module.custom_model",
+        ]
+        args = parse_worker_args(arguments)
+
+        worker = Worker(args)
         worker.set_model(model_inst)
         worker._stub = InProcessMaster(master)
 
diff --git a/elasticdl/python/tests/test_utils.py b/elasticdl/python/tests/test_utils.py
--- a/elasticdl/python/tests/test_utils.py
+++ b/elasticdl/python/tests/test_utils.py
@@ -8,6 +8,7 @@
 from odps import ODPS
 
 from elasticdl.proto import elasticdl_pb2
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import JobType, ODPSConfig
 from elasticdl.python.common.model_utils import (
     get_module_file_path,
@@ -165,16 +166,24 @@ def distributed_train_and_evaluate(
         else JobType.EVALUATION_ONLY
     )
     batch_size = 8 if dataset_name == DatasetName.IMAGENET else 16
-    worker = Worker(
-        1,
+    arguments = [
+        "--worker_id",
+        "1",
+        "--job_type",
         job_type,
+        "--minibatch_size",
         batch_size,
+        "--model_zoo",
         model_zoo_path,
-        model_def=model_def,
-        model_params=model_params,
-        channel=None,
-        get_model_steps=get_model_steps,
-    )
+        "--model_def",
+        model_def,
+        "--model_params",
+        model_params,
+        "--get_model_steps",
+        get_model_steps,
+    ]
+    args = parse_worker_args(arguments)
+    worker = Worker(args)
 
     if dataset_name in [DatasetName.IMAGENET, DatasetName.FRAPPE]:
         record_num = batch_size
diff --git a/elasticdl/python/tests/worker_ps_interaction_test.py b/elasticdl/python/tests/worker_ps_interaction_test.py
--- a/elasticdl/python/tests/worker_ps_interaction_test.py
+++ b/elasticdl/python/tests/worker_ps_interaction_test.py
@@ -7,6 +7,7 @@
 import tensorflow as tf
 
 from elasticdl.proto import elasticdl_pb2
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import GRPC
 from elasticdl.python.common.hash_utils import int_to_id, string_to_id
 from elasticdl.python.common.model_utils import get_model_spec
@@ -89,15 +90,22 @@ def _restart_pserver(self):
         )
 
     def test_worker_pull_embedding(self):
-        worker = Worker(
-            worker_id=0,
-            job_type=elasticdl_pb2.TRAINING,
-            minibatch_size=self._batch_size,
-            model_zoo=self._model_zoo_path,
-            model_def=self._model_def,
-            ps_channels=self._channel,
-            distribution_strategy="ParameterServerStrategy",
-        )
+        arguments = [
+            "--worker_id",
+            0,
+            "--job_type",
+            elasticdl_pb2.TRAINING,
+            "--minibatch_size",
+            self._batch_size,
+            "--model_zoo",
+            self._model_zoo_path,
+            "--model_def",
+            self._model_def,
+            "--distribution_strategy",
+            "ParameterServerStrategy",
+        ]
+        args = parse_worker_args(arguments)
+        worker = Worker(args, ps_channels=self._channel)
 
         # Test lookup embedding vectors that do not exist
         layers = ["test-2", "test-2-slot"]
@@ -133,17 +141,23 @@ def test_compare_onebatch_train(self):
         # TODO(yunjian.lmh): test optimizer wrapper
         tf.keras.backend.clear_session()
         tf.random.set_seed(22)
-        worker = Worker(
-            worker_id=0,
-            job_type=elasticdl_pb2.TRAINING,
-            minibatch_size=self._batch_size,
-            model_zoo=self._model_zoo_path,
-            model_def=(
-                "mnist_functional_api.mnist_functional_api.custom_model"
-            ),
-            ps_channels=self._channel,
-            distribution_strategy="ParameterServerStrategy",
-        )
+        arguments = [
+            "--worker_id",
+            0,
+            "--job_type",
+            elasticdl_pb2.TRAINING,
+            "--minibatch_size",
+            self._batch_size,
+            "--model_zoo",
+            self._model_zoo_path,
+            "--model_def",
+            "mnist_functional_api.mnist_functional_api.custom_model",
+            "--distribution_strategy",
+            "ParameterServerStrategy",
+        ]
+        args = parse_worker_args(arguments)
+
+        worker = Worker(args, ps_channels=self._channel)
         worker._run_model_call_before_training(images)
         worker.get_model(0, elasticdl_pb2.MINIMUM)
         w_loss, w_grads = worker.training_process_eagerly(images, labels)
@@ -196,15 +210,23 @@ def _worker_train(self, train_db, test_db, dataset, stop_step):
             )
         else:
             raise ValueError("dataset %s is not supported", dataset)
-        worker = Worker(
-            worker_id=0,
-            job_type=elasticdl_pb2.TRAINING,
-            minibatch_size=self._batch_size,
-            model_zoo=self._model_zoo_path,
-            model_def=model_def,
-            ps_channels=self._channel,
-            distribution_strategy="ParameterServerStrategy",
-        )
+        arguments = [
+            "--worker_id",
+            0,
+            "--job_type",
+            elasticdl_pb2.TRAINING,
+            "--minibatch_size",
+            self._batch_size,
+            "--model_zoo",
+            self._model_zoo_path,
+            "--model_def",
+            model_def,
+            "--distribution_strategy",
+            "ParameterServerStrategy",
+        ]
+        args = parse_worker_args(arguments)
+
+        worker = Worker(args, ps_channels=self._channel)
         acc_meter = tf.keras.metrics.Accuracy()
         worker_results = []
         for step, (x, y) in enumerate(train_db):
@@ -351,15 +373,23 @@ def test_restart_ps(self):
             self._restart_pserver()
             tf.keras.backend.clear_session()
             tf.random.set_seed(22)
-            worker = Worker(
-                worker_id=0,
-                job_type=elasticdl_pb2.TRAINING,
-                minibatch_size=self._batch_size,
-                model_zoo=self._model_zoo_path,
-                model_def=self._model_def,
-                ps_channels=self._channel,
-                distribution_strategy="ParameterServerStrategy",
-            )
+            arguments = [
+                "--worker_id",
+                0,
+                "--job_type",
+                elasticdl_pb2.TRAINING,
+                "--minibatch_size",
+                self._batch_size,
+                "--model_zoo",
+                self._model_zoo_path,
+                "--model_def",
+                self._model_def,
+                "--distribution_strategy",
+                "ParameterServerStrategy",
+            ]
+            args = parse_worker_args(arguments)
+
+            worker = Worker(args, ps_channels=self._channel)
             workers.append(worker)
             worker._run_model_call_before_training(training_data[0][0])
             for i in range(num_data):
diff --git a/elasticdl/python/tests/worker_test.py b/elasticdl/python/tests/worker_test.py
--- a/elasticdl/python/tests/worker_test.py
+++ b/elasticdl/python/tests/worker_test.py
@@ -7,6 +7,7 @@
 import tensorflow as tf
 from tensorflow.keras.layers import Concatenate, Dense, Flatten
 
+from elasticdl.python.common.args import parse_worker_args
 from elasticdl.python.common.constants import JobType
 from elasticdl.python.elasticdl.layers.embedding import Embedding
 from elasticdl.python.master.embedding_service import EmbeddingService
@@ -134,14 +135,20 @@ def test_distributed_train_get_model_steps(self):
         )
 
     def test_embedding_layer(self):
-        worker = Worker(
+        arguments = [
+            "--worker_id",
             1,
+            "--job_type",
             JobType.TRAINING_ONLY,
+            "--minibatch_size",
             32,
+            "--model_zoo",
             _model_zoo_path,
-            model_def="embedding_test_module.EdlEmbeddingModel",
-            channel=None,
-        )
+            "--model_def",
+            "embedding_test_module.EdlEmbeddingModel",
+        ]
+        args = parse_worker_args(arguments)
+        worker = Worker(args)
         self.assertTrue(len(worker._embedding_layers) == 2)
 
     def test_train_acceleration_with_embedding(self):
@@ -157,14 +164,20 @@ def test_train_acceleration_with_embedding(self):
             checkpoint_service=None,
             evaluation_service=None,
         )
-        worker = Worker(
+        arguments = [
+            "--worker_id",
             1,
+            "--job_type",
             JobType.TRAINING_ONLY,
+            "--minibatch_size",
             32,
+            "--model_zoo",
             _model_zoo_path,
-            model_def="embedding_test_module.EdlEmbeddingModel",
-            channel=None,
-        )
+            "--model_def",
+            "embedding_test_module.EdlEmbeddingModel",
+        ]
+        args = parse_worker_args(arguments)
+        worker = Worker(args)
         worker._stub = InProcessMaster(master)
 
         inputs_list = [