We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Remove manually create arguments in master/worker
In ElasticDL's master and worker. We manually create the arguments collection for them. It is a little hacky as we are adding more and more arguments.
After #1210 is merged, we can also update these two places where we manually add arguments.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `elasticdl/python/master/main.py`
Content:
```
1 import os
2 import time
3 from concurrent import futures
4 
5 import grpc
6 from kubernetes.client import V1EnvVar
7 
8 from elasticdl.proto import elasticdl_pb2_grpc
9 from elasticdl.python.common.args import parse_envs, parse_master_args
10 from elasticdl.python.common.constants import (
11     GRPC,
12     JobType,
13     WorkerManagerStatus,
14 )
15 from elasticdl.python.common.data_reader import RecordIODataReader
16 from elasticdl.python.common.k8s_tensorboard_client import TensorBoardClient
17 from elasticdl.python.common.log_util import get_logger
18 from elasticdl.python.common.model_helper import (
19     find_layer,
20     get_module_file_path,
21     load_model_from_module,
22     load_module,
23 )
24 from elasticdl.python.elasticdl.layers.embedding import Embedding
25 from elasticdl.python.master.checkpoint_service import CheckpointService
26 from elasticdl.python.master.embedding_service import EmbeddingService
27 from elasticdl.python.master.evaluation_service import EvaluationService
28 from elasticdl.python.master.k8s_worker_manager import WorkerManager
29 from elasticdl.python.master.servicer import MasterServicer
30 from elasticdl.python.master.task_dispatcher import _TaskDispatcher
31 from elasticdl.python.master.tensorboard_service import TensorboardService
32 
33 
34 def _make_task_dispatcher(
35     training_data_dir,
36     evaluation_data_dir,
37     prediction_data_dir,
38     records_per_task,
39     num_epochs,
40 ):
41     # TODO: Support any subclasses of `AbstractDataReader`
42     # and support passing specified parameters to the constructor
43     prediction_f_records = RecordIODataReader(
44         data_dir=prediction_data_dir
45     ).create_shards()
46 
47     return _TaskDispatcher(
48         RecordIODataReader(data_dir=training_data_dir).create_shards(),
49         RecordIODataReader(data_dir=evaluation_data_dir).create_shards(),
50         prediction_f_records,
51         records_per_task,
52         # Only generate prediction tasks for 1 epoch
53         1 if prediction_f_records else num_epochs,
54     )
55 
56 
57 def main():
58     args = parse_master_args()
59     logger = get_logger("master", level=args.log_level.upper())
60 
61     # Master addr
62     master_ip = os.getenv("MY_POD_IP", "localhost")
63     master_addr = "%s:%d" % (master_ip, args.port)
64 
65     # Start TensorBoard service if requested
66     if args.tensorboard_log_dir:
67         logger.info(
68             "Starting TensorBoard service with log directory %s",
69             args.tensorboard_log_dir,
70         )
71         # Start TensorBoard CLI
72         tb_service = TensorboardService(args.tensorboard_log_dir, master_ip)
73         tb_service.start()
74     else:
75         tb_service = None
76 
77     # Start task queue
78     logger.debug(
79         "Starting task queue with training data directory %s, "
80         "evaluation data directory %s, "
81         "and prediction data directory %s",
82         args.training_data_dir,
83         args.evaluation_data_dir,
84         args.prediction_data_dir,
85     )
86     task_d = _make_task_dispatcher(
87         args.training_data_dir,
88         args.evaluation_data_dir,
89         args.prediction_data_dir,
90         args.records_per_task,
91         args.num_epochs,
92     )
93     model_module = load_module(
94         get_module_file_path(args.model_zoo, args.model_def)
95     ).__dict__
96     model_inst = load_model_from_module(
97         args.model_def, model_module, args.model_params
98     )
99     optimizer = model_module[args.optimizer]()
100 
101     if all(
102         (
103             args.training_data_dir,
104             args.evaluation_data_dir,
105             args.evaluation_throttle_secs or args.evaluation_steps,
106         )
107     ):
108         job_type = JobType.TRAINING_WITH_EVALUATION
109     elif all(
110         (
111             args.evaluation_data_dir,
112             not args.training_data_dir,
113             not args.prediction_data_dir,
114         )
115     ):
116         job_type = JobType.EVALUATION_ONLY
117     elif all(
118         (
119             args.prediction_data_dir,
120             not args.evaluation_data_dir,
121             not args.training_data_dir,
122         )
123     ):
124         job_type = JobType.PREDICTION_ONLY
125     else:
126         job_type = JobType.TRAINING_ONLY
127 
128     # Initialize checkpoint service
129     if args.checkpoint_steps or job_type == JobType.TRAINING_WITH_EVALUATION:
130         logger.info("Starting checkpoint service")
131         checkpoint_service = CheckpointService(
132             args.checkpoint_dir,
133             args.checkpoint_steps,
134             args.keep_checkpoint_max,
135             job_type == JobType.TRAINING_WITH_EVALUATION,
136         )
137     else:
138         checkpoint_service = None
139 
140     # Initialize evaluation service
141     evaluation_service = None
142     if (
143         job_type == JobType.TRAINING_WITH_EVALUATION
144         or job_type == JobType.EVALUATION_ONLY
145     ):
146         logger.info(
147             "Starting evaluation service with throttle seconds %d "
148             " and evaluation steps %d",
149             args.evaluation_throttle_secs,
150             args.evaluation_steps,
151         )
152         evaluation_service = EvaluationService(
153             checkpoint_service,
154             tb_service,
155             task_d,
156             args.evaluation_start_delay_secs,
157             args.evaluation_throttle_secs,
158             args.evaluation_steps,
159             job_type == JobType.EVALUATION_ONLY,
160         )
161         evaluation_service.start()
162         task_d.set_evaluation_service(evaluation_service)
163 
164     embedding_service_endpoint = None
165     embedding_dims = {}
166     # Search for embedding layers in the model,
167     # if found, initialize embedding service
168     layers = find_layer(model_inst, Embedding)
169     if layers:
170         embedding_service = EmbeddingService()
171         embedding_service_endpoint = embedding_service.start_embedding_service(
172             job_name=args.job_name,
173             image_name=args.worker_image,
174             namespace=args.namespace,
175             resource_request=args.master_resource_request,
176             resource_limit=args.master_resource_limit,
177             pod_priority=args.worker_pod_priority,
178             volume=args.volume,
179             image_pull_policy=args.image_pull_policy,
180             restart_policy=args.restart_policy,
181             cluster_spec=args.cluster_spec,
182         )
183         logger.info(
184             "Embedding service start succeeded. The endpoint is %s."
185             % str(embedding_service_endpoint)
186         )
187         embedding_dims = dict(
188             [(layer.name, layer.output_dim) for layer in layers]
189         )
190 
191     # The master service
192     logger.info("Starting master service")
193     server = grpc.server(
194         futures.ThreadPoolExecutor(max_workers=64),
195         options=[
196             ("grpc.max_send_message_length", GRPC.MAX_SEND_MESSAGE_LENGTH),
197             (
198                 "grpc.max_receive_message_length",
199                 GRPC.MAX_RECEIVE_MESSAGE_LENGTH,
200             ),
201         ],
202     )
203     master_servicer = MasterServicer(
204         args.grads_to_wait,
205         args.minibatch_size,
206         optimizer,
207         task_d,
208         init_var=model_inst.trainable_variables if model_inst.built else [],
209         embedding_dims=embedding_dims,
210         checkpoint_filename_for_init=args.checkpoint_filename_for_init,
211         checkpoint_service=checkpoint_service,
212         evaluation_service=evaluation_service,
213         embedding_service_endpoint=embedding_service_endpoint,
214         lr_staleness_modulation=args.lr_staleness_modulation,
215         use_async=args.use_async,
216     )
217     elasticdl_pb2_grpc.add_MasterServicer_to_server(master_servicer, server)
218     server.add_insecure_port("[::]:{}".format(args.port))
219     server.start()
220     logger.info("Server started at port: %d", args.port)
221 
222     worker_manager = None
223     if args.num_workers:
224         assert args.worker_image, "Worker image cannot be empty"
225 
226         worker_command = ["python"]
227         worker_args = [
228             "-m",
229             "elasticdl.python.worker.main",
230             "--model_zoo",
231             args.model_zoo,
232             "--master_addr",
233             master_addr,
234             "--log_level",
235             args.log_level,
236             "--dataset_fn",
237             args.dataset_fn,
238             "--loss",
239             args.loss,
240             "--optimizer",
241             args.optimizer,
242             "--eval_metrics_fn",
243             args.eval_metrics_fn,
244             "--model_def",
245             args.model_def,
246             "--job_type",
247             job_type,
248             "--minibatch_size",
249             str(args.minibatch_size),
250             "--embedding_service_endpoint",
251             str(embedding_service_endpoint),
252             "--get_model_steps",
253             str(args.get_model_steps),
254         ]
255 
256         env_dict = parse_envs(args.envs)
257         env = []
258         for key in env_dict:
259             env.append(V1EnvVar(name=key, value=env_dict[key]))
260 
261         worker_manager = WorkerManager(
262             task_d,
263             job_name=args.job_name,
264             image_name=args.worker_image,
265             command=worker_command,
266             args=worker_args,
267             namespace=args.namespace,
268             num_workers=args.num_workers,
269             worker_resource_request=args.worker_resource_request,
270             worker_resource_limit=args.worker_resource_limit,
271             pod_priority=args.worker_pod_priority,
272             volume=args.volume,
273             image_pull_policy=args.image_pull_policy,
274             restart_policy=args.restart_policy,
275             cluster_spec=args.cluster_spec,
276             envs=env,
277         )
278         worker_manager.update_status(WorkerManagerStatus.PENDING)
279         logger.info("Launching %d workers", args.num_workers)
280         worker_manager.start_workers()
281         worker_manager.update_status(WorkerManagerStatus.RUNNING)
282 
283     # Start TensorBoard k8s Service if requested
284     if tb_service:
285         TensorBoardClient(
286             job_name=args.job_name,
287             image_name=args.worker_image,
288             namespace=args.namespace,
289         ).start_tensorboard_service()
290 
291     try:
292         while True:
293             if task_d.finished():
294                 if worker_manager:
295                     worker_manager.update_status(WorkerManagerStatus.FINISHED)
296                 if args.output:
297                     master_servicer.save_latest_checkpoint(args.output)
298                 break
299             time.sleep(30)
300     except KeyboardInterrupt:
301         logger.warning("Server stopping")
302 
303     if evaluation_service:
304         logger.info("Stopping evaluation service")
305         evaluation_service.stop()
306 
307     logger.info("Stopping RPC server")
308     server.stop(0)
309 
310     # Keep TensorBoard running when all the tasks are finished
311     if tb_service:
312         logger.info(
313             "All tasks finished. Keeping TensorBoard service running..."
314         )
315         while True:
316             if tb_service.is_active():
317                 time.sleep(10)
318             else:
319                 logger.warning(
320                     "Unable to keep TensorBoard running. "
321                     "It has already terminated"
322                 )
323                 break
324     logger.info("Master stopped")
325 
326 
327 if __name__ == "__main__":
328     main()
329
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/elasticdl/python/master/main.py b/elasticdl/python/master/main.py
--- a/elasticdl/python/master/main.py
+++ b/elasticdl/python/master/main.py
@@ -6,7 +6,11 @@
 from kubernetes.client import V1EnvVar
 
 from elasticdl.proto import elasticdl_pb2_grpc
-from elasticdl.python.common.args import parse_envs, parse_master_args
+from elasticdl.python.common.args import (
+    build_arguments_from_parsed_result,
+    parse_envs,
+    parse_master_args,
+)
 from elasticdl.python.common.constants import (
     GRPC,
     JobType,
@@ -227,31 +231,14 @@
         worker_args = [
             "-m",
             "elasticdl.python.worker.main",
-            "--model_zoo",
-            args.model_zoo,
             "--master_addr",
             master_addr,
-            "--log_level",
-            args.log_level,
-            "--dataset_fn",
-            args.dataset_fn,
-            "--loss",
-            args.loss,
-            "--optimizer",
-            args.optimizer,
-            "--eval_metrics_fn",
-            args.eval_metrics_fn,
-            "--model_def",
-            args.model_def,
             "--job_type",
             job_type,
-            "--minibatch_size",
-            str(args.minibatch_size),
             "--embedding_service_endpoint",
             str(embedding_service_endpoint),
-            "--get_model_steps",
-            str(args.get_model_steps),
         ]
+        worker_args.extend(build_arguments_from_parsed_result(args))
 
         env_dict = parse_envs(args.envs)
         env = []

{"golden_diff": "diff --git a/elasticdl/python/master/main.py b/elasticdl/python/master/main.py\n--- a/elasticdl/python/master/main.py\n+++ b/elasticdl/python/master/main.py\n@@ -6,7 +6,11 @@\n from kubernetes.client import V1EnvVar\n \n from elasticdl.proto import elasticdl_pb2_grpc\n-from elasticdl.python.common.args import parse_envs, parse_master_args\n+from elasticdl.python.common.args import (\n+    build_arguments_from_parsed_result,\n+    parse_envs,\n+    parse_master_args,\n+)\n from elasticdl.python.common.constants import (\n     GRPC,\n     JobType,\n@@ -227,31 +231,14 @@\n         worker_args = [\n             \"-m\",\n             \"elasticdl.python.worker.main\",\n-            \"--model_zoo\",\n-            args.model_zoo,\n             \"--master_addr\",\n             master_addr,\n-            \"--log_level\",\n-            args.log_level,\n-            \"--dataset_fn\",\n-            args.dataset_fn,\n-            \"--loss\",\n-            args.loss,\n-            \"--optimizer\",\n-            args.optimizer,\n-            \"--eval_metrics_fn\",\n-            args.eval_metrics_fn,\n-            \"--model_def\",\n-            args.model_def,\n             \"--job_type\",\n             job_type,\n-            \"--minibatch_size\",\n-            str(args.minibatch_size),\n             \"--embedding_service_endpoint\",\n             str(embedding_service_endpoint),\n-            \"--get_model_steps\",\n-            str(args.get_model_steps),\n         ]\n+        worker_args.extend(build_arguments_from_parsed_result(args))\n \n         env_dict = parse_envs(args.envs)\n         env = []\n", "issue": "Remove manually create arguments in master/worker\nIn ElasticDL's master and worker. We manually create the arguments collection for them. It is a little hacky as we are adding more and more arguments.\r\nAfter #1210 is merged, we can also update these two places where we manually add arguments.\n", "before_files": [{"content": "import os\nimport time\nfrom concurrent import futures\n\nimport grpc\nfrom kubernetes.client import V1EnvVar\n\nfrom elasticdl.proto import elasticdl_pb2_grpc\nfrom elasticdl.python.common.args import parse_envs, parse_master_args\nfrom elasticdl.python.common.constants import (\n    GRPC,\n    JobType,\n    WorkerManagerStatus,\n)\nfrom elasticdl.python.common.data_reader import RecordIODataReader\nfrom elasticdl.python.common.k8s_tensorboard_client import TensorBoardClient\nfrom elasticdl.python.common.log_util import get_logger\nfrom elasticdl.python.common.model_helper import (\n    find_layer,\n    get_module_file_path,\n    load_model_from_module,\n    load_module,\n)\nfrom elasticdl.python.elasticdl.layers.embedding import Embedding\nfrom elasticdl.python.master.checkpoint_service import CheckpointService\nfrom elasticdl.python.master.embedding_service import EmbeddingService\nfrom elasticdl.python.master.evaluation_service import EvaluationService\nfrom elasticdl.python.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.python.master.servicer import MasterServicer\nfrom elasticdl.python.master.task_dispatcher import _TaskDispatcher\nfrom elasticdl.python.master.tensorboard_service import TensorboardService\n\n\ndef _make_task_dispatcher(\n    training_data_dir,\n    evaluation_data_dir,\n    prediction_data_dir,\n    records_per_task,\n    num_epochs,\n):\n    # TODO: Support any subclasses of `AbstractDataReader`\n    # and support passing specified parameters to the constructor\n    prediction_f_records = RecordIODataReader(\n        data_dir=prediction_data_dir\n    ).create_shards()\n\n    return _TaskDispatcher(\n        RecordIODataReader(data_dir=training_data_dir).create_shards(),\n        RecordIODataReader(data_dir=evaluation_data_dir).create_shards(),\n        prediction_f_records,\n        records_per_task,\n        # Only generate prediction tasks for 1 epoch\n        1 if prediction_f_records else num_epochs,\n    )\n\n\ndef main():\n    args = parse_master_args()\n    logger = get_logger(\"master\", level=args.log_level.upper())\n\n    # Master addr\n    master_ip = os.getenv(\"MY_POD_IP\", \"localhost\")\n    master_addr = \"%s:%d\" % (master_ip, args.port)\n\n    # Start TensorBoard service if requested\n    if args.tensorboard_log_dir:\n        logger.info(\n            \"Starting TensorBoard service with log directory %s\",\n            args.tensorboard_log_dir,\n        )\n        # Start TensorBoard CLI\n        tb_service = TensorboardService(args.tensorboard_log_dir, master_ip)\n        tb_service.start()\n    else:\n        tb_service = None\n\n    # Start task queue\n    logger.debug(\n        \"Starting task queue with training data directory %s, \"\n        \"evaluation data directory %s, \"\n        \"and prediction data directory %s\",\n        args.training_data_dir,\n        args.evaluation_data_dir,\n        args.prediction_data_dir,\n    )\n    task_d = _make_task_dispatcher(\n        args.training_data_dir,\n        args.evaluation_data_dir,\n        args.prediction_data_dir,\n        args.records_per_task,\n        args.num_epochs,\n    )\n    model_module = load_module(\n        get_module_file_path(args.model_zoo, args.model_def)\n    ).__dict__\n    model_inst = load_model_from_module(\n        args.model_def, model_module, args.model_params\n    )\n    optimizer = model_module[args.optimizer]()\n\n    if all(\n        (\n            args.training_data_dir,\n            args.evaluation_data_dir,\n            args.evaluation_throttle_secs or args.evaluation_steps,\n        )\n    ):\n        job_type = JobType.TRAINING_WITH_EVALUATION\n    elif all(\n        (\n            args.evaluation_data_dir,\n            not args.training_data_dir,\n            not args.prediction_data_dir,\n        )\n    ):\n        job_type = JobType.EVALUATION_ONLY\n    elif all(\n        (\n            args.prediction_data_dir,\n            not args.evaluation_data_dir,\n            not args.training_data_dir,\n        )\n    ):\n        job_type = JobType.PREDICTION_ONLY\n    else:\n        job_type = JobType.TRAINING_ONLY\n\n    # Initialize checkpoint service\n    if args.checkpoint_steps or job_type == JobType.TRAINING_WITH_EVALUATION:\n        logger.info(\"Starting checkpoint service\")\n        checkpoint_service = CheckpointService(\n            args.checkpoint_dir,\n            args.checkpoint_steps,\n            args.keep_checkpoint_max,\n            job_type == JobType.TRAINING_WITH_EVALUATION,\n        )\n    else:\n        checkpoint_service = None\n\n    # Initialize evaluation service\n    evaluation_service = None\n    if (\n        job_type == JobType.TRAINING_WITH_EVALUATION\n        or job_type == JobType.EVALUATION_ONLY\n    ):\n        logger.info(\n            \"Starting evaluation service with throttle seconds %d \"\n            \" and evaluation steps %d\",\n            args.evaluation_throttle_secs,\n            args.evaluation_steps,\n        )\n        evaluation_service = EvaluationService(\n            checkpoint_service,\n            tb_service,\n            task_d,\n            args.evaluation_start_delay_secs,\n            args.evaluation_throttle_secs,\n            args.evaluation_steps,\n            job_type == JobType.EVALUATION_ONLY,\n        )\n        evaluation_service.start()\n        task_d.set_evaluation_service(evaluation_service)\n\n    embedding_service_endpoint = None\n    embedding_dims = {}\n    # Search for embedding layers in the model,\n    # if found, initialize embedding service\n    layers = find_layer(model_inst, Embedding)\n    if layers:\n        embedding_service = EmbeddingService()\n        embedding_service_endpoint = embedding_service.start_embedding_service(\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            namespace=args.namespace,\n            resource_request=args.master_resource_request,\n            resource_limit=args.master_resource_limit,\n            pod_priority=args.worker_pod_priority,\n            volume=args.volume,\n            image_pull_policy=args.image_pull_policy,\n            restart_policy=args.restart_policy,\n            cluster_spec=args.cluster_spec,\n        )\n        logger.info(\n            \"Embedding service start succeeded. The endpoint is %s.\"\n            % str(embedding_service_endpoint)\n        )\n        embedding_dims = dict(\n            [(layer.name, layer.output_dim) for layer in layers]\n        )\n\n    # The master service\n    logger.info(\"Starting master service\")\n    server = grpc.server(\n        futures.ThreadPoolExecutor(max_workers=64),\n        options=[\n            (\"grpc.max_send_message_length\", GRPC.MAX_SEND_MESSAGE_LENGTH),\n            (\n                \"grpc.max_receive_message_length\",\n                GRPC.MAX_RECEIVE_MESSAGE_LENGTH,\n            ),\n        ],\n    )\n    master_servicer = MasterServicer(\n        args.grads_to_wait,\n        args.minibatch_size,\n        optimizer,\n        task_d,\n        init_var=model_inst.trainable_variables if model_inst.built else [],\n        embedding_dims=embedding_dims,\n        checkpoint_filename_for_init=args.checkpoint_filename_for_init,\n        checkpoint_service=checkpoint_service,\n        evaluation_service=evaluation_service,\n        embedding_service_endpoint=embedding_service_endpoint,\n        lr_staleness_modulation=args.lr_staleness_modulation,\n        use_async=args.use_async,\n    )\n    elasticdl_pb2_grpc.add_MasterServicer_to_server(master_servicer, server)\n    server.add_insecure_port(\"[::]:{}\".format(args.port))\n    server.start()\n    logger.info(\"Server started at port: %d\", args.port)\n\n    worker_manager = None\n    if args.num_workers:\n        assert args.worker_image, \"Worker image cannot be empty\"\n\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.python.worker.main\",\n            \"--model_zoo\",\n            args.model_zoo,\n            \"--master_addr\",\n            master_addr,\n            \"--log_level\",\n            args.log_level,\n            \"--dataset_fn\",\n            args.dataset_fn,\n            \"--loss\",\n            args.loss,\n            \"--optimizer\",\n            args.optimizer,\n            \"--eval_metrics_fn\",\n            args.eval_metrics_fn,\n            \"--model_def\",\n            args.model_def,\n            \"--job_type\",\n            job_type,\n            \"--minibatch_size\",\n            str(args.minibatch_size),\n            \"--embedding_service_endpoint\",\n            str(embedding_service_endpoint),\n            \"--get_model_steps\",\n            str(args.get_model_steps),\n        ]\n\n        env_dict = parse_envs(args.envs)\n        env = []\n        for key in env_dict:\n            env.append(V1EnvVar(name=key, value=env_dict[key]))\n\n        worker_manager = WorkerManager(\n            task_d,\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=args.namespace,\n            num_workers=args.num_workers,\n            worker_resource_request=args.worker_resource_request,\n            worker_resource_limit=args.worker_resource_limit,\n            pod_priority=args.worker_pod_priority,\n            volume=args.volume,\n            image_pull_policy=args.image_pull_policy,\n            restart_policy=args.restart_policy,\n            cluster_spec=args.cluster_spec,\n            envs=env,\n        )\n        worker_manager.update_status(WorkerManagerStatus.PENDING)\n        logger.info(\"Launching %d workers\", args.num_workers)\n        worker_manager.start_workers()\n        worker_manager.update_status(WorkerManagerStatus.RUNNING)\n\n    # Start TensorBoard k8s Service if requested\n    if tb_service:\n        TensorBoardClient(\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            namespace=args.namespace,\n        ).start_tensorboard_service()\n\n    try:\n        while True:\n            if task_d.finished():\n                if worker_manager:\n                    worker_manager.update_status(WorkerManagerStatus.FINISHED)\n                if args.output:\n                    master_servicer.save_latest_checkpoint(args.output)\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if evaluation_service:\n        logger.info(\"Stopping evaluation service\")\n        evaluation_service.stop()\n\n    logger.info(\"Stopping RPC server\")\n    server.stop(0)\n\n    # Keep TensorBoard running when all the tasks are finished\n    if tb_service:\n        logger.info(\n            \"All tasks finished. Keeping TensorBoard service running...\"\n        )\n        while True:\n            if tb_service.is_active():\n                time.sleep(10)\n            else:\n                logger.warning(\n                    \"Unable to keep TensorBoard running. \"\n                    \"It has already terminated\"\n                )\n                break\n    logger.info(\"Master stopped\")\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "elasticdl/python/master/main.py"}], "after_files": [{"content": "import os\nimport time\nfrom concurrent import futures\n\nimport grpc\nfrom kubernetes.client import V1EnvVar\n\nfrom elasticdl.proto import elasticdl_pb2_grpc\nfrom elasticdl.python.common.args import (\n    build_arguments_from_parsed_result,\n    parse_envs,\n    parse_master_args,\n)\nfrom elasticdl.python.common.constants import (\n    GRPC,\n    JobType,\n    WorkerManagerStatus,\n)\nfrom elasticdl.python.common.data_reader import RecordIODataReader\nfrom elasticdl.python.common.k8s_tensorboard_client import TensorBoardClient\nfrom elasticdl.python.common.log_util import get_logger\nfrom elasticdl.python.common.model_helper import (\n    find_layer,\n    get_module_file_path,\n    load_model_from_module,\n    load_module,\n)\nfrom elasticdl.python.elasticdl.layers.embedding import Embedding\nfrom elasticdl.python.master.checkpoint_service import CheckpointService\nfrom elasticdl.python.master.embedding_service import EmbeddingService\nfrom elasticdl.python.master.evaluation_service import EvaluationService\nfrom elasticdl.python.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.python.master.servicer import MasterServicer\nfrom elasticdl.python.master.task_dispatcher import _TaskDispatcher\nfrom elasticdl.python.master.tensorboard_service import TensorboardService\n\n\ndef _make_task_dispatcher(\n    training_data_dir,\n    evaluation_data_dir,\n    prediction_data_dir,\n    records_per_task,\n    num_epochs,\n):\n    # TODO: Support any subclasses of `AbstractDataReader`\n    # and support passing specified parameters to the constructor\n    prediction_f_records = RecordIODataReader(\n        data_dir=prediction_data_dir\n    ).create_shards()\n\n    return _TaskDispatcher(\n        RecordIODataReader(data_dir=training_data_dir).create_shards(),\n        RecordIODataReader(data_dir=evaluation_data_dir).create_shards(),\n        prediction_f_records,\n        records_per_task,\n        # Only generate prediction tasks for 1 epoch\n        1 if prediction_f_records else num_epochs,\n    )\n\n\ndef main():\n    args = parse_master_args()\n    logger = get_logger(\"master\", level=args.log_level.upper())\n\n    # Master addr\n    master_ip = os.getenv(\"MY_POD_IP\", \"localhost\")\n    master_addr = \"%s:%d\" % (master_ip, args.port)\n\n    # Start TensorBoard service if requested\n    if args.tensorboard_log_dir:\n        logger.info(\n            \"Starting TensorBoard service with log directory %s\",\n            args.tensorboard_log_dir,\n        )\n        # Start TensorBoard CLI\n        tb_service = TensorboardService(args.tensorboard_log_dir, master_ip)\n        tb_service.start()\n    else:\n        tb_service = None\n\n    # Start task queue\n    logger.debug(\n        \"Starting task queue with training data directory %s, \"\n        \"evaluation data directory %s, \"\n        \"and prediction data directory %s\",\n        args.training_data_dir,\n        args.evaluation_data_dir,\n        args.prediction_data_dir,\n    )\n    task_d = _make_task_dispatcher(\n        args.training_data_dir,\n        args.evaluation_data_dir,\n        args.prediction_data_dir,\n        args.records_per_task,\n        args.num_epochs,\n    )\n    model_module = load_module(\n        get_module_file_path(args.model_zoo, args.model_def)\n    ).__dict__\n    model_inst = load_model_from_module(\n        args.model_def, model_module, args.model_params\n    )\n    optimizer = model_module[args.optimizer]()\n\n    if all(\n        (\n            args.training_data_dir,\n            args.evaluation_data_dir,\n            args.evaluation_throttle_secs or args.evaluation_steps,\n        )\n    ):\n        job_type = JobType.TRAINING_WITH_EVALUATION\n    elif all(\n        (\n            args.evaluation_data_dir,\n            not args.training_data_dir,\n            not args.prediction_data_dir,\n        )\n    ):\n        job_type = JobType.EVALUATION_ONLY\n    elif all(\n        (\n            args.prediction_data_dir,\n            not args.evaluation_data_dir,\n            not args.training_data_dir,\n        )\n    ):\n        job_type = JobType.PREDICTION_ONLY\n    else:\n        job_type = JobType.TRAINING_ONLY\n\n    # Initialize checkpoint service\n    if args.checkpoint_steps or job_type == JobType.TRAINING_WITH_EVALUATION:\n        logger.info(\"Starting checkpoint service\")\n        checkpoint_service = CheckpointService(\n            args.checkpoint_dir,\n            args.checkpoint_steps,\n            args.keep_checkpoint_max,\n            job_type == JobType.TRAINING_WITH_EVALUATION,\n        )\n    else:\n        checkpoint_service = None\n\n    # Initialize evaluation service\n    evaluation_service = None\n    if (\n        job_type == JobType.TRAINING_WITH_EVALUATION\n        or job_type == JobType.EVALUATION_ONLY\n    ):\n        logger.info(\n            \"Starting evaluation service with throttle seconds %d \"\n            \" and evaluation steps %d\",\n            args.evaluation_throttle_secs,\n            args.evaluation_steps,\n        )\n        evaluation_service = EvaluationService(\n            checkpoint_service,\n            tb_service,\n            task_d,\n            args.evaluation_start_delay_secs,\n            args.evaluation_throttle_secs,\n            args.evaluation_steps,\n            job_type == JobType.EVALUATION_ONLY,\n        )\n        evaluation_service.start()\n        task_d.set_evaluation_service(evaluation_service)\n\n    embedding_service_endpoint = None\n    embedding_dims = {}\n    # Search for embedding layers in the model,\n    # if found, initialize embedding service\n    layers = find_layer(model_inst, Embedding)\n    if layers:\n        embedding_service = EmbeddingService()\n        embedding_service_endpoint = embedding_service.start_embedding_service(\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            namespace=args.namespace,\n            resource_request=args.master_resource_request,\n            resource_limit=args.master_resource_limit,\n            pod_priority=args.worker_pod_priority,\n            volume=args.volume,\n            image_pull_policy=args.image_pull_policy,\n            restart_policy=args.restart_policy,\n            cluster_spec=args.cluster_spec,\n        )\n        logger.info(\n            \"Embedding service start succeeded. The endpoint is %s.\"\n            % str(embedding_service_endpoint)\n        )\n        embedding_dims = dict(\n            [(layer.name, layer.output_dim) for layer in layers]\n        )\n\n    # The master service\n    logger.info(\"Starting master service\")\n    server = grpc.server(\n        futures.ThreadPoolExecutor(max_workers=64),\n        options=[\n            (\"grpc.max_send_message_length\", GRPC.MAX_SEND_MESSAGE_LENGTH),\n            (\n                \"grpc.max_receive_message_length\",\n                GRPC.MAX_RECEIVE_MESSAGE_LENGTH,\n            ),\n        ],\n    )\n    master_servicer = MasterServicer(\n        args.grads_to_wait,\n        args.minibatch_size,\n        optimizer,\n        task_d,\n        init_var=model_inst.trainable_variables if model_inst.built else [],\n        embedding_dims=embedding_dims,\n        checkpoint_filename_for_init=args.checkpoint_filename_for_init,\n        checkpoint_service=checkpoint_service,\n        evaluation_service=evaluation_service,\n        embedding_service_endpoint=embedding_service_endpoint,\n        lr_staleness_modulation=args.lr_staleness_modulation,\n        use_async=args.use_async,\n    )\n    elasticdl_pb2_grpc.add_MasterServicer_to_server(master_servicer, server)\n    server.add_insecure_port(\"[::]:{}\".format(args.port))\n    server.start()\n    logger.info(\"Server started at port: %d\", args.port)\n\n    worker_manager = None\n    if args.num_workers:\n        assert args.worker_image, \"Worker image cannot be empty\"\n\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.python.worker.main\",\n            \"--master_addr\",\n            master_addr,\n            \"--job_type\",\n            job_type,\n            \"--embedding_service_endpoint\",\n            str(embedding_service_endpoint),\n        ]\n        worker_args.extend(build_arguments_from_parsed_result(args))\n\n        env_dict = parse_envs(args.envs)\n        env = []\n        for key in env_dict:\n            env.append(V1EnvVar(name=key, value=env_dict[key]))\n\n        worker_manager = WorkerManager(\n            task_d,\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=args.namespace,\n            num_workers=args.num_workers,\n            worker_resource_request=args.worker_resource_request,\n            worker_resource_limit=args.worker_resource_limit,\n            pod_priority=args.worker_pod_priority,\n            volume=args.volume,\n            image_pull_policy=args.image_pull_policy,\n            restart_policy=args.restart_policy,\n            cluster_spec=args.cluster_spec,\n            envs=env,\n        )\n        worker_manager.update_status(WorkerManagerStatus.PENDING)\n        logger.info(\"Launching %d workers\", args.num_workers)\n        worker_manager.start_workers()\n        worker_manager.update_status(WorkerManagerStatus.RUNNING)\n\n    # Start TensorBoard k8s Service if requested\n    if tb_service:\n        TensorBoardClient(\n            job_name=args.job_name,\n            image_name=args.worker_image,\n            namespace=args.namespace,\n        ).start_tensorboard_service()\n\n    try:\n        while True:\n            if task_d.finished():\n                if worker_manager:\n                    worker_manager.update_status(WorkerManagerStatus.FINISHED)\n                if args.output:\n                    master_servicer.save_latest_checkpoint(args.output)\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if evaluation_service:\n        logger.info(\"Stopping evaluation service\")\n        evaluation_service.stop()\n\n    logger.info(\"Stopping RPC server\")\n    server.stop(0)\n\n    # Keep TensorBoard running when all the tasks are finished\n    if tb_service:\n        logger.info(\n            \"All tasks finished. Keeping TensorBoard service running...\"\n        )\n        while True:\n            if tb_service.is_active():\n                time.sleep(10)\n            else:\n                logger.warning(\n                    \"Unable to keep TensorBoard running. \"\n                    \"It has already terminated\"\n                )\n                break\n    logger.info(\"Master stopped\")\n\n\nif __name__ == \"__main__\":\n    main()\n", "path": "elasticdl/python/master/main.py"}]}