paasta-tools 1.21.3__py3-none-any.whl

paasta_tools/spark_tools.py

@@ -0,0 +1,283 @@
+import logging
+import re
+import socket
+import sys
+from typing import Any
+from typing import cast
+from typing import Dict
+from typing import List
+from typing import Mapping
+from typing import Set
+
+from mypy_extensions import TypedDict
+
+from paasta_tools.utils import DockerVolume
+from paasta_tools.utils import PaastaColors
+
+KUBERNETES_NAMESPACE = "paasta-spark"
+DEFAULT_SPARK_SERVICE = "spark"
+DEFAULT_SPARK_RUNTIME_TIMEOUT = "12h"
+SPARK_AWS_CREDS_PROVIDER = "com.amazonaws.auth.WebIdentityTokenCredentialsProvider"
+SPARK_EXECUTOR_NAMESPACE = "paasta-spark"
+SPARK_DRIVER_POOL = "stable"
+SPARK_TRON_JOB_USER = "TRON"
+SPARK_PROMETHEUS_SHARD = "ml-compute"
+SPARK_DNS_POD_TEMPLATE = "/nail/srv/configs/spark_dns_pod_template.yaml"
+MEM_MULTIPLIER = {"k": 1024, "m": 1024**2, "g": 1024**3, "t": 1024**4}
+SPARK_DRIVER_DEFAULT_DISK_MB = 5120  # 5GB
+
+log = logging.getLogger(__name__)
+
+
+class SparkConfig(TypedDict):
+    account_id: str
+    default_event_log_dir: str
+    history_server: str
+
+
+class SparkEnvironmentConfig(TypedDict):
+    """
+    There's no set schema for the spark configuration from DEFAULT_SPARK_RUN_CONFIG,
+    but at the time of writing, this file looks like (with the addition of a `prod` key):
+    environments:
+        dev:
+            account_id: 'SOME_ACCOUNT_ID'
+            default_event_log_dir: s3a://SOME_BUCKET  # currently the only thing paasta reads
+            history_server: http://SOME_URL/
+    """
+
+    environments: Dict[str, SparkConfig]
+
+
+SparkEventLogConfiguration = TypedDict(
+    "SparkEventLogConfiguration",
+    {"spark.eventLog.enabled": str, "spark.eventLog.dir": str},
+    total=False,
+)
+
+
+def get_webui_url(port: str) -> str:
+    return f"http://{socket.getfqdn()}:{port}"
+
+
+def get_volumes_from_spark_mesos_configs(spark_conf: Mapping[str, str]) -> List[str]:
+    return (
+        spark_conf.get("spark.mesos.executor.docker.volumes", "").split(",")
+        if spark_conf.get("spark.mesos.executor.docker.volumes", "") != ""
+        else []
+    )
+
+
+def get_volumes_from_spark_k8s_configs(spark_conf: Mapping[str, str]) -> List[str]:
+    volume_names = []
+    for key in list(spark_conf.keys()):
+        if (
+            "spark.kubernetes.executor.volumes.hostPath." in key
+            and ".mount.path" in key
+        ):
+            v_name = re.match(
+                r"spark.kubernetes.executor.volumes.hostPath.([a-z0-9]([-a-z0-9]*[a-z0-9])?).mount.path",
+                key,
+            )
+            if v_name:
+                volume_names.append(v_name.group(1))
+            else:
+                log.error(
+                    f"Volume names must consist of lower case alphanumeric characters or '-', "
+                    f"and must start and end with an alphanumeric character. Config -> '{key}' must be fixed."
+                )
+                # Failing here because the k8s pod fails to start if the volume names
+                # don't follow the lowercase RFC 1123 standard.
+                sys.exit(1)
+
+    volumes = []
+    for volume_name in volume_names:
+        read_only = (
+            "ro"
+            if spark_conf.get(
+                f"spark.kubernetes.executor.volumes.hostPath.{volume_name}.mount.readOnly"
+            )
+            == "true"
+            else "rw"
+        )
+        container_path = spark_conf.get(
+            f"spark.kubernetes.executor.volumes.hostPath.{volume_name}.mount.path"
+        )
+        host_path = spark_conf.get(
+            f"spark.kubernetes.executor.volumes.hostPath.{volume_name}.options.path"
+        )
+        volumes.append(f"{host_path}:{container_path}:{read_only}")
+    return volumes
+
+
+def setup_volume_mounts(volumes: List[DockerVolume]) -> Dict[str, str]:
+    """
+    Returns Docker volume mount configurations in the format expected by Spark.
+    """
+    conf = {}
+
+    # XXX: why are these necessary?
+    extra_volumes: List[DockerVolume] = cast(
+        "List[DockerVolume]",
+        [
+            {"containerPath": "/etc/passwd", "hostPath": "/etc/passwd", "mode": "RO"},
+            {"containerPath": "/etc/group", "hostPath": "/etc/group", "mode": "RO"},
+        ],
+    )
+    seen_paths: Set[str] = set()  # dedupe volumes, just in case
+    for index, volume in enumerate(volumes + extra_volumes):
+        host_path, container_path, mode = (
+            volume["hostPath"],
+            volume["containerPath"],
+            volume["mode"],
+        )
+
+        if host_path in seen_paths:
+            log.warn(f"Skipping {host_path} - already added a binding for it.")
+            continue
+        seen_paths.add(host_path)
+
+        # the names here don't matter too much, so we just use the index in the volume
+        # list as an arbitrary name
+        conf[
+            f"spark.kubernetes.executor.volumes.hostPath.{index}.mount.path"
+        ] = container_path
+        conf[
+            f"spark.kubernetes.executor.volumes.hostPath.{index}.options.path"
+        ] = host_path
+        conf[
+            f"spark.kubernetes.executor.volumes.hostPath.{index}.mount.readOnly"
+        ] = str(mode.lower() == "ro").lower()
+
+    return conf
+
+
+def create_spark_config_str(spark_config_dict: Dict[str, Any], is_mrjob: bool) -> str:
+    conf_option = "--jobconf" if is_mrjob else "--conf"
+    spark_config_entries = list()
+
+    if is_mrjob:
+        spark_master = spark_config_dict["spark.master"]
+        spark_config_entries.append(f"--spark-master={spark_master}")
+        spark_config_dict.pop("spark.master", None)
+
+    for opt, val in spark_config_dict.items():
+        # Process Spark configs with multiple space separated values to be in single quotes
+        if isinstance(val, str) and " " in val:
+            val = f"'{val}'"
+        spark_config_entries.append(f"{conf_option} {opt}={val}")
+    return " ".join(spark_config_entries)
+
+
+def inject_spark_conf_str(original_cmd: str, spark_conf_str: str) -> str:
+    for base_cmd in ("pyspark", "spark-shell", "spark-submit"):
+        if base_cmd in original_cmd:
+            return original_cmd.replace(base_cmd, base_cmd + " " + spark_conf_str, 1)
+    return original_cmd
+
+
+def auto_add_timeout_for_spark_job(
+    cmd: str, timeout_job_runtime: str, silent: bool = False
+) -> str:
+    # Timeout only to be added for spark-submit commands
+    # TODO: Add timeout for jobs using mrjob with spark-runner
+    if "spark-submit" not in cmd:
+        return cmd
+    try:
+        # This is not an exhaustive regex, matches the invalid ones also, where as the invalid
+        # timeout command will fail during execution
+        options_regex = r"(--?[a-z][a-z-]*((\s+|=)[\w\d-]+)?\s+)*"
+        duration_regex = r"\d+\.?\d*[smhd]?"
+
+        timeout_present = re.match(
+            rf"^.*timeout\s+{options_regex}{duration_regex}\s+spark-submit .*$", cmd
+        )
+        if not timeout_present:
+            split_cmd = cmd.split("spark-submit")
+            # split_cmd[0] will always be an empty string or end with a space
+            cmd = f"{split_cmd[0]}timeout {timeout_job_runtime} spark-submit{split_cmd[1]}"
+            if not silent:
+                log.info(
+                    PaastaColors.blue(
+                        f"NOTE: Job will exit in given time {timeout_job_runtime}. "
+                        f"Adjust timeout value using --timeout-job-runtime. "
+                        f"New Updated Command with timeout: {cmd}"
+                    ),
+                )
+    except Exception as e:
+        err_msg = (
+            f"'timeout' could not be added to spark command: '{cmd}' due to error '{e}'. "
+            "Please report to #spark."
+        )
+        log.warn(err_msg)
+        print(PaastaColors.red(err_msg))
+    return cmd
+
+
+def build_spark_command(
+    original_cmd: str,
+    spark_config_dict: Dict[str, Any],
+    is_mrjob: bool,
+    timeout_job_runtime: str,
+    silent: bool = False,
+) -> str:
+    command = inject_spark_conf_str(
+        original_cmd, create_spark_config_str(spark_config_dict, is_mrjob=is_mrjob)
+    )
+    return auto_add_timeout_for_spark_job(command, timeout_job_runtime, silent=silent)
+
+
+def get_spark_ports_from_config(spark_conf: Dict[str, str]) -> List[int]:
+    ports = [int(v) for k, v in spark_conf.items() if k.endswith(".port")]
+    return ports
+
+
+# TODO: Reuse by ad-hoc Spark-driver-on-k8s
+def get_spark_driver_monitoring_annotations(
+    spark_config: Dict[str, str],
+) -> Dict[str, str]:
+    """
+    Returns Spark driver pod annotations - currently used for Prometheus metadata.
+    """
+    annotations: Dict[str, str] = {}
+
+    ui_port_str = spark_config.get("spark.ui.port")
+    if ui_port_str:
+        annotations.update(
+            {
+                "prometheus.io/port": ui_port_str,
+                "prometheus.io/path": "/metrics/prometheus",
+            }
+        )
+
+    paasta_service_non_truncated = spark_config.get(
+        "spark.kubernetes.executor.annotation.paasta.yelp.com/service"
+    )
+    paasta_instance_non_truncated = spark_config.get(
+        "spark.kubernetes.executor.annotation.paasta.yelp.com/instance"
+    )
+    if paasta_service_non_truncated and paasta_instance_non_truncated:
+        annotations.update(
+            {
+                "paasta.yelp.com/service": paasta_service_non_truncated,
+                "paasta.yelp.com/instance": paasta_instance_non_truncated,
+            }
+        )
+
+    return annotations
+
+
+def get_spark_driver_monitoring_labels(
+    spark_config: Dict[str, str],
+    user: str,
+) -> Dict[str, str]:
+    """
+    Returns Spark driver pod labels - generally for Prometheus metric relabeling.
+    """
+    ui_port_str = str(spark_config.get("spark.ui.port", ""))
+    labels = {
+        "paasta.yelp.com/prometheus_shard": SPARK_PROMETHEUS_SHARD,
+        "spark.yelp.com/user": user,
+        "spark.yelp.com/driver_ui_port": ui_port_str,
+    }
+    return labels

paasta_tools/synapse_srv_namespaces_fact.py

@@ -0,0 +1,42 @@
+#!/usr/bin/env python
+# 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.
+"""Usage: ./synapse_srv_namespaces_fact.py
+
+A simple script to enumerate all namespaces as a sorted comma separated
+string to stdout, with each entry in the form of full_name:proxy_port.
+
+If a proxy_port isn't defined for a namespace, that namespace is skipped.
+
+Example output: mumble.canary:5019,mumble.main:111,zookeeper.hab:4921
+
+This is nice to use as a facter fact for Synapse stuff!
+"""
+import sys
+
+from paasta_tools import long_running_service_tools
+
+
+def main():
+    strings = []
+    for full_name, config in long_running_service_tools.get_all_namespaces():
+        if "proxy_port" in config:
+            strings.append("{}:{}".format(full_name, config["proxy_port"]))
+    strings = sorted(strings)
+    print("synapse_srv_namespaces=" + ",".join(strings))
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

paasta_tools/tron/client.py

@@ -0,0 +1,158 @@
+# Copyright 2015-2018 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 logging
+import os
+from typing import Dict
+from urllib.parse import urljoin
+
+import requests
+
+from paasta_tools import yaml_tools as yaml
+from paasta_tools.cli.authentication import get_service_auth_token
+from paasta_tools.utils import get_user_agent
+
+
+log = logging.getLogger(__name__)
+
+
+class TronRequestError(Exception):
+    pass
+
+
+class TronClient:
+    """
+    Client for interacting with a Tron master.
+    """
+
+    def __init__(self, url):
+        self.master_url = url
+
+    def _request(self, method, url, data):
+        headers = {"User-Agent": get_user_agent()}
+        kwargs = {"url": urljoin(self.master_url, url), "headers": headers}
+        if method == "GET":
+            kwargs["params"] = data
+            response = requests.get(**kwargs)
+        elif method == "POST":
+            kwargs["data"] = data
+            if os.getenv("TRONCTL_API_AUTH"):
+                token = get_service_auth_token()
+                kwargs["headers"]["Authorization"] = f"Bearer {token}"
+            response = requests.post(**kwargs)
+        else:
+            raise ValueError(f"Unrecognized method: {method}")
+
+        return self._get_response_or_error(response)
+
+    def _get_response_or_error(self, response):
+        try:
+            result = response.json()
+            if "error" in result:
+                raise TronRequestError(result["error"])
+            return result
+        except ValueError:  # Not JSON
+            if not response.ok:
+                raise TronRequestError(
+                    "Status code {status_code} for {url}: {reason}".format(
+                        status_code=response.status_code,
+                        url=response.url,
+                        reason=response.reason,
+                    )
+                )
+            return response.text
+
+    def _get(self, url, data=None):
+        return self._request("GET", url, data)
+
+    def _post(self, url, data=None):
+        return self._request("POST", url, data)
+
+    def update_namespace(self, namespace, new_config, skip_if_unchanged=True):
+        """Updates the configuration for a namespace.
+
+        :param namespace: str
+        :param new_config: str, should be valid YAML.
+        :param skip_if_unchanged: boolean. If False, will send the update
+            even if the current config matches the new config.
+        """
+        current_config = self._get("/api/config", {"name": namespace, "no_header": 1})
+
+        if skip_if_unchanged:
+            if yaml.safe_load(new_config) == yaml.safe_load(current_config["config"]):
+                log.debug("No change in config, skipping update.")
+                return
+
+        return self._post(
+            "/api/config",
+            data={
+                "name": namespace,
+                "config": new_config,
+                "hash": current_config["hash"],
+                "check": 0,
+            },
+        )
+
+    def update_namespaces(
+        self, new_configs: Dict[str, str], skip_if_unchanged: bool = True
+    ):
+        """Updates the configuration for a namespace.
+
+        :param namespace: str
+        :param new_config: str, should be valid YAML.
+        :param skip_if_unchanged: boolean. If False, will send the update
+            even if the current config matches the new config.
+        """
+        current_configs: Dict[str, Dict[str, str]] = self._get("/api/config")  # type: ignore  # we don't have a good way to share types between tron/paasta
+        responses: Dict[str, str] = {}
+        for namespace, new_config in new_configs.items():
+            current_config = current_configs.get(namespace, {})
+            if skip_if_unchanged:
+                if yaml.safe_load(new_config) == yaml.safe_load(
+                    current_config["config"]
+                ):
+                    log.debug("No change in config, skipping update.")
+                    continue
+
+            responses[namespace] = self._post(
+                "/api/config",
+                data={
+                    "name": namespace,
+                    "config": new_config,
+                    "hash": current_config["hash"],
+                    "check": 0,
+                },
+            )
+        return responses
+
+    def list_namespaces(self):
+        """Gets the namespaces that are currently configured."""
+        response = self._get("/api")
+        return response.get("namespaces", [])
+
+    def get_job_content(self, job: str) -> dict:
+        return self._get(f"/api/jobs/{job}/")
+
+    def get_latest_job_run_id(self, job_content: dict) -> str:
+        job_runs = sorted(
+            job_content.get("runs", []),
+            key=lambda k: (k["state"] != "scheduled", k["run_num"]),
+            reverse=True,
+        )
+        if not job_runs:
+            return None
+        return job_runs[0]["run_num"]
+
+    def get_action_run(self, job: str, action: str, run_id: str) -> dict:
+        return self._get(
+            f"/api/jobs/{job}/{run_id}/{action}?include_stderr=1&include_stdout=1&num_lines=10"
+        )

paasta_tools/tron/tron_command_context.py

@@ -0,0 +1,194 @@
+"""Command Context is how we construct the command line for a command which may
+have variables that need to be rendered.
+This is a COPY of https://github.com/Yelp/Tron/blob/master/tron/command_context.py.
+"""
+import functools
+import operator
+
+from paasta_tools.tron import tron_timeutils
+
+
+def build_context(object, parent):
+    """Construct a CommandContext for object. object must have a property
+    'context_class'.
+    """
+    return CommandContext(object.context_class(object), parent)
+
+
+def build_filled_context(*context_objects):
+    """Create a CommandContext chain from context_objects, using a Filler
+    object to pass to each CommandContext. Can be used to validate a format
+    string.
+    """
+    if not context_objects:
+        return CommandContext()
+
+    filler = Filler()
+
+    def build(current, next):
+        return CommandContext(next(filler), current)
+
+    return functools.reduce(build, context_objects, None)
+
+
+class CommandContext:
+    """A CommandContext object is a wrapper around any object which has values
+    to be used to render a command for execution.  It looks up values by name.
+
+    It's lookup order is:
+        base[name],
+        base.__getattr__(name),
+        next[name],
+        next.__getattr__(name)
+    """
+
+    def __init__(self, base=None, next=None):
+        """
+        base - Object to look for attributes in
+        next - Next place to look for more pieces of context
+               Generally this will be another instance of CommandContext
+        """
+        self.base = base or {}
+        self.next = next or {}
+
+    def get(self, name, default=None):
+        try:
+            return self.__getitem__(name)
+        except KeyError:
+            return default
+
+    def __getitem__(self, name):
+        getters = [operator.itemgetter(name), operator.attrgetter(name)]
+        for target in [self.base, self.next]:
+            for getter in getters:
+                try:
+                    return getter(target)
+                except (KeyError, TypeError, AttributeError):
+                    pass
+
+        raise KeyError(name)
+
+    def __eq__(self, other):
+        return self.base == other.base and self.next == other.next
+
+    def __ne__(self, other):
+        return not self == other
+
+
+class JobContext:
+    """A class which exposes properties for rendering commands."""
+
+    def __init__(self, job):
+        self.job = job
+
+    @property
+    def name(self):
+        return self.job.name
+
+    def __getitem__(self, item):
+        date_name, date_spec = self._get_date_spec_parts(item)
+        if not date_spec:
+            raise KeyError(item)
+
+        if date_name == "last_success":
+            last_success = self.job.runs.last_success
+            last_success = last_success.run_time if last_success else None
+
+            time_value = tron_timeutils.DateArithmetic.parse(date_spec, last_success)
+            if time_value:
+                return time_value
+
+        raise KeyError(item)
+
+    def _get_date_spec_parts(self, name):
+        parts = name.rsplit(":", 1)
+        if len(parts) != 2:
+            return name, None
+        return parts
+
+
+class JobRunContext:
+    def __init__(self, job_run):
+        self.job_run = job_run
+
+    @property
+    def runid(self):
+        return self.job_run.id
+
+    @property
+    def cleanup_job_status(self):
+        """Provide 'SUCCESS' or 'FAILURE' to a cleanup action context based on
+        the status of the other steps
+        """
+        if self.job_run.action_runs.is_failed:
+            return "FAILURE"
+        elif self.job_run.action_runs.is_complete_without_cleanup:
+            return "SUCCESS"
+        return "UNKNOWN"
+
+    def __getitem__(self, name):
+        """Attempt to parse date arithmetic syntax and apply to run_time."""
+        run_time = self.job_run.run_time
+        time_value = tron_timeutils.DateArithmetic.parse(name, run_time)
+        if time_value:
+            return time_value
+
+        raise KeyError(name)
+
+
+class ActionRunContext:
+    """Context object that gives us access to data about the action run."""
+
+    def __init__(self, action_run):
+        self.action_run = action_run
+
+    @property
+    def actionname(self):
+        return self.action_run.action_name
+
+    @property
+    def node(self):
+        return self.action_run.node.hostname
+
+
+class ServiceInstancePidContext:
+    def __init__(self, service_instance):
+        self.service_instance = service_instance
+
+    @property
+    def instance_number(self):
+        return self.service_instance.instance_number
+
+    @property
+    def node(self):
+        return self.service_instance.node.hostname
+
+    @property
+    def name(self):
+        return self.service_instance.config.name
+
+
+class ServiceInstanceContext(ServiceInstancePidContext):
+    @property
+    def pid_file(self):
+        context = CommandContext(self, self.service_instance.parent_context)
+        return self.service_instance.config.pid_file % context
+
+
+class Filler:
+    """Filler object for using CommandContext during config parsing. This class
+    is used as a substitute for objects that would be passed to Context objects.
+    This allows the Context objects to be used directly for config validation.
+    """
+
+    def __getattr__(self, _):
+        return self
+
+    def __str__(self):
+        return "%(...)s"
+
+    def __mod__(self, _):
+        return self
+
+    def __nonzero__(self):
+        return False