FedML Adaptation for K8S Pod Solution

devops/k8s/fedml-k8s-agent/Dockerfile

@@ -0,0 +1,93 @@
+# docker build -t fedml/fedml-k8s-agent:202412271540 -f /root/fedml-k8s-agent/Dockerfile .
+FROM nvidia/cuda:12.1.0-cudnn8-devel-ubuntu22.04
+
+# Set environment variables
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PATH=/root/miniconda3/bin:$PATH
+
+# Set default environment variables (these can be overridden in K8s deployment)
+ENV FEDML_API_KEY=""
+ENV FEDML_ENV="release"
+ENV FEDML_PROVIDER="false"
+ENV FEDML_DEVICE_ID=""
+ENV INFERENCE_GATEWAY_PORT=""
+ENV INFERENCE_PROXY_PORT=""
+ENV FEDML_CONNECTION_TYPE=""
+ENV SSH_PUBLIC_KEY=""
+
+# Add build argument to bypass cache
+ARG CACHEBUST=1
+
+# Install system dependencies
+RUN apt-get update \
+    && apt-get install -y \
+    vim \
+    openssh-server \
+    net-tools \
+    wget \
+    curl \
+    iproute2 \
+    ca-certificates \
+    gnupg \
+    sudo \
+    lsb-release \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Configure SSH base setup
+RUN mkdir -p /var/run/sshd \
+    && mkdir -p /root/.ssh \
+    && chmod 700 /root/.ssh \
+    && echo 'Port 22\n\
+PermitRootLogin prohibit-password\n\
+PubkeyAuthentication yes\n\
+AuthorizedKeysFile .ssh/authorized_keys\n\
+PasswordAuthentication no\n\
+ChallengeResponseAuthentication no\n\
+UsePAM yes\n\
+X11Forwarding yes\n\
+PrintMotd no\n\
+AcceptEnv LANG LC_*\n\
+Subsystem sftp /usr/lib/openssh/sftp-server\n\
+LogLevel INFO' > /etc/ssh/sshd_config \
+    && chmod 600 /etc/ssh/sshd_config
+
+# Install Redis
+RUN curl -fsSL https://packages.redis.io/gpg | gpg --dearmor -o /usr/share/keyrings/redis-archive-keyring.gpg \
+    && echo "deb [signed-by=/usr/share/keyrings/redis-archive-keyring.gpg] https://packages.redis.io/deb $(lsb_release -cs) main" \
+    | tee /etc/apt/sources.list.d/redis.list \
+    && apt-get update \
+    && apt-get install -y redis \
+    && rm -rf /var/lib/apt/lists/*
+
+# Install crictl
+RUN wget https://github.com/kubernetes-sigs/cri-tools/releases/download/v1.31.1/crictl-v1.31.1-linux-amd64.tar.gz \
+    && tar -zxvf crictl-v1.31.1-linux-amd64.tar.gz \
+    && mv crictl /usr/local/bin/ \
+    && rm crictl-v1.31.1-linux-amd64.tar.gz
+
+# Install Miniconda
+RUN wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O /tmp/miniconda.sh \
+    && bash /tmp/miniconda.sh -b -p /root/miniconda3 \
+    && rm /tmp/miniconda.sh
+
+# Setup crictl configuration
+RUN echo 'runtime-endpoint: unix:///run/containerd/containerd.sock\n\
+image-endpoint: unix:///run/containerd/containerd.sock\n\
+timeout: 30\n\
+debug: false' > /etc/crictl.yaml
+
+# Copy entrypoint script
+COPY entrypoint.sh /entrypoint.sh
+RUN chmod +x /entrypoint.sh
+
+# Activate conda and install FedML
+RUN echo "Cache bust: $CACHEBUST" \
+    && conda init bash \
+    && . ~/.bashrc \
+    && conda create -y -n fedml python=3.10 \
+    && conda activate fedml \
+    && pip install --upgrade gevent multiprocess \
+    && pip install "git+https://github.com/FedML-AI/FedML.git@k8s/dev/v0.7.0#egg=fedml&subdirectory=python"
+
+ENTRYPOINT ["/entrypoint.sh"]

devops/k8s/fedml-k8s-agent/entrypoint.sh

@@ -0,0 +1,67 @@
+#!/bin/bash
+
+# Initialize conda
+eval "$(/root/miniconda3/bin/conda shell.bash hook)"
+conda init bash
+source ~/.bashrc
+
+# Activate conda environment
+conda activate fedml
+
+# Function to handle FedML login
+fedml_login() {
+    if [ "$FEDML_PROVIDER" = "true" ]; then
+        if [ -z "$INFERENCE_GATEWAY_PORT" ] && [ -z "$INFERENCE_PROXY_PORT" ] && [ -z "$FEDML_CONNECTION_TYPE" ]; then
+            fedml login -p "$FEDML_API_KEY" -v "$FEDML_ENV"
+        else
+            fedml login -p "$FEDML_API_KEY" -v "$FEDML_ENV" -mgp "$INFERENCE_GATEWAY_PORT" -wpp "$INFERENCE_PROXY_PORT" -wct "$FEDML_CONNECTION_TYPE"
+        fi
+    else
+        if [ -z "$INFERENCE_GATEWAY_PORT" ] && [ -z "$INFERENCE_PROXY_PORT" ] && [ -z "$FEDML_CONNECTION_TYPE" ]; then
+            fedml login "$FEDML_API_KEY" -v "$FEDML_ENV"
+        else
+            fedml login "$FEDML_API_KEY" -v "$FEDML_ENV" -mgp "$INFERENCE_GATEWAY_PORT" -wpp "$INFERENCE_PROXY_PORT" -wct "$FEDML_CONNECTION_TYPE"
+        fi
+    fi
+}
+
+# Print all environment variables
+echo "=== Environment Variables ==="
+printenv | sort
+echo "==========================="
+
+# Create devices.id file if FEDML_DEVICE_ID is provided
+if [ -n "$FEDML_DEVICE_ID" ]; then
+    mkdir -p "/root/.fedml/fedml-client/fedml/data/runner_infos"
+    echo -n "$FEDML_DEVICE_ID" > "/root/.fedml/fedml-client/fedml/data/runner_infos/devices.id"
+fi
+
+# Setup SSH key if provided
+sshDir="/root/.ssh"
+mkdir -p $sshDir
+if [ -n "$SSH_PUBLIC_KEY_SYSTEM" ]; then
+    if ! grep -qF "$SSH_PUBLIC_KEY_SYSTEM" $sshDir/authorized_keys 2>/dev/null; then
+        echo "$SSH_PUBLIC_KEY_SYSTEM" >> $sshDir/authorized_keys
+        chmod 600 $sshDir/authorized_keys
+    fi
+fi
+if [ -n "$SSH_PUBLIC_KEY_USER" ]; then
+    if ! grep -qF "$SSH_PUBLIC_KEY_USER" $sshDir/authorized_keys 2>/dev/null; then
+        echo "$SSH_PUBLIC_KEY_USER" >> $sshDir/authorized_keys
+        chmod 600 $sshDir/authorized_keys
+    fi
+fi
+
+# Start SSH server
+systemctl start ssh || service ssh start || /usr/sbin/sshd
+
+# Start Redis server
+redis-server --daemonize yes
+
+# Login to FedML
+fedml_login
+
+
+
+# Keep container running
+tail -f /dev/null