vantage6-node 5.0.2rc2__tar.gz → 5.0.3__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.
Files changed (30) hide show
  1. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/PKG-INFO +6 -6
  2. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/pyproject.toml +5 -5
  3. vantage6_node-5.0.3/tests_node/test_container_manager.py +219 -0
  4. vantage6_node-5.0.3/tests_node/test_network_isolation.py +124 -0
  5. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/tests_node/test_pod_status_mapping.py +0 -1
  6. vantage6_node-5.0.3/tests_node/test_share_algorithm_logs.py +141 -0
  7. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/__init__.py +110 -91
  8. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/cli/node.py +4 -4
  9. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/globals.py +17 -0
  10. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/container_manager.py +114 -42
  11. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/data_classes.py +8 -0
  12. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/exceptions.py +0 -12
  13. vantage6_node-5.0.3/vantage6/node/k8s/network_isolation.py +149 -0
  14. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/run_io.py +7 -7
  15. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/session_manager.py +3 -3
  16. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/proxy_server.py +8 -6
  17. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/socket.py +13 -3
  18. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/util/colorer.py +2 -5
  19. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/.gitignore +0 -0
  20. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/Makefile +0 -0
  21. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/README.md +0 -0
  22. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/dev_start.sh +0 -0
  23. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/tests/docker/__init__.py +0 -0
  24. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/tests_node/__init__.py +0 -0
  25. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/utest.py +0 -0
  26. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/dev_start.py +0 -0
  27. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/cli/__init__.py +0 -0
  28. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/enum.py +0 -0
  29. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/k8s/jobpod_state_to_run_status_mapper.py +0 -0
  30. {vantage6_node-5.0.2rc2 → vantage6_node-5.0.3}/vantage6/node/util/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.4
1
+ Metadata-Version: 2.5
2
2
  Name: vantage6-node
3
- Version: 5.0.2rc2
3
+ Version: 5.0.3
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.15.0
17
+ Requires-Dist: python-socketio==5.16.2
18
18
  Requires-Dist: requests==2.33.1
19
- Requires-Dist: vantage6-client==5.0.2rc2
20
- Requires-Dist: vantage6-common==5.0.2rc2
21
- Requires-Dist: vantage6==5.0.2rc2
19
+ Requires-Dist: vantage6-client==5.0.3
20
+ Requires-Dist: vantage6-common==5.0.3
21
+ Requires-Dist: vantage6==5.0.3
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.2rc2"
7
+ version = "5.0.3"
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.15.0",
28
+ "python-socketio==5.16.2",
29
29
  "requests==2.33.1",
30
- "vantage6==5.0.2rc2",
31
- "vantage6-client==5.0.2rc2",
32
- "vantage6-common==5.0.2rc2",
30
+ "vantage6==5.0.3",
31
+ "vantage6-client==5.0.3",
32
+ "vantage6-common==5.0.3",
33
33
  ]
34
34
 
35
35
  [project.optional-dependencies]
@@ -0,0 +1,219 @@
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
+
11
+ def _get_null_logger(name="null_logger_container_manager"):
12
+ logger = logging.getLogger(name)
13
+ logger.addHandler(logging.NullHandler())
14
+ logger.propagate = False
15
+ return logger
16
+
17
+
18
+ class TestIsImageAllowed(unittest.TestCase):
19
+ """
20
+ Tests for ContainerManager.is_image_allowed().
21
+
22
+ ContainerManager.__init__ talks to the Kubernetes API, so these tests bypass it
23
+ (via __new__) and set up only the attributes is_image_allowed() actually reads:
24
+ `_policies`, `client` and `log`. This keeps `_is_regex_pattern` (a real
25
+ staticmethod on the class) working, unlike a fully mocked `self` would.
26
+ """
27
+
28
+ def setUp(self):
29
+ self.manager = ContainerManager.__new__(ContainerManager)
30
+ self.manager.log = _get_null_logger()
31
+ self.manager.client = MagicMock()
32
+ self.manager._policies = {}
33
+
34
+ @staticmethod
35
+ def _task_info(store_id: int | None = None) -> dict:
36
+ task_info = {"init_org": {"id": 1}, "init_user": {"id": 1}}
37
+ if store_id is not None:
38
+ task_info["algorithm_store"] = {"id": store_id}
39
+ return task_info
40
+
41
+ def test_allowed_algorithms_exact_match(self):
42
+ """Baseline regression check for the (unchanged) allowed_algorithms policy."""
43
+ self.manager._policies = {
44
+ NodePolicy.ALLOWED_ALGORITHMS.value: ["some/image:tag"]
45
+ }
46
+
47
+ self.assertTrue(
48
+ self.manager.is_image_allowed("some/image:tag", self._task_info())
49
+ )
50
+ self.assertFalse(
51
+ self.manager.is_image_allowed("some/other-image:tag", self._task_info())
52
+ )
53
+
54
+ def test_allowed_algorithm_stores_verified(self):
55
+ """
56
+ The store URL is whitelisted and the store itself confirms the image is a
57
+ registered, approved algorithm there - the image is allowed.
58
+ """
59
+ self.manager._policies = {
60
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"]
61
+ }
62
+ self.manager.client.algorithm_store.get.return_value = {
63
+ "url": "https://store.example",
64
+ "api_path": "/api",
65
+ }
66
+ self.manager.client.algorithm_store.get_algorithm.return_value = {
67
+ "image": "some/image:tag",
68
+ "digest": "abc123",
69
+ }
70
+
71
+ self.assertTrue(
72
+ self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
73
+ )
74
+ self.manager.client.algorithm_store.get_algorithm.assert_called_once_with(
75
+ "some/image:tag"
76
+ )
77
+
78
+ def test_allowed_algorithm_stores_not_confirmed_by_store(self):
79
+ """
80
+ The store URL is whitelisted, but the store does not confirm the image is
81
+ registered/approved there (e.g. it genuinely isn't, the store hasn't been
82
+ upgraded to support node requests, or it could not be reached) - the image
83
+ is denied. This is the fix for GHSA-j43v-wwjr-9grj: previously the store URL
84
+ match alone was enough.
85
+ """
86
+ self.manager._policies = {
87
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"]
88
+ }
89
+ self.manager.client.algorithm_store.get.return_value = {
90
+ "url": "https://store.example",
91
+ "api_path": "/api",
92
+ }
93
+ self.manager.client.algorithm_store.get_algorithm.return_value = None
94
+
95
+ self.assertFalse(
96
+ self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
97
+ )
98
+
99
+ def test_allowed_algorithm_stores_url_not_whitelisted(self):
100
+ """
101
+ The store's URL isn't on the whitelist at all - the image is denied without
102
+ even asking the store to verify the image (no wasted network call).
103
+ """
104
+ self.manager._policies = {
105
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://trusted.example"]
106
+ }
107
+ self.manager.client.algorithm_store.get.return_value = {
108
+ "url": "https://untrusted.example",
109
+ "api_path": "/api",
110
+ }
111
+
112
+ self.assertFalse(
113
+ self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
114
+ )
115
+ self.manager.client.algorithm_store.get_algorithm.assert_not_called()
116
+
117
+ def test_allowed_algorithm_stores_localhost_translated_in_k8s(self):
118
+ """
119
+ A store registered under a `localhost` URL (e.g. running on the same local
120
+ dev machine as the node) is translated to an address reachable from inside
121
+ the node's pod before the live verification call, when V6_K8S_NODE_NAME
122
+ indicates a local k8s dev deployment.
123
+ """
124
+ self.manager._policies = {
125
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["http://localhost:7601"]
126
+ }
127
+ self.manager.client.algorithm_store.get.return_value = {
128
+ "url": "http://localhost:7601",
129
+ "api_path": "/api",
130
+ }
131
+ self.manager.client.algorithm_store.get_algorithm.return_value = {
132
+ "image": "some/image:tag",
133
+ "digest": "abc123",
134
+ }
135
+
136
+ with patch.dict(os.environ, {"V6_K8S_NODE_NAME": "docker-desktop"}):
137
+ self.assertTrue(
138
+ self.manager.is_image_allowed(
139
+ "some/image:tag", self._task_info(store_id=1)
140
+ )
141
+ )
142
+
143
+ self.assertEqual(
144
+ self.manager.client.algorithm_store.url,
145
+ "http://host.docker.internal:7601/api",
146
+ )
147
+
148
+ def test_allowed_algorithm_stores_localhost_unchanged_without_k8s_node(self):
149
+ """
150
+ Without V6_K8S_NODE_NAME set (the default outside local dev deployments), a
151
+ `localhost` store URL is used as-is - a real deployment never has the node
152
+ and the store on the same machine, so there is nothing to translate.
153
+ """
154
+ self.manager._policies = {
155
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["http://localhost:7601"]
156
+ }
157
+ self.manager.client.algorithm_store.get.return_value = {
158
+ "url": "http://localhost:7601",
159
+ "api_path": "/api",
160
+ }
161
+ self.manager.client.algorithm_store.get_algorithm.return_value = {
162
+ "image": "some/image:tag",
163
+ "digest": "abc123",
164
+ }
165
+
166
+ with patch.dict(os.environ, {}, clear=False):
167
+ os.environ.pop("V6_K8S_NODE_NAME", None)
168
+ self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
169
+
170
+ self.assertEqual(
171
+ self.manager.client.algorithm_store.url, "http://localhost:7601/api"
172
+ )
173
+
174
+ def test_allow_either_whitelist_or_store(self):
175
+ """
176
+ Baseline regression check: with allow_either_whitelist_or_store, an image
177
+ that fails the allowed_algorithms whitelist but is confirmed by the
178
+ (whitelisted) store is still allowed.
179
+ """
180
+ self.manager._policies = {
181
+ NodePolicy.ALLOWED_ALGORITHMS.value: ["some/other-image:tag"],
182
+ NodePolicy.ALLOWED_ALGORITHM_STORES.value: ["https://store.example"],
183
+ "allow_either_whitelist_or_store": True,
184
+ }
185
+ self.manager.client.algorithm_store.get.return_value = {
186
+ "url": "https://store.example",
187
+ "api_path": "/api",
188
+ }
189
+ self.manager.client.algorithm_store.get_algorithm.return_value = {
190
+ "image": "some/image:tag",
191
+ "digest": "abc123",
192
+ }
193
+
194
+ self.assertTrue(
195
+ self.manager.is_image_allowed("some/image:tag", self._task_info(store_id=1))
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
+ if __name__ == "__main__":
219
+ 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
@@ -63,7 +63,6 @@ class TestPodStatus(unittest.TestCase):
63
63
  )
64
64
 
65
65
  self.mock_pod = V1Pod(metadata={"name": ""}, status=self.mock_pod_status)
66
- pass
67
66
 
68
67
  def test_running_job_pod(self):
69
68
  self.mock_pod.status.phase = "Running"
@@ -0,0 +1,141 @@
1
+ import unittest
2
+ from unittest.mock import Mock, patch
3
+
4
+ from vantage6.common.enum import RunStatus
5
+
6
+ from vantage6.node import Node
7
+ from vantage6.node.k8s.container_manager import ContainerManager
8
+ from vantage6.node.k8s.data_classes import KilledRun, Result
9
+
10
+
11
+ class TestContainerManagerLogSharing(unittest.TestCase):
12
+ def _make_manager(self, share_algorithm_logs: bool) -> ContainerManager:
13
+ manager = ContainerManager.__new__(ContainerManager)
14
+ manager.log = Mock()
15
+ manager.share_algorithm_logs = share_algorithm_logs
16
+ manager.socket_io = Mock()
17
+ manager.client = Mock()
18
+ manager.client.collaboration_id = 7
19
+ manager.task_namespace = "test-namespace"
20
+ manager.core_api = Mock()
21
+ pod = Mock()
22
+ pod.metadata = Mock()
23
+ pod.metadata.name = "pod-1"
24
+ manager.core_api.list_namespaced_pod.return_value.items = [pod]
25
+ return manager
26
+
27
+ def test_log_stream_does_not_emit_when_sharing_disabled(self):
28
+ manager = self._make_manager(False)
29
+ run_io = Mock(run_id=42, container_name="algo-container")
30
+ fake_watch = Mock()
31
+ fake_watch.stream.return_value = ["hello from algorithm"]
32
+
33
+ with patch(
34
+ "vantage6.node.k8s.container_manager.watch.Watch", return_value=fake_watch
35
+ ):
36
+ manager._ContainerManager__log_stream(run_io=run_io, task_id=99)
37
+
38
+ manager.socket_io.emit.assert_not_called()
39
+
40
+ def test_log_stream_emits_when_sharing_enabled(self):
41
+ manager = self._make_manager(True)
42
+ run_io = Mock(run_id=42, container_name="algo-container")
43
+ fake_watch = Mock()
44
+ fake_watch.stream.return_value = ["hello from algorithm"]
45
+
46
+ with patch(
47
+ "vantage6.node.k8s.container_manager.watch.Watch", return_value=fake_watch
48
+ ):
49
+ manager._ContainerManager__log_stream(run_io=run_io, task_id=99)
50
+
51
+ manager.socket_io.emit.assert_called_once_with(
52
+ "algorithm_log",
53
+ {
54
+ "collaboration_id": 7,
55
+ "run_id": 42,
56
+ "task_id": 99,
57
+ "log": "hello from algorithm",
58
+ },
59
+ namespace="/tasks",
60
+ )
61
+
62
+
63
+ class TestFinishedTaskLogSharing(unittest.TestCase):
64
+ def _make_node(self, share_algorithm_logs: bool) -> Node:
65
+ node = Node.__new__(Node)
66
+ node.log = Mock()
67
+ node.config = {"share_algorithm_logs": share_algorithm_logs}
68
+ node.k8s_container_manager = Mock()
69
+ node.k8s_container_manager.process_next_completed_run.return_value = Result(
70
+ run_id=1,
71
+ task_id=2,
72
+ logs="sensitive algorithm logs",
73
+ data="output-data",
74
+ status=RunStatus.COMPLETED,
75
+ parent_id=None,
76
+ init_org_id=5,
77
+ )
78
+ node.client = Mock()
79
+ node.client.whoami = Mock(id_=1, organization_id=1)
80
+ node.client.collaboration_id = 1
81
+ node.socketIO = Mock()
82
+ return node
83
+
84
+ def _run_one_iteration(self, node: Node) -> None:
85
+ # __send_updates_finished_tasks loops forever with a trailing
86
+ # time.sleep(1); make sleep raise so the loop stops after one pass.
87
+ with (
88
+ patch("vantage6.node.time.sleep", side_effect=InterruptedError),
89
+ self.assertRaises(InterruptedError),
90
+ ):
91
+ node._Node__send_updates_finished_tasks()
92
+
93
+ def test_logs_redacted_when_sharing_disabled(self):
94
+ node = self._make_node(share_algorithm_logs=False)
95
+ self._run_one_iteration(node)
96
+
97
+ patch_call = node.client.run.patch.call_args
98
+ self.assertEqual(
99
+ patch_call.kwargs["data"]["log"],
100
+ "Node does not allow sharing algorithm logs",
101
+ )
102
+
103
+ def test_logs_shared_when_sharing_enabled(self):
104
+ node = self._make_node(share_algorithm_logs=True)
105
+ self._run_one_iteration(node)
106
+
107
+ patch_call = node.client.run.patch.call_args
108
+ self.assertEqual(patch_call.kwargs["data"]["log"], "sensitive algorithm logs")
109
+
110
+
111
+ class TestKillContainersLogSharing(unittest.TestCase):
112
+ def _make_node(self, share_algorithm_logs: bool) -> Node:
113
+ node = Node.__new__(Node)
114
+ node.log = Mock()
115
+ node.config = {"share_algorithm_logs": share_algorithm_logs}
116
+ node.client = Mock()
117
+ node.client.collaboration_id = 7
118
+ node.client.whoami = Mock(id_=1)
119
+ node.k8s_container_manager = Mock()
120
+ node.k8s_container_manager.kill_algorithm_runs.return_value = [
121
+ KilledRun(run_id=1, task_id=2, parent_id=None, logs="sensitive kill logs")
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()