aws / aws/sagemaker-python-sdk
Importing sagemaker library in airflow, makes airflow fail to upload logs to s3.
- Dominant language
- Python
- Stars
- 2.3k
- Forks
- 1.3k
- Avg merge
- 1d 22h
- Merged PRs (30d)
- 35
Description
**Describe the bug**
Importing the SageMaker library at the top of the Airflow plugin or hook declaration file causes issues with uploading logs to S3.
**To reproduce**
Setup a sample dag with the plugin as mentioned below. main problem is around importing sagemaker library.
### Sample dag
```python
from datetime import datetime
from airflow import DAG
from airflow.operators.python import PythonOperator
import logging
LOGGER = logging.getLogger(__name__)
def test_function(**_):
LOGGER.info("Success_test_dag")
with DAG(
dag_id="test_dag",
description=__doc__,
schedule_interval=None, # Every month on 6th at 6:00 PM IST
default_args={
"depends_on_past": False,
"owner": "Data Platform",
"start_date": datetime(2025, 4, 29),
"retries": 1,
},
catchup=False,
max_active_runs=1,
) as dag:
PythonOperator(
task_id="test_task",
python_callable=test_function,
provide_context=True,
)
```
### sample plugin
plugins/sample_hook.py
```python
from airflow.hooks.base import BaseHook
from airflow.plugins_manager import AirflowPlugin
import sagemaker
class SampleHook(BaseHook):
...
class SamplePlugin(AirflowPlugin):
name = "sample_hook"
hooks = [SampleHook]
```
Env variables:
AIRFLOW__LOGGING__REMOTE_BASE_LOG_FOLDER=s3://airflow/logs
AIRFLOW__LOGGING__REMOTE_LOGGING=True
AIRFLOW__LOGGING__REMOTE_LOG_CONN_ID=s3_airflow_logging
AIRFLOW__LOGGING__REMOTE_LOGGING=True
**Expected behavior**
When you run the DAG, it shows as successful, but Airflow fails to retrieve the logs from S3, even though the s3_airflow_logging connection is configured correctly. However, if you remove the line that imports the SageMaker library, the logs are successfully stored and retrieved from S3 as expected.
**Screenshots or logs**
When sagemaker is not imported
```logs
da08261f9df2
▼ Log message source details
*** Found logs in s3:
*** * s3://airflow/logs/dag_id=test_dag/run_id=manual__2025-06-16T10:10:39.397088+00:00/task_id=test_task/attempt=1.log
▲▲▲ Log group end
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{local_task_job_runner.py:123}} ▼ Pre task execution logs
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{taskinstance.py:2613}} INFO - Dependencies all met for dep_context=non-requeueable deps ti=
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{taskinstance.py:2613}} INFO - Dependencies all met for dep_context=requeueable deps ti=
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{taskinstance.py:2866}} INFO - Starting attempt 1 of 2
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{taskinstance.py:2889}} INFO - Executing on 2025-06-16 10:10:39.397088+00:00
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{standard_task_runner.py:72}} INFO - Started process 29 to run task
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{standard_task_runner.py:104}} INFO - Running: ['***', 'tasks', 'run', 'test_dag', 'test_task', 'manual__2025-06-16T10:10:39.397088+00:00', '--job-id', '65', '--raw', '--subdir', 'DAGS_FOLDER/test.py', '--cfg-path', '/tmp/tmpmn1u32bm']
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{standard_task_runner.py:105}} INFO - Job 65: Subtask test_task
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{task_command.py:467}} INFO - Running on host da08261f9df2
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{taskinstance.py:3132}} INFO - Exporting env vars: AIRFLOW_CTX_DAG_OWNER='Data Platform' AIRFLOW_CTX_DAG_ID='test_dag' AIRFLOW_CTX_TASK_ID='test_task' AIRFLOW_CTX_EXECUTION_DATE='2025-06-16T10:10:39.397088+00:00' AIRFLOW_CTX_TRY_NUMBER='1' AIRFLOW_CTX_DAG_RUN_ID='manual__2025-06-16T10:10:39.397088+00:00'
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{taskinstance.py:731}} ▲▲▲ Log group end
[2025-06-16, 15:40:42 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{test.py:10}} INFO - Success_test_dag
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{python.py:240}} INFO - Done. Returned value was: None
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{taskinstance.py:340}} ▼ Post task execution logs
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:29, thread = MainThread:256226136576032}} {{taskinstance.py:352}} INFO - Marking task as SUCCESS. dag_id=test_dag, task_id=test_task, run_id=manual__2025-06-16T10:10:39.397088+00:00, execution_date=20250616T101039, start_date=20250616T101042, end_date=20250616T101043
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{local_task_job_runner.py:266}} INFO - Task exited with return code 0
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{taskinstance.py:3895}} INFO - 0 downstream tasks scheduled from follow-on schedule check
[2025-06-16, 15:40:43 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:256226136576032}} {{local_task_job_runner.py:245}} ▲▲▲ Log group end
```
When sagemaker is imported
```logs
da08261f9df2
▼ Log message source details
*** No logs found on s3 for ti=
*** Found local files:
*** * /opt/airflow/logs/dag_id=test_dag/run_id=manual__2025-06-16T10:18:26.179356+00:00/task_id=test_task/attempt=1.log
▲▲▲ Log group end
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{local_task_job_runner.py:123}} ▼ Pre task execution logs
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{taskinstance.py:2613}} INFO - Dependencies all met for dep_context=non-requeueable deps ti=
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{taskinstance.py:2613}} INFO - Dependencies all met for dep_context=requeueable deps ti=
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{taskinstance.py:2866}} INFO - Starting attempt 1 of 2
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{taskinstance.py:2889}} INFO - Executing on 2025-06-16 10:18:26.179356+00:00
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{standard_task_runner.py:72}} INFO - Started process 38 to run task
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{standard_task_runner.py:104}} INFO - Running: ['***', 'tasks', 'run', 'test_dag', 'test_task', 'manual__2025-06-16T10:18:26.179356+00:00', '--job-id', '71', '--raw', '--subdir', 'DAGS_FOLDER/test.py', '--cfg-path', '/tmp/tmpegkx_0gj']
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{standard_task_runner.py:105}} INFO - Job 71: Subtask test_task
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{task_command.py:467}} INFO - Running on host da08261f9df2
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{taskinstance.py:3132}} INFO - Exporting env vars: AIRFLOW_CTX_DAG_OWNER='Data Platform' AIRFLOW_CTX_DAG_ID='test_dag' AIRFLOW_CTX_TASK_ID='test_task' AIRFLOW_CTX_EXECUTION_DATE='2025-06-16T10:18:26.179356+00:00' AIRFLOW_CTX_TRY_NUMBER='1' AIRFLOW_CTX_DAG_RUN_ID='manual__2025-06-16T10:18:26.179356+00:00'
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{taskinstance.py:731}} ▲▲▲ Log group end
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{test.py:10}} INFO - Success_test_dag
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{python.py:240}} INFO - Done. Returned value was: None
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{taskinstance.py:340}} ▼ Post task execution logs
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:38, thread = MainThread:248620777738272}} {{taskinstance.py:352}} INFO - Marking task as SUCCESS. dag_id=test_dag, task_id=test_task, run_id=manual__2025-06-16T10:18:26.179356+00:00, execution_date=20250616T101826, start_date=20250616T101829, end_date=20250616T101829
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{local_task_job_runner.py:266}} INFO - Task exited with return code 0
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{taskinstance.py:3895}} INFO - 0 downstream tasks scheduled from follow-on schedule check
[2025-06-16, 15:48:29 IST] {{process = ForkPoolWorker-3:26, thread = MainThread:248620777738272}} {{local_task_job_runner.py:245}} ▲▲▲ Log group end
```
**System information**
Version:
- airflow: 2.10.4
- python: 3.9
- sagemaker: 2.233.0
Contributor guide
Research direction
Reproduce the issue with plugins/sample_hook.py and the sample DAG, comparing runs with and without the import sagemaker line. Start by examining the SageMaker import's effect on Airflow's S3 remote logging, and consider the issue done when the DAG logs are stored and retrieved from s3://airflow/logs in both cases.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- aws, python
- Domain
- cloud, observability
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Stale
- Clarity
- Needs clarification
- Newbie friendliness
- 25/100