move custom exit code and signal handlers away from retry decorator

Author: saikonen

metaflow/plugins/argo/argo_workflows.py

@@ -54,7 +54,8 @@
 from metaflow.plugins.kubernetes.kube_utils import qos_requests_and_limits
 
 from metaflow.plugins.kubernetes.kubernetes_jobsets import KubernetesArgoJobSet
-from metaflow.plugins.retry_decorator import PLATFORM_EVICTED_EXITCODE, RetryEvents
+from metaflow.plugins.kubernetes.kubernetes import SPOT_INTERRUPT_EXITCODE
+from metaflow.plugins.retry_decorator import RetryEvents
 from metaflow.unbounded_foreach import UBF_CONTROL, UBF_TASK
 from metaflow.user_configs.config_options import ConfigInput
 from metaflow.util import (
@@ -1553,7 +1554,7 @@ def _container_templates(self):
             event_to_expr = {
                 RetryEvents.STEP: "asInt(lastRetry.exitCode) == 1",
                 RetryEvents.PREEMPT: "asInt(lastRetry.exitCode) == %s"
-                % PLATFORM_EVICTED_EXITCODE,
+                % SPOT_INTERRUPT_EXITCODE,
             }
             retry_expr = None
             if retry_conditions:

metaflow/plugins/aws/batch/batch.py

@@ -43,6 +43,8 @@
 STDOUT_PATH = os.path.join(LOGS_DIR, STDOUT_FILE)
 STDERR_PATH = os.path.join(LOGS_DIR, STDERR_FILE)
 
+SPOT_INTERRUPT_EXITCODE = 234
+
 
 class BatchException(MetaflowException):
     headline = "AWS Batch error"
@@ -488,7 +490,7 @@ def wait_for_launch(job, child_jobs):
         if self.job.is_crashed:
 
             # Custom exception for spot instance terminations
-            if self.job.status_code == 234:
+            if self.job.status_code == SPOT_INTERRUPT_EXITCODE:
                 raise BatchSpotInstanceTerminated()
 
             msg = next(

metaflow/plugins/aws/batch/batch_decorator.py

@@ -1,6 +1,8 @@
 import os
 import platform
+import signal
 import sys
+import threading
 import time
 
 from metaflow import R, current
@@ -24,7 +26,7 @@
     get_docker_registry,
     get_ec2_instance_metadata,
 )
-from .batch import BatchException
+from .batch import SPOT_INTERRUPT_EXITCODE, BatchException
 
 
 class BatchDecorator(StepDecorator):
@@ -298,6 +300,29 @@ def task_pre_step(
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Set up signal handling for spot termination
+            main_pid = os.getpid()
+
+            def _termination_timer():
+                time.sleep(30)
+                os.kill(main_pid, signal.SIGALRM)
+
+            def _spot_term_signal_handler(*args, **kwargs):
+                if os.path.isfile(current.spot_termination_notice):
+                    print(
+                        "Spot instance termination detected. Starting a timer to end the Metaflow task."
+                    )
+                    timer_thread = threading.Thread(
+                        target=_termination_timer, daemon=True
+                    )
+                    timer_thread.start()
+
+            def _curtain_call(*args, **kwargs):
+                # custom exit code in case of Spot termination
+                sys.exit(SPOT_INTERRUPT_EXITCODE)
+
+            signal.signal(signal.SIGUSR1, _spot_term_signal_handler)
+            signal.signal(signal.SIGALRM, _curtain_call)
             # Start spot termination monitor sidecar.
             # TODO: A nicer way to pass the main process id to a Sidecar, in order to allow sidecars to send signals back to the main process.
             os.environ["MF_MAIN_PID"] = str(os.getpid())

metaflow/plugins/aws/step_functions/step_functions.py

@@ -18,11 +18,11 @@
     SFN_S3_DISTRIBUTED_MAP_OUTPUT_PATH,
 )
 from metaflow.parameters import deploy_time_eval
-from metaflow.plugins.retry_decorator import PLATFORM_EVICTED_EXITCODE, RetryEvents
+from metaflow.plugins.retry_decorator import RetryEvents
 from metaflow.user_configs.config_options import ConfigInput
 from metaflow.util import dict_to_cli_options, to_pascalcase
 
-from ..batch.batch import Batch
+from ..batch.batch import Batch, SPOT_INTERRUPT_EXITCODE
 from .event_bridge_client import EventBridgeClient
 from .step_functions_client import StepFunctionsClient
 
@@ -842,7 +842,7 @@ def _batch(self, node):
             RetryEvents.STEP: {"action": "RETRY", "onExitCode": "1"},
             RetryEvents.PREEMPT: {
                 "action": "RETRY",
-                "onExitCode": str(PLATFORM_EVICTED_EXITCODE),
+                "onExitCode": str(SPOT_INTERRUPT_EXITCODE),
             },
         }
         retry_expr = None

metaflow/plugins/kubernetes/kubernetes.py

@@ -63,6 +63,8 @@
     "{METAFLOW_PARALLEL_STEP_CLI_OPTIONS_TEMPLATE}"
 )
 
+SPOT_INTERRUPT_EXITCODE = 234
+
 
 class KubernetesException(MetaflowException):
     headline = "Kubernetes error"
@@ -768,7 +770,7 @@ def _has_updates():
                     )
                 if int(exit_code) == 134:
                     raise KubernetesException("%s (exit code %s)" % (msg, exit_code))
-                if int(exit_code) == 234:
+                if int(exit_code) == SPOT_INTERRUPT_EXITCODE:
                     # NOTE. K8S exit codes are mod 256
                     raise KubernetesSpotInstanceTerminated()
                 else:

metaflow/plugins/kubernetes/kubernetes_decorator.py

@@ -1,7 +1,9 @@
 import json
 import os
 import platform
+import signal
 import sys
+import threading
 import time
 
 from metaflow import current
@@ -37,7 +39,7 @@
 from metaflow.unbounded_foreach import UBF_CONTROL
 
 from ..aws.aws_utils import get_docker_registry, get_ec2_instance_metadata
-from .kubernetes import KubernetesException
+from .kubernetes import KubernetesException, SPOT_INTERRUPT_EXITCODE
 from .kube_utils import validate_kube_labels, parse_kube_keyvalue_list
 
 try:
@@ -548,6 +550,29 @@ def task_pre_step(
             self._save_logs_sidecar = Sidecar("save_logs_periodically")
             self._save_logs_sidecar.start()
 
+            # Set up signal handling for spot termination
+            main_pid = os.getpid()
+
+            def _termination_timer():
+                time.sleep(30)
+                os.kill(main_pid, signal.SIGALRM)
+
+            def _spot_term_signal_handler(*args, **kwargs):
+                if os.path.isfile(current.spot_termination_notice):
+                    print(
+                        "Spot instance termination detected. Starting a timer to end the Metaflow task."
+                    )
+                    timer_thread = threading.Thread(
+                        target=_termination_timer, daemon=True
+                    )
+                    timer_thread.start()
+
+            def _curtain_call(*args, **kwargs):
+                # custom exit code in case of Spot termination
+                sys.exit(SPOT_INTERRUPT_EXITCODE)
+
+            signal.signal(signal.SIGUSR1, _spot_term_signal_handler)
+            signal.signal(signal.SIGALRM, _curtain_call)
             # Start spot termination monitor sidecar.
             # TODO: A nicer way to pass the main process id to a Sidecar, in order to allow sidecars to send signals back to the main process.
             os.environ["MF_MAIN_PID"] = str(os.getpid())

metaflow/plugins/retry_decorator.py

@@ -1,24 +1,15 @@
 from enum import Enum
-import os
-import signal
-import sys
-import threading
-from time import sleep
 
 from metaflow.decorators import StepDecorator
 from metaflow.exception import MetaflowException
 from metaflow.metaflow_config import MAX_ATTEMPTS
-from metaflow import current
 
 
 class RetryEvents(Enum):
     STEP = "step"
     PREEMPT = "instance-preemption"
 
 
-PLATFORM_EVICTED_EXITCODE = 234
-
-
 class RetryDecorator(StepDecorator):
     """
     Specifies the number of times the task corresponding
@@ -76,40 +67,5 @@ def _known_event(event: str):
                     % ", ".join("*%s*" % event for event in unsupported_events)
                 )
 
-    def task_pre_step(
-        self,
-        step_name,
-        task_datastore,
-        metadata,
-        run_id,
-        task_id,
-        flow,
-        graph,
-        retry_count,
-        max_user_code_retries,
-        ubf_context,
-        inputs,
-    ):
-        pid = os.getpid()
-
-        def _termination_timer():
-            sleep(30)
-            os.kill(pid, signal.SIGALRM)
-
-        def _spot_term_signal_handler(*args, **kwargs):
-            if os.path.isfile(current.spot_termination_notice):
-                print(
-                    "Spot instance termination detected. Starting a timer to end the Metaflow task."
-                )
-                timer_thread = threading.Thread(target=_termination_timer, daemon=True)
-                timer_thread.start()
-
-        def _curtain_call(*args, **kwargs):
-            # custom exit code in case of Spot termination
-            sys.exit(PLATFORM_EVICTED_EXITCODE)
-
-        signal.signal(signal.SIGUSR1, _spot_term_signal_handler)
-        signal.signal(signal.SIGALRM, _curtain_call)
-
     def step_task_retry_count(self):
         return int(self.attributes["times"]), 0