Skip to content

Commit a2b3e9a

Browse files
Add node selector (#99)
* add node selectors * lint * add node selector to local config * update snpashots * ui test * remove comment * update unit tests * update readme.md * update readme * update readme
1 parent 5678759 commit a2b3e9a

File tree

20 files changed

+224
-37
lines changed

20 files changed

+224
-37
lines changed

README.md

Lines changed: 44 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -218,7 +218,6 @@ c.JupyterHub.services.extend(
218218
'{"label": "Medium", "cpu": 4, "memory": 4}',
219219
"--machine_profiles",
220220
'{"label": "Large", "cpu": 8, "memory": 8}'
221-
222221
],
223222
"oauth_no_confirm": True,
224223
}
@@ -228,6 +227,50 @@ c.JupyterHub.services.extend(
228227

229228
![image](https://github.com/plasmabio/tljh-repo2docker/assets/4451292/c1f0231e-a02d-41dc-85e0-97a97ffa0311)
230229

230+
### Node Selector
231+
232+
`tljh-repo2docker` allows specifying node selectors to control which Kubernetes nodes user environments are scheduled on. This can be useful for assigning workloads to specific nodes based on hardware characteristics like GPUs, SSD storage, or other node labels.
233+
234+
## Configuring Node Selectors
235+
236+
To configure node selectors, add the `--node_selector` argument in the service definition:
237+
238+
```python
239+
c.JupyterHub.services.extend(
240+
[
241+
{
242+
"name": "tljh_repo2docker",
243+
"url": "http://127.0.0.1:6789",
244+
"command": [
245+
sys.executable,
246+
"-m",
247+
"tljh_repo2docker",
248+
"--ip",
249+
"127.0.0.1",
250+
"--port",
251+
"6789",
252+
"--node_selector",
253+
'{"gpu": {"description": "GPU availability", "values": ["yes", "no"]},'
254+
' "ssd": {"description": "SSD availability", "values": ["yes", "no"]}}'
255+
],
256+
"oauth_no_confirm": True,
257+
}
258+
]
259+
)
260+
```
261+
262+
This ensures that workloads are scheduled only on nodes that meet the specified criteria.
263+
264+
## Accessing Node Selector in Spawner
265+
266+
The node selector information is passed through the metadata field of `user_options` and can be accessed in the `start` method of the spawner:
267+
268+
```python
269+
user_options["metadata"]["node_selector"]
270+
```
271+
272+
![node_selector](https://github.com/user-attachments/assets/046bee93-2c7c-4e42-a9a0-94ade6f191d9)
273+
231274
### Extra documentation
232275

233276
`tljh-repo2docker` is currently developed as part of the [Plasma project](https://github.com/plasmabio/plasma).

src/environments/App.tsx

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,11 @@ import { ThemeProvider, createTheme } from '@mui/material/styles';
44

55
import { IEnvironmentData } from './types';
66
import { EnvironmentList } from './EnvironmentList';
7-
import { IMachineProfile, NewEnvironmentDialog } from './NewEnvironmentDialog';
7+
import {
8+
IMachineProfile,
9+
INodeSelector,
10+
NewEnvironmentDialog
11+
} from './NewEnvironmentDialog';
812
import { AxiosContext } from '../common/AxiosContext';
913
import { useEffect, useMemo, useState } from 'react';
1014
import { AxiosClient } from '../common/axiosclient';
@@ -16,6 +20,7 @@ export interface IAppProps {
1620
default_cpu_limit: string;
1721
default_mem_limit: string;
1822
machine_profiles: IMachineProfile[];
23+
node_selector: INodeSelector;
1924
use_binderhub: boolean;
2025
repo_providers?: { label: string; value: string }[];
2126
}
@@ -75,6 +80,7 @@ export default function App(props: IAppProps) {
7580
default_cpu_limit={props.default_cpu_limit}
7681
default_mem_limit={props.default_mem_limit}
7782
machine_profiles={props.machine_profiles}
83+
node_selector={props.node_selector}
7884
use_binderhub={props.use_binderhub}
7985
repo_providers={props.repo_providers}
8086
/>

src/environments/NewEnvironmentDialog.tsx

Lines changed: 73 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import {
1111
MenuItem,
1212
OutlinedTextFieldProps,
1313
Select,
14+
TextField,
1415
Typography
1516
} from '@mui/material';
1617
import {
@@ -31,10 +32,20 @@ export interface IMachineProfile {
3132
cpu: string;
3233
memory: string;
3334
}
35+
interface INodeSelectorOption {
36+
description: string;
37+
values: string[];
38+
}
39+
40+
export interface INodeSelector {
41+
[key: string]: INodeSelectorOption;
42+
}
43+
3444
export interface INewEnvironmentDialogProps {
3545
default_cpu_limit: string;
3646
default_mem_limit: string;
3747
machine_profiles: IMachineProfile[];
48+
node_selector: INodeSelector;
3849
use_binderhub: boolean;
3950
repo_providers?: { label: string; value: string }[];
4051
}
@@ -49,6 +60,7 @@ interface IFormValues {
4960
buildargs?: string;
5061
username?: string;
5162
password?: string;
63+
node_selector?: { [key: string]: string | undefined };
5264
}
5365
const commonInputProps: OutlinedTextFieldProps = {
5466
autoFocus: true,
@@ -76,10 +88,14 @@ function _NewEnvironmentDialog(props: INewEnvironmentDialogProps) {
7688

7789
const [formValues, setFormValues] = useState<IFormValues>({});
7890
const updateFormValue = useCallback(
79-
(key: keyof IFormValues, value: string | number) => {
80-
setFormValues(old => {
81-
return { ...old, [key]: value };
82-
});
91+
(
92+
key: keyof IFormValues,
93+
value: string | number | { [key: string]: string | undefined }
94+
) => {
95+
setFormValues(old => ({
96+
...old,
97+
[key]: value
98+
}));
8399
},
84100
[setFormValues]
85101
);
@@ -89,6 +105,15 @@ function _NewEnvironmentDialog(props: INewEnvironmentDialogProps) {
89105

90106
const [selectedProfile, setSelectedProfile] = useState<number>(0);
91107
const [selectedProvider, setSelectedProvider] = useState<number>(0);
108+
const [selectedNodeSelectors, setSelectedNodeSelectors] = useState<{
109+
[key: string]: string;
110+
}>(() => {
111+
const initialSelected: { [key: string]: string } = {};
112+
Object.entries(props.node_selector).forEach(([key, option]) => {
113+
initialSelected[key] = option.values[0] || '';
114+
});
115+
return initialSelected;
116+
});
92117

93118
const onMachineProfileChange = useCallback(
94119
(value?: string | number) => {
@@ -118,18 +143,39 @@ function _NewEnvironmentDialog(props: INewEnvironmentDialogProps) {
118143
},
119144
[props.repo_providers, updateFormValue]
120145
);
146+
147+
const onNodeSelectorChange = useCallback(
148+
(key: string, value: string) => {
149+
if (value !== undefined) {
150+
setSelectedNodeSelectors(prevState => {
151+
const newState = { ...prevState, [key]: value };
152+
updateFormValue('node_selector', newState);
153+
return newState;
154+
});
155+
}
156+
},
157+
[updateFormValue]
158+
);
159+
121160
useEffect(() => {
122161
if (props.machine_profiles.length > 0) {
123162
onMachineProfileChange(0);
124163
}
125164
if (props.repo_providers && props.repo_providers.length > 0) {
126165
onRepoProviderChange(0);
127166
}
167+
if (props.node_selector) {
168+
Object.entries(props.node_selector).forEach(([key, option]) => {
169+
onNodeSelectorChange(key, option.values[0]);
170+
});
171+
}
128172
}, [
129173
props.machine_profiles,
130174
props.repo_providers,
175+
props.node_selector,
131176
onMachineProfileChange,
132-
onRepoProviderChange
177+
onRepoProviderChange,
178+
onNodeSelectorChange
133179
]);
134180
const MemoryCpuSelector = useMemo(() => {
135181
return (
@@ -186,6 +232,27 @@ function _NewEnvironmentDialog(props: INewEnvironmentDialogProps) {
186232
);
187233
}, [props.machine_profiles, selectedProfile, onMachineProfileChange]);
188234

235+
const NodeSelectorDropdown = useMemo(() => {
236+
return Object.entries(props.node_selector).map(([key, option]) => (
237+
<FormControl key={key} fullWidth sx={{ marginTop: '8px' }}>
238+
<TextField
239+
id={`${key}-select`}
240+
value={selectedNodeSelectors[key]}
241+
label={key + option.description && `(${option.description})`}
242+
size="small"
243+
select
244+
onChange={e => onNodeSelectorChange(key, e.target.value)}
245+
>
246+
{option.values.map((val: string) => (
247+
<MenuItem key={val} value={val}>
248+
{val}
249+
</MenuItem>
250+
))}
251+
</TextField>
252+
</FormControl>
253+
));
254+
}, [props.node_selector, selectedNodeSelectors, onNodeSelectorChange]);
255+
189256
return (
190257
<Fragment>
191258
<Box sx={{ display: 'flex', flexDirection: 'row-reverse' }}>
@@ -292,6 +359,7 @@ function _NewEnvironmentDialog(props: INewEnvironmentDialogProps) {
292359
{props.machine_profiles.length > 0
293360
? MachineProfileSelector
294361
: MemoryCpuSelector}
362+
{props.node_selector && NodeSelectorDropdown}
295363
{!props.use_binderhub && (
296364
<Fragment>
297365
<Divider

src/environments/main.tsx

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ if (rootElement) {
1717
default_cpu_limit: '2',
1818
default_mem_limit: '2G',
1919
machine_profiles: [],
20+
node_selector: {},
2021
use_binderhub: false
2122
};
2223
if (dataElement) {

tljh_repo2docker/app.py

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,6 +108,19 @@ def _logo_file_default(self):
108108
def _default_log_level(self):
109109
return logging.INFO
110110

111+
node_selector = Dict(
112+
config=True,
113+
help="""
114+
The dictionary Selector labels used to match the Nodes where Pods will be launched.
115+
116+
Default is None and means it will be launched in any available Node.
117+
118+
For example to match the Nodes that have a label of `disktype: ssd` use::
119+
120+
c.KubeSpawner.node_selector = {'disktype': 'ssd'}
121+
""",
122+
)
123+
111124
machine_profiles = List(
112125
default_value=[], trait=Dict, config=True, help="Pre-defined machine profiles"
113126
)
@@ -162,6 +175,7 @@ def _default_log_level(self):
162175
"default_memory_limit": "TljhRepo2Docker.default_memory_limit",
163176
"default_cpu_limit": "TljhRepo2Docker.default_cpu_limit",
164177
"machine_profiles": "TljhRepo2Docker.machine_profiles",
178+
"node_selector": "TljhRepo2Docker.node_selector",
165179
"binderhub_url": "TljhRepo2Docker.binderhub_url",
166180
"db_url": "TljhRepo2Docker.db_url",
167181
}
@@ -193,6 +207,7 @@ def init_settings(self) -> tp.Dict:
193207
default_mem_limit=self.default_memory_limit,
194208
default_cpu_limit=self.default_cpu_limit,
195209
machine_profiles=self.machine_profiles,
210+
node_selector=self.node_selector,
196211
binderhub_url=self.binderhub_url,
197212
repo_providers=self.repo_providers,
198213
)

tljh_repo2docker/binderhub_builder.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ async def post(self):
7777
memory = data["memory"]
7878
cpu = data["cpu"]
7979
provider = data["provider"]
80+
node_selector = data["node_selector"]
8081
if len(repo) == 0:
8182
raise web.HTTPError(400, "Repository is empty")
8283

@@ -115,7 +116,7 @@ async def post(self):
115116
status=BuildStatusType.BUILDING,
116117
log="",
117118
image_meta=ImageMetadataType(
118-
display_name=name, repo=repo, ref=ref, cpu_limit=cpu, mem_limit=memory
119+
display_name=name, repo=repo, ref=ref, cpu_limit=cpu, mem_limit=memory, node_selector= node_selector
119120
),
120121
)
121122
self.set_status(200)

tljh_repo2docker/builder.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ async def post(self):
3939
name = data["name"].lower()
4040
memory = data["memory"]
4141
cpu = data["cpu"]
42+
node_selector = data.get("node_selector", {})
4243
buildargs = data.get("buildargs", None)
4344
username = data.get("username", None)
4445
password = data.get("password", None)
@@ -71,7 +72,7 @@ async def post(self):
7172
raise web.HTTPError(400, "Invalid build argument format")
7273
extra_buildargs.append(barg)
7374
await build_image(
74-
repo, ref, name, memory, cpu, username, password, extra_buildargs
75+
repo, ref, node_selector, name, memory, cpu, username, password, extra_buildargs
7576
)
7677

7778
self.set_status(200)

tljh_repo2docker/database/schemas.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ class ImageMetadataType(BaseModel):
1616
ref: str
1717
cpu_limit: str
1818
mem_limit: str
19+
node_selector: dict
1920

2021

2122
class DockerImageCreateSchema(BaseModel):

tljh_repo2docker/docker.py

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from urllib.parse import urlparse
33

44
from aiodocker import Docker
5+
from tornado import web
56

67

78
async def list_images():
@@ -20,6 +21,7 @@ async def list_images():
2021
"display_name": image["Labels"]["tljh_repo2docker.display_name"],
2122
"mem_limit": image["Labels"]["tljh_repo2docker.mem_limit"],
2223
"cpu_limit": image["Labels"]["tljh_repo2docker.cpu_limit"],
24+
"node_selector": image["Labels"]["tljh_repo2docker.node_selector"],
2325
"status": "built",
2426
}
2527
for image in r2d_images
@@ -45,6 +47,7 @@ async def list_containers():
4547
"display_name": container["Labels"]["tljh_repo2docker.display_name"],
4648
"mem_limit": container["Labels"]["tljh_repo2docker.mem_limit"],
4749
"cpu_limit": container["Labels"]["tljh_repo2docker.cpu_limit"],
50+
"node_selector": container["Labels"]["tljh_repo2docker.node_selector"],
4851
"status": "building",
4952
}
5053
for container in r2d_containers
@@ -53,9 +56,32 @@ async def list_containers():
5356
return containers
5457

5558

59+
async def get_image_metadata(image_name):
60+
"""
61+
Retrieve metadata of a specific locally built Docker image.
62+
"""
63+
async with Docker() as docker:
64+
images = await docker.images.list(
65+
filters=json.dumps({"reference": [image_name]})
66+
)
67+
if not images:
68+
raise web.HTTPError(404, "Image not found")
69+
70+
image = images[0]
71+
return {
72+
"repo": image["Labels"].get("repo2docker.repo", ""),
73+
"ref": image["Labels"].get("repo2docker.ref", ""),
74+
"display_name": image["Labels"].get("tljh_repo2docker.display_name", ""),
75+
"mem_limit": image["Labels"].get("tljh_repo2docker.mem_limit", ""),
76+
"cpu_limit": image["Labels"].get("tljh_repo2docker.cpu_limit", ""),
77+
"node_selector": image["Labels"].get("tljh_repo2docker.node_selector", ""),
78+
}
79+
80+
5681
async def build_image(
5782
repo,
5883
ref,
84+
node_selector={},
5985
name="",
6086
memory=None,
6187
cpu=None,
@@ -86,6 +112,7 @@ async def build_image(
86112
f"tljh_repo2docker.image_name={image_name}",
87113
f"tljh_repo2docker.mem_limit={memory}",
88114
f"tljh_repo2docker.cpu_limit={cpu}",
115+
f"tljh_repo2docker.node_selector={node_selector}",
89116
]
90117
cmd = [
91118
"jupyter-repo2docker",
@@ -118,6 +145,7 @@ async def build_image(
118145
"tljh_repo2docker.display_name": name,
119146
"tljh_repo2docker.mem_limit": memory,
120147
"tljh_repo2docker.cpu_limit": cpu,
148+
"tljh_repo2docker.node_selector": json.dumps(node_selector),
121149
},
122150
"Volumes": {
123151
"/var/run/docker.sock": {

0 commit comments

Comments
 (0)