From c994635af669ba390dc73aec449a9cd462f993f3 Mon Sep 17 00:00:00 2001
From: Andy Lee <andylizf@outlook.com>
Date: Wed, 13 Aug 2025 14:18:48 -0700
Subject: [PATCH] sky: expand leann-build.yaml with configurable params and
 flags (backend, recompute, compact, embedding options)

---
 sky/leann-build.yaml | 46 ++++++++++++++++++++++++++++++++++++++------
 1 file changed, 40 insertions(+), 6 deletions(-)

diff --git a/sky/leann-build.yaml b/sky/leann-build.yaml
index f5e04e7..9e762ca 100644
--- a/sky/leann-build.yaml
+++ b/sky/leann-build.yaml
@@ -7,10 +7,27 @@ resources:
   # cloud: aws
   disk_size: 100
 
+env:
+  # Build parameters (override with: sky launch -c leann-gpu sky/leann-build.yaml -e key=value)
+  index_name: my-index
+  docs: ./data
+  backend: hnsw               # hnsw | diskann
+  complexity: 64
+  graph_degree: 32
+  num_threads: 8
+  # Embedding selection
+  embedding_mode: sentence-transformers   # sentence-transformers | openai | mlx | ollama
+  embedding_model: facebook/contriever
+  # Storage/latency knobs
+  recompute: true             # true => selective recomputation; false => store full embeddings
+  compact: true               # for HNSW only: false when recompute=false
+  # Optional pass-through
+  extra_args: ""
+
 # Sync local paths to the remote VM. Adjust as needed.
 file_mounts:
   # Example: mount your local data directory used for building
-  ~/leann-data: ./data
+  ~/leann-data: ${docs}
 
 setup: |
   set -e
@@ -21,8 +38,25 @@ setup: |
   # Install the LEANN CLI globally on the remote machine
   uv tool install leann
 
-# Optional: you can immediately kick off a build here, or use `sky exec` later.
-# run: |
-#   export PATH="$HOME/.local/bin:$PATH"
-#   # Example build using the mounted data directory
-#   leann build my-index --docs ~/leann-data --backend hnsw --complexity 64 --graph-degree 32
+run: |
+  export PATH="$HOME/.local/bin:$PATH"
+  # Derive flags from env
+  recompute_flag=""
+  if [ "${recompute}" = "false" ] || [ "${recompute}" = "0" ]; then
+    recompute_flag="--no-recompute"
+  fi
+  compact_flag=""
+  if [ "${compact}" = "false" ] || [ "${compact}" = "0" ]; then
+    compact_flag="--no-compact"
+  fi
+
+  # Build command
+  leann build ${index_name} \
+    --docs ~/leann-data \
+    --backend ${backend} \
+    --complexity ${complexity} \
+    --graph-degree ${graph_degree} \
+    --num-threads ${num_threads} \
+    --embedding-mode ${embedding_mode} \
+    --embedding-model ${embedding_model} \
+    ${recompute_flag} ${compact_flag} ${extra_args}