Merge pull request #651 from AnguseZhang/fix/unsupported_archive

feat: enhance logging and environment flexibility by adding session_i…

Author: Rasic2

agents/matmaster_agent/constant.py

@@ -53,11 +53,13 @@
 if CURRENT_ENV == 'test':
     DFLOW_HOST = 'https://lbg-workflow-mlops.test.dp.tech'
     DFLOW_K8S_API_SERVER = 'https://lbg-workflow-mlops.test.dp.tech'
+    TIEFBLUE_NAS_HOST = 'https://tiefblue-nas-acs-bj.test.bohrium.com'
 elif CURRENT_ENV == 'uat':
-    pass
+    TIEFBLUE_NAS_HOST = 'https://tiefblue-nas-acs-bj.test.bohrium.com'
 elif CURRENT_ENV == 'prod':
     DFLOW_HOST = 'https://workflows.deepmodeling.com'
     DFLOW_K8S_API_SERVER = 'https://workflows.deepmodeling.com'
+    TIEFBLUE_NAS_HOST = 'https://tiefblue-nas-acs-bj.bohrium.com'
 
 # API URL
 OPENAPI_SANDBOX = f'{OPENAPI_HOST}/openapi/v1/sandbox'

agents/matmaster_agent/core_agents/public_agents/job_agents/result_core_agent/agent.py

@@ -141,7 +141,9 @@ async def _run_events(self, ctx: InvocationContext) -> AsyncGenerator[Event, Non
                     logger.info(f"{ctx.session.id} dict_result = {dict_result}")
 
                     if self.enable_tgz_unpack:
-                        tgz_flag, new_tool_result = await update_tgz_dict(dict_result)
+                        tgz_flag, new_tool_result = await update_tgz_dict(
+                            dict_result, session_id=ctx.session.id
+                        )
                     else:
                         new_tool_result = dict_result
                     parsed_tool_result = await parse_result(ctx, new_tool_result)

agents/matmaster_agent/services/job.py

@@ -11,6 +11,7 @@
     MATMASTER_AGENT_NAME,
     OPENAPI_FILE_TOKEN_API,
     OPENAPI_HOST,
+    TIEFBLUE_NAS_HOST,
     OpenAPIJobAPI,
 )
 from agents.matmaster_agent.logger import PrefixFilter
@@ -153,7 +154,7 @@ async def get_token_and_download_dir(dir_path, job_id, access_key):
     }
     async with aiohttp.ClientSession() as session:
         async with session.post(
-            'https://tiefblue-nas-acs-bj.test.bohrium.com/api/downloadr',
+            f'{TIEFBLUE_NAS_HOST}/api/downloadr',
             json=request_body,
             headers={
                 'Authorization': f"Bearer {token}",
@@ -197,7 +198,9 @@ async def get_iterate_files(
     return prefix, iterate_json
 
 
-async def parse_and_prepare_results(job_id: str = '', access_key: str = ''):
+async def parse_and_prepare_results(
+    job_id: str = '', access_key: str = '', session_id: str = ''
+):
     def norm(p: str) -> str:
         p = p.replace('\\', '/')
         if p.endswith('/') and p != '/':
@@ -294,7 +297,7 @@ async def safe_cleanup_file(file_download_path: Path):
         results_txt_parsed = jsonpickle.loads(
             f.read().replace('pathlib._local.PosixPath', 'pathlib.PosixPath')
         )
-        logger.info(f"results_txt_parsed = {results_txt_parsed}")
+        logger.info(f"{session_id} results_txt_parsed = {results_txt_parsed}")
     os.remove(RESULTS_TXT)
 
     final_results = {}
@@ -318,7 +321,7 @@ async def safe_cleanup_file(file_download_path: Path):
         obj = await find_obj(rel)
 
         if obj is None:
-            logger.warning(f"`{rel}` is not exist")
+            logger.warning(f"{session_id} `{rel}` is not exist")
             continue
 
         is_dir = bool(obj.get('isDir'))

agents/matmaster_agent/utils/io_oss.py

@@ -162,13 +162,14 @@ async def upload_base64_to_oss(data: str, oss_path: str) -> dict:
 
 
 async def extract_convert_and_upload(
-    compressed_url: str, temp_dir_path: str = './tmp'
+    compressed_url: str, temp_dir_path: str = './tmp', session_id: str = ''
 ) -> dict:
     """
     下载 TGZ → 解压 → JPG 转 Base64 → 上传 OSS → 自动清理
     """
     async with temp_dir(temp_dir_path) as temp_path:
         # 获取所有JPG文件的Base64编码
+        logger.info(f"{session_id} compressed_url = {compressed_url}")
         files = await extract_files_from_compressed_file_url(compressed_url, temp_path)
         conversion_tasks = [file_to_base64(file) for file in files]
         base64_results = await asyncio.gather(*conversion_tasks)
@@ -187,14 +188,16 @@ async def extract_convert_and_upload(
         }
 
 
-async def update_tgz_dict(tool_result: dict):
+async def update_tgz_dict(tool_result: dict, session_id: str = ''):
     new_tool_result = {}
     compressed_flag = False
     for k, v in tool_result.items():
         new_tool_result[k] = v
         if isinstance(v, str) and v.startswith('https') and v.endswith(('tgz', 'zip')):
             compressed_flag = True
-            new_tool_result.update(**await extract_convert_and_upload(v))
+            new_tool_result.update(
+                **await extract_convert_and_upload(v, session_id=session_id)
+            )
 
     return compressed_flag, new_tool_result
 
@@ -212,6 +215,7 @@ async def extract_file_content(file_url: str, temp_dir_path: str = './tmp') -> d
 
 
 if __name__ == '__main__':
-    invalid_zip = 'https://bohrium-agent-test.oss-cn-zhangjiakou.aliyuncs.com/agent/jobs/337612/72f4337ea51448edacabf9559bc630b3/outputs.zip'
+    invalid_zip = 'https://bohrium-agent-test.oss-cn-zhangjiakou.aliyuncs.com/agent/jobs/337612/fc6c0cc25d8847c8acb226c097557ed7/outputs.zip'
     valid_zip = 'https://bohrium-agent-test.oss-cn-zhangjiakou.aliyuncs.com/agent/jobs/110680/c15113bb106b46f1a49800e776a2b8fc/outputs.zip'
-    asyncio.run(extract_convert_and_upload(invalid_zip, './tmp'))
+    failed_zip = 'https://bohrium-agents.oss-cn-zhangjiakou.aliyuncs.com/agent/jobs/337612/263edcfa12a2460abd222f6deeb19969/outputs.zip'
+    asyncio.run(extract_convert_and_upload(failed_zip, './tmp'))

scripts/sandbox_cli.py

@@ -1,10 +1,12 @@
 import argparse
+import asyncio
 import json
 import os
 import sys
 import time
 from datetime import datetime
 
+import aiohttp
 import jsonpickle
 import requests
 from dotenv import find_dotenv, load_dotenv
@@ -100,7 +102,7 @@ def poll_job_status(job_id, interval=10):
             time.sleep(interval)
 
 
-def main():
+async def main():
     parser = argparse.ArgumentParser(description='Sandbox job cli')
     subparsers = parser.add_subparsers(
         dest='command', help='Available commands', required=True
@@ -138,7 +140,7 @@ def main():
     )
 
     args = parser.parse_args()
-
+    access_key = os.getenv('MATERIALS_ACCESS_KEY')
     # 调用API获取job信息
     logger.info(f"Job ID: {args.job_id}")
 
@@ -150,7 +152,7 @@ def main():
 
         try:
             response = requests.get(
-                f"{OpenAPIJobAPI}/{args.job_id}?accessKey={os.getenv('BOHRIUM_ACCESS_KEY')}"
+                f"{OpenAPIJobAPI}/{args.job_id}?accessKey={access_key}"
             )
             response.raise_for_status()
             job_info = response.json()
@@ -181,23 +183,26 @@ def main():
         logger.info(f"{job_name}[{job_status}] -- {duration}")
 
         # download log
-        get_token_and_download_file('log', args.job_id)
+        await get_token_and_download_file('log', args.job_id, access_key)
 
         if job_status in ['Running']:
             return
         elif job_status == 'Finished':
             # download result.txt
             results_txt = 'results.txt'
-            get_token_and_download_file(results_txt, args.job_id)
+            await get_token_and_download_file(results_txt, args.job_id, access_key)
             with open(results_txt) as f:
                 logger.info(jsonpickle.loads(f.read()))
             os.remove(results_txt)
 
             if args.download_output:
                 # 下载结果文件
                 if result_url and result_url != 'null':
-                    result_response = requests.get(result_url, stream=True)
-                    check_status_and_download_file(result_response, args.output)
+                    async with aiohttp.ClientSession() as session:
+                        async with session.get(result_url) as result_response:
+                            await check_status_and_download_file(
+                                result_response, args.output
+                            )
                 else:
                     logger.error('No resultUrl found or resultUrl is empty')
     elif args.command == 'kill':
@@ -207,4 +212,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    asyncio.run(main())