Skip to content

Commit 94bce5b

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

2 files changed

Lines changed: 607 additions & 1 deletion

File tree

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)