You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by gu...@apache.org on 2026/03/19 03:30:56 UTC
(spark) branch master updated: [SPARK-56062][PYTHON] Isolate memory_profiler to improve import time
This is an automated email from the ASF dual-hosted git repository.
gurwls223 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 811c7b780f33 [SPARK-56062][PYTHON] Isolate memory_profiler to improve import time
811c7b780f33 is described below
commit 811c7b780f33b6285f975496d1ca70a4538106a4
Author: Tian Gao <ga...@hotmail.com>
AuthorDate: Thu Mar 19 12:30:44 2026 +0900
[SPARK-56062][PYTHON] Isolate memory_profiler to improve import time
### What changes were proposed in this pull request?
* Moved `memory_profiler` related code to another file
* Make that module lazy so importing it won't trigger `import memory_profiler`
* Updated the code that uses `memory_profiler` related code to use the new module
### Why are the changes needed?
We should make `import pyspark` fast. Bootstrap time is important to user experience, and our CI speed. Also, we should make worker spawn faster so simple worker can be faster to boot. We have a pending SPIP about using simpler worker for UDF work and import time is critical.
`import pyspark` takes **220 ms** on master and will be **90 ms** after this PR. Similar acceleration will be achieved for `pyspark.worker`.
This would probably make a lot of our tests faster as a side effect. Including the newly added python data source simple worker test which timeout a lot in coverage run - https://github.com/apache/spark/actions/runs/23046212656/job/66936106272 .
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
I ran some profiler tests locally and they passed. CI is needed.
### Was this patch authored or co-authored using generative AI tooling?
No.
Closes #54895 from gaogaotiantian/isolate-memory-profiler.
Authored-by: Tian Gao <ga...@hotmail.com>
Signed-off-by: Hyukjin Kwon <gu...@apache.org>
---
python/pyspark/memory_profiler_ext.py | 190 +++++++++++++++++++++
python/pyspark/profiler.py | 128 +-------------
python/pyspark/sql/profiler.py | 8 +-
python/pyspark/sql/tests/test_python_datasource.py | 2 +-
python/pyspark/tests/test_memory_profiler.py | 6 +-
python/pyspark/tests/test_profiler.py | 2 +-
python/pyspark/worker.py | 11 +-
7 files changed, 204 insertions(+), 143 deletions(-)
diff --git a/python/pyspark/memory_profiler_ext.py b/python/pyspark/memory_profiler_ext.py
new file mode 100644
index 000000000000..215e160907bb
--- /dev/null
+++ b/python/pyspark/memory_profiler_ext.py
@@ -0,0 +1,190 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from types import CodeType
+from typing import Any, Optional, List, Iterator, Tuple, Type, TYPE_CHECKING, Callable
+import inspect
+import warnings
+
+if TYPE_CHECKING:
+ has_memory_profiler: bool
+ try:
+ from memory_profiler import CodeMap, LineProfiler
+
+ CodeMapForUDF: Type[CodeMap]
+ CodeMapForUDFV2: Type[CodeMap]
+ UDFLineProfiler: Type[LineProfiler]
+ UDFLineProfilerV2: Type[LineProfiler]
+ except Exception:
+ pass
+
+
+__all__ = [
+ "has_memory_profiler",
+ "CodeMapForUDF",
+ "CodeMapForUDFV2",
+ "UDFLineProfiler",
+ "UDFLineProfilerV2",
+]
+
+_module_initialized = False
+_has_memory_profiler = None
+
+
+def __getattr__(name: str) -> Any:
+ if name not in __all__:
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+ if not _module_initialized:
+ _init_module()
+ if name == "has_memory_profiler":
+ return _has_memory_profiler
+ elif name in globals():
+ return globals()[name]
+ else:
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+def _init_module() -> None:
+ global _has_memory_profiler
+ global _module_initialized
+
+ try:
+ from memory_profiler import CodeMap, LineProfiler
+
+ _has_memory_profiler = True
+ except Exception:
+ _has_memory_profiler = False
+
+ if not _has_memory_profiler:
+ _module_initialized = True
+ return
+
+ class CodeMapForUDF(CodeMap):
+ def add(
+ self,
+ code: Any,
+ toplevel_code: Optional[Any] = None,
+ *,
+ sub_lines: Optional[List] = None,
+ start_line: Optional[int] = None,
+ ) -> None:
+ if code in self:
+ return
+
+ if toplevel_code is None:
+ toplevel_code = code
+ filename = code.co_filename
+ if sub_lines is None or start_line is None:
+ sub_lines, start_line = inspect.getsourcelines(code)
+ linenos = range(start_line, start_line + len(sub_lines))
+ self._toplevel.append((filename, code, linenos))
+ self[code] = {}
+ else:
+ self[code] = self[toplevel_code]
+ for subcode in filter(inspect.iscode, code.co_consts):
+ self.add(subcode, toplevel_code=toplevel_code)
+
+ class CodeMapForUDFV2(CodeMap):
+ def add(
+ self,
+ code: Any,
+ toplevel_code: Optional[Any] = None,
+ ) -> None:
+ if code in self:
+ return
+
+ if toplevel_code is None:
+ toplevel_code = code
+ filename = code.co_filename
+ self._toplevel.append((filename, code))
+ self[code] = {}
+ else:
+ self[code] = self[toplevel_code]
+ for subcode in filter(inspect.iscode, code.co_consts):
+ self.add(subcode, toplevel_code=toplevel_code)
+
+ def items(self) -> Iterator[Tuple[str, Iterator[Tuple[int, Any]]]]:
+ """Iterate on the toplevel code blocks."""
+ for filename, code in self._toplevel:
+ measures = self[code]
+ if not measures:
+ continue # skip if no measurement
+ line_iterator = ((line, measures[line]) for line in measures.keys())
+ yield (filename, line_iterator)
+
+ class UDFLineProfiler(LineProfiler):
+ def __init__(self, **kw: Any) -> None:
+ super().__init__(**kw)
+ include_children = kw.get("include_children", False)
+ backend = kw.get("backend", "psutil")
+ self.code_map = CodeMapForUDF(include_children=include_children, backend=backend)
+
+ def __call__(
+ self,
+ func: Optional[Callable[..., Any]] = None,
+ precision: int = 1,
+ *,
+ sub_lines: Optional[List] = None,
+ start_line: Optional[int] = None,
+ ) -> Callable[..., Any]:
+ if func is not None:
+ self.add_function(func, sub_lines=sub_lines, start_line=start_line)
+ f = self.wrap_function(func)
+ f.__module__ = func.__module__
+ f.__name__ = func.__name__
+ f.__doc__ = func.__doc__
+ f.__dict__.update(getattr(func, "__dict__", {}))
+ return f
+ else:
+
+ def inner_partial(f: Callable[..., Any]) -> Any:
+ return self.__call__(f, precision=precision)
+
+ return inner_partial
+
+ def add_function(
+ self,
+ func: Callable[..., Any],
+ *,
+ sub_lines: Optional[List] = None,
+ start_line: Optional[int] = None,
+ ) -> None:
+ """Record line profiling information for the given Python function."""
+ try:
+ # func_code does not exist in Python3
+ code = func.__code__
+ except AttributeError:
+ warnings.warn("Could not extract a code object for the object %r" % func)
+ else:
+ self.code_map.add(code, sub_lines=sub_lines, start_line=start_line)
+
+ class UDFLineProfilerV2(LineProfiler):
+ def __init__(self, **kw: Any) -> None:
+ super().__init__(**kw)
+ include_children = kw.get("include_children", False)
+ backend = kw.get("backend", "psutil")
+ self.code_map = CodeMapForUDFV2(include_children=include_children, backend=backend)
+
+ def add_code(self, code: CodeType) -> None:
+ """Record line profiling information for the given code object."""
+ self.code_map.add(code)
+
+ for name in __all__:
+ if name != "has_memory_profiler":
+ globals()[name] = locals()[name]
+
+ _module_initialized = True
diff --git a/python/pyspark/profiler.py b/python/pyspark/profiler.py
index 28669d5343c3..8f797b5f21fc 100644
--- a/python/pyspark/profiler.py
+++ b/python/pyspark/profiler.py
@@ -15,12 +15,10 @@
# limitations under the License.
#
-from types import CodeType
from typing import (
Any,
Callable,
Dict,
- Iterator,
List,
Optional,
Tuple,
@@ -30,20 +28,11 @@ from typing import (
cast,
)
import cProfile
-import inspect
import pstats
import linecache
import os
import atexit
import sys
-import warnings
-
-try:
- from memory_profiler import CodeMap, LineProfiler
-
- has_memory_profiler = True
-except Exception:
- has_memory_profiler = False
import pyspark
from pyspark.accumulators import AccumulatorParam
@@ -172,119 +161,6 @@ class Profiler:
raise NotImplementedError
-if has_memory_profiler:
-
- class CodeMapForUDF(CodeMap):
- def add(
- self,
- code: Any,
- toplevel_code: Optional[Any] = None,
- *,
- sub_lines: Optional[List] = None,
- start_line: Optional[int] = None,
- ) -> None:
- if code in self:
- return
-
- if toplevel_code is None:
- toplevel_code = code
- filename = code.co_filename
- if sub_lines is None or start_line is None:
- sub_lines, start_line = inspect.getsourcelines(code)
- linenos = range(start_line, start_line + len(sub_lines))
- self._toplevel.append((filename, code, linenos))
- self[code] = {}
- else:
- self[code] = self[toplevel_code]
- for subcode in filter(inspect.iscode, code.co_consts):
- self.add(subcode, toplevel_code=toplevel_code)
-
- class CodeMapForUDFV2(CodeMap):
- def add(
- self,
- code: Any,
- toplevel_code: Optional[Any] = None,
- ) -> None:
- if code in self:
- return
-
- if toplevel_code is None:
- toplevel_code = code
- filename = code.co_filename
- self._toplevel.append((filename, code))
- self[code] = {}
- else:
- self[code] = self[toplevel_code]
- for subcode in filter(inspect.iscode, code.co_consts):
- self.add(subcode, toplevel_code=toplevel_code)
-
- def items(self) -> Iterator[Tuple[str, Iterator[Tuple[int, Any]]]]:
- """Iterate on the toplevel code blocks."""
- for filename, code in self._toplevel:
- measures = self[code]
- if not measures:
- continue # skip if no measurement
- line_iterator = ((line, measures[line]) for line in measures.keys())
- yield (filename, line_iterator)
-
- class UDFLineProfiler(LineProfiler):
- def __init__(self, **kw: Any) -> None:
- super().__init__(**kw)
- include_children = kw.get("include_children", False)
- backend = kw.get("backend", "psutil")
- self.code_map = CodeMapForUDF(include_children=include_children, backend=backend)
-
- def __call__(
- self,
- func: Optional[Callable[..., Any]] = None,
- precision: int = 1,
- *,
- sub_lines: Optional[List] = None,
- start_line: Optional[int] = None,
- ) -> Callable[..., Any]:
- if func is not None:
- self.add_function(func, sub_lines=sub_lines, start_line=start_line)
- f = self.wrap_function(func)
- f.__module__ = func.__module__
- f.__name__ = func.__name__
- f.__doc__ = func.__doc__
- f.__dict__.update(getattr(func, "__dict__", {}))
- return f
- else:
-
- def inner_partial(f: Callable[..., Any]) -> Any:
- return self.__call__(f, precision=precision)
-
- return inner_partial
-
- def add_function(
- self,
- func: Callable[..., Any],
- *,
- sub_lines: Optional[List] = None,
- start_line: Optional[int] = None,
- ) -> None:
- """Record line profiling information for the given Python function."""
- try:
- # func_code does not exist in Python3
- code = func.__code__
- except AttributeError:
- warnings.warn("Could not extract a code object for the object %r" % func)
- else:
- self.code_map.add(code, sub_lines=sub_lines, start_line=start_line)
-
- class UDFLineProfilerV2(LineProfiler):
- def __init__(self, **kw: Any) -> None:
- super().__init__(**kw)
- include_children = kw.get("include_children", False)
- backend = kw.get("backend", "psutil")
- self.code_map = CodeMapForUDFV2(include_children=include_children, backend=backend)
-
- def add_code(self, code: CodeType) -> None:
- """Record line profiling information for the given code object."""
- self.code_map.add(code)
-
-
class PStatsParam(AccumulatorParam[Optional[pstats.Stats]]):
"""PStatsParam is used to merge pstats.Stats"""
@@ -439,8 +315,8 @@ class MemoryProfiler(Profiler):
**kwargs: Any,
) -> Any:
"""Runs and profiles the method func passed in. A profile object is returned."""
- if has_memory_profiler:
- profiler = UDFLineProfiler()
+ if pyspark.memory_profiler_ext.has_memory_profiler:
+ profiler = pyspark.memory_profiler_ext.UDFLineProfiler()
wrapped = profiler(func, sub_lines=sub_lines, start_line=start_line)
ret = wrapped(*args, **kwargs)
codemap_dict = {
diff --git a/python/pyspark/sql/profiler.py b/python/pyspark/sql/profiler.py
index 9b4786c8c01c..88f48070f205 100644
--- a/python/pyspark/sql/profiler.py
+++ b/python/pyspark/sql/profiler.py
@@ -35,6 +35,7 @@ from typing import (
)
import warnings
+import pyspark.memory_profiler_ext
from pyspark.accumulators import (
Accumulator,
AccumulatorParam,
@@ -47,7 +48,6 @@ from pyspark.profiler import (
MemoryProfiler,
MemUsageParam,
PStatsParam,
- has_memory_profiler,
)
if TYPE_CHECKING:
@@ -138,7 +138,7 @@ class WorkerMemoryProfiler:
result_key: Union[int, str],
func_or_code: Union[Callable, CodeType],
) -> None:
- from pyspark.profiler import UDFLineProfilerV2
+ from pyspark.memory_profiler_ext import UDFLineProfilerV2
self._accumulator = accumulator
self._profiler = UDFLineProfilerV2()
@@ -245,7 +245,7 @@ class ProfilerCollector(ABC):
with self._lock:
code_map = self._memory_profile_results
- if not has_memory_profiler and not code_map:
+ if not pyspark.memory_profiler_ext.has_memory_profiler and not code_map:
warnings.warn(
"Install the 'memory_profiler' library in the cluster to enable memory profiling",
UserWarning,
@@ -331,7 +331,7 @@ class ProfilerCollector(ABC):
with self._lock:
code_map = self._memory_profile_results
- if not has_memory_profiler and not code_map:
+ if not pyspark.memory_profiler_ext.has_memory_profiler and not code_map:
warnings.warn(
"Install the 'memory_profiler' library in the cluster to enable memory profiling",
UserWarning,
diff --git a/python/pyspark/sql/tests/test_python_datasource.py b/python/pyspark/sql/tests/test_python_datasource.py
index d34611ac36e2..292d64725008 100644
--- a/python/pyspark/sql/tests/test_python_datasource.py
+++ b/python/pyspark/sql/tests/test_python_datasource.py
@@ -28,7 +28,7 @@ from decimal import Decimal
from typing import Callable, Iterable, List, Union, Iterator, Tuple
from pyspark.errors import AnalysisException, PythonException
-from pyspark.profiler import has_memory_profiler
+from pyspark.memory_profiler_ext import has_memory_profiler
from pyspark.sql.datasource import (
CaseInsensitiveDict,
DataSource,
diff --git a/python/pyspark/tests/test_memory_profiler.py b/python/pyspark/tests/test_memory_profiler.py
index e260e5997990..8664688c3b8e 100644
--- a/python/pyspark/tests/test_memory_profiler.py
+++ b/python/pyspark/tests/test_memory_profiler.py
@@ -27,7 +27,7 @@ from typing import Iterator
from unittest import mock
from pyspark import SparkConf
-from pyspark.profiler import has_memory_profiler
+from pyspark.memory_profiler_ext import has_memory_profiler
from pyspark.sql import SparkSession
from pyspark.sql.functions import col, pandas_udf, udf
from pyspark.sql.window import Window
@@ -70,7 +70,7 @@ class MemoryProfilerTests(PySparkTestCase):
self.spark = SparkSession(sparkContext=self.sc)
def test_code_map(self):
- from pyspark.profiler import CodeMapForUDF
+ from pyspark.memory_profiler_ext import CodeMapForUDF
code_map = CodeMapForUDF(include_children=False, backend="psutil")
@@ -85,7 +85,7 @@ class MemoryProfilerTests(PySparkTestCase):
self.assertEqual(len(code_map._toplevel), 1)
def test_udf_line_profiler(self):
- from pyspark.profiler import UDFLineProfiler
+ from pyspark.memory_profiler_ext import UDFLineProfiler
profiler = UDFLineProfiler()
diff --git a/python/pyspark/tests/test_profiler.py b/python/pyspark/tests/test_profiler.py
index baa71ebefd74..56fd0f473308 100644
--- a/python/pyspark/tests/test_profiler.py
+++ b/python/pyspark/tests/test_profiler.py
@@ -22,7 +22,7 @@ import unittest
from io import StringIO
from pyspark import SparkConf, SparkContext, BasicProfiler
-from pyspark.profiler import has_memory_profiler
+from pyspark.memory_profiler_ext import has_memory_profiler
from pyspark.sql import SparkSession
from pyspark.sql.functions import udf
from pyspark.errors import PythonException, PySparkRuntimeError
diff --git a/python/pyspark/worker.py b/python/pyspark/worker.py
index 6c8afb160b99..f01555d90e32 100644
--- a/python/pyspark/worker.py
+++ b/python/pyspark/worker.py
@@ -108,13 +108,6 @@ from pyspark.worker_util import (
)
from pyspark.logger.worker_io import capture_outputs
-try:
- import memory_profiler # noqa: F401
-
- has_memory_profiler = True
-except Exception:
- has_memory_profiler = False
-
class RunnerConf(Conf):
@property
@@ -1317,7 +1310,9 @@ def wrap_perf_profiler(f, eval_type, result_id):
def wrap_memory_profiler(f, eval_type, result_id):
from pyspark.sql.profiler import ProfileResultsParam, WorkerMemoryProfiler
- if not has_memory_profiler:
+ import pyspark.memory_profiler_ext
+
+ if not pyspark.memory_profiler_ext.has_memory_profiler:
return f
accumulator = _deserialize_accumulator(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org