# Use a Python base image with necessary libraries
FROM python:3.9-slim

# Set environment variables to avoid prompts
ENV PYTHONUNBUFFERED 1

# Set the working directory inside the container
WORKDIR /app

# Install system dependencies
RUN apt-get update && apt-get install -y \
    git \
    && rm -rf /var/lib/apt/lists/*

# Install necessary Python packages
RUN pip install --upgrade pip
RUN pip install git+https://github.com/bigscience-workshop/petals

# Copy the Python script into the container
COPY setup_petal_swarm.py /setup_petal_swarm.py

# Expose the necessary port (31337 is the default port for Petals)
EXPOSE 31337

# Set the command to run the Python script on container startup
CMD ["python", "/setup_petal_swarm.py"]