# 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 importlib.metadata
import importlib.util
import os
import platform
from rich.table import Table
import oumi.cli.cli_utils as cli_utils
def _get_package_version(package_name: str, version_fallback: str) -> str:
"""Gets the version of the specified package.
Args:
package_name: The name of the package.
version_fallback: The fallback version string.
Returns:
str: The version of the package, or a fallback string if the package is not
installed.
"""
try:
return importlib.metadata.version(package_name)
except importlib.metadata.PackageNotFoundError:
return version_fallback
[docs]
def env():
"""Prints information about the current environment."""
# Delayed imports
from oumi.utils.torch_utils import format_cudnn_version
# End imports
version_fallback = "<not installed>"
env_var_fallback = "<not set>"
# All relevant environment vars.
env_vars = sorted(
[
"ACCELERATE_DYNAMO_BACKEND",
"ACCELERATE_DYNAMO_MODE",
"ACCELERATE_DYNAMO_USE_FULLGRAPH",
"ACCELERATE_DYNAMO_USE_DYNAMIC",
"ACCELERATE_USE_FSDP",
"CUDA_VISIBLE_DEVICES",
"LOCAL_RANK",
"LOCAL_WORLD_SIZE",
"OUMI_EXTRA_DEPS_FILE",
"OUMI_SLURM_CONNECTIONS",
"OUMI_FORCE_EDITABLE_INSTALL",
"OUMI_USE_SPOT_VM",
"RANK",
"WORLD_SIZE",
]
)
# All deps, excluding dev, docs, and gcp.
core_packages = sorted(
[
"accelerate",
"aiohttp",
"bitsandbytes",
"datasets",
"diffusers",
"einops",
"jsonlines",
"llama-cpp-python",
"liger-kernel",
"lm-eval",
"numpy",
"nvidia-ml-py",
"omegaconf",
"open_clip_torch",
"pandas",
"peft",
"pexpect",
"pillow",
"pydantic",
"responses",
"sglang",
"skypilot",
"tensorboard",
"timm",
"torch",
"torchdata",
"torchvision",
"tqdm",
"transformers",
"trl",
"typer",
"vllm",
"wandb",
"mlflow",
]
)
package_versions = {
package: _get_package_version(package, version_fallback)
for package in core_packages
}
env_values = {env_var: os.getenv(env_var, env_var_fallback) for env_var in env_vars}
cli_utils.section_header("Oumi environment information:")
env_table = Table(show_header=False, show_lines=False)
env_table.add_row("Oumi version", _get_package_version("oumi", version_fallback))
env_table.add_row("Python version", platform.python_version())
env_table.add_row("Platform", platform.platform())
cli_utils.CONSOLE.print(env_table)
cli_utils.section_header("Installed dependencies:")
deps_table = Table(show_header=True, show_lines=False)
deps_table.add_column("PACKAGE", justify="left")
deps_table.add_column("VERSION", justify="left")
for package, version in package_versions.items():
deps_table.add_row(package, version)
cli_utils.CONSOLE.print(deps_table)
if env_vars:
cli_utils.section_header("Environment variables:")
env_var_table = Table(show_header=True, show_lines=False)
env_var_table.add_column("VARIABLE", justify="left")
env_var_table.add_column("VALUE", justify="left")
for var in env_vars:
env_var_table.add_row(var, env_values[var])
cli_utils.CONSOLE.print(env_var_table)
if importlib.util.find_spec("torch") is not None:
torch = importlib.import_module("torch")
cli_utils.section_header("PyTorch information:")
cuda_table = Table(show_header=False, show_lines=False)
cuda_table.add_row("CUDA available", str(torch.cuda.is_available()))
if torch.cuda.is_available():
cuda_table.add_row("CUDA version", str(torch.version.cuda))
cuda_table.add_row(
"cuDNN version", format_cudnn_version(torch.backends.cudnn.version())
)
cuda_table.add_row("Number of GPUs", str(torch.cuda.device_count()))
cuda_table.add_row("GPU type", torch.cuda.get_device_name())
total_memory_gb = float(torch.cuda.mem_get_info()[1]) / float(
1024 * 1024 * 1024
)
cuda_table.add_row("GPU memory", f"{total_memory_gb:.1f}GB")
cli_utils.CONSOLE.print(cuda_table)
