Skip to content

Dask client handling in TPI/SS: unguarded scatter breaks serial runs; SS re-scatters params every residual evaluation #1211

Description

@vahid-ahmadi

1. run_TPI raises when client=None

client.scatter at TPI.py:969 is called unconditionally at function-body indentation:

    # Scatter the parameters
    scattered_p_future = client.scatter(p, broadcast=True)

SS.inner_loop guards the equivalent block with if client: (SS.py:292), but TPI does not. So run_TPI(p, client=None) fails with AttributeError: 'NoneType' object has no attribute 'scatter' before the loop is entered — which makes the serial fallback at TPI.py:1039 unreachable in exactly the case it exists for.

Fix is to wrap the scatter/restore block in if client: and pass p directly when there is no client, mirroring SS.

2. SS.inner_loop re-scatters the parameters on every residual evaluation

client.scatter(p, broadcast=True) sits inside inner_loop (SS.py:304), which is called once per outer residual evaluation by SS_fsolve. The Specifications object is therefore re-serialised and re-broadcast to every worker dozens of times per steady-state solve, while p does not change across those calls.

run_TPI already does this correctly — scatter once before the loop, reuse the future. Hoisting the SS scatter into run_SS (or caching the future on the client) would match that pattern and remove the repeated broadcast.

I can open a PR for (1) straight away; happy to do (2) in the same PR or separately, whichever you prefer.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions