diff --git a/docker/Dockerfile.multi b/docker/Dockerfile.multi index 2b421e45b0b2..105418a5b0e2 100644 --- a/docker/Dockerfile.multi +++ b/docker/Dockerfile.multi @@ -175,20 +175,37 @@ ARG TRT_LLM_VER ARG GITHUB_MIRROR="" ARG BUILD_WHEEL_ARGS="--clean" ARG BUILD_WHEEL_SCRIPT="scripts/build_wheel.py" +# Used when BUILD_WHEEL_SCRIPT takes the wheel from a packaged build +# (scripts/get_wheel_from_package.py): the location read first (aws or urm; empty +# means aws) and the name of the credential behind the artifactory_netrc secret, +# which appears only in error messages. +ARG LLM_ARTIFACTS_PRIMARY="" +ARG ARTIFACTORY_CREDENTIAL_ID="" RUN --mount=type=cache,target=/root/.cache/pip --mount=type=cache,target=${CCACHE_DIR} \ --mount=type=secret,id=github_clone_token \ + --mount=type=secret,id=artifactory_netrc \ + --mount=type=bind,source=jenkins/scripts/fetch_build_artifact.sh,target=/src/tensorrt_llm/jenkins/scripts/fetch_build_artifact.sh \ if [ -n "$GITHUB_MIRROR" ]; then \ export PIP_INDEX_URL="https://artifactory.pdx.nvidia.com/artifactory/api/pypi/pypi-remote/simple"; \ fi && \ # Authenticates `git submodule update` and the CMake FetchContent clones # (3rdparty/fetch_content.json) that this build pulls from github.com. . docker/common/github_auth.sh && \ + # get_wheel_from_package.py reads the packaged build through the bind-mounted + # jenkins/scripts/fetch_build_artifact.sh, which takes the Artifactory access token + # for artifactory.nvidia.com from this netrc. The secret is present only when the + # build passes it (ARTIFACTORY_NETRC in docker/Makefile). + if [ -f /run/secrets/artifactory_netrc ]; then \ + export ARTIFACTORY_NETRC=/run/secrets/artifactory_netrc; \ + fi && \ set -- ${BUILD_WHEEL_ARGS} && \ if [ "${BUILD_WHEEL_SCRIPT}" = "scripts/build_wheel.py" ] && [ -n "${TRT_LLM_VER}" ]; then \ set -- "$@" --version-override "${TRT_LLM_VER}"; \ fi && \ TRTLLM_BUILD_SOURCE_COMMIT=${GIT_COMMIT} \ GITHUB_MIRROR=${GITHUB_MIRROR} \ + LLM_ARTIFACTS_PRIMARY=${LLM_ARTIFACTS_PRIMARY} \ + ARTIFACTORY_CREDENTIAL_ID=${ARTIFACTORY_CREDENTIAL_ID} \ python3 ${BUILD_WHEEL_SCRIPT} "$@" FROM ${DEVEL_IMAGE} AS release diff --git a/docker/Makefile b/docker/Makefile index cff393c4e728..bd59c447715e 100644 --- a/docker/Makefile +++ b/docker/Makefile @@ -56,8 +56,19 @@ GITHUB_MIRROR ?= # environment, never passed as a build arg -- build-arg values stay readable in # `docker history`. Empty means anonymous, rate-limited access. GITHUB_CLONE_TOKEN ?= +# Optional netrc file with an Artifactory access token for artifactory.nvidia.com, +# needed when the wheel stage takes the wheel from a packaged build +# (BUILD_WHEEL_SCRIPT=scripts/get_wheel_from_package.py, which reads it through +# jenkins/scripts/fetch_build_artifact.sh). Passed as a BuildKit secret, never as a +# build arg, for the same reason as GITHUB_CLONE_TOKEN. +ARTIFACTORY_NETRC ?= +# The location the packaged build is read from first (aws or urm; empty means aws), +# and the name of the CI credential behind ARTIFACTORY_NETRC, used only in messages. +LLM_ARTIFACTS_PRIMARY ?= +ARTIFACTORY_CREDENTIAL_ID ?= COMMA := , GITHUB_SECRET_ARG := $(if $(GITHUB_CLONE_TOKEN),--secret id=github_clone_token$(COMMA)env=GITHUB_CLONE_TOKEN) +ARTIFACTORY_SECRET_ARG := $(if $(ARTIFACTORY_NETRC),--secret id=artifactory_netrc$(COMMA)src=$(ARTIFACTORY_NETRC)) PYTHON_VERSION ?= NGC_STAGING_REPO ?= nvcr.io/nvstaging/tensorrt-llm NGC_REPO ?= nvcr.io/nvidia/tensorrt-llm @@ -96,6 +107,7 @@ base_pull: docker buildx build $(DOCKER_BUILD_OPTS) $(DOCKER_BUILD_ARGS) \ --progress $(DOCKER_PROGRESS) \ $(GITHUB_SECRET_ARG) \ + $(ARTIFACTORY_SECRET_ARG) \ $(if $(BASE_IMAGE), --build-arg BASE_IMAGE=$(BASE_IMAGE)) \ $(if $(BASE_TAG), --build-arg BASE_TAG=$(BASE_TAG)) \ $(if $(TRITON_IMAGE), --build-arg TRITON_IMAGE=$(TRITON_IMAGE)) \ @@ -103,6 +115,8 @@ base_pull: $(if $(CUDA_COMPAT_TAG), --build-arg CUDA_COMPAT_TAG=$(CUDA_COMPAT_TAG)) \ $(if $(BUILD_WHEEL_ARGS), --build-arg BUILD_WHEEL_ARGS="$(BUILD_WHEEL_ARGS)") \ $(if $(BUILD_WHEEL_SCRIPT), --build-arg BUILD_WHEEL_SCRIPT="$(BUILD_WHEEL_SCRIPT)") \ + $(if $(LLM_ARTIFACTS_PRIMARY), --build-arg LLM_ARTIFACTS_PRIMARY="$(LLM_ARTIFACTS_PRIMARY)") \ + $(if $(ARTIFACTORY_CREDENTIAL_ID), --build-arg ARTIFACTORY_CREDENTIAL_ID="$(ARTIFACTORY_CREDENTIAL_ID)") \ $(if $(TORCH_INSTALL_TYPE), --build-arg TORCH_INSTALL_TYPE="$(TORCH_INSTALL_TYPE)") \ $(if $(CUDA_VERSION), --build-arg CUDA_VER="$(CUDA_VERSION)") \ $(if $(CUDNN_VERSION), --build-arg CUDNN_VER="$(CUDNN_VERSION)") \ diff --git a/jenkins/BoltProfileGen.groovy b/jenkins/BoltProfileGen.groovy index 740d6c9fe756..0067017fdad7 100644 --- a/jenkins/BoltProfileGen.groovy +++ b/jenkins/BoltProfileGen.groovy @@ -42,13 +42,67 @@ import com.nvidia.bloom.SlurmCluster import com.nvidia.bloom.SlurmPartition import com.nvidia.bloom.Utils +// LLM repository configuration +withCredentials([string(credentialsId: 'default-llm-repo', variable: 'DEFAULT_LLM_REPO')]) { + LLM_REPO = env.gitlabSourceRepoHttpUrl ? env.gitlabSourceRepoHttpUrl : "${DEFAULT_LLM_REPO}" +} LLM_ROOT = "llm" // Input tarball location (passed from the parent via artifactPath, same as -// Build/L0_Test). The promoted bundle goes to a stable BRANCH-keyed path. +// Build/L0_Test). The promoted bundle goes to a stable BRANCH-keyed path on +// URM_ARTIFACTORY_BASE (PROFILE_PROMOTE_DIR below). ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" URM_ARTIFACTORY_BASE = "https://urm.nvidia.com/artifactory" +// Build artifacts (build_info.txt, packaged builds, waive_list/) are published to the +// primary location and read from it first; LLM_ARTIFACTS_PRIMARY ("aws", the default, or "urm") +// selects it. Reads fall back to the other location only when a file is absent. +// credential is the identity that publishes to the location; reads from aws use it too. +LLM_ARTIFACTS_LOCATIONS = [ + aws: [root: "https://artifactory.nvidia.com/artifactory", repo: "sw-tensorrt-llm-generic-local", + credential: "aws-artifactory-credentials"], + urm: [root: "https://urm.nvidia.com/artifactory", repo: "sw-tensorrt-generic", + credential: "urm-artifactory-creds"], +] + +// Exit code of jenkins/scripts/fetch_build_artifact.sh that this pipeline acts on, by the name +// the helper gives it. tests/unittest/scripts/test_fetch_build_artifact.py checks that it +// matches. +FETCH_ABSENT = 2 + +// The primary location a parent pipeline passed in globalVars, which this job uses instead of its +// own controller's LLM_ARTIFACTS_PRIMARY. Empty when the job was started on its own. +LLM_ARTIFACTS_PRIMARY_KEY = "llm_artifacts_primary" +LLM_ARTIFACTS_PRIMARY_FROM_PARENT = "" + +def llmArtifactsPrimary() +{ + def primary = LLM_ARTIFACTS_PRIMARY_FROM_PARENT ?: env.LLM_ARTIFACTS_PRIMARY ?: "aws" + if (!(primary in ["aws", "urm"])) { + error("LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '${primary}'") + } + return primary +} + +def llmArtifactsLayoutPath(String legacyPath) +{ + def layoutPath = legacyPath.replaceFirst(/^sw-tensorrt-generic(-local)?\//, "") + if (!layoutPath.startsWith("llm-artifacts/")) { + error("Not an llm-artifacts path: ${legacyPath}") + } + return layoutPath +} + +def llmArtifactsRepoPath(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].repo}/${llmArtifactsLayoutPath(legacyPath)}" +} + +def llmArtifactsUrl(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].root}/${llmArtifactsRepoPath(location, legacyPath)}" +} + X86_64_TRIPLE = "x86_64-linux-gnu" AARCH64_TRIPLE = "aarch64-linux-gnu" @@ -264,15 +318,30 @@ def b64BashRemoteCmdStdin(String script, String scriptPath) return "\"echo ${b64} | base64 -d > ${scriptPath}; trap 'rm -f ${scriptPath}' EXIT; bash ${scriptPath}\"" } +// For scripts too long for the two above, such as those carrying fetch_verified: +// a csh/tcsh login shell rejects a word that long ("Word too long"). Copies the +// script to the cluster as a file and returns the sshUserCmd remotecmd that runs it +// under bash, with stdin left connected to whatever ssh forwards. The script must +// hold no secret: it stays at remotePath, so a retried step can run it again. +def bashRemoteScriptFile(pipeline, remote, String script, String localPath, String remotePath) +{ + pipeline.writeFile(file: localPath, text: script) + Utils.copyFileToRemoteHost(pipeline, remote, localPath, remotePath) + return "\"bash ${remotePath}\"" +} + // --------------------------------------------------------------------------- // SLURM profile generation: fan out perf-harness runs (one per workload, BOLT // hook enabled) then a single cross-workload merge (slurm_merge.sh). // --------------------------------------------------------------------------- -def submitProfileGen(pipeline) +// readHelper: local copy of jenkins/scripts/fetch_build_artifact.sh from the run's +// commit; the tarball download below runs it on the cluster frontend. +def submitProfileGen(pipeline, String readHelper) { - // The tarball to profile: - def llmTarfile = "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}/${BOLT_TARNAME}" - pipeline.echo("Input tarball: ${llmTarfile}") + // The tarball to profile. The primary location is resolved once here and used + // both to read it and to publish the BOLTed build back. + def primary = llmArtifactsPrimary() + pipeline.echo("Input tarball: ${llmArtifactsLayoutPath(ARTIFACT_PATH)}/${BOLT_TARNAME} (read from ${primary} first)") SlurmPartition partition = SlurmConfig.resolvePlatform(SLURM_PLATFORM) SlurmCluster cluster = SlurmConfig.clusterConfig[partition.clusterName] @@ -289,9 +358,14 @@ def submitProfileGen(pipeline) def outDir = "${fdataRoot}/_bundle" def bundle = "${outDir}/bolt-profile-${BOLT_REF}-${TRIPLE}.tar.gz" - // Bootstrap on the frontend entirely from that tarball (no agent-> - // cluster file copy): extract the full source tree + wheel (the perf harness - // runs from the checkout, install_mode=wheel), and stage llvm-bolt once + // Scripts too long for the ssh command line are written here, beside the read + // helper, before bashRemoteScriptFile copies them to the cluster. + def localScriptDir = readHelper.substring(0, readHelper.lastIndexOf('/')) + + // Bootstrap on the frontend from that tarball (the only agent->cluster file + // copies are the read helper and the bootstrap scripts): extract the full + // source tree + wheel (the perf harness runs from the checkout, + // install_mode=wheel), and stage llvm-bolt once // (shared, reused by the per-node instrument hook + merge). The merge job // mounts scripts/bolt directly from the extracted tree (no separate copy). // Then fan out one perf-harness run per workload, and merge once. @@ -311,9 +385,24 @@ def submitProfileGen(pipeline) // xargs + coreutils): HEAD-probe for size + range support, fan out PARTS // ranged `curl -r` GETs via `xargs -P`, reassemble, verify. Falls back to a // single-stream curl when the server doesn't advertise byte ranges or on - // ANY failure, so nothing regresses. Defined once here, reused for BOTH - // downloads. NO resume on retry (append is a corruption vector); the - // per-file verify (gzip -t / tar -tJf) catches truncation/corruption. + // ANY failure, so nothing regresses. artifactory.nvidia.com and GitHub both + // redirect downloads to a CDN on another host. Each single-stream attempt + // follows the redirect; only when the CDN answers HTTP 403 does the attempt + // ask again, once, with X-JFrog-Download-Redirect-To: None, which makes + // Artifactory send the file itself (other servers ignore the header). Any + // other failure fails the attempt: a connection error, a stall, any other + // status (a 403 from the URL's own host too), or a failed request without + // the redirect. A failed attempt is followed by another from the URL, up to + // 10 more, 15 s apart. The header goes out only right after a CDN 403, at + // most once per attempt (11 times in all). The other Artifactory location + // is never tried. Logs name hosts only: a redirect target can be a signed + // CDN URL, whose query string is a short-lived credential. Defined once + // here, reused for BOTH downloads. NO resume on retry (append is a + // corruption vector); the per-file verify (gzip -t + SHA-1 / tar -tJf) + // catches truncation/corruption. + // Credentials come only from FETCH_NETRC, which the tarball stage sets for + // a URL on the aws location, so the llvm-bolt download stays anonymous. curl + // applies a netrc entry only to the host it names, so the CDN never gets them. def boltFetchLib = ''' _verify_txz() { tar -tJf "$1" >/dev/null 2>&1; } # _bolt_range_download URL DEST PARTS -> 0 if the whole file was fetched @@ -324,7 +413,8 @@ def submitProfileGen(pipeline) # HEAD (not a ranged GET) so a server that ignores Range can't make # us stream the whole file just to probe; -L follows the redirect to # the storage/CDN backend; || true so a probe miss doesn't trip set -e. - hdr=$(curl -fsSL -I --retry 5 --retry-all-errors "$url" 2>/dev/null | tr -d '\\r' || true) + hdr=$(curl -fsSL -I --retry 5 --retry-all-errors ${FETCH_NETRC:+--netrc-file "$FETCH_NETRC"} \\ + "$url" 2>/dev/null | tr -d '\\r' || true) total=$(printf '%s\\n' "$hdr" | awk 'tolower($1)=="content-length:"{v=$2} END{print v}') accept=$(printf '%s\\n' "$hdr" | awk 'tolower($1)=="accept-ranges:"{v=tolower($2)} END{print v}') case "$total" in ''|*[!0-9]*) return 1 ;; esac @@ -341,10 +431,11 @@ def submitProfileGen(pipeline) printf '%s %s %s\\n' "$i" "$start" "$end" >> "$ranges" done # One ranged curl per chunk, PARTS at a time. $0/$1/$2 = i/start/end - # (xargs -n 3); URL/DEST come from the inline-exported environment. - if ! BOLT_FETCH_URL="$url" BOLT_FETCH_DEST="$dest" \\ + # (xargs -n 3); URL/DEST/NETRC come from the inline-exported environment. + if ! BOLT_FETCH_URL="$url" BOLT_FETCH_DEST="$dest" BOLT_FETCH_NETRC="${FETCH_NETRC:-}" \\ xargs -P "$parts" -n 3 bash -c ' curl -fsSL --retry 10 --retry-all-errors --retry-delay 5 \\ + ${BOLT_FETCH_NETRC:+--netrc-file "$BOLT_FETCH_NETRC"} \\ -r "$1-$2" -o "$BOLT_FETCH_DEST.part.$0" "$BOLT_FETCH_URL" ' < "$ranges"; then rm -f "$ranges" "$dest".part.* @@ -364,8 +455,68 @@ def submitProfileGen(pipeline) [ "$final" -eq "$total" ] || return 1 return 0 } - # fetch_verified URL DEST VERIFY... : parallel fetch then `VERIFY DEST`; - # single-stream curl fallback on any failure. Reads PARTS (default 16). + # _bolt_url_host URL -> the lower-case host[:port] of URL. + _bolt_url_host() { + local rest="${1#*://}" + rest="${rest%%[/?#]*}" + rest="${rest##*@}" + printf '%s\\n' "${rest,,}" + } + # _bolt_get URL DEST [CURL OPTION...] -> one GET of URL into DEST that follows + # redirects, with FETCH_NETRC's entry for URL's host. curl does not retry it, + # and without -f it exits 0 whatever the HTTP status. Sets BOLT_STATUS (the + # last answer's status, 000 when none arrived) and BOLT_FINAL_HOST (the host + # of the last URL requested: that URL can be a signed CDN URL, so only its + # host is kept), and returns curl's exit status. + _bolt_get() { + local url="$1" dest="$2" out rc=0 + shift 2 + rm -f "$dest" + out=$(curl -sSL --connect-timeout 60 --speed-time 120 --speed-limit 10000 \\ + ${FETCH_NETRC:+--netrc-file "$FETCH_NETRC"} "$@" \\ + -w '%{http_code} %{url_effective}' -o "$dest" "$url") || rc=$? + BOLT_STATUS=${out%% *} + BOLT_FINAL_HOST=$(_bolt_url_host "${out#* }") + return "$rc" + } + # _bolt_single_stream URL DEST HOST -> one attempt at downloading URL, whose + # host is HOST, into DEST: 0 when DEST holds a 200 answer, else 1 with the + # reason in BOLT_FETCH_FAILURE. + _bolt_single_stream() { + local url="$1" dest="$2" host="$3" rc=0 + _bolt_get "$url" "$dest" || rc=$? + if [ "$rc" -ne 0 ]; then + # The status can be an earlier hop's, so only the last host is named. + BOLT_FETCH_FAILURE="download from ${BOLT_FINAL_HOST:-$host} failed (curl exit $rc)" + return 1 + fi + if [ "$BOLT_STATUS" = 200 ]; then + return 0 + fi + if [ "$BOLT_STATUS" != 403 ] || [ "$BOLT_FINAL_HOST" = "$host" ]; then + # Any other status, a 403 from URL's own host among them. + BOLT_FETCH_FAILURE="download answered HTTP $BOLT_STATUS from $BOLT_FINAL_HOST" + return 1 + fi + # The CDN that URL's host redirected to rejected the request, which says + # nothing about the credential. Ask URL's host for the file itself: + # Artifactory honors X-JFrog-Download-Redirect-To: None. + echo "[INFO] $BOLT_FINAL_HOST rejected the redirected download (HTTP 403); downloading from $host without the redirect." + _bolt_get "$url" "$dest" -H 'X-JFrog-Download-Redirect-To: None' || rc=$? + if [ "$rc" -ne 0 ]; then + BOLT_FETCH_FAILURE="download from $host without the redirect failed (curl exit $rc)" + return 1 + fi + if [ "$BOLT_STATUS" = 200 ]; then + return 0 + fi + BOLT_FETCH_FAILURE="download from $host without the redirect answered HTTP $BOLT_STATUS from $BOLT_FINAL_HOST" + return 1 + } + # fetch_verified URL DEST VERIFY... : parallel fetch, else a single-stream + # download, then `VERIFY DEST`. Reads PARTS (default 16) and FETCH_NETRC + # (netrc for URL's host; unset or empty sends none). On any failure it + # removes DEST and returns 1, also when the caller runs under set -e. fetch_verified() { local url="$1" dest="$2"; shift 2 local parts="${PARTS:-16}" @@ -374,31 +525,89 @@ def submitProfileGen(pipeline) return 0 fi echo "[INFO] parallel fetch unavailable/failed; single-stream curl: $dest" - rm -f "$dest" - curl -fSL --retry 10 --retry-all-errors --retry-delay 15 \\ - --connect-timeout 60 --speed-time 120 --speed-limit 10000 \\ - -o "$dest" "$url" || { rm -f "$dest"; return 1; } - "$@" "$dest" || { echo "[ERROR] verification failed: $dest"; rm -f "$dest"; return 1; } + # A failed attempt, whatever the cause, is followed by another from URL, + # through the redirect: up to `retries` more, `retry_delay` seconds apart. + # A file that fails VERIFY is not downloaded again. + local host attempt=1 retries=10 retry_delay=15 + local attempts=$((retries + 1)) + host=$(_bolt_url_host "$url") + until _bolt_single_stream "$url" "$dest" "$host"; do + if [ "$attempt" -ge "$attempts" ]; then + echo "[ERROR] single-stream attempt $attempt of $attempts: $BOLT_FETCH_FAILURE: $dest" >&2 + rm -f "$dest" + return 1 + fi + echo "[INFO] single-stream attempt $attempt of $attempts: $BOLT_FETCH_FAILURE; retrying in ${retry_delay}s." + attempt=$((attempt + 1)) + sleep "$retry_delay" + done + if ! "$@" "$dest"; then + echo "[ERROR] verification failed: $dest" >&2 + rm -f "$dest" + return 1 + fi return 0 } '''.stripIndent() - // If a prior run against THIS artifactPath already published the BOLTed build - // as canonical, publishBoltedCanonical preserved the original as - // unbolted-. Prefer it so we profile/BOLT the un-BOLTed input rather - // than double-BOLTing. Inert for a fresh artifactPath (no unbolted- object), - // where canonical IS the un-BOLTed build. - def unboltedTarUrl = "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}/unbolted-${BOLT_TARNAME}" + // The read helper (copied into the workspace by the download stage) locates + // the tarball: `--exists` looks on the primary location first and on the + // other one only when the file is absent there, and prints the URL and SHA-1 + // of the copy it found, which fetch_verified then downloads and checks. + // Jenkins environment variables do not reach the cluster, so the script + // exports what the helper reads. If a prior run against THIS artifactPath + // already published the BOLTed build as canonical, publishBoltedCanonical + // preserved the original as unbolted-. Prefer it so we profile/BOLT + // the un-BOLTed input rather than double-BOLTing. Inert for a fresh + // artifactPath (no unbolted- object), where canonical IS the un-BOLTed build. + // Reads use the publishing credential, the one the pipelines publish build artifacts with. + def awsCredentialId = LLM_ARTIFACTS_LOCATIONS.aws.credential + def readNetrc = "${ws}/builds/${BOLT_TARNAME}.netrc" def tarStage = """ - URL='${llmTarfile}' - if curl -fsI '${unboltedTarUrl}' >/dev/null 2>&1; then - echo '[INFO] found unbolted-${BOLT_TARNAME}; using it as the un-BOLTed input (avoids double-BOLT on reuse)' - URL='${unboltedTarUrl}' - fi + INPUT_DIR='${ARTIFACT_PATH}' + TARNAME='${BOLT_TARNAME}' + FETCH_HELPER='${ws}/fetch_build_artifact.sh' + FETCH_ABSENT=${FETCH_ABSENT} + AWS_ROOT='${LLM_ARTIFACTS_LOCATIONS.aws.root}/' DEST='${ws}/builds/${BOLT_TARNAME}' PARTS=16 + export LLM_ARTIFACTS_PRIMARY='${primary}' + export ARTIFACTORY_NETRC='${readNetrc}' + export ARTIFACTORY_CREDENTIAL_ID='${awsCredentialId}' """.stripIndent() + boltFetchLib + ''' set -e - fetch_verified "$URL" "$DEST" gzip -t + # The netrc arrives on stdin; the trap removes it however this exits. + trap 'rm -f "$ARTIFACTORY_NETRC"' EXIT + (umask 077 && cat > "$ARTIFACTORY_NETRC") + chmod 600 "$ARTIFACTORY_NETRC" + # FETCH_ABSENT from the helper means the file is on neither location. Other + # failures (a rejected credential, or a transfer that failed after its own + # retries) end the stage with the helper's exit code; the helper has logged why. + rc=0 + RESOLVED=$(bash "$FETCH_HELPER" --exists "$INPUT_DIR/unbolted-$TARNAME") || rc=$? + if [ "$rc" -eq 0 ]; then + echo "[INFO] found unbolted-$TARNAME; using it as the un-BOLTed input (avoids double-BOLT on reuse)" + elif [ "$rc" -eq "$FETCH_ABSENT" ]; then + echo "[INFO] no unbolted-$TARNAME; using $TARNAME as the un-BOLTed input" + rc=0 + RESOLVED=$(bash "$FETCH_HELPER" --exists "$INPUT_DIR/$TARNAME") || rc=$? + if [ "$rc" -eq "$FETCH_ABSENT" ]; then + echo "[ERROR] input tarball missing: neither unbolted-$TARNAME nor $TARNAME exists on either location" >&2 + fi + fi + [ "$rc" -eq 0 ] || exit "$rc" + URL=${RESOLVED%% *} + SHA1=${RESOLVED##* } + # Credentials go to the aws location only; urm is read anonymously. + FETCH_NETRC= + case "$URL" in "$AWS_ROOT"*) FETCH_NETRC=$ARTIFACTORY_NETRC ;; esac + # The archive check, then the SHA-1 that the serving location reported. + _verify_input() { + gzip -t "$1" || return 1 + [ "$(sha1sum "$1" | cut -d' ' -f1)" = "$SHA1" ] && return 0 + echo "[ERROR] $1 does not match SHA-1 $SHA1 reported for $URL" >&2 + return 1 + } + fetch_verified "$URL" "$DEST" _verify_input '''.stripIndent() // Stage llvm-bolt into the shared ${ws}/builds/llvm (the per-node instrument // hook via BOLT_LLVM_DIR and the merge job both reuse it). The GitHub @@ -453,8 +662,20 @@ def submitProfileGen(pipeline) // is lost -- marginal, especially on a cache HIT (a quick local extract). // bashWrappedRemoteCmd: not all clusters default to bash, so wrap the scripts. stage("Bootstrap: download tarball") { - Utils.exec(pipeline, timeout: false, numRetries: 2, - script: Utils.sshUserCmd(remote, b64BashRemoteCmd(tarStage))) + Utils.copyFileToRemoteHost(pipeline, remote, readHelper, "${ws}/fetch_build_artifact.sh") + def runDownload = bashRemoteScriptFile(pipeline, remote, tarStage, + "${localScriptDir}/download_tarball.sh", "${ws}/builds/${BOLT_TARNAME}.fetch.sh") + // The credential reaches the cluster as in promoteBundle(): a netrc + // rendered from shell variables on the agent and fed over ssh stdin, + // never part of the script text. set +x keeps the printf arguments out of + // the step's trace. + def awsHost = LLM_ARTIFACTS_LOCATIONS.aws.root.replaceFirst(/^https?:\/\//, "").tokenize('/').first() + pipeline.withCredentials([pipeline.usernamePassword(credentialsId: awsCredentialId, + usernameVariable: 'ART_USER', passwordVariable: 'ART_PASS')]) { + def feed = "set +x\nprintf 'machine ${awsHost} login %s password %s\\n' \"\$ART_USER\" \"\$ART_PASS\" | " + Utils.exec(pipeline, timeout: false, numRetries: 2, noNVDFEvent: true, + script: feed + Utils.sshUserCmd(remote, runDownload)) + } } // Extract the full source tree + wheel (packs the build commit's // TensorRT-LLM/src). The perf harness runs from this checkout @@ -465,8 +686,10 @@ def submitProfileGen(pipeline) script: Utils.sshUserCmd(remote, "\"tar -xf ${ws}/builds/${BOLT_TARNAME} -C ${ws} TensorRT-LLM\"")) } stage("Bootstrap: stage llvm-bolt") { + def runLlvmStage = bashRemoteScriptFile(pipeline, remote, llvmStage, + "${localScriptDir}/stage_llvm_bolt.sh", "${ws}/builds/stage_llvm_bolt.sh") Utils.exec(pipeline, timeout: false, numRetries: 2, - script: Utils.sshUserCmd(remote, b64BashRemoteCmd(llvmStage))) + script: Utils.sshUserCmd(remote, runLlvmStage)) } // 2) Fan-out: ONE perf-sanity run per workload (Jenkins parallel{}). @@ -547,7 +770,7 @@ def submitProfileGen(pipeline) publishBoltedCanonical(pipeline, remote, "${outDir}/bolt-${BOLT_TARNAME}", // bolted build from the merge job "${ws}/builds/${BOLT_TARNAME}", // un-BOLTed input on the cluster - "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}", BOLT_TARNAME) + primary, BOLT_TARNAME, localScriptDir) } } else if (!PUBLISH_BOLTED_CANONICAL) { pipeline.echo("PUBLISH_BOLTED_CANONICAL=false: not repushing a BOLTed canonical tarball.") @@ -794,14 +1017,23 @@ def promoteBundle(pipeline, remote, String bundle) // reused artifactPath: the unbolted- snapshot is skipped if it already exists, so a // re-run never overwrites the real unbolted with an already-BOLTed canonical. // -// Credentials use the same stdin-fed netrc pattern as promoteBundle(): the secret -// travels over ssh stdin (read by `cat > netrc`), never embedded in the base64'd -// script, so it can't leak via the decoded remote command or `ps`. +// `location` is the primary location ("aws" or "urm"), which decides the target +// URL and the publisher credential. Credentials use the same stdin-fed netrc +// pattern as promoteBundle(): the secret travels over ssh stdin (read by +// `cat > netrc`), never embedded in the script, so it can't leak via the script +// file, the remote command or `ps`. The script goes to the cluster as a file +// (bashRemoteScriptFile), written first under localScriptDir on the agent. // --------------------------------------------------------------------------- -def publishBoltedCanonical(pipeline, remote, String boltedLocal, String unboltedLocal, String artifactBase, String tarName) +def publishBoltedCanonical(pipeline, remote, String boltedLocal, String unboltedLocal, String location, String tarName, + String localScriptDir) { - def host = URM_ARTIFACTORY_BASE.replaceFirst(/^https?:\/\//, "").tokenize('/').first() + def artifactBase = llmArtifactsUrl(location, ARTIFACT_PATH) + def host = artifactBase.replaceFirst(/^https?:\/\//, "").tokenize('/').first() + def credentialId = LLM_ARTIFACTS_LOCATIONS[location].credential + def repo = LLM_ARTIFACTS_LOCATIONS[location].repo + def layoutDir = llmArtifactsLayoutPath(ARTIFACT_PATH) def netrc = "${boltedLocal}.netrc" + def response = "${boltedLocal}.response" def canonicalUrl = "${artifactBase}/${tarName}" def unboltedUrl = "${artifactBase}/unbolted-${tarName}" pipeline.echo("Publishing BOLTed build as canonical ${canonicalUrl} (un-BOLTed preserved at ${unboltedUrl})") @@ -811,34 +1043,56 @@ def publishBoltedCanonical(pipeline, remote, String boltedLocal, String unbolted def publish = """ set -e umask 077 - trap 'rm -f "${netrc}"' EXIT + trap 'rm -f "${netrc}" "${response}"' EXIT test -f "${boltedLocal}" || { echo "[ERROR] bolted tarball missing: ${boltedLocal} (did the merge run with BOLT_APPLY=1?)" >&2; exit 1; } test -f "${unboltedLocal}" || { echo "[ERROR] un-BOLTed tarball missing: ${unboltedLocal}" >&2; exit 1; } cat > "${netrc}" chmod 600 "${netrc}" + # http_status CURL_ARGS...: print the final HTTP status of one request (000 + # if no response arrived). Without -f, curl retries connection errors and + # transient statuses (408, 429, 5xx) but returns a 401/403 at once. The + # body goes to a regular file: curl cannot retry a transfer whose output + # it cannot truncate, such as /dev/null. + http_status() { + curl -sS --netrc-file "${netrc}" --retry 5 --retry-all-errors --connect-timeout 30 \\ + -o "${response}" -w '%{http_code}' "\$@" || true + } + # fail_status NAME STATUS: report the failed request for NAME and exit. + fail_status() { + case "\$2" in + 401|403) echo "[ERROR] ${host} rejected the credentials for ${layoutDir}/\$1 (HTTP \$2): the Artifactory access token in credential '${credentialId}' needs renewal or lacks deploy permission on ${repo}." >&2 ;; + *) echo "[ERROR] request for ${layoutDir}/\$1 on ${host} failed (HTTP \$2)" >&2 ;; + esac + exit 1 + } # 1) Preserve the un-BOLTed build as unbolted-${tarName}, unless already there. - # A PUT-upload, NOT /api/copy: the target repo is virtual, so a server-side - # copy returns 409. Idempotent via the existence check: on a re-run the - # canonical is ALREADY the BOLTed build, and re-publishing would otherwise - # overwrite the real unbolted with a bolted copy. - if curl -fsI --netrc-file "${netrc}" "${unboltedUrl}" >/dev/null 2>&1; then - echo "[INFO] unbolted-${tarName} already exists; skipping upload" - else - echo "[INFO] uploading un-BOLTed tarball -> unbolted-${tarName}" - curl -fS --netrc-file "${netrc}" --retry 5 --retry-all-errors \\ - --connect-timeout 30 --speed-time 300 --speed-limit 1024 \\ - -T "${unboltedLocal}" "${unboltedUrl}" - fi + # A PUT-upload, NOT /api/copy: on urm the target repo is virtual, so a + # server-side copy returns 409. Idempotent via the existence check: on a + # re-run the canonical is ALREADY the BOLTed build, and re-publishing would + # otherwise overwrite the real unbolted with a bolted copy. Only a 404 + # counts as absent: a rejected credential or a failed check stops here. + status=\$(http_status -I "${unboltedUrl}") + case "\$status" in + 2??|3??) echo "[INFO] unbolted-${tarName} already exists; skipping upload" ;; + 404) + echo "[INFO] uploading un-BOLTed tarball -> unbolted-${tarName}" + status=\$(http_status --speed-time 300 --speed-limit 1024 -T "${unboltedLocal}" "${unboltedUrl}") + case "\$status" in 2??) ;; *) fail_status "unbolted-${tarName}" "\$status" ;; esac + ;; + *) fail_status "unbolted-${tarName}" "\$status" ;; + esac # 2) Overwrite canonical ${tarName} with the BOLTed build (the release name). - curl -fS --netrc-file "${netrc}" --retry 5 --retry-all-errors \\ - --connect-timeout 30 --speed-time 300 --speed-limit 1024 \\ - -T "${boltedLocal}" "${canonicalUrl}" + status=\$(http_status --speed-time 300 --speed-limit 1024 -T "${boltedLocal}" "${canonicalUrl}") + case "\$status" in 2??) ;; *) fail_status "${tarName}" "\$status" ;; esac """.stripIndent() - pipeline.withCredentials([pipeline.usernamePassword(credentialsId: 'urm-artifactory-creds', + def runPublish = bashRemoteScriptFile(pipeline, remote, publish, + "${localScriptDir}/publish_bolted_canonical.sh", "${boltedLocal}.publish.sh") + pipeline.withCredentials([pipeline.usernamePassword(credentialsId: credentialId, usernameVariable: 'ART_USER', passwordVariable: 'ART_PASS')]) { - def feed = "printf 'machine ${host} login %s password %s\\n' \"\$ART_USER\" \"\$ART_PASS\" | " + // set +x keeps the printf arguments out of the step's trace. + def feed = "set +x\nprintf 'machine ${host} login %s password %s\\n' \"\$ART_USER\" \"\$ART_PASS\" | " Utils.exec(pipeline, timeout: false, numRetries: 2, noNVDFEvent: true, - script: feed + Utils.sshUserCmd(remote, b64BashRemoteCmdStdin(publish, "${boltedLocal}.publish.sh"))) + script: feed + Utils.sshUserCmd(remote, runPublish)) } pipeline.echo("Published. canonical = ${canonicalUrl}") } @@ -949,6 +1203,11 @@ pipeline { defaultValue: "", description: "SLURM CPU partition for the merge job. The merge is 100% CPU. An empty string maps to the value 'cpu'. Set this if a cluster names its CPU partition differently." ) + string( + name: "globalVars", + defaultValue: "", + description: "JSON settings passed by the parent pipeline. This job reads llm_artifacts_primary from it: the location to read the input tarball from first and to publish the BOLTed build to. Empty for a standalone run, which uses this controller's LLM_ARTIFACTS_PRIMARY." + ) } options { skipDefaultCheckout() @@ -959,8 +1218,22 @@ pipeline { stage("BOLT Profile Generation") { steps { script { + def parentVars = trtllm_utils.updateMapWithJson(this, + [(LLM_ARTIFACTS_PRIMARY_KEY): ""], env.globalVars, "globalVars") + if (parentVars[LLM_ARTIFACTS_PRIMARY_KEY]) { + LLM_ARTIFACTS_PRIMARY_FROM_PARENT = parentVars[LLM_ARTIFACTS_PRIMARY_KEY].toString() + echo "Build artifacts: primary location ${llmArtifactsPrimary()}, passed by the parent pipeline" + } + // The cluster runs jenkins/scripts/fetch_build_artifact.sh from the + // run's commit (the branch for a standalone run without one). + // checkoutFile needs git, which the pod's default (jnlp) container + // has and alpine lacks; the pod's containers share the workspace + // volume, so submitProfileGen copies it to the cluster from there. + def readHelperDir = "${pwd(tmp: true)}/fetch-build-artifact" + trtllm_utils.checkoutFile(LLM_REPO, env.gitlabCommit ?: BRANCH, + "jenkins/scripts/fetch_build_artifact.sh", readHelperDir) container('alpine') { - def bundle = submitProfileGen(this) + def bundle = submitProfileGen(this, "${readHelperDir}/fetch_build_artifact.sh") echo("BOLT profile bundle: ${bundle}") } } diff --git a/jenkins/Build.groovy b/jenkins/Build.groovy index 608fa8f1d776..aba7df0e5be8 100644 --- a/jenkins/Build.groovy +++ b/jenkins/Build.groovy @@ -31,6 +31,66 @@ LLM_ROOT = "llm" ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" +// Build artifacts (build_info.txt, packaged builds, waive_list/) are published to the +// primary location and read from it first; LLM_ARTIFACTS_PRIMARY ("aws", the default, or "urm") +// selects it. Reads fall back to the other location only when a file is absent. +// localRepo is the local repository that the copy API works on, and credential is the identity +// that publishes to the location. +LLM_ARTIFACTS_LOCATIONS = [ + aws: [root: "https://artifactory.nvidia.com/artifactory", repo: "sw-tensorrt-llm-generic-local", + localRepo: "sw-tensorrt-llm-generic-local", credential: "aws-artifactory-credentials"], + urm: [root: "https://urm.nvidia.com/artifactory", repo: "sw-tensorrt-generic", + localRepo: "sw-tensorrt-generic-local", credential: "urm-artifactory-creds"], +] + +// The primary location a parent pipeline passed in globalVars, which this job uses instead of its +// own controller's LLM_ARTIFACTS_PRIMARY. Empty when the job was started on its own. +LLM_ARTIFACTS_PRIMARY_FROM_PARENT = "" + +def llmArtifactsPrimary() +{ + def primary = LLM_ARTIFACTS_PRIMARY_FROM_PARENT ?: env.LLM_ARTIFACTS_PRIMARY ?: "aws" + if (!(primary in ["aws", "urm"])) { + error("LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '${primary}'") + } + return primary +} + +def llmArtifactsSecondary(String primary) +{ + return primary == "aws" ? "urm" : "aws" +} + +def llmArtifactsLayoutPath(String legacyPath) +{ + def layoutPath = legacyPath.replaceFirst(/^sw-tensorrt-generic(-local)?\//, "") + if (!layoutPath.startsWith("llm-artifacts/")) { + error("Not an llm-artifacts path: ${legacyPath}") + } + return layoutPath +} + +def llmArtifactsRepoPath(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].repo}/${llmArtifactsLayoutPath(legacyPath)}" +} + +def llmArtifactsUrl(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].root}/${llmArtifactsRepoPath(location, legacyPath)}" +} + +def llmArtifactsLocalRepoPath(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].localRepo}/${llmArtifactsLayoutPath(legacyPath)}" +} + +// The host name of a location, as a netrc machine entry and curl name it. +def llmArtifactsHost(String location) +{ + return LLM_ARTIFACTS_LOCATIONS[location].root.replaceFirst(/^[a-z]+:\/\//, "").replaceFirst(/[:\/].*$/, "") +} + X86_64_TRIPLE = "x86_64-linux-gnu" AARCH64_TRIPLE = "aarch64-linux-gnu" @@ -169,6 +229,8 @@ def BOLT_PUBLISH_VARIANT_KEY = "bolt_publish_variant" def BOLT_PROFILE_BRANCH = "bolt_profile_branch" @Field def BOLT_CONSUME_TRIPLES_KEY = "bolt_consume_triples" +@Field +def LLM_ARTIFACTS_PRIMARY_KEY = "llm_artifacts_primary" def globalVars = [ (GITHUB_PR_API_URL): null, (CACHED_CHANGED_FILE_LIST): null, @@ -186,6 +248,9 @@ def globalVars = [ (BOLT_PUBLISH_VARIANT_KEY): false, (BOLT_PROFILE_BRANCH): "", (BOLT_CONSUME_TRIPLES_KEY): "", + // Pre-declared so updateMapWithJson() populates it from the parent (see + // LLM_ARTIFACTS_PRIMARY_FROM_PARENT). + (LLM_ARTIFACTS_PRIMARY_KEY): "", ] // TODO: Move common variables to an unified location @@ -338,46 +403,157 @@ def echoNodeAndGpuInfo(pipeline, stageName) pipeline.echo "HOST_NODE_NAME = ${hostNodeName} ; GPU_UUIDS = ${gpuUuids} ; STAGE_NAME = ${stageName}" } +// Adds the publishing credential of a location to the netrc file at netrcPath, which only its +// owner can read. The shell expands the secret, so it never appears in a Groovy string or a +// command line. +def writeArtifactoryNetrc(String netrcPath, String location) +{ + withCredentials([usernamePassword(credentialsId: LLM_ARTIFACTS_LOCATIONS[location].credential, + usernameVariable: 'ARTIFACTORY_USER', passwordVariable: 'ARTIFACTORY_PASSWORD')]) { + withEnv(["ARTIFACTORY_NETRC=${netrcPath}", "ARTIFACTORY_HOST=${llmArtifactsHost(location)}"]) { + sh(label: "Write Artifactory netrc", script: ''' + set +x + umask 077 + printf 'machine %s login %s password %s\\n' "$ARTIFACTORY_HOST" "$ARTIFACTORY_USER" "$ARTIFACTORY_PASSWORD" >> "$ARTIFACTORY_NETRC" + ''') + } + } +} + +// Sends a HEAD or POST request and returns its HTTP status, or "000" when no response arrived. +// curl takes the credentials for the request's host from netrcPath when one is given, and sends +// none otherwise. HEAD asks Artifactory to answer itself rather than redirect to a CDN. +def artifactoryRequestStatus(String method, String url, String netrcPath) +{ + def status = "" + withEnv([ + "ARTIFACTORY_REQUEST_METHOD=${method}", + "ARTIFACTORY_REQUEST_URL=${url}", + "ARTIFACTORY_REQUEST_NETRC=${netrcPath}", + ]) { + status = sh(label: "Artifactory ${method}", returnStdout: true, script: ''' + set +x + set -- -sS -o /dev/null -w '%{http_code}' --max-time 300 + if [ -n "$ARTIFACTORY_REQUEST_NETRC" ]; then + set -- "$@" --netrc-file "$ARTIFACTORY_REQUEST_NETRC" + fi + if [ "$ARTIFACTORY_REQUEST_METHOD" = HEAD ]; then + set -- "$@" --head -H 'X-JFrog-Download-Redirect-To: None' + else + set -- "$@" --request "$ARTIFACTORY_REQUEST_METHOD" + fi + curl "$@" "$ARTIFACTORY_REQUEST_URL" || true + ''').trim() + } + return status +} + +// Describes a failed Artifactory request. status is the HTTP status, or "000" when no response +// arrived. A 401 or 403 names the credential that was sent, if any, and the permission that the +// request needs ("read" or "deploy"). +def artifactoryFailureMessage(String location, String path, String status, boolean sentCredentials, String permission) +{ + def host = llmArtifactsHost(location) + if (status in ["401", "403"]) { + if (!sentCredentials) { + return "${host} rejected the anonymous request for ${path} (HTTP ${status})." + } + def credential = LLM_ARTIFACTS_LOCATIONS[location].credential + def localRepo = LLM_ARTIFACTS_LOCATIONS[location].localRepo + return "${host} rejected the credentials for ${path} (HTTP ${status}): the Artifactory " + + "access token in credential '${credential}' needs renewal or lacks ${permission} permission on ${localRepo}." + } + if (!status || status == "000") { + return "No response from ${host} for ${path}." + } + return "${host} answered HTTP status ${status} for ${path}." +} + +// Whether a file of an earlier build is on a location: true on HTTP 200, false on 404. Any other +// answer raises an error. aws is read with the publishing credential in netrcPath, and urm +// anonymously, as every other reader of urm does. +def cachedArtifactExists(String location, String legacyPath, String netrcPath) +{ + def sendCredentials = (location == "aws") + def url = llmArtifactsUrl(location, legacyPath) + def status = artifactoryRequestStatus("HEAD", url, sendCredentials ? netrcPath : "") + echo "Cached artifact ${url}: HTTP status ${status}" + if (status == "200") { + return true + } + if (status == "404") { + return false + } + error(artifactoryFailureMessage(location, llmArtifactsRepoPath(location, legacyPath), status, sendCredentials, "read")) +} + +// Copies an earlier build's packaged builds into this run's directory on the primary location, +// with the copy API of that instance. When the primary lacks them and the secondary holds them, +// the run rebuilds instead, because a copy cannot cross instances. def copyCachedArtifacts(stageName, reuseArtifactPath, artifacts) { def reused = false stage(stageName) { - // The copy API does not support virtual repositories. Both paths point - // at the default deployment repository behind sw-tensorrt-generic. - def sourceRoot = reuseArtifactPath.replaceFirst( - /^sw-tensorrt-generic\//, 'sw-tensorrt-generic-local/') - def targetRoot = UPLOAD_PATH.replaceFirst( - /^sw-tensorrt-generic\//, 'sw-tensorrt-generic-local/') - - catchError( - buildResult: 'SUCCESS', - stageResult: 'UNSTABLE', - catchInterruptions: false) { - withCredentials([usernamePassword( - credentialsId: 'urm-artifactory-creds', - usernameVariable: 'ART_USER', - passwordVariable: 'ART_PASS')]) { + def primary = llmArtifactsPrimary() + def secondary = llmArtifactsSecondary(primary) + def sourceRoot = reuseArtifactPath + def sourceOnSecondary = false + def netrcPath = "" + try { + catchError( + buildResult: 'SUCCESS', + stageResult: 'UNSTABLE', + catchInterruptions: false) { + sourceRoot = llmArtifactsLocalRepoPath(primary, reuseArtifactPath) + def targetRoot = llmArtifactsLocalRepoPath(primary, UPLOAD_PATH) + // The primary's credential copies. + netrcPath = "${pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" + writeArtifactoryNetrc(netrcPath, primary) + def missing = [] for (artifact in artifacts) { - def artifactPath = artifact.key - def sourcePath = "${sourceRoot}/${artifactPath}" - def targetPath = "${targetRoot}/${artifactPath}" - echo "Copying cached artifact ${sourcePath} to ${targetPath}" - withEnv([ - "ARTIFACTORY_COPY_SOURCE=${sourcePath}", - "ARTIFACTORY_COPY_TARGET=${targetPath}", - ]) { - sh( - script: '''artifactory_url="https://urm.nvidia.com/artifactory" - copy_url="${artifactory_url}/api/copy/$ARTIFACTORY_COPY_SOURCE" - curl -fsSL --max-time 300 --user "$ART_USER:$ART_PASS" --request POST \ - "${copy_url}?to=/$ARTIFACTORY_COPY_TARGET"''' - ) + def legacyPath = "${reuseArtifactPath}/${artifact.key}" + if (!cachedArtifactExists(primary, legacyPath, netrcPath)) { + missing.add(legacyPath) } } + if (missing) { + // Reading aws as the secondary needs its credential too. + if (secondary == "aws") { + writeArtifactoryNetrc(netrcPath, "aws") + } + def absent = [] + for (legacyPath in missing) { + if (!cachedArtifactExists(secondary, legacyPath, netrcPath)) { + absent.add(legacyPath) + } + } + if (absent) { + error("Cached artifacts are on neither ${primary} nor ${secondary}: ${absent.join(', ')}") + } + sourceOnSecondary = true + } else { + def copyApi = "${LLM_ARTIFACTS_LOCATIONS[primary].root}/api/copy" + for (artifact in artifacts) { + def sourcePath = "${sourceRoot}/${artifact.key}" + def targetPath = "${targetRoot}/${artifact.key}" + echo "Copying cached artifact ${sourcePath} to ${targetPath}" + def status = artifactoryRequestStatus( + "POST", "${copyApi}/${sourcePath}?to=/${targetPath}", netrcPath) + if (!(status ==~ /2\d\d/)) { + error(artifactoryFailureMessage(primary, targetPath, status, true, "deploy")) + } + } + reused = true + } + } + } finally { + if (netrcPath) { + sh(label: "Remove Artifactory netrc", script: "rm -f '${netrcPath}'") } - reused = true } - if (!reused) { + if (sourceOnSecondary) { + echo "reuse source is on ${secondary}; rebuilding" + } else if (!reused) { echo "Failed to copy cached artifacts from ${sourceRoot}; need rebuild." } } @@ -413,6 +589,27 @@ def uploadArtifacts(artifacts, prefix = UPLOAD_PATH, retryTimes = 2, serverId = } } +// Publishes the packaged builds to the primary location: UPLOAD_PATH through server +// 'Artifactory' on urm, or the same layout path in sw-tensorrt-llm-generic-local through +// 'ArtifactoryAws' on aws. Both servers must be defined first. +def uploadToPrimary(artifacts) +{ + def primary = llmArtifactsPrimary() + if (primary == "urm") { + uploadArtifacts(artifacts) + return + } + try { + uploadArtifacts(artifacts, llmArtifactsRepoPath("aws", UPLOAD_PATH), 2, 'ArtifactoryAws') + } catch (InterruptedException e) { + throw e + } catch (Exception e) { + // The plugin reports an HTTP 401 or 403 only in its own log, so name the credential here. + echo "Upload to artifactory.nvidia.com failed using credential 'aws-artifactory-credentials'. If the upload log above reports HTTP 401 or 403, the Artifactory access token in that credential needs renewal or lacks deploy permission on sw-tensorrt-llm-generic-local." + throw e + } +} + def buildOrCache(pipeline, key, reuseArtifactPath, artifacts, image, k8s_cpu, runner) { if (reuseArtifactPath) { @@ -443,7 +640,14 @@ def buildOrCache(pipeline, key, reuseArtifactPath, artifacts, image, k8s_cpu, ru // The default value (if not configured) is 300 seconds: timeout: 300 ) - uploadArtifacts(artifacts) + rtServer ( + id: 'ArtifactoryAws', + url: 'https://artifactory.nvidia.com/artifactory', + credentialsId: 'aws-artifactory-credentials', + bypassProxy: true, + timeout: 300 + ) + uploadToPrimary(artifacts) } } }) @@ -775,6 +979,10 @@ def launchStages(pipeline, cpu_arch, enableFailFast, globalVars) globalVars = trtllm_utils.updateMapWithJson(pipeline, globalVars, env.globalVars, "globalVars") globalVars = trtllm_utils.initializeCiBudget(pipeline, globalVars, 24, 'HOURS', "Build-${cpu_arch}") globalVars[ACTION_INFO] = trtllm_utils.setupPipelineDescription(pipeline, globalVars[ACTION_INFO]) + if (globalVars[LLM_ARTIFACTS_PRIMARY_KEY]) { + LLM_ARTIFACTS_PRIMARY_FROM_PARENT = globalVars[LLM_ARTIFACTS_PRIMARY_KEY].toString() + echo "Build artifacts: primary location ${llmArtifactsPrimary()}, passed by the parent pipeline" + } // BOLT consume flag: source it from globalVars (set by L0_MergeRequest as // bolt_consume_build), which is reliably propagated. The standalone // boltConsume job parameter is NOT registered on the remote build jobs, so @@ -840,6 +1048,13 @@ def launchStages(pipeline, cpu_arch, enableFailFast, globalVars) // The default value (if not configured) is 300 seconds: timeout: 300 ) + rtServer ( + id: 'ArtifactoryAws', + url: 'https://artifactory.nvidia.com/artifactory', + credentialsId: 'aws-artifactory-credentials', + bypassProxy: true, + timeout: 300 + ) def reuseArtifactPath = env.reuseArtifactPath // copyCachedArtifacts walks the artifact map as it stands BEFORE the build // runs, and bolted- is added to that map by applyLatestBolt during diff --git a/jenkins/BuildDockerImage.groovy b/jenkins/BuildDockerImage.groovy index efcdc1deff9c..b7be8797b8f5 100644 --- a/jenkins/BuildDockerImage.groovy +++ b/jenkins/BuildDockerImage.groovy @@ -33,6 +33,36 @@ withCredentials([string(credentialsId: 'default-llm-repo', variable: 'DEFAULT_LL ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" +// Build artifacts (build_info.txt, packaged builds, waive_list/) are published to the +// primary location and read from it first; LLM_ARTIFACTS_PRIMARY ("aws", the default, or "urm") +// selects it. Reads fall back to the other location only when a file is absent. This job reads +// packaged builds with jenkins/scripts/fetch_build_artifact.sh, which maps the legacy paths above. + +// The primary location a parent pipeline passed in globalVars, which this job uses instead of its +// own controller's LLM_ARTIFACTS_PRIMARY. Empty when the job was started on its own. +LLM_ARTIFACTS_PRIMARY_FROM_PARENT = "" + +def llmArtifactsPrimary() +{ + def primary = LLM_ARTIFACTS_PRIMARY_FROM_PARENT ?: env.LLM_ARTIFACTS_PRIMARY ?: "aws" + if (!(primary in ["aws", "urm"])) { + error("LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '${primary}'") + } + return primary +} + +// Credential with which the read helper reads packaged builds on artifactory.nvidia.com: the +// publishing credential, the one the pipelines publish build artifacts with. Image pushes also +// use it. +LLM_ARTIFACTS_CREDENTIALS_ID = "aws-artifactory-credentials" + +// Exit codes of jenkins/scripts/fetch_build_artifact.sh that this pipeline acts on, by the +// names the helper gives them. tests/unittest/scripts/test_fetch_build_artifact.py checks +// that they match. +FETCH_ABSENT = 2 +FETCH_REJECTED = 3 +FETCH_TRANSFER_FAILED = 4 + LLM_ROOT = "llm" LLM_BRANCH = env.gitlabBranch ?: params.branch @@ -99,6 +129,8 @@ def MAINTENANCE_ENTRIES = "maintenance_entries" def BOLT_PROFILE_REF = "bolt_profile_ref" @Field def BOLT_PROFILE_BRANCH = "bolt_profile_branch" +@Field +def LLM_ARTIFACTS_PRIMARY_KEY = "llm_artifacts_primary" def globalVars = [ (GITHUB_PR_API_URL): null, (CACHED_CHANGED_FILE_LIST): null, @@ -112,6 +144,9 @@ def globalVars = [ // without saying so. (BOLT_PROFILE_REF): "", (BOLT_PROFILE_BRANCH): "", + // Pre-declared so updateMapWithJson() populates it from the parent (see + // LLM_ARTIFACTS_PRIMARY_FROM_PARENT). + (LLM_ARTIFACTS_PRIMARY_KEY): "", ] @Field @@ -290,6 +325,24 @@ def createKubernetesPodConfig(type, arch = "amd64", build_wheel = false) } +// Writes a netrc for LLM_ARTIFACTS_CREDENTIALS_ID, readable only by its owner, to netrcPath +// for jenkins/scripts/fetch_build_artifact.sh. Callers put it under pwd(tmp: true), outside the +// source tree and any docker build context, and remove it in a finally block. +def writeArtifactoryNetrc(String netrcPath) +{ + withCredentials([usernamePassword(credentialsId: LLM_ARTIFACTS_CREDENTIALS_ID, + usernameVariable: 'ARTIFACTORY_USER', passwordVariable: 'ARTIFACTORY_PASSWORD')]) { + withEnv(["ARTIFACTORY_NETRC=${netrcPath}"]) { + sh(label: "Write Artifactory netrc", script: ''' + set +x + umask 077 + printf 'machine artifactory.nvidia.com login %s password %s\\n' "$ARTIFACTORY_USER" "$ARTIFACTORY_PASSWORD" > "$ARTIFACTORY_NETRC" + ''') + } + } +} + + def prepareWheelFromBuildStage(dockerfileStage, arch) { // Whether THIS image has to ship a BOLT-optimized wheel. Answered before the // gates below, and deliberately not subject to them. @@ -677,59 +730,76 @@ def buildImage(config, imageKeyToTag, versionOverride) build_jobs = BUILD_JOBS_RELEASE_SBSA } } - stage ("make ${target}_${action} (${arch})") { - sh "env | sort" - def randomSleep = (Math.random() * 600 + 600).toInteger() - trtllm_utils.llmExecStepWithRetry(this, script: "docker pull ${TRITON_IMAGE}:${TRITON_BASE_TAG}", sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) - try { - trtllm_utils.llmExecStepWithRetry(this, script: """ - cd ${LLM_ROOT} && make -C docker ${target}_${action} \ - BASE_IMAGE=${BASE_IMAGE} \ - TRITON_IMAGE=${TRITON_IMAGE} \ - TORCH_INSTALL_TYPE=${torchInstallType} \ - IMAGE_WITH_TAG=${rawImageTag} \ - STAGE=${dockerfileStage} \ - BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} - """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) - } catch (InterruptedException ex) { - throw ex - } catch (Exception ex) { - if (buildWheelArgs.trim().isEmpty()) { + // With buildWheelArgs, the image build takes its wheel from the packaged build, which + // jenkins/scripts/fetch_build_artifact.sh reads with this netrc. docker/Makefile + // hands the netrc to `docker buildx build` as a BuildKit secret, never a build argument; + // it is written outside the build context (${LLM_ROOT}) and removed after the builds below. + def wheelNetrcPath = buildWheelArgs.trim() ? "${pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" : "" + try { + if (wheelNetrcPath) { + def primary = llmArtifactsPrimary() + writeArtifactoryNetrc(wheelNetrcPath) + buildWheelArgs += " ARTIFACTORY_NETRC='${wheelNetrcPath}' LLM_ARTIFACTS_PRIMARY=${primary}" + + " ARTIFACTORY_CREDENTIAL_ID=${LLM_ARTIFACTS_CREDENTIALS_ID}" + } + stage ("make ${target}_${action} (${arch})") { + sh "env | sort" + def randomSleep = (Math.random() * 600 + 600).toInteger() + trtllm_utils.llmExecStepWithRetry(this, script: "docker pull ${TRITON_IMAGE}:${TRITON_BASE_TAG}", sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) + try { + trtllm_utils.llmExecStepWithRetry(this, script: """ + cd ${LLM_ROOT} && make -C docker ${target}_${action} \ + BASE_IMAGE=${BASE_IMAGE} \ + TRITON_IMAGE=${TRITON_IMAGE} \ + TORCH_INSTALL_TYPE=${torchInstallType} \ + IMAGE_WITH_TAG=${rawImageTag} \ + STAGE=${dockerfileStage} \ + BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} + """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) + } catch (InterruptedException ex) { throw ex + } catch (Exception ex) { + if (buildWheelArgs.trim().isEmpty()) { + throw ex + } + if (!mayRetryWithoutBuildStageWheel(arch)) { + echo "Build failed with wheel arguments and the BOLT-optimized wheel is required for ${arch}; " + + "NOT retrying from source (that would publish an unoptimized release image)" + throw ex + } + echo "Build failed with wheel arguments, retrying without them" + buildWheelArgs = "" + trtllm_utils.llmExecStepWithRetry(this, script: """ + cd ${LLM_ROOT} && make -C docker ${target}_${action} \ + BASE_IMAGE=${BASE_IMAGE} \ + TRITON_IMAGE=${TRITON_IMAGE} \ + TORCH_INSTALL_TYPE=${torchInstallType} \ + IMAGE_WITH_TAG=${rawImageTag} \ + STAGE=${dockerfileStage} \ + BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} + """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) } - if (!mayRetryWithoutBuildStageWheel(arch)) { - echo "Build failed with wheel arguments and the BOLT-optimized wheel is required for ${arch}; " + - "NOT retrying from source (that would publish an unoptimized release image)" - throw ex + if (target == "ngc-release") { + imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag } - echo "Build failed with wheel arguments, retrying without them" - buildWheelArgs = "" - trtllm_utils.llmExecStepWithRetry(this, script: """ - cd ${LLM_ROOT} && make -C docker ${target}_${action} \ - BASE_IMAGE=${BASE_IMAGE} \ - TRITON_IMAGE=${TRITON_IMAGE} \ - TORCH_INSTALL_TYPE=${torchInstallType} \ - IMAGE_WITH_TAG=${rawImageTag} \ - STAGE=${dockerfileStage} \ - BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} - """, sleepInSecs: randomSleep, numRetries: 6, shortCommondRunTimeMax: 7200) } - if (target == "ngc-release") { - imageKeyToTag["NGC Release Image ${config.arch}"] = imageWithTag - } - } - if (customTag) { - stage ("Custom Tag: ${customTag} (${arch})") { - sh """ - cd ${LLM_ROOT} && make -C docker ${target}_${action} \ - BASE_IMAGE=${BASE_IMAGE} \ - TRITON_IMAGE=${TRITON_IMAGE} \ - TORCH_INSTALL_TYPE=${torchInstallType} \ - IMAGE_WITH_TAG=${rawCustomTag} \ - STAGE=${dockerfileStage} \ - BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} - """ + if (customTag) { + stage ("Custom Tag: ${customTag} (${arch})") { + sh """ + cd ${LLM_ROOT} && make -C docker ${target}_${action} \ + BASE_IMAGE=${BASE_IMAGE} \ + TRITON_IMAGE=${TRITON_IMAGE} \ + TORCH_INSTALL_TYPE=${torchInstallType} \ + IMAGE_WITH_TAG=${rawCustomTag} \ + STAGE=${dockerfileStage} \ + BUILD_WHEEL_OPTS='-j ${build_jobs}' ${args} ${buildWheelArgs} + """ + } + } + } finally { + if (wheelNetrcPath) { + sh(label: "Remove Artifactory netrc", script: "rm -f '${wheelNetrcPath}'") } } @@ -1023,6 +1093,10 @@ pipeline { sh "env | sort" globalVars = trtllm_utils.updateMapWithJson(this, globalVars, env.globalVars, "globalVars") globalVars[ACTION_INFO] = trtllm_utils.setupPipelineDescription(this, globalVars[ACTION_INFO]) + if (globalVars[LLM_ARTIFACTS_PRIMARY_KEY]) { + LLM_ARTIFACTS_PRIMARY_FROM_PARENT = globalVars[LLM_ARTIFACTS_PRIMARY_KEY].toString() + echo "Build artifacts: primary location ${llmArtifactsPrimary()}, passed by the parent pipeline" + } } } } @@ -1054,11 +1128,17 @@ pipeline { script { catchError(buildResult: 'FAILURE', stageResult: 'FAILURE') { container("python3") { - // Install wget - trtllm_utils.llmExecStepWithRetry(this, script: "apt-get update && apt-get install -y wget") + // git for checkoutFile, curl for jenkins/scripts/fetch_build_artifact.sh + trtllm_utils.llmExecStepWithRetry(this, script: "apt-get update && apt-get install -y wget git curl") + + // The read helper from this run's commit looks for each file on the + // primary location first; it needs a netrc for artifactory.nvidia.com. + def helperDir = "${pwd(tmp: true)}/fetch-build-artifact" + trtllm_utils.checkoutFile(LLM_REPO, LLM_COMMIT_OR_BRANCH, "jenkins/scripts/fetch_build_artifact.sh", helperDir) + def helperPath = "${helperDir}/fetch_build_artifact.sh" + def primary = llmArtifactsPrimary() // Poll for build artifacts - def artifactBaseUrl = "https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/" def requiredFiles = [ "TensorRT-LLM-GH200.tar.gz", "TensorRT-LLM.tar.gz" @@ -1072,34 +1152,57 @@ pipeline { def startTime = System.currentTimeMillis() def maxWaitMs = maxWaitMinutes * 60 * 1000 - while ((System.currentTimeMillis() - startTime) < maxWaitMs) { - def missingFiles = [] - - for (file in requiredFiles) { - def fileUrl = "${artifactBaseUrl}${file}" - def exitCode = sh( - script: "wget --spider --quiet --timeout=30 --tries=1 '${fileUrl}'", - returnStatus: true - ) - - if (exitCode != 0) { - missingFiles.add(file) + def netrcPath = "${pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" + try { + writeArtifactoryNetrc(netrcPath) + withEnv([ + "LLM_ARTIFACTS_PRIMARY=${primary}", + "ARTIFACTORY_NETRC=${netrcPath}", + "ARTIFACTORY_CREDENTIAL_ID=${LLM_ARTIFACTS_CREDENTIALS_ID}", + ]) { + while ((System.currentTimeMillis() - startTime) < maxWaitMs) { + def missingFiles = [] + + for (file in requiredFiles) { + // FETCH_ABSENT: not published yet. FETCH_TRANSFER_FAILED: the check + // failed after the helper's own retries. Both are checked again on + // the next poll. + def exitCode = sh( + script: "bash '${helperPath}' --exists '${UPLOAD_PATH}/${file}' > /dev/null", + returnStatus: true + ) + + if (exitCode == FETCH_REJECTED) { + // The helper has printed which request was rejected and why. + error(primary == "aws" ? + "Credential '${LLM_ARTIFACTS_CREDENTIALS_ID}' was rejected while checking for ${file}: its Artifactory access token needs renewal or lacks read permission on sw-tensorrt-llm-generic-local" : + "urm rejected the anonymous request while checking for ${file}") + } + if (!(exitCode in [0, FETCH_ABSENT, FETCH_TRANSFER_FAILED])) { + error "fetch_build_artifact.sh exited with ${exitCode} while checking for ${file}" + } + if (exitCode != 0) { + missingFiles.add(file) + } + } + + if (missingFiles.isEmpty()) { + echo "All build artifacts are ready!" + return + } + + def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) + echo "Waiting... (${elapsedMinutes.intValue()} minutes elapsed)" + echo "Missing files: ${missingFiles}" + sleep(pollIntervalSeconds) } - } - if (missingFiles.isEmpty()) { - echo "All build artifacts are ready!" - return + def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) + error "Timeout waiting for build artifacts (${elapsedMinutes.intValue()} minutes)" } - - def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) - echo "Waiting... (${elapsedMinutes.intValue()} minutes elapsed)" - echo "Missing files: ${missingFiles}" - sleep(pollIntervalSeconds) + } finally { + sh(label: "Remove Artifactory netrc", script: "rm -f '${netrcPath}'") } - - def elapsedMinutes = (System.currentTimeMillis() - startTime) / (60 * 1000) - error "Timeout waiting for build artifacts (${elapsedMinutes.intValue()} minutes)" } } } @@ -1115,6 +1218,7 @@ pipeline { script { catchError(buildResult: 'FAILURE', stageResult: 'FAILURE') { globalVars[IMAGE_KEY_TO_TAG] = imageKeyToTag + globalVars[LLM_ARTIFACTS_PRIMARY_KEY] = llmArtifactsPrimary() String globalVarsJson = writeJSON returnText: true, json: globalVars def parameters = getCommonParameters() parameters += [ diff --git a/jenkins/L0_MergeRequest.groovy b/jenkins/L0_MergeRequest.groovy index 7b10f964642e..8ac07c94f3c7 100644 --- a/jenkins/L0_MergeRequest.groovy +++ b/jenkins/L0_MergeRequest.groovy @@ -52,6 +52,42 @@ SCAN_ROOT = "scan" ARTIFACT_PATH = env.artifactPath ? env.artifactPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifacts/${JOB_NAME}/${BUILD_NUMBER}" +// Build artifacts (build_info.txt, packaged builds, waive_list/) are published to the +// primary location and read from it first; LLM_ARTIFACTS_PRIMARY ("aws", the default, or "urm") +// selects it. Reads fall back to the other location only when a file is absent. +LLM_ARTIFACTS_LOCATIONS = [ + aws: [root: "https://artifactory.nvidia.com/artifactory", repo: "sw-tensorrt-llm-generic-local"], + urm: [root: "https://urm.nvidia.com/artifactory", repo: "sw-tensorrt-generic"], +] + +def llmArtifactsPrimary() +{ + def primary = env.LLM_ARTIFACTS_PRIMARY ?: "aws" + if (!(primary in ["aws", "urm"])) { + error("LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '${primary}'") + } + return primary +} + +def llmArtifactsLayoutPath(String legacyPath) +{ + def layoutPath = legacyPath.replaceFirst(/^sw-tensorrt-generic(-local)?\//, "") + if (!layoutPath.startsWith("llm-artifacts/")) { + error("Not an llm-artifacts path: ${legacyPath}") + } + return layoutPath +} + +def llmArtifactsRepoPath(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].repo}/${llmArtifactsLayoutPath(legacyPath)}" +} + +def llmArtifactsUrl(String location, String legacyPath) +{ + return "${LLM_ARTIFACTS_LOCATIONS[location].root}/${llmArtifactsRepoPath(location, legacyPath)}" +} + // GitLab project id for TRT-LLM repo, used to set GitLab status. Only needed when triggering by timer or by hand. withCredentials([string(credentialsId: 'gitlab-llm-repo-id', variable: 'GITLAB_PROJECT_ID')]) { GITLAB_PROJECT_ID = env.gitlabProjectId ? env.gitlabProjectId : "${GITLAB_PROJECT_ID}" @@ -332,6 +368,10 @@ def BOLT_CONSUME_TRIPLES = "bolt_consume_triples" def RELEASE_TARGET = "release_target" @Field def JOB_TYPE = "job_type" +// The primary location of build artifacts (see llmArtifactsPrimary), passed to every child job. +// Child jobs run on other Jenkins controllers, whose own LLM_ARTIFACTS_PRIMARY can differ. +@Field +def LLM_ARTIFACTS_PRIMARY_KEY = "llm_artifacts_primary" def globalVars = [ (GITHUB_PR_API_URL): gitlabParamsFromBot.get('github_pr_api_url', null), (BASE_COMMIT): gitlabParamsFromBot.get('base_commit', null), @@ -637,8 +677,12 @@ def mergeWaiveList(pipeline, globalVars) --diff-file=diff_content.txt \ --output-file=waives.txt """ - trtllm_utils.uploadArtifacts("waives*.txt", "${UPLOAD_PATH}/waive_list/") - echo "New merged test waive list: https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/waive_list/waives.txt" + def primary = llmArtifactsPrimary() + def waiveListTarget = primary == "aws" ? + "${llmArtifactsRepoPath('aws', UPLOAD_PATH)}/waive_list/" : "${UPLOAD_PATH}/waive_list/" + trtllm_utils.uploadArtifacts("waives*.txt", waiveListTarget, primary) + def waiveListUrl = llmArtifactsUrl(primary, "${UPLOAD_PATH}/waive_list/waives.txt") + echo "New merged test waive list: ${waiveListUrl}" } catch (InterruptedException e) { throw e } catch (Exception e) { @@ -743,8 +787,12 @@ def preparation(pipeline, testFilter, globalVars) "date=${new Date().format('yyyy-MM-dd HH:mm:ss z', TimeZone.getTimeZone('UTC'))}\n" + "jenkins_url=${env.BUILD_URL}" writeFile file: 'build_info.txt', text: buildInfo - trtllm_utils.uploadArtifacts("build_info.txt", "${UPLOAD_PATH}/") - pipeline.echo "Build info: https://urm.nvidia.com/artifactory/${UPLOAD_PATH}/build_info.txt" + def primary = llmArtifactsPrimary() + def buildInfoTarget = primary == "aws" ? + "${llmArtifactsRepoPath('aws', UPLOAD_PATH)}/" : "${UPLOAD_PATH}/" + trtllm_utils.uploadArtifacts("build_info.txt", buildInfoTarget, primary) + def buildInfoUrl = llmArtifactsUrl(primary, "${UPLOAD_PATH}/build_info.txt") + pipeline.echo "Build info: ${buildInfoUrl}" } catch (InterruptedException e) { throw e } catch (Exception e) { @@ -1444,6 +1492,7 @@ def _cbtsCoverageDecline(boolean pilotEligible, String reason) def _cbtsCoverageAudit(pipeline, globalVars, String residualPath) { def pilotEligible = false + def artifactoryNetrc = "" try { // artifact.py resolves, downloads and merges the x86/SBSA DBs; paths come back // ${LLM_ROOT}-relative, matching the main.py caller's `cd ${LLM_ROOT}`. @@ -1456,6 +1505,30 @@ def _cbtsCoverageAudit(pipeline, globalVars, String residualPath) pilotEligible, "PR identity unavailable for coverage DB pin") } + // artifact.py reads build_info.txt with jenkins/scripts/fetch_build_artifact.sh, which + // looks in the primary location first and authenticates to aws with this netrc. It holds + // the publishing credential, the one the pipelines publish build artifacts with. + // When the credential is rejected the helper prints a message naming it, artifact.py + // exits 1, and the coverage tier declines like after any other failure. + def artifactoryCredentialId = 'aws-artifactory-credentials' + def artifactsPrimary = llmArtifactsPrimary() + artifactoryNetrc = "${pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" + withCredentials([usernamePassword(credentialsId: artifactoryCredentialId, + usernameVariable: 'ARTIFACTORY_USER', passwordVariable: 'ARTIFACTORY_PASSWORD')]) { + withEnv(["ARTIFACTORY_NETRC=${artifactoryNetrc}"]) { + sh(label: "Write Artifactory netrc", script: ''' + set +x + umask 077 + printf 'machine artifactory.nvidia.com login %s password %s\\n' "$ARTIFACTORY_USER" "$ARTIFACTORY_PASSWORD" > "$ARTIFACTORY_NETRC" + ''') + } + } + def artifactEnv = [ + "LLM_ARTIFACTS_PRIMARY=${artifactsPrimary}", + "ARTIFACTORY_NETRC=${artifactoryNetrc}", + "ARTIFACTORY_CREDENTIAL_ID=${artifactoryCredentialId}", + ] + def readyJson = "" def prAuthor = "" def pinPlanJson = "" @@ -1471,13 +1544,15 @@ def _cbtsCoverageAudit(pipeline, globalVars, String residualPath) ).trim() pilotEligible = prAuthor && CBTS_COVERAGE_PILOT_USERS.any { it.equalsIgnoreCase(prAuthor) } pipeline.echo("CBTS coverage pilot: pr_author=${prAuthor ?: 'unknown'}, eligible=${pilotEligible}") - pinPlanJson = sh( - script: "cd ${LLM_ROOT} && PYTHONPATH=jenkins/scripts/cbts " + - "python3 -m cbts.command coverage selection artifact " + - "--resolve-pin cbts_db_pin.json --pr-number ${prNumber} " + - "--pr-head ${prHead} || true", - returnStdout: true, - ).trim() + withEnv(artifactEnv) { + pinPlanJson = sh( + script: "cd ${LLM_ROOT} && PYTHONPATH=jenkins/scripts/cbts " + + "python3 -m cbts.command coverage selection artifact " + + "--resolve-pin cbts_db_pin.json --pr-number ${prNumber} " + + "--pr-head ${prHead} || true", + returnStdout: true, + ).trim() + } } if (!pinPlanJson) { pipeline.echo("CBTS audit: coverage DB pin could not be resolved") @@ -1511,15 +1586,17 @@ def _cbtsCoverageAudit(pipeline, globalVars, String residualPath) usernameVariable: 'NOT_USED_YET', passwordVariable: 'GITHUB_API_TOKEN'), ]) { - readyJson = sh( - script: "cd ${LLM_ROOT} && PYTHONPATH=jenkins/scripts/cbts " + - "python3 -m cbts.command coverage selection artifact " + - "--prepare cbts_cov --paths-json cbts_coverage_residual.json" + - " --pr-head ${prHead}" + - " --build ${pinPlan.build}" + - " --expected-commit ${pinPlan.commit} || true", - returnStdout: true, - ).trim() + withEnv(artifactEnv) { + readyJson = sh( + script: "cd ${LLM_ROOT} && PYTHONPATH=jenkins/scripts/cbts " + + "python3 -m cbts.command coverage selection artifact " + + "--prepare cbts_cov --paths-json cbts_coverage_residual.json" + + " --pr-head ${prHead}" + + " --build ${pinPlan.build}" + + " --expected-commit ${pinPlan.commit} || true", + returnStdout: true, + ).trim() + } } if (!readyJson) { pipeline.echo("CBTS audit: no coverage DB could be prepared — Tier 2 decision unavailable") @@ -1539,6 +1616,10 @@ def _cbtsCoverageAudit(pipeline, globalVars, String residualPath) } catch (Exception e) { pipeline.echo("CBTS audit: skipped (non-fatal): ${e.message}") return [db: null, pilotEligible: pilotEligible] + } finally { + if (artifactoryNetrc) { + sh(label: "Remove Artifactory netrc", script: "rm -f '${artifactoryNetrc}'") + } } } @@ -2347,6 +2428,7 @@ def launchJob(pipeline, jobName, reuseBuild, enableFailFast, globalVars, platfor // In particular, CACHED_CHANGED_FILE_LIST can become very large and may // trigger "Argument list too long" when passed to downstream jobs. def globalVarsToPass = globalVars.findAll { key, value -> key != CACHED_CHANGED_FILE_LIST } + globalVarsToPass[LLM_ARTIFACTS_PRIMARY_KEY] = llmArtifactsPrimary() String globalVarsJson = writeJSON returnText: true, json: globalVarsToPass parameters += [ 'enableFailFast': enableFailFast, diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 86d2137085f6..eb6953ab34eb 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -50,6 +50,32 @@ UPLOAD_PATH = env.uploadPath ? env.uploadPath : "sw-tensorrt-generic/llm-artifac URM_ARTIFACTORY_BASE = "https://urm.nvidia.com/artifactory" ENABLE_UPLOAD_TEST_RESULTS = params.enableUploadTestResults != null ? params.enableUploadTestResults : true +// The primary location a parent pipeline passed in globalVars, which this job uses instead of its +// own controller's LLM_ARTIFACTS_PRIMARY. Empty when the job was started on its own. +LLM_ARTIFACTS_PRIMARY_FROM_PARENT = "" + +// Build artifacts (build_info.txt, packaged builds, waive_list/) are published to the +// primary location and read from it first; LLM_ARTIFACTS_PRIMARY ("aws", the default, or "urm") +// selects it. Reads fall back to the other location only when a file is absent. +def llmArtifactsPrimary() +{ + def primary = LLM_ARTIFACTS_PRIMARY_FROM_PARENT ?: env.LLM_ARTIFACTS_PRIMARY ?: "aws" + if (!(primary in ["aws", "urm"])) { + error("LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '${primary}'") + } + return primary +} + +// "Setup Environment" stashes jenkins/scripts/fetch_build_artifact.sh from the run's commit under +// this name, because test stages read build artifacts before they have a source tree. +FETCH_BUILD_ARTIFACT_STASH = "fetch-build-artifact" + +// Exit codes of jenkins/scripts/fetch_build_artifact.sh that this pipeline acts on, by the +// names the helper gives them. tests/unittest/scripts/test_fetch_build_artifact.py checks +// that they match. +FETCH_ABSENT = 2 +FETCH_TRANSFER_FAILED = 4 + X86_64_TRIPLE = "x86_64-linux-gnu" AARCH64_TRIPLE = "aarch64-linux-gnu" @@ -70,6 +96,9 @@ ARTIFACTORY_IMAGE_PULL_SECRET = "trtllm-artifactory" ARTIFACTORY_DOCKER_HOST = "artifactory.nvidia.com" // Read-only artifactory pull credentials ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials" +// Credential with which the read helper reads build artifacts on artifactory.nvidia.com: the +// publishing credential, the one the pipelines publish build artifacts with. +LLM_ARTIFACTS_CREDENTIALS_ID = "aws-artifactory-credentials" // DLFW torch image DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.09-py3" @@ -906,15 +935,17 @@ def cleanUpNodeResources(def pipeline, SlurmCluster cluster, String clusterName, // login node. If that pod dies mid-run (eviction, container error, agent // offline), the in-pod cleanup can no longer reach the controller, so the SLURM // job and any Jenkins agent node leak. State + reconciliation live in the shared -// `resourceLedger` (trtllm-jenkins-shared-lib). Each SLURM stage registers two +// `resourceLedger` (trtllm-jenkins-shared-lib). Each SLURM stage registers // sibling entries -- "/dispatcher-pod" (the pod spec, from the pod -// wrapper) and "/slurm" (the per-attempt SLURM job / Jenkins node, from -// the stage body) -- so freeing the job/node after a successful in-pod cleanup -// (markReclaimed) never disturbs the pod spec a *later* attempt's off-pod -// reconciliation needs. The ledger stores only serializable primitives (it -// rejects anything else) so pipeline persistence is unaffected; the SlurmCluster -// is rebuilt from clusterName at reconcile time. The functions below are thin -// SLURM-specific adapters over that generic ledger. +// wrapper) and one "/slurm/" per attempt (that +// attempt's SLURM job / Jenkins node, from the stage body) -- so freeing the +// job/node after a successful in-pod cleanup (markReclaimed) never disturbs the +// pod spec a *later* attempt's off-pod reconciliation needs, and an attempt +// whose cleanup failed stays registered beside the next one. The ledger stores +// only serializable primitives (it rejects anything else) so pipeline +// persistence is unaffected; the SlurmCluster is rebuilt from clusterName at +// reconcile time. The functions below are thin SLURM-specific adapters over +// that generic ledger. // // Shared-library contract used below (trtllm-jenkins-shared-lib), documented here // for equivalence with the pre-shared inlined finalizer this replaced: @@ -936,8 +967,9 @@ def cleanUpNodeResources(def pipeline, SlurmCluster cluster, String clusterName, // Register a live SLURM resource. The pod wrapper passes a podSpec (recorded // under "/dispatcher-pod"); the stage body passes the SLURM job / Jenkins -// node identity (recorded under "/slurm"). Keyed separately so they are -// reclaimed independently. Absent fields are skipped by the ledger. +// node identity (recorded under "/slurm/", one entry +// per attempt). Keyed separately so they are reclaimed independently. Absent +// fields are skipped by the ledger. void registerSlurmResource(String stageName, Map fields) { if (!stageName) { return @@ -947,7 +979,7 @@ void registerSlurmResource(String stageName, Map fields) { fields: [podSpec: fields.podSpec, containerName: fields.containerName], ownerBuildTag: env.BUILD_TAG) } else { - resourceLedger.register(id: "${stageName}/slurm", + resourceLedger.register(id: "${stageName}/slurm/${fields.jobUID ?: fields.nodeName}", type: fields.usedSbatch ? 'slurmJob' : 'slurmNode', fields: [clusterName: fields.clusterName, nodeName: fields.nodeName, jobUID: fields.jobUID, slurmJobId: fields.slurmJobId, usedSbatch: fields.usedSbatch], @@ -955,14 +987,15 @@ void registerSlurmResource(String stageName, Map fields) { } } -// Called once an attempt's resources are actually torn down: drop the per-attempt -// SLURM job/node entry. The stage's dispatcher-pod entry is a separate id and is -// left intact (a later attempt's off-pod reconciliation still needs its pod spec). -void deregisterSlurmResource(String stageName) { +// Called once an attempt's resources are actually torn down: drop that attempt's +// SLURM job/node entry (`attempt` is its jobUID or nodeName). The stage's +// dispatcher-pod entry is a separate id and is left intact (a later attempt's +// off-pod reconciliation still needs its pod spec). +void deregisterSlurmResource(String stageName, String attempt) { if (!stageName) { return } - resourceLedger.markReclaimed("${stageName}/slurm") + resourceLedger.markReclaimed("${stageName}/slurm/${attempt}") } // Reclaim one orphaned SLURM entry off the (dead) dispatcher pod: launch a fresh @@ -986,7 +1019,7 @@ def reconcileSlurmResource(pipeline, def entry, def podSpecOverride = null) { if (!entry.jobUID && !entry.nodeName) { return true } - def podEntry = resourceLedger.get(((entry.id ?: "") as String).replaceFirst(/\/slurm$/, "/dispatcher-pod")) + def podEntry = resourceLedger.get(((entry.id ?: "") as String).replaceFirst(/\/slurm\/[^\/]+$/, "/dispatcher-pod")) def podSpec = podSpecOverride ?: podEntry?.podSpec def containerName = podEntry?.containerName ?: "trt-llm" def cluster = entry.clusterName ? SlurmConfig.clusterConfig[entry.clusterName] : null @@ -1014,15 +1047,15 @@ def reconcileSlurmResource(pipeline, def entry, def podSpecOverride = null) { } } -// Reconcile this stage's orphaned SLURM job/node off-pod (in-catch, on a detected -// dispatcher-pod death). podSpecOverride is the failing attempt's pod spec so -// reconciliation never depends on the sibling entry surviving. +// Reconcile this stage's orphaned SLURM jobs/nodes off-pod, every attempt's (in-catch, +// on a detected dispatcher-pod death). podSpecOverride is the failing attempt's pod +// spec so reconciliation never depends on the sibling entry surviving. def finalizeOrphanedSlurmResource(pipeline, String stageName, def podSpecOverride = null) { if (!stageName) { return } resourceLedger.reconcile(pipeline, - { it.id == "${stageName}/slurm" }, + { (it.id as String).startsWith("${stageName}/slurm/") }, { p, entry -> reconcileSlurmResource(p, entry, podSpecOverride) }) } @@ -1539,7 +1572,7 @@ def runLLMTestlistWithAgent(pipeline, platform, testList, config=VANILLA_CONFIG, } // Reaching here means cleanup succeeded within its retry budget. If all // attempts fail, the throw skips deregistration and leaves the entry for sweep. - deregisterSlurmResource(stageName) + deregisterSlurmResource(stageName, nodeName) } } } @@ -1760,7 +1793,87 @@ def getMountListForSlurmTest(SlurmCluster cluster, boolean useSbatch = false) return mounts } -// Which build of this stage tests. +// Writes a netrc for artifactory.nvidia.com with LLM_ARTIFACTS_CREDENTIALS_ID, readable only by its +// owner. The shell's printf builtin moves the secret from the bound variables into the file, so it +// never appears on a command line or in the trace. Callers remove the file in a finally block. +def writeArtifactoryNetrc(pipeline, String netrcPath) +{ + pipeline.withCredentials([pipeline.usernamePassword(credentialsId: LLM_ARTIFACTS_CREDENTIALS_ID, + usernameVariable: 'ARTIFACTORY_USER', passwordVariable: 'ARTIFACTORY_PASSWORD')]) { + pipeline.withEnv(["ARTIFACTORY_NETRC=${netrcPath}"]) { + pipeline.sh(label: "Write Artifactory netrc", script: ''' + set +x + umask 077 + printf 'machine artifactory.nvidia.com login %s password %s\\n' "$ARTIFACTORY_USER" "$ARTIFACTORY_PASSWORD" > "$ARTIFACTORY_NETRC" + ''') + } + } +} + +// Returns the path of the read helper, jenkins/scripts/fetch_build_artifact.sh from the run's +// commit, in the node's temporary directory, which workspace resets and package extraction leave +// alone. Unstashes it there the first time. +def readHelperPath(pipeline) +{ + def helperDir = "${pipeline.pwd(tmp: true)}/fetch-build-artifact" + def helper = "${helperDir}/fetch_build_artifact.sh" + if (!pipeline.fileExists(helper)) { + pipeline.dir(helperDir) { + pipeline.unstash(FETCH_BUILD_ARTIFACT_STASH) + } + } + return helper +} + +// Reads one in-scope build artifact with the read helper, which tries the primary location first +// and the other one only when the file is absent there, and checks each download against the +// SHA-1 that the serving location reports. The helper's netrc lives in the node's temporary +// directory beside the helper, and is removed before this returns. +// +// "download" writes dest and returns null. "exists" writes nothing and returns " " of +// the location that holds the file. Both fail the step once the helper's and +// llmExecStepWithRetry's retries are spent. "properties" writes the properties JSON to dest and +// returns the helper's exit status instead of failing: 0 when found, otherwise one of the exit +// codes the helper names, such as FETCH_ABSENT when neither location has properties for it. +def fetchBuildArtifact(pipeline, String artifactPath, String dest, String mode = "download") +{ + if (!(mode in ["download", "exists", "properties"])) { + pipeline.error("fetchBuildArtifact: unknown mode '${mode}'") + } + def primary = llmArtifactsPrimary() + def helper = readHelperPath(pipeline) + def command = "bash '${helper}'" + (mode == "download" ? "" : " --${mode}") + " '${artifactPath}'" + + (mode == "exists" ? "" : " '${dest}'") + def netrcPath = "${pipeline.pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" + try { + writeArtifactoryNetrc(pipeline, netrcPath) + return pipeline.withEnv([ + "LLM_ARTIFACTS_PRIMARY=${primary}", + "ARTIFACTORY_NETRC=${netrcPath}", + "ARTIFACTORY_CREDENTIAL_ID=${LLM_ARTIFACTS_CREDENTIALS_ID}", + ]) { + if (mode == "properties") { + // The helper retries transient errors itself. + return pipeline.sh(label: "Fetch properties of ${artifactPath}", script: command, returnStatus: true) + } + def stdout = trtllm_utils.llmExecStepWithRetry(pipeline, label: "Fetch ${artifactPath}", script: command, + returnStdout: mode == "exists") + if (mode == "download") { + return null + } + def location = stdout.trim() + if (!(location ==~ /\S+ [0-9a-f]{40}/)) { + pipeline.error("fetchBuildArtifact: unexpected location '${location}' for ${artifactPath}") + } + return location + } + } finally { + pipeline.sh(label: "Remove Artifactory netrc", script: "rm -f '${netrcPath}'") + } +} + +// Which build of this stage tests, as the artifact path to read it +// from with fetchBuildArtifact. // // Deterministic, not discovered. The pipeline already decided whether the build // published a BOLTed variant, so there is nothing to probe for: fetch @@ -1773,35 +1886,45 @@ def getMountListForSlurmTest(SlurmCluster cluster, boolean useSbatch = false) // by name proves only that something is there; the property proves it is the // right something, which is what keeps the tested build and the released wheel // on one profile set. Absent property means an artifact predating the label: -// reported and allowed rather than failed. -def boltedTarUrl(String config, String tarName) +// reported and allowed rather than failed, as is a property lookup that keeps +// failing. A rejected credential fails the stage. +def boltedTarArtifactPath(pipeline, String config, String tarName) { - def base = "${URM_ARTIFACTORY_BASE}/${ARTIFACT_PATH}" // Only the SBSA build is BOLTed, so only its tarball has a bolted- variant. // Asking for one on any other config would 404 the stage rather than fail // honestly, so the arch gate has to match the build's exactly. if (!BOLT_TEST_BOLTED_TARBALL || config != LINUX_AARCH64_CONFIG) { - return "${base}/${tarName}" + return "${ARTIFACT_PATH}/${tarName}" } if (BOLT_PINNED_REF) { - def propsUrl = "${URM_ARTIFACTORY_BASE}/api/storage/${ARTIFACT_PATH}/bolted-${tarName}?properties" - // -L: Artifactory redirects metadata reads, and without it curl returns - // the redirect rather than the document. - def got = sh(returnStdout: true, script: """ - curl -fsSL --retry 3 --retry-all-errors --connect-timeout 30 '${propsUrl}' 2>/dev/null \\ - | tr -d ' \\n' \\ - | sed -n 's/.*"bolt\\.ref":\\["\\([^"]*\\)"\\].*/\\1/p' || true - """).trim() + def propsFile = "${pipeline.pwd(tmp: true)}/bolted-${tarName}.properties.json" + def propsStatus = fetchBuildArtifact(pipeline, "${ARTIFACT_PATH}/bolted-${tarName}", propsFile, "properties") + // FETCH_ABSENT: no properties on either location. FETCH_TRANSFER_FAILED: the lookup kept + // failing. + if (!(propsStatus in [0, FETCH_ABSENT, FETCH_TRANSFER_FAILED])) { + error("[BOLT] Could not read the properties of bolted-${tarName}: " + + "fetch_build_artifact.sh exited ${propsStatus} (see its message above).") + } + def got = "" + if (propsStatus == 0) { + got = sh(returnStdout: true, script: """ + tr -d ' \\n' < '${propsFile}' \\ + | sed -n 's/.*"bolt\\.ref":\\["\\([^"]*\\)"\\].*/\\1/p' + rm -f '${propsFile}' + """).trim() + } if (got && got != BOLT_PINNED_REF) { error("[BOLT] bolted-${tarName} carries bolt.ref=${got}, but this pipeline pinned ${BOLT_PINNED_REF}. " + "Refusing to test a build optimized with a different profile bundle than the release artifacts use.") } - if (!got) { + if (propsStatus == FETCH_TRANSFER_FAILED) { + echo "[BOLT] the bolt.ref lookup for bolted-${tarName} failed (fetch_build_artifact.sh exited ${propsStatus}, see its message above); proceeding on pin ${BOLT_PINNED_REF}" + } else if (!got) { echo "[BOLT] bolted-${tarName} carries no bolt.ref (artifact predates the label); proceeding on pin ${BOLT_PINNED_REF}" } } echo "[BOLT] testing the BOLTed build: bolted-${tarName}" - return "${base}/bolted-${tarName}" + return "${ARTIFACT_PATH}/bolted-${tarName}" } def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG, perfMode=false, stageName="Undefined", splitId=1, splits=1, gpuCount=1, nodeCount=1, skipInstallWheel=false, cpver="cp312", String postTag="", boolean useClusterDurations=false, Map placementContext=null, Map retryContext=null) @@ -1847,7 +1970,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG ]) { CloudManager.withSlurmFrontendFailover(pipeline, partition.clusterName, cluster) { remote -> def tarName = BUILD_CONFIGS[config][TARNAME] - def llmTarfile = boltedTarUrl(config, tarName) + def llmTarArtifactPath = boltedTarArtifactPath(pipeline, config, tarName) def llmPath = sh (script: "realpath .", returnStdout: true).trim() def resourcePathNode = "/tmp" def llmSrcNode = "${resourcePathNode}/TensorRT-LLM/src" @@ -1858,6 +1981,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG def scriptInstallPathNode = "${jobWorkspace}/${jobUID}-slurm_install.sh" def scriptBashUtilsLocalPath = "${llmSrcLocal}/jenkins/scripts/bash_utils.sh" def scriptBashUtilsPathNode = "${jobWorkspace}/${jobUID}-bash_utils.sh" + def scriptFetchHelperPathNode = "${jobWorkspace}/${jobUID}-fetch_build_artifact.sh" def testListPathNode = "${jobWorkspace}/${testList}.txt" def waivesListPathNode = "${jobWorkspace}/waives.txt" def waivesListPathLocal = infraDryRun @@ -1875,6 +1999,9 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG def s3SecretKeyPathNode = "${jobWorkspace}/s3_secret_key" // Per-job enroot config (never shared ~/.config/enroot — that races across jobs). def enrootConfigDirNode = "${jobWorkspace}/enroot-config" + // Per-job netrc for the read helper that slurm_install.sh runs. It is not kept in + // enrootConfigDirNode, which the ENROOT srun prologue deletes before slurm_install.sh runs. + def artifactoryNetrcPathNode = "${jobWorkspace}/artifactory-netrc" def coverageConfigFile = "${jobWorkspace}/.coveragerc" stage("Initialize Test") { @@ -1882,11 +2009,16 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG // Create Job Workspace folder in Frontend Node Utils.exec(pipeline, script: Utils.sshUserCmd(remote, "\"mkdir -p ${jobWorkspace}\""), numRetries: 3) - // Download and Unzip Tar File + // Download and Unzip Tar File. The node downloads it again itself, with the read + // helper. The URL that --exists resolves is for the older slurm_install.sh inside a + // package built before the helper existed, which downloads that URL instead. + def llmTarfileLocation = "" timeout(time: 30, unit: 'MINUTES') { - trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && wget -nv -O '${BUILD_CONFIGS[config][TARNAME]}' '${llmTarfile}'") + llmTarfileLocation = fetchBuildArtifact(pipeline, llmTarArtifactPath, null, "exists") + fetchBuildArtifact(pipeline, llmTarArtifactPath, "${llmPath}/${tarName}") } sh "cd ${llmPath} && tar -zxf ${BUILD_CONFIGS[config][TARNAME]}" + def llmTarfile = llmTarfileLocation.tokenize(" ")[0] Utils.exec(pipeline, script: "echo \"Script for Slurm srun job to submit: \" && cat ${scriptRunLocalPath}") Utils.copyFileToRemoteHost( @@ -1913,6 +2045,16 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG scriptBashUtilsPathNode, true ) + // The read helper from the run's commit, which slurm_install.sh runs to download + // the packaged build. + def scriptFetchHelperLocalPath = readHelperPath(pipeline) + Utils.copyFileToRemoteHost( + pipeline, + remote, + scriptFetchHelperLocalPath, + scriptFetchHelperPathNode, + true + ) // Generate Test List and Upload to Frontend Node def makoArgs = getMakoArgsFromStageName(stageName, true) @@ -2000,6 +2142,27 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG } } + // Record this attempt's job before the token reaches its workspace, so that + // the off-pod cleanup removes the workspace even if this pod dies before the + // job is submitted. The registration after submission adds the job ID to the + // same entry; an earlier attempt's entry stays for its own cleanup. + registerSlurmResource(stageName, [clusterName: partition.clusterName, jobUID: jobUID, usedSbatch: true]) + // Every job gets the helper's netrc: the read helper on the node contacts + // artifactory.nvidia.com whenever aws is the primary, wherever the package is. + def artifactoryNetrcPathLocal = "${pipeline.pwd(tmp: true)}/artifactory-netrc-${UUID.randomUUID().toString()}" + try { + writeArtifactoryNetrc(pipeline, artifactoryNetrcPathLocal) + // scp -p keeps the netrc's 600 mode on the node. + Utils.copyFileToRemoteHost( + pipeline, + remote, + artifactoryNetrcPathLocal, + artifactoryNetrcPathNode + ) + } finally { + sh(label: "Remove Artifactory netrc", script: "rm -f '${artifactoryNetrcPathLocal}'") + } + // Generate .coveragerc: CBTS stages render coveragerc.template; all other stages get an empty rcfile (no coverage). if (isCbtsStage(stageName)) { // @TRTLLM_WHEEL_PATH@ stays a placeholder; slurm_run.sh substitutes it on the worker. @@ -2234,6 +2397,10 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG echo "Starting Slurm job \$SLURM_JOB_ID on \$SLURM_NODELIST" export jobWorkspace=$jobWorkspace export tarName=$tarName + export llmTarArtifactPath=$llmTarArtifactPath + export LLM_ARTIFACTS_PRIMARY=${llmArtifactsPrimary()} + export ARTIFACTORY_NETRC=$artifactoryNetrcPathNode + export ARTIFACTORY_CREDENTIAL_ID=$LLM_ARTIFACTS_CREDENTIALS_ID export llmTarfile=$llmTarfile export llmSrcNode=$llmSrcNode export stageName=$stageName @@ -2317,6 +2484,8 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG scriptRunPathNode, scriptInstallPathNode, scriptBashUtilsPathNode, + scriptFetchHelperPathNode, + artifactoryNetrcPathNode, scriptLaunchPathNode, scriptSubmitPathNode, scriptTrackPathNode, @@ -2817,7 +2986,7 @@ def runLLMTestlistWithSbatch(pipeline, platform, testList, config=VANILLA_CONFIG } // Reaching here means cleanup succeeded within its retry budget. If all // attempts fail, the throw skips deregistration and leaves the entry for sweep. - deregisterSlurmResource(stageName) + deregisterSlurmResource(stageName, jobUID) } } } @@ -3239,6 +3408,8 @@ def BOLT_PROFILE_REF = "bolt_profile_ref" def BOLT_PUBLISH_VARIANT = "bolt_publish_variant" @Field def BOLT_PROFILE_BRANCH = "bolt_profile_branch" +@Field +def LLM_ARTIFACTS_PRIMARY_KEY = "llm_artifacts_primary" // Hoisted out of globalVars into BINDING variables in launchTestJobs. globalVars // is not a parameter of runLLMTestlistWithSbatch or runLLMTestlistOnPlatformImpl, // so reading it inside them throws MissingPropertyException; Build.groovy handles @@ -3262,6 +3433,9 @@ def globalVars = [ (BOLT_PUBLISH_VARIANT): false, (BOLT_PROFILE_BRANCH): "", (JOB_TYPE): null, + // Pre-declared so updateMapWithJson() populates it from the parent (see + // LLM_ARTIFACTS_PRIMARY_FROM_PARENT). + (LLM_ARTIFACTS_PRIMARY_KEY): "", ] class GlobalState { @@ -4147,9 +4321,9 @@ def runLLMDocBuild(pipeline, config) def llmSrc = "${llmPath}/TensorRT-LLM/src" // Step 2: download TRT-LLM tarfile - def llmTarfile = "https://urm.nvidia.com/artifactory/${ARTIFACT_PATH}/${BUILD_CONFIGS[config][TARNAME]}" - trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && wget -nv -O '${BUILD_CONFIGS[config][TARNAME]}' '${llmTarfile}'") - sh "cd ${llmPath} && tar -zxf ${BUILD_CONFIGS[config][TARNAME]}" + def tarName = BUILD_CONFIGS[config][TARNAME] + fetchBuildArtifact(pipeline, "${ARTIFACT_PATH}/${tarName}", "${llmPath}/${tarName}") + sh "cd ${llmPath} && tar -zxf ${tarName}" // install python package if (env.alternativeTRT) { sh "cd ${llmSrc} && sed -i 's#tensorrt~=.*\$#tensorrt#g' requirements.txt && cat requirements.txt" @@ -4251,8 +4425,9 @@ def launchTestListCheck(pipeline, globalVars) sh "nvidia-smi && nvidia-smi -q && nvidia-smi topo -m" // download TRT-LLM tarfile def tarName = BUILD_CONFIGS[VANILLA_CONFIG][TARNAME] - def llmTarfile = "https://urm.nvidia.com/artifactory/${ARTIFACT_PATH}/${tarName}" - trtllm_utils.llmExecStepWithRetry(pipeline, script: "pwd && wget -nv -O '${tarName}' '${llmTarfile}' && ls -alh") + sh "pwd" + fetchBuildArtifact(pipeline, "${ARTIFACT_PATH}/${tarName}", tarName) + sh "ls -alh" sh "tar -zxf ${tarName}" def llmPath = sh (script: "realpath .", returnStdout: true).trim() def llmSrc = "${llmPath}/TensorRT-LLM/src" @@ -4905,9 +5080,8 @@ def generateRerunReport(stageName, llmSrc) { } def mergeWaivesTxt(pipeline, llmSrc, stageName) { - def waivesTxt = "https://urm.nvidia.com/artifactory/${ARTIFACT_PATH}/waive_list/waives.txt" try { - trtllm_utils.llmExecStepWithRetry(pipeline, script: "wget -nv ${waivesTxt}") + fetchBuildArtifact(pipeline, "${ARTIFACT_PATH}/waive_list/waives.txt", "waives.txt") if (!fileExists("waives.txt")) { error "There is no merged waives.txt file, use the default waives.txt." } @@ -5225,9 +5399,9 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO sh "rm -rf results-${stageName}.tar.gz ${stageName}/*" // download TRT-LLM tarfile def tarName = BUILD_CONFIGS[config][TARNAME] - def llmTarfile = boltedTarUrl(config, tarName) + def llmTarArtifactPath = boltedTarArtifactPath(pipeline, config, tarName) timeout(time: 30, unit: 'MINUTES') { - trtllm_utils.llmExecStepWithRetry(pipeline, script: "cd ${llmPath} && wget -nv -O '${tarName}' '${llmTarfile}'") + fetchBuildArtifact(pipeline, llmTarArtifactPath, "${llmPath}/${tarName}") } sh "cd ${llmPath} && tar -zxf ${tarName}" @@ -7599,6 +7773,16 @@ pipeline { globalVars = trtllm_utils.updateMapWithJson(this, globalVars, env.globalVars, "globalVars") globalVars = trtllm_utils.initializeCiBudget(this, globalVars, 24, 'HOURS', 'L0_Test') globalVars[ACTION_INFO] = trtllm_utils.setupPipelineDescription(this, globalVars[ACTION_INFO]) + if (globalVars[LLM_ARTIFACTS_PRIMARY_KEY]) { + LLM_ARTIFACTS_PRIMARY_FROM_PARENT = globalVars[LLM_ARTIFACTS_PRIMARY_KEY].toString() + echo "Build artifacts: primary location ${llmArtifactsPrimary()}, passed by the parent pipeline" + } + // The read helper that fetchBuildArtifact unstashes on every node. + def fetchHelperDir = "${pwd(tmp: true)}/fetch-build-artifact" + trtllm_utils.checkoutFile(LLM_REPO, env.gitlabCommit, "jenkins/scripts/fetch_build_artifact.sh", fetchHelperDir) + dir(fetchHelperDir) { + stash(name: FETCH_BUILD_ARTIFACT_STASH, includes: "fetch_build_artifact.sh") + } } } } diff --git a/jenkins/scripts/cbts/cbts/command/coverage/selection/artifact.py b/jenkins/scripts/cbts/cbts/command/coverage/selection/artifact.py index ba85449b5285..0545580f44f0 100644 --- a/jenkins/scripts/cbts/cbts/command/coverage/selection/artifact.py +++ b/jenkins/scripts/cbts/cbts/command/coverage/selection/artifact.py @@ -46,7 +46,7 @@ import urllib.request from functools import lru_cache from pathlib import Path -from typing import Literal, Optional +from typing import Literal, NoReturn, Optional import click @@ -80,6 +80,17 @@ _URM = "https://urm.nvidia.com/artifactory" _GITHUB_COMPARE = "https://api.github.com/repos/NVIDIA/TensorRT-LLM/compare" _JENKINS_BASE = "https://prod.blsm.nvidia.com/sw-tensorrt-top-1/job/LLM/job/main/job/L0_PostMerge" + +# `build_info.txt` is read with the read helper, jenkins/scripts/fetch_build_artifact.sh, which +# looks for it in the primary artifact location (LLM_ARTIFACTS_PRIMARY) first and in the other one +# only when it is absent there, authenticating to artifactory.nvidia.com with the ARTIFACTORY_NETRC +# file. Coverage archives and pins are read from `_URM` only. +_FETCH_BUILD_ARTIFACT = Path(__file__).resolve().parents[5] / "fetch_build_artifact.sh" +# Exit codes of the read helper that `_fetch_build_artifact` acts on, by the names the helper gives +# them; tests/unittest/scripts/test_fetch_build_artifact.py checks that they match. +_FETCH_ABSENT = 2 +_FETCH_REJECTED = 3 +_FETCH_TRANSFER_FAILED = 4 # Cover missing or unsuccessful recent builds. _MAX_PROBE = 50 # Per-request timeout in seconds, for the small JSON/metadata calls. @@ -96,6 +107,18 @@ _PR_NUMBER_RE = re.compile(r"[1-9][0-9]*") +class ArtifactError(Exception): + """An in-scope build artifact could not be read; the message is written for the job log.""" + + +class ArtifactAuthError(ArtifactError): + """The read helper exited with FETCH_REJECTED: the primary location rejected the request.""" + + +class ArtifactTransferError(ArtifactError): + """The read helper exited with FETCH_TRANSFER_FAILED: the transfer failed after its retries.""" + + def _get(url: str, headers: Optional[dict] = None) -> tuple[Optional[int], Optional[bytes]]: req = urllib.request.Request(url, headers=headers or {}) try: @@ -141,14 +164,55 @@ def tarball_urls(build: int, artifact_base: str = ARTIFACT_BASE) -> list[str]: return [tarball_url(build, name, artifact_base) for name in ARCH_TARBALL_NAMES] -def build_info_url(build: int, artifact_base: str = ARTIFACT_BASE) -> str: - return f"{_URM}/{artifact_base}/{build}/{BUILD_INFO_NAME}" +# Exceptions for the read helper's failure exit statuses, with what each status means. +_HELPER_FAILURES = { + _FETCH_REJECTED: (ArtifactAuthError, "the location rejected the request"), + _FETCH_TRANSFER_FAILED: (ArtifactTransferError, "the transfer failed after retries"), +} + + +def _fetch_build_artifact(artifact_path: str) -> Optional[bytes]: + """Content of a build artifact read with the read helper, or None when neither location has it. + + The helper's stdout is discarded, since this command's stdout carries JSON; its stderr, which + names the location that served the file or says why it could not, goes to the job log. + + Raises: + ArtifactAuthError: The helper exited with FETCH_REJECTED. + ArtifactTransferError: The helper exited with FETCH_TRANSFER_FAILED. + ArtifactError: The helper could not be run, or exited with another status. + """ + with tempfile.TemporaryDirectory(prefix="cbts-artifact-") as tmp: + dest = Path(tmp) / Path(artifact_path).name + try: + status = subprocess.run( + ["bash", str(_FETCH_BUILD_ARTIFACT), artifact_path, str(dest)], + stdout=subprocess.DEVNULL, + check=False, + ).returncode + content = dest.read_bytes() if status == 0 else None + except OSError as e: + raise ArtifactError( + f"{_FETCH_BUILD_ARTIFACT.name} could not read {artifact_path}: {e}" + ) from e + if status in (0, _FETCH_ABSENT): + return content + error, reason = _HELPER_FAILURES.get(status, (ArtifactError, None)) + detail = f": {reason}" if reason else "" + raise error( + f"{_FETCH_BUILD_ARTIFACT.name} could not read {artifact_path}{detail} " + f"(exit status {status}); its message above says why." + ) def build_commit(build: int, artifact_base: str = ARTIFACT_BASE) -> Optional[str]: - """Revision a build ran, from its `build_info.txt`, or None when unavailable.""" - status, data = _get(build_info_url(build, artifact_base)) - if status != 200 or not data: + """Revision a build ran, from its `build_info.txt`, or None when unavailable. + + Raises: + ArtifactError: As `_fetch_build_artifact` raises it; `main` reports it and exits 1. + """ + data = _fetch_build_artifact(f"{artifact_base}/{build}/{BUILD_INFO_NAME}") + if not data: return None for line in data.decode("utf-8", "replace").splitlines(): key, _, value = line.partition("=") @@ -770,7 +834,35 @@ def main( ) if expected_commit is not None and build is None: raise click.UsageError("--expected-commit requires --build") + try: + _run_entry_point( + resolve_pin_file, + resolve_build, + build, + expected_commit, + prepare_dir, + pr_head, + pr_number, + paths_json, + ) + except ArtifactError as e: + # Reported like the other failures: a message on stderr, no JSON, exit status 1. + for line in str(e).splitlines(): + print(f"[artifact] {line}", file=sys.stderr) + sys.exit(1) + +def _run_entry_point( + resolve_pin_file: Optional[str], + resolve_build: bool, + build: Optional[int], + expected_commit: Optional[str], + prepare_dir: Optional[str], + pr_head: Optional[str], + pr_number: Optional[str], + paths_json: Optional[str], +) -> NoReturn: + """Run the entry point `main` selected and exit with its status.""" if prepare_dir: relevant_paths = _load_paths_json(paths_json) if relevant_paths is None: diff --git a/jenkins/scripts/fetch_build_artifact.sh b/jenkins/scripts/fetch_build_artifact.sh new file mode 100755 index 000000000000..6aa0a7337b62 --- /dev/null +++ b/jenkins/scripts/fetch_build_artifact.sh @@ -0,0 +1,481 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Reads one file under llm-artifacts/ from the primary Artifactory location, and from the secondary +# location only when the primary does not have it. CI reads a build's build_info.txt, packaged +# builds, and waive_list/ with it, and it reads any other file there the same way. Run it with +# --help for the interface and the exit codes. +# +# Credentials go to the aws location only, through curl's --netrc-file; urm is read anonymously. +# The script turns off xtrace that a caller passes on, because a trace would show response URLs, +# which can carry a CDN's signature. Never pass a secret on a command line. + +# The steps run through with_retries and cleanup runs from a trap, so shellcheck cannot see that +# they are reachable. +# shellcheck disable=SC2317,SC2329 + +set +x +set -uo pipefail + +readonly PROG="fetch_build_artifact" +readonly MAX_ATTEMPTS=3 +readonly AWS_BASE_URL="https://artifactory.nvidia.com/artifactory/sw-tensorrt-llm-generic-local" +readonly URM_BASE_URL="https://urm.nvidia.com/artifactory/sw-tensorrt-generic" +# Exit codes other than success. Readers that act on one define the same name for it, and +# tests/unittest/scripts/test_fetch_build_artifact.py checks that the values match. +readonly FETCH_ABSENT=2 +readonly FETCH_REJECTED=3 +readonly FETCH_TRANSFER_FAILED=4 +readonly FETCH_USAGE=64 + +usage() { + cat <<'EOF' +Usage: fetch_build_artifact.sh [--properties | --exists] [] + +Reads a build artifact from the primary Artifactory location, and from the secondary location +only when the primary does not have it. A rejected request, a server or connection error, or a +checksum mismatch never falls back. + +Arguments: + A legacy path (sw-tensorrt-generic/llm-artifacts/... or + sw-tensorrt-generic-local/llm-artifacts/...) or a layout path + (llm-artifacts/...). + Output file, written as .part and renamed when complete. + Required unless --exists is given. + --properties Write the artifact's properties (storage API JSON) instead of the file. + --exists Write no file; print " " for the location that holds the artifact. + +Environment: + LLM_ARTIFACTS_PRIMARY aws (the default) or urm; the other location is the secondary. + ARTIFACTORY_NETRC netrc file holding an Artifactory access token for + artifactory.nvidia.com; required whenever aws is contacted. + ARTIFACTORY_CREDENTIAL_ID Name of the CI credential behind ARTIFACTORY_NETRC, for messages. + LLM_ARTIFACTS_AWS_URL, LLM_ARTIFACTS_URM_URL + Replace a location's base URL, /. + +Exit codes: + 0 Success; stderr gets "fetch_build_artifact: served-by= path=". + 2 Neither location has the artifact (with --properties: has properties for it). + 3 The primary rejected the request; the secondary was not contacted. + 4 The transfer failed after retries: a connection or server error, a persistent CDN + rejection, a checksum mismatch, or a failure other than "absent" on the secondary. + 64 Usage error, unmappable path, bad LLM_ARTIFACTS_PRIMARY, missing ARTIFACTORY_NETRC, + or a missing tool. +EOF +} + +log() { + printf '%s: %s\n' "$PROG" "$*" >&2 +} + +# fail ... +fail() { + local code="$1" + shift + log "$@" + exit "$code" +} + +usage_error() { + log "$@" + usage >&2 + exit "$FETCH_USAGE" +} + +# Prints the layout path (llm-artifacts/...) of a legacy or layout artifact path. Fails for any +# other path, for an empty, "." or ".." segment, and for characters that would change the meaning +# of the request URL. +map_path() { + local path="$1" unsafe='[[:space:];?#]' + case "$path" in + sw-tensorrt-generic/*) path="${path#sw-tensorrt-generic/}" ;; + sw-tensorrt-generic-local/*) path="${path#sw-tensorrt-generic-local/}" ;; + *) ;; + esac + [[ "$path" == llm-artifacts/?* ]] || return 1 + case "/${path}/" in + *//* | */./* | */../*) return 1 ;; + *) ;; + esac + [[ ! "$path" =~ $unsafe ]] || return 1 + printf '%s\n' "$path" +} + +# Prints the lower-case host[:port] of a URL. +url_host() { + local rest="${1#*://}" + rest="${rest%%[/?#]*}" + rest="${rest##*@}" + printf '%s\n' "${rest,,}" +} + +# use_location : sets LOC, LOC_BASE (/), LOC_ROOT, LOC_REPO, and +# LOC_HOST. Fails when the base URL does not have that shape. +use_location() { + LOC="$1" + if [[ "$LOC" == aws ]]; then + LOC_BASE="${LLM_ARTIFACTS_AWS_URL:-$AWS_BASE_URL}" + else + LOC_BASE="${LLM_ARTIFACTS_URM_URL:-$URM_BASE_URL}" + fi + LOC_BASE="${LOC_BASE%/}" + LOC_ROOT="${LOC_BASE%/*}" + LOC_REPO="${LOC_BASE##*/}" + LOC_HOST="$(url_host "$LOC_BASE")" + [[ "$LOC_ROOT" =~ ^https?://[^/]+ && -n "$LOC_REPO" ]] +} + +file_url() { + printf '%s/%s\n' "$LOC_BASE" "$LAYOUT_PATH" +} + +storage_url() { + printf '%s/api/storage/%s/%s\n' "$LOC_ROOT" "$LOC_REPO" "$LAYOUT_PATH" +} + +require_netrc() { + local need="reading $LAYOUT_PATH from aws needs a netrc file with an Artifactory access token" + need+=" for artifactory.nvidia.com." + if [[ -z "${ARTIFACTORY_NETRC:-}" ]]; then + fail "$FETCH_USAGE" "ARTIFACTORY_NETRC is not set; $need" + fi + if [[ ! -f "$ARTIFACTORY_NETRC" || ! -r "$ARTIFACTORY_NETRC" ]]; then + fail "$FETCH_USAGE" "ARTIFACTORY_NETRC ($ARTIFACTORY_NETRC) is not a readable file; $need" + fi +} + +# http_get [...] +# Sends one GET and follows redirects, with the netrc only when the location is aws; curl applies +# a netrc entry only to the host it names, so a redirect to another host carries no credential. +# Sets HTTP_STATUS (000 when no response arrived) and FINAL_HOST (host[:port] of the last URL +# requested), and returns curl's exit status. +http_get() { + local output="$1" url="$2" result rc + shift 2 + if [[ "$LOC" == aws ]]; then + set -- --netrc-file "$ARTIFACTORY_NETRC" "$@" + fi + # Start from an empty file, so that it never holds the body of an earlier attempt. + : >"$output" + # -q must come first. It keeps curl from reading a curlrc, which could, for example, turn on + # verbose output that prints the Authorization header. + result="$(curl -q --silent --show-error --location --connect-timeout 30 \ + --speed-limit 1 --speed-time 300 "$@" \ + --output "$output" --write-out '%{http_code} %{url_effective}' "$url")" + rc=$? + HTTP_STATUS="${result%% *}" + FINAL_HOST="$(url_host "${result#* }")" + return "$rc" +} + +# Prints checksums.sha1 from a storage API answer. Text tools keep the helper usable on hosts +# without python or jq. The pattern matches the "checksums" key only, never "originalChecksums". +extract_sha1() { + tr -d ' \t\r\n' <"$1" \ + | sed -n 's/.*"checksums":{[^}]*"sha1":"\([0-9a-fA-F]\{40\}\)".*/\1/p' \ + | tr 'A-F' 'a-f' +} + +# Steps make one attempt each and return 0 (done), 1 (transient failure, reason in FAILURE), +# FETCH_ABSENT, or FETCH_REJECTED (status in REJECTED_STATUS). + +# Reads the artifact's checksums.sha1 from the storage API into SHA1. +lookup_step() { + local rc + http_get "$WORK_DIR/storage.json" "$(storage_url)" + rc=$? + if (( rc != 0 )); then + FAILURE="the storage lookup failed (curl exit $rc)" + return 1 + fi + case "$HTTP_STATUS" in + 200) + SHA1="$(extract_sha1 "$WORK_DIR/storage.json")" + if [[ -n "$SHA1" ]]; then + return 0 + fi + FAILURE="the storage lookup answered without checksums.sha1" + ;; + 404) return "$FETCH_ABSENT" ;; + 401 | 403) + REJECTED_STATUS="$HTTP_STATUS" + return "$FETCH_REJECTED" + ;; + *) FAILURE="the storage lookup answered HTTP $HTTP_STATUS" ;; + esac + return 1 +} + +# Downloads the file to PART, and renames it to DEST when its SHA-1 matches SHA1. +download_step() { + local rc cdn_host="" + http_get "$PART" "$(file_url)" + rc=$? + if (( rc == 0 )) && [[ "$HTTP_STATUS" == 403 && "$FINAL_HOST" != "$LOC_HOST" ]]; then + # The CDN that the location redirected the download to rejected it, which says nothing + # about the credential. Ask the location to send the file itself. + cdn_host="$FINAL_HOST" + log "$cdn_host rejected the redirected download (HTTP 403);" \ + "downloading from $LOC_HOST without the redirect." + http_get "$PART" "$(file_url)" -H 'X-JFrog-Download-Redirect-To: None' + rc=$? + fi + if (( rc == 0 )) && [[ "$HTTP_STATUS" == 200 ]]; then + verify_download + return + fi + rm -f -- "$PART" + if (( rc != 0 )); then + FAILURE="the download failed (curl exit $rc)" + elif [[ -n "$cdn_host" ]]; then + FAILURE="the download without the CDN redirect answered HTTP $HTTP_STATUS" + elif [[ "$FINAL_HOST" == "$LOC_HOST" && "$HTTP_STATUS" =~ ^40[13]$ ]]; then + REJECTED_STATUS="$HTTP_STATUS" + return "$FETCH_REJECTED" + else + FAILURE="the download answered HTTP $HTTP_STATUS from $FINAL_HOST" + fi + return 1 +} + +verify_download() { + local actual + actual="$(sha1sum <"$PART")" + actual="${actual%% *}" + if [[ "$actual" != "$SHA1" ]]; then + rm -f -- "$PART" + FAILURE="checksum mismatch: the download has SHA-1 $actual, the storage API reports $SHA1" + return 1 + fi + if ! mv -f -- "$PART" "$DEST"; then + FAILURE="could not rename $PART to $DEST" + return 1 + fi +} + +# Writes the artifact's properties (storage API JSON) to DEST. +properties_step() { + local rc + http_get "$PART" "$(storage_url)?properties" + rc=$? + if (( rc == 0 )) && [[ "$HTTP_STATUS" == 200 ]] && mv -f -- "$PART" "$DEST"; then + return 0 + fi + rm -f -- "$PART" + if (( rc != 0 )); then + FAILURE="the properties lookup failed (curl exit $rc)" + return 1 + fi + case "$HTTP_STATUS" in + 404) return "$FETCH_ABSENT" ;; + 401 | 403) + REJECTED_STATUS="$HTTP_STATUS" + return "$FETCH_REJECTED" + ;; + 200) FAILURE="could not rename $PART to $DEST" ;; + *) FAILURE="the properties lookup answered HTTP $HTTP_STATUS" ;; + esac + return 1 +} + +# with_retries : repeats a step while it fails transiently, at most MAX_ATTEMPTS times in +# total, sleeping 2 s and then 4 s in between. Returns the step's result, or +# FETCH_TRANSFER_FAILED when its last attempt failed transiently. +with_retries() { + local attempt rc delay + for (( attempt = 1; ; attempt++ )); do + "$1" + rc=$? + if (( rc != 1 )); then + return "$rc" + fi + if (( attempt >= MAX_ATTEMPTS )); then + log "attempt $attempt of $MAX_ATTEMPTS on $LOC ($LOC_HOST) failed: $FAILURE;" \ + "giving up on $LAYOUT_PATH." + return "$FETCH_TRANSFER_FAILED" + fi + delay=$(( 2 ** attempt )) + log "attempt $attempt of $MAX_ATTEMPTS on $LOC ($LOC_HOST) failed: $FAILURE;" \ + "retrying in ${delay}s." + sleep "$delay" + done +} + +# read_from : reads the artifact from one location. Returns 0 (served), FETCH_ABSENT, +# FETCH_REJECTED, or FETCH_TRANSFER_FAILED (failed after retries). +read_from() { + local rc + use_location "$1" + if [[ "$LOC" == aws ]]; then + require_netrc + fi + if [[ "$MODE" == properties ]]; then + with_retries properties_step + return + fi + with_retries lookup_step + rc=$? + if (( rc != 0 )) || [[ "$MODE" == exists ]]; then + return "$rc" + fi + # The lookup found the file here, so a failed download is a transfer failure and never + # sends the read to the other location. + with_retries download_step +} + +# The message names both causes for a 401 and a 403 alike, so that a missing permission is never +# taken for an expired token. +report_rejection() { + if [[ "$LOC" == aws ]]; then + log "$LOC_HOST rejected the credentials for $LAYOUT_PATH (HTTP $REJECTED_STATUS)." + log "the Artifactory access token in credential" \ + "'${ARTIFACTORY_CREDENTIAL_ID:-}' needs renewal" \ + "or lacks read permission on $LOC_REPO;" \ + "an Artifactory access token for artifactory.nvidia.com is required." + else + log "$LOC_HOST rejected the anonymous request for $LAYOUT_PATH (HTTP $REJECTED_STATUS)." + fi +} + +served() { + if [[ "$MODE" == exists ]]; then + printf '%s %s\n' "$(file_url)" "$SHA1" + fi + log "served-by=$LOC path=$LAYOUT_PATH" + exit 0 +} + +cleanup() { + rm -rf -- "$WORK_DIR" + if [[ -n "$PART" ]]; then + rm -f -- "$PART" + fi +} + +MODE=download +ARGS=() +while (( $# > 0 )); do + case "$1" in + -h | --help) + usage + exit 0 + ;; + --properties | --exists) + if [[ "$MODE" != download && "$MODE" != "${1#--}" ]]; then + usage_error "--properties and --exists cannot be combined." + fi + MODE="${1#--}" + ;; + --) + shift + ARGS+=("$@") + break + ;; + -?*) usage_error "unknown option: $1" ;; + *) ARGS+=("$1") ;; + esac + shift +done + +case "$MODE:${#ARGS[@]}" in + download:2 | properties:2 | exists:1 | exists:2) ;; + *) usage_error "wrong number of arguments." ;; +esac + +if ! LAYOUT_PATH="$(map_path "${ARGS[0]}")"; then + usage_error "not a file path under llm-artifacts/: '${ARGS[0]}'" +fi + +# --exists writes no file, so it ignores . +DEST="" +PART="" +if [[ "$MODE" != exists ]]; then + DEST="${ARGS[1]}" + if [[ -z "$DEST" || "$DEST" == */ || -d "$DEST" ]]; then + fail "$FETCH_USAGE" " must name a file: '$DEST'" + fi + if [[ ! -d "$(dirname -- "$DEST")" ]]; then + fail "$FETCH_USAGE" "the directory of does not exist: '$DEST'" + fi + PART="$DEST.part" +fi + +PRIMARY="${LLM_ARTIFACTS_PRIMARY:-aws}" +case "$PRIMARY" in + aws) SECONDARY=urm ;; + urm) SECONDARY=aws ;; + *) fail "$FETCH_USAGE" "LLM_ARTIFACTS_PRIMARY must be 'aws' or 'urm', not '$PRIMARY'." ;; +esac + +for location in aws urm; do + override="LLM_ARTIFACTS_${location^^}_URL" + # Messages show the base URL, so it must not carry credentials. + authority="${!override:-}" + authority="${authority#*://}" + if [[ "${authority%%/*}" == *@* ]]; then + fail "$FETCH_USAGE" "$override must not contain credentials; ARTIFACTORY_NETRC provides them." + fi + if ! use_location "$location"; then + fail "$FETCH_USAGE" "$override must be an http(s) URL of the form /, not '$LOC_BASE'." + fi +done + +command -v curl >/dev/null 2>&1 || fail "$FETCH_USAGE" "curl is required." +if [[ "$MODE" == download ]]; then + command -v sha1sum >/dev/null 2>&1 || fail "$FETCH_USAGE" "sha1sum is required." +fi + +LOC="" +SHA1="" +FAILURE="" +REJECTED_STATUS="" +HTTP_STATUS="" +FINAL_HOST="" +WORK_DIR="$(mktemp -d "${TMPDIR:-/tmp}/fetch_build_artifact.XXXXXX")" \ + || fail "$FETCH_USAGE" "cannot create a temporary directory." +trap cleanup EXIT +trap 'exit 130' INT +trap 'exit 143' TERM + +read_from "$PRIMARY" +case $? in + 0) served ;; + "$FETCH_ABSENT") ;; + "$FETCH_REJECTED") + report_rejection + exit "$FETCH_REJECTED" + ;; + *) exit "$FETCH_TRANSFER_FAILED" ;; +esac + +read_from "$SECONDARY" +case $? in + 0) served ;; + "$FETCH_ABSENT") + if [[ "$MODE" == properties ]]; then + log "neither $PRIMARY nor $SECONDARY has properties for $LAYOUT_PATH." + else + log "neither $PRIMARY nor $SECONDARY has $LAYOUT_PATH." + fi + exit "$FETCH_ABSENT" + ;; + "$FETCH_REJECTED") + # A rejection by the secondary is reported like any other failure there. + report_rejection + exit "$FETCH_TRANSFER_FAILED" + ;; + *) exit "$FETCH_TRANSFER_FAILED" ;; +esac diff --git a/jenkins/scripts/slurm_install.sh b/jenkins/scripts/slurm_install.sh index aff14faaceba..578b208f2869 100644 --- a/jenkins/scripts/slurm_install.sh +++ b/jenkins/scripts/slurm_install.sh @@ -22,9 +22,15 @@ trap 'rc=$?; echo "Error in file ${BASH_SOURCE[0]} on line $LINENO: $BASH_COMMAN bashUtilsPath="$(dirname "${BASH_SOURCE[0]}")/$(basename "${BASH_SOURCE[0]}" | sed 's/slurm_install\.sh/bash_utils.sh/')" source "$bashUtilsPath" +# The read helper (jenkins/scripts/fetch_build_artifact.sh) sits beside this script, under the same +# per-job prefix. LLM_ARTIFACTS_PRIMARY, ARTIFACTORY_NETRC and ARTIFACTORY_CREDENTIAL_ID reach it +# from the job script. fetchHelperPath is used after slurm_install_setup changes into +# $resourcePathNode, so run or source this script by an absolute path. +fetchHelperPath="$(dirname "${BASH_SOURCE[0]}")/$(basename "${BASH_SOURCE[0]}" | sed 's/slurm_install\.sh/fetch_build_artifact.sh/')" + slurm_install_setup() { : "${tarName:?tarName is required}" - : "${llmTarfile:?llmTarfile is required}" + : "${llmTarArtifactPath:?llmTarArtifactPath is required}" : "${resourcePathNode:?resourcePathNode is required}" : "${stageName:?stageName is required}" : "${pytestCommand:?pytestCommand is required}" @@ -53,12 +59,11 @@ slurm_install_setup() { # Job/node-specific tmp path to avoid collisions on concurrent jobs archive_tmp="${archive_path}.tmp.${SLURM_JOB_ID:-local}.${SLURM_NODEID:-0}" rm -f "$archive_path" "$archive_tmp" - # Download the artifact idempotently with retry. A bare "retry_command wget " will - # save artifact as $tarName.1 when the first attempt fails in the middle of downloading. - # Here we download to the tmp path and only promote it on success - if ! retry_command --timeout 1800 bash -c 'wget -nv "$1" -O "$2" && mv -f "$2" "$3"' _ "$llmTarfile" "$archive_tmp" "$archive_path"; then + # Download the artifact idempotently with retry: each attempt writes to the tmp path, and + # only a successful one promotes it. + if ! retry_command --timeout 1800 bash -c 'bash "$1" "$2" "$3" && mv -f "$3" "$4"' _ "$fetchHelperPath" "$llmTarArtifactPath" "$archive_tmp" "$archive_path"; then rm -f "$archive_tmp" - echo "Artifact download failed after retries: $llmTarfile" + echo "Artifact download failed after retries: $llmTarArtifactPath" return 1 fi if [ ! -f "$archive_path" ]; then diff --git a/scripts/bolt/internal/artifactory.sh b/scripts/bolt/internal/artifactory.sh index 903d6ac06ffa..61a5e636a9b2 100755 --- a/scripts/bolt/internal/artifactory.sh +++ b/scripts/bolt/internal/artifactory.sh @@ -36,16 +36,16 @@ # Artifactory layout (override via env): # ////bolt-profile--.tar.gz # ////latest.tar.gz -# Defaults match the repo used by jenkins/Build.groovy (sw-tensorrt-generic). +# Defaults match the BOLT profile store on urm.nvidia.com (sw-tensorrt-generic), +# which jenkins/BoltProfileGen.groovy promotes to (PROFILE_PROMOTE_DIR). # # Auth: # promote -> `jf` (JFrog) CLI (write). If JF_URL + JF_ACCESS_TOKEN are set # this configures jf automatically; else assumes jf is # preconfigured. (The CI postmerge path in BoltProfileGen.groovy # instead uploads cluster-side via curl + urm-artifactory-creds.) -# pull-latest -> anonymous curl (read), matching Build.groovy's tarball -# download from the same repo; needs no creds and no jf, so it -# runs in the premerge build pod as-is. +# pull-latest -> anonymous curl or wget (read) from the profile store on urm; +# needs no creds and no jf, so it runs in a build pod as-is. set -euo pipefail @@ -185,9 +185,8 @@ cmd_pull_latest() { fi mkdir -p "$dest" log "Pulling $url -> $dest" - # Anonymous download (like jenkins/Build.groovy's tarball from the same - # sw-tensorrt-generic repo): the read path needs no creds, so this works in the - # build pod without a configured `jf`. Use curl if present, else fall back to + # Anonymous download: reading the profile store needs no creds, so this works in + # a build pod without a configured `jf`. Use curl if present, else fall back to # wget -- the image-overlay build pod (BuildDockerImage.overlayBoltBundle) has no # curl, so hard-requiring curl made the overlay fail (curl: command not found) # even though the bundle existed. --retry rides out flaky links; a missing object diff --git a/scripts/get_wheel_from_package.py b/scripts/get_wheel_from_package.py index 9fdf1607c8c0..45b8027e2b32 100644 --- a/scripts/get_wheel_from_package.py +++ b/scripts/get_wheel_from_package.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -117,6 +117,83 @@ def bolt_optimize_wheels(build_dir, arch, bolt_branch, bolt_profile_ref=None): f"refusing to build an unoptimized release image") +# Exit codes of jenkins/scripts/fetch_build_artifact.sh that download_package +# acts on, by the names the helper gives them; any other non-zero code stops +# it at once. tests/unittest/scripts/test_fetch_build_artifact.py checks that +# they match. +_FETCH_ABSENT = 2 +_FETCH_REJECTED = 3 +_FETCH_TRANSFER_FAILED = 4 + + +def _rejected_message(package_path: str) -> str: + """Describe why the primary location refused to serve `package_path`.""" + if (os.environ.get("LLM_ARTIFACTS_PRIMARY") or "aws") == "urm": + return ("The primary location (urm) rejected the anonymous request " + f"for {package_path}.") + credential = (os.environ.get("ARTIFACTORY_CREDENTIAL_ID") + or "") + return ( + f"Artifactory rejected the credentials for {package_path}: the " + f"Artifactory access token in credential '{credential}' needs " + "renewal or lacks read permission on sw-tensorrt-llm-generic-local; " + "an Artifactory access token for artifactory.nvidia.com is required.") + + +def download_package(artifact_path: str, tarfile_name: str, + timeout: int) -> None: + """Download a packaged build, polling until it has been published. + + jenkins/scripts/fetch_build_artifact.sh reads the primary location + (LLM_ARTIFACTS_PRIMARY) first, falls back to the other location only when + the file is absent there, and logs which location served it. It writes to + a temporary file and renames it only after the SHA-1 matches, so a failed + attempt never leaves a partial `tarfile_name` for the extract to read. + + Args: + artifact_path: The build's artifact directory in the legacy layout, + a repository name followed by llm-artifacts//. The + helper maps it to each location. + tarfile_name: File name of the packaged build, also used for the local + copy. + timeout: Maximum number of attempts, one minute apart. + + Raises: + PermissionError: The primary location rejected the request. + RuntimeError: The helper failed for a reason other than the package + being absent or a transfer failure, such as a usage error. + TimeoutError: The package was not downloaded within `timeout` + attempts. + """ + helper = (get_project_dir() / "jenkins" / "scripts" / + "fetch_build_artifact.sh") + package_path = f"{artifact_path}/{tarfile_name}" + for attempt in range(1, timeout + 1): + # The environment carries ARTIFACTORY_NETRC, LLM_ARTIFACTS_PRIMARY and + # ARTIFACTORY_CREDENTIAL_ID to the helper. + rc = subprocess.run( + ["bash", str(helper), package_path, tarfile_name], + env=os.environ.copy()).returncode + if rc == 0: + return + if rc == _FETCH_REJECTED: + raise PermissionError(_rejected_message(package_path)) + if rc not in (_FETCH_ABSENT, _FETCH_TRANSFER_FAILED): + raise RuntimeError( + f"fetch_build_artifact.sh failed with exit code {rc} for " + f"{package_path}") + if attempt < timeout: + reason = ("not published yet" + if rc == _FETCH_ABSENT else "transfer failed") + print( + f"Tarfile {reason}, waiting 60 seconds... " + f"(attempt {attempt}/{timeout})", + flush=True) + time.sleep(60) + raise TimeoutError( + f"Failed to download file after {timeout} attempts: {package_path}") + + def get_wheel_from_package(arch, artifact_path, timeout, @@ -127,25 +204,7 @@ def get_wheel_from_package(arch, else: tarfile_name = "TensorRT-LLM-GH200.tar.gz" - tarfile_link = f"https://urm.nvidia.com/artifactory/{artifact_path}/{tarfile_name}" - for attempt in range(timeout): - try: - # -O pins the output name: without it wget falls back to - # .1 when a previous attempt left a partial file behind, - # and the extract below would then read stale bytes. - subprocess.run(["wget", "-nv", "-O", tarfile_name, tarfile_link], - check=True) - print(f"Tarfile is available at {tarfile_link}") - break - except Exception: - if attempt == timeout - 1: - raise TimeoutError( - f"Failed to download file after {timeout} attempts: {tarfile_link}" - ) - print( - f"Tarfile not ready yet, waiting 60 seconds... (attempt {attempt + 1}/{timeout})" - ) - time.sleep(60) + download_package(artifact_path, tarfile_name, timeout) llm_root = get_project_dir() tmp_dir = llm_root / "tmp" diff --git a/tests/unittest/scripts/cbts/test_coverage_artifact.py b/tests/unittest/scripts/cbts/test_coverage_artifact.py index 602d6f0bf4b7..a58007479014 100644 --- a/tests/unittest/scripts/cbts/test_coverage_artifact.py +++ b/tests/unittest/scripts/cbts/test_coverage_artifact.py @@ -17,12 +17,17 @@ from __future__ import annotations +import hashlib import json +import shlex import shutil import sqlite3 import subprocess +import sys import tempfile +import threading import unittest +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from unittest import mock @@ -403,5 +408,251 @@ def extract(tarball: Path, destination: Path) -> bool: self.assertEqual(json.loads(Path(ready["meta"]).read_text()), selection) +# build_info.txt is read with jenkins/scripts/fetch_build_artifact.sh, which applies the read rule +# and is tested in tests/unittest/scripts/test_fetch_build_artifact.py. These tests check how +# artifact.py runs it and what it makes of each exit status. +_BUILD = 88 +_COMMIT = "a" * 40 +_BUILD_INFO = f"commit={_COMMIT}\nbranch=main\ndate=2026-10-16 00:00:00 UTC\n" +_ARTIFACT_PATH = f"{artifact.ARTIFACT_BASE}/{_BUILD}/{artifact.BUILD_INFO_NAME}" +_LAYOUT_PATH = f"llm-artifacts/LLM/main/L0_PostMerge/{_BUILD}/{artifact.BUILD_INFO_NAME}" +_CREDENTIAL_ID = "aws-artifactory-credentials" +_HELPER_ENVIRONMENT = ("LLM_ARTIFACTS_PRIMARY", "ARTIFACTORY_NETRC", "ARTIFACTORY_CREDENTIAL_ID") + +# Stands in for fetch_build_artifact.sh. It records its arguments and the environment it sees, +# writes to stdout and stderr, and exits with HELPER_STUB_STATUS; on 0 it first writes +# HELPER_STUB_CONTENT to . +_HELPER_STUB = r""" +import json +import os +import sys + +with open(os.environ["HELPER_STUB_LOG"], "a", encoding="utf-8") as log: + names = json.loads(os.environ["HELPER_STUB_ENV"]) + log.write( + json.dumps({"args": sys.argv[1:], "env": {name: os.environ.get(name) for name in names}}) + + "\n" + ) +print("helper stdout") +print("helper stderr", file=sys.stderr) +status = int(os.environ["HELPER_STUB_STATUS"]) +if status == 0: + with open(sys.argv[2], "w", encoding="utf-8") as dest: + dest.write(os.environ["HELPER_STUB_CONTENT"]) +sys.exit(status) +""" + + +class _HelperStub: + """Replaces the read helper that artifact.py runs, and records each call.""" + + def __init__(self, monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None: + self._log = tmp_path / "helper_calls.jsonl" + stub_script = tmp_path / "helper_stub.py" + stub_script.write_text(_HELPER_STUB, encoding="utf-8") + helper = tmp_path / "fetch_build_artifact.sh" + helper.write_text( + f'exec {shlex.quote(sys.executable)} -I -S {shlex.quote(str(stub_script))} "$@"\n', + encoding="utf-8", + ) + monkeypatch.setattr(artifact, "_FETCH_BUILD_ARTIFACT", helper) + monkeypatch.setenv("HELPER_STUB_LOG", str(self._log)) + monkeypatch.setenv("HELPER_STUB_ENV", json.dumps(_HELPER_ENVIRONMENT)) + monkeypatch.setenv("HELPER_STUB_CONTENT", _BUILD_INFO) + monkeypatch.setenv("LLM_ARTIFACTS_PRIMARY", "urm") + monkeypatch.setenv("ARTIFACTORY_NETRC", str(tmp_path / "artifactory-netrc")) + monkeypatch.setenv("ARTIFACTORY_CREDENTIAL_ID", _CREDENTIAL_ID) + self._monkeypatch = monkeypatch + self.exits_with(0) + + def exits_with(self, status: int) -> None: + self._monkeypatch.setenv("HELPER_STUB_STATUS", str(status)) + + @property + def calls(self) -> list[dict]: + if not self._log.exists(): + return [] + return [json.loads(line) for line in self._log.read_text(encoding="utf-8").splitlines()] + + +@pytest.fixture +def helper_stub(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> _HelperStub: + return _HelperStub(monkeypatch, tmp_path) + + +def test_build_commit_reads_build_info_through_the_helper( + helper_stub: _HelperStub, tmp_path: Path, capfd: pytest.CaptureFixture[str] +) -> None: + assert artifact.build_commit(_BUILD) == _COMMIT + + [call] = helper_stub.calls + artifact_path, dest = call["args"] + assert artifact_path == _ARTIFACT_PATH + assert Path(dest).name == artifact.BUILD_INFO_NAME + # A temporary file, removed once read. + assert not Path(dest).parent.exists() + assert call["env"] == { + "LLM_ARTIFACTS_PRIMARY": "urm", + "ARTIFACTORY_NETRC": str(tmp_path / "artifactory-netrc"), + "ARTIFACTORY_CREDENTIAL_ID": _CREDENTIAL_ID, + } + captured = capfd.readouterr() + # stdout carries the command's JSON, so the helper's stdout is dropped; its stderr is the + # job log's record of which location served the file, or why it could not. + assert captured.out == "" + assert "helper stderr" in captured.err + + +def test_build_commit_returns_none_when_neither_location_has_build_info( + helper_stub: _HelperStub, +) -> None: + helper_stub.exits_with(2) + + assert artifact.build_commit(_BUILD) is None + assert len(helper_stub.calls) == 1 + + +@pytest.mark.parametrize( + ("status", "error"), + ( + (3, artifact.ArtifactAuthError), + (4, artifact.ArtifactTransferError), + (64, artifact.ArtifactError), + (127, artifact.ArtifactError), + ), + ids=("rejected", "transfer-failed", "usage-error", "other"), +) +def test_build_commit_raises_on_helper_failures( + helper_stub: _HelperStub, status: int, error: type[artifact.ArtifactError] +) -> None: + helper_stub.exits_with(status) + + with pytest.raises(artifact.ArtifactError) as raised: + artifact.build_commit(_BUILD) + + assert type(raised.value) is error + assert _ARTIFACT_PATH in str(raised.value) + assert f"exit status {status}" in str(raised.value) + # One call: the helper retries transient failures itself. + assert len(helper_stub.calls) == 1 + + +def test_build_commit_reports_a_helper_that_cannot_start( + helper_stub: _HelperStub, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + empty_dir = tmp_path / "empty-path" + empty_dir.mkdir() + monkeypatch.setenv("PATH", str(empty_dir)) + + with pytest.raises(artifact.ArtifactError) as raised: + artifact.build_commit(_BUILD) + + assert type(raised.value) is artifact.ArtifactError + assert _ARTIFACT_PATH in str(raised.value) + assert helper_stub.calls == [] + + +@pytest.mark.parametrize("status", (3, 4, 64), ids=("rejected", "transfer-failed", "usage-error")) +def test_main_reports_helper_failures_and_exits_1( + helper_stub: _HelperStub, monkeypatch: pytest.MonkeyPatch, status: int +) -> None: + helper_stub.exits_with(status) + monkeypatch.setattr(artifact, "merge_base", lambda head: "c" * 40) + monkeypatch.setattr(artifact, "_exists", lambda url: True) + + result = CliRunner().invoke( + artifact.main, + ["--resolve-build", "--pr-head", "b" * 40, "--build", str(_BUILD)], + catch_exceptions=False, + ) + + assert result.exit_code == 1 + assert result.stdout == "" + errors = [line for line in result.stderr.splitlines() if line.startswith("[artifact] ")] + assert any(_ARTIFACT_PATH in line and f"exit status {status}" in line for line in errors) + + +class _ArtifactHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + body = self.server.routes.get(self.path) + self.send_response(404 if body is None else 200) + payload = b'{"errors": [{"status": 404}]}' if body is None else body + self.send_header("Content-Length", str(len(payload))) + self.end_headers() + self.wfile.write(payload) + + def log_message(self, format: str, *args: object) -> None: + """Keep access logs out of the test output.""" + + +class _ArtifactServer(ThreadingHTTPServer): + """A loopback Artifactory stand-in that answers GETs from a path map, and 404 otherwise.""" + + daemon_threads = True + + def __init__(self, host: str, routes: dict[str, bytes]) -> None: + super().__init__((host, 0), _ArtifactHandler) + self.routes = routes + + def base_url(self, repo: str) -> str: + host, port = self.server_address[:2] + return f"http://{host}:{port}/artifactory/{repo}" + + +def test_build_commit_with_the_real_helper_falls_back_to_urm( + monkeypatch: pytest.MonkeyPatch, tmp_path: Path, capfd: pytest.CaptureFixture[str] +) -> None: + content = _BUILD_INFO.encode() + storage_info = {"checksums": {"sha1": hashlib.sha1(content, usedforsecurity=False).hexdigest()}} + urm_routes = { + f"/artifactory/api/storage/sw-tensorrt-generic/{_LAYOUT_PATH}": json.dumps( + storage_info + ).encode(), + f"/artifactory/sw-tensorrt-generic/{_LAYOUT_PATH}": content, + } + # curl matches netrc entries by host name, so each location has its own loopback address. + servers = [] + try: + for host, routes in (("127.0.0.1", {}), ("127.0.0.2", urm_routes)): + servers.append(_ArtifactServer(host, routes)) + except OSError as error: + for server in servers: + server.server_close() + pytest.skip(f"cannot listen on the 127.0.0.0/8 loopback range: {error}") + aws, urm = servers + for server in servers: + threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start() + netrc = tmp_path / "artifactory-netrc" + netrc.write_text("machine 127.0.0.1 login ci-reader password reader-token-value\n") + netrc.chmod(0o600) + for name in ( + "http_proxy", + "https_proxy", + "all_proxy", + "HTTP_PROXY", + "HTTPS_PROXY", + "ALL_PROXY", + ): + monkeypatch.delenv(name, raising=False) + monkeypatch.delenv("LLM_ARTIFACTS_PRIMARY", raising=False) + monkeypatch.setenv("LLM_ARTIFACTS_AWS_URL", aws.base_url("sw-tensorrt-llm-generic-local")) + monkeypatch.setenv("LLM_ARTIFACTS_URM_URL", urm.base_url("sw-tensorrt-generic")) + monkeypatch.setenv("ARTIFACTORY_NETRC", str(netrc)) + monkeypatch.setenv("ARTIFACTORY_CREDENTIAL_ID", _CREDENTIAL_ID) + + try: + commit = artifact.build_commit(_BUILD) + finally: + for server in servers: + server.shutdown() + server.server_close() + + assert commit == _COMMIT + captured = capfd.readouterr() + assert captured.out == "" + assert f"fetch_build_artifact: served-by=urm path={_LAYOUT_PATH}" in captured.err + assert "reader-token-value" not in captured.err + + if __name__ == "__main__": unittest.main() diff --git a/tests/unittest/scripts/test_bolt_fetch_verified.py b/tests/unittest/scripts/test_bolt_fetch_verified.py new file mode 100644 index 000000000000..f0e01e1ec786 --- /dev/null +++ b/tests/unittest/scripts/test_bolt_fetch_verified.py @@ -0,0 +1,840 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for ``fetch_verified`` in ``jenkins/BoltProfileGen.groovy``. + +BOLT profile generation downloads its input build and llvm-bolt on a Slurm cluster with shell +functions that the pipeline keeps in the Groovy string ``boltFetchLib``. The tests build that +string from the Groovy source the way Groovy does, source it in bash, and call ``fetch_verified`` +the way the stages do: as a plain command in a ``set -e`` script. Two loopback HTTP servers stand +in for the ``aws`` Artifactory location and for the CDN that it redirects downloads to. Each +listens on its own host, because curl matches netrc entries by host name, and no request leaves +the loopback interface. The tests cut the retry counts and delays in the extracted text, so that +each test takes a few seconds at most. +""" + +from __future__ import annotations + +import base64 +import hashlib +import os +import re +import socket +import subprocess +import threading +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import cast + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +BOLT_PROFILE_GEN = REPO_ROOT / "jenkins" / "BoltProfileGen.groovy" +# The statement that defines the library, through its opening quotes, and the text after its +# closing quotes. +LIBRARY_START = "def boltFetchLib = '''" +LIBRARY_END = "'''.stripIndent()" + +HOSTS = {"aws": "127.0.0.1", "cdn": "127.0.0.3"} +FILE_PATH = ( + "/artifactory/sw-tensorrt-llm-generic-local/llm-artifacts/LLM/main/L0_PostMerge/88/" + "TensorRT-LLM.tar.gz" +) +# A signed CDN URL: the signature is a credential of its own and must not reach the log. +CDN_PATH = f"/cdn{FILE_PATH}?Expires=1798761600&Signature=c2lnbmVkLWNkbi11cmw&Key-Pair-Id=KTEST" +CONTENT = b"packaged build " * 64 + +NETRC_LOGIN = "ci-reader" +NETRC_PASSWORD = "Pw-5e0b7a14-must-never-be-printed" +BASIC_TOKEN = base64.b64encode(f"{NETRC_LOGIN}:{NETRC_PASSWORD}".encode()).decode() +AUTHORIZATION = f"Basic {BASIC_TOKEN}" + +# Asks Artifactory to send a file itself instead of redirecting the download to its CDN. +REDIRECT_OPT_OUT = "x-jfrog-download-redirect-to" + +# The retry counts and delays in the library, and the values the tests use instead: curl's options +# in the parallel download, and fetch_verified's own retries of a failed single-stream attempt. +# curl waits --retry-delay seconds before each retry; without that option it waits 1 s, then twice +# as long. +FAST_RETRIES = ( + (r"--retry \d+", "--retry 1"), + (r"--retry-delay \d+", "--retry-delay 1"), + (r"\bretries=\d+", "retries=1"), + (r"\bretry_delay=\d+", "retry_delay=1"), +) +# The single-stream attempts that fetch_verified makes with FAST_RETRIES before it gives up. +ATTEMPTS = 2 + +# The escapes of a single-quoted Groovy string. A backslash at the end of a line joins the next +# line to it. +_GROOVY_ESCAPES = { + "b": "\b", + "f": "\f", + "n": "\n", + "r": "\r", + "t": "\t", + "\\": "\\", + "'": "'", + '"': '"', + "$": "$", + "\n": "", +} +_GROOVY_ESCAPE = re.compile(r"\\(u[0-9A-Fa-f]{4}|[0-3][0-7]{0,2}|[4-7][0-7]?|.)", re.DOTALL) + +_PROXY_VARIABLES = {"http_proxy", "https_proxy", "all_proxy", "ftp_proxy", "no_proxy"} + +# Calls fetch_verified the way the stages do: as a plain command in a `set -e` script. The EXIT +# trap records the shell functions that were running when the script ended. There are none when +# fetch_verified returned, and there are some when `set -e` cut it short. +_SCRIPT = """\ +set -e +trap 'printf "%s" "${FUNCNAME[*]:-}" > "$EXITED_IN"' EXIT +. "$BOLT_FETCH_LIB" +_verify_test() { [ "$(sha1sum < "$1" | cut -d' ' -f1)" = "$EXPECTED_SHA1" ]; } +fetch_verified "$URL" "$DEST" _verify_test +echo "fetch_verified returned 0" +""" + + +# --------------------------------------------------------------------------- +# The library as Groovy builds it +# --------------------------------------------------------------------------- + + +def _groovy_unescape(literal: str) -> str: + """Return the value of the text between the quotes of a single-quoted Groovy string.""" + + def replace(match: re.Match[str]) -> str: + escape = match.group(1) + if len(escape) == 5: + return chr(int(escape[1:], 16)) + if escape[0] in "01234567": + return chr(int(escape, 8)) + if escape not in _GROOVY_ESCAPES: + pytest.fail(f"boltFetchLib contains \\{escape}, which is not a Groovy escape") + return _GROOVY_ESCAPES[escape] + + return _GROOVY_ESCAPE.sub(replace, literal) + + +def _strip_indent(text: str) -> str: + """Return ``text`` as Groovy's ``String.stripIndent()`` does. + + The smallest indent among the lines that are not blank is removed from every line, and blank + lines become empty. + """ + lines = text.split("\n") + indent = min((len(line) - len(line.lstrip()) for line in lines if line.strip()), default=0) + return "\n".join(line[indent:] if line.strip() else "" for line in lines) + + +def _bolt_fetch_lib() -> str: + """Return the shell text of ``boltFetchLib`` as the pipeline builds it.""" + source = BOLT_PROFILE_GEN.read_text() + start = source.find(LIBRARY_START) + end = source.find(LIBRARY_END, start) if start != -1 else -1 + if end == -1 or source.count(LIBRARY_START) != 1: + pytest.fail( + f"cannot find the boltFetchLib string in {BOLT_PROFILE_GEN}: expected one " + f"{LIBRARY_START!r} and then {LIBRARY_END!r}" + ) + literal = source[start + len(LIBRARY_START) : end] + # Groovy ends the string at the first ''' in it. + assert "'''" not in literal, "boltFetchLib contains ''', which ends the Groovy string early" + return _strip_indent(_groovy_unescape(literal)) + + +def _with_fast_retries(library: str) -> str: + """Return ``library`` with every curl retry count and delay replaced as ``FAST_RETRIES`` says. + + Each pattern must match, and no other retry value may remain, so that a change to how the + library passes these options fails here instead of slowing the tests down. + """ + for pattern, replacement in FAST_RETRIES: + library, count = re.subn(rf"{pattern}\b", replacement, library) + assert count, f"boltFetchLib has no {pattern!r} to replace" + values = re.findall(r"(?:--retry(?:-delay)?(?:\s+|=)|\bretries=|\bretry_delay=)(\S+)", library) + assert set(values) == {"1"}, f"retry values that FAST_RETRIES misses: {values}" + return library + + +# --------------------------------------------------------------------------- +# Fake servers +# --------------------------------------------------------------------------- + + +@dataclass +class Reply: + """An HTTP response for a fake server to send; a HEAD request gets it without the body.""" + + status: int + body: bytes = b"" + headers: dict[str, str] = field(default_factory=dict) + # A Content-Length above the body's length drops the connection partway through the body. + length: int | None = None + # After the body, keep the connection open without sending more until the test releases it. + stall: bool = False + + +@dataclass +class Request: + """A request that a fake server received; header names are lower case.""" + + server: str + method: str + path: str + headers: dict[str, str] + + @property + def opted_out(self) -> bool: + return self.headers.get(REDIRECT_OPT_OUT) == "None" + + +Route = Reply | Callable[[Request], Reply] + + +class _Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + self._answer(with_body=True) + + def do_HEAD(self) -> None: + self._answer(with_body=False) + + def _answer(self, *, with_body: bool) -> None: + server = cast(_Server, self.server) + headers = {name.lower(): value for name, value in self.headers.items()} + reply = server.fake.answer(Request(server.name, self.command, self.path, headers)) + self.send_response(reply.status) + for name, value in reply.headers.items(): + self.send_header(name, value) + length = len(reply.body) if reply.length is None else reply.length + self.send_header("Content-Length", str(length)) + self.end_headers() + if with_body: + self.wfile.write(reply.body) + if reply.stall: + self.wfile.flush() + server.fake.hold() + + def log_message(self, format: str, *args: object) -> None: + """Keep access logs out of the test output.""" + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, name: str, fake: FakeServers) -> None: + super().__init__((HOSTS[name], 0), _Handler) + self.name = name + self.fake = fake + + +class FakeServers: + """The aws location and its CDN, each a loopback HTTP server on its own host. + + Routes map ``(server, path with query)`` to a reply. A request without a route gets 404. + """ + + def __init__(self) -> None: + self.routes: dict[tuple[str, str], Route] = {} + self.requests: list[Request] = [] + self._lock = threading.Lock() + self._released = threading.Event() + self._held: list[threading.Thread] = [] + self._servers: dict[str, _Server] = {} + try: + for name in HOSTS: + self._servers[name] = _Server(name, self) + except OSError as error: + for server in self._servers.values(): + server.server_close() + pytest.skip(f"cannot listen on the 127.0.0.0/8 loopback range: {error}") + # A short poll interval keeps shutdown(), which waits for the next poll, fast. + self._threads = [ + threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True) + for server in self._servers.values() + ] + for thread in self._threads: + thread.start() + + def close(self) -> None: + self.release() + for server in self._servers.values(): + server.shutdown() + server.server_close() + for thread in self._threads: + thread.join() + + def hold(self) -> None: + """Keep the calling handler's connection open until ``release()``.""" + with self._lock: + self._held.append(threading.current_thread()) + self._released.wait(timeout=60) + + def release(self) -> None: + """Let the handlers in ``hold()`` finish, and wait until they have. + + A test that stalls a transfer calls this before it returns, because the thread leak check + runs before fixtures are torn down. + """ + self._released.set() + with self._lock: + held = list(self._held) + for thread in held: + thread.join(timeout=10) + + def host(self, name: str) -> str: + """Return the server's ``host:port``.""" + host, port = self._servers[name].server_address[:2] + return f"{host}:{port}" + + def url(self, name: str, path: str) -> str: + return f"http://{self.host(name)}{path}" + + def answer(self, request: Request) -> Reply: + with self._lock: + self.requests.append(request) + route = self.routes.get((request.server, request.path)) + if route is None: + return Reply(404) + return route if isinstance(route, Reply) else route(request) + + def received(self, name: str) -> list[Request]: + with self._lock: + return [request for request in self.requests if request.server == name] + + def gets(self) -> list[tuple[str, bool]]: + """Return the GET requests in the order received, as ``(server, opted out)``.""" + with self._lock: + return [(r.server, r.opted_out) for r in self.requests if r.method == "GET"] + + +def _redirect_to(fake: FakeServers, location: str, *, direct: Reply) -> None: + """Make aws redirect downloads to ``location``, and answer ``direct`` to those that opt out.""" + redirect = Reply(302, headers={"Location": location}) + fake.routes[("aws", FILE_PATH)] = lambda request: direct if request.opted_out else redirect + + +def _redirect_to_cdn(fake: FakeServers, *, cdn: Route, direct: Reply) -> None: + """Make aws redirect downloads to the CDN, which answers with ``cdn``.""" + _redirect_to(fake, fake.url("cdn", CDN_PATH), direct=direct) + fake.routes[("cdn", CDN_PATH)] = cdn + + +def _ranged(content: bytes) -> Callable[[Request], Reply]: + """Build a route that serves ``content`` and answers byte-range requests.""" + + def answer(request: Request) -> Reply: + headers = {"Accept-Ranges": "bytes"} + match = re.fullmatch(r"bytes=(\d+)-(\d+)", request.headers.get("range", "")) + if match is None: + return Reply(200, content, headers) + first, last = int(match[1]), int(match[2]) + headers["Content-Range"] = f"bytes {first}-{last}/{len(content)}" + return Reply(206, content[first : last + 1], headers) + + return answer + + +def _in_turn(*replies: Reply) -> Callable[[Request], Reply]: + """Build a route that answers successive GET requests with ``replies``, then repeats the last. + + A HEAD request gets the reply that the next GET request will get. + """ + gets = 0 + + def answer(request: Request) -> Reply: + nonlocal gets + reply = replies[min(gets, len(replies) - 1)] + if request.method == "GET": + gets += 1 + return reply + + return answer + + +def _unused_port(host: str) -> int: + with socket.socket() as sock: + sock.bind((host, 0)) + return sock.getsockname()[1] + + +# --------------------------------------------------------------------------- +# Running fetch_verified +# --------------------------------------------------------------------------- + + +@dataclass +class Result: + """The outcome of a script that called ``fetch_verified``.""" + + returncode: int + stdout: str + stderr: str + # The shell functions that were running when the script exited: none when fetch_verified + # returned, and fetch_verified among them when `set -e` cut it short. + exited_in: str + + +class FetchVerified: + """Runs ``fetch_verified`` from the extracted library against the fake servers.""" + + def __init__(self, tmp_path: Path, fake: FakeServers, library: str) -> None: + self.fake = fake + self.dest = tmp_path / "out" / "TensorRT-LLM.tar.gz" + self._cwd = tmp_path / "cwd" + self._tmpdir = tmp_path / "tmp" + self._exited_in = tmp_path / "exited_in" + home = tmp_path / "home" + for directory in (self.dest.parent, self._cwd, self._tmpdir, home): + directory.mkdir() + lib = tmp_path / "bolt_fetch_lib.sh" + lib.write_text(library) + self._script = tmp_path / "fetch.sh" + self._script.write_text(_SCRIPT) + + entry = f"machine {HOSTS['aws']} login {NETRC_LOGIN} password {NETRC_PASSWORD}\n" + netrc = tmp_path / "netrc" + # curl reads ~/.netrc only when told to, so a download without FETCH_NETRC sends nothing. + for path in (netrc, home / ".netrc"): + path.write_text(entry) + path.chmod(0o600) + + env = { + name: value + for name, value in os.environ.items() + if name.lower() not in _PROXY_VARIABLES + and name not in ("CURL_HOME", "XDG_CONFIG_HOME", "FETCH_NETRC", "PARTS") + } + env.update( + { + "HOME": str(home), + "TMPDIR": str(self._tmpdir), + "BOLT_FETCH_LIB": str(lib), + "EXITED_IN": str(self._exited_in), + "URL": fake.url("aws", FILE_PATH), + "DEST": str(self.dest), + "FETCH_NETRC": str(netrc), + } + ) + self._env = env + + def run(self, expected: bytes = CONTENT, **variables: str | None) -> Result: + """Download the aws file to ``dest``, with a check that it holds ``expected``. + + Every run checks that the output has neither a signed CDN URL nor the netrc password, and + that no file is left anywhere but at ``dest``. + + Args: + expected: The content that the check accepts. + **variables: Variables to set for the script, such as FETCH_NETRC; ``None`` unsets one. + """ + env = dict(self._env, EXPECTED_SHA1=hashlib.sha1(expected).hexdigest()) + for name, value in variables.items(): + if value is None: + env.pop(name, None) + else: + env[name] = value + self._exited_in.unlink(missing_ok=True) + + proc = subprocess.run( + ["bash", str(self._script)], + env=env, + cwd=self._cwd, + capture_output=True, + text=True, + timeout=60, + ) + + output = proc.stdout + proc.stderr + assert "Signature=" not in output, "a signed CDN URL reached the output" + assert NETRC_PASSWORD not in output + assert BASIC_TOKEN not in output + assert list(self._tmpdir.iterdir()) == [], "temporary files were left behind" + assert list(self._cwd.iterdir()) == [], "files were written outside DEST" + leftovers = [path.name for path in self.dest.parent.iterdir() if path != self.dest] + assert leftovers == [], "partial files were left next to DEST" + assert self._exited_in.exists(), "the script's EXIT trap did not run" + return Result(proc.returncode, proc.stdout, proc.stderr, self._exited_in.read_text()) + + +@pytest.fixture(scope="module") +def library() -> str: + return _with_fast_retries(_bolt_fetch_lib()) + + +@pytest.fixture +def fake() -> Iterator[FakeServers]: + servers = FakeServers() + yield servers + servers.close() + + +@pytest.fixture +def fetch(tmp_path: Path, fake: FakeServers, library: str) -> FetchVerified: + return FetchVerified(tmp_path, fake, library) + + +# --------------------------------------------------------------------------- +# The extracted library +# --------------------------------------------------------------------------- + + +def test_groovy_string_processing() -> None: + literal = "\n a \\\\\n b\\\n c \\$x \\'q\\' \\u0041\\t\n\n " + + assert _groovy_unescape(literal) == "\n a \\\n b c $x 'q' A\t\n\n " + assert _strip_indent(_groovy_unescape(literal)) == "\na \\\n b c $x 'q' A\t\n\n" + + +def test_library_is_valid_bash(library: str, tmp_path: Path) -> None: + lib = tmp_path / "bolt_fetch_lib.sh" + lib.write_text(library) + + proc = subprocess.run(["bash", "-n", str(lib)], capture_output=True, text=True, timeout=30) + + assert proc.returncode == 0, proc.stderr + # Groovy turned each doubled backslash into one. + assert "printf '%s\\n'" in library + assert "fetch_verified() {" in library + + +def test_long_scripts_are_not_sent_on_the_ssh_command_line() -> None: + """Long stage scripts run from a file copied to the cluster. + + ``b64BashRemoteCmd`` and ``b64BashRemoteCmdStdin`` put the whole script, base64-encoded, into + one word of the remote command, and a csh/tcsh login shell rejects a word that long with + "Word too long". Every script that carries the library is long, and so is ``publish``, the + script that publishes the BOLTed build. + """ + source = BOLT_PROFILE_GEN.read_text() + carriers = re.findall( + r'def (\w+) = """(?:(?!""").)*"""\.stripIndent\(\) \+ boltFetchLib \+', source, re.S + ) + + assert sorted(carriers) == ["llvmStage", "tarStage"] + for name in [*carriers, "publish"]: + assert not re.search(rf"b64BashRemoteCmd(?:Stdin)?\(\s*{name}\b", source), name + assert re.search(rf"bashRemoteScriptFile\(pipeline, remote, {name},", source), name + + +@pytest.mark.parametrize( + ("url", "host"), + [ + ("https://artifactory.nvidia.com/artifactory/repo/a.tar.gz", "artifactory.nvidia.com"), + ("http://127.0.0.3:8080/cdn/a.tar.gz?Signature=c2ln", "127.0.0.3:8080"), + ("https://reader@GitHub.com/llvm/a.tar.xz#part", "github.com"), + ("https://d1.cloudfront.net?Signature=c2ln", "d1.cloudfront.net"), + ], +) +def test_url_host(library: str, tmp_path: Path, url: str, host: str) -> None: + lib = tmp_path / "bolt_fetch_lib.sh" + lib.write_text(library) + + proc = subprocess.run( + ["bash", "-c", '. "$1" && _bolt_url_host "$2"', "bash", str(lib), url], + capture_output=True, + text=True, + timeout=30, + ) + + assert proc.returncode == 0, proc.stderr + assert proc.stdout == f"{host}\n" + + +# --------------------------------------------------------------------------- +# The single-stream download +# --------------------------------------------------------------------------- + + +def _retrying(reason: str) -> str: + """Return the line that fetch_verified logs when its first single-stream attempt fails.""" + return f"[INFO] single-stream attempt 1 of {ATTEMPTS}: {reason}; retrying in 1s." + + +def _gave_up(reason: str) -> str: + """Return the start of the line that fetch_verified logs when its last attempt fails.""" + return f"[ERROR] single-stream attempt {ATTEMPTS} of {ATTEMPTS}: {reason}: " + + +def test_cdn_rejection_is_retried_once_without_the_redirect(fetch: FetchVerified) -> None: + fake = fetch.fake + rejected = Reply(403, b"AccessDenied") + _redirect_to_cdn(fake, cdn=rejected, direct=Reply(200, CONTENT)) + + result = fetch.run() + + assert result.returncode == 0, result.stderr + assert result.exited_in == "" + assert fetch.dest.read_bytes() == CONTENT + assert fake.gets() == [("aws", False), ("cdn", False), ("aws", True)] + # Exactly the requests after the CDN rejected the download opt out of the redirect. + requests = fake.requests + cdn_get = next(i for i, r in enumerate(requests) if r.server == "cdn" and r.method == "GET") + assert [r.opted_out for r in requests] == [i > cdn_get for i in range(len(requests))] + assert ( + f"[INFO] {fake.host('cdn')} rejected the redirected download (HTTP 403); " + f"downloading from {fake.host('aws')} without the redirect." + ) in result.stdout.splitlines() + assert "single-stream attempt" not in result.stdout + result.stderr + + +@pytest.mark.parametrize("status", [401, 404, 500, 520]) +def test_cdn_failure_other_than_403_is_retried_through_the_redirect_then_fails( + fetch: FetchVerified, status: int +) -> None: + fake = fetch.fake + _redirect_to_cdn(fake, cdn=Reply(status), direct=Reply(200, CONTENT)) + + result = fetch.run() + + assert result.returncode == 1 + assert result.exited_in == "" + assert not fetch.dest.exists() + assert not any(request.opted_out for request in fake.requests) + # Each attempt starts again from the aws URL, which redirects to the CDN again. + assert fake.gets() == [("aws", False), ("cdn", False)] * ATTEMPTS + reason = f"download answered HTTP {status} from {fake.host('cdn')}" + assert _retrying(reason) in result.stdout.splitlines() + assert _gave_up(reason) in result.stderr + + +def test_attempt_after_a_failed_one_can_succeed(fetch: FetchVerified) -> None: + fake = fetch.fake + cdn = _in_turn(Reply(520), Reply(200, CONTENT)) + _redirect_to_cdn(fake, cdn=cdn, direct=Reply(500)) + + result = fetch.run() + + assert result.returncode == 0, result.stderr + assert result.exited_in == "" + assert fetch.dest.read_bytes() == CONTENT + assert fake.gets() == [("aws", False), ("cdn", False)] * 2 + reason = f"download answered HTTP 520 from {fake.host('cdn')}" + assert _retrying(reason) in result.stdout.splitlines() + + +def test_cdn_stall_is_retried_through_the_redirect_then_fails( + tmp_path: Path, fake: FakeServers, library: str +) -> None: + # curl ends a transfer that moves fewer than --speed-limit bytes a second for --speed-time + # seconds. + stalling_library, count = re.subn(r"--speed-time \d+\b", "--speed-time 1", library) + assert count, "boltFetchLib has no --speed-time to shorten" + fetch = FetchVerified(tmp_path, fake, stalling_library) + stalled = Reply(200, CONTENT[:100], length=len(CONTENT), stall=True) + _redirect_to_cdn(fake, cdn=stalled, direct=Reply(200, CONTENT)) + + try: + result = fetch.run() + finally: + fake.release() + + assert result.returncode == 1 + assert result.exited_in == "" + assert not fetch.dest.exists() + assert not any(request.opted_out for request in fake.requests) + assert fake.gets() == [("aws", False), ("cdn", False)] * ATTEMPTS + assert _gave_up(f"download from {fake.host('cdn')} failed (curl exit 28)") in result.stderr + + +def test_rejection_by_the_location_itself_is_not_retried_without_the_redirect( + fetch: FetchVerified, +) -> None: + fake = fetch.fake + fake.routes[("aws", FILE_PATH)] = lambda request: ( + Reply(200, CONTENT) if request.opted_out else Reply(403) + ) + + result = fetch.run() + + assert result.returncode == 1 + assert result.exited_in == "" + assert not fetch.dest.exists() + assert not any(request.opted_out for request in fake.requests) + # It fails the attempt like any other status, and the next attempt asks the same way. + assert fake.gets() == [("aws", False)] * ATTEMPTS + assert fake.received("cdn") == [] + assert _gave_up(f"download answered HTTP 403 from {fake.host('aws')}") in result.stderr + + +def test_download_without_a_redirect_does_not_opt_out(fetch: FetchVerified) -> None: + fake = fetch.fake + fake.routes[("aws", FILE_PATH)] = Reply(200, CONTENT) + + result = fetch.run() + + assert result.returncode == 0, result.stderr + assert result.exited_in == "" + assert fetch.dest.read_bytes() == CONTENT + assert "single-stream curl" in result.stdout + assert fake.gets() == [("aws", False)] + assert not any(request.opted_out for request in fake.requests) + assert fake.received("cdn") == [] + + +def test_credentials_go_to_the_location_only(fetch: FetchVerified) -> None: + fake = fetch.fake + _redirect_to_cdn(fake, cdn=Reply(403), direct=Reply(200, CONTENT)) + + result = fetch.run() + + assert result.returncode == 0, result.stderr + assert {request.headers.get("authorization") for request in fake.received("aws")} == { + AUTHORIZATION + } + assert {request.headers.get("authorization") for request in fake.received("cdn")} == {None} + + +@pytest.mark.parametrize("fetch_netrc", ["", None], ids=["empty", "unset"]) +def test_download_without_fetch_netrc_is_anonymous( + fetch: FetchVerified, fetch_netrc: str | None +) -> None: + fake = fetch.fake + _redirect_to_cdn(fake, cdn=Reply(200, CONTENT), direct=Reply(500)) + + result = fetch.run(FETCH_NETRC=fetch_netrc) + + assert result.returncode == 0, result.stderr + assert fetch.dest.read_bytes() == CONTENT + assert fake.gets() == [("aws", False), ("cdn", False)] + assert {request.headers.get("authorization") for request in fake.requests} == {None} + + +def test_parallel_download_follows_the_redirect(fetch: FetchVerified) -> None: + fake = fetch.fake + _redirect_to_cdn(fake, cdn=_ranged(CONTENT), direct=Reply(500)) + + result = fetch.run() + + assert result.returncode == 0, result.stderr + assert fetch.dest.read_bytes() == CONTENT + assert "[INFO] fetched via 16-way parallel: " in result.stdout + assert not any(request.opted_out for request in fake.requests) + assert {request.headers.get("authorization") for request in fake.received("aws")} == { + AUTHORIZATION + } + assert {request.headers.get("authorization") for request in fake.received("cdn")} == {None} + + +# --------------------------------------------------------------------------- +# Failures under set -e +# --------------------------------------------------------------------------- + + +@dataclass +class Failure: + """How a download fails: the start of the line it logs, and its GET requests in order.""" + + message: str + gets: list[tuple[str, bool]] + + +def _cdn_server_error(fake: FakeServers) -> Failure: + _redirect_to_cdn(fake, cdn=Reply(503), direct=Reply(200, CONTENT)) + return Failure( + _gave_up(f"download answered HTTP 503 from {fake.host('cdn')}"), + [("aws", False), ("cdn", False)] * ATTEMPTS, + ) + + +def _cdn_refuses_connections(fake: FakeServers) -> Failure: + closed = f"{HOSTS['cdn']}:{_unused_port(HOSTS['cdn'])}" + _redirect_to(fake, f"http://{closed}{CDN_PATH}", direct=Reply(200, CONTENT)) + return Failure( + _gave_up(f"download from {closed} failed (curl exit 7)"), [("aws", False)] * ATTEMPTS + ) + + +def _cdn_drops_the_connection(fake: FakeServers) -> Failure: + dropped = Reply(200, CONTENT[:100], length=len(CONTENT)) + _redirect_to_cdn(fake, cdn=dropped, direct=Reply(200, CONTENT)) + return Failure( + _gave_up(f"download from {fake.host('cdn')} failed (curl exit 18)"), + [("aws", False), ("cdn", False)] * ATTEMPTS, + ) + + +def _location_rejects(fake: FakeServers) -> Failure: + fake.routes[("aws", FILE_PATH)] = Reply(403) + return Failure( + _gave_up(f"download answered HTTP 403 from {fake.host('aws')}"), + [("aws", False)] * ATTEMPTS, + ) + + +def _download_without_the_redirect_fails(fake: FakeServers) -> Failure: + _redirect_to_cdn(fake, cdn=Reply(403), direct=Reply(500)) + aws = fake.host("aws") + return Failure( + _gave_up(f"download from {aws} without the redirect answered HTTP 500 from {aws}"), + [("aws", False), ("cdn", False), ("aws", True)] * ATTEMPTS, + ) + + +def _download_without_the_redirect_is_redirected(fake: FakeServers) -> Failure: + redirect = Reply(302, headers={"Location": fake.url("cdn", CDN_PATH)}) + _redirect_to_cdn(fake, cdn=Reply(403), direct=redirect) + return Failure( + _gave_up( + f"download from {fake.host('aws')} without the redirect answered HTTP 403 " + f"from {fake.host('cdn')}" + ), + [("aws", False), ("cdn", False), ("aws", True), ("cdn", True)] * ATTEMPTS, + ) + + +def _download_fails_the_check(fake: FakeServers) -> Failure: + _redirect_to_cdn(fake, cdn=Reply(200, b"corrupted package"), direct=Reply(200, CONTENT)) + return Failure("[ERROR] verification failed: ", [("aws", False), ("cdn", False)]) + + +def _download_without_the_redirect_fails_the_check(fake: FakeServers) -> Failure: + _redirect_to_cdn(fake, cdn=Reply(403), direct=Reply(200, b"corrupted package")) + return Failure("[ERROR] verification failed: ", [("aws", False), ("cdn", False), ("aws", True)]) + + +@pytest.mark.parametrize( + "failure", + [ + _cdn_server_error, + _cdn_refuses_connections, + _cdn_drops_the_connection, + _location_rejects, + _download_without_the_redirect_fails, + _download_without_the_redirect_is_redirected, + _download_fails_the_check, + _download_without_the_redirect_fails_the_check, + ], + ids=lambda failure: failure.__name__.strip("_").replace("_", "-"), +) +def test_failure_returns_1_under_set_e_and_removes_dest( + fetch: FetchVerified, failure: Callable[[FakeServers], Failure] +) -> None: + expected = failure(fetch.fake) + fetch.dest.write_bytes(b"earlier download") + + result = fetch.run() + + assert result.returncode == 1, result.stderr + assert result.exited_in == "", f"set -e ended the script inside {result.exited_in}" + assert "fetch_verified returned 0" not in result.stdout + assert expected.message in result.stderr + assert fetch.fake.gets() == expected.gets + assert not fetch.dest.exists() diff --git a/tests/unittest/scripts/test_fetch_build_artifact.py b/tests/unittest/scripts/test_fetch_build_artifact.py new file mode 100644 index 000000000000..0563d108e9ae --- /dev/null +++ b/tests/unittest/scripts/test_fetch_build_artifact.py @@ -0,0 +1,1301 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for ``jenkins/scripts/fetch_build_artifact.sh``. + +The helper reads one build artifact from the primary Artifactory location and falls back to the +secondary location only when the primary does not have it. Three loopback HTTP servers stand in +for the ``aws`` and ``urm`` locations and for the CDN that ``aws`` redirects downloads to. Each +listens on its own host name, because curl matches netrc entries by host name. A stub ``sleep`` on +``PATH`` records the back-off instead of waiting, and no request leaves the loopback interface. +""" + +from __future__ import annotations + +import ast +import base64 +import hashlib +import json +import os +import re +import socket +import subprocess +import threading +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import cast + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +FETCH_BUILD_ARTIFACT = REPO_ROOT / "jenkins" / "scripts" / "fetch_build_artifact.sh" + +HOSTS = {"aws": "127.0.0.1", "urm": "127.0.0.2", "cdn": "127.0.0.3"} +REPOS = {"aws": "sw-tensorrt-llm-generic-local", "urm": "sw-tensorrt-generic"} + +BUILD_DIR = "llm-artifacts/LLM/main/L0_PostMerge/88" +LAYOUT_PATH = f"{BUILD_DIR}/TensorRT-LLM.tar.gz" +# Not one of the three in-scope kinds: the helper reads any file under llm-artifacts/. +OTHER_FILE = "tensorrt_llm-1.2.0-cp312-cp312-linux_x86_64.whl" +LEGACY_PATH = f"sw-tensorrt-generic/{LAYOUT_PATH}" +SYNOPSIS = "fetch_build_artifact.sh [--properties | --exists] []" + +NETRC_LOGIN = "ci-reader" +NETRC_PASSWORD = "Pw-6c1f0e93-must-never-be-printed" +URM_NETRC_PASSWORD = "Pw-urm-41d8b2-must-never-be-sent" +CREDENTIAL_ID = "test-reader-credential" +BASIC_TOKEN = base64.b64encode(f"{NETRC_LOGIN}:{NETRC_PASSWORD}".encode()).decode() +AUTHORIZATION = f"Basic {BASIC_TOKEN}" + +# Asks Artifactory to send a file itself instead of redirecting the download to its CDN. +REDIRECT_OPT_OUT = "x-jfrog-download-redirect-to" +# A signed CDN URL: the signature is a credential of its own and must not reach the log. +CDN_PATH = f"/cdn/{LAYOUT_PATH}?Expires=1798761600&Signature=c2lnbmVkLWNkbi11cmw&Key-Pair-Id=KTEST" + +PROPERTIES = json.dumps( + {"properties": {"bolt.ref": ["0123456789abcdef"]}, "uri": "storage-uri"}, indent=2 +).encode() + +# Log phrases that make the shared Jenkins library retry a step as a network fault. The message +# for a rejected credential must contain none of them. +RETRY_SIGNATURES = ( + "403 Forbidden", + "URL returned error", + "HTTP response code said error", + "404 Not Found", + "ERROR 404", + "HTTP 50", + "ERROR 50", + "Server Error", + "Error downloading", + "Connection", + "Read timed out", + "Too Many Requests", +) + +_PROXY_VARIABLES = {"http_proxy", "https_proxy", "all_proxy", "ftp_proxy"} + + +@dataclass +class Reply: + """An HTTP response for a fake server to send.""" + + status: int + body: bytes = b"" + headers: dict[str, str] = field(default_factory=dict) + + +@dataclass +class Request: + """A request that a fake server received; header names are lower case.""" + + server: str + path: str + headers: dict[str, str] + + @property + def authorization(self) -> str | None: + return self.headers.get("authorization") + + +Route = Reply | Callable[[Request], Reply] + + +def _sha1(content: bytes) -> str: + return hashlib.sha1(content).hexdigest() + + +def _storage_info( + location: str, content: bytes, sha1: str, layout_path: str = LAYOUT_PATH +) -> bytes: + """Build a storage API answer for the artifact at ``layout_path``. + + ``originalChecksums`` comes first and holds a different SHA-1, so only a parser that reads the + ``checksums`` object gets the right value. + """ + info = { + "repo": REPOS[location], + "path": f"/{layout_path}", + "created": "2026-10-16T08:00:00.000Z", + "mimeType": "application/x-gzip", + "size": str(len(content)), + "originalChecksums": {"sha1": "f" * 40, "sha256": "e" * 64}, + "checksums": {"sha1": sha1, "sha256": hashlib.sha256(content).hexdigest()}, + "uri": f"storage-uri/{layout_path}", + } + return json.dumps(info, indent=2).encode() + + +def _in_turn(*replies: Reply) -> Callable[[Request], Reply]: + """Build a route that answers with ``replies`` in order, then repeats the last one.""" + remaining = list(replies) + + def answer(request: Request) -> Reply: + return remaining.pop(0) if len(remaining) > 1 else remaining[0] + + return answer + + +def _unused_port(host: str) -> int: + with socket.socket() as sock: + sock.bind((host, 0)) + return sock.getsockname()[1] + + +class _Handler(BaseHTTPRequestHandler): + def do_GET(self) -> None: + server = cast(_Server, self.server) + headers = {name.lower(): value for name, value in self.headers.items()} + reply = server.artifactory.answer(Request(server.name, self.path, headers)) + self.send_response(reply.status) + for name, value in reply.headers.items(): + self.send_header(name, value) + self.send_header("Content-Length", str(len(reply.body))) + self.end_headers() + self.wfile.write(reply.body) + + def log_message(self, format: str, *args: object) -> None: + """Keep access logs out of the test output.""" + + +class _Server(ThreadingHTTPServer): + daemon_threads = True + + def __init__(self, name: str, artifactory: FakeArtifactory) -> None: + super().__init__((HOSTS[name], 0), _Handler) + self.name = name + self.artifactory = artifactory + + +class FakeArtifactory: + """The two locations and the CDN, each a loopback HTTP server on its own host. + + Routes map ``(server, path with query)`` to a reply. A request without a route gets 404, which + is how Artifactory answers for a missing file. + """ + + def __init__(self) -> None: + self.routes: dict[tuple[str, str], Route] = {} + self.requests: list[Request] = [] + self._lock = threading.Lock() + self._servers: dict[str, _Server] = {} + try: + for name in HOSTS: + self._servers[name] = _Server(name, self) + except OSError as error: + self.close() + pytest.skip(f"cannot listen on the 127.0.0.0/8 loopback range: {error}") + for server in self._servers.values(): + # A short poll interval keeps shutdown(), which waits for the next poll, fast. + threading.Thread(target=server.serve_forever, args=(0.01,), daemon=True).start() + + def close(self) -> None: + for server in self._servers.values(): + server.shutdown() + server.server_close() + + def origin(self, name: str) -> str: + host, port = self._servers[name].server_address[:2] + return f"http://{host}:{port}" + + def host(self, name: str) -> str: + return self.origin(name).removeprefix("http://") + + def base_url(self, location: str) -> str: + """Return the location's base URL, ``/``.""" + return f"{self.origin(location)}/artifactory/{REPOS[location]}" + + def file_url(self, location: str) -> str: + return f"{self.base_url(location)}/{LAYOUT_PATH}" + + @staticmethod + def file_path(location: str, layout_path: str = LAYOUT_PATH) -> str: + return f"/artifactory/{REPOS[location]}/{layout_path}" + + @staticmethod + def storage_path(location: str, layout_path: str = LAYOUT_PATH) -> str: + return f"/artifactory/api/storage/{REPOS[location]}/{layout_path}" + + @classmethod + def properties_path(cls, location: str) -> str: + return f"{cls.storage_path(location)}?properties" + + def publish(self, location: str, content: bytes, layout_path: str = LAYOUT_PATH) -> None: + """Serve ``content`` as the artifact on a location: its storage API entry and download.""" + info = _storage_info(location, content, _sha1(content), layout_path) + self.routes[(location, self.storage_path(location, layout_path))] = Reply(200, info) + self.routes[(location, self.file_path(location, layout_path))] = Reply(200, content) + + def answer(self, request: Request) -> Reply: + with self._lock: + self.requests.append(request) + route = self.routes.get((request.server, request.path)) + if route is None: + return Reply(404, b'{"errors": [{"status": 404, "message": "File not found."}]}') + return route if isinstance(route, Reply) else route(request) + + def received(self, name: str) -> list[Request]: + with self._lock: + return [request for request in self.requests if request.server == name] + + def paths(self, name: str) -> list[str]: + return [request.path for request in self.received(name)] + + +@dataclass +class Result: + returncode: int + stdout: str + stderr: str + sleeps: list[str] + + +class Helper: + """Runs the helper against the fake servers, with a stub ``sleep`` and a test netrc.""" + + def __init__(self, tmp_path: Path, artifactory: FakeArtifactory) -> None: + self.artifactory = artifactory + self.dest = tmp_path / "out" / "TensorRT-LLM.tar.gz" + self.part = self.dest.with_name(f"{self.dest.name}.part") + self.cwd = tmp_path / "cwd" + self.tmpdir = tmp_path / "tmp" + stub_dir = tmp_path / "bin" + home = tmp_path / "home" + for directory in (self.dest.parent, self.cwd, self.tmpdir, stub_dir, home): + directory.mkdir() + + self._sleep_log = tmp_path / "sleep.log" + sleep = stub_dir / "sleep" + sleep.write_text(f'#!/bin/sh\necho "$*" >> "{self._sleep_log}"\n') + sleep.chmod(0o755) + + # curl must not read the caller's curlrc: a "verbose" line there would print the + # Authorization header that the netrc produces. + (home / ".curlrc").write_text("verbose\n") + + # The urm entry, which a developer's netrc may hold, must never be used: urm is read + # anonymously, so the helper may pass the netrc to curl only for aws requests. + netrc = tmp_path / "netrc" + netrc.write_text( + f"machine {HOSTS['aws']} login {NETRC_LOGIN} password {NETRC_PASSWORD}\n" + f"machine {HOSTS['urm']} login urm-reader password {URM_NETRC_PASSWORD}\n" + ) + netrc.chmod(0o600) + + env = { + name: value + for name, value in os.environ.items() + if name.lower() not in _PROXY_VARIABLES + and not name.startswith(("LLM_ARTIFACTS_", "ARTIFACTORY_")) + and name not in ("CURL_HOME", "XDG_CONFIG_HOME") + } + env.update( + { + "PATH": f"{stub_dir}{os.pathsep}{os.environ['PATH']}", + "HOME": str(home), + "TMPDIR": str(self.tmpdir), + "LLM_ARTIFACTS_AWS_URL": artifactory.base_url("aws"), + "LLM_ARTIFACTS_URM_URL": artifactory.base_url("urm"), + "ARTIFACTORY_NETRC": str(netrc), + "ARTIFACTORY_CREDENTIAL_ID": CREDENTIAL_ID, + } + ) + self.env = env + + def run(self, *args: str, **env_overrides: str | None) -> Result: + """Run the helper; an override of ``None`` removes that variable from its environment.""" + env = dict(self.env) + for name, value in env_overrides.items(): + if value is None: + env.pop(name, None) + else: + env[name] = value + self._sleep_log.unlink(missing_ok=True) + + proc = subprocess.run( + ["bash", str(FETCH_BUILD_ARTIFACT), *args], + env=env, + cwd=self.cwd, + capture_output=True, + text=True, + timeout=60, + ) + + output = proc.stdout + proc.stderr + assert NETRC_PASSWORD not in output + assert URM_NETRC_PASSWORD not in output + assert BASIC_TOKEN not in output + assert "Signature=" not in output + assert list(self.tmpdir.iterdir()) == [], "temporary files were left behind" + assert list(self.cwd.iterdir()) == [], "files were written outside " + sleeps = self._sleep_log.read_text().split() if self._sleep_log.exists() else [] + return Result(proc.returncode, proc.stdout, proc.stderr, sleeps) + + +@pytest.fixture +def artifactory() -> Iterator[FakeArtifactory]: + fake = FakeArtifactory() + yield fake + fake.close() + + +@pytest.fixture +def helper(tmp_path: Path, artifactory: FakeArtifactory) -> Helper: + return Helper(tmp_path, artifactory) + + +def _served_by(location: str, layout_path: str = LAYOUT_PATH) -> str: + return f"fetch_build_artifact: served-by={location} path={layout_path}" + + +def _rejection_message(credential_id: str) -> str: + """Return the line naming a credential that aws rejected, the same for HTTP 401 and 403.""" + return ( + f"fetch_build_artifact: the Artifactory access token in credential '{credential_id}' " + "needs renewal or lacks read permission on sw-tensorrt-llm-generic-local; an Artifactory " + "access token for artifactory.nvidia.com is required." + ) + + +def _assert_no_retry_signature(text: str) -> None: + for signature in RETRY_SIGNATURES: + assert signature not in text + + +def _assert_no_credential_message(text: str) -> None: + assert "needs renewal" not in text + assert CREDENTIAL_ID not in text + + +# --------------------------------------------------------------------------- +# Path mapping and the primary location +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("artifact_path", "layout_path"), + [ + (LEGACY_PATH, LAYOUT_PATH), + (f"sw-tensorrt-generic-local/{LAYOUT_PATH}", LAYOUT_PATH), + (LAYOUT_PATH, LAYOUT_PATH), + (f"sw-tensorrt-generic/{BUILD_DIR}/build_info.txt", f"{BUILD_DIR}/build_info.txt"), + ( + f"sw-tensorrt-generic/{BUILD_DIR}/waive_list/waives.txt", + f"{BUILD_DIR}/waive_list/waives.txt", + ), + (f"sw-tensorrt-generic/{BUILD_DIR}/{OTHER_FILE}", f"{BUILD_DIR}/{OTHER_FILE}"), + ], + ids=["legacy", "legacy-local", "layout", "build-info", "waive-list", "other-file"], +) +def test_reads_the_primary_location(helper: Helper, artifact_path: str, layout_path: str) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package", layout_path) + + result = helper.run(artifact_path, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert not helper.part.exists() + assert result.stderr.splitlines() == [_served_by("aws", layout_path)] + assert result.stdout == "" + assert artifactory.paths("aws") == [ + artifactory.storage_path("aws", layout_path), + artifactory.file_path("aws", layout_path), + ] + assert [request.authorization for request in artifactory.received("aws")] == [AUTHORIZATION] * 2 + assert artifactory.paths("urm") == [] + assert result.sleeps == [] + + +@pytest.mark.parametrize( + "artifact_path", + [ + f"sw-tensorrt-llm-generic-local/{LAYOUT_PATH}", + "sw-tensorrt-generic/other-artifacts/TensorRT-LLM.tar.gz", + f"/{LAYOUT_PATH}", + "llm-artifacts-old/88/TensorRT-LLM.tar.gz", + f"https://urm.nvidia.com/artifactory/{LEGACY_PATH}", + "llm-artifacts/LLM/../../TensorRT-LLM.tar.gz", + "llm-artifacts/LLM/main/L0_PostMerge/88/", + "llm-artifacts/LLM/main//88/TensorRT-LLM.tar.gz", + "llm-artifacts/LLM/main/./88/TensorRT-LLM.tar.gz", + "llm-artifacts/LLM/main/L0 PostMerge/88/TensorRT-LLM.tar.gz", + "llm-artifacts/LLM/main/L0_PostMerge/88/TensorRT-LLM.tar.gz\t", + f"{LAYOUT_PATH};bolt.ref=0123456789abcdef", + f"{LAYOUT_PATH}?properties", + f"{LAYOUT_PATH}#fragment", + ], +) +def test_unmappable_path_is_a_usage_error(helper: Helper, artifact_path: str) -> None: + helper.artifactory.publish("aws", b"aws package") + + result = helper.run(artifact_path, str(helper.dest)) + + assert result.returncode == 64 + assert helper.artifactory.requests == [] + assert not helper.dest.exists() + + +@pytest.mark.parametrize("location", ["aws", "urm"]) +def test_base_url_with_credentials_is_a_usage_error(helper: Helper, location: str) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + variable = f"LLM_ARTIFACTS_{location.upper()}_URL" + base_url = artifactory.base_url(location).replace("http://", "http://someone:url-secret@") + + result = helper.run(LEGACY_PATH, str(helper.dest), **{variable: base_url}) + + assert result.returncode == 64 + assert f"{variable} must not contain credentials" in result.stderr + assert "url-secret" not in result.stdout + result.stderr + assert artifactory.requests == [] + + +@pytest.mark.parametrize("primary", [None, "", "aws"], ids=["unset", "empty", "aws"]) +def test_aws_is_the_default_primary(helper: Helper, primary: str | None) -> None: + helper.artifactory.publish("aws", b"aws package") + helper.artifactory.publish("urm", b"urm package") + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY=primary) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert _served_by("aws") in result.stderr.splitlines() + assert helper.artifactory.paths("urm") == [] + + +def test_verified_download_replaces_the_destination(helper: Helper) -> None: + helper.artifactory.publish("aws", b"aws package") + helper.dest.write_bytes(b"older download") + helper.part.write_bytes(b"partial download left by an interrupted run") + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert not helper.part.exists() + + +# --------------------------------------------------------------------------- +# Fallback to the secondary location +# --------------------------------------------------------------------------- + + +def test_falls_back_to_urm_when_aws_lacks_the_file(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("urm", b"urm package") + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"urm package" + assert result.stderr.splitlines() == [_served_by("urm")] + assert artifactory.paths("aws") == [artifactory.storage_path("aws")] + assert artifactory.paths("urm") == [ + artifactory.storage_path("urm"), + artifactory.file_path("urm"), + ] + assert [request.authorization for request in artifactory.received("aws")] == [AUTHORIZATION] + assert [request.authorization for request in artifactory.received("urm")] == [None, None] + + +def test_absent_from_both_locations_exits_2(helper: Helper) -> None: + artifactory = helper.artifactory + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 2 + assert not helper.dest.exists() + assert not helper.part.exists() + assert artifactory.paths("aws") == [artifactory.storage_path("aws")] + assert artifactory.paths("urm") == [artifactory.storage_path("urm")] + assert result.sleeps == [] + _assert_no_retry_signature(result.stderr) + + +def test_rejection_by_the_secondary_exits_4(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.routes[("aws", artifactory.storage_path("aws"))] = Reply(401) + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm") + + assert result.returncode == 4 + assert _rejection_message(CREDENTIAL_ID) in result.stderr.splitlines() + assert artifactory.paths("aws") == [artifactory.storage_path("aws")] + assert result.sleeps == [] + + +def test_errors_from_the_secondary_are_retried_then_exit_4(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.routes[("urm", artifactory.storage_path("urm"))] = Reply(500) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + assert artifactory.paths("urm") == [artifactory.storage_path("urm")] * 3 + assert result.sleeps == ["2", "4"] + assert not helper.dest.exists() + + +# --------------------------------------------------------------------------- +# Rejected requests never fall back +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("status", [401, 403]) +@pytest.mark.parametrize("step", ["lookup", "download"]) +def test_rejected_credentials_exit_3_without_fallback( + helper: Helper, step: str, status: int +) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + rejected = artifactory.storage_path("aws") if step == "lookup" else artifactory.file_path("aws") + artifactory.routes[("aws", rejected)] = Reply(status, b'{"errors": [{"status": %d}]}' % status) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 3 + assert result.stderr.splitlines() == [ + f"fetch_build_artifact: {artifactory.host('aws')} rejected the credentials for " + f"{LAYOUT_PATH} (HTTP {status}).", + _rejection_message(CREDENTIAL_ID), + ] + _assert_no_retry_signature(result.stdout + result.stderr) + assert artifactory.paths("aws").count(rejected) == 1 + assert artifactory.paths("urm") == [] + assert result.sleeps == [] + assert not helper.dest.exists() + assert not helper.part.exists() + + +def test_rejection_message_marks_an_unset_credential_id(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.routes[("aws", artifactory.storage_path("aws"))] = Reply(401) + + result = helper.run(LEGACY_PATH, str(helper.dest), ARTIFACTORY_CREDENTIAL_ID=None) + + assert result.returncode == 3 + assert _rejection_message("") in result.stderr.splitlines() + + +# --------------------------------------------------------------------------- +# Downloads redirected to the CDN +# --------------------------------------------------------------------------- + + +def _redirect_downloads_to_cdn( + artifactory: FakeArtifactory, *, direct: Reply, cdn_origin: str | None = None +) -> None: + """Make aws redirect downloads to the CDN, and answer ``direct`` to requests that opt out. + + ``cdn_origin`` replaces the CDN's origin in the redirect, for example with a port that nothing + listens on. + """ + origin = cdn_origin or artifactory.origin("cdn") + redirect = Reply(302, headers={"Location": origin + CDN_PATH}) + + def download(request: Request) -> Reply: + return direct if request.headers.get(REDIRECT_OPT_OUT) == "None" else redirect + + artifactory.routes[("aws", artifactory.file_path("aws"))] = download + + +def test_download_through_the_cdn_carries_no_credential(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + _redirect_downloads_to_cdn(artifactory, direct=Reply(500)) + artifactory.routes[("cdn", CDN_PATH)] = Reply(200, b"aws package") + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert _served_by("aws") in result.stderr.splitlines() + assert artifactory.paths("aws").count(artifactory.file_path("aws")) == 1 + assert [request.authorization for request in artifactory.received("cdn")] == [None] + + +def test_inherited_xtrace_is_turned_off(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + _redirect_downloads_to_cdn(artifactory, direct=Reply(500)) + artifactory.routes[("cdn", CDN_PATH)] = Reply(200, b"aws package") + + # A caller running under `set -x` that exports SHELLOPTS passes xtrace on to the helper. A + # BASH_ENV startup file would run, traced, before the helper's first command. + result = helper.run(LEGACY_PATH, str(helper.dest), SHELLOPTS="xtrace", BASH_ENV=None) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + # Only the command that turns xtrace off is traced. Helper.run checks that neither the netrc + # password nor the CDN URL's signature reached the output. + assert [line for line in result.stderr.splitlines() if not line.endswith("set +x")] == [ + _served_by("aws") + ] + + +def test_cdn_rejection_is_retried_without_the_redirect(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + _redirect_downloads_to_cdn(artifactory, direct=Reply(200, b"aws package")) + artifactory.routes[("cdn", CDN_PATH)] = Reply(403, b"AccessDenied") + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert _served_by("aws") in result.stderr.splitlines() + downloads = [ + request + for request in artifactory.received("aws") + if request.path == artifactory.file_path("aws") + ] + assert [request.headers.get(REDIRECT_OPT_OUT) for request in downloads] == [None, "None"] + assert [request.authorization for request in downloads] == [AUTHORIZATION] * 2 + assert [request.authorization for request in artifactory.received("cdn")] == [None] + assert artifactory.paths("urm") == [] + assert result.sleeps == [] + _assert_no_credential_message(result.stderr) + + +@pytest.mark.parametrize("direct", ["rejected", "redirected-again"]) +def test_persistent_cdn_rejection_exits_4(helper: Helper, direct: str) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + redirect = Reply(302, headers={"Location": artifactory.origin("cdn") + CDN_PATH}) + _redirect_downloads_to_cdn(artifactory, direct=Reply(403) if direct == "rejected" else redirect) + artifactory.routes[("cdn", CDN_PATH)] = Reply(403) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + _assert_no_credential_message(result.stderr) + # Three attempts, each retried once without the redirect. + assert artifactory.paths("aws").count(artifactory.file_path("aws")) == 6 + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + assert not helper.dest.exists() + assert not helper.part.exists() + + +@pytest.mark.parametrize("failure", ["http-500", "connection-refused"]) +def test_cdn_failure_other_than_403_is_retried_through_the_redirect( + helper: Helper, failure: str +) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + if failure == "http-500": + cdn_origin = artifactory.origin("cdn") + artifactory.routes[("cdn", CDN_PATH)] = Reply(500) + else: + cdn_origin = f"http://{HOSTS['cdn']}:{_unused_port(HOSTS['cdn'])}" + # A request that opts out of the redirect would get the file, so only staying on the redirect + # can make the download fail. + _redirect_downloads_to_cdn( + artifactory, direct=Reply(200, b"aws package"), cdn_origin=cdn_origin + ) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + _assert_no_credential_message(result.stderr) + downloads = [ + request + for request in artifactory.received("aws") + if request.path == artifactory.file_path("aws") + ] + assert [request.headers.get(REDIRECT_OPT_OUT) for request in downloads] == [None] * 3 + if failure == "http-500": + assert artifactory.paths("cdn") == [CDN_PATH] * 3 + else: + assert artifactory.received("cdn") == [] + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + assert not helper.dest.exists() + assert not helper.part.exists() + + +# --------------------------------------------------------------------------- +# Checksums +# --------------------------------------------------------------------------- + + +def test_checksum_mismatch_is_retried(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.routes[("aws", artifactory.file_path("aws"))] = _in_turn( + Reply(200, b"corrupted package"), Reply(200, b"aws package") + ) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert not helper.part.exists() + assert artifactory.paths("aws").count(artifactory.file_path("aws")) == 2 + assert result.sleeps == ["2"] + + +def test_persistent_checksum_mismatch_exits_4_without_fallback(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + artifactory.routes[("aws", artifactory.file_path("aws"))] = Reply(200, b"corrupted package") + helper.dest.write_bytes(b"older download") + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + assert "checksum" in result.stderr + assert _sha1(b"corrupted package") in result.stderr + assert helper.dest.read_bytes() == b"older download" + assert not helper.part.exists() + assert artifactory.paths("aws").count(artifactory.file_path("aws")) == 3 + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + + +def test_lookup_without_checksum_is_retried_then_exit_4(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + info = json.dumps({"repo": REPOS["aws"], "path": f"/{LAYOUT_PATH}", "children": []}) + artifactory.routes[("aws", artifactory.storage_path("aws"))] = Reply(200, info.encode()) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + assert artifactory.paths("aws") == [artifactory.storage_path("aws")] * 3 + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + + +# --------------------------------------------------------------------------- +# Transient failures are retried and never fall back +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + ("step", "status"), + [("lookup", 500), ("lookup", 429), ("download", 502), ("download", 404)], +) +def test_transient_errors_are_retried_then_exit_4(helper: Helper, step: str, status: int) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + failing = artifactory.storage_path("aws") if step == "lookup" else artifactory.file_path("aws") + artifactory.routes[("aws", failing)] = Reply(status) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + assert artifactory.paths("aws").count(failing) == 3 + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + _assert_no_credential_message(result.stderr) + assert not helper.dest.exists() + assert not helper.part.exists() + + +def test_transient_error_then_success(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + storage = ("aws", artifactory.storage_path("aws")) + stored = artifactory.routes[storage] + assert isinstance(stored, Reply) + artifactory.routes[storage] = _in_turn(Reply(503), stored) + + result = helper.run(LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert _served_by("aws") in result.stderr.splitlines() + assert artifactory.paths("aws").count(artifactory.storage_path("aws")) == 2 + assert result.sleeps == ["2"] + + +def test_unreachable_primary_is_retried_then_exit_4(helper: Helper) -> None: + helper.artifactory.publish("urm", b"urm package") + closed = f"http://{HOSTS['aws']}:{_unused_port(HOSTS['aws'])}/artifactory/{REPOS['aws']}" + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_AWS_URL=closed) + + assert result.returncode == 4 + assert helper.artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + assert not helper.dest.exists() + + +# --------------------------------------------------------------------------- +# --exists and --properties +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("location", ["aws", "urm"]) +@pytest.mark.parametrize("with_dest", [False, True], ids=["no-dest", "dest-ignored"]) +def test_exists_prints_url_and_sha1(helper: Helper, location: str, with_dest: bool) -> None: + artifactory = helper.artifactory + artifactory.publish(location, b"package") + dest_args = [str(helper.dest)] if with_dest else [] + + result = helper.run("--exists", LEGACY_PATH, *dest_args) + + assert result.returncode == 0, result.stderr + assert result.stdout == f"{artifactory.file_url(location)} {_sha1(b'package')}\n" + assert result.stderr.splitlines() == [_served_by(location)] + assert list(helper.dest.parent.iterdir()) == [] + requested = artifactory.paths("aws") + artifactory.paths("urm") + assert requested and all(path.startswith("/artifactory/api/storage/") for path in requested) + + +@pytest.mark.parametrize(("status", "returncode"), [(404, 2), (403, 3)]) +def test_exists_reports_absent_and_rejected(helper: Helper, status: int, returncode: int) -> None: + artifactory = helper.artifactory + artifactory.routes[("aws", artifactory.storage_path("aws"))] = Reply(status) + + result = helper.run("--exists", LEGACY_PATH) + + assert result.returncode == returncode + assert result.stdout == "" + expected_urm = [artifactory.storage_path("urm")] if returncode == 2 else [] + assert artifactory.paths("urm") == expected_urm + + +@pytest.mark.parametrize("location", ["aws", "urm"]) +def test_properties_follow_the_read_rule(helper: Helper, location: str) -> None: + artifactory = helper.artifactory + artifactory.routes[(location, artifactory.properties_path(location))] = Reply(200, PROPERTIES) + + result = helper.run("--properties", LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == PROPERTIES + assert not helper.part.exists() + assert result.stderr.splitlines() == [_served_by(location)] + assert artifactory.paths("aws") == [artifactory.properties_path("aws")] + expected_urm = [artifactory.properties_path("urm")] if location == "urm" else [] + assert artifactory.paths("urm") == expected_urm + assert [request.authorization for request in artifactory.received("aws")] == [AUTHORIZATION] + assert all(request.authorization is None for request in artifactory.received("urm")) + + +def test_properties_absent_from_both_exit_2(helper: Helper) -> None: + artifactory = helper.artifactory + + result = helper.run("--properties", LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 2 + assert not helper.dest.exists() + assert not helper.part.exists() + assert artifactory.paths("aws") == [artifactory.properties_path("aws")] + assert artifactory.paths("urm") == [artifactory.properties_path("urm")] + + +def test_properties_rejected_exit_3_without_fallback(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.routes[("aws", artifactory.properties_path("aws"))] = Reply(403) + artifactory.routes[("urm", artifactory.properties_path("urm"))] = Reply(200, PROPERTIES) + + result = helper.run("--properties", LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 3 + assert _rejection_message(CREDENTIAL_ID) in result.stderr.splitlines() + _assert_no_retry_signature(result.stderr) + assert artifactory.paths("urm") == [] + assert not helper.dest.exists() + + +def test_properties_server_error_is_retried_then_exit_4(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.routes[("aws", artifactory.properties_path("aws"))] = Reply(500) + artifactory.routes[("urm", artifactory.properties_path("urm"))] = Reply(200, PROPERTIES) + + result = helper.run("--properties", LEGACY_PATH, str(helper.dest)) + + assert result.returncode == 4 + assert artifactory.paths("aws") == [artifactory.properties_path("aws")] * 3 + assert artifactory.paths("urm") == [] + assert result.sleeps == ["2", "4"] + + +# --------------------------------------------------------------------------- +# LLM_ARTIFACTS_PRIMARY=urm (rollback) +# --------------------------------------------------------------------------- + + +def test_primary_urm_reads_urm_first(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.publish("urm", b"urm package") + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm") + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"urm package" + assert result.stderr.splitlines() == [_served_by("urm")] + assert artifactory.paths("aws") == [] + + +def test_primary_urm_falls_back_to_aws(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm") + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"aws package" + assert result.stderr.splitlines() == [_served_by("aws")] + assert artifactory.paths("urm") == [artifactory.storage_path("urm")] + assert [request.authorization for request in artifactory.received("urm")] == [None] + assert [request.authorization for request in artifactory.received("aws")] == [AUTHORIZATION] * 2 + + +def test_primary_urm_rejection_exits_3_without_fallback(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + artifactory.routes[("urm", artifactory.storage_path("urm"))] = Reply(403) + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm") + + assert result.returncode == 3 + assert result.stderr.splitlines() == [ + f"fetch_build_artifact: {artifactory.host('urm')} rejected the anonymous request for " + f"{LAYOUT_PATH} (HTTP 403)." + ] + assert artifactory.paths("aws") == [] + + +@pytest.mark.parametrize("primary", ["bogus", "AWS", "urm "]) +def test_unknown_primary_is_a_usage_error(helper: Helper, primary: str) -> None: + helper.artifactory.publish("aws", b"aws package") + + result = helper.run(LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY=primary) + + assert result.returncode == 64 + assert "LLM_ARTIFACTS_PRIMARY" in result.stderr + assert helper.artifactory.requests == [] + + +# --------------------------------------------------------------------------- +# Credentials +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("netrc", ["unset", "empty", "missing-file"]) +def test_aws_needs_a_netrc(helper: Helper, tmp_path: Path, netrc: str) -> None: + helper.artifactory.publish("aws", b"aws package") + value = {"unset": None, "empty": "", "missing-file": str(tmp_path / "no-such-netrc")}[netrc] + + result = helper.run(LEGACY_PATH, str(helper.dest), ARTIFACTORY_NETRC=value) + + assert result.returncode == 64 + assert "ARTIFACTORY_NETRC" in result.stderr + assert helper.artifactory.requests == [] + + +def test_primary_urm_needs_no_netrc_for_a_file_on_urm(helper: Helper) -> None: + helper.artifactory.publish("urm", b"urm package") + + result = helper.run( + LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm", ARTIFACTORY_NETRC=None + ) + + assert result.returncode == 0, result.stderr + assert helper.dest.read_bytes() == b"urm package" + + +def test_primary_urm_needs_a_netrc_to_fall_back_to_aws(helper: Helper) -> None: + artifactory = helper.artifactory + artifactory.publish("aws", b"aws package") + + result = helper.run( + LEGACY_PATH, str(helper.dest), LLM_ARTIFACTS_PRIMARY="urm", ARTIFACTORY_NETRC=None + ) + + assert result.returncode == 64 + assert "ARTIFACTORY_NETRC" in result.stderr + assert artifactory.paths("urm") == [artifactory.storage_path("urm")] + assert artifactory.paths("aws") == [] + + +# --------------------------------------------------------------------------- +# Command line +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize("flag", ["-h", "--help"]) +def test_help(helper: Helper, flag: str) -> None: + result = helper.run(flag) + + assert result.returncode == 0 + assert SYNOPSIS in result.stdout + assert helper.artifactory.requests == [] + + +@pytest.mark.parametrize( + "args", + [ + (), + (LEGACY_PATH,), + ("--properties", LEGACY_PATH), + ("--exists", "--properties", LEGACY_PATH), + ("--verbose", LEGACY_PATH, "out"), + (LEGACY_PATH, "out", "extra"), + ], + ids=["no-args", "no-dest", "properties-no-dest", "two-modes", "unknown-option", "extra-arg"], +) +def test_usage_errors_exit_64(helper: Helper, args: tuple[str, ...]) -> None: + result = helper.run(*args) + + assert result.returncode == 64 + assert SYNOPSIS in result.stderr + assert helper.artifactory.requests == [] + + +@pytest.mark.parametrize("dest", ["missing-dir/TensorRT-LLM.tar.gz", "out"]) +def test_destination_must_be_a_file_in_an_existing_directory( + helper: Helper, tmp_path: Path, dest: str +) -> None: + helper.artifactory.publish("aws", b"aws package") + + result = helper.run(LEGACY_PATH, str(tmp_path / dest)) + + assert result.returncode == 64 + assert helper.artifactory.requests == [] + + +# --------------------------------------------------------------------------- +# Exit-code names shared with the readers +# --------------------------------------------------------------------------- + +# The helper's exit codes other than success, by the names that the helper defines for them. +EXIT_CODES = {"FETCH_ABSENT": 2, "FETCH_REJECTED": 3, "FETCH_TRANSFER_FAILED": 4, "FETCH_USAGE": 64} + + +@dataclass(frozen=True) +class ExitCodeReader: + """A reader that acts on particular exit codes of the helper.""" + + path: str + # Matches the reader's definitions of the helper's names, as (name, value). + definition: str + # Picks the code that checks the helper's exit status out of the reader's source. + checks: Callable[[str], str] + + +def _defined_exit_codes(text: str, definition: str) -> dict[str, int]: + return {name: int(value) for name, value in re.findall(definition, text, re.MULTILINE)} + + +def _helper_exit_codes() -> dict[str, int]: + return _defined_exit_codes( + FETCH_BUILD_ARTIFACT.read_text(), r"^readonly (FETCH_[A-Z_]+)=(\d+)$" + ) + + +def _braced_block(text: str, start: str) -> str: + """The code from ``start`` through the brace that closes the first block after it.""" + begin = text.index(start) + depth = 0 + for i in range(text.index("{", begin), len(text)): + if text[i] == "{": + depth += 1 + elif text[i] == "}": + depth -= 1 + if depth == 0: + return text[begin : i + 1] + raise AssertionError(f"no block closes after {start!r}") + + +def _between(text: str, start: str, end: str) -> str: + begin = text.index(start) + return text[begin : text.index(end, begin)] + + +def _python_definitions(text: str, names: set[str]) -> str: + """The source of the module-level functions and assignments called ``names``.""" + found = {} + for node in ast.parse(text).body: + if isinstance(node, ast.FunctionDef): + name = node.name + elif isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name): + name = node.targets[0].id + else: + continue + if name in names: + found[name] = ast.get_source_segment(text, node) or "" + assert set(found) == names + return "\n".join(found.values()) + + +def _bare_exit_codes(code: str, exit_codes: set[int]) -> list[int]: + """The numbers in ``code`` that are among ``exit_codes``. + + ``code`` is the part of a reader that checks the helper's exit status, so each of these + numbers is treated as an exit code checked by number, whatever the form of the check: a + comparison either way round, a shell test in ``[ ]`` or ``[[ ]]`` with or without quotes, a + ``case`` arm, or an item or key of a list, range, set, dict, or map. A file descriptor in a + shell redirection (``>&2``, ``2>``), a positional parameter (``$2``), and digits within a + word, version, or decimal are not counted. + """ + numbers = re.findall(r"(?]&)\d+(?![\w<>])(?!\.\d)", code) + return [int(number) for number in numbers if int(number) in exit_codes] + + +_GROOVY_DEFINITION = r"^(FETCH_[A-Z_]+) = (\d+)$" +_PYTHON_DEFINITION = r"^_(FETCH_[A-Z_]+) = (\d+)$" + +EXIT_CODE_READERS = [ + ExitCodeReader( + "jenkins/L0_Test.groovy", + _GROOVY_DEFINITION, + lambda text: _braced_block(text, "def boltedTarArtifactPath("), + ), + ExitCodeReader( + "jenkins/BuildDockerImage.groovy", + _GROOVY_DEFINITION, + lambda text: _braced_block(text, "for (file in requiredFiles)"), + ), + # The cluster script that resolves the input package. + ExitCodeReader( + "jenkins/BoltProfileGen.groovy", + _GROOVY_DEFINITION, + lambda text: _between(text, "def tarStage =", "'''.stripIndent()"), + ), + ExitCodeReader( + "jenkins/scripts/cbts/cbts/command/coverage/selection/artifact.py", + _PYTHON_DEFINITION, + lambda text: _python_definitions(text, {"_HELPER_FAILURES", "_fetch_build_artifact"}), + ), + ExitCodeReader( + "scripts/get_wheel_from_package.py", + _PYTHON_DEFINITION, + lambda text: _python_definitions(text, {"download_package"}), + ), +] + + +def test_helper_names_its_exit_codes() -> None: + text = FETCH_BUILD_ARTIFACT.read_text() + + assert _helper_exit_codes() == EXIT_CODES + bare = re.findall(r"\b(?:exit|fail|return)\s+(\d+)\b", text) + assert [int(code) for code in bare if int(code) in EXIT_CODES.values()] == [] + + +@pytest.mark.parametrize("reader", EXIT_CODE_READERS, ids=lambda reader: Path(reader.path).name) +def test_readers_use_the_helpers_exit_code_names(reader: ExitCodeReader) -> None: + helper_codes = _helper_exit_codes() + text = (REPO_ROOT / reader.path).read_text() + + defined = _defined_exit_codes(text, reader.definition) + assert defined, f"{reader.path} defines none of {sorted(helper_codes)}" + assert defined == {name: helper_codes.get(name) for name in defined} + checks = reader.checks(text) + assert any(name in checks for name in defined) + bare = _bare_exit_codes(checks, set(helper_codes.values())) + assert bare == [], f"{reader.path} uses {bare} as numbers where it checks the exit status" + + +@pytest.mark.parametrize( + ("code", "numbers"), + [ + ("if rc == 2:", [2]), + ("if (3 == exitCode) {", [3]), + ('elif [ "$rc" -eq 2 ]; then', [2]), + ("if [ $rc -eq 2 ]; then", [2]), + ('if [[ "$rc" == "3" ]]; then', [3]), + ('elif [ 2 -eq "$rc" ]; then', [2]), + ("if (!(propsStatus in [0, 2, 4])) {", [2, 4]), + ("if status in (0, 2):", [2]), + ("if status not in {0, 2, 4}:", [2, 4]), + ("if (exitCode in 2..4 || [64].contains(exitCode)) {", [2, 4, 64]), + ('case "$rc" in 2) echo absent ;; 3|4) exit "$rc" ;; esac', [2, 3, 4]), + ("switch (exitCode) {\n case 64:", [64]), + ("match status:\n case 2 | 4:", [2, 4]), + ("_HELPER_FAILURES = {\n 3: (AuthError, None),\n}", [3]), + ("_HELPER_FAILURES = {3: (AuthError, None), 4: (TransferError, None)}", [3, 4]), + ('def reasons = [3: "rejected", 4: "failed"]', [3, 4]), + ], + ids=[ + "comparison", + "reversed-comparison", + "quoted-test", + "unquoted-test", + "double-bracket-test", + "reversed-test", + "list", + "tuple", + "set", + "range-and-contains", + "shell-case-arms", + "groovy-case-arm", + "python-case-arm", + "dict-key-per-line", + "dict-keys-on-one-line", + "map-keys-on-one-line", + ], +) +def test_bare_number_check_finds_codes_in_any_form(code: str, numbers: list[int]) -> None: + assert _bare_exit_codes(code, set(EXIT_CODES.values())) == numbers + + +@pytest.mark.parametrize( + "code", + [ + "if (!(propsStatus in [0, FETCH_ABSENT, FETCH_TRANSFER_FAILED])) {", + "if (exitCode == FETCH_REJECTED) {", + "FETCH_ABSENT=${FETCH_ABSENT}", + 'elif [ "$rc" -eq "$FETCH_ABSENT" ]; then', + '[ "$rc" -eq 0 ] || exit "$rc"', + "_HELPER_FAILURES = {\n _FETCH_REJECTED: (ArtifactAuthError, None),\n}", + "if status in (0, _FETCH_ABSENT):", + "if rc not in (_FETCH_ABSENT, _FETCH_TRANSFER_FAILED):", + ], + ids=[ + "groovy-list", + "groovy-comparison", + "shell-definition", + "shell-test", + "shell-success", + "python-dict", + "python-tuple", + "python-not-in", + ], +) +def test_bare_number_check_passes_the_readers_named_forms(code: str) -> None: + assert _bare_exit_codes(code, set(EXIT_CODES.values())) == [] + + +@pytest.mark.parametrize( + "code", + [ + 'echo "[ERROR] input tarball missing" >&2', + 'gzip -t "$DEST" 2>/dev/null', + 'cp "$1" "$2"', + "if (config != LINUX_AARCH64_CONFIG || arch == 'x86_64' || version == '2.3.4') {", + ], + ids=["redirect-to-stderr", "redirect-stderr", "positional-parameter", "words-and-versions"], +) +def test_bare_number_check_ignores_digits_that_are_not_numbers(code: str) -> None: + assert _bare_exit_codes(code, set(EXIT_CODES.values())) == [] diff --git a/tests/unittest/scripts/test_get_wheel_from_package.py b/tests/unittest/scripts/test_get_wheel_from_package.py new file mode 100644 index 000000000000..86318311dc1d --- /dev/null +++ b/tests/unittest/scripts/test_get_wheel_from_package.py @@ -0,0 +1,265 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the packaged-build download in scripts/get_wheel_from_package.py. + +The script polls jenkins/scripts/fetch_build_artifact.sh until a run's packaged build is +published. subprocess.run and time.sleep are replaced, so no helper, network, or wait is involved. +""" + +import argparse +import importlib.util +import os +import re +import subprocess +from pathlib import Path +from types import ModuleType + +import pytest + +pytestmark = pytest.mark.cpu_only + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent.parent +SCRIPT_PATH = REPO_ROOT / "scripts" / "get_wheel_from_package.py" +HELPER_PATH = REPO_ROOT / "jenkins" / "scripts" / "fetch_build_artifact.sh" +# The legacy path the image build passes as --artifact_path; the helper maps it to each location. +ARTIFACT_PATH = "sw-tensorrt-generic/llm-artifacts/LLM/main/L0_PostMerge/88" +TARBALL = "TensorRT-LLM.tar.gz" +NETRC = "/run/secrets/artifactory_netrc" +CREDENTIAL_ID = "aws-artifactory-credentials" +# Log phrases that make the CI shared library retry a step as a network fault. +RETRY_SIGNATURES = ( + "403 Forbidden", + "URL returned error", + "HTTP response code said error", + "404 Not Found", + "ERROR 404", + "HTTP 50", + "ERROR 50", + "Server Error", + "Error downloading", + "Connection", + "Read timed out", + "Too Many Requests", +) + + +class _StopBeforeExtract(Exception): + """Raised for the first command after the download, so a test ends before extraction.""" + + +class _FakeRun: + """Stands in for subprocess.run and answers each helper call with the next exit code.""" + + def __init__(self, returncodes: list[int]) -> None: + self._returncodes = list(returncodes) + self.calls: list[tuple[list[str], dict[str, object]]] = [] + + def __call__(self, cmd: list[str], **kwargs: object) -> subprocess.CompletedProcess[bytes]: + self.calls.append((list(cmd), kwargs)) + if cmd[0] != "bash": + raise _StopBeforeExtract(cmd) + result: subprocess.CompletedProcess[bytes] = subprocess.CompletedProcess( + cmd, self._returncodes.pop(0) + ) + if kwargs.get("check"): + result.check_returncode() + return result + + +@pytest.fixture(scope="module") +def wheel_script() -> ModuleType: + spec = importlib.util.spec_from_file_location("get_wheel_from_package", SCRIPT_PATH) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +@pytest.fixture(autouse=True) +def read_helper_env(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setenv("ARTIFACTORY_NETRC", NETRC) + monkeypatch.setenv("ARTIFACTORY_CREDENTIAL_ID", CREDENTIAL_ID) + monkeypatch.delenv("LLM_ARTIFACTS_PRIMARY", raising=False) + + +@pytest.fixture(autouse=True) +def sleeps(wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch) -> list[float]: + recorded: list[float] = [] + monkeypatch.setattr(wheel_script.time, "sleep", recorded.append) + return recorded + + +def _fake_run( + wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch, returncodes: list[int] +) -> _FakeRun: + fake = _FakeRun(returncodes) + monkeypatch.setattr(wheel_script.subprocess, "run", fake) + return fake + + +def _download(wheel_script: ModuleType, timeout: int = 60) -> None: + wheel_script.download_package(ARTIFACT_PATH, TARBALL, timeout) + + +def test_download_runs_the_read_helper_with_the_environment( + wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch, sleeps: list[float] +) -> None: + monkeypatch.setenv("LLM_ARTIFACTS_PRIMARY", "aws") + run = _fake_run(wheel_script, monkeypatch, [0]) + + _download(wheel_script) + + assert len(run.calls) == 1 + cmd, kwargs = run.calls[0] + assert cmd == ["bash", str(HELPER_PATH), f"{ARTIFACT_PATH}/{TARBALL}", TARBALL] + env = kwargs["env"] + assert isinstance(env, dict) + assert env == dict(os.environ) + assert env["ARTIFACTORY_NETRC"] == NETRC + assert env["ARTIFACTORY_CREDENTIAL_ID"] == CREDENTIAL_ID + assert env["LLM_ARTIFACTS_PRIMARY"] == "aws" + # The helper's served-by line and its error messages belong in the build log. + assert not kwargs.get("capture_output") + assert kwargs.get("stderr") is None + assert sleeps == [] + + +def test_download_polls_until_the_package_is_published( + wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch, sleeps: list[float] +) -> None: + run = _fake_run(wheel_script, monkeypatch, [2, 4, 2, 0]) + + _download(wheel_script) + + assert len(run.calls) == 4 + assert sleeps == [60, 60, 60] + + +@pytest.mark.parametrize("returncode", (2, 4)) +def test_download_times_out_after_timeout_attempts( + returncode: int, + wheel_script: ModuleType, + monkeypatch: pytest.MonkeyPatch, + sleeps: list[float], +) -> None: + run = _fake_run(wheel_script, monkeypatch, [returncode] * 3) + + with pytest.raises(TimeoutError, match="after 3 attempts"): + _download(wheel_script, timeout=3) + + assert len(run.calls) == 3 + assert sleeps == [60, 60] + + +@pytest.mark.parametrize("primary", (None, "", "aws")) +def test_download_stops_when_the_credentials_are_rejected( + primary: str | None, + wheel_script: ModuleType, + monkeypatch: pytest.MonkeyPatch, + sleeps: list[float], +) -> None: + if primary is not None: + monkeypatch.setenv("LLM_ARTIFACTS_PRIMARY", primary) + run = _fake_run(wheel_script, monkeypatch, [3]) + + with pytest.raises(PermissionError) as excinfo: + _download(wheel_script) + + message = str(excinfo.value) + assert f"credential '{CREDENTIAL_ID}'" in message + assert "needs renewal or lacks read permission on sw-tensorrt-llm-generic-local" in message + assert "artifactory.nvidia.com" in message + assert f"{ARTIFACT_PATH}/{TARBALL}" in message + assert [signature for signature in RETRY_SIGNATURES if signature in message] == [] + assert len(run.calls) == 1 + assert sleeps == [] + + +def test_rejected_credentials_name_a_placeholder_when_the_credential_id_is_unset( + wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch, sleeps: list[float] +) -> None: + monkeypatch.delenv("ARTIFACTORY_CREDENTIAL_ID") + run = _fake_run(wheel_script, monkeypatch, [3]) + + with pytest.raises( + PermissionError, match=re.escape("credential ''") + ): + _download(wheel_script) + + assert len(run.calls) == 1 + assert sleeps == [] + + +def test_rejected_anonymous_request_when_urm_is_primary( + wheel_script: ModuleType, monkeypatch: pytest.MonkeyPatch, sleeps: list[float] +) -> None: + monkeypatch.setenv("LLM_ARTIFACTS_PRIMARY", "urm") + run = _fake_run(wheel_script, monkeypatch, [3]) + + with pytest.raises(PermissionError) as excinfo: + _download(wheel_script) + + message = str(excinfo.value) + assert "rejected the anonymous request" in message + assert f"{ARTIFACT_PATH}/{TARBALL}" in message + assert CREDENTIAL_ID not in message + assert [signature for signature in RETRY_SIGNATURES if signature in message] == [] + assert len(run.calls) == 1 + assert sleeps == [] + + +@pytest.mark.parametrize("returncode", (64, 1, 127)) +def test_download_stops_on_any_other_helper_failure( + returncode: int, + wheel_script: ModuleType, + monkeypatch: pytest.MonkeyPatch, + sleeps: list[float], +) -> None: + run = _fake_run(wheel_script, monkeypatch, [returncode]) + + with pytest.raises(RuntimeError, match=rf"exit code {returncode}\b"): + _download(wheel_script) + + assert len(run.calls) == 1 + assert sleeps == [] + + +@pytest.mark.parametrize( + ("arch", "tarball"), + (("x86_64", "TensorRT-LLM.tar.gz"), ("sbsa", "TensorRT-LLM-GH200.tar.gz")), +) +def test_command_line_downloads_the_package_for_its_arch( + arch: str, + tarball: str, + wheel_script: ModuleType, + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, +) -> None: + # Extraction writes under the project directory, so point it at tmp_path. + monkeypatch.setattr(wheel_script, "get_project_dir", lambda: tmp_path) + run = _fake_run(wheel_script, monkeypatch, [0]) + parser = argparse.ArgumentParser() + wheel_script.add_arguments(parser) + args = parser.parse_args(["--arch", arch, "--timeout", "60", "--artifact_path", ARTIFACT_PATH]) + + with pytest.raises(_StopBeforeExtract): + wheel_script.get_wheel_from_package(**vars(args)) + + helper = tmp_path / "jenkins" / "scripts" / "fetch_build_artifact.sh" + assert [cmd for cmd, _ in run.calls] == [ + ["bash", str(helper), f"{ARTIFACT_PATH}/{tarball}", tarball], + ["tar", "-zxf", tarball, "-C", str(tmp_path / "tmp")], + ] diff --git a/tests/unittest/scripts/test_slurm_install.py b/tests/unittest/scripts/test_slurm_install.py index 510028ff86ac..a4add5ef2398 100644 --- a/tests/unittest/scripts/test_slurm_install.py +++ b/tests/unittest/scripts/test_slurm_install.py @@ -14,8 +14,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import os +import re +import shlex +import shutil import subprocess +import sys +from dataclasses import dataclass from pathlib import Path import pytest @@ -26,6 +32,59 @@ BASH_UTILS = REPO_ROOT / "jenkins" / "scripts" / "bash_utils.sh" SLURM_INSTALL = REPO_ROOT / "jenkins" / "scripts" / "slurm_install.sh" L0_TEST = REPO_ROOT / "jenkins" / "L0_Test.groovy" +BUILD_GROOVY = REPO_ROOT / "jenkins" / "Build.groovy" +# The pipelines that run the read helper against the aws location with a netrc they write. +BUILD_ARTIFACT_READERS = ( + L0_TEST, + REPO_ROOT / "jenkins" / "BuildDockerImage.groovy", + REPO_ROOT / "jenkins" / "BoltProfileGen.groovy", + REPO_ROOT / "jenkins" / "L0_MergeRequest.groovy", +) +AWS_HOST = "artifactory.nvidia.com" + +# The launcher copies the scripts into the job workspace under this per-job prefix. +JOB_PREFIX = "cluster-multi_node_test-0a1b2c-" +ARTIFACT_PATH = ( + "sw-tensorrt-generic/llm-artifacts/LLM/main/L0_MergeRequest_PR/1234/TensorRT-LLM.tar.gz" +) +CREDENTIAL_ID = "test-reader-credential" +NETRC_PASSWORD = "netrc-password-5d1e9c0a" +PAYLOAD = "packaged build contents\n" +# What the job script exports for the read helper. +HELPER_ENVIRONMENT = ("LLM_ARTIFACTS_PRIMARY", "ARTIFACTORY_NETRC", "ARTIFACTORY_CREDENTIAL_ID") + +# Stands in for fetch_build_artifact.sh. It records its arguments and the environment it sees, and +# plays back the planned exit statuses, one per call; the last one repeats. A successful call +# writes the packaged build to ; a failed one leaves a partial file there, which must +# never be moved into place. +_FETCH_STUB = r""" +import json +import os +import sys + +with open(os.environ["FETCH_STUB_LOG"], "a", encoding="utf-8") as log: + log.write( + json.dumps( + { + "args": sys.argv[1:], + "env": {name: os.environ.get(name) for name in json.loads(os.environ["FETCH_STUB_ENV"])}, + } + ) + + "\n" + ) +with open(os.environ["FETCH_STUB_LOG"], encoding="utf-8") as log: + call_index = len(log.readlines()) - 1 +statuses = json.loads(os.environ["FETCH_STUB_STATUSES"]) +status = statuses[min(call_index, len(statuses) - 1)] +with open(sys.argv[2], "w", encoding="utf-8") as dest: + dest.write(os.environ["FETCH_STUB_PAYLOAD"] if status == 0 else "partial") +sys.exit(status) +""" + + +def _write_executable(path: Path, content: str) -> None: + path.write_text(content, encoding="utf-8") + path.chmod(0o755) @pytest.mark.parametrize("with_timeout", (False, True)) @@ -75,14 +134,14 @@ def test_retry_command_returns_terminal_status(with_timeout: bool) -> None: @pytest.mark.parametrize( - "missing", ("tarName", "llmTarfile", "resourcePathNode", "stageName", "pytestCommand") + "missing", ("tarName", "llmTarArtifactPath", "resourcePathNode", "stageName", "pytestCommand") ) def test_slurm_install_requires_artifact_inputs(missing: str) -> None: env = os.environ.copy() env.update( { "tarName": "/tmp/TensorRT-LLM.tar.gz", - "llmTarfile": "https://example.invalid/TensorRT-LLM.tar.gz", + "llmTarArtifactPath": ARTIFACT_PATH, "resourcePathNode": "/tmp", "stageName": "test-stage", "pytestCommand": "pytest", @@ -102,12 +161,475 @@ def test_slurm_install_requires_artifact_inputs(missing: str) -> None: assert f"{missing} is required" in result.stderr -def test_l0_artifact_download_retries_overwrite_archive() -> None: - download_commands = [ +@dataclass +class _Run: + """Outcome of running slurm_install.sh from a job workspace.""" + + returncode: int + stdout: str + stderr: str + fetch_calls: list[dict[str, object]] + archive: Path + + @property + def output(self) -> str: + return self.stdout + self.stderr + + +def _run_slurm_install(tmp_path: Path, statuses: list[int]) -> _Run: + """Run slurm_install.sh as a Slurm job does, up to extracting the downloaded package. + + The script, bash_utils.sh, and a stub read helper sit in a job workspace under one per-job + prefix, as the launcher copies them. A stub ``tar`` records its arguments and exits 42, which + ends the setup right after the download; a stub ``sleep`` makes the retry interval instant, and + a stub ``timeout`` records each attempt's limit before running the attempt. + """ + workspace = tmp_path / "job-workspace" + workspace.mkdir() + install = workspace / f"{JOB_PREFIX}slurm_install.sh" + shutil.copy(SLURM_INSTALL, install) + shutil.copy(BASH_UTILS, workspace / f"{JOB_PREFIX}bash_utils.sh") + stub_script = tmp_path / "fetch_stub.py" + stub_script.write_text(_FETCH_STUB, encoding="utf-8") + _write_executable( + workspace / f"{JOB_PREFIX}fetch_build_artifact.sh", + f'exec {shlex.quote(sys.executable)} -I -S {shlex.quote(str(stub_script))} "$@"\n', + ) + + bin_dir = tmp_path / "stub_bin" + bin_dir.mkdir() + _write_executable( + bin_dir / "tar", '#!/bin/bash\nprintf "%s\\n" "$@" > "$TAR_STUB_LOG"\nexit 42\n' + ) + _write_executable(bin_dir / "sleep", '#!/bin/bash\necho "$@" >> "$SLEEP_STUB_LOG"\n') + _write_executable( + bin_dir / "timeout", '#!/bin/bash\necho "$1" >> "$TIMEOUT_STUB_LOG"\nshift\nexec "$@"\n' + ) + + netrc = workspace / "artifactory-netrc" + netrc.write_text( + f"machine artifactory.nvidia.com login ci-reader password {NETRC_PASSWORD}\n", + encoding="utf-8", + ) + netrc.chmod(0o600) + + resource_dir = tmp_path / "resource" + resource_dir.mkdir() + fetch_log = tmp_path / "fetch_calls.jsonl" + env = { + name: value + for name, value in os.environ.items() + if name not in (*HELPER_ENVIRONMENT, "llmTarfile", "GITHUB_CLONE_TOKEN", "BASH_ENV") + } + env.update( + { + "PATH": os.pathsep.join((str(bin_dir), env.get("PATH", os.defpath))), + "HOME": str(tmp_path), + "tarName": "TensorRT-LLM.tar.gz", + "llmTarArtifactPath": ARTIFACT_PATH, + "resourcePathNode": str(resource_dir), + "stageName": "test-stage", + "pytestCommand": "pytest", + "SLURM_LOCALID": "0", + "SLURM_JOB_ID": "4321", + "SLURM_NODEID": "2", + "LLM_ARTIFACTS_PRIMARY": "aws", + "ARTIFACTORY_NETRC": str(netrc), + "ARTIFACTORY_CREDENTIAL_ID": CREDENTIAL_ID, + "FETCH_STUB_LOG": str(fetch_log), + "FETCH_STUB_ENV": json.dumps(HELPER_ENVIRONMENT), + "FETCH_STUB_STATUSES": json.dumps(statuses), + "FETCH_STUB_PAYLOAD": PAYLOAD, + "TAR_STUB_LOG": str(tmp_path / "tar_args.txt"), + "SLEEP_STUB_LOG": str(tmp_path / "sleep_args.txt"), + "TIMEOUT_STUB_LOG": str(tmp_path / "timeout_args.txt"), + } + ) + result = subprocess.run( + ["bash", str(install)], env=env, capture_output=True, text=True, timeout=60 + ) + fetch_calls = ( + [json.loads(line) for line in fetch_log.read_text(encoding="utf-8").splitlines()] + if fetch_log.exists() + else [] + ) + run = _Run( + result.returncode, + result.stdout, + result.stderr, + fetch_calls, + resource_dir / "TensorRT-LLM.tar.gz", + ) + + # The job script traces its commands, and the netrc reaches the helper only by path. + assert any(line.startswith("+") for line in run.stderr.splitlines()) + assert NETRC_PASSWORD not in run.output + for call in run.fetch_calls: + assert call["args"] == [ARTIFACT_PATH, f"{run.archive}.tmp.4321.2"] + assert call["env"] == { + "LLM_ARTIFACTS_PRIMARY": "aws", + "ARTIFACTORY_NETRC": str(netrc), + "ARTIFACTORY_CREDENTIAL_ID": CREDENTIAL_ID, + } + assert list(resource_dir.glob("*.tmp.*")) == [] + return run + + +def test_slurm_install_setup_downloads_package_through_retries(tmp_path: Path) -> None: + run = _run_slurm_install(tmp_path, [4, 0]) + + assert run.returncode == 42, run.output + assert (tmp_path / "tar_args.txt").read_text(encoding="utf-8").split() == [ + "-zxf", + str(run.archive), + ] + assert run.archive.read_text(encoding="utf-8") == PAYLOAD + assert len(run.fetch_calls) == 2 + assert (tmp_path / "timeout_args.txt").read_text(encoding="utf-8").split() == ["1800", "1800"] + assert len((tmp_path / "sleep_args.txt").read_text(encoding="utf-8").splitlines()) == 1 + + +@pytest.mark.parametrize("status", (2, 3, 4), ids=("absent", "rejected", "transfer-failed")) +def test_slurm_install_setup_fails_after_retries(tmp_path: Path, status: int) -> None: + run = _run_slurm_install(tmp_path, [status]) + + assert run.returncode == 1, run.output + assert len(run.fetch_calls) == 3 + assert f"Artifact download failed after retries: {ARTIFACT_PATH}" in run.stdout + assert not run.archive.exists() + assert not (tmp_path / "tar_args.txt").exists() + + +def _l0_test_code_lines() -> list[str]: + """Return the code lines of jenkins/L0_Test.groovy, as _groovy_code_lines does.""" + return _groovy_code_lines(L0_TEST) + + +def _groovy_code(path: Path) -> str: + """Return the code lines of a Groovy file, as _groovy_code_lines does, joined by newlines.""" + return "\n".join(_groovy_code_lines(path)) + + +def _groovy_code_lines(path: Path) -> list[str]: + """Return the code lines of a Groovy file. + + Whole-line and block comments are dropped, and backslash-continued lines are joined; a + statement that spans lines without a backslash stays split across them. + """ + code = [] + in_block_comment = False + continued = "" + for line in path.read_text(encoding="utf-8").splitlines(): + stripped = line.strip() + if in_block_comment: + in_block_comment = "*/" not in stripped + continue + if stripped.startswith("/*"): + in_block_comment = "*/" not in stripped + continue + if stripped.startswith("//"): + continue + if line.rstrip().endswith("\\"): + continued += line.rstrip().rstrip("\\") + " " + continue + code.append(continued + line) + continued = "" + return code + + +def _groovy_function_body(lines: list[str], name: str) -> str: + """Return the top-level Groovy function ``name``, up to the next top-level definition.""" + starts = [index for index, line in enumerate(lines) if re.match(rf"(def|void) {name}\(", line)] + assert len(starts) == 1, f"expected one top-level definition of {name}" + start = starts[0] + end = next( + ( + index + for index in range(start + 1, len(lines)) + if lines[index].startswith(("def ", "void ")) + ), + len(lines), + ) + return "\n".join(lines[start:end]) + + +def _enclosing_blocks(lines: list[str], index: int) -> list[str]: + """Return the openers of the blocks around ``lines[index]``, innermost first, by indentation.""" + openers = [] + indent = len(lines[index]) - len(lines[index].lstrip()) + for line in reversed(lines[:index]): + stripped = line.strip() + line_indent = len(line) - len(stripped) if stripped else indent + if line_indent < indent and stripped.endswith("{"): + openers.append(stripped) + indent = line_indent + return openers + + +def test_l0_reads_packaged_builds_through_fetch_helper() -> None: + lines = _l0_test_code_lines() + definitions = [line for line in lines if re.match(r"\s*def fetchBuildArtifact\(", line)] + assert len(definitions) == 1 + body = _groovy_function_body(lines, "fetchBuildArtifact") + assert "llmExecStepWithRetry" in body + assert "readHelperPath(pipeline)" in body + helper_body = _groovy_function_body(lines, "readHelperPath") + assert "fetch_build_artifact.sh" in helper_body + assert "unstash(FETCH_BUILD_ARTIFACT_STASH)" in helper_body + + +@pytest.mark.parametrize( + "reader", + [ + "runLLMTestlistWithSbatch", + "runLLMDocBuild", + "launchTestListCheck", + "runLLMTestlistOnPlatformImpl", + ], +) +def test_l0_packaged_build_reader_downloads_through_fetch_helper(reader: str) -> None: + body = _groovy_function_body(_l0_test_code_lines(), reader) + + # A download whose destination names the packaged build; an --exists probe has none. + assert re.search(r"\bfetchBuildArtifact\(pipeline, [^,]+, [^,)]*tarName[^,)]*\)", body) + + +def test_l0_has_no_direct_download_of_packaged_builds() -> None: + package_names = ("llmTarfile", "llmTarArtifactPath", "TARNAME", "tarName") + package_downloads = [ line.strip() - for line in L0_TEST.read_text().splitlines() - if "wget -nv" in line and "llmTarfile" in line + for line in _l0_test_code_lines() + if re.search(r"\b(wget|curl)\b", line) and any(name in line for name in package_names) + ] + + assert package_downloads == [] + + +def test_l0_slurm_job_gets_the_read_helper_and_its_netrc() -> None: + lines = _l0_test_code_lines() + body = _groovy_function_body(lines, "runLLMTestlistWithSbatch") + body_lines = body.splitlines() + + # The helper from the run's commit goes beside slurm_install.sh, which finds it by its own + # file name. + assert ( + 'def scriptFetchHelperPathNode = "${jobWorkspace}/${jobUID}-fetch_build_artifact.sh"' + in body + ) + assert "def scriptFetchHelperLocalPath = readHelperPath(pipeline)" in body + assert re.search( + r"Utils\.copyFileToRemoteHost\(\s*pipeline,\s*remote,\s*scriptFetchHelperLocalPath," + r"\s*scriptFetchHelperPathNode,\s*true\s*\)", + body, + ) + + # Every job gets the helper's netrc, since the helper on the node contacts + # artifactory.nvidia.com whenever aws is the primary, wherever the package is. + netrc_writes = [ + index + for index, line in enumerate(body_lines) + if "writeArtifactoryNetrc(pipeline, artifactoryNetrcPathLocal)" in line + ] + assert len(netrc_writes) == 1 + openers = _enclosing_blocks(body_lines, netrc_writes[0]) + assert 'stage("Initialize Test") {' in openers + within_stage = openers[: openers.index('stage("Initialize Test") {')] + assert [opener for opener in within_stage if opener.startswith("if")] == [] + assert re.search( + r"Utils\.copyFileToRemoteHost\(\s*pipeline,\s*remote,\s*artifactoryNetrcPathLocal," + r"\s*artifactoryNetrcPathNode\s*\)", + body, + ) + + # Before the netrc reaches the job workspace, the attempt's job is registered for the off-pod + # cleanup under an entry of its own, whatever an earlier attempt left registered. + registrations = [ + index + for index, line in enumerate(body_lines) + if "registerSlurmResource(stageName, [clusterName: partition.clusterName, jobUID: jobUID," + " usedSbatch: true])" in line + ] + assert len(registrations) == 1 + assert registrations[0] < netrc_writes[0] + openers = _enclosing_blocks(body_lines, registrations[0]) + within_stage = openers[: openers.index('stage("Initialize Test") {')] + assert [opener for opener in within_stage if opener.startswith("if")] == [] + register = _groovy_function_body(lines, "registerSlurmResource") + assert 'id: "${stageName}/slurm/${fields.jobUID ?: fields.nodeName}"' in register + + # The job script exports what the helper and slurm_install.sh read, with no condition. + # llmTarfile is the URL that --exists resolved on the agent, for the older slurm_install.sh + # inside a package built before the helper existed. + exports = [line.strip() for line in body_lines if re.search(r"\bexport \w+=", line)] + for name, value in ( + ("llmTarArtifactPath", "$llmTarArtifactPath"), + ("LLM_ARTIFACTS_PRIMARY", "${llmArtifactsPrimary()}"), + ("ARTIFACTORY_NETRC", "$artifactoryNetrcPathNode"), + ("ARTIFACTORY_CREDENTIAL_ID", "$LLM_ARTIFACTS_CREDENTIALS_ID"), + ("llmTarfile", "$llmTarfile"), + ): + assert [line for line in exports if re.search(rf"\bexport {name}=", line)] == [ + f"export {name}={value}" + ] + assert "llmTarfileSha1" not in body + assert "llmTarfileOnAws" not in body + + # A retry that reuses the job workspace keeps both files. + keep_list = re.search(r"def filesToKeepWhenRetry = \[(?P[^\]]*)\]", body) + assert keep_list + kept = {item.strip() for item in keep_list.group("items").split(",")} + assert {"scriptFetchHelperPathNode", "artifactoryNetrcPathNode"} <= kept + + # A netrc "default" entry would send the token to every host, including the CDN. + netrc_writer = _groovy_function_body(lines, "writeArtifactoryNetrc") + assert "printf 'machine artifactory.nvidia.com login %s password %s" in netrc_writer + assert not re.search(r"\bdefault\b", netrc_writer) + + +def test_l0_reads_build_artifacts_and_pulls_images_with_separate_credential_settings() -> None: + lines = _l0_test_code_lines() + + # The helper's netrc and the credential name that the helper prints come from one setting. + netrc_writer = _groovy_function_body(lines, "writeArtifactoryNetrc") + assert "credentialsId: LLM_ARTIFACTS_CREDENTIALS_ID," in netrc_writer + fetch = _groovy_function_body(lines, "fetchBuildArtifact") + assert '"ARTIFACTORY_CREDENTIAL_ID=${LLM_ARTIFACTS_CREDENTIALS_ID}",' in fetch + # fetchBuildArtifact gives the helper the netrc that writeArtifactoryNetrc writes, as the Slurm + # launcher does. + assert "writeArtifactoryNetrc(pipeline, netrcPath)" in fetch + assert '"ARTIFACTORY_NETRC=${netrcPath}",' in fetch + + # The enroot credentials of Slurm jobs and the docker login before an image pull keep the + # image-pull setting. + logins = [ + index + for index, line in enumerate(lines) + if "machine ${ARTIFACTORY_DOCKER_HOST} login" in line + or "docker login ${ARTIFACTORY_DOCKER_HOST}" in line ] + assert logins + for index in logins: + credential = next(line for line in reversed(lines[:index]) if "credentialsId:" in line) + assert credential.strip() == "credentialsId: ARTIFACTORY_CREDENTIALS_ID,", lines[index] + + # That setting holds the image-pull identity, not the credential that publishes build artifacts. + publishing = _artifact_location_field(_groovy_code(BUILD_GROOVY), "aws", "credential") + assert _groovy_string("\n".join(lines), "ARTIFACTORY_CREDENTIALS_ID") != publishing + + +# A netrc entry that a pipeline writes with printf, which keeps the secret off the command line. +# The two arguments after the format are the login and password variables. +_NETRC_ENTRY = re.compile(r"printf 'machine (?P\S+) login %s password %s") +_USERNAME_PASSWORD_BINDING = re.compile( + r"usernamePassword\(\s*credentialsId:\s*(?P[^,]+?),\s*" + r"usernameVariable:\s*'\w+',\s*passwordVariable:\s*'(?P\w+)'\s*\)" +) + + +def _assignments(code: str, name: str) -> list[str]: + """Return the values that a pipeline assigns to a variable, with or without def.""" + return re.findall(rf"^[ \t]*(?:def[ \t]+)?{name}[ \t]*=[ \t]*(.+?)[ \t]*$", code, re.M) + + +def _artifact_location_field(code: str, location: str, field: str) -> str: + """Return a field of one location in a pipeline's LLM_ARTIFACTS_LOCATIONS map.""" + locations = re.search(r"^LLM_ARTIFACTS_LOCATIONS = \[$(?P.*?)^\]$", code, re.M | re.S) + assert locations, "no LLM_ARTIFACTS_LOCATIONS map" + entry = re.search(rf"^[ \t]*{location}: \[(?P[^\]]*)\]", locations["entries"], re.M) + assert entry, f"no {location} entry in LLM_ARTIFACTS_LOCATIONS" + value = re.search(rf'\b{field}: "(?P[^"]*)"', entry["fields"]) + assert value, f"no {field} in the {location} entry of LLM_ARTIFACTS_LOCATIONS" + return value["value"] + + +def _groovy_string(code: str, expression: str) -> str: + """Return the string that a Groovy expression in a pipeline holds. + + The expression is a string literal, a field of LLM_ARTIFACTS_LOCATIONS, or a variable that the + pipeline assigns once. + """ + expression = expression.strip() + literal = re.fullmatch(r"""(["'])(?P[^"'$]*)\1""", expression) + if literal: + return literal["value"] + field = re.fullmatch(r"LLM_ARTIFACTS_LOCATIONS\.(?P\w+)\.(?P\w+)", expression) + if field: + return _artifact_location_field(code, field["location"], field["field"]) + assert re.fullmatch(r"\w+", expression), f"cannot evaluate {expression!r}" + values = _assignments(code, expression) + assert len(values) == 1, f"expected one assignment of {expression}, found {values}" + return _groovy_string(code, values[0]) + + +def _names_aws_host(code: str, host: str) -> bool: + """Tell whether a netrc entry's host is the aws location's. + + That is the literal artifactory.nvidia.com, or a variable whose every assignment derives from + LLM_ARTIFACTS_LOCATIONS.aws.root. ARTIFACTORY_DOCKER_HOST names the container registry instead. + """ + variable = re.fullmatch(r"\$\{?(?P\w+)\}?", host) + if variable is None: + return host == AWS_HOST + values = _assignments(code, variable["name"]) + return bool(values) and all( + value.startswith("LLM_ARTIFACTS_LOCATIONS.aws.root") for value in values + ) + + +def _aws_netrc_credential(code: str) -> str: + """Return the credentialsId expression behind the one netrc entry for the aws host. + + The entry prints its password from a variable, which the nearest usernamePassword binding of + that variable before the entry fills. + """ + entries = [ + entry for entry in _NETRC_ENTRY.finditer(code) if _names_aws_host(code, entry["host"]) + ] + assert len(entries) == 1, f"expected one netrc entry for {AWS_HOST}, found {len(entries)}" + (entry,) = entries + arguments = code[entry.end() :].splitlines()[0] + _, password = re.findall(r"\$\{?(\w+)", arguments)[:2] + bindings = [ + binding + for binding in _USERNAME_PASSWORD_BINDING.finditer(code, 0, entry.start()) + if binding["password"] == password + ] + assert bindings, f"no usernamePassword binding of {password} before the netrc entry" + return bindings[-1]["credential"] + + +def _helper_credential_ids(code: str) -> list[str]: + """Return the ARTIFACTORY_CREDENTIAL_ID values that a pipeline gives the read helper.""" + names = [] + for value in re.findall(r"ARTIFACTORY_CREDENTIAL_ID=([^\s\"]+)", code): + value = value.strip("'") + interpolation = re.fullmatch(r"\$\{?(?P[\w.]+)\}?", value) + names.append(_groovy_string(code, interpolation["expression"]) if interpolation else value) + return names + + +def test_build_uploads_to_aws_with_the_credential_of_its_location_map() -> None: + code = _groovy_code(BUILD_GROOVY) + + # Build.groovy uploads packaged builds through its ArtifactoryAws server, so the location map's + # aws credential, which the readers are compared with, is the one those uploads use. + servers = re.findall( + r"rtServer\s*\(\s*id:\s*'ArtifactoryAws',[^)]*?credentialsId:\s*'(?P[^']*)'", + code, + ) + assert servers + assert set(servers) == {_artifact_location_field(code, "aws", "credential")} + + +@pytest.mark.parametrize("pipeline", BUILD_ARTIFACT_READERS, ids=lambda path: path.name) +def test_build_artifacts_are_read_with_the_publishing_credential(pipeline: Path) -> None: + publishing = _artifact_location_field(_groovy_code(BUILD_GROOVY), "aws", "credential") + code = _groovy_code(pipeline) + + # The netrc that the read helper gets holds the credential that Build.groovy publishes with. + assert _groovy_string(code, _aws_netrc_credential(code)) == publishing - assert len(download_commands) == 4 - assert all("wget -nv -O" in command for command in download_commands) + # When aws rejects it, the read helper names that credential. + named = _helper_credential_ids(code) + assert named + assert set(named) == {publishing}