Add GPU KubeRay benchmark workflow

Author: Ubuntu

.dockerignore

@@ -0,0 +1,24 @@
+.git
+.github
+.pytest_cache
+.ruff_cache
+.mypy_cache
+__pycache__
+*.py[cod]
+*.log
+.coverage
+htmlcov
+build
+dist
+*.egg-info
+
+# Datasets and benchmark output are mounted or downloaded at runtime.
+data
+dataset
+benchmark/results
+results
+
+# Local walkthrough material is not needed by the runtime image.
+docs
+kuberay
+ray_cluster_configs

.pre-commit-config.yaml

@@ -6,6 +6,7 @@ repos:
     hooks:
       - id: check-merge-conflict
       - id: check-yaml
+        args: [--allow-multiple-documents]
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/PyCQA/isort

Dockerfile.gpu

@@ -0,0 +1,61 @@
+ARG RAY_VERSION=2.56.0
+FROM rayproject/ray:${RAY_VERSION}-py312-cu121
+
+USER root
+
+RUN apt-get update \
+    && apt-get install -y --no-install-recommends \
+        build-essential \
+        cmake \
+        git \
+        wget \
+    && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /app
+
+ARG TORCH_VERSION=2.5.1
+RUN python -m pip install --upgrade pip setuptools wheel \
+    && python -m pip install \
+        "torch==${TORCH_VERSION}" \
+        --index-url https://download.pytorch.org/whl/cu121 \
+    && python -m pip install \
+        "pyg-lib==0.4.0+pt25cu121" \
+        "torch-scatter==2.1.2+pt25cu121" \
+        "torch-sparse==0.6.18+pt25cu121" \
+        "torch-cluster==1.6.3+pt25cu121" \
+        "torch-spline-conv==1.2.2+pt25cu121" \
+        --find-links https://data.pyg.org/whl/torch-2.5.1+cu121.html \
+    && python -m pip install \
+        torch-geometric==2.8.0 \
+        numpy==2.4.4 \
+        scipy==1.18.0 \
+        pandas==3.0.3 \
+        scikit-learn==1.9.0 \
+        omegaconf==2.3.0 \
+        attridict==0.0.9 \
+        torchmetrics==1.9.0 \
+        tensorboard==2.21.0 \
+        dtaidistance==2.4.0 \
+        gdown==6.1.0 \
+        tenseal==0.3.16 \
+        ogb==1.3.6 \
+        openfhe==1.2.3.0.24.4 \
+        psutil==7.2.2 \
+        matplotlib==3.11.0
+
+COPY --chown=1000:100 setup.py README.md ./
+COPY --chown=1000:100 fedgraph ./fedgraph
+COPY --chown=1000:100 benchmark ./benchmark
+
+# The source remains in /app so FedGraph subpackages omitted by the legacy
+# setup.py package list are still importable.
+RUN mkdir -p /app/data \
+    && chown -R 1000:100 /app \
+    && python -m pip install --no-deps --editable . \
+    && python -c "import ray, torch, torch_geometric, fedgraph; print(ray.__version__, torch.__version__, torch_geometric.__version__)"
+
+ENV PYTHONPATH=/app \
+    PYTHONUNBUFFERED=1
+
+USER 1000
+CMD ["python", "-c", "import ray, torch; print('Ray', ray.__version__, 'CUDA build', torch.version.cuda)"]

Dockerfile.gpu.draft

@@ -0,0 +1,57 @@
+# Draft GPU image for FedGraph on KubeRay.
+#
+# This file is intentionally not used by setup_cluster.sh yet. Review the base
+# image tag and CUDA/PyTorch/PyG versions before building or pushing.
+#
+# The existing Dockerfile installs CPU PyTorch. For GPU Ray worker pods, the
+# image must include CUDA-compatible PyTorch and matching PyG extension wheels.
+
+ARG RAY_VERSION=2.56.0
+FROM rayproject/ray-ml:${RAY_VERSION}-gpu
+
+WORKDIR /app
+
+USER root
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    build-essential \
+    cmake \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY docker_requirements.txt .
+
+# Align these with the CUDA runtime in the selected rayproject/ray-ml image.
+# The local EC2 validation environment used torch 2.5.1+cu121 successfully.
+RUN python -m pip install --no-cache-dir --upgrade pip && \
+    python -m pip install --no-cache-dir \
+      torch==2.5.1+cu121 \
+      --index-url https://download.pytorch.org/whl/cu121 && \
+    python -m pip install --no-cache-dir \
+      pyg-lib \
+      torch-scatter \
+      torch-sparse \
+      torch-cluster \
+      torch-spline-conv \
+      -f https://data.pyg.org/whl/torch-2.5.1+cu121.html && \
+    python -m pip install --no-cache-dir torch-geometric
+
+RUN grep -v "tenseal" docker_requirements.txt | \
+    grep -v "torch" | \
+    grep -v "torch-cluster" | \
+    grep -v "torch-scatter" | \
+    grep -v "torch-sparse" | \
+    grep -v "torch-spline-conv" | \
+    grep -v "torch-geometric" > requirements_filtered.txt && \
+    python -m pip install --no-cache-dir -r requirements_filtered.txt
+
+RUN python -m pip install --no-cache-dir openfhe==1.2.3.0.24.4 || \
+    echo "Warning: OpenFHE installation may need manual verification"
+
+COPY fedgraph /app/fedgraph
+COPY benchmark /app/benchmark
+COPY setup.py README.md ./
+
+RUN python -m pip install --no-deps -e .
+
+CMD ["python", "-c", "import torch, ray, fedgraph; print('cuda_available=', torch.cuda.is_available()); print('ray=', ray.__version__)"]

benchmark/benchmark_NC_batch_size_convergence.py

@@ -82,6 +82,7 @@ class ExperimentConfig:
     num_layers: int
     num_hops: int
     gpu: bool
+    server_device: Optional[str]
     num_cpus_per_trainer: int
     num_gpus_per_trainer: float
     use_ogb_load_patch: bool
@@ -339,6 +340,7 @@ def to_fedgraph_args(
             "num_layers": config.num_layers,
             "num_hops": config.num_hops,
             "gpu": config.gpu,
+            "server_device": config.server_device,
             "num_cpus_per_trainer": config.num_cpus_per_trainer,
             "num_gpus_per_trainer": config.num_gpus_per_trainer,
             "logdir": str(logdir),
@@ -981,6 +983,7 @@ def build_configs(args) -> List[ExperimentConfig]:
                         num_layers=args.num_layers,
                         num_hops=args.num_hops,
                         gpu=args.gpu,
+                        server_device=args.server_device,
                         num_cpus_per_trainer=args.num_cpus_per_trainer,
                         num_gpus_per_trainer=args.num_gpus_per_trainer,
                         use_ogb_load_patch=use_ogb_load_patch,
@@ -1021,6 +1024,15 @@ def parse_args():
     parser.add_argument("--num-hops", type=int, default=0)
     parser.add_argument("--method", default=None)
     parser.add_argument("--gpu", action="store_true")
+    parser.add_argument(
+        "--server-device",
+        choices=("cpu", "cuda"),
+        default=None,
+        help=(
+            "Temporary NC server-device override. By default, the server uses "
+            "the same device selected by --gpu."
+        ),
+    )
     parser.add_argument("--num-cpus-per-trainer", type=int, default=1)
     parser.add_argument("--num-gpus-per-trainer", type=float, default=0.0)
     parser.add_argument(

deploy/kuberay/cluster-autoscaler-values.yaml

@@ -0,0 +1,30 @@
+cloudProvider: aws
+
+rbac:
+  serviceAccount:
+    create: false
+    name: cluster-autoscaler
+
+nodeSelector:
+  workload: system
+
+extraArgs:
+  expander: least-waste
+  scale-down-delay-after-add: 5m
+  scale-down-unneeded-time: 5m
+  scan-interval: 20s
+
+resources:
+  requests:
+    cpu: 100m
+    memory: 300Mi
+  limits:
+    cpu: 300m
+    memory: 600Mi
+
+serviceMonitor:
+  enabled: true
+  interval: 5s
+  namespace: prometheus-system
+  selector:
+    release: prometheus

deploy/kuberay/dcgm-values.yaml

@@ -0,0 +1,35 @@
+arguments:
+  - -f
+  - /etc/dcgm-exporter/default-counters.csv
+  - -c
+  - "5000"
+
+nodeSelector:
+  ray-node-type: gpu-worker
+
+tolerations:
+  - key: nvidia.com/gpu
+    operator: Equal
+    value: "true"
+    effect: NoSchedule
+
+serviceMonitor:
+  enabled: true
+  interval: 5s
+  scrapeTimeout: 4s
+  additionalLabels:
+    release: prometheus
+
+kubernetes:
+  enablePodLabels: true
+  enablePodUID: true
+  podLabelAllowlistRegex:
+    - ^ray\.io/(cluster|group|node-type)$
+
+resources:
+  requests:
+    cpu: 100m
+    memory: 128Mi
+  limits:
+    cpu: 300m
+    memory: 512Mi

deploy/kuberay/eks-cluster.yaml

@@ -0,0 +1,91 @@
+apiVersion: eksctl.io/v1alpha5
+kind: ClusterConfig
+
+metadata:
+  name: fedgraph-gpu
+  region: eu-central-1
+  version: "1.34"
+
+iam:
+  withOIDC: true
+  serviceAccounts:
+    - metadata:
+        name: cluster-autoscaler
+        namespace: kube-system
+      attachPolicy:
+        Version: "2012-10-17"
+        Statement:
+          - Effect: Allow
+            Action:
+              - autoscaling:SetDesiredCapacity
+              - autoscaling:TerminateInstanceInAutoScalingGroup
+            Resource: "*"
+            Condition:
+              StringEquals:
+                aws:ResourceTag/k8s.io/cluster-autoscaler/enabled: "true"
+                aws:ResourceTag/k8s.io/cluster-autoscaler/fedgraph-gpu: owned
+          - Effect: Allow
+            Action:
+              - autoscaling:DescribeAutoScalingGroups
+              - autoscaling:DescribeAutoScalingInstances
+              - autoscaling:DescribeLaunchConfigurations
+              - autoscaling:DescribeScalingActivities
+              - autoscaling:DescribeTags
+              - ec2:DescribeImages
+              - ec2:DescribeInstanceTypes
+              - ec2:DescribeLaunchTemplateVersions
+              - ec2:GetInstanceTypesFromInstanceRequirements
+              - eks:DescribeNodegroup
+            Resource: "*"
+
+addons:
+  - name: aws-ebs-csi-driver
+    wellKnownPolicies:
+      ebsCSIController: true
+
+managedNodeGroups:
+  # System services, monitoring, and the CPU-only Ray head share this node.
+  - name: system-head-cpu
+    instanceType: m7i.xlarge
+    amiFamily: AmazonLinux2023
+    minSize: 1
+    desiredCapacity: 1
+    maxSize: 1
+    volumeSize: 100
+    volumeType: gp3
+    volumeEncrypted: true
+    labels:
+      workload: system
+      ray-node-type: head
+    tags:
+      workload: system-head
+      k8s.io/cluster-autoscaler/enabled: "true"
+      k8s.io/cluster-autoscaler/fedgraph-gpu: owned
+    propagateASGTags: true
+
+  # GPU workers begin at zero. Ray creates pending worker pods on demand, then
+  # Cluster Autoscaler adds up to five EC2 nodes. Each node has one T4 GPU, so
+  # one one-GPU Ray worker pod and trainer actor can run on each node.
+  - name: ray-worker-gpu
+    instanceType: g4dn.xlarge
+    amiFamily: AmazonLinux2023
+    minSize: 0
+    desiredCapacity: 0
+    maxSize: 5
+    volumeSize: 100
+    volumeType: gp3
+    volumeEncrypted: true
+    labels:
+      workload: ray
+      ray-node-type: gpu-worker
+      accelerator: nvidia-tesla-t4
+      k8s.amazonaws.com/accelerator: nvidia-tesla-t4
+    taints:
+      - key: nvidia.com/gpu
+        value: "true"
+        effect: NoSchedule
+    tags:
+      workload: ray-gpu-worker
+      k8s.io/cluster-autoscaler/enabled: "true"
+      k8s.io/cluster-autoscaler/fedgraph-gpu: owned
+    propagateASGTags: true