# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import json import os import subprocess import sys import tempfile import unittest import uuid from op_test import is_custom_device import paddle from paddle import base MiB = 1 << 20 def _run_test_case(plan, flags, cuda_visible_devices="0"): script = os.path.join( os.path.dirname(__file__), "auto_growth_allocator_gpu.py" ) env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = cuda_visible_devices env["FLAGS_JSON"] = json.dumps(flags) env.setdefault("PYTHONUNBUFFERED", "1") keep = os.environ.get("AG_KEEP_OUT", "").strip() if keep: if keep == "1": out_dir = os.path.join(os.getcwd(), "_ag_out") os.makedirs(out_dir, exist_ok=True) out_path = os.path.join( out_dir, f"ag_{os.getpid()}_{uuid.uuid4().hex}.json" ) elif keep.endswith(".json"): os.makedirs( os.path.dirname(os.path.abspath(keep)) or ".", exist_ok=True ) out_path = os.path.abspath(keep) else: out_dir = os.path.abspath(keep) os.makedirs(out_dir, exist_ok=True) out_path = os.path.join( out_dir, f"ag_{os.getpid()}_{uuid.uuid4().hex}.json" ) else: fd, out_path = tempfile.mkstemp(prefix="ag_", suffix=".json") os.close(fd) log_path = out_path + ".log" if keep else None cmd = [ sys.executable, script, "--plan", json.dumps(plan), "--out", out_path, ] if log_path: cmd += ["--log", log_path] if env.get("AG_TEE", "") == "1": p = subprocess.run(cmd, env=env, text=True) else: p = subprocess.run(cmd, env=env, capture_output=True, text=True) if p.returncode != 0: raise RuntimeError( f"probe failed:\nSTDOUT:\n{p.stdout}\nSTDERR:\n{p.stderr}" ) with open(out_path, "r", encoding="utf-8") as f: data = json.load(f) if not keep: try: os.remove(out_path) if log_path: os.remove(log_path) except Exception: pass else: sys.stderr.write(f"[AG_KEEP_OUT] {out_path}\n") if log_path: sys.stderr.write(f"[AG_KEEP_OUT] {log_path}\n") return data class TestAllocatorFlagsWithSubprocess(unittest.TestCase): def setUp(self): if base.is_compiled_with_cuda() or is_custom_device(): paddle.set_flags( { 'FLAGS_allocator_strategy': 'auto_growth', 'FLAGS_use_cuda_malloc_async_allocator': 0, } ) def test_memory_pool_flags(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_small_pool_size_in_mb": 1, "FLAGS_auto_growth_chunk_size_in_mb": 10, # ignored because FLAGS_small_pool_size_in_mb > 0 "FLAGS_small_pool_auto_growth_chunk_size_in_mb": 2, "FLAGS_large_pool_auto_growth_chunk_size_in_mb": 16, "FLAGS_small_pool_pre_alloc_in_mb": 2, "FLAGS_large_pool_pre_alloc_in_mb": 20, } plan = [ {"op": "init"}, {"op": "alloc_small", "mb_per_block": 0.5, "blocks": 7}, ] out = _run_test_case(plan, flags) a0, a1 = out["allocated"][0], out["allocated"][1] r0, r1 = out["reserved"][0], out["reserved"][1] self.assertEqual(a1, int(3.5 * MiB)) self.assertEqual(r0, int(22 * MiB)) self.assertEqual(r1, r0 + int(2 * MiB), msg=f"r0={r0}, r1={r1}") def test_large_pool_growth_override_16mb(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_small_pool_size_in_mb": 1, "FLAGS_small_pool_auto_growth_chunk_size_in_mb": 0, "FLAGS_large_pool_auto_growth_chunk_size_in_mb": 16, "FLAGS_small_pool_pre_alloc_in_mb": 0, "FLAGS_large_pool_pre_alloc_in_mb": 6, } plan = [ {"op": "init"}, {"op": "alloc_large", "mb": 8}, ] out = _run_test_case(plan, flags) r0, r1 = out["reserved"][0], out["reserved"][1] self.assertEqual(r1, r0 + int(16 * MiB), msg=f"r0={r0}, r1={r1}") def test_single_pool(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_small_pool_size_in_mb": 0, "FLAGS_small_pool_auto_growth_chunk_size_in_mb": 2, "FLAGS_large_pool_auto_growth_chunk_size_in_mb": 4, "FLAGS_auto_growth_chunk_size_in_mb": 10, "FLAGS_small_pool_pre_alloc_in_mb": 2, "FLAGS_large_pool_pre_alloc_in_mb": 6, } plan = [ {"op": "init"}, {"op": "alloc_small", "mb_per_block": 0.5, "blocks": 1}, {"op": "alloc_large", "mb": 10}, ] out = _run_test_case(plan, flags) a0, a1, a2 = ( out["allocated"][0], out["allocated"][1], out["allocated"][2], ) r0, r1, r2 = out["reserved"][0], out["reserved"][1], out["reserved"][2] self.assertEqual(a1, int(0.5 * MiB)) self.assertEqual(a2, int(10.5 * MiB)) self.assertEqual(r0, int(10 * MiB), msg=f"r0={r0}") self.assertEqual(r1, int(10 * MiB), msg=f"r1={r1}") self.assertEqual(r2, int(20 * MiB), msg=f"r2={r2}") def test_all_block_info_without_vmm(self): if not base.is_compiled_with_cuda(): return flags = { "FLAGS_allocator_strategy": "auto_growth", "FLAGS_use_cuda_malloc_async_allocator": 0, "FLAGS_use_auto_growth_v2": False, "FLAGS_use_virtual_memory_auto_growth": False, "FLAGS_small_pool_size_in_mb": 0, "FLAGS_auto_growth_chunk_size_in_mb": 4, } plan = [ {"op": "init"}, {"op": "alloc_large", "mb": 1}, {"op": "all_block_info"}, ] out = _run_test_case(plan, flags) all_block_info = out["all_block_info"][0] if all_block_info is None: self.skipTest("all_block_info is not supported in this build") self.assertGreater(len(all_block_info), 0) self.assertTrue( any(len(chunk) > 0 for chunk in all_block_info), msg=f"all_block_info={all_block_info}", ) self.assertTrue( any( len(block) == 3 and block[0] > 0 for chunk in all_block_info for block in chunk ), msg=f"all_block_info={all_block_info}", ) def test_memory_limit(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_gpu_memory_limit_mb": 10, } plan = [ {"op": "try_alloc", "mb": 5}, {"op": "try_alloc", "mb": 20}, ] out = _run_test_case(plan, flags) self.assertEqual(out["try_alloc_ok"][0], True) self.assertEqual(out["try_alloc_ok"][1], False) def test_auto_growth_allocator_v2(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_use_auto_growth_v2": True, "FLAGS_large_pool_pre_alloc_in_mb": 6, } plan = [ {"op": "init"}, {"op": "alloc_large", "mb": 20}, ] out = _run_test_case(plan, flags) r0 = out["reserved"][0] self.assertLessEqual(r0, int(6 * MiB), msg=f"r0={r0}") def test_trace_flag(self): if not (base.is_compiled_with_cuda() or is_custom_device()): return flags = { "FLAGS_small_pool_size_in_mb": 1, "FLAGS_large_pool_pre_alloc_in_mb": 5, "FLAGS_free_idle_chunk": True, "FLAGS_free_when_no_cache_hit": True, "FLAGS_print_allocator_trace_info": True, } plan = [ {"op": "init"}, {"op": "alloc_small", "mb": 1}, ] out = _run_test_case(plan, flags) if __name__ == "__main__": unittest.main()