Fix reforge workspace isolation and durable result publication

README.md

@@ -84,6 +84,20 @@ sasskit analyze kernel.cubin -k MyKernel -t 64
 sasskit disassemble kernel.cubin -k MyKernel
 ```
 
+Each reforge call uses a private temporary workspace. The input is read-only;
+`--work-dir` selects the existing parent directory for temporary files. The best
+result is published outside that workspace and remains available after cleanup:
+
+```bash
+sasskit reforge kernel.cubin -k MyKernel -n 300 --work-dir /tmp -o best.cubin
+```
+
+Existing output files (including input aliases) are never overwritten. Without
+`-o`, each call creates `./reforge-results/<unique-run>/best.cubin`.
+Python callers can use the keyword-only `output_path` and `work_dir` arguments;
+`state.best_cubin_path` identifies the published result. A run with no accepted
+mutation publishes the baseline if its benchmark is valid.
+
 ## Architecture
 
 ```

docs/reforge-workspace-validation.md

@@ -0,0 +1,38 @@
+# Reforge workspace validation (SM120)
+
+Baseline: `kacper-daftcode/sasskit@913c13451285e07e189b8f89440bd848251a50e1`.
+Python 3.12.13; tests ran in a dedicated environment under `/home/gongji/0z5a`.
+Cubit was built from `1f7a5aa6cb0096223f4054930b58c9c0208251e5`;
+pyelftools 0.32. No system environment packages were changed.
+
+| Check | Baseline | Patched | Speedup |
+|---|---|---|---|
+| Barrier-synchronized 2-process shared-path regression | Cross-run token overwritten | Pass | N/A: correctness |
+| Barrier-synchronized 8-process shared-path regression | Cross-run token overwritten | Pass | N/A: correctness |
+| Workspace/CLI/publication suite | New output API absent | 11 passed | N/A |
+| Existing suite | 7 skipped | Same 7 skipped | N/A |
+| GPU workers sharing one workspace parent | Not run on unsafe baseline | 1 / 2 / 4 / 8 passed | No kernel speedup claimed |
+
+The CPU baseline reproducer relocates the **two shared filenames only** into a
+task-owned directory; it preserves the old shared-path behavior and forces all
+writers to reach a barrier before reading. The fixed test uses the actual
+`reforge` function, not a replacement lifecycle implementation.
+
+Tests cover repeated calls, spaces, explicit output collisions, atomic
+no-clobber publication, input symlink/hardlink aliases, benchmark exceptions,
+KeyboardInterrupt, publication failure, durable returned paths and legacy CLI
+arguments. Existing fixture-dependent tests remain skipped because their
+original four-kernel test binary is unavailable; they are not counted as passes.
+
+GPU checks used an explicit `(uint32_t* input, uint32_t* output, uint32_t n)` PTX
+fixture, 65,536 elements, 256 blocks × 256 threads, no dynamic shared memory,
+and full modulo-2^32 output comparison against `input * 3 + 7`. Each process
+verified its actual device UUID against its assigned UUID. Published outputs
+were reopened after cleanup and matched the input hash. These runs used
+`max_iters=0` and the integrated A1/A2/A3 tree with a test-only numerical runner.
+They establish file ownership and numerical survival, not search quality.
+
+RTX 5090 (SM120), driver 580.82.07, CUDA 13.0.88. The existing fixed-ABI
+`sass_test` harness was not treated as a generic numerical oracle. Validation
+subprocesses exited naturally; their timeout/termination paths were disabled in
+the test adapter. Production timeout behavior was not changed by this PR.

src/sasskit/recolor/cli.py

@@ -297,6 +297,8 @@ def cmd_reforge(args: argparse.Namespace) -> int:
         bench_smem=args.smem,
         temperature=args.temperature,
         seed=args.seed,
+        output_path=args.output,
+        work_dir=args.work_dir,
     )
 
     return 0
@@ -761,6 +763,9 @@ def main() -> int:
     p_reforge.add_argument('--smem', type=int, default=28672)
     p_reforge.add_argument('--temperature', type=float, default=0.1)
     p_reforge.add_argument('--seed', type=int, default=42)
+    p_reforge.add_argument('-o', '--output',
+                           help='Output cubin (no overwrite; default: ./reforge-results/<unique-run>/best.cubin)')
+    p_reforge.add_argument('--work-dir', help='Parent directory for private temporary workspaces')
 
     # --- forge commands ---
     try:

src/sasskit/schedule/reforge.py

@@ -16,7 +16,11 @@
 
 from __future__ import annotations
 
+import math
+import os
 import random
+import shutil
+import tempfile
 import struct
 import subprocess
 import sys
@@ -71,6 +75,7 @@ class ReforgeState:
     accepted: int = 0
     rejected: int = 0
     history: list[tuple[int, str, float]] = field(default_factory=list)
+    best_cubin_path: str | None = None
 
 
 def _is_memory_load(inst: Instruction) -> bool:
@@ -311,14 +316,70 @@ def gpu_bench(cubin_path: str, kernel_name: str,
 ]
 
 
+def _publish_best(source: Path, output: Path) -> str:
+    """Publish complete bytes atomically without replacing an existing result."""
+    with tempfile.NamedTemporaryFile(dir=output.parent, prefix='.reforge-',
+                                     delete=False) as staging:
+        staging_path = Path(staging.name)
+    try:
+        shutil.copyfile(source, staging_path)
+        os.link(staging_path, output)
+    finally:
+        staging_path.unlink()
+    return str(output.absolute())
+
+
 def reforge(cubin_path: str, kernel_name: str,
             max_iters: int = 200,
             bench_blocks: int = 1, bench_threads: int = 256,
             bench_smem: int = 28672,
             temperature: float = 0.1,
             cooling: float = 0.995,
             seed: int = 42,
-            verbose: bool = True) -> ReforgeState:
+            verbose: bool = True, *,
+            output_path: str | Path | None = None,
+            work_dir: str | Path | None = None) -> ReforgeState:
+    """Optimize in a private workspace and return a durable, non-overwritten best.
+
+    ``work_dir`` is the parent of a unique temporary directory. By default,
+    results are published in ``./reforge-results/<unique-run>/best.cubin``.
+    Explicit output parents must already exist. Input files are never modified.
+    """
+    output = Path(output_path).absolute() if output_path is not None else None
+    if output is not None:
+        if output.resolve() == Path(cubin_path).resolve() or (
+                output.exists() and output.samefile(cubin_path)):
+            raise ValueError('Output must not alias the input cubin')
+        if output.exists():
+            raise FileExistsError(output)
+
+    with tempfile.TemporaryDirectory(prefix='reforge-', dir=work_dir) as workspace:
+        current = Path(workspace) / 'current.cubin'
+        best = Path(workspace) / 'best.cubin'
+        state = _reforge(cubin_path, kernel_name, max_iters, bench_blocks,
+                         bench_threads, bench_smem, temperature, cooling, seed,
+                         verbose, tmp_path=str(current), best_path=str(best))
+        if not math.isfinite(state.best_time_ms) or state.best_time_ms <= 0:
+            raise ValueError('No valid benchmark result to publish')
+        if output is None:
+            results = Path.cwd() / 'reforge-results'
+            results.mkdir(exist_ok=True)
+            output = Path(tempfile.mkdtemp(prefix='run-', dir=results)) / 'best.cubin'
+        state.best_cubin_path = _publish_best(best, output)
+    if verbose:
+        print(f"  Saved:    {state.best_cubin_path}", file=sys.stderr)
+    return state
+
+
+def _reforge(cubin_path: str, kernel_name: str,
+            max_iters: int = 200,
+            bench_blocks: int = 1, bench_threads: int = 256,
+            bench_smem: int = 28672,
+            temperature: float = 0.1,
+            cooling: float = 0.995,
+            seed: int = 42,
+            verbose: bool = True, *,
+            tmp_path: str, best_path: str) -> ReforgeState:
     """Run the SASS-to-SASS optimization loop.
 
     Args:
@@ -348,8 +409,6 @@ def reforge(cubin_path: str, kernel_name: str,
     )
 
     # Baseline benchmark
-    tmp_path = '/tmp/reforge_current.cubin'
-    best_path = '/tmp/reforge_best.cubin'
     cubin.save(tmp_path)
     cubin.save(best_path)
 
@@ -444,6 +503,5 @@ def reforge(cubin_path: str, kernel_name: str,
         print(f"  Baseline: {baseline:.4f} ms", file=sys.stderr)
         print(f"  Best:     {state.best_time_ms:.4f} ms "
               f"({speedup:.3f}x)", file=sys.stderr)
-        print(f"  Saved:    {best_path}", file=sys.stderr)
 
     return state

tests/test_reforge_workspace.py

@@ -0,0 +1,180 @@
+"""File ownership tests; no CUDA runtime is required."""
+import multiprocessing
+from pathlib import Path
+from unittest.mock import patch
+
+import pytest
+
+from sasskit.schedule import reforge as rf
+
+
+class TokenCubin:
+    def __init__(self, path):
+        self.path = Path(path)
+        self.token = self.path.read_bytes()
+
+    @classmethod
+    def from_file(cls, path):
+        return cls(path)
+
+    def get_kernel(self, name):
+        return name
+
+    def save(self, path):
+        Path(path).write_bytes(self.token)
+
+
+def run_token(source, barrier=None, **kwargs):
+    token = Path(source).read_bytes()
+    saved = []
+
+    def save(obj, path):
+        saved.append(Path(path))
+        Path(path).write_bytes(obj.token)
+
+    def bench(path, *args):
+        if barrier is not None:
+            barrier.wait()
+        assert Path(path).read_bytes() == token, 'another run overwrote this cubin'
+        return 1.0
+
+    with patch.object(rf, 'Cubin', TokenCubin), \
+         patch.object(TokenCubin, 'save', save), \
+         patch.object(rf, 'decode_kernel', return_value=[]), \
+         patch.object(rf, 'gpu_bench', bench):
+        state = rf.reforge(str(source), 'test', max_iters=0, verbose=False, **kwargs)
+    return state, saved
+
+
+def worker(source, cwd, barrier):
+    import os
+    os.chdir(cwd)
+    state, saved = run_token(source, barrier)
+    assert Path(state.best_cubin_path).read_bytes() == Path(source).read_bytes()
+    assert not saved[0].parent.exists()
+
+
+@pytest.mark.parametrize('count', [2, 8])
+def test_concurrent_calls_share_cwd_and_temp_root(tmp_path, count):
+    ctx = multiprocessing.get_context('spawn')
+    barrier = ctx.Barrier(count)
+    processes = []
+    for i in range(count):
+        source = tmp_path / f'input-{i}.cubin'
+        source.write_bytes(str(i).encode())
+        process = ctx.Process(target=worker, args=(source, tmp_path, barrier))
+        process.start()
+        processes.append(process)
+    for process in processes:
+        process.join()
+    assert [p.exitcode for p in processes] == [0] * count
+    assert len(list((tmp_path / 'reforge-results').glob('*/best.cubin'))) == count
+
+
+def test_repeated_calls_and_spaces(tmp_path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    source = tmp_path / 'input with spaces.cubin'
+    source.write_bytes(b'original')
+    root = tmp_path / 'temporary root'
+    root.mkdir()
+    first, first_saved = run_token(source, work_dir=root)
+    second, second_saved = run_token(source, work_dir=root)
+    assert first_saved[0].parent != second_saved[0].parent
+    assert all(p.parent == first_saved[0].parent for p in first_saved)
+    assert list(root.iterdir()) == []
+    for state in (first, second):
+        assert Path(state.best_cubin_path).read_bytes() == b'original'
+        assert state.cubin_path == str(source)
+    assert first.best_cubin_path != second.best_cubin_path
+    assert source.read_bytes() == b'original'
+
+
+@pytest.mark.parametrize('alias', ['same', 'symlink', 'hardlink'])
+def test_input_alias_rejected(tmp_path, alias):
+    source = tmp_path / 'input'
+    source.write_bytes(b'original')
+    output = tmp_path / 'output'
+    if alias == 'same':
+        output = source
+    elif alias == 'symlink':
+        output.symlink_to(source)
+    else:
+        output.hardlink_to(source)
+    with pytest.raises(ValueError, match='alias'):
+        run_token(source, output_path=output)
+    assert source.read_bytes() == b'original'
+
+
+@pytest.mark.parametrize('error', [RuntimeError('benchmark failed'), KeyboardInterrupt()])
+def test_failure_cleans_only_owned_workspace(tmp_path, error):
+    source = tmp_path / 'input'
+    source.write_bytes(b'original')
+    other = tmp_path / 'reforge-other'
+    other.mkdir()
+    (other / 'current.cubin').write_bytes(b'other')
+    with patch.object(rf, 'Cubin', TokenCubin), \
+         patch.object(rf, 'decode_kernel', return_value=[]), \
+         patch.object(rf, 'gpu_bench', side_effect=error):
+        with pytest.raises(type(error)):
+            rf.reforge(str(source), 'test', work_dir=tmp_path, verbose=False)
+    assert sorted(p.name for p in tmp_path.iterdir()) == ['input', 'reforge-other']
+    assert (other / 'current.cubin').read_bytes() == b'other'
+
+
+def test_publish_failure_and_collision(tmp_path):
+    source = tmp_path / 'input'
+    source.write_bytes(b'original')
+    output = tmp_path / 'output'
+    with patch.object(rf.os, 'link', side_effect=OSError('publication failed')):
+        with pytest.raises(OSError, match='publication failed'):
+            run_token(source, output_path=output, work_dir=tmp_path)
+    assert sorted(p.name for p in tmp_path.iterdir()) == ['input']
+    output.write_bytes(b'preexisting')
+    with pytest.raises(FileExistsError):
+        run_token(source, output_path=output)
+    assert output.read_bytes() == b'preexisting'
+    assert source.read_bytes() == b'original'
+
+
+def publish_worker(source, output, barrier, results):
+    barrier.wait()
+    try:
+        rf._publish_best(Path(source), Path(output))
+        results.put('published')
+    except FileExistsError:
+        results.put('conflict')
+
+
+def test_atomic_publish_collision(tmp_path):
+    ctx = multiprocessing.get_context('spawn')
+    barrier, results = ctx.Barrier(2), ctx.Queue()
+    output = tmp_path / 'shared-output'
+    processes = []
+    for i in range(2):
+        source = tmp_path / f'input-{i}'
+        source.write_bytes(str(i).encode() * 65536)
+        process = ctx.Process(target=publish_worker, args=(source, output, barrier, results))
+        process.start()
+        processes.append(process)
+    for process in processes:
+        process.join()
+    assert [p.exitcode for p in processes] == [0, 0]
+    assert sorted(results.get() for _ in processes) == ['conflict', 'published']
+    assert output.read_bytes() in (b'0' * 65536, b'1' * 65536)
+    assert not list(tmp_path.glob('.reforge-*'))
+
+
+def test_cli_defaults_and_options(tmp_path, monkeypatch, capsys):
+    from sasskit.recolor import cli
+    source = tmp_path / 'input'
+    source.write_bytes(b'original')
+    with patch.object(cli.Cubin, 'from_file') as load, patch.object(rf, 'reforge') as run:
+        load.return_value.kernels = {'test': object()}
+        monkeypatch.setattr('sys.argv', ['sasskit', 'reforge', str(source)])
+        assert cli.main() == 0
+        assert run.call_args.kwargs['output_path'] is None
+        assert run.call_args.kwargs['work_dir'] is None
+    monkeypatch.setattr('sys.argv', ['sasskit', 'reforge', '--help'])
+    with pytest.raises(SystemExit, match='0'):
+        cli.main()
+    assert 'reforge-results' in ''.join(capsys.readouterr().out.split())