Merge pull request #164 from Yelp/u/sids/CLOUD-816/dry-run-metrics

dry run mode for --jira-ticket parmeter in paasta spark-run

Author: sidsatyelp

mypy.ini

@@ -5,3 +5,4 @@ warn_incomplete_stub = True
 follow_imports = silent
 ignore_missing_imports = True
 mypy_path = stubs
+strict_optional = False

service_configuration_lib/spark_config.py

@@ -79,6 +79,7 @@
 TICKET_NOT_REQUIRED_USERS = {
     'batch',  # non-human spark-run from batch boxes
     'TRON',  # tronjobs that run commands like paasta mark-for-deployment
+    'jenkins',  # username for jenkins pipeline jobs
     None,  # placeholder for being unable to determine user
 }
 USER_LABEL_UNSPECIFIED = 'UNSPECIFIED'
@@ -993,9 +994,79 @@ def _get_valid_jira_ticket(self, jira_ticket: Optional[str]) -> Optional[str]:
         if ticket and JIRA_TICKET_PATTERN.match(ticket):
             log.info(f'Valid Jira ticket provided: {ticket}')
             return ticket
-        log.warning(f'Jira ticket missing or invalid format: {ticket}')
+        log.info(f'Jira ticket missing or invalid format: {ticket}')
         return None
 
+    def _handle_jira_ticket_validation(
+        self,
+        cluster_manager: str,
+        user: Optional[str],
+        jira_ticket: Optional[str],
+        paasta_cluster: str,
+        paasta_service: str,
+        paasta_instance: str,
+    ) -> Optional[str]:
+        """
+        This method checks if Jira ticket validation is enabled, if the user needs
+        to provide a ticket, and validates the ticket if needed.
+
+        Returns:
+            The validated Jira ticket string if valid, otherwise None.
+
+        Args:
+            cluster_manager: The cluster manager being used
+            user: The user running the job
+            jira_ticket: The Jira ticket provided by the user
+        """
+        flag_enabled = self.mandatory_default_spark_srv_conf.get('spark.yelp.jira_ticket.enabled', 'false')
+        valid_ticket = self._get_valid_jira_ticket(jira_ticket)
+
+        # Skip validation for local cluster manager or exempt users
+        if cluster_manager == 'local' or user in TICKET_NOT_REQUIRED_USERS:
+            log.debug('Jira ticket check not required for this job configuration.')
+            # If exempt, we still pass through the original ticket if it's valid,
+            # otherwise None. This allows exempt users like tron to still have their valid tickets
+            # (if provided) attached as labels, without forcing validation.
+            return valid_ticket
+
+        if valid_ticket is None:
+            log_payload = {
+                'timestamp': int(time.time()),
+                'event': 'jira_ticket_validation_warning',
+                'level': 'WARNING',
+                'reason': 'Ticket missing or invalid',
+                'user': user,
+                'jira_ticket_provided': jira_ticket,
+                'paasta_cluster': paasta_cluster,
+                'paasta_service': paasta_service,
+                'paasta_instance': paasta_instance,
+            }
+            if flag_enabled == 'true':
+                error_msg = (
+                    f'Job requires a valid Jira ticket (format PROJ-1234).\n'
+                    f'Jira ticket check is enabled, but ticket "{jira_ticket}" is '
+                    f'missing or invalid for user "{user}".\n'
+                    'Please pass the parameter as: paasta spark-run --jira-ticket=PROJ-1234 \n'
+                    'For more information: http://y/spark-jira-ticket-param \n'
+                    'If you have questions, please reach out to #spark on Slack.\n'
+                    f'paasta_cluster={paasta_cluster}, paasta_service={paasta_service}\n'
+                    f'paasta_instance={paasta_instance}'
+                )
+                utils.log_to_clog('spark_jira_ticket', log_payload, error_msg, log)
+                raise RuntimeError(error_msg)
+            else:
+                warning_message = (
+                    f'Jira ticket check is configured, but ticket is missing or invalid for user "{user}". '
+                    f'Proceeding with job execution. Original ticket value: "{jira_ticket}". '
+                    'Please pass the parameter as: paasta spark-run --jira-ticket=PROJ-1234 '
+                    'For more information: http://y/spark-jira-ticket-param '
+                    'If you have questions, please reach out to #spark on Slack. '
+                    f'paasta_cluster={paasta_cluster}, paasta_service={paasta_service}\n'
+                    f'paasta_instance={paasta_instance}'
+                )
+                utils.log_to_clog('spark_jira_ticket', log_payload, warning_message, log)
+        return valid_ticket
+
     def get_spark_conf(
         self,
         cluster_manager: str,
@@ -1040,6 +1111,7 @@ def get_spark_conf(
             to launch the batch, and inside the batch use `spark_tools.paasta` to create
             spark session.
         :param aws_region: The default aws region to use
+        :param jira_ticket: The jira project that this spark job is related to.
         :param service_account_name: The k8s service account to use for spark k8s authentication.
         :param force_spark_resource_configs: skip the resource/instances recalculation.
             This is strongly not recommended.
@@ -1058,20 +1130,10 @@ def get_spark_conf(
         # Get user from environment variables if it's not set
         user = user or os.environ.get('USER', None)
 
-        if self.mandatory_default_spark_srv_conf.get('spark.yelp.jira_ticket.enabled') == 'true':
-            needs_jira_check = cluster_manager != 'local' and user not in TICKET_NOT_REQUIRED_USERS
-            if needs_jira_check:
-                valid_ticket = self._get_valid_jira_ticket(jira_ticket)
-                if valid_ticket is None:
-                    error_msg = (
-                        'Job requires a valid Jira ticket (format PROJ-1234).\n'
-                        'Please pass the parameter as: paasta spark-run --jira-ticket=PROJ-1234 \n'
-                        'For more information: https://yelpwiki.yelpcorp.com/spaces/AML/pages/402885641 \n'
-                        f'If you have questions, please reach out to #spark on Slack. (user={user})\n'
-                    )
-                    raise RuntimeError(error_msg)
-            else:
-                log.debug('Jira ticket check not required for this job configuration.')
+        # Handle Jira ticket validation if enabled
+        validated_jira_ticket = self._handle_jira_ticket_validation(
+            cluster_manager, user, jira_ticket, paasta_cluster, paasta_service, paasta_instance,
+        )
 
         app_base_name = (
             user_spark_opts.get('spark.app.name') or
@@ -1160,7 +1222,7 @@ def get_spark_conf(
                 include_self_managed_configs=not use_eks,
                 k8s_server_address=k8s_server_address,
                 user=user,
-                jira_ticket=jira_ticket,
+                jira_ticket=validated_jira_ticket,
             ))
         elif cluster_manager == 'local':
             spark_conf.update(_get_local_spark_env(

service_configuration_lib/utils.py

@@ -2,6 +2,7 @@
 import contextlib
 import errno
 import hashlib
+import json
 import logging
 import os
 import random
@@ -19,6 +20,12 @@
 import yaml
 from typing_extensions import Literal
 
+# Only works inside yelpy environments
+try:
+    import clog
+except ImportError:
+    clog = None
+
 DEFAULT_SPARK_RUN_CONFIG = '/nail/srv/configs/spark.yaml'
 POD_TEMPLATE_PATH = '/nail/tmp/spark-pt-{file_uuid}.yaml'
 SPARK_EXECUTOR_POD_TEMPLATE = '/nail/srv/configs/spark_executor_pod_template.yaml'
@@ -33,16 +40,16 @@
 SPARK_DRIVER_MEM_OVERHEAD_FACTOR_DEFAULT = 0.1
 
 
-log = logging.Logger(__name__)
+log = logging.getLogger(__name__)
 log.setLevel(logging.INFO)
 
 
 def load_spark_srv_conf(preset_values=None) -> Tuple[
-    Dict[str, Any],
-    Dict[str, Any],
-    Dict[str, Any],
-    Dict[str, Any],
-    Dict[str, Dict[str, float]],
+    Dict[str, Any],  # spark_srv_conf
+    Dict[str, Any],  # spark_constants
+    Dict[str, Any],  # default_spark_srv_conf
+    Dict[str, Any],  # mandatory_default_spark_srv_conf
+    Dict[str, Dict[str, float]],  # spark_costs
 ]:
     if preset_values is None:
         preset_values = dict()
@@ -220,3 +227,34 @@ def get_spark_driver_memory_overhead_mb(spark_conf: Dict[str, str]) -> float:
         )
         driver_mem_overhead_mb = driver_mem_mb * driver_mem_overhead_factor
     return round(driver_mem_overhead_mb, 5)
+
+
+def log_to_clog(log_stream: str, log_payload: dict, warning_message: str, log_instance=None):
+    """
+    Log a message to clog if available, otherwise log a warning.
+
+    Args:
+        log_stream: The clog stream name to log to
+        log_payload: Dictionary containing the log payload data
+        warning_message: Warning message to log if clog fails or is unavailable
+        log_instance: Logger instance to use for warnings (optional)
+    """
+    if clog:
+        try:
+            clog.config.configure(
+                scribe_host='169.254.255.254',
+                scribe_port='1463',
+                monk_disable=False,
+                scribe_disable=False,
+            )
+            clog.log_line(log_stream, json.dumps(log_payload))
+        except Exception as e:
+            if log_instance:
+                log_instance.warning(f'{warning_message} Clog operation failed with: {e}')
+            else:
+                log.warning(f'{warning_message} Clog operation failed with: {e}')
+    else:
+        if log_instance:
+            log_instance.warning(warning_message)
+        else:
+            log.warning(warning_message)

setup.py

@@ -43,4 +43,7 @@
         'scripts/services_needing_puppet_help',
         'scripts/services_that_run_here',
     ],
+    extras_require={
+        'yelp': ['yelp-clog>=7'],
+    },
 )

tests/conftest.py

@@ -1,3 +1,4 @@
+import logging
 from unittest.mock import MagicMock
 from unittest.mock import Mock
 from unittest.mock import patch
@@ -80,3 +81,20 @@ def yaml_configs_file_watcher(
         configs_suffixes=['.yaml'],
         configs_folder=mock_soa_dir,
     )
+
+
+@pytest.fixture(autouse=True)
+def mock_clog_logging(monkeypatch):
+    """
+    Autouse fixture to prevent clog logging during tests.
+    This mocks the log_to_clog function to avoid actual clog operations.
+    """
+    def mock_log_to_clog(log_stream, log_payload, warning_message, log_instance=None):
+        # During tests, just log the warning message without trying to use clog
+        if log_instance:
+            log_instance.warning(warning_message)
+        else:
+            logger = logging.getLogger(__name__)
+            logger.warning(warning_message)
+
+    monkeypatch.setattr('service_configuration_lib.utils.log_to_clog', mock_log_to_clog)