paasta-tools 1.2.0__tar.gz → 1.4.1__tar.gz

{paasta-tools-1.2.0/paasta_tools.egg-info → paasta-tools-1.4.1}/PKG-INFO

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

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/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.2.0"
+__version__ = "1.4.1"

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools/generate_authenticating_services.py

@@ -13,7 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-A simple script to enumerate all services partecipating in authenticated
+A simple script to enumerate all services participating in authenticated
 communications, and list them in a YAML/JSON file.
 """
 import argparse
@@ -27,6 +27,7 @@ from typing import Set
 import yaml
 
 from paasta_tools.utils import DEFAULT_SOA_DIR
+from paasta_tools.utils import load_system_paasta_config
 from paasta_tools.utils import write_json_configuration_file
 from paasta_tools.utils import write_yaml_configuration_file
 
@@ -53,6 +54,7 @@ def enumerate_authenticating_services() -> Dict[str, List[str]]:
     config_path_pattern = os.path.join(DEFAULT_SOA_DIR, "*", AUTHORIZATION_CONFIG_FILE)
     for authz_config in glob.glob(config_path_pattern):
         result.update(list_services_in_authz_config(authz_config))
+    result.update(load_system_paasta_config().get_always_authenticating_services())
     return {"services": sorted(result)}
 
 

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools/kubernetes_tools.py

@@ -21,6 +21,7 @@ import os
 import re
 from datetime import datetime
 from enum import Enum
+from functools import lru_cache
 from inspect import currentframe
 from pathlib import Path
 from typing import Any
@@ -2552,6 +2553,13 @@ class KubernetesDeploymentConfig(LongRunningServiceConfig):
             "topology_spread_constraints", default_pod_topology_spread_constraints
         )
 
+    def get_projected_sa_volumes(self) -> List[ProjectedSAVolume]:
+        return add_volumes_for_authenticating_services(
+            service_name=self.service,
+            config_volumes=super().get_projected_sa_volumes(),
+            soa_dir=self.soa_dir,
+        )
+
 
 def get_kubernetes_secret_hashes(
     environment_variables: Mapping[str, str], service: str, namespace: str
@@ -4406,3 +4414,35 @@ def get_kubernetes_secret_volumes(
                 ] = secret_contents
 
     return secret_volumes
+
+
+@lru_cache()
+def get_authenticating_services(soa_dir: str = DEFAULT_SOA_DIR) -> Set[str]:
+    """Load list of services participating in authenticated traffic"""
+    authenticating_services_conf_path = os.path.join(soa_dir, "authenticating.yaml")
+    config = service_configuration_lib.read_yaml_file(authenticating_services_conf_path)
+    return set(config.get("services", []))
+
+
+def add_volumes_for_authenticating_services(
+    service_name: str,
+    config_volumes: List[ProjectedSAVolume],
+    soa_dir: str = DEFAULT_SOA_DIR,
+) -> List[ProjectedSAVolume]:
+    """Add projected service account volume to the list of volumes if service
+    participates in authenticated traffic. In case of changes, a new list is returned,
+    no updates in-place.
+
+    :param str service_name: name of the service
+    :param List[ProjectedSAVolume] config_volumes: existing projected volumes from service config
+    :param str soa_dir: path to SOA configurations directory
+    :return: updated list of projected service account volumes
+    """
+    token_config = load_system_paasta_config().get_service_auth_token_volume_config()
+    if (
+        token_config
+        and service_name in get_authenticating_services(soa_dir)
+        and not any(volume == token_config for volume in config_volumes)
+    ):
+        config_volumes = [token_config, *config_volumes]
+    return config_volumes

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools/spark_tools.py

@@ -6,6 +6,7 @@ from typing import Any
 from typing import cast
 from typing import Dict
 from typing import List
+from typing import Literal
 from typing import Mapping
 from typing import Set
 
@@ -23,6 +24,7 @@ SPARK_DRIVER_POOL = "stable"
 SPARK_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}
 
 log = logging.getLogger(__name__)
 
@@ -247,3 +249,22 @@ def get_spark_driver_monitoring_labels(
         "spark.yelp.com/driver_ui_port": ui_port_str,
     }
     return labels
+
+
+def get_spark_memory_in_unit(mem: str, unit: Literal["k", "m", "g", "t"]) -> float:
+    """
+    Converts Spark memory to the desired unit.
+    mem is the same format as JVM memory strings: just number or number followed by 'k', 'm', 'g' or 't'.
+    unit can be 'k', 'm', 'g' or 't'.
+    Returns memory as a float converted to the desired unit.
+    """
+    try:
+        memory_bytes = float(mem)
+    except ValueError:
+        try:
+            memory_bytes = float(mem[:-1]) * MEM_MULTIPLIER[mem[-1]]
+        except (ValueError, IndexError):
+            print(f"Unable to parse memory value {mem}. Defaulting to 2 GB.")
+            memory_bytes = 2147483648  # default to 2 GB
+    memory_unit = memory_bytes / MEM_MULTIPLIER[unit]
+    return memory_unit

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools/tron_tools.py

@@ -56,6 +56,7 @@ from paasta_tools.utils import get_k8s_url_for_cluster
 from paasta_tools.utils import validate_pool
 from paasta_tools.utils import PoolsNotConfiguredError
 from paasta_tools.utils import DockerVolume
+from paasta_tools.utils import ProjectedSAVolume
 
 from paasta_tools import spark_tools
 
@@ -71,6 +72,7 @@ from paasta_tools.secret_tools import is_shared_secret
 from paasta_tools.secret_tools import is_shared_secret_from_secret_name
 from paasta_tools.secret_tools import get_secret_name_from_ref
 from paasta_tools.kubernetes_tools import get_paasta_secret_name
+from paasta_tools.kubernetes_tools import add_volumes_for_authenticating_services
 from paasta_tools.secret_tools import SHARED_SECRET_SERVICE
 
 from paasta_tools import monitoring_tools
@@ -280,6 +282,35 @@ class TronActionConfig(InstanceConfig):
         self.job, self.action = decompose_instance(instance)
         # Indicate whether this config object is created for validation
         self.for_validation = for_validation
+        self.action_spark_config = None
+        if self.get_executor() == "spark":
+            # build the complete Spark configuration
+            # TODO: add conditional check for Spark specific commands spark-submit, pyspark etc ?
+            self.action_spark_config = self.build_spark_config()
+
+    def get_cpus(self) -> float:
+        # set Spark driver pod CPU if it is specified by Spark arguments
+        if (
+            self.action_spark_config
+            and "spark.driver.cores" in self.action_spark_config
+        ):
+            return float(self.action_spark_config["spark.driver.cores"])
+        # we fall back to this default if there's no spark.driver.cores config
+        return super().get_cpus()
+
+    def get_mem(self) -> float:
+        # set Spark driver pod memory if it is specified by Spark arguments
+        if (
+            self.action_spark_config
+            and "spark.driver.memory" in self.action_spark_config
+        ):
+            return int(
+                spark_tools.get_spark_memory_in_unit(
+                    self.action_spark_config["spark.driver.memory"], "m"
+                )
+            )
+        # we fall back to this default if there's no spark.driver.memory config
+        return super().get_mem()
 
     def build_spark_config(self) -> Dict[str, str]:
         system_paasta_config = load_system_paasta_config()
@@ -438,7 +469,6 @@ class TronActionConfig(InstanceConfig):
         system_paasta_config: Optional["SystemPaastaConfig"] = None,
     ) -> Dict[str, str]:
         env = super().get_env(system_paasta_config=system_paasta_config)
-
         if self.get_executor() == "spark":
             # Required by some sdks like boto3 client. Throws NoRegionError otherwise.
             # AWS_REGION takes precedence if set.
@@ -603,6 +633,20 @@ class TronActionConfig(InstanceConfig):
             error_msgs.append(
                 f"{self.get_job_name()}.{self.get_action_name()} must have a deploy_group set"
             )
+        # We are not allowing users to specify `cpus` and `mem` configuration if the action is a Spark job
+        # with driver running on k8s (executor: spark), because we derive these values from `spark.driver.cores`
+        # and `spark.driver.memory` in order to avoid confusion.
+        if self.get_executor() == "spark":
+            if "cpus" in self.config_dict:
+                error_msgs.append(
+                    f"{self.get_job_name()}.{self.get_action_name()} is a Spark job. `cpus` config is not allowed. "
+                    f"Please specify the driver cores using `spark.driver.cores`."
+                )
+            if "mem" in self.config_dict:
+                error_msgs.append(
+                    f"{self.get_job_name()}.{self.get_action_name()} is a Spark job. `mem` config is not allowed. "
+                    f"Please specify the driver memory using `spark.driver.memory`."
+                )
         return error_msgs
 
     def get_pool(self) -> str:
@@ -627,6 +671,14 @@ class TronActionConfig(InstanceConfig):
     def get_service_account_name(self) -> Optional[str]:
         return self.config_dict.get("service_account_name")
 
+    def get_projected_sa_volumes(self) -> Optional[List[ProjectedSAVolume]]:
+        projected_volumes = add_volumes_for_authenticating_services(
+            service_name=self.service,
+            config_volumes=super().get_projected_sa_volumes(),
+            soa_dir=self.soa_dir,
+        )
+        return projected_volumes if projected_volumes else None
+
 
 class TronJobConfig:
     """Represents a job in Tron, consisting of action(s) and job-level configuration values."""
@@ -958,25 +1010,33 @@ def format_tron_action_dict(action_config: TronActionConfig):
                 dry_run=action_config.for_validation,
             )
 
+        # service account token volumes for service authentication
+        result["projected_sa_volumes"] = action_config.get_projected_sa_volumes()
+
         extra_volumes = action_config.get_extra_volumes()
         if executor == "spark":
             is_mrjob = action_config.config_dict.get("mrjob", False)
             system_paasta_config = load_system_paasta_config()
-            # inject spark configs to the original spark-submit command
-            spark_config = action_config.build_spark_config()
+            # inject additional Spark configs in case of Spark commands
             result["command"] = spark_tools.build_spark_command(
                 result["command"],
-                spark_config,
+                action_config.action_spark_config,
                 is_mrjob,
                 action_config.config_dict.get(
                     "max_runtime", spark_tools.DEFAULT_SPARK_RUNTIME_TIMEOUT
                 ),
             )
+            # point to the KUBECONFIG needed by Spark driver
             result["env"]["KUBECONFIG"] = system_paasta_config.get_spark_kubeconfig()
+
             # spark, unlike normal batches, needs to expose several ports for things like the spark
             # ui and for executor->driver communication
             result["ports"] = list(
-                set(spark_tools.get_spark_ports_from_config(spark_config))
+                set(
+                    spark_tools.get_spark_ports_from_config(
+                        action_config.action_spark_config
+                    )
+                )
             )
             # mount KUBECONFIG file for Spark drivers to communicate with EKS cluster
             extra_volumes.append(
@@ -990,10 +1050,12 @@ def format_tron_action_dict(action_config: TronActionConfig):
             )
             # Add pod annotations and labels for Spark monitoring metrics
             monitoring_annotations = (
-                spark_tools.get_spark_driver_monitoring_annotations(spark_config)
+                spark_tools.get_spark_driver_monitoring_annotations(
+                    action_config.action_spark_config
+                )
             )
             monitoring_labels = spark_tools.get_spark_driver_monitoring_labels(
-                spark_config
+                action_config.action_spark_config
             )
             result["annotations"].update(monitoring_annotations)
             result["labels"].update(monitoring_labels)

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools/utils.py

@@ -75,7 +75,7 @@ import ldap3
 import requests_cache
 import service_configuration_lib
 import yaml
-from docker import Client
+from docker import APIClient
 from docker.utils import kwargs_from_env
 from kazoo.client import KazooClient
 from mypy_extensions import TypedDict
@@ -2020,6 +2020,8 @@ class SystemPaastaConfigDict(TypedDict, total=False):
     eks_cluster_aliases: Dict[str, str]
     secret_sync_delay_seconds: float
     use_multiple_log_readers: Optional[List[str]]
+    service_auth_token_settings: ProjectedSAVolume
+    always_authenticating_services: List[str]
 
 
 def load_system_paasta_config(
@@ -2725,6 +2727,12 @@ class SystemPaastaConfig:
     def get_kube_clusters(self) -> Dict:
         return self.config_dict.get("kube_clusters", {})
 
+    def get_service_auth_token_volume_config(self) -> ProjectedSAVolume:
+        return self.config_dict.get("service_auth_token_settings", {})
+
+    def get_always_authenticating_services(self) -> List[str]:
+        return self.config_dict.get("always_authenticating_services", [])
+
 
 def _run(
     command: Union[str, List[str]],
@@ -3337,12 +3345,12 @@ def get_docker_host() -> str:
     return os.environ.get("DOCKER_HOST", "unix://var/run/docker.sock")
 
 
-def get_docker_client() -> Client:
+def get_docker_client() -> APIClient:
     client_opts = kwargs_from_env(assert_hostname=False)
     if "base_url" in client_opts:
-        return Client(**client_opts)
+        return APIClient(**client_opts)
     else:
-        return Client(base_url=get_docker_host(), **client_opts)
+        return APIClient(base_url=get_docker_host(), **client_opts)
 
 
 def get_running_mesos_docker_containers() -> List[Dict]:

{paasta-tools-1.2.0 → paasta-tools-1.4.1/paasta_tools.egg-info}/PKG-INFO

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

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/paasta_tools.egg-info/requires.txt

@@ -10,7 +10,7 @@ certifi
 choice>=0.1
 cookiecutter>=1.4.0
 croniter
-docker-py>=1.2.3
+docker
 dulwich>=0.17.3
 ephemeral-port-reserve>=1.0.1
 graphviz

{paasta-tools-1.2.0 → paasta-tools-1.4.1}/requirements-minimal.txt

@@ -10,9 +10,7 @@ certifi
 choice >= 0.1
 cookiecutter >= 1.4.0
 croniter
-# Don't update this unless you have confirmed the client works with
-# the Docker version deployed on PaaSTA servers
-docker-py >= 1.2.3
+docker
 dulwich >= 0.17.3
 ephemeral-port-reserve >= 1.0.1
 graphviz