Merge pull request #45 from UKGovernmentBEIS/craig/uninstall-all-rele…

…ases

Add support for `inspect sandbox cleanup k8s`

CHANGELOG.md

@@ -2,5 +2,6 @@
 
 ## Unreleased
 
+- Add support for `inspect sandbox cleanup k8s` command to uninstall all Inspect Helm charts.
 - Remove use of Inspect's deleted `SANDBOX` log level in favour of `trace_action()` and `trace_message()` functions.
 - Initial release.

src/k8s_sandbox/_helm.py

@@ -70,7 +70,7 @@ async def install(self) -> None:
                         await asyncio.sleep(INSTALL_RETRY_DELAY_SECONDS)
 
     async def uninstall(self, quiet: bool) -> None:
-        await uninstall(self.release_name, quiet)
+        await uninstall(self.release_name, self._namespace, quiet)
 
     async def get_sandbox_pods(self) -> dict[str, Pod]:
         client = k8s_client()
@@ -121,6 +121,9 @@ async def _install(self, upgrade: bool) -> None:
                 # Annotation do not have strict length reqs. Quoting/escaping
                 # handled by asyncio.create_subprocess_exec.
                 f"annotations.inspectTaskName={self.task_name}",
+                # Include a label to identify releases created by Inspect.
+                "--labels",
+                "inspectSandbox=true",
             ]
             + values,
             capture_output=True,
@@ -148,8 +151,7 @@ def _raise_install_error(self, result: ExecResult[str]) -> NoReturn:
         )
 
 
-async def uninstall(release_name: str, quiet: bool) -> None:
-    namespace = get_current_context_namespace()
+async def uninstall(release_name: str, namespace: str, quiet: bool) -> None:
     async with _uninstall_semaphore():
         with inspect_trace_action(
             "K8s uninstall Helm chart", release=release_name, namespace=namespace
@@ -177,6 +179,24 @@ async def uninstall(release_name: str, quiet: bool) -> None:
                 )
 
 
+async def get_all_release_names(namespace: str) -> list[str]:
+    result = await _run_subprocess(
+        "helm",
+        [
+            "list",
+            "--namespace",
+            namespace,
+            "-q",
+            "--selector",
+            "inspectSandbox=true",
+            "--max",
+            "0",
+        ],
+        capture_output=True,
+    )
+    return result.stdout.splitlines()
+
+
 def _raise_runtime_error(
     message: str, from_exception: Exception | None = None, **kwargs: Any
 ) -> NoReturn:

src/k8s_sandbox/_manager.py

@@ -5,10 +5,12 @@
 
 from rich import box, print
 from rich.panel import Panel
+from rich.prompt import Confirm
 from rich.table import Table
 
-from k8s_sandbox._helm import Release
+from k8s_sandbox._helm import Release, get_all_release_names
 from k8s_sandbox._helm import uninstall as helm_uninstall
+from k8s_sandbox._kubernetes_api import get_current_context_namespace
 
 
 class HelmReleaseManager:
@@ -78,13 +80,13 @@ async def uninstall_all(self, print_only: bool) -> None:
 
     def _print_cleanup_instructions(self) -> None:
         table = Table(
-            title="K8s Sandbox Environments (not yet cleaned up):",
+            title="K8s Sandbox Releases (not yet cleaned up):",
             box=box.SQUARE_DOUBLE_HEAD,
             show_lines=True,
             title_style="bold",
             title_justify="left",
         )
-        table.add_column("Container(s)", no_wrap=True)
+        table.add_column("Release(s)", no_wrap=True)
         table.add_column("Cleanup")
         for release in self._installed_releases:
             table.add_row(
@@ -93,11 +95,10 @@ def _print_cleanup_instructions(self) -> None:
             )
         print("")
         print(table)
-        # TODO: Once supported, tell user how to cleanup all environments.
-        # print(
-        #     "\nCleanup all environments with: "
-        #     "[blue]inspect sandbox cleanup k8s[/blue]\n"
-        # )
+        print(
+            "\nCleanup all sandbox releases with: "
+            "[blue]inspect sandbox cleanup k8s[/blue]\n"
+        )
 
 
 async def uninstall_unmanaged_release(release_name: str) -> None:
@@ -108,7 +109,40 @@ async def uninstall_unmanaged_release(release_name: str) -> None:
       release_name (str): The name of the release to uninstall (e.g. "lsphdyup").
     """
     _print_do_not_interrupt()
-    await helm_uninstall(release_name, quiet=False)
+    namespace = get_current_context_namespace()
+    await helm_uninstall(release_name, namespace, quiet=False)
+
+
+async def uninstall_all_unmanaged_releases():
+    def _print_table(releases: list[str]) -> None:
+        print("Releases to be uninstalled:")
+        table = Table(
+            box=box.SQUARE,
+            show_lines=False,
+            title_style="bold",
+            title_justify="left",
+        )
+        table.add_column("Release")
+        for release in releases:
+            table.add_row(f"[red]{release}[/red]")
+        print(table)
+
+    namespace = get_current_context_namespace()
+    releases = await get_all_release_names(namespace)
+    if len(releases) == 0:
+        print(f"No Inspect sandbox releases found in '{namespace}' namespace.")
+        return
+    _print_table(releases)
+    if not Confirm.ask(
+        f"Are you sure you want to uninstall ALL {len(releases)} Inspect sandbox "
+        f"release(s) in '{namespace}' namespace? If this is a shared namespace, "
+        "this may affect other users.",
+    ):
+        print("Cancelled.")
+        return
+    tasks = [helm_uninstall(release, namespace, quiet=False) for release in releases]
+    await asyncio.gather(*tasks, return_exceptions=True)
+    print("Complete.")
 
 
 def _print_do_not_interrupt() -> None:

src/k8s_sandbox/_sandbox_environment.py

@@ -21,6 +21,7 @@
 )
 from k8s_sandbox._manager import (
     HelmReleaseManager,
+    uninstall_all_unmanaged_releases,
     uninstall_unmanaged_release,
 )
 from k8s_sandbox._pod import Pod
@@ -58,13 +59,10 @@ async def task_cleanup(
 
     @classmethod
     async def cli_cleanup(cls, id: str | None) -> None:
-        if id is None:
-            # TODO: How can we avoid uninstalling others' releases? Uninstall all
-            # releases in namespace regardless?
-            raise NotImplementedError(
-                "Cleanup all for k8s is not supported. Please specify a release name."
-            )
-        await uninstall_unmanaged_release(id)
+        if id is not None:
+            await uninstall_unmanaged_release(id)
+        else:
+            await uninstall_all_unmanaged_releases()
 
     @classmethod
     async def sample_init(
@@ -102,8 +100,8 @@ async def sample_cleanup(
         environments: dict[str, SandboxEnvironment],
         interrupted: bool,
     ) -> None:
-        # If we were interrupted, wait unil the end of the task to cleanup (this enables
-        # us to show output for the cleanup operation).
+        # If we were interrupted, wait until the end of the task to cleanup (this
+        # enables us to show output for the cleanup operation).
         if interrupted:
             return
         sandbox: K8sSandboxEnvironment = cast(

test/k8s_sandbox/inspect_integration/test_cleanup.py

@@ -2,8 +2,12 @@
 from unittest.mock import patch
 
 import pytest
+from inspect_ai import Task
+from inspect_ai.model import Model
 
 from k8s_sandbox._helm import uninstall
+from k8s_sandbox._kubernetes_api import get_current_context_namespace
+from k8s_sandbox._sandbox_environment import K8sSandboxEnvironment
 from test.k8s_sandbox.inspect_integration.testing_utils.mock_model import (
     MockToolCallModel,
 )
@@ -17,19 +21,26 @@
 pytestmark = pytest.mark.req_k8s
 
 
-def test_with_cleanup() -> None:
-    model = MockToolCallModel([tool_call("bash", {"cmd": "echo 'success'"})])
-    task = create_task(__file__, target="success")
+@pytest.fixture
+def model() -> Model:
+    return MockToolCallModel.from_tool_call(
+        tool_call("bash", {"cmd": "echo 'success'"})
+    )
 
+
+@pytest.fixture
+def task() -> Task:
+    return create_task(__file__, target="success")
+
+
+def test_with_cleanup(model: Model, task: Task) -> None:
     with patch("k8s_sandbox._helm.uninstall", wraps=uninstall) as spy:
         run_and_verify_inspect_eval(task=task, model=model)
 
     assert spy.call_count == 1
 
 
-def test_without_cleanup() -> None:
-    model = MockToolCallModel([tool_call("bash", {"cmd": "echo 'success'"})])
-    task = create_task(__file__, target="success")
+def test_without_cleanup(model: Model, task: Task) -> None:
     release = "no-clean"
 
     with patch(
@@ -40,4 +51,23 @@ def test_without_cleanup() -> None:
             run_and_verify_inspect_eval(task=task, model=model, sandbox_cleanup=False)
 
     assert spy.call_count == 0
-    asyncio.run(uninstall(release, quiet=False))
+    asyncio.run(uninstall(release, get_current_context_namespace(), quiet=False))
+
+
+def test_cli_cleanup_all_gets_user_confirmation(model: Model, task: Task) -> None:
+    release = "no-clean"
+    with patch(
+        "k8s_sandbox._helm.Release._generate_release_name",
+        return_value=release,
+    ):
+        run_and_verify_inspect_eval(task=task, model=model, sandbox_cleanup=False)
+
+    with patch("k8s_sandbox._helm.uninstall", wraps=uninstall) as spy:
+        # We don't want to actually uninstall all releases in this test (the test could
+        # be run on a production cluster).
+        with patch("rich.prompt.Confirm.ask", return_value=False) as confirm:
+            asyncio.run(K8sSandboxEnvironment.cli_cleanup(id=None))
+
+    assert "Are you sure you want to uninstall ALL" in confirm.call_args.args[0]
+    assert spy.call_count == 0
+    asyncio.run(uninstall(release, get_current_context_namespace(), quiet=False))