#!/bin/bash
# Ray Worker Entrypoint
# Connects to Ray head node and registers custom resources

set -e

# Ensure Ray is in PATH (works across all base images)
export PATH="/home/ray/.local/bin:/home/ray/anaconda3/bin:${PATH}"

# Get Ray head address from environment or default
RAY_HEAD_ADDRESS="${RAY_HEAD_SVC:-ray-head-svc}:6379"

# Get custom resources from environment
GPU_RESOURCE="${GPU_RESOURCE:-gpu_amd}"
NUM_GPUS="${NUM_GPUS:-1}"

echo "Starting Ray worker..."
echo "  Head address: $RAY_HEAD_ADDRESS"
echo "  GPU resource: $GPU_RESOURCE"
echo "  Num GPUs: $NUM_GPUS"

# Start Ray worker with custom resources
exec ray start \
    --address="$RAY_HEAD_ADDRESS" \
    --num-gpus="$NUM_GPUS" \
    --resources="{\"$GPU_RESOURCE\": 1}" \
    --block