# Copyright 2025 - Oumi
#
# 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.
import copy
import enum
import os
import sys
import time
from subprocess import Popen
from sys import stderr, stdout
from typing import Any, Final, NamedTuple, Optional
import typer
import oumi.cli.cli_utils as cli_utils
from oumi.utils.logging import logger
# Port range [1024, 65535] is generally available
# for application use w/o root permissions (non-privileged)
_MASTER_PORT_MIN_VALID_VALUE: Final[int] = 1024
_MASTER_PORT_MAX_VALID_VALUE: Final[int] = 65535
_SKY_ENV_VARS = {
"SKYPILOT_NODE_RANK",
"SKYPILOT_NODE_IPS",
"SKYPILOT_NUM_GPUS_PER_NODE",
}
_POLARIS_ENV_VARS = {
"PBS_NODEFILE",
"PBS_JOBID",
}
_MASTER_ADDR_ENV = "MASTER_ADDRESS"
_MASTER_PORT_ENV = "MASTER_PORT"
_DEFAULT_MASTER_ADDR = "127.0.0.1"
_DEFAULT_MASTER_PORT = 8007
class _RunBackend(str, enum.Enum):
SKYPILOT = "SkyPilot"
POLARIS = "Polaris"
LOCAL_MACHINE = "LocalMachine"
class _WorldInfo(NamedTuple):
num_nodes: int
"""Total number of nodes (machines)."""
gpus_per_node: int
"""Number of GPU-s per node."""
class _ProcessRunInfo:
def __init__(
self,
node_rank: int,
world_info: _WorldInfo,
master_address: str,
master_port: int,
node_ips: list[str],
):
"""Initializes run info, and validates arguments."""
if not (world_info.num_nodes > 0 and world_info.gpus_per_node > 0):
raise ValueError(
f"Non-positive number of nodes or GPUs per node: {world_info}"
)
elif not (node_rank >= 0 and node_rank < world_info.num_nodes):
raise ValueError(
f"Node rank {node_rank} is out of range: [0, {world_info.num_nodes})."
)
elif len(master_address) == 0:
raise ValueError(f"Empty master address: {master_address}.")
elif not (
master_port >= _MASTER_PORT_MIN_VALID_VALUE
and master_port <= _MASTER_PORT_MAX_VALID_VALUE
):
raise ValueError(
f"Master port: {master_port} is outside of valid range: "
f"[{_MASTER_PORT_MIN_VALID_VALUE}, {_MASTER_PORT_MAX_VALID_VALUE}]."
)
self._world_info = world_info
self._node_rank = int(node_rank)
self._master_address = master_address
self._master_port = master_port
self._node_ips = node_ips
@property
def node_rank(self) -> int:
"""Node rank in the [0, num_nodes) range."""
return self._node_rank
@property
def num_nodes(self) -> int:
"""Total number of nodes (machines)."""
return self._world_info.num_nodes
@property
def gpus_per_node(self) -> int:
"""Number of GPU-s per node."""
return self._world_info.gpus_per_node
@property
def total_gpus(self) -> int:
"""Total number of nodes (machines)."""
return self._world_info.num_nodes * self._world_info.gpus_per_node
@property
def master_address(self) -> str:
"""Master address."""
return self._master_address
@property
def node_ips(self) -> list[str]:
"""List of node IPs."""
return self._node_ips
@property
def master_port(self) -> int:
"""Master port."""
return self._master_port
def __repr__(self) -> str:
"""Defines how this class is properly printed."""
fields_dict: dict[str, Any] = {
"node_rank": self.node_rank,
"num_nodes": self.num_nodes,
"gpus_per_node": self.gpus_per_node,
"total_gpus": self.total_gpus,
"master_address": self.master_address,
"master_port": self.master_port,
"node_ips": self.node_ips,
}
return repr(fields_dict)
#
# Comamnds
#
[docs]
def torchrun(
ctx: typer.Context,
level: cli_utils.LOG_LEVEL_TYPE = None,
) -> None:
"""Starts `torchrun` sub-process w/ automatically configured common params.
Args:
ctx: The Typer context object.
level: The logging level for the specified command.
"""
try:
run_info: _ProcessRunInfo = _detect_process_run_info(os.environ.copy())
except (ValueError, RuntimeError):
logger.exception("Failed to detect process run info!")
raise
try:
cmds: list[str] = [
"torchrun",
f"--nnodes={run_info.num_nodes}",
f"--node-rank={run_info.node_rank}",
f"--nproc-per-node={run_info.gpus_per_node}",
f"--master-addr={run_info.master_address}",
f"--master-port={run_info.master_port}",
]
cmds.extend(ctx.args)
_run_subprocess(cmds, rank=run_info.node_rank)
except Exception:
logger.exception(f"`torchrun` failed (Rank: {run_info.node_rank})!")
raise
[docs]
def accelerate(
ctx: typer.Context,
level: cli_utils.LOG_LEVEL_TYPE = None,
) -> None:
"""Starts `accelerate` sub-process w/ automatically configured common params.
Args:
ctx: The Typer context object.
level: The logging level for the specified command.
"""
try:
run_info: _ProcessRunInfo = _detect_process_run_info(os.environ.copy())
except (ValueError, RuntimeError):
logger.exception("Failed to detect process run info!")
raise
try:
accelerate_subcommand: Optional[str] = None
extra_args = copy.deepcopy(ctx.args)
if (
len(extra_args) > 0
and len(extra_args[0]) > 0
and not extra_args[0].startswith("-")
):
# Copy sub-commands like "launch" to insert them right after `accelerate`
# ("accelerate launch ...")
accelerate_subcommand = extra_args.pop(0)
cmds: list[str] = (
["accelerate"]
+ ([accelerate_subcommand] if accelerate_subcommand is not None else [])
+ [
f"--num_machines={run_info.num_nodes}",
f"--machine_rank={run_info.node_rank}",
f"--num_processes={run_info.total_gpus}",
f"--main_process_ip={run_info.master_address}",
f"--main_process_port={run_info.master_port}",
]
)
cmds.extend(extra_args)
_run_subprocess(cmds, rank=run_info.node_rank)
except Exception:
logger.exception(f"`accelerate` failed (Rank: {run_info.node_rank})!")
raise
#
# Helper functions
#
def _detect_process_run_info(env: dict[str, str]) -> _ProcessRunInfo:
"""Detects process run info.
Uses known environment variables to detect common runtime parameters.
Args:
env: All environment variables.
Returns:
Process run info.
Raises:
ValueError: If any of the required environment variables are missing or invalid.
RuntimeError: If the node list is empty, or there are issues with backend
detection.
"""
# Detect the process run info depending on the runtime environment.
# Each runtime environment is checked in the order of priority.
process_run_info = _detect_skypilot_process_run_info(env)
if process_run_info is None:
process_run_info = _detect_polaris_process_run_info(env)
if process_run_info is None:
process_run_info = _detect_local_machine_process_run_info(env)
if process_run_info is None:
raise RuntimeError("Failed to detect process run info!")
# Extra verification logic to make sure that the detected process run info is
# consistent with the environment variables.
# Will raise an exception if the detected process run info is not consistent.
_verify_process_run_info(process_run_info, env)
return process_run_info
def _run_subprocess(cmds: list[str], *, rank: int) -> None:
env_copy = os.environ.copy()
start_time = time.perf_counter()
logger.info(f"Running the command: {cmds}")
p = Popen(
cmds,
env=env_copy,
stdout=stdout,
stderr=stderr,
bufsize=1,
universal_newlines=True,
)
rc = p.wait()
duration_sec = time.perf_counter() - start_time
duration_str = f"Duration: {duration_sec:.1f} sec"
if rc != 0:
logger.error(
f"{cmds[0]} failed with exit code: {rc} ({duration_str}). Command: {cmds}"
)
sys.exit(rc)
logger.info(f"Successfully completed! (Rank: {rank}. {duration_str})")
def _verify_process_run_info(run_info: _ProcessRunInfo, env: dict[str, str]) -> None:
oumi_total_gpus: Optional[int] = _get_optional_int_env_var(
"OUMI_TOTAL_NUM_GPUS", env
)
oumi_num_nodes: Optional[int] = _get_optional_int_env_var("OUMI_NUM_NODES", env)
oumi_master_address: Optional[str] = env.get("OUMI_MASTER_ADDR", None)
if oumi_master_address is not None and len(oumi_master_address) == 0:
raise ValueError("Empty master address in 'OUMI_MASTER_ADDR'!")
assert len(run_info.node_ips) > 0, "Empty list of nodes!"
assert run_info.node_rank is not None
if oumi_num_nodes is not None and oumi_num_nodes != run_info.num_nodes:
raise ValueError(
"Inconsistent number of nodes: "
f"{run_info.num_nodes} vs {oumi_num_nodes} in 'OUMI_NUM_NODES'."
)
elif oumi_total_gpus is not None and (oumi_total_gpus != run_info.total_gpus):
raise ValueError(
"Inconsistent total number of GPUs: "
f"{run_info.total_gpus} vs {oumi_total_gpus} "
"in 'OUMI_TOTAL_NUM_GPUS'. "
f"Nodes: {run_info.num_nodes}. GPU-s per node: {run_info.gpus_per_node}."
)
elif oumi_master_address and oumi_master_address not in run_info.node_ips:
raise ValueError(
f"Master address '{oumi_master_address}' not found in the list of nodes."
)
#
# Parse environment variables
#
def _detect_polaris_process_run_info(env: dict[str, str]) -> Optional[_ProcessRunInfo]:
polaris_node_file = env.get("PBS_NODEFILE", None)
if polaris_node_file is None:
return None
logger.debug("Running in Polaris environment!")
for env_var_name in _POLARIS_ENV_VARS:
if env.get(env_var_name, None) is None:
raise ValueError(
f"Polaris environment variable '{env_var_name}' is not defined!"
)
if not polaris_node_file:
raise ValueError("Empty value in the 'PBS_NODEFILE' environment variable!")
with open(polaris_node_file) as f:
nodes_str = f.read()
node_ips = _parse_nodes_str(nodes_str)
if len(node_ips) == 0:
raise RuntimeError("Empty list of nodes in 'PBS_NODEFILE'!")
gpus_per_node = 4 # Per Polaris spec.
node_rank = _get_optional_int_env_var("PMI_RANK", env)
if node_rank is None:
node_rank = 0
return _ProcessRunInfo(
node_rank=node_rank,
world_info=_WorldInfo(num_nodes=len(node_ips), gpus_per_node=gpus_per_node),
master_address=node_ips[0],
master_port=_DEFAULT_MASTER_PORT,
node_ips=node_ips,
)
def _detect_skypilot_process_run_info(env: dict[str, str]) -> Optional[_ProcessRunInfo]:
node_rank: Optional[int] = _get_optional_int_env_var("SKYPILOT_NODE_RANK", env)
if node_rank is None:
return None
logger.debug("Running in SkyPilot environment!")
for env_var_name in _SKY_ENV_VARS:
if env.get(env_var_name, None) is None:
raise ValueError(
f"SkyPilot environment variable '{env_var_name}' is not defined!"
)
node_ips = _parse_nodes_str(env.get("SKYPILOT_NODE_IPS", ""))
if len(node_ips) == 0:
raise RuntimeError("Empty list of nodes in 'SKYPILOT_NODE_IPS'!")
gpus_per_node = _get_positive_int_env_var("SKYPILOT_NUM_GPUS_PER_NODE", env)
return _ProcessRunInfo(
node_rank=node_rank,
world_info=_WorldInfo(num_nodes=len(node_ips), gpus_per_node=gpus_per_node),
master_address=node_ips[0],
master_port=_DEFAULT_MASTER_PORT,
node_ips=node_ips,
)
def _detect_local_machine_process_run_info(env: dict[str, str]) -> _ProcessRunInfo:
import torch # Importing torch takes time so only load it in this scenario.
# Attempt to produce a local configuration
if not torch.cuda.is_available():
raise RuntimeError(
"No supported distributed backends found and no GPUs on local machine!"
)
num_gpus_available = torch.cuda.device_count()
if num_gpus_available > 0:
oumi_num_nodes = 1
oumi_master_address = env.get(_MASTER_ADDR_ENV, _DEFAULT_MASTER_ADDR)
oumi_master_port = int(env.get(_MASTER_PORT_ENV, _DEFAULT_MASTER_PORT))
node_rank = 0
gpus_per_node = num_gpus_available
node_ips = [oumi_master_address]
cli_utils.configure_common_env_vars()
else:
raise RuntimeError("CUDA available but no GPUs found on local machine!")
return _ProcessRunInfo(
node_rank=node_rank,
world_info=_WorldInfo(num_nodes=oumi_num_nodes, gpus_per_node=gpus_per_node),
master_address=oumi_master_address,
master_port=oumi_master_port,
node_ips=node_ips,
)
#
# Private helper functions to parse environment variables
#
def _get_optional_int_env_var(var_name: str, env: dict[str, str]) -> Optional[int]:
str_value = env.get(var_name, None)
if str_value is None:
return None
try:
int_value = int(str_value)
except ValueError as e:
raise ValueError(f"Environment variable '{var_name}' is not an integer!") from e
return int_value
def _get_int_env_var(var_name: str, env: dict[str, str]) -> int:
int_value = _get_optional_int_env_var(var_name, env)
if int_value is None:
raise ValueError(f"Environment variable '{var_name}' is not defined!")
return int_value
def _get_positive_int_env_var(var_name: str, env: dict[str, str]) -> int:
int_value = _get_int_env_var(var_name, env)
if not (int_value > 0):
raise ValueError(
f"Environment variable '{var_name}' is not positive: {int_value}!"
)
return int_value
def _parse_nodes_str(nodes_str: str) -> list[str]:
node_ips = [x.strip() for x in nodes_str.split("\n")]
node_ips = [x for x in node_ips if len(x) > 0]
return node_ips
