chore: unify Python function validation between TrainerClient and SparkClient

kubeflow/common/utils.py

@@ -11,7 +11,11 @@
 # 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.
+import ast
+from collections.abc import Callable
+import inspect
 import os
+import textwrap
 
 from kubernetes import config
 
@@ -61,3 +65,47 @@ def validate_wait_for_job_status(polling_interval: int, timeout: int) -> None:
             "Polling interval must be strictly less than timeout. "
             f"Received polling_interval={polling_interval}, timeout={timeout}"
         )
+
+
+def validate_python_function(func: Callable) -> None:
+    """Validate a Python function.
+
+    Args:
+        func: The Python function to validate.
+
+    Raises:
+        ValueError: If the function is not a Python function, is asynchronous,
+            is a lambda function, is decorated, or its source cannot be
+            inspected or parsed.
+    """
+    if not inspect.isfunction(func):
+        raise ValueError("Function must be a Python function.")
+
+    if inspect.iscoroutinefunction(func):
+        raise ValueError("Async functions are not supported.")
+
+    if func.__name__ == "<lambda>":
+        raise ValueError("Lambda functions are not supported.")
+
+    try:
+        func_source = textwrap.dedent(inspect.getsource(func))
+    except TypeError as e:
+        raise ValueError(
+            "Function must be a pure-Python function; built-in or "
+            "C-implemented callables are not supported."
+        ) from e
+    except OSError as e:
+        raise ValueError(
+            "Function source could not be read. Functions defined interactively are not supported."
+        ) from e
+
+    try:
+        func_tree = ast.parse(func_source)
+    except SyntaxError as e:
+        raise ValueError("Function source could not be parsed.") from e
+
+    if any(
+        isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.decorator_list
+        for node in func_tree.body
+    ):
+        raise ValueError("Decorated functions are not supported.")

kubeflow/common/utils_test.py

@@ -16,6 +16,36 @@
 from kubeflow.common import utils
 from kubeflow.trainer.test.common import SUCCESS, TestCase
 
+# --------------------------
+# Test Helpers
+# --------------------------
+
+
+def sample_function() -> None:
+    """Sample function for testing."""
+    print("Hello World")
+
+
+async def sample_async_function() -> None:
+    """Sample async function for testing."""
+    pass
+
+
+def sample_decorator(func):
+    """Pass-through decorator for testing."""
+    return func
+
+
+@sample_decorator
+def sample_decorated_function() -> None:
+    """Sample decorated function."""
+    pass
+
+
+# --------------------------
+# Tests
+# --------------------------
+
 
 @pytest.mark.parametrize(
     "test_case",
@@ -67,3 +97,54 @@ def test_validate_wait_for_job_status(test_case):
             utils.validate_wait_for_job_status(polling_interval, timeout)
     else:
         utils.validate_wait_for_job_status(polling_interval, timeout)
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="valid function",
+            expected_status=SUCCESS,
+            config={"func": sample_function},
+        ),
+        TestCase(
+            name="not a function",
+            expected_status=SUCCESS,
+            config={"func": "not a function"},
+            expected_error=ValueError,
+            expected_output="Function must be a Python function.",
+        ),
+        TestCase(
+            name="lambda function",
+            expected_status=SUCCESS,
+            config={"func": lambda: None},
+            expected_error=ValueError,
+            expected_output="Lambda functions are not supported.",
+        ),
+        TestCase(
+            name="async function",
+            expected_status=SUCCESS,
+            config={"func": sample_async_function},
+            expected_error=ValueError,
+            expected_output="Async functions are not supported.",
+        ),
+        TestCase(
+            name="decorated function",
+            expected_status=SUCCESS,
+            config={"func": sample_decorated_function},
+            expected_error=ValueError,
+            expected_output="Decorated functions are not supported.",
+        ),
+    ],
+)
+def test_validate_python_function(test_case):
+    """Test validate_python_function with valid and invalid functions."""
+    print("Executing test:", test_case.name)
+
+    func = test_case.config["func"]
+
+    if test_case.expected_error:
+        with pytest.raises(test_case.expected_error, match=test_case.expected_output):
+            utils.validate_python_function(func)
+    else:
+        utils.validate_python_function(func)

kubeflow/spark/backends/kubernetes/backend.py

@@ -14,7 +14,6 @@
 
 """Kubernetes backend for Spark operations."""
 
-import ast
 from collections.abc import Iterator
 import contextlib
 import inspect
@@ -37,6 +36,7 @@
 
 from kubeflow.common import constants as common_constants
 from kubeflow.common.types import KubernetesBackendConfig
+from kubeflow.common.utils import validate_python_function
 from kubeflow.spark.backends.base import RuntimeBackend
 from kubeflow.spark.backends.kubernetes import constants
 from kubeflow.spark.backends.kubernetes.utils import (
@@ -853,29 +853,13 @@ def _validate_func_job(
                 If the function or function arguments are invalid.
         """
 
-        if not inspect.isfunction(job.func):
-            raise ValueError("`job.func` must be a Python function.")
+        # Validate generic Python function properties.
+        validate_python_function(job.func)
 
-        if inspect.iscoroutinefunction(job.func):
-            raise ValueError("Async functions are not supported.")
-
-        if job.func.__name__ == "<lambda>":
-            raise ValueError("Lambda functions are not supported.")
-
-        try:
-            func_source = textwrap.dedent(inspect.getsource(job.func))
-        except TypeError as e:
-            raise ValueError(
-                "`job.func` must be a pure-Python function; built-in or "
-                "C-implemented callables are not supported."
-            ) from e
-        except OSError as e:
-            raise ValueError(
-                "`job.func` source could not be read. Functions defined "
-                "interactively (REPL/Jupyter) or generated dynamically are "
-                "not supported; define it in a Python module."
-            ) from e
+        # Get the function source for Spark-specific validation.
+        func_source = textwrap.dedent(inspect.getsource(job.func))
 
+        # Ensure the function source does not contain the reserved heredoc delimiter.
         if any(
             line.strip() == constants.FUNC_JOB_SCRIPT_DELIMITER for line in func_source.splitlines()
         ):
@@ -884,17 +868,6 @@ def _validate_func_job(
                 f"{constants.FUNC_JOB_SCRIPT_DELIMITER!r}, which is not supported."
             )
 
-        try:
-            func_tree = ast.parse(func_source)
-        except SyntaxError as e:
-            raise ValueError("`job.func` source could not be parsed.") from e
-
-        if any(
-            isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.decorator_list
-            for node in func_tree.body
-        ):
-            raise ValueError("Decorated functions are not supported.")
-
         if job.func_args is not None:
             if not isinstance(job.func_args, dict):
                 raise ValueError("`job.func_args` must be a dictionary.")

kubeflow/spark/backends/kubernetes/backend_test.py

@@ -15,7 +15,6 @@
 """Unit tests for KubernetesBackend."""
 
 from datetime import datetime
-from functools import wraps
 import multiprocessing
 from unittest.mock import Mock, patch
 
@@ -372,11 +371,8 @@ def sample_func_with_args(
 
 
 def _decorator(func: callable) -> callable:
-    @wraps(func)
-    def wrapper(*args, **kwargs):
-        return func(*args, **kwargs)
-
-    return wrapper
+    """Pass-through decorator for testing."""
+    return func
 
 
 @_decorator

kubeflow/spark/backends/kubernetes/utils.py

@@ -30,6 +30,7 @@
 from kubernetes import client
 
 from kubeflow.common import constants as common_constants
+from kubeflow.common.utils import validate_python_function
 from kubeflow.spark.backends.kubernetes import constants
 from kubeflow.spark.types.types import (
     Driver,
@@ -673,10 +674,9 @@ def get_command_using_spark_func(
 
     Raises:
         ValueError:
-            If ``func`` is not callable.
+            If ``func`` is not a supported Python function.
     """
-    if not callable(func):
-        raise ValueError(f"Expected a callable function, got {type(func)}.")
+    validate_python_function(func)
 
     func_code = textwrap.dedent(
         inspect.getsource(func),

kubeflow/spark/backends/kubernetes/utils_test.py

@@ -1122,7 +1122,7 @@ def test_get_func_job_init_container(test_case: TestCase) -> None:
                 "func_args": None,
             },
             expected_error=ValueError,
-            expected_output="Expected a callable function",
+            expected_output="Function must be a Python function.",
         ),
     ],
 )

kubeflow/trainer/backends/container/utils.py

@@ -23,6 +23,7 @@
 import shlex
 
 from kubeflow.common.constants import UNKNOWN
+from kubeflow.common.utils import validate_python_function
 from kubeflow.trainer.constants import constants
 from kubeflow.trainer.types import types
 
@@ -64,6 +65,7 @@ def get_training_script_code(trainer: types.CustomTrainer) -> str:
     import inspect
     import textwrap
 
+    validate_python_function(trainer.func)
     code = inspect.getsource(trainer.func)
     code = textwrap.dedent(code)
     if trainer.func_args is None:

kubeflow/trainer/backends/container/utils_test.py

@@ -164,3 +164,45 @@ def test_aggregate_status_from_containers(test_case: TestCase):
     result = container_utils.aggregate_status_from_containers(test_case.config["statuses"])
     assert result == test_case.expected_output
     print("test execution complete")
+
+
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        TestCase(
+            name="valid function",
+            expected_status=SUCCESS,
+            config={"func": simple_train_func},
+        ),
+        TestCase(
+            name="non python function",
+            expected_status=FAILED,
+            config={"func": "not_a_function"},
+            expected_error=ValueError,
+            expected_output="Function must be a Python function.",
+        ),
+        TestCase(
+            name="lambda function",
+            expected_status=FAILED,
+            config={"func": lambda: None},
+            expected_error=ValueError,
+            expected_output="Lambda functions are not supported.",
+        ),
+    ],
+)
+def test_get_training_script_code(test_case: TestCase):
+    """Test training script generation."""
+
+    trainer = types.CustomTrainer(func=test_case.config["func"])
+
+    if test_case.expected_status == SUCCESS:
+        code = container_utils.get_training_script_code(trainer)
+
+        assert "def simple_train_func" in code
+        assert "simple_train_func()" in code
+    else:
+        with pytest.raises(
+            test_case.expected_error,
+            match=test_case.expected_output,
+        ):
+            container_utils.get_training_script_code(trainer)

kubeflow/trainer/backends/kubernetes/backend_test.py

@@ -83,6 +83,11 @@
 # --------------------------
 
 
+def sample_train_func() -> None:
+    """Sample training function."""
+    print("Hello World")
+
+
 @pytest.fixture
 def kubernetes_backend(request):
     """Provide a KubernetesBackend with mocked Kubernetes APIs."""
@@ -289,8 +294,10 @@ def get_custom_trainer(
     # with torchrun as the entrypoint and a fixed lambda for deterministic tests.
     func_script = (
         "\nread -r -d '' SCRIPT << EOM\n\n"
-        'func=lambda: print("Hello World"),\n\n'
-        "<lambda>(**{'learning_rate': 0.001, 'batch_size': 32})\n\n"
+        "def sample_train_func() -> None:\n"
+        '    """Sample training function."""\n'
+        '    print("Hello World")\n\n'
+        "sample_train_func(**{'learning_rate': 0.001, 'batch_size': 32})\n\n"
         'EOM\nprintf "%s" "$SCRIPT" > "backend_test.py"\n'
         'torchrun "backend_test.py"'
     )
@@ -1235,7 +1242,7 @@ def test_get_runtime_packages(kubernetes_backend, test_case):
             expected_status=SUCCESS,
             config={
                 "trainer": types.CustomTrainer(
-                    func=lambda: print("Hello World"),
+                    func=sample_train_func,
                     func_args={"learning_rate": 0.001, "batch_size": 32},
                     packages_to_install=["torch", "numpy"],
                     pip_index_urls=constants.DEFAULT_PIP_INDEX_URLS,
@@ -1256,7 +1263,7 @@ def test_get_runtime_packages(kubernetes_backend, test_case):
             expected_status=SUCCESS,
             config={
                 "trainer": types.CustomTrainer(
-                    func=lambda: print("Hello World"),
+                    func=sample_train_func,
                     func_args={"learning_rate": 0.001, "batch_size": 32},
                     packages_to_install=["torch", "numpy"],
                     pip_index_urls=constants.DEFAULT_PIP_INDEX_URLS,
@@ -1340,7 +1347,7 @@ def test_get_runtime_packages(kubernetes_backend, test_case):
             expected_status=FAILED,
             config={
                 "trainer": types.CustomTrainer(
-                    func=lambda: print("Hello World"),
+                    func=sample_train_func,
                     num_nodes=2,
                 ),
                 "runtime": TORCH_TUNE_RUNTIME,