Misleading Warning: "Custom Operator Inheriting from SageMakerTrainingOperator Incorrectly Triggers execute cannot be called outside TaskInstance"

[luancaarvalho]

Apache Airflow Provider(s)

amazon

Versions of Apache Airflow Providers

8.28.0

Apache Airflow version

2.10.2

Operating System

ubuntu

Deployment

Docker-Compose

Deployment details

No response

What happened

We have created a custom operator, CustomSageMakerTrainingOperator, that inherits from SageMakerTrainingOperator. Although the operator is executed within a proper TaskInstance context, we receive a persistent and seemingly incorrect warning in the logs:

[2025-01-03, 22:57:16 UTC] {baseoperator.py:405} WARNING - CustomSageMakerTrainingOperator.execute cannot be called outside TaskInstance!

This warning appears even though our operator’s execute method is only called as part of a TaskInstance execution. The operator code, which extends the functionality of SageMakerTrainingOperator, confirms that no misuse exists. This suggests that the warning is a false positive, potentially due to an issue within Airflow’s internals handling of such inherited operators.

What you think should happen instead

The custom operator should run the SageMaker training job without generating the misleading warning, since the execute is indeed called in a valid TaskInstance context.

How to reproduce

The custom operator should run the SageMaker training job without generating the misleading warning. The warning should not appear when the execute method is correctly invoked within a TaskInstance context. Ideally, Airflow should correctly assess the context for inherited operators like CustomSageMakerTrainingOperator and refrain from showing this warning.

Anything else

Below is the core of the custom operator code (anonymized). Full code snippet is in the original text:

import logging
from typing import Optional, Dict, Any

from airflow.providers.amazon.aws.operators.sagemaker import SageMakerTrainingOperator
from airflow.exceptions import AirflowSkipException

class CustomSageMakerTrainingOperator(SageMakerTrainingOperator):
    ...

I've tried to confirm that the execute method is only invoked via a TaskInstance, so I'm not sure why the warning triggers. Possibly there's a check in baseoperator.py or within SageMakerTrainingOperator that incorrectly flags this scenario.

Potential Evaluated Solutions:

Investigate if the super().init call with config=self.config is triggering unexpected logic in SageMakerTrainingOperator.
Check if there's an internal method or property in SageMakerTrainingOperator that must be overridden or called differently.
Possibly a bug in the logic that triggers the “outside TaskInstance” warning for operators inheriting from SageMakerTrainingOperator.

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[eladkal]

cc @vincbeck @ferruzzi

[insomnes]

I've met this issue with my extended operators too. I believe any child class with super().execute call would lead to this warning.

In #41470 comment says:

The fix is now available in Airflow 2.10.3