the remote-run missing bits (#4028)

This was meant to be included in the last release but I (luisp) don't know how to read and merged things in the "wrong" order. 

This correctly merges in the following PRs:

* cleanup for remote-run resources (#4024)

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

* new remote-run cli (#4025)

Last portion of #3986, with the updated logic to reflect the updated API conventions introduced in #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/cli/cmds/remote_run.py

@@ -13,21 +13,193 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import argparse
+import pty
+import shlex
+import time
 
-from paasta_tools.utils import PaastaColors
+from paasta_tools.cli.utils import get_paasta_oapi_api_clustername
+from paasta_tools.cli.utils import get_paasta_oapi_client_with_auth
+from paasta_tools.cli.utils import lazy_choices_completer
+from paasta_tools.paastaapi.model.remote_run_start import RemoteRunStart
+from paasta_tools.paastaapi.model.remote_run_stop import RemoteRunStop
+from paasta_tools.utils import get_username
+from paasta_tools.utils import list_clusters
+from paasta_tools.utils import list_services
+from paasta_tools.utils import load_system_paasta_config
+from paasta_tools.utils import SystemPaastaConfig
 
 
-def add_subparser(subparsers: argparse._SubParsersAction):
-    subparsers.add_parser(
+KUBECTL_CMD_TEMPLATE = (
+    "kubectl-eks-{cluster} --token {token} exec -it -n {namespace} {pod} -- /bin/bash"
+)
+
+
+def paasta_remote_run_start(
+    args: argparse.Namespace,
+    system_paasta_config: SystemPaastaConfig,
+) -> int:
+    client = get_paasta_oapi_client_with_auth(
+        cluster=get_paasta_oapi_api_clustername(cluster=args.cluster, is_eks=True),
+        system_paasta_config=system_paasta_config,
+    )
+    if not client:
+        print("Cannot get a paasta-api client")
+        return 1
+
+    user = get_username()
+    start_response = client.remote_run.remote_run_start(
+        args.service,
+        args.instance,
+        RemoteRunStart(
+            user=user,
+            interactive=args.interactive,
+            recreate=args.recreate,
+            max_duration=args.max_duration,
+        ),
+    )
+    if start_response.status >= 300:
+        print(f"Error from PaaSTA APIs while starting job: {start_response.message}")
+        return 1
+
+    start_time = time.time()
+    while time.time() - start_time < args.timeout:
+        poll_response = client.remote_run.remote_run_poll(
+            args.service,
+            args.instance,
+            start_response.job_name,
+        )
+        if poll_response.status == 200:
+            break
+        time.sleep(10)
+    else:
+        print("Timed out while waiting for job to start")
+        return 1
+
+    if not args.interactive:
+        print("Successfully started remote-run job")
+        return 0
+
+    token_response = client.remote_run.remote_run_token(
+        args.service, args.instance, user
+    )
+
+    exec_command = KUBECTL_CMD_TEMPLATE.format(
+        cluster=args.cluster,
+        namespace=poll_response.namespace,
+        pod=poll_response.pod_name,
+        token=token_response.token,
+    )
+    pty.spawn(shlex.split(exec_command))
+    return 0
+
+
+def paasta_remote_run_stop(
+    args: argparse.Namespace,
+    system_paasta_config: SystemPaastaConfig,
+) -> int:
+    client = get_paasta_oapi_client_with_auth(
+        cluster=get_paasta_oapi_api_clustername(cluster=args.cluster, is_eks=True),
+        system_paasta_config=system_paasta_config,
+    )
+    if not client:
+        print("Cannot get a paasta-api client")
+        return 1
+    response = client.remote_run.remote_run_stop(
+        args.service, args.instance, RemoteRunStop(user=get_username())
+    )
+    print(response.message)
+    return 0 if response.status < 300 else 1
+
+
+def add_common_args_to_parser(parser: argparse.ArgumentParser):
+    service_arg = parser.add_argument(
+        "-s",
+        "--service",
+        help="The name of the service you wish to inspect. Required.",
+        required=True,
+    )
+    service_arg.completer = lazy_choices_completer(list_services)  # type: ignore
+    parser.add_argument(
+        "-i",
+        "--instance",
+        help=(
+            "Simulate a docker run for a particular instance of the "
+            "service, like 'main' or 'canary'. Required."
+        ),
+        required=True,
+    )
+    cluster_arg = parser.add_argument(
+        "-c",
+        "--cluster",
+        help="The name of the cluster you wish to run your task on. Required.",
+        required=True,
+    )
+    cluster_arg.completer = lazy_choices_completer(list_clusters)  # type: ignore
+
+
+def add_subparser(subparsers: argparse._SubParsersAction) -> None:
+    remote_run_parser = subparsers.add_parser(
         "remote-run",
-        help="Schedule adhoc service sandbox on PaaSTA cluster",
-        description=(
-            "`paasta remote-run` is useful for running adhoc commands in "
-            "context of a service's Docker image."
+        help="Run services / jobs remotely",
+        description="'paasta remote-run' runs services / jobs remotely",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    subparsers = remote_run_parser.add_subparsers(dest="remote_run_command")
+    start_parser = subparsers.add_parser(
+        "start",
+        help="Start or connect to a remote-run job",
+        description="Starts or connects to a remote-run-job",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    start_parser.add_argument(
+        "-I",
+        "--interactive",
+        help=(
+            "Run container in interactive mode. If interactive is set the "
+            'default command will be "bash" unless otherwise set by the "--cmd" flag'
+        ),
+        action="store_true",
+        default=False,
+    )
+    start_parser.add_argument(
+        "-m",
+        "--max-duration",
+        help=(
+            "Amount of time in seconds after which the job is "
+            "automatically stopped (capped by the API backend)"
         ),
+        type=int,
+        default=1800,
+    )
+    start_parser.add_argument(
+        "-r",
+        "--recreate",
+        help="Recreate remote-run job if already existing",
+        action="store_true",
+        default=False,
+    )
+    start_parser.add_argument(
+        "-t",
+        "--timeout",
+        help="Maximum time to wait for a job to start, in seconds",
+        type=int,
+        default=600,
+    )
+    stop_parser = subparsers.add_parser(
+        "stop",
+        help="Stop your remote-run job if it exists",
+        description="Stop your remote-run job if it exists",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
     )
+    add_common_args_to_parser(start_parser)
+    add_common_args_to_parser(stop_parser)
+    remote_run_parser.set_defaults(command=paasta_remote_run)
 
 
-def paasta_remote_run(args: argparse.Namespace):
-    print(PaastaColors.red("Error: functionality under construction"))
-    return 1
+def paasta_remote_run(args: argparse.Namespace) -> int:
+    system_paasta_config = load_system_paasta_config()
+    if args.remote_run_command == "start":
+        return paasta_remote_run_start(args, system_paasta_config)
+    elif args.remote_run_command == "stop":
+        return paasta_remote_run_stop(args, system_paasta_config)
+    raise ValueError(f"Unsupported subcommand: {args.remote_run_command}")

paasta_tools/cli/utils.py

@@ -43,6 +43,8 @@
 
 from paasta_tools import remote_git
 from paasta_tools.adhoc_tools import load_adhoc_job_config
+from paasta_tools.api.client import get_paasta_oapi_client
+from paasta_tools.api.client import PaastaOApiClient
 from paasta_tools.cassandracluster_tools import load_cassandracluster_instance_config
 from paasta_tools.eks_tools import EksDeploymentConfig
 from paasta_tools.eks_tools import load_eks_service_config
@@ -1135,3 +1137,16 @@ def get_sso_service_auth_token() -> str:
     """Generate an authentication token for the calling user from the Single Sign On provider"""
     client_id = load_system_paasta_config().get_service_auth_sso_oidc_client_id()
     return get_and_cache_jwt_default(client_id)
+
+
+def get_paasta_oapi_client_with_auth(
+    cluster: str = None,
+    system_paasta_config: SystemPaastaConfig = None,
+    http_res: bool = False,
+) -> Optional[PaastaOApiClient]:
+    return get_paasta_oapi_client(
+        cluster=cluster,
+        system_paasta_config=system_paasta_config,
+        http_res=http_res,
+        auth_token=get_sso_service_auth_token(),
+    )

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()