Skip to content

Commit

Permalink
Fix open symlinks (#3098)
Browse files Browse the repository at this point in the history
- Cherry-pick of #3096
- Cherry-pick of #3097

test_examples/622/ - passed
  • Loading branch information
KodiaqQ authored Nov 19, 2024
1 parent f6f1513 commit f143e1c
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions nncf/common/tensor_statistics/statistics_serializer.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,8 @@
from typing import Any, Dict, Optional, Tuple, cast

import nncf
from nncf.common.utils.os import fail_if_symlink
from nncf.common.utils.os import safe_open

METADATA_FILE = "statistics_metadata.json"

Expand All @@ -35,7 +37,7 @@ def load_metadata(dir_path: Path) -> Dict[str, Any]:
"""
metadata_file = dir_path / METADATA_FILE
if metadata_file.exists():
with open(metadata_file, "r") as f:
with safe_open(metadata_file, "r") as f:
return cast(Dict[str, Any], json.load(f))
return {"mapping": {}, "metadata": {}}

Expand All @@ -47,7 +49,7 @@ def save_metadata(metadata: Dict[str, Any], dir_path: Path) -> None:
:param dir_path: The directory where the metadata file will be stored.
"""
metadata_file = dir_path / METADATA_FILE
with open(metadata_file, "w") as f:
with safe_open(metadata_file, "w") as f:
json.dump(metadata, f, indent=4)


Expand All @@ -70,6 +72,7 @@ def load_from_dir(dir_path: str) -> Tuple[Dict[str, Any], Dict[str, str]]:
continue # Skip the metadata file

try:
fail_if_symlink(statistics_file)
with gzip.open(statistics_file, "rb") as f:
sanitized_name = statistics_file.name
original_name = mapping.get(sanitized_name, sanitized_name)
Expand Down Expand Up @@ -101,6 +104,7 @@ def dump_to_dir(
mapping[sanitized_name] = original_name

try:
fail_if_symlink(file_path)
with gzip.open(file_path, "wb") as f:
pickle.dump(statistics_value, f)
except (IOError, pickle.PicklingError) as e:
Expand Down

0 comments on commit f143e1c

Please sign in to comment.