vantage6-node 5.0.2__tar.gz → 5.0.3rc5__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/PKG-INFO +6 -6
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/pyproject.toml +5 -5
- vantage6_node-5.0.3rc5/tests_node/test_container_manager.py +220 -0
- vantage6_node-5.0.3rc5/tests_node/test_network_isolation.py +124 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/tests_node/test_pod_status_mapping.py +0 -1
- vantage6_node-5.0.3rc5/tests_node/test_share_algorithm_logs.py +141 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/__init__.py +110 -91
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/cli/node.py +4 -4
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/globals.py +17 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/container_manager.py +114 -42
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/data_classes.py +8 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/exceptions.py +0 -12
- vantage6_node-5.0.3rc5/vantage6/node/k8s/network_isolation.py +149 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/run_io.py +7 -7
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/session_manager.py +3 -3
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/proxy_server.py +8 -6
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/socket.py +13 -3
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/util/colorer.py +2 -5
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/.gitignore +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/Makefile +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/README.md +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/dev_start.sh +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/tests/docker/__init__.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/tests_node/__init__.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/utest.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/dev_start.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/cli/__init__.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/enum.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/k8s/jobpod_state_to_run_status_mapper.py +0 -0
- {vantage6_node-5.0.2 → vantage6_node-5.0.3rc5}/vantage6/node/util/__init__.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
2
|
Name: vantage6-node
|
|
3
|
-
Version: 5.0.
|
|
3
|
+
Version: 5.0.3rc5
|
|
4
4
|
Summary: vantage6 node
|
|
5
5
|
Author: Vantage6 Team
|
|
6
6
|
Maintainer-email: Frank Martin <f.martin@iknl.nl>, Bart van Beusekom <b.vanbeusekom@iknl.nl>
|
|
@@ -14,11 +14,11 @@ Requires-Dist: kubernetes==36.0.2
|
|
|
14
14
|
Requires-Dist: nvidia-ml-py==12.535.133
|
|
15
15
|
Requires-Dist: parquet-tools==0.2.16
|
|
16
16
|
Requires-Dist: psutil==7.0.0
|
|
17
|
-
Requires-Dist: python-socketio==5.
|
|
17
|
+
Requires-Dist: python-socketio==5.16.2
|
|
18
18
|
Requires-Dist: requests==2.33.1
|
|
19
|
-
Requires-Dist: vantage6-client==5.0.
|
|
20
|
-
Requires-Dist: vantage6-common==5.0.
|
|
21
|
-
Requires-Dist: vantage6==5.0.
|
|
19
|
+
Requires-Dist: vantage6-client==5.0.3rc5
|
|
20
|
+
Requires-Dist: vantage6-common==5.0.3rc5
|
|
21
|
+
Requires-Dist: vantage6==5.0.3rc5
|
|
22
22
|
Provides-Extra: dev
|
|
23
23
|
Requires-Dist: black; extra == 'dev'
|
|
24
24
|
Requires-Dist: coverage==7.10.2; extra == 'dev'
|
|
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "vantage6-node"
|
|
7
|
-
version = "5.0.
|
|
7
|
+
version = "5.0.3rc5"
|
|
8
8
|
description = "vantage6 node"
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
license = {text = "MIT"}
|
|
@@ -25,11 +25,11 @@ dependencies = [
|
|
|
25
25
|
"nvidia-ml-py==12.535.133",
|
|
26
26
|
"parquet-tools==0.2.16",
|
|
27
27
|
"psutil==7.0.0",
|
|
28
|
-
"python-socketio==5.
|
|
28
|
+
"python-socketio==5.16.2",
|
|
29
29
|
"requests==2.33.1",
|
|
30
|
-
"vantage6==5.0.
|
|
31
|
-
"vantage6-client==5.0.
|
|
32
|
-
"vantage6-common==5.0.
|
|
30
|
+
"vantage6==5.0.3rc5",
|
|
31
|
+
"vantage6-client==5.0.3rc5",
|
|
32
|
+
"vantage6-common==5.0.3rc5",
|
|
33
33
|
]
|
|
34
34
|
|
|
35
35
|
[project.optional-dependencies]
|
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
import unittest
|
|
4
|
+
from unittest.mock import MagicMock, patch
|
|
5
|
+
|
|
6
|
+
from vantage6.common.globals import NodePolicy
|
|
7
|
+
|
|
8
|
+
from vantage6.node.k8s.container_manager import ContainerManager
|
|
9
|
+
|
|
10
|
+
def _get_null_logger(name="null_logger_container_manager"):
|
|
11
|
+
logger = logging.getLogger(name)
|
|
12
|
+
logger.addHandler(logging.NullHandler())
|
|
13
|
+
logger.propagate = False
|
|
14
|
+
return logger
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class TestIsImageAllowed(unittest.TestCase):
|
|
18
|
+
"""
|
|
19
|
+
Tests for ContainerManager.is_image_allowed().
|
|
20
|
+
|
|
21
|
+
ContainerManager.__init__ talks to the Kubernetes API, so these tests bypass it
|
|
22
|
+
(via __new__) and set up only the attributes is_image_allowed() actually reads:
|
|
23
|
+
`_policies`, `client` and `log`. This keeps `_is_regex_pattern` (a real
|
|
24
|
+
staticmethod on the class) working, unlike a fully mocked `self` would.
|
|
25
|
+
"""
|
|
26
|
+
|
|
27
|
+
def setUp(self):
|
|
28
|
+
self.manager = ContainerManager.__new__(ContainerManager)
|
|
29
|
+
self.manager.log = _get_null_logger()
|
|
30
|
+
self.manager.client = MagicMock()
|
|
31
|
+
self.manager._policies = {}
|
|
32
|
+
|
|
33
|
+
@staticmethod
|
|
34
|
+
def _task_info(store_id: int | None = None) -> dict:
|
|
35
|
+
task_info = {"init_org": {"id": 1}, "init_user": {"id": 1}}
|
|
36
|
+
if store_id is not None:
|
|
37
|
+
task_info["algorithm_store"] = {"id": store_id}
|
|
38
|
+
return task_info
|
|
39
|
+
|
|
40
|
+
def test_allowed_algorithms_exact_match(self):
|
|
41
|
+
"""Baseline regression check for the (unchanged) allowed_algorithms policy."""
|
|
42
|
+
self.manager._policies = {
|
|
43
|
+
NodePolicy.ALLOWED_ALGORITHMS.value: ["some/image:tag"]
|
|
44
|
+
}
|
|
45
|
+
|
|
46
|
+
self.assertTrue(
|
|
47
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info())
|
|
48
|
+
)
|
|
49
|
+
self.assertFalse(
|
|
50
|
+
self.manager.is_image_allowed("some/other-image:tag", self._task_info())
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
def test_allowed_algorithm_stores_verified(self):
|
|
54
|
+
"""
|
|
55
|
+
The store URL is whitelisted and the store itself confirms the image is a
|
|
56
|
+
registered, approved algorithm there - the image is allowed.
|
|
57
|
+
"""
|
|
58
|
+
self.manager._policies = {
|
|
59
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"]
|
|
60
|
+
}
|
|
61
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
62
|
+
"url": "https://store.example",
|
|
63
|
+
"api_path": "/api",
|
|
64
|
+
}
|
|
65
|
+
self.manager.client.algorithm_store.get_algorithm.return_value = {
|
|
66
|
+
"image": "some/image:tag",
|
|
67
|
+
"digest": "abc123",
|
|
68
|
+
}
|
|
69
|
+
|
|
70
|
+
self.assertTrue(
|
|
71
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
|
|
72
|
+
)
|
|
73
|
+
self.manager.client.algorithm_store.get_algorithm.assert_called_once_with(
|
|
74
|
+
"some/image:tag"
|
|
75
|
+
)
|
|
76
|
+
|
|
77
|
+
def test_allowed_algorithm_stores_not_confirmed_by_store(self):
|
|
78
|
+
"""
|
|
79
|
+
The store URL is whitelisted, but the store does not confirm the image is
|
|
80
|
+
registered/approved there (e.g. it genuinely isn't, the store hasn't been
|
|
81
|
+
upgraded to support node requests, or it could not be reached) - the image
|
|
82
|
+
is denied. This is the fix for GHSA-j43v-wwjr-9grj: previously the store URL
|
|
83
|
+
match alone was enough.
|
|
84
|
+
"""
|
|
85
|
+
self.manager._policies = {
|
|
86
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"]
|
|
87
|
+
}
|
|
88
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
89
|
+
"url": "https://store.example",
|
|
90
|
+
"api_path": "/api",
|
|
91
|
+
}
|
|
92
|
+
self.manager.client.algorithm_store.get_algorithm.return_value = None
|
|
93
|
+
|
|
94
|
+
self.assertFalse(
|
|
95
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
|
|
96
|
+
)
|
|
97
|
+
|
|
98
|
+
def test_allowed_algorithm_stores_url_not_whitelisted(self):
|
|
99
|
+
"""
|
|
100
|
+
The store's URL isn't on the whitelist at all - the image is denied without
|
|
101
|
+
even asking the store to verify the image (no wasted network call).
|
|
102
|
+
"""
|
|
103
|
+
self.manager._policies = {
|
|
104
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://trusted.example"]
|
|
105
|
+
}
|
|
106
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
107
|
+
"url": "https://untrusted.example",
|
|
108
|
+
"api_path": "/api",
|
|
109
|
+
}
|
|
110
|
+
|
|
111
|
+
self.assertFalse(
|
|
112
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
|
|
113
|
+
)
|
|
114
|
+
self.manager.client.algorithm_store.get_algorithm.assert_not_called()
|
|
115
|
+
|
|
116
|
+
def test_allowed_algorithm_stores_localhost_translated_in_k8s(self):
|
|
117
|
+
"""
|
|
118
|
+
A store registered under a `localhost` URL (e.g. running on the same local
|
|
119
|
+
dev machine as the node) is translated to an address reachable from inside
|
|
120
|
+
the node's pod before the live verification call, when V6_K8S_NODE_NAME
|
|
121
|
+
indicates a local k8s dev deployment.
|
|
122
|
+
"""
|
|
123
|
+
self.manager._policies = {
|
|
124
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["http://localhost:7601"]
|
|
125
|
+
}
|
|
126
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
127
|
+
"url": "http://localhost:7601",
|
|
128
|
+
"api_path": "/api",
|
|
129
|
+
}
|
|
130
|
+
self.manager.client.algorithm_store.get_algorithm.return_value = {
|
|
131
|
+
"image": "some/image:tag",
|
|
132
|
+
"digest": "abc123",
|
|
133
|
+
}
|
|
134
|
+
|
|
135
|
+
with patch.dict(os.environ, {"V6_K8S_NODE_NAME": "docker-desktop"}):
|
|
136
|
+
self.assertTrue(
|
|
137
|
+
self.manager.is_image_allowed(
|
|
138
|
+
"some/image:tag", self._task_info(store_id=1)
|
|
139
|
+
)
|
|
140
|
+
)
|
|
141
|
+
|
|
142
|
+
self.assertEqual(
|
|
143
|
+
self.manager.client.algorithm_store.url,
|
|
144
|
+
"http://host.docker.internal:7601/api",
|
|
145
|
+
)
|
|
146
|
+
|
|
147
|
+
def test_allowed_algorithm_stores_localhost_unchanged_without_k8s_node(self):
|
|
148
|
+
"""
|
|
149
|
+
Without V6_K8S_NODE_NAME set (the default outside local dev deployments), a
|
|
150
|
+
`localhost` store URL is used as-is - a real deployment never has the node
|
|
151
|
+
and the store on the same machine, so there is nothing to translate.
|
|
152
|
+
"""
|
|
153
|
+
self.manager._policies = {
|
|
154
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["http://localhost:7601"]
|
|
155
|
+
}
|
|
156
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
157
|
+
"url": "http://localhost:7601",
|
|
158
|
+
"api_path": "/api",
|
|
159
|
+
}
|
|
160
|
+
self.manager.client.algorithm_store.get_algorithm.return_value = {
|
|
161
|
+
"image": "some/image:tag",
|
|
162
|
+
"digest": "abc123",
|
|
163
|
+
}
|
|
164
|
+
|
|
165
|
+
with patch.dict(os.environ, {}, clear=False):
|
|
166
|
+
os.environ.pop("V6_K8S_NODE_NAME", None)
|
|
167
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
|
|
168
|
+
|
|
169
|
+
self.assertEqual(
|
|
170
|
+
self.manager.client.algorithm_store.url, "http://localhost:7601/api"
|
|
171
|
+
)
|
|
172
|
+
|
|
173
|
+
def test_allow_either_whitelist_or_store(self):
|
|
174
|
+
"""
|
|
175
|
+
Baseline regression check: with allow_either_whitelist_or_store, an image
|
|
176
|
+
that fails the allowed_algorithms whitelist but is confirmed by the
|
|
177
|
+
(whitelisted) store is still allowed.
|
|
178
|
+
"""
|
|
179
|
+
self.manager._policies = {
|
|
180
|
+
NodePolicy.ALLOWED_ALGORITHMS.value: ["some/other-image:tag"],
|
|
181
|
+
NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"],
|
|
182
|
+
"allow_either_whitelist_or_store": True,
|
|
183
|
+
}
|
|
184
|
+
self.manager.client.algorithm_store.get.return_value = {
|
|
185
|
+
"url": "https://store.example",
|
|
186
|
+
"api_path": "/api",
|
|
187
|
+
}
|
|
188
|
+
self.manager.client.algorithm_store.get_algorithm.return_value = {
|
|
189
|
+
"image": "some/image:tag",
|
|
190
|
+
"digest": "abc123",
|
|
191
|
+
}
|
|
192
|
+
|
|
193
|
+
self.assertTrue(
|
|
194
|
+
self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
|
|
195
|
+
)
|
|
196
|
+
|
|
197
|
+
|
|
198
|
+
|
|
199
|
+
class TestGetImagePullPolicy:
|
|
200
|
+
def _manager_with_policies(self, policies: dict) -> ContainerManager:
|
|
201
|
+
manager = ContainerManager.__new__(ContainerManager)
|
|
202
|
+
manager._policies = policies
|
|
203
|
+
return manager
|
|
204
|
+
|
|
205
|
+
def test_require_algorithm_pull_false_means_if_not_present(self):
|
|
206
|
+
manager = self._manager_with_policies({"require_algorithm_pull": False})
|
|
207
|
+
assert manager._get_image_pull_policy() == "IfNotPresent"
|
|
208
|
+
|
|
209
|
+
def test_require_algorithm_pull_true_means_always(self):
|
|
210
|
+
manager = self._manager_with_policies({"require_algorithm_pull": True})
|
|
211
|
+
assert manager._get_image_pull_policy() == "Always"
|
|
212
|
+
|
|
213
|
+
def test_require_algorithm_pull_absent_defaults_to_always(self):
|
|
214
|
+
manager = self._manager_with_policies({})
|
|
215
|
+
assert manager._get_image_pull_policy() == "Always"
|
|
216
|
+
|
|
217
|
+
|
|
218
|
+
|
|
219
|
+
if __name__ == "__main__":
|
|
220
|
+
unittest.main()
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
from unittest.mock import MagicMock, patch
|
|
2
|
+
|
|
3
|
+
import pytest
|
|
4
|
+
from kubernetes import client as k8s_client
|
|
5
|
+
|
|
6
|
+
from vantage6.node import Node
|
|
7
|
+
from vantage6.node.k8s.network_isolation import (
|
|
8
|
+
_probe_pod_reached_target,
|
|
9
|
+
validate_algorithm_isolation,
|
|
10
|
+
)
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class TestProbePodReachedTarget:
|
|
14
|
+
def _pod_with_exit_code(self, exit_code: int) -> k8s_client.V1Pod:
|
|
15
|
+
return k8s_client.V1Pod(
|
|
16
|
+
status=k8s_client.V1PodStatus(
|
|
17
|
+
container_statuses=[
|
|
18
|
+
k8s_client.V1ContainerStatus(
|
|
19
|
+
name="probe",
|
|
20
|
+
image="curlimages/curl:latest",
|
|
21
|
+
image_id="curlimages/curl@sha256:abc",
|
|
22
|
+
ready=False,
|
|
23
|
+
restart_count=0,
|
|
24
|
+
state=k8s_client.V1ContainerState(
|
|
25
|
+
terminated=k8s_client.V1ContainerStateTerminated(
|
|
26
|
+
exit_code=exit_code
|
|
27
|
+
)
|
|
28
|
+
),
|
|
29
|
+
)
|
|
30
|
+
]
|
|
31
|
+
)
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
def test_exit_code_zero_means_not_isolated(self):
|
|
35
|
+
assert _probe_pod_reached_target(self._pod_with_exit_code(0)) is True
|
|
36
|
+
|
|
37
|
+
def test_non_zero_exit_code_means_isolated(self):
|
|
38
|
+
assert _probe_pod_reached_target(self._pod_with_exit_code(28)) is False
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
class TestValidateAlgorithmIsolation:
|
|
42
|
+
@patch("vantage6.node.k8s.network_isolation._wait_for_probe_pod")
|
|
43
|
+
def test_reachable_target_returns_not_isolated(self, mock_wait):
|
|
44
|
+
mock_wait.return_value = TestProbePodReachedTarget()._pod_with_exit_code(0)
|
|
45
|
+
core_api = MagicMock()
|
|
46
|
+
|
|
47
|
+
isolated, message = validate_algorithm_isolation(
|
|
48
|
+
core_api=core_api,
|
|
49
|
+
task_namespace="vantage6-tasks",
|
|
50
|
+
log=MagicMock(),
|
|
51
|
+
)
|
|
52
|
+
|
|
53
|
+
assert isolated is False
|
|
54
|
+
assert "example.com" in message
|
|
55
|
+
core_api.delete_namespaced_pod.assert_called_once()
|
|
56
|
+
|
|
57
|
+
@patch("vantage6.node.k8s.network_isolation._wait_for_probe_pod")
|
|
58
|
+
def test_blocked_target_returns_isolated(self, mock_wait):
|
|
59
|
+
mock_wait.return_value = TestProbePodReachedTarget()._pod_with_exit_code(28)
|
|
60
|
+
core_api = MagicMock()
|
|
61
|
+
|
|
62
|
+
isolated, message = validate_algorithm_isolation(
|
|
63
|
+
core_api=core_api,
|
|
64
|
+
task_namespace="vantage6-tasks",
|
|
65
|
+
log=MagicMock(),
|
|
66
|
+
)
|
|
67
|
+
|
|
68
|
+
assert isolated is True
|
|
69
|
+
assert "verified" in message
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
class TestNodeProductionGating:
|
|
73
|
+
def test_production_true_exits_on_isolation_failure(self):
|
|
74
|
+
ctx = MagicMock()
|
|
75
|
+
ctx.config = {"production": True}
|
|
76
|
+
|
|
77
|
+
with (
|
|
78
|
+
patch("vantage6.node.validate_required_env_vars"),
|
|
79
|
+
patch("vantage6.node.ContainerManager") as mock_cm_cls,
|
|
80
|
+
):
|
|
81
|
+
mock_cm = mock_cm_cls.return_value
|
|
82
|
+
mock_cm.ensure_task_namespace.return_value = True
|
|
83
|
+
mock_cm.validate_algorithm_isolation.return_value = (
|
|
84
|
+
False,
|
|
85
|
+
"probe failed",
|
|
86
|
+
)
|
|
87
|
+
with patch("vantage6.node.exit", side_effect=SystemExit(1)) as mock_exit:
|
|
88
|
+
with (
|
|
89
|
+
patch.object(Node, "_setup_node_client", return_value=MagicMock()),
|
|
90
|
+
pytest.raises(SystemExit),
|
|
91
|
+
):
|
|
92
|
+
Node(ctx)
|
|
93
|
+
mock_exit.assert_called_once_with(1)
|
|
94
|
+
|
|
95
|
+
def test_production_false_warns_on_isolation_failure(self, caplog):
|
|
96
|
+
ctx = MagicMock()
|
|
97
|
+
ctx.config = {"production": False, "debug": {}}
|
|
98
|
+
|
|
99
|
+
def _connect_socket(self):
|
|
100
|
+
self.socketIO = MagicMock()
|
|
101
|
+
|
|
102
|
+
with (
|
|
103
|
+
patch("vantage6.node.validate_required_env_vars"),
|
|
104
|
+
patch("vantage6.node.ContainerManager") as mock_cm_cls,
|
|
105
|
+
):
|
|
106
|
+
mock_cm = mock_cm_cls.return_value
|
|
107
|
+
mock_cm.ensure_task_namespace.return_value = True
|
|
108
|
+
mock_cm.validate_algorithm_isolation.return_value = (
|
|
109
|
+
False,
|
|
110
|
+
"probe failed",
|
|
111
|
+
)
|
|
112
|
+
with patch("vantage6.node.exit") as mock_exit:
|
|
113
|
+
with (
|
|
114
|
+
patch.object(Node, "_setup_node_client", return_value=MagicMock()),
|
|
115
|
+
patch.object(Node, "authenticate"),
|
|
116
|
+
patch.object(Node, "setup_encryption"),
|
|
117
|
+
patch.object(Node, "connect_to_socket", _connect_socket),
|
|
118
|
+
patch.object(Node, "start_processing_threads"),
|
|
119
|
+
patch("vantage6.node.Thread"),
|
|
120
|
+
caplog.at_level("WARNING"),
|
|
121
|
+
):
|
|
122
|
+
Node(ctx)
|
|
123
|
+
mock_exit.assert_not_called()
|
|
124
|
+
assert "probe failed" in caplog.text
|
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
import unittest
|
|
2
|
+
from unittest.mock import Mock, patch
|
|
3
|
+
|
|
4
|
+
from vantage6.common.enum import RunStatus
|
|
5
|
+
from vantage6.node import Node
|
|
6
|
+
from vantage6.node.k8s.container_manager import ContainerManager
|
|
7
|
+
from vantage6.node.k8s.data_classes import KilledRun, Result
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class TestContainerManagerLogSharing(unittest.TestCase):
|
|
11
|
+
def _make_manager(self, share_algorithm_logs: bool) -> ContainerManager:
|
|
12
|
+
manager = ContainerManager.__new__(ContainerManager)
|
|
13
|
+
manager.log = Mock()
|
|
14
|
+
manager.share_algorithm_logs = share_algorithm_logs
|
|
15
|
+
manager.socket_io = Mock()
|
|
16
|
+
manager.client = Mock()
|
|
17
|
+
manager.client.collaboration_id = 7
|
|
18
|
+
manager.task_namespace = "test-namespace"
|
|
19
|
+
manager.core_api = Mock()
|
|
20
|
+
pod = Mock()
|
|
21
|
+
pod.metadata = Mock()
|
|
22
|
+
pod.metadata.name = "pod-1"
|
|
23
|
+
manager.core_api.list_namespaced_pod.return_value.items = [pod]
|
|
24
|
+
return manager
|
|
25
|
+
|
|
26
|
+
def test_log_stream_does_not_emit_when_sharing_disabled(self):
|
|
27
|
+
manager = self._make_manager(False)
|
|
28
|
+
run_io = Mock(run_id=42, container_name="algo-container")
|
|
29
|
+
fake_watch = Mock()
|
|
30
|
+
fake_watch.stream.return_value = ["hello from algorithm"]
|
|
31
|
+
|
|
32
|
+
with patch("vantage6.node.k8s.container_manager.watch.Watch", return_value=fake_watch):
|
|
33
|
+
manager._ContainerManager__log_stream(run_io=run_io, task_id=99)
|
|
34
|
+
|
|
35
|
+
manager.socket_io.emit.assert_not_called()
|
|
36
|
+
|
|
37
|
+
def test_log_stream_emits_when_sharing_enabled(self):
|
|
38
|
+
manager = self._make_manager(True)
|
|
39
|
+
run_io = Mock(run_id=42, container_name="algo-container")
|
|
40
|
+
fake_watch = Mock()
|
|
41
|
+
fake_watch.stream.return_value = ["hello from algorithm"]
|
|
42
|
+
|
|
43
|
+
with patch("vantage6.node.k8s.container_manager.watch.Watch", return_value=fake_watch):
|
|
44
|
+
manager._ContainerManager__log_stream(run_io=run_io, task_id=99)
|
|
45
|
+
|
|
46
|
+
manager.socket_io.emit.assert_called_once_with(
|
|
47
|
+
"algorithm_log",
|
|
48
|
+
{
|
|
49
|
+
"collaboration_id": 7,
|
|
50
|
+
"run_id": 42,
|
|
51
|
+
"task_id": 99,
|
|
52
|
+
"log": "hello from algorithm",
|
|
53
|
+
},
|
|
54
|
+
namespace="/tasks",
|
|
55
|
+
)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
class TestFinishedTaskLogSharing(unittest.TestCase):
|
|
59
|
+
def _make_node(self, share_algorithm_logs: bool) -> Node:
|
|
60
|
+
node = Node.__new__(Node)
|
|
61
|
+
node.log = Mock()
|
|
62
|
+
node.config = {"share_algorithm_logs": share_algorithm_logs}
|
|
63
|
+
node.k8s_container_manager = Mock()
|
|
64
|
+
node.k8s_container_manager.process_next_completed_run.return_value = Result(
|
|
65
|
+
run_id=1,
|
|
66
|
+
task_id=2,
|
|
67
|
+
logs="sensitive algorithm logs",
|
|
68
|
+
data="output-data",
|
|
69
|
+
status=RunStatus.COMPLETED,
|
|
70
|
+
parent_id=None,
|
|
71
|
+
)
|
|
72
|
+
node.client = Mock()
|
|
73
|
+
node.client.request.side_effect = [
|
|
74
|
+
{"task": {"id": 2}},
|
|
75
|
+
{"init_org": {"id": 5}},
|
|
76
|
+
]
|
|
77
|
+
node.client.whoami = Mock(id_=1, organization_id=1)
|
|
78
|
+
node.client.collaboration_id = 1
|
|
79
|
+
node.socketIO = Mock()
|
|
80
|
+
return node
|
|
81
|
+
|
|
82
|
+
def _run_one_iteration(self, node: Node) -> None:
|
|
83
|
+
# __send_updates_finished_tasks loops forever with a trailing
|
|
84
|
+
# time.sleep(1); make sleep raise so the loop stops after one pass.
|
|
85
|
+
with patch("vantage6.node.time.sleep", side_effect=InterruptedError):
|
|
86
|
+
with self.assertRaises(InterruptedError):
|
|
87
|
+
node._Node__send_updates_finished_tasks()
|
|
88
|
+
|
|
89
|
+
def test_logs_redacted_when_sharing_disabled(self):
|
|
90
|
+
node = self._make_node(share_algorithm_logs=False)
|
|
91
|
+
self._run_one_iteration(node)
|
|
92
|
+
|
|
93
|
+
patch_call = node.client.run.patch.call_args
|
|
94
|
+
self.assertEqual(
|
|
95
|
+
patch_call.kwargs["data"]["log"],
|
|
96
|
+
"Node does not allow sharing algorithm logs",
|
|
97
|
+
)
|
|
98
|
+
|
|
99
|
+
def test_logs_shared_when_sharing_enabled(self):
|
|
100
|
+
node = self._make_node(share_algorithm_logs=True)
|
|
101
|
+
self._run_one_iteration(node)
|
|
102
|
+
|
|
103
|
+
patch_call = node.client.run.patch.call_args
|
|
104
|
+
self.assertEqual(
|
|
105
|
+
patch_call.kwargs["data"]["log"], "sensitive algorithm logs"
|
|
106
|
+
)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
class TestKillContainersLogSharing(unittest.TestCase):
|
|
110
|
+
def _make_node(self, share_algorithm_logs: bool) -> Node:
|
|
111
|
+
node = Node.__new__(Node)
|
|
112
|
+
node.log = Mock()
|
|
113
|
+
node.config = {"share_algorithm_logs": share_algorithm_logs}
|
|
114
|
+
node.client = Mock()
|
|
115
|
+
node.client.collaboration_id = 7
|
|
116
|
+
node.client.whoami = Mock(id_=1)
|
|
117
|
+
node.k8s_container_manager = Mock()
|
|
118
|
+
node.k8s_container_manager.kill_algorithm_runs.return_value = [
|
|
119
|
+
KilledRun(
|
|
120
|
+
run_id=1, task_id=2, parent_id=None, logs="sensitive kill logs"
|
|
121
|
+
)
|
|
122
|
+
]
|
|
123
|
+
return node
|
|
124
|
+
|
|
125
|
+
def test_kill_logs_not_sent_when_sharing_disabled(self):
|
|
126
|
+
node = self._make_node(share_algorithm_logs=False)
|
|
127
|
+
node.kill_containers({"collaboration_id": 7})
|
|
128
|
+
|
|
129
|
+
node.client.run.patch.assert_not_called()
|
|
130
|
+
|
|
131
|
+
def test_kill_logs_sent_when_sharing_enabled(self):
|
|
132
|
+
node = self._make_node(share_algorithm_logs=True)
|
|
133
|
+
node.kill_containers({"collaboration_id": 7})
|
|
134
|
+
|
|
135
|
+
node.client.run.patch.assert_called_once_with(
|
|
136
|
+
id_=1, data={"log": "sensitive kill logs"}
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
|
|
140
|
+
if __name__ == "__main__":
|
|
141
|
+
unittest.main()
|