update skypilot yaml and us uv instead of pip (sgl-project#279)

Author: sii-xinglong

docs/basic_usage/qwen.md

@@ -7,7 +7,7 @@ Qwen is Alibaba's family of large language models optimized for diverse applicat
 Launch a Qwen-7B-Chat server on TPU:
 
 ```bash
-JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server \
+JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache uv run python -u -m sgl_jax.launch_server \
     --model-path Qwen/Qwen-7B-Chat \
     --trust-remote-code \
     --dist-init-addr=0.0.0.0:10011 \
@@ -56,7 +56,7 @@ Set `--mem-fraction-static 0.8` for optimal TPU memory utilization. For larger m
 
 ### Throughput Testing
 ```bash
-python3 -m sgl_jax.bench_serving \
+uv run python -m sgl_jax.bench_serving \
     --backend sgl-jax \
     --dataset-name random \
     --num-prompts 100 \

docs/get_started/install.md

@@ -4,9 +4,12 @@ You can install SGLang-Jax using one of the methods below.
 
 This page is mainly applicable to TPU devices running through JAX.
 
-## Method 1: With pip or uv
+## Method 1: With uv
 
-🚧 **Under Construction** 🚧
+```bash
+uv venv --python 3.12 && source .venv/bin/activate
+uv pip install sglang-jax
+```
 
 ## Method 2: From source
 
@@ -16,11 +19,11 @@ git clone https://github.com/sgl-project/sglang-jax
 cd sglang-jax
 
 # Install the python packages
-pip install --upgrade pip setuptools packaging
-pip install -e "python[all]"
+uv venv --python 3.12 && source .venv/bin/activate
+uv pip install -e python/
 
 # Run Qwen-7B Model
-JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server --model-path Qwen/Qwen-7B-Chat --trust-remote-code  --dist-init-addr=0.0.0.0:10011 --nnodes=1  --tp-size=4 --device=tpu --random-seed=3 --node-rank=0 --mem-fraction-static=0.8 --max-prefill-tokens=8192 --download-dir=/tmp --dtype=bfloat16  --skip-server-warmup --host 0.0.0.0 --port 30000
+JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache uv run python -u -m sgl_jax.launch_server --model-path Qwen/Qwen-7B-Chat --trust-remote-code  --dist-init-addr=0.0.0.0:10011 --nnodes=1  --tp-size=4 --device=tpu --random-seed=3 --node-rank=0 --mem-fraction-static=0.8 --max-prefill-tokens=8192 --download-dir=/tmp --dtype=bfloat16  --skip-server-warmup --host 0.0.0.0 --port 30000
 ```
 
 ## Method 3: Using docker
@@ -31,11 +34,7 @@ JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server --m
 
 🚧 **Under Construction** 🚧
 
-## Method 5: Using docker compose
-
-🚧 **Under Construction** 🚧
-
-## Method 6: Run on Cloud TPU with SkyPilot
+## Method 5: Run on Cloud TPU with SkyPilot
 
 <details>
 <summary>More</summary>
@@ -54,25 +53,19 @@ resources:
    accelerator_args:
       tpu_vm: True
       runtime_version: v2-alpha-tpuv6e
-file_mounts:
-  ~/.ssh/id_rsa: ~/.ssh/id_rsa
-setup: |
-  chmod 600 ~/.ssh/id_rsa
-  rm ~/.ssh/config
-  GIT_SSH_COMMAND="ssh -o UserKnownHostsFile=/dev/null -o StrictHostKeyChecking=no" git clone https://github.com/sgl-project/sglang-jax
 run: |
-  cd sglang-jax
-  pip install -e "python[all]"
-  JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 -u -m sgl_jax.launch_server --model-path Qwen/Qwen-7B-Chat --trust-remote-code  --dist-init-addr=0.0.0.0:10011 --nnodes=1  --tp-size=4 --device=tpu --random-seed=3 --node-rank=0 --mem-fraction-static=0.8 --max-prefill-tokens=8192 --download-dir=/tmp --dtype=bfloat16  --skip-server-warmup --attention-backend=fa --host 0.0.0.0 --port 30000
+  git clone https://github.com/sgl-project/sglang-jax.git
+  cd sglang-jax && git fetch origin $REF:$REF && git checkout $REF
+  uv venv --python 3.12
+  source .venv/bin/activate
+  uv pip install -e python/
 ```
 
 </details>
 
 ```bash
 sky launch -c sglang-jax sglang.yaml --infra=gcp
 
-# Get the HTTP API endpoint
-sky status --endpoint 30000 sglang-jax
 ```
 - For debugging and testing purposes, you can use spot instances to reduce costs by adding the `--use-spot` flag to your SkyPilot commands:
   ```bash

scripts/launch_tpu.sh

@@ -36,26 +36,15 @@ sanitize_ref() {
     ref=$(echo "$ref" | sed 's/^[^a-zA-Z]*//')
     # Ensure it doesn't end with dash, dot or underscore
     ref=$(echo "$ref" | sed 's/[-._]*$//')
-    # Truncate to reasonable length (max 40 chars to leave room for prefix/suffix)
-    ref=$(echo "$ref" | cut -c1-40)
+    # Truncate to reasonable length (max 20 chars to leave room for prefix/suffix)
+    ref=$(echo "$ref" | cut -c1-20)
     # If empty after sanitization, use "default"
     if [ -z "$ref" ]; then
         ref="default"
     fi
     echo "$ref"
 }
 
-# Check environment variables
-if [ -z "$USERNAME" ]; then
-    echo "Error: USERNAME environment variable is not set"
-    exit 1
-fi
-
-if [ -z "$GIT_TOKEN" ]; then
-    echo "Error: GIT_TOKEN environment variable is not set"
-    exit 1
-fi
-
 # Sanitize ref name
 SANITIZED_REF=$(sanitize_ref "$REF")
 
@@ -73,9 +62,9 @@ sed -e "s|\$ACCELERATOR|${ACCELERATOR}|g" \
 
 # Create cluster name with ref
 if [ -n "$TEST_TYPE" ]; then
-    CLUSTER_NAME="sgl-jax-ci-$ACCELERATOR-$SANITIZED_REF-$TEST_TYPE"
+    CLUSTER_NAME="sgl-jax-$ACCELERATOR-$SANITIZED_REF-$TEST_TYPE-$RANDOM"
 else
-    CLUSTER_NAME="sgl-jax-ci-$ACCELERATOR-$SANITIZED_REF"
+    CLUSTER_NAME="sgl-jax-$ACCELERATOR-$SANITIZED_REF-$RANDOM"
 fi
 
 # Execute sky launch command
@@ -93,9 +82,7 @@ sky launch "$TEMP_YAML" \
     --infra=gcp \
     -i 10 \
     --down \
-    -y \
-    --secret USERNAME=${USERNAME} \
-    --secret GIT_TOKEN=${GIT_TOKEN}
+    -y
 
 # Store the exit code
 EXIT_CODE=$?

scripts/tpu_resource.yaml

@@ -1,19 +1,16 @@
 resources:
   accelerators: $ACCELERATOR
-  use_spot: True
-  disk_size: 512  # Ensure model checkpoints can fit.
-  disk_tier: best
-
-secrets:
-  USERNAME: null
-  GIT_TOKEN: null
-
+  accelerator_args:
+      tpu_vm: True
+      runtime_version: v2-alpha-tpuv6e
+  use_spot: True # The default instance is spot.
 
 setup: |
   if [ -d "sglang-jax" ]; then pip uninstall -y sgl-jax || true && rm -rf sglang-jax; fi
-  git clone https://${USERNAME}:${GIT_TOKEN}@github.com/sgl-project/sglang-jax.git
+  git clone https://github.com/sgl-project/sglang-jax.git
   cd sglang-jax && git fetch origin $REF:$REF && git checkout $REF
   uv venv --python 3.12
   source .venv/bin/activate
   uv pip install -e python/
   bash scripts/killall_sglang.sh
+  uv tool install evalscope