Skip to content

Commit 9e7765a

Browse files
Merge branch 'main' into feature/try-to-beat-the-limitation-of-ee-in-terms-of-singular-elements-pushed-into-batch-inputs
2 parents b72719c + 4aa303b commit 9e7765a

File tree

11 files changed

+1285
-12
lines changed

11 files changed

+1285
-12
lines changed

inference/core/cache/model_artifacts.py

Lines changed: 40 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import errno
2+
import json
23
import os.path
34
import re
45
import shutil
@@ -7,12 +8,16 @@
78

89
from filelock import FileLock
910

10-
from inference.core.env import MODEL_CACHE_DIR
11+
from inference.core.env import ATOMIC_CACHE_WRITES_ENABLED, MODEL_CACHE_DIR
12+
from inference.core.exceptions import ModelArtefactError
1113
from inference.core.logger import logger
1214
from inference.core.utils.file_system import (
1315
dump_bytes,
16+
dump_bytes_atomic,
1417
dump_json,
18+
dump_json_atomic,
1519
dump_text_lines,
20+
dump_text_lines_atomic,
1621
read_json,
1722
read_text_file,
1823
)
@@ -67,7 +72,10 @@ def load_json_from_cache(
6772
file: str, model_id: Optional[str] = None, **kwargs
6873
) -> Optional[Union[dict, list]]:
6974
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
70-
return read_json(path=cached_file_path, **kwargs)
75+
try:
76+
return read_json(path=cached_file_path, **kwargs)
77+
except json.JSONDecodeError as e:
78+
raise ModelArtefactError(f"Error loading JSON from cache: {e}")
7179

7280

7381
def save_bytes_in_cache(
@@ -77,7 +85,14 @@ def save_bytes_in_cache(
7785
allow_override: bool = True,
7886
) -> None:
7987
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
80-
dump_bytes(path=cached_file_path, content=content, allow_override=allow_override)
88+
if ATOMIC_CACHE_WRITES_ENABLED:
89+
dump_bytes_atomic(
90+
path=cached_file_path, content=content, allow_override=allow_override
91+
)
92+
else:
93+
dump_bytes(
94+
path=cached_file_path, content=content, allow_override=allow_override
95+
)
8196

8297

8398
def save_json_in_cache(
@@ -88,9 +103,20 @@ def save_json_in_cache(
88103
**kwargs,
89104
) -> None:
90105
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
91-
dump_json(
92-
path=cached_file_path, content=content, allow_override=allow_override, **kwargs
93-
)
106+
if ATOMIC_CACHE_WRITES_ENABLED:
107+
dump_json_atomic(
108+
path=cached_file_path,
109+
content=content,
110+
allow_override=allow_override,
111+
**kwargs,
112+
)
113+
else:
114+
dump_json(
115+
path=cached_file_path,
116+
content=content,
117+
allow_override=allow_override,
118+
**kwargs,
119+
)
94120

95121

96122
def save_text_lines_in_cache(
@@ -100,9 +126,14 @@ def save_text_lines_in_cache(
100126
allow_override: bool = True,
101127
) -> None:
102128
cached_file_path = get_cache_file_path(file=file, model_id=model_id)
103-
dump_text_lines(
104-
path=cached_file_path, content=content, allow_override=allow_override
105-
)
129+
if ATOMIC_CACHE_WRITES_ENABLED:
130+
dump_text_lines_atomic(
131+
path=cached_file_path, content=content, allow_override=allow_override
132+
)
133+
else:
134+
dump_text_lines(
135+
path=cached_file_path, content=content, allow_override=allow_override
136+
)
106137

107138

108139
def get_cache_file_path(file: str, model_id: Optional[str] = None) -> str:

inference/core/env.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@
5555

5656
MD5_VERIFICATION_ENABLED = str2bool(os.getenv("MD5_VERIFICATION_ENABLED", False))
5757

58+
ATOMIC_CACHE_WRITES_ENABLED = str2bool(os.getenv("ATOMIC_CACHE_WRITES_ENABLED", False))
59+
5860
# Base URL for metrics collector
5961
METRICS_COLLECTOR_BASE_URL = os.getenv(
6062
"METRICS_COLLECTOR_BASE_URL",

inference/core/models/roboflow.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -748,9 +748,9 @@ def __init__(
748748
expanded_execution_providers.append(ep)
749749
self.onnxruntime_execution_providers = expanded_execution_providers
750750

751-
self.initialize_model()
752751
self.image_loader_threadpool = ThreadPoolExecutor(max_workers=None)
753752
try:
753+
self.initialize_model()
754754
self.validate_model()
755755
except ModelArtefactError as e:
756756
logger.error(f"Unable to validate model artifacts, clearing cache: {e}")

inference/core/utils/file_system.py

Lines changed: 104 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,70 @@
11
import json
2+
import os
23
import os.path
34
import re
5+
import tempfile
46
from typing import List, Optional, Union
57

68

9+
class AtomicPath:
10+
"""Context manager for atomic file writes.
11+
12+
Ensures that files are either written completely or not at all,
13+
preventing partial/corrupted files from power failures or crashes.
14+
15+
Usage:
16+
with AtomicPath(target_path, allow_override=False) as temp_path:
17+
# Write to temp_path
18+
with open(temp_path, 'w') as f:
19+
f.write(data)
20+
# File is atomically moved to target_path on successful exit
21+
"""
22+
23+
def __init__(self, target_path: str, allow_override: bool = False):
24+
self.target_path = target_path
25+
self.allow_override = allow_override
26+
self.temp_path: Optional[str] = None
27+
self.temp_file = None
28+
29+
def __enter__(self) -> str:
30+
ensure_write_is_allowed(
31+
path=self.target_path, allow_override=self.allow_override
32+
)
33+
ensure_parent_dir_exists(path=self.target_path)
34+
35+
dir_name = os.path.dirname(os.path.abspath(self.target_path))
36+
base_name = os.path.basename(self.target_path)
37+
self.temp_file = tempfile.NamedTemporaryFile(
38+
dir=dir_name, prefix=".tmp_", suffix="_" + base_name, delete=False
39+
)
40+
self.temp_path = self.temp_file.name
41+
self.temp_file.close()
42+
return self.temp_path
43+
44+
def __exit__(self, exc_type, exc_val, exc_tb):
45+
if exc_type is None:
46+
try:
47+
if os.name == "nt": # Windows
48+
if os.path.exists(self.target_path):
49+
os.remove(self.target_path)
50+
os.rename(self.temp_path, self.target_path)
51+
else: # POSIX
52+
os.replace(self.temp_path, self.target_path)
53+
except Exception:
54+
try:
55+
os.unlink(self.temp_path)
56+
except OSError:
57+
pass
58+
raise
59+
else:
60+
# Error occurred - clean up temp file
61+
try:
62+
os.unlink(self.temp_path)
63+
except OSError:
64+
pass
65+
return False # Don't suppress exceptions
66+
67+
768
def read_text_file(
869
path: str,
970
split_lines: bool = False,
@@ -28,31 +89,72 @@ def read_json(path: str, **kwargs) -> Optional[Union[dict, list]]:
2889

2990

3091
def dump_json(
31-
path: str, content: Union[dict, list], allow_override: bool = False, **kwargs
92+
path: str,
93+
content: Union[dict, list],
94+
allow_override: bool = False,
95+
fsync: bool = False,
96+
**kwargs,
3297
) -> None:
3398
ensure_write_is_allowed(path=path, allow_override=allow_override)
3499
ensure_parent_dir_exists(path=path)
35100
with open(path, "w") as f:
36101
json.dump(content, fp=f, **kwargs)
102+
if fsync:
103+
os.fsync(f.fileno())
104+
105+
106+
def dump_json_atomic(
107+
path: str, content: Union[dict, list], allow_override: bool = False, **kwargs
108+
) -> None:
109+
with AtomicPath(path, allow_override=allow_override) as temp_path:
110+
dump_json(temp_path, content, allow_override=True, fsync=True, **kwargs)
37111

38112

39113
def dump_text_lines(
40114
path: str,
41115
content: List[str],
42116
allow_override: bool = False,
43117
lines_connector: str = "\n",
118+
fsync: bool = False,
44119
) -> None:
45120
ensure_write_is_allowed(path=path, allow_override=allow_override)
46121
ensure_parent_dir_exists(path=path)
47122
with open(path, "w") as f:
48123
f.write(lines_connector.join(content))
124+
if fsync:
125+
os.fsync(f.fileno())
49126

50127

51-
def dump_bytes(path: str, content: bytes, allow_override: bool = False) -> None:
128+
def dump_text_lines_atomic(
129+
path: str,
130+
content: List[str],
131+
allow_override: bool = False,
132+
lines_connector: str = "\n",
133+
) -> None:
134+
with AtomicPath(path, allow_override=allow_override) as temp_path:
135+
dump_text_lines(
136+
temp_path,
137+
content,
138+
allow_override=True,
139+
lines_connector=lines_connector,
140+
fsync=True,
141+
)
142+
143+
144+
def dump_bytes(
145+
path: str, content: bytes, allow_override: bool = False, fsync: bool = False
146+
) -> None:
52147
ensure_write_is_allowed(path=path, allow_override=allow_override)
53148
ensure_parent_dir_exists(path=path)
54149
with open(path, "wb") as f:
55150
f.write(content)
151+
if fsync:
152+
os.fsync(f.fileno())
153+
154+
155+
def dump_bytes_atomic(path: str, content: bytes, allow_override: bool = False) -> None:
156+
with AtomicPath(path, allow_override=allow_override) as temp_path:
157+
dump_bytes(temp_path, content, allow_override=True, fsync=True)
56158

57159

58160
def ensure_parent_dir_exists(path: str) -> None:

inference/core/workflows/core_steps/loader.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -341,6 +341,9 @@
341341
from inference.core.workflows.core_steps.transformations.perspective_correction.v1 import (
342342
PerspectiveCorrectionBlockV1,
343343
)
344+
from inference.core.workflows.core_steps.transformations.qr_code_generator.v1 import (
345+
QRCodeGeneratorBlockV1,
346+
)
344347
from inference.core.workflows.core_steps.transformations.relative_static_crop.v1 import (
345348
RelativeStaticCropBlockV1,
346349
)
@@ -667,6 +670,7 @@ def load_blocks() -> List[Type[WorkflowBlock]]:
667670
Moondream2BlockV1,
668671
OverlapBlockV1,
669672
ONVIFSinkBlockV1,
673+
QRCodeGeneratorBlockV1,
670674
]
671675

672676

inference/core/workflows/core_steps/transformations/qr_code_generator/__init__.py

Whitespace-only changes.

0 commit comments

Comments
 (0)