Skip to content

Return value of worker_init_fn is ignored by MPRS #1083

@sehoffmann

Description

@sehoffmann

🐛 Describe the bug

def abc(x):
    return -1

def worker_init(dp, worker_info):
    return dp.map(abc)
    
pipe = dp.iter.IterableWrapper(range(10))

rs = MultiProcessingReadingService(num_workers=1, worker_init_fn=worker_init)
dl = DataLoader2(pipe, reading_service=rs)

for x in dl:
    print(x)

Output:

0
1
2
3
4
5
6
7
8
9

Expected:

-1
-1
-1
-1
-1
-1
-1
-1
-1
-1

Notice that worker_init_fn MUST return a datapipe, and process_init_fn also accounts for that, but MPRS simply ignores it. The signature of the function heavily implies that worker_init_fn should be able to modify the datapipe in a per-worker manner; and arguably this would be a useful feature.

Versions

main & torch nightly

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions