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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bohr-agent-sdk
-Version: 0.1.19
+Version: 0.1.20
 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.20.dist-info}/RECORD

@@ -25,7 +25,7 @@ dp/agent/server/preprocessor.py,sha256=XUWu7QOwo_sIDMYS2b1OTrM33EXEVH_73vk-ju1Ok
 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/dispatcher_executor.py,sha256=urpzmKH_tBOgblBdJEa3y8eEhXqUDrdcdWCnUdJpfZk,9420
 dp/agent/server/executor/local_executor.py,sha256=alNpEPicZQOQMtXc1vqGSo4El_9agXAJ5ldqi-OJGL4,3887
 dp/agent/server/storage/__init__.py,sha256=Sgsyp5hb0_hhIGugAPfQFzBHt_854rS_MuMuE3sn8Gs,389
 dp/agent/server/storage/base_storage.py,sha256=728-oNG6N8isV95gZVnyi4vTznJPJhSjxw9Gl5Y_y5o,2356
@@ -33,8 +33,8 @@ dp/agent/server/storage/bohrium_storage.py,sha256=EsKX4dWWvZTn2TEhZv4zsvihfDK0mm
 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.20.dist-info/METADATA,sha256=ZEoxrydj541JG9txn0Wr0iARoxZ_w2KA6E0aJFFGlAY,6329
+bohr_agent_sdk-0.1.20.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+bohr_agent_sdk-0.1.20.dist-info/entry_points.txt,sha256=5n5kneF5IbDQtoQ2WfF-QuBjDtsimJte9Rv9baSGgc0,86
+bohr_agent_sdk-0.1.20.dist-info/top_level.txt,sha256=87xLUDhu_1nQHoGLwlhJ6XlO7OsjILh6i1nX6ljFzDo,3
+bohr_agent_sdk-0.1.20.dist-info/RECORD,,

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