matrice-compute 0.1.1__py3-none-any.whl

matrice_compute/instance_utils.py

@@ -0,0 +1,707 @@
+"""Module providing instance utilities functionality."""
+
+import os
+import socket
+import urllib.request
+import subprocess
+import logging
+import base64
+from datetime import datetime
+import psutil
+from cryptography.hazmat.primitives.ciphers import (
+    Cipher,
+    algorithms,
+    modes,
+)
+from cryptography.hazmat.backends import default_backend
+from matrice_common.utils import log_errors
+
+
+def get_instance_info(service_provider: str = None, instance_id: str = None) -> tuple:
+    """
+    Get instance provider and ID information.
+
+    Returns:
+        tuple: (service_provider, instance_id) strings
+    """
+    auto_service_provider = service_provider or os.environ.get("SERVICE_PROVIDER") or "LOCAL"
+    auto_instance_id = instance_id or os.environ.get("INSTANCE_ID") or ""
+    try:
+        gcp_check = subprocess.run(
+            "curl -s -m 1 -H 'Metadata-Flavor: Google' 'http://metadata.google.internal/computeMetadata/v1/instance/id'",
+            shell=True,
+            capture_output=True,
+            check=True,
+        )
+        if gcp_check.returncode == 0:
+            auto_service_provider = "GCP"
+            auto_instance_id = gcp_check.stdout.decode().strip()
+    except subprocess.CalledProcessError:
+        pass
+    try:
+        azure_check = subprocess.run(
+            "curl -s -m 1 -H Metadata:true 'http://169.254.169.254/metadata/instance?api-version=2020-09-01'",
+            shell=True,
+            capture_output=True,
+            check=True,
+        )
+        if azure_check.returncode == 0:
+            auto_service_provider = "AZURE"
+            azure_id = subprocess.run(
+                "curl -s -H Metadata:true 'http://169.254.169.254/metadata/instance/compute/vmId?api-version=2017-08-01&format=text'",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            auto_instance_id = azure_id.stdout.decode().strip()
+    except subprocess.CalledProcessError:
+        pass
+    try:
+        oci_check = subprocess.run(
+            "curl -s -m 1 -H 'Authorization: Bearer OracleCloud' 'http://169.254.169.254/opc/v1/instance/'",
+            shell=True,
+            capture_output=True,
+            check=True,
+        )
+        if oci_check.returncode == 0:
+            auto_service_provider = "OCI"
+            oci_id = subprocess.run(
+                "curl -s http://169.254.169.254/opc/v1/instance/id",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            auto_instance_id = oci_id.stdout.decode().strip()
+    except subprocess.CalledProcessError:
+        pass
+    try:
+        aws_check = subprocess.run(
+            "curl -s -m 1 http://169.254.169.254/latest/meta-data/",
+            shell=True,
+            capture_output=True,
+            check=True,
+        )
+        if aws_check.returncode == 0:
+            auto_service_provider = "AWS"
+            aws_id = subprocess.run(
+                "curl -s http://169.254.169.254/latest/meta-data/instance-id",
+                shell=True,
+                capture_output=True,
+                check=True,
+            )
+            auto_instance_id = aws_id.stdout.decode().strip()
+    except subprocess.CalledProcessError:
+        pass
+    return str(auto_service_provider), str(auto_instance_id)
+
+
+@log_errors(default_return=0, raise_exception=False, log_error=False)
+def calculate_time_difference(start_time_str: str, finish_time_str: str) -> int:
+    """
+    Calculate time difference between start and finish times.
+
+    Args:
+        start_time_str (str): Start time string
+        finish_time_str (str): Finish time string
+
+    Returns:
+        int: Time difference in seconds
+    """
+    if os.environ["SERVICE_PROVIDER"] in [
+        "AWS",
+        "OCI",
+        "LAMBDA",
+    ]:
+        start_time = datetime.fromisoformat(start_time_str.split(".")[0] + "+00:00")
+        finish_time = datetime.fromisoformat(finish_time_str.split(".")[0] + "+00:00")
+    else:
+        start_time = datetime.fromisoformat(start_time_str.replace("Z", "+00:00"))
+        finish_time = datetime.fromisoformat(finish_time_str.replace("Z", "+00:00"))
+    return int((finish_time - start_time).total_seconds())
+
+
+@log_errors(default_return=False, raise_exception=False, log_error=False)
+def has_gpu() -> bool:
+    """
+    Check if the system has a GPU.
+
+    Returns:
+        bool: True if GPU is present, False otherwise
+    """
+    subprocess.run("nvidia-smi", check=True)
+    return True
+
+
+@log_errors(default_return=0, raise_exception=False)
+def get_gpu_memory_usage() -> float:
+    """
+    Get GPU memory usage percentage.
+
+    Returns:
+        float: Memory usage between 0 and 1
+    """
+    command = "nvidia-smi --query-gpu=memory.used,memory.total --format=csv,nounits,noheader"
+    output = subprocess.check_output(command.split()).decode("ascii").strip().split("\n")
+    memory_percentages = []
+    for line in output:
+        used, total = map(int, line.split(","))
+        usage_percentage = used / total
+        memory_percentages.append(usage_percentage)
+    return min(memory_percentages)
+
+
+@log_errors(default_return=0, raise_exception=False)
+def get_cpu_memory_usage() -> float:
+    """
+    Get CPU memory usage.
+
+    Returns:
+        float: Memory usage between 0 and 1
+    """
+    memory = psutil.virtual_memory()
+    return memory.percent / 100
+
+
+@log_errors(default_return=0, raise_exception=False)
+def get_mem_usage() -> float:
+    """
+    Get memory usage for either GPU or CPU.
+
+    Returns:
+        float: Memory usage between 0 and 1
+    """
+    if has_gpu():
+        try:
+            mem_usage = get_gpu_memory_usage()
+        except Exception as err:
+            logging.error(
+                "Error getting GPU memory usage: %s",
+                err,
+            )
+            mem_usage = get_cpu_memory_usage()
+    else:
+        mem_usage = get_cpu_memory_usage()
+    if mem_usage is None:
+        mem_usage = 0
+    return mem_usage
+
+
+@log_errors(default_return=[], raise_exception=False)
+def get_gpu_info() -> list:
+    """
+    Get GPU information.
+
+    Returns:
+        list: GPU information strings
+    """
+    with subprocess.Popen(
+        [
+            "nvidia-smi",
+            "--query-gpu=index,uuid,utilization.gpu,memory.total,memory.used,memory.free,driver_version,name,gpu_serial,display_active,display_mode,temperature.gpu",
+            "--format=csv,noheader,nounits",
+        ],
+        stdout=subprocess.PIPE,
+    ) as proc:
+        stdout, _ = proc.communicate()
+        output = stdout.decode("UTF-8")
+        return output.split("\n")[:-1]
+
+
+@log_errors(default_return="", raise_exception=False)
+def get_instance_id() -> str:
+    """
+    Get instance ID.
+
+    Returns:
+        str: Instance ID or empty string
+    """
+    return os.environ["INSTANCE_ID"]
+
+
+@log_errors(default_return=False, raise_exception=False, log_error=False)
+def is_docker_running() -> bool:
+    """
+    Check if Docker is running.
+
+    Returns:
+        bool: True if Docker containers are running
+    """
+    command = "docker ps"
+    docker_images = (
+        subprocess.check_output(command.split()).decode("ascii").split("\n")[:-1][1:]
+    )
+    return bool(docker_images)
+
+
+@log_errors(default_return=None, raise_exception=False)
+def prune_docker_images() -> None:
+    """Prune Docker images."""
+    subprocess.run(
+        [
+            "docker",
+            "image",
+            "prune",
+            "-a",
+            "-f",
+        ],
+        check=True,
+    )
+    logging.info("Docker images pruned successfully.")
+
+
+@log_errors(default_return=0.0, raise_exception=False)
+def _normalize_disk_usage_to_gb(disk_space: str) -> float:
+    """
+    Normalize disk usage to GB.
+
+    Args:
+        disk_space (str): Disk space with unit
+
+    Returns:
+        float: Disk space in GB
+    """
+    if disk_space.endswith("G"):
+        result = float(disk_space[:-1])
+    elif disk_space.endswith("T"):
+        result = float(disk_space[:-1]) * 1024
+    elif disk_space.endswith("M"):
+        result = float(disk_space[:-1]) / 1024
+    elif disk_space.endswith("K"):
+        result = float(disk_space[:-1]) / (1024 * 1024)
+    else:
+        result = float(disk_space)
+    logging.debug(
+        "Normalized disk space value to %f GB",
+        result,
+    )
+    return result
+
+
+@log_errors(default_return=None, raise_exception=False)
+def _parse_disk_usage_info(line: str) -> dict:
+    """
+    Parse disk usage information.
+
+    Args:
+        line (str): Disk usage line from df command
+
+    Returns:
+        dict: Parsed disk usage information
+    """
+    parts = line.split()
+    parsed_info = {
+        "filesystem": parts[0],
+        "size": _normalize_disk_usage_to_gb(parts[1]),
+        "used": _normalize_disk_usage_to_gb(parts[2]),
+        "available": _normalize_disk_usage_to_gb(parts[3]),
+        "use_percentage": float(parts[4].rstrip("%")),
+        "mounted_on": parts[5],
+    }
+    logging.debug(
+        "Successfully parsed disk usage info: %s",
+        parsed_info,
+    )
+    return parsed_info
+
+
+@log_errors(default_return=None, raise_exception=False)
+def get_disk_space_usage() -> list:
+    """
+    Get disk space usage for all filesystems.
+
+    Returns:
+        list: List of disk usage information dictionaries
+    """
+    logging.info("Getting disk space usage information")
+    result = subprocess.run(
+        ["df", "-h"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    lines = result.stdout.strip().split("\n")[1:]
+    disk_usage = []
+    for line in lines:
+        disk = _parse_disk_usage_info(line)
+        if disk:
+            disk_usage.append(disk)
+    logging.info(
+        "Found disk usage info for %d filesystems",
+        len(disk_usage),
+    )
+    return disk_usage
+
+
+@log_errors(default_return=None, raise_exception=False)
+def get_max_file_system() -> str:
+    """
+    Get filesystem with maximum available space.
+
+    Returns:
+        str: Path to filesystem with most space or None
+    """
+    logging.info("Finding filesystem with maximum available space")
+    disk_usage = get_disk_space_usage()
+    if not disk_usage:
+        logging.warning("No disk usage information available")
+        return None
+    filtered_disks = [
+        disk
+        for disk in disk_usage
+        if disk["mounted_on"] != "/boot/efi"
+        and "overlay" not in disk["filesystem"]
+        and disk["available"] > 0
+    ]
+    if not filtered_disks:
+        logging.warning("No suitable filesystems found after filtering")
+        max_available_filesystem = ""
+    else:
+        max_disk = max(
+            filtered_disks,
+            key=lambda x: x["available"],
+        )
+        max_available_filesystem = max_disk["mounted_on"]
+        logging.info(
+            "Found filesystem with maximum space: %s (%f GB available)",
+            max_available_filesystem,
+            max_disk["available"],
+        )
+    if max_available_filesystem in ["/", ""]:
+        home_dir = os.path.expanduser("~")
+        if not os.environ.get("WORKSPACE_DIR"):
+            logging.error("WORKSPACE_DIR environment variable not set")
+            return None
+        workspace_dir = os.path.join(
+            home_dir,
+            os.environ["WORKSPACE_DIR"],
+        )
+        os.makedirs(workspace_dir, exist_ok=True)
+        logging.info(
+            "Created workspace directory at: %s",
+            workspace_dir,
+        )
+        return workspace_dir
+    return max_available_filesystem
+
+
+@log_errors(default_return=None, raise_exception=False)
+def get_docker_disk_space_usage() -> dict:
+    """
+    Get disk space usage for Docker storage.
+
+    Returns:
+        dict: Docker disk usage information
+    """
+    result = subprocess.run(
+        ["docker", "info"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    docker_info = result.stdout
+    docker_root_dir = None
+    for line in docker_info.split("\n"):
+        if line.strip().startswith("Docker Root Dir"):
+            docker_root_dir = line.split(":")[1].strip()
+            break
+    if docker_root_dir is None:
+        logging.error("Unable to find Docker root directory")
+        raise ValueError("Unable to find Docker root directory")
+    logging.debug(
+        "Found Docker root directory: %s",
+        docker_root_dir,
+    )
+    result = subprocess.run(
+        ["df", "-h", docker_root_dir],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    lines = result.stdout.strip().split("\n")[1:]
+    if not lines:
+        logging.error("No disk usage information found for Docker root directory")
+        raise ValueError("No disk usage information found for Docker root directory")
+    docker_disk_usage = _parse_disk_usage_info(lines[0])
+    if docker_disk_usage is None:
+        logging.error("Failed to parse Docker disk usage information")
+        raise ValueError("Failed to parse Docker disk usage information")
+    logging.info(
+        "Successfully retrieved Docker disk usage: %s",
+        docker_disk_usage,
+    )
+    return docker_disk_usage
+
+
+@log_errors(raise_exception=False)
+def cleanup_docker_storage() -> None:
+    """Clean up Docker storage if space is low."""
+    docker_disk_usage = get_docker_disk_space_usage()
+    if docker_disk_usage is None:
+        logging.error("Failed to get Docker disk space usage, skipping cleanup")
+        return
+    if docker_disk_usage["use_percentage"] >= 90 or docker_disk_usage["available"] <= 30:
+        logging.info(
+            "Pruning Docker images. Disk space is low: %s",
+            docker_disk_usage,
+        )
+        prune_docker_images()
+
+
+@log_errors(default_return=0, raise_exception=False)
+def get_required_gpu_memory(action_details: dict) -> int:
+    """
+    Get required GPU memory from action details.
+
+    Args:
+        action_details (dict): Action details
+
+    Returns:
+        int: Required GPU memory
+    """
+    try:
+        return action_details["actionDetails"]["expectedResources"]["gpuMemory"]
+    except KeyError:
+        return 0
+
+
+@log_errors(default_return=True, raise_exception=False)
+def is_allowed_gpu_device(gpu_index: int) -> bool:
+    """Check if GPU device is allowed.
+
+    Args:
+        gpu_index (int): GPU device index
+
+    Returns:
+        bool: True if GPU is allowed
+    """
+    gpus = os.environ.get("GPUS")
+    if not gpus:
+        return True
+    allowed_gpus = [int(x) for x in gpus.split(",") if x.strip()]
+    return int(gpu_index) in allowed_gpus
+
+
+@log_errors(raise_exception=True)
+def get_gpu_with_sufficient_memory_for_action(
+    action_details: dict,
+) -> list:
+    """
+    Get GPUs with sufficient memory for action.
+
+    Args:
+        action_details (dict): Action details
+
+    Returns:
+        list: List of GPU indices
+
+    Raises:
+        ValueError: If insufficient GPU memory
+    """
+    required_gpu_memory = get_required_gpu_memory(action_details)
+    command = "nvidia-smi --query-gpu=memory.free --format=csv"
+    memory_free_info = subprocess.check_output(command.split()).decode("ascii").split("\n")
+    if len(memory_free_info) < 2:
+        raise ValueError("No GPU information available from nvidia-smi")
+    memory_free_values = [int(x.split()[0]) for x in memory_free_info[1:-1]]
+    if required_gpu_memory < 80000:
+        try:
+            return get_single_gpu_with_sufficient_memory_for_action(action_details)
+        except ValueError:
+            pass
+    selected_gpus = []
+    total_memory = 0
+    for i, mem in enumerate(memory_free_values):
+        if not is_allowed_gpu_device(i):
+            continue
+        if total_memory >= required_gpu_memory:
+            break
+        selected_gpus.append(i)
+        total_memory += mem
+    if total_memory >= required_gpu_memory:
+        return selected_gpus
+    raise ValueError(
+        f"Insufficient GPU memory available. Required: {required_gpu_memory}, Available: {total_memory}"
+    )
+
+
+@log_errors(raise_exception=True)
+def get_single_gpu_with_sufficient_memory_for_action(
+    action_details: dict,
+) -> list:
+    """
+    Get single GPU with sufficient memory.
+
+    Args:
+        action_details (dict): Action details
+
+    Returns:
+        list: List with single GPU index
+
+    Raises:
+        ValueError: If no GPU has sufficient memory
+    """
+    required_gpu_memory = get_required_gpu_memory(action_details)
+    command = "nvidia-smi --query-gpu=memory.free --format=csv"
+    memory_free_info = subprocess.check_output(command.split()).decode("ascii").split("\n")
+    if len(memory_free_info) < 2:
+        raise ValueError("No GPU information available from nvidia-smi")
+    memory_free_values = [int(x.split()[0]) for x in memory_free_info[1:-1]]
+    best_fit_gpu = None
+    best_fit_memory = float("inf")
+    for i, mem in enumerate(memory_free_values):
+        if not is_allowed_gpu_device(i):
+            continue
+        if mem >= required_gpu_memory and mem < best_fit_memory:
+            best_fit_gpu = i
+            best_fit_memory = mem
+    if best_fit_gpu is not None:
+        return [best_fit_gpu]
+    raise ValueError(
+        f"No single GPU with sufficient memory ({required_gpu_memory}MB) available"
+    )
+
+
+@log_errors(default_return=(None, None), raise_exception=False)
+def get_decrypted_access_key_pair(
+    enc_access_key: str,
+    enc_secret_key: str,
+    encryption_key: str = "",
+) -> tuple:
+    """
+    Get decrypted access key pair.
+
+    Args:
+        enc_access_key (str): Encrypted access key
+        enc_secret_key (str): Encrypted secret key
+        encryption_key (str): Encryption key
+
+    Returns:
+        tuple: (access_key, secret_key) strings
+    """
+    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY")
+    if not encryption_key:
+        logging.warning("Encryption key is not set, Will assume that the keys are not encrypted")
+        return enc_access_key, enc_secret_key
+    encrypted_access_key = base64.b64decode(enc_access_key)
+    encrypted_secret_key = base64.b64decode(enc_secret_key)
+    nonce = encrypted_access_key[:12]
+    tag = encrypted_access_key[-16:]
+    ciphertext = encrypted_access_key[12:-16]
+    cipher = Cipher(
+        algorithms.AES(encryption_key.encode()),
+        modes.GCM(nonce, tag),
+        backend=default_backend(),
+    )
+    decryptor = cipher.decryptor()
+    decrypted_access_key = decryptor.update(ciphertext) + decryptor.finalize()
+    nonce = encrypted_secret_key[:12]
+    tag = encrypted_secret_key[-16:]
+    ciphertext = encrypted_secret_key[12:-16]
+    cipher = Cipher(
+        algorithms.AES(encryption_key.encode()),
+        modes.GCM(nonce, tag),
+        backend=default_backend(),
+    )
+    decryptor = cipher.decryptor()
+    decrypted_secret_key = decryptor.update(ciphertext) + decryptor.finalize()
+    access_key = decrypted_access_key.decode("utf-8", errors="replace")
+    secret_key = decrypted_secret_key.decode("utf-8", errors="replace")
+    return access_key, secret_key
+
+@log_errors(default_return=(None, None), raise_exception=False)
+def get_encrypted_access_key_pair(
+    access_key: str,
+    secret_key: str,
+    encryption_key: str = "",
+) -> tuple:
+    """
+    Get encrypted access key pair.
+
+    Args:
+        access_key (str):  access key
+        secret_key (str):  secret key
+        encryption_key (str): Encryption key
+
+    Returns:
+        tuple: (encrypted_access_key, encrypted_secret_key) strings
+    """
+    encryption_key = encryption_key or os.environ.get("MATRICE_ENCRYPTION_KEY")
+    if not encryption_key:
+        logging.warning("Encryption key is not set, returning unencrypted keys")
+        return access_key, secret_key
+    
+    # Convert encryption key to bytes
+    key = encryption_key.encode()
+    
+    # Encrypt access key
+    nonce = os.urandom(12)
+    cipher = Cipher(
+        algorithms.AES(key),
+        modes.GCM(nonce),
+        backend=default_backend()
+    )
+    encryptor = cipher.encryptor()
+    encrypted_access_key = encryptor.update(access_key.encode()) + encryptor.finalize()
+    encrypted_access_key_with_nonce = nonce + encrypted_access_key + encryptor.tag
+    
+    # Encrypt secret key
+    nonce = os.urandom(12)
+    cipher = Cipher(
+        algorithms.AES(key),
+        modes.GCM(nonce),
+        backend=default_backend()
+    )
+    encryptor = cipher.encryptor()
+    encrypted_secret_key = encryptor.update(secret_key.encode()) + encryptor.finalize()
+    encrypted_secret_key_with_nonce = nonce + encrypted_secret_key + encryptor.tag
+    
+    # Encode to base64 for storage
+    encoded_access_key = base64.b64encode(encrypted_access_key_with_nonce).decode()
+    encoded_secret_key = base64.b64encode(encrypted_secret_key_with_nonce).decode()
+    
+    return encoded_access_key, encoded_secret_key
+
+@log_errors(default_return=False, raise_exception=False)
+def check_public_port_exposure(port: int) -> bool:
+    """
+    Check if port is publicly accessible.
+
+    Args:
+        port (int): Port number to check
+
+    Returns:
+        bool: True if port is publicly accessible
+    """
+    is_public_exposed = False
+    is_locally_available = False
+    # Check if port is publicly accessible
+    public_ip = urllib.request.urlopen("https://ident.me", timeout=10).read().decode("utf8")
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as conn_sock:
+        conn_sock.settimeout(3)
+        result = conn_sock.connect_ex((public_ip, port))
+        is_public_exposed = result == 0
+    
+    # Check if port is locally available
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as bind_sock:
+        bind_sock.setsockopt(
+            socket.SOL_SOCKET,
+            socket.SO_REUSEADDR,
+            1,
+        )
+        bind_sock.bind(("", port))
+        bind_sock.listen(1)
+        is_locally_available = True
+
+    if not is_public_exposed:
+        logging.debug(
+            "Port %d is not publicly exposed",
+            port,
+        )
+        return False
+    if not is_locally_available:
+        logging.debug(
+            "Port %d is not locally available",
+            port,
+        )
+        return False
+    return True