-
Notifications
You must be signed in to change notification settings - Fork 47
/
messages_compat.py
80 lines (67 loc) · 2.71 KB
/
messages_compat.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
from __future__ import annotations
import logging
import pickle
import uuid
from globus_compute_common.messagepack import Message as OutgoingMessage
from globus_compute_common.messagepack import pack
from globus_compute_common.messagepack.message_types import (
EPStatusReport as OutgoingEPStatusReport,
)
from globus_compute_common.messagepack.message_types import Result as OutgoingResult
from globus_compute_common.messagepack.message_types import (
ResultErrorDetails as OutgoingResultErrorDetails,
)
from globus_compute_common.messagepack.message_types import Task as OutgoingTask
from globus_compute_common.messagepack.message_types import TaskTransition
from globus_compute_endpoint.executors.high_throughput.messages import (
EPStatusReport as InternalEPStatusReport,
)
from globus_compute_endpoint.executors.high_throughput.messages import (
Task as InternalTask,
)
logger = logging.getLogger(__name__)
def try_convert_to_messagepack(message: bytes) -> bytes:
try:
unpacked = pickle.loads(message)
except pickle.UnpicklingError:
# message isn't pickled; assume that it's already in messagepack format
return message
messagepack_msg: OutgoingMessage | None = None
if isinstance(unpacked, InternalEPStatusReport):
messagepack_msg = OutgoingEPStatusReport(
endpoint_id=unpacked._header,
global_state=unpacked.global_state,
task_statuses=unpacked.task_statuses,
)
elif isinstance(unpacked, dict):
kwargs: dict[
str, str | uuid.UUID | OutgoingResultErrorDetails | list[TaskTransition]
] = {
"task_id": uuid.UUID(unpacked["task_id"]),
}
if "task_statuses" in unpacked:
kwargs["task_statuses"] = unpacked["task_statuses"]
if "exception" in unpacked:
kwargs["data"] = unpacked["exception"]
code, user_message = unpacked.get("error_details", ("Unknown", "Unknown"))
kwargs["error_details"] = OutgoingResultErrorDetails(
code=code, user_message=user_message
)
else:
kwargs["data"] = unpacked["data"]
messagepack_msg = OutgoingResult(**kwargs)
if messagepack_msg:
message = pack(messagepack_msg)
return message
def convert_to_internaltask(message: OutgoingTask, container_type: str | None) -> bytes:
container_loc = "RAW"
if message.container:
for img in message.container.images:
if img.image_type == container_type:
container_loc = img.location
break
return InternalTask(
task_id=str(message.task_id),
container_id=container_loc,
task_buffer=message.task_buffer,
).pack()