[AIR] Raise ValueError if TorchCheckpoint can't serialize model

[bveeramani]

Depends on:

• [AIR] Add TorchCheckpoint.from_state_dict #27970
• [AIR] Correct error doesn't propagate #27922

Why are these changes needed?

If a model is defined in the top-level directory, then Torch can't serialize the model. This PR adds a clearer error message.

Before:
If you're using more than one worker:

RuntimeError: Some workers returned results while others didn't. Make sure that session.report() (legacy API:train.report() and train.save_checkpoint()) are called the same number of times on all workers.

If you're using one worker:

_pickle.PicklingError: Can't pickle <class 'main.Identity'>: attribute lookup Identity on main failed

After:

ValueError: TorchCheckpoint can't serialize model of type Identity because Identity is defined in the top-level environment. To work around this error, call TorchCheckpoint.from_state_dict instead of TorchCheckpoint.from_model. Alternatively, move the definition of Identity to a different module.

Related issue number

Closes #27922

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[bveeramani]

This test makes assumptions about the implementation of from_model, but I wasn't sure how else to test it.

[richardliaw]

This particular catch-warn mechanism seems very niche.

Can we figure out how to get the _pickle.PicklingError: Can't pickle <class 'main.Identity'>: attribute lookup Identity on main failed error to raise even on multiple machines workers?

[bveeramani]

This particular catch-warn mechanism seems very niche.

Can we figure out how to get the _pickle.PicklingError: Can't pickle <class 'main.Identity'>: attribute lookup Identity on main failed error to raise even on multiple machines?

What error did you get when you ran it on multiple machines?

[richardliaw]

Sorry typo -- meant workers.

[bveeramani]

This particular catch-warn mechanism seems very niche.

@richardliaw can you elaborate on this? The try-catch mechanism works with any number of workers, and it works with both Jupyter notebooks and Python programs.

Also, the error _pickle.PicklingError: Can't pickle <class 'main.Identity'>: attribute lookup Identity on main failed isn't actionable. It isn't clear what the user needs to do.

[matthewdeng]

@bveeramani I think Richard's point is that we should start with generalizable problem of surfacing the correct error and not

RuntimeError: Some workers returned results while others didn't. Make sure that session.report() (legacy API:train.report() and train.save_checkpoint()) are called the same number of times on all workers.

For this particular error, it is niche but I'm wondering if this should instead be handled as part of a try/catch.

[bveeramani]

@bveeramani I think Richard's point is that we should start with generalizable problem of surfacing the correct error and not

RuntimeError: Some workers returned results while others didn't. Make sure that session.report() (legacy API:train.report() and train.save_checkpoint()) are called the same number of times on all workers.

I agree we should fix the general problem. That being said, I still think we should raise a ValueError. The "correct error" is confusing and unactionable.

For this particular error, it is niche but I'm wondering if this should instead be handled as part of a try/catch.

How could we handle this as part of a try/catch?

[matthewdeng]

If we can fix the general problem, then we have access to the original error. We then catch this particular error and raise the ValueError.

[stale]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

• If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[bveeramani]

Closing for now. Can re-open when #27922 is fixed