Repository: Lightning-AI/pytorch-lightning
File: src/lightning/fabric/cli.py (Python)

Code before:
        args = Namespace(checkpoint_folder=checkpoint_folder, output_file=output_file)
        config = _process_cli_args(args)
        checkpoint = _load_distributed_checkpoint(config.checkpoint_folder)
        torch.save(checkpoint, config.output_file)


def _set_env_variables(args: Namespace) -> None:
    """Set the environment variables for the new processes.

    The Fabric connector will parse the arguments set here.

    """
    os.environ["LT_CLI_USED"] = "1"
    if args.accelerator is not None:
        os.environ["LT_ACCELERATOR"] = str(args.accelerator)
    if args.strategy is not None:
        os.environ["LT_STRATEGY"] = str(args.strategy)
    os.environ["LT_DEVICES"] = str(args.devices)
    os.environ["LT_NUM_NODES"] = str(args.num_nodes)
    if args.precision is not None:
        os.environ["LT_PRECISION"] = str(args.precision)


def _get_num_processes(accelerator: str, devices: str) -> int:
    """Parse the `devices` argument to determine how many processes need to be launched on the current machine."""
    if accelerator == "auto" or accelerator is None:
        if torch.cuda.is_available() and torch.cuda.device_count() > 0:
            accelerator = "cuda"
        elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
            accelerator = "mps"
        else:
            accelerator = "cpu"

    if devices == "auto":
        if accelerator == "cuda" and torch.cuda.device_count() > 0 or accelerator == "mps" or accelerator == "cpu":
            devices = "1"
        else:
            raise ValueError(f"Cannot default to '1' device for accelerator='{accelerator}'")
    if accelerator == "gpu":
        parsed_devices = _parse_gpu_ids(devices, include_cuda=True, include_mps=True)
    elif accelerator == "cuda":
        parsed_devices = CUDAAccelerator.parse_devices(devices)
    elif accelerator == "mps":
        parsed_devices = MPSAccelerator.parse_devices(devices)
    elif accelerator == "tpu":
        raise ValueError("Launching processes for TPU through the CLI is not supported.")
    else:
        return CPUAccelerator.parse_devices(devices)
    return len(parsed_devices) if parsed_devices is not None else 0



Review comment (refactor):
lets not duplicate the logic and use already written functions like:
https://github.com/Lightning-AI/pytorch-lightning/blob/c85660c0da7777c667309a3ef937276353aacc27/src/lightning/pytorch/trainer/connectors/accelerator_connector.py#L333

Code after:
        args = Namespace(checkpoint_folder=checkpoint_folder, output_file=output_file)
        config = _process_cli_args(args)
        checkpoint = _load_distributed_checkpoint(config.checkpoint_folder)
        torch.save(checkpoint, config.output_file)


def _set_env_variables(args: Namespace) -> None:
    """Set the environment variables for the new processes.

    The Fabric connector will parse the arguments set here.

    """
    os.environ["LT_CLI_USED"] = "1"
    if args.accelerator is not None:
        os.environ["LT_ACCELERATOR"] = str(args.accelerator)
    if args.strategy is not None:
        os.environ["LT_STRATEGY"] = str(args.strategy)
    os.environ["LT_DEVICES"] = str(args.devices)
    os.environ["LT_NUM_NODES"] = str(args.num_nodes)
    if args.precision is not None:
        os.environ["LT_PRECISION"] = str(args.precision)


def _get_num_processes(accelerator: str, devices: str) -> int:
    """Parse the `devices` argument to determine how many processes need to be launched on the current machine."""

    if accelerator == "auto" or accelerator is None:
        accelerator = _select_auto_accelerator()
    if devices == "auto":
        if accelerator == "cuda" or accelerator == "mps" or accelerator == "cpu":
            devices = "1"
        else:
            raise ValueError(f"Cannot default to '1' device for accelerator='{accelerator}'")
    if accelerator == "gpu":
        parsed_devices = _parse_gpu_ids(devices, include_cuda=True, include_mps=True)
    elif accelerator == "cuda":
        parsed_devices = CUDAAccelerator.parse_devices(devices)
    elif accelerator == "mps":
        parsed_devices = MPSAccelerator.parse_devices(devices)
    elif accelerator == "tpu":
        raise ValueError("Launching processes for TPU through the CLI is not supported.")
    else:
        return CPUAccelerator.parse_devices(devices)
    return len(parsed_devices) if parsed_devices is not None else 0


def _torchrun_launch(args: Namespace, script_args: list[str]) -> None:
    """This will invoke `torchrun` programmatically to launch the given script in new processes."""
    import torch.distributed.run as torchrun

    num_processes = 1 if args.strategy == "dp" else _get_num_processes(args.accelerator, args.devices)