Skip to content

Commit 9c2f9cb

Browse files
dingzhaohandp-dingzhaohanpre-commit-ci[bot]njzjz
authored
update bohrium-sdk version 0.6.0 (#520)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added utility functions for creating and extracting zip files using standard Python tools. - **Bug Fixes** - Improved error handling for missing credentials by validating environment variables and profile settings. - **Refactor** - Updated integration to use the latest Bohrium SDK. - Replaced custom zip utilities with standard Python implementations for file compression and extraction. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: dingzhaohan <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Jinzhe Zeng <[email protected]>
1 parent d81c9f6 commit 9c2f9cb

File tree

2 files changed

+105
-19
lines changed

2 files changed

+105
-19
lines changed

‎dpdispatcher/contexts/openapi_context.py

Lines changed: 67 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,20 @@
1+
import glob
12
import os
23
import shutil
34
import uuid
5+
from zipfile import ZipFile
46

57
import tqdm
68

79
try:
8-
from bohriumsdk.client import Client
9-
from bohriumsdk.job import Job
10-
from bohriumsdk.storage import Storage
11-
from bohriumsdk.util import Util
12-
except ModuleNotFoundError:
10+
from bohrium import Bohrium
11+
from bohrium.resources import Job, Tiefblue
12+
except ModuleNotFoundError as e:
1313
found_bohriumsdk = False
14+
import_bohrium_error = e
1415
else:
1516
found_bohriumsdk = True
17+
import_bohrium_error = None
1618

1719
from dpdispatcher.base_context import BaseContext
1820
from dpdispatcher.dlog import dlog
@@ -23,6 +25,36 @@
2325
)
2426

2527

28+
def unzip_file(zip_file, out_dir="./"):
29+
obj = ZipFile(zip_file, "r")
30+
for item in obj.namelist():
31+
obj.extract(item, out_dir)
32+
33+
34+
def zip_file_list(root_path, zip_filename, file_list=[]):
35+
out_zip_file = os.path.join(root_path, zip_filename)
36+
# print('debug: file_list', file_list)
37+
zip_obj = ZipFile(out_zip_file, "w")
38+
for f in file_list:
39+
matched_files = os.path.join(root_path, f)
40+
for ii in glob.glob(matched_files):
41+
# print('debug: matched_files:ii', ii)
42+
if os.path.isdir(ii):
43+
arcname = os.path.relpath(ii, start=root_path)
44+
zip_obj.write(ii, arcname)
45+
for root, dirs, files in os.walk(ii):
46+
for file in files:
47+
filename = os.path.join(root, file)
48+
arcname = os.path.relpath(filename, start=root_path)
49+
# print('debug: filename:arcname:root_path', filename, arcname, root_path)
50+
zip_obj.write(filename, arcname)
51+
else:
52+
arcname = os.path.relpath(ii, start=root_path)
53+
zip_obj.write(ii, arcname)
54+
zip_obj.close()
55+
return out_zip_file
56+
57+
2658
class OpenAPIContext(BaseContext):
2759
def __init__(
2860
self,
@@ -35,15 +67,39 @@ def __init__(
3567
if not found_bohriumsdk:
3668
raise ModuleNotFoundError(
3769
"bohriumsdk not installed. Install dpdispatcher with `pip install dpdispatcher[bohrium]`"
38-
)
70+
) from import_bohrium_error
3971
self.init_local_root = local_root
4072
self.init_remote_root = remote_root
4173
self.temp_local_root = os.path.abspath(local_root)
4274
self.remote_profile = remote_profile
43-
self.client = Client()
44-
self.storage = Storage(client=self.client)
75+
access_key = (
76+
remote_profile.get("access_key", None)
77+
or os.getenv("BOHRIUM_ACCESS_KEY", None)
78+
or os.getenv("ACCESS_KEY", None)
79+
)
80+
project_id = (
81+
remote_profile.get("project_id", None)
82+
or os.getenv("BOHRIUM_PROJECT_ID", None)
83+
or os.getenv("PROJECT_ID", None)
84+
)
85+
app_key = (
86+
remote_profile.get("app_key", None)
87+
or os.getenv("BOHRIUM_APP_KEY", None)
88+
or os.getenv("APP_KEY", None)
89+
)
90+
if access_key is None:
91+
raise ValueError(
92+
"remote_profile must contain 'access_key' or set environment variable 'BOHRIUM_ACCESS_KEY'"
93+
)
94+
if project_id is None:
95+
raise ValueError(
96+
"remote_profile must contain 'project_id' or set environment variable 'BOHRIUM_PROJECT_ID'"
97+
)
98+
self.client = Bohrium(
99+
access_key=access_key, project_id=project_id, app_key=app_key
100+
)
101+
self.storage = Tiefblue()
45102
self.job = Job(client=self.client)
46-
self.util = Util()
47103
self.jgid = None
48104

49105
@classmethod
@@ -97,7 +153,7 @@ def upload_job(self, job, common_files=None):
97153
for file in task.forward_files:
98154
upload_file_list.append(os.path.join(task.task_work_path, file))
99155

100-
upload_zip = Util.zip_file_list(
156+
upload_zip = zip_file_list(
101157
self.local_root, zip_task_file, file_list=upload_file_list
102158
)
103159
project_id = self.remote_profile.get("project_id", 0)
@@ -189,7 +245,7 @@ def download(self, submission):
189245
):
190246
continue
191247
self.storage.download_from_url(info["resultUrl"], target_result_zip)
192-
Util.unzip_file(target_result_zip, out_dir=self.local_root)
248+
unzip_file(target_result_zip, out_dir=self.local_root)
193249
self._backup(self.local_root, target_result_zip)
194250
self._clean_backup(
195251
self.local_root, keep_backup=self.remote_profile.get("keep_backup", True)

‎dpdispatcher/machines/openapi.py

Lines changed: 38 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,13 @@
11
import os
22
import shutil
33
import time
4+
from zipfile import ZipFile
45

56
from dpdispatcher.utils.utils import customized_script_header_template
67

78
try:
8-
from bohriumsdk.client import Client
9-
from bohriumsdk.job import Job
10-
from bohriumsdk.storage import Storage
11-
from bohriumsdk.util import Util
9+
from bohrium import Bohrium
10+
from bohrium.resources import Job, Tiefblue
1211
except ModuleNotFoundError:
1312
found_bohriumsdk = False
1413
else:
@@ -23,6 +22,12 @@
2322
"""
2423

2524

25+
def unzip_file(zip_file, out_dir="./"):
26+
obj = ZipFile(zip_file, "r")
27+
for item in obj.namelist():
28+
obj.extract(item, out_dir)
29+
30+
2631
class OpenAPI(Machine):
2732
def __init__(self, context):
2833
if not found_bohriumsdk:
@@ -35,9 +40,35 @@ def __init__(self, context):
3540
self.grouped = self.remote_profile.get("grouped", True)
3641
self.retry_count = self.remote_profile.get("retry_count", 3)
3742
self.ignore_exit_code = context.remote_profile.get("ignore_exit_code", True)
38-
self.client = Client()
43+
44+
access_key = (
45+
self.remote_profile.get("access_key", None)
46+
or os.getenv("BOHRIUM_ACCESS_KEY", None)
47+
or os.getenv("ACCESS_KEY", None)
48+
)
49+
project_id = (
50+
self.remote_profile.get("project_id", None)
51+
or os.getenv("BOHRIUM_PROJECT_ID", None)
52+
or os.getenv("PROJECT_ID", None)
53+
)
54+
app_key = (
55+
self.remote_profile.get("app_key", None)
56+
or os.getenv("BOHRIUM_APP_KEY", None)
57+
or os.getenv("APP_KEY", None)
58+
)
59+
if access_key is None:
60+
raise ValueError(
61+
"remote_profile must contain 'access_key' or set environment variable 'BOHRIUM_ACCESS_KEY'"
62+
)
63+
if project_id is None:
64+
raise ValueError(
65+
"remote_profile must contain 'project_id' or set environment variable 'BOHRIUM_PROJECT_ID'"
66+
)
67+
self.client = Bohrium(
68+
access_key=access_key, project_id=project_id, app_key=app_key
69+
)
70+
self.storage = Tiefblue()
3971
self.job = Job(client=self.client)
40-
self.storage = Storage(client=self.client)
4172
self.group_id = None
4273

4374
def gen_script(self, job):
@@ -102,7 +133,6 @@ def do_submit(self, job):
102133
}
103134
if job.job_state == JobStatus.unsubmitted:
104135
openapi_params["job_id"] = job.job_id
105-
106136
data = self.job.insert(**openapi_params)
107137

108138
job.job_id = data.get("jobId", 0) # type: ignore
@@ -170,7 +200,7 @@ def _download_job(self, job):
170200
result_filename = job_hash + "_back.zip"
171201
target_result_zip = os.path.join(self.context.local_root, result_filename)
172202
self.storage.download_from_url(job_url, target_result_zip)
173-
Util.unzip_file(target_result_zip, out_dir=self.context.local_root)
203+
unzip_file(target_result_zip, out_dir=self.context.local_root)
174204
try:
175205
os.makedirs(os.path.join(self.context.local_root, "backup"), exist_ok=True)
176206
shutil.move(

0 commit comments

Comments
 (0)