From 5d4c371c794df450ae87e5257c0dd542fdb40bd0 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Thu, 22 Jan 2026 12:13:14 -0800 Subject: [PATCH 1/9] add set_ir function to change each worker's ir --- python/mneme/async_executor.py | 21 +++++++++++++++++++++ python/mneme/replay_executor.py | 27 ++++++++++++++++++++++++++- 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index 19c94174..c823f481 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -5,6 +5,7 @@ from multiprocessing import Event as ProcessEvent from multiprocessing import Process from multiprocessing import Queue as ProcessQueue +from pathlib import Path from queue import Queue as ThreadQueue from threading import Event as ThreadEvent from typing import Callable, Dict, Optional @@ -467,6 +468,26 @@ def _fail_pending_futures(self, error: str): return future.set_error(error) + def set_ir(self, ir: str | Path): + """ Sets the LLVM IR for all workers to this IR. + Async, returns before the IR is set; But is guaranteed to set the IR + before any submit or evaluate calls after set_ir is called. + + Parameters + ---------- + ir : str | Path + Path to the LLVM IR file (.bc or .ll) or the IR as a string. + """ + if isinstance(ir, Path): + ir_data = str(ir.absolute()) + else: + ir_data = ir + + msg = {"payload": "set_ir", "data": ir_data} + for w in self.workers: + w._ipc_write_q.put(msg) + + # ------------------------------------------------------------------ # Submit new job (non-blocking) # ------------------------------------------------------------------ diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 1084ccb9..46d88645 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -34,6 +34,7 @@ import os from datetime import datetime, timezone from multiprocessing import Event, Queue +from pathlib import Path from typing import Optional, Tuple from mneme.device import ( @@ -44,7 +45,7 @@ set_device, ) from mneme.llvm.buffer import MemBufferRef -from mneme.llvm.module import ModuleRef +from mneme.llvm.module import ModuleRef, parse_assembly, parse_bitcode from mneme.mneme_logging import logger from mneme.mneme_types import ExperimentConfiguration, ExperimentResult from mneme.page_manager import PageManagerRef @@ -178,6 +179,23 @@ def __exit__(self, exc_type, exc_val, exc_tb): def link_ir(self): return self.records.link_llvm_modules(prune=True, internalize=True) + def set_new_ir(self, ir_path_or_asm: str): + if isinstance(ir_path_or_asm, str) and (ir_path_or_asm.endswith('.ll') or ir_path_or_asm.endswith('.bc')): + ir_path = Path(ir_path_or_asm) + if ir_path.suffix == '.bc': + with open(ir_path, 'rb') as f: + new_ir = parse_bitcode(f.read()) + else: # .ll + with open(ir_path, 'r') as f: + new_ir = parse_assembly(f.read()) + else: + # assume str is text IR + new_ir = parse_assembly(ir_path_or_asm) + + # apply internalization and pruning + jit.internalize(new_ir, self.kernel_descr.kernel_name) + jit.pruneIR(new_ir) + @cond_time("preprocess_ir_time") def _preprocess_ir( self, @@ -817,6 +835,13 @@ def run( f"Worker {worker.device_id} received terminate request, exiting ..." ) break + elif msg["payload"] == "set_ir": + # update the root_ir for subsequent requests + logger.debug(f"Worker {worker.device_id} received set_ir request") + new_ir_data = msg["data"] + + root_ir = worker.set_new_ir(new_ir_data) + elif msg["payload"] == "process": logger.debug( f"Worker {worker.device_id} received processing request {msg['exp_id']}" From 762028adef2e626976353a6bab906a489fb03021 Mon Sep 17 00:00:00 2001 From: koparasy Date: Mon, 26 Jan 2026 09:01:28 -0800 Subject: [PATCH 2/9] Fix typing --- python/mneme/async_executor.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index c823f481..c410e0bb 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -8,7 +8,7 @@ from pathlib import Path from queue import Queue as ThreadQueue from threading import Event as ThreadEvent -from typing import Callable, Dict, Optional +from typing import Callable, Dict, Optional, Union from mneme.futures import EvalFuture from mneme.mneme_logging import logger @@ -468,7 +468,7 @@ def _fail_pending_futures(self, error: str): return future.set_error(error) - def set_ir(self, ir: str | Path): + def set_ir(self, ir: Union[str, Path]): """ Sets the LLVM IR for all workers to this IR. Async, returns before the IR is set; But is guaranteed to set the IR before any submit or evaluate calls after set_ir is called. From a1d3d0088a2728ff2db145d3007de8a6db1a06ed Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Tue, 17 Feb 2026 17:45:34 -0800 Subject: [PATCH 3/9] moving commits from other branch to this one; return ir fix --- python/mneme/replay_executor.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 46d88645..32f756b8 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -196,6 +196,8 @@ def set_new_ir(self, ir_path_or_asm: str): jit.internalize(new_ir, self.kernel_descr.kernel_name) jit.pruneIR(new_ir) + return new_ir + @cond_time("preprocess_ir_time") def _preprocess_ir( self, From e55bc36e376df8271e9152d502db0c506d9b13e0 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Tue, 17 Feb 2026 17:46:40 -0800 Subject: [PATCH 4/9] add guard and status message --- python/mneme/replay_executor.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/python/mneme/replay_executor.py b/python/mneme/replay_executor.py index 32f756b8..e597d0ae 100644 --- a/python/mneme/replay_executor.py +++ b/python/mneme/replay_executor.py @@ -844,6 +844,9 @@ def run( root_ir = worker.set_new_ir(new_ir_data) + if root_ir is None: + logger.error(f"Worker {worker.device_id} failed to set new IR") + elif msg["payload"] == "process": logger.debug( f"Worker {worker.device_id} received processing request {msg['exp_id']}" From ba6be3b9c8af2ba1e239bc3fadade13f8abb6d10 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Mon, 4 May 2026 17:39:53 -0700 Subject: [PATCH 5/9] update tests for async executor --- python/tests/test_async_executor.py | 59 +++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/python/tests/test_async_executor.py b/python/tests/test_async_executor.py index d84ae8ab..8b8489af 100644 --- a/python/tests/test_async_executor.py +++ b/python/tests/test_async_executor.py @@ -393,6 +393,65 @@ def test_async_executor_evaluate(monkeypatch): assert out.executed is True +def test_async_executor_set_ir_sends_ir_to_all_workers(monkeypatch): + workers = [] + + class FakeHandle: + def __init__(self): + self._ipc_write_q = queue.Queue() + workers.append(self) + + monkeypatch.setattr( + "mneme.async_executor.TuneWorkerHandle", lambda *a, **k: FakeHandle() + ) + + exe = AsyncReplayExecutor( + record_db="db", + record_id="rid", + iterations=3, + results_db_dir="/tmp", + num_workers=2, + ) + + exe.set_ir("define void @kernel() { ret void }") + + for worker in workers: + assert worker._ipc_write_q.get_nowait() == { + "payload": "set_ir", + "data": "define void @kernel() { ret void }", + } + assert worker._ipc_write_q.empty() + + +def test_async_executor_set_ir_normalizes_path(monkeypatch, tmp_path): + workers = [] + + class FakeHandle: + def __init__(self): + self._ipc_write_q = queue.Queue() + workers.append(self) + + monkeypatch.setattr( + "mneme.async_executor.TuneWorkerHandle", lambda *a, **k: FakeHandle() + ) + + exe = AsyncReplayExecutor( + record_db="db", + record_id="rid", + iterations=3, + results_db_dir="/tmp", + num_workers=1, + ) + ir_path = tmp_path / "kernel.ll" + + exe.set_ir(ir_path) + + assert workers[0]._ipc_write_q.get_nowait() == { + "payload": "set_ir", + "data": str(ir_path.absolute()), + } + + def test_async_executor_shutdown(monkeypatch): joined = [] From ba7b4904aa066aeabebb5af3400e988b2fd62690 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Mon, 4 May 2026 18:18:47 -0700 Subject: [PATCH 6/9] patch lower level replay executor APIs to ensure set_new_ir is covered during tests --- python/tests/test_replay_executor.py | 154 +++++++++++++++++++++++++++ 1 file changed, 154 insertions(+) diff --git a/python/tests/test_replay_executor.py b/python/tests/test_replay_executor.py index 519dd5da..116c3c8d 100644 --- a/python/tests/test_replay_executor.py +++ b/python/tests/test_replay_executor.py @@ -360,6 +360,85 @@ def test_link_ir_forwards_prune_and_internalize(monkeypatch): assert rec.link_calls == [(True, True)] +@pytest.mark.parametrize( + ("ir_input", "file_mode", "file_contents", "expected_parser", "expected_data"), + [ + ( + "define void @K() { ret void }", + None, + None, + "asm", + "define void @K() { ret void }", + ), + ( + "kernel.ll", + "w", + "define void @K() { ret void }", + "asm", + "define void @K() { ret void }", + ), + ("kernel.bc", "wb", b"bitcode", "bitcode", b"bitcode"), + ], +) +def test_set_new_ir_parses_and_prunes_replacement_ir( + monkeypatch, + tmp_path, + ir_input, + file_mode, + file_contents, + expected_parser, + expected_data, +): + mod = _reload_with_identity_decorators(monkeypatch) + + kernel = FakeKernelDescr(kernel_name="K") + rec = FakeRecordedExecution(kernel) + + monkeypatch.setattr( + mod.RecordedExecution, "from_json", staticmethod(lambda _: rec), raising=True + ) + monkeypatch.setattr(mod, "set_device", lambda _: None, raising=True) + monkeypatch.setattr(mod, "get_device_arch", lambda: "sm", raising=True) + monkeypatch.setattr(mod, "get_device_count", lambda: 1, raising=True) + + parsed_module = FakeModule("parsed") + calls = [] + + def fake_parse_assembly(data): + calls.append(("asm", data)) + return parsed_module + + def fake_parse_bitcode(data): + calls.append(("bitcode", data)) + return parsed_module + + def fake_internalize(ir, kernel_name): + calls.append(("internalize", ir, kernel_name)) + + def fake_prune(ir): + calls.append(("prune", ir)) + + monkeypatch.setattr(mod, "parse_assembly", fake_parse_assembly, raising=True) + monkeypatch.setattr(mod, "parse_bitcode", fake_parse_bitcode, raising=True) + monkeypatch.setattr(mod.jit, "internalize", fake_internalize, raising=True) + monkeypatch.setattr(mod.jit, "pruneIR", fake_prune, raising=True) + + if file_mode is not None: + ir_path = tmp_path / ir_input + with open(ir_path, file_mode) as f: + f.write(file_contents) + ir_input = str(ir_path) + + ex = mod.BaseExecutor(record_db="x", record_id="rid") + + assert ex.set_new_ir(ir_input) is parsed_module + assert calls == [ + (expected_parser, expected_data), + ("internalize", parsed_module, "K"), + ("prune", parsed_module), + ] + + def test_preprocess_ir_calls_jit_hooks_based_on_config(monkeypatch): mod = _reload_with_identity_decorators(monkeypatch) @@ -633,6 +712,9 @@ def test_tuneworker_run_process_and_terminate(monkeypatch, tmp_path): monkeypatch.setattr(mod.os, "open", lambda *a, **k: 999, raising=True) monkeypatch.setattr(mod.os, "dup2", lambda *a, **k: None, raising=True) + set_ir_calls = [] + process_ir_names = [] + class FakeWorker: def __init__(self, record_db, record_id, device_id, iterations, warmup): self.record_db = record_db @@ -644,6 +726,10 @@ def __init__(self, record_db, record_id, device_id, iterations, warmup): def link_ir(self): return FakeModule("root_ir") + def set_new_ir(self, ir_data): + set_ir_calls.append(ir_data) + return FakeModule("replacement_ir") + def __enter__(self): return self @@ -651,6 +737,7 @@ def __exit__(self, exc_type, exc, tb): return False def process_payload(self, ir_module, config): + process_ir_names.append(ir_module.name) res = mod.ExperimentResult(executed=True, verified=True) return res, FakeModule("ir_out") @@ -671,6 +758,7 @@ def set(self): state = FakeEvent() + req_q.put({"payload": "set_ir", "data": "replacement asm"}) req_q.put( { "payload": "process", @@ -693,8 +781,74 @@ def set(self): ) assert state.set_called == 1 + assert set_ir_calls == ["replacement asm"] + assert process_ir_names == ["replacement_ir_clone"] msg = resp_q.get_nowait() assert msg["payload"] == "result" assert msg["exp_id"] == 7 assert isinstance(msg["data"], dict) assert msg["llvm_ir"] == "" + + +def test_tuneworker_run_logs_failed_set_ir(monkeypatch, tmp_path): + mod = _reload_with_identity_decorators(monkeypatch) + + real_run = mod.TuneWorker.run + + monkeypatch.setattr(mod.os, "open", lambda *a, **k: 999, raising=True) + monkeypatch.setattr(mod.os, "dup2", lambda *a, **k: None, raising=True) + + errors = [] + monkeypatch.setattr( + mod.logger, "error", lambda msg: errors.append(msg), raising=True + ) + + class FakeWorker: + def __init__(self, record_db, record_id, device_id, iterations): + self.device_id = device_id + + def link_ir(self): + return FakeModule("root_ir") + + def set_new_ir(self, ir_data): + return None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + monkeypatch.setattr( + mod, "TuneWorker", lambda *a, **k: FakeWorker(*a, **k), raising=True + ) + + req_q = queue.Queue() + resp_q = queue.Queue() + + class FakeEvent: + def __init__(self): + self.set_called = 0 + + def set(self): + self.set_called += 1 + + state = FakeEvent() + + req_q.put({"payload": "set_ir", "data": "bad asm"}) + req_q.put({"payload": "terminate"}) + + real_run( + request_q=req_q, + response_q=resp_q, + record_db="db.json", + record_id="rid", + device_id=3, + iterations=3, + results_db_dir=str(tmp_path), + state=state, + ) + + assert state.set_called == 1 + assert errors == ["Worker 3 failed to set new IR"] + assert resp_q.empty() From 0878e660715265d14efd2adb9f492383592bece0 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Tue, 21 Jul 2026 16:02:39 -0700 Subject: [PATCH 7/9] add revision system to fix issue --- python/mneme/async_executor.py | 20 ++++++++++++-------- python/mneme/futures.py | 18 +++++++++++++++++- 2 files changed, 29 insertions(+), 9 deletions(-) diff --git a/python/mneme/async_executor.py b/python/mneme/async_executor.py index c410e0bb..24052265 100644 --- a/python/mneme/async_executor.py +++ b/python/mneme/async_executor.py @@ -194,6 +194,7 @@ def _spawn_process(self): ) self._process.start() self._action = self.StateMachine.SUBMIT + self._ir_revision = 0 def _startup_failure_error(self): return ( @@ -276,6 +277,10 @@ def _submit(self): return self.current = future + if future.ir_revision != self._ir_revision: + self._ipc_write_q.put({"payload": "set_ir", "data": future.ir_data}) + self._ir_revision = future.ir_revision + msg = { "payload": "process", "data": future.config.to_dict(), @@ -428,6 +433,8 @@ def __init__( self._futures: Dict[int, EvalFuture] = {} self._next_id = 0 self._lock = threading.Lock() + self._ir_revision = 0 + self._ir_data = None self.iterations = iterations self.warmup = warmup self.max_startup_failures = max_startup_failures @@ -469,9 +476,7 @@ def _fail_pending_futures(self, error: str): future.set_error(error) def set_ir(self, ir: Union[str, Path]): - """ Sets the LLVM IR for all workers to this IR. - Async, returns before the IR is set; But is guaranteed to set the IR - before any submit or evaluate calls after set_ir is called. + """Use this LLVM IR for evaluations submitted after this call. Parameters ---------- @@ -483,10 +488,9 @@ def set_ir(self, ir: Union[str, Path]): else: ir_data = ir - msg = {"payload": "set_ir", "data": ir_data} - for w in self.workers: - w._ipc_write_q.put(msg) - + with self._lock: + self._ir_revision += 1 + self._ir_data = ir_data # ------------------------------------------------------------------ # Submit new job (non-blocking) @@ -513,7 +517,7 @@ def submit(self, config: ExperimentConfiguration) -> EvalFuture: job_id = self._next_id self._next_id += 1 logger.debug(f"[{self.__class__.__name__}] Submitting job {job_id}") - future = EvalFuture(job_id, config) + future = EvalFuture(job_id, config, self._ir_revision, self._ir_data) self._futures[job_id] = future if self._broken_error is not None: future.set_error(self._broken_error) diff --git a/python/mneme/futures.py b/python/mneme/futures.py index 44cbd723..669849c4 100644 --- a/python/mneme/futures.py +++ b/python/mneme/futures.py @@ -41,9 +41,19 @@ class EvalFuture: Stable identifier assigned by the submitting executor. config : ExperimentConfiguration Configuration associated with this evaluation request. + ir_revision : int + IR revision captured when this evaluation was submitted. + ir_data : str | None + Replacement IR associated with ``ir_revision``. """ - def __init__(self, job_id: int, config: ExperimentConfiguration): + def __init__( + self, + job_id: int, + config: ExperimentConfiguration, + ir_revision: int = 0, + ir_data: Optional[str] = None, + ): """ Parameters ---------- @@ -51,9 +61,15 @@ def __init__(self, job_id: int, config: ExperimentConfiguration): Unique identifier for this evaluation. config : ExperimentConfiguration Configuration to be evaluated. + ir_revision : int, optional + IR revision captured for this evaluation. + ir_data : str | None, optional + Replacement IR associated with the captured revision. """ self.job_id = job_id self.config = config # small dict of input params + self.ir_revision = ir_revision + self.ir_data = ir_data self._cond = threading.Condition() self._done = False From fdfa014a602e35de274ff3c7b8c25a2fd1b539f7 Mon Sep 17 00:00:00 2001 From: Daniel Nichols Date: Wed, 22 Jul 2026 08:42:15 -0700 Subject: [PATCH 8/9] update tests --- python/tests/test_async_executor.py | 63 ++++++++++++++--------------- 1 file changed, 30 insertions(+), 33 deletions(-) diff --git a/python/tests/test_async_executor.py b/python/tests/test_async_executor.py index 8b8489af..a9b7a5ff 100644 --- a/python/tests/test_async_executor.py +++ b/python/tests/test_async_executor.py @@ -144,6 +144,23 @@ def test_submit_moves_state_and_sets_current(handle): assert handle.current is future assert handle._action == handle.StateMachine.RECEIVE + assert handle._ipc_write_q.get_nowait()["payload"] == "process" + + +def test_submit_installs_future_ir_revision(handle): + future = EvalFuture( + 3, ExperimentConfiguration(), ir_revision=1, ir_data="replacement ir" + ) + handle.global_q.put(future) + + handle._submit() + + assert handle._ipc_write_q.get_nowait() == { + "payload": "set_ir", + "data": "replacement ir", + } + assert handle._ipc_write_q.get_nowait()["payload"] == "process" + assert handle._ir_revision == 1 def test_try_receive_no_message(handle): @@ -393,63 +410,43 @@ def test_async_executor_evaluate(monkeypatch): assert out.executed is True -def test_async_executor_set_ir_sends_ir_to_all_workers(monkeypatch): - workers = [] - - class FakeHandle: - def __init__(self): - self._ipc_write_q = queue.Queue() - workers.append(self) - - monkeypatch.setattr( - "mneme.async_executor.TuneWorkerHandle", lambda *a, **k: FakeHandle() - ) +def test_async_executor_set_ir_only_affects_subsequent_work(monkeypatch): + monkeypatch.setattr("mneme.async_executor.TuneWorkerHandle", lambda *a, **k: None) exe = AsyncReplayExecutor( record_db="db", record_id="rid", iterations=3, results_db_dir="/tmp", - num_workers=2, + num_workers=0, ) + old_future = exe.submit(ExperimentConfiguration()) exe.set_ir("define void @kernel() { ret void }") + new_future = exe.submit(ExperimentConfiguration()) - for worker in workers: - assert worker._ipc_write_q.get_nowait() == { - "payload": "set_ir", - "data": "define void @kernel() { ret void }", - } - assert worker._ipc_write_q.empty() + assert old_future.ir_revision == 0 + assert old_future.ir_data is None + assert new_future.ir_revision == 1 + assert new_future.ir_data == "define void @kernel() { ret void }" def test_async_executor_set_ir_normalizes_path(monkeypatch, tmp_path): - workers = [] - - class FakeHandle: - def __init__(self): - self._ipc_write_q = queue.Queue() - workers.append(self) - - monkeypatch.setattr( - "mneme.async_executor.TuneWorkerHandle", lambda *a, **k: FakeHandle() - ) + monkeypatch.setattr("mneme.async_executor.TuneWorkerHandle", lambda *a, **k: None) exe = AsyncReplayExecutor( record_db="db", record_id="rid", iterations=3, results_db_dir="/tmp", - num_workers=1, + num_workers=0, ) ir_path = tmp_path / "kernel.ll" exe.set_ir(ir_path) + future = exe.submit(ExperimentConfiguration()) - assert workers[0]._ipc_write_q.get_nowait() == { - "payload": "set_ir", - "data": str(ir_path.absolute()), - } + assert future.ir_data == str(ir_path.absolute()) def test_async_executor_shutdown(monkeypatch): From ae36ce4cb1262cc97cc0f1b004109dcba3c8ec9c Mon Sep 17 00:00:00 2001 From: Zane Fink Date: Mon, 27 Jul 2026 11:33:20 -0700 Subject: [PATCH 9/9] fix missing param --- python/tests/test_replay_executor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tests/test_replay_executor.py b/python/tests/test_replay_executor.py index 116c3c8d..0d09f828 100644 --- a/python/tests/test_replay_executor.py +++ b/python/tests/test_replay_executor.py @@ -804,7 +804,7 @@ def test_tuneworker_run_logs_failed_set_ir(monkeypatch, tmp_path): ) class FakeWorker: - def __init__(self, record_db, record_id, device_id, iterations): + def __init__(self, record_db, record_id, device_id, iterations, warmup): self.device_id = device_id def link_ir(self):