nshtrainer 0.9.1__py3-none-any.whl → 0.10.0__py3-none-any.whl

nshtrainer/model/_environment.py

@@ -0,0 +1,777 @@
+import getpass
+import inspect
+import logging
+import os
+import platform
+import socket
+import sys
+from datetime import timedelta
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, cast
+
+import git
+import nshconfig as C
+import psutil
+import torch
+from typing_extensions import Self
+
+from ..util.slurm import parse_slurm_node_list
+
+if TYPE_CHECKING:
+    from .base import LightningModuleBase
+    from .config import BaseConfig
+
+
+log = logging.getLogger(__name__)
+
+
+class EnvironmentClassInformationConfig(C.Config):
+    """Configuration for class information in the environment."""
+
+    name: str | None
+    """The name of the class."""
+
+    module: str | None
+    """The module where the class is defined."""
+
+    full_name: str | None
+    """The fully qualified name of the class."""
+
+    file_path: Path | None
+    """The file path where the class is defined."""
+
+    source_file_path: Path | None
+    """The source file path of the class, if available."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            name=None,
+            module=None,
+            full_name=None,
+            file_path=None,
+            source_file_path=None,
+        )
+
+    @classmethod
+    def from_class(cls, cls_: type):
+        name = cls_.__name__
+        module = cls_.__module__
+        full_name = f"{cls_.__module__}.{cls_.__qualname__}"
+
+        file_path = inspect.getfile(cls_)
+        source_file_path = inspect.getsourcefile(cls_)
+        return cls(
+            name=name,
+            module=module,
+            full_name=full_name,
+            file_path=Path(file_path),
+            source_file_path=Path(source_file_path) if source_file_path else None,
+        )
+
+    @classmethod
+    def from_instance(cls, instance: object):
+        return cls.from_class(type(instance))
+
+
+class EnvironmentSLURMInformationConfig(C.Config):
+    """Configuration for SLURM environment information."""
+
+    hostname: str | None
+    """The hostname of the current node."""
+
+    hostnames: list[str] | None
+    """List of hostnames for all nodes in the job."""
+
+    job_id: str | None
+    """The SLURM job ID."""
+
+    raw_job_id: str | None
+    """The raw SLURM job ID."""
+
+    array_job_id: str | None
+    """The SLURM array job ID, if applicable."""
+
+    array_task_id: str | None
+    """The SLURM array task ID, if applicable."""
+
+    num_tasks: int | None
+    """The number of tasks in the SLURM job."""
+
+    num_nodes: int | None
+    """The number of nodes in the SLURM job."""
+
+    node: str | int | None
+    """The node ID or name."""
+
+    global_rank: int | None
+    """The global rank of the current process."""
+
+    local_rank: int | None
+    """The local rank of the current process within its node."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            hostname=None,
+            hostnames=None,
+            job_id=None,
+            raw_job_id=None,
+            array_job_id=None,
+            array_task_id=None,
+            num_tasks=None,
+            num_nodes=None,
+            node=None,
+            global_rank=None,
+            local_rank=None,
+        )
+
+    @classmethod
+    def from_current_environment(cls):
+        try:
+            from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
+
+            if not SLURMEnvironment.detect():
+                return None
+
+            hostname = socket.gethostname()
+            hostnames = [hostname]
+            if node_list := os.environ.get("SLURM_JOB_NODELIST", ""):
+                hostnames = parse_slurm_node_list(node_list)
+
+            raw_job_id = os.environ["SLURM_JOB_ID"]
+            job_id = raw_job_id
+            array_job_id = os.environ.get("SLURM_ARRAY_JOB_ID")
+            array_task_id = os.environ.get("SLURM_ARRAY_TASK_ID")
+            if array_job_id and array_task_id:
+                job_id = f"{array_job_id}_{array_task_id}"
+
+            num_tasks = int(os.environ["SLURM_NTASKS"])
+            num_nodes = int(os.environ["SLURM_JOB_NUM_NODES"])
+
+            node_id = os.environ.get("SLURM_NODEID")
+
+            global_rank = int(os.environ["SLURM_PROCID"])
+            local_rank = int(os.environ["SLURM_LOCALID"])
+
+            return cls(
+                hostname=hostname,
+                hostnames=hostnames,
+                job_id=job_id,
+                raw_job_id=raw_job_id,
+                array_job_id=array_job_id,
+                array_task_id=array_task_id,
+                num_tasks=num_tasks,
+                num_nodes=num_nodes,
+                node=node_id,
+                global_rank=global_rank,
+                local_rank=local_rank,
+            )
+        except (ImportError, RuntimeError, ValueError, KeyError):
+            return None
+
+
+class EnvironmentLSFInformationConfig(C.Config):
+    """Configuration for LSF environment information."""
+
+    hostname: str | None
+    """The hostname of the current node."""
+
+    hostnames: list[str] | None
+    """List of hostnames for all nodes in the job."""
+
+    job_id: str | None
+    """The LSF job ID."""
+
+    array_job_id: str | None
+    """The LSF array job ID, if applicable."""
+
+    array_task_id: str | None
+    """The LSF array task ID, if applicable."""
+
+    num_tasks: int | None
+    """The number of tasks in the LSF job."""
+
+    num_nodes: int | None
+    """The number of nodes in the LSF job."""
+
+    node: str | int | None
+    """The node ID or name."""
+
+    global_rank: int | None
+    """The global rank of the current process."""
+
+    local_rank: int | None
+    """The local rank of the current process within its node."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            hostname=None,
+            hostnames=None,
+            job_id=None,
+            array_job_id=None,
+            array_task_id=None,
+            num_tasks=None,
+            num_nodes=None,
+            node=None,
+            global_rank=None,
+            local_rank=None,
+        )
+
+    @classmethod
+    def from_current_environment(cls):
+        try:
+            import os
+            import socket
+
+            hostname = socket.gethostname()
+            hostnames = [hostname]
+            if node_list := os.environ.get("LSB_HOSTS", ""):
+                hostnames = node_list.split()
+
+            job_id = os.environ["LSB_JOBID"]
+            array_job_id = os.environ.get("LSB_JOBINDEX")
+            array_task_id = os.environ.get("LSB_JOBINDEX")
+
+            num_tasks = int(os.environ.get("LSB_DJOB_NUMPROC", 1))
+            num_nodes = len(set(hostnames))
+
+            node_id = (
+                os.environ.get("LSB_HOSTS", "").split().index(hostname)
+                if "LSB_HOSTS" in os.environ
+                else None
+            )
+
+            # LSF doesn't have direct equivalents for global_rank and local_rank
+            # You might need to calculate these based on your specific setup
+            global_rank = int(os.environ.get("PMI_RANK", 0))
+            local_rank = int(os.environ.get("LSB_RANK", 0))
+
+            return cls(
+                hostname=hostname,
+                hostnames=hostnames,
+                job_id=job_id,
+                array_job_id=array_job_id,
+                array_task_id=array_task_id,
+                num_tasks=num_tasks,
+                num_nodes=num_nodes,
+                node=node_id,
+                global_rank=global_rank,
+                local_rank=local_rank,
+            )
+        except (ImportError, RuntimeError, ValueError, KeyError):
+            return None
+
+
+class EnvironmentLinuxEnvironmentConfig(C.Config):
+    """Configuration for Linux environment information."""
+
+    user: str | None
+    """The current user."""
+
+    hostname: str | None
+    """The hostname of the machine."""
+
+    system: str | None
+    """The operating system name."""
+
+    release: str | None
+    """The operating system release."""
+
+    version: str | None
+    """The operating system version."""
+
+    machine: str | None
+    """The machine type."""
+
+    processor: str | None
+    """The processor type."""
+
+    cpu_count: int | None
+    """The number of CPUs."""
+
+    memory: int | None
+    """The total system memory in bytes."""
+
+    uptime: timedelta | None
+    """The system uptime."""
+
+    boot_time: float | None
+    """The system boot time as a timestamp."""
+
+    load_avg: tuple[float, float, float] | None
+    """The system load average (1, 5, and 15 minutes)."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            user=None,
+            hostname=None,
+            system=None,
+            release=None,
+            version=None,
+            machine=None,
+            processor=None,
+            cpu_count=None,
+            memory=None,
+            uptime=None,
+            boot_time=None,
+            load_avg=None,
+        )
+
+    @classmethod
+    def from_current_environment(cls):
+        return cls(
+            user=getpass.getuser(),
+            hostname=platform.node(),
+            system=platform.system(),
+            release=platform.release(),
+            version=platform.version(),
+            machine=platform.machine(),
+            processor=platform.processor(),
+            cpu_count=os.cpu_count(),
+            memory=psutil.virtual_memory().total,
+            uptime=timedelta(seconds=psutil.boot_time()),
+            boot_time=psutil.boot_time(),
+            load_avg=os.getloadavg(),
+        )
+
+
+class EnvironmentSnapshotConfig(C.Config):
+    """Configuration for environment snapshot information."""
+
+    snapshot_dir: Path | None
+    """The directory where the snapshot is stored."""
+
+    modules: list[str] | None
+    """List of modules included in the snapshot."""
+
+    @classmethod
+    def empty(cls):
+        return cls(snapshot_dir=None, modules=None)
+
+    @classmethod
+    def from_current_environment(cls):
+        draft = cls.draft()
+        if snapshot_dir := os.environ.get("NSHRUNNER_SNAPSHOT_DIR"):
+            draft.snapshot_dir = Path(snapshot_dir)
+        if modules := os.environ.get("NSHRUNNER_SNAPSHOT_MODULES"):
+            draft.modules = modules.split(",")
+        return draft.finalize()
+
+
+class EnvironmentPackageConfig(C.Config):
+    """Configuration for Python package information."""
+
+    name: str | None
+    """The name of the package."""
+
+    version: str | None
+    """The version of the package."""
+
+    path: Path | None
+    """The installation path of the package."""
+
+    summary: str | None
+    """A brief summary of the package."""
+
+    author: str | None
+    """The author of the package."""
+
+    license: str | None
+    """The license of the package."""
+
+    requires: list[str] | None
+    """List of package dependencies."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            name=None,
+            version=None,
+            path=None,
+            summary=None,
+            author=None,
+            license=None,
+            requires=None,
+        )
+
+    @classmethod
+    def from_current_environment(cls):
+        # Add Python package information
+        python_packages: dict[str, Self] = {}
+        try:
+            import pkg_resources
+
+            for package in pkg_resources.working_set:
+                python_packages[package.key] = cls(
+                    name=package.project_name,
+                    version=package.version,
+                    path=Path(package.location) if package.location else None,
+                    summary=package.summary,
+                    author=package.author,
+                    license=package.license,
+                    requires=[str(req) for req in package.requires()],
+                )
+        except ImportError:
+            log.warning("pkg_resources not available, skipping package information")
+
+        return python_packages
+
+
+class EnvironmentGPUConfig(C.Config):
+    """Configuration for individual GPU information."""
+
+    name: str | None
+    """Name of the GPU."""
+
+    total_memory: int | None
+    """Total memory of the GPU in bytes."""
+
+    major: int | None
+    """Major version of CUDA capability."""
+
+    minor: int | None
+    """Minor version of CUDA capability."""
+
+    multi_processor_count: int | None
+    """Number of multiprocessors on the GPU."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            name=None,
+            total_memory=None,
+            major=None,
+            minor=None,
+            multi_processor_count=None,
+        )
+
+
+class EnvironmentCUDAConfig(C.Config):
+    """Configuration for CUDA environment information."""
+
+    is_available: bool | None
+    """Whether CUDA is available."""
+
+    version: str | None
+    """CUDA version."""
+
+    cudnn_version: int | None
+    """cuDNN version."""
+
+    @classmethod
+    def empty(cls):
+        return cls(is_available=None, version=None, cudnn_version=None)
+
+
+class EnvironmentHardwareConfig(C.Config):
+    """Configuration for hardware information."""
+
+    cpu_count_physical: int | None
+    """Number of physical CPU cores."""
+
+    cpu_count_logical: int | None
+    """Number of logical CPU cores."""
+
+    cpu_frequency_current: float | None
+    """Current CPU frequency in MHz."""
+
+    cpu_frequency_min: float | None
+    """Minimum CPU frequency in MHz."""
+
+    cpu_frequency_max: float | None
+    """Maximum CPU frequency in MHz."""
+
+    ram_total: int | None
+    """Total RAM in bytes."""
+
+    ram_available: int | None
+    """Available RAM in bytes."""
+
+    disk_total: int | None
+    """Total disk space in bytes."""
+
+    disk_used: int | None
+    """Used disk space in bytes."""
+
+    disk_free: int | None
+    """Free disk space in bytes."""
+
+    gpu_count: int | None
+    """Number of GPUs available."""
+
+    gpus: list[EnvironmentGPUConfig] | None
+    """List of GPU configurations."""
+
+    cuda: EnvironmentCUDAConfig | None
+    """CUDA environment configuration."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            cpu_count_physical=None,
+            cpu_count_logical=None,
+            cpu_frequency_current=None,
+            cpu_frequency_min=None,
+            cpu_frequency_max=None,
+            ram_total=None,
+            ram_available=None,
+            disk_total=None,
+            disk_used=None,
+            disk_free=None,
+            gpu_count=None,
+            gpus=None,
+            cuda=None,
+        )
+
+    @classmethod
+    def from_current_environment(cls):
+        draft = cls.draft()
+
+        # CPU information
+        draft.cpu_count_physical = psutil.cpu_count(logical=False)
+        draft.cpu_count_logical = psutil.cpu_count(logical=True)
+        cpu_freq = psutil.cpu_freq()
+        if cpu_freq:
+            draft.cpu_frequency_current = cpu_freq.current
+            draft.cpu_frequency_min = cpu_freq.min
+            draft.cpu_frequency_max = cpu_freq.max
+
+        # RAM information
+        ram = psutil.virtual_memory()
+        draft.ram_total = ram.total
+        draft.ram_available = ram.available
+
+        # Disk information
+        disk = psutil.disk_usage("/")
+        draft.disk_total = disk.total
+        draft.disk_used = disk.used
+        draft.disk_free = disk.free
+
+        # GPU and CUDA information
+        draft.cuda = EnvironmentCUDAConfig(
+            is_available=torch.cuda.is_available(),
+            version=cast(Any, torch).version.cuda,
+            cudnn_version=torch.backends.cudnn.version()
+            if torch.backends.cudnn.is_available()
+            else None,
+        )
+
+        if draft.cuda.is_available:
+            draft.gpu_count = torch.cuda.device_count()
+            draft.gpus = []
+            for i in range(draft.gpu_count):
+                gpu_props = torch.cuda.get_device_properties(i)
+                gpu_config = EnvironmentGPUConfig(
+                    name=gpu_props.name,
+                    total_memory=gpu_props.total_memory,
+                    major=gpu_props.major,
+                    minor=gpu_props.minor,
+                    multi_processor_count=gpu_props.multi_processor_count,
+                )
+                draft.gpus.append(gpu_config)
+
+        return draft.finalize()
+
+
+class GitRepositoryConfig(C.Config):
+    """Configuration for Git repository information."""
+
+    is_git_repo: bool | None
+    """Whether the current directory is a Git repository."""
+
+    branch: str | None
+    """The current Git branch."""
+
+    commit_hash: str | None
+    """The current commit hash."""
+
+    commit_message: str | None
+    """The current commit message."""
+
+    author: str | None
+    """The author of the current commit."""
+
+    commit_date: str | None
+    """The date of the current commit."""
+
+    remote_url: str | None
+    """The URL of the remote repository."""
+
+    is_dirty: bool | None
+    """Whether there are uncommitted changes."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            is_git_repo=None,
+            branch=None,
+            commit_hash=None,
+            commit_message=None,
+            author=None,
+            commit_date=None,
+            remote_url=None,
+            is_dirty=None,
+        )
+
+    @classmethod
+    def from_current_directory(cls):
+        draft = cls.draft()
+        try:
+            repo = git.Repo(os.getcwd(), search_parent_directories=True)
+            draft.is_git_repo = True
+            draft.branch = repo.active_branch.name
+            commit = repo.head.commit
+            draft.commit_hash = commit.hexsha
+
+            # Handle both str and bytes for commit message
+            if isinstance(commit.message, str):
+                draft.commit_message = commit.message.strip()
+            elif isinstance(commit.message, bytes):
+                draft.commit_message = commit.message.decode(
+                    "utf-8", errors="replace"
+                ).strip()
+            else:
+                draft.commit_message = str(commit.message).strip()
+
+            draft.author = f"{commit.author.name} <{commit.author.email}>"
+            draft.commit_date = commit.committed_datetime.isoformat()
+            if repo.remotes:
+                draft.remote_url = repo.remotes.origin.url
+            draft.is_dirty = repo.is_dirty()
+        except git.InvalidGitRepositoryError:
+            draft.is_git_repo = False
+        except Exception as e:
+            log.warning(f"Failed to get Git repository information: {e}")
+            draft.is_git_repo = None
+
+        return draft.finalize()
+
+
+class EnvironmentConfig(C.Config):
+    """Configuration for the overall environment."""
+
+    cwd: Path | None
+    """The current working directory."""
+
+    snapshot: EnvironmentSnapshotConfig | None
+    """The environment snapshot configuration."""
+
+    python_executable: Path | None
+    """The path to the Python executable."""
+
+    python_path: list[Path] | None
+    """The Python path."""
+
+    python_version: str | None
+    """The Python version."""
+
+    python_packages: dict[str, EnvironmentPackageConfig] | None
+    """A mapping of package names to their configurations."""
+
+    config: EnvironmentClassInformationConfig | None
+    """The configuration class information."""
+
+    model: EnvironmentClassInformationConfig | None
+    """The Lightning module class information."""
+
+    linux: EnvironmentLinuxEnvironmentConfig | None
+    """The Linux environment information."""
+
+    hardware: EnvironmentHardwareConfig | None
+    """Hardware configuration information."""
+
+    slurm: EnvironmentSLURMInformationConfig | None
+    """The SLURM environment information."""
+
+    lsf: EnvironmentLSFInformationConfig | None
+    """The LSF environment information."""
+
+    base_dir: Path | None
+    """The base directory for the run."""
+
+    log_dir: Path | None
+    """The directory for logs."""
+
+    checkpoint_dir: Path | None
+    """The directory for checkpoints."""
+
+    stdio_dir: Path | None
+    """The directory for standard input/output files."""
+
+    seed: int | None
+    """The global random seed."""
+
+    seed_workers: bool | None
+    """Whether to seed workers."""
+
+    git: GitRepositoryConfig | None
+    """Git repository information."""
+
+    @classmethod
+    def empty(cls):
+        return cls(
+            cwd=None,
+            snapshot=None,
+            python_executable=None,
+            python_path=None,
+            python_version=None,
+            python_packages=None,
+            config=None,
+            model=None,
+            linux=None,
+            hardware=None,
+            slurm=None,
+            lsf=None,
+            base_dir=None,
+            log_dir=None,
+            checkpoint_dir=None,
+            stdio_dir=None,
+            seed=None,
+            seed_workers=None,
+            git=None,
+        )
+
+    @classmethod
+    def from_current_environment(
+        cls,
+        root_config: "BaseConfig",
+        model: "LightningModuleBase",
+    ):
+        draft = cls.draft()
+        draft.cwd = Path(os.getcwd())
+        draft.python_executable = Path(sys.executable)
+        draft.python_path = [Path(path) for path in sys.path]
+        draft.python_version = sys.version
+        draft.python_packages = EnvironmentPackageConfig.from_current_environment()
+        draft.config = EnvironmentClassInformationConfig.from_instance(root_config)
+        draft.model = EnvironmentClassInformationConfig.from_instance(model)
+        draft.linux = EnvironmentLinuxEnvironmentConfig.from_current_environment()
+        draft.hardware = EnvironmentHardwareConfig.from_current_environment()
+        draft.slurm = EnvironmentSLURMInformationConfig.from_current_environment()
+        draft.lsf = EnvironmentLSFInformationConfig.from_current_environment()
+        draft.base_dir = root_config.directory.resolve_run_root_directory(
+            root_config.id
+        )
+        draft.log_dir = root_config.directory.resolve_subdirectory(
+            root_config.id, "log"
+        )
+        draft.checkpoint_dir = root_config.directory.resolve_subdirectory(
+            root_config.id, "checkpoint"
+        )
+        draft.stdio_dir = root_config.directory.resolve_subdirectory(
+            root_config.id, "stdio"
+        )
+        draft.seed = (
+            int(seed_str) if (seed_str := os.environ.get("PL_GLOBAL_SEED")) else None
+        )
+        draft.seed_workers = (
+            bool(int(seed_everything))
+            if (seed_everything := os.environ.get("PL_SEED_WORKERS"))
+            else None
+        )
+        draft.snapshot = EnvironmentSnapshotConfig.from_current_environment()
+        draft.git = GitRepositoryConfig.from_current_directory()
+        return draft.finalize()