bohr-agent-sdk 0.1.19__py3-none-any.whl → 0.1.21__py3-none-any.whl

{bohr_agent_sdk-0.1.19.dist-info → bohr_agent_sdk-0.1.21.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.19
+Version: 0.1.21
 Summary: SDK for scientific agents
 Home-page: https://github.com/dptech-corp/bohr-agent-sdk/
 Author: DP Technology

{bohr_agent_sdk-0.1.19.dist-info → bohr_agent_sdk-0.1.21.dist-info}/RECORD

@@ -20,21 +20,21 @@ dp/agent/device/device/__init__.py,sha256=w7_1S16S1vWUq0RGl0GFgjq2vFkc5oNvy8cQTn
 dp/agent/device/device/device.py,sha256=9ZRIJth-4qMO-i-u_b_cO3d6a4eTbTQjPaxFsV_zEkc,9643
 dp/agent/device/device/types.py,sha256=JuxB-hjf1CjjvfBxCLwRAXVFlYS-nPEdiJpBWLFVCzo,1924
 dp/agent/server/__init__.py,sha256=rckaYd8pbYyB4ENEhgjXKeGMXjdnrgcJpdM1gu5u1Wc,508
-dp/agent/server/calculation_mcp_server.py,sha256=omPcrjVT3k8ek3BbB5Dc3TRjdSC_0IEqjCrK2Ka0YqE,10888
+dp/agent/server/calculation_mcp_server.py,sha256=oumU2xtqrrgp8BKiuF9YsqMNevXba2JACyYsvNy4nQo,11008
 dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok8A,1264
 dp/agent/server/utils.py,sha256=8jgYZEW4XBp86AF2Km6QkwHltBmrnS-soTpHov7ZEJw,4501
 dp/agent/server/executor/__init__.py,sha256=s95M5qKQk39Yi9qaVJZhk_nfj54quSf7EDghR3OCFUA,248
 dp/agent/server/executor/base_executor.py,sha256=EFJBsYVYAvuRbiLAbLOwLTw3h7ScjN025xnSP4uJHrQ,2052
-dp/agent/server/executor/dispatcher_executor.py,sha256=zp9RDUoT2RYZGbKXsUsk4HbS0lyM9IEt1zTwOTCQ1YY,9239
-dp/agent/server/executor/local_executor.py,sha256=alNpEPicZQOQMtXc1vqGSo4El_9agXAJ5ldqi-OJGL4,3887
+dp/agent/server/executor/dispatcher_executor.py,sha256=urpzmKH_tBOgblBdJEa3y8eEhXqUDrdcdWCnUdJpfZk,9420
+dp/agent/server/executor/local_executor.py,sha256=9JEtZTEbKsmHAY5nsdLZ36gH5M_rSG6En0W9IrCUwm0,5803
 dp/agent/server/storage/__init__.py,sha256=Sgsyp5hb0_hhIGugAPfQFzBHt_854rS_MuMuE3sn8Gs,389
 dp/agent/server/storage/base_storage.py,sha256=728-oNG6N8isV95gZVnyi4vTznJPJhSjxw9Gl5Y_y5o,2356
 dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mmPFecrln-Ytk40,10488
 dp/agent/server/storage/http_storage.py,sha256=KiySq7g9-iJr12XQCKKyJLn8wJoDnSRpQAR5_qPJ1ZU,1471
 dp/agent/server/storage/local_storage.py,sha256=t1wfjByjXew9ws3PuUxWxmZQ0-Wt1a6t4wmj3fW62GI,1352
 dp/agent/server/storage/oss_storage.py,sha256=pgjmi7Gir3Y5wkMDCvU4fvSls15fXT7Ax-h9MYHFPK0,3359
-bohr_agent_sdk-0.1.19.dist-info/METADATA,sha256=tDUixjAnwCjzdtcuXXc-71c2ZhUiJgKmRlXVMQguDEE,6329
-bohr_agent_sdk-0.1.19.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-bohr_agent_sdk-0.1.19.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
-bohr_agent_sdk-0.1.19.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
-bohr_agent_sdk-0.1.19.dist-info/RECORD,,
+bohr_agent_sdk-0.1.21.dist-info/METADATA,sha256=nHfB7MVnAuzZs1EUND-29vdVA3KDSzpuJmQxTswFT_I,6329
+bohr_agent_sdk-0.1.21.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.21.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.21.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.21.dist-info/RECORD,,

dp/agent/server/calculation_mcp_server.py

@@ -29,7 +29,7 @@ def parse_uri(uri):
 
 
 def init_storage(storage_config: Optional[dict] = None):
-    if storage_config is None:
+    if not storage_config:
         storage_config = {"type": "local"}
     storage_config = deepcopy(storage_config)
     storage_type = storage_config.pop("type")
@@ -38,7 +38,7 @@ def init_storage(storage_config: Optional[dict] = None):
 
 
 def init_executor(executor_config: Optional[dict] = None):
-    if executor_config is None:
+    if not executor_config:
         executor_config = {"type": "local"}
     executor_config = deepcopy(executor_config)
     executor_type = executor_config.pop("type")
@@ -227,7 +227,11 @@ class CalculationMCPServer:
                     exec_id = res["job_id"]
                     job_id = "%s/%s" % (trace_id, exec_id)
                     logger.info("Job submitted (ID: %s)" % job_id)
-                return convert_to_content({"job_id": job_id}, job_info={
+                result = {
+                    "job_id": job_id,
+                    "extra_info": res.get("extra_info"),
+                }
+                return convert_to_content(result, job_info={
                     "trace_id": trace_id,
                     "executor_type": executor_type,
                     "job_id": job_id,

dp/agent/server/executor/dispatcher_executor.py

@@ -12,12 +12,15 @@ from dpdispatcher import Machine, Resources, Task, Submission
 
 from .base_executor import BaseExecutor
 from .... import __path__
+
 config = {
     "username": os.environ.get("BOHRIUM_USERNAME", ""),
     "password": os.environ.get("BOHRIUM_PASSWORD", ""),
     "project_id": os.environ.get("BOHRIUM_PROJECT_ID", ""),
     "access_key": os.environ.get("BOHRIUM_ACCESS_KEY", ""),
     "app_key": os.environ.get("BOHRIUM_APP_KEY", "agent"),
+    "bohrium_url": os.environ.get("BOHRIUM_BOHRIUM_URL",
+                                  "https://bohrium.dp.tech"),
 }
 logger = logging.getLogger(__name__)
 
@@ -26,17 +29,17 @@ def get_source_code(fn):
     source_lines, start_line = inspect.getsourcelines(fn)
     source_file = inspect.getsourcefile(fn)
     with open(source_file, "r", encoding="utf-8") as fd:
-        pre_lines = fd.readlines()[:start_line-1]
+        pre_lines = fd.readlines()[:start_line - 1]
     return "".join(pre_lines + source_lines) + "\n"
 
 
 class DispatcherExecutor(BaseExecutor):
     def __init__(
-        self,
-        machine=None,
-        resources=None,
-        python_packages=None,
-        python_executable="python3",
+            self,
+            machine=None,
+            resources=None,
+            python_packages=None,
+            python_executable="python3",
     ):
         """Use DPDispatcher to execute the tool
         Refer to https://docs.deepmodeling.com/projects/dpdispatcher.
@@ -107,7 +110,7 @@ class DispatcherExecutor(BaseExecutor):
         script += "if __name__ == \"__main__\":\n"
         script += "    cwd = os.getcwd()\n"
         script += "    kwargs = jsonpickle.loads(r'''%s''')\n" % \
-            jsonpickle.dumps(kwargs)
+                  jsonpickle.dumps(kwargs)
         script += "    try:\n"
         if import_func_line is not None:
             script += "        " + import_func_line
@@ -118,8 +121,8 @@ class DispatcherExecutor(BaseExecutor):
             script += "        if isinstance(results, dict):\n"
             script += "            for name in results:\n"
             script += "                if isinstance(results[name], Path):\n"
-            script += "                    results[name] = "\
-                "results[name].absolute().relative_to(cwd)\n"
+            script += "                    results[name] = " \
+                      "results[name].absolute().relative_to(cwd)\n"
         script += "    except Exception as e:\n"
         script += "        os.chdir(cwd)\n"
         script += "        with open('err', 'w') as f:\n"
@@ -173,8 +176,8 @@ class DispatcherExecutor(BaseExecutor):
             extra_info = {
                 "bohr_job_id": bohr_job_id,
                 "bohr_group_id": bohr_group_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info
@@ -182,8 +185,8 @@ class DispatcherExecutor(BaseExecutor):
             bohr_job_id = submission.belonging_jobs[0].job_id
             extra_info = {
                 "bohr_job_id": bohr_job_id,
-                "job_link": "https://bohrium.dp.tech/jobs/detail/%s" %
-                bohr_job_id,
+                "job_link": f"{config['bohrium_url']}/jobs/detail/%s" %
+                            bohr_job_id,
             }
             logger.info(extra_info)
             res["extra_info"] = extra_info

dp/agent/server/executor/local_executor.py

@@ -1,19 +1,43 @@
 import asyncio
 import importlib
 import inspect
+import io
 import jsonpickle
 import os
 import psutil
+import re
 import sys
+import time
 import uuid
 from multiprocessing import Process
 from typing import Dict, Optional
 
 from .base_executor import BaseExecutor
+from ..utils import get_logger
 
+DFLOW_ID_PATTERN = r"Workflow has been submitted \(ID: ([^,]*), UID: ([^)]*)\)"
+DFLOW_LINK_PATTERN = r"Workflow link: (.*)"
+logger = get_logger(__name__)
 
-def wrapped_fn(fn, kwargs):
+
+class Tee(io.TextIOBase):
+    def __init__(self, file, stdout):
+        self.file = file
+        self.stdout = stdout
+
+    def write(self, text):
+        self.stdout.write(text)
+        self.file.write(text)
+        self.file.flush()
+        return len(text)
+
+
+def wrapped_fn(fn, kwargs, redirect_file=None):
     pid = os.getpid()
+    if redirect_file:
+        stdout = sys.stdout
+        flog = open(redirect_file, "w")
+        sys.stdout = Tee(flog, stdout)
     try:
         if inspect.iscoroutinefunction(fn):
             result = asyncio.run(fn(**kwargs))
@@ -23,6 +47,10 @@ def wrapped_fn(fn, kwargs):
         with open("err", "w") as f:
             f.write(str(e))
         raise e
+    finally:
+        if redirect_file:
+            sys.stdout = stdout
+            flog.close()
     with open("%s.txt" % pid, "w") as f:
         f.write(jsonpickle.dumps(result))
 
@@ -47,13 +75,16 @@ def reload_dflow_config():
 
 
 class LocalExecutor(BaseExecutor):
-    def __init__(self, env: Optional[Dict[str, str]] = None):
+    def __init__(self, env: Optional[Dict[str, str]] = None,
+                 dflow: bool = False):
         """
         Execute the tool locally
         Args:
             env: The environmental variables at run time
+            dflow: Wait until workflow submitted in submit method
         """
         self.env = env or {}
+        self.dflow = dflow
 
     def set_env(self):
         old_env = {}
@@ -73,10 +104,34 @@ class LocalExecutor(BaseExecutor):
     def submit(self, fn, kwargs):
         os.environ["DP_AGENT_RUNNING_MODE"] = "1"
         old_env = self.set_env()
-        p = Process(target=wrapped_fn, kwargs={"fn": fn, "kwargs": kwargs})
+        params = {"fn": fn, "kwargs": kwargs}
+        if self.dflow:
+            params["redirect_file"] = "log.txt"
+        p = Process(target=wrapped_fn, kwargs=params)
         p.start()
+        extra_info = {}
+        if self.dflow:
+            while True:
+                alive = p.is_alive()
+                if os.path.isfile("log.txt"):
+                    with open("log.txt", "r") as f:
+                        log = f.read()
+                    match_id = re.search(DFLOW_ID_PATTERN, log)
+                    match_link = re.search(DFLOW_LINK_PATTERN, log)
+                    if match_id and match_link:
+                        wf_id = match_id.group(1)
+                        wf_uid = match_id.group(2)
+                        wf_link = match_link.group(1)
+                        extra_info["workflow_id"] = wf_id
+                        extra_info["workflow_uid"] = wf_uid
+                        extra_info["workflow_link"] = wf_link
+                        break
+                if not alive:
+                    break
+                logger.info("Waiting workflow to be submitted")
+                time.sleep(1)
         self.recover_env(old_env)
-        return {"job_id": str(p.pid)}
+        return {"job_id": str(p.pid), "extra_info": extra_info}
 
     def query_status(self, job_id):
         try: