cleanup for remote-run resources (#4024)

Clean up logic for the (meant to be) ephemeral resources which are created in remote-run invocations (#4022).

Author: piax93

debian/paasta-tools.links

@@ -16,6 +16,7 @@ opt/venvs/paasta-tools/bin/generate_authenticating_services.py usr/bin/generate_
 opt/venvs/paasta-tools/bin/kubernetes_remove_evicted_pods.py usr/bin/kubernetes_remove_evicted_pods
 opt/venvs/paasta-tools/bin/paasta-api usr/bin/paasta-api
 opt/venvs/paasta-tools/bin/paasta-fsm usr/bin/paasta-fsm
+opt/venvs/paasta-tools/bin/paasta_cleanup_remote_run_resources.py usr/bin/paasta_cleanup_remote_run_resources
 opt/venvs/paasta-tools/bin/paasta_cleanup_stale_nodes.py usr/bin/paasta_cleanup_stale_nodes
 opt/venvs/paasta-tools/bin/paasta_prune_completed_pods usr/bin/paasta_prune_completed_pods
 opt/venvs/paasta-tools/bin/paasta_cleanup_tron_namespaces usr/bin/paasta_cleanup_tron_namespaces

paasta_tools/kubernetes/bin/paasta_cleanup_remote_run_resources.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+# Copyright 2015-2019 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import argparse
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+from typing import Any
+from typing import Callable
+from typing import Sequence
+from typing import Tuple
+
+from paasta_tools.kubernetes.remote_run import get_remote_run_role_bindings
+from paasta_tools.kubernetes.remote_run import get_remote_run_roles
+from paasta_tools.kubernetes.remote_run import get_remote_run_service_accounts
+from paasta_tools.kubernetes_tools import get_all_managed_namespaces
+from paasta_tools.kubernetes_tools import KubeClient
+
+
+ListingFuncType = Callable[[KubeClient, str], Sequence[Any]]
+DeletionFuncType = Callable[[str, str], Any]
+
+
+def clean_namespace(kube_client: KubeClient, namespace: str, age_limit: datetime):
+    """Clean ephemeral remote-run resource in a namespace
+
+    :param KubeClient kube_client: kubernetes client
+    :param str namepsace: kubernetes namespace
+    :param datetime age_limit: expiration time for resources
+    """
+    cleanup_actions: Sequence[Tuple[DeletionFuncType, ListingFuncType]] = (
+        (
+            kube_client.core.delete_namespaced_service_account,
+            get_remote_run_service_accounts,
+        ),
+        (kube_client.rbac.delete_namespaced_role, get_remote_run_roles),
+        (kube_client.rbac.delete_namespaced_role_binding, get_remote_run_role_bindings),
+    )
+    for delete_func, list_func in cleanup_actions:
+        for entity in list_func(kube_client, namespace):
+            if (
+                not entity.metadata.name.startswith("remote-run-")
+                or entity.metadata.creation_timestamp > age_limit
+            ):
+                continue
+            delete_func(entity.metadata.name, namespace)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Clean ephemeral Kubernetes resources created by remote-run invocations",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--max-age",
+        type=int,
+        default=600,
+        help="Maximum age, in seconds, resources are allowed to have",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    kube_client = KubeClient()
+    age_limit = datetime.now(tzinfo=timezone.utc) - timedelta(seconds=args.max_age)
+    for namespace in get_all_managed_namespaces(kube_client):
+        clean_namespace(kube_client, namespace, age_limit)
+
+
+if __name__ == "__main__":
+    main()

paasta_tools/kubernetes/remote_run.py

@@ -14,7 +14,9 @@
 import hashlib
 import logging
 from time import sleep
+from typing import List
 from typing import Optional
+from typing import Sequence
 from typing import TypedDict
 
 from kubernetes.client import AuthenticationV1TokenRequest
@@ -236,7 +238,7 @@ def remote_run_token(
         kube_client, namespace, pod_name, user
     )
     role = create_pod_scoped_role(kube_client, namespace, pod_name, user)
-    bind_role_to_service_account(kube_client, namespace, service_account, role)
+    bind_role_to_service_account(kube_client, namespace, service_account, role, user)
     return create_temp_exec_token(kube_client, namespace, service_account)
 
 
@@ -293,6 +295,23 @@ def create_temp_exec_token(
     return response.status.token
 
 
+def get_remote_run_service_accounts(
+    kube_client: KubeClient, namespace: str, user: str = ""
+) -> Sequence[V1ServiceAccount]:
+    """List all temporary service account related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :param str user: optionally filter by owning user
+    :return: list of service accounts
+    """
+    return get_all_service_accounts(
+        kube_client,
+        namespace=namespace,
+        label_selector=(f"{POD_OWNER_LABEL}={user}" if user else POD_OWNER_LABEL),
+    )
+
+
 def create_remote_run_service_account(
     kube_client: KubeClient,
     namespace: str,
@@ -308,11 +327,7 @@ def create_remote_run_service_account(
     """
     pod_name_hash = hashlib.sha1(pod_name.encode("utf-8")).hexdigest()[:12]
     service_account_name = limit_size_with_hash(f"remote-run-{user}-{pod_name_hash}")
-    service_accounts = get_all_service_accounts(
-        kube_client,
-        namespace=namespace,
-        label_selector=f"{POD_OWNER_LABEL}={user}",
-    )
+    service_accounts = get_remote_run_service_accounts(kube_client, namespace, user)
     if any(item.metadata.name == service_account_name for item in service_accounts):
         return service_account_name
     service_account = V1ServiceAccount(
@@ -366,18 +381,21 @@ def bind_role_to_service_account(
     namespace: str,
     service_account: str,
     role: str,
+    user: str,
 ) -> None:
     """Bind service account to role
 
     :param KubeClient kube_client: Kubernetes client
     :param str namespace: service account namespace
     :param str service_account: service account name
     :param str role: role name
+    :param str user: user requiring the role
     """
     role_binding = V1RoleBinding(
         metadata=V1ObjectMeta(
-            name=limit_size_with_hash(f"binding-{role}"),
+            name=limit_size_with_hash(f"remote-run-binding-{role}"),
             namespace=namespace,
+            labels={POD_OWNER_LABEL: user},
         ),
         role_ref=V1RoleRef(
             api_group="rbac.authorization.k8s.io",
@@ -395,3 +413,29 @@ def bind_role_to_service_account(
         namespace=namespace,
         body=role_binding,
     )
+
+
+def get_remote_run_roles(kube_client: KubeClient, namespace: str) -> List[V1Role]:
+    """List all temporary roles related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :return: list of roles
+    """
+    return kube_client.rbac.list_namespaced_role(
+        namespace, label_selector=POD_OWNER_LABEL
+    ).items
+
+
+def get_remote_run_role_bindings(
+    kube_client: KubeClient, namespace: str
+) -> List[V1RoleBinding]:
+    """List all temporary role bindings related to remote-run
+
+    :param KubeClient kube_client: Kubernetes client
+    :param str namespace: pod namespace
+    :return: list of roles
+    """
+    return kube_client.rbac.list_namespaced_role_binding(
+        namespace, label_selector=POD_OWNER_LABEL
+    ).items

setup.py

@@ -59,6 +59,7 @@ def get_install_requires():
         "paasta_tools/generate_services_yaml.py",
         "paasta_tools/generate_authenticating_services.py",
         "paasta_tools/kubernetes/bin/kubernetes_remove_evicted_pods.py",
+        "paasta_tools/kubernetes/bin/paasta_cleanup_remote_run_resources.py",
         "paasta_tools/kubernetes/bin/paasta_cleanup_stale_nodes.py",
         "paasta_tools/kubernetes/bin/paasta_secrets_sync.py",
         "paasta_tools/paasta_deploy_tron_jobs",

tests/kubernetes/bin/test_paasta_cleanup_remote_run_resources.py

@@ -0,0 +1,100 @@
+# Copyright 2015-2016 Yelp Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from datetime import datetime
+from unittest.mock import call
+from unittest.mock import MagicMock
+from unittest.mock import patch
+
+from paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources import (
+    clean_namespace,
+)
+from paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources import main
+
+
+def _create_mock_kube_resource(name: str, creation_time: datetime):
+    mock_resource = MagicMock()
+    mock_resource.metadata.name = name
+    mock_resource.metadata.creation_timestamp = creation_time
+    return mock_resource
+
+
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.get_remote_run_role_bindings",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.get_remote_run_roles",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.get_remote_run_service_accounts",
+    autospec=True,
+)
+def test_clean_namespace(mock_get_sa, mock_get_roles, mock_get_bindings):
+    mock_client = MagicMock()
+    mock_get_sa.return_value = [
+        _create_mock_kube_resource("foobar", datetime(2025, 1, 1, 0, 0, 0)),
+        _create_mock_kube_resource("remote-run-abc", datetime(2025, 1, 1, 0, 0, 0)),
+    ]
+    mock_get_roles.return_value = [
+        _create_mock_kube_resource("remote-run-abc", datetime(2025, 1, 1, 2, 0, 0)),
+        _create_mock_kube_resource("remote-run-def", datetime(2025, 1, 1, 0, 0, 0)),
+    ]
+    mock_get_bindings.return_value = [
+        _create_mock_kube_resource("whatever", datetime(2025, 1, 1, 0, 0, 0)),
+    ]
+    clean_namespace(mock_client, "abc", datetime(2025, 1, 1, 1, 1, 1))
+    mock_client.core.delete_namespaced_service_account.assert_has_calls(
+        [call("remote-run-abc", "abc")]
+    )
+    mock_client.rbac.delete_namespaced_role.assert_has_calls(
+        [call("remote-run-def", "abc")]
+    )
+    mock_client.rbac.delete_namespaced_role_binding.assert_not_called()
+
+
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.get_all_managed_namespaces",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.clean_namespace",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.KubeClient",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.parse_args",
+    autospec=True,
+)
+@patch(
+    "paasta_tools.kubernetes.bin.paasta_cleanup_remote_run_resources.datetime",
+    autospec=True,
+)
+def test_main(
+    mock_datetime, mock_parse_args, mock_kube, mock_clean, mock_get_namespaces
+):
+    mock_parse_args.return_value.max_age = 60
+    mock_datetime.now.return_value = datetime(2025, 1, 1, 0, 1, 0)
+    mock_get_namespaces.return_value = ["a", "b", "c"]
+    main()
+    mock_clean.assert_has_calls(
+        [
+            call(mock_kube.return_value, "a", datetime(2025, 1, 1, 0, 0, 0)),
+            call(mock_kube.return_value, "b", datetime(2025, 1, 1, 0, 0, 0)),
+            call(mock_kube.return_value, "c", datetime(2025, 1, 1, 0, 0, 0)),
+        ]
+    )

tests/kubernetes/test_remote_run.py

@@ -32,6 +32,8 @@
 from paasta_tools.kubernetes.remote_run import create_remote_run_service_account
 from paasta_tools.kubernetes.remote_run import create_temp_exec_token
 from paasta_tools.kubernetes.remote_run import find_job_pod
+from paasta_tools.kubernetes.remote_run import get_remote_run_role_bindings
+from paasta_tools.kubernetes.remote_run import get_remote_run_roles
 from paasta_tools.kubernetes.remote_run import remote_run_ready
 from paasta_tools.kubernetes.remote_run import remote_run_start
 from paasta_tools.kubernetes.remote_run import remote_run_stop
@@ -193,7 +195,11 @@ def test_remote_run_token(
         mock_client, "namespace", "remote-run-someuser-somejob-112233", "someuser"
     )
     mock_bind_role.assert_called_once_with(
-        mock_client, "namespace", "somesa", "somerole"
+        mock_client,
+        "namespace",
+        "somesa",
+        "somerole",
+        "someuser",
     )
     mock_create_token.assert_called_once_with(mock_client, "namespace", "somesa")
     # job not found
@@ -320,13 +326,16 @@ def test_create_pod_scoped_role():
 
 def test_bind_role_to_service_account():
     mock_client = MagicMock()
-    bind_role_to_service_account(mock_client, "namespace", "somesa", "somerole")
+    bind_role_to_service_account(
+        mock_client, "namespace", "somesa", "somerole", "someuser"
+    )
     mock_client.rbac.create_namespaced_role_binding.assert_called_once_with(
         namespace="namespace",
         body=V1RoleBinding(
             metadata=V1ObjectMeta(
-                name=f"binding-somerole",
+                name=f"remote-run-binding-somerole",
                 namespace="namespace",
+                labels={"paasta.yelp.com/pod_owner": "someuser"},
             ),
             role_ref=V1RoleRef(
                 api_group="rbac.authorization.k8s.io",
@@ -341,3 +350,19 @@ def test_bind_role_to_service_account():
             ],
         ),
     )
+
+
+def test_get_remote_run_roles():
+    mock_client = MagicMock()
+    get_remote_run_roles(mock_client, "namespace")
+    mock_client.rbac.list_namespaced_role.assert_called_once_with(
+        "namespace", label_selector="paasta.yelp.com/pod_owner"
+    )
+
+
+def test_get_remote_run_role_bindings():
+    mock_client = MagicMock()
+    get_remote_run_role_bindings(mock_client, "namespace")
+    mock_client.rbac.list_namespaced_role_binding.assert_called_once_with(
+        "namespace", label_selector="paasta.yelp.com/pod_owner"
+    )