rdmatop 0.1.31

A real-time TUI monitor for RDMA network interfaces
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101

import os
import shutil
import subprocess
import sys
from pathlib import Path

import torch
from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CppExtension

ROOT = os.path.abspath(os.path.dirname(__file__))
TARGET_DIR = os.path.join(ROOT, "target")
STATICLIB = os.path.join(TARGET_DIR, "release", "librdmatop.a")
KINETO_DIR = os.path.join(ROOT, "kineto")
CAPTURE_HEADER = os.path.join(KINETO_DIR, "rdmatop_capture.h")
SHIM_SOURCE = "kineto/rdmatop_kineto.cpp"
TORCH_DIR = os.path.dirname(torch.__file__)
TORCH_LIB = os.path.join(TORCH_DIR, "lib")
KINETO_INCLUDE = os.path.join(TORCH_DIR, "include", "kineto")
RUST_RUNTIME_LIBS = ["-ldl", "-lgcc_s", "-lutil", "-lrt", "-lpthread", "-lm"]


def kineto_capabilities():
    headers = Path(KINETO_INCLUDE)
    activity_types = (headers / "ActivityType.h").read_text()
    trace_activity = (headers / "GenericTraceActivity.h").read_text()
    # Older wheels do not export fmt symbols used by inline Kineto metadata APIs.
    return [
        ("FMT_HEADER_ONLY", "1"),
        ("RDMATOP_NATIVE_COUNTERS", str(int("MTIA_COUNTERS" in activity_types))),
        ("RDMATOP_TYPED_COUNTERS", str(int("addCounterValue(" in trace_activity))),
    ]


class BuildRustThenExt(BuildExtension):
    def run(self):
        if sys.platform != "linux":
            raise SystemExit("rdmatop: only Linux is supported")
        if shutil.which("cargo") is None:
            raise SystemExit("rdmatop: cargo not found on PATH; install Rust first")
        env = {**os.environ, "CARGO_TARGET_DIR": TARGET_DIR}
        subprocess.check_call(
            ["cargo", "build", "--release", "--lib", "--bin", "rdmatop"],
            cwd=ROOT,
            env=env,
        )
        super().run()
        binary_dir = (
            Path(self.get_ext_fullpath("rdmatop._rdmatop_kineto")).parent / "bin"
        )
        binary_dir.mkdir(parents=True, exist_ok=True)
        shutil.copy2(Path(TARGET_DIR) / "release" / "rdmatop", binary_dir / "rdmatop")
        shutil.copy2(
            Path(TARGET_DIR) / "release" / "librdmatop.so",
            binary_dir.parent / "librdmatop.so",
        )

    def get_outputs(self):
        binary = (
            Path(self.get_ext_fullpath("rdmatop._rdmatop_kineto")).parent
            / "bin"
            / "rdmatop"
        )
        return [
            *super().get_outputs(),
            str(binary),
            str(binary.parent.parent / "librdmatop.so"),
        ]


def cargo_version():
    with open(os.path.join(ROOT, "Cargo.toml")) as manifest:
        lines = manifest.read().splitlines()
    for line in lines:
        if line.startswith("version = "):
            return line.split('"')[1]
    raise SystemExit("rdmatop: version not found in Cargo.toml")


setup(
    version=cargo_version(),
    # Kineto's C++ ABI is tied to the PyTorch version used for compilation.
    install_requires=[f"torch=={torch.__version__.split('+')[0]}"],
    ext_modules=[
        CppExtension(
            name="rdmatop._rdmatop_kineto",
            sources=[SHIM_SOURCE],
            include_dirs=[KINETO_INCLUDE, KINETO_DIR],
            depends=[STATICLIB, CAPTURE_HEADER],
            extra_objects=[STATICLIB],
            define_macros=kineto_capabilities(),
            extra_link_args=[
                f"-Wl,-rpath,{TORCH_LIB}",
                "-Wl,--exclude-libs,ALL",
            ]
            + RUST_RUNTIME_LIBS,
            libraries=["torch_cpu"],
        )
    ],
    cmdclass={"build_ext": BuildRustThenExt},
)