Add base container image Jax TCPX

sample_workloads/jax-tcpx-base-image/Dockerfile

@@ -0,0 +1,42 @@
+FROM python:3.10-slim
+
+ENV VERSION=0.4.25
+
+RUN pip install --no-cache-dir --upgrade pip
+
+RUN pip install --no-cache-dir --upgrade "jax[cuda12_pip]==${VERSION}" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+
+WORKDIR /workspace/
+
+ADD entrypoint.sh .
+
+ENTRYPOINT ["bash", "entrypoint.sh"]
+
+# Environment variables required by TCPX
+ENV NCCL_NVLS_ENABLE=0
+ENV NCCL_CROSS_NIC=0
+ENV NCCL_ALGO=Ring
+ENV NCCL_PROTO=Simple
+ENV NCCL_DEBUG=INFO
+ENV NCCL_NET_GDR_LEVEL=PIX
+ENV NCCL_P2P_PXN_LEVEL=0
+ENV NCCL_DEBUG_SUBSYS=INIT,GRAPH,ENV,TUNING,NET,VERSION
+ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:/usr/local/tcpx/lib64"
+ENV NCCL_GPUDIRECTTCPX_FORCE_ACK=0
+ENV NCCL_GPUDIRECTTCPX_TX_COMPLETION_NANOSLEEP=1000
+ENV NCCL_DYNAMIC_CHUNK_SIZE=524288
+ENV NCCL_P2P_NET_CHUNKSIZE=524288
+ENV NCCL_P2P_PCI_CHUNKSIZE=524288
+ENV NCCL_P2P_NVL_CHUNKSIZE=1048576
+ENV NCCL_NSOCKS_PERTHREAD=4
+ENV NCCL_SOCKET_NTHREADS=1
+ENV NCCL_MAX_NCHANNELS=12
+ENV NCCL_MIN_NCHANNELS=12
+ENV NCCL_GPUDIRECTTCPX_PROGRAM_FLOW_STEERING_WAIT_MICROS=1000000
+ENV NCCL_SOCKET_IFNAME=eth0
+ENV NCCL_GPUDIRECTTCPX_TX_BINDINGS="eth1:8-21,112-125;eth2:8-21,112-125;eth3:60-73,164-177;eth4:60-73,164-177"
+ENV NCCL_GPUDIRECTTCPX_RX_BINDINGS="eth1:22-35,124-139;eth2:22-35,124-139;eth3:74-87,178-191;eth4:74-87,178-191"
+ENV NCCL_GPUDIRECTTCPX_SOCKET_IFNAME=eth1,eth2,eth3,eth4
+ENV NCCL_GPUDIRECTTCPX_CTRL_DEV=eth0
+# Might require adjusting based on where the TCPX socket is mounted
+ENV NCCL_GPUDIRECTTCPX_UNIX_CLIENT_PREFIX=/run/tcpx

sample_workloads/jax-tcpx-base-image/README.md

@@ -0,0 +1,3 @@
+# Base image with JAX and TCPX config
+
+A base image to use with Jax

sample_workloads/jax-tcpx-base-image/entrypoint.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+function on_script_completion {
+  # semaphore to cleanly exit hardware utilization monitor
+  touch /run/tcpx/workload_terminated
+}
+trap on_script_completion EXIT
+
+exec "$@"