Skip to content

Commit 2e57d11

Browse files
nboveeRajeshwariGangoni
nbovee
and
RajeshwariGangoni
authored
feat(encryption): Implemented AES-CBC and AES-CTR encryptions (#13)
* Implemented AES-CBC and AES-CTR encryptions --------- Co-authored-by: Rajeshwari Gangoni <rajeshwarig.5639@gmail.com>
1 parent 83e82c4 commit 2e57d11

30 files changed

+3459
-1779
lines changed

‎config/alexnetsplit.yaml‎

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@
44
# DEFAULT CONFIGURATIONS
55
# ================================================================
66
default:
7-
device: cuda
8-
save_layer_images: false
9-
collect_metrics: false
7+
device: cpu
8+
save_layer_images: true
9+
collect_metrics: true
1010

1111
# ================================================================
1212
# LOGGING CONFIGURATIONS
@@ -60,3 +60,12 @@ compression:
6060
clevel: 3
6161
filter: SHUFFLE
6262
codec: ZSTD
63+
64+
# ================================================================
65+
# ENCRYPTION CONFIGURATIONS
66+
# ================================================================
67+
encryption:
68+
enabled: false
69+
algorithm: AES-CBC
70+
key_size: 256
71+
test_key: "12345678921234597890123456789012"

‎config/alexnetsplit_ctr.yaml‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
# config/alexnetsplit_ctr.yaml
2+
3+
# ================================================================
4+
# DEFAULT CONFIGURATIONS
5+
# ================================================================
6+
default:
7+
device: cpu
8+
save_layer_images: true
9+
collect_metrics: true
10+
11+
# ================================================================
12+
# LOGGING CONFIGURATIONS
13+
# ================================================================
14+
logging:
15+
log_level: DEBUG
16+
log_file: logs/app.log
17+
18+
# ================================================================
19+
# MODEL CONFIGURATIONS
20+
# ================================================================
21+
model:
22+
model_name: alexnet
23+
version: null
24+
pretrained: true
25+
weight_path: null
26+
input_size: [3, 224, 224]
27+
split_layer: 12
28+
save_layers: [2, 4, 6]
29+
num_classes: 1000
30+
mode: eval
31+
depth: 2
32+
flush_buffer_size: 100
33+
warmup_iterations: 10
34+
log_file: logs/alexnetsplit_ctr.log
35+
36+
# ================================================================
37+
# DATASET CONFIGURATIONS
38+
# ================================================================
39+
dataset:
40+
name: "imagenet"
41+
root: "data/imagenet"
42+
class_names: "data/imagenet/imagenet_classes.txt"
43+
img_directory: "data/imagenet/sample_images"
44+
transform: null
45+
max_samples: -1
46+
47+
# ================================================================
48+
# DATALOADER CONFIGURATIONS
49+
# ================================================================
50+
dataloader:
51+
batch_size: 1
52+
shuffle: false
53+
num_workers: 4
54+
collate_fn: imagenet_collate
55+
56+
# ================================================================
57+
# COMPRESSION CONFIGURATIONS
58+
# ================================================================
59+
compression:
60+
clevel: 3
61+
filter: SHUFFLE
62+
codec: ZSTD
63+
64+
# ================================================================
65+
# ENCRYPTION CONFIGURATIONS - AES-CTR MODE
66+
# ================================================================
67+
encryption:
68+
enabled: true
69+
algorithm: AES-CTR # Using AES-CTR instead of AES-CBC
70+
key_size: 256
71+
test_key: "12345678921234597890123456789012"

‎host.py‎

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -460,7 +460,7 @@ def main() -> None:
460460

461461
host = None
462462
try:
463-
print(f"Initializing experiment with config from {config_path}...")
463+
logger.info(f"Initializing experiment with config from {config_path}...")
464464

465465
# Load config to check if we should modify it
466466
config = read_yaml_file(str(config_path))
@@ -470,34 +470,34 @@ def main() -> None:
470470
config["experiment"] = {}
471471
if "type" not in config["experiment"]:
472472
config["experiment"]["type"] = "networked"
473-
print("Setting experiment type to 'networked'")
473+
logger.info("Setting experiment type to 'networked'")
474474

475475
host = ExperimentHost(str(config_path))
476476

477-
print("Starting experiment...")
477+
logger.info("Starting experiment...")
478478
host.run_experiment()
479479

480480
if args.copy_results and host:
481-
print("Copying results to server...")
481+
logger.info("Copying results to server...")
482482
success = host._copy_results_to_server()
483483
if success:
484-
print("Results successfully copied to server")
484+
logger.info("Results successfully copied to server")
485485
else:
486-
print("Failed to copy results to server")
486+
logger.error("Failed to copy results to server")
487487

488488
except KeyboardInterrupt:
489-
print("\nExperiment interrupted by user")
489+
logger.info("\nExperiment interrupted by user")
490490
logger.info("Experiment interrupted by user")
491491
except Exception as e:
492-
print(f"Error: {e}")
492+
logger.error(f"Error: {e}")
493493
logger.error(f"Experiment failed: {e}", exc_info=True)
494494
finally:
495495
if host:
496-
print("Cleaning up...")
496+
logger.info("Cleaning up...")
497497
host.cleanup()
498-
print("Done.")
498+
logger.info("Done.")
499499
else:
500-
print("Exiting without cleanup (host was not initialized)")
500+
logger.warning("Exiting without cleanup (host was not initialized)")
501501

502502

503503
if __name__ == "__main__":

‎pyproject.toml‎

Lines changed: 34 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[project]
22
name = "tracr"
3-
version = "0.4.2"
3+
version = "0.5.0"
44
description = "An experimental framework for computational offloading and distributed neural network inference through split computing"
55
readme = "README.md"
66
requires-python = ">=3.10"
@@ -19,6 +19,9 @@ dependencies = [
1919
"pyyaml>=6.0.2",
2020
"pandas>=2.2.3",
2121
"tqdm>=4.67.0",
22+
"cryptography>=43.0.0",
23+
"psutil>=7.0.0",
24+
"openpyxl>=3.1.5",
2225
]
2326

2427
[project.optional-dependencies]
@@ -43,6 +46,16 @@ cu124 = [
4346
"torchvision>=0.20",
4447
"torchaudio>=2.5",
4548
]
49+
cu126 = [
50+
"torch>=2.5",
51+
"torchvision>=0.20",
52+
"torchaudio>=2.5",
53+
]
54+
cu128 = [
55+
"torch>=2.5",
56+
"torchvision>=0.20",
57+
"torchaudio>=2.5",
58+
]
4659
full = [
4760
"ultralytics>=8.3.28",
4861
]
@@ -55,11 +68,13 @@ conflicts = [
5568
{ extra = "cu118" },
5669
{ extra = "cu121" },
5770
{ extra = "cu124" },
71+
{ extra = "cu126" },
72+
{ extra = "cu128" },
5873
],
5974
]
6075

6176
[tool.ruff.lint]
62-
ignore = ["F841"]
77+
# ignore = ["F841"]
6378

6479
[dependency-groups]
6580
dev = [
@@ -72,19 +87,26 @@ torch = [
7287
{ index = "pytorch-cu118", extra = "cu118" },
7388
{ index = "pytorch-cu121", extra = "cu121" },
7489
{ index = "pytorch-cu124", extra = "cu124" },
90+
{ index = "pytorch-cu126", extra = "cu126" },
91+
{ index = "pytorch-cu128", extra = "cu128" },
7592
]
7693
torchvision = [
7794
{ index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
7895
{ index = "pytorch-cu118", extra = "cu118" },
7996
{ index = "pytorch-cu121", extra = "cu121" },
8097
{ index = "pytorch-cu124", extra = "cu124" },
98+
{ index = "pytorch-cu126", extra = "cu126" },
99+
{ index = "pytorch-cu128", extra = "cu128" },
81100
]
82101

83102
torchaudio = [
84103
{ index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
85104
{ index = "pytorch-cu118", extra = "cu118" },
86105
{ index = "pytorch-cu121", extra = "cu121" },
87106
{ index = "pytorch-cu124", extra = "cu124" },
107+
{ index = "pytorch-cu126", extra = "cu126" },
108+
{ index = "pytorch-cu128", extra = "cu128" },
109+
88110
]
89111

90112
[[tool.uv.index]]
@@ -107,33 +129,16 @@ name = "pytorch-cu124"
107129
url = "https://download.pytorch.org/whl/cu124"
108130
explicit = true
109131

132+
[[tool.uv.index]]
133+
name = "pytorch-cu126"
134+
url = "https://download.pytorch.org/whl/cu126"
135+
explicit = true
136+
137+
[[tool.uv.index]]
138+
name = "pytorch-cu128"
139+
url = "https://download.pytorch.org/whl/cu128"
140+
explicit = true
141+
110142
[tool.ruff]
111143
line-length = 88
112144
indent-width = 4
113-
114-
# [build-system]
115-
# requires = ["hatchling"]
116-
# build-backend = "hatchling.build"
117-
118-
# [project.scripts]
119-
# tracr-server = "server:main"
120-
# tracr-client = "host:main"
121-
122-
# [tool.hatch.build]
123-
# packages = ["src"]
124-
# include = [
125-
# "server.py",
126-
# "host.py",
127-
# ]
128-
# exclude = [
129-
# "data/",
130-
# "tests/",
131-
# "venv/",
132-
# "logs/",
133-
# "*.pyc",
134-
# "__pycache__",
135-
# "*.egg-info",
136-
# ]
137-
138-
# [tool.hatch.build.targets.wheel]
139-
# packages = ["src/"]

‎requirements-cu118.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,18 @@
22
# uv pip compile pyproject.toml -o requirements-cu118.txt --no-deps --extra cu118 --extra full
33
blosc2==3.2.0
44
# via tracr (pyproject.toml)
5+
cryptography==46.0.1
6+
# via tracr (pyproject.toml)
57
loguru==0.7.3
68
# via tracr (pyproject.toml)
9+
openpyxl==3.1.5
10+
# via tracr (pyproject.toml)
711
pandas==2.2.3
812
# via tracr (pyproject.toml)
913
paramiko==3.5.1
1014
# via tracr (pyproject.toml)
15+
psutil==7.1.0
16+
# via tracr (pyproject.toml)
1117
pyyaml==6.0.2
1218
# via tracr (pyproject.toml)
1319
rich==13.9.4

‎requirements-cu121.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,18 @@
22
# uv pip compile pyproject.toml -o requirements-cu121.txt --no-deps --extra cu121 --extra full
33
blosc2==3.2.0
44
# via tracr (pyproject.toml)
5+
cryptography==46.0.1
6+
# via tracr (pyproject.toml)
57
loguru==0.7.3
68
# via tracr (pyproject.toml)
9+
openpyxl==3.1.5
10+
# via tracr (pyproject.toml)
711
pandas==2.2.3
812
# via tracr (pyproject.toml)
913
paramiko==3.5.1
1014
# via tracr (pyproject.toml)
15+
psutil==7.1.0
16+
# via tracr (pyproject.toml)
1117
pyyaml==6.0.2
1218
# via tracr (pyproject.toml)
1319
rich==13.9.4

‎requirements-cu124.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,18 @@
22
# uv pip compile pyproject.toml -o requirements-cu124.txt --no-deps --extra cu124 --extra full
33
blosc2==3.2.0
44
# via tracr (pyproject.toml)
5+
cryptography==46.0.1
6+
# via tracr (pyproject.toml)
57
loguru==0.7.3
68
# via tracr (pyproject.toml)
9+
openpyxl==3.1.5
10+
# via tracr (pyproject.toml)
711
pandas==2.2.3
812
# via tracr (pyproject.toml)
913
paramiko==3.5.1
1014
# via tracr (pyproject.toml)
15+
psutil==7.1.0
16+
# via tracr (pyproject.toml)
1117
pyyaml==6.0.2
1218
# via tracr (pyproject.toml)
1319
rich==13.9.4

‎requirements-cu126.txt‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
# This file was autogenerated by uv via the following command:
2+
# uv pip compile pyproject.toml -o requirements-cu126.txt --no-deps --extra cu126 --extra full
3+
blosc2==3.9.0
4+
# via tracr (pyproject.toml)
5+
cryptography==46.0.1
6+
# via tracr (pyproject.toml)
7+
loguru==0.7.3
8+
# via tracr (pyproject.toml)
9+
openpyxl==3.1.5
10+
# via tracr (pyproject.toml)
11+
pandas==2.3.2
12+
# via tracr (pyproject.toml)
13+
paramiko==4.0.0
14+
# via tracr (pyproject.toml)
15+
psutil==7.1.0
16+
# via tracr (pyproject.toml)
17+
pyyaml==6.0.3
18+
# via tracr (pyproject.toml)
19+
rich==14.1.0
20+
# via tracr (pyproject.toml)
21+
rpyc==6.0.2
22+
# via tracr (pyproject.toml)
23+
tomli==2.2.1
24+
# via tracr (pyproject.toml)
25+
torch==2.8.0+cu126
26+
# via tracr (pyproject.toml)
27+
torchaudio==2.8.0+cu126
28+
# via tracr (pyproject.toml)
29+
torchinfo==1.8.0
30+
# via tracr (pyproject.toml)
31+
torchvision==0.23.0+cu126
32+
# via tracr (pyproject.toml)
33+
tqdm==4.67.1
34+
# via tracr (pyproject.toml)
35+
ultralytics==8.3.203
36+
# via tracr (pyproject.toml)

0 commit comments

Comments
 (0)