PKG Add jax

[rth]

Looking at the list of dependencies jax should be easier to build than torch or tensorflow #175, and it would be useful to include it.

[horaceg]

Being able to have demos of JAX ML models running in the browser would be truly amazing !

[josephrocca]

Question: Given that the jaxlib repo is in a sub-folder within the jax github repo, and that jaxlib isn't a standalone pypi package, I'm assuming that I should be setting build.library:true in the meta.yaml, but the build instructions for jaxlib are simply:

# install requirements:
sudo apt-get update
sudo apt-get install -y python3-dev
pip install numpy wheel
# build CPU-only jaxlib:
python build/build.py
pip install dist/*.whl  # installs jaxlib (includes XLA)

This is a bit confusing for me because all the build scripts for other build.library:true use emconfigure and emmake directly (example), whereas in this case it seems like we just want to use python build/build.py but with all the build tools "swapped out" with emscripten versions, just like when building a normal package. So that makes me think that maybe I shouldn't be using build.library:true.

But my problem is that in build.script, the build tooling haven't been swapped out for the emscripten versions, so I'm not sure how to run this:

python build/build.py
pip install dist/*.whl  # installs jaxlib (includes XLA)

in such a way that it's all built with emscripten.

Here are the conda-forge meta.yaml files:

• https://github.com/conda-forge/jaxlib-feedstock/blob/main/recipe/meta.yaml
• https://github.com/conda-forge/jax-feedstock/blob/main/recipe/meta.yaml

For install_requires, jaxlib just has:

['scipy', 'numpy>=1.19', 'absl-py', 'flatbuffers >= 1.12, < 3.0']

and for jax itself:

['absl-py','numpy>=1.19','opt_einsum','scipy>=1.2.1','typing_extensions','etils[epath]']

But as of a few days ago the jax team is removing the absl-py requirement (making it optional).

opt_einsum.contract seems to be an API-compatible sped-up version of numpy.einsum (if I'm reading this correctly), so it sounds like that might be easy to swap out for now? Is there a simple way to "remap" the opt_einsum import? E.g. create a dummy "local" package with that name that overrides the pypi version and that exports numpy.einsum as contract? I know this would be simple in JS land, but I'm really clueless when it comes to Python packaging stuff.

I'm not sure what typing_extensions and etils[epath] are for. These might cause some trouble? Edit: actually, I think I was confusing build-requirements for actual requirements - it seems like typing_extensions is a pure Python lib? And I think etils[epath] just needs zipp (pure?) and importlib_resources (pure?).

CC @hoodmane

[hoodmane]

You certainly shouldn't be setting library: true because that's for static libraries. We might need to expand the build system a bit for this one. You could hypothetically build libjax in the build.script, this happens before the package itself is built. But in that case you need to ensure that libjax uses emcc with all the proper config.

[josephrocca]

Yeah I was wondering if there's a way to trigger the swap to the emscripten tools before running build.script rather than after? That said I don't really know how the whole build pipeline fits together. I tried this (doing emmake/emcc/etc. symlinking in built.script):

meta.yaml

package:
  name: jax
  version: 0.3.13
source:
  url: https://github.com/google/jax/archive/refs/tags/jax-v0.3.13.tar.gz
  sha256: d7ea74f8fd4ee69698f7350078f7987fe66333a0ac84b057dc792aa72b5f1813

build:
  script: |
    # install requirements:
    sudo apt-get update
    sudo apt-get install -y python3-dev
    pip install numpy wheel
    # build CPU-only jaxlib:
    mkdir -p /src/symlink-tools
    ln -sf "$(which emcc)" /src/symlink-tools/gcc
    ln -sf "$(which emmake)" /src/symlink-tools/make
    ln -sf "$(which emconfigure)" /src/symlink-tools/configure
    export PATH="/src/symlink-tools:$PATH"
    python build/build.py
    pip install dist/*.whl  # installs jaxlib (includes XLA)

test:
  imports:
    - jax
about:
  home: https://github.com/google/jax
  PyPI: https://pypi.org/project/jax
  summary: Differentiate, compile, and transform Numpy code.
  license: Apache-2.0

Results in this error when running python -m pyodide_build buildpkg meta.yaml:

...
...
...
Building XLA and installing it in the jaxlib source tree...
./bazel-5.1.1-linux-x86_64 run --verbose_failures=true --config=avx_posix --config=mkl_open_source_only :build_wheel -- --output_path=/src/packages/jax/build/jax-0.3.13/dist --cpu=x86_64
b''
Traceback (most recent call last):
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 528, in <module>
    main()
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 523, in main
    shell(command)
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/usr/local/lib/python3.10/subprocess.py", line 420, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/usr/local/lib/python3.10/subprocess.py", line 524, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.1-linux-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/src/packages/jax/build/jax-0.3.13/dist', '--cpu=x86_64']' returned non-zero exit status 1.

But that gives an error as shown in the above folded section. It builds fine if I leave out the symlink stuff (meta.yaml and full logs here), but obviously the resulting wheel wouldn't work in Pyodide because emscripten wasn't involved in the compilation of jaxlib. I'm guessing that the above error is at least partly because buildpkg does more than just symlink the emscripten versions of the build tools - at least going by this comment in the compile function in pyodide_build/buildpkg.py:

Runs pywasmcross for the package. The effect of this is to first run setup.py with compiler wrappers subbed in, which don't actually build the package but write the compile commands to build.log. Then we walk over build log and invoke the same set of commands but with some flags munged around or removed to make it work with emcc.

Is there anything that I could investigate/"prod at" here as a newbie? If someone (with some more experience/understanding of this stuff) could provide a high level dot-point plan that they think is the best approach, that would be super useful!

[hoodmane]

So there are quite a few layers of junk. The innermost layer is here:
https://github.com/pyodide/pyodide/blob/d3544a06a852f72c6a4343ece46625456cccc2e8/pyodide-build/pyodide_build/pypabuild.py#L97
The with replace_env(build_env) substitutes environment variables so that 1) when the compiler tools are invoked, they call the pywasmcross wrappers (which do some rewriting and then forward to emcc) and 2) when python is invoked as a build tool in this context, it sees the emscripten target sysconfigdata.

If you modify it to invoke subprocess.call(sys.executable, "build/build.py") inside of that, then it will run build.py in the cross build environment. Who knows if it will work, but it will at least be going through the full pywasmcross machinery.

I've been thinking about adding a variant of build.script that is run inside the cross build environment, and I think that would be the thing to do to add jax.

[josephrocca]

@hoodmane Thanks! Which set of environment variables should I be using for the subprocess.call? I've tried env, build_env and os.environ.copy(), but they all result in errors.

If I use build_env, os.environ.copy() or omit the env parameter, then I get this error:

File "/src/tools/symlinks/cc", line 18, in <module>
  PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
  raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'

If I use env (created in this line: with _IsolatedEnvBuilder() as env) then I get this error:

Traceback (most recent call last):
  File "/src/pyodide-build/pyodide_build/pypabuild.py", line 116, in build
    built = _build_in_isolated_env(
  File "/src/pyodide-build/pyodide_build/pypabuild.py", line 101, in _build_in_isolated_env
    subprocess.call([sys.executable, "build/build.py"], env=env)
  File "/usr/local/lib/python3.10/subprocess.py", line 345, in call
    with Popen(*popenargs, **kwargs) as p:
  File "/usr/local/lib/python3.10/subprocess.py", line 966, in __init__
    self._execute_child(args, executable, preexec_fn, close_fds,
  File "/usr/local/lib/python3.10/subprocess.py", line 1758, in _execute_child
    for k, v in env.items():
AttributeError: '_IsolatedEnvVenvPip' object has no attribute 'items'

Just to be clear:

subprocess.call([sys.executable, "build/build.py"], env=what_do_i_put_here)

[hoodmane]

It should be build_env, I'm not sure why it's missing PYWASMCROSS_ARGS. Could you push what you've done to a branch so I can look at it?

[hoodmane]

But maybe I should just make a branch adding a meta.yaml key for you and then you won't have to worry about mucking with internals.

[josephrocca]

But maybe I should just make a branch adding a meta.yaml key for you and then you won't have to worry about mucking with internals.

That would be great!

Could you push what you've done to a branch so I can look at it?

The only code change I've done is to add this:

import subprocess
subprocess.call([sys.executable, "build/build.py"], env=build_env)

below with replace_env(build_env): in pyodide/pyodide-build/pyodide_build/pypabuild.py.

Here's a branch with that change, and the packages/jax/meta.yaml that I'm using: https://github.com/josephrocca/pyodide/tree/patch-2

[hoodmane]

Hmmm yeah I don't know why you are getting the no PYWASMCROSS_ARGS error then it's assigned here:
https://github.com/pyodide/pyodide/blob/a47e9085b8b5d53413bfadd95fbc9e2a24a1ccc0/pyodide-build/pyodide_build/pywasmcross.py#L108
If it's inside of with replace_env you shouldn't need to pass the env argument or do anything it should just work ...

[josephrocca]

One really strange thing that might be a clue here is that if I have this (in pyodide/pyodide-build/pyodide_build/pypabuild.py):

with replace_env(build_env):
    import subprocess
    print("1-START-######################################################################")
    subprocess.call([sys.executable, "build/build.py"], env=build_env.copy())
    print("2-END-######################################################################")

and if I add this at the top of that build/build.py file:

print("@@@@@@@@@@@@")
import os
print(os.environ["PYWASMCROSS_ARGS"])
print("@@@@@@@@@@@@")

then I get this:

Logs

1-START-######################################################################
Starting local Bazel server and connecting to it...
INFO: Options provided by the client:
  Inherited 'common' options: --isatty=0 --terminal_columns=80
INFO: Reading rc options for 'run' from /src/packages/jax/build/jax-0.3.13/.bazelrc:
  Inherited 'common' options: --experimental_repo_remote_exec
INFO: Reading rc options for 'run' from /src/packages/jax/build/jax-0.3.13/.bazelrc:
  Inherited 'build' options: --apple_platform_type=macos --macos_minimum_os=10.9 --announce_rc --define open_source_build=true --spawn_strategy=standalone --enable_platform_specific_config --experimental_cc_shared_library --define=no_aws_support=true --define=no_gcp_support=true --define=no_hdfs_support=true --define=no_kafka_support=true --define=no_ignite_support=true --define=grpc_no_ares=true -c opt --config=short_logs --copt=-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.
INFO: Reading rc options for 'run' from /src/packages/jax/build/jax-0.3.13/.jax_configure.bazelrc:
  Inherited 'build' options: --strategy=Genrule=standalone --repo_env PYTHON_BIN_PATH=/usr/local/bin/python --action_env=PYENV_ROOT --python_path=/usr/local/bin/python --distinct_host_configuration=false
INFO: Found applicable config definition build:short_logs in file /src/packages/jax/build/jax-0.3.13/.bazelrc: --output_filter=DONT_MATCH_ANYTHING
INFO: Found applicable config definition build:avx_posix in file /src/packages/jax/build/jax-0.3.13/.bazelrc: --copt=-mavx --host_copt=-mavx
INFO: Found applicable config definition build:mkl_open_source_only in file /src/packages/jax/build/jax-0.3.13/.bazelrc: --define=tensorflow_mkldnn_contraction_kernel=1
INFO: Found applicable config definition build:linux in file /src/packages/jax/build/jax-0.3.13/.bazelrc: --config=posix --copt=-Wno-stringop-truncation --copt=-Wno-array-parameter
INFO: Found applicable config definition build:posix in file /src/packages/jax/build/jax-0.3.13/.bazelrc: --copt=-fvisibility=hidden --copt=-Wno-sign-compare --cxxopt=-std=c++14 --host_cxxopt=-std=c++14
Loading: 
Loading: 0 packages loaded
Loading: 0 packages loaded
Analyzing: target //build:build_wheel (1 packages loaded, 0 targets configured)
Analyzing: target //build:build_wheel (68 packages loaded, 301 targets configured)
Analyzing: target //build:build_wheel (182 packages loaded, 2810 targets configured)
Analyzing: target //build:build_wheel (211 packages loaded, 15224 targets configured)
INFO: Analyzed target //build:build_wheel (229 packages loaded, 16044 targets configured).
INFO: Found 1 target...

[4 / 9] [Prepa] BazelWorkspaceStatusAction stable-status.txt
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling inffast.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inffast.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inffast.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/inffast.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inffast.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling infback.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/infback.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/infback.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/infback.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/infback.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling gzwrite.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzwrite.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzwrite.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/gzwrite.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzwrite.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling adler32.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/adler32.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/adler32.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/adler32.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/adler32.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling zutil.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/zutil.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/zutil.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/zutil.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/zutil.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling deflate.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/deflate.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/deflate.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/deflate.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/deflate.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling inflate.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inflate.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inflate.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/inflate.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/inflate.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling crc32.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/crc32.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/crc32.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/crc32.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/crc32.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
ERROR: /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/external/zlib/BUILD.bazel:5:11: Compiling gzlib.c failed: (Exit 1): cc failed: error executing command 
  (cd /src/.docker_home/.cache/bazel/_bazel_joe/7e454c473e82025e727812785b12ccc7/execroot/__main__ && \
  exec env - \
    PATH=/src/tools/symlinks:/src/emsdk/emsdk:/src/emsdk/emsdk/upstream/emscripten:/src/emsdk/emsdk/node/14.18.2_64bit/bin:/src/node_modules/.bin/:/usr/local/bin:/usr/local/sbin:/usr/sbin:/usr/bin:/sbin:/bin:/src/packages/.artifacts/bin/ \
    PWD=/proc/self/cwd \
  /src/tools/symlinks/cc -U_FORTIFY_SOURCE -fstack-protector -Wall -Wthread-safety -Wself-assign -Wunused-but-set-parameter -Wno-free-nonheap-object -fcolor-diagnostics -fno-omit-frame-pointer -g0 -O2 '-D_FORTIFY_SOURCE=1' -DNDEBUG -ffunction-sections -fdata-sections -MD -MF bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzlib.d '-frandom-seed=bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzlib.o' -iquote external/zlib -iquote bazel-out/k8-opt/bin/external/zlib -isystem external/zlib -isystem bazel-out/k8-opt/bin/external/zlib '-fvisibility=hidden' -Wno-sign-compare -Wno-stringop-truncation -Wno-array-parameter '-DMLIR_PYTHON_PACKAGE_PREFIX=jaxlib.mlir.' -mavx -Wno-shift-negative-value -DZ_HAVE_UNISTD_H -no-canonical-prefixes -Wno-builtin-macro-redefined '-D__DATE__="redacted"' '-D__TIMESTAMP__="redacted"' '-D__TIME__="redacted"' -c external/zlib/gzlib.c -o bazel-out/k8-opt/bin/external/zlib/_objs/zlib/gzlib.o)
# Configuration: 2d7bf1a4c5451c37d2ec070281f978338f5ebc0cdd1b6b3c88a48fe0a0d34362
# Execution platform: @local_execution_config_platform//:platform
Traceback (most recent call last):
  File "/src/tools/symlinks/cc", line 18, in <module>
    PYWASMCROSS_ARGS = json.loads(os.environ["PYWASMCROSS_ARGS"])
  File "/usr/local/lib/python3.10/os.py", line 679, in __getitem__
    raise KeyError(key) from None
KeyError: 'PYWASMCROSS_ARGS'
Target //build:build_wheel failed to build
INFO: Elapsed time: 8.608s, Critical Path: 0.15s
INFO: 23 processes: 23 internal.
FAILED: Build did NOT complete successfully
ERROR: Build failed. Not running target
FAILED: Build did NOT complete successfully
@@@@@@@@@@@@
{"pkgname": "jax", "cflags": " -O2 -g0 -fPIC", "cxxflags": " ", "ldflags": " -O2 -g0  -s MODULARIZE=1 -s LINKABLE=1 -s EXPORT_ALL=1 -s WASM=1 -std=c++14 -s LZ4=1 -L /src/cpython/installs/python-3.10.2/lib/ -s WASM_BIGINT  -s SIDE_MODULE=1", "target_install_dir": "/src/cpython/installs/python-3.10.2", "replace_libs": "", "builddir": "/src/packages/jax/build/jax-0.3.13", "PYTHONPATH": ["/src/packages/jax", "/usr/local/lib/python310.zip", "/usr/local/lib/python3.10", "/usr/local/lib/python3.10/lib-dynload", "/src/.docker_home/.local/lib/python3.10/site-packages", "/src/pyodide-build", "/usr/local/lib/python3.10/site-packages", "/src/cpython/installs/python-3.10.2/sysconfigdata"], "orig__name__": "pyodide_build.pywasmcross"}
@@@@@@@@@@@@

     _   _  __  __
    | | / \ \ \/ /
 _  | |/ _ \ \  /
| |_| / ___ \/  \
 \___/_/   \/_/\_\


Downloading bazel from: https://github.com/bazelbuild/bazel/releases/download/5.1.1/bazel-5.1.1-linux-x86_64

Bazel binary path: ./bazel-5.1.1-linux-x86_64
Bazel version: 5.1.1
Python binary path: /usr/local/bin/python
Python version: 3.10
NumPy version: 1.22.4
MKL-DNN enabled: yes
Target CPU: x86_64
Target CPU features: release
CUDA enabled: no
TPU enabled: no
ROCm enabled: no

Building XLA and installing it in the jaxlib source tree...
./bazel-5.1.1-linux-x86_64 run --verbose_failures=true --config=avx_posix --config=mkl_open_source_only :build_wheel -- --output_path=/src/packages/jax/build/jax-0.3.13/dist --cpu=x86_64
b''
Traceback (most recent call last):
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 532, in <module>
    main()
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 527, in main
    shell(command)
  File "/src/packages/jax/build/jax-0.3.13/build/build.py", line 57, in shell
    output = subprocess.check_output(cmd)
  File "/usr/local/lib/python3.10/subprocess.py", line 420, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/usr/local/lib/python3.10/subprocess.py", line 524, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-5.1.1-linux-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', ':build_wheel', '--', '--output_path=/src/packages/jax/build/jax-0.3.13/dist', '--cpu=x86_64']' returned non-zero exit status 1.
2-END-######################################################################

Notice that the PYWASMCROSS_ARGS error actually occurs before the build/build.py file is executed (unless I severely misunderstand the order in which Python files are executed). Somehow the Bazel server starts and then a while later the @@@@@@@@@@@@ lines with the actual PYWASMCROSS_ARGS dictionary are logged, even though that code is at the very top of build/build.py.

Any idea how that could be happening? I was expecting the logs to start like this:

1-START-######################################################################
@@@@@@@@@@@@
{"pkgname": "jax", "cflags": " -O2 -g0 -fPIC", "cxxflags": " ", "ldflags": " -O2 -g0  -s MODULARIZE=1 -s LINKABLE=1 -s EXPORT_ALL=1 -s WASM=1 -std=c++14 -s LZ4=1 -L /src/cpython/installs/python-3.10.2/lib/ -s WASM_BIGINT  -s SIDE_MODULE=1", "target_install_dir": "/src/cpython/installs/python-3.10.2", "replace_libs": "", "builddir": "/src/packages/jax/build/jax-0.3.13", "PYTHONPATH": ["/src/packages/jax", "/usr/local/lib/python310.zip", "/usr/local/lib/python3.10", "/usr/local/lib/python3.10/lib-dynload", "/src/.docker_home/.local/lib/python3.10/site-packages", "/src/pyodide-build", "/usr/local/lib/python3.10/site-packages", "/src/cpython/installs/python-3.10.2/sysconfigdata"], "orig__name__": "pyodide_build.pywasmcross"}
@@@@@@@@@@@@
...
...
...

[josephrocca]

Okay, so I didn't work out the above-mentioned weirdness, but I bypassed it by realising that Bazel ignores "caller" environment variables by default. To inherit environment variables from caller environment, you need to add --action_env=VAR_NAME for each env variable as mentioned here. So I've added the following code here:

for e in ["--action_env="+k for k in dict.keys(dict(os.environ))]:
  config_args += [e]

So this can be done with a patch. I'm still using the pypabuild.py hack mentioned here until we get a new meta.yaml property to run some pre-build code with the environment swapped.

I've also switched to jax's new logging setup (not yet merged into main) which allows me to remove the absl dependency:

git clone -b logging-setup https://github.com/nicholasjng/jax jax-no-absl

So with that branch I can comment out absl-py in the jax/setup.py install_requires array (by default it's not used).

I've also had to add this to meta.yaml:

build:
  cxxflags: |
    -msimd128
  cflags: |
    -msimd128

I'm going to keep trying to chip away at this - I'm currently stuck on this - fatal error: 'linux/futex.h' file not found.