Skip to content

Commit 214c6b3

Browse files
authored
Remove pickle, minor variable rename (#1924)
Removing pickle usage to avoid security issue.
1 parent c54dbd3 commit 214c6b3

File tree

1 file changed

+15
-16
lines changed

1 file changed

+15
-16
lines changed

apex/contrib/sparsity/permutation_search_kernels/exhaustive_search.py

Lines changed: 15 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -59,36 +59,35 @@ def generate_unique_combinations(built_permutation, remaining_columns, full_perm
5959
# remove the most recent column and put it back on the remaining column list where we found it (sorted)
6060
remaining_columns.insert(c, built_permutation.pop(-1))
6161

62-
import pickle
6362
import os
6463
from os import environ, path
65-
master_unique_permutation_list = {}
64+
unique_permutation_list = {}
6665
def generate_all_unique_combinations(C, M, must_use_all_groups = False):
6766

6867
cache_dir_path = ASP_CACHE_DIR_DEFAULT
6968
# The user is allowed to set the cache directory via an environment variable.
7069
if environ.get(ASP_CACHE_DIR_ENV_VAR) is not None:
7170
cache_dir_path = environ.get(ASP_CACHE_DIR_ENV_VAR)
72-
cache_file_path = path.join(cache_dir_path, "master_list.pkl")
71+
cache_file_path = path.join(cache_dir_path, f"permutations_{C}_{M}.npy")
7372

74-
global master_unique_permutation_list
75-
if len(master_unique_permutation_list) == 0 and path.exists(cache_file_path):
76-
with open(cache_file_path,"rb") as cache:
77-
master_unique_permutation_list = pickle.load(cache)
73+
global unique_permutation_list
74+
if (C,M) not in unique_permutation_list:
75+
if path.exists(cache_file_path):
76+
unique_permutation_list[(C,M)] = np.load(cache_file_path, allow_pickle=False)
7877

79-
if (C,M) not in master_unique_permutation_list:
80-
full_permutation_list = []
81-
generate_unique_combinations([0], [c for c in range(1,C)], full_permutation_list, M)
82-
master_unique_permutation_list[(C,M)] = full_permutation_list
83-
if not path.exists(cache_dir_path):
84-
os.makedirs(cache_dir_path)
85-
with open(cache_file_path, "wb") as cache:
86-
pickle.dump(master_unique_permutation_list, cache)
78+
else:
79+
full_permutation_list = []
80+
generate_unique_combinations([0], [c for c in range(1,C)], full_permutation_list, M)
81+
unique_permutation_list[(C,M)] = full_permutation_list
82+
if not path.exists(cache_dir_path):
83+
os.makedirs(cache_dir_path)
84+
np.save(cache_file_path, full_permutation_list, allow_pickle=False)
8785

88-
unique_permutations = master_unique_permutation_list[(C,M)]
86+
unique_permutations = unique_permutation_list[(C,M)]
8987

9088
return unique_permutations
9189

90+
9291
# analytical solution
9392
import math
9493
def predict_unique_combinations(C, M):

0 commit comments

Comments
 (0)