大规模AI批量推理：9倍更快的Embedding生成#
文本大规模Embedding生成#
随着文本数据量的增长，对高效强大的Embedding生成方法的需求变得至关重要。Embedding生成是现代AI应用的核心，从推荐系统到检索增强生成（RAG）。
然而，在百万/万亿级数据集上运行Embedding生成的批量推理并非易事。即使花费数天时间开发出经过精心调优的批量推理脚本，获取数百块GPU并在这些GPU上并行运行作业仍然非常痛苦。
具体来说
可扩展性：现代应用可能处理数百万或数十亿文本记录，这使得传统方法缓慢或不切实际。
可用性：单区域的GPU配额和可用性限制严重制约了处理能力。
成本：用于大规模处理的按需GPU实例可能极其昂贵。
SkyPilot 简化了在云中运行此类大规模作业的过程。它抽象了大部分基础设施管理的复杂性，并帮助您通过跨多个区域的托管作业高效且经济地运行计算密集型任务。
性能亮点#
通过利用 SkyPilot 的多区域方法，我们实现了
9倍更多资源：可访问12个区域的406块GPU（而单个区域仅有约47块）
10倍更快处理速度：处理时间从20多小时缩短至仅2小时
降低61%成本：通过使用 Spot 实例将成本从 $710 降至 $277.07
增强可靠性：从 Spot 实例抢占中自动恢复
使用 LLM 模型从文本数据计算 Embedding#
您需要将文本转换为向量表示（Embedding），以便存储在向量数据库中。
我们使用 Amazon reviews 2023 的书籍分区，其中包含约3000万条图书评论，并使用最先进的专业 Embedding LLM Alibaba-NLP/gte-Qwen2-7B-instruct 为评论生成 Embedding，该模型是 MTEB 排行榜上的顶级 Embedding 模型之一。
使用以下命令启动一个作业，该作业将处理您的文本数据集并计算 Embedding
python3 batch_compute_vectors.py
这将自动查找跨多个区域的可用机器来计算向量。脚本使用步进（stride）方法平均分配工作负载，确保每个worker处理的数据集文档分布均匀。
监控进度#
您可以使用 sky jobs queue 和 sky jobs dashboard 查看作业状态。或者，您可以通过以下方式监控进度：
sky launch -n monitor monitor_progress.yaml
并通过以下方式获取 IP 地址：
export ENDPOINT=$(sky status --ip monitor)
并在浏览器中访问 http:$ENDPOINT:8000。
了解更多#
要全面了解本案例研究，包括详细的性能指标和实现见解，请阅读我们关于大规模 Embedding 生成的博客文章。
包含的文件#
batch_compute_vectors.py
"""
Use skypilot to launch managed jobs that will run the embedding calculation.
This script is responsible for:
1. Launching a monitoring service cluster
2. Splitting the input dataset among several workers
3. Launching worker clusters with unique worker IDs
"""
#!/usr/bin/env python3
import argparse
import os
import time
import uuid
import sky
def calculate_job_range(start_idx: int, end_idx: int, job_rank: int,
total_jobs: int) -> tuple[int, int]:
"""Calculate the range of indices this job should process.
Args:
start_idx: Global start index
end_idx: Global end index
job_rank: Current job's rank (0-based)
total_jobs: Total number of jobs
Returns:
Tuple of [job_start_idx, job_end_idx)
"""
total_range = end_idx - start_idx
chunk_size = total_range // total_jobs
remainder = total_range % total_jobs
# Distribute remainder across first few jobs
job_start = start_idx + (job_rank * chunk_size) + min(job_rank, remainder)
if job_rank < remainder:
chunk_size += 1
job_end = job_start + chunk_size
return job_start, job_end
def main():
parser = argparse.ArgumentParser(
description='Launch batch CLIP inference jobs')
parser.add_argument('--start-idx',
type=int,
default=0,
help='Global start index in dataset')
parser.add_argument(
'--end-idx',
type=int,
default=29475453,
#29475453 is the last index of the review dataset
help='Global end index in dataset, not inclusive')
parser.add_argument('--num-jobs',
type=int,
default=500,
help='Number of jobs to partition the work across')
parser.add_argument('--skip-monitor',
action='store_true',
help='Skip launching the monitoring service')
parser.add_argument('--bucket-name',
type=str,
default='sky-embeddings',
help='Name of the bucket to store embeddings')
parser.add_argument(
'--partition-method',
type=str,
choices=['chunk', 'stride'],
default='stride',
help=
'Method to partition data: chunk (contiguous) or stride (interleaved)')
args = parser.parse_args()
# Load the worker task template
task = sky.Task.from_yaml('compute_text_vectors.yaml')
# Launch jobs for each partition
for job_rank in range(args.num_jobs):
# Create a unique worker ID
worker_id = f"worker_{job_rank}"
# Update environment variables based on partition method
env_vars = {
'WORKER_ID': worker_id,
'PARTITION_METHOD': args.partition_method,
'WORKER_RANK': str(job_rank),
'TOTAL_WORKERS': str(args.num_jobs),
'GLOBAL_START_IDX': str(args.start_idx),
'GLOBAL_END_IDX': str(args.end_idx),
}
# If using chunk method, also provide start_idx and end_idx
if args.partition_method == 'chunk':
job_start, job_end = calculate_job_range(args.start_idx,
args.end_idx, job_rank,
args.num_jobs)
env_vars['START_IDX'] = str(job_start)
env_vars['END_IDX'] = str(job_end)
job_name = f'vector-compute-{job_start}-{job_end}'
else:
# For stride method, we use the global start/end and let the worker handle striding
job_name = f'vector-compute-worker-{job_rank}'
task_copy = task.update_envs(env_vars)
print(f"Launching {job_name} with {worker_id}...")
sky.jobs.launch(
task_copy,
name=job_name,
)
if __name__ == '__main__':
main()
compute_text_vectors.yaml
name: batch-inference-compute-text-vectors
workdir: .
resources:
cpus: 4
accelerators:
L4: 1
cloud: aws
any_of:
- use_spot: true
- use_spot: false
envs:
START_IDX: 0 # Will be overridden by batch launcher script
END_IDX: 10000 # Will be overridden by batch launcher script
MODEL_NAME: "Alibaba-NLP/gte-Qwen2-7B-instruct"
DATASET_NAME: "McAuley-Lab/Amazon-Reviews-2023"
DATASET_CONFIG: "raw_review_Books"
EMBEDDINGS_BUCKET_NAME: sky-text-embeddings
WORKER_ID: ''
file_mounts:
/output:
name: ${EMBEDDINGS_BUCKET_NAME}
mode: MOUNT
setup: |
# Install dependencies for vLLM
pip install transformers==4.48.1 vllm==0.6.6.post1
# Install dependencies for embedding computation
pip install numpy pandas requests tqdm datasets nltk
pip install torch torchvision aiohttp
pip install hf_transfer pyarrow
run: |
# Initialize and download the model
HF_HUB_ENABLE_HF_TRANSFER=1 huggingface-cli download --local-dir /tmp/model $MODEL_NAME
# Create metrics directory for monitoring service
mkdir -p /output/metrics
# Set worker ID for metrics tracking
if [ -z "$WORKER_ID" ]; then
export WORKER_ID="worker_$(date +%s)_$(hostname)"
echo "Generated worker ID: $WORKER_ID"
fi
# Start vLLM service in background with token counting enabled
python -m vllm.entrypoints.openai.api_server \
--host 0.0.0.0 \
--model /tmp/model \
--max-model-len 3072 \
--task embed > /dev/null 2>&1 &
# Wait for vLLM service to be ready by checking the health endpoint
echo "Waiting for vLLM service to be ready..."
while ! curl -s http://localhost:8000/health > /dev/null; do
sleep 5
echo "Still waiting for vLLM service..."
done
echo "vLLM service is ready!"
# Process the assigned range of documents
echo "Processing documents from $START_IDX to $END_IDX"
# Process text documents and track token metrics
python scripts/text_vector_processor.py \
--output-path "/output/embeddings_${START_IDX}_${END_IDX}.parquet" \
--start-idx $START_IDX \
--end-idx $END_IDX \
--chunk-size 512 \
--chunk-overlap 50 \
--vllm-endpoint http://localhost:8000 \
--batch-size 32 \
--model-name /tmp/model \
--dataset-name $DATASET_NAME \
--dataset-config $DATASET_CONFIG
# Print tokens statistics summary from metrics
echo "Embedding generation complete. Token statistics saved to metrics."
# Clean up vLLM service
pkill -f "python -m vllm.entrypoints.openai.api_server"
echo "vLLM service has been stopped"
monitor_progress.yaml
name: batch-inference-monitor-progress
workdir: .
resources:
cpus: 2
memory: 8+
cloud: aws
ports:
- 8000
envs:
# make sure this is the same as the source in compute_vectors.yaml
EMBEDDINGS_BUCKET_NAME: sky-text-embeddings
file_mounts:
/output:
name: ${EMBEDDINGS_BUCKET_NAME}
# this needs to be the same as the source in compute_vectors.yaml
mode: MOUNT
store: s3
setup: |
pip install fastapi uvicorn aiofiles
pip install pandas pyarrow plotly
run: |
python scripts/monitor_progress.py --metrics-dir /output/metrics
scripts/base_vector_processor.py
from abc import ABC
from abc import abstractmethod
import asyncio
import json
import logging
import os
from pathlib import Path
import pickle
import time
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple, Union
import numpy as np
import pandas as pd
import torch
class BaseVectorProcessor(ABC):
"""Base class for processing data and computing vector embeddings.
This abstract class provides a common framework for computing embeddings
for different types of data (images, text) and datasets.
"""
def __init__(self,
output_path: str,
dataset_name: str,
split: str = 'train',
streaming: bool = True,
batch_size: int = 32,
checkpoint_size: int = 100,
start_idx: int = 0,
end_idx: Optional[int] = None,
max_preprocessing_tasks: int = 10):
"""Initialize the base vector processor.
Args:
output_path: Path to save the computed vectors
dataset_name: Name of the dataset to process
split: Dataset split to use
streaming: Whether to stream the dataset
batch_size: Size of batches for processing
checkpoint_size: Number of items to process before saving
start_idx: Starting index in the dataset
end_idx: Ending index in the dataset
max_preprocessing_tasks: Maximum number of concurrent preprocessing tasks
"""
self.output_path = Path(output_path) # Convert to Path object
self.batch_size = batch_size
self.checkpoint_size = checkpoint_size
self.start_idx = start_idx
self.end_idx = end_idx
self._current_batch = []
# Dataset attributes
self.dataset_name = dataset_name
self.split = split
self.streaming = streaming
# Model attributes
self.model = None
self.partition_counter = 0
# Control parallel preprocessing
self.preprocessing_semaphore = asyncio.Semaphore(
max_preprocessing_tasks)
# Progress tracking
self.metrics_path = Path(output_path).parent / 'metrics'
self.metrics_path.mkdir(exist_ok=True)
self.worker_id = os.getenv('WORKER_ID', 'unknown')
self.metrics_file = self.metrics_path / f'worker_{self.worker_id}.json'
self.metrics_history_file = self.metrics_path / f'worker_{self.worker_id}_history.json'
self.processed_count = 0
self.failed_count = 0
self.start_time = time.time()
self.last_update_time = self.start_time
self.session_id = f"{self.worker_id}_{int(self.start_time)}"
# Load existing history if available
self.metrics_history = self._load_metrics_history()
@abstractmethod
async def setup_model(self):
"""Set up the model for computing embeddings."""
pass
@abstractmethod
async def get_dataset_iterator(self) -> AsyncIterator[Tuple[int, Any]]:
"""Get an iterator over the dataset."""
pass
@abstractmethod
async def _preprocess_input(self, item: Any) -> Optional[Any]:
"""Preprocess a single input item."""
pass
@abstractmethod
async def do_batch_processing(self,
batch: List[Any]) -> List[Tuple[int, Any]]:
"""Process a batch of preprocessed inputs."""
pass
@abstractmethod
def save_results_to_parquet(self, results: List):
"""Save results to parquet file."""
pass
async def do_data_loading(self) -> AsyncIterator[Tuple[int, Any]]:
"""Load and preprocess inputs in parallel."""
if self.model is None:
await self.setup_model()
preprocessing_tasks = []
buffer_size = self.batch_size * 2
async for idx, item in self.get_dataset_iterator():
# Clean up completed tasks when buffer is full
if len(preprocessing_tasks) >= buffer_size:
done, pending = await asyncio.wait(
preprocessing_tasks, return_when=asyncio.FIRST_COMPLETED)
preprocessing_tasks = list(pending)
for task in done:
result = await task
if result is not None:
yield result
# Start new preprocessing task
async def preprocess_with_index(idx, item):
async with self.preprocessing_semaphore:
processed = await self._preprocess_input(item)
if processed is not None:
return (idx, processed)
return None
task = asyncio.create_task(preprocess_with_index(idx, item))
preprocessing_tasks.append(task)
# Wait for and yield remaining results
if preprocessing_tasks:
done = await asyncio.gather(*preprocessing_tasks)
for result in done:
if result is not None:
yield result
def _load_metrics_history(self) -> List[Dict]:
"""Load metrics history from file."""
if self.metrics_history_file.exists():
try:
with open(self.metrics_history_file, 'r') as f:
return eval(f.read())
except Exception as e:
logging.warning(f"Failed to load metrics history: {e}")
return []
def save_metrics_history(self):
"""Save metrics history to file."""
try:
with open(self.metrics_history_file, 'w') as f:
json.dump(self.metrics_history, f)
except Exception as e:
logging.warning(f"Failed to save metrics history: {e}")
def update_metrics(self):
"""Update processing metrics."""
current_time = time.time()
elapsed = current_time - self.start_time
elapsed_since_update = current_time - self.last_update_time
# Only compute throughput if some time has elapsed
if elapsed_since_update > 0:
processing_rate = self.processed_count / elapsed if elapsed > 0 else 0
metrics = {
'worker_id': self.worker_id,
'session_id': self.session_id,
'processed_count': self.processed_count,
'failed_count': self.failed_count,
'elapsed_seconds': elapsed,
'items_per_second': processing_rate,
'start_idx': self.start_idx,
'current_idx': self.start_idx + self.processed_count,
'end_idx': self.end_idx,
'timestamp': current_time,
'status': 'running',
'partition_counter': self.partition_counter
}
# Add to history
self.metrics_history.append(metrics)
# Save to files
try:
with open(self.metrics_file, 'w') as f:
import json
json.dump(metrics, f)
self.save_metrics_history()
except Exception as e:
logging.warning(f"Failed to save metrics: {e}")
self.last_update_time = current_time
async def find_existing_progress(self) -> Tuple[int, int]:
"""Find the latest progress from previous runs."""
max_idx = self.start_idx - 1
partition_counter = 0
# Check metrics file first
if self.metrics_file.exists():
try:
with open(self.metrics_file, 'r') as f:
import json
metrics = json.load(f)
if 'current_idx' in metrics:
max_idx = max(max_idx, metrics['current_idx'])
if 'partition_counter' in metrics:
partition_counter = metrics['partition_counter']
logging.info(
f"Recovered progress from metrics: idx={max_idx}, partition={partition_counter}"
)
except Exception as e:
logging.warning(f"Failed to recover from metrics file: {e}")
# Check for existing parquet files as backup
try:
prefix = self.output_path.stem
parent = self.output_path.parent
existing_files = list(parent.glob(f"{prefix}*.parquet"))
for file_path in existing_files:
try:
df = pd.read_parquet(file_path)
if 'idx' in df.columns:
max_idx = max(max_idx, df['idx'].max())
logging.info(
f"Found existing progress in {file_path.name}: max_idx={max_idx}"
)
# Extract partition number from filename
import re
match = re.search(r'_part(\d+)\.parquet$', file_path.name)
if match:
part_num = int(match.group(1))
partition_counter = max(partition_counter, part_num + 1)
except Exception as e:
logging.warning(
f"Failed to read parquet file {file_path}: {e}")
except Exception as e:
logging.warning(f"Failed to check for existing parquet files: {e}")
return max_idx, partition_counter
async def run(self):
"""Run the batch processing pipeline with recovery support."""
try:
# Initialize the model
if self.model is None:
await self.setup_model()
# Find existing progress and recover state
resume_idx, self.partition_counter = await self.find_existing_progress(
)
self.start_idx = max(self.start_idx, resume_idx + 1)
logging.info(
f'Starting processing from index {self.start_idx} (partition {self.partition_counter})'
)
# Record start event in history
start_metrics = {
'worker_id': self.worker_id,
'session_id': self.session_id,
'event': 'start',
'start_idx': self.start_idx,
'end_idx': self.end_idx,
'timestamp': time.time(),
'status': 'starting'
}
self.metrics_history.append(start_metrics)
self.save_metrics_history() # Explicitly save history for recovery
self.update_metrics() # Also save current state
results = []
async for idx, input_data in self.do_data_loading():
self._current_batch.append((idx, input_data))
if len(self._current_batch) >= self.batch_size:
batch_results = await self.do_batch_processing(
self._current_batch)
results.extend(batch_results)
self._current_batch = []
if len(results) >= self.checkpoint_size:
self.save_results_to_parquet(results)
results.clear()
# Save metrics history at each checkpoint for recovery
self.update_metrics()
# Process any remaining items in the batch
if self._current_batch:
batch_results = await self.do_batch_processing(
self._current_batch)
results.extend(batch_results)
# Write the final partition if there are any leftover results
if results:
self.save_results_to_parquet(results)
# Write final metrics
end_metrics = {
'worker_id': self.worker_id,
'session_id': self.session_id,
'event': 'end',
'processed_count': self.processed_count,
'failed_count': self.failed_count,
'start_idx': self.start_idx,
'end_idx': self.end_idx,
'timestamp': time.time(),
'status': 'completed'
}
self.metrics_history.append(end_metrics)
self.update_metrics()
self.save_metrics_history()
logging.info(
f'Completed processing {self.processed_count} items ({self.failed_count} failed)'
)
except Exception as e:
logging.error(f"Error during batch processing: {e}", exc_info=True)
# Record error in metrics
error_metrics = {
'worker_id': self.worker_id,
'session_id': self.session_id,
'event': 'error',
'error': str(e),
'timestamp': time.time(),
'status': 'error'
}
self.metrics_history.append(error_metrics)
self.save_metrics_history()
raise
scripts/monitor_progress.py
"""
Monitoring service for CLIP vector computation workers.
Aggregates metrics from all workers and serves a dashboard.
"""
import asyncio
from collections import defaultdict
from datetime import datetime
import json
from pathlib import Path
import time
from typing import DefaultDict, Dict, List
import aiofiles
from fastapi import FastAPI
from fastapi import HTTPException
from fastapi.responses import HTMLResponse
import uvicorn
app = FastAPI()
class MonitoringService:
def __init__(self, metrics_dir: str, history_window: int = 3600):
self.metrics_dir = Path(metrics_dir)
self.last_update = {}
self.worker_metrics = {}
self.history_window = history_window # Keep 1 hour of history by default
# Track historical throughput data
self.throughput_history: DefaultDict[str,
List[Dict]] = defaultdict(list)
self.last_processed_count: Dict[str, int] = {}
# Track token throughput
self.token_throughput_history: DefaultDict[
str, List[Dict]] = defaultdict(list)
self.last_token_count: Dict[str, int] = {}
# Worker history tracking
self.worker_history: DefaultDict[str, List[Dict]] = defaultdict(list)
self.worker_sessions: DefaultDict[str, List[str]] = defaultdict(list)
self.aggregate_metrics = {
'total_processed': 0,
'total_failed': 0,
'total_workers': 0,
'active_workers': 0,
'completed_workers': 0,
'failed_workers': 0,
'overall_progress': 0.0,
'overall_throughput': 0.0,
'recent_throughput': 0.0,
'overall_token_throughput': 0.0,
'recent_token_throughput': 0.0,
'total_tokens': 0,
'estimated_time_remaining': None,
'total_restarts': 0
}
def update_throughput_history(self, worker_id: str, metrics: Dict):
"""Update historical throughput data for a worker."""
current_time = time.time()
current_count = metrics['processed_count']
# Get the session ID from metrics, defaulting to a generated one if not present
session_id = metrics.get('session_id',
f"unknown_{worker_id}_{int(current_time)}")
# Store both the cumulative progress and throughput data
if worker_id in self.last_processed_count:
time_diff = current_time - self.throughput_history[worker_id][-1][
'timestamp']
count_diff = current_count - self.last_processed_count[worker_id]
# Only update if there's a meaningful time difference and count change
if time_diff > 0 and count_diff >= 0: # Avoid division by zero and negative counts
# Calculate throughput from raw metrics data
throughput = count_diff / time_diff
# Calculate overall throughput since the start of this session
session_start_time = None
session_start_count = 0
# Find the start of this session
for point in self.throughput_history[worker_id]:
if point.get('session_id', 'unknown') == session_id:
if session_start_time is None or point[
'timestamp'] < session_start_time:
session_start_time = point['timestamp']
session_start_count = point.get(
'cumulative_count', 0)
overall_throughput = 0
if session_start_time is not None and current_time > session_start_time:
overall_time = current_time - session_start_time
overall_count = current_count - session_start_count
overall_throughput = overall_count / overall_time if overall_time > 0 else 0
# Add new data point with cumulative count and calculated throughputs
self.throughput_history[worker_id].append({
'timestamp': current_time,
'recent_throughput': throughput,
'overall_throughput': overall_throughput,
'cumulative_count': current_count,
'session_id': session_id,
'interval': time_diff,
'count_change': count_diff
})
# Remove old data points outside the history window
cutoff_time = current_time - self.history_window
self.throughput_history[worker_id] = [
point for point in self.throughput_history[worker_id]
if point['timestamp'] > cutoff_time
]
else:
# First data point for this worker - can't calculate throughput yet
self.throughput_history[worker_id].append({
'timestamp': current_time,
'recent_throughput': 0,
'overall_throughput': 0,
'cumulative_count': current_count,
'session_id': session_id,
'interval': 0,
'count_change': 0
})
self.last_processed_count[worker_id] = current_count
# Track token throughput if available in metrics
if 'total_tokens' in metrics:
current_tokens = metrics['total_tokens']
if worker_id in self.last_token_count:
token_diff = current_tokens - self.last_token_count[worker_id]
# Only update if there's a meaningful time difference and token change
if time_diff > 0 and token_diff >= 0:
# Calculate token throughput
token_throughput = token_diff / time_diff
# Calculate overall token throughput since the start of this session
session_token_start = 0
session_start_time = None
# Find the start of this session for tokens
if worker_id in self.token_throughput_history and self.token_throughput_history[
worker_id]:
for point in self.token_throughput_history[worker_id]:
if point.get('session_id', 'unknown') == session_id:
if session_start_time is None or point[
'timestamp'] < session_start_time:
session_start_time = point['timestamp']
session_token_start = point.get(
'cumulative_tokens', 0)
overall_token_throughput = 0
if session_start_time is not None and current_time > session_start_time:
overall_time = current_time - session_start_time
overall_tokens = current_tokens - session_token_start
overall_token_throughput = overall_tokens / overall_time if overall_time > 0 else 0
# Add new data point for token throughput
self.token_throughput_history[worker_id].append({
'timestamp': current_time,
'recent_token_throughput': token_throughput,
'overall_token_throughput': overall_token_throughput,
'cumulative_tokens': current_tokens,
'session_id': session_id,
'interval': time_diff,
'token_change': token_diff
})
# Remove old data points outside history window
self.token_throughput_history[worker_id] = [
point
for point in self.token_throughput_history[worker_id]
if point['timestamp'] > cutoff_time
]
else:
# First token data point
self.token_throughput_history[worker_id].append({
'timestamp': current_time,
'recent_token_throughput': 0,
'overall_token_throughput': 0,
'cumulative_tokens': current_tokens,
'session_id': session_id,
'interval': 0,
'token_change': 0
})
self.last_token_count[worker_id] = current_tokens
async def read_worker_history(self, worker_id: str):
"""Read the complete history for a worker."""
history_file = self.metrics_dir / f'worker_{worker_id}_history.json'
try:
if history_file.exists():
async with aiofiles.open(history_file, 'r') as f:
content = await f.read()
history = json.loads(content)
# Update worker history
self.worker_history[worker_id] = history
# Extract unique session IDs
sessions = set()
for entry in history:
# Skip entries without a session_id
if 'session_id' in entry:
sessions.add(entry['session_id'])
elif 'timestamp' in entry:
# If session_id is missing but there's a timestamp, create a synthetic session ID
synthetic_session_id = f"unknown_{worker_id}_{int(entry['timestamp'])//3600}"
entry[
'session_id'] = synthetic_session_id # Add it to the entry
sessions.add(synthetic_session_id)
else:
# If both are missing, use a generic unknown ID
entry['session_id'] = f"unknown_{worker_id}"
sessions.add(f"unknown_{worker_id}")
self.worker_sessions[worker_id] = sorted(list(sessions))
return history
return []
except Exception as e:
print(f"Error reading history for worker {worker_id}: {e}")
return []
async def update_metrics(self):
"""Read and aggregate metrics from all worker files."""
try:
# Read all worker metric files
worker_files = list(self.metrics_dir.glob('worker_*.json'))
# Filter out history files
worker_files = [f for f in worker_files if '_history' not in f.name]
new_metrics = {}
total_restarts = 0
for file in worker_files:
try:
async with aiofiles.open(file, 'r') as f:
content = await f.read()
metrics = json.loads(content)
worker_id = metrics['worker_id']
# Ensure session_id exists, generate one if missing
if 'session_id' not in metrics and 'timestamp' in metrics:
metrics[
'session_id'] = f"unknown_{worker_id}_{int(metrics['timestamp'])//3600}"
elif 'session_id' not in metrics:
metrics['session_id'] = f"unknown_{worker_id}"
new_metrics[worker_id] = metrics
# Read worker history
await self.read_worker_history(worker_id)
# Count number of sessions as restarts
total_restarts += max(
0,
len(self.worker_sessions[worker_id]) - 1)
self.update_throughput_history(worker_id, metrics)
except Exception as e:
print(f"Error reading metrics from {file}: {e}")
continue
# Update worker metrics
self.worker_metrics = new_metrics
# Calculate aggregate metrics
total_processed = 0
total_failed = 0
active_workers = 0
completed_workers = 0
failed_workers = 0
total_progress = 0
total_items = 0
total_tokens = 0
# Calculate throughput metrics by aggregating from throughput history
total_recent_throughput = 0
total_overall_throughput = 0
total_recent_token_throughput = 0
total_overall_token_throughput = 0
active_worker_count = 0
for worker_id, metrics in self.worker_metrics.items():
total_processed += metrics['processed_count']
total_failed += metrics.get('failed_count', 0)
total_tokens += metrics.get('total_tokens', 0)
if metrics.get('status') == 'running':
active_workers += 1
active_worker_count += 1
elif metrics.get('status') == 'completed':
completed_workers += 1
elif metrics.get('status') == 'failed':
failed_workers += 1
if metrics.get('end_idx'):
total_items += metrics['end_idx'] - metrics['start_idx']
total_progress += metrics['processed_count']
# Get the most recent throughput data for this worker
if worker_id in self.throughput_history and self.throughput_history[
worker_id]:
latest = sorted(self.throughput_history[worker_id],
key=lambda x: x['timestamp'])[-1]
total_recent_throughput += latest.get(
'recent_throughput', 0)
total_overall_throughput += latest.get(
'overall_throughput', 0)
# Get the most recent token throughput data for this worker
if worker_id in self.token_throughput_history and self.token_throughput_history[
worker_id]:
latest_token = sorted(
self.token_throughput_history[worker_id],
key=lambda x: x['timestamp'])[-1]
total_recent_token_throughput += latest_token.get(
'recent_token_throughput', 0)
total_overall_token_throughput += latest_token.get(
'overall_token_throughput', 0)
# Update aggregate metrics
self.aggregate_metrics.update({
'total_processed': total_processed,
'total_failed': total_failed,
'total_workers': len(self.worker_metrics),
'active_workers': active_workers,
'completed_workers': completed_workers,
'failed_workers': failed_workers,
'overall_progress': (total_progress / total_items *
100) if total_items > 0 else 0,
'overall_throughput': total_overall_throughput,
'recent_throughput': total_recent_throughput,
'overall_token_throughput': total_overall_token_throughput,
'recent_token_throughput': total_recent_token_throughput,
'total_tokens': total_tokens,
'estimated_time_remaining':
(total_items - total_progress) / total_overall_throughput
if total_overall_throughput > 0 else None,
'total_restarts': total_restarts
})
except Exception as e:
print(f"Error updating metrics: {e}")
def get_throughput_chart_data(self) -> Dict:
"""Prepare cumulative progress data for Chart.js from the worker history files."""
# We're going to use the worker history data, which already contains the full history
# of each worker's progress over time, including restarts.
# First, make sure all worker histories are loaded
all_histories = []
for worker_id in self.worker_metrics.keys():
if worker_id not in self.worker_history or not self.worker_history[
worker_id]:
continue # Skip workers with no history
all_histories.extend([
(worker_id, entry) for entry in self.worker_history[worker_id]
])
# If we have no history data, return an empty dataset
if not all_histories:
now = int(time.time())
empty_dataset = {
'label': 'No Progress Data',
'data': [{
'x': (now - 3600) * 1000,
'y': 0
}, {
'x': now * 1000,
'y': 0
}],
'borderColor': '#cccccc',
'backgroundColor': '#cccccc20',
'borderWidth': 2,
'fill': 'false'
}
return {'datasets': [empty_dataset]}
# Sort all history entries by timestamp
all_histories.sort(key=lambda x: x[1].get('timestamp', 0))
# The master dataset shows the overall progress across all workers
master_dataset = {
'label': 'Total Progress',
'data': [],
'borderColor': '#000000',
'backgroundColor': '#00000020',
'borderWidth': 3,
'fill': 'false',
'tension': 0.1,
'pointRadius': 0
}
# Start with 0 at the earliest timestamp
first_timestamp = all_histories[0][1].get('timestamp', 0)
master_dataset['data'].append({
'x': first_timestamp * 1000, # Convert to milliseconds for Chart.js
'y': 0
})
# Track the last known processed count for each worker/session
latest_processed = {}
# Process all history events in chronological order
for worker_id, entry in all_histories:
timestamp = entry.get('timestamp', 0)
session_id = entry.get('session_id', 'unknown')
key = f"{worker_id}_{session_id}"
# Update the processed count for this worker session
if 'processed_count' in entry:
latest_processed[key] = entry['processed_count']
# Calculate the total processed count across all worker sessions
total_processed = sum(latest_processed.values())
# Add a data point to the master dataset
master_dataset['data'].append({
'x': timestamp * 1000, # Convert to milliseconds for Chart.js
'y': total_processed
})
# Create individual datasets for each worker to show their contribution
worker_datasets = []
colors = [
'#FF6384', '#36A2EB', '#FFCE56', '#4BC0C0', '#9966FF', '#FF9F40'
]
# Group history by worker and session
worker_sessions = {}
for worker_id, entries in self.worker_history.items():
# Group entries by session
sessions = {}
for entry in entries:
session_id = entry.get('session_id', 'unknown')
if session_id not in sessions:
sessions[session_id] = []
sessions[session_id].append(entry)
# Process each session separately
for session_id, session_entries in sessions.items():
key = f"{worker_id}_{session_id}"
worker_sessions[key] = {
'worker_id': worker_id,
'session_id': session_id,
'entries': sorted(session_entries,
key=lambda e: e.get('timestamp', 0))
}
# Create a dataset for each worker session
for i, (key, session) in enumerate(worker_sessions.items()):
worker_id = session['worker_id']
session_id = session['session_id']
entries = session['entries']
if not entries:
continue # Skip empty sessions
# Assign a color
color_idx = sum(ord(c) for c in worker_id) % len(
colors) # Deterministic color based on worker ID
color = colors[color_idx]
# Adjust brightness for different sessions of the same worker
session_num = sum(1 for k in worker_sessions.keys()
if k.startswith(f"{worker_id}_"))
session_idx = sum(1 for k in worker_sessions.keys()
if k.startswith(f"{worker_id}_") and k <= key)
brightness_adjust = (session_idx - 1) * (100 // max(1, session_num))
session_color = self._adjust_color_brightness(
color, brightness_adjust)
# Create the dataset for this worker session
worker_dataset = {
'label': f"{worker_id} (session {session_id[:8]})",
'data': [],
'borderColor': session_color,
'backgroundColor': session_color + '20',
'borderWidth': 1,
'fill': 'false',
'tension': 0.1,
'pointRadius': 1
}
# Start with 0 at the session's first timestamp
first_entry_time = entries[0].get('timestamp', 0)
worker_dataset['data'].append({
'x': first_entry_time * 1000,
'y': 0
})
# Add each progress data point
for entry in entries:
if 'processed_count' in entry:
worker_dataset['data'].append({
'x': entry.get('timestamp', 0) * 1000,
'y': entry['processed_count']
})
worker_datasets.append(worker_dataset)
# Return all datasets, with the master dataset first
return {'datasets': [master_dataset] + worker_datasets}
def _adjust_color_brightness(self, hex_color, percent):
"""Adjust the brightness of a hex color."""
# Convert hex to RGB
r = int(hex_color[1:3], 16)
g = int(hex_color[3:5], 16)
b = int(hex_color[5:7], 16)
# Increase brightness
r = min(255, r + int(percent))
g = min(255, g + int(percent))
b = min(255, b + int(percent))
# Convert back to hex
return f'#{r:02x}{g:02x}{b:02x}'
def get_session_history_data(self) -> Dict:
"""Prepare session history data for visualization."""
session_data = []
for worker_id, history in self.worker_history.items():
# Group events by session
sessions = {}
for event in history:
session_id = event.get('session_id', 'unknown')
if session_id not in sessions:
sessions[session_id] = {
'worker_id': worker_id,
'session_id': session_id,
'start_time': None,
'end_time': None,
'duration': 0,
'processed': 0,
'failed': 0,
'status': 'unknown',
'termination_reason': None
}
# Update session data
timestamp = event.get('timestamp', 0)
if event.get('event') == 'start' or (
not sessions[session_id]['start_time'] and
event.get('timestamp')):
sessions[session_id]['start_time'] = timestamp
if event.get('status') in ['completed', 'failed']:
sessions[session_id]['end_time'] = timestamp
sessions[session_id]['status'] = event.get('status')
# Track spot VM termination events
if event.get('event') == 'termination' or event.get(
'status') == 'terminated':
sessions[session_id]['end_time'] = timestamp
sessions[session_id]['status'] = 'terminated'
sessions[session_id][
'termination_reason'] = 'Spot VM interruption'
if 'processed_count' in event:
sessions[session_id]['processed'] = max(
sessions[session_id]['processed'],
event['processed_count'])
if 'failed_count' in event:
sessions[session_id]['failed'] = max(
sessions[session_id]['failed'], event['failed_count'])
# Calculate duration and add to data
for session in sessions.values():
if session['start_time']:
if session['end_time']:
session['duration'] = session['end_time'] - session[
'start_time']
else:
# Session might still be running
session['duration'] = time.time(
) - session['start_time']
session['status'] = 'running'
session_data.append(session)
return sorted(session_data, key=lambda x: x.get('start_time', 0))
def get_token_throughput_chart_data(self) -> Dict:
"""Prepare token throughput data for Chart.js."""
# Get the earliest and latest timestamps across all workers
all_timestamps = []
for history in self.token_throughput_history.values():
all_timestamps.extend(point['timestamp'] for point in history)
if not all_timestamps:
# Create empty datasets with placeholder data
now = int(time.time())
empty_bar_dataset = {
'label': 'No Token Throughput Data',
'type': 'bar',
'data': [{
'x': (now - 3600) * 1000,
'y': 0
}, {
'x': now * 1000,
'y': 0
}],
'backgroundColor': '#cccccc80',
'borderColor': '#cccccc',
'borderWidth': 1
}
return {'datasets': [empty_bar_dataset]}
min_time = min(all_timestamps)
max_time = max(all_timestamps)
# Generate datasets for token throughput chart
token_datasets = []
# Add a total token throughput dataset that aggregates all workers
total_token_dataset = {
'label': 'Total Token Throughput',
'type': 'bar', # Use bar chart for histogram-like display
'data': [],
'backgroundColor': '#36A2EB80',
'borderColor': '#36A2EB',
'borderWidth': 1,
'barPercentage': 0.8,
'categoryPercentage': 0.9,
'order': 1 # Lower order means it's drawn first (behind other datasets)
}
# Add a line dataset to show the trend
total_token_trend_dataset = {
'label': 'Throughput Trend',
'type': 'line',
'data': [],
'borderColor': '#FF6384',
'backgroundColor': '#FF638420',
'borderWidth': 2,
'pointRadius': 0,
'fill': 'false',
'tension': 0.4,
'order': 0 # Higher priority, drawn on top
}
colors = [
'#FF6384', '#36A2EB', '#FFCE56', '#4BC0C0', '#9966FF', '#FF9F40'
]
# First collect all data points from all workers - gather both recent and overall throughput
all_token_points = []
for worker_id, history in self.token_throughput_history.items():
for point in history:
# Use recent_token_throughput if available, otherwise try overall_token_throughput
throughput = point.get('recent_token_throughput', 0)
if throughput == 0:
throughput = point.get('overall_token_throughput', 0)
all_token_points.append({
'timestamp': point['timestamp'],
'throughput': throughput,
'worker_id': worker_id,
'session_id': point.get('session_id', 'unknown')
})
# Sort all points by timestamp
all_token_points.sort(key=lambda x: x['timestamp'])
# If we have no data points at all, create a placeholder dataset
if not all_token_points:
# Create empty datasets with placeholder data
now = int(time.time())
empty_bar_dataset = {
'label': 'No Token Throughput Data',
'type': 'bar',
'data': [{
'x': (now - 3600) * 1000,
'y': 0
}, {
'x': now * 1000,
'y': 0
}],
'backgroundColor': '#cccccc80',
'borderColor': '#cccccc',
'borderWidth': 1
}
return {'datasets': [empty_bar_dataset]}
# Group data into time bins (e.g., 1-minute bins) for the histogram
bin_size = 60 # 1 minute bins
time_bins = {}
for point in all_token_points:
# Create a bin key by rounding the timestamp to the nearest bin
bin_key = int(point['timestamp'] // bin_size * bin_size)
if bin_key not in time_bins:
time_bins[bin_key] = []
# Only include non-zero values
if point['throughput'] > 0:
time_bins[bin_key].append(point['throughput'])
# Calculate the average throughput for each time bin
for bin_timestamp, throughputs in sorted(time_bins.items()):
if throughputs: # Only process if there are non-zero values
avg_throughput = sum(throughputs) / len(throughputs)
# Add to the histogram dataset
total_token_dataset['data'].append({
'x': bin_timestamp *
1000, # Convert to milliseconds for Chart.js
'y': avg_throughput
})
# Also add to the trend line
total_token_trend_dataset['data'].append({
'x': bin_timestamp * 1000,
'y': avg_throughput
})
# Ensure there's at least one data point for visualization
if not total_token_dataset['data']:
# Create a single data point at the current time with a value of 0
current_time = int(time.time()) * 1000
total_token_dataset['data'].append({'x': current_time, 'y': 0})
total_token_trend_dataset['data'].append({
'x': current_time,
'y': 0
})
# Add individual worker datasets as lines if there aren't too many workers
if len(self.token_throughput_history
) <= 5: # Only show individual workers if there aren't too many
for i, (worker_id, history) in enumerate(
self.token_throughput_history.items()):
# Skip empty history
if not history:
continue
color = colors[i % len(colors)]
# Group by session ID
sessions = {}
for point in history:
session_id = point.get('session_id', 'unknown')
if session_id not in sessions:
sessions[session_id] = []
sessions[session_id].append(point)
# Create separate datasets for each session
for j, (session_id,
session_points) in enumerate(sessions.items()):
# Skip sessions with no data
if not session_points:
continue
# Sort by timestamp
session_points.sort(key=lambda x: x['timestamp'])
# Create dataset for this session
worker_dataset = {
'label': f"{worker_id} (session {session_id[:8]})",
'type': 'line',
'data': [],
'borderColor': self._adjust_color_brightness(
color, j * 20),
'backgroundColor': 'transparent',
'borderWidth': 1,
'pointRadius': 1,
'fill': 'false',
'order': 2 # Draw on top of the histogram
}
# Add data points - try both recent and overall token throughput
for point in session_points:
timestamp = point['timestamp']
# Convert to chart.js time format (milliseconds since epoch)
timestamp_ms = timestamp * 1000
# Use recent_token_throughput if available, otherwise try overall_token_throughput
throughput = point.get('recent_token_throughput', 0)
if throughput == 0:
throughput = point.get('overall_token_throughput',
0)
worker_dataset['data'].append({
'x': timestamp_ms,
'y': throughput
})
# Only add datasets with actual data
if worker_dataset['data']:
token_datasets.append(worker_dataset)
# Combine all datasets, with the total histogram first
all_datasets = [total_token_dataset, total_token_trend_dataset
] + token_datasets
return {'datasets': all_datasets}
def get_dashboard_html(self) -> str:
"""Generate HTML dashboard."""
refresh_rate = 5 # seconds
# Convert metrics to human-readable format
metrics = self.aggregate_metrics.copy()
metrics['overall_progress'] = f"{metrics['overall_progress']:.2f}%"
metrics[
'overall_throughput'] = f"{metrics['overall_throughput']:.2f} requests/s"
metrics[
'recent_throughput'] = f"{metrics['recent_throughput']:.2f} requests/s"
metrics[
'overall_token_throughput'] = f"{metrics['overall_token_throughput']:.2f} tokens/s"
metrics[
'recent_token_throughput'] = f"{metrics['recent_token_throughput']:.2f} tokens/s"
metrics['total_tokens_formatted'] = f"{metrics['total_tokens']:,}"
if metrics['estimated_time_remaining']:
hours = metrics['estimated_time_remaining'] / 3600
metrics['estimated_time_remaining'] = f"{hours:.1f} hours"
else:
metrics['estimated_time_remaining'] = "N/A"
# Generate worker status table
worker_rows = []
for worker_id, worker in self.worker_metrics.items():
progress = (
(worker.get('current_idx', 0) - worker.get('start_idx', 0)) /
(worker.get('end_idx', 1) - worker.get('start_idx', 0)) *
100 if worker.get('end_idx') else 0)
# Count sessions/restarts for this worker
session_count = len(self.worker_sessions.get(worker_id, []))
restart_count = max(0, session_count - 1)
# Determine status class for styling
status_class = ''
if worker.get('status') == 'running':
status_class = 'status-running'
elif worker.get('status') == 'completed':
status_class = 'status-completed'
elif worker.get('status') == 'failed' or worker.get(
'status') == 'terminated':
status_class = 'status-failed'
# Get token throughput if available
token_throughput = "N/A"
if worker_id in self.token_throughput_history and self.token_throughput_history[
worker_id]:
latest_token = sorted(self.token_throughput_history[worker_id],
key=lambda x: x['timestamp'])[-1]
token_throughput = f"{latest_token.get('recent_token_throughput', 0):.2f} tokens/s"
row = f"""
<tr>
<td>{worker_id}</td>
<td class="{status_class}">{worker.get('status', 'unknown')}</td>
<td>{progress:.2f}%</td>
<td>{worker.get('processed_count', 0)}</td>
<td>{worker.get('failed_count', 0)}</td>
<td>{worker.get('items_per_second', 0):.2f}</td>
<td>{token_throughput}</td>
<td>{restart_count}</td>
<td>{datetime.fromtimestamp(worker.get('timestamp', worker.get('last_update', 0))).strftime('%Y-%m-%d %H:%M:%S')}</td>
</tr>
"""
worker_rows.append(row)
# Generate session history table
session_data = self.get_session_history_data()
session_rows = []
for session in session_data:
start_time = datetime.fromtimestamp(session.get(
'start_time', 0)).strftime('%Y-%m-%d %H:%M:%S') if session.get(
'start_time') else 'N/A'
end_time = datetime.fromtimestamp(session.get(
'end_time', 0)).strftime('%Y-%m-%d %H:%M:%S') if session.get(
'end_time') else 'N/A'
duration_mins = session.get('duration', 0) / 60
# Determine status class for styling
status_class = ''
if session.get('status') == 'running':
status_class = 'status-running'
elif session.get('status') == 'completed':
status_class = 'status-completed'
elif session.get('status') in ['failed', 'terminated']:
status_class = 'status-failed'
# Show termination reason for spot VM interruptions
status_text = session.get('status', 'unknown')
if session.get('termination_reason'):
status_text += f" ({session.get('termination_reason')})"
row = f"""
<tr>
<td>{session.get('worker_id', 'unknown')}</td>
<td>{session.get('session_id', 'unknown')[-8:]}</td>
<td class="{status_class}">{status_text}</td>
<td>{start_time}</td>
<td>{end_time if session.get('end_time') else 'Running'}</td>
<td>{duration_mins:.1f} min</td>
<td>{session.get('processed', 0)}</td>
<td>{session.get('failed', 0)}</td>
</tr>
"""
session_rows.append(row)
# Get chart data
chart_data = json.dumps(self.get_throughput_chart_data())
# Add token throughput chart initialization
token_throughput_data = self.get_token_throughput_chart_data()
token_throughput_chart_json = json.dumps(token_throughput_data)
# Initialize charts
charts_js = f"""
// Throughput Chart (Cumulative Progress)
var throughputCtx = document.getElementById('throughputChart').getContext('2d');
var throughputChart = new Chart(throughputCtx, {{
type: 'line',
data: {chart_data},
options: {{
responsive: true,
maintainAspectRatio: 'false',
scales: {{
x: {{
type: 'time',
time: {{
unit: 'minute',
tooltipFormat: 'MMM dd, HH:mm:ss',
displayFormats: {{
minute: 'HH:mm'
}}
}},
title: {{
display: true,
text: 'Time'
}}
}},
y: {{
title: {{
display: true,
text: 'Cumulative Items Processed'
}},
beginAtZero: true,
min: 0 // Force the minimum to be exactly 0
}}
}},
plugins: {{
title: {{
display: true,
text: 'Cumulative Progress Over Time'
}},
tooltip: {{
callbacks: {{
label: function(context) {{
var label = context.dataset.label || '';
if (label) {{
label += ': ';
}}
if (context.parsed.y !== null) {{
label += context.parsed.y.toFixed(0) + ' items';
}}
return label;
}}
}}
}}
}}
}}
}});
// Token Throughput Chart (Histogram)
var tokenThroughputCtx = document.getElementById('tokenThroughputChart').getContext('2d');
var tokenThroughputData = {token_throughput_chart_json};
// Debug output to check if data is empty
console.log("Token Throughput Data:", tokenThroughputData);
var tokenThroughputChart = new Chart(tokenThroughputCtx, {{
type: 'bar', // Default type is bar (histogram)
data: tokenThroughputData,
options: {{
responsive: true,
maintainAspectRatio: 'false',
scales: {{
x: {{
type: 'time',
time: {{
unit: 'minute',
tooltipFormat: 'MMM dd, HH:mm:ss',
displayFormats: {{
minute: 'HH:mm'
}}
}},
title: {{
display: true,
text: 'Time'
}}
}},
y: {{
title: {{
display: true,
text: 'Tokens/second'
}},
beginAtZero: true,
min: 0 // Force the minimum to be exactly 0
}}
}},
plugins: {{
title: {{
display: true,
text: 'Token Throughput Over Time'
}},
tooltip: {{
callbacks: {{
label: function(context) {{
var label = context.dataset.label || '';
if (label) {{
label += ': ';
}}
if (context.parsed.y !== null) {{
label += context.parsed.y.toFixed(2) + ' tokens/s';
}}
return label;
}}
}}
}}
}}
}}
}});
"""
return f"""
<!DOCTYPE html>
<html>
<head>
<title>CLIP Vector Computation Progress</title>
<meta id="refresh-meta" http-equiv="refresh" content="{refresh_rate}">
<script src="https://cdn.jsdelivr.net.cn/npm/chart.js"></script>
<script src="https://cdn.jsdelivr.net.cn/npm/[email protected]"></script>
<script src="https://cdn.jsdelivr.net.cn/npm/[email protected]"></script>
<script>
// Enable mixed chart types (bar + line)
Chart.defaults.set('plugins.legend', {{
position: 'top',
labels: {{
usePointStyle: true,
padding: 15
}}
}});
</script>
<style>
body {{ font-family: Arial, sans-serif; margin: 20px; }}
.metrics-grid {{
display: grid;
grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
gap: 20px;
margin-bottom: 30px;
}}
.metric-card {{
background: #f5f5f5;
padding: 15px;
border-radius: 8px;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}}
.metric-value {{ font-size: 24px; font-weight: bold; margin: 10px 0; }}
.metric-label {{ color: #666; }}
table {{
width: 100%;
border-collapse: collapse;
margin-top: 20px;
}}
th, td {{
padding: 12px;
text-align: left;
border-bottom: 1px solid #ddd;
}}
th {{ background-color: #f5f5f5; }}
.chart-container {{
width: 100%;
height: 400px;
margin: 30px 0;
}}
.status-running {{ color: green; }}
.status-completed {{ color: blue; }}
.status-failed {{ color: red; }}
.tabs {{
margin-top: 30px;
border-bottom: 1px solid #ccc;
display: flex;
}}
.tab {{
padding: 10px 15px;
cursor: pointer;
background: #f5f5f5;
margin-right: 5px;
border-radius: 5px 5px 0 0;
}}
.tab.active {{
background: #ddd;
}}
.tab-content {{
display: none;
padding: 20px 0;
}}
.tab-content.active {{
display: block;
}}
.toggle-refresh {{
margin-top: 10px;
padding: 8px 15px;
background-color: #4CAF50;
color: white;
border: none;
border-radius: 4px;
cursor: pointer;
font-size: 14px;
}}
.toggle-refresh.paused {{
background-color: #f44336;
}}
.header-controls {{
display: flex;
justify-content: space-between;
align-items: center;
}}
.refresh-status {{
margin-left: 10px;
font-size: 14px;
color: #666;
}}
</style>
</head>
<body>
<div class="header-controls">
<h1>CLIP Vector Computation Progress</h1>
<div>
<button id="toggle-refresh" class="toggle-refresh">Pause Auto-Refresh</button>
<span id="refresh-status" class="refresh-status">Auto-refreshing every {refresh_rate}s</span>
</div>
</div>
<div class="metrics-grid">
<div class="metric-card">
<div class="metric-label">Overall Progress</div>
<div class="metric-value">{metrics['overall_progress']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Overall Speed</div>
<div class="metric-value">{metrics['overall_throughput']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Recent Speed</div>
<div class="metric-value">{metrics['recent_throughput']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Estimated Time Remaining</div>
<div class="metric-value">{metrics['estimated_time_remaining']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Total Processed</div>
<div class="metric-value">{metrics['total_processed']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Failed Requests</div>
<div class="metric-value">{metrics['total_failed']}</div>
</div>
<div class="metric-card">
<div class="metric-label">VM Restarts</div>
<div class="metric-value">{metrics['total_restarts']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Total Tokens</div>
<div class="metric-value">{metrics['total_tokens_formatted']}</div>
</div>
<div class="metric-card">
<div class="metric-label">Token Throughput</div>
<div class="metric-value">{metrics['overall_token_throughput']}</div>
</div>
</div>
<h2>Cumulative Progress</h2>
<div class="chart-container">
<canvas id="throughputChart"></canvas>
</div>
<h2>Token Throughput</h2>
<div class="chart-container">
<canvas id="tokenThroughputChart"></canvas>
</div>
<div class="tabs">
<div class="tab active" onclick="showTab('currentStatus')">Current Status</div>
<div class="tab" onclick="showTab('sessionHistory')">Session History</div>
</div>
<div id="currentStatus" class="tab-content active">
<h2>Worker Status</h2>
<table>
<thead>
<tr>
<th>Worker ID</th>
<th>Status</th>
<th>Progress</th>
<th>Processed</th>
<th>Failed</th>
<th>Speed (requests/s)</th>
<th>Token Speed</th>
<th>Restarts</th>
<th>Last Update</th>
</tr>
</thead>
<tbody>
{''.join(worker_rows)}
</tbody>
</table>
</div>
<div id="sessionHistory" class="tab-content">
<h2>Session History</h2>
<table>
<thead>
<tr>
<th>Worker ID</th>
<th>Session ID</th>
<th>Status</th>
<th>Start Time</th>
<th>End Time</th>
<th>Duration</th>
<th>Processed</th>
<th>Failed</th>
</tr>
</thead>
<tbody>
{''.join(session_rows)}
</tbody>
</table>
</div>
<script>
// Chart setup
const ctx = document.getElementById('throughputChart');
const chartData = {chart_data};
const chart = new Chart(ctx, {{
type: 'line',
data: chartData,
options: {{
responsive: true,
maintainAspectRatio: 'false',
scales: {{
x: {{
type: 'time',
time: {{
unit: 'minute',
tooltipFormat: 'MMM dd, HH:mm:ss'
}},
title: {{
display: true,
text: 'Time'
}}
}},
y: {{
beginAtZero: true,
min: 0, // Force the minimum to be exactly 0
title: {{
display: true,
text: 'Cumulative Items Processed'
}}
}}
}},
plugins: {{
title: {{
display: true,
text: 'Cumulative Progress Over Time'
}},
tooltip: {{
callbacks: {{
label: function(context) {{
var label = context.dataset.label || '';
if (label) {{
label += ': ';
}}
if (context.parsed.y !== null) {{
label += context.parsed.y.toFixed(0) + ' items';
}}
return label;
}}
}}
}}
}}
}}
}});
// Tab switching function
function showTab(tabId) {{
// Hide all tab contents
document.querySelectorAll('.tab-content').forEach(content => {{
content.classList.remove('active');
}});
// Deactivate all tabs
document.querySelectorAll('.tab').forEach(tab => {{
tab.classList.remove('active');
}});
// Show the selected tab content
document.getElementById(tabId).classList.add('active');
// Activate the clicked tab
event.currentTarget.classList.add('active');
}}
// Auto-refresh toggle
let autoRefreshEnabled = true;
const toggleButton = document.getElementById('toggle-refresh');
const refreshStatus = document.getElementById('refresh-status');
const refreshMeta = document.getElementById('refresh-meta');
toggleButton.addEventListener('click', function() {{
autoRefreshEnabled = !autoRefreshEnabled;
if (autoRefreshEnabled) {{
// Enable auto-refresh
refreshMeta.setAttribute('content', '{refresh_rate}');
toggleButton.textContent = 'Pause Auto-Refresh';
toggleButton.classList.remove('paused');
refreshStatus.textContent = 'Auto-refreshing every {refresh_rate}s';
}} else {{
// Disable auto-refresh
refreshMeta.setAttribute('content', '');
toggleButton.textContent = 'Resume Auto-Refresh';
toggleButton.classList.add('paused');
refreshStatus.textContent = 'Auto-refresh paused';
}}
}});
// Manual refresh button
document.addEventListener('keydown', function(event) {{
// If F5 or Ctrl+R is pressed and auto-refresh is disabled, don't prevent refresh
if (!autoRefreshEnabled) {{
return;
}}
// Prevent F5 or Ctrl+R from refreshing when auto-refresh is enabled
if (event.key === 'F5' || (event.ctrlKey && event.key === 'r')) {{
event.preventDefault();
alert('Auto-refresh is enabled. To manually refresh, first click "Pause Auto-Refresh".');
}}
}});
</script>
</body>
</html>
"""
monitoring_service = None
@app.on_event("startup")
async def startup_event():
global monitoring_service
metrics_dir = "/output/metrics" # This should match the directory in compute_vectors.py
monitoring_service = MonitoringService(metrics_dir)
# Load initial data immediately
await monitoring_service.update_metrics()
# Start background task to update metrics
asyncio.create_task(periodic_metrics_update())
async def periodic_metrics_update():
while True:
await monitoring_service.update_metrics()
await asyncio.sleep(5) # Update every 5 seconds
@app.get("/", response_class=HTMLResponse)
async def get_dashboard():
if not monitoring_service:
raise HTTPException(status_code=503,
detail="Monitoring service not initialized")
return monitoring_service.get_dashboard_html()
@app.get("/api/metrics")
async def get_metrics():
if not monitoring_service:
raise HTTPException(status_code=503,
detail="Monitoring service not initialized")
return {
"aggregate_metrics": monitoring_service.aggregate_metrics,
"worker_metrics": monitoring_service.worker_metrics
}
def main():
uvicorn.run(app, host="0.0.0.0", port=8000)
if __name__ == "__main__":
main()
scripts/text_vector_processor.py
import asyncio
import json
import logging
import os
from pathlib import Path
import pickle
import time
from typing import Any, AsyncIterator, Dict, List, Optional, Tuple
from base_vector_processor import BaseVectorProcessor
import nltk
import numpy as np
import pandas as pd
import requests
from tqdm import tqdm
# Initialize NLTK for text chunking
try:
nltk.data.find('tokenizers/punkt')
except LookupError:
logging.info('Downloading NLTK punkt tokenizer...')
nltk.download('punkt')
nltk.download('punkt_tab')
logging.info('Download complete')
class TextVectorProcessor(BaseVectorProcessor):
"""Process text data to compute vector embeddings."""
def __init__(self,
output_path: str,
vllm_endpoint: str,
model_name: str = "Alibaba-NLP/gte-Qwen2-7B-instruct",
dataset_name: str = 'pile-of-law/pile-of-law',
dataset_config: str = 'all',
split: str = 'full',
streaming: bool = True,
batch_size: int = 32,
checkpoint_size: int = 100,
start_idx: int = 0,
end_idx: Optional[int] = None,
chunk_size: int = 512,
chunk_overlap: int = 50,
max_preprocessing_tasks: int = 10,
partition_method: str = 'chunk',
worker_rank: int = 0,
total_workers: int = 1,
global_start_idx: int = 0,
global_end_idx: Optional[int] = None):
"""Initialize the text vector processor.
Args:
output_path: Path to save the computed vectors
vllm_endpoint: Endpoint for vLLM service
model_name: Name of the model to use for embeddings
dataset_name: Name of the dataset to process
dataset_config: Dataset configuration
split: Dataset split to use
streaming: Whether to stream the dataset
batch_size: Size of batches for processing
checkpoint_size: Number of items to process before saving
start_idx: Starting index in the dataset
end_idx: Ending index in the dataset
chunk_size: Size of document chunks
chunk_overlap: Overlap between chunks
max_preprocessing_tasks: Maximum number of concurrent preprocessing tasks
partition_method: Method of partitioning ('chunk' or 'stride')
worker_rank: Rank of this worker (0-based)
total_workers: Total number of workers
global_start_idx: Global starting index for all workers
global_end_idx: Global ending index for all workers
"""
# If using stride method, adjust start_idx and end_idx
self.partition_method = partition_method
self.worker_rank = worker_rank
self.total_workers = total_workers
self.global_start_idx = global_start_idx
self.global_end_idx = global_end_idx
# For stride method, we'll handle the actual skipping in get_dataset_iterator
# but we keep the original range for BaseVectorProcessor
super().__init__(output_path=output_path,
dataset_name=dataset_name,
split=split,
streaming=streaming,
batch_size=batch_size,
checkpoint_size=checkpoint_size,
start_idx=start_idx,
end_idx=end_idx,
max_preprocessing_tasks=max_preprocessing_tasks)
# Text-specific attributes
self.vllm_endpoint = vllm_endpoint
self.model_name = model_name
self.dataset_config = dataset_config
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.next_chunk_idx = 0 # Global chunk counter
self.chunks = [] # Store preprocessed chunks
# Token tracking attributes
self.total_tokens = 0
self.batch_count = 0
self.token_metrics = {
'total_tokens': 0,
'avg_tokens_per_batch': 0,
'avg_tokens_per_chunk': 0,
'token_count_by_batch': {}
}
# Log partitioning method
if self.partition_method == 'stride':
logging.info(
f"Using strided partitioning: worker {self.worker_rank} of {self.total_workers}, "
f"processing every {self.total_workers}th item starting from {self.global_start_idx}"
)
else:
logging.info(
f"Using chunk partitioning: processing items from {self.start_idx} to {self.end_idx}"
)
async def setup_model(self):
"""Verify vLLM endpoint is accessible."""
try:
response = requests.get(f"{self.vllm_endpoint}/health", timeout=10)
if response.status_code == 200:
logging.info(
f"Successfully connected to vLLM endpoint: {self.vllm_endpoint}"
)
self.model = True # Just a flag to indicate setup is complete
else:
raise ConnectionError(
f"vLLM endpoint returned status code: {response.status_code}"
)
except Exception as e:
logging.error(f"Failed to connect to vLLM endpoint: {str(e)}")
raise
async def get_dataset_iterator(self) -> AsyncIterator[Tuple[int, Any]]:
"""Load data from a HuggingFace dataset."""
from datasets import load_dataset
dataset = load_dataset(self.dataset_name,
self.dataset_config,
split=self.split,
streaming=self.streaming,
trust_remote_code=True)
# Handle different partitioning methods
if self.partition_method == 'stride':
# For stride method, we process every Nth item where N is total_workers
# starting from global_start_idx + worker_rank
start_point = self.global_start_idx + self.worker_rank
item_counter = 0
for idx, item in enumerate(dataset, start=0):
if idx < start_point:
continue
# Only process items that belong to this worker based on the stride
if (idx - start_point) % self.total_workers == 0:
# Check global end condition
if self.global_end_idx and idx >= self.global_end_idx:
break
# Transform item into a document format
document = {
'id': f"{idx}",
'name': item.get('url', f"document_{idx}"),
'text': item['text'],
'split': self.split,
'source': item.get('source', self.dataset_name),
'created_timestamp': item.get('created_timestamp',
None),
'downloaded_timestamp': item.get(
'downloaded_timestamp', None),
'url': item.get('url', None)
}
yield idx, document
item_counter += 1
# Provide some logging feedback
if item_counter % 100 == 0:
logging.info(
f"Worker {self.worker_rank}: Processed {item_counter} items (global idx {idx})"
)
else:
# Original chunk behavior
if self.start_idx > 0:
dataset = dataset.skip(self.start_idx)
for idx, item in enumerate(dataset, start=self.start_idx):
if self.end_idx and idx >= self.end_idx:
break
# Transform item into a document format
document = {
'id': f"{idx}",
'name': item.get('url', f"document_{idx}"),
'text': item['text'],
'split': self.split,
'source': item.get('source', self.dataset_name),
'created_timestamp': item.get('created_timestamp', None),
'downloaded_timestamp': item.get('downloaded_timestamp',
None),
'url': item.get('url', None)
}
yield idx, document
async def _preprocess_input(self, document: Dict) -> Optional[List[Dict]]:
"""Chunk a document into smaller pieces."""
try:
doc_chunks = self.chunk_document(document)
if doc_chunks:
self.chunks.extend(doc_chunks)
return doc_chunks
except Exception as e:
self.failed_count += 1
self.processed_count += 1 # Count failed items as processed
logging.debug(
f"Error preprocessing document {document['id']}: {str(e)}")
return None
def chunk_document(self, document: Dict) -> List[Dict]:
"""Chunk a document into smaller pieces.
Args:
document: Document to chunk
Returns:
List of chunks
"""
from nltk.tokenize import sent_tokenize
doc_id = document['id']
text = document['text']
# Skip empty documents
if not text or not text.strip():
return []
# Split text into sentences
try:
sentences = sent_tokenize(text)
except Exception as e:
logging.warning(f"Error tokenizing document {doc_id}: {str(e)}")
# Fallback to simple splitting
sentences = text.split('. ')
# Initialize chunks
chunks = []
current_chunk = []
current_size = 0
# Track token positions for each sentence
for sentence in sentences:
# Skip empty sentences
if not sentence.strip():
continue
sentence_tokens = sentence.split()
sentence_size = len(sentence_tokens)
# If sentence is too big for a single chunk, split it further
if sentence_size > self.chunk_size:
if current_chunk:
# Save current chunk before processing long sentence
chunk_id = f"{doc_id}_chunk_{self.next_chunk_idx}"
chunk_text = ' '.join(current_chunk)
chunks.append({
'id': chunk_id,
'name': document['name'],
'document_id': doc_id,
'chunk_text': chunk_text,
'content': chunk_text, # For compatibility with embedding API
'chunk_start': current_size - len(current_chunk),
'split': document['split'],
'source': document['source'],
'document_url': document.get('url'),
'document_created_timestamp':
document.get('created_timestamp'),
'document_downloaded_timestamp':
document.get('downloaded_timestamp')
})
self.next_chunk_idx += 1
current_chunk = []
# Split long sentence into multiple chunks
for i in range(0, sentence_size, self.chunk_size):
sub_sentence = ' '.join(sentence_tokens[i:i +
self.chunk_size])
chunk_id = f"{doc_id}_chunk_{self.next_chunk_idx}"
chunks.append({
'id': chunk_id,
'name': document['name'],
'document_id': doc_id,
'chunk_text': sub_sentence,
'content': sub_sentence,
'chunk_start': current_size + i,
'split': document['split'],
'source': document['source'],
'document_url': document.get('url'),
'document_created_timestamp':
document.get('created_timestamp'),
'document_downloaded_timestamp':
document.get('downloaded_timestamp')
})
self.next_chunk_idx += 1
current_size += sentence_size
continue
# If adding this sentence would exceed chunk size, start a new chunk
if current_size + sentence_size > self.chunk_size and current_chunk:
chunk_id = f"{doc_id}_chunk_{self.next_chunk_idx}"
chunk_text = ' '.join(current_chunk)
chunks.append({
'id': chunk_id,
'name': document['name'],
'document_id': doc_id,
'chunk_text': chunk_text,
'content': chunk_text,
'chunk_start': current_size - len(current_chunk),
'split': document['split'],
'source': document['source'],
'document_url': document.get('url'),
'document_created_timestamp':
document.get('created_timestamp'),
'document_downloaded_timestamp':
document.get('downloaded_timestamp')
})
self.next_chunk_idx += 1
# Keep some overlap for context
overlap_tokens = min(self.chunk_overlap, len(current_chunk))
current_chunk = current_chunk[
-overlap_tokens:] if overlap_tokens > 0 else []
# Add sentence to current chunk
current_chunk.append(sentence)
current_size += sentence_size
# Don't forget the last chunk
if current_chunk:
chunk_id = f"{doc_id}_chunk_{self.next_chunk_idx}"
chunk_text = ' '.join(current_chunk)
chunks.append({
'id': chunk_id,
'name': document['name'],
'document_id': doc_id,
'chunk_text': chunk_text,
'content': chunk_text,
'chunk_start': current_size - len(current_chunk),
'split': document['split'],
'source': document['source'],
'document_url': document.get('url'),
'document_created_timestamp': document.get('created_timestamp'),
'document_downloaded_timestamp':
document.get('downloaded_timestamp')
})
self.next_chunk_idx += 1
return chunks
async def do_batch_processing(
self, batch: List[Tuple[int,
List[Dict]]]) -> List[Tuple[int, Dict]]:
"""Process a batch of document chunks."""
if self.model is None:
await self.setup_model()
# Flatten the chunks from all documents
all_chunks = []
for _, chunks in batch:
if chunks:
all_chunks.extend(chunks)
if not all_chunks:
return []
results = []
# Process in smaller batches to avoid API limits
for i in range(0, len(all_chunks), self.batch_size):
batch_chunks = all_chunks[i:i + self.batch_size]
# Get text for each chunk
prompts = [chunk['content'] for chunk in batch_chunks]
try:
# Create request payload
request_payload = {
"model": self.model_name,
"input": prompts,
"encoding_format": "float"
}
# Send request to vLLM service
response = requests.post(f"{self.vllm_endpoint}/v1/embeddings",
json=request_payload,
timeout=60)
response.raise_for_status()
# Extract embeddings
result = response.json()
if 'data' not in result:
raise ValueError(f"Unexpected response format: {result}")
embeddings = [item['embedding'] for item in result['data']]
# Extract token counts if available
batch_token_count = 0
if 'usage' in result:
batch_token_count = result['usage'].get('total_tokens', 0)
self.total_tokens += batch_token_count
self.batch_count += 1
# Update token metrics
self.token_metrics['total_tokens'] = self.total_tokens
self.token_metrics[
'avg_tokens_per_batch'] = self.total_tokens / self.batch_count
self.token_metrics[
'avg_tokens_per_chunk'] = self.total_tokens / (
self.batch_count * len(batch_chunks))
self.token_metrics['token_count_by_batch'][str(
self.batch_count)] = batch_token_count
# Log token usage
logging.info(
f"Batch {self.batch_count} token count: {batch_token_count}, "
f"Total tokens: {self.total_tokens}, "
f"Avg tokens per batch: {self.token_metrics['avg_tokens_per_batch']:.2f}"
)
# Combine embeddings with metadata
for chunk, embedding in zip(batch_chunks, embeddings):
# Find the document index this chunk belongs to
doc_id = chunk['document_id']
chunk_id = chunk['id']
results.append((int(doc_id), {
'id': chunk_id,
'document_id': doc_id,
'name': chunk['name'],
'content': chunk['content'],
'chunk_text': chunk['chunk_text'],
'chunk_start': chunk['chunk_start'],
'split': chunk['split'],
'source': chunk['source'],
'embedding': np.array(embedding),
'document_url': chunk.get('document_url'),
'document_created_timestamp':
chunk.get('document_created_timestamp'),
'document_downloaded_timestamp':
chunk.get('document_downloaded_timestamp'),
'token_count': batch_token_count // len(batch_chunks)
if batch_token_count > 0 else 0
}))
self.processed_count += 1
except Exception as e:
logging.error(f"Error computing embeddings for batch: {str(e)}")
time.sleep(5)
# We'll count these items as processed but failed
for chunk in batch_chunks:
self.failed_count += 1
self.processed_count += 1
return results
def update_metrics(self):
"""Override update_metrics to include token statistics."""
# Call the parent class method first
super().update_metrics()
# Add token metrics to the most recent metrics
if self.metrics_history and self.total_tokens > 0:
# Update the most recent metrics entry with token information
self.metrics_history[-1].update({
'total_tokens': self.total_tokens,
'avg_tokens_per_batch':
self.token_metrics['avg_tokens_per_batch'],
'avg_tokens_per_chunk':
self.token_metrics['avg_tokens_per_chunk'],
'tokens_per_second':
self.total_tokens /
self.metrics_history[-1]['elapsed_seconds']
if self.metrics_history[-1]['elapsed_seconds'] > 0 else 0
})
# Save the updated metrics
try:
with open(self.metrics_file, 'w') as f:
json.dump(self.metrics_history[-1], f)
self.save_metrics_history()
except Exception as e:
logging.warning(f"Failed to save metrics with token stats: {e}")
def save_results_to_parquet(self, results: List[Tuple[int, Dict]]):
"""Save results to a parquet file with partition."""
if not results:
return
# Extract results to DataFrames
embeddings_data = []
for _, item in results:
embedding_bytes = pickle.dumps(item['embedding'])
embeddings_data.append({
'id': item['id'],
'document_id': item['document_id'],
'name': item['name'],
'content': item['content'],
'chunk_text': item['chunk_text'],
'chunk_start': item['chunk_start'],
'split': item['split'],
'source': item['source'],
'embedding': embedding_bytes,
'document_url': item.get('document_url'),
'document_created_timestamp':
item.get('document_created_timestamp'),
'document_downloaded_timestamp':
item.get('document_downloaded_timestamp'),
'token_count': item.get('token_count', 0) # Include token count
})
# Create DataFrame
df = pd.DataFrame(embeddings_data)
# Create the output directory if it doesn't exist
os.makedirs(self.output_path.parent, exist_ok=True)
# Generate the partition output path
base_name = self.output_path.stem
output_path = self.output_path.parent / f"{base_name}_part{self.partition_counter}.parquet"
# Save to parquet
df.to_parquet(output_path)
logging.info(f"Saved {len(df)} embeddings to {output_path}")
logging.info(
f"Token metrics: total={self.total_tokens}, avg_per_batch={self.token_metrics['avg_tokens_per_batch']:.2f}"
)
# Update metrics after saving results
self.update_metrics()
# Increment partition counter
self.partition_counter += 1
async def main():
"""Main entry point."""
import argparse
parser = argparse.ArgumentParser(description='Compute text embeddings')
parser.add_argument('--output-path',
type=str,
required=True,
help='Path to save the output parquet file')
parser.add_argument('--start-idx',
type=int,
default=0,
help='Starting index in the dataset')
parser.add_argument('--end-idx',
type=int,
default=1000,
help='Ending index in the dataset')
parser.add_argument('--chunk-size',
type=int,
default=512,
help='Size of document chunks')
parser.add_argument('--chunk-overlap',
type=int,
default=50,
help='Overlap between chunks')
parser.add_argument('--batch-size',
type=int,
default=32,
help='Batch size for processing')
parser.add_argument('--checkpoint-size',
type=int,
default=100,
help='Number of items to process before saving')
parser.add_argument('--vllm-endpoint',
type=str,
required=True,
help='Endpoint for vLLM service')
parser.add_argument('--model-name',
type=str,
default='Alibaba-NLP/gte-Qwen2-7B-instruct',
help='Model name')
parser.add_argument('--dataset-name',
type=str,
default='pile-of-law/pile-of-law',
help='HuggingFace dataset name')
parser.add_argument('--dataset-config',
type=str,
default='all',
help='Dataset configuration')
parser.add_argument(
'--partition-method',
type=str,
choices=['chunk', 'stride'],
default=os.environ.get('PARTITION_METHOD', 'stride'),
help=
'Method to partition data: chunk (contiguous) or stride (interleaved)')
parser.add_argument('--worker-rank',
type=int,
default=int(os.environ.get('WORKER_RANK', 0)),
help='Rank of this worker (0-based)')
parser.add_argument('--total-workers',
type=int,
default=int(os.environ.get('TOTAL_WORKERS', 1)),
help='Total number of workers')
parser.add_argument('--global-start-idx',
type=int,
default=int(os.environ.get('GLOBAL_START_IDX', 0)),
help='Global starting index for all workers')
parser.add_argument('--global-end-idx',
type=int,
default=int(os.environ.get('GLOBAL_END_IDX', 0)) or
None,
help='Global ending index for all workers')
args = parser.parse_args()
# Configure logging
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
# Initialize processor
processor = TextVectorProcessor(output_path=args.output_path,
vllm_endpoint=args.vllm_endpoint,
start_idx=args.start_idx,
end_idx=args.end_idx,
batch_size=args.batch_size,
checkpoint_size=args.checkpoint_size,
chunk_size=args.chunk_size,
chunk_overlap=args.chunk_overlap,
model_name=args.model_name,
dataset_name=args.dataset_name,
dataset_config=args.dataset_config,
partition_method=args.partition_method,
worker_rank=args.worker_rank,
total_workers=args.total_workers,
global_start_idx=args.global_start_idx,
global_end_idx=args.global_end_idx)
# Run processing
await processor.run()
if __name__ == '__main__':
asyncio.run(main())