flax_model/alphafold3/data/tools/mmseqs.py

"""Library to run Mmseqs from Python."""

import os
import tempfile
import subprocess
import time
from typing import List, Dict, Optional


from absl import logging
from flax_model.alphafold3.data import parsers
from flax_model.alphafold3.data.tools import msa_tool
from flax_model.alphafold3.data.tools import subprocess_utils
import shlex

import shutil
import re


class Mmseqs(msa_tool.MsaTool):
    
    _database_types = ["mgnify", "uniprot_cluster_annot", "uniref90", "small_bfd"]

    def __init__(
        self,
        *,
        binary_path: str,
        database_path: str,
        n_cpu: int = 8,
        e_value: float = 1e-4,
        max_sequences: int = 5000,
        use_gpu: int = 1,
        msa_format_mode: int = 4,
        mmseqs_options: str = "",
        result2msa_options: str = "",
    ):
        self.binary_path = binary_path
        self.database_path = database_path
        
        subprocess_utils.check_binary_exists(path=self.binary_path, name='MMseqs')
        
        if not os.path.exists(self.database_path):
            raise ValueError(f'Database not found: {database_path}')
        
        self.n_cpu = n_cpu
        self.e_value = e_value
        self.max_sequences = max_sequences
        self.use_gpu = use_gpu
        self.msa_format_mode = msa_format_mode
        self.mmseqs_options = mmseqs_options
        self.result2msa_options = result2msa_options

    def _get_gpu_memory_gb(self) -> float:
        hip_devices = os.environ.get('HIP_VISIBLE_DEVICES', '')
        cuda_devices = os.environ.get('CUDA_VISIBLE_DEVICES', '')

        hy_smi = shutil.which('hy-smi')
        nvidia_smi = shutil.which('nvidia-smi')

        prefer_hip = bool(hip_devices or hy_smi)

        device_id: Optional[str] = None
        if prefer_hip and hip_devices:
            ids = [d.strip() for d in hip_devices.split(',') if d.strip()]
            device_id = ids[0] if ids else None
        elif cuda_devices:
            ids = [d.strip() for d in cuda_devices.split(',') if d.strip()]
            device_id = ids[0] if ids else None

        # 1) hy-smi (HIP path)
        if hy_smi:
            try:
                mem = self._get_gpu_memory_gb_via_hy_smi(hy_smi, device_id)
                if mem is not None:
                    logging.info(f"GPU memory (hy-smi): {mem:.1f}GB")
                    return mem
            except Exception as e:  # noqa: BLE001
                logging.warning(f"hy-smi query failed: {e}")

        # 2) NVIDIA fallback when CUDA devices are provided
        if nvidia_smi and cuda_devices:
            try:
                ids = [d.strip() for d in cuda_devices.split(',') if d.strip()]
                if ids:
                    cmd = [
                        nvidia_smi,
                        '--query-gpu=memory.total',
                        '--format=csv,noheader,nounits',
                        f'--id={ids[0]}',
                    ]
                    result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                    memory_mb = float(result.stdout.strip())
                    return memory_mb / 1024.0
            except (subprocess.CalledProcessError, ValueError, FileNotFoundError) as e:
                logging.warning(f"nvidia-smi query failed: {e}")

        logging.info("Falling back to default GPU memory: 40GB")
        return 40.0

    def _parse_memory_gb_from_text(self, text: str) -> Optional[float]:
        units = {
            'gib': 1.0,
            'gb': 1.0,
            'mib': 1.0 / 1024.0,
            'mb': 1.0 / 1024.0,
            'kib': 1.0 / (1024.0 * 1024.0),
            'kb': 1.0 / (1024.0 * 1024.0),
            'b': 1.0 / (1024.0 * 1024.0 * 1024.0),
        }
        pattern = re.compile(r"(\d+(?:\.\d+)?)\s*(GiB|GB|MiB|MB|KiB|KB|B)", re.IGNORECASE)
        candidates = []
        for match in pattern.finditer(text):
            value = float(match.group(1))
            unit = match.group(2).lower()
            factor = units.get(unit, None)
            if factor is not None:
                candidates.append(value * factor)
        if candidates:
            return max(candidates)
        # rocm-smi often prints bytes as: Total Memory (B): 17163091968
        bytes_pattern = re.compile(r"Total\s*Memory\s*\(B\)\s*[:=]\s*(\d{6,})", re.IGNORECASE)
        m = bytes_pattern.search(text)
        if m:
            try:
                return float(m.group(1)) / (1024.0 * 1024.0 * 1024.0)
            except ValueError:
                return None
        # generic Total: <bytes>
        generic_bytes = re.compile(r"Total\s*[:=]\s*(\d{6,})", re.IGNORECASE)
        m2 = generic_bytes.search(text)
        if m2:
            try:
                return float(m2.group(1)) / (1024.0 * 1024.0 * 1024.0)
            except ValueError:
                return None
        return None

    def _get_gpu_memory_gb_via_hy_smi(self, hy_smi_path: str, device_id: Optional[str]) -> Optional[float]:
        candidate_cmds = []
        # 1) Plain text output exactly as user provided; most reliable for MiB lines
        if device_id is not None:
            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '-d', str(device_id)])
        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram'])
        # 2) CSV/JSON (may not include units next to values, kept as fallback)
        if device_id is not None:
            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--csv', '-d', str(device_id)])
        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--csv'])
        if device_id is not None:
            candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--json', '-d', str(device_id)])
        candidate_cmds.append([hy_smi_path, '--showmeminfo', 'vram', '--json'])
        # 3) All-in-one info as a last resort
        candidate_cmds.append([hy_smi_path, '-a'])

        for cmd in candidate_cmds:
            try:
                result = subprocess.run(cmd, capture_output=True, text=True, check=True)
                # Prefer exact vram parser for hy-smi
                parsed = self._parse_hysmi_vram_total_gb(result.stdout, device_id)
                if parsed is None:
                    parsed = self._parse_memory_gb_from_text(result.stdout)
                if parsed is not None and parsed > 0:
                    return parsed
            except subprocess.CalledProcessError:
                continue
        return None

    def _parse_hysmi_vram_total_gb(self, text: str, device_id: Optional[str]) -> Optional[float]:
        # Match lines like: HCU[0]          : vram Total Memory (MiB): 65520
        pattern = re.compile(r"(?:DCU|HCU)\[(\d+)\].*?vram\s+Total\s+Memory\s*\(MiB\)\s*:\s*(\d+)", re.IGNORECASE)
        matches = pattern.findall(text)
        if not matches:
            return None
        values_gb = []
        for dev, mib_str in matches:
            try:
                if device_id is not None and str(device_id) != str(dev):
                    continue
                mib = float(mib_str)
                values_gb.append(mib / 1024.0)
            except ValueError:
                continue
        if not values_gb:
            return None
        return max(values_gb)

    
    
    def _get_database_type(self, database_path: str) -> str:
        path_lower = database_path.lower()
        for db_type in self._database_types:
            if db_type in path_lower:
                return db_type
        return "unknown"
    
    def _calculate_memory_allocation(self, database_path: str) -> int:
        db_type = self._get_database_type(database_path)
        if db_type == "unknown":
            return 8
        
        if not self.use_gpu:
            return 8
        
        gpu_memory_gb = self._get_gpu_memory_gb()
        
        available_memory = max(0, gpu_memory_gb)
        
        if db_type == "mgnify":
            small_dbs_memory = 3 * 2
            allocated_memory = int(available_memory - small_dbs_memory)
            allocated_memory = max(2, allocated_memory)
        else:
            allocated_memory = 2
        
        logging.info(f"Database {db_type}: allocated {allocated_memory}GB GPU memory")
        return allocated_memory

    def query(self, target_sequence: str) -> msa_tool.MsaToolResult:
        with tempfile.TemporaryDirectory() as tmp_dir:
            input_fasta = os.path.join(tmp_dir, 'query.fasta')
            subprocess_utils.create_query_fasta_file(target_sequence, input_fasta)
            
            query_db = os.path.join(tmp_dir, 'queryDB')
            self._run_createdb(input_fasta, query_db)
            
            result_db = os.path.join(tmp_dir, 'resultDB')
            self._run_search(query_db, result_db)  
            
            output_sto = os.path.join(tmp_dir, 'output.sto')
            self._run_result2msa(query_db, result_db, output_sto)
            
            with open(output_sto) as f:
                a3m = self._parse_output(f)
        
        return msa_tool.MsaToolResult(
            target_sequence=target_sequence,
            a3m=a3m,
            e_value=self.e_value
        )

    def _run_search(self, input_db: str, output_db: str):
        with tempfile.TemporaryDirectory() as tmp_dir:
            cmd = [
                self.binary_path,
                'search',
                input_db,
                self.database_path,
                output_db,
                tmp_dir,
                '--threads', str(self.n_cpu),
                '-e', str(self.e_value),
                '--gpu', str(self.use_gpu),
                *shlex.split(self.mmseqs_options)
            ]
            
            if self.use_gpu:
                max_gpu_mem = self._calculate_memory_allocation(self.database_path)
                cmd.extend(['--max-gpu-mem', f'{max_gpu_mem}G'])
                logging.info(f"Running MMseqs search with max GPU memory: {max_gpu_mem}GB")
            else:
                logging.info("Running MMseqs search in CPU mode")
            
            subprocess_utils.run(
                cmd=cmd,
                cmd_name='MMseqs2 search',
                log_stderr=True
            )

    def _run_createdb(self, input_fasta: str, output_db: str):
        cmd = [self.binary_path, 'createdb', input_fasta, output_db]
        subprocess_utils.run(cmd, 'MMseqs2 createdb')

    def _run_result2msa(self, query_db: str, result_db: str, output_sto: str):
        cmd = [
            self.binary_path,
            'result2msa',
            query_db,
            self.database_path,
            result_db,
            output_sto,
            '--msa-format-mode', str(self.msa_format_mode),
            *shlex.split(self.result2msa_options)
        ]
        subprocess_utils.run(cmd, 'MMseqs2 result2msa')

    def _parse_output(self, file_handle):
        try:
            return parsers.convert_stockholm_to_a3m(file_handle, self.max_sequences)
        except Exception as e:
            logging.warning(f"Stockholm parse failed: {e}")
            
            file_handle.seek(0)
            content = file_handle.read()
            debug_file = f"/tmp/stockholm_debug_{int(time.time())}.sto"
            with open(debug_file, 'w') as f:
                f.write(content)
            logging.info(f"Saved debug file: {debug_file}")
            
            file_handle.seek(0)
            return parsers.convert_mmseqs_stockholm_to_a3m(file_handle, self.max_sequences)