prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Add check for mandatory functions in model files
  Users are required to implement some functions(e.g, `loss()`) in the model file. We can add some checks for the model file to see if these required functions are implemented correctly. Otherwise, if the functions are not implemented (correctly), it won't be found until the job starts to run on k8s cluster.
  
  --- 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/common/model_helper.py`
  Content:
  ```
  1 import importlib.util
  2 import os
  3 
  4 from elasticdl.python.common.log_util import default_logger as logger
  5 from elasticdl.python.worker.prediction_outputs_processor import (
  6     BasePredictionOutputsProcessor,
  7 )
  8 
  9 
  10 def load_module(module_file):
  11     spec = importlib.util.spec_from_file_location(module_file, module_file)
  12     module = importlib.util.module_from_spec(spec)
  13     spec.loader.exec_module(module)
  14     return module
  15 
  16 
  17 # TODO: Discuss whether we need to support default model
  18 # function/class names such as `custom_model()`
  19 # or `CustomModel()`
  20 def load_model_from_module(model_def, model_module, model_params):
  21     model_def_name = model_def.split(".")[-1]
  22     if model_def_name in model_module:
  23         custom_model_name = model_def_name
  24     else:
  25         raise ValueError(
  26             "Cannot find the custom model function/class "
  27             "in model definition files"
  28         )
  29     if model_params:
  30         kvs = model_params.split(",")
  31         model_params_dict = {}
  32         for kv in kvs:
  33             k, v = kv.split("=")
  34             model_params_dict[k] = eval(v)
  35         return model_module[custom_model_name](**model_params_dict)
  36     else:
  37         return model_module[custom_model_name]()
  38 
  39 
  40 def get_module_file_path(model_zoo, spec_key):
  41     """Get the path to module file from model zoo and the spec string.
  42 
  43     For example, if `model_zoo = "model_zoo"` and
  44     `spec_key = "test_module.custom_model"`, the function returns
  45     "model_zoo/test_module.py".
  46     """
  47     return os.path.join(model_zoo, "/".join(spec_key.split(".")[:-1]) + ".py")
  48 
  49 
  50 def _get_spec_value(spec_key, model_zoo, default_module):
  51     """Get the value to the given spec key.
  52 
  53     Notes:
  54 
  55     * If the dot-splitted spec key (e.g. "test_module.custom_model"
  56       is splitted into "test_module" and "custom_model") is of length 1
  57       (e.g. `spec_key` is "custom_model"), return the value in the
  58       specified `default_module`.
  59     * If the spec key does not exist in the module, return `None`.
  60     """
  61     spec_key_items = spec_key.split(".")
  62     spec_key_base = spec_key_items[-1]
  63     if len(spec_key_items) == 1:
  64         spec_key_module = default_module
  65     else:
  66         spec_key_module = load_module(
  67             get_module_file_path(model_zoo, spec_key)
  68         ).__dict__
  69     return (
  70         spec_key_module[spec_key_base]
  71         if spec_key_base in spec_key_module
  72         else None
  73     )
  74 
  75 
  76 def get_model_spec(
  77     model_zoo,
  78     model_def,
  79     model_params,
  80     dataset_fn,
  81     loss,
  82     optimizer,
  83     eval_metrics_fn,
  84     prediction_outputs_processor,
  85 ):
  86     """Get the model spec items in a tuple.
  87 
  88     The model spec tuple contains the following items in order:
  89 
  90     * The model object instantiated with parameters specified
  91       in `model_params`,
  92     * The `dataset_fn`,
  93     * The `loss`,
  94     * The `optimizer`,
  95     * The `eval_metrics_fn`,
  96     * The `prediction_outputs_processor`. Note that it will print
  97       warning if it's not inherited from `BasePredictionOutputsProcessor`.
  98     """
  99     model_def_module_file = get_module_file_path(model_zoo, model_def)
  100     default_module = load_module(model_def_module_file).__dict__
  101     model = load_model_from_module(model_def, default_module, model_params)
  102     prediction_outputs_processor = _get_spec_value(
  103         prediction_outputs_processor, model_zoo, default_module
  104     )
  105     if prediction_outputs_processor and not isinstance(
  106         prediction_outputs_processor, BasePredictionOutputsProcessor
  107     ):
  108         logger.warning(
  109             "prediction_outputs_processor is not "
  110             "inherited from BasePredictionOutputsProcessor. "
  111             "Prediction outputs may not be processed correctly."
  112         )
  113     return (
  114         model,
  115         _get_spec_value(dataset_fn, model_zoo, default_module),
  116         _get_spec_value(loss, model_zoo, default_module),
  117         _get_spec_value(optimizer, model_zoo, default_module),
  118         _get_spec_value(eval_metrics_fn, model_zoo, default_module),
  119         prediction_outputs_processor,
  120     )
  121 
  122 
  123 def save_checkpoint_to_file(pb_model, file_name):
  124     encoded_model = pb_model.SerializeToString()
  125     with open(file_name, "wb") as f:
  126         f.write(encoded_model)
  127 
  128 
  129 def load_from_checkpoint_file(file_name):
  130     from elasticdl.proto import elasticdl_pb2
  131 
  132     pb_model = elasticdl_pb2.Model()
  133     with open(file_name, "rb") as f:
  134         pb_model.ParseFromString(f.read())
  135     return pb_model
  136 
  137 
  138 def find_layer(model, layer_class):
  139     """
  140     Find all layers in model that are instances of layer_class
  141     """
  142     layers = []
  143     for layer in model.layers:
  144         if isinstance(layer, layer_class):
  145             layers.append(layer)
  146         elif hasattr(layer, "layers"):
  147             # search in nested layers
  148             layers += find_layer(layer, layer_class)
  149     return layers
  150
  ```
  --- 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/common/model_helper.py b/elasticdl/python/common/model_helper.py
  --- a/elasticdl/python/common/model_helper.py
  +++ b/elasticdl/python/common/model_helper.py
  @@ -14,9 +14,6 @@
       return module
   
   
  -# TODO: Discuss whether we need to support default model
  -# function/class names such as `custom_model()`
  -# or `CustomModel()`
   def load_model_from_module(model_def, model_module, model_params):
       model_def_name = model_def.split(".")[-1]
       if model_def_name in model_module:
  @@ -47,7 +44,7 @@
       return os.path.join(model_zoo, "/".join(spec_key.split(".")[:-1]) + ".py")
   
   
  -def _get_spec_value(spec_key, model_zoo, default_module):
  +def _get_spec_value(spec_key, model_zoo, default_module, required=False):
       """Get the value to the given spec key.
   
       Notes:
  @@ -66,11 +63,17 @@
           spec_key_module = load_module(
               get_module_file_path(model_zoo, spec_key)
           ).__dict__
  -    return (
  +    spec_value = (
           spec_key_module[spec_key_base]
           if spec_key_base in spec_key_module
           else None
       )
  +    if required and spec_value is None:
  +        raise Exception(
  +            "Missing required spec key %s in the module: %s"
  +            % (spec_key_base, spec_key)
  +        )
  +    return spec_value
   
   
   def get_model_spec(
  @@ -112,10 +115,12 @@
           )
       return (
           model,
  -        _get_spec_value(dataset_fn, model_zoo, default_module),
  -        _get_spec_value(loss, model_zoo, default_module),
  -        _get_spec_value(optimizer, model_zoo, default_module),
  -        _get_spec_value(eval_metrics_fn, model_zoo, default_module),
  +        _get_spec_value(dataset_fn, model_zoo, default_module, required=True),
  +        _get_spec_value(loss, model_zoo, default_module, required=True),
  +        _get_spec_value(optimizer, model_zoo, default_module, required=True),
  +        _get_spec_value(
  +            eval_metrics_fn, model_zoo, default_module, required=True
  +        ),
           prediction_outputs_processor,
       )
  
verification_info: {"golden_diff": "diff --git a/elasticdl/python/common/model_helper.py b/elasticdl/python/common/model_helper.py\n--- a/elasticdl/python/common/model_helper.py\n+++ b/elasticdl/python/common/model_helper.py\n@@ -14,9 +14,6 @@\n     return module\n \n \n-# TODO: Discuss whether we need to support default model\n-# function/class names such as `custom_model()`\n-# or `CustomModel()`\n def load_model_from_module(model_def, model_module, model_params):\n     model_def_name = model_def.split(\".\")[-1]\n     if model_def_name in model_module:\n@@ -47,7 +44,7 @@\n     return os.path.join(model_zoo, \"/\".join(spec_key.split(\".\")[:-1]) + \".py\")\n \n \n-def _get_spec_value(spec_key, model_zoo, default_module):\n+def _get_spec_value(spec_key, model_zoo, default_module, required=False):\n     \"\"\"Get the value to the given spec key.\n \n     Notes:\n@@ -66,11 +63,17 @@\n         spec_key_module = load_module(\n             get_module_file_path(model_zoo, spec_key)\n         ).__dict__\n-    return (\n+    spec_value = (\n         spec_key_module[spec_key_base]\n         if spec_key_base in spec_key_module\n         else None\n     )\n+    if required and spec_value is None:\n+        raise Exception(\n+            \"Missing required spec key %s in the module: %s\"\n+            % (spec_key_base, spec_key)\n+        )\n+    return spec_value\n \n \n def get_model_spec(\n@@ -112,10 +115,12 @@\n         )\n     return (\n         model,\n-        _get_spec_value(dataset_fn, model_zoo, default_module),\n-        _get_spec_value(loss, model_zoo, default_module),\n-        _get_spec_value(optimizer, model_zoo, default_module),\n-        _get_spec_value(eval_metrics_fn, model_zoo, default_module),\n+        _get_spec_value(dataset_fn, model_zoo, default_module, required=True),\n+        _get_spec_value(loss, model_zoo, default_module, required=True),\n+        _get_spec_value(optimizer, model_zoo, default_module, required=True),\n+        _get_spec_value(\n+            eval_metrics_fn, model_zoo, default_module, required=True\n+        ),\n         prediction_outputs_processor,\n     )\n", "issue": "Add check for mandatory functions in model files\nUsers are required to implement some functions(e.g, `loss()`) in the model file. We can add some checks for the model file to see if these required functions are implemented correctly. Otherwise, if the functions are not implemented (correctly), it won't be found until the job starts to run on k8s cluster.\n", "before_files": [{"content": "import importlib.util\nimport os\n\nfrom elasticdl.python.common.log_util import default_logger as logger\nfrom elasticdl.python.worker.prediction_outputs_processor import (\n    BasePredictionOutputsProcessor,\n)\n\n\ndef load_module(module_file):\n    spec = importlib.util.spec_from_file_location(module_file, module_file)\n    module = importlib.util.module_from_spec(spec)\n    spec.loader.exec_module(module)\n    return module\n\n\n# TODO: Discuss whether we need to support default model\n# function/class names such as `custom_model()`\n# or `CustomModel()`\ndef load_model_from_module(model_def, model_module, model_params):\n    model_def_name = model_def.split(\".\")[-1]\n    if model_def_name in model_module:\n        custom_model_name = model_def_name\n    else:\n        raise ValueError(\n            \"Cannot find the custom model function/class \"\n            \"in model definition files\"\n        )\n    if model_params:\n        kvs = model_params.split(\",\")\n        model_params_dict = {}\n        for kv in kvs:\n            k, v = kv.split(\"=\")\n            model_params_dict[k] = eval(v)\n        return model_module[custom_model_name](**model_params_dict)\n    else:\n        return model_module[custom_model_name]()\n\n\ndef get_module_file_path(model_zoo, spec_key):\n    \"\"\"Get the path to module file from model zoo and the spec string.\n\n    For example, if `model_zoo = \"model_zoo\"` and\n    `spec_key = \"test_module.custom_model\"`, the function returns\n    \"model_zoo/test_module.py\".\n    \"\"\"\n    return os.path.join(model_zoo, \"/\".join(spec_key.split(\".\")[:-1]) + \".py\")\n\n\ndef _get_spec_value(spec_key, model_zoo, default_module):\n    \"\"\"Get the value to the given spec key.\n\n    Notes:\n\n    * If the dot-splitted spec key (e.g. \"test_module.custom_model\"\n      is splitted into \"test_module\" and \"custom_model\") is of length 1\n      (e.g. `spec_key` is \"custom_model\"), return the value in the\n      specified `default_module`.\n    * If the spec key does not exist in the module, return `None`.\n    \"\"\"\n    spec_key_items = spec_key.split(\".\")\n    spec_key_base = spec_key_items[-1]\n    if len(spec_key_items) == 1:\n        spec_key_module = default_module\n    else:\n        spec_key_module = load_module(\n            get_module_file_path(model_zoo, spec_key)\n        ).__dict__\n    return (\n        spec_key_module[spec_key_base]\n        if spec_key_base in spec_key_module\n        else None\n    )\n\n\ndef get_model_spec(\n    model_zoo,\n    model_def,\n    model_params,\n    dataset_fn,\n    loss,\n    optimizer,\n    eval_metrics_fn,\n    prediction_outputs_processor,\n):\n    \"\"\"Get the model spec items in a tuple.\n\n    The model spec tuple contains the following items in order:\n\n    * The model object instantiated with parameters specified\n      in `model_params`,\n    * The `dataset_fn`,\n    * The `loss`,\n    * The `optimizer`,\n    * The `eval_metrics_fn`,\n    * The `prediction_outputs_processor`. Note that it will print\n      warning if it's not inherited from `BasePredictionOutputsProcessor`.\n    \"\"\"\n    model_def_module_file = get_module_file_path(model_zoo, model_def)\n    default_module = load_module(model_def_module_file).__dict__\n    model = load_model_from_module(model_def, default_module, model_params)\n    prediction_outputs_processor = _get_spec_value(\n        prediction_outputs_processor, model_zoo, default_module\n    )\n    if prediction_outputs_processor and not isinstance(\n        prediction_outputs_processor, BasePredictionOutputsProcessor\n    ):\n        logger.warning(\n            \"prediction_outputs_processor is not \"\n            \"inherited from BasePredictionOutputsProcessor. \"\n            \"Prediction outputs may not be processed correctly.\"\n        )\n    return (\n        model,\n        _get_spec_value(dataset_fn, model_zoo, default_module),\n        _get_spec_value(loss, model_zoo, default_module),\n        _get_spec_value(optimizer, model_zoo, default_module),\n        _get_spec_value(eval_metrics_fn, model_zoo, default_module),\n        prediction_outputs_processor,\n    )\n\n\ndef save_checkpoint_to_file(pb_model, file_name):\n    encoded_model = pb_model.SerializeToString()\n    with open(file_name, \"wb\") as f:\n        f.write(encoded_model)\n\n\ndef load_from_checkpoint_file(file_name):\n    from elasticdl.proto import elasticdl_pb2\n\n    pb_model = elasticdl_pb2.Model()\n    with open(file_name, \"rb\") as f:\n        pb_model.ParseFromString(f.read())\n    return pb_model\n\n\ndef find_layer(model, layer_class):\n    \"\"\"\n    Find all layers in model that are instances of layer_class\n    \"\"\"\n    layers = []\n    for layer in model.layers:\n        if isinstance(layer, layer_class):\n            layers.append(layer)\n        elif hasattr(layer, \"layers\"):\n            # search in nested layers\n            layers += find_layer(layer, layer_class)\n    return layers\n", "path": "elasticdl/python/common/model_helper.py"}], "after_files": [{"content": "import importlib.util\nimport os\n\nfrom elasticdl.python.common.log_util import default_logger as logger\nfrom elasticdl.python.worker.prediction_outputs_processor import (\n    BasePredictionOutputsProcessor,\n)\n\n\ndef load_module(module_file):\n    spec = importlib.util.spec_from_file_location(module_file, module_file)\n    module = importlib.util.module_from_spec(spec)\n    spec.loader.exec_module(module)\n    return module\n\n\ndef load_model_from_module(model_def, model_module, model_params):\n    model_def_name = model_def.split(\".\")[-1]\n    if model_def_name in model_module:\n        custom_model_name = model_def_name\n    else:\n        raise ValueError(\n            \"Cannot find the custom model function/class \"\n            \"in model definition files\"\n        )\n    if model_params:\n        kvs = model_params.split(\",\")\n        model_params_dict = {}\n        for kv in kvs:\n            k, v = kv.split(\"=\")\n            model_params_dict[k] = eval(v)\n        return model_module[custom_model_name](**model_params_dict)\n    else:\n        return model_module[custom_model_name]()\n\n\ndef get_module_file_path(model_zoo, spec_key):\n    \"\"\"Get the path to module file from model zoo and the spec string.\n\n    For example, if `model_zoo = \"model_zoo\"` and\n    `spec_key = \"test_module.custom_model\"`, the function returns\n    \"model_zoo/test_module.py\".\n    \"\"\"\n    return os.path.join(model_zoo, \"/\".join(spec_key.split(\".\")[:-1]) + \".py\")\n\n\ndef _get_spec_value(spec_key, model_zoo, default_module, required=False):\n    \"\"\"Get the value to the given spec key.\n\n    Notes:\n\n    * If the dot-splitted spec key (e.g. \"test_module.custom_model\"\n      is splitted into \"test_module\" and \"custom_model\") is of length 1\n      (e.g. `spec_key` is \"custom_model\"), return the value in the\n      specified `default_module`.\n    * If the spec key does not exist in the module, return `None`.\n    \"\"\"\n    spec_key_items = spec_key.split(\".\")\n    spec_key_base = spec_key_items[-1]\n    if len(spec_key_items) == 1:\n        spec_key_module = default_module\n    else:\n        spec_key_module = load_module(\n            get_module_file_path(model_zoo, spec_key)\n        ).__dict__\n    spec_value = (\n        spec_key_module[spec_key_base]\n        if spec_key_base in spec_key_module\n        else None\n    )\n    if required and spec_value is None:\n        raise Exception(\n            \"Missing required spec key %s in the module: %s\"\n            % (spec_key_base, spec_key)\n        )\n    return spec_value\n\n\ndef get_model_spec(\n    model_zoo,\n    model_def,\n    model_params,\n    dataset_fn,\n    loss,\n    optimizer,\n    eval_metrics_fn,\n    prediction_outputs_processor,\n):\n    \"\"\"Get the model spec items in a tuple.\n\n    The model spec tuple contains the following items in order:\n\n    * The model object instantiated with parameters specified\n      in `model_params`,\n    * The `dataset_fn`,\n    * The `loss`,\n    * The `optimizer`,\n    * The `eval_metrics_fn`,\n    * The `prediction_outputs_processor`. Note that it will print\n      warning if it's not inherited from `BasePredictionOutputsProcessor`.\n    \"\"\"\n    model_def_module_file = get_module_file_path(model_zoo, model_def)\n    default_module = load_module(model_def_module_file).__dict__\n    model = load_model_from_module(model_def, default_module, model_params)\n    prediction_outputs_processor = _get_spec_value(\n        prediction_outputs_processor, model_zoo, default_module\n    )\n    if prediction_outputs_processor and not isinstance(\n        prediction_outputs_processor, BasePredictionOutputsProcessor\n    ):\n        logger.warning(\n            \"prediction_outputs_processor is not \"\n            \"inherited from BasePredictionOutputsProcessor. \"\n            \"Prediction outputs may not be processed correctly.\"\n        )\n    return (\n        model,\n        _get_spec_value(dataset_fn, model_zoo, default_module, required=True),\n        _get_spec_value(loss, model_zoo, default_module, required=True),\n        _get_spec_value(optimizer, model_zoo, default_module, required=True),\n        _get_spec_value(\n            eval_metrics_fn, model_zoo, default_module, required=True\n        ),\n        prediction_outputs_processor,\n    )\n\n\ndef save_checkpoint_to_file(pb_model, file_name):\n    encoded_model = pb_model.SerializeToString()\n    with open(file_name, \"wb\") as f:\n        f.write(encoded_model)\n\n\ndef load_from_checkpoint_file(file_name):\n    from elasticdl.proto import elasticdl_pb2\n\n    pb_model = elasticdl_pb2.Model()\n    with open(file_name, \"rb\") as f:\n        pb_model.ParseFromString(f.read())\n    return pb_model\n\n\ndef find_layer(model, layer_class):\n    \"\"\"\n    Find all layers in model that are instances of layer_class\n    \"\"\"\n    layers = []\n    for layer in model.layers:\n        if isinstance(layer, layer_class):\n            layers.append(layer)\n        elif hasattr(layer, \"layers\"):\n            # search in nested layers\n            layers += find_layer(layer, layer_class)\n    return layers\n", "path": "elasticdl/python/common/model_helper.py"}]}