aws / aws/sagemaker-python-sdk

Importing sagemaker library in airflow, makes airflow fail to upload logs to s3.

Open
#5,207 0 comments 0 reactions 0 assignees View on GitHub
component: pysdk-team type: bug
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.