Repository navigation
fix(trainer): use writable generated script paths - #604
danish9039 wants to merge 5 commits into
Conversation
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
🎉 Welcome to the Kubeflow SDK! 🎉 Thanks for opening your first PR! We're happy to have you as part of our community 🚀 Here's what happens next:
Join the community:
Feel free to ask questions in the comments if you need any help or clarification! |
lntutor
left a comment
There was a problem hiding this comment.
Reviewed the generated-script path change and its behavioral coverage. The non-MPI path keeps the image working directory first for imports while writing generated files and install logs under /tmp; MPI retains its existing user-home behavior. The read-only, safe-path, shadowing, PYTHONPATH, concurrent, and MPI cases cover the relevant execution paths. Approved.
|
@lntutor: changing LGTM is restricted to collaborators DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes/test-infra repository. |
|
Two updates since the approval.
Scoping note: this pull request fixes the generated script and package-installation log paths. It does not make a training function's own relative writes work; the Fashion-MNIST example still downloads into Separately, |
@andreyvelich for /approve |
There was a problem hiding this comment.
🟡 Changes recommended
Safe-path mode can incorrectly replace a user-provided /tmp import path with the working directory.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Relocates non-MPI generated scripts and installation logs to /tmp while preserving import behavior.
Changes:
- Uses writable
/tmppaths for non-MPI execution. - Preserves working-directory import precedence.
- Adds read-only and safe-path tests.
File summaries
| File | Description |
|---|---|
utils.py |
Relocates generated artifacts and adjusts imports. |
utils_test.py |
Tests paths and import behavior. |
backend_test.py |
Updates expected trainer commands. |
Review details
- Files reviewed: 3/3 changed files
- Comments generated: 1
- Review effort level: Balanced
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
6b8cd2a to
e3e4ca2
Compare
e3e4ca2 to
21369ae
Compare
| 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", |
There was a problem hiding this comment.
What if /tmp directory doesn't exist in the image?
Any thoughts on this @briangallagher @robert-bell @kubeflow/kubeflow-sdk-team @tariq-hasan ?
There was a problem hiding this comment.
I am wondering if we can pass in something along the lines of "${TMPDIR:-/tmp}" and have the resolution occur within the container shell. I say this because Unix tools would be expected to recognize TMPDIR and there's an added fallback to /tmp otherwise. I suppose if neither works then we could just fail with an error.
There was a problem hiding this comment.
/tmp is really the default in limux so "${TEMPORARY_DIRECTORY:-/tmp} might already be a bit of overengineering but either way it is good enough for me to merge
There was a problem hiding this comment.
I'm fine assuming tmp exists.
An alternative would be for the sdk to add and mount an emptydir volume through runtimePatches api. I think that's overkill and probably more likely to confuse users who look at the manifest.
There was a problem hiding this comment.
I like @tariq-hasan's approach with ${TMPDIR:-/tmp}/pip_install.py.
WDYT @danish9039 @juliusvonkohout @robert-bell ?
There was a problem hiding this comment.
@andreyvelich @juliusvonkohout @robert-bell I agree with @tariq-hasan's suggestion.
The container shell now resolves ${TMPDIR:-/tmp} for the generated
training script and the pip install log, so an image that sets TMPDIR
is honoured and /tmp stays the default.
The behavioural test points TMPDIR at a per-test directory, which
proves the resolution and keeps the test out of the shared /tmp.
https: //github.com/kubeflow/pull/604#discussion_r4209700197
Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
The prelude restores the working directory as the first import root after the generated script moved to /tmp. Its `else` branch inserted that directory even when Python had deliberately left it out, defeating `-P` and `PYTHONSAFEPATH=1`, whose whole purpose is to keep a writable directory off the import path. Insert only when Python itself would have added a script directory. This matches how CPython fixed `pdb` in python/cpython#111762 and how `ipykernel` and `pyright` resolved the same problem. `sys.flags.safe_path` exists only on Python 3.11+, and this package supports 3.10, so it is read through `getattr` with a `False` default; 3.10 has no safe path mode, so the historical behaviour is preserved there. Verified on Python 3.12 and 3.10: - normal mode resolves the working directory, with existing PYTHONPATH entries behind it - `-P` and `PYTHONSAFEPATH=1` no longer import from the working directory - under safe path an existing PYTHONPATH entry still resolves, which is what Python does on its own - Python 3.10 raises no AttributeError and keeps the normal-mode behaviour Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
Python 3.10 has no safe path mode and ignores PYTHONSAFEPATH, so the generated prelude correctly restores the working directory there. The safe-path test asserted the 3.11+ behaviour on every version. Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
The container shell now resolves ${TMPDIR:-/tmp} for the generated
training script and the pip install log, so an image that sets TMPDIR
is honoured and /tmp stays the default.
The behavioural test points TMPDIR at a per-test directory, which
proves the resolution and keeps the test out of the shared /tmp.
https: //github.com/kubeflow/pull/604#discussion_r4209700197
Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
108f688 to
db8182c
Compare
|
@andreyvelich @juliusvonkohout @tariq-hasan @robert-bell generated script and the pip install log now go to The behavioural test now sets |
|
/lgtm |
andreyvelich
left a comment
There was a problem hiding this comment.
Thanks for the fix, @danish9039! This review was done with the help of AI tools.
| 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(constants.DEFAULT_TEMPORARY_DIRECTORY, func_file) | ||
| install_log_file = os.path.join(constants.DEFAULT_TEMPORARY_DIRECTORY, "pip_install.log") |
There was a problem hiding this comment.
These are container paths, so join them with /. On a Windows client, os.path.join produces ${TMPDIR:-/tmp}\file.py. Also, please update the comment on L383: it now sits above both branches but only describes MPI.
| 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(constants.DEFAULT_TEMPORARY_DIRECTORY, func_file) | |
| install_log_file = os.path.join(constants.DEFAULT_TEMPORARY_DIRECTORY, "pip_install.log") | |
| if is_mpi: | |
| func_file = f"{constants.DEFAULT_MPI_USER_HOME}/{func_file}" | |
| install_log_file = f"{constants.DEFAULT_MPI_USER_HOME}/pip_install.log" | |
| else: | |
| func_file = f"{constants.DEFAULT_TEMPORARY_DIRECTORY}/{func_file}" | |
| install_log_file = f"{constants.DEFAULT_TEMPORARY_DIRECTORY}/pip_install.log" |
There was a problem hiding this comment.
Good catch on Windows, switched to / and fixed the comment.
| # Preserve the working directory as Python's first import root instead of the | ||
| # temporary directory. | ||
| # | ||
| # Python puts the generated script's directory first on `sys.path`. Because the script now | ||
| # lives in the temporary directory, that entry is replaced with the working directory to | ||
| # keep the import precedence the training function had before the script was relocated. | ||
| # | ||
| # Under `-P` or `PYTHONSAFEPATH=1` Python deliberately omits that entry, because the point | ||
| # of safe path mode is to keep a writable directory off the import path. In that mode | ||
| # nothing is inserted, matching how CPython fixed `pdb` and how `ipykernel` and `pyright` | ||
| # handle the same problem. `sys.flags.safe_path` only exists on Python 3.11+, and this | ||
| # package supports 3.10, so it is read defensively. | ||
| 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" | ||
| 'elif not getattr(sys.flags, "safe_path", False):\n' | ||
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" | ||
| ) | ||
| func_code = python_import_path_prelude + func_code |
There was a problem hiding this comment.
Move the prelude into constants (see the suggestion there) and drop the elif. Under python or torchrun, sys.path[0] is the script directory unless safe path is on, so the if alone covers it. The elif only fires for -I on Python <=3.10, where adding the working directory works against isolation.
| # Preserve the working directory as Python's first import root instead of the | |
| # temporary directory. | |
| # | |
| # Python puts the generated script's directory first on `sys.path`. Because the script now | |
| # lives in the temporary directory, that entry is replaced with the working directory to | |
| # keep the import precedence the training function had before the script was relocated. | |
| # | |
| # Under `-P` or `PYTHONSAFEPATH=1` Python deliberately omits that entry, because the point | |
| # of safe path mode is to keep a writable directory off the import path. In that mode | |
| # nothing is inserted, matching how CPython fixed `pdb` and how `ipykernel` and `pyright` | |
| # handle the same problem. `sys.flags.safe_path` only exists on Python 3.11+, and this | |
| # package supports 3.10, so it is read defensively. | |
| 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" | |
| 'elif not getattr(sys.flags, "safe_path", False):\n' | |
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" | |
| ) | |
| func_code = python_import_path_prelude + func_code | |
| # Keep the working directory, not the temporary directory, as the first import root. | |
| func_code = constants.PYTHON_IMPORT_PATH_PRELUDE + func_code |
There was a problem hiding this comment.
Makes sense, dropped it. The swap covers every case now that the -I on 3.10 path is gone.
| # The directory for the generated training script and the pip install log on non-MPI runtimes. | ||
| # The container shell resolves it at runtime: `TMPDIR` when the image sets one, otherwise /tmp. | ||
| # Both are writable for a non-root user, unlike the image working directory. | ||
| DEFAULT_TEMPORARY_DIRECTORY = "${TMPDIR:-/tmp}" |
There was a problem hiding this comment.
Make it clear that this is a shell expression, and keep the prelude next to EXEC_FUNC_SCRIPT so tests can reuse it:
| # The directory for the generated training script and the pip install log on non-MPI runtimes. | |
| # The container shell resolves it at runtime: `TMPDIR` when the image sets one, otherwise /tmp. | |
| # Both are writable for a non-root user, unlike the image working directory. | |
| DEFAULT_TEMPORARY_DIRECTORY = "${TMPDIR:-/tmp}" | |
| # Shell expression for the writable directory that holds the generated training script and the | |
| # pip install log on non-MPI runtimes. The container shell expands it to `TMPDIR` if set, | |
| # otherwise /tmp, so it must only be used inside double-quoted shell strings. | |
| DEFAULT_TEMPORARY_DIRECTORY = "${TMPDIR:-/tmp}" | |
| # Prepended to the non-MPI generated script. Python puts the script's directory (the temporary | |
| # directory) first on sys.path, so swap it for the working directory. Under safe path mode | |
| # (`-P` / `PYTHONSAFEPATH`) that entry is absent and sys.path is left untouched. | |
| PYTHON_IMPORT_PATH_PRELUDE = ( | |
| "import os\n" | |
| "import sys\n\n" | |
| "if sys.path and os.path.realpath(sys.path[0]) == " | |
| "os.path.realpath(os.path.dirname(__file__)):\n" | |
| " sys.path[0] = os.getcwd()\n\n" | |
| ) |
| 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" | ||
| 'elif not getattr(sys.flags, "safe_path", False):\n' | ||
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" | ||
| ) |
There was a problem hiding this comment.
Reuse the constant instead of keeping a copy:
| 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" | |
| 'elif not getattr(sys.flags, "safe_path", False):\n' | |
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" | |
| ) | |
| NON_MPI_IMPORT_PATH_PRELUDE = constants.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" | ||
| 'elif not getattr(sys.flags, "safe_path", False):\n' | ||
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" |
There was a problem hiding this comment.
Reuse the constant instead of keeping a copy:
| "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" | |
| 'elif not getattr(sys.flags, "safe_path", False):\n' | |
| " sys.path.insert(0, _kubeflow_working_directory)\n\n" | |
| f"{constants.PYTHON_IMPORT_PATH_PRELUDE}" |
| tmp_path, | ||
| safe_path_enabled, | ||
| ): | ||
| unique_name_suffix = uuid.uuid4().hex |
There was a problem hiding this comment.
uuid4 isn't needed here. spec_from_file_location and exec_module never register the module in sys.modules, and tmp_path is already unique per run. Please use fixed module names and drop import uuid, since CLAUDE.md asks to avoid non-determinism. A nested def training_function() would also remove the need for the importlib scaffolding.
There was a problem hiding this comment.
Went with the nested def, that removed uuid and the importlib bits together.
| 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() | ||
|
|
||
| if safe_path_enabled and sys.version_info >= (3, 11): |
There was a problem hiding this comment.
pytest already cleans up tmp_path. Root also ignores chmod 0o555, so as root this test passes even if the script is written to the working directory. Assert the script's location directly:
| 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() | |
| if safe_path_enabled and sys.version_info >= (3, 11): | |
| finally: | |
| working_directory.chmod(0o755) | |
| # chmod does not stop root, so assert directly that the script landed in TMPDIR. | |
| assert generated_function_file.is_file() | |
| if safe_path_enabled and sys.version_info >= (3, 11): |
There was a problem hiding this comment.
Right, root ignores the chmod. Asserting is_file() on the generated path now.
| # Python 3.10 has no safe path mode and ignores PYTHONSAFEPATH, so the prelude's | ||
| # `getattr(sys.flags, "safe_path", False)` fallback must restore the working directory | ||
| # exactly as it does on the normal path. |
There was a problem hiding this comment.
On 3.10, sys.path[0] is the script directory, so the if swap runs and the getattr fallback is never reached:
| # Python 3.10 has no safe path mode and ignores PYTHONSAFEPATH, so the prelude's | |
| # `getattr(sys.flags, "safe_path", False)` fallback must restore the working directory | |
| # exactly as it does on the normal path. | |
| # Normal path, or Python 3.10 which ignores PYTHONSAFEPATH: sys.path[0] is the generated | |
| # script's directory, so the prelude swaps in the working directory. |
Join container paths with / so a Windows client cannot produce a backslash, keep the prelude next to EXEC_FUNC_SCRIPT so the tests reuse it, and drop the elif branch: python and torchrun always put the script directory first on sys.path unless safe path mode is on, so the swap alone covers every case. The behavioural test uses fixed module names and a nested training function, and asserts the generated script's location directly because chmod does not stop root. https: //github.com/kubeflow/pull/604#pullrequestreview-5548386367 Signed-off-by: danish9039 <danishsiddiqui040@gmail.com>
|
New changes are detected. LGTM label has been removed. |
|
@andreyvelich @juliusvonkohout Addressed all the suggestions PTAL |
| "if sys.path and os.path.realpath(sys.path[0]) == " | ||
| "os.path.realpath(os.path.dirname(__file__)):\n" |


What this PR does / why we need it:
Writes non-MPI CustomTrainer generated Python files and package-installation logs under
/tmpinstead of the image working directory.Restricted non-root containers can have a root-owned working directory, which prevents the SDK from creating its generated files. The generated script restores the image working directory as the first Python import root while preserving existing import paths, including Python safe-path mode. MPI continues to use
DEFAULT_MPI_USER_HOME.This is required by kubeflow/trainer#3702, which intentionally preserves each image's configured working directory.
Which issue(s) this PR fixes:
Follow-up to kubeflow/trainer#3702.
Validation:
/tmpmodulePYTHONPATHpreservationgit diff --checkChecklist: