amoghrajesh commented on code in PR #76413:
URL: https://github.com/apache/airflow/pull/76413#discussion_r4508686333
##########
providers/amazon/src/airflow/providers/amazon/aws/triggers/sagemaker.py:
##########
@@ -123,17 +124,27 @@ def hook(self) -> AwsGenericHook:
config=self.botocore_config,
)
+ STOP_METHODS_BY_JOB_TYPE: ClassVar[dict[str, str]] = {
+ "training": "stop_training_job",
+ "processing": "stop_processing_job",
+ }
+
async def on_kill(self) -> None:
- """Stop the training job when the user acts on the deferred task."""
- if not self.cancel_on_kill or self.job_type.lower() != "training":
+ """Stop the job when the user acts on the deferred task."""
+ if not self.cancel_on_kill:
+ return
+ # Absent job types are a no-op, not a call against another type's stop
API.
+ stop_method = self.STOP_METHODS_BY_JOB_TYPE.get(self.job_type.lower())
+ if stop_method is None:
return
- self.log.info("Stopping SageMaker training job %s.", self.job_name)
+ self.log.info("Stopping SageMaker %s job %s.", self.job_type,
self.job_name)
hook: SageMakerHook = self.hook() # type: ignore[assignment]
try:
- await sync_to_async(hook.stop_training_job)(self.job_name)
+ await sync_to_async(getattr(hook, stop_method))(self.job_name)
Review Comment:
The map is now `_STOP_METHODS_BY_JOB_TYPE` and its values are spelled
`SageMakerHook.stop_training_job.__name__`, so a typo fails at import rather
than at kill time.
--
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]
