#!/bin/bash
# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# EOS container wrapper for nvshmem4py tests.
#
# Handles container-specific setup (apt packages, MPI/CUDA paths),
# then delegates to run_tests.sh for venv creation, dependency
# installation, library overlays, and test execution.
#
# Invoked via gitlab_ci_launcher.sh:
#   sbatch ... python_test_wrapper.sh $CI_PROJECT_DIR $TEST_SUITE_NAME

#SBATCH --job-name=coreai_libraries_nvshmem-ci-jobs.nvshmem4py
#SBATCH --ntasks=8
#SBATCH --time=00:30:00
#SBATCH --account coreai_libraries_nvshmem
#SBATCH --partition=batch
#SBATCH --exclusive
#SBATCH --output=nvshmem4py_test.out
#SBATCH --error=nvshmem4py_test.out
#SBATCH --container-name=nvshmem_ci_container

set -x

########################################
# Parse arguments
########################################
if [ -z "$1" ]; then
    echo "Usage: sbatch python_test_wrapper.sh CI_PROJECT_HOME [TESTSUITE_NAME]"
    exit 1
fi
CI_PROJECT_HOME="$1"
TEST_SUITE="${2:-all}"

########################################
# Container-specific: install rdma-core
########################################
apt update && apt install -y librdmacm-dev rdma-core ibverbs-providers

########################################
# Set environment for run_tests.sh
########################################
export NVSHMEM_HOME="$CI_PROJECT_HOME/nvshmem_pkg"
export CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
export VENV_DIR="/workspace/nvshmem/nvshmem4py_venv"
export MPI_RUN="${MPI_RUN:-mpirun --oversubscribe --allow-run-as-root --mca pml ob1 --mca btl self,tcp,vader --mca btl_openib_warn_no_device_params_found 0}"
export RDMA_CORE_HOME="/usr"
export NP="${NP:-2}"
export NP_TEAM="${NP_TEAM:-8}"    # EOS nodes have 8 GPUs
if [ -z "${MPI_HOME:-}" ] || [ ! -x "$MPI_HOME/bin/mpirun" ]; then
    MPI_HOME="/usr/local/openmpi"
fi
export PATH="$MPI_HOME/bin:$CUDA_HOME/bin:$PATH"
export LD_LIBRARY_PATH="$MPI_HOME/lib:${LD_LIBRARY_PATH:-}"

########################################
# Delegate to run_tests.sh
########################################
exec bash "$CI_PROJECT_HOME/nvshmem4py/test/run_tests.sh" "$TEST_SUITE"
