parallelize docker smoke tests

Author: PastelStorm

CHANGELOG.md

@@ -11,7 +11,7 @@
 * Split lint tools into a lightweight dependency group so the CI lint step no longer installs heavy runtime dependencies
 * Add explicit dependencies for `backoff`, `pandas`, `psutil`, `pypdf`, and `requests` (previously only transitive via `unstructured[all-docs]`)
 * Pre-download NLTK models before parallel test runs to prevent race conditions
-* Parallelize Docker smoke tests with pytest-xdist
+* Parallelize Docker smoke tests by running one container per xdist worker on dedicated ports
 * Remove unused `ARCH` variable from Makefile
 
 ## 0.0.93

scripts/docker-smoke-test.sh

@@ -11,24 +11,20 @@
 
 set -e
 
-CONTAINER_NAME=unstructured-api-smoke-test
+CONTAINER_NAME_PREFIX=unstructured-api-smoke-test
 CONTAINER_NAME_PARALLEL=unstructured-api-smoke-test-parallel
 PIPELINE_FAMILY=${PIPELINE_FAMILY:-"general"}
 DOCKER_IMAGE="${DOCKER_IMAGE:-pipeline-family-${PIPELINE_FAMILY}-dev:latest}"
 SKIP_INFERENCE_TESTS="${SKIP_INFERENCE_TESTS:-false}"
+NUM_WORKERS="${NUM_WORKERS:-4}"
+BASE_PORT=8000
 
 start_container() {
+    local port=$1
+    local name=$2
+    local use_parallel_mode=${3:-false}
 
-    port=$1
-    use_parallel_mode=$2
-
-    if [ "$use_parallel_mode" = "true" ]; then
-        name=$CONTAINER_NAME_PARALLEL
-    else
-        name=$CONTAINER_NAME
-    fi
-
-    echo Starting container "$name"
+    echo Starting container "$name" on port "$port"
     docker run --platform "$DOCKER_PLATFORM" \
            -p "$port":"$port" \
            --entrypoint uvicorn \
@@ -42,8 +38,8 @@ start_container() {
 }
 
 await_server_ready() {
-    port=$1
-    url=localhost:$port/healthcheck
+    local port=$1
+    local url=localhost:$port/healthcheck
 
     # NOTE(rniko): Increasing the timeout to 120 seconds because emulated arm tests are slow to start
     for _ in {1..120}; do
@@ -60,37 +56,52 @@ await_server_ready() {
     exit 1
 }
 
-stop_container() {
-    echo Stopping container "$CONTAINER_NAME"
-    # Note (austin) - if you're getting an error from the api, try dumping the logs
-    # docker logs $CONTAINER_NAME 2> docker_logs.txt
-    docker stop "$CONTAINER_NAME" 2> /dev/null || true
+stop_all_containers() {
+    for i in $(seq 0 $((NUM_WORKERS-1))); do
+        local name="${CONTAINER_NAME_PREFIX}-${i}"
+        echo Stopping container "$name"
+        docker stop "$name" 2> /dev/null || true
+    done
 
     echo Stopping container "$CONTAINER_NAME_PARALLEL"
     docker stop "$CONTAINER_NAME_PARALLEL" 2> /dev/null || true
 }
 
-# Always clean up the container
-trap stop_container EXIT
+# Always clean up all containers
+trap stop_all_containers EXIT
+
+#######################
+# Start worker containers
+#######################
+for i in $(seq 0 $((NUM_WORKERS-1))); do
+    port=$((BASE_PORT + i))
+    start_container "$port" "${CONTAINER_NAME_PREFIX}-${i}" "false"
+done
 
-start_container 8000 "false"
-await_server_ready 8000
+for i in $(seq 0 $((NUM_WORKERS-1))); do
+    port=$((BASE_PORT + i))
+    await_server_ready "$port"
+done
 
 #######################
 # Smoke Tests
 #######################
-echo Running smoke tests with SKIP_INFERENCE_TESTS: "$SKIP_INFERENCE_TESTS"
-PYTHONPATH=. SKIP_INFERENCE_TESTS=$SKIP_INFERENCE_TESTS uv run pytest -n auto -vv scripts/smoketest.py
+echo "Running smoke tests with SKIP_INFERENCE_TESTS: $SKIP_INFERENCE_TESTS, NUM_WORKERS: $NUM_WORKERS"
+PYTHONPATH=. SKIP_INFERENCE_TESTS=$SKIP_INFERENCE_TESTS SMOKE_TEST_BASE_PORT=$BASE_PORT \
+    uv run pytest -n "$NUM_WORKERS" -vv scripts/smoketest.py
 
 #######################
 # Test parallel vs single mode
 #######################
 if ! $SKIP_INFERENCE_TESTS; then
-    start_container 9000 true
-    await_server_ready 9000
+    # Reuse the first container on BASE_PORT for single mode,
+    # start a new one for parallel mode on a non-conflicting port
+    parallel_port=$((BASE_PORT + NUM_WORKERS))
+    start_container "$parallel_port" "$CONTAINER_NAME_PARALLEL" "true"
+    await_server_ready "$parallel_port"
 
     echo Running parallel mode test
-    ./scripts/parallel-mode-test.sh localhost:8000 localhost:9000
+    ./scripts/parallel-mode-test.sh "localhost:$BASE_PORT" "localhost:$parallel_port"
 fi
 
 result=$?

scripts/smoketest.py

@@ -11,11 +11,23 @@
 import pytest
 import requests
 
-API_URL = "http://localhost:8000/general/v0/general"
 # NOTE(rniko): Skip inference tests if we're running on an emulated architecture
 skip_inference_tests = os.getenv("SKIP_INFERENCE_TESTS", "").lower() in {"true", "yes", "y", "1"}
 
 
+def _get_api_url() -> str:
+    """Determine the API URL for the current xdist worker.
+
+    Each pytest-xdist worker (gw0, gw1, ...) is assigned its own container
+    on a dedicated port (base_port + worker_number). When running without
+    xdist, defaults to base_port.
+    """
+    base_port = int(os.environ.get("SMOKE_TEST_BASE_PORT", "8000"))
+    worker_id = os.environ.get("PYTEST_XDIST_WORKER", "gw0")
+    worker_num = int(worker_id.replace("gw", ""))
+    return f"http://localhost:{base_port + worker_num}/general/v0/general"
+
+
 def send_document(
     filenames: List[str],
     filenames_gzipped: Optional[List[str]] = None,
@@ -42,7 +54,7 @@ def send_document(
         options["gz_uncompressed_content_type"] = uncompressed_content_type
 
     return requests.post(
-        API_URL,
+        _get_api_url(),
         files=files,
         data=options,
     )