We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
"No such file or directory" occasionally happens when the PS save checkpoint.
``` 
File "/elasticdl/python/worker/worker.py", line 579, in report_gradient_to_ps
    res = report_future.result()
  File "/usr/local/lib/python3.6/dist-packages/grpc/_channel.py", line 625, in result
    raise self
grpc._channel._MultiThreadedRendezvous: <_MultiThreadedRendezvous of RPC that terminated with:
        status = StatusCode.UNKNOWN
        details = "Exception calling application: [Errno 2] No such file or directory: '/testdata/ckpt/20200309-19/version-26000/variables-0-of-3.ckpt'"
        debug_error_string = "{"created":"@1583757719.258002676","description":"Error received from peer ipv4:11.149.168.3:2222","file":"src/core/lib/surface/call.cc","file_line":1056,"grpc_message":"Exception calling application: [Errno 2] No such file or directory: '/testdata/elasticdl/kept/dban_lookup/20200309-19/version-26000/variables-0-of-3.ckpt'","grpc_status":2}"
```

--- 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/save_utils.py`
Content:
```
1 import contextlib
2 import os
3 import tempfile
4 
5 import tensorflow as tf
6 
7 from elasticdl.proto import elasticdl_pb2
8 from elasticdl.python.common.hash_utils import int_to_id, string_to_id
9 from elasticdl.python.common.tensor_utils import (
10     pb_to_indexed_slices,
11     pb_to_ndarray,
12 )
13 from elasticdl.python.ps.embedding_table import create_embedding_table
14 from elasticdl.python.ps.parameters import Parameters
15 
16 
17 def save_pb_to_file(pb_obj, file_name):
18     """Save a protobuf object to file"""
19     encoded_model = pb_obj.SerializeToString()
20     with open(file_name, "wb") as f:
21         f.write(encoded_model)
22 
23 
24 def load_pb_from_file(pb_obj, file_name):
25     """Load a protobuf object from a file"""
26     with open(file_name, "rb") as f:
27         pb_obj.ParseFromString(f.read())
28     return pb_obj
29 
30 
31 def _get_params_shard_from_pb(model_pb, shard_index, shard_num):
32     """Get parameters including variables values and embedding table
33     from a model protobuf.
34     Args:
35         model_pb: A Model protobuf instance.
36         shard_index: Model shard index.
37         shard_num: The total number of model shards.
38     Return:
39         non_embedding_vars: A Python dict in which the key is a variable
40             name and the value is a `tf.Variable` object.
41         embedding_table_values: A Python dict in which the key is an embedding
42             table name and the value is a tuple with 2 elements. The value[0]
43             is indices and value[1] is the corresponding embedding vector.
44     """
45     non_embedding_vars = {}
46     embedding_table_values = {}
47 
48     for name, pb in model_pb.dense_parameters.items():
49         if string_to_id(name, shard_num) == shard_index:
50             non_embedding_vars[name] = tf.Variable(
51                 initial_value=pb_to_ndarray(pb), trainable=True
52             )
53     for name, pb in model_pb.embedding_tables.items():
54         embedding_table_values.setdefault(name, ([], []))
55         t = pb_to_indexed_slices(pb)
56         for embedding_id, vector in zip(t.indices, t.values):
57             if int_to_id(embedding_id, shard_num) == shard_index:
58                 embedding_table_values[name][0].append(embedding_id)
59                 embedding_table_values[name][1].append(vector)
60     return non_embedding_vars, embedding_table_values
61 
62 
63 class Checkpoint(object):
64     def __init__(self, version, file):
65         self.version = version
66         self.file = file
67 
68 
69 class CheckpointSaver(object):
70     """Checkpoint Saver implementation"""
71 
72     def __init__(
73         self,
74         checkpoint_dir,
75         checkpoint_steps,
76         keep_checkpoint_max,
77         include_evaluation,
78     ):
79         """
80         Arguments:
81             checkpoint_dir: The directory to store the checkpoint files.
82                             Directory will be created if not exist.
83             checkpoint_steps: Save checkpoint every this many steps.
84             keep_checkpoint_max: The maximum number of recent checkpoint
85                                  files to keep.
86         """
87         self._directory = checkpoint_dir
88         self._steps = checkpoint_steps
89         self._max_versions = keep_checkpoint_max
90         if not self._directory:
91             self._directory = os.getcwd() + "/checkpoint_dir"
92         if self._steps:
93             os.makedirs(self._directory, exist_ok=True)
94         self._checkpoint_list = []
95         self._include_evaluation = include_evaluation
96         self._eval_checkpoint_dir = (
97             tempfile.mkdtemp() if include_evaluation else ""
98         )
99 
100     def _get_checkpoint_file(
101         self, version, is_eval_checkpoint=False, shard_index=0, shard_num=1
102     ):
103         checkpoint_dir = (
104             self._eval_checkpoint_dir
105             if is_eval_checkpoint
106             else self._directory
107         )
108         checkpoint_version_dir = os.path.join(
109             checkpoint_dir, "version-%s" % str(version)
110         )
111         with contextlib.suppress(FileExistsError):
112             os.makedirs(checkpoint_version_dir, exist_ok=True)
113         return "%s/variables-%s-of-%s.ckpt" % (
114             checkpoint_version_dir,
115             str(shard_index),
116             str(shard_num),
117         )
118 
119     def is_enabled(self):
120         """Checkpoint is enabled or not"""
121         return self._steps
122 
123     def need_to_checkpoint(self, version):
124         """Check if the given model version needs to be checkpointed"""
125         return self.is_enabled() and version % self._steps == 0
126 
127     def save(
128         self, version, model, is_eval_checkpoint, shard_index=0, shard_num=1
129     ):
130         """Checkpoint the given model
131 
132         Args:
133             version (int): iteration steps
134             model: a model protobuf instance
135             is_eval_checkpoint (bool): if True, the model will be saved to
136                 a temporary directory.
137             shard_index (int): default 0. The shard index in all
138                 model shard files, e.g. the shard_index is PS instance index
139                 using ParameterServerStrategy.
140             shard_number (int): default 1. The number of model shards,
141                 e.g. shard_number is the number of PS instances using
142                 ParameterServerStrategy.
143         """
144         file = self._get_checkpoint_file(
145             version, is_eval_checkpoint, shard_index, shard_num
146         )
147         save_pb_to_file(model, file)
148         if not is_eval_checkpoint:
149             self._checkpoint_list.append(Checkpoint(version, file))
150             if self._max_versions:
151                 self._delete_old_checkpoints_if_needed()
152 
153     def _delete_old_checkpoints_if_needed(self):
154         """Delete the oldest checkpoint files and keep the number of
155         checkpoints is not beyond max_version.
156         """
157         while len(self._checkpoint_list) > self._max_versions:
158             checkpoint_file = self._checkpoint_list.pop(0).file
159             checkpoint_version_dir = os.path.dirname(checkpoint_file)
160             with contextlib.suppress(FileNotFoundError):
161                 os.remove(checkpoint_file)
162                 # Remove the directory if empty
163                 if not os.listdir(checkpoint_version_dir):
164                     os.rmdir(checkpoint_version_dir)
165 
166     def get_latest_checkpoint_version(self):
167         """Get the latest checkpointed model version"""
168         if not self._checkpoint_list:
169             raise RuntimeError("No model checkpoint available")
170         return self._checkpoint_list[-1].version
171 
172     @staticmethod
173     def get_valid_lastest_version_dir(checkpoint_dir):
174         """Get the valid and lastest version checkpoint directory"""
175         if not checkpoint_dir or not os.path.exists(checkpoint_dir):
176             return None
177 
178         version_folders = os.listdir(checkpoint_dir)
179         if not version_folders:
180             return None
181         version_num = [int(v.split("-")[-1]) for v in version_folders]
182         version_folder_pairs = sorted(
183             zip(version_num, version_folders), reverse=True
184         )
185         for version, folder in version_folder_pairs:
186             folder_dir = os.path.join(checkpoint_dir, folder)
187             if CheckpointSaver.check_checkpoint_valid(folder_dir):
188                 return folder_dir
189         return None
190 
191     @staticmethod
192     def check_checkpoint_valid(checkpoint_dir):
193         """Check whether the checkpoint directory is valid. The filename template
194         in the checkpoint directory like "variables-{i}-of-{N}.ckpt". We will
195         parse any filename to get N which is the total number of parameters
196         shards. It is valid if the number of files in the directory N.
197         """
198         if not os.path.exists(checkpoint_dir):
199             return False
200 
201         shard_files = os.listdir(checkpoint_dir)
202         if not shard_files:
203             return False
204 
205         shard_file_prefix = shard_files[0].split(".")[0]
206         expected_shard_num = int(shard_file_prefix.split("-")[-1])
207         return expected_shard_num == len(shard_files)
208 
209     @staticmethod
210     def restore_params_from_checkpoint(checkpoint_dir, shard_index, shard_num):
211         """Restore a shard parameters from the checkpoint directory.
212         If shard_num=1, a entire model parameters will be restored.
213 
214         Args:
215             checkpoint_dir: a directory with checkpoint files.
216             shard_index: Model shard index, e.g. the PS instance index
217                 using ParameterServerStrategy with multiple PS instances.
218             shard_num: The total number of model shards, e.g. the total PS
219                 instancecount using ParameterServerStrategy with multiple
220                 PS instances.
221 
222         Return:
223             parameters: A Parameter object which contains model version,
224                 non-embedding parameters and embedding tables for the
225                 PS instance with ps_id.
226         """
227 
228         variable_shard_files = os.listdir(checkpoint_dir)
229         non_embedding_vars = {}
230         embedding_tables = {}
231         version = None
232         for shard_file in variable_shard_files:
233             shard_file_path = os.path.join(checkpoint_dir, shard_file)
234             model_pb = elasticdl_pb2.Model()
235             model_pb = load_pb_from_file(model_pb, shard_file_path)
236             if version is None:
237                 version = model_pb.version
238             elif version != model_pb.version:
239                 raise ValueError(
240                     "The versions in model shards are not consistent"
241                 )
242 
243             for embedding_info_pb in model_pb.embedding_table_infos:
244                 embedding_table = create_embedding_table(embedding_info_pb)
245                 embedding_tables.setdefault(
246                     embedding_table.name, embedding_table
247                 )
248 
249             (
250                 shard_non_embedding_vars,
251                 shard_embedding_table_values,
252             ) = _get_params_shard_from_pb(model_pb, shard_index, shard_num)
253 
254             non_embedding_vars.update(shard_non_embedding_vars)
255             for name, pair in shard_embedding_table_values.items():
256                 embedding_tables[name].set(pair[0], pair[1])
257 
258         parameters = Parameters()
259         parameters.non_embedding_params.update(non_embedding_vars)
260         parameters.embedding_params.update(embedding_tables)
261         parameters.version = version
262         return parameters
263 
264     @staticmethod
265     def get_version_from_checkpoint(checkpoint_dir):
266         """Get model version from the checkpoint. There may be several shard
267         files in the checkpoint directory. The model versions of shard files
268         are same, so we only need to read one shard file to get model version.
269         """
270         variable_shard_files = os.listdir(checkpoint_dir)
271         shard_file_path = os.path.join(checkpoint_dir, variable_shard_files[0])
272         model_pb = elasticdl_pb2.Model()
273         model_pb = load_pb_from_file(model_pb, shard_file_path)
274         return model_pb.version
275
```
--- 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/common/save_utils.py b/elasticdl/python/common/save_utils.py
--- a/elasticdl/python/common/save_utils.py
+++ b/elasticdl/python/common/save_utils.py
@@ -1,5 +1,6 @@
 import contextlib
 import os
+import shutil
 import tempfile
 
 import tensorflow as tf
@@ -91,7 +92,7 @@
             self._directory = os.getcwd() + "/checkpoint_dir"
         if self._steps:
             os.makedirs(self._directory, exist_ok=True)
-        self._checkpoint_list = []
+        self._checkpoint_dir_list = []
         self._include_evaluation = include_evaluation
         self._eval_checkpoint_dir = (
             tempfile.mkdtemp() if include_evaluation else ""
@@ -141,12 +142,12 @@
                 e.g. shard_number is the number of PS instances using
                 ParameterServerStrategy.
         """
-        file = self._get_checkpoint_file(
+        filename = self._get_checkpoint_file(
             version, is_eval_checkpoint, shard_index, shard_num
         )
-        save_pb_to_file(model, file)
+        save_pb_to_file(model, filename)
         if not is_eval_checkpoint:
-            self._checkpoint_list.append(Checkpoint(version, file))
+            self._checkpoint_dir_list.append(os.path.dirname(filename))
             if self._max_versions:
                 self._delete_old_checkpoints_if_needed()
 
@@ -154,20 +155,16 @@
         """Delete the oldest checkpoint files and keep the number of
         checkpoints is not beyond max_version.
         """
-        while len(self._checkpoint_list) > self._max_versions:
-            checkpoint_file = self._checkpoint_list.pop(0).file
-            checkpoint_version_dir = os.path.dirname(checkpoint_file)
-            with contextlib.suppress(FileNotFoundError):
-                os.remove(checkpoint_file)
-                # Remove the directory if empty
-                if not os.listdir(checkpoint_version_dir):
-                    os.rmdir(checkpoint_version_dir)
-
-    def get_latest_checkpoint_version(self):
-        """Get the latest checkpointed model version"""
-        if not self._checkpoint_list:
-            raise RuntimeError("No model checkpoint available")
-        return self._checkpoint_list[-1].version
+        if len(self._checkpoint_dir_list) > self._max_versions:
+            old_version_dir = self._checkpoint_dir_list[0]
+
+            # Some PS instances have not saved checkpoint shard files of
+            # the version if invalid. And the slowest PS will remove the
+            # old version checkpoint.
+            if self.check_checkpoint_valid(old_version_dir):
+                self._checkpoint_dir_list.pop(0)
+                with contextlib.suppress(FileNotFoundError):
+                    shutil.rmtree(old_version_dir)
 
     @staticmethod
     def get_valid_lastest_version_dir(checkpoint_dir):

{"golden_diff": "diff --git a/elasticdl/python/common/save_utils.py b/elasticdl/python/common/save_utils.py\n--- a/elasticdl/python/common/save_utils.py\n+++ b/elasticdl/python/common/save_utils.py\n@@ -1,5 +1,6 @@\n import contextlib\n import os\n+import shutil\n import tempfile\n \n import tensorflow as tf\n@@ -91,7 +92,7 @@\n             self._directory = os.getcwd() + \"/checkpoint_dir\"\n         if self._steps:\n             os.makedirs(self._directory, exist_ok=True)\n-        self._checkpoint_list = []\n+        self._checkpoint_dir_list = []\n         self._include_evaluation = include_evaluation\n         self._eval_checkpoint_dir = (\n             tempfile.mkdtemp() if include_evaluation else \"\"\n@@ -141,12 +142,12 @@\n                 e.g. shard_number is the number of PS instances using\n                 ParameterServerStrategy.\n         \"\"\"\n-        file = self._get_checkpoint_file(\n+        filename = self._get_checkpoint_file(\n             version, is_eval_checkpoint, shard_index, shard_num\n         )\n-        save_pb_to_file(model, file)\n+        save_pb_to_file(model, filename)\n         if not is_eval_checkpoint:\n-            self._checkpoint_list.append(Checkpoint(version, file))\n+            self._checkpoint_dir_list.append(os.path.dirname(filename))\n             if self._max_versions:\n                 self._delete_old_checkpoints_if_needed()\n \n@@ -154,20 +155,16 @@\n         \"\"\"Delete the oldest checkpoint files and keep the number of\n         checkpoints is not beyond max_version.\n         \"\"\"\n-        while len(self._checkpoint_list) > self._max_versions:\n-            checkpoint_file = self._checkpoint_list.pop(0).file\n-            checkpoint_version_dir = os.path.dirname(checkpoint_file)\n-            with contextlib.suppress(FileNotFoundError):\n-                os.remove(checkpoint_file)\n-                # Remove the directory if empty\n-                if not os.listdir(checkpoint_version_dir):\n-                    os.rmdir(checkpoint_version_dir)\n-\n-    def get_latest_checkpoint_version(self):\n-        \"\"\"Get the latest checkpointed model version\"\"\"\n-        if not self._checkpoint_list:\n-            raise RuntimeError(\"No model checkpoint available\")\n-        return self._checkpoint_list[-1].version\n+        if len(self._checkpoint_dir_list) > self._max_versions:\n+            old_version_dir = self._checkpoint_dir_list[0]\n+\n+            # Some PS instances have not saved checkpoint shard files of\n+            # the version if invalid. And the slowest PS will remove the\n+            # old version checkpoint.\n+            if self.check_checkpoint_valid(old_version_dir):\n+                self._checkpoint_dir_list.pop(0)\n+                with contextlib.suppress(FileNotFoundError):\n+                    shutil.rmtree(old_version_dir)\n \n     @staticmethod\n     def get_valid_lastest_version_dir(checkpoint_dir):\n", "issue": "\"No such file or directory\" occasionally happens when the PS save checkpoint.\n``` \r\nFile \"/elasticdl/python/worker/worker.py\", line 579, in report_gradient_to_ps\r\n    res = report_future.result()\r\n  File \"/usr/local/lib/python3.6/dist-packages/grpc/_channel.py\", line 625, in result\r\n    raise self\r\ngrpc._channel._MultiThreadedRendezvous: <_MultiThreadedRendezvous of RPC that terminated with:\r\n        status = StatusCode.UNKNOWN\r\n        details = \"Exception calling application: [Errno 2] No such file or directory: '/testdata/ckpt/20200309-19/version-26000/variables-0-of-3.ckpt'\"\r\n        debug_error_string = \"{\"created\":\"@1583757719.258002676\",\"description\":\"Error received from peer ipv4:11.149.168.3:2222\",\"file\":\"src/core/lib/surface/call.cc\",\"file_line\":1056,\"grpc_message\":\"Exception calling application: [Errno 2] No such file or directory: '/testdata/elasticdl/kept/dban_lookup/20200309-19/version-26000/variables-0-of-3.ckpt'\",\"grpc_status\":2}\"\r\n```\n", "before_files": [{"content": "import contextlib\nimport os\nimport tempfile\n\nimport tensorflow as tf\n\nfrom elasticdl.proto import elasticdl_pb2\nfrom elasticdl.python.common.hash_utils import int_to_id, string_to_id\nfrom elasticdl.python.common.tensor_utils import (\n    pb_to_indexed_slices,\n    pb_to_ndarray,\n)\nfrom elasticdl.python.ps.embedding_table import create_embedding_table\nfrom elasticdl.python.ps.parameters import Parameters\n\n\ndef save_pb_to_file(pb_obj, file_name):\n    \"\"\"Save a protobuf object to file\"\"\"\n    encoded_model = pb_obj.SerializeToString()\n    with open(file_name, \"wb\") as f:\n        f.write(encoded_model)\n\n\ndef load_pb_from_file(pb_obj, file_name):\n    \"\"\"Load a protobuf object from a file\"\"\"\n    with open(file_name, \"rb\") as f:\n        pb_obj.ParseFromString(f.read())\n    return pb_obj\n\n\ndef _get_params_shard_from_pb(model_pb, shard_index, shard_num):\n    \"\"\"Get parameters including variables values and embedding table\n    from a model protobuf.\n    Args:\n        model_pb: A Model protobuf instance.\n        shard_index: Model shard index.\n        shard_num: The total number of model shards.\n    Return:\n        non_embedding_vars: A Python dict in which the key is a variable\n            name and the value is a `tf.Variable` object.\n        embedding_table_values: A Python dict in which the key is an embedding\n            table name and the value is a tuple with 2 elements. The value[0]\n            is indices and value[1] is the corresponding embedding vector.\n    \"\"\"\n    non_embedding_vars = {}\n    embedding_table_values = {}\n\n    for name, pb in model_pb.dense_parameters.items():\n        if string_to_id(name, shard_num) == shard_index:\n            non_embedding_vars[name] = tf.Variable(\n                initial_value=pb_to_ndarray(pb), trainable=True\n            )\n    for name, pb in model_pb.embedding_tables.items():\n        embedding_table_values.setdefault(name, ([], []))\n        t = pb_to_indexed_slices(pb)\n        for embedding_id, vector in zip(t.indices, t.values):\n            if int_to_id(embedding_id, shard_num) == shard_index:\n                embedding_table_values[name][0].append(embedding_id)\n                embedding_table_values[name][1].append(vector)\n    return non_embedding_vars, embedding_table_values\n\n\nclass Checkpoint(object):\n    def __init__(self, version, file):\n        self.version = version\n        self.file = file\n\n\nclass CheckpointSaver(object):\n    \"\"\"Checkpoint Saver implementation\"\"\"\n\n    def __init__(\n        self,\n        checkpoint_dir,\n        checkpoint_steps,\n        keep_checkpoint_max,\n        include_evaluation,\n    ):\n        \"\"\"\n        Arguments:\n            checkpoint_dir: The directory to store the checkpoint files.\n                            Directory will be created if not exist.\n            checkpoint_steps: Save checkpoint every this many steps.\n            keep_checkpoint_max: The maximum number of recent checkpoint\n                                 files to keep.\n        \"\"\"\n        self._directory = checkpoint_dir\n        self._steps = checkpoint_steps\n        self._max_versions = keep_checkpoint_max\n        if not self._directory:\n            self._directory = os.getcwd() + \"/checkpoint_dir\"\n        if self._steps:\n            os.makedirs(self._directory, exist_ok=True)\n        self._checkpoint_list = []\n        self._include_evaluation = include_evaluation\n        self._eval_checkpoint_dir = (\n            tempfile.mkdtemp() if include_evaluation else \"\"\n        )\n\n    def _get_checkpoint_file(\n        self, version, is_eval_checkpoint=False, shard_index=0, shard_num=1\n    ):\n        checkpoint_dir = (\n            self._eval_checkpoint_dir\n            if is_eval_checkpoint\n            else self._directory\n        )\n        checkpoint_version_dir = os.path.join(\n            checkpoint_dir, \"version-%s\" % str(version)\n        )\n        with contextlib.suppress(FileExistsError):\n            os.makedirs(checkpoint_version_dir, exist_ok=True)\n        return \"%s/variables-%s-of-%s.ckpt\" % (\n            checkpoint_version_dir,\n            str(shard_index),\n            str(shard_num),\n        )\n\n    def is_enabled(self):\n        \"\"\"Checkpoint is enabled or not\"\"\"\n        return self._steps\n\n    def need_to_checkpoint(self, version):\n        \"\"\"Check if the given model version needs to be checkpointed\"\"\"\n        return self.is_enabled() and version % self._steps == 0\n\n    def save(\n        self, version, model, is_eval_checkpoint, shard_index=0, shard_num=1\n    ):\n        \"\"\"Checkpoint the given model\n\n        Args:\n            version (int): iteration steps\n            model: a model protobuf instance\n            is_eval_checkpoint (bool): if True, the model will be saved to\n                a temporary directory.\n            shard_index (int): default 0. The shard index in all\n                model shard files, e.g. the shard_index is PS instance index\n                using ParameterServerStrategy.\n            shard_number (int): default 1. The number of model shards,\n                e.g. shard_number is the number of PS instances using\n                ParameterServerStrategy.\n        \"\"\"\n        file = self._get_checkpoint_file(\n            version, is_eval_checkpoint, shard_index, shard_num\n        )\n        save_pb_to_file(model, file)\n        if not is_eval_checkpoint:\n            self._checkpoint_list.append(Checkpoint(version, file))\n            if self._max_versions:\n                self._delete_old_checkpoints_if_needed()\n\n    def _delete_old_checkpoints_if_needed(self):\n        \"\"\"Delete the oldest checkpoint files and keep the number of\n        checkpoints is not beyond max_version.\n        \"\"\"\n        while len(self._checkpoint_list) > self._max_versions:\n            checkpoint_file = self._checkpoint_list.pop(0).file\n            checkpoint_version_dir = os.path.dirname(checkpoint_file)\n            with contextlib.suppress(FileNotFoundError):\n                os.remove(checkpoint_file)\n                # Remove the directory if empty\n                if not os.listdir(checkpoint_version_dir):\n                    os.rmdir(checkpoint_version_dir)\n\n    def get_latest_checkpoint_version(self):\n        \"\"\"Get the latest checkpointed model version\"\"\"\n        if not self._checkpoint_list:\n            raise RuntimeError(\"No model checkpoint available\")\n        return self._checkpoint_list[-1].version\n\n    @staticmethod\n    def get_valid_lastest_version_dir(checkpoint_dir):\n        \"\"\"Get the valid and lastest version checkpoint directory\"\"\"\n        if not checkpoint_dir or not os.path.exists(checkpoint_dir):\n            return None\n\n        version_folders = os.listdir(checkpoint_dir)\n        if not version_folders:\n            return None\n        version_num = [int(v.split(\"-\")[-1]) for v in version_folders]\n        version_folder_pairs = sorted(\n            zip(version_num, version_folders), reverse=True\n        )\n        for version, folder in version_folder_pairs:\n            folder_dir = os.path.join(checkpoint_dir, folder)\n            if CheckpointSaver.check_checkpoint_valid(folder_dir):\n                return folder_dir\n        return None\n\n    @staticmethod\n    def check_checkpoint_valid(checkpoint_dir):\n        \"\"\"Check whether the checkpoint directory is valid. The filename template\n        in the checkpoint directory like \"variables-{i}-of-{N}.ckpt\". We will\n        parse any filename to get N which is the total number of parameters\n        shards. It is valid if the number of files in the directory N.\n        \"\"\"\n        if not os.path.exists(checkpoint_dir):\n            return False\n\n        shard_files = os.listdir(checkpoint_dir)\n        if not shard_files:\n            return False\n\n        shard_file_prefix = shard_files[0].split(\".\")[0]\n        expected_shard_num = int(shard_file_prefix.split(\"-\")[-1])\n        return expected_shard_num == len(shard_files)\n\n    @staticmethod\n    def restore_params_from_checkpoint(checkpoint_dir, shard_index, shard_num):\n        \"\"\"Restore a shard parameters from the checkpoint directory.\n        If shard_num=1, a entire model parameters will be restored.\n\n        Args:\n            checkpoint_dir: a directory with checkpoint files.\n            shard_index: Model shard index, e.g. the PS instance index\n                using ParameterServerStrategy with multiple PS instances.\n            shard_num: The total number of model shards, e.g. the total PS\n                instancecount using ParameterServerStrategy with multiple\n                PS instances.\n\n        Return:\n            parameters: A Parameter object which contains model version,\n                non-embedding parameters and embedding tables for the\n                PS instance with ps_id.\n        \"\"\"\n\n        variable_shard_files = os.listdir(checkpoint_dir)\n        non_embedding_vars = {}\n        embedding_tables = {}\n        version = None\n        for shard_file in variable_shard_files:\n            shard_file_path = os.path.join(checkpoint_dir, shard_file)\n            model_pb = elasticdl_pb2.Model()\n            model_pb = load_pb_from_file(model_pb, shard_file_path)\n            if version is None:\n                version = model_pb.version\n            elif version != model_pb.version:\n                raise ValueError(\n                    \"The versions in model shards are not consistent\"\n                )\n\n            for embedding_info_pb in model_pb.embedding_table_infos:\n                embedding_table = create_embedding_table(embedding_info_pb)\n                embedding_tables.setdefault(\n                    embedding_table.name, embedding_table\n                )\n\n            (\n                shard_non_embedding_vars,\n                shard_embedding_table_values,\n            ) = _get_params_shard_from_pb(model_pb, shard_index, shard_num)\n\n            non_embedding_vars.update(shard_non_embedding_vars)\n            for name, pair in shard_embedding_table_values.items():\n                embedding_tables[name].set(pair[0], pair[1])\n\n        parameters = Parameters()\n        parameters.non_embedding_params.update(non_embedding_vars)\n        parameters.embedding_params.update(embedding_tables)\n        parameters.version = version\n        return parameters\n\n    @staticmethod\n    def get_version_from_checkpoint(checkpoint_dir):\n        \"\"\"Get model version from the checkpoint. There may be several shard\n        files in the checkpoint directory. The model versions of shard files\n        are same, so we only need to read one shard file to get model version.\n        \"\"\"\n        variable_shard_files = os.listdir(checkpoint_dir)\n        shard_file_path = os.path.join(checkpoint_dir, variable_shard_files[0])\n        model_pb = elasticdl_pb2.Model()\n        model_pb = load_pb_from_file(model_pb, shard_file_path)\n        return model_pb.version\n", "path": "elasticdl/python/common/save_utils.py"}], "after_files": [{"content": "import contextlib\nimport os\nimport shutil\nimport tempfile\n\nimport tensorflow as tf\n\nfrom elasticdl.proto import elasticdl_pb2\nfrom elasticdl.python.common.hash_utils import int_to_id, string_to_id\nfrom elasticdl.python.common.tensor_utils import (\n    pb_to_indexed_slices,\n    pb_to_ndarray,\n)\nfrom elasticdl.python.ps.embedding_table import create_embedding_table\nfrom elasticdl.python.ps.parameters import Parameters\n\n\ndef save_pb_to_file(pb_obj, file_name):\n    \"\"\"Save a protobuf object to file\"\"\"\n    encoded_model = pb_obj.SerializeToString()\n    with open(file_name, \"wb\") as f:\n        f.write(encoded_model)\n\n\ndef load_pb_from_file(pb_obj, file_name):\n    \"\"\"Load a protobuf object from a file\"\"\"\n    with open(file_name, \"rb\") as f:\n        pb_obj.ParseFromString(f.read())\n    return pb_obj\n\n\ndef _get_params_shard_from_pb(model_pb, shard_index, shard_num):\n    \"\"\"Get parameters including variables values and embedding table\n    from a model protobuf.\n    Args:\n        model_pb: A Model protobuf instance.\n        shard_index: Model shard index.\n        shard_num: The total number of model shards.\n    Return:\n        non_embedding_vars: A Python dict in which the key is a variable\n            name and the value is a `tf.Variable` object.\n        embedding_table_values: A Python dict in which the key is an embedding\n            table name and the value is a tuple with 2 elements. The value[0]\n            is indices and value[1] is the corresponding embedding vector.\n    \"\"\"\n    non_embedding_vars = {}\n    embedding_table_values = {}\n\n    for name, pb in model_pb.dense_parameters.items():\n        if string_to_id(name, shard_num) == shard_index:\n            non_embedding_vars[name] = tf.Variable(\n                initial_value=pb_to_ndarray(pb), trainable=True\n            )\n    for name, pb in model_pb.embedding_tables.items():\n        embedding_table_values.setdefault(name, ([], []))\n        t = pb_to_indexed_slices(pb)\n        for embedding_id, vector in zip(t.indices, t.values):\n            if int_to_id(embedding_id, shard_num) == shard_index:\n                embedding_table_values[name][0].append(embedding_id)\n                embedding_table_values[name][1].append(vector)\n    return non_embedding_vars, embedding_table_values\n\n\nclass Checkpoint(object):\n    def __init__(self, version, file):\n        self.version = version\n        self.file = file\n\n\nclass CheckpointSaver(object):\n    \"\"\"Checkpoint Saver implementation\"\"\"\n\n    def __init__(\n        self,\n        checkpoint_dir,\n        checkpoint_steps,\n        keep_checkpoint_max,\n        include_evaluation,\n    ):\n        \"\"\"\n        Arguments:\n            checkpoint_dir: The directory to store the checkpoint files.\n                            Directory will be created if not exist.\n            checkpoint_steps: Save checkpoint every this many steps.\n            keep_checkpoint_max: The maximum number of recent checkpoint\n                                 files to keep.\n        \"\"\"\n        self._directory = checkpoint_dir\n        self._steps = checkpoint_steps\n        self._max_versions = keep_checkpoint_max\n        if not self._directory:\n            self._directory = os.getcwd() + \"/checkpoint_dir\"\n        if self._steps:\n            os.makedirs(self._directory, exist_ok=True)\n        self._checkpoint_dir_list = []\n        self._include_evaluation = include_evaluation\n        self._eval_checkpoint_dir = (\n            tempfile.mkdtemp() if include_evaluation else \"\"\n        )\n\n    def _get_checkpoint_file(\n        self, version, is_eval_checkpoint=False, shard_index=0, shard_num=1\n    ):\n        checkpoint_dir = (\n            self._eval_checkpoint_dir\n            if is_eval_checkpoint\n            else self._directory\n        )\n        checkpoint_version_dir = os.path.join(\n            checkpoint_dir, \"version-%s\" % str(version)\n        )\n        with contextlib.suppress(FileExistsError):\n            os.makedirs(checkpoint_version_dir, exist_ok=True)\n        return \"%s/variables-%s-of-%s.ckpt\" % (\n            checkpoint_version_dir,\n            str(shard_index),\n            str(shard_num),\n        )\n\n    def is_enabled(self):\n        \"\"\"Checkpoint is enabled or not\"\"\"\n        return self._steps\n\n    def need_to_checkpoint(self, version):\n        \"\"\"Check if the given model version needs to be checkpointed\"\"\"\n        return self.is_enabled() and version % self._steps == 0\n\n    def save(\n        self, version, model, is_eval_checkpoint, shard_index=0, shard_num=1\n    ):\n        \"\"\"Checkpoint the given model\n\n        Args:\n            version (int): iteration steps\n            model: a model protobuf instance\n            is_eval_checkpoint (bool): if True, the model will be saved to\n                a temporary directory.\n            shard_index (int): default 0. The shard index in all\n                model shard files, e.g. the shard_index is PS instance index\n                using ParameterServerStrategy.\n            shard_number (int): default 1. The number of model shards,\n                e.g. shard_number is the number of PS instances using\n                ParameterServerStrategy.\n        \"\"\"\n        filename = self._get_checkpoint_file(\n            version, is_eval_checkpoint, shard_index, shard_num\n        )\n        save_pb_to_file(model, filename)\n        if not is_eval_checkpoint:\n            self._checkpoint_dir_list.append(os.path.dirname(filename))\n            if self._max_versions:\n                self._delete_old_checkpoints_if_needed()\n\n    def _delete_old_checkpoints_if_needed(self):\n        \"\"\"Delete the oldest checkpoint files and keep the number of\n        checkpoints is not beyond max_version.\n        \"\"\"\n        if len(self._checkpoint_dir_list) > self._max_versions:\n            old_version_dir = self._checkpoint_dir_list[0]\n\n            # Some PS instances have not saved checkpoint shard files of\n            # the version if invalid. And the slowest PS will remove the\n            # old version checkpoint.\n            if self.check_checkpoint_valid(old_version_dir):\n                self._checkpoint_dir_list.pop(0)\n                with contextlib.suppress(FileNotFoundError):\n                    shutil.rmtree(old_version_dir)\n\n    @staticmethod\n    def get_valid_lastest_version_dir(checkpoint_dir):\n        \"\"\"Get the valid and lastest version checkpoint directory\"\"\"\n        if not checkpoint_dir or not os.path.exists(checkpoint_dir):\n            return None\n\n        version_folders = os.listdir(checkpoint_dir)\n        if not version_folders:\n            return None\n        version_num = [int(v.split(\"-\")[-1]) for v in version_folders]\n        version_folder_pairs = sorted(\n            zip(version_num, version_folders), reverse=True\n        )\n        for version, folder in version_folder_pairs:\n            folder_dir = os.path.join(checkpoint_dir, folder)\n            if CheckpointSaver.check_checkpoint_valid(folder_dir):\n                return folder_dir\n        return None\n\n    @staticmethod\n    def check_checkpoint_valid(checkpoint_dir):\n        \"\"\"Check whether the checkpoint directory is valid. The filename template\n        in the checkpoint directory like \"variables-{i}-of-{N}.ckpt\". We will\n        parse any filename to get N which is the total number of parameters\n        shards. It is valid if the number of files in the directory N.\n        \"\"\"\n        if not os.path.exists(checkpoint_dir):\n            return False\n\n        shard_files = os.listdir(checkpoint_dir)\n        if not shard_files:\n            return False\n\n        shard_file_prefix = shard_files[0].split(\".\")[0]\n        expected_shard_num = int(shard_file_prefix.split(\"-\")[-1])\n        return expected_shard_num == len(shard_files)\n\n    @staticmethod\n    def restore_params_from_checkpoint(checkpoint_dir, shard_index, shard_num):\n        \"\"\"Restore a shard parameters from the checkpoint directory.\n        If shard_num=1, a entire model parameters will be restored.\n\n        Args:\n            checkpoint_dir: a directory with checkpoint files.\n            shard_index: Model shard index, e.g. the PS instance index\n                using ParameterServerStrategy with multiple PS instances.\n            shard_num: The total number of model shards, e.g. the total PS\n                instancecount using ParameterServerStrategy with multiple\n                PS instances.\n\n        Return:\n            parameters: A Parameter object which contains model version,\n                non-embedding parameters and embedding tables for the\n                PS instance with ps_id.\n        \"\"\"\n\n        variable_shard_files = os.listdir(checkpoint_dir)\n        non_embedding_vars = {}\n        embedding_tables = {}\n        version = None\n        for shard_file in variable_shard_files:\n            shard_file_path = os.path.join(checkpoint_dir, shard_file)\n            model_pb = elasticdl_pb2.Model()\n            model_pb = load_pb_from_file(model_pb, shard_file_path)\n            if version is None:\n                version = model_pb.version\n            elif version != model_pb.version:\n                raise ValueError(\n                    \"The versions in model shards are not consistent\"\n                )\n\n            for embedding_info_pb in model_pb.embedding_table_infos:\n                embedding_table = create_embedding_table(embedding_info_pb)\n                embedding_tables.setdefault(\n                    embedding_table.name, embedding_table\n                )\n\n            (\n                shard_non_embedding_vars,\n                shard_embedding_table_values,\n            ) = _get_params_shard_from_pb(model_pb, shard_index, shard_num)\n\n            non_embedding_vars.update(shard_non_embedding_vars)\n            for name, pair in shard_embedding_table_values.items():\n                embedding_tables[name].set(pair[0], pair[1])\n\n        parameters = Parameters()\n        parameters.non_embedding_params.update(non_embedding_vars)\n        parameters.embedding_params.update(embedding_tables)\n        parameters.version = version\n        return parameters\n\n    @staticmethod\n    def get_version_from_checkpoint(checkpoint_dir):\n        \"\"\"Get model version from the checkpoint. There may be several shard\n        files in the checkpoint directory. The model versions of shard files\n        are same, so we only need to read one shard file to get model version.\n        \"\"\"\n        variable_shard_files = os.listdir(checkpoint_dir)\n        shard_file_path = os.path.join(checkpoint_dir, variable_shard_files[0])\n        model_pb = elasticdl_pb2.Model()\n        model_pb = load_pb_from_file(model_pb, shard_file_path)\n        return model_pb.version\n", "path": "elasticdl/python/common/save_utils.py"}]}