Reduce Python memory use while preparing offline pipeline archives

dm/games/01006FE013472000/README.md

@@ -42,6 +42,8 @@ Everything else needed was general: see [Found on this game](#found-on-this-game
 
 ## Found on this game
 
+- [Memory during offline archive preparation](../../../docs/reference/shaders.md#python-memory-during-offline-pipeline-archive-preparation): the 1.1.1 report has 140,036 pipeline bases. The compiler failure remains open ([#1022](https://github.com/seedds/nxnative/issues/1022)).
+
 General changes first needed here; each applies to every game:
 
 - **Texture and sampler pools for handles traced to no slot** (`Program.cs`, vendored `BindlessElimination.cs`, `dmsh_build.py` `pool_arrays`, `nvn_mtl.m` `pool_gpu`): [Texture handles from game memory index the pools](../../../docs/reference/shaders.md#texture-handles-from-game-memory-index-the-pools).

dm/gfx/tools/pso_offline.py

@@ -26,7 +26,7 @@
 #   --os-min: the archives' system (default: for iOS the SDK's major version, e.g. 27.0; for macOS the major version this Mac runs, since the Mac app is built for the Mac that builds it).
 #   A device opens only an archive of its own major version (macOS 27.2 refused archives made for 26.0 and 26.6 and opened those for 27.0 and 27.2), and builds the bases itself
 #   otherwise: the SDK's version on a Mac whose system is older than its Xcode's SDK (macOS 26.6 with Xcode 27.1) left that Mac without them (GitHub issue #673)
-import argparse, hashlib, json, os, shutil, struct, subprocess, sys, time
+import argparse, filecmp, hashlib, json, os, shutil, struct, subprocess, sys, tempfile, time
 from pathlib import Path
 
 LAYER = b'render_target_array_index'
@@ -134,6 +134,16 @@ def fd(lib, entry):
             'pipeline_descriptors': {'render_pipeline_descriptors': pipes}}
 
 
+def write_script(path, configuration):
+    """Stream the script to disk. Keep the previous file until a complete replacement is ready."""
+    with tempfile.TemporaryDirectory(dir=path.parent) as tmp:
+        candidate = Path(tmp) / path.name
+        with candidate.open('w', encoding='utf-8', newline='\n') as out:
+            json.dump(configuration, out, indent=1)
+        if not path.is_file() or not filecmp.cmp(path, candidate, shallow=False):
+            os.replace(candidate, path)
+
+
 def library_digest(lib_dir, bases):
     """The bytes metal-tt consumes: a profile/chain stamp does not change when the same shaders are rebuilt by another Metal compiler."""
     h = hashlib.sha256()
@@ -167,7 +177,10 @@ def archive_key(configuration, libraries, arch, platform, osmin, toolchain):
     config = dict(configuration, libraries=[dict(lib, path=lib['label'] + '.metallib') for lib in configuration['libraries']])
     inputs = {'format': 1, 'configuration': config, 'libraries': libraries,
               'arch': arch, 'platform': platform, 'os_min': osmin, 'toolchain': toolchain}
-    return hashlib.sha256(json.dumps(inputs, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
+    digest = hashlib.sha256()
+    for chunk in json.JSONEncoder(sort_keys=True, separators=(',', ':')).iterencode(inputs):
+        digest.update(chunk.encode())
+    return digest.hexdigest()
 
 
 def sdk(platform):
@@ -220,8 +233,7 @@ def main(argv=None):
     work.mkdir(parents=True, exist_ok=True); out.mkdir(parents=True, exist_ok=True)
     js = work / 'bases.mtl4-json'   # metal-tt takes the format from the extension: .mtl4-json is Metal 4's (.mtlp-json is the Metal 3 format)
     configuration = script(lib_dir, bases)
-    body = json.dumps(configuration, indent=1).encode()
-    if not js.is_file() or js.read_bytes() != body: js.write_bytes(body)
+    write_script(js, configuration)
     st = stamp(chain, zero, text)
     print(f'pso_offline: {len(pairs)} pairs: {len(bases)} bases, {len(rest)} left to the device (function constants), left out as the host leaves them: '
           + ', '.join(f'{k} {v}' for k, v in left.items()) + f' ({time.time() - t0:.1f} s)')
@@ -231,8 +243,9 @@ def main(argv=None):
     if not bases: archs = []   # no pair, or every eligible pair needs device specialization; clear old archives and still write the stamp and the remainder list
     osmin = a.os_min or os_min_default(a.platform); tv = tool_version(a.platform) if bases else 'not used (function constants only)'
     libraries = library_digest(lib_dir, bases)
-    for arch in archs:
-        key = archive_key(configuration, libraries, arch, a.platform, osmin, tv)
+    keys = {arch: archive_key(configuration, libraries, arch, a.platform, osmin, tv) for arch in archs}
+    del configuration   # the compiler reads the script from disk
+    for arch, key in keys.items():
         dst, kf = out / f'{arch}.mtl4a', work / f'key.{arch}'
         if dst.is_file() and kf.is_file() and kf.read_text().strip() == key: print(f'pso_offline: {arch}: reused ({dst.stat().st_size / 2**20:.0f} MB)'); continue
         kf.unlink(missing_ok=True); dst.unlink(missing_ok=True)

dm/gfx/tools/shader_split.py

@@ -550,20 +550,18 @@ def linkdata_index(path):
 
 def linkdata_jobs(path, rel, parts=64):
     """A LINKDATA archive's index (linkdata_index; Three Houses' DATA1.bin is 6.9 GB) has 32-byte entries {u64 offset, u64 inflated length, u64 stored length, u64 compressed}. The entries are split into
-    ranges of about equal stored size, one scan job each. → [(path, rel, (first entry, end entry)), …], or None when there is no index beside the file"""
+    ranges of about equal stored size, one scan job each. → [(path, rel, (first entry, end entry), index path), …], or None when there is no index beside the file"""
     idx = linkdata_index(path)
     if not idx: return None
     n = os.path.getsize(idx) // 32; step = os.path.getsize(path) / parts; jobs, lo = [], 0
     with open(idx, 'rb') as f: ents = [struct.unpack('<QQQQ', f.read(32)) for _ in range(n)]
     for i in range(1, n + 1):
-        if i == n or ents[i][0] >= (len(jobs) + 1) * step: jobs.append((path, rel, (lo, i))); lo = i
+        if i == n or ents[i][0] >= (len(jobs) + 1) * step: jobs.append((path, rel, (lo, i), idx)); lo = i
     return jobs
 
-def linkdata_shaders(path, rel, rng, out):
+def linkdata_shaders(path, rel, rng, index_path, out):
     from romfs import mmap
-    index = linkdata_index(path)
-    if not index: return red(f'{path}: the LINKDATA index its scan jobs were planned from is gone')
-    with open(index, 'rb') as f: idx = f.read()
+    with open(index_path, 'rb') as f: idx = f.read()
     with open(path, 'rb') as f, mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as d:
         for i in range(*rng):
             off, n, stored, comp = struct.unpack_from('<QQQQ', idx, i * 32)
@@ -2587,7 +2585,7 @@ def scan_file_in(args):
         bdt_shaders(path, rel, args[2], found)
         return rel, uniq_keys(found), var_pairs(found), trees
     if len(args) > 2:   # a range of LINKDATA entries (linkdata_jobs)
-        linkdata_shaders(path, rel, args[2], found)
+        linkdata_shaders(path, rel, args[2], args[3], found)
         return rel, uniq_keys(found), var_pairs(found), trees
     if head[:4] == BND2:   # a Criterion Bundle 2 (Need for Speed: Hot Pursuit Remastered's SHADERS.BNDL): GLSLC outputs in its zlib resources (bnd2_shaders)
         found += bnd2_shaders(open(path, 'rb').read(), '/' + rel)

dm/tools/tests/test_pso_offline.py

@@ -16,6 +16,7 @@
 import sys
 import tempfile
 import unittest
+import tracemalloc
 from unittest import mock
 
 REPO = Path(__file__).resolve().parents[3]
@@ -160,6 +161,36 @@ def main(self, *extra):
         with mock.patch.object(po.subprocess, 'run', self.fake_run), mock.patch.dict(os.environ, DM_GAMES_DIR=str(self.prof.parent)), contextlib.redirect_stdout(io.StringIO()):
             return po.main(args)
 
+    def test_large_script_memory_and_compatible_archive_key(self):
+        # Share the two synthetic library files but give every base a distinct entry pair.
+        bases = [('v_bp', f'vertex{i}', 'f_b', 'main0') for i in range(512)]
+        at_compile = []
+        compiler = self.fake_run
+        def measure(cmd, **kwargs):
+            if '-o' in cmd:
+                at_compile.append(tracemalloc.get_traced_memory()[0])
+            return compiler(cmd, **kwargs)
+        with mock.patch.object(po, 'plan', return_value=(bases, [], {})), mock.patch.object(self, 'fake_run', measure):
+            tracemalloc.start()
+            try:
+                self.main('--os-min', '27.0')
+                peak = tracemalloc.get_traced_memory()[1]
+            finally:
+                tracemalloc.stop()
+        path = self.work / 'bases.mtl4-json'
+        self.assertLess(peak, 3 * path.stat().st_size, 'serializing descriptors must not retain multiple full JSON buffers')
+        self.assertEqual(len(at_compile), 1)
+        self.assertLess(at_compile[0], path.stat().st_size, 'the compiler only needs the script on disk')
+        config = json.loads(path.read_bytes())
+        self.assertEqual(len(config['pipeline_descriptors']['render_pipeline_descriptors']), len(bases))
+        # Keys from earlier builds must continue to hit. Only library locations are normalized.
+        for lib in config['libraries']:
+            lib['path'] = lib['label'] + '.metallib'
+        inputs = dict(format=1, configuration=config, libraries=po.library_digest(self.lib, bases),
+                      arch='applegpu_g18p', platform='ios', os_min='27.0', toolchain=self.compiler)
+        expected = hashlib.sha256(json.dumps(inputs, sort_keys=True, separators=(',', ':')).encode()).hexdigest()
+        self.assertEqual((self.work / 'key.applegpu_g18p').read_text().strip(), expected)
+
     def test_build_reuse_and_stamp(self):
         self.assertEqual(self.main('--gpu', 'g18p'), 0)
         self.assertEqual(len(self.calls), 1)

dm/tools/tests/test_split_unreadable.py

@@ -46,6 +46,27 @@ def scan(self, path, rel):
         with contextlib.redirect_stdout(out): r = ss.scan_file((str(path), rel))
         return r, out.getvalue()
 
+    def test_linkdata_jobs_keep_their_index(self):
+        for archive_name, index_name in [('DATA1.bin', 'DATA0.bin'), ('LD_C0_R.BIN', 'LI_C0_R.IDX')]:
+            with self.subTest(archive=archive_name):
+                archive = self.tmp / archive_name
+                companion = self.tmp / index_name
+                archive.write_bytes(b'abcdwxyz')
+                self.assertIsNone(ss.linkdata_jobs(str(archive), archive_name))
+                companion.write_bytes(struct.pack('<8Q', 0, 4, 0, 0, 4, 4, 0, 0))
+                jobs = ss.linkdata_jobs(str(archive), archive_name, parts=2)
+                if jobs is None: self.fail(f'no LINKDATA jobs for {archive_name}')
+                self.assertEqual(len(jobs), 2)
+                with mock.patch.object(ss, 'kt_shaders') as read:
+                    for job in jobs: ss.scan_file_in(job)
+                self.assertEqual([(data, label) for data, label, _ in (c.args for c in read.call_args_list)],
+                                 [(b'abcd', f'/{archive_name}#0'), (b'wxyz', f'/{archive_name}#1')])
+                companion.unlink()
+                with contextlib.redirect_stdout(io.StringIO()):
+                    _, _, _, extra = ss.scan_file(jobs[0])
+                self.assertEqual(extra[0], 'unreadable')
+                self.assertIn(index_name, extra[1])
+
     def test_red_line_of_a_reader(self):
         p = self.tmp / 'a.dict'; p.write_bytes(index(5, 2)[:0x28])
         (rel, keys, pairs, extra), out = self.scan(p, 'sub/a.dict')

docs/agents/symptoms.md

@@ -739,6 +739,7 @@ One line per symptom in the [subsystem references](../reference/): what you see
 - rebuilding an unchanged NSP after a shader-chain update runs `metal-tt` for minutes even when all referenced Metal libraries and pipeline settings remain identical. → [Offline pipeline archive rebuilt after a chain-only update](../reference/shaders.md#offline-pipeline-archive-rebuilt-after-a-chain-only-update)
 - the host logs `pipeline archive bundled with the app: made for ... this run needs ...`, naming different depths, and pipeline bases compile during play despite a successful archive build. There need not be a `[RED]` line. → [Offline pipeline archive and profile depth disagree](../reference/shaders.md#offline-pipeline-archive-and-profile-depth-disagree)
 - `pipeline archive bundled with the app: this system cannot open …/applegpu_<gpu>.mtl4a (deployment target for architecture … found in archive … is not compatible with current OS …): the bases are built on the device`. On macOS 26.6 with Xcode 27.1, `t_pso_bundle` failed 10 checks, from `bundle: opened` on, and a Mac app packed there compiled every base on its first start, as without the archive. → [Offline pipeline archive made for another major version](../reference/shaders.md#offline-pipeline-archive-made-for-another-major-version)
+- preparing a large offline archive increases the builder's memory use before the compiler starts. → [Python memory during offline pipeline archive preparation](../reference/shaders.md#python-memory-during-offline-pipeline-archive-preparation)
 - `[RED] pso_offline: metal-tt failed for applegpu_<gpu> (exit 1)` with only `applegpu-nt command failed` in `metal-tt.<arch>.log`. → [Offline pipeline archive on a full disk](../reference/shaders.md#offline-pipeline-archive-on-a-full-disk)
 - `[RED] pipeline prewarm: seed list does not exist: …/pso_manifest.tsv` at every start, no `pipeline prewarm list …/pso_user.tsv` line, and the same compile stalls every session in places already visited, while `pso_user.tsv` keeps growing; the pipeline archive gains files every start and is never compacted (Donkey Kong Bananza: 88 files, 7 s to open them). → [The device's prewarm list unread without a seed list](../reference/shaders.md#the-devices-prewarm-list-unread-without-a-seed-list)
 - the first frame that draws with new programs loads their shader libraries on the submit thread (`[longframe] ... shader libraries N, X ms`): tens of milliseconds with the files cached, seconds when the drive is slow. → [Shader libraries loaded when the program is made](../reference/shaders.md#shader-libraries-loaded-when-the-program-is-made)

docs/development.md

@@ -768,7 +768,9 @@ The settings:
 - `DM_PSO_IOS_MIN`: the archives' minimum iOS version. The default is the SDK's major version (`27.0` with the iOS 27 SDK). A device opens only an archive made for its own major version, and builds the bases itself otherwise: an iPhone on iOS 26 cannot open a 27.0 archive, nor one on iOS 27 a 26.0 archive. The Mac app's archive (`--platform mac-app`) takes the major version of the Mac that builds it, not the SDK's ([Offline pipeline archive made for another major version](reference/shaders.md#offline-pipeline-archive-made-for-another-major-version)).
 - The profile key `depth_mode` (`gl`, the default, or `zero`) picks the vertex entry (`main0` or `main0_z`). The host learns the game's depth mode only when the game sets it, so the build cannot derive it. A wrong value costs only the device's own prebuild, as without the archive.
 
-An archive is made again only when its inputs change: the consumed libraries' bytes and stamp (`inputs.tsv`), the stamp above, the architecture, the minimum OS, `metal-tt`'s version and `pso_offline.py` itself (keys in `$W/build/ios/pso/work`). A shader rebuild with another Metal compiler can change the library bytes without changing the profile or chain stamp, so the stamp alone is insufficient. Archive validation and fallback regressions are recorded under [Offline pipeline archive reuse and fallback](reference/shaders.md#offline-pipeline-archive-reuse-and-fallback).
+The builder streams the pipeline script into the work directory and computes the archive key incrementally. During replacement the work directory can briefly contain both script files. The Python descriptor table is released before `metal-tt` starts. The compiler still receives one complete script per architecture.
+
+An archive is made again only when what `metal-tt` compiles from changes: the pipeline script without its library locations, the consumed libraries' names and bytes, the architecture, the platform, the minimum OS and `metal-tt`'s version (keys in `$W/build/ios/pso/work`). A chain update that leaves those libraries byte-identical keeps the archive; the stamp above and the other metadata are written on every run, and the reflection is still checked against the current chain first ([Offline pipeline archive rebuilt after a chain-only update](reference/shaders.md#offline-pipeline-archive-rebuilt-after-a-chain-only-update)). A shader rebuild with another Metal compiler can change the library bytes without changing the profile or chain stamp, so the stamp alone is insufficient. Archive validation and fallback regressions are recorded under [Offline pipeline archive reuse and fallback](reference/shaders.md#offline-pipeline-archive-reuse-and-fallback).
 
 ## Power
 

docs/reference/shaders.md

@@ -1289,6 +1289,13 @@ Finding, keying and translating shaders (`dm/gfx/tools/shader_split.py`, `dmsh_b
 - **Found on:** no game: [GitHub issue #673](https://github.com/seedds/nxnative/issues/673), `t_pso_bundle` on an M4 Pro with macOS 26.6.2 and Xcode 27.1.
 - **Check:** `dm/tools/tests/test_pso_offline.py` (`test_macos_takes_this_macs_major_version`). `t_pso_bundle` builds its archive with the default and opens it on macOS 27.2: 52 checks pass with its libraries built for 27.0 and for 26.6 (`MACOSX_DEPLOYMENT_TARGET=26.6`), the mirror of the reporter's Mac; the same libraries in an archive made for 26.6 failed the 10 checks. Open: no run on macOS 26.x here.
 
+### Python memory during offline pipeline archive preparation
+
+- **Symptom:** preparing a large offline archive increases the builder's memory use before the compiler starts.
+- **Cause:** the builder materialized the pipeline script as a JSON string and as bytes, read the previous script in full to compare it, and serialized the entire cache-key input again. The script bytes and descriptors remained live during compilation.
+- **Fix:** stream the script and the key hash, compare the on-disk scripts in chunks and release the descriptors before starting the compiler.
+- **Found on:** [Mario Party Superstars](../../dm/games/01006FE013472000/README.md). The Python allocation pattern is reproducible with synthetic descriptors. The cause of the reported compiler exit and the compiler processes' memory use remain open.
+
 ### Offline pipeline archive on a full disk
 
 - **Symptom:** `[RED] pso_offline: metal-tt failed for applegpu_<gpu> (exit 1)` with only `applegpu-nt command failed` in `metal-tt.<arch>.log`.