Skip to content

Commit b6392ae

Browse files
hussein-awalao-nikolas
hussein-awala
and
o-nikolas
authored
Support deleting the local log files when using remote logging (#29772)
* add a new config * add remote_task_handler_kwargs conf and add its content as kwargs for remote task handlers * add delete_local_logs to logging tasks doc Co-authored-by: Niko Oliveira <onikolas@amazon.com> --------- Co-authored-by: Niko Oliveira <onikolas@amazon.com>
1 parent c405ecb commit b6392ae

13 files changed

Lines changed: 305 additions & 29 deletions

File tree

‎airflow/config_templates/airflow_local_settings.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,7 @@
203203
# WASB buckets should start with "wasb"
204204
# just to help Airflow select correct handler
205205
REMOTE_BASE_LOG_FOLDER: str = conf.get_mandatory_value("logging", "REMOTE_BASE_LOG_FOLDER")
206+
REMOTE_TASK_HANDLER_KWARGS = conf.getjson("logging", "REMOTE_TASK_HANDLER_KWARGS", fallback={})
206207

207208
if REMOTE_BASE_LOG_FOLDER.startswith("s3://"):
208209
S3_REMOTE_HANDLERS: dict[str, dict[str, str | None]] = {
@@ -252,7 +253,6 @@
252253
"wasb_log_folder": REMOTE_BASE_LOG_FOLDER,
253254
"wasb_container": "airflow-logs",
254255
"filename_template": FILENAME_TEMPLATE,
255-
"delete_local_copy": False,
256256
},
257257
}
258258

@@ -315,3 +315,4 @@
315315
"section 'elasticsearch' if you are using Elasticsearch. In the other case, "
316316
"'remote_base_log_folder' option in the 'logging' section."
317317
)
318+
DEFAULT_LOGGING_CONFIG["handlers"]["task"].update(REMOTE_TASK_HANDLER_KWARGS)

‎airflow/config_templates/config.yml‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -606,6 +606,14 @@ logging:
606606
type: string
607607
example: ~
608608
default: ""
609+
delete_local_logs:
610+
description: |
611+
Whether the local log files for GCS, S3, WASB and OSS remote logging should be deleted after
612+
they are uploaded to the remote location.
613+
version_added: 2.6.0
614+
type: string
615+
example: ~
616+
default: "False"
609617
google_key_path:
610618
description: |
611619
Path to Google Credential JSON file. If omitted, authorization based on `the Application Default
@@ -628,6 +636,16 @@ logging:
628636
type: string
629637
example: ~
630638
default: ""
639+
remote_task_handler_kwargs:
640+
description: |
641+
The remote_task_handler_kwargs param is loaded into a dictionary and passed to __init__ of remote
642+
task handler and it overrides the values provided by Airflow config. For example if you set
643+
`delete_local_logs=False` and you provide ``{{"delete_local_copy": true}}``, then the local
644+
log files will be deleted after they are uploaded to remote location.
645+
version_added: 2.6.0
646+
type: string
647+
example: '{"delete_local_copy": true}'
648+
default: ""
631649
encrypt_s3_logs:
632650
description: |
633651
Use server-side encryption for logs stored in S3

‎airflow/config_templates/default_airflow.cfg‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,10 @@ remote_logging = False
345345
# reading logs, not writing them.
346346
remote_log_conn_id =
347347

348+
# Whether the local log files for GCS, S3, WASB and OSS remote logging should be deleted after
349+
# they are uploaded to the remote location.
350+
delete_local_logs = False
351+
348352
# Path to Google Credential JSON file. If omitted, authorization based on `the Application Default
349353
# Credentials
350354
# <https://cloud.google.com/docs/authentication/production#finding_credentials_automatically>`__ will
@@ -359,6 +363,13 @@ google_key_path =
359363
# Stackdriver logs should start with "stackdriver://"
360364
remote_base_log_folder =
361365

366+
# The remote_task_handler_kwargs param is loaded into a dictionary and passed to __init__ of remote
367+
# task handler and it overrides the values provided by Airflow config. For example if you set
368+
# `delete_local_logs=False` and you provide ``{{"delete_local_copy": true}}``, then the local
369+
# log files will be deleted after they are uploaded to remote location.
370+
# Example: remote_task_handler_kwargs = {{"delete_local_copy": true}}
371+
remote_task_handler_kwargs =
372+
362373
# Use server-side encryption for logs stored in S3
363374
encrypt_s3_logs = False
364375

‎airflow/providers/alibaba/cloud/log/oss_task_handler.py‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,9 @@
2020
import contextlib
2121
import os
2222
import pathlib
23+
import shutil
24+
25+
from packaging.version import Version
2326

2427
from airflow.compat.functools import cached_property
2528
from airflow.configuration import conf
@@ -28,21 +31,35 @@
2831
from airflow.utils.log.logging_mixin import LoggingMixin
2932

3033

34+
def get_default_delete_local_copy():
35+
"""Load delete_local_logs conf if Airflow version > 2.6 and return False if not
36+
TODO: delete this function when min airflow version >= 2.6
37+
"""
38+
from airflow.version import version
39+
40+
if Version(version) < Version("2.6"):
41+
return False
42+
return conf.getboolean("logging", "delete_local_logs")
43+
44+
3145
class OSSTaskHandler(FileTaskHandler, LoggingMixin):
3246
"""
3347
OSSTaskHandler is a python log handler that handles and reads
3448
task instance logs. It extends airflow FileTaskHandler and
3549
uploads to and reads from OSS remote storage.
3650
"""
3751

38-
def __init__(self, base_log_folder, oss_log_folder, filename_template=None):
52+
def __init__(self, base_log_folder, oss_log_folder, filename_template=None, **kwargs):
3953
self.log.info("Using oss_task_handler for remote logging...")
4054
super().__init__(base_log_folder, filename_template)
4155
(self.bucket_name, self.base_folder) = OSSHook.parse_oss_url(oss_log_folder)
4256
self.log_relative_path = ""
4357
self._hook = None
4458
self.closed = False
4559
self.upload_on_close = True
60+
self.delete_local_copy = (
61+
kwargs["delete_local_copy"] if "delete_local_copy" in kwargs else get_default_delete_local_copy()
62+
)
4663

4764
@cached_property
4865
def hook(self):
@@ -92,7 +109,9 @@ def close(self):
92109
if os.path.exists(local_loc):
93110
# read log and remove old logs to get just the latest additions
94111
log = pathlib.Path(local_loc).read_text()
95-
self.oss_write(log, remote_loc)
112+
oss_write = self.oss_write(log, remote_loc)
113+
if oss_write and self.delete_local_copy:
114+
shutil.rmtree(os.path.dirname(local_loc))
96115

97116
# Mark closed so we don't double write if close is called twice
98117
self.closed = True
@@ -154,15 +173,16 @@ def oss_read(self, remote_log_location, return_error=False):
154173
if return_error:
155174
return msg
156175

157-
def oss_write(self, log, remote_log_location, append=True):
176+
def oss_write(self, log, remote_log_location, append=True) -> bool:
158177
"""
159-
Writes the log to the remote_log_location. Fails silently if no hook
160-
was created.
178+
Writes the log to the remote_log_location and return `True` when done. Fails silently
179+
and return `False` if no log was created.
161180
162181
:param log: the log to write to the remote_log_location
163182
:param remote_log_location: the log's location in remote storage
164183
:param append: if False, any existing log file is overwritten. If True,
165184
the new log is appended to any existing logs.
185+
:return: whether the log is successfully written to remote location or not.
166186
"""
167187
oss_remote_log_location = f"{self.base_folder}/{remote_log_location}"
168188
pos = 0
@@ -180,3 +200,5 @@ def oss_write(self, log, remote_log_location, append=True):
180200
str(pos),
181201
str(append),
182202
)
203+
return False
204+
return True

‎airflow/providers/amazon/aws/log/s3_task_handler.py‎

Lines changed: 30 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,9 @@
1919

2020
import os
2121
import pathlib
22+
import shutil
23+
24+
from packaging.version import Version
2225

2326
from airflow.compat.functools import cached_property
2427
from airflow.configuration import conf
@@ -27,6 +30,17 @@
2730
from airflow.utils.log.logging_mixin import LoggingMixin
2831

2932

33+
def get_default_delete_local_copy():
34+
"""Load delete_local_logs conf if Airflow version > 2.6 and return False if not
35+
TODO: delete this function when min airflow version >= 2.6
36+
"""
37+
from airflow.version import version
38+
39+
if Version(version) < Version("2.6"):
40+
return False
41+
return conf.getboolean("logging", "delete_local_logs")
42+
43+
3044
class S3TaskHandler(FileTaskHandler, LoggingMixin):
3145
"""
3246
S3TaskHandler is a python log handler that handles and reads
@@ -36,13 +50,18 @@ class S3TaskHandler(FileTaskHandler, LoggingMixin):
3650

3751
trigger_should_wrap = True
3852

39-
def __init__(self, base_log_folder: str, s3_log_folder: str, filename_template: str | None = None):
53+
def __init__(
54+
self, base_log_folder: str, s3_log_folder: str, filename_template: str | None = None, **kwargs
55+
):
4056
super().__init__(base_log_folder, filename_template)
4157
self.remote_base = s3_log_folder
4258
self.log_relative_path = ""
4359
self._hook = None
4460
self.closed = False
4561
self.upload_on_close = True
62+
self.delete_local_copy = (
63+
kwargs["delete_local_copy"] if "delete_local_copy" in kwargs else get_default_delete_local_copy()
64+
)
4665

4766
@cached_property
4867
def hook(self):
@@ -84,7 +103,9 @@ def close(self):
84103
if os.path.exists(local_loc):
85104
# read log and remove old logs to get just the latest additions
86105
log = pathlib.Path(local_loc).read_text()
87-
self.s3_write(log, remote_loc)
106+
write_to_s3 = self.s3_write(log, remote_loc)
107+
if write_to_s3 and self.delete_local_copy:
108+
shutil.rmtree(os.path.dirname(local_loc))
88109

89110
# Mark closed so we don't double write if close is called twice
90111
self.closed = True
@@ -164,23 +185,25 @@ def s3_read(self, remote_log_location: str, return_error: bool = False) -> str:
164185
return msg
165186
return ""
166187

167-
def s3_write(self, log: str, remote_log_location: str, append: bool = True, max_retry: int = 1):
188+
def s3_write(self, log: str, remote_log_location: str, append: bool = True, max_retry: int = 1) -> bool:
168189
"""
169-
Writes the log to the remote_log_location. Fails silently if no hook
170-
was created.
190+
Writes the log to the remote_log_location and return `True` when done. Fails silently
191+
and return `False` if no log was created.
171192
172193
:param log: the log to write to the remote_log_location
173194
:param remote_log_location: the log's location in remote storage
174195
:param append: if False, any existing log file is overwritten. If True,
175196
the new log is appended to any existing logs.
176197
:param max_retry: Maximum number of times to retry on upload failure
198+
:return: whether the log is successfully written to remote location or not.
177199
"""
178200
try:
179201
if append and self.s3_log_exists(remote_log_location):
180202
old_log = self.s3_read(remote_log_location)
181203
log = "\n".join([old_log, log]) if old_log else log
182204
except Exception:
183205
self.log.exception("Could not verify previous log to append")
206+
return False
184207

185208
# Default to a single retry attempt because s3 upload failures are
186209
# rare but occasionally occur. Multiple retry attempts are unlikely
@@ -199,3 +222,5 @@ def s3_write(self, log: str, remote_log_location: str, append: bool = True, max_
199222
self.log.warning("Failed attempt to write logs to %s, will retry", remote_log_location)
200223
else:
201224
self.log.exception("Could not write logs to %s", remote_log_location)
225+
return False
226+
return True

‎airflow/providers/google/cloud/log/gcs_task_handler.py‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,11 +19,13 @@
1919

2020
import logging
2121
import os
22+
import shutil
2223
from pathlib import Path
2324
from typing import Collection
2425

2526
# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
2627
from google.cloud import storage # type: ignore[attr-defined]
28+
from packaging.version import Version
2729

2830
from airflow.compat.functools import cached_property
2931
from airflow.configuration import conf
@@ -43,6 +45,17 @@
4345
logger = logging.getLogger(__name__)
4446

4547

48+
def get_default_delete_local_copy():
49+
"""Load delete_local_logs conf if Airflow version > 2.6 and return False if not
50+
TODO: delete this function when min airflow version >= 2.6
51+
"""
52+
from airflow.version import version
53+
54+
if Version(version) < Version("2.6"):
55+
return False
56+
return conf.getboolean("logging", "delete_local_logs")
57+
58+
4659
class GCSTaskHandler(FileTaskHandler, LoggingMixin):
4760
"""
4861
GCSTaskHandler is a python log handler that handles and reads
@@ -63,6 +76,8 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
6376
:param gcp_scopes: Comma-separated string containing OAuth2 scopes
6477
:param project_id: Project ID to read the secrets from. If not passed, the project ID from credentials
6578
will be used.
79+
:param delete_local_copy: Whether local log files should be deleted after they are downloaded when using
80+
remote logging
6681
"""
6782

6883
trigger_should_wrap = True
@@ -77,6 +92,7 @@ def __init__(
7792
gcp_keyfile_dict: dict | None = None,
7893
gcp_scopes: Collection[str] | None = _DEFAULT_SCOPESS,
7994
project_id: str | None = None,
95+
**kwargs,
8096
):
8197
super().__init__(base_log_folder, filename_template)
8298
self.remote_base = gcs_log_folder
@@ -87,6 +103,9 @@ def __init__(
87103
self.gcp_keyfile_dict = gcp_keyfile_dict
88104
self.scopes = gcp_scopes
89105
self.project_id = project_id
106+
self.delete_local_copy = (
107+
kwargs["delete_local_copy"] if "delete_local_copy" in kwargs else get_default_delete_local_copy()
108+
)
90109

91110
@cached_property
92111
def hook(self) -> GCSHook | None:
@@ -147,7 +166,9 @@ def close(self):
147166
# read log and remove old logs to get just the latest additions
148167
with open(local_loc) as logfile:
149168
log = logfile.read()
150-
self.gcs_write(log, remote_loc)
169+
gcs_write = self.gcs_write(log, remote_loc)
170+
if gcs_write and self.delete_local_copy:
171+
shutil.rmtree(os.path.dirname(local_loc))
151172

152173
# Mark closed so we don't double write if close is called twice
153174
self.closed = True
@@ -207,13 +228,14 @@ def _read(self, ti, try_number, metadata=None):
207228

208229
return "".join([f"*** {x}\n" for x in messages]) + "\n".join(logs), {"end_of_log": True}
209230

210-
def gcs_write(self, log, remote_log_location):
231+
def gcs_write(self, log, remote_log_location) -> bool:
211232
"""
212-
Writes the log to the remote_log_location. Fails silently if no log
213-
was created.
233+
Writes the log to the remote_log_location and return `True` when done. Fails silently
234+
and return `False` if no log was created.
214235
215236
:param log: the log to write to the remote_log_location
216237
:param remote_log_location: the log's location in remote storage
238+
:return: whether the log is successfully written to remote location or not.
217239
"""
218240
try:
219241
blob = storage.Blob.from_string(remote_log_location, self.client)
@@ -232,6 +254,8 @@ def gcs_write(self, log, remote_log_location):
232254
blob.upload_from_string(log, content_type="text/plain")
233255
except Exception as e:
234256
self.log.error("Could not write logs to %s: %s", remote_log_location, e)
257+
return False
258+
return True
235259

236260
@staticmethod
237261
def no_log_found(exc):

0 commit comments

Comments
 (0)