agentscope.manager._manager 源代码

# -*- coding: utf-8 -*-
"""A manager for AgentScope."""
import os
from typing import Union, Any
from copy import deepcopy

from loguru import logger

from ._monitor import MonitorManager
from ._file import FileManager
from ._model import ModelManager
from ..logging import LOG_LEVEL, setup_logger
from .._version import __version__
from ..utils.common import (
    _generate_random_code,
    _get_process_creation_time,
    _get_timestamp,
)
from ..constants import _RUNTIME_ID_FORMAT, _RUNTIME_TIMESTAMP_FORMAT
from ..studio._client import _studio_client


[文档] class ASManager: """A manager for AgentScope.""" _instance = None __serialized_attrs = [ "project", "name", "disable_saving", "run_id", "pid", "timestamp", ] def __new__(cls, *args: Any, **kwargs: Any) -> "ASManager": if cls._instance is None: cls._instance = super(ASManager, cls).__new__(cls) return cls._instance
[文档] @classmethod def get_instance(cls) -> "ASManager": """Get the instance of the singleton class.""" if cls._instance is None: raise ValueError( "AgentScope hasn't been initialized. Please call " "`agentscope.init` function first.", ) return cls._instance
[文档] def __init__(self) -> None: """Initialize the manager. Note we initialize the managers by default arguments to avoid unnecessary errors when user doesn't call `agentscope.init` function""" self.project = "" self.name = "" self.run_id = "" self.pid = -1 self.timestamp = "" self.disable_saving = True self.file = FileManager() self.model = ModelManager() self.monitor = MonitorManager() # TODO: unified with logger and studio self.logger_level: LOG_LEVEL = "INFO"
[文档] def initialize( self, model_configs: Union[dict, str, list, None], project: Union[str, None], name: Union[str, None], disable_saving: bool, save_dir: str, save_log: bool, save_code: bool, save_api_invoke: bool, cache_dir: str, use_monitor: bool, logger_level: LOG_LEVEL, run_id: Union[str, None], studio_url: Union[str, None], ) -> None: """Initialize the package.""" # =============== Init the runtime =============== self.project = project or _generate_random_code() self.name = name or _generate_random_code(uppercase=False) self.pid = os.getpid() timestamp = _get_process_creation_time() self.timestamp = timestamp.strftime(_RUNTIME_TIMESTAMP_FORMAT) self.run_id = run_id or _get_timestamp( _RUNTIME_ID_FORMAT, timestamp, ).format(self.name) self.disable_saving = disable_saving # =============== Init the file manager =============== if disable_saving: save_log = False save_code = False save_api_invoke = False use_monitor = False run_dir = None else: run_dir = os.path.abspath(os.path.join(save_dir, self.run_id)) self.file.initialize( run_dir=run_dir, save_log=save_log, save_code=save_code, save_api_invoke=save_api_invoke, cache_dir=cache_dir, ) # Save the python code here to avoid duplicated saving in the child # process (when calling deserialize function) if save_code: self.file.save_python_code() if not disable_saving: # Save the runtime information in .config file self.file.save_runtime_information(self.state_dict()) # =============== Init the logger =============== # TODO: unified with studio and gradio self.logger_level = logger_level # run_dir will be None if save_log is False setup_logger(self.file.run_dir, logger_level) # =============== Init the model manager =============== self.model.initialize(model_configs) # =============== Init the monitor manager =============== self.monitor.initialize(use_monitor) # =============== Init the studio =============== # TODO: unified with studio and gradio # Init studio client, which will push messages to web ui and fetch user # inputs from web ui if studio_url is not None: _studio_client.initialize(self.run_id, studio_url) # Register in AgentScope Studio _studio_client.register_running_instance( project=self.project, name=self.name, timestamp=self.timestamp, run_dir=self.file.run_dir, pid=self.pid, )
[文档] def state_dict(self) -> dict: """Serialize the runtime information.""" serialized_data = { k: getattr(self, k) for k in self.__serialized_attrs } serialized_data["agentscope_version"] = __version__ serialized_data["file"] = self.file.state_dict() serialized_data["model"] = self.model.state_dict() serialized_data["logger"] = { "level": self.logger_level, } serialized_data["studio"] = _studio_client.state_dict() serialized_data["monitor"] = self.monitor.state_dict() return deepcopy(serialized_data)
[文档] def load_dict(self, data: dict) -> None: """Load the runtime information from a dictionary""" for k in self.__serialized_attrs: assert k in data, f"Key {k} not found in data." setattr(self, k, data[k]) self.file.load_dict(data["file"]) # TODO: unified the logger with studio and gradio self.logger_level = data["logger"]["level"] setup_logger(self.file.run_dir, self.logger_level) self.model.load_dict(data["model"]) _studio_client.load_dict(data["studio"]) self.monitor.load_dict(data["monitor"])
[文档] def flush(self) -> None: """Flush the runtime information.""" self.project = "" self.name = "" self.run_id = "" self.pid = -1 self.timestamp = "" self.disable_saving = True self.file.flush() self.model.flush() self.monitor.flush() logger.remove() _studio_client.flush() self.logger_level = "INFO"