| ############################################################################ |
| # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| ############################################################################ |
| import pytest |
| |
| import atf |
| |
| sbcast_src = "sbcast_src" |
| sbcast_tgt = "sbcast_tgt" |
| job_marker = "job_marker" |
| |
| |
| def setup_module(): |
| atf.require_nodes(1, [("CPUs", 1)]) |
| atf.require_slurm_running() |
| |
| |
| @pytest.fixture(scope="function") |
| def sbcast_job(request): |
| """Submit a job that waits for the sbcast target file to appear, then writes |
| a separate marker file. Returns the job identifier based on request.param |
| (JobId or SLUID) for use with sbcast -j.""" |
| |
| if request.param == "SLUID": |
| atf.require_version( |
| (26, 5), |
| "bin/sbcast", |
| reason="Ticket 22180: SLUID availability added in 26.05+", |
| ) |
| |
| script = "sbcast_wait.sh" |
| atf.make_bash_script( |
| script, |
| f"""for i in $(seq 1 30); do |
| if [ -f {sbcast_tgt} ]; then |
| echo "sbcast_tgt found" > {job_marker} |
| exit 0 |
| fi |
| sleep 1 |
| done |
| exit 1""", |
| ) |
| |
| # Create a source file to broadcast |
| with open(sbcast_src, "w") as f: |
| f.write("sbcast_ok\n") |
| |
| job_id = atf.submit_job_sbatch(f"-N1 -n1 {script}", fatal=True) |
| atf.wait_for_job_state(job_id, "RUNNING", fatal=True) |
| |
| identifier = str(atf.get_job_parameter(job_id, request.param)) |
| |
| return identifier, job_id |
| |
| |
| @pytest.mark.parametrize("sbcast_job", ["JobId", "SLUID"], indirect=True) |
| def test_sbcast(sbcast_job): |
| """Verify sbcast -j <identifier> broadcasts a file to the job's nodes.""" |
| |
| identifier, job_id = sbcast_job |
| |
| atf.run_command(f"sbcast -j {identifier} -f {sbcast_src} {sbcast_tgt}", fatal=True) |
| |
| # Job should find the broadcast file and exit successfully |
| atf.wait_for_job_state(job_id, "DONE", fatal=True) |
| |
| # Verify the job saw the broadcast file and wrote the marker |
| output = atf.run_command_output(f"cat {job_marker}", fatal=True) |
| assert ( |
| "sbcast_tgt found" in output |
| ), f"Expected 'sbcast_tgt found' in marker file, got: {output}" |
| atf.run_command(f"rm -f {job_marker} {sbcast_tgt}") |