diff --git a/.buildkite/scripts/ci-otel/ci_otel.sh b/.buildkite/scripts/ci-otel/ci_otel.sh index c9897e732e..7f1b700fe6 100755 --- a/.buildkite/scripts/ci-otel/ci_otel.sh +++ b/.buildkite/scripts/ci-otel/ci_otel.sh @@ -65,6 +65,39 @@ _ci_otel_python() { fi } +# The Docker test commands run as root, while /workdir is owned by the +# Buildkite checkout user. Run the GPU sampler as that unprivileged identity so +# any future project-code import cannot leave root-owned files in the checkout. +_ci_otel_prepare_gpu_user() { + _CI_INFRA_GPU_UID="" + _CI_INFRA_GPU_GID="" + [ "$(id -u 2>/dev/null || :)" = "0" ] || return 0 + command -v setpriv >/dev/null 2>&1 || return 1 + + _CI_INFRA_OTEL_WORKSPACE_DIR="${CI_INFRA_OTEL_WORKSPACE_DIR:-/workdir}" + [ -d "${_CI_INFRA_OTEL_WORKSPACE_DIR}" ] || return 1 + # The owner fields are numeric and whitespace-free on GNU stat. + # shellcheck disable=SC2046 + set -- $(stat -c '%u %g' -- "${_CI_INFRA_OTEL_WORKSPACE_DIR}" 2>/dev/null || :) + [ "$#" -eq 2 ] || return 1 + case "$1:$2" in + *[!0-9:]* | 0:* | *:0) return 1 ;; + esac + + _CI_INFRA_GPU_UID="$1" + _CI_INFRA_GPU_GID="$2" + chown "${_CI_INFRA_GPU_UID}:${_CI_INFRA_GPU_GID}" \ + "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + chmod 0700 "${CI_INFRA_OTEL_SPOOL_DIR}" || return 1 + if [ "${_CI_INFRA_OTEL_OWNS_RUNTIME}" = "1" ]; then + # Permit traversal to the private child spool, but not directory listing. + chmod 0711 "${CI_INFRA_OTEL_RUNTIME_DIR}" || return 1 + fi + setpriv --reuid="${_CI_INFRA_GPU_UID}" \ + --regid="${_CI_INFRA_GPU_GID}" --clear-groups --no-new-privs \ + -- test -w "${CI_INFRA_OTEL_SPOOL_DIR}" +} + ci_otel_start() { _CI_INFRA_OTEL_COMMAND_INDEX="$1" _CI_INFRA_OTEL_COMMAND_LABEL="${2:-command ${_CI_INFRA_OTEL_COMMAND_INDEX}}" @@ -87,9 +120,23 @@ ci_otel_start() { if [ "${CI_INFRA_GPU_SAMPLING:-1}" != "0" ] && [ -f "${CI_INFRA_OTEL_DIR}/ci_gpu.py" ] && command -v nvidia-smi >/dev/null 2>&1; then - PYTHONDONTWRITEBYTECODE=1 \ - "${_CI_INFRA_OTEL_PYTHON}" "${CI_INFRA_OTEL_DIR}/ci_gpu.py" "$$" &2 || : + fi fi } diff --git a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py index ae51010f9e..c66530989f 100644 --- a/.buildkite/scripts/ci-otel/tests/test_ci_otel.py +++ b/.buildkite/scripts/ci-otel/tests/test_ci_otel.py @@ -987,6 +987,85 @@ def test_gpu_periodic_batches_are_bounded_even_when_upload_fails(monkeypatch): assert [len(batch[0].events) for batch in batches] == [1, 1, 1] +def _write_executable(path: Path, body: str) -> None: + path.write_text(f"#!/bin/sh\n{body}") + path.chmod(0o755) + + +@pytest.mark.parametrize( + "workspace_owner,expect_sampler", + [("997 996", True), ("0 0", False)], +) +def test_root_shell_runs_gpu_sampler_as_checkout_owner( + tmp_path, workspace_owner, expect_sampler +): + bin_dir = tmp_path / "bin" + workspace = tmp_path / "workdir" + runtime = tmp_path / "runtime" + log = tmp_path / "calls.log" + bin_dir.mkdir() + workspace.mkdir() + + _write_executable( + bin_dir / "id", + '[ "$1" = "-u" ] && { echo 0; exit 0; }; exec /usr/bin/id "$@"\n', + ) + _write_executable( + bin_dir / "stat", + f"printf '%s\\n' '{workspace_owner}'\n", + ) + for command in ("chown", "chmod"): + _write_executable(bin_dir / command, "exit 0\n") + _write_executable( + bin_dir / "setpriv", + 'printf "setpriv %s\\n" "$*" >> "$CI_OTEL_TEST_LOG"\n' + 'while [ "$#" -gt 0 ] && [ "$1" != "--" ]; do shift; done\n' + '[ "$#" -gt 0 ] && shift\n' + 'exec "$@"\n', + ) + _write_executable( + bin_dir / "python3", + '[ "$1" = "-c" ] && exit 0\n' + '[ "$2" = "new-context" ] && {\n' + ' echo "01010101010101010101010101010101 0202020202020202 - 1"\n' + " exit 0\n" + "}\n" + 'case "$1" in *ci_gpu.py) echo gpu >> "$CI_OTEL_TEST_LOG";; esac\n' + "exit 0\n", + ) + _write_executable(bin_dir / "nvidia-smi", "exit 0\n") + shell = ( + f'. "{SCRIPTS_DIR / "ci_otel.sh"}"; ' + f"ci_otel_start 1 {_quoted('true')}; ci_otel_finish 0" + ) + + result = subprocess.run( + ["/bin/sh", "-c", shell], + check=False, + capture_output=True, + text=True, + env={ + **os.environ, + "PATH": f"{bin_dir}:{os.environ['PATH']}", + "CI_INFRA_OTEL_DIR": str(SCRIPTS_DIR), + "CI_INFRA_OTEL_RUNTIME_DIR": str(runtime), + "CI_INFRA_OTEL_WORKSPACE_DIR": str(workspace), + "CI_OTEL_TEST_LOG": str(log), + }, + ) + + assert result.returncode == 0, result.stderr + calls = log.read_text() if log.exists() else "" + if expect_sampler: + expected = "--reuid=997 --regid=996 --clear-groups --no-new-privs --" + assert calls.count("setpriv") == 2 + assert calls.count(expected) == 2 + assert "gpu" in calls + else: + assert calls == "" + assert "non-root GPU sampler unavailable" in result.stderr + + @pytest.mark.parametrize("sampling", ["1", "0"]) def test_gpu_sampler_stops_with_command_and_preserves_failure_status( tmp_path, sampling