pycityagent 2.0.0a20__py3-none-any.whl → 2.0.0a22__py3-none-any.whl

pycityagent/environment/utils/protobuf.py

@@ -1,13 +1,15 @@
-from typing import Any, Awaitable, TypeVar, Union, Dict
-from google.protobuf.message import Message
+from collections.abc import Awaitable
+from typing import Any, TypeVar, Union
+
 from google.protobuf.json_format import MessageToDict
+from google.protobuf.message import Message
 
 __all__ = ["parse", "async_parse"]
 
 T = TypeVar("T", bound=Message)
 
 
-def parse(res: T, dict_return: bool) -> Union[Dict[str, Any], T]:
+def parse(res: T, dict_return: bool) -> Union[dict[str, Any], T]:
     """
     将Protobuf返回值转换为dict或者原始值
     Convert Protobuf return value to dict or original value
@@ -23,7 +25,7 @@ def parse(res: T, dict_return: bool) -> Union[Dict[str, Any], T]:
         return res
 
 
-async def async_parse(res: Awaitable[T], dict_return: bool) -> Union[Dict[str, Any], T]:
+async def async_parse(res: Awaitable[T], dict_return: bool) -> Union[dict[str, Any], T]:
     """
     将Protobuf await返回值转换为dict或者原始值
     Convert Protobuf await return value to dict or original value

pycityagent/memory/memory_base.py

@@ -6,7 +6,8 @@ import asyncio
 import logging
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Optional, Sequence, Union
+from collections.abc import Callable, Sequence
+from typing import Any, Optional, Union
 
 from .const import *
 
@@ -16,8 +17,8 @@ logger = logging.getLogger("pycityagent")
 class MemoryUnit:
     def __init__(
         self,
-        content: Optional[Dict] = None,
-        required_attributes: Optional[Dict] = None,
+        content: Optional[dict] = None,
+        required_attributes: Optional[dict] = None,
         activate_timestamp: bool = False,
     ) -> None:
         self._content = {}
@@ -52,7 +53,7 @@ class MemoryUnit:
         else:
             setattr(self, f"{SELF_DEFINE_PREFIX}{property_name}", property_value)
 
-    async def update(self, content: Dict) -> None:
+    async def update(self, content: dict) -> None:
         await self._lock.acquire()
         for k, v in content.items():
             if k in self._content:
@@ -111,14 +112,14 @@ class MemoryUnit:
 
     async def dict_values(
         self,
-    ) -> Dict[Any, Any]:
+    ) -> dict[Any, Any]:
         return self._content
 
 
 class MemoryBase(ABC):
 
     def __init__(self) -> None:
-        self._memories: Dict[Any, Dict] = {}
+        self._memories: dict[Any, dict] = {}
         self._lock = asyncio.Lock()
 
     @abstractmethod

pycityagent/memory/profile.py

@@ -2,8 +2,9 @@
 Agent Profile
 """
 
+from collections.abc import Callable, Sequence
 from copy import deepcopy
-from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
+from typing import Any, Optional, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *
@@ -14,7 +15,7 @@ from .utils import convert_msg_to_sequence
 class ProfileMemoryUnit(MemoryUnit):
     def __init__(
         self,
-        content: Optional[Dict] = None,
+        content: Optional[dict] = None,
         activate_timestamp: bool = False,
     ) -> None:
         super().__init__(
@@ -28,7 +29,7 @@ class ProfileMemory(MemoryBase):
     def __init__(
         self,
         msg: Optional[
-            Union[ProfileMemoryUnit, Sequence[ProfileMemoryUnit], Dict, Sequence[Dict]]
+            Union[ProfileMemoryUnit, Sequence[ProfileMemoryUnit], dict, Sequence[dict]]
         ] = None,
         activate_timestamp: bool = False,
     ) -> None:
@@ -74,7 +75,7 @@ class ProfileMemory(MemoryBase):
     @lock_decorator
     async def load(
         self,
-        snapshots: Union[Dict, Sequence[Dict]],
+        snapshots: Union[dict, Sequence[dict]],
         reset_memory: bool = False,
     ) -> None:
         if reset_memory:
@@ -91,7 +92,7 @@ class ProfileMemory(MemoryBase):
     @lock_decorator
     async def export(
         self,
-    ) -> Sequence[Dict]:
+    ) -> Sequence[dict]:
         _res = []
         for m in self._memories.keys():
             m = cast(ProfileMemoryUnit, m)
@@ -145,7 +146,7 @@ class ProfileMemory(MemoryBase):
                     self._memories[unit] = {}
 
     @lock_decorator
-    async def update_dict(self, to_update_dict: Dict, store_snapshot: bool = False):
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
         _latest_memories = self._fetch_recent_memory()
         _latest_memory: ProfileMemoryUnit = _latest_memories[-1]
         if not store_snapshot:

pycityagent/memory/self_define.py

@@ -2,8 +2,9 @@
 Self Define Data
 """
 
+from collections.abc import Callable, Sequence
 from copy import deepcopy
-from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
+from typing import Any, Optional, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *
@@ -14,8 +15,8 @@ from .utils import convert_msg_to_sequence
 class DynamicMemoryUnit(MemoryUnit):
     def __init__(
         self,
-        content: Optional[Dict] = None,
-        required_attributes: Optional[Dict] = None,
+        content: Optional[dict] = None,
+        required_attributes: Optional[dict] = None,
         activate_timestamp: bool = False,
     ) -> None:
         super().__init__(
@@ -29,7 +30,7 @@ class DynamicMemory(MemoryBase):
 
     def __init__(
         self,
-        required_attributes: Dict[Any, Any],
+        required_attributes: dict[Any, Any],
         activate_timestamp: bool = False,
     ) -> None:
         super().__init__()
@@ -69,7 +70,7 @@ class DynamicMemory(MemoryBase):
     @lock_decorator
     async def load(
         self,
-        snapshots: Union[Dict, Sequence[Dict]],
+        snapshots: Union[dict, Sequence[dict]],
         reset_memory: bool = False,
     ) -> None:
         if reset_memory:
@@ -86,7 +87,7 @@ class DynamicMemory(MemoryBase):
     @lock_decorator
     async def export(
         self,
-    ) -> Sequence[Dict]:
+    ) -> Sequence[dict]:
         _res = []
         for m in self._memories.keys():
             m = cast(DynamicMemoryUnit, m)
@@ -143,7 +144,7 @@ class DynamicMemory(MemoryBase):
                     self._memories[unit] = {}
 
     @lock_decorator
-    async def update_dict(self, to_update_dict: Dict, store_snapshot: bool = False):
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
         _latest_memories = self._fetch_recent_memory()
         _latest_memory: DynamicMemoryUnit = _latest_memories[-1]
         if not store_snapshot:

pycityagent/memory/state.py

@@ -2,8 +2,9 @@
 Agent State
 """
 
+from collections.abc import Callable, Sequence
 from copy import deepcopy
-from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
+from typing import Any, Optional, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *
@@ -14,7 +15,7 @@ from .utils import convert_msg_to_sequence
 class StateMemoryUnit(MemoryUnit):
     def __init__(
         self,
-        content: Optional[Dict] = None,
+        content: Optional[dict] = None,
         activate_timestamp: bool = False,
     ) -> None:
         super().__init__(
@@ -28,7 +29,7 @@ class StateMemory(MemoryBase):
     def __init__(
         self,
         msg: Optional[
-            Union[MemoryUnit, Sequence[MemoryUnit], Dict, Sequence[Dict]]
+            Union[MemoryUnit, Sequence[MemoryUnit], dict, Sequence[dict]]
         ] = None,
         activate_timestamp: bool = False,
     ) -> None:
@@ -73,7 +74,7 @@ class StateMemory(MemoryBase):
     @lock_decorator
     async def load(
         self,
-        snapshots: Union[Dict, Sequence[Dict]],
+        snapshots: Union[dict, Sequence[dict]],
         reset_memory: bool = False,
     ) -> None:
 
@@ -91,7 +92,7 @@ class StateMemory(MemoryBase):
     @lock_decorator
     async def export(
         self,
-    ) -> Sequence[Dict]:
+    ) -> Sequence[dict]:
 
         _res = []
         for m in self._memories.keys():
@@ -151,7 +152,7 @@ class StateMemory(MemoryBase):
                     self._memories[unit] = {}
 
     @lock_decorator
-    async def update_dict(self, to_update_dict: Dict, store_snapshot: bool = False):
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
 
         _latest_memories = self._fetch_recent_memory()
         _latest_memory: StateMemoryUnit = _latest_memories[-1]

pycityagent/memory/utils.py

@@ -1,4 +1,5 @@
-from typing import Any, Sequence, Union
+from collections.abc import Sequence
+from typing import Any, Union
 
 from .memory_base import MemoryUnit
 

pycityagent/simulation/__init__.py

@@ -3,5 +3,6 @@
 """
 
 from .simulation import AgentSimulation
+from .storage.pg import PgWriter, create_pg_tables
 
-__all__ = ["AgentSimulation"]
+__all__ = ["AgentSimulation", "PgWriter", "create_pg_tables"]

pycityagent/simulation/agentgroup.py

@@ -3,7 +3,7 @@ import json
 import logging
 import time
 import uuid
-from datetime import datetime
+from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any
 from uuid import UUID
@@ -35,7 +35,7 @@ class AgentGroup:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_copy_writer: ray.ObjectRef,
+        pgsql_writer: ray.ObjectRef,
         mlflow_run_id: str,
         logging_level: int,
     ):
@@ -45,6 +45,7 @@ class AgentGroup:
         self.config = config
         self.exp_id = exp_id
         self.enable_avro = enable_avro
+        self.enable_pgsql = enable_pgsql
         if enable_avro:
             self.avro_path = avro_path / f"{self._uuid}"
             self.avro_path.mkdir(parents=True, exist_ok=True)
@@ -54,6 +55,8 @@ class AgentGroup:
                 "status": self.avro_path / f"status.avro",
                 "survey": self.avro_path / f"survey.avro",
             }
+        if self.enable_pgsql:
+            pass
 
         self.messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
@@ -61,6 +64,8 @@ class AgentGroup:
             username=config["simulator_request"]["mqtt"].get("username", None),
             password=config["simulator_request"]["mqtt"].get("password", None),
         )
+        self._pgsql_writer = pgsql_writer
+        self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
         self.initialized = False
         self.id2agent = {}
         # Step:1 prepare LLM client
@@ -105,6 +110,8 @@ class AgentGroup:
             agent.set_messager(self.messager)
             if self.enable_avro:
                 agent.set_avro_file(self.avro_file)  # type: ignore
+            if self.enable_pgsql:
+                agent.set_pgsql_writer(self._pgsql_writer)
 
     async def init_agents(self):
         logger.debug(f"-----Initializing Agents in AgentGroup {self._uuid} ...")
@@ -161,6 +168,20 @@ class AgentGroup:
             with open(filename, "wb") as f:
                 surveys = []
                 fastavro.writer(f, SURVEY_SCHEMA, surveys)
+
+        if self.enable_pgsql:
+            if not issubclass(type(self.agents[0]), InstitutionAgent):
+                profiles: list[Any] = []
+                for agent in self.agents:
+                    profile = await agent.memory._profile.export()
+                    profile = profile[0]
+                    profile["id"] = agent._uuid
+                    profiles.append(
+                        (agent._uuid, profile.get("name", ""), json.dumps(profile))
+                    )
+                await self._pgsql_writer.async_write_profile.remote(  # type:ignore
+                    profiles
+                )
         self.initialized = True
         logger.debug(f"-----AgentGroup {self._uuid} initialized")
 
@@ -218,11 +239,13 @@ class AgentGroup:
             await asyncio.sleep(0.5)
 
     async def save_status(self):
+        _statuses_time_list: list[tuple[dict, datetime]] = []
         if self.enable_avro:
             logger.debug(f"-----Saving status for group {self._uuid}")
             avros = []
             if not issubclass(type(self.agents[0]), InstitutionAgent):
                 for agent in self.agents:
+                    _date_time = datetime.now(timezone.utc)
                     position = await agent.memory.get("position")
                     lng = position["longlat_position"]["longitude"]
                     lat = position["longlat_position"]["latitude"]
@@ -248,13 +271,15 @@ class AgentGroup:
                         "tired": needs["tired"],
                         "safe": needs["safe"],
                         "social": needs["social"],
-                        "created_at": int(datetime.now().timestamp() * 1000),
+                        "created_at": int(_date_time.timestamp() * 1000),
                     }
                     avros.append(avro)
+                    _statuses_time_list.append((avro, _date_time))
                 with open(self.avro_file["status"], "a+b") as f:
                     fastavro.writer(f, STATUS_SCHEMA, avros, codec="snappy")
             else:
                 for agent in self.agents:
+                    _date_time = datetime.now(timezone.utc)
                     avro = {
                         "id": agent._uuid,
                         "day": await self.simulator.get_simulator_day(),
@@ -274,8 +299,109 @@ class AgentGroup:
                         "customers": await agent.memory.get("customers"),
                     }
                     avros.append(avro)
+                    _statuses_time_list.append((avro, _date_time))
                 with open(self.avro_file["status"], "a+b") as f:
                     fastavro.writer(f, INSTITUTION_STATUS_SCHEMA, avros, codec="snappy")
+        if self.enable_pgsql:
+            # data already acquired from Avro part
+            if len(_statuses_time_list) > 0:
+                for _status_dict, _date_time in _statuses_time_list:
+                    for key in ["lng", "lat", "parent_id"]:
+                        if key not in _status_dict:
+                            _status_dict[key] = -1
+                    for key in [
+                        "action",
+                    ]:
+                        if key not in _status_dict:
+                            _status_dict[key] = ""
+                    _status_dict["created_at"] = _date_time
+            else:
+                if not issubclass(type(self.agents[0]), InstitutionAgent):
+                    for agent in self.agents:
+                        _date_time = datetime.now(timezone.utc)
+                        position = await agent.memory.get("position")
+                        lng = position["longlat_position"]["longitude"]
+                        lat = position["longlat_position"]["latitude"]
+                        if "aoi_position" in position:
+                            parent_id = position["aoi_position"]["aoi_id"]
+                        elif "lane_position" in position:
+                            parent_id = position["lane_position"]["lane_id"]
+                        else:
+                            # BUG: 需要处理
+                            parent_id = -1
+                        needs = await agent.memory.get("needs")
+                        action = await agent.memory.get("current_step")
+                        action = action["intention"]
+                        _status_dict = {
+                            "id": agent._uuid,
+                            "day": await self.simulator.get_simulator_day(),
+                            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                            "lng": lng,
+                            "lat": lat,
+                            "parent_id": parent_id,
+                            "action": action,
+                            "hungry": needs["hungry"],
+                            "tired": needs["tired"],
+                            "safe": needs["safe"],
+                            "social": needs["social"],
+                            "created_at": _date_time,
+                        }
+                        _statuses_time_list.append((_status_dict, _date_time))
+                else:
+                    for agent in self.agents:
+                        _date_time = datetime.now(timezone.utc)
+                        _status_dict = {
+                            "id": agent._uuid,
+                            "day": await self.simulator.get_simulator_day(),
+                            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+                            "lng": -1,
+                            "lat": -1,
+                            "parent_id": -1,
+                            "action": "",
+                            "type": await agent.memory.get("type"),
+                            "nominal_gdp": await agent.memory.get("nominal_gdp"),
+                            "real_gdp": await agent.memory.get("real_gdp"),
+                            "unemployment": await agent.memory.get("unemployment"),
+                            "wages": await agent.memory.get("wages"),
+                            "prices": await agent.memory.get("prices"),
+                            "inventory": await agent.memory.get("inventory"),
+                            "price": await agent.memory.get("price"),
+                            "interest_rate": await agent.memory.get("interest_rate"),
+                            "bracket_cutoffs": await agent.memory.get(
+                                "bracket_cutoffs"
+                            ),
+                            "bracket_rates": await agent.memory.get("bracket_rates"),
+                            "employees": await agent.memory.get("employees"),
+                            "customers": await agent.memory.get("customers"),
+                            "created_at": _date_time,
+                        }
+                        _statuses_time_list.append((_status_dict, _date_time))
+        to_update_statues: list[tuple] = []
+        for _status_dict, _ in _statuses_time_list:
+            BASIC_KEYS = [
+                "id",
+                "day",
+                "t",
+                "lng",
+                "lat",
+                "parent_id",
+                "action",
+                "created_at",
+            ]
+            _data = [_status_dict[k] for k in BASIC_KEYS if k != "created_at"]
+            _other_dict = json.dumps(
+                {k: v for k, v in _status_dict.items() if k not in BASIC_KEYS}
+            )
+            _data.append(_other_dict)
+            _data.append(_status_dict["created_at"])
+            to_update_statues.append(tuple(_data))
+        if self._last_asyncio_pg_task is not None:
+            await self._last_asyncio_pg_task
+        self._last_asyncio_pg_task = (
+            self._pgsql_writer.async_write_status.remote(  # type:ignore
+                to_update_statues
+            )
+        )
 
     async def step(self):
         if not self.initialized:

pycityagent/simulation/simulation.py

@@ -23,6 +23,7 @@ from ..message.messager import Messager
 from ..metrics import init_mlflow_connection
 from ..survey import Survey
 from .agentgroup import AgentGroup
+from .storage.pg import PgWriter, create_pg_tables
 
 logger = logging.getLogger("pycityagent")
 
@@ -63,6 +64,7 @@ class AgentSimulation:
         self._user_survey_topics: dict[uuid.UUID, str] = {}
         self._user_interview_topics: dict[uuid.UUID, str] = {}
         self._loop = asyncio.get_event_loop()
+        # self._last_asyncio_pg_task = None  # 将SQL写入的IO隐藏到计算任务后
 
         self._messager = Messager(
             hostname=config["simulator_request"]["mqtt"]["server"],
@@ -89,22 +91,13 @@ class AgentSimulation:
         self._enable_pgsql = _pgsql_config.get("enabled", False)
         if not self._enable_pgsql:
             logger.warning("PostgreSQL is not enabled, NO POSTGRESQL DATABASE STORAGE")
-            self._pgsql_args = ("", "", "", "", "")
+            self._pgsql_dsn = ""
         else:
-            self._pgsql_host = _pgsql_config["host"]
-            self._pgsql_port = _pgsql_config["port"]
-            self._pgsql_database = _pgsql_config["database"]
-            self._pgsql_user = _pgsql_config.get("user", None)
-            self._pgsql_password = _pgsql_config.get("password", None)
-            self._pgsql_args: tuple[str, str, str, str, str] = (
-                self._pgsql_host,
-                self._pgsql_port,
-                self._pgsql_database,
-                self._pgsql_user,
-                self._pgsql_password,
-            )
+            self._pgsql_dsn = _pgsql_config["data_source_name"]
 
         # 添加实验信息相关的属性
+        self._exp_created_time = datetime.now(timezone.utc)
+        self._exp_updated_time = datetime.now(timezone.utc)
         self._exp_info = {
             "id": self.exp_id,
             "name": exp_name,
@@ -114,7 +107,8 @@ class AgentSimulation:
             "cur_t": 0.0,
             "config": json.dumps(config),
             "error": "",
-            "created_at": datetime.now(timezone.utc).isoformat(),
+            "created_at": self._exp_created_time.isoformat(),
+            "updated_at": self._exp_updated_time.isoformat(),
         }
 
         # 创建异步任务保存实验信息
@@ -168,7 +162,7 @@ class AgentSimulation:
         enable_avro: bool,
         avro_path: Path,
         enable_pgsql: bool,
-        pgsql_copy_writer: ray.ObjectRef,
+        pgsql_writer: ray.ObjectRef,
         mlflow_run_id: str = None,  # type: ignore
         logging_level: int = logging.WARNING,
     ):
@@ -181,7 +175,7 @@ class AgentSimulation:
             enable_avro,
             avro_path,
             enable_pgsql,
-            pgsql_copy_writer,
+            pgsql_writer,
             mlflow_run_id,
             logging_level,
         )
@@ -191,6 +185,7 @@ class AgentSimulation:
         self,
         agent_count: Union[int, list[int]],
         group_size: int = 1000,
+        pg_sql_writers: int = 32,
         memory_config_func: Optional[Union[Callable, list[Callable]]] = None,
     ) -> None:
         """初始化智能体
@@ -251,8 +246,8 @@ class AgentSimulation:
                     memory=memory,
                 )
 
-                self._agents[agent._uuid] = agent
-                self._agent_uuids.append(agent._uuid)
+                self._agents[agent._uuid] = agent  # type:ignore
+                self._agent_uuids.append(agent._uuid)  # type:ignore
 
             # 计算需要的组数,向上取整以处理不足一组的情况
             num_group = (agent_count_i + group_size - 1) // group_size
@@ -282,9 +277,23 @@ class AgentSimulation:
             )
         else:
             mlflow_run_id = None
+        # 建表
+        if self.enable_pgsql:
+            _num_workers = min(1, pg_sql_writers)
+            create_pg_tables(
+                exp_id=self.exp_id,
+                dsn=self._pgsql_dsn,
+            )
+            self._pgsql_writers = _workers = [
+                PgWriter.remote(self.exp_id, self._pgsql_dsn)
+                for _ in range(_num_workers)
+            ]
+        else:
+            _num_workers = 1
+            self._pgsql_writers = _workers = [None for _ in range(_num_workers)]
         # 收集所有创建组的参数
         creation_tasks = []
-        for group_name, agents in group_creation_params:
+        for i, (group_name, agents) in enumerate(group_creation_params):
             # 直接创建异步任务
             group = AgentGroup.remote(
                 agents,
@@ -294,10 +303,8 @@ class AgentSimulation:
                 self.enable_avro,
                 self.avro_path,
                 self.enable_pgsql,
-                # TODO:
-                # self._pgsql_copy_writer, # type:ignore
-                None,
-                mlflow_run_id,
+                _workers[i % _num_workers],  # type:ignore
+                mlflow_run_id,  # type:ignore
                 self.logging_level,
             )
             creation_tasks.append((group_name, group, agents))
@@ -469,11 +476,13 @@ class AgentSimulation:
         survey_dict = survey.to_dict()
         if agent_uuids is None:
             agent_uuids = self._agent_uuids
+        _date_time = datetime.now(timezone.utc)
         payload = {
             "from": "none",
             "survey_id": survey_dict["id"],
-            "timestamp": int(datetime.now().timestamp() * 1000),
+            "timestamp": int(_date_time.timestamp() * 1000),
             "data": survey_dict,
+            "_date_time": _date_time,
         }
         for uuid in agent_uuids:
             topic = self._user_survey_topics[uuid]
@@ -483,10 +492,12 @@ class AgentSimulation:
         self, content: str, agent_uuids: Union[uuid.UUID, list[uuid.UUID]]
     ):
         """发送面试消息"""
+        _date_time = datetime.now(timezone.utc)
         payload = {
             "from": "none",
             "content": content,
-            "timestamp": int(datetime.now().timestamp() * 1000),
+            "timestamp": int(_date_time.timestamp() * 1000),
+            "_date_time": _date_time,
         }
         if not isinstance(agent_uuids, Sequence):
             agent_uuids = [agent_uuids]
@@ -515,15 +526,29 @@ class AgentSimulation:
             logger.error(f"Avro保存实验信息失败: {str(e)}")
         try:
             if self.enable_pgsql:
-                #   TODO
-                pass
+                worker: ray.ObjectRef = self._pgsql_writers[0]  # type:ignore
+                # if self._last_asyncio_pg_task is not None:
+                #     await self._last_asyncio_pg_task
+                # self._last_asyncio_pg_task = (
+                #     worker.async_update_exp_info.remote(  # type:ignore
+                #         pg_exp_info
+                #     )
+                # )
+                pg_exp_info = {k: v for k, v in self._exp_info.items()}
+                pg_exp_info["created_at"] = self._exp_created_time
+                pg_exp_info["updated_at"] = self._exp_updated_time
+                await worker.async_update_exp_info.remote(  # type:ignore
+                    pg_exp_info
+                )
         except Exception as e:
             logger.error(f"PostgreSQL保存实验信息失败: {str(e)}")
 
     async def _update_exp_status(self, status: int, error: str = "") -> None:
+        self._exp_updated_time = datetime.now(timezone.utc)
         """更新实验状态并保存"""
         self._exp_info["status"] = status
         self._exp_info["error"] = error
+        self._exp_info["updated_at"] = self._exp_updated_time.isoformat()
         await self._save_exp_info()
 
     async def _monitor_exp_status(self, stop_event: asyncio.Event):