diff --git a/nvflow/lib/gpu_layout.py b/nvflow/lib/gpu_layout.py index 4afe6ea..2c3dd02 100644 --- a/nvflow/lib/gpu_layout.py +++ b/nvflow/lib/gpu_layout.py @@ -66,15 +66,24 @@ def resolve_gpu_layout( Resolved :class:`GpuLayout`. Raises: - ValueError: If ``total_gpus`` is not evenly divisible by - ``gpus_per_node``. + ValueError: If ``gpus_per_node`` is not a positive integer, if + ``total_gpus`` is not a positive integer evenly divisible by + ``gpus_per_node``, or if the legacy ``num_nodes``/``num_gpus`` + fields are not positive integers. """ gpus_per_node = _DEFAULT_GPUS_PER_NODE if cluster_config: gpus_per_node = cluster_config.get("gpus_per_node", _DEFAULT_GPUS_PER_NODE) + if gpus_per_node <= 0: + raise ValueError( + f"gpus_per_node ({gpus_per_node}) must be a positive integer. " + f"Check the gpus_per_node value in the cluster config." + ) if "total_gpus" in config: total = config["total_gpus"] + if total <= 0: + raise ValueError(f"total_gpus ({total}) must be a positive integer.") if total % gpus_per_node != 0: raise ValueError( f"total_gpus ({total}) is not evenly divisible by " @@ -86,7 +95,11 @@ def resolve_gpu_layout( if "num_nodes" in config: num_nodes = config["num_nodes"] + if num_nodes <= 0: + raise ValueError(f"num_nodes ({num_nodes}) must be a positive integer.") num_gpus = config.get("num_gpus", gpus_per_node) + if num_gpus <= 0: + raise ValueError(f"num_gpus ({num_gpus}) must be a positive integer.") LOG.debug( "Using legacy num_nodes=%d / num_gpus=%d (not portable across clusters)", num_nodes, diff --git a/tests/test_gpu_layout.py b/tests/test_gpu_layout.py new file mode 100644 index 0000000..c2bf189 --- /dev/null +++ b/tests/test_gpu_layout.py @@ -0,0 +1,87 @@ +# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Tests for ``nvflow.lib.gpu_layout``.""" + +from __future__ import annotations + +import pytest + +from nvflow.lib.gpu_layout import GpuLayout, resolve_gpu_layout + + +class TestGpuLayout: + def test_total_gpus_property(self) -> None: + layout = GpuLayout(num_nodes=2, gpus_per_node=8) + assert layout.total_gpus == 16 + + +class TestResolveGpuLayoutTotalGpus: + def test_divides_evenly_with_cluster_config(self) -> None: + layout = resolve_gpu_layout({"total_gpus": 16}, {"gpus_per_node": 4}) + assert layout == GpuLayout(num_nodes=4, gpus_per_node=4) + + def test_uses_default_gpus_per_node_without_cluster_config(self) -> None: + layout = resolve_gpu_layout({"total_gpus": 16}) + assert layout == GpuLayout(num_nodes=2, gpus_per_node=8) + + def test_not_evenly_divisible_raises(self) -> None: + with pytest.raises(ValueError, match="not evenly divisible"): + resolve_gpu_layout({"total_gpus": 10}, {"gpus_per_node": 8}) + + def test_zero_total_gpus_raises(self) -> None: + with pytest.raises(ValueError, match="must be a positive integer"): + resolve_gpu_layout({"total_gpus": 0}, {"gpus_per_node": 8}) + + def test_negative_total_gpus_raises(self) -> None: + with pytest.raises(ValueError, match="must be a positive integer"): + resolve_gpu_layout({"total_gpus": -8}, {"gpus_per_node": 8}) + + +class TestResolveGpuLayoutGpusPerNode: + def test_zero_gpus_per_node_raises_value_error_not_zero_division(self) -> None: + with pytest.raises(ValueError, match="gpus_per_node.*must be a positive integer"): + resolve_gpu_layout({"total_gpus": 16}, {"gpus_per_node": 0}) + + def test_negative_gpus_per_node_raises(self) -> None: + with pytest.raises(ValueError, match="gpus_per_node.*must be a positive integer"): + resolve_gpu_layout({"total_gpus": 16}, {"gpus_per_node": -4}) + + +class TestResolveGpuLayoutLegacyNumNodes: + def test_num_nodes_with_explicit_num_gpus(self) -> None: + layout = resolve_gpu_layout({"num_nodes": 3, "num_gpus": 4}) + assert layout == GpuLayout(num_nodes=3, gpus_per_node=4) + + def test_num_nodes_without_num_gpus_falls_back_to_gpus_per_node(self) -> None: + layout = resolve_gpu_layout({"num_nodes": 3}, {"gpus_per_node": 2}) + assert layout == GpuLayout(num_nodes=3, gpus_per_node=2) + + def test_zero_num_nodes_raises(self) -> None: + with pytest.raises(ValueError, match="num_nodes.*must be a positive integer"): + resolve_gpu_layout({"num_nodes": 0}) + + def test_zero_num_gpus_raises(self) -> None: + with pytest.raises(ValueError, match="num_gpus.*must be a positive integer"): + resolve_gpu_layout({"num_nodes": 2, "num_gpus": 0}) + + +class TestResolveGpuLayoutDefault: + def test_neither_field_defaults_to_single_node(self) -> None: + layout = resolve_gpu_layout({}) + assert layout == GpuLayout(num_nodes=1, gpus_per_node=8) + + def test_neither_field_uses_cluster_gpus_per_node(self) -> None: + layout = resolve_gpu_layout({}, {"gpus_per_node": 4}) + assert layout == GpuLayout(num_nodes=1, gpus_per_node=4)