Skip to content

Commit 0a23a23

Browse files
committed
Automatically mirror connections in user's sidecar conainters
1 parent 87bc4cc commit 0a23a23

8 files changed

Lines changed: 285 additions & 4 deletions

File tree

cli/polyaxon/_auxiliaries/sidecar.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -143,6 +143,15 @@ class V1PolyaxonSidecarContainer(BaseSchemaModel):
143143
>>> sidecar:
144144
>>> monitorSpec: true
145145
```
146+
147+
### noConnections
148+
149+
Whether to prevent mounting connections and context volumes to sidecars, default `None` (false).
150+
151+
```yaml
152+
>>> sidecar:
153+
>>> noConnections: true
154+
```
146155
"""
147156

148157
_IDENTIFIER = "polyaxon_sidecar"
@@ -156,6 +165,7 @@ class V1PolyaxonSidecarContainer(BaseSchemaModel):
156165
sync_interval: Optional[IntOrRef] = Field(alias="syncInterval", default=None)
157166
monitor_logs: Optional[BoolOrRef] = Field(alias="monitorLogs", default=None)
158167
monitor_spec: Optional[BoolOrRef] = Field(alias="monitorSpec", default=None)
168+
no_connections: Optional[BoolOrRef] = Field(alias="noConnections", default=None)
159169
resources: Optional[Union[Dict[str, Any], RefField]] = None
160170

161171
def get_image(self):

cli/polyaxon/_docker/converter/base/base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,11 @@ def get_replica_resource(
238238
plugins=plugins,
239239
artifacts_store=artifacts_store,
240240
sidecar_containers=sidecars,
241+
connections=connections,
242+
connection_by_names=connection_by_names,
243+
secrets=secrets,
244+
config_maps=config_maps,
245+
kv_env_vars=kv_env_vars,
241246
log_level=plugins.log_level,
242247
volumes=volumes,
243248
)

cli/polyaxon/_flow/run/ray/autoscaler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -80,4 +80,4 @@ class V1RayAutoscalerOptions(BaseSchemaModel):
8080

8181
upscaling_mode: Optional[str] = Field(alias="upscalingMode", default=None)
8282
image_pull_policy: Optional[str] = Field(alias="imagePullPolicy", default=None)
83-
resources: Optional[Union[k8s_schemas.V1ResourceRequirements, RefField]] = None
83+
resources: Optional[Union[k8s_schemas.V1ResourceRequirements, RefField]] = None

cli/polyaxon/_k8s/converter/base/base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,11 @@ def get_replica_resource(
185185
plugins=plugins,
186186
artifacts_store=artifacts_store,
187187
sidecar_containers=sidecars,
188+
connections=connections,
189+
connection_by_names=connection_by_names,
190+
secrets=secrets,
191+
config_maps=config_maps,
192+
kv_env_vars=kv_env_vars,
188193
log_level=plugins.log_level,
189194
)
190195

cli/polyaxon/_k8s/converter/base/sidecar.py

Lines changed: 92 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
1-
from typing import List, Optional
1+
from typing import Dict, Iterable, List, Optional
22

33
from clipped.utils.lists import to_list
44

55
from polyaxon._auxiliaries import V1PolyaxonSidecarContainer
6-
from polyaxon._connections import V1Connection
6+
from polyaxon._connections import V1Connection, V1ConnectionResource
77
from polyaxon._containers.names import SIDECAR_CONTAINER
88
from polyaxon._env_vars.keys import ENV_KEYS_ARTIFACTS_STORE_NAME, ENV_KEYS_CONTAINER_ID
99
from polyaxon._flow import V1Plugins
@@ -165,3 +165,93 @@ def _get_sidecar_container(
165165
)
166166

167167
return cls._patch_container(container)
168+
169+
def _get_user_sidecar_container(
170+
self,
171+
sidecar: k8s_schemas.V1Container,
172+
plugins: V1Plugins,
173+
artifacts_store: Optional[V1Connection],
174+
connections: Optional[List[str]],
175+
connection_by_names: Optional[Dict[str, V1Connection]],
176+
secrets: Optional[Iterable[V1ConnectionResource]],
177+
config_maps: Optional[Iterable[V1ConnectionResource]],
178+
kv_env_vars: Optional[List[List]],
179+
) -> k8s_schemas.V1Container:
180+
"""Apply connections and context volumes to a user-defined sidecar container."""
181+
if plugins and plugins.sidecar and plugins.sidecar.no_connections:
182+
return sidecar
183+
184+
connections = connections or []
185+
connection_by_names = connection_by_names or {}
186+
secrets = secrets or []
187+
config_maps = config_maps or []
188+
189+
if artifacts_store and (
190+
not plugins.collect_artifacts or plugins.mount_artifacts_store
191+
):
192+
if artifacts_store.name not in connection_by_names:
193+
connection_by_names[artifacts_store.name] = artifacts_store
194+
if artifacts_store.name not in connections:
195+
connections.append(artifacts_store.name)
196+
197+
requested_connections = [connection_by_names[c] for c in connections]
198+
requested_config_maps = V1Connection.get_requested_resources(
199+
resources=config_maps,
200+
connections=requested_connections,
201+
resource_key="config_map",
202+
)
203+
requested_secrets = V1Connection.get_requested_resources(
204+
resources=secrets, connections=requested_connections, resource_key="secret"
205+
)
206+
207+
# Volume mounts
208+
volume_mounts = (
209+
self._get_mounts(
210+
use_auth_context=plugins.auth,
211+
use_artifacts_context=False,
212+
use_docker_context=plugins.docker,
213+
use_shm_context=plugins.shm,
214+
run_path=self.run_path,
215+
)
216+
if plugins
217+
else []
218+
)
219+
volume_mounts = volume_mounts + self._get_main_volume_mounts(
220+
plugins=plugins,
221+
init=[],
222+
connections=requested_connections,
223+
secrets=requested_secrets,
224+
config_maps=requested_config_maps,
225+
run_path=self.run_path,
226+
)
227+
228+
# Env vars
229+
env = self._get_main_env_vars(
230+
plugins=plugins,
231+
kv_env_vars=kv_env_vars,
232+
artifacts_store_name=artifacts_store.name if artifacts_store else None,
233+
connections=requested_connections,
234+
secrets=requested_secrets,
235+
config_maps=requested_config_maps,
236+
)
237+
env += self._get_resources_env_vars(sidecar.resources)
238+
if sidecar.env:
239+
sidecar.env = list(sidecar.env) + env
240+
else:
241+
sidecar.env = env
242+
243+
# Env from
244+
env_from = self._get_env_from_k8s_resources(
245+
secrets=requested_secrets, config_maps=requested_config_maps
246+
)
247+
if env_from:
248+
if sidecar.env_from:
249+
sidecar.env_from = list(sidecar.env_from) + env_from
250+
else:
251+
sidecar.env_from = env_from
252+
if sidecar.volume_mounts:
253+
sidecar.volume_mounts = list(sidecar.volume_mounts) + volume_mounts
254+
else:
255+
sidecar.volume_mounts = volume_mounts
256+
257+
return sidecar

cli/polyaxon/_local_process/converter/base/base.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -120,6 +120,11 @@ def get_replica_resource(
120120
plugins=plugins,
121121
artifacts_store=artifacts_store,
122122
sidecar_containers=sidecars,
123+
connections=connections,
124+
connection_by_names=connection_by_names,
125+
secrets=secrets,
126+
config_maps=config_maps,
127+
kv_env_vars=kv_env_vars,
123128
log_level=plugins.log_level,
124129
volumes=volumes,
125130
)

cli/polyaxon/_runner/converter/converter.py

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -968,6 +968,11 @@ def get_sidecar_containers(
968968
plugins: V1Plugins,
969969
artifacts_store: V1Connection,
970970
sidecar_containers: List[Container],
971+
connections: Optional[List[str]] = None,
972+
connection_by_names: Optional[Dict[str, V1Connection]] = None,
973+
secrets: Optional[Iterable[V1ConnectionResource]] = None,
974+
config_maps: Optional[Iterable[V1ConnectionResource]] = None,
975+
kv_env_vars: Optional[List[List]] = None,
971976
log_level: Optional[str] = None,
972977
volumes: Optional[List[k8s_schemas.V1Volume]] = None,
973978
) -> List[Container]:
@@ -987,9 +992,37 @@ def get_sidecar_containers(
987992
run_path=self.run_path,
988993
)
989994
containers = to_list(polyaxon_sidecar_container, check_none=True)
990-
containers += sidecar_containers
995+
996+
# Handle user sidecars with connections
997+
for sidecar in sidecar_containers:
998+
_sidecar = self._get_user_sidecar_container(
999+
sidecar=sidecar,
1000+
plugins=plugins,
1001+
artifacts_store=artifacts_store,
1002+
connections=connections,
1003+
connection_by_names=connection_by_names,
1004+
secrets=secrets,
1005+
config_maps=config_maps,
1006+
kv_env_vars=kv_env_vars,
1007+
)
1008+
if _sidecar:
1009+
containers.append(sidecar)
1010+
9911011
return [self._ensure_container(c, volumes) for c in containers]
9921012

1013+
def _get_user_sidecar_container(
1014+
self,
1015+
sidecar: Container,
1016+
plugins: V1Plugins,
1017+
artifacts_store: Optional[V1Connection],
1018+
connections: Optional[List[str]],
1019+
connection_by_names: Optional[Dict[str, V1Connection]],
1020+
secrets: Optional[Iterable[V1ConnectionResource]],
1021+
config_maps: Optional[Iterable[V1ConnectionResource]],
1022+
kv_env_vars: Optional[List[List]],
1023+
) -> Optional[Container]:
1024+
return None
1025+
9931026
def get_main_container(
9941027
self,
9951028
main_container: Container,
Lines changed: 133 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,133 @@
1+
from polyaxon._auxiliaries import V1PolyaxonSidecarContainer
2+
from polyaxon._connections import (
3+
V1ClaimConnection,
4+
V1Connection,
5+
V1ConnectionKind,
6+
)
7+
from polyaxon._flow import V1Plugins
8+
from polyaxon._k8s import k8s_schemas
9+
from tests.test_k8s.test_converters.base import BaseConverterTest
10+
11+
12+
class TestSidecarConnections(BaseConverterTest):
13+
def test_get_sidecars_with_connections(self):
14+
store = V1Connection(
15+
name="test_claim",
16+
kind=V1ConnectionKind.VOLUME_CLAIM,
17+
schema_=V1ClaimConnection(
18+
mount_path="/claim/path", volume_claim="claim", read_only=True
19+
),
20+
)
21+
plugins = V1Plugins.get_or_create(
22+
V1Plugins(
23+
collect_logs=False,
24+
collect_artifacts=False,
25+
auth=True,
26+
sidecar=V1PolyaxonSidecarContainer(no_connections=False),
27+
)
28+
)
29+
sidecar = k8s_schemas.V1Container(name="sidecar", image="foo")
30+
31+
containers = self.converter.get_sidecar_containers(
32+
plugins=plugins,
33+
artifacts_store=None,
34+
sidecar_containers=[sidecar],
35+
polyaxon_sidecar=V1PolyaxonSidecarContainer(image="sidecar/sidecar"),
36+
connections=[store.name],
37+
connection_by_names={store.name: store},
38+
)
39+
40+
# Check that sidecar has volume mounts
41+
# Since collect_artifacts=False and collect_logs=False, polyaxon sidecar is not created
42+
# Only user sidecar is returned
43+
assert len(containers) == 1
44+
user_sidecar = containers[0]
45+
assert "sidecar" in user_sidecar.name
46+
47+
# Check mounts
48+
mount_paths = [m.mount_path for m in user_sidecar.volume_mounts]
49+
assert "/claim/path" in mount_paths
50+
51+
# Check auth context (since plugins.auth=True)
52+
# Auth context mount path is /plx-context/configs
53+
from polyaxon._contexts import paths as ctx_paths
54+
55+
assert ctx_paths.CONTEXT_MOUNT_CONFIGS in mount_paths
56+
57+
def test_get_sidecars_with_multiple_connections(self):
58+
"""Test that sidecars get volume mounts for multiple connections."""
59+
store1 = V1Connection(
60+
name="test_claim",
61+
kind=V1ConnectionKind.VOLUME_CLAIM,
62+
schema_=V1ClaimConnection(
63+
mount_path="/claim/path", volume_claim="claim", read_only=True
64+
),
65+
)
66+
store2 = V1Connection(
67+
name="test_claim2",
68+
kind=V1ConnectionKind.VOLUME_CLAIM,
69+
schema_=V1ClaimConnection(
70+
mount_path="/claim/path2", volume_claim="claim2", read_only=False
71+
),
72+
)
73+
plugins = V1Plugins.get_or_create(
74+
V1Plugins(collect_logs=False, collect_artifacts=False, auth=False)
75+
)
76+
sidecar = k8s_schemas.V1Container(name="sidecar", image="foo")
77+
78+
containers = self.converter.get_sidecar_containers(
79+
plugins=plugins,
80+
artifacts_store=None,
81+
sidecar_containers=[sidecar],
82+
polyaxon_sidecar=V1PolyaxonSidecarContainer(image="sidecar/sidecar"),
83+
connections=[store1.name, store2.name],
84+
connection_by_names={store1.name: store1, store2.name: store2},
85+
)
86+
87+
assert len(containers) == 1
88+
user_sidecar = containers[0]
89+
90+
# Check both connections are mounted
91+
mount_paths = [m.mount_path for m in user_sidecar.volume_mounts]
92+
assert "/claim/path" in mount_paths
93+
assert "/claim/path2" in mount_paths
94+
95+
def test_get_sidecars_without_mirroring(self):
96+
"""Test that sidecars do NOT get volume mounts when mirror_connections is False or default."""
97+
store = V1Connection(
98+
name="test_claim",
99+
kind=V1ConnectionKind.VOLUME_CLAIM,
100+
schema_=V1ClaimConnection(
101+
mount_path="/claim/path", volume_claim="claim", read_only=True
102+
),
103+
)
104+
# plugins without sidecar config (defaults to mirror_connections=None/False)
105+
plugins = V1Plugins.get_or_create(
106+
V1Plugins(
107+
collect_logs=False,
108+
collect_artifacts=False,
109+
auth=True,
110+
sidecar=V1PolyaxonSidecarContainer(no_connections=True),
111+
)
112+
)
113+
sidecar = k8s_schemas.V1Container(name="sidecar", image="foo")
114+
115+
containers = self.converter.get_sidecar_containers(
116+
plugins=plugins,
117+
artifacts_store=None,
118+
sidecar_containers=[sidecar],
119+
polyaxon_sidecar=V1PolyaxonSidecarContainer(image="sidecar/sidecar"),
120+
connections=[store.name],
121+
connection_by_names={store.name: store},
122+
)
123+
124+
assert len(containers) == 1
125+
user_sidecar = containers[0]
126+
127+
# Check NO connection mounts are present
128+
mount_paths = (
129+
[m.mount_path for m in user_sidecar.volume_mounts]
130+
if user_sidecar.volume_mounts
131+
else []
132+
)
133+
assert "/claim/path" not in mount_paths

0 commit comments

Comments
 (0)