[air] Fix NoneType error loading TorchCheckpoint through from_uri.

[xwjiang2010]

• encode/decode logic of TorchCheckpoint is only relevant in the context of training
• it's currently intercepting __getstate__ and __setstate__
• training worker first sends TorchCheckpoint to Trainable (through ser/deser). The TorchCheckpoint that Trainable receives is already "decoded"
• the last statement may not hold true anymore if training worker uploads TorchCheckpoint directly with to_uri, at which point we should revisit where to put the encode/decode logic.

Signed-off-by: xwjiang2010 xwjiang2010@gmail.com

Why are these changes needed?

Related issue number

Closes #32284

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 :(

[Yard1]

I don't think we need this, the model should be already encoded when it was put into the URI. I think the only thing that needs to be fixed is the missing check for the case where _data_dict is None:

    def __setstate__(self, state: dict):
        if state.get("_data_dict", None):
            state = state.copy()
            state["_data_dict"] = self._decode_data_dict(state["_data_dict"])
        super().__setstate__(state)

[krfricke]

This looks good!

[xwjiang2010]

I don't think we need this, the model should be already encoded when it was put into the URI. I think the only thing that needs to be fixed is the missing check for the case where _data_dict is None:

    def __setstate__(self, state: dict):
        if state.get("_data_dict", None):
            state = state.copy()
            state["_data_dict"] = self._decode_data_dict(state["_data_dict"])
        super().__setstate__(state)

isn't this for decoding ? We encode it when putting into the URI and decode it when retrieving?

[xwjiang2010]

waiting for https://github.com/ray-project/ray/pull/32479/files to be landed first...