crash instead of trying to set the available option

Author: Noé Pion

nerfstudio/configs/base_config.py

@@ -68,8 +68,8 @@ class MachineConfig(PrintableConfig):
     """current machine's rank (for DDP)"""
     dist_url: str = "auto"
     """distributed connection point (for DDP)"""
-    device_type: Literal["cpu", "cuda", "mps"] | None = None
-    """device type to use for training. If none set, script will do its best to set the value."""
+    device_type: Literal["cpu", "cuda", "mps"]
+    """device type to use for training"""
 
 
 @dataclass

nerfstudio/scripts/train.py

@@ -62,7 +62,7 @@
 from nerfstudio.configs.method_configs import AnnotatedBaseConfigUnion
 from nerfstudio.engine.trainer import TrainerConfig
 from nerfstudio.utils import comms, profiler
-from nerfstudio.utils.best_device import get_best_device
+from nerfstudio.utils.available_devices import get_available_devices
 from nerfstudio.utils.rich_utils import CONSOLE
 
 DEFAULT_TIMEOUT = timedelta(minutes=30)
@@ -227,6 +227,15 @@ def launch(
 def main(config: TrainerConfig) -> None:
     """Main function."""
 
+    # Check if the specified device type is available
+    available_device_types = get_available_devices()
+    if config.machine.device_type not in available_device_types:
+        raise RuntimeError(
+            f"Specified device type '{config.machine.device_type}' is not available. "
+            f"Available device types: {available_device_types}. "
+            "Please specify a valid device type using the CLI option: --machine.device_type [cuda|mps|cpu]"
+        )
+
     if config.data:
         CONSOLE.log("Using --data alias for --data.pipeline.datamanager.data")
         config.pipeline.datamanager.data = config.data
@@ -239,12 +248,6 @@ def main(config: TrainerConfig) -> None:
         CONSOLE.log(f"Loading pre-set config from: {config.load_config}")
         config = yaml.load(config.load_config.read_text(), Loader=yaml.Loader)
 
-    if not hasattr(config.machine, "device_type") or config.machine.device_type is None:
-        device_type, reason = get_best_device()
-        config.machine.device_type = device_type
-        CONSOLE.log(f"No device specified: {reason}")
-        CONSOLE.log("You can force a certain device type with --machine.device_type [cuda|mps|cpu]")
-
     config.set_timestamp()
 
     # print and save config

nerfstudio/utils/best_device.py nerfstudio/utils/available_devices.pynerfstudio/utils/best_device.py renamed to nerfstudio/utils/available_devices.py

@@ -12,21 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Literal, Tuple
+from typing import List, Literal
 
 import torch
 
 
-def get_best_device() -> Tuple[Literal["cpu", "cuda", "mps"], str]:
-    """Determine the best available device to run nerfstudio inference.
+def get_available_devices() -> List[Literal["cpu", "cuda", "mps"]]:
+    """Determine the available devices on the machine
 
     Returns:
-        tuple: (device_type, reason) where device_type is the selected device and
-        reason is an explanation of why it was chosen
+        list: List of available device types
     """
+    available_devices: List[Literal["cpu", "cuda", "mps"]] = []
     if torch.cuda.is_available():
-        return "cuda", "CUDA GPU available - using for optimal performance"
+        available_devices.append("cuda")
     elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
-        return "mps", "Apple Metal (MPS) available - using for accelerated performance"
-    else:
-        return "cpu", "No GPU/MPS detected - falling back to CPU"
+        available_devices.append("mps")
+    available_devices.append("cpu")
+    return available_devices