Skip to content

Commit 9062e8a

Browse files
authored
fix key error in dpgen2 download (#297)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **Bug Fixes** - Improved artifact retrieval to prevent errors when accessing missing artifacts, resulting in a more robust download process. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: zjgemi <[email protected]>
1 parent 53d93f9 commit 9062e8a

File tree

3 files changed

+8
-2
lines changed

3 files changed

+8
-2
lines changed

dpgen2/utils/download_dpgen2_artifacts.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -359,9 +359,9 @@ def _dl_step_item(
359359
[step_key, io, name] = item.split(global_step_def_split)
360360
pref = _item_path(prefix, item)
361361
if io in ["input"]:
362-
target = step.inputs.artifacts[name]
362+
target = step.inputs.artifacts.get(name)
363363
elif io in ["output"]:
364-
target = step.outputs.artifacts[name]
364+
target = step.outputs.artifacts.get(name)
365365
else:
366366
raise RuntimeError("unknown io style {io}")
367367
try:

tests/utils/test_dl_dpgen2_arti.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,9 @@
2828

2929

3030
class MockedArti:
31+
def get(self, key):
32+
return self.__getitem__(key)
33+
3134
def __getitem__(
3235
self,
3336
key,

tests/utils/test_dl_dpgen2_arti_by_def.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,9 @@
3030

3131

3232
class MockedArti:
33+
def get(self, key):
34+
return self.__getitem__(key)
35+
3336
def __getitem__(
3437
self,
3538
key,

0 commit comments

Comments
 (0)