maas-base/gpustack/utils/envs.py

import os
import platform
import subprocess
from pathlib import Path
from typing import Dict, List, Optional


def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
    """
    Extracts the environment variables from a source-able script on Unix.
    On Windows, uses PowerShell to source the script and capture env changes.
    """
    for script_path in script_paths:
        if not script_path.is_file():
            raise Exception(
                f"The file '{script_path}' does not exist or is not a file."
            )

    def parse_env(env_str):
        env = {}
        for line in env_str.splitlines():
            if '=' in line:
                key, value = line.split('=', 1)
                env[key] = value
        return env

    system = platform.system().lower()

    if system == "windows":
        return _extract_env_via_powershell(script_paths)
    else:
        return _extract_env_via_bash(script_paths)


def _extract_env_via_bash(script_paths: List[Path]) -> Dict[str, str]:
    def parse_env(env_str):
        env = {}
        for line in env_str.splitlines():
            if '=' in line:
                key, value = line.split('=', 1)
                env[key] = value
        return env

    try:
        original_env_output = subprocess.check_output(
            ['bash', '-c', 'env'],
            stderr=subprocess.PIPE,
            text=True,
        )
        original = parse_env(original_env_output)

        source_command = ' && '.join(
            [f'source {script_path}' for script_path in script_paths]
        )

        sourced_env_output = subprocess.check_output(
            ['bash', '-c', f'{source_command} && env'],
            stderr=subprocess.PIPE,
            text=True,
        )
        sourced = parse_env(sourced_env_output)

        diff = {
            k: v
            for k, v in sourced.items()
            if k not in original or original.get(k) != v
        }

        return diff
    except subprocess.CalledProcessError as e:
        raise Exception(
            f"Failed to extract environment variables from [{script_paths}]: {e.stderr}"
        )


def _extract_env_via_powershell(script_paths: List[Path]) -> Dict[str, str]:
    """
    Uses PowerShell to source a script (e.g. .ps1 or env-setup script)
    and capture environment variable changes.
    """
    try:
        # Build a PowerShell command that captures env before and after sourcing
        script_args = ' '.join(
            [f'"{str(p)}"' for p in script_paths]
        )
        ps_command = (
            '$before = Get-ChildItem Env: | ForEach-Object { "$($_.Name)=$($_.Value)" }; '
            f'{script_args}; '
            '$after = Get-ChildItem Env: | ForEach-Object { "$($_.Name)=$($_.Value)" }; '
            'Write-Output "---BEFORE---"; '
            'Write-Output $before; '
            'Write-Output "---AFTER---"; '
            'Write-Output $after'
        )

        output = subprocess.check_output(
            ['powershell', '-NoProfile', '-NonInteractive', '-Command', ps_command],
            stderr=subprocess.PIPE,
            text=True,
        )

        # Split output into before and after sections
        parts = output.split('---BEFORE---')
        if len(parts) < 2:
            return {}
        after_section = parts[1].split('---AFTER---')
        before_str = after_section[0]
        after_str = after_section[1] if len(after_section) > 1 else ''

        before = parse_env(before_str)
        after = parse_env(after_str)

        diff = {
            k: v
            for k, v in after.items()
            if k not in before or before.get(k) != v
        }

        return diff
    except subprocess.CalledProcessError as e:
        raise Exception(
            f"Failed to extract environment variables from [{script_paths}] via PowerShell: {e.stderr}"
        )


def get_gpustack_env(env_var: str) -> Optional[str]:
    env_name = "GPUSTACK_" + env_var
    return os.getenv(env_name)


def get_gpustack_env_bool(env_var: str) -> Optional[bool]:
    env_name = "GPUSTACK_" + env_var
    env_value = os.getenv(env_name)
    if env_value is not None:
        return env_value.lower() in ["true", "1"]
    return None


def is_docker_env() -> bool:
    return os.path.exists("/.dockerenv")


def sanitize_env(env: Dict[str, str]) -> Dict[str, str]:
    """
    Sanitize the environment variables by removing any keys that are not valid
    environment variable names.
    """
    prefixes = ("GPUSTACK_",)
    suffixes = (
        "_KEY",
        "_key",
        "_TOKEN",
        "_token",
        "_SECRET",
        "_secret",
        "_PASSWORD",
        "_password",
        "_PASS",
        "_pass",
    )
    return {
        k: v
        for k, v in env.items()
        if not k.startswith(prefixes) and not k.endswith(suffixes)
    }


def filter_env_vars(env_dict: dict) -> dict:
    """
    Filter out environment variables that should not be passed to the model instance or container.
    """
    return {
        k: v
        for k, v in env_dict.items()
        if not (
            k.startswith(
                (
                    "GPUSTACK_",
                    "PIP_",
                    "PIPX_",
                    "POETRY_",
                    "UV_",
                    "S6_",
                    "PGCONFIG_",
                    "POSTGRES_",
                )
            )
            or k.endswith(
                (
                    "_VISIBLE_DEVICES",
                    "_DISABLE_REQUIRE",
                    "_DRIVER_CAPABILITIES",
                    "_PATH",
                    "_HOME",
                )
            )
            or (
                k
                in (
                    "DEBIAN_FRONTEND",
                    "LANG",
                    "LANGUAGE",
                    "LC_ALL",
                    "PYTHON_VERSION",
                    "HOME",
                    "HOSTNAME",
                    "PWD",
                    "_",
                    "TERM",
                    "SHLVL",
                    "LS_COLORS",
                    "PATH",
                )
            )
        )
    }