icrdr/AnyImage
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Implement unified memory management for ONNX model inference

Browse Source 
- Introduced a new memory management policy for ONNX Runtime to release temporary GPU memory after inference completion.
- Established a shared inference entry point for all ONNX models, differentiating between long-lived sessions and temporary memory arenas.
- Updated CUDA sessions to utilize a shrinkable memory arena, reclaiming GPU memory post-inference.
- Ensured consistent memory release semantics across BEN2, DA3, MoGe-2, and Upscale models while reusing existing sessions.
- Added lifecycle tests to validate temporary memory release and session reuse.
- Modified model adapters to integrate the new memory management features, ensuring backward compatibility with existing inference behavior.
- Enhanced unit tests to cover new memory release scenarios and verify correct behavior across different execution providers.
This commit is contained in:
马南
2026-09-01 13:22:35 +08:00
parent 4d4b892d64
commit 2ff7b04f56
24 changed files with 767 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
tests/test_model_ben2_onnx.py
View File

@@ -2,6 +2,7 @@ import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from PIL import Image
@@ -20,6 +21,9 @@ import onnx_ben2
class FakeSession:
def get_providers(self):
return ["CPUExecutionProvider"]
def get_inputs(self):
return [SimpleNamespace(name="input")]
@@ -34,6 +38,25 @@ class FakeSession:
class Ben2OnnxTest(unittest.TestCase):
def test_infer_forwards_memory_release_boundary(self):
session = FakeSession()
with patch.object(
onnx_ben2,
"run_session",
side_effect=lambda model, names, feeds, **_options: model.run(
names, feeds
),
) as run:
result = onnx_ben2.infer(
session,
Image.new("RGB", (40, 20)),
release_memory=False,
)
self.assertEqual(result.size, (40, 20))
self.assertFalse(run.call_args.kwargs["release_memory"])
def test_preprocess_matches_model_layout(self):
image = Image.new("RGB", (40, 20), (64, 128, 192))
result = onnx_ben2.preprocess(image)

25
tests/test_model_da3_onnx.py
View File

@@ -3,6 +3,7 @@ import tempfile
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from PIL import Image
@@ -21,6 +22,9 @@ import onnx_da3
class FakeSession:
def get_providers(self):
return ["CPUExecutionProvider"]
def get_inputs(self):
return [SimpleNamespace(name="images"), SimpleNamespace(name="strategy")]
@@ -36,6 +40,27 @@ class FakeSession:
class Da3OnnxTest(unittest.TestCase):
def test_infer_forwards_memory_release_boundary(self):
with tempfile.TemporaryDirectory() as directory:
image_path = Path(directory) / "input.png"
Image.new("RGB", (80, 60)).save(image_path)
session = FakeSession()
with patch.object(
onnx_da3,
"run_session",
side_effect=lambda model, names, feeds, **_options: model.run(
names, feeds
),
) as run:
onnx_da3.infer(
session,
[image_path],
process_res=56,
release_memory=False,
)
self.assertFalse(run.call_args.kwargs["release_memory"])
def test_target_size_uses_patch_multiple(self):
self.assertEqual(
onnx_da3.target_size((800, 600), 504, "upper_bound_resize"),

32
tests/test_model_moge2_onnx.py
View File

@@ -56,6 +56,9 @@ class FakeSession:
self.calls.append((names, inputs))
return [self.outputs[name] for name in names]
def get_providers(self):
return ["CPUExecutionProvider"]
def get_outputs(self):
return [FakeInput(name) for name in self.outputs]
@@ -145,6 +148,35 @@ class Moge2OnnxTest(unittest.TestCase):
self.assertEqual(inputs["num_tokens"].dtype, np.int64)
self.assertEqual(inputs["num_tokens"], 3600)
def test_infer_forwards_memory_release_boundary(self):
height, width = 8, 12
raw = {
"points": np.zeros((1, height, width, 3), dtype=np.float32),
"normal": np.zeros((1, height, width, 3), dtype=np.float32),
"mask": np.ones((1, height, width), dtype=np.float32),
"scale": np.ones((1,), dtype=np.float32),
}
session = FakeSession(raw)
with (
patch.object(self.moge, "postprocess", return_value={}),
patch.object(
self.moge,
"run_session",
side_effect=lambda model, names, feeds, **_options: model.run(
names, feeds
),
) as run,
):
self.moge.infer(
session,
np.zeros((height, width, 3), dtype=np.float32),
0,
release_memory=False,
)
self.assertFalse(run.call_args.kwargs["release_memory"])
def test_postprocess_recovers_camera_projection_and_metric_scale(self):
height, width = 64, 80
focal = np.float32(0.9)

135
tests/test_model_onnx_runtime.py Normal file
View File

@@ -0,0 +1,135 @@
import sys
import unittest
from pathlib import Path
from types import ModuleType
from unittest.mock import patch
SOURCE_DIR = (
Path(__file__).resolve().parents[1]
/ "src"
/ "anyimage"
/ "server"
/ "models"
)
sys.path.insert(0, str(SOURCE_DIR))
import onnx_runtime
class FakeSession:
def __init__(self, providers):
self.providers = providers
self.calls = []
def get_providers(self):
return self.providers
def run(self, output_names, input_feed, run_options=None):
self.calls.append((output_names, input_feed, run_options))
return ["prediction"]
class OnnxRuntimeTest(unittest.TestCase):
def create_runtime(self, available):
calls = []
runtime = ModuleType("onnxruntime")
runtime.get_available_providers = lambda: available
runtime.InferenceSession = lambda path, providers: calls.append(
(path, providers)
)
return runtime, calls
def test_cuda_session_uses_shrinkable_arena(self):
runtime, calls = self.create_runtime(
["CUDAExecutionProvider", "CPUExecutionProvider"]
)
with patch.dict(sys.modules, {"onnxruntime": runtime}):
onnx_runtime.create_session("model.onnx", "cuda")
self.assertEqual(
calls[0],
(
"model.onnx",
[
(
"CUDAExecutionProvider",
{"arena_extend_strategy": "kSameAsRequested"},
)
],
),
)
def test_coreml_and_cpu_session_options_are_preserved(self):
runtime, calls = self.create_runtime(
["CoreMLExecutionProvider", "CPUExecutionProvider"]
)
with patch.dict(sys.modules, {"onnxruntime": runtime}):
onnx_runtime.create_session("model.onnx", "auto")
self.assertEqual(
calls[0][1],
[
(
"CoreMLExecutionProvider",
dict(onnx_runtime.COREML_PROVIDER_OPTIONS),
),
"CPUExecutionProvider",
],
)
def test_completing_cuda_run_shrinks_memory_arena(self):
configs = []
class FakeRunOptions:
def add_run_config_entry(self, key, value):
configs.append((key, value))
runtime = ModuleType("onnxruntime")
runtime.RunOptions = FakeRunOptions
session = FakeSession(["CUDAExecutionProvider", "CPUExecutionProvider"])
with patch.dict(sys.modules, {"onnxruntime": runtime}):
result = onnx_runtime.run_session(session, ["output"], {"input": 1})
self.assertEqual(result, ["prediction"])
self.assertIsInstance(session.calls[0][2], FakeRunOptions)
self.assertEqual(
configs,
[("memory.enable_memory_arena_shrinkage", "gpu:0")],
)
def test_intermediate_and_non_cuda_runs_use_plain_session_run(self):
cuda = FakeSession(["CUDAExecutionProvider"])
directml = FakeSession(["DmlExecutionProvider"])
onnx_runtime.run_session(
cuda,
["output"],
{"input": 1},
release_memory=False,
)
result = onnx_runtime.run_session(
directml,
["output"],
{"input": 2},
)
self.assertIsNone(cuda.calls[0][2])
self.assertIsNone(directml.calls[0][2])
self.assertEqual(result, ["prediction"])
def test_installed_runtime_accepts_memory_arena_shrinkage_option(self):
import onnxruntime
run_options = onnxruntime.RunOptions()
run_options.add_run_config_entry(
"memory.enable_memory_arena_shrinkage",
"gpu:0",
)
if __name__ == "__main__":
unittest.main()

38
tests/test_model_upscale_onnx.py
View File

@@ -2,6 +2,7 @@ import sys
import unittest
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch
import numpy as np
from PIL import Image
@@ -20,6 +21,9 @@ import onnx_upscale
class FakeSession:
def get_providers(self):
return ["CPUExecutionProvider"]
def get_inputs(self):
return [SimpleNamespace(name="input")]
@@ -121,6 +125,40 @@ class UpscaleOnnxTest(unittest.TestCase):
self.assertEqual(session.calls, 6)
self.assertEqual(session.feeds["input"].shape, (1, 3, 256, 256))
def test_only_final_tile_releases_memory(self):
releases = []
session = FakeSession()
def run(model, output_names, feeds, *, release_memory):
releases.append(release_memory)
return model.run(output_names, feeds)
with patch.object(onnx_upscale, "run_session", side_effect=run):
result = onnx_upscale.infer(
session,
Image.new("RGB", (449, 225), (1, 2, 3)),
)
self.assertEqual(result.size, (1796, 900))
self.assertEqual(releases, [False, False, False, False, False, True])
def test_non_final_image_does_not_release_its_final_tile(self):
releases = []
session = FakeSession()
def run(model, output_names, feeds, *, release_memory):
releases.append(release_memory)
return model.run(output_names, feeds)
with patch.object(onnx_upscale, "run_session", side_effect=run):
onnx_upscale.infer(
session,
Image.new("RGB", (449, 225), (1, 2, 3)),
release_memory=False,
)
self.assertEqual(releases, [False] * 6)
def test_provider_unicode_error_is_presented_as_resource_error(self):
class InvalidProviderSession(FakeSession):
def run(self, _output_names, _feeds):

105
tests/test_server_jobs.py
View File

@@ -504,8 +504,8 @@ class ServerJobsTest(unittest.TestCase):
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("server.models.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
fake_ben2.infer = lambda _model, image, **options: (
inference_options.append((image, options)) or FakeImage()
)
fake_pil = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(
@@ -537,7 +537,7 @@ class ServerJobsTest(unittest.TestCase):
self.assertEqual(
inference_options[0][1],
{},
{"release_memory": True},
)
self.assertEqual(
save_options[0],
@@ -556,7 +556,7 @@ class ServerJobsTest(unittest.TestCase):
model_result = Image.new("RGBA", (2, 1), (10, 20, 30, 128))
received = []
def infer(_model, image):
def infer(_model, image, **_options):
received.append(np.asarray(image.convert("RGB"), dtype=np.uint8))
return model_result
@@ -644,8 +644,8 @@ class ServerJobsTest(unittest.TestCase):
fake_torch.inference_mode = nullcontext
fake_ben2 = ModuleType("server.models.onnx_ben2")
fake_ben2.create_session = lambda _model_dir, _device: object()
fake_ben2.infer = lambda _model, image: (
inference_options.append((image, {})) or FakeImage()
fake_ben2.infer = lambda _model, image, **options: (
inference_options.append((image, options)) or FakeImage()
)
fake_pil = ModuleType("PIL")
fake_pil.Image = SimpleNamespace(
@@ -689,8 +689,10 @@ class ServerJobsTest(unittest.TestCase):
)
self.assertEqual(len(inference_options), 2)
for _image, options in inference_options:
self.assertEqual(options, {})
self.assertEqual(
[options["release_memory"] for _image, options in inference_options],
[False, True],
)
def test_da3_debug_saves_only_direct_api_prediction(self):
prediction = SimpleNamespace(
@@ -754,6 +756,52 @@ class ServerJobsTest(unittest.TestCase):
self.assertFalse(any(context.directory.glob("*.png")))
self.assertEqual(calls[0][0], [str(input_path)])
self.assertEqual(calls[0][1]["process_res"], 518)
self.assertTrue(calls[0][1]["release_memory"])
def test_da3_per_frame_releases_only_after_final_frame(self):
releases = []
def prediction():
return SimpleNamespace(
depth=np.ones((1, 2, 2), dtype=np.float32),
conf=np.ones((1, 2, 2), dtype=np.float32),
extrinsics=np.eye(4, dtype=np.float32)[None, :3],
intrinsics=np.eye(3, dtype=np.float32)[None],
processed_images=np.zeros((1, 2, 2, 3), dtype=np.uint8),
)
fake_da3 = ModuleType("server.models.onnx_da3")
def infer(_session, _images, **options):
releases.append(options["release_memory"])
return prediction()
fake_da3.infer = infer
fake_runtime = ModuleType("server.models.onnx_runtime")
fake_runtime.provider_label = lambda _session: "CPU"
parameters = {
"model": "DA3",
"device": "cpu",
"process_res": 56,
"process_res_method": "upper_bound_resize",
"ref_view_strategy": "middle",
"inference_mode": "per_frame",
}
with patch.dict(
sys.modules,
{
"server.models.onnx_da3": fake_da3,
"server.models.onnx_runtime": fake_runtime,
},
):
self.da3_model.infer(
FakeModels(),
parameters,
[Path("a.png"), Path("b.png")],
)
self.assertEqual(releases, [False, True])
def test_moge2_debug_saves_each_direct_api_dictionary(self):
prediction = {
@@ -786,12 +834,42 @@ class ServerJobsTest(unittest.TestCase):
},
)
self.assertEqual(infer_one.call_count, 2)
self.assertEqual(
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
[False, True],
)
archives = sorted(context.directory.glob("prediction.*.npz"))
self.assertEqual(len(archives), 2)
with np.load(archives[0]) as result:
self.assertEqual(set(result.files), set(prediction))
self.assertFalse(any(context.directory.glob("*.glb")))
def test_moge2_batch_releases_only_after_final_image(self):
prediction = {
"points": np.ones((2, 2, 3), dtype=np.float32),
"depth": np.ones((2, 2), dtype=np.float32),
"normal": np.ones((2, 2, 3), dtype=np.float32),
"mask": np.ones((2, 2), dtype=bool),
"intrinsics": np.eye(3, dtype=np.float32),
}
with patch.object(
self.moge2_model,
"infer_one",
return_value=prediction,
) as infer_one:
frames = self.moge2_model.infer(
FakeModels(),
{"model": "MOGE2_VITS_NORMAL", "device": "cpu"},
[Path("a.png"), Path("b.png")],
include_points=True,
)
self.assertEqual(len(frames), 2)
self.assertEqual(
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
[False, True],
)
def test_ben2_debug_saves_direct_api_foreground(self):
class Foreground:
def save(self, path):
@@ -817,6 +895,9 @@ class ServerJobsTest(unittest.TestCase):
},
)
infer_foreground.assert_called_once()
self.assertTrue(
infer_foreground.call_args.kwargs["release_memory"]
)
self.assertEqual(
(context.directory / "foreground.png").read_bytes(),
b"foreground",
@@ -848,6 +929,10 @@ class ServerJobsTest(unittest.TestCase):
)
self.assertEqual(infer_one.call_count, 2)
self.assertEqual(
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
[False, True],
)
predictions = sorted(context.directory.glob("prediction.*.png"))
self.assertEqual(len(predictions), 2)
with Image.open(predictions[0]) as prediction:
@@ -880,6 +965,10 @@ class ServerJobsTest(unittest.TestCase):
)
self.assertEqual(infer_one.call_count, 2)
self.assertEqual(
[call.kwargs["release_memory"] for call in infer_one.call_args_list],
[False, True],
)
self.assertEqual(result, {"upscale": "."})
colors = sorted(context.directory.glob("upscale.*.png"))
self.assertEqual(

84
tests/test_server_runtime.py
View File

@@ -46,10 +46,12 @@ class ServerRuntimeTest(unittest.TestCase):
importlib.import_module("server.jobs.cutout")
from blendjob import JobServer
from server.model_manager import ModelManager, ModelSessionCache
from server.models import onnx_runtime
cls.JobServer = JobServer
cls.ModelManager, cls.ModelSessionCache = ModelManager, ModelSessionCache
cls.onnx_runtime = onnx_runtime
cls.server = cls.server_module.server
cls.app_module = sys.modules["server.app"]
@@ -224,6 +226,88 @@ class ServerRuntimeTest(unittest.TestCase):
self.assertEqual(loads, [("model", "cpu")])
self.assertEqual(cache.snapshot()["upscale"], "model")
def test_inference_cleanup_keeps_every_cached_session(self):
class Session:
def get_providers(self):
return ["CUDAExecutionProvider"]
def run(self, _names, _feeds, run_options=None):
self.run_options = run_options
return [object()]
class RunOptions:
def add_run_config_entry(self, _key, _value):
return None
runtime = ModuleType("onnxruntime")
runtime.RunOptions = RunOptions
sessions = {
"ben2": Session(),
"da3": Session(),
"moge2": Session(),
"upscale": Session(),
}
ben2 = ModuleType("server.models.onnx_ben2")
ben2.create_session = lambda _directory, _device: sessions["ben2"]
da3 = ModuleType("server.models.onnx_da3")
da3.create_session = lambda _directory, _device: sessions["da3"]
moge2 = ModuleType("server.models.onnx_moge2")
moge2.create_session = lambda _directory, _device: sessions["moge2"]
upscale = ModuleType("server.models.onnx_upscale")
upscale.create_session = lambda _directory, _device: sessions["upscale"]
cache = self.create_cache()
def assert_reused(first, getter, release_name):
with (
patch.dict(sys.modules, {"onnxruntime": runtime}),
patch.object(
cache,
release_name,
wraps=getattr(cache, release_name),
) as release,
):
self.onnx_runtime.run_session(first, ["output"], {"input": 1})
second = getter()
self.assertIs(first, second)
release.assert_not_called()
with patch.dict(
sys.modules,
{
"server.models.onnx_ben2": ben2,
"server.models.onnx_da3": da3,
"server.models.onnx_moge2": moge2,
"server.models.onnx_upscale": upscale,
},
):
first_ben2, _load_ms = cache.get_ben2("ben2", "cuda")
assert_reused(
first_ben2,
lambda: cache.get_ben2("ben2", "cuda")[0],
"release_ben2",
)
first_da3 = cache.get_da3("da3", "cuda")
assert_reused(
first_da3,
lambda: cache.get_da3("da3", "cuda"),
"release_geometry",
)
first_moge2 = cache.get_moge2("moge2", "cuda")
assert_reused(
first_moge2,
lambda: cache.get_moge2("moge2", "cuda"),
"release_geometry",
)
first_upscale = cache.get_upscale("upscale", "cuda")
assert_reused(
first_upscale,
lambda: cache.get_upscale("upscale", "cuda"),
"release_upscale",
)
def test_dedicated_server_runs_only_one_job_at_a_time(self):
server = self.JobServer("Test", storage_root=self.storage_root)
first_started = threading.Event()