setup_environment.py

import os
import subprocess
import sys 
import gdown
import zipfile


# pip 업데이트
def upgrade_pip():
    """pip를 최신 버전으로 업그레이드"""
    print("Upgrading pip...")
    subprocess.run([sys.executable, "-m", "pip", "install", "--upgrade", "pip"], check=True)


def set_pythonpath():
    """PYTHONPATH 환경 변수를 설정 (stf-tools의 설치 경로를 기반으로)"""
    
    # pip show stf-tools 명령어 실행
    try:
        result = subprocess.run(['pip', 'show', 'stf-alternative'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        
        if result.returncode != 0:
            print("stf-tools 패키지를 찾을 수 없습니다.")
            return
        
        # 출력에서 Location 찾기
        location_line = next((line for line in result.stdout.splitlines() if line.startswith("Location:")), None)
        
        if location_line:
            # Location에서 경로 추출
            location_path = location_line.split("Location: ")[1].strip()
            print('location_path===',location_path)
        else:
            print("Location 정보를 찾을 수 없습니다.")
            return
    
    except Exception as e:
        print(f"패키지 경로를 가져오는 중 에러가 발생했습니다: {e}")
        return
    
    # 현재 PYTHONPATH 확인
    current_pythonpath = os.environ.get("PYTHONPATH", "")
    
    # 새로 추가할 PYTHONPATH 경로 (stf-tools 위치 기반)
    new_pythonpath = location_path
    
    # 기존 PYTHONPATH가 있으면 추가하고, 없으면 새로 설정
    if current_pythonpath:
        new_pythonpath = f"{new_pythonpath}:{current_pythonpath}"
    
    # PYTHONPATH 설정
    os.environ["PYTHONPATH"] = new_pythonpath
    print(f"PYTHONPATH set to: {os.environ['PYTHONPATH']}")


# 현재 디렉토리 저장
original_dir = os.getcwd()

def run_pip_install(path):
    """해당 경로로 이동한 후 pip install . 실행"""
    if os.path.exists(os.path.join(path, 'setup.py')) or os.path.exists(os.path.join(path, 'pyproject.toml')):
        print(f"Installing dependencies in {path} using pip install .")
        os.chdir(path)
        try:
            result = subprocess.run([sys.executable, "-m", "pip", "install", "--user", "."], check=True, stderr=subprocess.PIPE, text=True)
            print(result.stdout)
        except subprocess.CalledProcessError as e:
            print(f"Error occurred while installing in {path}: {e.stderr}")
    else:
        print(f"No setup.py or pyproject.toml found in {path}, skipping...")


def check_stf_alternative_installed():
    # pip list | grep stf-alternative 명령을 실행
    result = subprocess.run(['pip', 'list'], stdout=subprocess.PIPE, text=True)
    
    # 출력에서 stf-alternative 패키지가 있는지 확인
    if "stf-alternative" in result.stdout:
        print("stf-alternative 패키지가 설치되어 있습니다.")
    else:
        print("stf-alternative 패키지가 설치되지 않았습니다.")



# 추가된 함수들: 'libcublasLt.so.11'의 경로를 찾아 LD_LIBRARY_PATH에 추가하는 함수
def find_library(library_name):
    try:
        result = subprocess.run(['find', '/usr', '-name', library_name], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        library_path = result.stdout.decode('utf-8').strip()

        if library_path:
            library_dir = os.path.dirname(library_path)
            return library_dir
        else:
            print(f"라이브러리 {library_name}을(를) 찾을 수 없습니다.")
            return None
    except Exception as e:
        print(f"오류 발생: {e}")
        return None

def add_to_ld_library_path(library_dir):
    if library_dir:
        ld_library_path = os.environ.get('LD_LIBRARY_PATH', '')
        if library_dir not in ld_library_path:
            os.environ['LD_LIBRARY_PATH'] = library_dir + ':' + ld_library_path
            print(f"{library_dir}이(가) LD_LIBRARY_PATH에 추가되었습니다.")
        else:
            print(f"{library_dir}이(가) 이미 LD_LIBRARY_PATH에 포함되어 있습니다.")
    else:
        print("유효한 라이브러리 경로가 없습니다.")




def initialize_environment():
    """환경 설정을 위한 전체 프로세스 실행"""
    print('aaaaaaaaaaaa')

    # pip 업그레이드 실행
    upgrade_pip()
    
    # 함수 호출
    print('11111111')
    check_stf_alternative_installed()


    # stf-api-alternative에서 pip install . 실행
    run_pip_install("/tmp/stf/stf-api-alternative")

    # PYTHONPATH 설정
    set_pythonpath()

    # 'libcublasLt.so.11' 라이브러리 경로 찾기 및 LD_LIBRARY_PATH에 추가
    library_name = 'libcublasLt.so.*'
    library_dir = find_library(library_name)
    add_to_ld_library_path(library_dir)

    # 함수 호출
    print('222222')
    check_stf_alternative_installed()

    # 원래 작업 디렉토리로 돌아옴
    os.chdir(original_dir)