diff --git a/kubeflow/common/utils.py b/kubeflow/common/utils.py index 7faa5e961..ef2efd9c0 100644 --- a/kubeflow/common/utils.py +++ b/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__ == "": + 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.") diff --git a/kubeflow/common/utils_test.py b/kubeflow/common/utils_test.py index ab0281e06..408a16123 100644 --- a/kubeflow/common/utils_test.py +++ b/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) diff --git a/kubeflow/spark/backends/kubernetes/backend.py b/kubeflow/spark/backends/kubernetes/backend.py index bdfc2adf5..b1cbd4583 100644 --- a/kubeflow/spark/backends/kubernetes/backend.py +++ b/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__ == "": - 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.") diff --git a/kubeflow/spark/backends/kubernetes/backend_test.py b/kubeflow/spark/backends/kubernetes/backend_test.py index f954dc8c8..77bdf5ca1 100644 --- a/kubeflow/spark/backends/kubernetes/backend_test.py +++ b/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 diff --git a/kubeflow/spark/backends/kubernetes/utils.py b/kubeflow/spark/backends/kubernetes/utils.py index ca7fdc4b6..5cecdf43d 100644 --- a/kubeflow/spark/backends/kubernetes/utils.py +++ b/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), diff --git a/kubeflow/spark/backends/kubernetes/utils_test.py b/kubeflow/spark/backends/kubernetes/utils_test.py index 344db527c..a04fcf343 100644 --- a/kubeflow/spark/backends/kubernetes/utils_test.py +++ b/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.", ), ], ) diff --git a/kubeflow/trainer/backends/container/utils.py b/kubeflow/trainer/backends/container/utils.py index 5ee4bbc3c..80588a219 100644 --- a/kubeflow/trainer/backends/container/utils.py +++ b/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: diff --git a/kubeflow/trainer/backends/container/utils_test.py b/kubeflow/trainer/backends/container/utils_test.py index e929ded84..adaa1aa15 100644 --- a/kubeflow/trainer/backends/container/utils_test.py +++ b/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) diff --git a/kubeflow/trainer/backends/kubernetes/backend_test.py b/kubeflow/trainer/backends/kubernetes/backend_test.py index 01e2fdf22..d5967d2a3 100644 --- a/kubeflow/trainer/backends/kubernetes/backend_test.py +++ b/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' - "(**{'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, diff --git a/kubeflow/trainer/backends/kubernetes/utils.py b/kubeflow/trainer/backends/kubernetes/utils.py index 66fdc481e..a60dca2fb 100644 --- a/kubeflow/trainer/backends/kubernetes/utils.py +++ b/kubeflow/trainer/backends/kubernetes/utils.py @@ -30,6 +30,7 @@ from kubeflow_trainer_api import models import requests +from kubeflow.common.utils import validate_python_function from kubeflow.trainer.constants import constants from kubeflow.trainer.types import types @@ -352,10 +353,7 @@ def get_command_using_train_func( raise ValueError(f"Runtime must have a trainer: {runtime}") # Check if training function is callable. - if not callable(train_func): - raise ValueError( - f"Training function must be callable, got function type: {type(train_func)}" - ) + validate_python_function(train_func) # Extract the function implementation. func_code = inspect.getsource(train_func) diff --git a/kubeflow/trainer/backends/kubernetes/utils_test.py b/kubeflow/trainer/backends/kubernetes/utils_test.py index 13b5d35c9..f8adaaa26 100644 --- a/kubeflow/trainer/backends/kubernetes/utils_test.py +++ b/kubeflow/trainer/backends/kubernetes/utils_test.py @@ -25,6 +25,10 @@ from kubeflow.trainer.test.common import FAILED, SUCCESS, TestCase from kubeflow.trainer.types import types +# -------------------------- +# Test Helpers +# -------------------------- + def _build_runtime() -> types.Runtime: runtime_trainer = types.RuntimeTrainer( @@ -42,6 +46,21 @@ def _build_runtime() -> types.Runtime: ) +def sample_train_func() -> None: + """Sample training function.""" + print("Hello World") + + +def sample_train_func_kwargs(a: int, b: str, c: float) -> str: + """Sample training function with kwargs.""" + return "ok" + + +# -------------------------- +# Tests +# -------------------------- + + @pytest.mark.parametrize( "test_case", [ @@ -390,7 +409,7 @@ def test_get_script_for_python_packages(test_case): name="with args dict always unpacks kwargs", expected_status=SUCCESS, config={ - "func": (lambda: print("Hello World")), + "func": sample_train_func, "func_args": {"batch_size": 128, "learning_rate": 0.001, "epochs": 20}, "runtime": _build_runtime(), }, @@ -399,8 +418,10 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" - '"func": (lambda: print("Hello World")),\n\n' - "(**{'batch_size': 128, 'learning_rate': 0.001, 'epochs': 20})\n\n" + "def sample_train_func() -> None:\n" + ' """Sample training function."""\n' + ' print("Hello World")\n\n' + "sample_train_func(**{'batch_size': 128, 'learning_rate': 0.001, 'epochs': 20})\n\n" "EOM\n" 'printf "%s" "$SCRIPT" > "utils_test.py"\n' 'python "utils_test.py"' @@ -411,7 +432,7 @@ def test_get_script_for_python_packages(test_case): name="without args calls function with no params", expected_status=SUCCESS, config={ - "func": (lambda: print("Hello World")), + "func": sample_train_func, "func_args": None, "runtime": _build_runtime(), }, @@ -420,8 +441,10 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" - '"func": (lambda: print("Hello World")),\n\n' - "()\n\n" + "def sample_train_func() -> None:\n" + ' """Sample training function."""\n' + ' print("Hello World")\n\n' + "sample_train_func()\n\n" "EOM\n" 'printf "%s" "$SCRIPT" > "utils_test.py"\n' 'python "utils_test.py"' @@ -432,7 +455,7 @@ def test_get_script_for_python_packages(test_case): name="raises when runtime has no trainer", expected_status=FAILED, config={ - "func": (lambda: print("Hello World")), + "func": sample_train_func, "func_args": None, "runtime": types.Runtime( name="no-trainer", @@ -456,7 +479,7 @@ def test_get_script_for_python_packages(test_case): name="single dict param also unpacks kwargs", expected_status=SUCCESS, config={ - "func": (lambda: print("Hello World")), + "func": sample_train_func, "func_args": {"a": 1, "b": 2}, "runtime": _build_runtime(), }, @@ -465,8 +488,10 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" - '"func": (lambda: print("Hello World")),\n\n' - "(**{'a': 1, 'b': 2})\n\n" + "def sample_train_func() -> None:\n" + ' """Sample training function."""\n' + ' print("Hello World")\n\n' + "sample_train_func(**{'a': 1, 'b': 2})\n\n" "EOM\n" 'printf "%s" "$SCRIPT" > "utils_test.py"\n' 'python "utils_test.py"' @@ -477,7 +502,7 @@ def test_get_script_for_python_packages(test_case): name="multi-param function uses kwargs-unpacking", expected_status=SUCCESS, config={ - "func": (lambda **kwargs: "ok"), + "func": sample_train_func_kwargs, "func_args": {"a": 3, "b": "hi", "c": 0.2}, "runtime": _build_runtime(), }, @@ -486,8 +511,10 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" - '"func": (lambda **kwargs: "ok"),\n\n' - "(**{'a': 3, 'b': 'hi', 'c': 0.2})\n\n" + "def sample_train_func_kwargs(a: int, b: str, c: float) -> str:\n" + ' """Sample training function with kwargs."""\n' + ' return "ok"\n\n' + "sample_train_func_kwargs(**{'a': 3, 'b': 'hi', 'c': 0.2})\n\n" "EOM\n" 'printf "%s" "$SCRIPT" > "utils_test.py"\n' 'python "utils_test.py"' @@ -498,7 +525,7 @@ def test_get_script_for_python_packages(test_case): name="with packages to install", expected_status=SUCCESS, config={ - "func": (lambda: print("Hello World")), + "func": sample_train_func, "func_args": None, "runtime": _build_runtime(), "packages_to_install": ["requests"], @@ -528,8 +555,10 @@ def test_get_script_for_python_packages(test_case): " exit 1\n" "fi\n\n" "\nread -r -d '' SCRIPT << EOM\n\n" - '"func": (lambda: print("Hello World")),\n\n' - "()\n\n" + "def sample_train_func() -> None:\n" + ' """Sample training function."""\n' + ' print("Hello World")\n\n' + "sample_train_func()\n\n" "EOM\n" 'printf "%s" "$SCRIPT" > "utils_test.py"\n' 'python "utils_test.py"' diff --git a/kubeflow/trainer/backends/localprocess/utils.py b/kubeflow/trainer/backends/localprocess/utils.py index ef01fb501..365aadec6 100644 --- a/kubeflow/trainer/backends/localprocess/utils.py +++ b/kubeflow/trainer/backends/localprocess/utils.py @@ -8,6 +8,7 @@ import textwrap from typing import Any +from kubeflow.common.utils import validate_python_function from kubeflow.trainer.backends.localprocess import constants as local_exec_constants from kubeflow.trainer.backends.localprocess.types import LocalRuntimeTrainer from kubeflow.trainer.constants import constants @@ -192,11 +193,8 @@ def get_command_using_train_func( if not runtime.trainer: raise ValueError(f"Runtime must have a trainer: {runtime}") - # Check if training function is callable. - if not callable(train_func): - raise ValueError( - f"Training function must be callable, got function type: {type(train_func)}" - ) + # Validate the training function. + validate_python_function(train_func) # Extract the function implementation. func_code = inspect.getsource(train_func)