Skip to content

fix(trainer): use writable generated script paths - #604

Open
danish9039 wants to merge 5 commits into
kubeflow:mainfrom
danish9039:pr3487/sdk-writable-script-path
Open

danish9039 wants to merge 5 commits into
kubeflow:mainfrom
danish9039:pr3487/sdk-writable-script-path

Conversation

@danish9039

Copy link
Copy Markdown
Member

What this PR does / why we need it:

Writes non-MPI CustomTrainer generated Python files and package-installation logs under /tmp instead 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:

  • 127 focused Kubernetes backend tests
  • Read-only working-directory execution in normal and safe-path modes
  • Working-directory precedence over a same-named /tmp module
  • Existing PYTHONPATH preservation
  • Concurrent behavioral execution
  • MPI path preservation
  • Ruff lint and formatting checks
  • git diff --check

Checklist:

  • Documentation included if any changes are user facing

@google-oss-prow

Copy link
Copy Markdown
Contributor

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign electronic-waste for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details Needs approval from an approver in each of these files:

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

@github-actions

Copy link
Copy Markdown
Contributor

🎉 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:

  • If you haven't already, please check out our Contributing Guide for repo-specific guidelines and the Kubeflow Contributor Guide for general community standards
  • Our team will review your PR soon! cc @kubeflow/kubeflow-sdk-team

Join the community:

Feel free to ask questions in the comments if you need any help or clarification!
Thanks again for contributing to Kubeflow! 🙏

@lntutor lntutor left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@google-oss-prow

Copy link
Copy Markdown
Contributor

@lntutor: changing LGTM is restricted to collaborators

Details

In response to this:

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.

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.

@danish9039

Copy link
Copy Markdown
Member Author

Two updates since the approval.

  1. 20501cb makes the generated-script prelude honour Python safe path mode. 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 under -P or PYTHONSAFEPATH=1, whose entire purpose is to keep a writable directory off the import path. The branch is now elif not getattr(sys.flags, "safe_path", False):, so the working directory is added only when Python itself would have added a script directory. CPython made the same fix to pdb in gh-99367: Do not mangle sys.path[0] in pdb if safe_path is set python/cpython#111762; ipykernel (Respect sys.flags.safe_path ipython/ipykernel#1311) and pyright (Respect PYTHONSAFEPATH if set microsoft/pyright#8168) resolved the same problem the same way. getattr rather than sys.flags.safe_path because the attribute only exists on 3.11+ and this package declares requires-python = ">=3.10".
  2. 6b8cd2a fixes the Test (Python 3.10) failure that the first commit introduced: Python 3.10 has no safe path mode and ignores PYTHONSAFEPATH, so the prelude correctly falls back there and the safe-path test now asserts that fallback on 3.10 instead of the 3.11+ behaviour. Verified in python:3.10-slim and python:3.12-slim; ruff check and ruff format are clean.
Case Before After
normal working directory working directory (unchanged)
normal with PYTHONPATH working directory, PYTHONPATH behind it unchanged
-P or PYTHONSAFEPATH=1 (3.11+) working directory, safe path defeated ModuleNotFoundError, honoured
-P with PYTHONPATH (3.11+) working directory ahead of PYTHONPATH resolves from PYTHONPATH
Python 3.10, any mode working directory working directory (unchanged)

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 ./data under the image working directory, which kubeflow/trainer#3702 now handles on the runtime side.

Separately, kubeflow/trainer/constants/constants.py uses read -r -d '' SCRIPT << EOM with an unquoted heredoc delimiter, so bash expands $ and backticks inside the user's Python source before it is written. That is pre-existing and unrelated, so it is deliberately left out here; happy to raise it separately.

@danish9039
danish9039 marked this pull request as ready for review August 30, 2026 14:52
@google-oss-prow
google-oss-prow Bot requested a review from tariq-hasan August 30, 2026 14:53
@juliusvonkohout

Copy link
Copy Markdown
Member

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.

@andreyvelich for /approve
/lgtm

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 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 /tmp paths 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.

Comment thread kubeflow/trainer/backends/kubernetes/utils.py Outdated

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot review overview

🟡 Changes recommended

Fixed shared /tmp filenames can collide and may execute stale script contents after a failed write.

Review effort: Balanced
Findings: 1 High severity

Open (1)
Resolved since last review (1)

Comment thread kubeflow/trainer/backends/kubernetes/utils.py Outdated
@juliusvonkohout
juliusvonkohout force-pushed the pr3487/sdk-writable-script-path branch from e3e4ca2 to 21369ae Compare October 2, 2026 11:46
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",

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What if /tmp directory doesn't exist in the image?
Any thoughts on this @briangallagher @robert-bell @kubeflow/kubeflow-sdk-team @tariq-hasan ?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@juliusvonkohout juliusvonkohout Oct 7, 2026 •

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

/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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I like @tariq-hasan's approach with ${TMPDIR:-/tmp}/pip_install.py.
WDYT @danish9039 @juliusvonkohout @robert-bell ?

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sounds good

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

danish9039 added a commit to danish9039/sdk that referenced this pull request Oct 8, 2026
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>
@danish9039
danish9039 force-pushed the pr3487/sdk-writable-script-path branch from 108f688 to db8182c Compare October 8, 2026 17:00
@danish9039

Copy link
Copy Markdown
Member Author

@andreyvelich @juliusvonkohout @tariq-hasan @robert-bell generated script and the pip install log now go to ${TMPDIR:-/tmp}, resolved by the container shell, so an image that sets TMPDIR is honoured and /tmp remains the default. The value lives in one place,DEFAULT_TEMPORARY_DIRECTORY in constants.py. MPI runtimes are unchanged and keep /home/mpiuser.

The behavioural test now sets TMPDIR to a per-test directory, which proves the resolution end to end and keeps the test out of the shared /tmp.

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Approval recommended

The implementation is consistent with the stated compatibility requirements and has focused behavioral coverage.

0 open findings

1 resolved since last review

🧠 Review effort: Balanced

@juliusvonkohout

Copy link
Copy Markdown
Member

/lgtm

@andreyvelich andreyvelich left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the fix, @danish9039! This review was done with the help of AI tools.

Comment on lines +385 to +389
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")

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
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"

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch on Windows, switched to / and fixed the comment.

Comment on lines +390 to +414
# 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

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
# 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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Makes sense, dropped it. The swap covers every case now that the -I on 3.10 path is gone.

Comment thread kubeflow/trainer/constants/constants.py Outdated
Comment on lines +153 to +156
# 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}"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Make it clear that this is a shell expression, and keep the prelude next to EXEC_FUNC_SCRIPT so tests can reuse it:

Suggested change
# 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"
)

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Applied as is.

Comment on lines +38 to +47
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"
)

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reuse the constant instead of keeping a copy:

Suggested change
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

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

Comment on lines +298 to +306
"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"

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Reuse the constant instead of keeping a copy:

Suggested change
"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}"

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done here too.

tmp_path,
safe_path_enabled,
):
unique_name_suffix = uuid.uuid4().hex

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Went with the nested def, that removed uuid and the importlib bits together.

Comment on lines +758 to +767
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):

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

Suggested change
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):

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right, root ignores the chmod. Asserting is_file() on the generated path now.

Comment on lines +776 to +778
# 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.

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On 3.10, sys.path[0] is the script directory, so the if swap runs and the getattr fallback is never reached:

Suggested change
# 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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Updated.

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>
@google-oss-prow google-oss-prow Bot removed the lgtm label Oct 21, 2026
@google-oss-prow

Copy link
Copy Markdown
Contributor

New changes are detected. LGTM label has been removed.

@danish9039

Copy link
Copy Markdown
Member Author

@andreyvelich @juliusvonkohout Addressed all the suggestions PTAL

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The prelude can replace an explicit PYTHONPATH entry with the working directory when safe-path mode is enabled.

1 open issue

Review effort: Balanced

Comment on lines +164 to +165
"if sys.path and os.path.realpath(sys.path[0]) == "
"os.path.realpath(os.path.dirname(__file__)):\n"

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

Status: Ready to Merge

Development

Successfully merging this pull request may close these issues.

7 participants