利用importlib实现Python大型数组内存驻留及代码热更新

import pickle
import numpy as np

FILE_PATH_TO_PICKLED_ARRAY = "large_array.pkl" # 假设这是一个很大的数组文件

def some_task(arr):
    # arr 在此任务中是只读的
    # 例如：
    result = arr.sum()
    return result

if __name__ == "__main__":
    # 模拟生成一个大型数组并保存
    # arr = np.random.rand(10000, 10000) # 示例，实际可能更大
    # with open(FILE_PATH_TO_PICKLED_ARRAY, "wb") as f:
    #     pickle.dump(arr, f)

    arr = pickle.load(open(FILE_PATH_TO_PICKLED_ARRAY, "rb"))
    some_task(arr)

# the_other_script.py
import numpy as np

def some_task(arr: np.ndarray):
    """
    对传入的数组执行某个任务。
    这个函数会频繁修改和测试。
    """
    # arr 在此任务中是只读的
    # 示例：计算数组的和，并打印一个版本信息
    result = arr.sum()
    print(f"执行 some_task (版本 1.0): 数组总和为 {result}")
    return result

# 可以在这里添加更多需要测试的函数
def another_task(arr: np.ndarray):
    print(f"执行 another_task (版本 1.0): 数组均值为 {arr.mean()}")
    return arr.mean()

# main_script.py
import importlib.util
import sys
import pickle
import numpy as np
import os

# 假设大型数组文件路径
FILE_PATH_TO_PICKLED_ARRAY = "large_array.pkl"
# 任务脚本文件路径
FILE_PATH_TO_THE_OTHER_SCRIPT = "the_other_script.py"

def load_module(source: str, module_name: str):
    """
    从指定文件路径加载或重新加载一个模块。

    :param source: 模块文件的路径。
    :param module_name: 模块在 sys.modules 中注册的名称。
    :return: 加载或重新加载的模块对象。
    """
    # 检查模块是否已经存在于 sys.modules 中
    if module_name in sys.modules:
        # 如果模块已加载，则尝试重新加载
        print(f"重新加载模块: {module_name} from {source}")
        # 获取模块对象
        module = sys.modules[module_name]
        # 使用 reload() 函数重新加载模块
        importlib.reload(module)
        return module
    else:
        # 如果模块未加载，则从文件加载
        print(f"首次加载模块: {module_name} from {source}")
        spec = importlib.util.spec_from_file_location(module_name, source)
        if spec is None:
            raise ImportError(f"无法找到模块文件: {source}")
        module = importlib.util.module_from_spec(spec)
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
        return module

if __name__ == "__main__":
    # 模拟生成一个大型数组并保存（如果文件不存在）
    if not os.path.exists(FILE_PATH_TO_PICKLED_ARRAY):
        print(f"生成大型数组并保存到 {FILE_PATH_TO_PICKLED_ARRAY}...")
        arr_to_save = np.random.rand(5000, 5000) # 示例，约200MB
        with open(FILE_PATH_TO_PICKLED_ARRAY, "wb") as f:
            pickle.dump(arr_to_save, f)
        print("数组生成完毕。")

    print(f"正在加载大型数组从 {FILE_PATH_TO_PICKLED_ARRAY}...")
    arr = pickle.load(open(FILE_PATH_TO_PICKLED_ARRAY, "rb"))
    print(f"大型数组加载完成，形状：{arr.shape}，占用内存：{arr.nbytes / (1024**2):.2f} MB")

    while True:
        user_input = input("\n请按 Enter 键重新加载模块并运行 some_task (输入 'q' 退出): ")
        if user_input.lower() == 'q':
            print("程序退出。")
            break

        try:
            # 动态加载或重新加载 the_other_script 模块
            the_other_script_module = load_module(FILE_PATH_TO_THE_OTHER_SCRIPT, "the_other_script")

            # 调用模块中的 some_task 函数
            print("正在执行 some_task...")
            the_other_script_module.some_task(arr)

            # 也可以调用其他函数
            # print("正在执行 another_task...")
            # the_other_script_module.another_task(arr)

        except Exception as e:
            print(f"执行任务时发生错误: {e}")
            import traceback
            traceback.print_exc()