jscheffl commented on code in PR #65838: URL: https://github.com/apache/airflow/pull/65838#discussion_r3143429395
########## providers/edge3/src/airflow/providers/edge3/cli/worker.py: ########## Review Comment: Can you change the wording slightly in the PR? Because "Error starting" as a message after kill has been made sounds a bit strange. <img width="1056" height="121" alt="Image" src="https://github.com/user-attachments/assets/0112f61e-6410-48c5-ac6a-65926db3d220" /> ```suggestion log_chunk_data=f"Error executing job:\n{ex_txt}", ``` ########## providers/edge3/src/airflow/providers/edge3/cli/worker.py: ########## @@ -210,18 +217,112 @@ def signal_status(self): ) def signal_drain(self): - self.drain = True - logger.info("Request to shut down Edge Worker received, waiting for jobs to complete.") + if self._start_draining(): + logger.info( + "Request to shut down Edge Worker received, waiting for jobs to complete. %s", + self._drain_policy_description(), + ) def shutdown_handler(self): - if self.drain: + if not self._start_draining(): return - self.drain = True logger.info("SIGTERM received. Sending SIGTERM to all jobs and quit") + try: + loop = get_running_loop() + except RuntimeError: + pass + else: + task = loop.create_task( + self._push_drain_notice_to_all_jobs( + "Edge worker received external SIGTERM; terminating task supervisor." + ) + ) + self.background_tasks.add(task) + task.add_done_callback(self.background_tasks.discard) + self._terminate_jobs(signal.SIGTERM) + + def _start_draining(self) -> bool: + """Mark drain start. Returns ``True`` on the first call, ``False`` on subsequent calls.""" + if self.drain: + return False + self.drain = True + self.drain_started_at = time.monotonic() + return True + + def _drain_policy_description(self) -> str: + """One-line description of the configured drain-timeout policy for logging.""" + if self.drain_timeout_sec <= 0: + return "Drain timeout disabled; will wait indefinitely." + return ( + f"Drain timeout: {self.drain_timeout_sec}s (then SIGTERM), " + f"kill grace: {self.drain_kill_grace_sec}s (then SIGKILL)." + ) + + def _terminate_jobs(self, sig: int) -> None: + """Send ``sig`` to every running job process. Safe to call repeatedly.""" for job in self.jobs: if job.process.pid: with suppress(ProcessLookupError, PermissionError): - os.kill(job.process.pid, signal.SIGTERM) + os.kill(job.process.pid, sig) + + async def _push_drain_notice_to_all_jobs(self, message: str) -> None: + """Best-effort push of ``message`` into each running job's task log stream.""" + + async def push_one(job: Job) -> None: + try: + await logs_push( + task=job.edge_job.key, + log_chunk_time=timezone.utcnow(), + log_chunk_data=message, Review Comment: Can you add a line-feed to the message? Normally this is by logging formatter but here you post a single line. Ends-up missing line feed: <img width="1056" height="121" alt="Image" src="https://github.com/user-attachments/assets/5412ceb1-72ea-4224-ac04-6436a1170366" /> ```suggestion log_chunk_data=f"{message}\n", ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
