CODEWITHSUNDEEP
bazelbazel-rulesbazel-python
Jack Shi
Jack Shi

Reputation: 23

Custom py_binary bazel target that uses deepspeed instead of python3 toolchain

In order to use deepspeed, I need to use the cli to run deepspeed main.py instead of python3 main.py. My goal is to create a deepspeed_py_binary target that is 99% the same as py_binary except it uses the deepspeed toolchain instead of the python3 toolchain.

My approach is currently as follows:

First, I want to create a deepspeed toolchain:

load("@rules_python//python:defs.bzl", "py_runtime_pair")

py_runtime(
    name = "deepspeed_py_runtime",
    interpreter_path = "/usr/bin/deepspeed",
    python_version = "PY3"
)

py_runtime_pair(
    name = "deepspeed_py_runtime_pair",
    py2_runtime = None,
    py3_runtime = ":deepspeed_py_runtime",
)

toolchain(
    name = "deepspeed_py_toolchain",
    toolchain = ":deepspeed_py_runtime_pair",
    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
)

Then my plan is to extend the py_binary target to use this toolchain. I then got super lost trying to understand the py_binary starlark source code...

Anyone have any ideas on how to create a deepspeed_py_binary target that uses the deepspeed command to run main.py instead of python3?

Thank you!!

Upvotes: 0

Views: 39

Answers (0)

Related Questions