Skip to content

Commit f6d0788

Browse files
committed
Add unit tests to LeptonExecutor
Signed-Off-By: Robert Clark <roclark@nvidia.com>
1 parent 2071ad1 commit f6d0788

File tree

3 files changed

+608
-2
lines changed

3 files changed

+608
-2
lines changed

.vscode/settings.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66
"editor.defaultFormatter": "charliermarsh.ruff",
77
"editor.codeActionsOnSave": {
88
"source.fixAll": "always",
9-
"source.organizeImports": "always"
9+
"source.organizeImports": "never"
1010
},
1111
"python.analysis.typeCheckingMode": "standard",
1212
"ruff.lint.args": [

nemo_run/core/execution/lepton.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,8 +133,17 @@ def _node_group_id(self, client: APIClient) -> DedicatedNodeGroup:
133133
from the user with the list of node groups. Assumes there are no duplicate node groups.
134134
"""
135135
node_groups = client.nodegroup.list_all()
136+
if len(node_groups) < 1:
137+
raise RuntimeError(
138+
"No node groups found in cluster. Ensure Lepton workspace has at least one node group."
139+
)
136140
node_group_map = {ng.metadata.name: ng for ng in node_groups}
137-
node_group_id = node_group_map[self.node_group]
141+
try:
142+
node_group_id = node_group_map[self.node_group]
143+
except KeyError:
144+
raise RuntimeError(
145+
"Could not find node group that matches requested ID in the Lepton workspace. Ensure your requested node group exists."
146+
)
138147
return node_group_id
139148

140149
def _valid_node_ids(self, node_group_id: DedicatedNodeGroup, client: APIClient) -> Set:

0 commit comments

Comments
 (0)