fat-catTW commented on code in PR #74497:
URL: https://github.com/apache/airflow/pull/74497#discussion_r4306477767
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/job.py:
##########
@@ -243,12 +244,38 @@ def execute(self, context: Context):
# still looping. So unlike KubernetesPodOperator, this
path must
# fail loudly when the sidecar cannot be killed rather
than hang.
xcom_result.append(self.extract_xcom(pod=pod,
ignore_kill_failure=False))
- self.job = self.hook.wait_until_job_complete(
- job_name=self.job.metadata.name,
- namespace=self.job.metadata.namespace,
- job_poll_interval=self.job_poll_interval,
- )
+ while True:
+ self.job = self.hook.wait_until_job_complete(
+ job_name=self.job.metadata.name,
+ namespace=self.job.metadata.namespace,
+ job_poll_interval=self.job_poll_interval,
+ return_on_suspension=True,
+ )
+
+ if self.hook.is_job_complete(job=self.job):
+ break
+
+ if getattr(self.job.spec, "suspend", False) and not (
+ self.job.status and self.job.status.active
+ ):
+ self.log.info(
+ "Job '%s' was preempted and is suspended. Going
back to pod discovery.",
+ self.job.metadata.name,
+ )
+ # Wait until it is readmitted before pod discovery
+ while getattr(self.job.spec, "suspend", False) and not
(
+ self.job.status and self.job.status.active
+ ):
+ sleep(self.job_poll_interval)
+ self.job = self.hook.get_job_status(
+ job_name=self.job.metadata.name,
namespace=self.job.metadata.namespace
+ )
+ self.pods =
self.get_pods(pod_request_obj=self.pod_request_obj, context=context)
Review Comment:
I think this can race with pod creation after readmission too.
The loop exits as soon as spec.suspend becomes false, but the replacement
pod may not be discoverable yet, so `get_pods()` can raise No pods running.
Maybe we can wait or retry pod discovery here.
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/job.py:
##########
Review Comment:
I think this initial `get_pods()` can fail before the suspension-aware
polling begins.
A Kueue-managed Job can be suspended before a pod is created, or preempted
before discovery, so this can raise `No pods running` instead of waiting for
readmission.
Maybe we can check for suspension before pod discovery and wait for a pod to
become available?
##########
providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/operators/job.py:
##########
@@ -287,6 +314,15 @@ def execute_deferrable(self):
)
def execute_complete(self, context: Context, event: dict, **kwargs):
+ if event.get("status") == "readmitted":
+ self.log.info("Job preempted but now readmitted. Re-deferring to
restart pod discovery.")
+ self.job = self.hook.get_job_status(job_name=event["name"],
namespace=event["namespace"])
+ self.job_request_obj = self.build_job_request_obj(context)
+ self.pod_request_obj = self.build_pod_request_obj(context)
+ self.pods = self.get_pods(pod_request_obj=self.pod_request_obj,
context=context)
Review Comment:
Same as above.
--
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]
