Manage runtime environments (#452)

* Add `Env` class that manages all runtime environments.
* Changed `NPKIT_DUMP_DIR` to `MSCCLPP_NPKIT_DUMP_DIR`.
This commit is contained in:
Changho Hwang
2025-01-15 09:44:52 -08:00
committed by GitHub
parent 8ac50dc85d
commit 869cdba00c
19 changed files with 229 additions and 51 deletions

View File

@@ -8,6 +8,7 @@ from mscclpp import (
ExecutionPlan,
PacketType,
npkit,
env,
)
import mscclpp.comm as mscclpp_comm
from mscclpp.utils import KernelBuilder, GpuBuffer, pack
@@ -171,8 +172,8 @@ def main(
mscclpp_group = mscclpp_comm.CommGroup(MPI.COMM_WORLD)
cp.cuda.Device(mscclpp_group.my_rank % mscclpp_group.nranks_per_node).use()
executor = Executor(mscclpp_group.communicator)
npkit_dump_dir = os.getenv("NPKIT_DUMP_DIR")
if npkit_dump_dir is not None:
npkit_dump_dir = env().npkit_dump_dir
if npkit_dump_dir != "":
npkit.init(mscclpp_group.my_rank)
execution_plan = ExecutionPlan(execution_plan_path)
collective = execution_plan.collective()