import sky
from sky.serve.service_spec import SkyServiceSpec 

run_str = """
docker run  \
--runtime nvidia \
--gpus all \
--shm-size 1g \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HUGGING_FACE_HUB_TOKEN=${HUGGINGFACE_TOKEN}" \
--ipc=host \
-p 8000:8000 \
vllm/vllm-openai:latest \
--trust-remote-code \                                                                            
--tensor-parallel-size $SKYPILOT_NUM_GPUS_PER_NODE \
--max-model-len 18000 \
--chat-template '<s>[INST] You are a fair judge assistant tasked with providing clear, objective feedback based on specific criteria, ensuring each assessment reflects the absolute standards set for performance. </s>{% for message in messages %}{% if message["role"] == "user" %}{{ "[INST] " + message["content"] + " [/INST]" }}{% elif message["role"] == "assistant" %}{{ message["content"] + "</s> " }}{% endif %}{% endfor %}' \
--model prometheus-eval/prometheus-7b-v2.0
"""
task = sky.Task(
    envs={
        "HUGGINGFACE_TOKEN": "<token>",
    },
    run=run_str,
)

task.set_service(SkyServiceSpec(
    "/v1/models",
    1200,
    1,
))
resources = sky.Resources(
    cloud=sky.clouds.AWS(),
    ports="8080",
    region="us-west-2",
    accelerators="A10G:1",
    instance_type="g5.8xlarge",
)
task.set_resources(resources)
service_name, endpoint = sky.serve.core.up(task)