prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Make master exist when there are no tasks left.
  Currently, master exists when there are no tasks left AND all workers are gone. It might left hanging if a worker got preempted.
  
  --- 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/elasticdl/master/main.py`
  Content:
  ```
  1 import logging
  2 import time
  3 import argparse
  4 import os
  5 
  6 import grpc
  7 import tensorflow as tf
  8 
  9 tf.enable_eager_execution()
  10 
  11 from concurrent import futures
  12 from recordio import File
  13 from elasticdl.proto import master_pb2_grpc
  14 from elasticdl.master.servicer import MasterServicer
  15 from elasticdl.master.task_queue import _TaskQueue
  16 from elasticdl.master.k8s_worker_manager import WorkerManager
  17 from elasticdl.common.model_helper import load_user_model, build_model
  18 
  19 
  20 def _make_task_queue(data_dir, record_per_task, num_epoch):
  21     f_records = {}
  22     for f in os.listdir(data_dir):
  23         p = os.path.join(data_dir, f)
  24         with File(p, "r") as rio:
  25             f_records[p] = rio.count()
  26     return _TaskQueue(f_records, record_per_task, num_epoch)
  27 
  28 
  29 def _parse_args():
  30     parser = argparse.ArgumentParser(description="ElasticDL Master")
  31     parser.add_argument(
  32         "--model_file",
  33         help="Full file path of user defined neural model",
  34         required=True,
  35     )
  36     parser.add_argument(
  37         "--train_data_dir",
  38         help="Training data directory. Files should be in RecordIO format",
  39         required=True,
  40     )
  41     parser.add_argument("--record_per_task", type=int, required=True)
  42     parser.add_argument("--num_epoch", type=int, required=True)
  43     parser.add_argument(
  44         "--grads_to_wait",
  45         type=int,
  46         help="Number of gradients to wait before updating model",
  47         required=True,
  48     )
  49     parser.add_argument(
  50         "--minibatch_size",
  51         type=int,
  52         help="Minibatch size used by workers to compute gradients",
  53         required=True,
  54     )
  55     parser.add_argument(
  56         "--num_worker",
  57         type=int,
  58         help="the number of workers used in training",
  59         default=0,
  60     )
  61     parser.add_argument(
  62         "--worker_cpu_request",
  63         help="the minimal cpu required by worker in training",
  64         default="1000m",
  65     )
  66     parser.add_argument(
  67         "--worker_cpu_limit",
  68         help="the maximal cpu used by worker in training",
  69         default="1000m",
  70     )
  71     parser.add_argument(
  72         "--worker_memory_request",
  73         help="the minimal memory required by worker in training",
  74         default="4096Mi",
  75     )
  76     parser.add_argument(
  77         "--worker_memory_limit",
  78         help="the maximal memory used by worker in training",
  79         default="4096Mi",
  80     )
  81     parser.add_argument(
  82         "--worker_pod_priority",
  83         help="the requested priority of worker pod")
  84     parser.add_argument(
  85         "--worker_image", help="docker image for worker", default=None
  86     )
  87     parser.add_argument("--job_name", help="job name", required=True)
  88     parser.add_argument(
  89         "--codec_type",
  90         default="bytes",
  91         choices=["tf_example", "bytes"],
  92         help="Type of codec(tf_example or bytes)",
  93     )
  94     return parser.parse_args()
  95 
  96 
  97 def main():
  98     # TODO: pass port via flags.
  99     PORT = 50001
  100     logger = logging.getLogger("master")
  101     args = _parse_args()
  102     task_q = _make_task_queue(
  103         args.train_data_dir, args.record_per_task, args.num_epoch
  104     )
  105     model_module = load_user_model(args.model_file)
  106     model_inst = model_module.model
  107     build_model(model_inst, model_module.feature_columns())
  108     optimizer = model_module.optimizer()
  109 
  110     server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))
  111     master_pb2_grpc.add_MasterServicer_to_server(
  112         MasterServicer(
  113             logger,
  114             args.grads_to_wait,
  115             args.minibatch_size,
  116             optimizer,
  117             task_q,
  118             init_var=model_inst.trainable_variables,
  119         ),
  120         server,
  121     )
  122     server.add_insecure_port("[::]:{}".format(PORT))
  123     server.start()
  124     logger.warning("Server started at port: %d", PORT)
  125 
  126     if args.num_worker:
  127         master_addr = "%s:%d" % (os.getenv("MY_POD_IP", "localhost"), PORT)
  128         worker_command = ["python"]
  129         worker_args = [
  130             "-m",
  131             "elasticdl.worker.main",
  132             "--model_file",
  133             args.model_file,
  134             "--master_addr",
  135             master_addr,
  136             "--codec_type",
  137             args.codec_type
  138         ]
  139 
  140         worker_manager = WorkerManager(
  141             job_name=args.job_name,
  142             worker_image=args.worker_image,
  143             command=worker_command,
  144             args=worker_args,
  145             namespace="default",
  146             num_worker=args.num_worker,
  147             cpu_request=args.worker_cpu_request,
  148             cpu_limit=args.worker_cpu_limit,
  149             memory_request=args.worker_memory_request,
  150             memory_limit=args.worker_memory_limit,
  151             pod_priority=args.worker_pod_priority,
  152         )
  153         worker_manager.start_workers(restart_policy="Never")
  154 
  155     try:
  156         while True:
  157             if task_q.finished():
  158                 break
  159             time.sleep(30)
  160     except KeyboardInterrupt:
  161         logger.warning("Server stopping")
  162 
  163     if args.num_worker:
  164         # TODO: worker_manager.remove_workers supports synchronized call
  165         worker_manager.remove_workers()
  166         # wait for worker pod to be deleted
  167         max_check_num = 10
  168         for _ in range(max_check_num):
  169             time.sleep(3)
  170             counters = worker_manager.get_counters()
  171             if not counters:
  172                 break
  173     server.stop(0)
  174 
  175 
  176 if __name__ == "__main__":
  177     logging.basicConfig()
  178     main()
  179
  ```
  --- 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()
  
  ```
  
golden_diff: 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
  @@ -163,13 +163,7 @@
       if args.num_worker:
           # TODO: worker_manager.remove_workers supports synchronized call
           worker_manager.remove_workers()
  -        # wait for worker pod to be deleted
  -        max_check_num = 10
  -        for _ in range(max_check_num):
  -            time.sleep(3)
  -            counters = worker_manager.get_counters()
  -            if not counters:
  -                break
  +
       server.stop(0)
  
verification_info: {"golden_diff": "diff --git a/elasticdl/python/elasticdl/master/main.py b/elasticdl/python/elasticdl/master/main.py\n--- a/elasticdl/python/elasticdl/master/main.py\n+++ b/elasticdl/python/elasticdl/master/main.py\n@@ -163,13 +163,7 @@\n     if args.num_worker:\n         # TODO: worker_manager.remove_workers supports synchronized call\n         worker_manager.remove_workers()\n-        # wait for worker pod to be deleted\n-        max_check_num = 10\n-        for _ in range(max_check_num):\n-            time.sleep(3)\n-            counters = worker_manager.get_counters()\n-            if not counters:\n-                break\n+\n     server.stop(0)\n", "issue": "Make master exist when there are no tasks left.\nCurrently, master exists when there are no tasks left AND all workers are gone. It might left hanging if a worker got preempted.\n", "before_files": [{"content": "import logging\nimport time\nimport argparse\nimport os\n\nimport grpc\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom concurrent import futures\nfrom recordio import File\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.master.servicer import MasterServicer\nfrom elasticdl.master.task_queue import _TaskQueue\nfrom elasticdl.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.common.model_helper import load_user_model, build_model\n\n\ndef _make_task_queue(data_dir, record_per_task, num_epoch):\n    f_records = {}\n    for f in os.listdir(data_dir):\n        p = os.path.join(data_dir, f)\n        with File(p, \"r\") as rio:\n            f_records[p] = rio.count()\n    return _TaskQueue(f_records, record_per_task, num_epoch)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Master\")\n    parser.add_argument(\n        \"--model_file\",\n        help=\"Full file path of user defined neural model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--train_data_dir\",\n        help=\"Training data directory. Files should be in RecordIO format\",\n        required=True,\n    )\n    parser.add_argument(\"--record_per_task\", type=int, required=True)\n    parser.add_argument(\"--num_epoch\", type=int, required=True)\n    parser.add_argument(\n        \"--grads_to_wait\",\n        type=int,\n        help=\"Number of gradients to wait before updating model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--minibatch_size\",\n        type=int,\n        help=\"Minibatch size used by workers to compute gradients\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--num_worker\",\n        type=int,\n        help=\"the number of workers used in training\",\n        default=0,\n    )\n    parser.add_argument(\n        \"--worker_cpu_request\",\n        help=\"the minimal cpu required by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_cpu_limit\",\n        help=\"the maximal cpu used by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_memory_request\",\n        help=\"the minimal memory required by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_memory_limit\",\n        help=\"the maximal memory used by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_pod_priority\",\n        help=\"the requested priority of worker pod\")\n    parser.add_argument(\n        \"--worker_image\", help=\"docker image for worker\", default=None\n    )\n    parser.add_argument(\"--job_name\", help=\"job name\", required=True)\n    parser.add_argument(\n        \"--codec_type\",\n        default=\"bytes\",\n        choices=[\"tf_example\", \"bytes\"],\n        help=\"Type of codec(tf_example or bytes)\",\n    )\n    return parser.parse_args()\n\n\ndef main():\n    # TODO: pass port via flags.\n    PORT = 50001\n    logger = logging.getLogger(\"master\")\n    args = _parse_args()\n    task_q = _make_task_queue(\n        args.train_data_dir, args.record_per_task, args.num_epoch\n    )\n    model_module = load_user_model(args.model_file)\n    model_inst = model_module.model\n    build_model(model_inst, model_module.feature_columns())\n    optimizer = model_module.optimizer()\n\n    server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))\n    master_pb2_grpc.add_MasterServicer_to_server(\n        MasterServicer(\n            logger,\n            args.grads_to_wait,\n            args.minibatch_size,\n            optimizer,\n            task_q,\n            init_var=model_inst.trainable_variables,\n        ),\n        server,\n    )\n    server.add_insecure_port(\"[::]:{}\".format(PORT))\n    server.start()\n    logger.warning(\"Server started at port: %d\", PORT)\n\n    if args.num_worker:\n        master_addr = \"%s:%d\" % (os.getenv(\"MY_POD_IP\", \"localhost\"), PORT)\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.worker.main\",\n            \"--model_file\",\n            args.model_file,\n            \"--master_addr\",\n            master_addr,\n            \"--codec_type\",\n            args.codec_type\n        ]\n\n        worker_manager = WorkerManager(\n            job_name=args.job_name,\n            worker_image=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=\"default\",\n            num_worker=args.num_worker,\n            cpu_request=args.worker_cpu_request,\n            cpu_limit=args.worker_cpu_limit,\n            memory_request=args.worker_memory_request,\n            memory_limit=args.worker_memory_limit,\n            pod_priority=args.worker_pod_priority,\n        )\n        worker_manager.start_workers(restart_policy=\"Never\")\n\n    try:\n        while True:\n            if task_q.finished():\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if args.num_worker:\n        # TODO: worker_manager.remove_workers supports synchronized call\n        worker_manager.remove_workers()\n        # wait for worker pod to be deleted\n        max_check_num = 10\n        for _ in range(max_check_num):\n            time.sleep(3)\n            counters = worker_manager.get_counters()\n            if not counters:\n                break\n    server.stop(0)\n\n\nif __name__ == \"__main__\":\n    logging.basicConfig()\n    main()\n", "path": "elasticdl/python/elasticdl/master/main.py"}], "after_files": [{"content": "import logging\nimport time\nimport argparse\nimport os\n\nimport grpc\nimport tensorflow as tf\n\ntf.enable_eager_execution()\n\nfrom concurrent import futures\nfrom recordio import File\nfrom elasticdl.proto import master_pb2_grpc\nfrom elasticdl.master.servicer import MasterServicer\nfrom elasticdl.master.task_queue import _TaskQueue\nfrom elasticdl.master.k8s_worker_manager import WorkerManager\nfrom elasticdl.common.model_helper import load_user_model, build_model\n\n\ndef _make_task_queue(data_dir, record_per_task, num_epoch):\n    f_records = {}\n    for f in os.listdir(data_dir):\n        p = os.path.join(data_dir, f)\n        with File(p, \"r\") as rio:\n            f_records[p] = rio.count()\n    return _TaskQueue(f_records, record_per_task, num_epoch)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(description=\"ElasticDL Master\")\n    parser.add_argument(\n        \"--model_file\",\n        help=\"Full file path of user defined neural model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--train_data_dir\",\n        help=\"Training data directory. Files should be in RecordIO format\",\n        required=True,\n    )\n    parser.add_argument(\"--record_per_task\", type=int, required=True)\n    parser.add_argument(\"--num_epoch\", type=int, required=True)\n    parser.add_argument(\n        \"--grads_to_wait\",\n        type=int,\n        help=\"Number of gradients to wait before updating model\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--minibatch_size\",\n        type=int,\n        help=\"Minibatch size used by workers to compute gradients\",\n        required=True,\n    )\n    parser.add_argument(\n        \"--num_worker\",\n        type=int,\n        help=\"the number of workers used in training\",\n        default=0,\n    )\n    parser.add_argument(\n        \"--worker_cpu_request\",\n        help=\"the minimal cpu required by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_cpu_limit\",\n        help=\"the maximal cpu used by worker in training\",\n        default=\"1000m\",\n    )\n    parser.add_argument(\n        \"--worker_memory_request\",\n        help=\"the minimal memory required by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_memory_limit\",\n        help=\"the maximal memory used by worker in training\",\n        default=\"4096Mi\",\n    )\n    parser.add_argument(\n        \"--worker_pod_priority\",\n        help=\"the requested priority of worker pod\")\n    parser.add_argument(\n        \"--worker_image\", help=\"docker image for worker\", default=None\n    )\n    parser.add_argument(\"--job_name\", help=\"job name\", required=True)\n    parser.add_argument(\n        \"--codec_type\",\n        default=\"bytes\",\n        choices=[\"tf_example\", \"bytes\"],\n        help=\"Type of codec(tf_example or bytes)\",\n    )\n    return parser.parse_args()\n\n\ndef main():\n    # TODO: pass port via flags.\n    PORT = 50001\n    logger = logging.getLogger(\"master\")\n    args = _parse_args()\n    task_q = _make_task_queue(\n        args.train_data_dir, args.record_per_task, args.num_epoch\n    )\n    model_module = load_user_model(args.model_file)\n    model_inst = model_module.model\n    build_model(model_inst, model_module.feature_columns())\n    optimizer = model_module.optimizer()\n\n    server = grpc.server(futures.ThreadPoolExecutor(max_workers=64))\n    master_pb2_grpc.add_MasterServicer_to_server(\n        MasterServicer(\n            logger,\n            args.grads_to_wait,\n            args.minibatch_size,\n            optimizer,\n            task_q,\n            init_var=model_inst.trainable_variables,\n        ),\n        server,\n    )\n    server.add_insecure_port(\"[::]:{}\".format(PORT))\n    server.start()\n    logger.warning(\"Server started at port: %d\", PORT)\n\n    if args.num_worker:\n        master_addr = \"%s:%d\" % (os.getenv(\"MY_POD_IP\", \"localhost\"), PORT)\n        worker_command = [\"python\"]\n        worker_args = [\n            \"-m\",\n            \"elasticdl.worker.main\",\n            \"--model_file\",\n            args.model_file,\n            \"--master_addr\",\n            master_addr,\n            \"--codec_type\",\n            args.codec_type\n        ]\n\n        worker_manager = WorkerManager(\n            job_name=args.job_name,\n            worker_image=args.worker_image,\n            command=worker_command,\n            args=worker_args,\n            namespace=\"default\",\n            num_worker=args.num_worker,\n            cpu_request=args.worker_cpu_request,\n            cpu_limit=args.worker_cpu_limit,\n            memory_request=args.worker_memory_request,\n            memory_limit=args.worker_memory_limit,\n            pod_priority=args.worker_pod_priority,\n        )\n        worker_manager.start_workers(restart_policy=\"Never\")\n\n    try:\n        while True:\n            if task_q.finished():\n                break\n            time.sleep(30)\n    except KeyboardInterrupt:\n        logger.warning(\"Server stopping\")\n\n    if args.num_worker:\n        # TODO: worker_manager.remove_workers supports synchronized call\n        worker_manager.remove_workers()\n\n    server.stop(0)\n\n\nif __name__ == \"__main__\":\n    logging.basicConfig()\n    main()\n", "path": "elasticdl/python/elasticdl/master/main.py"}]}