diff --git a/kubeflow/trainer/backends/kubernetes/backend_test.py b/kubeflow/trainer/backends/kubernetes/backend_test.py index 8853ae259..e2c31f561 100644 --- a/kubeflow/trainer/backends/kubernetes/backend_test.py +++ b/kubeflow/trainer/backends/kubernetes/backend_test.py @@ -238,17 +238,26 @@ def get_custom_trainer( install_script = utils.get_script_for_python_packages( packages_to_install=packages_to_install, pip_index_urls=pip_index_urls, - install_log_file="pip_install.log", + install_log_file="/tmp/pip_install.log", ) # Append the embedded training function script that matches EXEC_FUNC_SCRIPT # with torchrun as the entrypoint and a fixed lambda for deterministic tests. func_script = ( "\nread -r -d '' SCRIPT << EOM\n\n" + "import os\n" + "import sys\n\n" + "_kubeflow_generated_script_directory = os.path.realpath(os.path.dirname(__file__))\n" + "_kubeflow_working_directory = os.getcwd()\n" + "if sys.path and os.path.realpath(sys.path[0]) == " + "_kubeflow_generated_script_directory:\n" + " sys.path[0] = _kubeflow_working_directory\n" + "else:\n" + " sys.path.insert(0, _kubeflow_working_directory)\n\n" 'func=lambda: print("Hello World"),\n\n' "(**{'learning_rate': 0.001, 'batch_size': 32})\n\n" - 'EOM\nprintf "%s" "$SCRIPT" > "backend_test.py"\n' - 'torchrun "backend_test.py"' + 'EOM\nprintf "%s" "$SCRIPT" > "/tmp/backend_test.py"\n' + 'torchrun "/tmp/backend_test.py"' ) full_command = install_script + func_script diff --git a/kubeflow/trainer/backends/kubernetes/utils.py b/kubeflow/trainer/backends/kubernetes/utils.py index e7b1f6e6a..9ed6e1f9b 100644 --- a/kubeflow/trainer/backends/kubernetes/utils.py +++ b/kubeflow/trainer/backends/kubernetes/utils.py @@ -381,7 +381,22 @@ def get_command_using_train_func( func_file = os.path.join(constants.DEFAULT_MPI_USER_HOME, func_file) install_log_file = os.path.join(constants.DEFAULT_MPI_USER_HOME, "pip_install.log") else: - install_log_file = "pip_install.log" + func_file = os.path.join("/tmp", func_file) + install_log_file = os.path.join("/tmp", "pip_install.log") + # Preserve the working directory as Python's first import root instead of /tmp. + python_import_path_prelude = ( + "import os\n" + "import sys\n\n" + "_kubeflow_generated_script_directory = " + "os.path.realpath(os.path.dirname(__file__))\n" + "_kubeflow_working_directory = os.getcwd()\n" + "if sys.path and os.path.realpath(sys.path[0]) == " + "_kubeflow_generated_script_directory:\n" + " sys.path[0] = _kubeflow_working_directory\n" + "else:\n" + " sys.path.insert(0, _kubeflow_working_directory)\n\n" + ) + func_code = python_import_path_prelude + func_code # Install Python packages if that is required. install_packages = "" diff --git a/kubeflow/trainer/backends/kubernetes/utils_test.py b/kubeflow/trainer/backends/kubernetes/utils_test.py index a634447a0..ef46ed5bf 100644 --- a/kubeflow/trainer/backends/kubernetes/utils_test.py +++ b/kubeflow/trainer/backends/kubernetes/utils_test.py @@ -12,10 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +import importlib.util import os +from pathlib import Path +import shlex +import subprocess +import sys import tempfile import time from unittest.mock import MagicMock, patch +import uuid from kubeflow_trainer_api import models import pytest @@ -25,8 +31,26 @@ from kubeflow.trainer.test.common import FAILED, SUCCESS, TestCase from kubeflow.trainer.types import types +MPI_FUNCTION_FILE = os.path.join(constants.DEFAULT_MPI_USER_HOME, "utils_test.py") +MPI_INSTALLATION_LOG_FILE = os.path.join( + constants.DEFAULT_MPI_USER_HOME, + "pip_install.log", +) +NON_MPI_IMPORT_PATH_PRELUDE = ( + "import os\n" + "import sys\n\n" + "_kubeflow_generated_script_directory = os.path.realpath(os.path.dirname(__file__))\n" + "_kubeflow_working_directory = os.getcwd()\n" + "if sys.path and os.path.realpath(sys.path[0]) == _kubeflow_generated_script_directory:\n" + " sys.path[0] = _kubeflow_working_directory\n" + "else:\n" + " sys.path.insert(0, _kubeflow_working_directory)\n\n" +) -def _build_runtime() -> types.Runtime: + +def _build_runtime( + command: tuple[str, ...] = constants.DEFAULT_COMMAND, +) -> types.Runtime: runtime_trainer = types.RuntimeTrainer( trainer_type=types.TrainerType.CUSTOM_TRAINER, framework="torch", @@ -34,7 +58,7 @@ def _build_runtime() -> types.Runtime: device_count="1", image="example.com/image", ) - runtime_trainer.set_command(constants.DEFAULT_COMMAND) + runtime_trainer.set_command(command) return types.Runtime( name="test-runtime", trainer=runtime_trainer, @@ -395,11 +419,12 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" + f"{NON_MPI_IMPORT_PATH_PRELUDE}" '"func": (lambda: print("Hello World")),\n\n' "(**{'batch_size': 128, 'learning_rate': 0.001, 'epochs': 20})\n\n" "EOM\n" - 'printf "%s" "$SCRIPT" > "utils_test.py"\n' - 'python "utils_test.py"' + 'printf "%s" "$SCRIPT" > "/tmp/utils_test.py"\n' + 'python "/tmp/utils_test.py"' ), ], ), @@ -416,11 +441,12 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" + f"{NON_MPI_IMPORT_PATH_PRELUDE}" '"func": (lambda: print("Hello World")),\n\n' "()\n\n" "EOM\n" - 'printf "%s" "$SCRIPT" > "utils_test.py"\n' - 'python "utils_test.py"' + 'printf "%s" "$SCRIPT" > "/tmp/utils_test.py"\n' + 'python "/tmp/utils_test.py"' ), ], ), @@ -461,11 +487,12 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" + f"{NON_MPI_IMPORT_PATH_PRELUDE}" '"func": (lambda: print("Hello World")),\n\n' "(**{'a': 1, 'b': 2})\n\n" "EOM\n" - 'printf "%s" "$SCRIPT" > "utils_test.py"\n' - 'python "utils_test.py"' + 'printf "%s" "$SCRIPT" > "/tmp/utils_test.py"\n' + 'python "/tmp/utils_test.py"' ), ], ), @@ -482,11 +509,12 @@ def test_get_script_for_python_packages(test_case): "-c", ( "\nread -r -d '' SCRIPT << EOM\n\n" + f"{NON_MPI_IMPORT_PATH_PRELUDE}" '"func": (lambda **kwargs: "ok"),\n\n' "(**{'a': 3, 'b': 'hi', 'c': 0.2})\n\n" "EOM\n" - 'printf "%s" "$SCRIPT" > "utils_test.py"\n' - 'python "utils_test.py"' + 'printf "%s" "$SCRIPT" > "/tmp/utils_test.py"\n' + 'python "/tmp/utils_test.py"' ), ], ), @@ -509,7 +537,7 @@ def test_get_script_for_python_packages(test_case): "fi\n\n\n" "PACKAGES=(requests)\n" "PIP_OPTS=(--index-url https://pypi.org/simple)\n" - 'LOG_FILE="pip_install.log"\n' + 'LOG_FILE="/tmp/pip_install.log"\n' 'rm -f "$LOG_FILE"\n' "\n" "if PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_BREAK_SYSTEM_PACKAGES=1 python -m pip install --quiet \\\n" @@ -524,11 +552,57 @@ def test_get_script_for_python_packages(test_case): " exit 1\n" "fi\n\n" "\nread -r -d '' SCRIPT << EOM\n\n" + f"{NON_MPI_IMPORT_PATH_PRELUDE}" '"func": (lambda: print("Hello World")),\n\n' "()\n\n" "EOM\n" - 'printf "%s" "$SCRIPT" > "utils_test.py"\n' - 'python "utils_test.py"' + 'printf "%s" "$SCRIPT" > "/tmp/utils_test.py"\n' + 'python "/tmp/utils_test.py"' + ), + ], + ), + TestCase( + name="with MPI preserves mpiuser paths", + expected_status=SUCCESS, + config={ + "func": (lambda: print("Hello World")), + "func_args": None, + "runtime": _build_runtime(constants.MPI_COMMAND), + "packages_to_install": ["requests"], + }, + expected_output=[ + "mpirun", + "--hostfile", + "/etc/mpi/hostfile", + "bash", + "-c", + ( + '\nif ! [ -x "$(command -v pip)" ]; then\n' + " python -m ensurepip || python -m ensurepip --user || " + "apt-get install python-pip\n" + "fi\n\n\n" + "PACKAGES=(requests)\n" + "PIP_OPTS=(--index-url https://pypi.org/simple)\n" + f'LOG_FILE="{MPI_INSTALLATION_LOG_FILE}"\n' + 'rm -f "$LOG_FILE"\n' + "\n" + "if PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_BREAK_SYSTEM_PACKAGES=1 python -m pip install --quiet \\\n" + ' --no-warn-script-location "${PIP_OPTS[@]}" --user "${PACKAGES[@]}" >"$LOG_FILE" 2>&1; then\n' + ' echo "Successfully installed Python packages (user): ${PACKAGES[*]}"\n' + "elif PIP_DISABLE_PIP_VERSION_CHECK=1 PIP_BREAK_SYSTEM_PACKAGES=1 python -m pip install --quiet \\\n" + ' --no-warn-script-location "${PIP_OPTS[@]}" "${PACKAGES[@]}" >>"$LOG_FILE" 2>&1; then\n' + ' echo "Successfully installed Python packages (system-wide): ${PACKAGES[*]}"\n' + "else\n" + ' echo "ERROR: Failed to install Python packages: ${PACKAGES[*]}" >&2\n' + ' cat "$LOG_FILE" >&2\n' + " exit 1\n" + "fi\n\n" + "\nread -r -d '' SCRIPT << EOM\n\n" + '"func": (lambda: print("Hello World")),\n\n' + "()\n\n" + "EOM\n" + f'printf "%s" "$SCRIPT" > "{MPI_FUNCTION_FILE}"\n' + f'python "{MPI_FUNCTION_FILE}"' ), ], ), @@ -551,6 +625,111 @@ def test_get_command_using_train_func(test_case: TestCase): assert type(e) is test_case.expected_error +@pytest.mark.parametrize( + "safe_path_enabled", + [False, True], + ids=["normal-path", "safe-path"], +) +def test_generated_command_runs_from_read_only_working_directory( + tmp_path, + safe_path_enabled, +): + unique_name_suffix = uuid.uuid4().hex + working_directory_module_name = f"working_directory_module_{unique_name_suffix}" + training_function_module_name = f"training_function_{unique_name_suffix}" + working_directory = tmp_path / "working-directory" + existing_python_path_directory = tmp_path / "existing-python-path" + training_function_source_file = tmp_path / f"{training_function_module_name}.py" + temporary_directory = Path("/tmp") + temporary_shadow_module_file = temporary_directory / f"{working_directory_module_name}.py" + generated_function_file = temporary_directory / training_function_source_file.name + working_directory.mkdir() + existing_python_path_directory.mkdir() + + try: + (working_directory / f"{working_directory_module_name}.py").write_text( + 'VALUE = "working-directory"\n', + encoding="utf-8", + ) + temporary_shadow_module_file.write_text( + 'VALUE = "temporary-directory"\n', + encoding="utf-8", + ) + (existing_python_path_directory / "existing_python_path_module.py").write_text( + 'VALUE = "existing-python-path"\n', + encoding="utf-8", + ) + training_function_source_file.write_text( + ( + "def training_function(expected_python_path):\n" + " import os\n" + " import existing_python_path_module\n" + f" import {working_directory_module_name}\n\n" + ' assert os.environ["PYTHONPATH"] == expected_python_path\n' + " print(\n" + f' f"{{{working_directory_module_name}.VALUE}}:' + '{existing_python_path_module.VALUE}"\n' + " )\n" + ), + encoding="utf-8", + ) + training_function_module_specification = importlib.util.spec_from_file_location( + training_function_module_name, + training_function_source_file, + ) + assert training_function_module_specification is not None + assert training_function_module_specification.loader is not None + training_function_module = importlib.util.module_from_spec( + training_function_module_specification + ) + training_function_module_specification.loader.exec_module(training_function_module) + + working_directory.chmod(0o555) + python_executable_command = ( + "bash", + "-c", + constants.EXEC_FUNC_SCRIPT.replace( + "__ENTRYPOINT__", + shlex.quote(sys.executable), + ), + ) + command = utils.get_command_using_train_func( + runtime=_build_runtime(python_executable_command), + train_func=training_function_module.training_function, + train_func_parameters={ + "expected_python_path": str(existing_python_path_directory), + }, + pip_index_urls=constants.DEFAULT_PIP_INDEX_URLS, + packages_to_install=[], + ) + environment = os.environ.copy() + environment["PYTHONPATH"] = str(existing_python_path_directory) + if safe_path_enabled: + environment["PYTHONSAFEPATH"] = "1" + else: + environment.pop("PYTHONSAFEPATH", None) + + result = subprocess.run( + command, + cwd=working_directory, + env=environment, + capture_output=True, + text=True, + check=False, + ) + finally: + working_directory.chmod(0o755) + generated_function_file.unlink(missing_ok=True) + temporary_shadow_module_file.unlink(missing_ok=True) + for owned_python_cache_file in (temporary_directory / "__pycache__").glob( + f"{working_directory_module_name}.*.pyc" + ): + owned_python_cache_file.unlink() + + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "working-directory:existing-python-path" + + @pytest.mark.parametrize( "test_case", [