paasta-tools 1.16.2__tar.gz → 1.16.4__tar.gz

{paasta-tools-1.16.2/paasta_tools.egg-info → paasta-tools-1.16.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.2
 Name: paasta-tools
-Version: 1.16.2
+Version: 1.16.4
 Summary: Tools for Yelps SOA infrastructure
 Home-page: UNKNOWN
 Author: Compute Infrastructure @ Yelp

{paasta-tools-1.16.2 → paasta-tools-1.16.4}/paasta_tools/__init__.py

@@ -17,4 +17,4 @@
 # setup phase, the dependencies may not exist on disk yet.
 #
 # Don't bump version manually. See `make release` docs in ./Makefile
-__version__ = "1.16.2"
+__version__ = "1.16.4"

{paasta-tools-1.16.2 → paasta-tools-1.16.4}/paasta_tools/cli/cmds/remote_run.py

@@ -13,13 +13,12 @@
 # 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.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.cli.utils import run_interactive_cli
 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
@@ -96,7 +95,7 @@ def paasta_remote_run_start(
         pod=poll_response.pod_name,
         token=token_response.token,
     )
-    pty.spawn(shlex.split(exec_command))
+    run_interactive_cli(exec_command)
     return 0
 
 

{paasta-tools-1.16.2 → paasta-tools-1.16.4}/paasta_tools/cli/cmds/spark_run.py

@@ -6,6 +6,7 @@ import re
 import shlex
 import socket
 import sys
+from configparser import ConfigParser
 from typing import Any
 from typing import cast
 from typing import Dict
@@ -84,6 +85,14 @@ DEPRECATED_OPTS = {
 
 SPARK_COMMANDS = {"pyspark", "spark-submit"}
 
+# config looks as follows:
+# [default]
+# aws_access_key_id = ...
+# aws_secret_access_key = ...
+SPARK_DRIVER_IAM_USER = (
+    "/nail/etc/spark_driver_k8s_role_assumer/spark_driver_k8s_role_assumer.ini"
+)
+
 log = logging.getLogger(__name__)
 
 
@@ -154,6 +163,11 @@ def add_subparser(subparsers):
         help=argparse.SUPPRESS,
         action=DeprecatedAction,
     )
+    list_parser.add_argument(
+        "--get-eks-token-via-iam-user",
+        help="Use IAM user to get EKS token for long running spark-run jobs",
+        action="store_true",
+    )
 
     group = list_parser.add_mutually_exclusive_group()
     group.add_argument(
@@ -648,7 +662,22 @@ def get_spark_env(
         spark_env["SPARK_DAEMON_CLASSPATH"] = "/opt/spark/extra_jars/*"
         spark_env["SPARK_NO_DAEMONIZE"] = "true"
 
-    spark_env["KUBECONFIG"] = system_paasta_config.get_spark_kubeconfig()
+    if args.get_eks_token_via_iam_user:
+        with open(SPARK_DRIVER_IAM_USER) as f:
+            config = ConfigParser()
+            config.read_file(f)
+
+        # these env variables are consumed by a script specified in the spark kubeconfig - and which will result in a tightly-scoped IAM identity being used for EKS cluster access
+        spark_env["GET_EKS_TOKEN_AWS_ACCESS_KEY_ID"] = config["default"][
+            "aws_access_key_id"
+        ]
+        spark_env["GET_EKS_TOKEN_AWS_SECRET_ACCESS_KEY"] = config["default"][
+            "aws_secret_access_key"
+        ]
+
+        spark_env["KUBECONFIG"] = system_paasta_config.get_spark_iam_user_kubeconfig()
+    else:
+        spark_env["KUBECONFIG"] = system_paasta_config.get_spark_kubeconfig()
 
     return spark_env
 
@@ -838,9 +867,9 @@ def configure_and_run_docker_container(
     if pod_template_path:
         volumes.append(f"{pod_template_path}:{pod_template_path}:rw")
 
-    volumes.append(
-        f"{system_paasta_config.get_spark_kubeconfig()}:{system_paasta_config.get_spark_kubeconfig()}:ro"
-    )
+    # NOTE: we mount a directory here since the kubeconfig we're transitioning to requires a helper script that will co-exist in the same directory
+    kubeconfig_dir = os.path.dirname(system_paasta_config.get_spark_kubeconfig())
+    volumes.append(f"{kubeconfig_dir}:{kubeconfig_dir}:ro")
 
     environment = instance_config.get_env_dictionary()  # type: ignore
     spark_conf_str = create_spark_config_str(spark_conf, is_mrjob=args.mrjob)
@@ -1148,6 +1177,12 @@ def update_args_from_tronfig(args: argparse.Namespace) -> Optional[Dict[str, str
 
 
 def paasta_spark_run(args: argparse.Namespace) -> int:
+    if args.get_eks_token_via_iam_user and os.getuid() != 0:
+        print("Re-executing paasta spark-run with sudo..", file=sys.stderr)
+        # argv[0] is treated as command name, so prepending "sudo"
+        os.execvp("sudo", ["sudo"] + sys.argv)
+        return  # will not reach unless above function is mocked
+
     driver_envs_from_tronfig: Dict[str, str] = dict()
     if args.tronfig is not None:
         if args.job_id is None:

{paasta-tools-1.16.2 → paasta-tools-1.16.4}/paasta_tools/cli/utils.py

@@ -18,8 +18,10 @@ import getpass
 import hashlib
 import logging
 import os
+import pty
 import random
 import re
+import shutil
 import socket
 import subprocess
 from collections import defaultdict
@@ -1150,3 +1152,20 @@ def get_paasta_oapi_client_with_auth(
         http_res=http_res,
         auth_token=get_sso_service_auth_token(),
     )
+
+
+def run_interactive_cli(cmd: str, shell: str = "bash", term: str = "xterm-256color"):
+    """Runs interactive command in a pseudo terminal, handling terminal size management
+
+    :param str cmd: shell command
+    :param str shell: shell utility to use as wrapper
+    :param str term: terminal type
+    """
+    cols, rows = shutil.get_terminal_size()
+    wrapped_cmd = (
+        f"export SHELL={shell};"
+        f"export TERM={term};"
+        f"stty columns {cols} rows {rows};"
+        f"exec {cmd}"
+    )
+    pty.spawn([shell, "-c", wrapped_cmd])

{paasta-tools-1.16.2 → paasta-tools-1.16.4}/paasta_tools/utils.py

@@ -2036,6 +2036,7 @@ class SystemPaastaConfigDict(TypedDict, total=False):
     skip_cpu_burst_validation: List[str]
     tron_default_pool_override: str
     spark_kubeconfig: str
+    spark_iam_user_kubeconfig: str
     kube_clusters: Dict
     spark_use_eks_default: bool
     sidecar_requirements_config: Dict[str, KubeContainerResourceRequest]
@@ -2758,6 +2759,11 @@ class SystemPaastaConfig:
     def get_spark_kubeconfig(self) -> str:
         return self.config_dict.get("spark_kubeconfig", "/etc/kubernetes/spark.conf")
 
+    def get_spark_iam_user_kubeconfig(self) -> str:
+        return self.config_dict.get(
+            "spark_iam_user_kubeconfig", "/etc/kubernetes/spark2.conf"
+        )
+
     def get_kube_clusters(self) -> Dict:
         return self.config_dict.get("kube_clusters", {})
 

{paasta-tools-1.16.2 → paasta-tools-1.16.4/paasta_tools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 1.2
 Name: paasta-tools
-Version: 1.16.2
+Version: 1.16.4
 Summary: Tools for Yelps SOA infrastructure
 Home-page: UNKNOWN
 Author: Compute Infrastructure @ Yelp