From 5fe46d2dcae17d66f5b89769a453f2820bc5fca8 Mon Sep 17 00:00:00 2001 From: Raja Mukerji Date: Mon, 7 Sep 2026 09:43:39 -0700 Subject: [PATCH 1/2] feat(cli): expose --layer-selection so layer strategy is reachable without the Python API --- obliteratus/cli.py | 17 ++++ obliteratus/remote.py | 2 + tests/test_cli_layer_selection.py | 125 ++++++++++++++++++++++++++++++ 3 files changed, 144 insertions(+) create mode 100644 tests/test_cli_layer_selection.py diff --git a/obliteratus/cli.py b/obliteratus/cli.py index 15feb5a..4b28faa 100644 --- a/obliteratus/cli.py +++ b/obliteratus/cli.py @@ -325,6 +325,15 @@ def main(argv: list[str] | None = None): ) p.add_argument("--regularization", type=float, default=None, help="Override: fraction to preserve (0.0-1.0)") p.add_argument("--refinement-passes", type=int, default=None, help="Override: number of iterative passes") + p.add_argument( + "--layer-selection", type=str, default=None, + choices=["knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"], + help=( + "Override which layers the method selects. Defaults to the " + "method's own setting (knee_cosmic for most). 'all' matches the " + "Heretic method and lets per-layer weights do the narrowing." + ), + ) p.add_argument( "--min-layer-fraction", type=float, default=None, help="Optional layer floor as fraction of depth; e.g. 0.75 keeps only the final quarter.", @@ -475,6 +484,10 @@ def main(argv: list[str] | None = None): si_parser.add_argument("--n-directions", type=int, default=None) si_parser.add_argument("--regularization", type=float, default=None) si_parser.add_argument("--refinement-passes", type=int, default=None) + si_parser.add_argument( + "--layer-selection", type=str, default=None, + choices=["knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"], + ) si_parser.add_argument("--min-layer-fraction", type=float, default=None) si_parser.add_argument("--max-layer-fraction", type=float, default=None) si_parser.add_argument("--harmless-pc-count", type=int, default=None) @@ -1001,6 +1014,7 @@ def _cmd_self_improve(args): direction_method=args.direction_method, regularization=regularization, refinement_passes=refinement_passes, + layer_selection=getattr(args, "layer_selection", None), min_layer_fraction=args.min_layer_fraction, max_layer_fraction=args.max_layer_fraction, harmless_pc_count=args.harmless_pc_count, @@ -1499,6 +1513,7 @@ def _cmd_abliterate(args): direction_method=getattr(args, "direction_method", None), regularization=args.regularization, refinement_passes=args.refinement_passes, + layer_selection=getattr(args, "layer_selection", None), min_layer_fraction=getattr(args, "min_layer_fraction", None), max_layer_fraction=getattr(args, "max_layer_fraction", None), harmless_pc_count=getattr(args, "harmless_pc_count", None), @@ -1825,6 +1840,8 @@ def _cmd_remote_abliterate(args): kwargs["regularization"] = args.regularization if args.refinement_passes is not None: kwargs["refinement_passes"] = args.refinement_passes + if getattr(args, "layer_selection", None) is not None: + kwargs["layer_selection"] = args.layer_selection if getattr(args, "min_layer_fraction", None) is not None: kwargs["min_layer_fraction"] = args.min_layer_fraction if getattr(args, "max_layer_fraction", None) is not None: diff --git a/obliteratus/remote.py b/obliteratus/remote.py index 96a9f7e..bcdd325 100644 --- a/obliteratus/remote.py +++ b/obliteratus/remote.py @@ -312,6 +312,7 @@ class RemoteRunner: refinement_passes: int | None = None, large_model: bool = False, verify_sample_size: int | None = None, + layer_selection: str | None = None, min_layer_fraction: float | None = None, max_layer_fraction: float | None = None, harmless_pc_count: int | None = None, @@ -351,6 +352,7 @@ class RemoteRunner: if verify_sample_size is not None: parts.extend(["--verify-sample-size", str(verify_sample_size)]) optional_values = ( + ("--layer-selection", layer_selection), ("--min-layer-fraction", min_layer_fraction), ("--max-layer-fraction", max_layer_fraction), ("--harmless-pc-count", harmless_pc_count), diff --git a/tests/test_cli_layer_selection.py b/tests/test_cli_layer_selection.py new file mode 100644 index 0000000..360142e --- /dev/null +++ b/tests/test_cli_layer_selection.py @@ -0,0 +1,125 @@ +"""CPU-safe contract tests for the ``--layer-selection`` CLI override. + +``AbliterationPipeline`` has accepted ``layer_selection`` since layer selection +was made configurable, and ``_distill`` dispatches on six distinct values, but no +command exposed it — reaching anything other than a method's built-in default +required constructing the pipeline in Python. These tests pin the flag to the +values the implementation actually dispatches on, so the two cannot drift. +""" + +from __future__ import annotations + +import re +import shlex + +import pytest + +from obliteratus import cli +from obliteratus.remote import RemoteConfig, RemoteRunner + +# The values `AbliterationPipeline._distill` branches on. "knee_cosmic" is the +# implicit default (the `else` arm) and is spelled out here so selecting it +# explicitly is possible rather than only reachable by omission. +EXPECTED_CHOICES = {"knee_cosmic", "knee", "all", "all_except_first", "middle60", "top_k"} + + +@pytest.mark.parametrize("command", ["obliterate", "abliterate"]) +def test_layer_selection_rejects_an_unknown_strategy(command, monkeypatch): + """An unrecognised value must fail at parse time, not fall through silently. + + `_distill` treats every unknown value as the default `knee_cosmic` arm, so + without constrained choices a typo would run a different strategy than the + one asked for and report success. + """ + monkeypatch.setattr(cli, "_cmd_abliterate", lambda _args: None) + + with pytest.raises(SystemExit) as excinfo: + cli.main([command, "org/model", "--layer-selection", "not-a-strategy"]) + + assert excinfo.value.code == 2 + + +@pytest.mark.parametrize("command", ["obliterate", "abliterate"]) +@pytest.mark.parametrize("strategy", sorted(EXPECTED_CHOICES)) +def test_layer_selection_parses_and_reaches_the_command(command, strategy, monkeypatch): + captured = {} + monkeypatch.setattr(cli, "_cmd_abliterate", lambda args: captured.update(vars(args))) + + cli.main([command, "org/model", "--layer-selection", strategy]) + + assert captured["layer_selection"] == strategy + + +def test_layer_selection_defaults_to_none_so_the_method_keeps_its_own_setting(monkeypatch): + """Omitting the flag must not override the method's configured strategy. + + The pipeline resolves `layer_selection or method_cfg[...]`, so passing + anything other than None here would silently flatten every method's default. + """ + captured = {} + monkeypatch.setattr(cli, "_cmd_abliterate", lambda args: captured.update(vars(args))) + + cli.main(["obliterate", "org/model"]) + + assert captured["layer_selection"] is None + + +def test_self_improve_accepts_the_same_strategies(monkeypatch, tmp_path): + captured = {} + monkeypatch.setattr(cli, "_cmd_self_improve", lambda args: captured.update(vars(args))) + + cli.main([ + "self-improve", "org/model", + "--audit", str(tmp_path / "audit.json"), + "--output-dir", str(tmp_path / "out"), + "--layer-selection", "all", + ]) + + assert captured["layer_selection"] == "all" + + +def test_remote_command_forwards_layer_selection(): + runner = RemoteRunner(RemoteConfig(host="example.invalid", user="operator")) + + command = runner.build_obliterate_command( + "org/model", + method="optimized", + layer_selection="all", + ) + + tokens = shlex.split(command) + assert "--layer-selection" in tokens + assert tokens[tokens.index("--layer-selection") + 1] == "all" + + +def test_remote_command_omits_the_flag_when_unset(): + runner = RemoteRunner(RemoteConfig(host="example.invalid", user="operator")) + + command = runner.build_obliterate_command("org/model", method="optimized") + + assert "--layer-selection" not in shlex.split(command) + + +def test_every_offered_strategy_is_one_distill_actually_dispatches_on(): + """Guard against the flag and the implementation drifting apart. + + A choice the CLI offers but `_distill` does not branch on would fall into the + default arm and silently run `knee_cosmic` instead — the failure this flag + exists to make impossible. + """ + import inspect + + from obliteratus.abliterate import AbliterationPipeline + + source = inspect.getsource(AbliterationPipeline._distill) + + # knee_cosmic is the implicit `else` arm and so has no equality branch. + for strategy in sorted(EXPECTED_CHOICES - {"knee_cosmic"}): + assert f'selection_method == "{strategy}"' in source, ( + f"CLI offers {strategy!r} but _distill has no branch for it" + ) + + branches = set(re.findall(r'selection_method == "([a-z_0-9]+)"', source)) + assert branches | {"knee_cosmic"} == EXPECTED_CHOICES, ( + f"_distill dispatches on {branches | {'knee_cosmic'}}, CLI offers {EXPECTED_CHOICES}" + ) From bbf06fd6c86a7cf40efa2daad06fdedc697a62e8 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Mon, 7 Sep 2026 16:20:38 -0400 Subject: [PATCH 2/2] test(cli): verify layer selection reaches local and remote execution --- tests/test_cli_boundaries.py | 17 ++++++++++++----- tests/test_cli_layer_selection.py | 22 ++++++++++++++++++++++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/tests/test_cli_boundaries.py b/tests/test_cli_boundaries.py index 0133107..ac2d606 100644 --- a/tests/test_cli_boundaries.py +++ b/tests/test_cli_boundaries.py @@ -550,7 +550,8 @@ def test_remote_cli_accepts_valid_ssh_port_and_dispatches(monkeypatch): assert dispatch.call_args.args[0].ssh_port == 2222 -def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp_path): +@pytest.mark.parametrize("layer_selection", [None, "all", "middle60"]) +def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp_path, layer_selection): import obliteratus.abliterate import obliteratus.community import obliteratus.hard_negative @@ -604,9 +605,11 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp quantization=None, gpu_memory_utilization=0.95, large_model=False, verify_sample_size=3, refusal_max_tokens=512, residue_file=["audit.json"], dataset="builtin", residue_weight=2, residue_max=4, - contribute=True, contribute_notes="fixture", + contribute=True, contribute_notes="fixture", layer_selection=layer_selection, ) cli._cmd_abliterate(args) + assert factory.call_args.kwargs["layer_selection"] == layer_selection + pipeline.run.assert_called_once_with() assert factory.call_args.kwargs["refusal_max_tokens"] == 512 assert factory.call_args.kwargs["gpu_memory_utilization"] == 0.95 assert factory.call_args.kwargs["trust_remote_code"] is False @@ -614,7 +617,8 @@ def test_abliterate_pipeline_callbacks_residue_and_contribution(monkeypatch, tmp telemetry.assert_called_once_with(pipeline) -def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path): +@pytest.mark.parametrize("layer_selection", [None, "all", "middle60"]) +def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path, layer_selection): import obliteratus.abliterate import obliteratus.hard_negative import obliteratus.model_profile @@ -651,7 +655,7 @@ def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path): direction_method="svd", min_layer_fraction=None, max_layer_fraction=None, harmless_pc_count=None, shield_concept_count=None, shield_ridge=None, shield_residualize=None, shield_layer_penalty=None, projection_target=None, - device="cpu", dry_run=True, + device="cpu", dry_run=True, layer_selection=layer_selection, ) cli._cmd_self_improve(args) assert (output / "self_improve_plan.json").is_file() @@ -660,7 +664,10 @@ def test_self_improve_dry_run_and_pipeline(monkeypatch, tmp_path): result.mkdir() pipeline = MagicMock() pipeline.run.return_value = str(result) - monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", Mock(return_value=pipeline)) + factory = Mock(return_value=pipeline) + monkeypatch.setattr(obliteratus.abliterate, "AbliterationPipeline", factory) args.dry_run = False cli._cmd_self_improve(args) + assert factory.call_args.kwargs["layer_selection"] == layer_selection + pipeline.run.assert_called_once_with() assert (result / "hard_negative_residue.json").is_file() diff --git a/tests/test_cli_layer_selection.py b/tests/test_cli_layer_selection.py index 360142e..97d3edc 100644 --- a/tests/test_cli_layer_selection.py +++ b/tests/test_cli_layer_selection.py @@ -123,3 +123,25 @@ def test_every_offered_strategy_is_one_distill_actually_dispatches_on(): assert branches | {"knee_cosmic"} == EXPECTED_CHOICES, ( f"_distill dispatches on {branches | {'knee_cosmic'}}, CLI offers {EXPECTED_CHOICES}" ) + + +@pytest.mark.parametrize("strategy", [None, "all", "middle60"]) +def test_remote_cli_preserves_override_and_method_default(strategy, monkeypatch): + from unittest.mock import Mock + + runner = Mock() + runner.run_obliterate.return_value = "result" + monkeypatch.setattr(cli, "_make_remote_runner", lambda args: runner) + argv = ["obliterate", "org/model", "--remote", "operator@example.invalid", + "--method", "optimized"] + if strategy is not None: + argv += ["--layer-selection", strategy] + + cli.main(argv) + + forwarded = runner.run_obliterate.call_args.kwargs + assert forwarded["method"] == "optimized" + if strategy is None: + assert "layer_selection" not in forwarded + else: + assert forwarded["layer_selection"] == strategy