fix: restrict policy var save for distributed setup

Author: jatinsharechat

tensorflow_recommenders_addons/dynamic_embedding/python/keras/models.py

@@ -98,34 +98,43 @@ def _check_saveable_and_redirect_new_de_dir(hvd_rank=0):
       # Redirect new de_dir
       if hasattr(de_var, 'saveable'):
         de_var.saveable._saver_config.save_path = de_dir
-
+  
+  def _save_de_var(de_var, proc_size=1, proc_rank=0):
+    a2a_emb = de_var._created_in_class
+    if de_var._saveable_object_creator is not None:
+      if not isinstance(de_var.kv_creator.saver, de.FileSystemSaver):
+        # This function only serves FileSystemSaver.
+        return
+      # save optimizer parameters of Dynamic Embedding
+      if include_optimizer is True:
+        de_opt_vars = a2a_emb.optimizer_vars.as_list() if hasattr(
+          a2a_emb.optimizer_vars, "as_list") else a2a_emb.optimizer_vars
+        for de_opt_var in de_opt_vars:
+          de_opt_var.save_to_file_system(dirpath=de_dir,
+                                         proc_size=proc_size,
+                                         proc_rank=proc_rank)
+      if proc_rank == 0:
+        # FileSystemSaver works well at rank 0.
+        return
+      # save Dynamic Embedding Parameters
+      de_var.save_to_file_system(dirpath=de_dir,
+                                 proc_size=proc_size,
+                                 proc_rank=proc_rank)
+  
+  def _maybe_save_restrict_policy_params(var, proc_size=1, proc_rank=0):
+    if var.restrict_policy is not None:
+      de_var = var.restrict_policy._restrict_var
+      _save_de_var(de_var, proc_size=proc_size, proc_rank=proc_rank)
+  
   def _traverse_emb_layers_and_save(proc_size=1, proc_rank=0):
     for var in model.variables:
       if not hasattr(var, "params"):
         continue
       if not hasattr(var.params, "_created_in_class"):
         continue
       de_var = var.params
-      a2a_emb = de_var._created_in_class
-      if de_var._saveable_object_creator is not None:
-        if not isinstance(de_var.kv_creator.saver, de.FileSystemSaver):
-          # This function only serves FileSystemSaver.
-          continue
-        # save optimizer parameters of Dynamic Embedding
-        if include_optimizer is True:
-          de_opt_vars = a2a_emb.optimizer_vars.as_list() if hasattr(
-              a2a_emb.optimizer_vars, "as_list") else a2a_emb.optimizer_vars
-          for de_opt_var in de_opt_vars:
-            de_opt_var.save_to_file_system(dirpath=de_dir,
-                                           proc_size=proc_size,
-                                           proc_rank=proc_rank)
-        if proc_rank == 0:
-          # FileSystemSaver works well at rank 0.
-          continue
-        # save Dynamic Embedding Parameters
-        de_var.save_to_file_system(dirpath=de_dir,
-                                   proc_size=proc_size,
-                                   proc_rank=proc_rank)
+      _save_de_var(de_var, proc_size=proc_size, proc_rank=proc_rank)
+      _maybe_save_restrict_policy_params(var, proc_size=proc_size, proc_rank=proc_rank)
 
   if hvd is None:
     call_original_save_func()