-
Notifications
You must be signed in to change notification settings - Fork 92
Expand file tree
/
Copy pathget_compose_hash.py
More file actions
289 lines (244 loc) · 10.2 KB
/
Copy pathget_compose_hash.py
File metadata and controls
289 lines (244 loc) · 10.2 KB
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
# SPDX-FileCopyrightText: © 2025 Phala Network <dstack@phala.network>
#
# SPDX-License-Identifier: Apache-2.0
"""Compose hash calculation module for dstack SDK.
Provides deterministic JSON serialization and SHA256 hashing of AppCompose configurations,
compatible with the TypeScript implementation.
"""
import hashlib
import json
from typing import Any
from typing import Dict
from typing import List
from typing import Literal
from typing import Optional
from typing import Union
KeyProviderKind = Literal["none", "kms", "local", "tpm"]
class DockerConfig:
"""Docker configuration for app compose."""
def __init__(
self,
registry: Optional[str] = None,
username: Optional[str] = None,
token_key: Optional[str] = None,
) -> None:
"""Initialize a new ``DockerConfig`` instance."""
self.registry = registry
self.username = username
self.token_key = token_key
def to_dict(self) -> Dict[str, Any]:
"""Return a dictionary representation excluding ``None`` fields."""
result: Dict[str, Any] = {}
if self.registry is not None:
result["registry"] = self.registry
if self.username is not None:
result["username"] = self.username
if self.token_key is not None:
result["token_key"] = self.token_key
return result
class Requirements:
"""Guest-side requirements for app compose."""
def __init__(
self,
os_version: Optional[str] = None,
platforms: Optional[List[str]] = None,
tdx_measure_acpi_tables: Optional[bool] = None,
launch_token_hash: Optional[str] = None,
) -> None:
"""Initialize a new ``Requirements`` instance."""
self.os_version = os_version
self.platforms = platforms
self.tdx_measure_acpi_tables = tdx_measure_acpi_tables
self.launch_token_hash = launch_token_hash
def to_dict(self) -> Dict[str, Any]:
"""Return a dictionary representation excluding ``None`` fields."""
result: Dict[str, Any] = {}
if self.os_version is not None:
result["os_version"] = self.os_version
if self.platforms is not None:
result["platforms"] = self.platforms
if self.tdx_measure_acpi_tables is not None:
result["tdx_measure_acpi_tables"] = self.tdx_measure_acpi_tables
if self.launch_token_hash is not None:
result["launch_token_hash"] = self.launch_token_hash
return result
class AppCompose:
"""App compose configuration."""
def __init__(
self,
runner: str,
manifest_version: Optional[Union[int, str]] = None,
name: Optional[str] = None,
features: Optional[List[str]] = None, # Deprecated
docker_compose_file: Optional[str] = None,
docker_config: Optional[DockerConfig] = None,
public_logs: Optional[bool] = None,
public_sysinfo: Optional[bool] = None,
public_tcbinfo: Optional[bool] = None,
kms_enabled: Optional[bool] = None,
gateway_enabled: Optional[bool] = None,
tproxy_enabled: Optional[bool] = None, # For backward compatibility
local_key_provider_enabled: Optional[bool] = None,
key_provider: Optional[KeyProviderKind] = None,
key_provider_id: Optional[str] = None,
allowed_envs: Optional[List[str]] = None,
no_instance_id: Optional[bool] = None,
secure_time: Optional[bool] = None,
requirements: Optional[Union[Requirements, Dict[str, Any]]] = None,
bash_script: Optional[str] = None, # Legacy
pre_launch_script: Optional[str] = None, # Legacy
snapshotter: Optional[str] = None,
**kwargs: Any,
) -> None:
"""Initialize a new ``AppCompose`` instance with arbitrary extra fields."""
self.runner = runner
self.snapshotter = snapshotter
self.manifest_version = manifest_version
self.name = name
self.features = features
self.docker_compose_file = docker_compose_file
self.docker_config = docker_config
self.public_logs = public_logs
self.public_sysinfo = public_sysinfo
self.public_tcbinfo = public_tcbinfo
self.kms_enabled = kms_enabled
self.gateway_enabled = gateway_enabled
self.tproxy_enabled = tproxy_enabled
self.local_key_provider_enabled = local_key_provider_enabled
self.key_provider = key_provider
self.key_provider_id = key_provider_id
self.allowed_envs = allowed_envs
self.no_instance_id = no_instance_id
self.secure_time = secure_time
self.requirements = requirements
self.bash_script = bash_script
self.pre_launch_script = pre_launch_script
# Add any additional fields
for key, value in kwargs.items():
setattr(self, key, value)
def to_dict(self) -> Dict[str, Any]:
"""Convert to dictionary representation."""
result: Dict[str, Any] = {}
# Add all attributes that are not None
for attr_name in dir(self):
if not attr_name.startswith("_") and not callable(getattr(self, attr_name)):
value = getattr(self, attr_name)
if value is not None:
if isinstance(value, (DockerConfig, Requirements)):
result[attr_name] = value.to_dict()
else:
# Handle special float values
if isinstance(value, float):
if value != value: # NaN check
result[attr_name] = None
elif value == float("inf") or value == float("-inf"):
result[attr_name] = None
else:
result[attr_name] = value
else:
result[attr_name] = value
return result
@classmethod
def from_dict(cls, data: Dict[str, Any]) -> "AppCompose":
"""Create AppCompose from dictionary."""
# Handle docker_config
docker_config: Optional[DockerConfig] = None
if "docker_config" in data and data["docker_config"] is not None:
dc = data.pop("docker_config")
docker_config = DockerConfig(**dc)
# Handle requirements
requirements: Optional[Requirements] = None
if "requirements" in data and data["requirements"] is not None:
req = data.pop("requirements")
requirements = Requirements(**req)
# Handle special float values
processed_data: Dict[str, Any] = {}
for key, value in data.items():
if isinstance(value, float):
if value != value: # NaN check
processed_data[key] = None
elif value == float("inf") or value == float("-inf"):
processed_data[key] = None
else:
processed_data[key] = value
else:
processed_data[key] = value
runner_value = processed_data.pop("runner")
runner: str = str(runner_value)
return cls(
runner=runner,
docker_config=docker_config,
requirements=requirements,
**processed_data,
)
def sort_object(obj: Any) -> Any:
"""Recursively sort object keys lexicographically.
This is crucial for deterministic JSON.stringify.
"""
if obj is None:
return obj
elif isinstance(obj, list):
return [sort_object(item) for item in obj]
elif isinstance(obj, dict):
return {key: sort_object(value) for key, value in sorted(obj.items())}
else:
return obj
def preprocess_app_compose(app_compose: AppCompose) -> AppCompose:
"""Preprocess app compose by removing conflicting fields based on runner."""
# Create a copy
data = app_compose.to_dict()
if data.get("runner") == "bash" and "docker_compose_file" in data:
del data["docker_compose_file"]
elif (
data.get("runner") in ("docker-compose", "nerdctl-compose")
and "bash_script" in data
):
del data["bash_script"]
if "pre_launch_script" in data and not data["pre_launch_script"]:
del data["pre_launch_script"]
return AppCompose.from_dict(data)
def to_deterministic_json(app_compose: AppCompose) -> str:
"""Serialize to deterministic JSON following cross-language standards.
- Recursively sorts object keys lexicographically
- Compact output (no spaces)
- Handles special values (NaN, Infinity) by converting them to null
- UTF-8 encoding (default in Python)
"""
data = sort_object(app_compose.to_dict())
def convert_special_values(obj: Any) -> Any:
"""Convert NaN and Infinity to null for deterministic output."""
if isinstance(obj, float):
if obj != obj: # NaN check
return None
if obj == float("inf") or obj == float("-inf"):
return None
return obj
def json_serializer(obj: Any) -> Any:
"""Handle special float values during JSON serialization."""
return convert_special_values(obj)
# Convert special values recursively
def process_data(obj: Any) -> Any:
if isinstance(obj, dict):
return {key: process_data(value) for key, value in obj.items()}
elif isinstance(obj, list):
return [process_data(item) for item in obj]
else:
return convert_special_values(obj)
processed_data = process_data(data)
return json.dumps(processed_data, separators=(",", ":"), ensure_ascii=False)
def get_compose_hash(
app_compose: Union[AppCompose, Dict[str, Any]], normalize: bool = False
) -> str:
"""Calculate SHA256 hash of app compose configuration.
Args:
app_compose: AppCompose object or dictionary
normalize: Whether to preprocess the compose (remove conflicting fields)
Returns:
str: SHA256 hash as hex string
"""
if isinstance(app_compose, dict):
app_compose = AppCompose.from_dict(app_compose)
if normalize:
app_compose = preprocess_app_compose(app_compose)
manifest_str = to_deterministic_json(app_compose)
return hashlib.sha256(manifest_str.encode("utf-8")).hexdigest()