Skip to content

fix: improve logging around packages_to_install #269

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
briangallagher wants to merge 2 commits into
base: main
Choose a base branch
from
+108 −59
Open

fix: improve logging around packages_to_install #269

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
363a68f
improve logging around packages_to_install
briangallagher Feb 6, 2026
f009677
exit when pip install fails, append errors from both attempts
briangallagher Feb 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 19 additions & 15 deletions kubeflow/trainer/backends/kubernetes/backend_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -223,26 +223,30 @@ def get_custom_trainer(
"""
Get the custom trainer for the TrainJob.
"""
pip_command = [f"--index-url {pip_index_urls[0]}"]
pip_command.extend([f"--extra-index-url {repo}" for repo in pip_index_urls[1:]])
pip_command = " ".join(pip_command)
# Use the same helper as production code to build the pip install script so
# tests stay in sync with the runtime behavior.
install_script = utils.get_script_for_python_packages(
packages_to_install=packages_to_install,
pip_index_urls=pip_index_urls,
)

# 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"
'func=lambda: print("Hello World"),\n\n'
"<lambda>(**{'learning_rate': 0.001, 'batch_size': 32})\n\n"
'EOM\nprintf "%s" "$SCRIPT" > "backend_test.py"\n'
'torchrun "backend_test.py"'
)

full_command = install_script + func_script

packages_command = " ".join(packages_to_install)
return models.TrainerV1alpha1Trainer(
command=[
"bash",
"-c",
'\nif ! [ -x "$(command -v pip)" ]; then\n python -m ensurepip '
"|| python -m ensurepip --user || apt-get install python-pip"
"\nfi\n\n"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet"
f" --no-warn-script-location {pip_command} --user {packages_command}"
" ||\nPIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet"
f" --no-warn-script-location {pip_command} {packages_command}"
"\n\nread -r -d '' SCRIPT << EOM\n\nfunc=lambda: "
'print("Hello World"),\n\n<lambda>(**'
"{'learning_rate': 0.001, 'batch_size': 32})\n\nEOM\nprintf \"%s\" "
'"$SCRIPT" > "backend_test.py"\ntorchrun "backend_test.py"',
full_command,
],
numNodes=2,
env=env,
Expand Down
36 changes: 24 additions & 12 deletions kubeflow/trainer/backends/kubernetes/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -268,6 +268,7 @@ def get_script_for_python_packages(
# first url will be the index-url.
options = [f"--index-url {pip_index_urls[0]}"]
options.extend(f"--extra-index-url {extra_index_url}" for extra_index_url in pip_index_urls[1:])
options_str = " ".join(options)
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

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

Command injection vulnerability: pip_index_urls are embedded in shell options without escaping. A malicious URL like 'https://evil.com"; rm -rf / #' could break out of the quotes and execute arbitrary commands. Apply shlex.quote() to each URL when building options: options = [f"--index-url {shlex.quote(pip_index_urls[0])}"] and similarly for extra-index-url.

Copilot uses AI. Check for mistakes.
Copy link
Member

@andreyvelich andreyvelich Feb 20, 2026 •
edited
Loading

Choose a reason for hiding this comment

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

Good catch, tho.
@kubeflow/kubeflow-sdk-team @kubeflow/wg-pipeline-leads Do we know if we have the same CVE in KFP Client upstream?


header_script = textwrap.dedent(
"""
Expand All @@ -278,18 +279,29 @@ def get_script_for_python_packages(
"""
)

script_for_python_packages = (
header_script
+ "PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
+ "--no-warn-script-location {} --user {}".format(
" ".join(options),
packages_str,
)
+ " ||\nPIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
+ "--no-warn-script-location {} {}\n".format(
" ".join(options),
packages_str,
)
# First try per-user installation, then fall back to system-wide installation.
# Pip output is captured to a log file and only printed when both attempts fail;
# on success we emit a single concise confirmation line.
script_for_python_packages = header_script + textwrap.dedent(
f"""
PACKAGES="{packages_str}"
PIP_OPTS="{options_str}"
Comment on lines +287 to +288
Copy link

Copilot AI Feb 18, 2026

Choose a reason for hiding this comment

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

Command injection vulnerability: packages_str and options_str are directly interpolated into double-quoted shell strings without escaping. A malicious package name like 'torch"; rm -rf / #' or URL with shell metacharacters could break out of the quotes and execute arbitrary commands. Use shlex.quote() to properly escape each package name and URL before joining them, or use shell arrays for safer variable handling.

Copilot uses AI. Check for mistakes.
LOG_FILE=/tmp/pip_install.log
Copy link
Member

@andreyvelich andreyvelich Feb 20, 2026

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, shall we just use ?
Do we know if KFP client write pip logs to a file?
cc @kubeflow/wg-pipeline-leads

rm -f "$LOG_FILE"

if PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet \\
--no-warn-script-location $PIP_OPTS --user $PACKAGES >"$LOG_FILE" 2>&1; then
echo "Successfully installed Python packages: $PACKAGES"
Copy link
Member

@andreyvelich andreyvelich Feb 20, 2026

Choose a reason for hiding this comment

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

Maybe we should say?

Suggested change
echo "Successfully installed Python packages: $PACKAGES"
echo "Successfully installed the user' Python packages: $PACKAGES"

elif PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet \\
--no-warn-script-location $PIP_OPTS $PACKAGES >>"$LOG_FILE" 2>&1; then
echo "Successfully installed Python packages: $PACKAGES"
Copy link
Member

@andreyvelich andreyvelich Feb 20, 2026

Choose a reason for hiding this comment

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

Suggested change
echo "Successfully installed Python packages: $PACKAGES"
echo "Successfully installed the system's Python packages: $PACKAGES"

else
echo "ERROR: Failed to install Python packages: $PACKAGES" >&2
cat "$LOG_FILE" >&2
Comment on lines +299 to +300
Copy link
Contributor

@kramaranya kramaranya Feb 11, 2026

Choose a reason for hiding this comment

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

shall we have exit 1 here?

Copy link
Contributor Author

@briangallagher briangallagher Feb 18, 2026

Choose a reason for hiding this comment

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

I think we should, updated.

Comment on lines 292 to 300
Copy link
Contributor

@kramaranya kramaranya Feb 11, 2026

Choose a reason for hiding this comment

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

Looks like we overwrite the first attempt's output with the second. Can we append the second output?

Copy link
Contributor Author

@briangallagher briangallagher Feb 18, 2026

Choose a reason for hiding this comment

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

Good catch, updated now

exit 1
fi

"""
)

return script_for_python_packages
Expand Down
97 changes: 65 additions & 32 deletions kubeflow/trainer/backends/kubernetes/utils_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -151,17 +151,23 @@ def test_get_resources_per_node(test_case: TestCase):
'\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="torch numpy custom-package"\n'
'PIP_OPTS="--index-url https://pypi.org/simple --extra-index-url https://private.repo.com/simple --extra-index-url https://internal.company.com/simple"\n'
"LOG_FILE=/tmp/pip_install.log\n"
'rm -f "$LOG_FILE"\n'
"\n"
"if PIP_DISABLE_PIP_VERSION_CHECK=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: $PACKAGES"\n'
"elif PIP_DISABLE_PIP_VERSION_CHECK=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: $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"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"--extra-index-url https://private.repo.com/simple "
"--extra-index-url https://internal.company.com/simple "
"--user torch numpy custom-package ||\n"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"--extra-index-url https://private.repo.com/simple "
"--extra-index-url https://internal.company.com/simple "
"torch numpy custom-package\n"
),
),
TestCase(
Expand All @@ -175,13 +181,23 @@ def test_get_resources_per_node(test_case: TestCase):
'\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="torch numpy custom-package"\n'
'PIP_OPTS="--index-url https://pypi.org/simple"\n'
"LOG_FILE=/tmp/pip_install.log\n"
'rm -f "$LOG_FILE"\n'
"\n"
"if PIP_DISABLE_PIP_VERSION_CHECK=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: $PACKAGES"\n'
"elif PIP_DISABLE_PIP_VERSION_CHECK=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: $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"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"--user torch numpy custom-package ||\n"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"torch numpy custom-package\n"
),
),
TestCase(
Expand All @@ -199,17 +215,23 @@ def test_get_resources_per_node(test_case: TestCase):
'\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="torch numpy custom-package"\n'
'PIP_OPTS="--index-url https://pypi.org/simple --extra-index-url https://private.repo.com/simple --extra-index-url https://internal.company.com/simple"\n'
"LOG_FILE=/tmp/pip_install.log\n"
'rm -f "$LOG_FILE"\n'
"\n"
"if PIP_DISABLE_PIP_VERSION_CHECK=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: $PACKAGES"\n'
"elif PIP_DISABLE_PIP_VERSION_CHECK=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: $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"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"--extra-index-url https://private.repo.com/simple "
"--extra-index-url https://internal.company.com/simple "
"--user torch numpy custom-package ||\n"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
"--no-warn-script-location --index-url https://pypi.org/simple "
"--extra-index-url https://private.repo.com/simple "
"--extra-index-url https://internal.company.com/simple "
"torch numpy custom-package\n"
),
),
TestCase(
Expand All @@ -223,13 +245,24 @@ def test_get_resources_per_node(test_case: TestCase):
'\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="torch numpy"\n'
"PIP_OPTS="
f'"--index-url {constants.DEFAULT_PIP_INDEX_URLS[0]}"\n'
"LOG_FILE=/tmp/pip_install.log\n"
'rm -f "$LOG_FILE"\n'
"\n"
"if PIP_DISABLE_PIP_VERSION_CHECK=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: $PACKAGES"\n'
"elif PIP_DISABLE_PIP_VERSION_CHECK=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: $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"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
f"--no-warn-script-location --index-url "
f"{constants.DEFAULT_PIP_INDEX_URLS[0]} --user torch numpy ||\n"
"PIP_DISABLE_PIP_VERSION_CHECK=1 python -m pip install --quiet "
f"--no-warn-script-location --index-url "
f"{constants.DEFAULT_PIP_INDEX_URLS[0]} torch numpy\n"
),
),
],
Expand Down
Loading