293 lines
11 KiB
Python
293 lines
11 KiB
Python
from typing import TypedDict
|
|
import copy
|
|
import json
|
|
import io
|
|
import base64
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
class AnyType(str):
|
|
def __ne__(self, __value: object) -> bool:
|
|
return False
|
|
|
|
|
|
class ContainsDynamicDict(dict):
|
|
"""
|
|
A custom dictionary that dynamically returns values for keys based on a pattern.
|
|
- If a key in the passed dictionary has a value with `{"_dynamic": "number"}` in the tuple's second position,
|
|
then any other key starting with the same string and ending with a number will return that value.
|
|
- For other keys, normal dictionary lookup behavior applies.
|
|
"""
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
# Store prefixes associated with `_dynamic` values for efficient lookup
|
|
self._dynamic_prefixes = {
|
|
key.rstrip("0123456789"): value
|
|
for key, value in self.items()
|
|
if isinstance(value, tuple) and len(value) > 1 and value[1].get("_dynamic") == "number"
|
|
}
|
|
|
|
def __contains__(self, key):
|
|
# Check if key matches a dynamically handled prefix or exists normally
|
|
return any(key.startswith(prefix) and key[len(prefix) :].isdigit() for prefix in self._dynamic_prefixes) or super().__contains__(
|
|
key
|
|
)
|
|
|
|
def __getitem__(self, key):
|
|
# Dynamically return the value for keys matching a `prefix<number>` pattern
|
|
for prefix, value in self._dynamic_prefixes.items():
|
|
if key.startswith(prefix) and key[len(prefix) :].isdigit():
|
|
return value
|
|
# Fallback to normal dictionary behavior for other keys
|
|
return super().__getitem__(key)
|
|
|
|
|
|
class Closure(TypedDict):
|
|
body: dict
|
|
captures: list
|
|
output: list
|
|
|
|
|
|
def create_graph_from_closure(closure: Closure, params, caller_unique_id=None):
|
|
body = copy.deepcopy(closure["body"])
|
|
graph = {}
|
|
|
|
def transform_input(spec):
|
|
if not isinstance(spec, list):
|
|
return spec
|
|
# Return normal spec if not a capture or param
|
|
prefix = f"{caller_unique_id}_" if caller_unique_id is not None else ""
|
|
if spec[0] != "__capture" and spec[0] != "__param":
|
|
return [f"{prefix}{spec[0]}", spec[1]]
|
|
if spec[0] == "__param":
|
|
if spec[1] >= len(params):
|
|
raise IndexError(f"Parameter index {spec[1]} out of range for function call. Only {len(params)} parameters were provided.")
|
|
value = params[spec[1]]
|
|
else:
|
|
value = closure["captures"][spec[1]]
|
|
if isinstance(value, list):
|
|
recover_id = f"{prefix}{spec[0]}_{spec[1]}"
|
|
if recover_id not in graph:
|
|
graph[recover_id] = {
|
|
"inputs": {
|
|
# Wrap the list in a tuple so it's treated as a literal, not a link
|
|
"values": (value,),
|
|
},
|
|
"class_type": "__RecoverList__",
|
|
}
|
|
return [recover_id, 0]
|
|
else:
|
|
return value
|
|
|
|
for node_id, node_data in body.items():
|
|
inputs = node_data["inputs"]
|
|
for key in inputs.keys():
|
|
inputs[key] = transform_input(inputs[key])
|
|
node_data["override_display_id"] = node_id
|
|
graph[f"{caller_unique_id}_{node_id}"] = node_data
|
|
|
|
output = transform_input(closure["output"])
|
|
return graph, output
|
|
|
|
|
|
def create_remote_workflow_from_closure(closure, params, use_shared_memory=False):
|
|
body = copy.deepcopy(closure["body"])
|
|
graph = {}
|
|
|
|
def transform_input(spec):
|
|
if not isinstance(spec, list):
|
|
return spec
|
|
# Return normal spec if not a capture or param
|
|
if spec[0] != "__capture" and spec[0] != "__param":
|
|
return spec
|
|
|
|
# Check if we have already created a deserialize node for this spec
|
|
deserialize_node_id = f"deserialize_{spec[0]}_{spec[1]}"
|
|
if deserialize_node_id in graph:
|
|
return [deserialize_node_id, 0]
|
|
|
|
# Try to create a deserialize node
|
|
if spec[0] == "__param":
|
|
if spec[1] >= len(params):
|
|
raise IndexError(f"Parameter index {spec[1]} out of range for function call. Only {len(params)} parameters were provided.")
|
|
value = params[spec[1]]
|
|
else:
|
|
value = closure["captures"][spec[1]]
|
|
|
|
try:
|
|
serialized_value, _ = serialize(value, use_shared_memory=use_shared_memory)
|
|
graph[deserialize_node_id] = {
|
|
"inputs": {
|
|
"data": serialized_value,
|
|
},
|
|
"class_type": "__Deserialize__",
|
|
}
|
|
return [deserialize_node_id, 0]
|
|
except TypeError as e:
|
|
if spec[0] == "__capture":
|
|
raise ValueError(
|
|
f"Cannot serialize captured value at index {spec[1]}. "
|
|
"Captured values must be serializable to be used in remote execution. "
|
|
"Consider disabling capture for this function to initialize the value remotely. "
|
|
f"Error: {e}"
|
|
)
|
|
else:
|
|
raise ValueError(
|
|
f"Cannot serialize parameter at index {spec[1]}. "
|
|
"Parameters must be serializable to be used in remote execution. "
|
|
f"Error: {e}"
|
|
)
|
|
|
|
for node_id, node_data in body.items():
|
|
inputs = node_data["inputs"]
|
|
for key in inputs.keys():
|
|
inputs[key] = transform_input(inputs[key])
|
|
graph[node_id] = node_data
|
|
|
|
output = transform_input(closure["output"])
|
|
graph["__output__"] = {
|
|
"inputs": {"data": output, "use_shared_memory": use_shared_memory},
|
|
"class_type": "__Serialize__",
|
|
"_meta": {"title": "%__output__%"},
|
|
}
|
|
|
|
return graph
|
|
|
|
|
|
class WhitelistEncoder(json.JSONEncoder):
|
|
def __init__(self, use_shared_memory=False, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.use_shared_memory = use_shared_memory
|
|
self.shared_tensors = set()
|
|
|
|
def default(self, obj):
|
|
# 区分 set 和 tuple
|
|
if isinstance(obj, set):
|
|
return {"__type__": "set", "data": list(obj)}
|
|
if isinstance(obj, tuple):
|
|
return {"__type__": "tuple", "data": list(obj)}
|
|
|
|
if isinstance(obj, bytes):
|
|
return {"__type__": "bytes", "data": base64.b64encode(obj).decode("utf-8")}
|
|
|
|
if isinstance(obj, np.ndarray):
|
|
buffer = io.BytesIO()
|
|
np.save(buffer, obj, allow_pickle=False)
|
|
return {
|
|
"__type__": "numpy.ndarray",
|
|
"data": base64.b64encode(buffer.getvalue()).decode("utf-8"),
|
|
}
|
|
|
|
if isinstance(obj, torch.Tensor):
|
|
obj = obj.cpu().contiguous()
|
|
if self.use_shared_memory:
|
|
# 启用共享内存优化
|
|
obj.share_memory_() # 确保移动到共享内存
|
|
storage = obj.untyped_storage()
|
|
manager_handle_bytes, handle_bytes, storage_size = storage._share_filename_cpu_()
|
|
self.shared_tensors.add(obj)
|
|
return {
|
|
"__type__": "torch.Tensor_shared",
|
|
"manager_handle": manager_handle_bytes.decode("utf-8"),
|
|
"handle": handle_bytes.decode("utf-8"),
|
|
"storage_size": storage_size,
|
|
"offset": obj.storage_offset(),
|
|
"shape": list(obj.shape),
|
|
"strides": list(obj.stride()),
|
|
"dtype": str(obj.dtype),
|
|
"requires_grad": obj.requires_grad,
|
|
}
|
|
else:
|
|
# 普通序列化（回退）
|
|
buffer = io.BytesIO()
|
|
torch.save(obj, buffer) # 确保 CPU
|
|
return {
|
|
"__type__": "torch.Tensor",
|
|
"data": base64.b64encode(buffer.getvalue()).decode("utf-8"),
|
|
}
|
|
|
|
if isinstance(obj, Image.Image):
|
|
buffer = io.BytesIO()
|
|
obj.save(buffer, format="PNG")
|
|
return {
|
|
"__type__": "PIL.Image",
|
|
"data": base64.b64encode(buffer.getvalue()).decode("utf-8"),
|
|
}
|
|
|
|
raise TypeError(
|
|
f"Object of type '{type(obj).__name__}' is not serializable. "
|
|
"Only basic types, set, tuple, bytes, numpy.ndarray, "
|
|
"torch.Tensor, and PIL.Image are allowed."
|
|
)
|
|
|
|
|
|
def whitelist_decoder(dct):
|
|
"""
|
|
用于 json.load/loads 的 object_hook，用于恢复自定义类型。
|
|
"""
|
|
if "__type__" in dct:
|
|
obj_type = dct["__type__"]
|
|
data = dct.get("data")
|
|
|
|
if obj_type == "set":
|
|
return set(data)
|
|
if obj_type == "tuple":
|
|
return tuple(data)
|
|
if obj_type == "bytes":
|
|
return base64.b64decode(data)
|
|
|
|
if obj_type == "numpy.ndarray":
|
|
buffer = io.BytesIO(base64.b64decode(data))
|
|
return np.load(buffer, allow_pickle=True)
|
|
|
|
if obj_type == "torch.Tensor":
|
|
buffer = io.BytesIO(base64.b64decode(data))
|
|
return torch.load(buffer)
|
|
|
|
if obj_type == "torch.Tensor_shared":
|
|
# 重建共享 Tensor
|
|
manager_handle_bytes = dct["manager_handle"].encode("utf-8")
|
|
handle_bytes = dct["handle"].encode("utf-8")
|
|
storage_size = dct["storage_size"]
|
|
offset = dct["offset"]
|
|
shape = tuple(dct["shape"])
|
|
strides = tuple(dct["strides"])
|
|
dtype_str = dct["dtype"].split(".")[-1] # 如 'float32'
|
|
dtype = getattr(torch, dtype_str)
|
|
requires_grad = dct["requires_grad"]
|
|
|
|
# 重建存储（attach 到共享内存）
|
|
storage = torch.UntypedStorage._new_shared_filename_cpu(manager_handle_bytes, handle_bytes, storage_size)
|
|
|
|
# 创建 Tensor 视图
|
|
tensor = torch.empty(shape, dtype=dtype).set_(storage, offset, shape, strides)
|
|
tensor.requires_grad_(requires_grad)
|
|
return tensor
|
|
|
|
if obj_type == "PIL.Image":
|
|
buffer = io.BytesIO(base64.b64decode(data))
|
|
img = Image.open(buffer)
|
|
return img.copy()
|
|
|
|
return dct
|
|
|
|
|
|
def serialize(obj, use_shared_memory=False):
|
|
"""
|
|
将对象序列化为 JSON 字符串。
|
|
仅支持基本类型、set、tuple、bytes、numpy.ndarray、torch.Tensor 和 PIL.Image。
|
|
use_shared_memory: 如果为 True，且 obj 是 CPU torch.Tensor，则使用共享内存优化（传输元数据）。
|
|
"""
|
|
encoder = WhitelistEncoder(use_shared_memory=use_shared_memory)
|
|
return encoder.encode(obj), encoder.shared_tensors
|
|
|
|
|
|
def deserialize(json_str):
|
|
"""
|
|
从 JSON 字符串反序列化对象。
|
|
"""
|
|
return json.loads(json_str, object_hook=whitelist_decoder)
|