`orbax-checkpoint` broken on Windows since 0.11.33

[JeppeKlitgaard]

I wanted to try on Flax for size, but am unfortunately hoping to run it on a Windows system.

Since Flax depends on orbax-export, which in turn depends on orbax-checkpoint, this no longer works on Windows due to the recently added dependency uvloop (see: ebffb79) , which is Unix-only. There is a fork that also supports Windows (Winloop).

Reproducer (on Windows):

uv pip install orbax-checkpoint
Resolved 22 packages in 276ms
  × Failed to build `uvloop==0.22.1`
  ├─▶ The build backend returned an error                                                                                                     
  ╰─▶ Call to `setuptools.build_meta.build_wheel` failed (exit code: 1)

      [stderr]
      Traceback (most recent call last):
        File "<string>", line 14, in <module>
          requires = get_requires_for_build({})
        File "C:\Users\jeppe\AppData\Local\uv\cache\builds-v0\.tmpXe8j9j\Lib\site-packages\setuptools\build_meta.py", line 333, in
      get_requires_for_build_wheel
          return self._get_build_requires(config_settings, requirements=[])
                 ~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        File "C:\Users\jeppe\AppData\Local\uv\cache\builds-v0\.tmpXe8j9j\Lib\site-packages\setuptools\build_meta.py", line 301, in
      _get_build_requires
          self.run_setup()
          ~~~~~~~~~~~~~~^^
        File "C:\Users\jeppe\AppData\Local\uv\cache\builds-v0\.tmpXe8j9j\Lib\site-packages\setuptools\build_meta.py", line 317, in run_setup  
          exec(code, locals())
          ~~~~^^^^^^^^^^^^^^^^
        File "<string>", line 8, in <module>
          import setuptools.build_meta as backend
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
      RuntimeError: uvloop does not support Windows at the moment

      hint: This usually indicates a problem with the package or the build environment.
  help: `uvloop` (v0.22.1) was included because `orbax-checkpoint` (v0.11.33) depends on `uvloop`

Workaround:

Installing or pinning to orbax-checkpoint<=0.11.32

[luzechao]

Can someone in google change orbax to use winloop for windows and uvloop for unix-like OS? Or do you rely on community effort?

[JustinPan-goog]

@JeppeKlitgaard @luzechao Thanks for bringing up the issue. We generally try to maintain the compatibility cross-platform, but unfortunately certain edge cases of orbax-checkpoint may not work on Windows, as noted in our documentation. We also do not currently accept contributions from the open source community.

Since we might not have the bandwidth to fully import and test out for winloop library, I think the alternative we could do is to fallback to nest_asyncio on win platform, which is not ideal but should unblock the workflow. I will start looking into this fix. cc @cpgaffney1

[JeppeKlitgaard]

@JustinPan-goog, thank you for taking a look. I appreciate that some edge-cases may not be supported on Windows, which seems increasingly nightmarish to deal with. That being said, the ramifications of orbax-checkpoint not being installable on Windows has far-reaching consequences for the whole Jax ML ecosystem, at least while orbax-export is a required dependency of Flax.

[JustinPan-goog]

Quick update, here is my draft for a fallback fix: #2984

Thanks for the input, and we completely see your point! While orbax-checkpoint doesn't officially accept external contributions rn, we are open to exceptional cases for unexpected breakage like this. If you’d like to fork the repo and implement a more robust solution for other platforms (such as winloop support), we’d be happy to review and help merging it into the main with full acknowledgment