Skip to content

Commit

Permalink
Move output path to be inside the wheel build command execution loop
Browse files Browse the repository at this point in the history
This was causing an issue when building multiple wheels in editable mode.

i.e instead of wheels being stored as:
```
# jax-cuda12-pjrt   0.4.36.dev20241125           ./dist/jax-cuda-pjrt
# jax-cuda12-plugin 0.4.36.dev20241125           ./dist/jax-cuda-plugin
# jaxlib            0.4.36.dev20241125           ./dist/jaxlib
```

they were being stored as:
```
# jaxlib            0.4.36.dev20241125           ./dist/jaxlib
# jax-cuda12-pjrt   0.4.36.dev20241125           ./dist/jaxlib/jax-cuda-pjrt
# jax-cuda12-plugin 0.4.36.dev20241125           ./dist/jaxlib/jax-cuda-plugin
```

PiperOrigin-RevId: 708468522
  • Loading branch information
nitins17 authored and Google-ML-Automation committed Dec 21, 2024
1 parent 0e0fc0a commit 0159bea
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions build/build.py
Original file line number Diff line number Diff line change
Expand Up @@ -580,11 +580,11 @@ async def main():
if args.configure_only:
logging.info("--configure_only is set so not running any Bazel commands.")
else:
output_path = args.output_path
logger.debug("Artifacts output directory: %s", output_path)

# Wheel build command execution
for wheel in args.wheels.split(","):
output_path = args.output_path
logger.debug("Artifacts output directory: %s", output_path)

# Allow CUDA/ROCm wheels without the "jax-" prefix.
if ("plugin" in wheel or "pjrt" in wheel) and "jax" not in wheel:
wheel = "jax-" + wheel
Expand Down

0 comments on commit 0159bea

Please sign in to comment.