Skip to content

Commit c8351c5

Browse files
abhijeet-dhumalabhijeet-dhumal
committed
fix(spark): pre-create S3A event log dir before SparkContext init to prevent silent materialize failure
Signed-off-by: abhijeet-dhumal <abhijeetdhumal652@gmail.com>
1 parent 835cda8 commit c8351c5

2 files changed

Lines changed: 141 additions & 1 deletion

File tree

‎sdk/python/feast/infra/compute_engines/spark/utils.py‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
import logging
2+
import os
13
from typing import Dict, Iterable, Literal, Optional
24

35
import pandas as pd
@@ -9,6 +11,68 @@
911
from feast.infra.common.serde import SerializedArtifacts
1012
from feast.utils import _convert_arrow_to_proto, _run_pyarrow_field_mapping
1113

14+
logger = logging.getLogger(__name__)
15+
16+
17+
def _ensure_s3a_event_log_dir(spark_config: Dict[str, str]) -> None:
18+
"""Pre-create the S3A event log prefix before SparkContext initialisation.
19+
20+
Spark's EventLogFileWriter.requireLogBaseDirAsDirectory() is called inside
21+
SparkContext.__init__ and crashes if the S3A path doesn't exist yet (S3 has no
22+
real directories, so an empty prefix returns a 404). This function writes a
23+
zero-byte placeholder so the prefix exists before SparkContext is built.
24+
25+
This is only attempted when:
26+
- spark.eventLog.enabled == "true"
27+
- spark.eventLog.dir starts with "s3a://"
28+
Failures are non-fatal: Spark will surface its own error if the dir is still missing.
29+
"""
30+
if spark_config.get("spark.eventLog.enabled", "false").lower() != "true":
31+
return
32+
event_dir = spark_config.get("spark.eventLog.dir", "")
33+
if not event_dir.startswith("s3a://"):
34+
return
35+
36+
path = event_dir[len("s3a://") :]
37+
bucket, _, prefix = path.partition("/")
38+
prefix = prefix.rstrip("/") + "/"
39+
placeholder_key = prefix + ".keep"
40+
41+
endpoint = spark_config.get(
42+
"spark.hadoop.fs.s3a.endpoint",
43+
os.environ.get("FEAST_S3A_ENDPOINT", ""),
44+
)
45+
access_key = os.environ.get("AWS_ACCESS_KEY_ID", "")
46+
secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY", "")
47+
48+
try:
49+
import boto3
50+
from botocore.client import Config
51+
52+
s3 = boto3.client(
53+
"s3",
54+
endpoint_url=endpoint if endpoint else None,
55+
aws_access_key_id=access_key or None,
56+
aws_secret_access_key=secret_key or None,
57+
config=Config(signature_version="s3v4"),
58+
)
59+
resp = s3.list_objects_v2(Bucket=bucket, Prefix=prefix, MaxKeys=1)
60+
if resp.get("KeyCount", 0) == 0:
61+
s3.put_object(Bucket=bucket, Key=placeholder_key, Body=b"")
62+
logger.debug(
63+
"Created S3A event log dir placeholder: s3a://%s/%s",
64+
bucket,
65+
placeholder_key,
66+
)
67+
except Exception as exc:
68+
logger.warning(
69+
"Could not pre-create S3A event log dir s3a://%s/%s — "
70+
"SparkContext may fail if the path still doesn't exist: %s",
71+
bucket,
72+
prefix,
73+
exc,
74+
)
75+
1276

1377
def get_or_create_new_spark_session(
1478
spark_config: Optional[Dict[str, str]] = None,
@@ -17,6 +81,10 @@ def get_or_create_new_spark_session(
1781
if not spark_session:
1882
spark_builder = SparkSession.builder
1983
if spark_config:
84+
# Spark's EventLogFileWriter.requireLogBaseDirAsDirectory() is called
85+
# during SparkContext.__init__ and will crash if the S3A event log
86+
# prefix doesn't exist yet. Ensure the prefix exists first.
87+
_ensure_s3a_event_log_dir(spark_config)
2088
spark_builder = spark_builder.config(
2189
conf=SparkConf().setAll([(k, v) for k, v in spark_config.items()])
2290
)

‎sdk/python/tests/component/spark/test_compute.py‎

Lines changed: 73 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
from datetime import timedelta
22
from typing import cast
3-
from unittest.mock import MagicMock
3+
from unittest.mock import MagicMock, patch
44

55
import pytest
66
from pyspark.sql import DataFrame
@@ -15,6 +15,7 @@
1515
from feast.infra.common.retrieval_task import HistoricalRetrievalTask
1616
from feast.infra.compute_engines.spark.compute import SparkComputeEngine
1717
from feast.infra.compute_engines.spark.job import SparkDAGRetrievalJob
18+
from feast.infra.compute_engines.spark.utils import _ensure_s3a_event_log_dir
1819
from feast.infra.offline_stores.contrib.spark_offline_store.spark import (
1920
SparkOfflineStore,
2021
)
@@ -192,5 +193,76 @@ def tqdm_builder(length):
192193
spark_environment.teardown()
193194

194195

196+
# ---------------------------------------------------------------------------
197+
# Unit tests for _ensure_s3a_event_log_dir — no Spark dependency needed
198+
# ---------------------------------------------------------------------------
199+
200+
201+
def _base_conf(event_log_dir: str) -> dict:
202+
return {
203+
"spark.eventLog.enabled": "true",
204+
"spark.eventLog.dir": event_log_dir,
205+
"spark.hadoop.fs.s3a.endpoint": "http://minio:9000",
206+
}
207+
208+
209+
@patch("feast.infra.compute_engines.spark.utils.boto3")
210+
def test_ensure_s3a_event_log_dir_creates_placeholder_when_empty(mock_boto3):
211+
"""S3A prefix doesn't exist → placeholder object is written."""
212+
s3 = MagicMock()
213+
mock_boto3.client.return_value = s3
214+
s3.list_objects_v2.return_value = {"KeyCount": 0}
215+
216+
_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))
217+
218+
s3.list_objects_v2.assert_called_once_with(
219+
Bucket="my-bucket", Prefix="spark-events/", MaxKeys=1
220+
)
221+
s3.put_object.assert_called_once_with(
222+
Bucket="my-bucket", Key="spark-events/.keep", Body=b""
223+
)
224+
225+
226+
@patch("feast.infra.compute_engines.spark.utils.boto3")
227+
def test_ensure_s3a_event_log_dir_skips_when_prefix_exists(mock_boto3):
228+
"""S3A prefix already has objects → no placeholder written."""
229+
s3 = MagicMock()
230+
mock_boto3.client.return_value = s3
231+
s3.list_objects_v2.return_value = {"KeyCount": 3}
232+
233+
_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))
234+
235+
s3.put_object.assert_not_called()
236+
237+
238+
@patch("feast.infra.compute_engines.spark.utils.boto3")
239+
def test_ensure_s3a_event_log_dir_noop_when_event_log_disabled(mock_boto3):
240+
"""spark.eventLog.enabled != true → boto3 never called."""
241+
_ensure_s3a_event_log_dir(
242+
{"spark.eventLog.enabled": "false", "spark.eventLog.dir": "s3a://b/p/"}
243+
)
244+
mock_boto3.client.assert_not_called()
245+
246+
247+
@patch("feast.infra.compute_engines.spark.utils.boto3")
248+
def test_ensure_s3a_event_log_dir_noop_for_non_s3a_path(mock_boto3):
249+
"""Non-S3A paths (hdfs://, file://, etc.) are left untouched."""
250+
_ensure_s3a_event_log_dir(
251+
{"spark.eventLog.enabled": "true", "spark.eventLog.dir": "hdfs:///spark-logs"}
252+
)
253+
mock_boto3.client.assert_not_called()
254+
255+
256+
@patch("feast.infra.compute_engines.spark.utils.boto3")
257+
def test_ensure_s3a_event_log_dir_non_fatal_on_s3_error(mock_boto3):
258+
"""boto3 errors are swallowed — SparkContext will surface its own error."""
259+
s3 = MagicMock()
260+
mock_boto3.client.return_value = s3
261+
s3.list_objects_v2.side_effect = Exception("connection refused")
262+
263+
# Must not raise
264+
_ensure_s3a_event_log_dir(_base_conf("s3a://my-bucket/spark-events/"))
265+
266+
195267
if __name__ == "__main__":
196268
test_spark_compute_engine_get_historical_features()

0 commit comments

Comments
 (0)