From ded49c1f05e010e237b513cf8198adeb76c7002f Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 03:50:01 +0800 Subject: [PATCH 01/17] =?UTF-8?q?=E2=9A=A1=20perf:=20activate=20only=20sel?= =?UTF-8?q?ected=20CUPTI=20sources?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cupti/src/cupti_agent.c | 288 +++++++++++++++++++++++-------- tests/integration/test_inject.py | 5 + 2 files changed, 224 insertions(+), 69 deletions(-) diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index de78d76..40789b2 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -99,7 +99,7 @@ static _Atomic int clock_alignment_warning_emitted; static CUpti_SubscriberHandle subscriber; static int subscriber_active; static int agent_initialized; -static size_t enabled_activity_count; +static uint32_t enabled_activity_mask; static char output_path[PATH_MAX]; static char snapshot_socket_path[sizeof(((struct sockaddr_un *)0)->sun_path)]; static pthread_mutex_t flush_mutex = PTHREAD_MUTEX_INITIALIZER; @@ -109,14 +109,12 @@ static _Atomic int snapshot_listener = -1; static int snapshot_thread_started; static struct xprobe_cupti_filter capture_filters[XPROBE_CUPTI_FILTER_COUNT]; static _Atomic int capture_filter_enabled; -static const CUpti_ActivityKind enabled_activity_kinds[] = { - CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL, - CUPTI_ACTIVITY_KIND_MEMCPY, - CUPTI_ACTIVITY_KIND_MEMSET, -}; +#define XPROBE_CUPTI_ACTIVITY_KERNEL (1U << 0) +#define XPROBE_CUPTI_ACTIVITY_MEMCPY (1U << 1) +#define XPROBE_CUPTI_ACTIVITY_MEMSET (1U << 2) static int shutdown_agent(void); -static int initialize_agent(void); +static int activate_capture(void); static int allocate_capture(uint64_t capacity) { @@ -303,10 +301,13 @@ static void copy_name(char destination[XPROBE_CUPTI_NAME_LENGTH], destination[index] = '\0'; } -static size_t bounded_name_length(const char name[XPROBE_CUPTI_NAME_LENGTH]) +static size_t bounded_name_length(const char *name) { size_t length = 0U; + if (name == NULL) { + return 0U; + } while (length < XPROBE_CUPTI_NAME_LENGTH && name[length] != '\0') { ++length; } @@ -314,7 +315,7 @@ static size_t bounded_name_length(const char name[XPROBE_CUPTI_NAME_LENGTH]) } static int name_matches(const struct xprobe_cupti_filter *filter, - const char record_name[XPROBE_CUPTI_NAME_LENGTH]) + const char *record_name) { size_t filter_length; size_t record_length; @@ -322,6 +323,9 @@ static int name_matches(const struct xprobe_cupti_filter *filter, if (filter->name_match == XPROBE_CUPTI_NAME_ANY) { return 1; } + if (record_name == NULL) { + record_name = ""; + } filter_length = bounded_name_length(filter->name); record_length = bounded_name_length(record_name); if (filter_length == XPROBE_CUPTI_NAME_LENGTH || @@ -378,21 +382,49 @@ static uint32_t semantic_memcpy_kind(uint32_t cupti_kind) } } -static int record_matches_filter(const struct xprobe_cupti_record *record, - const struct xprobe_cupti_filter *filter) +static int filter_matches_values(const struct xprobe_cupti_filter *filter, + uint32_t record_kind, uint32_t api_domain, + uint32_t memcpy_kind, const char *name) { - if (filter->record_kind == 0U || filter->record_kind != record->kind) { + if (filter->record_kind == 0U || filter->record_kind != record_kind) { return 0; } - if (filter->api_domain != 0U && - filter->api_domain != record->callback_domain) { + if (filter->api_domain != 0U && filter->api_domain != api_domain) { return 0; } - if (filter->memcpy_kind != 0U && - filter->memcpy_kind != semantic_memcpy_kind(record->grid_z)) { + if (filter->memcpy_kind != 0U && filter->memcpy_kind != memcpy_kind) { return 0; } - return name_matches(filter, record->name); + return name_matches(filter, name); +} + +static int record_matches_filter(const struct xprobe_cupti_record *record, + const struct xprobe_cupti_filter *filter) +{ + return filter_matches_values(filter, record->kind, record->callback_domain, + semantic_memcpy_kind(record->grid_z), + record->name); +} + +static int values_match_capture(uint32_t first_kind, uint32_t second_kind, + uint32_t api_domain, uint32_t memcpy_kind, + const char *name) +{ + if (atomic_load_explicit(&capture_filter_enabled, memory_order_relaxed) == 0) { + return 1; + } + for (size_t index = 0U; index < XPROBE_CUPTI_FILTER_COUNT; ++index) { + const struct xprobe_cupti_filter *filter = &capture_filters[index]; + + if (filter_matches_values(filter, first_kind, api_domain, memcpy_kind, + name) != 0 || + (second_kind != 0U && + filter_matches_values(filter, second_kind, api_domain, memcpy_kind, + name) != 0)) { + return 1; + } + } + return 0; } static int capture_filter_matches(const struct xprobe_cupti_record *record) @@ -478,6 +510,7 @@ static void CUPTIAPI api_callback(void *userdata, CUpti_CallbackDomain domain, const void *callback_data) { const CUpti_CallbackData *data = callback_data; + uint32_t kind; uint64_t timestamp_ns; (void)userdata; @@ -485,17 +518,17 @@ static void CUPTIAPI api_callback(void *userdata, CUpti_CallbackDomain domain, domain != CUPTI_CB_DOMAIN_DRIVER_API) { return; } + kind = data->callbackSite == CUPTI_API_ENTER ? XPROBE_CUPTI_CUDA_API_ENTRY + : XPROBE_CUPTI_CUDA_API_EXIT; + if (values_match_capture(kind, 0U, (uint32_t)domain, 0U, + data->functionName) == 0) { + return; + } if (monotonic_timestamp_ns(×tamp_ns) != XPROBE_CUPTI_AGENT_READY) { return; } - if (data->callbackSite == CUPTI_API_ENTER) { - enqueue_api_record(data, domain, callback_id, - XPROBE_CUPTI_CUDA_API_ENTRY, timestamp_ns); - } else { - enqueue_api_record(data, domain, callback_id, - XPROBE_CUPTI_CUDA_API_EXIT, timestamp_ns); - } + enqueue_api_record(data, domain, callback_id, kind, timestamp_ns); } static void CUPTIAPI activity_buffer_requested(uint8_t **buffer, size_t *size, @@ -552,6 +585,11 @@ static int activity_started_during_capture(uint64_t timestamp_ns) if (activity_started_during_capture((kernel)->start) == 0) { \ return; \ } \ + if (values_match_capture(XPROBE_CUPTI_GPU_KERNEL_START, \ + XPROBE_CUPTI_GPU_KERNEL_END, 0U, 0U, \ + (kernel)->name) == 0) { \ + return; \ + } \ enqueue_kernel_record( \ (kernel)->deviceId, (kernel)->contextId, (kernel)->streamId, \ (kernel)->correlationId, (kernel)->gridX, (kernel)->gridY, \ @@ -648,6 +686,10 @@ static void CUPTIAPI activity_buffer_completed(CUcontext context, uint32_t strea remember_cupti_error(result); break; } + if (atomic_load_explicit(&capture_state, memory_order_relaxed) != + XPROBE_CUPTI_CAPTURE_ACTIVE) { + continue; + } if (activity->kind == CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) { enqueue_kernel_activity(activity); } else if (activity->kind == CUPTI_ACTIVITY_KIND_MEMCPY) { @@ -656,6 +698,13 @@ static void CUPTIAPI activity_buffer_completed(CUcontext context, uint32_t strea if (activity_started_during_capture(memcpy_record->start) == 0) { continue; } + if (values_match_capture( + XPROBE_CUPTI_GPU_MEMCPY_START, + XPROBE_CUPTI_GPU_MEMCPY_END, 0U, + semantic_memcpy_kind((uint32_t)memcpy_record->copyKind), + NULL) == 0) { + continue; + } enqueue_memcpy_record(memcpy_record, XPROBE_CUPTI_GPU_MEMCPY_START, memcpy_record->start); enqueue_memcpy_record(memcpy_record, XPROBE_CUPTI_GPU_MEMCPY_END, @@ -666,6 +715,11 @@ static void CUPTIAPI activity_buffer_completed(CUcontext context, uint32_t strea if (activity_started_during_capture(memset_record->start) == 0) { continue; } + if (values_match_capture(XPROBE_CUPTI_GPU_MEMSET_START, + XPROBE_CUPTI_GPU_MEMSET_END, 0U, 0U, + NULL) == 0) { + continue; + } enqueue_memset_record(memset_record, XPROBE_CUPTI_GPU_MEMSET_START, memset_record->start); enqueue_memset_record(memset_record, XPROBE_CUPTI_GPU_MEMSET_END, @@ -981,7 +1035,7 @@ static int arm_capture(const struct xprobe_cupti_control_request *request) return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } } - if (subscriber_active != 0 || enabled_activity_count != 0U) { + if (subscriber_active != 0 || enabled_activity_mask != 0U) { status = shutdown_agent(); if (status != XPROBE_CUPTI_AGENT_READY) { return status; @@ -994,18 +1048,19 @@ static int arm_capture(const struct xprobe_cupti_control_request *request) reset_capture(); memcpy(capture_filters, request->filters, sizeof(capture_filters)); atomic_store_explicit(&capture_filter_enabled, 1, memory_order_release); - return initialize_agent(); + return activate_capture(); } static int stop_capture(void) { int status = XPROBE_CUPTI_AGENT_READY; - if (subscriber_active != 0 || enabled_activity_count != 0U) { + if (enabled_activity_mask != 0U) { status = flush_activity_buffers(1); - if (status == XPROBE_CUPTI_AGENT_READY) { - status = shutdown_agent(); - } + } + if (status == XPROBE_CUPTI_AGENT_READY && + (subscriber_active != 0 || enabled_activity_mask != 0U)) { + status = shutdown_agent(); } if (status == XPROBE_CUPTI_AGENT_READY && atomic_load_explicit(&capture_state, memory_order_acquire) == @@ -1181,17 +1236,30 @@ static void stop_snapshot_server(void) static int disable_activities(void) { CUptiResult result; - - while (enabled_activity_count > 0U) { - CUpti_ActivityKind kind = enabled_activity_kinds[enabled_activity_count - 1U]; - result = cuptiActivityDisable(kind); + int status = XPROBE_CUPTI_AGENT_READY; + static const struct { + uint32_t bit; + CUpti_ActivityKind kind; + } activities[] = { + {XPROBE_CUPTI_ACTIVITY_MEMSET, CUPTI_ACTIVITY_KIND_MEMSET}, + {XPROBE_CUPTI_ACTIVITY_MEMCPY, CUPTI_ACTIVITY_KIND_MEMCPY}, + {XPROBE_CUPTI_ACTIVITY_KERNEL, CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL}, + }; + + for (size_t index = 0U; index < sizeof(activities) / sizeof(activities[0]); + ++index) { + if ((enabled_activity_mask & activities[index].bit) == 0U) { + continue; + } + result = cuptiActivityDisable(activities[index].kind); if (result != CUPTI_SUCCESS) { report_cupti_error("cuptiActivityDisable", result); - return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + status = XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } else { + enabled_activity_mask &= ~activities[index].bit; } - --enabled_activity_count; } - return XPROBE_CUPTI_AGENT_READY; + return status; } static int shutdown_agent(void) @@ -1210,12 +1278,73 @@ static int shutdown_agent(void) return status; } -static int initialize_agent(void) +static int capture_needs_record_kind(uint32_t first_kind, uint32_t second_kind) +{ + if (atomic_load_explicit(&capture_filter_enabled, memory_order_relaxed) == 0) { + return 1; + } + for (size_t index = 0U; index < XPROBE_CUPTI_FILTER_COUNT; ++index) { + uint32_t kind = capture_filters[index].record_kind; + + if (kind == first_kind || kind == second_kind) { + return 1; + } + } + return 0; +} + +static int capture_needs_api_domain(uint32_t domain) +{ + if (capture_needs_record_kind(XPROBE_CUPTI_CUDA_API_ENTRY, + XPROBE_CUPTI_CUDA_API_EXIT) == 0) { + return 0; + } + if (atomic_load_explicit(&capture_filter_enabled, memory_order_relaxed) == 0) { + return 1; + } + for (size_t index = 0U; index < XPROBE_CUPTI_FILTER_COUNT; ++index) { + const struct xprobe_cupti_filter *filter = &capture_filters[index]; + + if ((filter->record_kind == XPROBE_CUPTI_CUDA_API_ENTRY || + filter->record_kind == XPROBE_CUPTI_CUDA_API_EXIT) && + (filter->api_domain == 0U || filter->api_domain == domain)) { + return 1; + } + } + return 0; +} + +static int enable_activity(uint32_t bit, CUpti_ActivityKind kind) +{ + CUptiResult result = cuptiActivityEnable(kind); + + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiActivityEnable", result); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + enabled_activity_mask |= bit; + return XPROBE_CUPTI_AGENT_READY; +} + +static int activate_capture(void) { #if CUPTI_API_VERSION >= 130000 CUpti_TimestampCallbackFunc timestamp_callback = activity_timestamp; #endif CUptiResult result; + int needs_kernel = + capture_needs_record_kind(XPROBE_CUPTI_GPU_KERNEL_START, + XPROBE_CUPTI_GPU_KERNEL_END); + int needs_memcpy = + capture_needs_record_kind(XPROBE_CUPTI_GPU_MEMCPY_START, + XPROBE_CUPTI_GPU_MEMCPY_END); + int needs_memset = + capture_needs_record_kind(XPROBE_CUPTI_GPU_MEMSET_START, + XPROBE_CUPTI_GPU_MEMSET_END); + int needs_activities = needs_kernel != 0 || needs_memcpy != 0 || + needs_memset != 0; + int needs_runtime = capture_needs_api_domain(CUPTI_CB_DOMAIN_RUNTIME_API); + int needs_driver = capture_needs_api_domain(CUPTI_CB_DOMAIN_DRIVER_API); result = cuptiGetVersion(&runtime_cupti_version); if (result != CUPTI_SUCCESS) { @@ -1226,52 +1355,73 @@ static int initialize_agent(void) XPROBE_CUPTI_AGENT_READY) { return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } - result = cuptiSubscribe(&subscriber, api_callback, NULL); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiSubscribe", result); - return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + if (needs_runtime != 0 || needs_driver != 0) { + result = cuptiSubscribe(&subscriber, api_callback, NULL); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiSubscribe", result); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + subscriber_active = 1; } - subscriber_active = 1; + if (needs_activities != 0) { #if CUPTI_API_VERSION < 130000 - int calibration_status = calibrate_activity_timestamp(); - if (calibration_status != XPROBE_CUPTI_AGENT_READY) { + int calibration_status = calibrate_activity_timestamp(); + if (calibration_status != XPROBE_CUPTI_AGENT_READY) { + shutdown_agent(); + return calibration_status; + } +#else + result = cuptiActivityRegisterTimestampCallback(timestamp_callback); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiActivityRegisterTimestampCallback", result); + shutdown_agent(); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } +#endif + result = cuptiActivityRegisterCallbacks(activity_buffer_requested, + activity_buffer_completed); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiActivityRegisterCallbacks", result); + shutdown_agent(); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + } + if (needs_kernel != 0 && + enable_activity(XPROBE_CUPTI_ACTIVITY_KERNEL, + CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL) != + XPROBE_CUPTI_AGENT_READY) { shutdown_agent(); - return calibration_status; + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } -#else - result = cuptiActivityRegisterTimestampCallback(timestamp_callback); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiActivityRegisterTimestampCallback", result); + if (needs_memcpy != 0 && + enable_activity(XPROBE_CUPTI_ACTIVITY_MEMCPY, + CUPTI_ACTIVITY_KIND_MEMCPY) != + XPROBE_CUPTI_AGENT_READY) { shutdown_agent(); return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } -#endif - result = cuptiActivityRegisterCallbacks(activity_buffer_requested, - activity_buffer_completed); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiActivityRegisterCallbacks", result); + if (needs_memset != 0 && + enable_activity(XPROBE_CUPTI_ACTIVITY_MEMSET, + CUPTI_ACTIVITY_KIND_MEMSET) != + XPROBE_CUPTI_AGENT_READY) { shutdown_agent(); return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } - for (size_t index = 0U; - index < sizeof(enabled_activity_kinds) / sizeof(enabled_activity_kinds[0]); - ++index) { - result = cuptiActivityEnable(enabled_activity_kinds[index]); + if (needs_runtime != 0) { + result = cuptiEnableDomain(1U, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API); if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiActivityEnable", result); + report_cupti_error("cuptiEnableDomain", result); shutdown_agent(); return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } - ++enabled_activity_count; } - result = cuptiEnableDomain(1U, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API); - if (result == CUPTI_SUCCESS) { + if (needs_driver != 0) { result = cuptiEnableDomain(1U, subscriber, CUPTI_CB_DOMAIN_DRIVER_API); - } - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiEnableDomain", result); - shutdown_agent(); - return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiEnableDomain", result); + shutdown_agent(); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } } atomic_store_explicit(&agent_status, XPROBE_CUPTI_AGENT_READY, memory_order_relaxed); @@ -1336,7 +1486,7 @@ int xprobe_cupti_agent_start(const char *configured_socket, uint64_t capacity) } } else { atomic_store_explicit(&capture_filter_enabled, 0, memory_order_release); - if (initialize_agent() != XPROBE_CUPTI_AGENT_READY) { + if (activate_capture() != XPROBE_CUPTI_AGENT_READY) { return xprobe_cupti_agent_status(); } } diff --git a/tests/integration/test_inject.py b/tests/integration/test_inject.py index bed6721..c4f3b8f 100755 --- a/tests/integration/test_inject.py +++ b/tests/integration/test_inject.py @@ -66,6 +66,11 @@ def main() -> None: assert result["ok"] is True assert result["status"] == "completed" assert result["measurement"]["samples"]["matched"] == 3 + assert result["collection"]["cuda_events"] >= 6 + assert result["collection"]["cuda_events"] % 2 == 0 + cupti = result["collection"]["cupti"] + assert cupti["observed_records"] == cupti["retained_records"] + assert cupti["retained_records"] == result["collection"]["cuda_events"] stderr = (output / f"{name}.stderr").read_text() assert "activating the CUPTI agent modifies target PID" in stderr first = json.loads((output / "first.json").read_text()) From 40bb7f8f2a1c88d76b16ec297424692fe7d915bd Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 03:55:23 +0800 Subject: [PATCH 02/17] =?UTF-8?q?=E2=9A=A1=20perf:=20enable=20only=20match?= =?UTF-8?q?ed=20CUDA=20callbacks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cupti/src/cupti_agent.c | 111 +++++++++++++++++++++++++-- tests/integration/run-inject-live.sh | 19 +++++ tests/integration/test_inject.py | 6 +- 3 files changed, 126 insertions(+), 10 deletions(-) diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index 40789b2..aa41bba 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -44,6 +44,8 @@ int xprobe_cupti_agent_flush(void) #include #include #include +#include +#include #include #include @@ -1326,6 +1328,99 @@ static int enable_activity(uint32_t bit, CUpti_ActivityKind kind) return XPROBE_CUPTI_AGENT_READY; } +static int api_filter_matches_callback(CUpti_CallbackDomain domain, + const char *callback_name) +{ + char semantic_name[XPROBE_CUPTI_NAME_LENGTH]; + size_t length; + + copy_name(semantic_name, callback_name); + length = bounded_name_length(semantic_name); + if (length < XPROBE_CUPTI_NAME_LENGTH) { + for (size_t index = length; index > 2U; --index) { + size_t suffix = index - 2U; + + if (semantic_name[suffix] != '_' || + semantic_name[suffix + 1U] != 'v') { + continue; + } + size_t digit = suffix + 2U; + while (digit < length && semantic_name[digit] >= '0' && + semantic_name[digit] <= '9') { + ++digit; + } + if (digit == length && digit > suffix + 2U) { + semantic_name[suffix] = '\0'; + } + break; + } + } + for (size_t index = 0U; index < XPROBE_CUPTI_FILTER_COUNT; ++index) { + const struct xprobe_cupti_filter *filter = &capture_filters[index]; + + if ((filter->record_kind == XPROBE_CUPTI_CUDA_API_ENTRY || + filter->record_kind == XPROBE_CUPTI_CUDA_API_EXIT) && + (filter->api_domain == 0U || + filter->api_domain == (uint32_t)domain) && + name_matches(filter, semantic_name) != 0) { + return 1; + } + } + return 0; +} + +static int enable_filtered_callbacks(CUpti_CallbackDomain domain, + uint32_t callback_count) +{ + uint32_t enabled = 0U; + + for (uint32_t callback_id = 1U; callback_id < callback_count; ++callback_id) { + const char *callback_name = NULL; + CUptiResult result = + cuptiGetCallbackName(domain, callback_id, &callback_name); + + if (result == CUPTI_ERROR_INVALID_PARAMETER) { + continue; + } + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiGetCallbackName", result); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + if (api_filter_matches_callback(domain, callback_name) == 0) { + continue; + } + result = cuptiEnableCallback(1U, subscriber, domain, callback_id); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiEnableCallback", result); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + ++enabled; + } + if (enabled == 0U) { + fprintf(stderr, + "xprobe CUPTI: no callback ID matched the requested API filter\n"); + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } + return XPROBE_CUPTI_AGENT_READY; +} + +static int enable_api_callbacks(CUpti_CallbackDomain domain, + uint32_t callback_count) +{ + CUptiResult result; + + if (atomic_load_explicit(&capture_filter_enabled, memory_order_relaxed) == 0) { + result = cuptiEnableDomain(1U, subscriber, domain); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiEnableDomain", result); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + return XPROBE_CUPTI_AGENT_READY; + } + return enable_filtered_callbacks(domain, callback_count); +} + static int activate_capture(void) { #if CUPTI_API_VERSION >= 130000 @@ -1408,19 +1503,19 @@ static int activate_capture(void) return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } if (needs_runtime != 0) { - result = cuptiEnableDomain(1U, subscriber, CUPTI_CB_DOMAIN_RUNTIME_API); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiEnableDomain", result); + if (enable_api_callbacks(CUPTI_CB_DOMAIN_RUNTIME_API, + CUPTI_RUNTIME_TRACE_CBID_SIZE) != + XPROBE_CUPTI_AGENT_READY) { shutdown_agent(); - return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + return xprobe_cupti_agent_status(); } } if (needs_driver != 0) { - result = cuptiEnableDomain(1U, subscriber, CUPTI_CB_DOMAIN_DRIVER_API); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiEnableDomain", result); + if (enable_api_callbacks(CUPTI_CB_DOMAIN_DRIVER_API, + CUPTI_DRIVER_TRACE_CBID_SIZE) != + XPROBE_CUPTI_AGENT_READY) { shutdown_agent(); - return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + return xprobe_cupti_agent_status(); } } atomic_store_explicit(&agent_status, XPROBE_CUPTI_AGENT_READY, diff --git a/tests/integration/run-inject-live.sh b/tests/integration/run-inject-live.sh index 719e297..c5fecbd 100755 --- a/tests/integration/run-inject-live.sh +++ b/tests/integration/run-inject-live.sh @@ -70,8 +70,27 @@ run_measure() { wait "${measurement_pid}" } +run_api_measure() { + agent_args=() + if [[ ${auto_agent} != 1 ]]; then + agent_args=(--agent "${agent}") + fi + "${xprobe_bin}" measure \ + --pid "${target_pid}" \ + "${agent_args[@]}" \ + --from 'cuda:runtime_api:cudaLaunchKernel:entry' \ + --to 'cuda:runtime_api:cudaLaunchKernel:exit' \ + --match exact \ + --samples 3 \ + --max-events 32 \ + --timeout-ms 10000 \ + --json --non-interactive --no-color \ + >"$1.json" 2>"$1.stderr" +} + run_measure "$1/first" run_measure "$1/second" +run_api_measure "$1/api" if [[ ${auto_agent} == 1 ]]; then mapped_agents=$(awk '$0 ~ /lib\/xprobe\/cuda(12|13)\/libxprobe-cupti.so/ {print $NF}' "/proc/${target_pid}/maps" | sort -u | wc -l) diff --git a/tests/integration/test_inject.py b/tests/integration/test_inject.py index c4f3b8f..2a2140b 100755 --- a/tests/integration/test_inject.py +++ b/tests/integration/test_inject.py @@ -61,7 +61,7 @@ def main() -> None: raise SystemExit(completed.returncode) output = pathlib.Path(output_dir) - for name in ("first", "second"): + for name in ("first", "second", "api"): result = json.loads((output / f"{name}.json").read_text()) assert result["ok"] is True assert result["status"] == "completed" @@ -75,8 +75,10 @@ def main() -> None: assert "activating the CUPTI agent modifies target PID" in stderr first = json.loads((output / "first.json").read_text()) second = json.loads((output / "second.json").read_text()) + api = json.loads((output / "api.json").read_text()) assert any(warning["code"] == "CUPTI_AGENT_INJECTED" for warning in first["warnings"]) assert all(warning["code"] != "CUPTI_AGENT_INJECTED" for warning in second["warnings"]) + assert all(warning["code"] != "CUPTI_AGENT_INJECTED" for warning in api["warnings"]) assert int((output / "mapped-agents.txt").read_text()) == 1 if package is not None: expected_major = 12 if ":12." in sys.argv[1] else 13 @@ -94,7 +96,7 @@ def main() -> None: "schema_version": "2.0", "ok": True, "gpu": "NVIDIA GeForce RTX 3060 Laptop GPU", - "measurements": 2, + "measurements": 3, }, sort_keys=True, ) From b690096f05a68435b0c51669acfc9f19f2f7c5c0 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 04:03:22 +0800 Subject: [PATCH 03/17] =?UTF-8?q?=E2=9A=A1=20perf:=20stream=20incremental?= =?UTF-8?q?=20CUPTI=20snapshots?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cupti/include/xprobe/cupti_agent.h | 6 +- cupti/src/cupti_agent.c | 36 ++++++--- cupti/tests/cupti_smoke.c | 10 ++- docs/architecture.md | 10 ++- docs/cupti-agent.md | 43 ++++++----- tests/integration/test_cupti.py | 10 ++- xprobe/cli/src/main.rs | 115 ++++++++++++++++++++++++++--- xprobe/cli/tests/cupti.rs | 6 +- xprobe/cli/tests/measure.rs | 10 +-- xprobe/collector/src/cupti.rs | 84 ++++++++++++++++----- 10 files changed, 247 insertions(+), 83 deletions(-) diff --git a/cupti/include/xprobe/cupti_agent.h b/cupti/include/xprobe/cupti_agent.h index 193b320..f1766db 100644 --- a/cupti/include/xprobe/cupti_agent.h +++ b/cupti/include/xprobe/cupti_agent.h @@ -7,10 +7,10 @@ extern "C" { #endif -#define XPROBE_CUPTI_AGENT_ABI_VERSION 2U +#define XPROBE_CUPTI_AGENT_ABI_VERSION 3U #define XPROBE_CUPTI_OUTPUT_MAGIC "XPCUPTI" #define XPROBE_CUPTI_CONTROL_MAGIC "XPCTRL\0" -#define XPROBE_CUPTI_CONTROL_VERSION 2U +#define XPROBE_CUPTI_CONTROL_VERSION 3U #define XPROBE_CUPTI_NAME_LENGTH 128U #define XPROBE_CUPTI_FILTER_COUNT 2U #define XPROBE_CUPTI_VALUE_UNKNOWN UINT32_MAX @@ -71,6 +71,7 @@ struct xprobe_cupti_control_request { uint32_t version; uint32_t command; uint64_t record_capacity; + uint64_t record_offset; struct xprobe_cupti_filter filters[XPROBE_CUPTI_FILTER_COUNT]; }; @@ -99,6 +100,7 @@ struct xprobe_cupti_output_header { uint64_t agent_dropped_records; uint64_t cupti_dropped_records; uint64_t unknown_records; + uint64_t record_offset; }; struct xprobe_cupti_record { diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index aa41bba..0a632b8 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -70,13 +70,13 @@ int xprobe_cupti_agent_flush(void) #define XPROBE_CUPTI_CLOCK_MARGIN_NS 1000000000U #define XPROBE_CUPTI_CORRELATION_CLOCK_MARGIN_NS 1000000U -_Static_assert(sizeof(struct xprobe_cupti_output_header) == 80U, +_Static_assert(sizeof(struct xprobe_cupti_output_header) == 88U, "unexpected CUPTI output header layout"); _Static_assert(sizeof(struct xprobe_cupti_record) == 200U, "unexpected CUPTI record layout"); _Static_assert(sizeof(struct xprobe_cupti_filter) == 144U, "unexpected CUPTI filter layout"); -_Static_assert(sizeof(struct xprobe_cupti_control_request) == 312U, +_Static_assert(sizeof(struct xprobe_cupti_control_request) == 320U, "unexpected CUPTI control request layout"); static struct xprobe_cupti_record *records; @@ -911,7 +911,7 @@ static int activity_timestamps_are_host_monotonic(uint64_t available) } static void initialize_output_header(struct xprobe_cupti_output_header *header, - uint64_t available) + uint64_t available, uint64_t record_offset) { memset(header, 0, sizeof(*header)); memcpy(header->magic, XPROBE_CUPTI_OUTPUT_MAGIC, sizeof(header->magic)); @@ -926,7 +926,7 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, (uint32_t)atomic_load_explicit(&capture_state, memory_order_acquire); header->stop_reason = (uint32_t)atomic_load_explicit(&stop_reason, memory_order_relaxed); - header->record_count = available; + header->record_count = available - record_offset; header->record_capacity = record_capacity; header->observed_records = atomic_load_explicit(&record_count, memory_order_relaxed); @@ -936,9 +936,10 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, atomic_load_explicit(&cupti_dropped_records, memory_order_relaxed); header->unknown_records = atomic_load_explicit(&unknown_records, memory_order_relaxed); + header->record_offset = record_offset; } -static int write_capture(int descriptor, int is_socket) +static int write_capture(int descriptor, int is_socket, uint64_t record_offset) { struct xprobe_cupti_output_header header; uint64_t available = @@ -947,11 +948,19 @@ static int write_capture(int descriptor, int is_socket) is_socket != 0 ? send_all : write_all; int result; - initialize_output_header(&header, available); + if (record_offset > available) { + fprintf(stderr, + "xprobe CUPTI: requested record offset %llu exceeds committed " + "record count %llu\n", + (unsigned long long)record_offset, + (unsigned long long)available); + return -1; + } + initialize_output_header(&header, available, record_offset); result = write_function(descriptor, &header, sizeof(header)); if (result == 0) { - result = write_function(descriptor, records, - available * sizeof(records[0])); + result = write_function(descriptor, records + record_offset, + (available - record_offset) * sizeof(records[0])); } return result; } @@ -968,7 +977,7 @@ static int write_output(void) remember_output_error(); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } - result = write_capture(descriptor, 0); + result = write_capture(descriptor, 0, 0U); if (close(descriptor) != 0 && result == 0) { result = -1; } @@ -1125,7 +1134,12 @@ static void *serve_snapshots(void *unused) } else { int command_status = XPROBE_CUPTI_AGENT_READY; - if (request.command == XPROBE_CUPTI_CONTROL_ARM) { + if (request.command == XPROBE_CUPTI_CONTROL_ARM && + request.record_offset != 0U) { + fprintf(stderr, + "xprobe CUPTI: ARM record offset must be zero\n"); + command_status = XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } else if (request.command == XPROBE_CUPTI_CONTROL_ARM) { command_status = arm_capture(&request); } else if (request.command == XPROBE_CUPTI_CONTROL_SNAPSHOT) { if (atomic_load_explicit(&capture_state, memory_order_acquire) == @@ -1138,7 +1152,7 @@ static void *serve_snapshots(void *unused) request.command == XPROBE_CUPTI_CONTROL_CLOSE; } (void)command_status; - if (write_capture(client, 1) != 0) { + if (write_capture(client, 1, request.record_offset) != 0) { fprintf(stderr, "xprobe CUPTI: failed to send snapshot: %s\n", strerror(errno)); } diff --git a/cupti/tests/cupti_smoke.c b/cupti/tests/cupti_smoke.c index f70f6f0..125d658 100644 --- a/cupti/tests/cupti_smoke.c +++ b/cupti/tests/cupti_smoke.c @@ -7,7 +7,7 @@ int main(void) if (xprobe_cupti_agent_abi_version() != XPROBE_CUPTI_AGENT_ABI_VERSION) { return 1; } - if (sizeof(struct xprobe_cupti_output_header) != 80U) { + if (sizeof(struct xprobe_cupti_output_header) != 88U) { return 2; } if (sizeof(struct xprobe_cupti_record) != 200U) { @@ -17,7 +17,8 @@ int main(void) return 4; } if (offsetof(struct xprobe_cupti_output_header, record_count) != 32U || - offsetof(struct xprobe_cupti_output_header, unknown_records) != 72U) { + offsetof(struct xprobe_cupti_output_header, unknown_records) != 72U || + offsetof(struct xprobe_cupti_output_header, record_offset) != 80U) { return 5; } if (offsetof(struct xprobe_cupti_record, grid_x) != 44U || @@ -28,11 +29,12 @@ int main(void) return 6; } if (sizeof(struct xprobe_cupti_filter) != 144U || - sizeof(struct xprobe_cupti_control_request) != 312U) { + sizeof(struct xprobe_cupti_control_request) != 320U) { return 7; } if (offsetof(struct xprobe_cupti_control_request, record_capacity) != 16U || - offsetof(struct xprobe_cupti_control_request, filters) != 24U || + offsetof(struct xprobe_cupti_control_request, record_offset) != 24U || + offsetof(struct xprobe_cupti_control_request, filters) != 32U || offsetof(struct xprobe_cupti_filter, name) != 16U) { return 8; } diff --git a/docs/architecture.md b/docs/architecture.md index 258a040..39c2dd8 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -89,10 +89,12 @@ began before the ARM epoch is excluded. The Agent verifies the retained window before setting the host-monotonic ABI feature, preserving GPU same-domain durations while preventing invalid cross-domain subtraction. -Snapshot flushes and reads the active capture. Stop flushes and disables CUPTI -but retains an externally managed socket for a later ARM. Automatically -injected collection closes its private socket after the final capture. Neither -path unloads the shared object. +Snapshot flushes and reads only records after a checked caller watermark. The +CLI rejects noncontiguous offsets and counter rollback while accumulating the +bounded capture. Stop flushes the final delta and disables CUPTI but retains an +externally managed socket for a later ARM. Automatically injected collection +closes its private socket after the final capture. Neither path unloads the +shared object. ## Online injection diff --git a/docs/cupti-agent.md b/docs/cupti-agent.md index f520bab..b0c3d3d 100644 --- a/docs/cupti-agent.md +++ b/docs/cupti-agent.md @@ -25,12 +25,14 @@ directions, and simple kernel-name exact/prefix/suffix/contains patterns are filtered before they consume capture capacity. Complex kernel regular expressions use a wider Agent filter and retain exact Rust-side matching. -`SNAPSHOT` flushes pending activity and returns an immutable capture. `STOP` -returns the final capture and disables CUPTI while retaining the socket, so a -preloaded Agent can service a later bounded measurement. An automatically -injected measurement uses `CLOSE`, which performs the same logical stop and -then removes its private socket. In both cases the shared object remains mapped; -it must not be passed to `dlclose` while CUDA may still reference callback code. +`SNAPSHOT` flushes pending activity and returns records after the caller's +checked record offset. The CLI accumulates only contiguous deltas rather than +retransferring and decoding the growing capture. `STOP` returns the final delta +and disables CUPTI while retaining the socket, so a preloaded Agent can service +a later bounded measurement. An automatically injected measurement uses +`CLOSE`, which performs the same logical stop and then removes its private +socket. In both cases the shared object remains mapped; it must not be passed to +`dlclose` while CUDA may still reference callback code. Activity records that began before the ARM epoch are excluded as a whole. This prevents CUPTI's boundary records, which can have an unavailable start @@ -43,18 +45,19 @@ paths. ## Control and capture ABI -Control version 2 uses one fixed 312-byte native-endian request defined by +Control version 3 uses one fixed 320-byte native-endian request defined by `cupti/include/xprobe/cupti_agent.h`. It contains magic, version, command, -record capacity, and two fixed endpoint filters. Commands are `ARM`, -`SNAPSHOT`, `STOP`, and `CLOSE`. Every accepted command returns one capture -followed by EOF. +record capacity, record offset, and two fixed endpoint filters. Commands are +`ARM`, `SNAPSHOT`, `STOP`, and `CLOSE`. ARM requires offset zero; later commands +return records at or after the requested offset followed by EOF. -Capture ABI v2 starts with an 80-byte header and zero or more 200-byte records. +Capture ABI v3 starts with an 88-byte header and zero or more 200-byte records. The header reports capture state and stop reason, configured capacity, observed and retained counts, Agent and CUPTI drops, unknown activity, record sizes, and -feature flags. Capacity comes from `--max-events`; there is no special 2^16 -limit. Reaching the configured limit freezes capture and causes measurement to -fail explicitly instead of returning partial success. +feature flags. It also reports the payload record offset so the caller can +reject gaps, replays, and counter rollback. Capacity comes from `--max-events`; +there is no special 2^16 limit. Reaching the configured limit freezes capture +and causes measurement to fail explicitly instead of returning partial success. Records preserve process/thread, device/context/stream, correlation IDs, callback domain/ID, dimensions or transfer metadata, and one bounded name. @@ -67,11 +70,13 @@ interpolation error bound, so normalized results emit ## Hot-path constraints -The API callback reads time and fixed metadata and writes one preallocated -record. Endpoint matching is bounded string/integer comparison. The path does -not allocate, perform file or socket I/O, symbolize names, or take blocking -locks. CUPTI activity buffer decoding is likewise bounded by the configured -capture capacity; flushing and response I/O run on the control thread. +ARM enables only activity families and Runtime/Driver callback IDs required by +the validated endpoints. The API callback checks its fixed filter before +reading time or constructing a record. Activity decoding checks event family, +name, and transfer direction before copying fixed metadata, and stops +converting records after the capture limit. These paths do not allocate, +perform file or socket I/O, symbolize names, or take blocking locks. Flushing +and incremental response I/O run on the control thread. ## Verification diff --git a/tests/integration/test_cupti.py b/tests/integration/test_cupti.py index 5b1dbc1..4e6d011 100755 --- a/tests/integration/test_cupti.py +++ b/tests/integration/test_cupti.py @@ -8,7 +8,7 @@ import tempfile -HEADER = struct.Struct("<8s6I6Q") +HEADER = struct.Struct("<8s6I7Q") RECORD = struct.Struct(" tuple[dict[str, int], list[dict[str, int if len(data) < HEADER.size: raise AssertionError("CUPTI capture is shorter than its header") fields = HEADER.unpack_from(data) - if fields[1] != 2: - raise AssertionError(f"expected capture ABI 2, found {fields[1]}") + if fields[1] != 3: + raise AssertionError(f"expected capture ABI 3, found {fields[1]}") header = { "abi_version": fields[1], "header_size": fields[2], @@ -64,9 +64,10 @@ def read_capture(path: pathlib.Path) -> tuple[dict[str, int], list[dict[str, int "agent_dropped_records": fields[10], "cupti_dropped_records": fields[11], "unknown_records": fields[12], + "record_offset": fields[13], } assert fields[0] == b"XPCUPTI\0" - assert header["abi_version"] == 2 + assert header["abi_version"] == 3 assert header["header_size"] == HEADER.size assert header["record_size"] == RECORD.size assert header["feature_flags"] in {2, 3} @@ -283,6 +284,7 @@ def main() -> None: assert header["agent_dropped_records"] == 0 assert header["cupti_dropped_records"] == 0 assert header["unknown_records"] == 0 + assert header["record_offset"] == 0 assert bool(header["feature_flags"] & 1) is expect_aligned assert all(record["timestamp_ns"] > 0 for record in records) assert all(record["pid"] > 0 and record["tid"] > 0 for record in records) diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index ae64aff..63742c0 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -1208,6 +1208,7 @@ struct LiveCollection { collectors: HostCollectors, host_captures: Option>, latest_cuda: Option, + cuda_record_offset: u64, managed_agent: bool, cupti_armed: bool, agent_injected: bool, @@ -1421,6 +1422,7 @@ fn prepare_live_collection(request: &LiveMeasureRequest) -> Result Err(failure), Err(cleanup) => Err(CommandFailure::new( ErrorCode::CleanupFailed, @@ -1801,16 +1803,26 @@ fn stop_cupti_agent( .expect("armed CUPTI agent has a socket"); collection.cupti_armed = false; let capture = if collection.managed_agent { - cupti::close(socket, request.timeout, &request.options.session_id) + cupti::close( + socket, + request.timeout, + &request.options.session_id, + collection.cuda_record_offset, + ) } else { - cupti::stop(socket, request.timeout, &request.options.session_id) + cupti::stop( + socket, + request.timeout, + &request.options.session_id, + collection.cuda_record_offset, + ) } .map_err(|error| { CommandFailure::new(ErrorCode::CleanupFailed, error.to_string(), true) .with_detail("cleanup_phase", "stop_cupti") .with_hint("verify the target state before starting another measurement") })?; - collection.latest_cuda = Some(capture); + append_cuda_capture(collection, capture)?; collection.managed_agent = false; Ok(()) } @@ -1820,9 +1832,19 @@ impl Drop for LiveCollection { if self.cupti_armed { if let Some(socket) = self.socket.as_deref() { let result = if self.managed_agent { - cupti::close(socket, Duration::from_secs(2), "xp_cleanup") + cupti::close( + socket, + Duration::from_secs(2), + "xp_cleanup", + self.cuda_record_offset, + ) } else { - cupti::stop(socket, Duration::from_secs(2), "xp_cleanup") + cupti::stop( + socket, + Duration::from_secs(2), + "xp_cleanup", + self.cuda_record_offset, + ) }; if let Err(error) = result { eprintln!("xprobe: failed to stop CUPTI agent during cleanup: {error}"); @@ -1932,11 +1954,80 @@ fn refresh_live_sources( collection.collectors.cancel(); return Ok(()); } - collection.latest_cuda = Some( - cupti::snapshot(path, remaining, &request.options.session_id).map_err(|error| { - CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true) - })?, - ); + let capture = cupti::snapshot( + path, + remaining, + &request.options.session_id, + collection.cuda_record_offset, + ) + .map_err(|error| CommandFailure::new(ErrorCode::CuptiNotAvailable, error.to_string(), true))?; + append_cuda_capture(collection, capture)?; + Ok(()) +} + +fn append_cuda_capture( + collection: &mut LiveCollection, + mut capture: cupti::CuptiCapture, +) -> Result<(), CommandFailure> { + if capture.record_offset != collection.cuda_record_offset { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "CUPTI incremental capture started at record {}, expected {}", + capture.record_offset, collection.cuda_record_offset + ), + false, + )); + } + let returned = u64::try_from(capture.events.len()).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("CUPTI incremental event count exceeds u64: {error}"), + false, + ) + })?; + let next_offset = capture.record_offset.checked_add(returned).ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI incremental record offset overflowed", + false, + ) + })?; + if let Some(accumulated) = collection.latest_cuda.as_mut() { + if accumulated.record_capacity != capture.record_capacity { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "CUPTI record capacity changed from {} to {} during capture", + accumulated.record_capacity, capture.record_capacity + ), + false, + )); + } + if capture.observed_records < accumulated.observed_records + || capture.agent_dropped_records < accumulated.agent_dropped_records + || capture.cupti_dropped_records < accumulated.cupti_dropped_records + || capture.unknown_records < accumulated.unknown_records + { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "CUPTI incremental counters moved backwards", + false, + )); + } + accumulated.state = capture.state; + accumulated.stop_reason = capture.stop_reason; + accumulated.observed_records = capture.observed_records; + accumulated.agent_dropped_records = capture.agent_dropped_records; + accumulated.cupti_dropped_records = capture.cupti_dropped_records; + accumulated.dropped_records = capture.dropped_records; + accumulated.unknown_records = capture.unknown_records; + accumulated.events.append(&mut capture.events); + } else { + capture.record_offset = 0; + collection.latest_cuda = Some(capture); + } + collection.cuda_record_offset = next_offset; Ok(()) } @@ -2496,7 +2587,7 @@ fn run_cupti(args: CuptiArgs) -> ExitCode { } } else { let socket = socket.expect("clap requires one CUPTI capture source"); - match cupti::snapshot(&socket, Duration::from_millis(timeout_ms), &session_id) { + match cupti::snapshot(&socket, Duration::from_millis(timeout_ms), &session_id, 0) { Ok(capture) => capture, Err(error) => { return emit_error(ErrorCode::TraceExportFailed, error.to_string(), true, json); diff --git a/xprobe/cli/tests/cupti.rs b/xprobe/cli/tests/cupti.rs index ddf1d11..10f8773 100644 --- a/xprobe/cli/tests/cupti.rs +++ b/xprobe/cli/tests/cupti.rs @@ -2,7 +2,7 @@ use std::{fs, path::PathBuf, process::Command}; use xprobe_protocol::{ErrorCode, ErrorResponse, Event, EventSource, EventType}; -const HEADER_SIZE: usize = 80; +const HEADER_SIZE: usize = 88; const RECORD_SIZE: usize = 200; fn capture_path(name: &str) -> PathBuf { @@ -12,8 +12,8 @@ fn capture_path(name: &str) -> PathBuf { fn write_capture(path: &PathBuf) { let mut bytes = vec![0_u8; HEADER_SIZE + RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&2_u32.to_le_bytes()); - bytes[12..16].copy_from_slice(&80_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[24..28].copy_from_slice(&3_u32.to_le_bytes()); bytes[28..32].copy_from_slice(&1_u32.to_le_bytes()); diff --git a/xprobe/cli/tests/measure.rs b/xprobe/cli/tests/measure.rs index 820173b..92c8c9e 100644 --- a/xprobe/cli/tests/measure.rs +++ b/xprobe/cli/tests/measure.rs @@ -8,7 +8,7 @@ use xprobe_protocol::{ SchemaVersion, SessionStatus, TargetIdentity, }; -const HEADER_SIZE: usize = 80; +const HEADER_SIZE: usize = 88; const RECORD_SIZE: usize = 200; fn capture_path(name: &str) -> PathBuf { @@ -35,8 +35,8 @@ fn record(kind: u32, timestamp: u64, correlation_id: u32, name: &str) -> [u8; RE fn write_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { let mut bytes = vec![0_u8; HEADER_SIZE + records.len() * RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&2_u32.to_le_bytes()); - bytes[12..16].copy_from_slice(&80_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[24..28].copy_from_slice(&3_u32.to_le_bytes()); bytes[28..32].copy_from_slice(&1_u32.to_le_bytes()); @@ -54,8 +54,8 @@ fn write_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { fn write_normalized_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { let mut bytes = vec![0_u8; HEADER_SIZE + records.len() * RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&2_u32.to_le_bytes()); - bytes[12..16].copy_from_slice(&80_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[20..24].copy_from_slice(&1_u32.to_le_bytes()); bytes[24..28].copy_from_slice(&3_u32.to_le_bytes()); diff --git a/xprobe/collector/src/cupti.rs b/xprobe/collector/src/cupti.rs index d18403c..895ea41 100644 --- a/xprobe/collector/src/cupti.rs +++ b/xprobe/collector/src/cupti.rs @@ -17,13 +17,13 @@ use xprobe_protocol::{ const OUTPUT_MAGIC: &[u8; 8] = b"XPCUPTI\0"; const CONTROL_MAGIC: &[u8; 8] = b"XPCTRL\0\0"; -const CONTROL_VERSION: u32 = 2; -const ABI_VERSION: u32 = 2; -const HEADER_SIZE: usize = 80; -const HEADER_SIZE_U32: u32 = 80; +const CONTROL_VERSION: u32 = 3; +const ABI_VERSION: u32 = 3; +const HEADER_SIZE: usize = 88; +const HEADER_SIZE_U32: u32 = 88; const RECORD_SIZE: usize = 200; const RECORD_SIZE_U32: u32 = 200; -const CONTROL_SIZE: usize = 312; +const CONTROL_SIZE: usize = 320; const FILTER_SIZE: usize = 144; const FILTER_COUNT: usize = 2; const FILTER_NAME_SIZE: usize = 128; @@ -89,6 +89,7 @@ pub struct CuptiArmConfig { #[derive(Debug, Clone, PartialEq)] pub struct CuptiCapture { + pub record_offset: u64, pub state: CuptiCaptureState, pub stop_reason: CuptiStopReason, pub record_capacity: u64, @@ -151,6 +152,7 @@ pub enum CuptiDecodeError { InvalidCaptureState(u32), InvalidStopReason(u32), InvalidCounters { + offset: u64, records: u64, capacity: u64, observed: u64, @@ -217,12 +219,13 @@ impl fmt::Display for CuptiDecodeError { write!(formatter, "CUPTI capture stop reason {reason} is invalid") } Self::InvalidCounters { + offset, records, capacity, observed, } => write!( formatter, - "CUPTI capture counters are inconsistent: records={records}, \ + "CUPTI capture counters are inconsistent: offset={offset}, records={records}, \ capacity={capacity}, observed={observed}" ), Self::CounterOverflow => { @@ -293,7 +296,7 @@ pub fn arm( session_id: &str, config: &CuptiArmConfig, ) -> Result { - request(socket_path, timeout, session_id, 1, Some(config)) + request(socket_path, timeout, session_id, 1, 0, Some(config)) } /// Request one immutable capture snapshot from a running CUPTI agent. @@ -309,8 +312,9 @@ pub fn snapshot( socket_path: &Path, timeout: Duration, session_id: &str, + record_offset: u64, ) -> Result { - request(socket_path, timeout, session_id, 2, None) + request(socket_path, timeout, session_id, 2, record_offset, None) } /// Request a final snapshot and logically disable a running CUPTI agent. @@ -324,8 +328,9 @@ pub fn stop( socket_path: &Path, timeout: Duration, session_id: &str, + record_offset: u64, ) -> Result { - request(socket_path, timeout, session_id, 3, None) + request(socket_path, timeout, session_id, 3, record_offset, None) } /// Stop capture and close the Agent control socket. @@ -337,8 +342,9 @@ pub fn close( socket_path: &Path, timeout: Duration, session_id: &str, + record_offset: u64, ) -> Result { - request(socket_path, timeout, session_id, 4, None) + request(socket_path, timeout, session_id, 4, record_offset, None) } fn request( @@ -346,9 +352,10 @@ fn request( timeout: Duration, session_id: &str, command: u32, + record_offset: u64, config: Option<&CuptiArmConfig>, ) -> Result { - let control = encode_control(command, config)?; + let control = encode_control(command, record_offset, config)?; let mut stream = UnixStream::connect(socket_path).map_err(|source| CuptiSnapshotError::Connect { path: socket_path.to_owned(), @@ -373,12 +380,14 @@ fn request( fn encode_control( command: u32, + record_offset: u64, config: Option<&CuptiArmConfig>, ) -> Result<[u8; CONTROL_SIZE], CuptiSnapshotError> { let mut control = [0_u8; CONTROL_SIZE]; control[..8].copy_from_slice(CONTROL_MAGIC); control[8..12].copy_from_slice(&CONTROL_VERSION.to_ne_bytes()); control[12..16].copy_from_slice(&command.to_ne_bytes()); + control[24..32].copy_from_slice(&record_offset.to_ne_bytes()); let Some(config) = config else { return Ok(control); }; @@ -393,9 +402,14 @@ fn encode_control( config.filters.len() ))); } + if record_offset != 0 { + return Err(CuptiSnapshotError::InvalidControl( + "CUPTI ARM record offset must be zero".to_owned(), + )); + } control[16..24].copy_from_slice(&config.record_capacity.to_ne_bytes()); for (index, filter) in config.filters.iter().enumerate() { - let offset = 24 + index * FILTER_SIZE; + let offset = 32 + index * FILTER_SIZE; control[offset..offset + 4].copy_from_slice(&(filter.record_kind as u32).to_ne_bytes()); control[offset + 4..offset + 8].copy_from_slice(&(filter.api_domain as u32).to_ne_bytes()); control[offset + 8..offset + 12] @@ -475,8 +489,13 @@ pub fn decode_capture(bytes: &[u8], session_id: &str) -> Result record_capacity || observed_records < record_count_raw { + let record_offset = read_u64(bytes, 80); + let committed_records = record_offset + .checked_add(record_count_raw) + .ok_or(CuptiDecodeError::CounterOverflow)?; + if committed_records > record_capacity || observed_records < committed_records { return Err(CuptiDecodeError::InvalidCounters { + offset: record_offset, records: record_count_raw, capacity: record_capacity, observed: observed_records, @@ -502,7 +521,7 @@ pub fn decode_capture(bytes: &[u8], session_id: &str) -> Result Result Date: Fri, 24 Jul 2026 04:04:43 +0800 Subject: [PATCH 04/17] =?UTF-8?q?=F0=9F=A7=AA=20test:=20align=20benchmark?= =?UTF-8?q?=20with=20CUPTI=20ABI=20v3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- benchmarks/cuda-callback/run.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/benchmarks/cuda-callback/run.py b/benchmarks/cuda-callback/run.py index 176d3fd..34f6159 100755 --- a/benchmarks/cuda-callback/run.py +++ b/benchmarks/cuda-callback/run.py @@ -8,7 +8,7 @@ import tempfile -HEADER = struct.Struct("<8s6I6Q") +HEADER = struct.Struct("<8s6I7Q") RECORD = struct.Struct(" tuple[dict[str, int], list[dict[str, int "agent_dropped_records": fields[10], "cupti_dropped_records": fields[11], "unknown_records": fields[12], + "record_offset": fields[13], } assert fields[0] == b"XPCUPTI\0", header - assert header["abi_version"] == 2, header + assert header["abi_version"] == 3, header assert header["header_size"] == HEADER.size, header assert header["record_size"] == RECORD.size, header assert len(data) == HEADER.size + header["record_count"] * RECORD.size, header @@ -122,6 +123,7 @@ def main() -> None: assert header["agent_dropped_records"] == 0, header assert header["cupti_dropped_records"] == 0, header assert header["unknown_records"] == 0, header + assert header["record_offset"] == 0, header cupti_ns = precision_duration(records) reference_ns = instrumented["precision_cuda_event_ns"] From 96375cb650012a939150a16b06d98266f6035eb8 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 04:48:55 +0800 Subject: [PATCH 05/17] =?UTF-8?q?=E2=9A=A1=20perf:=20add=20bounded=20aggre?= =?UTF-8?q?gate=20inventory?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- benchmarks/cuda-callback/run.py | 2 +- cupti/include/xprobe/cupti_agent.h | 27 +- cupti/src/cupti_agent.c | 310 +++++++- cupti/tests/cuda_multisource_fixture.cu | 9 + cupti/tests/cupti_smoke.c | 18 +- .../aggregate-inventory-result.schema.json | 279 ++++++++ schemas/measurement-spec.schema.json | 29 +- tests/integration/run-inject-live.sh | 47 +- tests/integration/test_cupti.py | 6 +- tests/integration/test_inject.py | 26 +- xprobe/cli/src/main.rs | 669 +++++++++++++++++- xprobe/cli/tests/cupti.rs | 2 +- xprobe/cli/tests/measure.rs | 60 +- xprobe/cli/tests/trace.rs | 11 +- xprobe/collector/src/cupti.rs | 301 +++++++- xprobe/protocol/src/lib.rs | 8 +- xprobe/protocol/src/measurement.rs | 81 ++- xprobe/protocol/src/schema.rs | 11 +- xprobe/protocol/tests/contracts.rs | 52 +- 19 files changed, 1841 insertions(+), 107 deletions(-) create mode 100644 schemas/aggregate-inventory-result.schema.json diff --git a/benchmarks/cuda-callback/run.py b/benchmarks/cuda-callback/run.py index 34f6159..d3279d3 100755 --- a/benchmarks/cuda-callback/run.py +++ b/benchmarks/cuda-callback/run.py @@ -35,7 +35,7 @@ def read_capture(path: pathlib.Path) -> tuple[dict[str, int], list[dict[str, int "record_offset": fields[13], } assert fields[0] == b"XPCUPTI\0", header - assert header["abi_version"] == 3, header + assert header["abi_version"] == 4, header assert header["header_size"] == HEADER.size, header assert header["record_size"] == RECORD.size, header assert len(data) == HEADER.size + header["record_count"] * RECORD.size, header diff --git a/cupti/include/xprobe/cupti_agent.h b/cupti/include/xprobe/cupti_agent.h index f1766db..a5c5fc3 100644 --- a/cupti/include/xprobe/cupti_agent.h +++ b/cupti/include/xprobe/cupti_agent.h @@ -7,17 +7,23 @@ extern "C" { #endif -#define XPROBE_CUPTI_AGENT_ABI_VERSION 3U +#define XPROBE_CUPTI_AGENT_ABI_VERSION 4U #define XPROBE_CUPTI_OUTPUT_MAGIC "XPCUPTI" #define XPROBE_CUPTI_CONTROL_MAGIC "XPCTRL\0" -#define XPROBE_CUPTI_CONTROL_VERSION 3U +#define XPROBE_CUPTI_CONTROL_VERSION 4U #define XPROBE_CUPTI_NAME_LENGTH 128U #define XPROBE_CUPTI_FILTER_COUNT 2U #define XPROBE_CUPTI_VALUE_UNKNOWN UINT32_MAX enum xprobe_cupti_feature { XPROBE_CUPTI_FEATURE_HOST_MONOTONIC_TIMESTAMPS = 1U << 0, - XPROBE_CUPTI_FEATURE_TRANSFER_RECORDS = 1U << 1 + XPROBE_CUPTI_FEATURE_TRANSFER_RECORDS = 1U << 1, + XPROBE_CUPTI_FEATURE_AGGREGATE_RECORDS = 1U << 2 +}; + +enum xprobe_cupti_capture_mode { + XPROBE_CUPTI_CAPTURE_EXACT = 0, + XPROBE_CUPTI_CAPTURE_AGGREGATE = 1 }; enum xprobe_cupti_agent_status { @@ -72,6 +78,8 @@ struct xprobe_cupti_control_request { uint32_t command; uint64_t record_capacity; uint64_t record_offset; + uint32_t capture_mode; + uint32_t reserved; struct xprobe_cupti_filter filters[XPROBE_CUPTI_FILTER_COUNT]; }; @@ -124,6 +132,19 @@ struct xprobe_cupti_record { char name[XPROBE_CUPTI_NAME_LENGTH]; }; +struct xprobe_cupti_aggregate_record { + uint64_t count; + uint64_t total_duration_ns; + uint64_t min_duration_ns; + uint64_t max_duration_ns; + uint64_t total_bytes; + uint32_t kind; + uint32_t device_id; + uint32_t memcpy_kind; + uint32_t reserved[5]; + char name[XPROBE_CUPTI_NAME_LENGTH]; +}; + /* * For memcpy and memset records, grid_x/grid_y hold the low/high halves of the * byte count. grid_z holds the CUpti_ActivityMemcpyKind for memcpy records, diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index 0a632b8..b1de857 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -74,14 +74,31 @@ _Static_assert(sizeof(struct xprobe_cupti_output_header) == 88U, "unexpected CUPTI output header layout"); _Static_assert(sizeof(struct xprobe_cupti_record) == 200U, "unexpected CUPTI record layout"); +_Static_assert(sizeof(struct xprobe_cupti_aggregate_record) == 200U, + "unexpected CUPTI aggregate record layout"); _Static_assert(sizeof(struct xprobe_cupti_filter) == 144U, "unexpected CUPTI filter layout"); -_Static_assert(sizeof(struct xprobe_cupti_control_request) == 320U, +_Static_assert(sizeof(struct xprobe_cupti_control_request) == 328U, "unexpected CUPTI control request layout"); +struct xprobe_cupti_aggregate_slot { + _Atomic uint64_t state; + _Atomic uint64_t count; + _Atomic uint64_t total_duration_ns; + _Atomic uint64_t min_duration_ns; + _Atomic uint64_t max_duration_ns; + _Atomic uint64_t total_bytes; + uint32_t kind; + uint32_t device_id; + uint32_t memcpy_kind; + char name[XPROBE_CUPTI_NAME_LENGTH]; +}; + static struct xprobe_cupti_record *records; static _Atomic unsigned char *record_ready; +static struct xprobe_cupti_aggregate_slot *aggregate_slots; static uint64_t record_capacity; +static uint32_t capture_mode; static _Atomic uint64_t record_count; static _Atomic uint64_t committed_record_count; static _Atomic uint64_t agent_dropped_records; @@ -122,6 +139,7 @@ static int allocate_capture(uint64_t capacity) { struct xprobe_cupti_record *new_records; _Atomic unsigned char *new_ready; + struct xprobe_cupti_aggregate_slot *new_aggregate_slots; if (capacity == 0U || capacity > SIZE_MAX / sizeof(*records) || capacity > SIZE_MAX / sizeof(*record_ready)) { @@ -129,22 +147,41 @@ static int allocate_capture(uint64_t capacity) (unsigned long long)capacity); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } - new_records = calloc((size_t)capacity, sizeof(*new_records)); - new_ready = malloc((size_t)capacity * sizeof(*new_ready)); - if (new_records == NULL || new_ready == NULL) { + new_records = NULL; + new_ready = NULL; + new_aggregate_slots = NULL; + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + if (capacity > SIZE_MAX / sizeof(*aggregate_slots)) { + fprintf(stderr, "xprobe CUPTI: aggregate capacity exceeds size_t\n"); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } + new_aggregate_slots = calloc((size_t)capacity, sizeof(*new_aggregate_slots)); + } else { + new_records = calloc((size_t)capacity, sizeof(*new_records)); + new_ready = malloc((size_t)capacity * sizeof(*new_ready)); + } + if ((capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE && + new_aggregate_slots == NULL) || + (capture_mode == XPROBE_CUPTI_CAPTURE_EXACT && + (new_records == NULL || new_ready == NULL))) { fprintf(stderr, "xprobe CUPTI: failed to allocate %llu capture records\n", (unsigned long long)capacity); free(new_records); free(new_ready); + free(new_aggregate_slots); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } - for (uint64_t index = 0U; index < capacity; ++index) { - atomic_init(&new_ready[index], 0U); + if (new_ready != NULL) { + for (uint64_t index = 0U; index < capacity; ++index) { + atomic_init(&new_ready[index], 0U); + } } free(records); free(record_ready); + free(aggregate_slots); records = new_records; record_ready = new_ready; + aggregate_slots = new_aggregate_slots; record_capacity = capacity; return XPROBE_CUPTI_AGENT_READY; } @@ -153,8 +190,15 @@ static void reset_capture(void) { uint64_t index; - for (index = 0U; index < record_capacity; ++index) { - atomic_store_explicit(&record_ready[index], 0U, memory_order_relaxed); + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + for (index = 0U; index < record_capacity; ++index) { + atomic_store_explicit(&aggregate_slots[index].state, 0U, + memory_order_relaxed); + } + } else { + for (index = 0U; index < record_capacity; ++index) { + atomic_store_explicit(&record_ready[index], 0U, memory_order_relaxed); + } } atomic_store_explicit(&record_count, 0U, memory_order_relaxed); atomic_store_explicit(&committed_record_count, 0U, memory_order_relaxed); @@ -477,6 +521,154 @@ static void enqueue_record(const struct xprobe_cupti_record *record) } } +static uint64_t aggregate_hash(uint32_t kind, uint32_t device_id, + uint32_t memcpy_kind, const char *name) +{ + uint64_t hash = 1469598103934665603ULL; + const unsigned char *bytes = (const unsigned char *)name; + uint32_t values[] = {kind, device_id, memcpy_kind}; + + for (size_t index = 0U; index < sizeof(values); ++index) { + hash ^= ((const unsigned char *)values)[index]; + hash *= 1099511628211ULL; + } + if (bytes != NULL) { + for (size_t index = 0U; + index + 1U < XPROBE_CUPTI_NAME_LENGTH && bytes[index] != '\0'; + ++index) { + hash ^= bytes[index]; + hash *= 1099511628211ULL; + } + } + if (hash < 2U) { + hash += 2U; + } + return hash; +} + +static int aggregate_key_matches( + const struct xprobe_cupti_aggregate_slot *slot, uint32_t kind, + uint32_t device_id, uint32_t memcpy_kind, const char *name) +{ + const char *key_name = name == NULL ? "" : name; + + return slot->kind == kind && slot->device_id == device_id && + slot->memcpy_kind == memcpy_kind && + strncmp(slot->name, key_name, XPROBE_CUPTI_NAME_LENGTH - 1U) == 0; +} + +static int atomic_add_checked(_Atomic uint64_t *value, uint64_t addend) +{ + uint64_t current = atomic_load_explicit(value, memory_order_relaxed); + + for (;;) { + if (UINT64_MAX - current < addend) { + remember_output_error(); + return 0; + } + if (atomic_compare_exchange_weak_explicit( + value, ¤t, current + addend, memory_order_relaxed, + memory_order_relaxed)) { + return 1; + } + } +} + +static void atomic_min(_Atomic uint64_t *value, uint64_t candidate) +{ + uint64_t current = atomic_load_explicit(value, memory_order_relaxed); + + while (candidate < current && + !atomic_compare_exchange_weak_explicit( + value, ¤t, candidate, memory_order_relaxed, + memory_order_relaxed)) { + } +} + +static void atomic_max(_Atomic uint64_t *value, uint64_t candidate) +{ + uint64_t current = atomic_load_explicit(value, memory_order_relaxed); + + while (candidate > current && + !atomic_compare_exchange_weak_explicit( + value, ¤t, candidate, memory_order_relaxed, + memory_order_relaxed)) { + } +} + +static void update_aggregate_slot(struct xprobe_cupti_aggregate_slot *slot, + uint64_t duration_ns, uint64_t bytes) +{ + if (atomic_add_checked(&slot->total_duration_ns, duration_ns) == 0 || + atomic_add_checked(&slot->total_bytes, bytes) == 0 || + atomic_add_checked(&slot->count, 1U) == 0) { + return; + } + atomic_min(&slot->min_duration_ns, duration_ns); + atomic_max(&slot->max_duration_ns, duration_ns); +} + +static void aggregate_activity(uint32_t kind, uint32_t device_id, + uint32_t memcpy_kind, const char *name, + uint64_t start_ns, uint64_t end_ns, uint64_t bytes) +{ + uint64_t hash; + uint64_t first; + + if (atomic_load_explicit(&capture_state, memory_order_relaxed) != + XPROBE_CUPTI_CAPTURE_ACTIVE) { + return; + } + if (start_ns == 0U || end_ns < start_ns) { + remember_output_error(); + return; + } + atomic_fetch_add_explicit(&record_count, 1U, memory_order_relaxed); + hash = aggregate_hash(kind, device_id, memcpy_kind, name); + first = hash % record_capacity; + for (uint64_t probe = 0U; probe < record_capacity; ++probe) { + uint64_t index = (first + probe) % record_capacity; + struct xprobe_cupti_aggregate_slot *slot = &aggregate_slots[index]; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + + if (state == hash && + aggregate_key_matches(slot, kind, device_id, memcpy_kind, name) != 0) { + update_aggregate_slot(slot, end_ns - start_ns, bytes); + return; + } + if (state == 0U) { + uint64_t expected = 0U; + if (!atomic_compare_exchange_strong_explicit( + &slot->state, &expected, 1U, memory_order_acq_rel, + memory_order_acquire)) { + continue; + } + slot->kind = kind; + slot->device_id = device_id; + slot->memcpy_kind = memcpy_kind; + copy_name(slot->name, name); + atomic_store_explicit(&slot->count, 0U, memory_order_relaxed); + atomic_store_explicit(&slot->total_duration_ns, 0U, + memory_order_relaxed); + atomic_store_explicit(&slot->min_duration_ns, UINT64_MAX, + memory_order_relaxed); + atomic_store_explicit(&slot->max_duration_ns, 0U, + memory_order_relaxed); + atomic_store_explicit(&slot->total_bytes, 0U, memory_order_relaxed); + atomic_store_explicit(&slot->state, hash, memory_order_release); + atomic_fetch_add_explicit(&committed_record_count, 1U, + memory_order_relaxed); + update_aggregate_slot(slot, end_ns - start_ns, bytes); + return; + } + } + atomic_fetch_add_explicit(&agent_dropped_records, 1U, memory_order_relaxed); + atomic_store_explicit(&stop_reason, XPROBE_CUPTI_STOP_RECORD_LIMIT, + memory_order_relaxed); + atomic_store_explicit(&capture_state, XPROBE_CUPTI_CAPTURE_LIMIT_REACHED, + memory_order_release); +} + static void initialize_record(struct xprobe_cupti_record *record, uint32_t kind, uint64_t timestamp_ns) { @@ -592,6 +784,13 @@ static int activity_started_during_capture(uint64_t timestamp_ns) (kernel)->name) == 0) { \ return; \ } \ + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { \ + aggregate_activity( \ + XPROBE_CUPTI_GPU_KERNEL_START, (kernel)->deviceId, 0U, \ + (kernel)->name, normalize_activity_timestamp((kernel)->start), \ + normalize_activity_timestamp((kernel)->end), 0U); \ + return; \ + } \ enqueue_kernel_record( \ (kernel)->deviceId, (kernel)->contextId, (kernel)->streamId, \ (kernel)->correlationId, (kernel)->gridX, (kernel)->gridY, \ @@ -707,6 +906,15 @@ static void CUPTIAPI activity_buffer_completed(CUcontext context, uint32_t strea NULL) == 0) { continue; } + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + aggregate_activity( + XPROBE_CUPTI_GPU_MEMCPY_START, memcpy_record->deviceId, + semantic_memcpy_kind((uint32_t)memcpy_record->copyKind), NULL, + normalize_activity_timestamp(memcpy_record->start), + normalize_activity_timestamp(memcpy_record->end), + memcpy_record->bytes); + continue; + } enqueue_memcpy_record(memcpy_record, XPROBE_CUPTI_GPU_MEMCPY_START, memcpy_record->start); enqueue_memcpy_record(memcpy_record, XPROBE_CUPTI_GPU_MEMCPY_END, @@ -722,6 +930,14 @@ static void CUPTIAPI activity_buffer_completed(CUcontext context, uint32_t strea NULL) == 0) { continue; } + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + aggregate_activity( + XPROBE_CUPTI_GPU_MEMSET_START, memset_record->deviceId, 0U, + NULL, normalize_activity_timestamp(memset_record->start), + normalize_activity_timestamp(memset_record->end), + memset_record->bytes); + continue; + } enqueue_memset_record(memset_record, XPROBE_CUPTI_GPU_MEMSET_START, memset_record->start); enqueue_memset_record(memset_record, XPROBE_CUPTI_GPU_MEMSET_END, @@ -919,7 +1135,9 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, header->header_size = sizeof(*header); header->record_size = sizeof(records[0]); header->feature_flags = XPROBE_CUPTI_FEATURE_TRANSFER_RECORDS; - if (activity_timestamps_are_host_monotonic(available) != 0) { + if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + header->feature_flags |= XPROBE_CUPTI_FEATURE_AGGREGATE_RECORDS; + } else if (activity_timestamps_are_host_monotonic(available) != 0) { header->feature_flags |= XPROBE_CUPTI_FEATURE_HOST_MONOTONIC_TIMESTAMPS; } header->capture_state = @@ -939,6 +1157,38 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, header->record_offset = record_offset; } +static int write_aggregate_records(int descriptor, + int (*write_function)(int, const void *, size_t)) +{ + for (uint64_t index = 0U; index < record_capacity; ++index) { + const struct xprobe_cupti_aggregate_slot *slot = &aggregate_slots[index]; + struct xprobe_cupti_aggregate_record output; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + + if (state < 2U) { + continue; + } + memset(&output, 0, sizeof(output)); + output.count = atomic_load_explicit(&slot->count, memory_order_relaxed); + output.total_duration_ns = + atomic_load_explicit(&slot->total_duration_ns, memory_order_relaxed); + output.min_duration_ns = + atomic_load_explicit(&slot->min_duration_ns, memory_order_relaxed); + output.max_duration_ns = + atomic_load_explicit(&slot->max_duration_ns, memory_order_relaxed); + output.total_bytes = + atomic_load_explicit(&slot->total_bytes, memory_order_relaxed); + output.kind = slot->kind; + output.device_id = slot->device_id; + output.memcpy_kind = slot->memcpy_kind; + memcpy(output.name, slot->name, sizeof(output.name)); + if (write_function(descriptor, &output, sizeof(output)) != 0) { + return -1; + } + } + return 0; +} + static int write_capture(int descriptor, int is_socket, uint64_t record_offset) { struct xprobe_cupti_output_header header; @@ -948,7 +1198,9 @@ static int write_capture(int descriptor, int is_socket, uint64_t record_offset) is_socket != 0 ? send_all : write_all; int result; - if (record_offset > available) { + if ((capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE && + record_offset != 0U) || + record_offset > available) { fprintf(stderr, "xprobe CUPTI: requested record offset %llu exceeds committed " "record count %llu\n", @@ -958,7 +1210,9 @@ static int write_capture(int descriptor, int is_socket, uint64_t record_offset) } initialize_output_header(&header, available, record_offset); result = write_function(descriptor, &header, sizeof(header)); - if (result == 0) { + if (result == 0 && capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { + result = write_aggregate_records(descriptor, write_function); + } else if (result == 0) { result = write_function(descriptor, records + record_offset, (available - record_offset) * sizeof(records[0])); } @@ -1035,6 +1289,27 @@ static int valid_filter(const struct xprobe_cupti_filter *filter) bounded_name_length(filter->name) < XPROBE_CUPTI_NAME_LENGTH; } +static int valid_aggregate_filters( + const struct xprobe_cupti_filter filters[XPROBE_CUPTI_FILTER_COUNT]) +{ + uint32_t start = filters[0].record_kind; + uint32_t end = filters[1].record_kind; + + if (!((start == XPROBE_CUPTI_GPU_KERNEL_START && + end == XPROBE_CUPTI_GPU_KERNEL_END) || + (start == XPROBE_CUPTI_GPU_MEMCPY_START && + end == XPROBE_CUPTI_GPU_MEMCPY_END) || + (start == XPROBE_CUPTI_GPU_MEMSET_START && + end == XPROBE_CUPTI_GPU_MEMSET_END))) { + return 0; + } + return filters[0].api_domain == filters[1].api_domain && + filters[0].memcpy_kind == filters[1].memcpy_kind && + filters[0].name_match == filters[1].name_match && + memcmp(filters[0].name, filters[1].name, + XPROBE_CUPTI_NAME_LENGTH) == 0; +} + static int arm_capture(const struct xprobe_cupti_control_request *request) { int status; @@ -1046,12 +1321,21 @@ static int arm_capture(const struct xprobe_cupti_control_request *request) return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } } + if (request->capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE && + valid_aggregate_filters(request->filters) == 0) { + fprintf(stderr, + "xprobe CUPTI: aggregate capture requires one matching " + "activity start/end pair\n"); + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } if (subscriber_active != 0 || enabled_activity_mask != 0U) { status = shutdown_agent(); if (status != XPROBE_CUPTI_AGENT_READY) { return status; } } + capture_mode = request->capture_mode; if (allocate_capture(request->record_capacity) != XPROBE_CUPTI_AGENT_READY) { remember_output_error(); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; @@ -1127,7 +1411,9 @@ static void *serve_snapshots(void *unused) sizeof(request.magic)) != 0 || request.version != XPROBE_CUPTI_CONTROL_VERSION || (request.command < XPROBE_CUPTI_CONTROL_ARM || - request.command > XPROBE_CUPTI_CONTROL_CLOSE)) { + request.command > XPROBE_CUPTI_CONTROL_CLOSE) || + request.capture_mode > XPROBE_CUPTI_CAPTURE_AGGREGATE || + request.reserved != 0U) { fprintf(stderr, "xprobe CUPTI: invalid snapshot control request\n"); } else if (pthread_mutex_lock(&flush_mutex) != 0) { fprintf(stderr, "xprobe CUPTI: failed to lock snapshot flush mutex\n"); diff --git a/cupti/tests/cuda_multisource_fixture.cu b/cupti/tests/cuda_multisource_fixture.cu index 5024aea..06114e9 100644 --- a/cupti/tests/cuda_multisource_fixture.cu +++ b/cupti/tests/cuda_multisource_fixture.cu @@ -8,6 +8,11 @@ __global__ void xprobe_multisource_kernel(int *output) *output += 1; } +__global__ void xprobe_auxiliary_kernel(int *output) +{ + *output += 2; +} + extern "C" __attribute__((noinline, visibility("default"))) void xprobe_request_marker() { @@ -55,6 +60,10 @@ int main(int argc, char **argv) xprobe_request_marker(); xprobe_multisource_kernel<<<1, 1>>>(device_output); result = cudaGetLastError(); + if (result == cudaSuccess) { + xprobe_auxiliary_kernel<<<1, 1>>>(device_output); + result = cudaGetLastError(); + } if (result == cudaSuccess) { result = cudaDeviceSynchronize(); } diff --git a/cupti/tests/cupti_smoke.c b/cupti/tests/cupti_smoke.c index 125d658..8f770c4 100644 --- a/cupti/tests/cupti_smoke.c +++ b/cupti/tests/cupti_smoke.c @@ -13,30 +13,34 @@ int main(void) if (sizeof(struct xprobe_cupti_record) != 200U) { return 3; } - if (offsetof(struct xprobe_cupti_output_header, feature_flags) != 20U) { + if (sizeof(struct xprobe_cupti_aggregate_record) != 200U) { return 4; } + if (offsetof(struct xprobe_cupti_output_header, feature_flags) != 20U) { + return 5; + } if (offsetof(struct xprobe_cupti_output_header, record_count) != 32U || offsetof(struct xprobe_cupti_output_header, unknown_records) != 72U || offsetof(struct xprobe_cupti_output_header, record_offset) != 80U) { - return 5; + return 6; } if (offsetof(struct xprobe_cupti_record, grid_x) != 44U || offsetof(struct xprobe_cupti_record, grid_z) != 52U || offsetof(struct xprobe_cupti_record, block_x) != 56U || offsetof(struct xprobe_cupti_record, runtime_correlation_id) != 68U || offsetof(struct xprobe_cupti_record, name) != 72U) { - return 6; + return 7; } if (sizeof(struct xprobe_cupti_filter) != 144U || - sizeof(struct xprobe_cupti_control_request) != 320U) { - return 7; + sizeof(struct xprobe_cupti_control_request) != 328U) { + return 8; } if (offsetof(struct xprobe_cupti_control_request, record_capacity) != 16U || offsetof(struct xprobe_cupti_control_request, record_offset) != 24U || - offsetof(struct xprobe_cupti_control_request, filters) != 32U || + offsetof(struct xprobe_cupti_control_request, capture_mode) != 32U || + offsetof(struct xprobe_cupti_control_request, filters) != 40U || offsetof(struct xprobe_cupti_filter, name) != 16U) { - return 8; + return 9; } return 0; } diff --git a/schemas/aggregate-inventory-result.schema.json b/schemas/aggregate-inventory-result.schema.json new file mode 100644 index 0000000..5345973 --- /dev/null +++ b/schemas/aggregate-inventory-result.schema.json @@ -0,0 +1,279 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "title": "AggregateInventoryResult", + "type": "object", + "properties": { + "collection": { + "$ref": "#/$defs/AggregateCollectionSummary" + }, + "inventory": { + "$ref": "#/$defs/AggregateInventory" + }, + "ok": { + "type": "boolean" + }, + "schema_version": { + "$ref": "#/$defs/SchemaVersion" + }, + "session_id": { + "type": "string" + }, + "status": { + "$ref": "#/$defs/SessionStatus" + }, + "warnings": { + "type": "array", + "default": [], + "items": { + "$ref": "#/$defs/Warning" + } + } + }, + "additionalProperties": false, + "required": [ + "schema_version", + "ok", + "session_id", + "status", + "inventory", + "collection" + ], + "$defs": { + "AggregateActivity": { + "type": "string", + "enum": [ + "kernel", + "memcpy", + "memset" + ] + }, + "AggregateCollectionSummary": { + "type": "object", + "properties": { + "completeness": { + "$ref": "#/$defs/CaptureCompleteness" + }, + "dropped_activities": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "group_capacity": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "grouped_activities": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "groups": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "observed_activities": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "occupied_slots": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "table_utilization": { + "type": "number", + "format": "double" + } + }, + "additionalProperties": false, + "required": [ + "completeness", + "observed_activities", + "grouped_activities", + "dropped_activities", + "group_capacity", + "groups", + "occupied_slots", + "table_utilization" + ] + }, + "AggregateDuration": { + "type": "object", + "properties": { + "max": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "mean": { + "type": "number", + "format": "double" + }, + "min": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "total": { + "type": "integer", + "format": "uint64", + "minimum": 0 + } + }, + "additionalProperties": false, + "required": [ + "min", + "mean", + "max", + "total" + ] + }, + "AggregateGroup": { + "type": "object", + "properties": { + "activity": { + "$ref": "#/$defs/AggregateActivity" + }, + "count": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "device_id": { + "type": [ + "integer", + "null" + ], + "format": "uint32", + "minimum": 0 + }, + "duration_ns": { + "$ref": "#/$defs/AggregateDuration" + }, + "end_selector_hint": { + "type": "string" + }, + "memcpy_kind": { + "anyOf": [ + { + "$ref": "#/$defs/MemcpyKind" + }, + { + "type": "null" + } + ] + }, + "name": { + "type": [ + "string", + "null" + ] + }, + "start_selector_hint": { + "type": "string" + }, + "total_bytes": { + "type": [ + "integer", + "null" + ], + "format": "uint64", + "minimum": 0 + } + }, + "additionalProperties": false, + "required": [ + "activity", + "start_selector_hint", + "end_selector_hint", + "count", + "duration_ns" + ] + }, + "AggregateInventory": { + "type": "object", + "properties": { + "end_selector": { + "type": "string" + }, + "groups": { + "type": "array", + "items": { + "$ref": "#/$defs/AggregateGroup" + } + }, + "name": { + "type": [ + "string", + "null" + ] + }, + "start_selector": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "start_selector", + "end_selector", + "groups" + ] + }, + "CaptureCompleteness": { + "type": "string", + "enum": [ + "complete" + ] + }, + "MemcpyKind": { + "type": "string", + "enum": [ + "HtoD", + "DtoH", + "DtoD", + "HtoH", + "PtoP", + "unknown" + ] + }, + "SchemaVersion": { + "type": "string", + "enum": [ + "2.0" + ] + }, + "SessionStatus": { + "type": "string", + "enum": [ + "completed", + "timed_out", + "cancelled", + "failed" + ] + }, + "Warning": { + "type": "object", + "properties": { + "code": { + "type": "string" + }, + "details": { + "type": "object", + "additionalProperties": true, + "default": {} + }, + "message": { + "type": "string" + } + }, + "additionalProperties": false, + "required": [ + "code", + "message" + ] + } + } +} diff --git a/schemas/measurement-spec.schema.json b/schemas/measurement-spec.schema.json index 12c56ba..28f333d 100644 --- a/schemas/measurement-spec.schema.json +++ b/schemas/measurement-spec.schema.json @@ -18,10 +18,27 @@ "$ref": "#/$defs/MatchPolicy" }, "max_events": { - "type": "integer", + "type": [ + "integer", + "null" + ], + "format": "uint64", + "default": null, + "minimum": 0 + }, + "max_groups": { + "type": [ + "integer", + "null" + ], "format": "uint64", + "default": null, "minimum": 0 }, + "measurement_mode": { + "$ref": "#/$defs/MeasurementMode", + "default": "exact" + }, "name": { "type": [ "string", @@ -58,8 +75,7 @@ "start_selector", "end_selector", "match_policy", - "timeout_ms", - "max_events" + "timeout_ms" ], "$defs": { "MatchPolicy": { @@ -72,6 +88,13 @@ "stream_order" ] }, + "MeasurementMode": { + "type": "string", + "enum": [ + "exact", + "aggregate" + ] + }, "SchemaVersion": { "type": "string", "enum": [ diff --git a/tests/integration/run-inject-live.sh b/tests/integration/run-inject-live.sh index c5fecbd..3eb451e 100755 --- a/tests/integration/run-inject-live.sh +++ b/tests/integration/run-inject-live.sh @@ -56,8 +56,8 @@ run_measure() { "${xprobe_bin}" measure \ --pid "${target_pid}" \ "${agent_args[@]}" \ - --from 'cuda:kernel_start:name~xprobe_multisource_kernel.*' \ - --to 'cuda:kernel_end:name~xprobe_multisource_kernel.*' \ + --from 'cuda:kernel_start:name~.*xprobe_multisource_kernel.*' \ + --to 'cuda:kernel_end:name~.*xprobe_multisource_kernel.*' \ --match exact \ --samples 3 \ --max-events 32 \ @@ -88,9 +88,52 @@ run_api_measure() { >"$1.json" 2>"$1.stderr" } +run_aggregate_measure() { + agent_args=() + if [[ ${auto_agent} != 1 ]]; then + agent_args=(--agent "${agent}") + fi + "${xprobe_bin}" measure \ + --pid "${target_pid}" \ + "${agent_args[@]}" \ + --from 'cuda:kernel_start:name~.*xprobe_multisource_kernel.*' \ + --to 'cuda:kernel_end:name~.*xprobe_multisource_kernel.*' \ + --match exact \ + --aggregate \ + --duration-ms 200 \ + --max-groups 1 \ + --timeout-ms 10000 \ + --json --non-interactive --no-color \ + >"$1.json" 2>"$1.stderr" +} + +run_aggregate_limit() { + agent_args=() + if [[ ${auto_agent} != 1 ]]; then + agent_args=(--agent "${agent}") + fi + if "${xprobe_bin}" measure \ + --pid "${target_pid}" \ + "${agent_args[@]}" \ + --from 'cuda:kernel_start' \ + --to 'cuda:kernel_end' \ + --match exact \ + --aggregate \ + --duration-ms 200 \ + --max-groups 1 \ + --timeout-ms 10000 \ + --json --non-interactive --no-color \ + >"$1.json" 2>"$1.stderr"; then + echo "aggregate capacity test unexpectedly succeeded" >&2 + return 1 + fi +} + run_measure "$1/first" run_measure "$1/second" run_api_measure "$1/api" +run_aggregate_measure "$1/aggregate" +run_aggregate_limit "$1/aggregate-limit" if [[ ${auto_agent} == 1 ]]; then mapped_agents=$(awk '$0 ~ /lib\/xprobe\/cuda(12|13)\/libxprobe-cupti.so/ {print $NF}' "/proc/${target_pid}/maps" | sort -u | wc -l) diff --git a/tests/integration/test_cupti.py b/tests/integration/test_cupti.py index 4e6d011..089db29 100755 --- a/tests/integration/test_cupti.py +++ b/tests/integration/test_cupti.py @@ -49,8 +49,8 @@ def read_capture(path: pathlib.Path) -> tuple[dict[str, int], list[dict[str, int if len(data) < HEADER.size: raise AssertionError("CUPTI capture is shorter than its header") fields = HEADER.unpack_from(data) - if fields[1] != 3: - raise AssertionError(f"expected capture ABI 3, found {fields[1]}") + if fields[1] != 4: + raise AssertionError(f"expected capture ABI 4, found {fields[1]}") header = { "abi_version": fields[1], "header_size": fields[2], @@ -67,7 +67,7 @@ def read_capture(path: pathlib.Path) -> tuple[dict[str, int], list[dict[str, int "record_offset": fields[13], } assert fields[0] == b"XPCUPTI\0" - assert header["abi_version"] == 3 + assert header["abi_version"] == 4 assert header["header_size"] == HEADER.size assert header["record_size"] == RECORD.size assert header["feature_flags"] in {2, 3} diff --git a/tests/integration/test_inject.py b/tests/integration/test_inject.py index 2a2140b..c6a2536 100755 --- a/tests/integration/test_inject.py +++ b/tests/integration/test_inject.py @@ -79,6 +79,30 @@ def main() -> None: assert any(warning["code"] == "CUPTI_AGENT_INJECTED" for warning in first["warnings"]) assert all(warning["code"] != "CUPTI_AGENT_INJECTED" for warning in second["warnings"]) assert all(warning["code"] != "CUPTI_AGENT_INJECTED" for warning in api["warnings"]) + aggregate = json.loads((output / "aggregate.json").read_text()) + assert aggregate["ok"] is True + assert aggregate["status"] == "completed" + assert aggregate["collection"]["observed_activities"] > 0 + assert ( + aggregate["collection"]["observed_activities"] + > aggregate["collection"]["group_capacity"] + ) + assert aggregate["collection"]["observed_activities"] == aggregate["collection"]["grouped_activities"] + assert aggregate["collection"]["dropped_activities"] == 0 + assert aggregate["collection"]["groups"] == 1 + assert aggregate["collection"]["occupied_slots"] == 1 + assert aggregate["inventory"]["groups"][0]["activity"] == "kernel" + assert aggregate["inventory"]["groups"][0]["count"] > 0 + assert aggregate["inventory"]["groups"][0]["duration_ns"]["min"] > 0 + assert all( + warning["code"] != "CUPTI_AGENT_INJECTED" + for warning in aggregate["warnings"] + ) + aggregate_limit = json.loads((output / "aggregate-limit.json").read_text()) + assert aggregate_limit["ok"] is False + assert aggregate_limit["error"]["code"] == "EVENT_RATE_TOO_HIGH" + assert aggregate_limit["error"]["details"]["group_capacity"] == 1 + assert aggregate_limit["error"]["details"]["observed_activities"] >= 2 assert int((output / "mapped-agents.txt").read_text()) == 1 if package is not None: expected_major = 12 if ":12." in sys.argv[1] else 13 @@ -96,7 +120,7 @@ def main() -> None: "schema_version": "2.0", "ok": True, "gpu": "NVIDIA GeForce RTX 3060 Laptop GPU", - "measurements": 3, + "measurements": 4, }, sort_keys=True, ) diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 63742c0..2c8567d 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -25,8 +25,10 @@ use xprobe_core::{cupti_compat, discover, doctor, inject, inspect, resolve, vali use xprobe_correlator::{MeasureError, MeasureOptions, measure}; use xprobe_exporter::{events_to_chrome_trace, events_to_jsonl}; use xprobe_protocol::{ - CapabilityReport, CheckResult, ClockDomain, CuptiCollectionSummary, DiscoveryResult, ErrorCode, - ErrorResponse, Event, EventSource, EventType, ExportFormat, HostCaptureResult, MatchPolicy, + AggregateActivity, AggregateCollectionSummary, AggregateDuration, AggregateGroup, + AggregateInventory, AggregateInventoryResult, CapabilityReport, CaptureCompleteness, + CheckResult, ClockDomain, CuptiCollectionSummary, DiscoveryResult, ErrorCode, ErrorResponse, + Event, EventSource, EventType, ExportFormat, HostCaptureResult, MatchPolicy, MeasurementMode, MeasurementResult, MeasurementSpec, MemcpyKind, ProcessReport, ResolvedCudaSelector, ResolvedProbe, SchemaVersion, SessionStatus, TargetIdentity, TraceExportResult, ValidationResult, Warning, XprobeError, @@ -54,7 +56,7 @@ enum Command { /// Validate two event selectors and their correlation policy. Validate(ValidateArgs), /// Measure latency between two events in files or a running process. - Measure(MeasureArgs), + Measure(Box), /// Run a live measurement from a versioned specification. #[command(hide = true)] Trace(TraceArgs), @@ -253,7 +255,7 @@ struct ValidateArgs { #[derive(Debug, Args)] struct MeasureArgs { /// Versioned `MeasurementSpec` JSON for a live target. - #[arg(long, conflicts_with_all = ["input", "pid", "from", "to", "match_policy", "samples", "duration_ms", "timeout_ms", "max_events", "name"])] + #[arg(long, conflicts_with_all = ["input", "pid", "from", "to", "match_policy", "samples", "duration_ms", "timeout_ms", "max_events", "aggregate", "max_groups", "name"])] spec: Option, /// Completed CUPTI binary, host capture JSON, or Event JSONL; repeat to merge. @@ -308,10 +310,28 @@ struct MeasureArgs { #[arg(long, default_value_t = 100_000)] max_events: usize, + /// Return bounded activity aggregates instead of exact event evidence. + #[arg(long, conflicts_with_all = ["input", "samples", "events_out", "format"])] + aggregate: bool, + + /// Bound distinct activity groups retained by --aggregate. + #[arg( + long, + requires = "aggregate", + default_value_if("aggregate", "true", "4096") + )] + max_groups: Option, + /// Optional measurement name. #[arg(long)] name: Option, + #[command(flatten)] + output: CommonOutputArgs, +} + +#[derive(Debug, Clone, Copy, Args)] +struct CommonOutputArgs { /// Emit only the versioned JSON result on stdout. #[arg(long)] json: bool, @@ -417,7 +437,7 @@ fn main() -> ExitCode { Command::Inspect(args) => run_inspect(args), Command::Resolve(args) => run_resolve(args), Command::Validate(args) => run_validate(args), - Command::Measure(args) => run_measure(args), + Command::Measure(args) => run_measure(*args), Command::Trace(args) => run_trace(args), Command::Export(args) => run_export(args), Command::Capture(args) | Command::Dev(args) => run_capture(args), @@ -461,6 +481,14 @@ fn run_trace(args: TraceArgs) -> ExitCode { ); } }; + if spec.measurement_mode != MeasurementMode::Exact || spec.max_groups.is_some() { + return emit_error( + ErrorCode::InvalidEventSelector, + "legacy trace accepts only exact MeasurementSpec requests".to_owned(), + true, + json, + ); + } let samples = match spec.samples.map(usize::try_from).transpose() { Ok(samples) => samples, Err(error) => { @@ -472,7 +500,15 @@ fn run_trace(args: TraceArgs) -> ExitCode { ); } }; - let max_events = match usize::try_from(spec.max_events) { + let Some(max_events) = spec.max_events else { + return emit_error( + ErrorCode::SessionLimitExceeded, + "exact MeasurementSpec requires max_events".to_owned(), + true, + json, + ); + }; + let max_events = match usize::try_from(max_events) { Ok(max_events) => max_events, Err(error) => { return emit_error( @@ -490,6 +526,7 @@ fn run_trace(args: TraceArgs) -> ExitCode { agent_path: None, timeout: Duration::from_millis(spec.timeout_ms), match_policy_text: match_policy_name(spec.match_policy), + mode: MeasurementMode::Exact, options: MeasureOptions { session_id: format!("xp_trace_{}", std::process::id()), name: spec.name, @@ -676,10 +713,15 @@ fn run_measure(args: MeasureArgs) -> ExitCode { samples, duration_ms, max_events, + aggregate, + max_groups, name, - json, - no_color: _, - non_interactive: _, + output: + CommonOutputArgs { + json, + no_color: _, + non_interactive: _, + }, } = args; if let Some(spec) = spec { return run_measure_spec( @@ -710,13 +752,26 @@ fn run_measure(args: MeasureArgs) -> ExitCode { dropped_events: 0, }; - if pid.is_some() && !input.is_empty() { - return emit_error( - ErrorCode::InvalidEventSelector, - "--pid and --input select different collection modes".to_owned(), - true, - json, - ); + if let Err(error) = validate_measure_source(pid, &input) { + return emit_command_failure(error, json); + } + if aggregate { + let request = match direct_aggregate_request( + pid, + cupti_socket, + agent, + timeout_ms, + duration_ms, + max_groups, + options, + ) { + Ok(request) => request, + Err(error) => return emit_command_failure(error, json), + }; + return match collect_live_aggregate(&request) { + Ok(result) => emit_aggregate_inventory(&result, json), + Err(error) => emit_command_failure(error, json), + }; } if let Some(pid) = pid { let request = LiveMeasureRequest { @@ -726,6 +781,7 @@ fn run_measure(args: MeasureArgs) -> ExitCode { agent_path: agent, timeout: Duration::from_millis(timeout_ms), match_policy_text: match_policy_name(match_policy), + mode: MeasurementMode::Exact, options, }; return match collect_live_measurement(&request) { @@ -744,6 +800,56 @@ fn run_measure(args: MeasureArgs) -> ExitCode { ) } +fn validate_measure_source(pid: Option, input: &[PathBuf]) -> Result<(), CommandFailure> { + if pid.is_some() && !input.is_empty() { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "--pid and --input select different collection modes", + true, + )); + } + Ok(()) +} + +fn direct_aggregate_request( + pid: Option, + cupti_socket: Option, + agent_path: Option, + timeout_ms: u64, + duration_ms: Option, + max_groups: Option, + options: MeasureOptions, +) -> Result { + let pid = pid.ok_or_else(|| { + CommandFailure::new( + ErrorCode::InvalidEventSelector, + "--aggregate requires --pid", + true, + ) + })?; + let duration_ms = duration_ms.ok_or_else(|| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "--aggregate requires --duration-ms", + true, + ) + })?; + Ok(LiveMeasureRequest { + pid, + expected_target: None, + cupti_socket, + agent_path, + timeout: Duration::from_millis(timeout_ms), + match_policy_text: match_policy_name(options.match_policy), + mode: MeasurementMode::Aggregate, + options: MeasureOptions { + duration: Some(Duration::from_millis(duration_ms)), + max_events: max_groups.expect("clap supplies aggregate group capacity"), + ..options + }, + }) +} + fn parse_measure_selection( from: Option, to: Option, @@ -861,16 +967,9 @@ fn run_measure_spec( ); } }; - let max_events = match usize::try_from(spec.max_events) { - Ok(max_events) => max_events, - Err(error) => { - return emit_error( - ErrorCode::SessionLimitExceeded, - format!("MeasurementSpec max_events exceed this platform: {error}"), - true, - json, - ); - } + let max_events = match measurement_spec_capacity(&spec) { + Ok(capacity) => capacity, + Err(error) => return emit_command_failure(error, json), }; let request = LiveMeasureRequest { pid: spec.target.pid, @@ -879,6 +978,7 @@ fn run_measure_spec( agent_path, timeout: Duration::from_millis(spec.timeout_ms), match_policy_text: match_policy_name(spec.match_policy), + mode: spec.measurement_mode, options: MeasureOptions { session_id: format!("xp_measure_{}", std::process::id()), name: spec.name, @@ -891,12 +991,61 @@ fn run_measure_spec( dropped_events: 0, }, }; - match collect_live_measurement(&request) { - Ok(execution) => finish_measurement(&execution, events_out, format, json), - Err(error) => finish_measurement_failure(error, events_out, format, json), + match request.mode { + MeasurementMode::Exact => match collect_live_measurement(&request) { + Ok(execution) => finish_measurement(&execution, events_out, format, json), + Err(error) => finish_measurement_failure(error, events_out, format, json), + }, + MeasurementMode::Aggregate => { + if events_out.is_some() || format.is_some() { + return emit_error( + ErrorCode::TraceExportFailed, + "aggregate measurement does not produce event evidence artifacts".to_owned(), + true, + json, + ); + } + match collect_live_aggregate(&request) { + Ok(result) => emit_aggregate_inventory(&result, json), + Err(error) => emit_command_failure(error, json), + } + } } } +fn measurement_spec_capacity(spec: &MeasurementSpec) -> Result { + let capacity = match spec.measurement_mode { + MeasurementMode::Exact if spec.max_groups.is_some() => { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "MeasurementSpec max_groups is only valid in aggregate mode", + true, + )); + } + MeasurementMode::Exact => spec.max_events.ok_or_else(|| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "exact MeasurementSpec requires max_events", + true, + ) + })?, + MeasurementMode::Aggregate => spec.max_groups.ok_or_else(|| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "aggregate MeasurementSpec requires max_groups", + true, + ) + })?, + }; + usize::try_from(capacity).map_err(|error| { + CommandFailure::new( + ErrorCode::SessionLimitExceeded, + format!("MeasurementSpec capacity exceeds this platform: {error}"), + true, + ) + }) +} + fn finish_measurement( execution: &MeasurementExecution, events_out: Option<&Path>, @@ -1151,6 +1300,7 @@ struct LiveMeasureRequest { agent_path: Option, timeout: Duration, match_policy_text: &'static str, + mode: MeasurementMode, options: MeasureOptions, } @@ -1220,6 +1370,7 @@ struct LiveCollection { fn collect_live_measurement( request: &LiveMeasureRequest, ) -> Result { + assert_eq!(request.mode, MeasurementMode::Exact); validate_live_limits(request)?; let mut collection = prepare_live_collection(request)?; @@ -1317,6 +1468,62 @@ fn collect_live_measurement( Ok(execution) } +fn collect_live_aggregate( + request: &LiveMeasureRequest, +) -> Result { + assert_eq!(request.mode, MeasurementMode::Aggregate); + validate_live_limits(request)?; + let mut collection = prepare_live_collection(request)?; + let collection_end = collection + .collection_end + .expect("aggregate validation requires duration"); + if collection_end > collection.deadline { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "aggregate duration must not exceed timeout", + true, + )); + } + + while Instant::now() < collection_end { + let remaining = collection_end.saturating_duration_since(Instant::now()); + thread::sleep(remaining.min(Duration::from_millis(10))); + } + let target_result = inspect::verify_target(&collection.report.target) + .map_err(|error| CommandFailure::new(error.code(), error.to_string(), error.recoverable())); + let stop_result = stop_cupti_agent(&mut collection, request); + match (target_result, stop_result) { + (Err(error), Ok(())) | (Ok(()), Err(error)) => return Err(error), + (Err(target), Err(cleanup)) => { + return Err(CommandFailure::new( + ErrorCode::CleanupFailed, + format!( + "{}; target verification also failed: {}", + cleanup.message, target.message + ), + cleanup.recoverable, + ) + .with_detail("cleanup_phase", "stop_cupti") + .with_detail("target_error_code", target.code.to_string()) + .with_hint("verify the target state before starting another measurement")); + } + (Ok(()), Ok(())) => {} + } + + let capture = collection.latest_cuda.as_ref().ok_or_else(|| { + CommandFailure::new( + ErrorCode::CuptiNotAvailable, + "CUPTI Agent returned no aggregate capture", + false, + ) + })?; + let mut result = aggregate_inventory_result(capture, request)?; + if collection.agent_injected { + result.warnings.push(injection_warning(&collection)); + } + Ok(result) +} + fn validate_live_limits(request: &LiveMeasureRequest) -> Result<(), CommandFailure> { if request.timeout.is_zero() { return Err(CommandFailure::new( @@ -1339,6 +1546,25 @@ fn validate_live_limits(request: &LiveMeasureRequest) -> Result<(), CommandFailu true, )); } + if request.mode == MeasurementMode::Aggregate { + if request.options.samples.is_some() + || request.options.duration.is_none() + || request.options.match_policy != MatchPolicy::Exact + { + return Err(CommandFailure::new( + ErrorCode::InvalidCorrelationPolicy, + "aggregate measurement requires duration, exact activity endpoints, and no samples limit", + true, + )); + } + if request.options.duration > Some(request.timeout) { + return Err(CommandFailure::new( + ErrorCode::SessionLimitExceeded, + "aggregate duration must not exceed timeout", + true, + )); + } + } Ok(()) } @@ -1377,9 +1603,12 @@ fn prepare_live_collection(request: &LiveMeasureRequest) -> Result Result Result<(), CommandFailure> { + let (Some(start), Some(end)) = (validation.start.cuda.as_ref(), validation.end.cuda.as_ref()) + else { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate measurement supports CUPTI activity endpoints only", + true, + )); + }; + let paired = matches!( + (&start.event_type, &end.event_type), + (&EventType::GpuKernelStart, &EventType::GpuKernelEnd) + | (&EventType::GpuMemcpyStart, &EventType::GpuMemcpyEnd) + | (&EventType::GpuMemsetStart, &EventType::GpuMemsetEnd) + ); + if !paired + || start.api_domain != end.api_domain + || start.api_name != end.api_name + || start.kernel_name_regex != end.kernel_name_regex + || start.memcpy_kind != end.memcpy_kind + { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate measurement requires matching kernel, memcpy, or memset start/end selectors", + true, + )); + } + if let Some(pattern) = start.kernel_name_regex.as_deref() + && pattern != ".*" + && matches!( + bounded_kernel_name_filter(pattern), + cupti::CuptiNameFilter::Any + ) + { + return Err(CommandFailure::new( + ErrorCode::InvalidEventSelector, + "aggregate kernel regex must be an exact, prefix, suffix, or contains pattern that the CUPTI Agent can apply", + true, + ) + .with_hint("use ^literal$, ^literal.*, .*literal$, or .*literal.*")); + } + Ok(()) +} + fn arm_cupti_capture( activation: &CuptiActivation, config: Option<&cupti::CuptiArmConfig>, @@ -1494,6 +1767,7 @@ fn cleanup_failed_arm( fn cupti_arm_config( validation: &ValidationResult, max_events: usize, + mode: MeasurementMode, ) -> Result, CommandFailure> { if !validation.requirements.needs_cupti { return Ok(None); @@ -1512,6 +1786,10 @@ fn cupti_arm_config( })?; Ok(Some(cupti::CuptiArmConfig { record_capacity, + capture_mode: match mode { + MeasurementMode::Exact => cupti::CuptiCaptureMode::Exact, + MeasurementMode::Aggregate => cupti::CuptiCaptureMode::Aggregate, + }, filters, })) } @@ -2031,6 +2309,307 @@ fn append_cuda_capture( Ok(()) } +fn aggregate_inventory_result( + capture: &cupti::CuptiCapture, + request: &LiveMeasureRequest, +) -> Result { + validate_aggregate_capture(capture)?; + let raw_groups = u64::try_from(capture.aggregate_groups.len()).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("aggregate group count exceeds u64: {error}"), + false, + ) + })?; + let groups = merge_aggregate_groups(capture.aggregate_groups.clone())?; + let grouped_activities = groups.iter().try_fold(0_u64, |total, group| { + total.checked_add(group.count).ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate activity count overflowed", + false, + ) + }) + })?; + if grouped_activities != capture.observed_records { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + format!( + "aggregate activity counters are inconsistent: grouped={grouped_activities}, observed={}", + capture.observed_records + ), + false, + )); + } + let distinct_groups = u64::try_from(groups.len()).map_err(|error| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + format!("distinct aggregate group count exceeds u64: {error}"), + false, + ) + })?; + let groups = groups + .into_iter() + .map(protocol_aggregate_group) + .collect::>(); + Ok(AggregateInventoryResult { + schema_version: SchemaVersion::current(), + ok: true, + session_id: request.options.session_id.clone(), + status: SessionStatus::Completed, + inventory: AggregateInventory { + name: request.options.name.clone(), + start_selector: request.options.start_selector.clone(), + end_selector: request.options.end_selector.clone(), + groups, + }, + collection: AggregateCollectionSummary { + completeness: CaptureCompleteness::Complete, + observed_activities: capture.observed_records, + grouped_activities, + dropped_activities: 0, + group_capacity: capture.record_capacity, + groups: distinct_groups, + occupied_slots: raw_groups, + table_utilization: utilization(raw_groups, capture.record_capacity), + }, + warnings: Vec::new(), + }) +} + +fn validate_aggregate_capture(capture: &cupti::CuptiCapture) -> Result<(), CommandFailure> { + if capture.state == cupti::CuptiCaptureState::LimitReached { + return Err(CommandFailure::new( + ErrorCode::EventRateTooHigh, + "aggregate group table reached its configured capacity", + true, + ) + .with_detail("group_capacity", capture.record_capacity) + .with_detail("observed_activities", capture.observed_records) + .with_hint("narrow the selector or explicitly increase max-groups")); + } + if capture.state != cupti::CuptiCaptureState::Stopped + || capture.stop_reason != cupti::CuptiStopReason::Requested + { + return Err(CommandFailure::new( + ErrorCode::CuptiNotAvailable, + format!( + "aggregate CUPTI capture ended in state {:?} with reason {:?}", + capture.state, capture.stop_reason + ), + true, + )); + } + if capture.dropped_records != 0 { + return Err(CommandFailure::new( + ErrorCode::EventsDropped, + "aggregate CUPTI capture dropped activities", + true, + ) + .with_detail("dropped_activities", capture.dropped_records) + .with_hint("reduce capture duration or narrow the activity selector")); + } + if capture.unknown_records != 0 { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate CUPTI capture contains unknown activities", + false, + ) + .with_detail("unknown_records", capture.unknown_records) + .with_hint("rebuild xprobe and the CUPTI Agent from the same release")); + } + if !capture.events.is_empty() { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate CUPTI capture unexpectedly contains exact events", + false, + )); + } + Ok(()) +} + +fn merge_aggregate_groups( + mut groups: Vec, +) -> Result, CommandFailure> { + groups.sort_by(|left, right| { + aggregate_activity_rank(left.activity) + .cmp(&aggregate_activity_rank(right.activity)) + .then_with(|| left.name.cmp(&right.name)) + .then_with(|| left.device_id.cmp(&right.device_id)) + .then_with(|| { + memcpy_kind_rank(left.memcpy_kind.as_ref()) + .cmp(&memcpy_kind_rank(right.memcpy_kind.as_ref())) + }) + }); + let mut merged: Vec = Vec::with_capacity(groups.len()); + for group in groups { + let Some(previous) = merged.last_mut() else { + merged.push(group); + continue; + }; + if previous.activity != group.activity + || previous.name != group.name + || previous.device_id != group.device_id + || previous.memcpy_kind != group.memcpy_kind + { + merged.push(group); + continue; + } + previous.count = previous.count.checked_add(group.count).ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate group count overflowed while merging", + false, + ) + })?; + previous.total_duration_ns = previous + .total_duration_ns + .checked_add(group.total_duration_ns) + .ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate duration overflowed while merging", + false, + ) + })?; + previous.min_duration_ns = previous.min_duration_ns.min(group.min_duration_ns); + previous.max_duration_ns = previous.max_duration_ns.max(group.max_duration_ns); + previous.total_bytes = match (previous.total_bytes, group.total_bytes) { + (Some(left), Some(right)) => Some(left.checked_add(right).ok_or_else(|| { + CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate byte count overflowed while merging", + false, + ) + })?), + (None, None) => None, + _ => { + return Err(CommandFailure::new( + ErrorCode::TraceExportFailed, + "aggregate byte counters are inconsistent", + false, + )); + } + }; + } + Ok(merged) +} + +const fn aggregate_activity_rank(activity: cupti::CuptiAggregateActivity) -> u8 { + match activity { + cupti::CuptiAggregateActivity::Kernel => 0, + cupti::CuptiAggregateActivity::Memcpy => 1, + cupti::CuptiAggregateActivity::Memset => 2, + } +} + +const fn memcpy_kind_rank(kind: Option<&MemcpyKind>) -> u8 { + match kind { + None => 0, + Some(&MemcpyKind::Unknown) => 1, + Some(&MemcpyKind::HostToDevice) => 2, + Some(&MemcpyKind::DeviceToHost) => 3, + Some(&MemcpyKind::DeviceToDevice) => 4, + Some(&MemcpyKind::HostToHost) => 5, + Some(&MemcpyKind::PeerToPeer) => 6, + } +} + +#[allow(clippy::cast_precision_loss)] +fn protocol_aggregate_group(group: cupti::CuptiAggregateGroup) -> AggregateGroup { + let (start_selector_hint, end_selector_hint) = aggregate_selector_hints(&group); + AggregateGroup { + activity: match group.activity { + cupti::CuptiAggregateActivity::Kernel => AggregateActivity::Kernel, + cupti::CuptiAggregateActivity::Memcpy => AggregateActivity::Memcpy, + cupti::CuptiAggregateActivity::Memset => AggregateActivity::Memset, + }, + name: group.name, + device_id: group.device_id, + memcpy_kind: group.memcpy_kind, + start_selector_hint, + end_selector_hint, + count: group.count, + duration_ns: AggregateDuration { + min: group.min_duration_ns, + mean: group.total_duration_ns as f64 / group.count as f64, + max: group.max_duration_ns, + total: group.total_duration_ns, + }, + total_bytes: group.total_bytes, + } +} + +fn aggregate_selector_hints(group: &cupti::CuptiAggregateGroup) -> (String, String) { + match group.activity { + cupti::CuptiAggregateActivity::Kernel => { + let name = group.name.as_deref().expect("kernel aggregate has a name"); + let escaped = escape_selector_regex(name); + ( + format!("cuda:kernel_start:name~^{escaped}$"), + format!("cuda:kernel_end:name~^{escaped}$"), + ) + } + cupti::CuptiAggregateActivity::Memcpy => { + let kind = match group.memcpy_kind.as_ref() { + Some(MemcpyKind::HostToDevice) => Some("HtoD"), + Some(MemcpyKind::DeviceToHost) => Some("DtoH"), + Some(MemcpyKind::DeviceToDevice) => Some("DtoD"), + Some(MemcpyKind::HostToHost) => Some("HtoH"), + Some(MemcpyKind::PeerToPeer) => Some("PtoP"), + Some(MemcpyKind::Unknown) | None => None, + }; + kind.map_or_else( + || ("cuda:memcpy_start".to_owned(), "cuda:memcpy_end".to_owned()), + |kind| { + ( + format!("cuda:memcpy_start:kind={kind}"), + format!("cuda:memcpy_end:kind={kind}"), + ) + }, + ) + } + cupti::CuptiAggregateActivity::Memset => { + ("cuda:memset_start".to_owned(), "cuda:memset_end".to_owned()) + } + } +} + +fn escape_selector_regex(name: &str) -> String { + let mut escaped = String::with_capacity(name.len()); + for character in name.chars() { + if "\\.^$|?*+()[]{}".contains(character) { + escaped.push('\\'); + } + escaped.push(character); + } + escaped +} + +fn injection_warning(collection: &LiveCollection) -> Warning { + let mut details = BTreeMap::new(); + if let Some(major) = collection.agent_major { + details.insert("cuda_major".to_owned(), serde_json::Value::from(major)); + } + if let Some(path) = collection.agent_path.as_deref() { + details.insert( + "agent_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + if let Some(path) = collection.cupti_path.as_deref() { + details.insert( + "cupti_path".to_owned(), + serde_json::Value::from(path.display().to_string()), + ); + } + Warning { + code: "CUPTI_AGENT_INJECTED".to_owned(), + message: "xprobe injected the CUPTI agent and left the shared object mapped".to_owned(), + details, + } +} + fn evaluate_live_result( collection: &LiveCollection, request: &LiveMeasureRequest, @@ -2323,6 +2902,32 @@ fn emit_measurement(result: &MeasurementResult, json: bool) -> ExitCode { ExitCode::SUCCESS } +fn emit_aggregate_inventory(result: &AggregateInventoryResult, json: bool) -> ExitCode { + if json { + println!( + "{}", + serde_json::to_string_pretty(result) + .expect("aggregate inventory result must serialize") + ); + } else { + println!( + "aggregate inventory: {} activities in {} groups", + result.collection.grouped_activities, result.collection.groups + ); + for group in &result.inventory.groups { + let label = group + .name + .as_deref() + .map_or_else(|| format!("{:?}", group.activity), str::to_owned); + println!( + " {label}: count={} mean={:.0}ns min={}ns max={}ns", + group.count, group.duration_ns.mean, group.duration_ns.min, group.duration_ns.max + ); + } + } + ExitCode::SUCCESS +} + fn run_validate(args: ValidateArgs) -> ExitCode { let ValidateArgs { pid, @@ -2916,6 +3521,10 @@ mod tests { bounded_kernel_name_filter("cudaLaunch"), CuptiNameFilter::Contains("cudaLaunch".to_owned()) ); + assert_eq!( + bounded_kernel_name_filter(".*xprobe_multisource_kernel.*"), + CuptiNameFilter::Contains("xprobe_multisource_kernel".to_owned()) + ); assert_eq!( bounded_kernel_name_filter("^(flash|attention)$"), CuptiNameFilter::Any diff --git a/xprobe/cli/tests/cupti.rs b/xprobe/cli/tests/cupti.rs index 10f8773..f3f3d01 100644 --- a/xprobe/cli/tests/cupti.rs +++ b/xprobe/cli/tests/cupti.rs @@ -12,7 +12,7 @@ fn capture_path(name: &str) -> PathBuf { fn write_capture(path: &PathBuf) { let mut bytes = vec![0_u8; HEADER_SIZE + RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&4_u32.to_le_bytes()); bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[24..28].copy_from_slice(&3_u32.to_le_bytes()); diff --git a/xprobe/cli/tests/measure.rs b/xprobe/cli/tests/measure.rs index 92c8c9e..182b061 100644 --- a/xprobe/cli/tests/measure.rs +++ b/xprobe/cli/tests/measure.rs @@ -35,7 +35,7 @@ fn record(kind: u32, timestamp: u64, correlation_id: u32, name: &str) -> [u8; RE fn write_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { let mut bytes = vec![0_u8; HEADER_SIZE + records.len() * RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&4_u32.to_le_bytes()); bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[24..28].copy_from_slice(&3_u32.to_le_bytes()); @@ -54,7 +54,7 @@ fn write_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { fn write_normalized_capture(path: &PathBuf, records: &[[u8; RECORD_SIZE]]) { let mut bytes = vec![0_u8; HEADER_SIZE + records.len() * RECORD_SIZE]; bytes[0..8].copy_from_slice(b"XPCUPTI\0"); - bytes[8..12].copy_from_slice(&3_u32.to_le_bytes()); + bytes[8..12].copy_from_slice(&4_u32.to_le_bytes()); bytes[12..16].copy_from_slice(&88_u32.to_le_bytes()); bytes[16..20].copy_from_slice(&200_u32.to_le_bytes()); bytes[20..24].copy_from_slice(&1_u32.to_le_bytes()); @@ -583,3 +583,59 @@ fn requires_exactly_one_measurement_source_mode() { let error: ErrorResponse = serde_json::from_slice(&both.stdout).expect("error JSON"); assert_eq!(error.error.code, ErrorCode::InvalidEventSelector); } + +#[test] +fn aggregate_requires_a_duration_bound() { + let output = Command::new(env!("CARGO_BIN_EXE_xprobe")) + .args([ + "measure", + "--pid", + &std::process::id().to_string(), + "--from", + "cuda:kernel_start", + "--to", + "cuda:kernel_end", + "--match", + "exact", + "--aggregate", + "--json", + ]) + .output() + .expect("xprobe measure must run"); + + assert_eq!(output.status.code(), Some(1)); + let error: ErrorResponse = + serde_json::from_slice(&output.stdout).expect("stdout must contain error JSON"); + assert_eq!(error.error.code, ErrorCode::SessionLimitExceeded); + assert!(error.error.message.contains("--duration-ms")); +} + +#[test] +fn aggregate_rejects_kernel_regex_that_cannot_run_in_the_agent() { + let output = Command::new(env!("CARGO_BIN_EXE_xprobe")) + .args([ + "measure", + "--pid", + &std::process::id().to_string(), + "--from", + "cuda:kernel_start:name~^(flash|attention)$", + "--to", + "cuda:kernel_end:name~^(flash|attention)$", + "--match", + "exact", + "--aggregate", + "--duration-ms", + "1", + "--max-groups", + "8", + "--json", + ]) + .output() + .expect("xprobe measure must run"); + + assert_eq!(output.status.code(), Some(1)); + let error: ErrorResponse = + serde_json::from_slice(&output.stdout).expect("stdout must contain error JSON"); + assert_eq!(error.error.code, ErrorCode::InvalidEventSelector); + assert!(error.error.message.contains("CUPTI Agent can apply")); +} diff --git a/xprobe/cli/tests/trace.rs b/xprobe/cli/tests/trace.rs index 66ad1c7..be8787f 100644 --- a/xprobe/cli/tests/trace.rs +++ b/xprobe/cli/tests/trace.rs @@ -1,7 +1,8 @@ use std::{fs, path::PathBuf, process::Command}; use xprobe_protocol::{ - ErrorCode, ErrorResponse, MatchPolicy, MeasurementSpec, SchemaVersion, TargetIdentity, + ErrorCode, ErrorResponse, MatchPolicy, MeasurementMode, MeasurementSpec, SchemaVersion, + TargetIdentity, }; fn spec_path(name: &str) -> PathBuf { @@ -27,7 +28,9 @@ fn trace_rejects_a_reused_target_from_the_spec() { samples: Some(1), duration_ms: None, timeout_ms: 1_000, - max_events: 100, + max_events: Some(100), + measurement_mode: MeasurementMode::Exact, + max_groups: None, }; fs::write(&path, serde_json::to_vec_pretty(&spec).unwrap()) .expect("trace spec must be written"); @@ -67,7 +70,9 @@ fn measure_accepts_a_versioned_live_spec() { samples: Some(1), duration_ms: None, timeout_ms: 1_000, - max_events: 100, + max_events: Some(100), + measurement_mode: MeasurementMode::Exact, + max_groups: None, }; fs::write(&path, serde_json::to_vec_pretty(&spec).unwrap()) .expect("measurement spec must be written"); diff --git a/xprobe/collector/src/cupti.rs b/xprobe/collector/src/cupti.rs index 895ea41..6af97d2 100644 --- a/xprobe/collector/src/cupti.rs +++ b/xprobe/collector/src/cupti.rs @@ -17,19 +17,21 @@ use xprobe_protocol::{ const OUTPUT_MAGIC: &[u8; 8] = b"XPCUPTI\0"; const CONTROL_MAGIC: &[u8; 8] = b"XPCTRL\0\0"; -const CONTROL_VERSION: u32 = 3; -const ABI_VERSION: u32 = 3; +const CONTROL_VERSION: u32 = 4; +const ABI_VERSION: u32 = 4; const HEADER_SIZE: usize = 88; const HEADER_SIZE_U32: u32 = 88; const RECORD_SIZE: usize = 200; const RECORD_SIZE_U32: u32 = 200; -const CONTROL_SIZE: usize = 320; +const CONTROL_SIZE: usize = 328; const FILTER_SIZE: usize = 144; const FILTER_COUNT: usize = 2; const FILTER_NAME_SIZE: usize = 128; const FEATURE_HOST_MONOTONIC_TIMESTAMPS: u32 = 1 << 0; const FEATURE_TRANSFER_RECORDS: u32 = 1 << 1; -const SUPPORTED_FEATURES: u32 = FEATURE_HOST_MONOTONIC_TIMESTAMPS | FEATURE_TRANSFER_RECORDS; +const FEATURE_AGGREGATE_RECORDS: u32 = 1 << 2; +const SUPPORTED_FEATURES: u32 = + FEATURE_HOST_MONOTONIC_TIMESTAMPS | FEATURE_TRANSFER_RECORDS | FEATURE_AGGREGATE_RECORDS; const UNKNOWN_U32: u32 = u32::MAX; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -64,6 +66,13 @@ pub enum CuptiMemcpyKind { PeerToPeer = 5, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[repr(u32)] +pub enum CuptiCaptureMode { + Exact = 0, + Aggregate = 1, +} + #[derive(Debug, Clone, PartialEq, Eq)] pub enum CuptiNameFilter { Any, @@ -84,9 +93,30 @@ pub struct CuptiEventFilter { #[derive(Debug, Clone, PartialEq, Eq)] pub struct CuptiArmConfig { pub record_capacity: u64, + pub capture_mode: CuptiCaptureMode, pub filters: Vec, } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CuptiAggregateActivity { + Kernel, + Memcpy, + Memset, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct CuptiAggregateGroup { + pub activity: CuptiAggregateActivity, + pub name: Option, + pub device_id: Option, + pub memcpy_kind: Option, + pub count: u64, + pub total_duration_ns: u64, + pub min_duration_ns: u64, + pub max_duration_ns: u64, + pub total_bytes: Option, +} + #[derive(Debug, Clone, PartialEq)] pub struct CuptiCapture { pub record_offset: u64, @@ -99,6 +129,7 @@ pub struct CuptiCapture { pub dropped_records: u64, pub unknown_records: u64, pub events: Vec, + pub aggregate_groups: Vec, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -149,6 +180,10 @@ pub enum CuptiDecodeError { index: usize, kind: u32, }, + InvalidAggregateRecord { + index: usize, + message: &'static str, + }, InvalidCaptureState(u32), InvalidStopReason(u32), InvalidCounters { @@ -212,6 +247,12 @@ impl fmt::Display for CuptiDecodeError { "CUPTI record {index} of kind {kind} has a zero timestamp" ) } + Self::InvalidAggregateRecord { index, message } => { + write!( + formatter, + "CUPTI aggregate record {index} is invalid: {message}" + ) + } Self::InvalidCaptureState(state) => { write!(formatter, "CUPTI capture state {state} is invalid") } @@ -408,8 +449,9 @@ fn encode_control( )); } control[16..24].copy_from_slice(&config.record_capacity.to_ne_bytes()); + control[32..36].copy_from_slice(&(config.capture_mode as u32).to_ne_bytes()); for (index, filter) in config.filters.iter().enumerate() { - let offset = 32 + index * FILTER_SIZE; + let offset = 40 + index * FILTER_SIZE; control[offset..offset + 4].copy_from_slice(&(filter.record_kind as u32).to_ne_bytes()); control[offset + 4..offset + 8].copy_from_slice(&(filter.api_domain as u32).to_ne_bytes()); control[offset + 8..offset + 12] @@ -516,19 +558,22 @@ pub fn decode_capture(bytes: &[u8], session_id: &str) -> Result Result Result<(Vec, Vec), CuptiDecodeError> { + if feature_flags & FEATURE_AGGREGATE_RECORDS != 0 { + let mut groups = Vec::with_capacity(record_count); + for (index, record) in bytes[HEADER_SIZE..].chunks_exact(RECORD_SIZE).enumerate() { + groups.push(decode_aggregate_record(record, index)?); + } + return Ok((Vec::new(), groups)); + } + + let mut events = Vec::with_capacity(record_count); + for (index, record) in bytes[HEADER_SIZE..].chunks_exact(RECORD_SIZE).enumerate() { + events.push(decode_record(record, index, session_id, feature_flags)?); + } + events.sort_by_key(|event| event.timestamp_ns); + let mut sequence = record_offset; + for event in &mut events { + event.sequence = sequence; + event.event_id = format!("evt_{sequence}"); + sequence = sequence + .checked_add(1) + .ok_or(CuptiDecodeError::CaptureLengthOverflow)?; + } + Ok((events, Vec::new())) +} + +fn decode_aggregate_record( + record: &[u8], + index: usize, +) -> Result { + let count = read_u64(record, 0); + let total_duration_ns = read_u64(record, 8); + let min_duration_ns = read_u64(record, 16); + let max_duration_ns = read_u64(record, 24); + if count == 0 { + return Err(CuptiDecodeError::InvalidAggregateRecord { + index, + message: "count is zero", + }); + } + if min_duration_ns > max_duration_ns + || total_duration_ns / count < min_duration_ns + || total_duration_ns / count > max_duration_ns + { + return Err(CuptiDecodeError::InvalidAggregateRecord { + index, + message: "duration counters are inconsistent", + }); + } + let kind = read_u32(record, 40); + let (activity, has_name, has_bytes) = match kind { + 3 => (CuptiAggregateActivity::Kernel, true, false), + 5 => (CuptiAggregateActivity::Memcpy, false, true), + 7 => (CuptiAggregateActivity::Memset, false, true), + _ => { + return Err(CuptiDecodeError::InvalidAggregateRecord { + index, + message: "activity kind is unsupported", + }); + } + }; + let name_bytes = &record[72..200]; + let name_length = name_bytes + .iter() + .position(|byte| *byte == 0) + .unwrap_or(name_bytes.len()); + let name = str::from_utf8(&name_bytes[..name_length]) + .map_err(|_| CuptiDecodeError::InvalidName { index })?; + if has_name && name.is_empty() { + return Err(CuptiDecodeError::InvalidAggregateRecord { + index, + message: "kernel name is empty", + }); + } + let memcpy_kind = (activity == CuptiAggregateActivity::Memcpy) + .then(|| decode_semantic_memcpy_kind(read_u32(record, 48))); + Ok(CuptiAggregateGroup { + activity, + name: has_name.then(|| name.to_owned()), + device_id: optional_unknown(read_u32(record, 44)), + memcpy_kind, + count, + total_duration_ns, + min_duration_ns, + max_duration_ns, + total_bytes: has_bytes.then(|| read_u64(record, 32)), }) } @@ -694,6 +834,17 @@ const fn decode_memcpy_kind(kind: u32) -> MemcpyKind { } } +const fn decode_semantic_memcpy_kind(kind: u32) -> MemcpyKind { + match kind { + 1 => MemcpyKind::HostToDevice, + 2 => MemcpyKind::DeviceToHost, + 3 => MemcpyKind::DeviceToDevice, + 4 => MemcpyKind::HostToHost, + 5 => MemcpyKind::PeerToPeer, + _ => MemcpyKind::Unknown, + } +} + fn decode_dim(record: &[u8], offset: usize) -> Option { let dimension = Dim3 { x: read_u32(record, offset), @@ -736,9 +887,9 @@ mod tests { use std::io::Cursor; use super::{ - CuptiApiDomain, CuptiArmConfig, CuptiDecodeError, CuptiEventFilter, CuptiMemcpyKind, - CuptiNameFilter, CuptiRecordKind, HEADER_SIZE, OUTPUT_MAGIC, RECORD_SIZE, decode_capture, - encode_control, read_snapshot, + CuptiApiDomain, CuptiArmConfig, CuptiCaptureMode, CuptiDecodeError, CuptiEventFilter, + CuptiMemcpyKind, CuptiNameFilter, CuptiRecordKind, HEADER_SIZE, OUTPUT_MAGIC, RECORD_SIZE, + decode_capture, encode_control, read_snapshot, }; use xprobe_protocol::{ClockDomain, EventSource, EventType, MemcpyKind}; @@ -769,10 +920,40 @@ mod tests { fn transfer_capture(records: &[[u8; RECORD_SIZE]]) -> Vec { let mut bytes = normalized_capture(records); - bytes[20..24].copy_from_slice(&super::SUPPORTED_FEATURES.to_le_bytes()); + let features = super::FEATURE_HOST_MONOTONIC_TIMESTAMPS | super::FEATURE_TRANSFER_RECORDS; + bytes[20..24].copy_from_slice(&features.to_le_bytes()); bytes } + fn aggregate_capture(records: &[[u8; RECORD_SIZE]], observed: u64) -> Vec { + let mut bytes = capture(records); + let features = super::FEATURE_TRANSFER_RECORDS | super::FEATURE_AGGREGATE_RECORDS; + bytes[20..24].copy_from_slice(&features.to_le_bytes()); + bytes[48..56].copy_from_slice(&observed.to_le_bytes()); + bytes + } + + fn aggregate_record( + kind: u32, + count: u64, + total_duration_ns: u64, + device_id: u32, + memcpy_kind: u32, + name: &str, + ) -> [u8; RECORD_SIZE] { + let mut record = [0_u8; RECORD_SIZE]; + record[0..8].copy_from_slice(&count.to_le_bytes()); + record[8..16].copy_from_slice(&total_duration_ns.to_le_bytes()); + record[16..24].copy_from_slice(&(total_duration_ns / count).to_le_bytes()); + record[24..32].copy_from_slice(&(total_duration_ns / count).to_le_bytes()); + record[32..40].copy_from_slice(&4096_u64.to_le_bytes()); + record[40..44].copy_from_slice(&kind.to_le_bytes()); + record[44..48].copy_from_slice(&device_id.to_le_bytes()); + record[48..52].copy_from_slice(&memcpy_kind.to_le_bytes()); + record[72..72 + name.len()].copy_from_slice(name.as_bytes()); + record + } + fn record(kind: u32, timestamp: u64, correlation: u32, name: &str) -> [u8; RECORD_SIZE] { let mut record = [0_u8; RECORD_SIZE]; record[0..8].copy_from_slice(×tamp.to_le_bytes()); @@ -930,10 +1111,10 @@ mod tests { #[test] fn rejects_unknown_feature_flags() { let mut bytes = capture(&[]); - bytes[20..24].copy_from_slice(&4_u32.to_le_bytes()); + bytes[20..24].copy_from_slice(&8_u32.to_le_bytes()); assert_eq!( decode_capture(&bytes, "xp_test"), - Err(CuptiDecodeError::UnsupportedFeatureFlags(4)) + Err(CuptiDecodeError::UnsupportedFeatureFlags(8)) ); } @@ -1000,6 +1181,36 @@ mod tests { assert_eq!(decoded.events[0].event_id, "evt_2"); } + #[test] + fn decodes_bounded_aggregate_groups() { + let bytes = aggregate_capture( + &[ + aggregate_record(3, 4, 400, 0, 0, "vector_add"), + aggregate_record(5, 2, 80, 0, 1, ""), + ], + 6, + ); + + let decoded = decode_capture(&bytes, "unused").expect("aggregate capture must decode"); + + assert!(decoded.events.is_empty()); + assert_eq!(decoded.aggregate_groups.len(), 2); + assert_eq!( + decoded.aggregate_groups[0].activity, + super::CuptiAggregateActivity::Kernel + ); + assert_eq!( + decoded.aggregate_groups[0].name.as_deref(), + Some("vector_add") + ); + assert_eq!(decoded.aggregate_groups[0].count, 4); + assert_eq!( + decoded.aggregate_groups[1].memcpy_kind, + Some(MemcpyKind::HostToDevice) + ); + assert_eq!(decoded.aggregate_groups[1].total_bytes, Some(4096)); + } + #[test] fn encodes_bounded_arm_filters() { let control = encode_control( @@ -1007,6 +1218,7 @@ mod tests { 0, Some(&CuptiArmConfig { record_capacity: 250_000, + capture_mode: CuptiCaptureMode::Exact, filters: vec![CuptiEventFilter { record_kind: CuptiRecordKind::GpuKernelStart, api_domain: CuptiApiDomain::Any, @@ -1017,7 +1229,7 @@ mod tests { ) .expect("ARM request must encode"); - assert_eq!(control.len(), 320); + assert_eq!(control.len(), 328); assert_eq!(&control[0..8], super::CONTROL_MAGIC); assert_eq!(u32::from_ne_bytes(control[12..16].try_into().unwrap()), 1); assert_eq!( @@ -1025,10 +1237,42 @@ mod tests { 250_000 ); assert_eq!(u64::from_ne_bytes(control[24..32].try_into().unwrap()), 0); - assert_eq!(u32::from_ne_bytes(control[32..36].try_into().unwrap()), 3); - assert_eq!(u32::from_ne_bytes(control[44..48].try_into().unwrap()), 2); - assert_eq!(&control[48..54], b"flash_"); - assert!(control[176..].iter().all(|byte| *byte == 0)); + assert_eq!(u32::from_ne_bytes(control[32..36].try_into().unwrap()), 0); + assert_eq!(u32::from_ne_bytes(control[40..44].try_into().unwrap()), 3); + assert_eq!(u32::from_ne_bytes(control[52..56].try_into().unwrap()), 2); + assert_eq!(&control[56..62], b"flash_"); + assert!(control[184..].iter().all(|byte| *byte == 0)); + } + + #[test] + fn encodes_aggregate_arm_mode() { + let control = encode_control( + 1, + 0, + Some(&CuptiArmConfig { + record_capacity: 4096, + capture_mode: CuptiCaptureMode::Aggregate, + filters: vec![ + CuptiEventFilter { + record_kind: CuptiRecordKind::GpuKernelStart, + api_domain: CuptiApiDomain::Any, + memcpy_kind: CuptiMemcpyKind::Any, + name: CuptiNameFilter::Any, + }, + CuptiEventFilter { + record_kind: CuptiRecordKind::GpuKernelEnd, + api_domain: CuptiApiDomain::Any, + memcpy_kind: CuptiMemcpyKind::Any, + name: CuptiNameFilter::Any, + }, + ], + }), + ) + .expect("aggregate ARM request must encode"); + + assert_eq!(u32::from_ne_bytes(control[32..36].try_into().unwrap()), 1); + assert_eq!(u32::from_ne_bytes(control[40..44].try_into().unwrap()), 3); + assert_eq!(u32::from_ne_bytes(control[184..188].try_into().unwrap()), 4); } #[test] @@ -1038,6 +1282,7 @@ mod tests { 0, Some(&CuptiArmConfig { record_capacity: 1, + capture_mode: CuptiCaptureMode::Exact, filters: Vec::new(), }), ) diff --git a/xprobe/protocol/src/lib.rs b/xprobe/protocol/src/lib.rs index df36c31..deb536a 100644 --- a/xprobe/protocol/src/lib.rs +++ b/xprobe/protocol/src/lib.rs @@ -25,9 +25,11 @@ pub use event::{ }; pub use export::{ExportFormat, TraceExportResult}; pub use measurement::{ - CaptureCompleteness, ClockQuality, CollectionSummary, CorrelationConfidence, - CorrelationSummary, CuptiCollectionSummary, LatencyStatistics, MatchPolicy, MatchedEventPair, - Measurement, MeasurementResult, MeasurementSpec, SampleSummary, SessionStatus, TargetIdentity, + AggregateActivity, AggregateCollectionSummary, AggregateDuration, AggregateGroup, + AggregateInventory, AggregateInventoryResult, CaptureCompleteness, ClockQuality, + CollectionSummary, CorrelationConfidence, CorrelationSummary, CuptiCollectionSummary, + LatencyStatistics, MatchPolicy, MatchedEventPair, Measurement, MeasurementMode, + MeasurementResult, MeasurementSpec, SampleSummary, SessionStatus, TargetIdentity, }; pub use process::{CgroupEntry, ProcessCredentials, ProcessCudaState, ProcessReport}; pub use resolve::{ElfObjectKind, ProcessMapping, ResolvedProbe}; diff --git a/xprobe/protocol/src/measurement.rs b/xprobe/protocol/src/measurement.rs index b76a966..3eaf4ea 100644 --- a/xprobe/protocol/src/measurement.rs +++ b/xprobe/protocol/src/measurement.rs @@ -20,6 +20,14 @@ pub enum MatchPolicy { StreamOrder, } +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum MeasurementMode { + #[default] + Exact, + Aggregate, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct MeasurementSpec { @@ -32,7 +40,12 @@ pub struct MeasurementSpec { pub samples: Option, pub duration_ms: Option, pub timeout_ms: u64, - pub max_events: u64, + #[serde(default)] + pub max_events: Option, + #[serde(default)] + pub measurement_mode: MeasurementMode, + #[serde(default)] + pub max_groups: Option, } #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -147,3 +160,69 @@ pub struct CuptiCollectionSummary { pub dropped_records: u64, pub buffer_utilization: f64, } + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AggregateInventoryResult { + pub schema_version: SchemaVersion, + pub ok: bool, + pub session_id: String, + pub status: SessionStatus, + pub inventory: AggregateInventory, + pub collection: AggregateCollectionSummary, + #[serde(default)] + pub warnings: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AggregateInventory { + pub name: Option, + pub start_selector: String, + pub end_selector: String, + pub groups: Vec, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AggregateGroup { + pub activity: AggregateActivity, + pub name: Option, + pub device_id: Option, + pub memcpy_kind: Option, + pub start_selector_hint: String, + pub end_selector_hint: String, + pub count: u64, + pub duration_ns: AggregateDuration, + pub total_bytes: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(rename_all = "snake_case")] +pub enum AggregateActivity { + Kernel, + Memcpy, + Memset, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AggregateDuration { + pub min: u64, + pub mean: f64, + pub max: u64, + pub total: u64, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct AggregateCollectionSummary { + pub completeness: CaptureCompleteness, + pub observed_activities: u64, + pub grouped_activities: u64, + pub dropped_activities: u64, + pub group_capacity: u64, + pub groups: u64, + pub occupied_slots: u64, + pub table_utilization: f64, +} diff --git a/xprobe/protocol/src/schema.rs b/xprobe/protocol/src/schema.rs index a2be216..a14969b 100644 --- a/xprobe/protocol/src/schema.rs +++ b/xprobe/protocol/src/schema.rs @@ -1,12 +1,13 @@ use schemars::{Schema, schema_for}; use crate::{ - CapabilityReport, DiscoveryResult, ErrorResponse, Event, HostCaptureResult, MeasurementResult, - MeasurementSpec, ProcessReport, ResolvedProbe, TraceExportResult, ValidationResult, + AggregateInventoryResult, CapabilityReport, DiscoveryResult, ErrorResponse, Event, + HostCaptureResult, MeasurementResult, MeasurementSpec, ProcessReport, ResolvedProbe, + TraceExportResult, ValidationResult, }; #[must_use] -pub fn generated_schemas() -> [(&'static str, Schema); 11] { +pub fn generated_schemas() -> [(&'static str, Schema); 12] { [ ("event.schema.json", schema_for!(Event)), ("error.schema.json", schema_for!(ErrorResponse)), @@ -15,6 +16,10 @@ pub fn generated_schemas() -> [(&'static str, Schema); 11] { "measurement-result.schema.json", schema_for!(MeasurementResult), ), + ( + "aggregate-inventory-result.schema.json", + schema_for!(AggregateInventoryResult), + ), ("capability.schema.json", schema_for!(CapabilityReport)), ("discover.schema.json", schema_for!(DiscoveryResult)), ("inspect.schema.json", schema_for!(ProcessReport)), diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index e698664..025d0f1 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -3,9 +3,9 @@ use std::{fs, path::PathBuf}; use serde::{Serialize, de::DeserializeOwned}; use serde_json::{Value, json}; use xprobe_protocol::{ - CapabilityReport, DiscoveryResult, ErrorResponse, Event, HostCaptureResult, MeasurementResult, - MeasurementSpec, ProcessReport, ResolvedProbe, TraceExportResult, ValidationResult, - schema::generated_schemas, + AggregateInventoryResult, CapabilityReport, DiscoveryResult, ErrorResponse, Event, + HostCaptureResult, MeasurementResult, MeasurementSpec, ProcessReport, ResolvedProbe, + TraceExportResult, ValidationResult, schema::generated_schemas, }; fn assert_round_trip(fixture: &Value) @@ -172,7 +172,51 @@ fn measurement_spec_contract_round_trips() { "samples": 100, "duration_ms": null, "timeout_ms": 30_000, - "max_events": 100_000 + "max_events": 100_000, + "measurement_mode": "exact", + "max_groups": null + })); +} + +#[test] +fn aggregate_inventory_contract_round_trips() { + assert_round_trip::(&json!({ + "schema_version": "2.0", + "ok": true, + "session_id": "xp_inventory", + "status": "completed", + "inventory": { + "name": "kernel_inventory", + "start_selector": "cuda:kernel_start", + "end_selector": "cuda:kernel_end", + "groups": [{ + "activity": "kernel", + "name": "vector_add", + "device_id": 0, + "memcpy_kind": null, + "start_selector_hint": "cuda:kernel_start:name~^vector_add$", + "end_selector_hint": "cuda:kernel_end:name~^vector_add$", + "count": 10, + "duration_ns": { + "min": 1000, + "mean": 1500.0, + "max": 2000, + "total": 15000 + }, + "total_bytes": null + }] + }, + "collection": { + "completeness": "complete", + "observed_activities": 10, + "grouped_activities": 10, + "dropped_activities": 0, + "group_capacity": 4096, + "groups": 1, + "occupied_slots": 1, + "table_utilization": 0.000_244_140_625 + }, + "warnings": [] })); } From f53dead5cbcdf3e2d0d32c041365f66b2cedb13a Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 04:49:04 +0800 Subject: [PATCH 06/17] =?UTF-8?q?=F0=9F=93=9D=20docs:=20guide=20aggregate-?= =?UTF-8?q?first=20profiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- README.md | 9 +++- docs/architecture.md | 6 +++ docs/cli-contract.md | 17 ++++++++ docs/cupti-agent.md | 30 ++++++++++---- skills/xprobe-measure-latency/SKILL.md | 18 ++++---- .../examples/coarse-kernel-inventory.json | 4 +- .../examples/coarse-memcpy-inventory.json | 4 +- .../references/cli-contract.md | 9 ++++ .../references/investigation.md | 41 ++++++++++--------- .../references/result-quality.md | 15 ++++--- tests/agent-contract/test_contract.py | 13 +++++- 11 files changed, 119 insertions(+), 47 deletions(-) diff --git a/README.md b/README.md index 770dcbf..e8e4939 100644 --- a/README.md +++ b/README.md @@ -48,6 +48,11 @@ xprobe validate --pid 4242 \ --to 'cuda:kernel_start:name~flash.*' \ --match exact --json --non-interactive --no-color +xprobe measure --pid 4242 \ + --from 'cuda:kernel_start' --to 'cuda:kernel_end' \ + --match exact --aggregate --duration-ms 1000 --max-groups 4096 \ + --json --non-interactive --no-color + xprobe measure --pid 4242 \ --from 'cuda:runtime_api:cudaLaunchKernel:exit' \ --to 'cuda:kernel_start:name~flash.*' \ @@ -58,7 +63,9 @@ xprobe measure --pid 4242 \ Kernel launch latency is only one event pair. The same workflow measures host function spans, CUDA API calls, GPU operation durations, transfers, and paths -across CPU and GPU events after selecting the correct CUDA worker. +across CPU and GPU events after selecting the correct CUDA worker. Aggregate +mode provides a bounded coarse inventory of GPU operations before an exact +evidence measurement narrows the question. `measure` also accepts completed `--input` captures and versioned live `--spec` files. Evidence can be exported as `jsonl` or `chrome`. JSON results diff --git a/docs/architecture.md b/docs/architecture.md index 39c2dd8..9388b95 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -83,6 +83,12 @@ Driver callbacks plus kernel, memcpy, and memset activity are filtered before fixed records consume capacity. Callback hot paths do not perform blocking I/O or allocation. +Broad GPU inventory uses the same `measure` primitive with aggregate mode. The +Agent updates a `--max-groups`-bounded table for matching kernel, memcpy, or +memset activity and returns only final count/duration/byte summaries. Exact +measurement remains the evidence path; aggregate output has a separate result +contract and cannot be exported as event JSONL. + CUPTI activity timestamps are normalized to `CLOCK_MONOTONIC` through its timestamp callback or an explicit CUDA 12 clock calibration. Activity that began before the ARM epoch is excluded. The Agent verifies the retained window diff --git a/docs/cli-contract.md b/docs/cli-contract.md index bbb0f60..3bad6cc 100644 --- a/docs/cli-contract.md +++ b/docs/cli-contract.md @@ -133,9 +133,26 @@ xprobe measure --spec measurement.json \ --json --non-interactive --no-color ``` +Broad GPU activity inventory: + +```bash +xprobe measure --pid 4242 \ + --from 'cuda:kernel_start' --to 'cuda:kernel_end' \ + --match exact --aggregate --duration-ms 1000 --max-groups 4096 \ + --json --non-interactive --no-color +``` + Exactly one source mode is used: `--pid`, one or more `--input`, or `--spec`. At least one positive `--samples` or `--duration-ms` bound is required in direct mode. `--timeout-ms` defaults to 30 seconds and `--max-events` to 100,000. +`--aggregate` is live-only, duration-bounded, and accepts one matching kernel, +memcpy, or memset activity start/end pair. It uses `--max-groups` (default +4,096), does not accept `--samples` or `--events-out`, and emits +`schemas/aggregate-inventory-result.schema.json`. The result is a coarse +inventory, not exact event evidence: it reports count, total/min/max/mean +duration, optional transferred bytes, table occupancy, and drop completeness. +Aggregate kernel regex must be reducible to an exact, prefix, suffix, or +contains filter because this mode intentionally has no Rust-side event pass. Live host endpoints attach PID-scoped eBPF probes. CUDA endpoints automatically activate the CUPTI agent. If it is absent, `--agent` or diff --git a/docs/cupti-agent.md b/docs/cupti-agent.md index b0c3d3d..9d1e98c 100644 --- a/docs/cupti-agent.md +++ b/docs/cupti-agent.md @@ -25,6 +25,13 @@ directions, and simple kernel-name exact/prefix/suffix/contains patterns are filtered before they consume capture capacity. Complex kernel regular expressions use a wider Agent filter and retain exact Rust-side matching. +An aggregate ARM uses a fixed-capacity hash table instead of event records. +Kernel activity is grouped by name and device, memcpy by direction and device, +and memset by device. The activity callback updates count, total/min/max +duration, and transfer bytes without materializing start/end events. Aggregate +collection uses one final CLI read instead of incremental snapshots, and table +saturation fails instead of returning a sampled inventory. + `SNAPSHOT` flushes pending activity and returns records after the caller's checked record offset. The CLI accumulates only contiguous deltas rather than retransferring and decoding the growing capture. `STOP` returns the final delta @@ -45,22 +52,27 @@ paths. ## Control and capture ABI -Control version 3 uses one fixed 320-byte native-endian request defined by +Control version 4 uses one fixed 328-byte native-endian request defined by `cupti/include/xprobe/cupti_agent.h`. It contains magic, version, command, -record capacity, record offset, and two fixed endpoint filters. Commands are -`ARM`, `SNAPSHOT`, `STOP`, and `CLOSE`. ARM requires offset zero; later commands -return records at or after the requested offset followed by EOF. +capacity, record offset, capture mode, and two fixed endpoint filters. Commands +are `ARM`, `SNAPSHOT`, `STOP`, and `CLOSE`. Exact ARM requires offset zero; +later exact commands return records at or after the requested offset followed +by EOF. Aggregate mode requires offset zero for every command. -Capture ABI v3 starts with an 88-byte header and zero or more 200-byte records. +Capture ABI v4 starts with an 88-byte header and zero or more 200-byte records. The header reports capture state and stop reason, configured capacity, observed and retained counts, Agent and CUPTI drops, unknown activity, record sizes, and feature flags. It also reports the payload record offset so the caller can -reject gaps, replays, and counter rollback. Capacity comes from `--max-events`; -there is no special 2^16 limit. Reaching the configured limit freezes capture -and causes measurement to fail explicitly instead of returning partial success. +reject gaps, replays, and counter rollback. Exact capacity comes from +`--max-events` and aggregate table capacity from `--max-groups`; there is no +special 2^16 limit. Reaching either configured limit freezes capture and causes +measurement to fail explicitly instead of returning partial success. -Records preserve process/thread, device/context/stream, correlation IDs, +Exact records preserve process/thread, device/context/stream, correlation IDs, callback domain/ID, dimensions or transfer metadata, and one bounded name. +Aggregate records contain a bounded group key and exact integer counters, not +event evidence or a latency distribution; they therefore do not report +percentiles or correlation confidence. Activity timestamps are normalized to `CLOCK_MONOTONIC` through the CUPTI timestamp callback or CUDA 12 clock calibration. If alignment cannot be established, GPU durations remain usable in the CUPTI domain but host/GPU diff --git a/skills/xprobe-measure-latency/SKILL.md b/skills/xprobe-measure-latency/SKILL.md index 62f85df..338d221 100644 --- a/skills/xprobe-measure-latency/SKILL.md +++ b/skills/xprobe-measure-latency/SKILL.md @@ -35,14 +35,16 @@ CLI and selector syntax is in [references/cli-contract.md](references/cli-contra the selected process PID. 5. Map GPU or mixed work before choosing a name. Validate broad kernel, memcpy, or memset activity endpoints, then collect one bounded, representative coarse - inventory per event family with `--events-out`. For CPU-only work, resolve and - validate the intended host-function boundary directly; do not require CUDA or - CUPTI. Scope breadth and collection duration are independent: keep the - selector broad where an activity inventory exists, but choose a duration that - covers the workload cycle being diagnosed. Give `--max-events` headroom. -6. For GPU artifacts, run `scripts/analyze_trace.py`. Use kernel names, selector - hints, duration aggregates, launch variants, stream distribution, busy union, - overlap factor, and adjacent gaps to form one narrow hypothesis. Read + inventory per event family with `measure --aggregate --duration-ms ...`. + For CPU-only work, resolve and validate the intended host-function boundary + directly; do not require CUDA or CUPTI. Scope breadth and collection duration + are independent: keep the selector broad where an activity inventory exists, + choose a duration that covers the workload cycle being diagnosed, and give + `--max-groups` headroom. +6. Use aggregate names, selector hints, counts, duration totals and bounds, and + transfer bytes to form one narrow hypothesis. For an exact GPU artifact, run + `scripts/analyze_trace.py` and use launch variants, stream distribution, busy + union, overlap factor, and adjacent gaps. Read [references/trace-analysis.md](references/trace-analysis.md) when interpreting the report. For CPU-only work, use resolved host selectors and result evidence to form the hypothesis instead. diff --git a/skills/xprobe-measure-latency/examples/coarse-kernel-inventory.json b/skills/xprobe-measure-latency/examples/coarse-kernel-inventory.json index e64b0a6..8cebf11 100644 --- a/skills/xprobe-measure-latency/examples/coarse-kernel-inventory.json +++ b/skills/xprobe-measure-latency/examples/coarse-kernel-inventory.json @@ -11,5 +11,7 @@ "samples": null, "duration_ms": 1000, "timeout_ms": 30000, - "max_events": 200000 + "max_events": null, + "measurement_mode": "aggregate", + "max_groups": 4096 } diff --git a/skills/xprobe-measure-latency/examples/coarse-memcpy-inventory.json b/skills/xprobe-measure-latency/examples/coarse-memcpy-inventory.json index f8ae56e..63afa08 100644 --- a/skills/xprobe-measure-latency/examples/coarse-memcpy-inventory.json +++ b/skills/xprobe-measure-latency/examples/coarse-memcpy-inventory.json @@ -11,5 +11,7 @@ "samples": null, "duration_ms": 1000, "timeout_ms": 30000, - "max_events": 200000 + "max_events": null, + "measurement_mode": "aggregate", + "max_groups": 64 } diff --git a/skills/xprobe-measure-latency/references/cli-contract.md b/skills/xprobe-measure-latency/references/cli-contract.md index c73f80e..da2c439 100644 --- a/skills/xprobe-measure-latency/references/cli-contract.md +++ b/skills/xprobe-measure-latency/references/cli-contract.md @@ -40,6 +40,15 @@ Direct `measure` calls require a positive `--samples` or `--duration-ms` bound. `--timeout-ms` defaults to 30 seconds and `--max-events` to 100,000. Use exactly one source mode: `--pid`, one or more `--input` files, or `--spec`. +Use `measure --aggregate --duration-ms ... --max-groups ...` for a live, coarse +kernel, memcpy, or memset inventory. Aggregate endpoints must be one matching +activity start/end pair with `--match exact`. The result contains bounded group +counts, total/min/max/mean duration, transfer bytes, selector hints, and table +quality. It contains no event evidence, percentiles, or correlation confidence; +`--samples`, `--input`, and `--events-out` are invalid in this mode. +Kernel regex must be an exact, prefix, suffix, or contains shape that the Agent +can apply before aggregation; other regex is rejected instead of widened. + Kernel and other GPU activity durations require separate start and end records, so `max-events` is record capacity rather than sample capacity. Sample completion is checked after bounded snapshots and does not reserve space in the CUPTI diff --git a/skills/xprobe-measure-latency/references/investigation.md b/skills/xprobe-measure-latency/references/investigation.md index 8e7dfb1..2e7b645 100644 --- a/skills/xprobe-measure-latency/references/investigation.md +++ b/skills/xprobe-measure-latency/references/investigation.md @@ -28,12 +28,9 @@ xprobe validate --pid "$PID" \ xprobe measure --pid "$PID" \ --from cuda:kernel_start --to cuda:kernel_end --match exact \ - --duration-ms "$REPRESENTATIVE_WINDOW_MS" --timeout-ms "$TIMEOUT_MS" \ - --max-events "$MAX_EVENTS" --events-out coarse-kernels.jsonl --format jsonl \ - --json --non-interactive --no-color - -skills/xprobe-measure-latency/scripts/analyze_trace.py coarse-kernels.jsonl \ - > coarse-kernels-analysis.json + --aggregate --duration-ms "$REPRESENTATIVE_WINDOW_MS" \ + --timeout-ms "$TIMEOUT_MS" --max-groups "$MAX_GROUPS" \ + --json --non-interactive --no-color > coarse-kernels.json ``` Map copies and memsets in separate bounded captures when they could matter: @@ -45,22 +42,21 @@ xprobe validate --pid "$PID" \ xprobe measure --pid "$PID" \ --from cuda:memcpy_start --to cuda:memcpy_end --match exact \ - --duration-ms "$REPRESENTATIVE_WINDOW_MS" --timeout-ms "$TIMEOUT_MS" \ - --max-events "$MAX_EVENTS" --events-out coarse-memcpy.jsonl --format jsonl \ - --json --non-interactive --no-color + --aggregate --duration-ms "$REPRESENTATIVE_WINDOW_MS" \ + --timeout-ms "$TIMEOUT_MS" --max-groups "$MAX_GROUPS" \ + --json --non-interactive --no-color > coarse-memcpy.json ``` Do not guess a named kernel or host function in this stage. CUDA API selectors require a concrete Runtime or Driver API name, so inventory activity first and choose API boundaries from application, framework, or trace evidence later. -Analyze each artifact separately; busy union only describes the capture window -inside that artifact. +Read each inventory separately. Use its selector hints to start the next exact +measurement; device-specific groups may still need workload-level GPU routing. -Treat capacity as a consequence of observed event rate, not a universal default. -On `EVENT_RATE_TOO_HIGH`, inspect the written artifact and error counters. First -split event families or reduce selector scope while retaining a representative -cycle; reduce duration only when the remaining window is still representative. -Preserve the artifact as evidence for the change. +Treat group capacity as a consequence of workload diversity, not event rate. +On `EVENT_RATE_TOO_HIGH`, split event families or reduce selector scope while +retaining a representative cycle; reduce duration only when the remaining +window is still representative. Aggregate mode never returns partial output. ## Derive CUDA selectors @@ -103,16 +99,21 @@ Choose one next boundary from evidence: - host function entry to return with `stack-nested` for CPU span; - host marker to GPU activity with `first-after` only as a disclosed heuristic. -After capture, analyze the artifact and all result quality fields. Re-correlate a -completed artifact with `measure --input` when only selectors or policy change; -do not attach again merely to recompute pairing. +After capture, analyze the artifact and all result quality fields. An aggregate +inventory cannot be re-correlated because it intentionally contains no events. +Collect one exact artifact for the selected hypothesis, then use +`measure --input` when only selectors or policy change. ```bash -xprobe measure --input coarse-kernels.jsonl \ +xprobe measure --pid "$PID" \ --from 'cuda:kernel_start:name~^selected_kernel$' \ --to 'cuda:kernel_end:name~^selected_kernel$' \ --match exact --samples 100 --max-events 200000 \ + --events-out selected-kernel.jsonl --format jsonl \ --json --non-interactive --no-color + +skills/xprobe-measure-latency/scripts/analyze_trace.py selected-kernel.jsonl \ + > selected-kernel-analysis.json ``` ## Escalate at the right boundary diff --git a/skills/xprobe-measure-latency/references/result-quality.md b/skills/xprobe-measure-latency/references/result-quality.md index f7e1391..bb88db6 100644 --- a/skills/xprobe-measure-latency/references/result-quality.md +++ b/skills/xprobe-measure-latency/references/result-quality.md @@ -13,6 +13,12 @@ changed, collection was incomplete, or clock alignment failed. Report unmatched and ambiguous counts with matched samples. `estimated_error_ns: null` means no quantified interpolation error bound. +Aggregate inventory is a different contract. Require `completeness: complete`, +zero dropped activities, equal observed and grouped activity counts, and +reasonable table utilization before using its group names or selector hints. +Its min/max/mean come from exact integer totals, but it has no event ordering, +percentiles, stream overlap, or correlation evidence. + ## Concurrency Group GPU evidence by device, context, and stream before interpreting order. @@ -41,11 +47,10 @@ minimum_records = samples * (start_records_per_sample + end_records_per_sample) max_events >= minimum_records + expected_unmatched_records ``` -Use at least 2x headroom for stable narrow selectors and 4-10x for high-rate or -broad inventories. For a duration inventory, size from a pilot artifact's -observed records per second. Keep enough duration to cover a representative -cycle; increasing max-events without narrowing a noisy selector only increases -profiler work. +Use at least 2x headroom for stable narrow selectors. For broad activity +inventory, use aggregate mode and size `max-groups` from expected operation +diversity rather than event rate. Keep enough duration to cover a representative +cycle; table saturation is an explicit failure. `duration-ms` limits correlation to a window beginning at the first selected event. In live mode it also sets a collection stop from ARM completion, so finish diff --git a/tests/agent-contract/test_contract.py b/tests/agent-contract/test_contract.py index 04b6691..45b05ec 100755 --- a/tests/agent-contract/test_contract.py +++ b/tests/agent-contract/test_contract.py @@ -101,11 +101,19 @@ def check_skill(workspace: pathlib.Path) -> None: examples = sorted((skill_root / "examples").glob("*.json")) assert len(examples) >= 5 policies = set() + modes = set() for example in examples: specification = json.loads(example.read_text()) assert specification["schema_version"] == "2.0", example - assert specification["max_events"] > 0, example + mode = specification.get("measurement_mode", "exact") + modes.add(mode) + if mode == "aggregate": + assert specification["max_events"] is None, example + assert specification["max_groups"] > 0, example + else: + assert specification["max_events"] > 0, example policies.add(specification["match_policy"]) + assert modes == {"exact", "aggregate"} assert {"exact", "first_after", "stack_nested", "stream_order"} <= policies investigation = (skill_root / "references/investigation.md").read_text() @@ -131,7 +139,8 @@ def check_skill(workspace: pathlib.Path) -> None: "first selected event", "ARM completion", "Summed kernel", - "profiler", + "Aggregate inventory", + "max-groups", ): assert required in normalized_quality for required in ( From f80f315c68ef0211f6c0264333aa63da18a03a50 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 05:12:13 +0800 Subject: [PATCH 07/17] =?UTF-8?q?=E2=9A=A1=20perf:=20bound=20aggregate=20b?= =?UTF-8?q?road=20capture?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/hardware.yml | 1 + AGENTS.md | 2 + .../cuda_aggregate_benchmark.cu | 63 ++++++++ benchmarks/cuda-aggregate/resource_runner.c | 150 ++++++++++++++++++ benchmarks/cuda-aggregate/run-container.sh | 83 ++++++++++ benchmarks/cuda-aggregate/run.py | 111 +++++++++++++ cupti/src/cupti_agent.c | 50 ++++-- docs/benchmarks.md | 19 +++ docs/cupti-agent.md | 4 +- docs/development.md | 7 + justfile | 3 + 11 files changed, 477 insertions(+), 16 deletions(-) create mode 100644 benchmarks/cuda-aggregate/cuda_aggregate_benchmark.cu create mode 100644 benchmarks/cuda-aggregate/resource_runner.c create mode 100755 benchmarks/cuda-aggregate/run-container.sh create mode 100755 benchmarks/cuda-aggregate/run.py diff --git a/.github/workflows/hardware.yml b/.github/workflows/hardware.yml index 3ef7815..a253bd0 100644 --- a/.github/workflows/hardware.yml +++ b/.github/workflows/hardware.yml @@ -24,3 +24,4 @@ jobs: - run: just test-injection-live - run: just test-multisource-live - run: just benchmark-gpu + - run: just benchmark-aggregate diff --git a/AGENTS.md b/AGENTS.md index f6bab9c..92a8a0b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -71,4 +71,6 @@ - Test bundled Skill scripts with deterministic fixtures and include them in `just test-agent-contract`. - Run `just benchmark-gpu` for callback hot-path changes. +- Run `just benchmark-aggregate` for aggregate inventory hot-path or capacity + changes. - Use emoji conventional commits, for example `🐛 fix: restore target registers`. diff --git a/benchmarks/cuda-aggregate/cuda_aggregate_benchmark.cu b/benchmarks/cuda-aggregate/cuda_aggregate_benchmark.cu new file mode 100644 index 0000000..54e771c --- /dev/null +++ b/benchmarks/cuda-aggregate/cuda_aggregate_benchmark.cu @@ -0,0 +1,63 @@ +#include + +#include +#include + +__global__ void xprobe_aggregate_primary(int *output) +{ + *output += 1; +} + +__global__ void xprobe_aggregate_secondary(int *output) +{ + *output += 2; +} + +static int report_cuda_error(const char *operation, cudaError_t result) +{ + std::fprintf(stderr, "%s failed: %s\n", operation, cudaGetErrorString(result)); + return 13; +} + +int main(int argc, char **argv) +{ + if (argc != 3) { + std::fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + cudaError_t result = cudaSetDevice(0); + if (result != cudaSuccess) { + return report_cuda_error("cudaSetDevice", result); + } + int *device_output = nullptr; + result = cudaMalloc(&device_output, sizeof(*device_output)); + if (result != cudaSuccess) { + return report_cuda_error("cudaMalloc", result); + } + + FILE *ready = std::fopen(argv[1], "w"); + if (ready == nullptr || std::fclose(ready) != 0) { + std::perror("ready file"); + return 5; + } + while (access(argv[2], F_OK) != 0) { + xprobe_aggregate_primary<<<1, 1>>>(device_output); + xprobe_aggregate_secondary<<<1, 1>>>(device_output); + result = cudaGetLastError(); + if (result == cudaSuccess) { + result = cudaDeviceSynchronize(); + } + if (result != cudaSuccess) { + return report_cuda_error("kernel workload", result); + } + } + + result = cudaFree(device_output); + if (result == cudaSuccess) { + result = cudaDeviceReset(); + } + if (result != cudaSuccess) { + return report_cuda_error("CUDA shutdown", result); + } + return 0; +} diff --git a/benchmarks/cuda-aggregate/resource_runner.c b/benchmarks/cuda-aggregate/resource_runner.c new file mode 100644 index 0000000..6907e3e --- /dev/null +++ b/benchmarks/cuda-aggregate/resource_runner.c @@ -0,0 +1,150 @@ +#define _GNU_SOURCE + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static uint64_t elapsed_ns(const struct timespec *start, const struct timespec *end) +{ + time_t seconds = end->tv_sec - start->tv_sec; + long nanoseconds = end->tv_nsec - start->tv_nsec; + + if (nanoseconds < 0) { + --seconds; + nanoseconds += 1000000000L; + } + return (uint64_t)seconds * 1000000000ULL + (uint64_t)nanoseconds; +} + +static uint64_t timeval_us(const struct timeval *value) +{ + return (uint64_t)value->tv_sec * 1000000ULL + (uint64_t)value->tv_usec; +} + +static uint64_t process_rss_kib(pid_t pid) +{ + char path[64]; + char line[256]; + FILE *status; + uint64_t rss = 0U; + unsigned long long parsed_rss; + + if (snprintf(path, sizeof(path), "/proc/%ld/status", (long)pid) < 0) { + return 0U; + } + status = fopen(path, "r"); + if (status == NULL) { + return 0U; + } + while (fgets(line, sizeof(line), status) != NULL) { + if (sscanf(line, "VmRSS: %llu kB", &parsed_rss) == 1) { + rss = (uint64_t)parsed_rss; + break; + } + } + if (ferror(status) != 0 || fclose(status) != 0) { + return 0U; + } + return rss; +} + +int main(int argc, char **argv) +{ + struct timespec started; + struct timespec finished; + struct timespec interval = {.tv_sec = 0, .tv_nsec = 10000000L}; + struct rusage usage; + pid_t target_pid; + pid_t child; + int status = 0; + uint64_t target_start_rss; + uint64_t target_peak_rss; + FILE *output; + + if (argc < 4) { + fprintf(stderr, + "usage: %s [args...]\n", + argv[0]); + return 2; + } + target_pid = (pid_t)strtol(argv[2], NULL, 10); + target_start_rss = process_rss_kib(target_pid); + if (target_pid <= 0 || target_start_rss == 0U) { + fprintf(stderr, "failed to read target process RSS\n"); + return 3; + } + target_peak_rss = target_start_rss; + if (clock_gettime(CLOCK_MONOTONIC, &started) != 0) { + perror("clock_gettime"); + return 4; + } + child = fork(); + if (child < 0) { + perror("fork"); + return 5; + } + if (child == 0) { + execvp(argv[3], &argv[3]); + perror("execvp"); + _exit(127); + } + + for (;;) { + pid_t result = wait4(child, &status, WNOHANG, &usage); + uint64_t target_rss = process_rss_kib(target_pid); + + if (target_rss > target_peak_rss) { + target_peak_rss = target_rss; + } + if (result == child) { + break; + } + if (result < 0) { + perror("wait4"); + return 6; + } + if (nanosleep(&interval, NULL) != 0 && errno != EINTR) { + perror("nanosleep"); + return 7; + } + } + if (clock_gettime(CLOCK_MONOTONIC, &finished) != 0) { + perror("clock_gettime"); + return 8; + } + output = fopen(argv[1], "w"); + if (output == NULL) { + perror("metrics output"); + return 9; + } + fprintf(output, "wall_ns %llu\n", + (unsigned long long)elapsed_ns(&started, &finished)); + fprintf(output, "user_us %llu\n", + (unsigned long long)timeval_us(&usage.ru_utime)); + fprintf(output, "system_us %llu\n", + (unsigned long long)timeval_us(&usage.ru_stime)); + fprintf(output, "max_rss_kib %ld\n", usage.ru_maxrss); + fprintf(output, "target_start_rss_kib %llu\n", + (unsigned long long)target_start_rss); + fprintf(output, "target_peak_rss_kib %llu\n", + (unsigned long long)target_peak_rss); + if (fclose(output) != 0) { + perror("metrics output"); + return 10; + } + if (WIFEXITED(status)) { + return WEXITSTATUS(status); + } + if (WIFSIGNALED(status)) { + raise(WTERMSIG(status)); + } + return 11; +} diff --git a/benchmarks/cuda-aggregate/run-container.sh b/benchmarks/cuda-aggregate/run-container.sh new file mode 100755 index 0000000..534bbbb --- /dev/null +++ b/benchmarks/cuda-aggregate/run-container.sh @@ -0,0 +1,83 @@ +#!/usr/bin/env bash +set -euo pipefail + +if [[ $# -ne 1 ]]; then + echo "usage: run-container.sh " >&2 + exit 2 +fi + +output_dir=$1 +build_dir=/tmp/xprobe-cuda-aggregate +cuda_root=/usr/local/cuda +agent="${build_dir}/libxprobe-cupti.so" +fixture="${build_dir}/xprobe-cuda-aggregate" +resource_runner="${build_dir}/resource-runner" +socket="${build_dir}/agent.sock" +ready="${build_dir}/ready" +stop="${build_dir}/stop" +xprobe_bin=/workspace/target/debug/xprobe +compute_capability=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed -n '1p') +compute_arch=${compute_capability//./} + +mkdir -p "${build_dir}" "${output_dir}" +rm -f "${socket}" "${ready}" "${stop}" + +gcc \ + -std=c11 -D_GNU_SOURCE -DXPROBE_HAS_CUPTI=1 \ + -fPIC -shared -pthread -O2 -Wall -Wextra -Wpedantic -Werror \ + -I/workspace/cupti/include -isystem "${cuda_root}/include" \ + /workspace/cupti/src/cupti_agent.c \ + -L"${cuda_root}/lib64" -Wl,-rpath,"${cuda_root}/lib64" -lcupti \ + -o "${agent}" + +nvcc \ + -std=c++17 -O2 \ + -gencode="arch=compute_${compute_arch},code=sm_${compute_arch}" \ + /workspace/benchmarks/cuda-aggregate/cuda_aggregate_benchmark.cu \ + -o "${fixture}" + +gcc \ + -std=c11 -O2 -Wall -Wextra -Wpedantic -Werror \ + /workspace/benchmarks/cuda-aggregate/resource_runner.c \ + -o "${resource_runner}" + +XPROBE_CUPTI_SOCKET="${socket}" CUDA_INJECTION64_PATH="${agent}" \ + "${fixture}" "${ready}" "${stop}" & +target_pid=$! +trap 'touch "${stop}"; kill "${target_pid}" 2>/dev/null || true; wait "${target_pid}" 2>/dev/null || true' EXIT + +for _ in $(seq 1 500); do + [[ -e "${ready}" && -S "${socket}" ]] && break + kill -0 "${target_pid}" 2>/dev/null || wait "${target_pid}" + sleep 0.01 +done +[[ -e "${ready}" && -S "${socket}" ]] || { + echo "aggregate benchmark readiness timed out" >&2 + exit 1 +} + +"${resource_runner}" "${output_dir}/exact-metrics.txt" "${target_pid}" \ + "${xprobe_bin}" measure \ + --pid "${target_pid}" --cupti-socket "${socket}" \ + --from cuda:kernel_start --to cuda:kernel_end --match exact \ + --duration-ms 1800 --max-events 500000 --timeout-ms 10000 \ + --events-out "${output_dir}/exact-events.jsonl" \ + --json --non-interactive --no-color \ + >"${output_dir}/exact-result.json" 2>"${output_dir}/exact.stderr" + +"${resource_runner}" "${output_dir}/aggregate-metrics.txt" "${target_pid}" \ + "${xprobe_bin}" measure \ + --pid "${target_pid}" --cupti-socket "${socket}" \ + --from cuda:kernel_start --to cuda:kernel_end --match exact \ + --aggregate --duration-ms 1800 --max-groups 4 --timeout-ms 10000 \ + --json --non-interactive --no-color \ + >"${output_dir}/aggregate-result.json" 2>"${output_dir}/aggregate.stderr" + +nvidia-smi \ + --query-gpu=name,driver_version,compute_cap \ + --format=csv,noheader >"${output_dir}/gpu.txt" + +touch "${stop}" +wait "${target_pid}" +trap - EXIT +chmod 0644 "${output_dir}"/* diff --git a/benchmarks/cuda-aggregate/run.py b/benchmarks/cuda-aggregate/run.py new file mode 100755 index 0000000..3f4eceb --- /dev/null +++ b/benchmarks/cuda-aggregate/run.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +import json +import pathlib +import subprocess +import sys +import tempfile + + +def read_metrics(path: pathlib.Path) -> dict[str, int]: + return { + key: int(value) + for key, value in (line.split() for line in path.read_text().splitlines()) + } + + +def main() -> None: + if len(sys.argv) != 2: + raise SystemExit("usage: run.py ") + + workspace = pathlib.Path(__file__).resolve().parents[2] + with tempfile.TemporaryDirectory(prefix="xprobe-cuda-aggregate-") as temporary: + output = pathlib.Path(temporary) + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--gpus", + "all", + "--volume", + f"{workspace}:/workspace:ro", + "--volume", + f"{output}:/output", + "--workdir", + "/workspace", + sys.argv[1], + "/workspace/benchmarks/cuda-aggregate/run-container.sh", + "/output", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + raise SystemExit(completed.returncode) + + exact = json.loads((output / "exact-result.json").read_text()) + aggregate = json.loads((output / "aggregate-result.json").read_text()) + exact_metrics = read_metrics(output / "exact-metrics.txt") + aggregate_metrics = read_metrics(output / "aggregate-metrics.txt") + exact_artifact_bytes = (output / "exact-events.jsonl").stat().st_size + aggregate_artifact_bytes = (output / "aggregate-result.json").stat().st_size + gpu = (output / "gpu.txt").read_text().strip() + + assert exact["ok"] is True + assert exact["status"] == "completed" + assert exact["collection"]["cuda_events"] > 10_000 + assert exact["collection"]["dropped_events"] == 0 + assert aggregate["ok"] is True + assert aggregate["status"] == "completed" + collection = aggregate["collection"] + assert collection["observed_activities"] > 10_000 + assert collection["observed_activities"] == collection["grouped_activities"] + assert collection["observed_activities"] > collection["group_capacity"] + assert collection["groups"] == 2 + assert collection["occupied_slots"] == 2 + assert collection["dropped_activities"] == 0 + assert exact_artifact_bytes > aggregate_artifact_bytes * 100 + exact_cpu_us = exact_metrics["user_us"] + exact_metrics["system_us"] + aggregate_cpu_us = aggregate_metrics["user_us"] + aggregate_metrics["system_us"] + exact_target_growth = ( + exact_metrics["target_peak_rss_kib"] - exact_metrics["target_start_rss_kib"] + ) + aggregate_target_growth = ( + aggregate_metrics["target_peak_rss_kib"] + - aggregate_metrics["target_start_rss_kib"] + ) + assert aggregate_cpu_us < exact_cpu_us + assert aggregate_metrics["max_rss_kib"] < exact_metrics["max_rss_kib"] + assert aggregate_target_growth < exact_target_growth + + print( + json.dumps( + { + "schema_version": "2.0", + "ok": True, + "gpu": gpu, + "resources": { + "exact": exact_metrics, + "aggregate": aggregate_metrics, + "exact_cpu_us": exact_cpu_us, + "aggregate_cpu_us": aggregate_cpu_us, + "exact_target_growth_kib": exact_target_growth, + "aggregate_target_growth_kib": aggregate_target_growth, + }, + "artifacts": { + "exact_bytes": exact_artifact_bytes, + "aggregate_bytes": aggregate_artifact_bytes, + "reduction": exact_artifact_bytes / aggregate_artifact_bytes, + }, + "collection": collection, + }, + sort_keys=True, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index b1de857..ba9bad3 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -559,19 +559,14 @@ static int aggregate_key_matches( static int atomic_add_checked(_Atomic uint64_t *value, uint64_t addend) { - uint64_t current = atomic_load_explicit(value, memory_order_relaxed); + uint64_t previous = + atomic_fetch_add_explicit(value, addend, memory_order_relaxed); - for (;;) { - if (UINT64_MAX - current < addend) { - remember_output_error(); - return 0; - } - if (atomic_compare_exchange_weak_explicit( - value, ¤t, current + addend, memory_order_relaxed, - memory_order_relaxed)) { - return 1; - } + if (UINT64_MAX - previous < addend) { + remember_output_error(); + return 0; } + return 1; } static void atomic_min(_Atomic uint64_t *value, uint64_t candidate) @@ -600,7 +595,7 @@ static void update_aggregate_slot(struct xprobe_cupti_aggregate_slot *slot, uint64_t duration_ns, uint64_t bytes) { if (atomic_add_checked(&slot->total_duration_ns, duration_ns) == 0 || - atomic_add_checked(&slot->total_bytes, bytes) == 0 || + (bytes != 0U && atomic_add_checked(&slot->total_bytes, bytes) == 0) || atomic_add_checked(&slot->count, 1U) == 0) { return; } @@ -623,7 +618,6 @@ static void aggregate_activity(uint32_t kind, uint32_t device_id, remember_output_error(); return; } - atomic_fetch_add_explicit(&record_count, 1U, memory_order_relaxed); hash = aggregate_hash(kind, device_id, memcpy_kind, name); first = hash % record_capacity; for (uint64_t probe = 0U; probe < record_capacity; ++probe) { @@ -1126,9 +1120,36 @@ static int activity_timestamps_are_host_monotonic(uint64_t available) return 0; } +static uint64_t aggregate_observed_records(void) +{ + uint64_t observed = 0U; + + for (uint64_t index = 0U; index < record_capacity; ++index) { + const struct xprobe_cupti_aggregate_slot *slot = &aggregate_slots[index]; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + uint64_t count; + + if (state < 2U) { + continue; + } + count = atomic_load_explicit(&slot->count, memory_order_relaxed); + if (UINT64_MAX - observed < count) { + remember_output_error(); + return UINT64_MAX; + } + observed += count; + } + return observed; +} + static void initialize_output_header(struct xprobe_cupti_output_header *header, uint64_t available, uint64_t record_offset) { + uint64_t observed = + capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE + ? aggregate_observed_records() + : atomic_load_explicit(&record_count, memory_order_relaxed); + memset(header, 0, sizeof(*header)); memcpy(header->magic, XPROBE_CUPTI_OUTPUT_MAGIC, sizeof(header->magic)); header->abi_version = XPROBE_CUPTI_AGENT_ABI_VERSION; @@ -1146,8 +1167,7 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, (uint32_t)atomic_load_explicit(&stop_reason, memory_order_relaxed); header->record_count = available - record_offset; header->record_capacity = record_capacity; - header->observed_records = - atomic_load_explicit(&record_count, memory_order_relaxed); + header->observed_records = observed; header->agent_dropped_records = atomic_load_explicit(&agent_dropped_records, memory_order_relaxed); header->cupti_dropped_records = diff --git a/docs/benchmarks.md b/docs/benchmarks.md index b8f968b..6af2ace 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -24,3 +24,22 @@ The output records the GPU name, driver, compute capability, workload shape, reference and CUPTI durations, error, baseline and instrumented medians, overhead ratio, ABI metadata, and drop counters. Keep complete JSON output with benchmark results; do not label a run by a GPU model inferred from machine documentation. + +For aggregate inventory changes, run: + +```bash +just benchmark-aggregate +``` + +This benchmark keeps a two-kernel CUDA workload at a high event rate and +captures it once as raw exact events and once into a four-slot aggregate table. +It fails unless both captures are complete and drop-free, the aggregate count +exceeds table capacity without saturation, and exactly two groups are retained. +It also requires aggregate capture to use less collector CPU, collector peak +RSS, target-process RSS growth, and artifact space than raw broad capture. + +The resource comparison uses `wait4` for collector CPU and peak RSS and samples +target RSS from procfs while the command runs. It deliberately does not gate on +short GPU throughput windows: device clock changes can make those samples move +opposite to profiler overhead. The JSON output includes every asserted resource +and collection value so regressions remain diagnosable. diff --git a/docs/cupti-agent.md b/docs/cupti-agent.md index 9d1e98c..35edc28 100644 --- a/docs/cupti-agent.md +++ b/docs/cupti-agent.md @@ -97,4 +97,6 @@ test-injection-live` performs first-load and reactivation measurements against one mapped Agent. `just test-multisource-live` covers host/GPU orchestration and repeated ARM/STOP windows. CUDA 12 variants exercise the same paths against the other linked major, and `just benchmark-gpu` checks callback overhead and -timestamp precision. +timestamp precision. `just benchmark-aggregate` drives a high-rate two-kernel +workload and compares bounded aggregation with raw broad capture for collection +completeness, process resources, and artifact size. diff --git a/docs/development.md b/docs/development.md index 0982e20..300519e 100644 --- a/docs/development.md +++ b/docs/development.md @@ -136,6 +136,13 @@ CUPTI timing or callback hot path: just benchmark-gpu ``` +Run the high-rate bounded-inventory benchmark after changing aggregate storage, +capacity, or activity accounting: + +```bash +just benchmark-aggregate +``` + See `docs/agent-integration.md` and `docs/benchmarks.md` for the tested contract and benchmark interpretation. diff --git a/justfile b/justfile index 0ee885d..3cbcc00 100644 --- a/justfile +++ b/justfile @@ -69,6 +69,9 @@ test-multisource-live-cuda12: build benchmark-gpu: python3 benchmarks/cuda-callback/run.py "{{cuda13_devel_image}}" +benchmark-aggregate: + python3 benchmarks/cuda-aggregate/run.py "{{cuda13_devel_image}}" + fmt: cargo fmt --all From 2fa7bcaa7974532f30db01ec40ab3f6d2e839993 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 05:42:17 +0800 Subject: [PATCH 08/17] =?UTF-8?q?=F0=9F=A4=96=20feat:=20orchestrate=20conc?= =?UTF-8?q?urrent=20worker=20captures?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .github/workflows/hardware.yml | 1 + AGENTS.md | 1 + .../cuda_multiprocess_worker.cu | 101 ++ benchmarks/cuda-multiprocess/run-container.pl | 999 ++++++++++++++++++ benchmarks/cuda-multiprocess/run.py | 68 ++ docs/agent-integration.md | 7 + docs/architecture.md | 4 +- docs/benchmarks.md | 18 + docs/development.md | 7 + justfile | 5 + skills/xprobe-measure-latency/SKILL.md | 23 +- .../references/investigation.md | 6 +- .../references/multi-process.md | 85 ++ .../references/result-quality.md | 7 + .../fixtures/multi-process.json | 88 ++ tests/agent-contract/test_contract.py | 13 + .../test_multi_process_workflow.py | 186 ++++ 17 files changed, 1609 insertions(+), 10 deletions(-) create mode 100644 benchmarks/cuda-multiprocess/cuda_multiprocess_worker.cu create mode 100755 benchmarks/cuda-multiprocess/run-container.pl create mode 100755 benchmarks/cuda-multiprocess/run.py create mode 100644 skills/xprobe-measure-latency/references/multi-process.md create mode 100644 tests/agent-contract/fixtures/multi-process.json create mode 100755 tests/agent-contract/test_multi_process_workflow.py diff --git a/.github/workflows/hardware.yml b/.github/workflows/hardware.yml index a253bd0..f81f655 100644 --- a/.github/workflows/hardware.yml +++ b/.github/workflows/hardware.yml @@ -25,3 +25,4 @@ jobs: - run: just test-multisource-live - run: just benchmark-gpu - run: just benchmark-aggregate + - run: just benchmark-multiprocess diff --git a/AGENTS.md b/AGENTS.md index 92a8a0b..ce79cde 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -73,4 +73,5 @@ - Run `just benchmark-gpu` for callback hot-path changes. - Run `just benchmark-aggregate` for aggregate inventory hot-path or capacity changes. +- Run `just benchmark-multiprocess` for concurrent worker orchestration changes. - Use emoji conventional commits, for example `🐛 fix: restore target registers`. diff --git a/benchmarks/cuda-multiprocess/cuda_multiprocess_worker.cu b/benchmarks/cuda-multiprocess/cuda_multiprocess_worker.cu new file mode 100644 index 0000000..cb94ac5 --- /dev/null +++ b/benchmarks/cuda-multiprocess/cuda_multiprocess_worker.cu @@ -0,0 +1,101 @@ +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +struct SharedControl { + std::uint64_t ready; + std::uint64_t stop; + std::uint64_t iterations; +}; + +static_assert(sizeof(SharedControl) == 24U); +static_assert(offsetof(SharedControl, ready) == 0U); +static_assert(offsetof(SharedControl, stop) == 8U); +static_assert(offsetof(SharedControl, iterations) == 16U); + +__global__ void xprobe_multiprocess_stable_kernel(unsigned long long *output) +{ + const unsigned long long started = clock64(); + while (clock64() - started < 50000ULL) { + } + atomicAdd(output, 1ULL); +} + +static int report_cuda_error(const char *operation, cudaError_t result) +{ + std::fprintf(stderr, "%s failed: %s\n", operation, cudaGetErrorString(result)); + return 13; +} + +int main(int argc, char **argv) +{ + if (argc != 2) { + std::fprintf(stderr, "usage: %s \n", argv[0]); + return 2; + } + + const int control_fd = open(argv[1], O_RDWR | O_CLOEXEC); + if (control_fd < 0) { + std::perror("open shared control"); + return 3; + } + void *mapping = mmap(nullptr, sizeof(SharedControl), PROT_READ | PROT_WRITE, MAP_SHARED, + control_fd, 0); + if (mapping == MAP_FAILED) { + std::perror("mmap shared control"); + close(control_fd); + return 4; + } + auto *control = static_cast(mapping); + + cudaError_t result = cudaSetDevice(0); + if (result != cudaSuccess) { + return report_cuda_error("cudaSetDevice", result); + } + unsigned long long *device_output = nullptr; + result = cudaMalloc(&device_output, sizeof(*device_output)); + if (result != cudaSuccess) { + return report_cuda_error("cudaMalloc", result); + } + result = cudaMemset(device_output, 0, sizeof(*device_output)); + if (result != cudaSuccess) { + return report_cuda_error("cudaMemset", result); + } + + __atomic_store_n(&control->ready, 1ULL, __ATOMIC_RELEASE); + while (__atomic_load_n(&control->stop, __ATOMIC_ACQUIRE) == 0ULL) { + xprobe_multiprocess_stable_kernel<<<1, 1>>>(device_output); + result = cudaGetLastError(); + if (result == cudaSuccess) { + result = cudaDeviceSynchronize(); + } + if (result != cudaSuccess) { + return report_cuda_error("kernel workload", result); + } + __atomic_fetch_add(&control->iterations, 1ULL, __ATOMIC_RELAXED); + } + + result = cudaFree(device_output); + if (result == cudaSuccess) { + result = cudaDeviceReset(); + } + if (result != cudaSuccess) { + return report_cuda_error("CUDA shutdown", result); + } + if (munmap(mapping, sizeof(SharedControl)) != 0) { + std::perror("munmap shared control"); + return 5; + } + if (close(control_fd) != 0) { + std::perror("close shared control"); + return 6; + } + return 0; +} diff --git a/benchmarks/cuda-multiprocess/run-container.pl b/benchmarks/cuda-multiprocess/run-container.pl new file mode 100755 index 0000000..9f36780 --- /dev/null +++ b/benchmarks/cuda-multiprocess/run-container.pl @@ -0,0 +1,999 @@ +#!/usr/bin/perl +use strict; +use warnings; + +use Errno qw(EINTR); +use Fcntl qw(O_CREAT O_EXCL O_RDONLY O_RDWR O_WRONLY SEEK_SET); +use File::Find qw(find); +use IO::Handle; +use JSON::PP (); +use POSIX qw(WNOHANG _exit); +use Time::HiRes qw(CLOCK_MONOTONIC clock_gettime usleep); + +use constant { + BASELINE_SECONDS => 0.35, + INVENTORY_DURATION_MS => 500, + MEASURE_DURATION_MS => 650, + COLLECTION_WARMUP_US => 150_000, + COLLECTION_SAMPLE_US => 350_000, + COMMAND_TIMEOUT_MS => 15_000, + CONTROL_SIZE => 24, + READY_OFFSET => 0, + STOP_OFFSET => 8, + ITERATIONS_OFFSET => 16, +}; + +my $JSON = JSON::PP->new->canonical(1); +my $OUTPUT_DIRECTORY; + +END { + if (defined($OUTPUT_DIRECTORY) && -d $OUTPUT_DIRECTORY + && defined($ENV{XPROBE_HOST_UID}) && defined($ENV{XPROBE_HOST_GID})) { + my @paths; + find({wanted => sub { push @paths, $File::Find::name }, no_chdir => 1}, + $OUTPUT_DIRECTORY); + my $changed = chown( + 0 + $ENV{XPROBE_HOST_UID}, 0 + $ENV{XPROBE_HOST_GID}, @paths); + warn "failed to restore benchmark artifact ownership\n" + if $changed != @paths; + } +} + +sub fail { + my ($message) = @_; + die "$message\n"; +} + +sub require_condition { + my ($condition, $message) = @_; + fail($message) unless $condition; +} + +sub monotonic_ns { + return int(clock_gettime(CLOCK_MONOTONIC) * 1_000_000_000); +} + +sub read_text { + my ($path) = @_; + open my $input, '<', $path or fail("unable to read $path: $!"); + local $/; + my $content = <$input>; + close $input or fail("unable to close $path: $!"); + return defined($content) ? $content : ''; +} + +sub write_text { + my ($path, $content) = @_; + open my $output, '>', $path or fail("unable to write $path: $!"); + print {$output} $content or fail("unable to write $path: $!"); + close $output or fail("unable to close $path: $!"); +} + +sub write_json { + my ($path, $value) = @_; + write_text($path, $JSON->encode($value) . "\n"); +} + +sub procfs_start_time { + my ($pid) = @_; + my $path = "/proc/$pid/stat"; + my $stat = read_text($path); + my $close_parenthesis = rindex($stat, ')'); + require_condition($close_parenthesis >= 0, "malformed procfs stat for PID $pid"); + my $tail = substr($stat, $close_parenthesis + 2); + my @fields = split /\s+/, $tail; + require_condition(@fields > 19 && $fields[19] =~ /^\d+$/, + "malformed procfs start time for PID $pid"); + return 0 + $fields[19]; +} + +sub verify_identity { + my ($worker, $phase) = @_; + require_condition(defined($worker->{pid}), + "worker $worker->{ordinal} has no PID during $phase"); + require_condition(defined($worker->{process_start_time}), + "worker $worker->{ordinal} has no procfs start time during $phase"); + my $current = procfs_start_time($worker->{pid}); + require_condition($current == $worker->{process_start_time}, + "worker $worker->{ordinal} PID $worker->{pid} identity mismatch during $phase: " + . "expected start time $worker->{process_start_time}, found $current"); +} + +sub worker_stem { + my ($worker) = @_; + require_condition(defined($worker->{pid}), 'worker has no PID for artifact naming'); + require_condition(defined($worker->{process_start_time}), + 'worker has no start time for artifact naming'); + return "worker-$worker->{pid}-$worker->{process_start_time}"; +} + +sub worker_path { + my ($batch_dir, $worker, $suffix) = @_; + return "$batch_dir/" . worker_stem($worker) . "-$suffix"; +} + +sub rename_worker_startup_files { + my ($batch_dir, $worker) = @_; + my $control_path = worker_path($batch_dir, $worker, 'control'); + my $stdout_path = worker_path($batch_dir, $worker, 'fixture.stdout'); + my $stderr_path = worker_path($batch_dir, $worker, 'fixture.stderr'); + rename $worker->{control_path}, $control_path + or fail("unable to rename worker control file: $!"); + rename $worker->{stdout_path}, $stdout_path + or fail("unable to rename worker stdout: $!"); + rename $worker->{stderr_path}, $stderr_path + or fail("unable to rename worker stderr: $!"); + $worker->{control_path} = $control_path; + $worker->{stdout_path} = $stdout_path; + $worker->{stderr_path} = $stderr_path; +} + +sub write_control { + my ($path, $offset, $value) = @_; + sysopen my $control, $path, O_RDWR + or fail("unable to open shared control $path: $!"); + sysseek($control, $offset, SEEK_SET) == $offset + or fail("unable to seek shared control $path: $!"); + my $bytes = pack('Q<', $value); + syswrite($control, $bytes, length($bytes)) == length($bytes) + or fail("unable to write shared control $path: $!"); + $control->sync or fail("unable to sync shared control $path: $!"); + close $control or fail("unable to close shared control $path: $!"); +} + +sub read_control { + my ($path, $offset) = @_; + sysopen my $control, $path, O_RDONLY + or fail("unable to open shared control $path: $!"); + sysseek($control, $offset, SEEK_SET) == $offset + or fail("unable to seek shared control $path: $!"); + my $bytes = ''; + my $count = sysread($control, $bytes, 8); + require_condition(defined($count) && $count == 8, + "unable to read shared control $path: $!"); + close $control or fail("unable to close shared control $path: $!"); + return unpack('Q<', $bytes); +} + +sub iteration_snapshot { + my ($worker) = @_; + return read_control($worker->{control_path}, ITERATIONS_OFFSET); +} + +sub sample_iterations { + my ($worker, $seconds) = @_; + my $before = iteration_snapshot($worker); + my $started = monotonic_ns(); + usleep(int($seconds * 1_000_000)); + my $after = iteration_snapshot($worker); + my $finished = monotonic_ns(); + my $elapsed_ns = $finished - $started; + require_condition($after >= $before, + "worker $worker->{ordinal} iteration counter moved backwards"); + require_condition($after > $before, + "worker $worker->{ordinal} made no baseline kernel progress"); + return { + start => $before, + end => $after, + delta => $after - $before, + wall_ns => $elapsed_ns, + iterations_per_second => ($after - $before) * 1_000_000_000 / $elapsed_ns, + }; +} + +sub artifact { + my ($path) = @_; + require_condition(-f $path, "missing artifact $path"); + return {path => $path, bytes => -s $path}; +} + +sub warning_present { + my ($result, $code) = @_; + require_condition(ref($result->{warnings}) eq 'ARRAY', + 'result warnings are malformed'); + for my $warning (@{$result->{warnings}}) { + return JSON::PP::true + if ref($warning) eq 'HASH' && defined($warning->{code}) + && $warning->{code} eq $code; + } + return JSON::PP::false; +} + +sub load_json { + my ($path, $description) = @_; + my $value = eval { $JSON->decode(read_text($path)) }; + fail("malformed $description JSON at $path: $@") if $@; + require_condition(ref($value) eq 'HASH', + "malformed $description JSON at $path: expected object"); + require_condition(defined($value->{schema_version}) + && $value->{schema_version} eq '2.0', + "unexpected schema in $description"); + return $value; +} + +sub decode_wait_status { + my ($status) = @_; + return 128 + ($status & 127) if ($status & 127); + return ($status >> 8) & 255; +} + +sub launch_command { + my ($command, $stdout_path, $stderr_path) = @_; + my $pid = fork(); + fail("fork failed: $!") unless defined($pid); + if ($pid == 0) { + open STDOUT, '>', $stdout_path or do { + print STDERR "unable to open $stdout_path: $!\n"; + _exit(127); + }; + open STDERR, '>', $stderr_path or do { + print STDERR "unable to open $stderr_path: $!\n"; + _exit(127); + }; + exec {$command->[0]} @{$command} or do { + print STDERR "exec $command->[0] failed: $!\n"; + _exit(127); + }; + } + return $pid; +} + +sub wait_for_command { + my ($pid) = @_; + while (1) { + my $result = waitpid($pid, 0); + next if $result < 0 && $! == EINTR; + fail("waitpid failed for PID $pid: $!") if $result < 0; + return decode_wait_status($?); + } +} + +sub run_command { + my ($command, $stdout_path, $stderr_path) = @_; + return wait_for_command(launch_command($command, $stdout_path, $stderr_path)); +} + +sub run_checked { + my ($description, @command) = @_; + my $status = system {$command[0]} @command; + fail("$description failed to execute: $!") if $status == -1; + my $exit_status = decode_wait_status($status); + require_condition($exit_status == 0, + "$description failed with exit status $exit_status"); +} + +sub capture_checked { + my ($description, @command) = @_; + open my $input, '-|', @command + or fail("$description failed to execute: $!"); + local $/; + my $output = <$input>; + my $closed = close $input; + my $status = $?; + require_condition($closed, + "$description failed with exit status " . decode_wait_status($status)); + return defined($output) ? $output : ''; +} + +sub command_for_aggregate { + my ($xprobe, $agent, $worker, $start_selector, $end_selector, $duration_ms) = @_; + require_condition(defined($worker->{pid}), + 'cannot construct measurement without a worker PID'); + return [ + $xprobe, 'measure', + '--pid', $worker->{pid}, + '--agent', $agent, + '--from', $start_selector, + '--to', $end_selector, + '--match', 'exact', + '--duration-ms', $duration_ms, + '--timeout-ms', COMMAND_TIMEOUT_MS, + '--json', '--non-interactive', '--no-color', + '--aggregate', '--max-groups', 8, + ]; +} + +sub write_measurement_spec { + my ($path, $worker, $start_selector, $end_selector) = @_; + require_condition(defined($worker->{pid}), + 'cannot write a MeasurementSpec without a worker PID'); + require_condition(defined($worker->{process_start_time}), + 'cannot write a MeasurementSpec without a start time'); + write_json($path, { + schema_version => '2.0', + name => 'cuda_multiprocess_kernel_duration', + target => { + pid => $worker->{pid}, + process_start_time => $worker->{process_start_time}, + }, + start_selector => $start_selector, + end_selector => $end_selector, + match_policy => 'exact', + samples => undef, + duration_ms => MEASURE_DURATION_MS, + timeout_ms => COMMAND_TIMEOUT_MS, + max_events => 200_000, + measurement_mode => 'exact', + }); +} + +sub command_for_spec_measure { + my ($xprobe, $agent, $spec_path, $events_path) = @_; + return [ + $xprobe, 'measure', + '--spec', $spec_path, + '--agent', $agent, + '--events-out', $events_path, + '--json', '--non-interactive', '--no-color', + ]; +} + +sub validate_result { + my ($result, $worker, $start_selector, $end_selector) = @_; + require_condition($result->{ok} && $result->{valid}, 'selector validation failed'); + require_condition(ref($result->{target}) eq 'HASH', + 'validation result has no target identity'); + require_condition($result->{target}{pid} == $worker->{pid}, + 'validation target PID does not match worker'); + require_condition( + $result->{target}{process_start_time} == $worker->{process_start_time}, + 'validation target procfs start time does not match worker'); + require_condition(ref($result->{start}) eq 'HASH' + && $result->{start}{selector} eq $start_selector, + 'validation start selector mismatch'); + require_condition(ref($result->{end}) eq 'HASH' + && $result->{end}{selector} eq $end_selector, + 'validation end selector mismatch'); + require_condition(ref($result->{requirements}) eq 'HASH', + 'validation requirements are malformed'); + my $activation = $result->{requirements}{agent_activation}; + require_condition(defined($activation) + && ($activation eq 'already_loaded' || $activation eq 'injection_required'), + 'unexpected CUDA agent activation requirement'); + my $mutation_warning = warning_present($result, 'TARGET_PROCESS_WILL_BE_MODIFIED'); + require_condition($mutation_warning, + 'validation omitted required target-mutation warning') + if $activation eq 'injection_required'; + return { + agent_activation => $activation, + target_mutation => $result->{requirements}{target_mutation}, + target_mutation_warning => $mutation_warning, + policy_recommendation => $result->{policy_recommendation}, + }; +} + +sub validate_exact_result { + my ($result, $worker) = @_; + require_condition($result->{ok}, + "worker $worker->{ordinal} measure result is not successful"); + require_condition($result->{status} eq 'completed', + "worker $worker->{ordinal} capture is incomplete"); + my $collection = $result->{collection}; + require_condition(ref($collection) eq 'HASH', + 'measure collection is malformed'); + require_condition($collection->{completeness} eq 'complete', + 'measure collection is incomplete'); + require_condition($collection->{dropped_events} == 0, + 'measure capture dropped events'); + require_condition($collection->{cuda_events} > 0, + 'measure capture retained no CUDA events'); + my $cupti = $collection->{cupti}; + require_condition(ref($cupti) eq 'HASH', + 'measure result lacks CUPTI quality fields'); + require_condition($cupti->{dropped_records} == 0, + 'measure CUPTI capture dropped records'); + require_condition($cupti->{observed_records} == $cupti->{retained_records}, + 'measure CUPTI retained record count is incomplete'); + require_condition($cupti->{retained_records} == $collection->{cuda_events}, + 'measure CUDA event count does not match retained records'); + my $measurement = $result->{measurement}; + require_condition(ref($measurement) eq 'HASH', + 'measure result lacks measurement quality'); + require_condition(ref($measurement->{samples}) eq 'HASH' + && $measurement->{samples}{matched} > 0, + 'measure capture has no matched kernel samples'); + return { + completeness => $collection->{completeness}, + cuda_events => $collection->{cuda_events}, + dropped_events => $collection->{dropped_events}, + cupti => $cupti, + samples => $measurement->{samples}, + correlation => $result->{correlation}, + clock => $result->{clock}, + }; +} + +sub make_worker { + my ($batch_dir, $ordinal, $fixture) = @_; + my $control_path = "$batch_dir/worker-$ordinal.control"; + sysopen my $control, $control_path, O_WRONLY | O_CREAT | O_EXCL, 0644 + or fail("unable to create $control_path: $!"); + my $empty = "\0" x CONTROL_SIZE; + syswrite($control, $empty, length($empty)) == length($empty) + or fail("unable to initialize $control_path: $!"); + close $control or fail("unable to close $control_path: $!"); + my $worker = { + ordinal => $ordinal, + control_path => $control_path, + stdout_path => "$batch_dir/worker-$ordinal.stdout", + stderr_path => "$batch_dir/worker-$ordinal.stderr", + launched_monotonic_ns => monotonic_ns(), + files => {}, + }; + $worker->{process_pid} = launch_command( + [$fixture, $control_path], + $worker->{stdout_path}, + $worker->{stderr_path}); + return $worker; +} + +sub poll_worker { + my ($worker) = @_; + return $worker->{exit_status} if defined($worker->{exit_status}); + my $result = waitpid($worker->{process_pid}, WNOHANG); + return undef if $result == 0; + fail("waitpid failed for worker $worker->{ordinal}: $!") if $result < 0; + $worker->{exit_status} = decode_wait_status($?); + return $worker->{exit_status}; +} + +sub wait_for_workers_ready { + my ($batch_dir, $workers) = @_; + my $deadline = monotonic_ns() + 10_000_000_000; + my %pending = map { $_->{ordinal} => $_ } @{$workers}; + while (%pending && monotonic_ns() < $deadline) { + for my $ordinal (keys %pending) { + my $worker = $pending{$ordinal}; + my $exit_status = poll_worker($worker); + fail("worker $ordinal exited before readiness with status $exit_status") + if defined($exit_status); + if (read_control($worker->{control_path}, READY_OFFSET) == 1) { + $worker->{ready_monotonic_ns} = monotonic_ns(); + $worker->{pid} = $worker->{process_pid}; + $worker->{process_start_time} = procfs_start_time($worker->{pid}); + verify_identity($worker, 'startup'); + rename_worker_startup_files($batch_dir, $worker); + delete $pending{$ordinal}; + } + } + usleep(10_000) if %pending; + } + require_condition(!%pending, + 'worker readiness timed out for indexes ' . join(',', sort keys %pending)); +} + +sub stop_workers { + my ($workers) = @_; + my @failures; + for my $worker (@{$workers}) { + next if defined(poll_worker($worker)); + eval { write_control($worker->{control_path}, STOP_OFFSET, 1); 1 } + or push @failures, + "worker $worker->{ordinal} stop signal failed: " . ($@ || 'unknown error'); + } + my $deadline = monotonic_ns() + 10_000_000_000; + while (monotonic_ns() < $deadline) { + my $running = 0; + for my $worker (@{$workers}) { + $running++ unless defined(poll_worker($worker)); + } + last unless $running; + usleep(10_000); + } + for my $worker (@{$workers}) { + unless (defined(poll_worker($worker))) { + kill 'KILL', $worker->{process_pid}; + waitpid($worker->{process_pid}, 0); + $worker->{exit_status} = decode_wait_status($?); + push @failures, + "worker $worker->{ordinal} did not stop after its stop signal"; + } + push @failures, + "worker $worker->{ordinal} exited with status $worker->{exit_status}" + if $worker->{exit_status} != 0; + } + fail(join('; ', @failures)) if @failures; +} + +sub discover_workers { + my ($batch_dir, $xprobe, $workers) = @_; + my $root_pid = $$; + my $root_start_time = procfs_start_time($root_pid); + my $stdout_path = "$batch_dir/root-$root_pid-$root_start_time-discover.json"; + my $stderr_path = "$batch_dir/root-$root_pid-$root_start_time-discover.stderr"; + my $exit_status = run_command([ + $xprobe, 'discover', + '--pid', $root_pid, + '--limit', scalar(@{$workers}) + 8, + '--json', '--non-interactive', '--no-color', + ], $stdout_path, $stderr_path); + require_condition($exit_status == 0, + 'CUDA worker discovery command failed'); + my $result = load_json($stdout_path, 'CUDA worker discovery'); + require_condition($result->{ok}, 'CUDA worker discovery was unsuccessful'); + require_condition(ref($result->{root}) eq 'HASH', + 'CUDA worker discovery root is malformed'); + require_condition($result->{root}{pid} == $root_pid, + 'CUDA worker discovery root PID mismatch'); + require_condition($result->{root}{process_start_time} == $root_start_time, + 'CUDA worker discovery root procfs start time mismatch'); + require_condition(!$result->{truncated}, + 'CUDA worker discovery candidates were truncated'); + require_condition(ref($result->{candidates}) eq 'ARRAY', + 'CUDA worker discovery candidates are malformed'); + my %discovered; + for my $candidate (@{$result->{candidates}}) { + next unless ref($candidate) eq 'HASH' + && ref($candidate->{target}) eq 'HASH'; + $discovered{"$candidate->{target}{pid}:$candidate->{target}{process_start_time}"} = 1; + } + for my $worker (@{$workers}) { + require_condition( + $discovered{"$worker->{pid}:$worker->{process_start_time}"}, + 'CUDA worker discovery did not return every launched worker identity'); + verify_identity($worker, 'after CUDA discovery'); + } + return { + exit_status => $exit_status, + root => { + pid => $root_pid, + process_start_time => $root_start_time, + }, + candidates => $result->{candidates}, + files => { + stdout => artifact($stdout_path), + stderr => artifact($stderr_path), + }, + }; +} + +sub compile_fixture { + my ($workspace, $build_dir) = @_; + my $cuda_root = '/usr/local/cuda'; + my $agent = "$build_dir/libxprobe-cupti.so"; + my $fixture = "$build_dir/xprobe-cuda-multiprocess-worker"; + my $capabilities = capture_checked('GPU compute-capability query', + 'nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader'); + my ($capability) = grep { length($_) } split /\n/, $capabilities; + require_condition(defined($capability), 'GPU compute-capability query returned no rows'); + $capability =~ s/^\s+|\s+$//g; + require_condition($capability =~ /^\d+\.\d+$/, + "invalid GPU compute capability $capability"); + (my $architecture = $capability) =~ s/\.//g; + run_checked('CUPTI Agent compilation', + 'gcc', + '-std=c11', '-D_GNU_SOURCE', '-DXPROBE_HAS_CUPTI=1', + '-fPIC', '-shared', '-pthread', '-O2', + '-Wall', '-Wextra', '-Wpedantic', '-Werror', + '-I/workspace/cupti/include', '-isystem', "$cuda_root/include", + "$workspace/cupti/src/cupti_agent.c", + "-L$cuda_root/lib64", "-Wl,-rpath,$cuda_root/lib64", '-lcupti', + '-o', $agent); + run_checked('CUDA worker compilation', + 'nvcc', + '-std=c++17', '-O2', + "-gencode=arch=compute_$architecture,code=sm_$architecture", + "$workspace/benchmarks/cuda-multiprocess/cuda_multiprocess_worker.cu", + '-o', $fixture); + return ($agent, $fixture); +} + +sub wait_for_concurrent_commands { + my ($entries, $description) = @_; + my $deadline = monotonic_ns() + (COMMAND_TIMEOUT_MS + 5_000) * 1_000_000; + my %pending = map { $_->{ordinal} => $_ } @{$entries}; + while (%pending) { + for my $ordinal (keys %pending) { + my $entry = $pending{$ordinal}; + my $result = waitpid($entry->{command_pid}, WNOHANG); + next if $result == 0; + fail("waitpid failed for $description worker $ordinal: $!") + if $result < 0; + $entry->{finished_monotonic_ns} = monotonic_ns(); + $entry->{exit_status} = decode_wait_status($?); + delete $pending{$ordinal}; + } + if (%pending && monotonic_ns() >= $deadline) { + for my $entry (values %pending) { + kill 'KILL', $entry->{command_pid}; + waitpid($entry->{command_pid}, 0); + $entry->{finished_monotonic_ns} = monotonic_ns(); + $entry->{exit_status} = decode_wait_status($?); + } + fail("per-worker $description commands exceeded the benchmark deadline"); + } + usleep(5_000) if %pending; + } +} + +sub run_batch_body { + my ($output_dir, $xprobe, $agent, $fixture, $worker_count, + $workers, $worker_reports) = @_; + my $batch_dir = "$output_dir/workers-$worker_count"; + mkdir $batch_dir or fail("unable to create $batch_dir: $!"); + for my $ordinal (0 .. $worker_count - 1) { + push @{$workers}, make_worker($batch_dir, $ordinal, $fixture); + } + wait_for_workers_ready($batch_dir, $workers); + my $discovery = discover_workers($batch_dir, $xprobe, $workers); + my %baselines; + for my $worker (@{$workers}) { + verify_identity($worker, 'before baseline'); + $baselines{$worker->{ordinal}} = + sample_iterations($worker, BASELINE_SECONDS); + } + + my $representative = $workers->[0]; + my $broad_validate_stdout = + worker_path($batch_dir, $representative, 'broad-validate.json'); + my $broad_validate_stderr = + worker_path($batch_dir, $representative, 'broad-validate.stderr'); + my $broad_validate_exit = run_command([ + $xprobe, 'validate', + '--pid', $representative->{pid}, + '--from', 'cuda:kernel_start', + '--to', 'cuda:kernel_end', + '--match', 'exact', + '--json', '--non-interactive', '--no-color', + ], $broad_validate_stdout, $broad_validate_stderr); + require_condition($broad_validate_exit == 0, + 'representative broad selector validation command failed'); + my $broad_validation = validate_result( + load_json($broad_validate_stdout, 'broad selector validation'), + $representative, 'cuda:kernel_start', 'cuda:kernel_end'); + $broad_validation->{exit_status} = $broad_validate_exit; + verify_identity($representative, 'after broad selector validation'); + + my $inventory_stdout = + worker_path($batch_dir, $representative, 'inventory.json'); + my $inventory_stderr = + worker_path($batch_dir, $representative, 'inventory.stderr'); + my $inventory_before = iteration_snapshot($representative); + my $inventory_started = monotonic_ns(); + my $inventory_exit = run_command(command_for_aggregate( + $xprobe, $agent, $representative, + 'cuda:kernel_start', 'cuda:kernel_end', INVENTORY_DURATION_MS), + $inventory_stdout, $inventory_stderr); + my $inventory_finished = monotonic_ns(); + my $inventory_after = iteration_snapshot($representative); + require_condition($inventory_exit == 0, + 'representative aggregate inventory command failed'); + verify_identity($representative, 'after aggregate inventory'); + my $inventory = load_json($inventory_stdout, 'aggregate inventory'); + require_condition($inventory->{ok} && $inventory->{status} eq 'completed', + 'aggregate inventory is not complete'); + my $inventory_collection = $inventory->{collection}; + require_condition(ref($inventory_collection) eq 'HASH', + 'aggregate inventory collection is malformed'); + require_condition($inventory_collection->{completeness} eq 'complete', + 'aggregate inventory is incomplete'); + require_condition($inventory_collection->{dropped_activities} == 0, + 'aggregate inventory dropped activities'); + require_condition( + $inventory_collection->{observed_activities} + == $inventory_collection->{grouped_activities}, + 'aggregate inventory grouped activity count is inconsistent'); + require_condition(ref($inventory->{inventory}) eq 'HASH' + && ref($inventory->{inventory}{groups}) eq 'ARRAY' + && @{$inventory->{inventory}{groups}} == 1, + 'homogeneous aggregate inventory did not return exactly one group'); + my $group = $inventory->{inventory}{groups}[0]; + require_condition(ref($group) eq 'HASH' && $group->{activity} eq 'kernel', + 'inventory group is not a kernel'); + require_condition(defined($group->{name}) + && index($group->{name}, 'xprobe_multiprocess_stable_kernel') >= 0, + 'inventory did not identify the stable worker kernel'); + require_condition($group->{count} > 0, + 'inventory kernel group has no activities'); + my $start_selector = $group->{start_selector_hint}; + my $end_selector = $group->{end_selector_hint}; + require_condition(defined($start_selector) + && index($start_selector, 'name~') >= 0, + 'inventory did not provide a narrow kernel start selector'); + require_condition(defined($end_selector) + && index($end_selector, 'name~') >= 0, + 'inventory did not provide a narrow kernel end selector'); + require_condition(warning_present($inventory, 'CUPTI_AGENT_INJECTED'), + 'aggregate inventory omitted the required injection warning'); + + my @validation_commands; + for my $worker (@{$workers}) { + verify_identity($worker, 'before selector validation'); + my $stdout_path = worker_path($batch_dir, $worker, 'validate.json'); + my $stderr_path = worker_path($batch_dir, $worker, 'validate.stderr'); + $worker->{files}{validate_stdout} = $stdout_path; + $worker->{files}{validate_stderr} = $stderr_path; + my $started = monotonic_ns(); + my $command_pid = launch_command([ + $xprobe, 'validate', + '--pid', $worker->{pid}, + '--from', $start_selector, + '--to', $end_selector, + '--match', 'exact', + '--json', '--non-interactive', '--no-color', + ], $stdout_path, $stderr_path); + push @validation_commands, { + ordinal => $worker->{ordinal}, + worker => $worker, + command_pid => $command_pid, + started_monotonic_ns => $started, + }; + } + wait_for_concurrent_commands(\@validation_commands, 'selector validation'); + my %validations; + for my $entry (@validation_commands) { + my $worker = $entry->{worker}; + require_condition($entry->{exit_status} == 0, + "worker $worker->{ordinal} selector validation command failed"); + my $validation = validate_result( + load_json($worker->{files}{validate_stdout}, 'selector validation'), + $worker, $start_selector, $end_selector); + $validation->{exit_status} = $entry->{exit_status}; + $validation->{started_monotonic_ns} = + $entry->{started_monotonic_ns}; + $validation->{finished_monotonic_ns} = + $entry->{finished_monotonic_ns}; + $validations{$worker->{ordinal}} = $validation; + verify_identity($worker, 'after selector validation'); + } + + my @measure_commands; + for my $worker (@{$workers}) { + verify_identity($worker, 'before exact measurement'); + my $spec = worker_path($batch_dir, $worker, 'measurement-spec.json'); + my $stdout_path = worker_path($batch_dir, $worker, 'measure.json'); + my $stderr_path = worker_path($batch_dir, $worker, 'measure.stderr'); + my $events = worker_path($batch_dir, $worker, 'events.jsonl'); + write_measurement_spec($spec, $worker, $start_selector, $end_selector); + $worker->{files}{spec} = $spec; + $worker->{files}{measure_stdout} = $stdout_path; + $worker->{files}{measure_stderr} = $stderr_path; + $worker->{files}{events} = $events; + my $started = monotonic_ns(); + my $command_pid = launch_command( + command_for_spec_measure($xprobe, $agent, $spec, $events), + $stdout_path, $stderr_path); + push @measure_commands, { + ordinal => $worker->{ordinal}, + worker => $worker, + command_pid => $command_pid, + started_monotonic_ns => $started, + }; + } + usleep(COLLECTION_WARMUP_US); + my $collection_sample_started_ns = monotonic_ns(); + for my $entry (@measure_commands) { + $entry->{collection_start} = iteration_snapshot($entry->{worker}); + } + usleep(COLLECTION_SAMPLE_US); + for my $entry (@measure_commands) { + $entry->{collection_end} = iteration_snapshot($entry->{worker}); + } + my $collection_sample_finished_ns = monotonic_ns(); + my $collection_sample_wall_ns = + $collection_sample_finished_ns - $collection_sample_started_ns; + wait_for_concurrent_commands(\@measure_commands, 'measure'); + my $latest_start = 0; + my $earliest_finish; + for my $entry (@measure_commands) { + $latest_start = $entry->{started_monotonic_ns} + if $entry->{started_monotonic_ns} > $latest_start; + $earliest_finish = $entry->{finished_monotonic_ns} + if !defined($earliest_finish) + || $entry->{finished_monotonic_ns} < $earliest_finish; + } + require_condition($latest_start < $earliest_finish, + 'per-worker measure command windows did not overlap'); + + my @all_paths = ( + $discovery->{files}{stdout}{path}, + $discovery->{files}{stderr}{path}, + $broad_validate_stdout, + $broad_validate_stderr, + $inventory_stdout, + $inventory_stderr, + ); + for my $entry (@measure_commands) { + my $worker = $entry->{worker}; + my $collection_end = $entry->{collection_end}; + require_condition($collection_end >= $entry->{collection_start}, + 'worker iteration counter moved backwards during measure'); + require_condition($entry->{exit_status} == 0, + "worker $worker->{ordinal} exact measurement command failed"); + verify_identity($worker, 'after exact measurement'); + my $exact = load_json($worker->{files}{measure_stdout}, + 'exact measurement'); + my $quality = validate_exact_result($exact, $worker); + my $injection_warning = + warning_present($exact, 'CUPTI_AGENT_INJECTED'); + require_condition($injection_warning, + 'measure omitted required injection warning') + if $validations{$worker->{ordinal}}{agent_activation} + eq 'injection_required'; + my $wall_ns = + $entry->{finished_monotonic_ns} - $entry->{started_monotonic_ns}; + my $collection_delta = + $collection_end - $entry->{collection_start}; + my %files = ( + worker_stdout => artifact($worker->{stdout_path}), + worker_stderr => artifact($worker->{stderr_path}), + ); + for my $key (keys %{$worker->{files}}) { + $files{$key} = artifact($worker->{files}{$key}); + } + push @{$worker_reports}, { + ordinal => $worker->{ordinal}, + target => { + pid => $worker->{pid}, + process_start_time => $worker->{process_start_time}, + }, + fixture_exit_status => undef, + startup => { + launched_monotonic_ns => $worker->{launched_monotonic_ns}, + ready_monotonic_ns => $worker->{ready_monotonic_ns}, + wall_ns => $worker->{ready_monotonic_ns} + - $worker->{launched_monotonic_ns}, + }, + baseline_iterations => $baselines{$worker->{ordinal}}, + collection_iterations => { + start => $entry->{collection_start}, + end => $collection_end, + delta => $collection_delta, + wall_ns => $collection_sample_wall_ns, + iterations_per_second => + $collection_delta * 1_000_000_000 + / $collection_sample_wall_ns, + retained_events_per_second => + $quality->{cuda_events} * 1_000_000_000 / $wall_ns, + }, + perturbation => { + baseline_iterations_per_second => + $baselines{$worker->{ordinal}}{iterations_per_second}, + collection_iterations_per_second => + $collection_delta * 1_000_000_000 + / $collection_sample_wall_ns, + relative_change => + ($collection_delta * 1_000_000_000 + / $collection_sample_wall_ns) + / $baselines{$worker->{ordinal}}{iterations_per_second} + - 1.0, + }, + command => { + exit_status => $entry->{exit_status}, + started_monotonic_ns => $entry->{started_monotonic_ns}, + finished_monotonic_ns => $entry->{finished_monotonic_ns}, + wall_ns => $wall_ns, + }, + validation => $validations{$worker->{ordinal}}, + injection_warning => $injection_warning, + quality => $quality, + files => \%files, + }; + push @all_paths, $worker->{stdout_path}, $worker->{stderr_path}; + push @all_paths, values %{$worker->{files}}; + } + my %unique_paths = map { $_ => 1 } @all_paths; + require_condition(@all_paths == keys(%unique_paths), + 'benchmark artifact path reuse detected'); + my $inventory_wall_ns = $inventory_finished - $inventory_started; + return { + workers => $worker_count, + representative => { + pid => $representative->{pid}, + process_start_time => $representative->{process_start_time}, + }, + discovery => $discovery, + selectors => { + from => $start_selector, + to => $end_selector, + }, + aggregate_inventory => { + broad_validation => { + %{$broad_validation}, + files => { + stdout => artifact($broad_validate_stdout), + stderr => artifact($broad_validate_stderr), + }, + }, + command => { + exit_status => $inventory_exit, + started_monotonic_ns => $inventory_started, + finished_monotonic_ns => $inventory_finished, + wall_ns => $inventory_wall_ns, + }, + collection_iterations => { + start => $inventory_before, + end => $inventory_after, + delta => $inventory_after - $inventory_before, + iterations_per_second => + ($inventory_after - $inventory_before) + * 1_000_000_000 / $inventory_wall_ns, + }, + injection_warning => JSON::PP::true, + collection => $inventory_collection, + group => $group, + files => { + stdout => artifact($inventory_stdout), + stderr => artifact($inventory_stderr), + }, + }, + measure_windows_overlap => JSON::PP::true, + workers_report => $worker_reports, + }; +} + +sub run_batch { + my ($output_dir, $xprobe, $agent, $fixture, $worker_count) = @_; + my @workers; + my @worker_reports; + my ($result, $primary_error, $cleanup_error); + eval { + $result = run_batch_body( + $output_dir, $xprobe, $agent, $fixture, $worker_count, + \@workers, \@worker_reports); + 1; + } or $primary_error = $@ || 'unknown benchmark failure'; + eval { + stop_workers(\@workers); + for my $report (@worker_reports) { + my $worker = $workers[$report->{ordinal}]; + require_condition(defined($worker->{exit_status}), + "worker $report->{ordinal} fixture did not produce an exit status"); + $report->{fixture_exit_status} = $worker->{exit_status}; + } + 1; + } or $cleanup_error = $@ || 'unknown cleanup failure'; + if ($primary_error) { + chomp $primary_error; + if ($cleanup_error) { + chomp $cleanup_error; + fail("$primary_error; cleanup failed: $cleanup_error"); + } + fail($primary_error); + } + fail($cleanup_error) if $cleanup_error; + return $result; +} + +sub main { + require_condition(@ARGV == 1, + 'usage: run-container.pl '); + my $output_dir = $ARGV[0]; + $OUTPUT_DIRECTORY = $output_dir; + mkdir $output_dir unless -d $output_dir; + require_condition(-d $output_dir, + "output directory is unavailable at $output_dir"); + my $workspace = '/workspace'; + my $xprobe = "$workspace/target/debug/xprobe"; + require_condition(-f $xprobe, "xprobe binary is missing at $xprobe"); + my $build_dir = '/tmp/xprobe-cuda-multiprocess'; + mkdir $build_dir unless -d $build_dir; + require_condition(-d $build_dir, + "build directory is unavailable at $build_dir"); + my ($agent, $fixture) = compile_fixture($workspace, $build_dir); + my $gpu_rows = capture_checked('GPU information query', + 'nvidia-smi', + '--query-gpu=name,driver_version,compute_cap', + '--format=csv,noheader'); + my ($gpu) = grep { length($_) } split /\n/, $gpu_rows; + require_condition(defined($gpu), 'GPU information query returned no rows'); + $gpu =~ s/^\s+|\s+$//g; + my @batches; + for my $count (1, 2, 4) { + push @batches, run_batch( + $output_dir, $xprobe, $agent, $fixture, $count); + } + write_json("$output_dir/report.json", { + schema_version => '2.0', + ok => JSON::PP::true, + gpu => $gpu, + batches => \@batches, + }); +} + +eval { main(); 1 } or do { + my $error = $@ || 'unknown failure'; + chomp $error; + print STDERR "cuda multiprocess benchmark failed: $error\n"; + exit 1; +}; diff --git a/benchmarks/cuda-multiprocess/run.py b/benchmarks/cuda-multiprocess/run.py new file mode 100755 index 0000000..78fd631 --- /dev/null +++ b/benchmarks/cuda-multiprocess/run.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +import json +import os +import pathlib +import subprocess +import sys +import tempfile + + +def dump_failure_artifacts(output: pathlib.Path) -> None: + for path in sorted(output.rglob("*")): + if not path.is_file() or path.suffix not in {".json", ".stderr", ".stdout"}: + continue + sys.stderr.write(f"\n--- {path.relative_to(output)} ---\n") + sys.stderr.write(path.read_text(errors="replace")) + + +def main() -> None: + if len(sys.argv) != 2: + raise SystemExit("usage: run.py ") + + workspace = pathlib.Path(__file__).resolve().parents[2] + with tempfile.TemporaryDirectory(prefix="xprobe-cuda-multiprocess-") as temporary: + output = pathlib.Path(temporary) + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--gpus", + "all", + "--cap-add", + "SYS_PTRACE", + "--security-opt", + "seccomp=unconfined", + "--env", + f"XPROBE_HOST_UID={os.getuid()}", + "--env", + f"XPROBE_HOST_GID={os.getgid()}", + "--volume", + f"{workspace}:/workspace:ro", + "--volume", + f"{output}:/output", + "--workdir", + "/workspace", + sys.argv[1], + "/usr/bin/perl", + "/workspace/benchmarks/cuda-multiprocess/run-container.pl", + "/output", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + dump_failure_artifacts(output) + raise SystemExit(completed.returncode) + + report = json.loads((output / "report.json").read_text()) + if report.get("schema_version") != "2.0" or report.get("ok") is not True: + raise RuntimeError("multiprocess benchmark report is malformed or unsuccessful") + print(json.dumps(report, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/docs/agent-integration.md b/docs/agent-integration.md index 9421eb9..f7832e6 100644 --- a/docs/agent-integration.md +++ b/docs/agent-integration.md @@ -34,6 +34,13 @@ installation with `skills` CLI 1.5.20 in isolated home directories. This pinned test protects released behavior while the documented `skills@1` selector receives compatible path updates. +For multi-process workloads, the Skill selects explicit PID/start-time +identities, inventories a representative worker when homogeneity is supported +by evidence, and asks the agent framework to launch one independent bounded +measurement per selected worker concurrently. Results, warnings, failures, and +artifacts remain per process; xprobe does not add a multi-process command or +claim cross-process causality. + ## Contract test ```bash diff --git a/docs/architecture.md b/docs/architecture.md index 9388b95..ba57631 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -46,7 +46,9 @@ that pair before and after metadata operations and attachment. PID reuse returns the requested process-tree root, and returns only confirmed CUDA context holders with stable process identity, parent PID, command line, and GPU UUIDs. It does not attach and does not choose a worker. That orchestration belongs to -the calling Agent or user. +the calling Agent or user. A multi-process investigation uses one bounded +command per selected PID/start-time identity; concurrency and per-worker failure +handling remain in the caller, and event correlation never crosses processes. Host selector resolution converts ELF virtual addresses through load segments to file offsets, then through `/proc//maps` to runtime addresses. This diff --git a/docs/benchmarks.md b/docs/benchmarks.md index 6af2ace..de13de5 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -43,3 +43,21 @@ target RSS from procfs while the command runs. It deliberately does not gate on short GPU throughput windows: device clock changes can make those samples move opposite to profiler overhead. The JSON output includes every asserted resource and collection value so regressions remain diagnosable. + +Run the agent-orchestrated worker benchmark with: + +```bash +just benchmark-multiprocess +``` + +Fresh homogeneous batches of one, two, and four CUDA workers are discovered by +PID plus start time. One representative worker supplies a broad aggregate +inventory; its selector hints are validated for every worker before independent +spec-based measurements start concurrently. The benchmark rejects identity +changes, missing injection warnings, path reuse, non-overlapping commands, +incomplete captures, and drops. + +Its JSON keeps validation, command timing, quality, artifact metadata, and +baseline-versus-collection throughput per worker. Perturbation is reported +without a fixed pass ratio because shared GPU scheduling and first injection +cost vary by workload and worker count. diff --git a/docs/development.md b/docs/development.md index 300519e..2dcdbf5 100644 --- a/docs/development.md +++ b/docs/development.md @@ -143,6 +143,13 @@ capacity, or activity accounting: just benchmark-aggregate ``` +Run the concurrent worker benchmark after changing multi-process Skill guidance +or per-worker orchestration: + +```bash +just benchmark-multiprocess +``` + See `docs/agent-integration.md` and `docs/benchmarks.md` for the tested contract and benchmark interpretation. diff --git a/justfile b/justfile index 3cbcc00..df60cda 100644 --- a/justfile +++ b/justfile @@ -19,10 +19,12 @@ test: build ctest --test-dir build --output-on-failure python3 tests/agent-contract/test_contract.py target/debug/xprobe python3 tests/agent-contract/test_trace_analysis.py + python3 tests/agent-contract/test_multi_process_workflow.py test-agent-contract: build python3 tests/agent-contract/test_contract.py target/debug/xprobe python3 tests/agent-contract/test_trace_analysis.py + python3 tests/agent-contract/test_multi_process_workflow.py test-skill-install: tests/agent-contract/test_skill_install.sh @@ -72,6 +74,9 @@ benchmark-gpu: benchmark-aggregate: python3 benchmarks/cuda-aggregate/run.py "{{cuda13_devel_image}}" +benchmark-multiprocess: + python3 benchmarks/cuda-multiprocess/run.py "{{cuda13_devel_image}}" + fmt: cargo fmt --all diff --git a/skills/xprobe-measure-latency/SKILL.md b/skills/xprobe-measure-latency/SKILL.md index 338d221..7fba978 100644 --- a/skills/xprobe-measure-latency/SKILL.md +++ b/skills/xprobe-measure-latency/SKILL.md @@ -12,6 +12,8 @@ to install or repair the CLI and this Skill. Read unknown workload. Read [references/result-quality.md](references/result-quality.md) before interpreting correlation, clocks, concurrency, or overhead. The exact CLI and selector syntax is in [references/cli-contract.md](references/cli-contract.md). +For more than one selected process, follow +[references/multi-process.md](references/multi-process.md). ## Workflow @@ -30,9 +32,10 @@ CLI and selector syntax is in [references/cli-contract.md](references/cli-contra 4. For GPU or mixed work, run `xprobe discover --pid ROOT_PID --limit 200 --json --non-interactive --no-color`. It returns NVML-confirmed CUDA context holders under that process tree. Choose a worker from workload, PID/start-time, - command line, and GPU UUID evidence. Measure workers separately when several - ranks are relevant. For CPU-only work, do not run `discover`; continue with - the selected process PID. + command line, and GPU UUID evidence. When several ranks are relevant, retain + every selected PID plus process start time and use the multi-process workflow. + For CPU-only work, do not run `discover`; continue with the selected process + PID. 5. Map GPU or mixed work before choosing a name. Validate broad kernel, memcpy, or memset activity endpoints, then collect one bounded, representative coarse inventory per event family with `measure --aggregate --duration-ms ...`. @@ -40,7 +43,9 @@ CLI and selector syntax is in [references/cli-contract.md](references/cli-contra directly; do not require CUDA or CUPTI. Scope breadth and collection duration are independent: keep the selector broad where an activity inventory exists, choose a duration that covers the workload cycle being diagnosed, and give - `--max-groups` headroom. + `--max-groups` headroom. For defensibly homogeneous workers, inventory one + representative worker and apply its evidence-derived narrow selector to all + selected workers. 6. Use aggregate names, selector hints, counts, duration totals and bounds, and transfer bytes to form one narrow hypothesis. For an exact GPU artifact, run `scripts/analyze_trace.py` and use launch variants, stream distribution, busy @@ -48,14 +53,18 @@ CLI and selector syntax is in [references/cli-contract.md](references/cli-contra [references/trace-analysis.md](references/trace-analysis.md) when interpreting the report. For CPU-only work, use resolved host selectors and result evidence to form the hypothesis instead. -7. Run `xprobe validate --pid WORKER_PID --from SELECTOR --to SELECTOR --match - POLICY --json --non-interactive --no-color`. Stop when `valid` is false. If +7. Run one read-only `xprobe validate` per selected worker. Compare every + response target with the PID plus process start time retained from discovery + before mutation, and stop that worker when `valid` is false. If `agent_activation` is `injection_required`, disclose that `measure` will ptrace the target and leave the CUPTI shared object mapped. Use `policy_recommendation` explicitly; xprobe never changes policy for the caller. 8. Run one bounded `xprobe measure` for that hypothesis. Set samples or duration, timeout, and max-events; write `--events-out` when the capture may need audit - or offline re-correlation. Use `--spec FILE` for a versioned live configuration. + or offline re-correlation. Use a versioned `--spec FILE` containing the stable + target identity. For multiple workers, launch the independent calls + concurrently, preserve per-worker outputs and failures, and never correlate + across process artifacts. 9. Check `status`, matched/unmatched/ambiguous/dropped counts, collection completeness, buffer utilization, clock alignment, estimated error, correlation method/confidence/score, warnings, and every evidence pair. diff --git a/skills/xprobe-measure-latency/references/investigation.md b/skills/xprobe-measure-latency/references/investigation.md index 2e7b645..07ad7ea 100644 --- a/skills/xprobe-measure-latency/references/investigation.md +++ b/skills/xprobe-measure-latency/references/investigation.md @@ -11,8 +11,10 @@ or while repeatable traffic is active. A quiet target commonly produces Keep launcher and worker roles separate. `discover --pid ROOT_PID` reports only CUDA context holders, so a wrapper PID should not be selected unless it appears as a candidate. Map GPU UUID and command/rank metadata to workload ownership. -Measure multiple relevant workers one at a time. After a restart, rerun discovery -and use PID plus procfs start time; never reuse an old PID-only choice. +For multiple relevant workers, follow [multi-process.md](multi-process.md): +inventory a representative worker where homogeneity is defensible, then run +independent narrow captures concurrently. After a restart, rerun discovery and +use PID plus procfs start time; never reuse an old PID-only choice. ## Map broadly before selecting diff --git a/skills/xprobe-measure-latency/references/multi-process.md b/skills/xprobe-measure-latency/references/multi-process.md new file mode 100644 index 0000000..30ad2e7 --- /dev/null +++ b/skills/xprobe-measure-latency/references/multi-process.md @@ -0,0 +1,85 @@ +# Multi-process workloads + +Keep process selection and concurrency in the agent framework. xprobe has no +multi-process command: each `validate` and `measure` remains an independent, +bounded invocation for one stable target identity. + +## Select workers + +Run one `discover` after workload warmup and retain the full candidate records +used for selection. Select workers from application ownership, command/rank +metadata, GPU UUIDs, and the intended experiment. Do not select every candidate +merely because it owns a CUDA context. + +Treat workers as homogeneous only when application configuration and observed +workload evidence support that claim. For a homogeneous rank set, use one +representative worker for each broad aggregate inventory. Derive narrow +selectors from that inventory, then apply those selectors to every selected +worker. For heterogeneous workers, inventory one representative per defensible +class or investigate workers separately. + +Record each selected candidate's exact `target.pid` and +`target.process_start_time`. Name every spec, result, stderr log, and event +artifact with both values. Never reuse an artifact path across workers or +attempts. + +## Revalidate identity + +Create one narrow Measurement Spec per selected worker. Copy its full discovery +identity into `target`; give every spec its own samples or duration, +`timeout_ms`, and `max_events`. Capacities are per worker, never a shared budget. + +Issue one read-only `validate` call per selected worker, concurrently when the +agent framework supports concurrent tool calls. Before any mutation, require +that each validation response is valid and that its `target` exactly matches +the selected PID and `process_start_time`. Preserve each validation JSON and its +warnings. A mismatch, exit, or failed validation removes that worker from the +measurement batch; do not replace it with a newly observed PID. + +When validation reports `injection_required`, disclose and retain that worker's +mutation warning. Do not preload or repeatedly broad-profile all ranks merely +to make injection look uniform. Injection completes before that worker's ARM +window, but its application-level startup perturbation still belongs in the +report. + +## Capture concurrently + +After every selected identity has passed validation, launch one bounded +`xprobe measure --spec WORKER_SPEC` per selected worker using native concurrent +tool calls from the agent framework. Start the calls in one concurrency batch +so their capture windows cover the same controlled request, batch, or iteration +cycle. Do not serialize the commands unless the investigation explicitly needs +different workload windows. + +Keep these outputs separate for every worker: + +- spec with the discovered target identity; +- validation JSON and stderr; +- measurement JSON and stderr; +- exact Event JSONL artifact when requested; +- command exit status and start/end wall timestamps. + +Wait for every command. Do not cancel sibling commands or discard successful +outputs when one worker fails. Preserve injection warnings, configured capacity, +collection quality, and the structured error for each worker. Mark the overall +experiment incomplete when any selected worker fails, while describing usable +single-worker evidence separately. + +## Interpret without merging + +Read every result's status, completeness, drops, utilization, ambiguity, clock, +correlation, and evidence independently. Compare per-worker distributions only +after checking that selectors, workload phase, device assignment, and quality +are comparable. + +Do not concatenate artifacts and rerun correlation across processes. CUPTI +correlation IDs, CUDA contexts, streams, host thread identity, and process +clocks are process-scoped evidence; overlapping wall timestamps do not establish +cross-process causality. Report variation or concurrent observation, not an +exact relationship between worker events. + +Measure an application baseline before the batch and again afterward. Report +per-worker command startup time and workload throughput during collection. +Concurrent captures can perturb shared CPU, GPU, storage, and driver resources, +so a result that changes the workload materially needs a smaller worker set, +narrower selectors, or a shorter but still representative window. diff --git a/skills/xprobe-measure-latency/references/result-quality.md b/skills/xprobe-measure-latency/references/result-quality.md index bb88db6..a7fed6d 100644 --- a/skills/xprobe-measure-latency/references/result-quality.md +++ b/skills/xprobe-measure-latency/references/result-quality.md @@ -31,6 +31,13 @@ for wall-clock busy time and `overlap_factor` to quantify concurrency. Compare per-stream gaps separately. A top kernel's `summed_kernel_time_share` describes its share of summed kernel work, not its exclusive share of wall time. +For multi-process captures, keep every result and artifact scoped to its PID +plus process start time. Compare worker summaries only after checking individual +quality and workload alignment. Do not merge Event JSONL files for correlation: +timestamp overlap, matching names, or matching correlation IDs across processes +does not establish causality. Preserve partial command failures instead of +silently reporting only successful workers. + ## Bounds and completion `completed` means either requested samples or duration was reached. `timed_out` diff --git a/tests/agent-contract/fixtures/multi-process.json b/tests/agent-contract/fixtures/multi-process.json new file mode 100644 index 0000000..4939678 --- /dev/null +++ b/tests/agent-contract/fixtures/multi-process.json @@ -0,0 +1,88 @@ +{ + "schema_version": "2.0", + "homogeneous": true, + "representative": { + "pid": 4101, + "process_start_time": 90001 + }, + "selected": [ + { + "pid": 4101, + "process_start_time": 90001 + }, + { + "pid": 4102, + "process_start_time": 90002 + } + ], + "discovery": { + "candidates": [ + { + "target": { + "pid": 4101, + "process_start_time": 90001 + }, + "command_line": ["python", "worker.py", "--rank", "0"], + "gpu_uuids": ["GPU-fixture"] + }, + { + "target": { + "pid": 4102, + "process_start_time": 90002 + }, + "command_line": ["python", "worker.py", "--rank", "1"], + "gpu_uuids": ["GPU-fixture"] + }, + { + "target": { + "pid": 4199, + "process_start_time": 90199 + }, + "command_line": ["python", "unrelated.py"], + "gpu_uuids": ["GPU-fixture"] + } + ] + }, + "validations": [ + { + "schema_version": "2.0", + "ok": true, + "valid": true, + "target": { + "pid": 4101, + "process_start_time": 90001 + }, + "warnings": [ + { + "code": "TARGET_PROCESS_WILL_BE_MODIFIED", + "message": "fixture mutation warning" + } + ] + }, + { + "schema_version": "2.0", + "ok": true, + "valid": true, + "target": { + "pid": 4102, + "process_start_time": 90002 + }, + "warnings": [ + { + "code": "TARGET_PROCESS_WILL_BE_MODIFIED", + "message": "fixture mutation warning" + } + ] + } + ], + "stale_validation": { + "schema_version": "2.0", + "ok": true, + "valid": true, + "target": { + "pid": 4102, + "process_start_time": 99002 + }, + "warnings": [] + } +} diff --git a/tests/agent-contract/test_contract.py b/tests/agent-contract/test_contract.py index 45b05ec..392a7bc 100755 --- a/tests/agent-contract/test_contract.py +++ b/tests/agent-contract/test_contract.py @@ -118,10 +118,12 @@ def check_skill(workspace: pathlib.Path) -> None: investigation = (skill_root / "references/investigation.md").read_text() quality = (skill_root / "references/result-quality.md").read_text() + multi_process = (skill_root / "references/multi-process.md").read_text() trace_analysis = (skill_root / "references/trace-analysis.md").read_text() setup = (skill_root / "references/setup.md").read_text() normalized_investigation = re.sub(r"\s+", " ", investigation) normalized_quality = re.sub(r"\s+", " ", quality) + normalized_multi_process = re.sub(r"\s+", " ", multi_process) normalized_trace_analysis = re.sub(r"\s+", " ", trace_analysis) normalized_setup = re.sub(r"\s+", " ", setup) for required in ( @@ -143,6 +145,17 @@ def check_skill(workspace: pathlib.Path) -> None: "max-groups", ): assert required in normalized_quality + for required in ( + "representative worker", + "target.process_start_time", + "target` exactly matches", + "native concurrent tool calls", + "Do not cancel sibling commands", + "Never reuse an artifact path", + "Do not concatenate artifacts", + "overall experiment incomplete", + ): + assert required in normalized_multi_process for required in ( "busy_union_ns", "summed_activity_ns", diff --git a/tests/agent-contract/test_multi_process_workflow.py b/tests/agent-contract/test_multi_process_workflow.py new file mode 100755 index 0000000..73e0721 --- /dev/null +++ b/tests/agent-contract/test_multi_process_workflow.py @@ -0,0 +1,186 @@ +#!/usr/bin/env python3 +import json +import pathlib +import subprocess +import sys +import tempfile +import time + + +def identity(target: dict) -> tuple[int, int]: + return target["pid"], target["process_start_time"] + + +def verified_selection(fixture: dict, validations: list[dict]) -> list[dict]: + candidates = { + identity(candidate["target"]): candidate + for candidate in fixture["discovery"]["candidates"] + } + selected = fixture["selected"] + assert len(selected) == len({identity(target) for target in selected}) + assert identity(fixture["representative"]) == identity(selected[0]) + for target, validation in zip(selected, validations, strict=True): + assert identity(target) in candidates + if not validation["valid"] or identity(validation["target"]) != identity(target): + raise ValueError("worker identity changed before measurement") + return selected + + +def fake_measure(spec_path: pathlib.Path, events_path: pathlib.Path) -> int: + spec = json.loads(spec_path.read_text()) + target = spec["target"] + started_ns = time.monotonic_ns() + time.sleep(0.15) + events_path.write_text( + json.dumps( + { + "schema_version": "2.0", + "pid": target["pid"], + "process_start_time": target["process_start_time"], + } + ) + + "\n" + ) + finished_ns = time.monotonic_ns() + if target["pid"] == 4102: + print("fixture worker failed", file=sys.stderr) + print( + json.dumps( + { + "schema_version": "2.0", + "ok": False, + "error": {"code": "FIXTURE_FAILURE"}, + "started_ns": started_ns, + "finished_ns": finished_ns, + } + ) + ) + return 7 + print("fixture injection warning", file=sys.stderr) + print( + json.dumps( + { + "schema_version": "2.0", + "ok": True, + "status": "completed", + "started_ns": started_ns, + "finished_ns": finished_ns, + } + ) + ) + return 0 + + +def run_contract() -> None: + workspace = pathlib.Path(__file__).resolve().parents[2] + fixture = json.loads( + (workspace / "tests/agent-contract/fixtures/multi-process.json").read_text() + ) + selected = verified_selection(fixture, fixture["validations"]) + try: + verified_selection( + fixture, [fixture["validations"][0], fixture["stale_validation"]] + ) + except ValueError: + pass + else: + raise AssertionError("stale worker identity was accepted") + + with tempfile.TemporaryDirectory(prefix="xprobe-multi-process-") as directory: + root = pathlib.Path(directory) + processes = [] + paths = [] + for target in selected: + suffix = f"{target['pid']}-{target['process_start_time']}" + spec_path = root / f"spec-{suffix}.json" + result_path = root / f"result-{suffix}.json" + stderr_path = root / f"stderr-{suffix}.log" + events_path = root / f"events-{suffix}.jsonl" + spec_path.write_text( + json.dumps( + { + "schema_version": "2.0", + "target": target, + "start_selector": "cuda:kernel_start:name~^fixture_kernel$", + "end_selector": "cuda:kernel_end:name~^fixture_kernel$", + "match_policy": "exact", + "samples": 10, + "duration_ms": None, + "timeout_ms": 5000, + "max_events": 100, + "measurement_mode": "exact", + "max_groups": None, + "name": suffix, + } + ) + ) + command = [ + sys.executable, + __file__, + "--fake-measure", + str(spec_path), + str(events_path), + ] + processes.append( + ( + target, + subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ), + result_path, + stderr_path, + ) + ) + paths.extend((spec_path, result_path, stderr_path, events_path)) + + outcomes = [] + for target, process, result_path, stderr_path in processes: + stdout, stderr = process.communicate(timeout=5) + result_path.write_text(stdout) + stderr_path.write_text(stderr) + outcomes.append( + { + "target": target, + "returncode": process.returncode, + "result": json.loads(stdout), + "stderr": stderr, + } + ) + + assert len(paths) == len(set(paths)) + assert {outcome["returncode"] for outcome in outcomes} == {0, 7} + assert any(outcome["result"]["ok"] for outcome in outcomes) + assert any(not outcome["result"]["ok"] for outcome in outcomes) + latest_start = max(outcome["result"]["started_ns"] for outcome in outcomes) + earliest_finish = min(outcome["result"]["finished_ns"] for outcome in outcomes) + assert latest_start < earliest_finish + assert all(path.is_file() for path in paths) + + print( + json.dumps( + { + "schema_version": "2.0", + "ok": True, + "selected": [identity(target) for target in selected], + "representative": identity(fixture["representative"]), + "partial_failure_preserved": True, + }, + sort_keys=True, + ) + ) + + +def main() -> int: + if len(sys.argv) == 4 and sys.argv[1] == "--fake-measure": + return fake_measure(pathlib.Path(sys.argv[2]), pathlib.Path(sys.argv[3])) + if len(sys.argv) != 1: + raise SystemExit(f"usage: {sys.argv[0]}") + run_contract() + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From b3fcf8b66f8cedbfde9963439f775883fafcf51e Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 05:52:09 +0800 Subject: [PATCH 09/17] =?UTF-8?q?=E2=9C=A8=20feat:=20define=20Linux=20even?= =?UTF-8?q?t=20selectors?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- schemas/validate.schema.json | 34 ++++ xprobe/core/src/validate.rs | 209 ++++++++++++++++++++- xprobe/correlator/src/lib.rs | 292 +++++++++++++++++++++++++++-- xprobe/protocol/src/lib.rs | 4 +- xprobe/protocol/src/validate.rs | 14 +- xprobe/protocol/tests/contracts.rs | 2 + 6 files changed, 528 insertions(+), 27 deletions(-) diff --git a/schemas/validate.schema.json b/schemas/validate.schema.json index a4bd0ee..80aa4d6 100644 --- a/schemas/validate.schema.json +++ b/schemas/validate.schema.json @@ -260,6 +260,30 @@ "event_type" ] }, + "ResolvedLinuxSelector": { + "type": "object", + "properties": { + "category": { + "type": "string" + }, + "event_type": { + "$ref": "#/$defs/EventType" + }, + "name": { + "type": "string" + }, + "probe_kind": { + "$ref": "#/$defs/HostProbeKind" + } + }, + "additionalProperties": false, + "required": [ + "event_type", + "probe_kind", + "category", + "name" + ] + }, "ResolvedProbe": { "type": "object", "properties": { @@ -395,6 +419,16 @@ } ] }, + "linux": { + "anyOf": [ + { + "$ref": "#/$defs/ResolvedLinuxSelector" + }, + { + "type": "null" + } + ] + }, "selector": { "type": "string" }, diff --git a/xprobe/core/src/validate.rs b/xprobe/core/src/validate.rs index 618aba2..b01d0ee 100644 --- a/xprobe/core/src/validate.rs +++ b/xprobe/core/src/validate.rs @@ -4,8 +4,8 @@ use regex::Regex; use xprobe_protocol::{ AgentActivation, EndpointSource, ErrorCode, EventType, HostProbeKind, MatchPolicy, MemcpyKind, PolicyRecommendation, PolicyRecommendationReason, ProcessReport, ResolvedCudaSelector, - SchemaVersion, ValidatedEndpoint, ValidationIssue, ValidationRequirements, ValidationResult, - Warning, + ResolvedLinuxSelector, SchemaVersion, ValidatedEndpoint, ValidationIssue, + ValidationRequirements, ValidationResult, Warning, }; use crate::{ @@ -81,6 +81,9 @@ impl From for ValidateError { enum EndpointKind { HostEntry, HostReturn, + SyscallEntry { name: String }, + SyscallExit { name: String }, + Tracepoint, CudaApi { domain: String, name: String }, Kernel, Memcpy, @@ -97,13 +100,25 @@ impl EndpointKind { } const fn is_host(&self) -> bool { - matches!(self, Self::HostEntry | Self::HostReturn) + matches!( + self, + Self::HostEntry + | Self::HostReturn + | Self::SyscallEntry { .. } + | Self::SyscallExit { .. } + | Self::Tracepoint + ) } const fn uses_host_clock(&self) -> bool { matches!( self, - Self::HostEntry | Self::HostReturn | Self::CudaApi { .. } + Self::HostEntry + | Self::HostReturn + | Self::SyscallEntry { .. } + | Self::SyscallExit { .. } + | Self::Tracepoint + | Self::CudaApi { .. } ) } } @@ -267,6 +282,22 @@ fn resolve_endpoint( event_type, collectable, host: Some(host), + linux: None, + cuda: None, + }, + kind, + }); + } + if selector.starts_with("syscall:") || selector.starts_with("tracepoint:") { + let (linux, kind) = parse_linux_selector(selector)?; + return Ok(Endpoint { + public: ValidatedEndpoint { + selector: selector.to_owned(), + source: EndpointSource::Host, + event_type: linux.event_type.clone(), + collectable: true, + host: None, + linux: Some(linux), cuda: None, }, kind, @@ -281,19 +312,87 @@ fn resolve_endpoint( event_type: cuda.event_type.clone(), collectable, host: None, + linux: None, cuda: Some(cuda), }, kind, }) } +fn parse_linux_selector( + selector: &str, +) -> Result<(ResolvedLinuxSelector, EndpointKind), ValidateError> { + let fields = selector.split(':').collect::>(); + match fields.as_slice() { + ["syscall", name, boundary] if valid_tracepoint_component(name) => { + let (event_type, kind) = match *boundary { + "entry" => ( + EventType::SyscallEntry, + EndpointKind::SyscallEntry { + name: (*name).to_owned(), + }, + ), + "exit" => ( + EventType::SyscallExit, + EndpointKind::SyscallExit { + name: (*name).to_owned(), + }, + ), + _ => { + return Err(ValidateError::InvalidSelector( + "syscall boundary must be entry or exit".to_owned(), + )); + } + }; + Ok(( + ResolvedLinuxSelector { + event_type, + probe_kind: HostProbeKind::Syscall, + category: "syscalls".to_owned(), + name: (*name).to_owned(), + }, + kind, + )) + } + ["syscall", ..] => Err(ValidateError::InvalidSelector( + "syscall selector must be syscall::".to_owned(), + )), + ["tracepoint", category, name] + if valid_tracepoint_component(category) && valid_tracepoint_component(name) => + { + Ok(( + ResolvedLinuxSelector { + event_type: EventType::Tracepoint, + probe_kind: HostProbeKind::Tracepoint, + category: (*category).to_owned(), + name: (*name).to_owned(), + }, + EndpointKind::Tracepoint, + )) + } + ["tracepoint", ..] => Err(ValidateError::InvalidSelector( + "tracepoint selector must be tracepoint::".to_owned(), + )), + _ => Err(ValidateError::InvalidSelector( + "expected uprobe:, syscall:, tracepoint:, or cuda: prefix".to_owned(), + )), + } +} + +fn valid_tracepoint_component(value: &str) -> bool { + !value.is_empty() + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_') +} + fn parse_cuda_selector( selector: &str, ) -> Result<(ResolvedCudaSelector, EndpointKind, bool), ValidateError> { let fields: Vec<&str> = selector.splitn(3, ':').collect(); if fields.first() != Some(&"cuda") || fields.len() < 2 { return Err(ValidateError::InvalidSelector( - "expected uprobe: or cuda: prefix".to_owned(), + "expected uprobe:, syscall:, tracepoint:, or cuda: prefix".to_owned(), )); } @@ -491,12 +590,24 @@ fn check_capabilities( || matches!(end.kind, EndpointKind::HostEntry); let needs_uretprobe = matches!(start.kind, EndpointKind::HostReturn) || matches!(end.kind, EndpointKind::HostReturn); + let needs_tracepoint = matches!( + start.kind, + EndpointKind::SyscallEntry { .. } + | EndpointKind::SyscallExit { .. } + | EndpointKind::Tracepoint + ) || matches!( + end.kind, + EndpointKind::SyscallEntry { .. } + | EndpointKind::SyscallExit { .. } + | EndpointKind::Tracepoint + ); if (needs_uprobe && !report.capabilities.uprobe) || (needs_uretprobe && !report.capabilities.uretprobe) + || (needs_tracepoint && !report.capabilities.tracepoint) { issues.push(issue( ErrorCode::PermissionDenied, - "required eBPF userspace probe capability is unavailable".to_owned(), + "required eBPF probe capability is unavailable".to_owned(), )); } } @@ -572,6 +683,10 @@ fn supports_exact(start: &EndpointKind, end: &EndpointKind) -> bool { (EndpointKind::Kernel, EndpointKind::Kernel) | (EndpointKind::Memcpy, EndpointKind::Memcpy) | (EndpointKind::Memset, EndpointKind::Memset) => true, + ( + EndpointKind::SyscallEntry { name: start_name }, + EndpointKind::SyscallExit { name: end_name }, + ) => start_name == end_name, (EndpointKind::CudaApi { domain, name }, EndpointKind::Kernel) | (EndpointKind::Kernel, EndpointKind::CudaApi { domain, name }) => { (domain == "runtime_api" && name == "cudaLaunchKernel") @@ -635,12 +750,13 @@ fn warning(code: &str, message: &str) -> Warning { mod tests { use xprobe_protocol::{ AgentActivation, ElfObjectKind, EndpointSource, EventType, HostProbeKind, MatchPolicy, - MemcpyKind, PolicyRecommendationReason, ProcessMapping, ResolvedProbe, SchemaVersion, - TargetIdentity, ValidatedEndpoint, + MemcpyKind, PolicyRecommendationReason, ProcessMapping, ResolvedLinuxSelector, + ResolvedProbe, SchemaVersion, TargetIdentity, ValidatedEndpoint, }; use super::{ - Endpoint, EndpointKind, parse_cuda_selector, parse_match_policy, recommend_policy, run, + Endpoint, EndpointKind, parse_cuda_selector, parse_linux_selector, parse_match_policy, + recommend_policy, run, }; use crate::inspect; @@ -679,6 +795,80 @@ mod tests { assert!(parse_match_policy("probably-near").is_err()); } + #[test] + fn parses_named_syscalls_and_tracepoints() { + let (syscall, kind) = parse_linux_selector("syscall:mmap:entry").unwrap(); + assert_eq!(syscall.event_type, EventType::SyscallEntry); + assert_eq!(syscall.probe_kind, HostProbeKind::Syscall); + assert_eq!(syscall.category, "syscalls"); + assert_eq!(syscall.name, "mmap"); + assert_eq!( + kind, + EndpointKind::SyscallEntry { + name: "mmap".to_owned() + } + ); + + let (tracepoint, kind) = parse_linux_selector("tracepoint:sched:sched_switch").unwrap(); + assert_eq!(tracepoint.event_type, EventType::Tracepoint); + assert_eq!(tracepoint.probe_kind, HostProbeKind::Tracepoint); + assert_eq!(tracepoint.category, "sched"); + assert_eq!(tracepoint.name, "sched_switch"); + assert_eq!(kind, EndpointKind::Tracepoint); + + assert!(parse_linux_selector("syscall:mmap:return").is_err()); + assert!(parse_linux_selector("tracepoint:sched:sched-switch").is_err()); + } + + #[test] + fn recommends_exact_matching_for_one_syscall_lifecycle() { + let endpoint = |selector: &str, event_type, kind, linux| Endpoint { + public: ValidatedEndpoint { + selector: selector.to_owned(), + source: EndpointSource::Host, + event_type, + collectable: true, + host: None, + linux: Some(linux), + cuda: None, + }, + kind, + }; + let start = endpoint( + "syscall:mmap:entry", + EventType::SyscallEntry, + EndpointKind::SyscallEntry { + name: "mmap".to_owned(), + }, + ResolvedLinuxSelector { + event_type: EventType::SyscallEntry, + probe_kind: HostProbeKind::Syscall, + category: "syscalls".to_owned(), + name: "mmap".to_owned(), + }, + ); + let end = endpoint( + "syscall:mmap:exit", + EventType::SyscallExit, + EndpointKind::SyscallExit { + name: "mmap".to_owned(), + }, + ResolvedLinuxSelector { + event_type: EventType::SyscallExit, + probe_kind: HostProbeKind::Syscall, + category: "syscalls".to_owned(), + name: "mmap".to_owned(), + }, + ); + + let recommendation = recommend_policy(&start, &end); + assert_eq!(recommendation.policy, MatchPolicy::Exact); + assert_eq!( + recommendation.reason, + PolicyRecommendationReason::DeterministicCorrelationKey + ); + } + #[test] fn reports_missing_cupti_without_failing_validation() { let report = inspect::run(std::process::id()).unwrap(); @@ -763,6 +953,7 @@ mod tests { file_offset: 0, }, }), + linux: None, cuda: None, }, kind, diff --git a/xprobe/correlator/src/lib.rs b/xprobe/correlator/src/lib.rs index 46ae9be..31674f1 100644 --- a/xprobe/correlator/src/lib.rs +++ b/xprobe/correlator/src/lib.rs @@ -97,6 +97,14 @@ enum Selector { binary_path: String, target: HostTarget, }, + Syscall { + event_type: EventType, + name: String, + }, + Tracepoint { + category: String, + name: String, + }, Api { event_type: EventType, domain: String, @@ -126,10 +134,17 @@ impl Selector { if text.starts_with("uprobe:") { return Self::parse_host(text); } + if text.starts_with("syscall:") { + return Self::parse_syscall(text); + } + if text.starts_with("tracepoint:") { + return Self::parse_tracepoint(text); + } let fields: Vec<&str> = text.splitn(3, ':').collect(); if fields.first() != Some(&"cuda") || fields.len() < 2 { return Err(MeasureError::InvalidSelector( - "completed captures require a uprobe: or cuda: selector".to_owned(), + "completed captures require a uprobe:, syscall:, tracepoint:, or cuda: selector" + .to_owned(), )); } match fields[1] { @@ -215,6 +230,52 @@ impl Selector { }) } + fn parse_syscall(text: &str) -> Result { + let fields = text.split(':').collect::>(); + let ["syscall", name, boundary] = fields.as_slice() else { + return Err(MeasureError::InvalidSelector( + "syscall selector must be syscall::".to_owned(), + )); + }; + if !valid_linux_component(name) { + return Err(MeasureError::InvalidSelector( + "syscall name must contain only letters, digits, and underscores".to_owned(), + )); + } + let event_type = match *boundary { + "entry" => EventType::SyscallEntry, + "exit" => EventType::SyscallExit, + _ => { + return Err(MeasureError::InvalidSelector( + "syscall boundary must be entry or exit".to_owned(), + )); + } + }; + Ok(Self::Syscall { + event_type, + name: (*name).to_owned(), + }) + } + + fn parse_tracepoint(text: &str) -> Result { + let fields = text.split(':').collect::>(); + let ["tracepoint", category, name] = fields.as_slice() else { + return Err(MeasureError::InvalidSelector( + "tracepoint selector must be tracepoint::".to_owned(), + )); + }; + if !valid_linux_component(category) || !valid_linux_component(name) { + return Err(MeasureError::InvalidSelector( + "tracepoint category and name must contain only letters, digits, and underscores" + .to_owned(), + )); + } + Ok(Self::Tracepoint { + category: (*category).to_owned(), + name: (*name).to_owned(), + }) + } + fn parse_kernel(fields: &[&str]) -> Result { let event_type = match fields[1] { "kernel_start" => EventType::GpuKernelStart, @@ -336,6 +397,25 @@ impl Selector { .is_some_and(|kernel| pattern.is_match(kernel)) }) } + Self::Syscall { event_type, name } => { + event.event_type == *event_type + && event.host.as_ref().is_some_and(|host| { + host.probe_kind == HostProbeKind::Syscall + && host.symbol.as_deref() == Some(name.as_str()) + }) + } + Self::Tracepoint { category, name } => { + event.event_type == EventType::Tracepoint + && event.host.as_ref().is_some_and(|host| { + host.probe_kind == HostProbeKind::Tracepoint + && host.symbol.as_deref() == Some(name.as_str()) + }) + && event + .attributes + .get("tracepoint_category") + .and_then(serde_json::Value::as_str) + == Some(category.as_str()) + } Self::Memcpy { event_type, kind } => { event.event_type == *event_type && kind.as_ref().is_none_or(|kind| { @@ -350,8 +430,38 @@ impl Selector { } } - const fn supports_exact(&self) -> bool { - !matches!(self, Self::Host { .. }) + fn supports_exact(&self, end: &Self) -> bool { + match (self, end) { + ( + Self::Syscall { + event_type: EventType::SyscallEntry, + name: start_name, + }, + Self::Syscall { + event_type: EventType::SyscallExit, + name: end_name, + }, + ) => start_name == end_name, + (Self::Host { .. } | Self::Tracepoint { .. }, _) + | (_, Self::Host { .. } | Self::Tracepoint { .. }) => false, + _ => true, + } + } + + fn is_syscall_lifecycle(&self, end: &Self) -> bool { + matches!( + (self, end), + ( + Self::Syscall { + event_type: EventType::SyscallEntry, + name: start_name, + }, + Self::Syscall { + event_type: EventType::SyscallExit, + name: end_name, + }, + ) if start_name == end_name + ) } fn supports_stack_nested(&self, end: &Self) -> bool { @@ -382,6 +492,13 @@ impl Selector { } } +fn valid_linux_component(value: &str) -> bool { + !value.is_empty() + && value + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_') +} + #[derive(Debug)] struct Pair<'a> { start: &'a Event, @@ -435,13 +552,14 @@ pub fn measure( let clock_domain = common_clock_domain(&starts, &ends)?; apply_duration_window(&mut starts, &mut ends, options.duration)?; - let outcome = match options.match_policy { - MatchPolicy::Exact => correlate_exact(&starts, &ends, options.samples), - MatchPolicy::FirstAfter => correlate_first_after(&starts, &ends, options.samples), - MatchPolicy::Nearest => correlate_nearest(&starts, &ends, options.samples), - MatchPolicy::StackNested => correlate_stack_nested(&starts, &ends, options.samples), - MatchPolicy::StreamOrder => correlate_stream_order(&starts, &ends, options.samples), - }; + let syscall_lifecycle = start_selector.is_syscall_lifecycle(&end_selector); + let outcome = correlate_selected( + &starts, + &ends, + options.match_policy, + options.samples, + syscall_lifecycle, + ); if outcome.pairs.is_empty() { return Err(MeasureError::NoMatchedSamples); } @@ -453,7 +571,7 @@ pub fn measure( .map(|pair| pair.latency_ns) .collect::>(); let denominator = matched + outcome.unmatched_start + outcome.unmatched_end + outcome.ambiguous; - let (method, confidence) = correlation_metadata(options.match_policy); + let (method, confidence) = correlation_metadata(options.match_policy, syscall_lifecycle); let warnings = measurement_warnings(options, &starts, &ends); let host_events = events @@ -511,14 +629,31 @@ pub fn measure( }) } +fn correlate_selected<'a>( + starts: &[&'a Event], + ends: &[&'a Event], + policy: MatchPolicy, + limit: Option, + syscall_lifecycle: bool, +) -> Outcome<'a> { + match policy { + MatchPolicy::Exact if syscall_lifecycle => correlate_syscall_lifecycle(starts, ends, limit), + MatchPolicy::Exact => correlate_exact(starts, ends, limit), + MatchPolicy::FirstAfter => correlate_first_after(starts, ends, limit), + MatchPolicy::Nearest => correlate_nearest(starts, ends, limit), + MatchPolicy::StackNested => correlate_stack_nested(starts, ends, limit), + MatchPolicy::StreamOrder => correlate_stream_order(starts, ends, limit), + } +} + fn validate_policy( start: &Selector, end: &Selector, policy: MatchPolicy, ) -> Result<(), MeasureError> { let message = match policy { - MatchPolicy::Exact if !start.supports_exact() || !end.supports_exact() => { - Some("exact matching requires CUDA endpoints with a deterministic correlation ID") + MatchPolicy::Exact if !start.supports_exact(end) => { + Some("exact matching requires endpoints with a deterministic correlation key") } MatchPolicy::StackNested if !start.supports_stack_nested(end) => { Some("stack-nested requires entry and return selectors for the same host function") @@ -534,8 +669,14 @@ fn validate_policy( Ok(()) } -const fn correlation_metadata(policy: MatchPolicy) -> (&'static str, CorrelationConfidence) { +const fn correlation_metadata( + policy: MatchPolicy, + syscall_lifecycle: bool, +) -> (&'static str, CorrelationConfidence) { match policy { + MatchPolicy::Exact if syscall_lifecycle => { + ("exact_syscall_tid_lifecycle", CorrelationConfidence::Exact) + } MatchPolicy::Exact => ("exact_cupti_correlation_id", CorrelationConfidence::Exact), MatchPolicy::FirstAfter => ("first_after", CorrelationConfidence::Heuristic), MatchPolicy::Nearest => ("nearest", CorrelationConfidence::Heuristic), @@ -544,6 +685,54 @@ const fn correlation_metadata(policy: MatchPolicy) -> (&'static str, Correlation } } +fn correlate_syscall_lifecycle<'a>( + starts: &[&'a Event], + ends: &[&'a Event], + limit: Option, +) -> Outcome<'a> { + let mut timeline = starts + .iter() + .map(|event| (event.timestamp_ns, 0_u8, *event)) + .chain(ends.iter().map(|event| (event.timestamp_ns, 1_u8, *event))) + .collect::>(); + timeline.sort_by_key(|(timestamp, boundary, event)| (*timestamp, *boundary, event.sequence)); + + let mut pending = BTreeMap::<(u32, u32, Option), &Event>::new(); + let mut pairs = Vec::new(); + let mut unmatched_end = 0_u64; + let mut ambiguous = 0_u64; + for (_, boundary, event) in timeline { + let key = (event.pid, event.tid, event.process_start_time); + if boundary == 0 { + if pending.insert(key, event).is_some() { + ambiguous += 1; + } + } else if let Some(start) = pending.remove(&key) { + if let Some(latency_ns) = event.timestamp_ns.checked_sub(start.timestamp_ns) { + pairs.push(Pair { + start, + end: event, + latency_ns, + }); + } else { + ambiguous += 1; + } + } else { + unmatched_end += 1; + } + } + pairs.sort_by_key(|pair| pair.start.timestamp_ns); + if let Some(limit) = limit { + pairs.truncate(limit); + } + Outcome { + pairs, + unmatched_start: pending.len() as u64, + unmatched_end, + ambiguous, + } +} + fn validate_options(options: &MeasureOptions) -> Result<(), MeasureError> { if options.samples.is_none() && options.duration.is_none() { return Err(MeasureError::InvalidLimit( @@ -1097,6 +1286,25 @@ mod tests { event } + fn syscall_event(event_type: EventType, timestamp: u64, tid: u32, name: &str) -> Event { + let mut event = event(event_type, timestamp, 0); + event.source = EventSource::Ebpf; + event.clock_domain = ClockDomain::HostMonotonic; + event.tid = tid; + event.process_start_time = Some(42); + event.cuda = None; + event.host = Some(HostEvent { + probe_kind: HostProbeKind::Syscall, + binary_path: None, + build_id: None, + symbol: Some(name.to_owned()), + offset: None, + return_value: (event.event_type == EventType::SyscallExit).then_some(0), + arguments: Vec::new(), + }); + event + } + #[test] fn exact_matching_uses_cupti_correlation_ids() { let events = vec![ @@ -1127,6 +1335,62 @@ mod tests { assert_eq!(result.measurement.latency_ns.max, 80); } + #[test] + fn exact_syscall_matching_follows_thread_lifecycle() { + let events = vec![ + syscall_event(EventType::SyscallEntry, 100, 10, "mmap"), + syscall_event(EventType::SyscallEntry, 110, 20, "mmap"), + syscall_event(EventType::SyscallExit, 150, 20, "mmap"), + syscall_event(EventType::SyscallExit, 180, 10, "mmap"), + ]; + let options = MeasureOptions { + session_id: "xp_syscall".to_owned(), + name: Some("mmap".to_owned()), + start_selector: "syscall:mmap:entry".to_owned(), + end_selector: "syscall:mmap:exit".to_owned(), + match_policy: MatchPolicy::Exact, + samples: Some(2), + duration: None, + max_events: 100, + dropped_events: 0, + }; + + let result = measure(&events, &options).unwrap(); + assert_eq!(result.measurement.samples.matched, 2); + assert_eq!(result.measurement.latency_ns.min, 40); + assert_eq!(result.measurement.latency_ns.max, 80); + assert_eq!(result.correlation.method, "exact_syscall_tid_lifecycle"); + assert_eq!(result.correlation.confidence, CorrelationConfidence::Exact); + assert_eq!(result.evidence[0].start.tid, 10); + assert_eq!(result.evidence[0].end.tid, 10); + } + + #[test] + fn exact_syscall_matching_does_not_cross_process_identity() { + let mut exit = syscall_event(EventType::SyscallExit, 150, 10, "munmap"); + exit.process_start_time = Some(43); + let events = vec![ + syscall_event(EventType::SyscallEntry, 100, 10, "munmap"), + exit, + ]; + let options = MeasureOptions { + session_id: "xp_syscall".to_owned(), + name: None, + start_selector: "syscall:munmap:entry".to_owned(), + end_selector: "syscall:munmap:exit".to_owned(), + match_policy: MatchPolicy::Exact, + samples: Some(1), + duration: None, + max_events: 100, + dropped_events: 0, + }; + + assert_eq!( + measure(&events, &options).unwrap_err(), + MeasureError::NoMatchedSamples + ); + } + #[test] fn exact_matching_never_pairs_across_cuda_contexts() { let mut events = vec![ diff --git a/xprobe/protocol/src/lib.rs b/xprobe/protocol/src/lib.rs index deb536a..99533f1 100644 --- a/xprobe/protocol/src/lib.rs +++ b/xprobe/protocol/src/lib.rs @@ -35,8 +35,8 @@ pub use process::{CgroupEntry, ProcessCredentials, ProcessCudaState, ProcessRepo pub use resolve::{ElfObjectKind, ProcessMapping, ResolvedProbe}; pub use validate::{ AgentActivation, EndpointSource, PolicyRecommendation, PolicyRecommendationReason, - ResolvedCudaSelector, ValidatedEndpoint, ValidationIssue, ValidationRequirements, - ValidationResult, + ResolvedCudaSelector, ResolvedLinuxSelector, ValidatedEndpoint, ValidationIssue, + ValidationRequirements, ValidationResult, }; pub use version::SchemaVersion; diff --git a/xprobe/protocol/src/validate.rs b/xprobe/protocol/src/validate.rs index 138d05b..90bec15 100644 --- a/xprobe/protocol/src/validate.rs +++ b/xprobe/protocol/src/validate.rs @@ -2,8 +2,8 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use crate::{ - ErrorCode, EventType, MatchPolicy, MemcpyKind, ResolvedProbe, SchemaVersion, TargetIdentity, - Warning, + ErrorCode, EventType, HostProbeKind, MatchPolicy, MemcpyKind, ResolvedProbe, SchemaVersion, + TargetIdentity, Warning, }; #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -31,6 +31,15 @@ pub struct ResolvedCudaSelector { pub memcpy_kind: Option, } +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct ResolvedLinuxSelector { + pub event_type: EventType, + pub probe_kind: HostProbeKind, + pub category: String, + pub name: String, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct ValidatedEndpoint { @@ -39,6 +48,7 @@ pub struct ValidatedEndpoint { pub event_type: EventType, pub collectable: bool, pub host: Option, + pub linux: Option, pub cuda: Option, } diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index 025d0f1..3257955 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -361,6 +361,7 @@ fn validation_result_contract_round_trips() { "event_type": "cuda_api_exit", "collectable": true, "host": null, + "linux": null, "cuda": { "event_type": "cuda_api_exit", "api_domain": "runtime_api", @@ -375,6 +376,7 @@ fn validation_result_contract_round_trips() { "event_type": "gpu_kernel_start", "collectable": true, "host": null, + "linux": null, "cuda": { "event_type": "gpu_kernel_start", "api_domain": null, From 9ea0594aa492c035199c81cc51a5c81e851473e7 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 06:20:43 +0800 Subject: [PATCH 10/17] =?UTF-8?q?=E2=9C=A8=20feat:=20collect=20bounded=20L?= =?UTF-8?q?inux=20events?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- bpf/xprobe.bpf.c | 223 +++++++++ justfile | 1 + schemas/host-capture.schema.json | 4 + schemas/validate.schema.json | 8 + tests/CMakeLists.txt | 3 + tests/fixtures/syscall_target.c | 16 + tests/integration/run-linux-live.sh | 74 +++ tests/integration/test_linux.py | 104 +++++ xprobe/cli/src/main.rs | 105 ++++- xprobe/cli/tests/measure.rs | 1 + xprobe/collector/src/completed.rs | 3 +- xprobe/collector/src/lib.rs | 1 + xprobe/collector/src/linux.rs | 675 ++++++++++++++++++++++++++++ xprobe/collector/src/uprobe.rs | 1 + xprobe/core/src/validate.rs | 116 +++++ xprobe/protocol/src/capture.rs | 2 + xprobe/protocol/src/validate.rs | 1 + xprobe/protocol/tests/contracts.rs | 1 + 18 files changed, 1316 insertions(+), 23 deletions(-) create mode 100644 tests/fixtures/syscall_target.c create mode 100755 tests/integration/run-linux-live.sh create mode 100755 tests/integration/test_linux.py create mode 100644 xprobe/collector/src/linux.rs diff --git a/bpf/xprobe.bpf.c b/bpf/xprobe.bpf.c index b9c4a23..bafa1ec 100644 --- a/bpf/xprobe.bpf.c +++ b/bpf/xprobe.bpf.c @@ -8,6 +8,9 @@ static void *(*bpf_map_lookup_elem)(void *map, const void *key) = (void *)BPF_FUNC_map_lookup_elem; static __u64 (*bpf_ktime_get_ns)(void) = (void *)BPF_FUNC_ktime_get_ns; static __u32 (*bpf_get_smp_processor_id)(void) = (void *)BPF_FUNC_get_smp_processor_id; +static long (*bpf_probe_read_kernel)(void *dst, __u32 size, + const void *unsafe_ptr) = + (void *)BPF_FUNC_probe_read_kernel; static long (*bpf_get_ns_current_pid_tgid)(__u64 dev, __u64 ino, struct bpf_pidns_info *nsdata, __u32 size) = @@ -32,6 +35,53 @@ struct xprobe_event { __u32 probe_id; }; +struct xprobe_linux_config { + __u64 pidns_dev; + __u64 pidns_ino; + __u32 target_pid; + __u32 reserved; + __s64 syscall_entry_numbers[2]; + __s64 syscall_exit_numbers[2]; +}; + +struct xprobe_linux_event { + __u64 timestamp_ns; + __u64 sequence; + __u64 values[6]; + __u32 pid; + __u32 tid; + __u32 cpu; + __u32 probe_id; +}; + +struct xprobe_raw_tracepoint_context { + __u64 arguments[2]; +}; + +struct xprobe_x86_64_registers { + __u64 r15; + __u64 r14; + __u64 r13; + __u64 r12; + __u64 bp; + __u64 bx; + __u64 r11; + __u64 r10; + __u64 r9; + __u64 r8; + __u64 ax; + __u64 cx; + __u64 dx; + __u64 si; + __u64 di; + __u64 orig_ax; + __u64 ip; + __u64 cs; + __u64 flags; + __u64 sp; + __u64 ss; +}; + struct { __uint(type, BPF_MAP_TYPE_ARRAY); __uint(max_entries, 1); @@ -58,6 +108,32 @@ struct { __uint(max_entries, 256 * 1024); } events SEC(".maps"); +struct { + __uint(type, BPF_MAP_TYPE_ARRAY); + __uint(max_entries, 1); + __type(key, __u32); + __type(value, struct xprobe_linux_config); +} linux_config SEC(".maps"); + +struct { + __uint(type, BPF_MAP_TYPE_ARRAY); + __uint(max_entries, 1); + __type(key, __u32); + __type(value, __u64); +} linux_sequence SEC(".maps"); + +struct { + __uint(type, BPF_MAP_TYPE_ARRAY); + __uint(max_entries, 1); + __type(key, __u32); + __type(value, __u64); +} linux_dropped SEC(".maps"); + +struct { + __uint(type, BPF_MAP_TYPE_RINGBUF); + __uint(max_entries, 256 * 1024); +} linux_events SEC(".maps"); + SEC("uprobe") int xprobe_handle_uprobe(void *context) { @@ -94,4 +170,151 @@ int xprobe_handle_uprobe(void *context) return 0; } +static __attribute__((always_inline)) int +xprobe_emit_linux_event(__u32 probe_id, const __u64 *values) +{ + __u32 key = 0; + struct xprobe_linux_config *current = + bpf_map_lookup_elem(&linux_config, &key); + struct bpf_pidns_info nsdata = {}; + struct xprobe_linux_event *event; + __u64 *counter; + int index; + + if (!current || !current->reserved || + bpf_get_ns_current_pid_tgid(current->pidns_dev, current->pidns_ino, + &nsdata, sizeof(nsdata)) != 0 || + current->target_pid != nsdata.tgid) + return 0; + + event = bpf_ringbuf_reserve(&linux_events, sizeof(*event), 0); + if (!event) { + counter = bpf_map_lookup_elem(&linux_dropped, &key); + if (counter) + __sync_fetch_and_add(counter, 1); + return 0; + } + + counter = bpf_map_lookup_elem(&linux_sequence, &key); + event->timestamp_ns = bpf_ktime_get_ns(); + event->sequence = counter ? __sync_fetch_and_add(counter, 1) + 1 : 0; +#pragma unroll + for (index = 0; index < 6; index++) + event->values[index] = values ? values[index] : 0; + event->pid = nsdata.tgid; + event->tid = nsdata.pid; + event->cpu = bpf_get_smp_processor_id(); + event->probe_id = probe_id; + bpf_ringbuf_submit(event, 0); + return 0; +} + +#define XPROBE_TRACEPOINT_PROGRAM(slot) \ + SEC("tracepoint") \ + int xprobe_handle_tracepoint_##slot(void *context) \ + { \ + (void)context; \ + return xprobe_emit_linux_event(slot, 0); \ + } + +XPROBE_TRACEPOINT_PROGRAM(1) +XPROBE_TRACEPOINT_PROGRAM(2) + +#define XPROBE_RAW_TRACEPOINT_PROGRAM(slot) \ + SEC("raw_tracepoint") \ + int xprobe_handle_raw_tracepoint_##slot(void *context) \ + { \ + (void)context; \ + return xprobe_emit_linux_event(slot, 0); \ + } + +XPROBE_RAW_TRACEPOINT_PROGRAM(1) +XPROBE_RAW_TRACEPOINT_PROGRAM(2) + +static __attribute__((always_inline)) __u32 +xprobe_syscall_probe_id(const __s64 numbers[2], __s64 syscall_number) +{ + if (numbers[0] == syscall_number) + return 1; + if (numbers[1] == syscall_number) + return 2; + return 0; +} + +static __attribute__((always_inline)) int +xprobe_target_linux_process(struct xprobe_linux_config *current) +{ + struct bpf_pidns_info nsdata = {}; + + return current && current->reserved && + bpf_get_ns_current_pid_tgid(current->pidns_dev, current->pidns_ino, + &nsdata, sizeof(nsdata)) == 0 && + current->target_pid == nsdata.tgid; +} + +static __attribute__((always_inline)) void xprobe_count_linux_drop(void) +{ + __u32 key = 0; + __u64 *counter = bpf_map_lookup_elem(&linux_dropped, &key); + + if (counter) + __sync_fetch_and_add(counter, 1); +} + +SEC("raw_tracepoint") +int xprobe_handle_syscall_entry(struct xprobe_raw_tracepoint_context *context) +{ + __u32 key = 0; + struct xprobe_linux_config *current = + bpf_map_lookup_elem(&linux_config, &key); + struct xprobe_x86_64_registers registers; + __u64 values[6]; + __s64 syscall_number = (__s64)context->arguments[1]; + __u32 probe_id; + + if (!xprobe_target_linux_process(current)) + return 0; + probe_id = xprobe_syscall_probe_id(current->syscall_entry_numbers, + syscall_number); + if (!probe_id) + return 0; + if (bpf_probe_read_kernel(®isters, sizeof(registers), + (void *)context->arguments[0]) != 0) { + xprobe_count_linux_drop(); + return 0; + } + values[0] = registers.di; + values[1] = registers.si; + values[2] = registers.dx; + values[3] = registers.r10; + values[4] = registers.r8; + values[5] = registers.r9; + return xprobe_emit_linux_event(probe_id, values); +} + +SEC("raw_tracepoint") +int xprobe_handle_syscall_exit(struct xprobe_raw_tracepoint_context *context) +{ + __u32 key = 0; + struct xprobe_linux_config *current = + bpf_map_lookup_elem(&linux_config, &key); + struct xprobe_x86_64_registers registers; + __u64 values[6] = {}; + __u32 probe_id; + + if (!xprobe_target_linux_process(current)) + return 0; + if (bpf_probe_read_kernel(®isters, sizeof(registers), + (void *)context->arguments[0]) != 0) { + xprobe_count_linux_drop(); + return 0; + } + probe_id = xprobe_syscall_probe_id(current->syscall_exit_numbers, + (__s64)registers.orig_ax); + if (!probe_id) + return 0; + values[0] = context->arguments[1]; + return xprobe_emit_linux_event(probe_id, values); +} + char _license[] SEC("license") = "GPL"; diff --git a/justfile b/justfile index df60cda..2c59437 100644 --- a/justfile +++ b/justfile @@ -41,6 +41,7 @@ test-bpf: test-bpf-live: build python3 tests/integration/test_uprobe.py "{{cuda_smoke_image}}" + python3 tests/integration/test_linux.py "{{cuda_smoke_image}}" test-cupti: cmake -S . -B build -G Ninja -DXPROBE_BUILD_CUPTI=ON diff --git a/schemas/host-capture.schema.json b/schemas/host-capture.schema.json index 5d36818..7857efd 100644 --- a/schemas/host-capture.schema.json +++ b/schemas/host-capture.schema.json @@ -27,6 +27,10 @@ "format": "uint32", "minimum": 0 }, + "record_limit_reached": { + "type": "boolean", + "default": false + }, "schema_version": { "$ref": "#/$defs/SchemaVersion" }, diff --git a/schemas/validate.schema.json b/schemas/validate.schema.json index 80aa4d6..b7a5b2c 100644 --- a/schemas/validate.schema.json +++ b/schemas/validate.schema.json @@ -274,6 +274,14 @@ }, "probe_kind": { "$ref": "#/$defs/HostProbeKind" + }, + "syscall_number": { + "type": [ + "integer", + "null" + ], + "format": "uint32", + "minimum": 0 } }, "additionalProperties": false, diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 43fd76a..aa34a7e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -1,2 +1,5 @@ add_executable(xprobe-uprobe-target fixtures/uprobe_target.c) target_compile_options(xprobe-uprobe-target PRIVATE -O0 -Wall -Wextra -Werror) + +add_executable(xprobe-syscall-target fixtures/syscall_target.c) +target_compile_options(xprobe-syscall-target PRIVATE -O0 -Wall -Wextra -Werror) diff --git a/tests/fixtures/syscall_target.c b/tests/fixtures/syscall_target.c new file mode 100644 index 0000000..34d412f --- /dev/null +++ b/tests/fixtures/syscall_target.c @@ -0,0 +1,16 @@ +#include +#include + +int main(void) +{ + for (;;) { + void *mapping = mmap(NULL, 4096, PROT_READ | PROT_WRITE, + MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + if (mapping == MAP_FAILED) + return 1; + *(volatile char *)mapping = 1; + if (munmap(mapping, 4096) != 0) + return 1; + usleep(10000); + } +} diff --git a/tests/integration/run-linux-live.sh b/tests/integration/run-linux-live.sh new file mode 100755 index 0000000..ed41e15 --- /dev/null +++ b/tests/integration/run-linux-live.sh @@ -0,0 +1,74 @@ +#!/usr/bin/env bash +set -euo pipefail + +target=/workspace/build/tests/xprobe-syscall-target +binary=/workspace/target/debug/xprobe +mmap_capture=/tmp/xprobe-mmap.json +munmap_capture=/tmp/xprobe-munmap.json +tracepoint_capture=/tmp/xprobe-tracepoint.json +capacity_capture=/tmp/xprobe-capacity.json +capacity_stderr=/tmp/xprobe-capacity.stderr + +"${target}" & +target_pid=$! +cleanup() { + kill "${target_pid}" 2>/dev/null || true + wait "${target_pid}" 2>/dev/null || true + rm -f "${mmap_capture}" "${munmap_capture}" "${tracepoint_capture}" \ + "${capacity_capture}" "${capacity_stderr}" +} +trap cleanup EXIT + +"${binary}" measure \ + --pid "${target_pid}" \ + --from syscall:mmap:entry \ + --to syscall:mmap:exit \ + --match exact \ + --samples 3 \ + --max-events 64 \ + --timeout-ms 5000 \ + --json --non-interactive --no-color >"${mmap_capture}" + +"${binary}" measure \ + --pid "${target_pid}" \ + --from syscall:munmap:entry \ + --to syscall:munmap:exit \ + --match exact \ + --samples 3 \ + --max-events 64 \ + --timeout-ms 5000 \ + --json --non-interactive --no-color >"${munmap_capture}" + +"${binary}" measure \ + --pid "${target_pid}" \ + --from tracepoint:raw_syscalls:sys_enter \ + --to tracepoint:raw_syscalls:sys_exit \ + --match first-after \ + --samples 3 \ + --max-events 64 \ + --timeout-ms 5000 \ + --json --non-interactive --no-color >"${tracepoint_capture}" + +if "${binary}" measure \ + --pid "${target_pid}" \ + --from syscall:mmap:entry \ + --to syscall:mmap:exit \ + --match exact \ + --duration-ms 500 \ + --max-events 4 \ + --timeout-ms 5000 \ + --json --non-interactive --no-color \ + >"${capacity_capture}" 2>"${capacity_stderr}"; then + echo "Linux capacity test unexpectedly succeeded" >&2 + exit 1 +fi + +printf '{"mmap":' +cat "${mmap_capture}" +printf ',"munmap":' +cat "${munmap_capture}" +printf ',"tracepoint":' +cat "${tracepoint_capture}" +printf ',"capacity":' +cat "${capacity_capture}" +printf '}\n' diff --git a/tests/integration/test_linux.py b/tests/integration/test_linux.py new file mode 100755 index 0000000..5437319 --- /dev/null +++ b/tests/integration/test_linux.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +import json +import pathlib +import subprocess +import sys + + +def main() -> None: + if len(sys.argv) != 2: + raise SystemExit("usage: test_linux.py ") + + workspace = pathlib.Path(__file__).resolve().parents[2] + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--cap-add", + "BPF", + "--cap-add", + "PERFMON", + "--cap-add", + "SYS_ADMIN", + "--cap-add", + "SYS_RESOURCE", + "--security-opt", + "seccomp=unconfined", + "--volume", + f"{workspace}:/workspace:ro", + "--workdir", + "/workspace", + sys.argv[1], + "/workspace/tests/integration/run-linux-live.sh", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + raise SystemExit(completed.returncode) + + captures = json.loads(completed.stdout) + assert_syscall_capture(captures["mmap"], "mmap") + assert_syscall_capture(captures["munmap"], "munmap") + assert_tracepoint_capture(captures["tracepoint"]) + assert_capacity_failure(captures["capacity"]) + print("captured mmap, munmap, and named tracepoint latency evidence") + + +def assert_syscall_capture(result: dict, name: str) -> None: + assert result["ok"] is True + assert result["status"] == "completed" + assert result["measurement"]["samples"]["matched"] == 3 + assert result["measurement"]["samples"]["dropped"] == 0 + assert result["correlation"]["method"] == "exact_syscall_tid_lifecycle" + assert result["correlation"]["confidence"] == "exact" + assert 0 < result["correlation"]["score"] <= 1 + assert result["collection"]["host_events"] == 8 + assert result["collection"]["cuda_events"] == 0 + assert result["collection"]["dropped_events"] == 0 + assert len(result["evidence"]) == 3 + for pair in result["evidence"]: + start = pair["start"] + end = pair["end"] + assert start["event_type"] == "syscall_entry" + assert end["event_type"] == "syscall_exit" + assert start["pid"] == end["pid"] + assert start["tid"] == end["tid"] + assert start["process_start_time"] == end["process_start_time"] + assert start["host"]["probe_kind"] == "syscall" + assert end["host"]["probe_kind"] == "syscall" + assert start["host"]["symbol"] == name + assert end["host"]["symbol"] == name + assert len(start["host"]["arguments"]) == 6 + assert all(argument["read_error"] is None for argument in start["host"]["arguments"]) + assert end["host"]["arguments"] == [] + assert end["host"]["return_value"] is not None + assert pair["latency_ns"] == end["timestamp_ns"] - start["timestamp_ns"] + + +def assert_tracepoint_capture(result: dict) -> None: + assert result["ok"] is True + assert result["measurement"]["samples"]["matched"] == 3 + assert result["correlation"]["method"] == "first_after" + assert result["correlation"]["confidence"] == "heuristic" + assert result["collection"]["dropped_events"] == 0 + for pair in result["evidence"]: + assert pair["start"]["host"]["probe_kind"] == "tracepoint" + assert pair["start"]["host"]["symbol"] == "sys_enter" + assert pair["end"]["host"]["symbol"] == "sys_exit" + assert pair["start"]["attributes"]["tracepoint_category"] == "raw_syscalls" + assert pair["end"]["attributes"]["tracepoint_category"] == "raw_syscalls" + + +def assert_capacity_failure(result: dict) -> None: + assert result["ok"] is False + assert result["error"]["code"] == "EVENT_RATE_TOO_HIGH" + assert result["error"]["details"]["record_capacity"] == 4 + + +if __name__ == "__main__": + main() diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 2c8567d..7b3e85f 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -19,6 +19,7 @@ static EXPORT_SEQUENCE: AtomicU64 = AtomicU64::new(0); use clap::{Args, Parser, Subcommand, ValueEnum}; use xprobe_collector::{ completed, cupti, + linux::{self, LinuxCaptureRequest}, uprobe::{self, UprobeRequest}, }; use xprobe_core::{cupti_compat, discover, doctor, inject, inspect, resolve, validate}; @@ -30,8 +31,8 @@ use xprobe_protocol::{ CheckResult, ClockDomain, CuptiCollectionSummary, DiscoveryResult, ErrorCode, ErrorResponse, Event, EventSource, EventType, ExportFormat, HostCaptureResult, MatchPolicy, MeasurementMode, MeasurementResult, MeasurementSpec, MemcpyKind, ProcessReport, ResolvedCudaSelector, - ResolvedProbe, SchemaVersion, SessionStatus, TargetIdentity, TraceExportResult, - ValidationResult, Warning, XprobeError, + ResolvedLinuxSelector, ResolvedProbe, SchemaVersion, SessionStatus, TargetIdentity, + TraceExportResult, ValidationResult, Warning, XprobeError, }; #[derive(Debug, Parser)] @@ -1304,7 +1305,7 @@ struct LiveMeasureRequest { options: MeasureOptions, } -type HostCaptureHandle = JoinHandle>; +type HostCaptureHandle = JoinHandle>; const SNAPSHOT_QUIET_PERIOD: Duration = Duration::from_millis(100); struct HostCollectors { @@ -1327,9 +1328,7 @@ impl HostCollectors { let result = handle.join().map_err(|_| { CommandFailure::new(ErrorCode::Internal, "host collector thread panicked", false) })?; - captures.push(result.map_err(|error| { - CommandFailure::new(error.code(), error.to_string(), error.recoverable()) - })?); + captures.push(result?); } Ok(captures) } @@ -2137,19 +2136,27 @@ fn start_host_collectors( validation: &ValidationResult, request: &LiveMeasureRequest, ) -> Result { + inspect::verify_target(&report.target).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + })?; let cancelled = Arc::new(AtomicBool::new(false)); let mut probes = Vec::new(); + let mut linux_probes = Vec::new(); for endpoint in [&validation.start, &validation.end] { - let Some(probe) = endpoint.host.as_ref() else { - continue; - }; - if probes - .iter() - .any(|existing: &ResolvedProbe| existing.selector == probe.selector) + if let Some(probe) = endpoint.host.as_ref() + && !probes + .iter() + .any(|existing: &ResolvedProbe| existing.selector == probe.selector) { - continue; + probes.push(probe.clone()); + } + if let Some(probe) = endpoint.linux.as_ref() + && !linux_probes + .iter() + .any(|existing: &ResolvedLinuxSelector| existing == probe) + { + linux_probes.push(probe.clone()); } - probes.push(probe.clone()); } let host_timeout = request .options @@ -2159,8 +2166,8 @@ fn start_host_collectors( .options .samples .unwrap_or(request.options.max_events); - let mut receivers = Vec::with_capacity(probes.len()); - let handles = probes + let mut receivers = Vec::with_capacity(probes.len() + usize::from(!linux_probes.is_empty())); + let mut handles = probes .into_iter() .enumerate() .map(|(index, probe)| { @@ -2183,12 +2190,66 @@ fn start_host_collectors( cancelled: Arc::clone(&cancelled), ready: Some(ready), }; - thread::spawn(move || uprobe::capture(&capture_request)) + thread::spawn(move || { + uprobe::capture(&capture_request).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + }) + }) }) - .collect(); + .collect::>(); + if !linux_probes.is_empty() { + let (ready, receiver) = mpsc::sync_channel(1); + receivers.push(receiver); + let endpoint_count = linux_probes.len(); + let event_limit = linux_capture_event_limit(&request.options, endpoint_count); + let capture_request = LinuxCaptureRequest { + target: report.target.clone(), + probes: linux_probes, + event_limit, + capacity_limit: request.options.samples.is_none(), + timeout: host_timeout, + cancelled: Arc::clone(&cancelled), + ready: Some(ready), + }; + handles.push(thread::spawn(move || { + linux::capture(&capture_request).map_err(|error| { + CommandFailure::new(error.code(), error.to_string(), error.recoverable()) + }) + })); + } let mut collectors = HostCollectors { cancelled, handles }; + wait_for_host_collectors(&mut collectors, receivers, request.timeout)?; + if let Err(error) = inspect::verify_target(&report.target) { + collectors.cancel(); + collectors.join()?; + return Err(CommandFailure::new( + error.code(), + error.to_string(), + error.recoverable(), + )); + } + Ok(collectors) +} + +fn linux_capture_event_limit(options: &MeasureOptions, endpoint_count: usize) -> usize { + options + .samples + .and_then(|samples| { + samples + .checked_mul(endpoint_count) + .and_then(|events| events.checked_add(endpoint_count)) + }) + .unwrap_or(options.max_events) + .min(options.max_events) +} + +fn wait_for_host_collectors( + collectors: &mut HostCollectors, + receivers: Vec>, + timeout: Duration, +) -> Result<(), CommandFailure> { for receiver in receivers { - match receiver.recv_timeout(request.timeout) { + match receiver.recv_timeout(timeout) { Ok(()) => {} Err(mpsc::RecvTimeoutError::Disconnected) => { collectors.cancel(); @@ -2212,7 +2273,7 @@ fn start_host_collectors( } } } - Ok(collectors) + Ok(()) } fn refresh_live_sources( @@ -2739,7 +2800,7 @@ fn completed_live_capture( .map(|capture| completed::CompletedCapture { dropped_records: capture.dropped, unknown_records: 0, - record_limit_reached: None, + record_limit_reached: capture.record_limit_reached.then_some(capture.captured), capture_failed: false, cupti: None, events: capture.events.clone(), @@ -2828,7 +2889,7 @@ fn reject_incomplete_capture( if let Some(capacity) = capture.record_limit_reached { let mut failure = CommandFailure::new( ErrorCode::EventRateTooHigh, - format!("CUPTI capture reached its configured limit of {capacity} records"), + format!("capture reached its configured limit of {capacity} records"), true, ) .with_detail("record_capacity", capacity) diff --git a/xprobe/cli/tests/measure.rs b/xprobe/cli/tests/measure.rs index 182b061..c1d8e3b 100644 --- a/xprobe/cli/tests/measure.rs +++ b/xprobe/cli/tests/measure.rs @@ -112,6 +112,7 @@ fn write_host_capture(path: &PathBuf, timestamp_ns: u64) { captured: 1, dropped: 2, timed_out: false, + record_limit_reached: false, events: vec![event], }; fs::write(path, serde_json::to_vec_pretty(&capture).unwrap()) diff --git a/xprobe/collector/src/completed.rs b/xprobe/collector/src/completed.rs index 2310eec..a147034 100644 --- a/xprobe/collector/src/completed.rs +++ b/xprobe/collector/src/completed.rs @@ -130,7 +130,7 @@ fn decode_json(bytes: &[u8]) -> Result return Ok(CompletedCapture { dropped_records: capture.dropped, unknown_records: 0, - record_limit_reached: None, + record_limit_reached: capture.record_limit_reached.then_some(capture.captured), capture_failed: false, cupti: None, events: capture.events, @@ -289,6 +289,7 @@ mod tests { captured: 1, dropped: 3, timed_out: false, + record_limit_reached: false, events: vec![event(12, 20)], }; let input = serde_json::to_vec_pretty(&capture).unwrap(); diff --git a/xprobe/collector/src/lib.rs b/xprobe/collector/src/lib.rs index 6f1eb75..7e354d2 100644 --- a/xprobe/collector/src/lib.rs +++ b/xprobe/collector/src/lib.rs @@ -2,4 +2,5 @@ pub mod completed; pub mod cupti; +pub mod linux; pub mod uprobe; diff --git a/xprobe/collector/src/linux.rs b/xprobe/collector/src/linux.rs new file mode 100644 index 0000000..edce901 --- /dev/null +++ b/xprobe/collector/src/linux.rs @@ -0,0 +1,675 @@ +use std::{ + cell::RefCell, + collections::BTreeMap, + error::Error, + ffi::OsStr, + fmt, fs, io, + os::unix::fs::MetadataExt, + path::PathBuf, + rc::Rc, + sync::{ + Arc, + atomic::{AtomicBool, AtomicU64, Ordering}, + mpsc::SyncSender, + }, + time::{Duration, Instant}, +}; + +use libbpf_rs::{ + ErrorKind, Link, MapCore, MapFlags, Object, ObjectBuilder, RingBufferBuilder, + TracepointCategory, +}; +use serde_json::Value; +use xprobe_protocol::{ + ArgumentValue, ClockDomain, ErrorCode, Event, EventSource, EventType, HostCaptureResult, + HostEvent, HostProbeKind, ResolvedLinuxSelector, SchemaVersion, TargetIdentity, +}; + +const BPF_OBJECT: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/xprobe.bpf.o")); +const RAW_EVENT_SIZE: usize = 80; +const MIN_RING_BYTES: usize = 256 * 1024; +const POLL_INTERVAL: Duration = Duration::from_millis(100); +static SESSION_SEQUENCE: AtomicU64 = AtomicU64::new(1); + +#[derive(Debug, Clone)] +pub struct LinuxCaptureRequest { + pub target: TargetIdentity, + pub probes: Vec, + pub event_limit: usize, + pub capacity_limit: bool, + pub timeout: Duration, + pub cancelled: Arc, + pub ready: Option>, +} + +#[derive(Debug)] +pub enum LinuxCaptureError { + InvalidRequest(String), + TargetNamespace { + path: PathBuf, + source: io::Error, + }, + MissingObjectMember { + kind: &'static str, + name: String, + }, + Libbpf { + operation: &'static str, + source: libbpf_rs::Error, + }, + MalformedEvent { + expected: usize, + actual: usize, + }, + UnknownProbeId(u32), +} + +impl LinuxCaptureError { + #[must_use] + pub fn code(&self) -> ErrorCode { + match self { + Self::InvalidRequest(_) => ErrorCode::InvalidEventSelector, + Self::TargetNamespace { source, .. } + if source.kind() == io::ErrorKind::PermissionDenied => + { + ErrorCode::PermissionDenied + } + Self::TargetNamespace { source, .. } if source.kind() == io::ErrorKind::NotFound => { + ErrorCode::TargetExited + } + Self::Libbpf { source, .. } if source.kind() == ErrorKind::PermissionDenied => { + ErrorCode::PermissionDenied + } + Self::Libbpf { + operation: "attach tracepoint", + source, + } if source.kind() == ErrorKind::NotFound => ErrorCode::InvalidEventSelector, + Self::TargetNamespace { .. } + | Self::MissingObjectMember { .. } + | Self::Libbpf { .. } + | Self::MalformedEvent { .. } + | Self::UnknownProbeId(_) => ErrorCode::Internal, + } + } + + #[must_use] + pub fn recoverable(&self) -> bool { + matches!( + self.code(), + ErrorCode::InvalidEventSelector | ErrorCode::PermissionDenied | ErrorCode::TargetExited + ) + } +} + +impl fmt::Display for LinuxCaptureError { + fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + Self::InvalidRequest(message) => formatter.write_str(message), + Self::TargetNamespace { path, source } => write!( + formatter, + "failed to inspect target PID namespace at {}: {source}", + path.display() + ), + Self::MissingObjectMember { kind, name } => { + write!(formatter, "BPF object is missing {kind} {name}") + } + Self::Libbpf { operation, source } => { + write!(formatter, "failed to {operation}: {source:#}") + } + Self::MalformedEvent { expected, actual } => write!( + formatter, + "BPF ring buffer record has size {actual}, expected {expected}" + ), + Self::UnknownProbeId(probe_id) => { + write!( + formatter, + "BPF record has unknown Linux probe ID {probe_id}" + ) + } + } + } +} + +impl Error for LinuxCaptureError { + fn source(&self) -> Option<&(dyn Error + 'static)> { + match self { + Self::TargetNamespace { source, .. } => Some(source), + Self::Libbpf { source, .. } => Some(source), + Self::InvalidRequest(_) + | Self::MissingObjectMember { .. } + | Self::MalformedEvent { .. } + | Self::UnknownProbeId(_) => None, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct RawEvent { + timestamp_ns: u64, + sequence: u64, + values: [u64; 6], + pid: u32, + tid: u32, + cpu: u32, + probe_id: u32, +} + +impl RawEvent { + fn decode(bytes: &[u8]) -> Result { + if bytes.len() != RAW_EVENT_SIZE { + return Err(LinuxCaptureError::MalformedEvent { + expected: RAW_EVENT_SIZE, + actual: bytes.len(), + }); + } + let mut values = [0_u64; 6]; + for (index, value) in values.iter_mut().enumerate() { + let start = 16 + index * 8; + *value = u64::from_ne_bytes(bytes[start..start + 8].try_into().expect("u64 field")); + } + Ok(Self { + timestamp_ns: u64::from_ne_bytes(bytes[0..8].try_into().expect("u64 field")), + sequence: u64::from_ne_bytes(bytes[8..16].try_into().expect("u64 field")), + values, + pid: u32::from_ne_bytes(bytes[64..68].try_into().expect("u32 field")), + tid: u32::from_ne_bytes(bytes[68..72].try_into().expect("u32 field")), + cpu: u32::from_ne_bytes(bytes[72..76].try_into().expect("u32 field")), + probe_id: u32::from_ne_bytes(bytes[76..80].try_into().expect("u32 field")), + }) + } +} + +/// Capture up to two PID-scoped Linux tracepoint or syscall endpoints. +/// +/// # Errors +/// +/// Returns [`LinuxCaptureError`] when the request is invalid, the BPF object +/// cannot be loaded or attached, or a ring-buffer record is malformed. +pub fn capture(request: &LinuxCaptureRequest) -> Result { + validate_request(request)?; + let deadline = Instant::now() + .checked_add(request.timeout) + .ok_or_else(|| LinuxCaptureError::InvalidRequest("timeout exceeds Instant range".into()))?; + let object = load_object(request.event_limit)?; + configure_object(&object, request)?; + let (raw_events, dropped) = collect_raw_events(&object, request, deadline)?; + let session_id = format!( + "xp_linux_{}_{}_{}", + request.target.pid, + std::process::id(), + SESSION_SEQUENCE.fetch_add(1, Ordering::Relaxed) + ); + let events = raw_events + .iter() + .map(|event| normalize_event(event, request, &session_id)) + .collect::, _>>()?; + + Ok(HostCaptureResult { + schema_version: SchemaVersion::current(), + ok: true, + session_id, + target: request.target.clone(), + probe_id: 0, + captured: events.len() as u64, + dropped, + timed_out: events.len() < request.event_limit, + record_limit_reached: request.capacity_limit && events.len() == request.event_limit, + events, + }) +} + +fn load_object(event_limit: usize) -> Result { + let ring_bytes = event_limit + .checked_mul(RAW_EVENT_SIZE) + .and_then(usize::checked_next_power_of_two) + .unwrap_or(usize::MAX) + .max(MIN_RING_BYTES); + let ring_bytes = u32::try_from(ring_bytes).map_err(|_| { + LinuxCaptureError::InvalidRequest( + "max-events requires a Linux ring buffer larger than u32".to_owned(), + ) + })?; + let mut builder = ObjectBuilder::default(); + builder + .name("xprobe_linux") + .map_err(|source| libbpf_error("name BPF object", source))?; + let mut open_object = builder + .open_memory(BPF_OBJECT) + .map_err(|source| libbpf_error("open BPF object", source))?; + open_object + .maps_mut() + .find(|map| map.name() == OsStr::new("linux_events")) + .ok_or_else(|| missing("map", "linux_events"))? + .set_max_entries(ring_bytes) + .map_err(|source| libbpf_error("size Linux ring buffer", source))?; + open_object + .load() + .map_err(|source| libbpf_error("load BPF object", source)) +} + +fn configure_object( + object: &Object, + request: &LinuxCaptureRequest, +) -> Result<(), LinuxCaptureError> { + let namespace_path = PathBuf::from(format!("/proc/{}/ns/pid", request.target.pid)); + let namespace = + fs::metadata(&namespace_path).map_err(|source| LinuxCaptureError::TargetNamespace { + path: namespace_path, + source, + })?; + let key = 0_u32.to_ne_bytes(); + let mut config = [0_u8; 56]; + config[0..8].copy_from_slice(&namespace.dev().to_ne_bytes()); + config[8..16].copy_from_slice(&namespace.ino().to_ne_bytes()); + config[16..20].copy_from_slice(&request.target.pid.to_ne_bytes()); + for slot in 0..4 { + let start = 24 + slot * 8; + config[start..start + 8].copy_from_slice(&(-1_i64).to_ne_bytes()); + } + for (index, probe) in request.probes.iter().enumerate() { + let Some(number) = probe.syscall_number else { + continue; + }; + let base = match probe.event_type { + EventType::SyscallEntry => 24, + EventType::SyscallExit => 40, + _ => continue, + }; + let start = base + index * 8; + config[start..start + 8].copy_from_slice(&i64::from(number).to_ne_bytes()); + } + object + .maps() + .find(|map| map.name() == OsStr::new("linux_config")) + .ok_or_else(|| missing("map", "linux_config"))? + .update(&key, &config, MapFlags::ANY) + .map_err(|source| libbpf_error("configure Linux BPF map", source)) +} + +fn collect_raw_events( + object: &Object, + request: &LinuxCaptureRequest, + deadline: Instant, +) -> Result<(Vec, u64), LinuxCaptureError> { + let raw_events = Rc::new(RefCell::new(Vec::with_capacity(request.event_limit))); + let callback_events = Rc::clone(&raw_events); + let callback_error = Rc::new(RefCell::new(None)); + let callback_error_slot = Rc::clone(&callback_error); + let event_limit = request.event_limit; + let events_map = object + .maps() + .find(|map| map.name() == OsStr::new("linux_events")) + .ok_or_else(|| missing("map", "linux_events"))?; + let mut ring_builder = RingBufferBuilder::new(); + ring_builder + .add(&events_map, move |bytes| { + if callback_events.borrow().len() >= event_limit { + return 0; + } + match RawEvent::decode(bytes) { + Ok(event) => callback_events.borrow_mut().push(event), + Err(error) => { + *callback_error_slot.borrow_mut() = Some(error); + return -1; + } + } + 0 + }) + .map_err(|source| libbpf_error("register Linux ring buffer", source))?; + let ring_buffer = ring_builder + .build() + .map_err(|source| libbpf_error("build Linux ring buffer", source))?; + let _links = attach_probes(object, &request.probes)?; + arm_object(object)?; + if let Some(ready) = &request.ready { + ready.send(()).map_err(|_| { + LinuxCaptureError::InvalidRequest( + "Linux collector readiness receiver closed".to_owned(), + ) + })?; + } + + while raw_events.borrow().len() < request.event_limit + && !request.cancelled.load(Ordering::Acquire) + { + let now = Instant::now(); + if now >= deadline { + break; + } + let wait = deadline.saturating_duration_since(now).min(POLL_INTERVAL); + if let Err(source) = ring_buffer.poll(wait) { + if let Some(error) = callback_error.borrow_mut().take() { + return Err(error); + } + return Err(libbpf_error("poll Linux ring buffer", source)); + } + } + + let key = 0_u32.to_ne_bytes(); + let dropped = read_counter(object, "linux_dropped", &key)?; + let collected = raw_events.borrow().clone(); + Ok((collected, dropped)) +} + +fn arm_object(object: &Object) -> Result<(), LinuxCaptureError> { + let key = 0_u32.to_ne_bytes(); + let map = object + .maps() + .find(|map| map.name() == OsStr::new("linux_config")) + .ok_or_else(|| missing("map", "linux_config"))?; + let mut config = map + .lookup(&key, MapFlags::ANY) + .map_err(|source| libbpf_error("read Linux BPF config", source))? + .ok_or_else(|| missing("map value", "linux_config"))?; + if config.len() != 56 { + return Err(LinuxCaptureError::MalformedEvent { + expected: 56, + actual: config.len(), + }); + } + config[20..24].copy_from_slice(&1_u32.to_ne_bytes()); + map.update(&key, &config, MapFlags::ANY) + .map_err(|source| libbpf_error("arm Linux BPF collection", source)) +} + +fn attach_probes( + object: &Object, + probes: &[ResolvedLinuxSelector], +) -> Result, LinuxCaptureError> { + let mut links = Vec::with_capacity(probes.len()); + for (event_type, program_name, tracepoint_name) in [ + ( + EventType::SyscallEntry, + "xprobe_handle_syscall_entry", + "sys_enter", + ), + ( + EventType::SyscallExit, + "xprobe_handle_syscall_exit", + "sys_exit", + ), + ] { + if probes.iter().any(|probe| probe.event_type == event_type) { + let program = object + .progs_mut() + .find(|program| program.name() == OsStr::new(program_name)) + .ok_or_else(|| missing("program", program_name))?; + links.push( + program + .attach_raw_tracepoint(tracepoint_name) + .map_err(|source| libbpf_error("attach tracepoint", source))?, + ); + } + } + for (index, probe) in probes.iter().enumerate() { + if probe.event_type != EventType::Tracepoint { + continue; + } + let raw = probe.category == "raw_syscalls"; + let stem = if raw { + "xprobe_handle_raw_tracepoint" + } else { + "xprobe_handle_tracepoint" + }; + let program_name = format!("{stem}_{}", index + 1); + let program = object + .progs_mut() + .find(|program| program.name() == OsStr::new(&program_name)) + .ok_or_else(|| missing("program", &program_name))?; + let link = if raw { + program + .attach_raw_tracepoint(&probe.name) + .map_err(|source| libbpf_error("attach tracepoint", source))? + } else { + program + .attach_tracepoint( + TracepointCategory::Custom(probe.category.clone()), + &probe.name, + ) + .map_err(|source| libbpf_error("attach tracepoint", source))? + }; + links.push(link); + } + Ok(links) +} + +fn validate_request(request: &LinuxCaptureRequest) -> Result<(), LinuxCaptureError> { + if request.probes.is_empty() || request.probes.len() > 2 { + return Err(LinuxCaptureError::InvalidRequest( + "Linux collector requires one or two endpoints".to_owned(), + )); + } + if request.event_limit == 0 { + return Err(LinuxCaptureError::InvalidRequest( + "max-events must be greater than zero".to_owned(), + )); + } + if request.timeout.is_zero() { + return Err(LinuxCaptureError::InvalidRequest( + "timeout must be greater than zero".to_owned(), + )); + } + for probe in &request.probes { + let valid = match probe.probe_kind { + HostProbeKind::Syscall => { + matches!( + probe.event_type, + EventType::SyscallEntry | EventType::SyscallExit + ) && probe.category == "syscalls" + && probe.syscall_number.is_some() + } + HostProbeKind::Tracepoint => { + probe.event_type == EventType::Tracepoint && probe.syscall_number.is_none() + } + _ => false, + }; + if !valid || probe.category.is_empty() || probe.name.is_empty() { + return Err(LinuxCaptureError::InvalidRequest( + "Linux endpoint metadata is inconsistent".to_owned(), + )); + } + } + Ok(()) +} + +fn normalize_event( + raw: &RawEvent, + request: &LinuxCaptureRequest, + session_id: &str, +) -> Result { + let probe_index = + usize::try_from(raw.probe_id.saturating_sub(1)).expect("u32 probe ID fits usize"); + let probe = request + .probes + .get(probe_index) + .ok_or(LinuxCaptureError::UnknownProbeId(raw.probe_id))?; + let arguments = if probe.event_type == EventType::SyscallEntry { + raw.values + .iter() + .enumerate() + .map(|(index, value)| ArgumentValue { + index: u16::try_from(index).expect("six arguments fit u16"), + abi_type: "u64".to_owned(), + value: Some(Value::from(*value)), + read_error: None, + }) + .collect() + } else { + Vec::new() + }; + let mut attributes = BTreeMap::new(); + attributes.insert( + "tracepoint_category".to_owned(), + Value::String(probe.category.clone()), + ); + Ok(Event { + schema_version: SchemaVersion::current(), + session_id: session_id.to_owned(), + event_id: format!("evt_{}", raw.sequence), + sequence: raw.sequence, + source: EventSource::Ebpf, + event_type: probe.event_type.clone(), + pid: raw.pid, + tid: raw.tid, + cpu: Some(raw.cpu), + timestamp_raw: raw.timestamp_ns, + timestamp_ns: raw.timestamp_ns, + clock_domain: ClockDomain::HostMonotonic, + timestamp_error_ns: None, + process_start_time: Some(request.target.process_start_time), + host: Some(HostEvent { + probe_kind: probe.probe_kind.clone(), + binary_path: None, + build_id: None, + symbol: Some(probe.name.clone()), + offset: None, + return_value: (probe.event_type == EventType::SyscallExit) + .then_some(i64::from_ne_bytes(raw.values[0].to_ne_bytes())), + arguments, + }), + cuda: None, + attributes, + }) +} + +fn read_counter(object: &Object, name: &'static str, key: &[u8]) -> Result { + let value = object + .maps() + .find(|map| map.name() == OsStr::new(name)) + .ok_or_else(|| missing("map", name))? + .lookup(key, MapFlags::ANY) + .map_err(|source| libbpf_error("read Linux BPF map", source))? + .ok_or_else(|| missing("map value", name))?; + let bytes: [u8; 8] = + value + .try_into() + .map_err(|value: Vec| LinuxCaptureError::MalformedEvent { + expected: 8, + actual: value.len(), + })?; + Ok(u64::from_ne_bytes(bytes)) +} + +fn missing(kind: &'static str, name: &str) -> LinuxCaptureError { + LinuxCaptureError::MissingObjectMember { + kind, + name: name.to_owned(), + } +} + +fn libbpf_error(operation: &'static str, source: libbpf_rs::Error) -> LinuxCaptureError { + LinuxCaptureError::Libbpf { operation, source } +} + +#[cfg(test)] +mod tests { + use std::{ + sync::{Arc, atomic::AtomicBool}, + time::Duration, + }; + + use xprobe_protocol::{EventType, HostProbeKind, ResolvedLinuxSelector, TargetIdentity}; + + use super::{ + LinuxCaptureError, LinuxCaptureRequest, RAW_EVENT_SIZE, RawEvent, normalize_event, + validate_request, + }; + + fn syscall(event_type: EventType) -> ResolvedLinuxSelector { + ResolvedLinuxSelector { + event_type, + probe_kind: HostProbeKind::Syscall, + category: "syscalls".to_owned(), + name: "mmap".to_owned(), + syscall_number: Some(9), + } + } + + fn request(probes: Vec) -> LinuxCaptureRequest { + LinuxCaptureRequest { + target: TargetIdentity { + pid: 1234, + process_start_time: 99, + }, + probes, + event_limit: 4, + capacity_limit: false, + timeout: Duration::from_secs(1), + cancelled: Arc::new(AtomicBool::new(false)), + ready: None, + } + } + + #[test] + fn decodes_native_ring_buffer_layout() { + let mut bytes = [0_u8; RAW_EVENT_SIZE]; + bytes[0..8].copy_from_slice(&1000_u64.to_ne_bytes()); + bytes[8..16].copy_from_slice(&9_u64.to_ne_bytes()); + bytes[16..24].copy_from_slice(&0x1234_u64.to_ne_bytes()); + bytes[64..68].copy_from_slice(&1234_u32.to_ne_bytes()); + bytes[68..72].copy_from_slice(&1235_u32.to_ne_bytes()); + bytes[72..76].copy_from_slice(&3_u32.to_ne_bytes()); + bytes[76..80].copy_from_slice(&1_u32.to_ne_bytes()); + + let raw = RawEvent::decode(&bytes).unwrap(); + assert_eq!(raw.timestamp_ns, 1000); + assert_eq!(raw.sequence, 9); + assert_eq!(raw.values[0], 0x1234); + assert_eq!(raw.pid, 1234); + assert_eq!(raw.tid, 1235); + assert_eq!(raw.cpu, 3); + assert_eq!(raw.probe_id, 1); + } + + #[test] + fn normalizes_syscall_scalars_without_pointer_reads() { + let request = request(vec![ + syscall(EventType::SyscallEntry), + syscall(EventType::SyscallExit), + ]); + let entry = RawEvent { + timestamp_ns: 100, + sequence: 1, + values: [0x1000, 4096, 3, 0x22, u64::MAX, 0], + pid: 1234, + tid: 1235, + cpu: 2, + probe_id: 1, + }; + let exit = RawEvent { + timestamp_ns: 150, + sequence: 2, + values: [u64::MAX, 0, 0, 0, 0, 0], + probe_id: 2, + ..entry + }; + + let entry = normalize_event(&entry, &request, "session").unwrap(); + let exit = normalize_event(&exit, &request, "session").unwrap(); + assert_eq!(entry.event_type, EventType::SyscallEntry); + assert_eq!( + entry.host.unwrap().arguments[0].value, + Some(serde_json::json!(4096_u64)) + ); + assert_eq!(exit.event_type, EventType::SyscallExit); + assert_eq!(exit.host.unwrap().return_value, Some(-1)); + } + + #[test] + fn rejects_inconsistent_or_unbounded_requests() { + let mut invalid = request(vec![ResolvedLinuxSelector { + event_type: EventType::Tracepoint, + probe_kind: HostProbeKind::Syscall, + category: "syscalls".to_owned(), + name: "mmap".to_owned(), + syscall_number: Some(9), + }]); + assert!(matches!( + validate_request(&invalid), + Err(LinuxCaptureError::InvalidRequest(_)) + )); + invalid.probes = vec![syscall(EventType::SyscallEntry)]; + invalid.event_limit = 0; + assert!(validate_request(&invalid).is_err()); + } +} diff --git a/xprobe/collector/src/uprobe.rs b/xprobe/collector/src/uprobe.rs index 9f580bc..73277e5 100644 --- a/xprobe/collector/src/uprobe.rs +++ b/xprobe/collector/src/uprobe.rs @@ -208,6 +208,7 @@ pub fn capture(request: &UprobeRequest) -> Result>(); match fields.as_slice() { ["syscall", name, boundary] if valid_tracepoint_component(name) => { + if !cfg!(target_arch = "x86_64") { + return Err(ValidateError::InvalidSelector( + "syscall selectors currently require x86_64".to_owned(), + )); + } + let syscall_number = syscall_number(name).ok_or_else(|| { + ValidateError::InvalidSelector(format!( + "syscall {name:?} is not supported on x86_64" + )) + })?; let (event_type, kind) = match *boundary { "entry" => ( EventType::SyscallEntry, @@ -350,6 +360,7 @@ fn parse_linux_selector( probe_kind: HostProbeKind::Syscall, category: "syscalls".to_owned(), name: (*name).to_owned(), + syscall_number: Some(syscall_number), }, kind, )) @@ -366,6 +377,7 @@ fn parse_linux_selector( probe_kind: HostProbeKind::Tracepoint, category: (*category).to_owned(), name: (*name).to_owned(), + syscall_number: None, }, EndpointKind::Tracepoint, )) @@ -379,6 +391,108 @@ fn parse_linux_selector( } } +fn syscall_number(name: &str) -> Option { + Some(match name { + "read" => 0, + "write" => 1, + "close" => 3, + "fstat" => 5, + "poll" => 7, + "lseek" => 8, + "mmap" => 9, + "mprotect" => 10, + "munmap" => 11, + "brk" => 12, + "ioctl" => 16, + "pread64" => 17, + "pwrite64" => 18, + "readv" => 19, + "writev" => 20, + "sched_yield" => 24, + "mremap" => 25, + "msync" => 26, + "mincore" => 27, + "madvise" => 28, + "nanosleep" => 35, + "getpid" => 39, + "socket" => 41, + "connect" => 42, + "accept" => 43, + "sendto" => 44, + "recvfrom" => 45, + "sendmsg" => 46, + "recvmsg" => 47, + "clone" => 56, + "fork" => 57, + "vfork" => 58, + "execve" => 59, + "exit" => 60, + "wait4" => 61, + "fcntl" => 72, + "flock" => 73, + "fsync" => 74, + "fdatasync" => 75, + "getdents" => 78, + "futex" => 202, + "sched_setaffinity" => 203, + "sched_getaffinity" => 204, + "gettid" => 186, + "epoll_wait" => 232, + "epoll_ctl" => 233, + "tgkill" => 234, + "openat" => 257, + "newfstatat" => 262, + "pselect6" => 270, + "ppoll" => 271, + "unshare" => 272, + "splice" => 275, + "epoll_pwait" => 281, + "accept4" => 288, + "eventfd2" => 290, + "epoll_create1" => 291, + "dup3" => 292, + "pipe2" => 293, + "preadv" => 295, + "pwritev" => 296, + "perf_event_open" => 298, + "recvmmsg" => 299, + "prlimit64" => 302, + "sendmmsg" => 307, + "getcpu" => 309, + "process_vm_readv" => 310, + "process_vm_writev" => 311, + "sched_setattr" => 314, + "sched_getattr" => 315, + "seccomp" => 317, + "getrandom" => 318, + "memfd_create" => 319, + "bpf" => 321, + "execveat" => 322, + "userfaultfd" => 323, + "membarrier" => 324, + "mlock2" => 325, + "copy_file_range" => 326, + "preadv2" => 327, + "pwritev2" => 328, + "statx" => 332, + "rseq" => 334, + "pidfd_send_signal" => 424, + "io_uring_setup" => 425, + "io_uring_enter" => 426, + "io_uring_register" => 427, + "pidfd_open" => 434, + "clone3" => 435, + "close_range" => 436, + "openat2" => 437, + "pidfd_getfd" => 438, + "faccessat2" => 439, + "process_madvise" => 440, + "epoll_pwait2" => 441, + "futex_waitv" => 449, + _ => return None, + }) +} + fn valid_tracepoint_component(value: &str) -> bool { !value.is_empty() && value @@ -845,6 +959,7 @@ mod tests { probe_kind: HostProbeKind::Syscall, category: "syscalls".to_owned(), name: "mmap".to_owned(), + syscall_number: Some(9), }, ); let end = endpoint( @@ -858,6 +973,7 @@ mod tests { probe_kind: HostProbeKind::Syscall, category: "syscalls".to_owned(), name: "mmap".to_owned(), + syscall_number: Some(9), }, ); diff --git a/xprobe/protocol/src/capture.rs b/xprobe/protocol/src/capture.rs index 71cc057..258072c 100644 --- a/xprobe/protocol/src/capture.rs +++ b/xprobe/protocol/src/capture.rs @@ -14,5 +14,7 @@ pub struct HostCaptureResult { pub captured: u64, pub dropped: u64, pub timed_out: bool, + #[serde(default)] + pub record_limit_reached: bool, pub events: Vec, } diff --git a/xprobe/protocol/src/validate.rs b/xprobe/protocol/src/validate.rs index 90bec15..5d49a36 100644 --- a/xprobe/protocol/src/validate.rs +++ b/xprobe/protocol/src/validate.rs @@ -38,6 +38,7 @@ pub struct ResolvedLinuxSelector { pub probe_kind: HostProbeKind, pub category: String, pub name: String, + pub syscall_number: Option, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index 3257955..5a1f305 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -59,6 +59,7 @@ fn host_capture_contract_round_trips() { "captured": 1, "dropped": 0, "timed_out": false, + "record_limit_reached": false, "events": [{ "schema_version": "2.0", "session_id": "xp_uprobe_1234_1000", From eedb920c94b72e49842d0a99b667ac9354089771 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 06:26:58 +0800 Subject: [PATCH 11/17] =?UTF-8?q?=F0=9F=93=9A=20docs:=20guide=20Linux=20ev?= =?UTF-8?q?ent=20profiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 7 +++- README.md | 20 +++++----- docs/architecture.md | 31 ++++++++++----- docs/cli-contract.md | 39 +++++++++++++++---- docs/development.md | 7 ++-- skills/xprobe-measure-latency/SKILL.md | 15 ++++--- .../examples/syscall-duration.json | 15 +++++++ .../references/cli-contract.md | 25 +++++++++--- .../references/investigation.md | 13 ++++++- .../references/result-quality.md | 11 +++--- 10 files changed, 134 insertions(+), 49 deletions(-) create mode 100644 skills/xprobe-measure-latency/examples/syscall-duration.json diff --git a/AGENTS.md b/AGENTS.md index ce79cde..7c3ae4c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,6 +19,9 @@ - Do not add a daemon or persistent session service without a concrete workflow that cannot be expressed by the four bounded commands. - Keep eBPF and CUPTI callback hot paths bounded and free of blocking I/O. +- Resolve architecture-specific syscall names in core. In BPF, filter process + identity and syscall number before reading scalar registers or reserving a + record. Arm multi-link collectors only after every link is attached. ## Failures and safety @@ -31,7 +34,9 @@ reports `injection_required`; log the mutation and include a JSON warning. - Stop CUPTI logically after collection. Do not `dlclose` the injected agent. - Never collect pointer-referenced payloads, environments, or GPU buffer data by - default, and never describe temporal correlation as exact causality. + default. Named tracepoints retain identity and timestamps unless a versioned + scalar payload is explicitly designed. Never describe temporal correlation + as exact causality. ## Agent workflow diff --git a/README.md b/README.md index e8e4939..d6f6115 100644 --- a/README.md +++ b/README.md @@ -10,9 +10,10 @@ `xprobe` is an AI harness for measuring latency between two observable events -in a process, on the CPU, NVIDIA GPU, or across both. Its bounded native profiler -combines eBPF uprobes and NVIDIA CUPTI with an agent-friendly CLI, strict JSON -contracts, explicit correlation quality, and no daemon or server lifecycle. +in a process, on the CPU, NVIDIA GPU, or across both. Its bounded native +profiler combines eBPF function, syscall, and tracepoint evidence with NVIDIA +CUPTI, an agent-friendly CLI, strict JSON contracts, explicit correlation +quality, and no daemon or server lifecycle. ## Install @@ -62,10 +63,11 @@ xprobe measure --pid 4242 \ ``` Kernel launch latency is only one event pair. The same workflow measures host -function spans, CUDA API calls, GPU operation durations, transfers, and paths -across CPU and GPU events after selecting the correct CUDA worker. Aggregate -mode provides a bounded coarse inventory of GPU operations before an exact -evidence measurement narrows the question. +function spans, syscall latency, named Linux events, CUDA API calls, GPU +operation durations, transfers, and paths across CPU and GPU events after +selecting the correct process. Aggregate mode provides a bounded coarse +inventory of GPU operations before an exact evidence measurement narrows the +question. `measure` also accepts completed `--input` captures and versioned live `--spec` files. Evidence can be exported as `jsonl` or `chrome`. JSON results @@ -90,10 +92,10 @@ for safe reactivation. ## Support -| Surface | 0.3.3 support | +| Surface | Current support | | --- | --- | | OS/architecture | Linux x86_64, glibc 2.34 or newer | -| Host events | ELF function entry/return through PID-scoped uprobes | +| Host events | PID-scoped ELF function, named syscall, and tracepoint boundaries | | CUDA callbacks | Runtime and Driver API entry/exit | | GPU activity | Kernel, memcpy, and memset start/end | | CUDA/CUPTI | 12.x and 13.x with automatic major selection | diff --git a/docs/architecture.md b/docs/architecture.md index ba57631..7807356 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -33,7 +33,7 @@ attachment, collection, correlation, logical CUPTI shutdown, and cleanup. | `xprobe/collector` | eBPF collection, CUPTI control protocol and ABI decoding | | `xprobe/correlator` | Selector matching, pair evidence, statistics and quality | | `xprobe/exporter` | Event JSONL and Chrome Trace Event Format | -| `bpf/` | PID-scoped uprobe/uretprobe programs | +| `bpf/` | PID-scoped uprobe, syscall, and named tracepoint programs | | `cupti/` | Reusable in-process CUDA callback/activity agent | ## Identity and discovery @@ -57,9 +57,10 @@ base is a symbol address. ## Validation -`validate` is read-only. It resolves host selectors, parses CUDA selectors, -checks correlation-policy compatibility, and reports eBPF, CUPTI, callback, -activity, and clock requirements. CUPTI activation is one of: +`validate` is read-only. It resolves ELF probes and Linux syscall numbers, +parses named tracepoint and CUDA selectors, checks correlation-policy +compatibility, and reports eBPF, CUPTI, callback, activity, and clock +requirements. CUPTI activation is one of: - `not_required` - `already_loaded` @@ -72,10 +73,20 @@ required host collection remain explicit issues or errors. ## Collection -Host events use libbpf-rs with one PID-scoped uprobe or uretprobe per unique -host endpoint. The BPF path only filters identity, records monotonic timestamps -and CPU/TID, updates bounded counters, and submits fixed records. Rust ownership -detaches links on every return path. +Host events use libbpf-rs. ELF functions attach one PID-scoped uprobe or +uretprobe per unique endpoint. Linux syscall endpoints share raw +`sys_enter`/`sys_exit` links and compare up to two configured x86_64 syscall +numbers after PID-namespace filtering. Only matching entries read the six +scalar ABI registers; exits retain the scalar return value. Pointer-referenced +memory is never read. + +Named tracepoints attach by category and name and retain timestamp, PID/TID, +CPU, and selector identity without copying payload fields. The two endpoint +links are armed together after attachment, so setup events cannot enter the +capture. Linux records use a `--max-events`-sized ring buffer and fixed record +layout. Ring exhaustion, scalar-read failure, malformed records, and +duration-capacity exhaustion remain explicit failures. Rust ownership detaches +all links on every return path. CUDA events use an in-process CUPTI Agent. CUDA 12 and CUDA 13 builds share the same source and capture ABI but link their matching CUPTI SONAME. Loading the @@ -126,7 +137,7 @@ All sources normalize into the versioned `Event` type. `measure` supports: | Policy | Pairing | Confidence | | --- | --- | --- | -| `exact` | CUPTI correlation ID | exact | +| `exact` | CUPTI correlation ID or same-thread syscall lifecycle | exact | | `first-after` | First unused end at or after each start | heuristic | | `nearest` | Nearest unused end by timestamp | heuristic | | `stack-nested` | Per-thread LIFO host entry/return | high | @@ -140,7 +151,7 @@ CUPTI event without a reported interpolation bound makes ## Contracts and failure model -Public records use schema version `1.0`; generated schemas are checked in and +Public records use schema version `2.0`; generated schemas are checked in and tested for drift. Unknown fields and unsupported versions fail deserialization. Expected capability absence is represented as available/restricted/unavailable data. Unexpected I/O, malformed procfs/ELF/capture data, target changes, diff --git a/docs/cli-contract.md b/docs/cli-contract.md index 3bad6cc..2653303 100644 --- a/docs/cli-contract.md +++ b/docs/cli-contract.md @@ -1,7 +1,7 @@ # CLI contract -xprobe 0.3.3 exposes exactly four public commands: `doctor`, `discover`, -`validate`, and `measure`. +xprobe exposes exactly four public commands: `doctor`, `discover`, `validate`, +and `measure`. ## Common behavior @@ -66,8 +66,18 @@ uprobe:::entry uprobe:::return uprobe::+0x:entry uprobe::+0x:return +syscall::entry +syscall::exit +tracepoint:: ``` +Named syscall selectors are Linux x86_64 endpoints resolved by `validate`. +Their eBPF path filters PID and syscall number before reserving an event, +records the six scalar syscall ABI registers on entry, and records the scalar +return value on exit. It never dereferences pointer arguments. Named +tracepoints record identity and timestamp only; they do not copy the tracepoint +payload. Unsupported syscall names and unavailable tracepoints fail explicitly. + CUDA forms include Runtime/Driver API entry/exit and kernel, memcpy, or memset activity start/end. Kernel selectors accept `name~REGEX`; memcpy selectors accept `kind=`. @@ -86,9 +96,10 @@ parses CUDA filters, and checks collection and correlation requirements. Results conform to `schemas/validate.schema.json`. Supported policies are `exact`, `first-after`, `nearest`, `stack-nested`, and -`stream-order`. Exact requires deterministic CUPTI correlation IDs. Nested -requires entry/return of the same host function. Stream order requires GPU -activity endpoints. Temporal policies always warn that they are heuristic. +`stream-order`. Exact uses deterministic CUPTI correlation IDs or one named +syscall's per-thread entry/exit lifecycle. Nested requires entry/return of the +same host function. Stream order requires GPU activity endpoints. Temporal +policies always warn that they are heuristic. `policy_recommendation` reports the strongest compatible policy, a stable machine-readable reason, and all compatible alternatives. The caller still chooses the policy; xprobe does not silently replace or retry it. @@ -116,6 +127,15 @@ xprobe measure --pid 4242 \ --json --non-interactive --no-color ``` +Linux syscall duration: + +```bash +xprobe measure --pid 4242 \ + --from 'syscall:mmap:entry' --to 'syscall:mmap:exit' \ + --match exact --samples 100 --max-events 1000 \ + --json --non-interactive --no-color +``` + Completed captures: ```bash @@ -154,8 +174,13 @@ duration, optional transferred bytes, table occupancy, and drop completeness. Aggregate kernel regex must be reducible to an exact, prefix, suffix, or contains filter because this mode intentionally has no Rust-side event pass. -Live host endpoints attach PID-scoped eBPF probes. CUDA endpoints automatically -activate the CUPTI agent. If it is absent, `--agent` or +Live host endpoints attach PID-scoped eBPF probes. Linux syscall endpoints use +raw tracepoints so they do not depend on tracingfs event IDs; ordinary named +tracepoints use their kernel category and name. A samples-bound Linux capture +allows bounded startup slack for a target already inside an event boundary. A +duration capture that fills `--max-events` returns `EVENT_RATE_TOO_HIGH` rather +than reporting partial success. CUDA endpoints automatically activate the CUPTI +agent. If it is absent, `--agent` or `XPROBE_CUPTI_AGENT_PATH` selects the shared object; otherwise xprobe searches `../lib/xprobe/cuda12` or `../lib/xprobe/cuda13` according to the target and then the matching development build path. An unobservable target major is diff --git a/docs/development.md b/docs/development.md index 2dcdbf5..4eff9cd 100644 --- a/docs/development.md +++ b/docs/development.md @@ -52,14 +52,15 @@ Compile the BPF object without attaching it: just test-bpf ``` -Run the real PID-scoped uprobe and uretprobe test in the pinned container: +Run the real PID-scoped ELF and Linux event tests in the pinned container: ```bash just test-bpf-live ``` -The live test captures function entry and return events from the same target. -It requires Docker daemon access and grants the container `BPF`, +The live suite captures function entry/return, mmap/munmap lifecycle, generic +raw tracepoint, and host-capacity failure paths from controlled targets. It +requires Docker daemon access and grants the container `BPF`, `PERFMON`, `SYS_ADMIN`, and `SYS_RESOURCE`, with seccomp disabled for BPF/perf syscalls. It does not use `--privileged`, does not require GPU access, mounts the workspace read-only, and removes the container after the test. diff --git a/skills/xprobe-measure-latency/SKILL.md b/skills/xprobe-measure-latency/SKILL.md index 7fba978..91f2432 100644 --- a/skills/xprobe-measure-latency/SKILL.md +++ b/skills/xprobe-measure-latency/SKILL.md @@ -39,10 +39,12 @@ For more than one selected process, follow 5. Map GPU or mixed work before choosing a name. Validate broad kernel, memcpy, or memset activity endpoints, then collect one bounded, representative coarse inventory per event family with `measure --aggregate --duration-ms ...`. - For CPU-only work, resolve and validate the intended host-function boundary - directly; do not require CUDA or CUPTI. Scope breadth and collection duration - are independent: keep the selector broad where an activity inventory exists, - choose a duration that covers the workload cycle being diagnosed, and give + For CPU-only work, use existing application evidence or a bounded system + summary to choose a function, named syscall, or tracepoint family before + detailed collection; do not require CUDA or CUPTI, and do not begin with an + unfiltered high-rate raw tracepoint. Scope breadth and collection duration are + independent: keep the selector broad where a bounded aggregate exists, choose + a duration that covers the workload cycle being diagnosed, and give `--max-groups` headroom. For defensibly homogeneous workers, inventory one representative worker and apply its evidence-derived narrow selector to all selected workers. @@ -51,8 +53,8 @@ For more than one selected process, follow `scripts/analyze_trace.py` and use launch variants, stream distribution, busy union, overlap factor, and adjacent gaps. Read [references/trace-analysis.md](references/trace-analysis.md) when interpreting - the report. For CPU-only work, use resolved host selectors and result evidence - to form the hypothesis instead. + the report. For CPU-only work, use resolved host selectors or filtered + syscall/tracepoint evidence to form the hypothesis instead. 7. Run one read-only `xprobe validate` per selected worker. Compare every response target with the PID plus process start time retained from discovery before mutation, and stop that worker when `valid` is false. If @@ -79,6 +81,7 @@ or [coarse memcpy inventory](examples/coarse-memcpy-inventory.json), then use th [kernel duration](examples/kernel-duration.json), [same-stream gap](examples/same-stream-kernel-gap.json), [host span](examples/host-function-span.json), and +[syscall duration](examples/syscall-duration.json), [memcpy duration](examples/memcpy-duration.json) specs, plus the [CUDA synchronization API](examples/cuda-api-duration.json) shape, after replacing target identity and selectors. Each bounded call answers one diff --git a/skills/xprobe-measure-latency/examples/syscall-duration.json b/skills/xprobe-measure-latency/examples/syscall-duration.json new file mode 100644 index 0000000..a5b956e --- /dev/null +++ b/skills/xprobe-measure-latency/examples/syscall-duration.json @@ -0,0 +1,15 @@ +{ + "schema_version": "2.0", + "name": "mmap_duration", + "target": { + "pid": 1234, + "process_start_time": 987654 + }, + "start_selector": "syscall:mmap:entry", + "end_selector": "syscall:mmap:exit", + "match_policy": "exact", + "samples": 100, + "duration_ms": null, + "timeout_ms": 30000, + "max_events": 1000 +} diff --git a/skills/xprobe-measure-latency/references/cli-contract.md b/skills/xprobe-measure-latency/references/cli-contract.md index da2c439..f32aea5 100644 --- a/skills/xprobe-measure-latency/references/cli-contract.md +++ b/skills/xprobe-measure-latency/references/cli-contract.md @@ -20,6 +20,14 @@ uprobe::+0x:entry uprobe::+0x:return ``` +Linux selectors use `syscall::entry|exit` and +`tracepoint::`. Syscall entry/exit for one name supports +`exact` per-thread lifecycle matching. Entry evidence contains scalar ABI +register values and exit evidence contains the scalar return value; xprobe +does not dereference pointers. Named tracepoints contain no payload fields. +Always let `validate` resolve the named syscall or tracepoint on the current +host before measurement. + CUDA selectors cover Runtime and Driver API entry/exit plus kernel, memcpy, and memset activity start/end. Kernel activity accepts `name~REGEX`; memcpy accepts `kind=`. `validate` must accept the complete selector @@ -27,12 +35,12 @@ and policy before measurement. ## Correlation -Use `exact` for CUDA events with the same CUPTI correlation ID, -`stack-nested` for entry/return of the same host function, and `stream-order` -for activity events on one CUDA stream. `first-after` and `nearest` are temporal -heuristics and cannot establish causality. Read `policy_recommendation.policy`, -its machine-readable `reason`, and `compatible_policies`; xprobe never silently -changes the requested policy. +Use `exact` for CUDA events with the same CUPTI correlation ID or entry/exit of +one named syscall, `stack-nested` for entry/return of the same host function, +and `stream-order` for activity events on one CUDA stream. `first-after` and +`nearest` are temporal heuristics and cannot establish causality. Read +`policy_recommendation.policy`, its machine-readable `reason`, and +`compatible_policies`; xprobe never silently changes the requested policy. ## Bounds and failures @@ -57,6 +65,11 @@ sets a live stop from ARM completion; either samples or duration may complete a call when both are present. Timeout bounds the complete foreground operation and cleanup. +Linux syscall filtering runs in BPF before event reservation. A duration-bound +host capture that reaches `max-events` fails with `EVENT_RATE_TOO_HIGH`; narrow +the selector or increase the explicit bound. Do not replace it with partial +evidence. + `--events-out PATH` atomically writes the bounded capture with mode `0600`, not only matched evidence. Collection completeness and CUPTI capacity, observed, retained, dropped, and buffer utilization fields describe capture integrity diff --git a/skills/xprobe-measure-latency/references/investigation.md b/skills/xprobe-measure-latency/references/investigation.md index 07ad7ea..cc0c74d 100644 --- a/skills/xprobe-measure-latency/references/investigation.md +++ b/skills/xprobe-measure-latency/references/investigation.md @@ -74,9 +74,10 @@ generated source, or application logs. xprobe does not read JIT cache contents and cannot select by grid/block. Long mangled names should be narrowed to a short, observed-unique literal instead of copied wholesale. -## Derive host selectors +## Derive CPU selectors -Resolve the mapped object in the target, then inspect its symbols: +Choose the narrowest observable boundary supported by existing evidence. For a +function, resolve the mapped object in the target and inspect its symbols: ```bash readlink -f "/proc/$PID/exe" @@ -91,6 +92,13 @@ stripped or local code, derive a file offset with `readelf`/`objdump` and use `validate`; do not infer a runtime virtual address from one process and reuse it as a file offset. +For kernel-facing latency, first use application logs, `/proc` state, or a +bounded syscall summary to identify a candidate. Then validate +`syscall:NAME:entry` to `syscall:NAME:exit` with `exact`. Use +`tracepoint:CATEGORY:NAME` only when the kernel event itself is the intended +boundary. Do not start an unknown high-rate workload with unfiltered raw +syscall tracepoints: select first, then collect detailed evidence. + ## Measure one narrow hypothesis Choose one next boundary from evidence: @@ -99,6 +107,7 @@ Choose one next boundary from evidence: - CUDA API exit to kernel start with `exact` for launch delay; - kernel end to next activity start with `stream-order` for one-stream gaps; - host function entry to return with `stack-nested` for CPU span; +- named syscall entry to exit with `exact` for kernel-facing latency; - host marker to GPU activity with `first-after` only as a disclosed heuristic. After capture, analyze the artifact and all result quality fields. An aggregate diff --git a/skills/xprobe-measure-latency/references/result-quality.md b/skills/xprobe-measure-latency/references/result-quality.md index a7fed6d..1b13f22 100644 --- a/skills/xprobe-measure-latency/references/result-quality.md +++ b/skills/xprobe-measure-latency/references/result-quality.md @@ -1,10 +1,11 @@ # Result quality -Prefer `exact` when CUDA endpoints carry the same CUPTI correlation ID. Prefer -`stack-nested` for entry/return pairs of the same host function. Use -`stream-order` only for GPU activity endpoints on the same device, context, and -stream. Treat `first-after` and `nearest` as temporal heuristics, never request -causality. +Prefer `exact` when CUDA endpoints carry the same CUPTI correlation ID or when +one named syscall has entry/exit records from the same thread and process +identity. Prefer `stack-nested` for entry/return pairs of the same host +function. Use `stream-order` only for GPU activity endpoints on the same device, +context, and stream. Treat `first-after` and `nearest` as temporal heuristics, +never request causality. Inspect every evidence pair for selector scope, process identity, correlation IDs, device/context/stream, timestamps, and clock domains. A result is not From 2979473f1fccf1887bfdd070941126b6b58d274d Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 06:48:13 +0800 Subject: [PATCH 12/17] =?UTF-8?q?=E2=9C=A8=20feat:=20resolve=20native=20fr?= =?UTF-8?q?amework=20symbols?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.lock | 10 + Cargo.toml | 1 + docs/architecture.md | 5 +- docs/cli-contract.md | 8 + docs/development.md | 14 ++ justfile | 6 + schemas/event.schema.json | 7 + schemas/host-capture.schema.json | 7 + schemas/measurement-result.schema.json | 7 + schemas/resolve.schema.json | 6 + schemas/validate.schema.json | 6 + .../references/cli-contract.md | 7 + .../references/investigation.md | 20 +- tests/fixtures/resolve_library.c | 15 +- tests/integration/test_pytorch.py | 55 +++++ tests/integration/test_pytorch_symbols.py | 185 +++++++++++++++++ xprobe/cli/src/main.rs | 7 +- xprobe/cli/tests/measure.rs | 1 + xprobe/cli/tests/resolve.rs | 16 ++ xprobe/collector/src/linux.rs | 1 + xprobe/collector/src/uprobe.rs | 5 + xprobe/core/Cargo.toml | 1 + xprobe/core/src/resolve.rs | 188 ++++++++++++++---- xprobe/core/src/validate.rs | 1 + xprobe/correlator/src/lib.rs | 43 +++- xprobe/protocol/src/event.rs | 2 + xprobe/protocol/src/resolve.rs | 1 + xprobe/protocol/tests/contracts.rs | 3 + 28 files changed, 580 insertions(+), 48 deletions(-) create mode 100644 tests/integration/test_pytorch.py create mode 100755 tests/integration/test_pytorch_symbols.py diff --git a/Cargo.lock b/Cargo.lock index 283a306..c178dee 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -135,6 +135,15 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" +[[package]] +name = "cpp_demangle" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2bb79cb74d735044c972aae58ed0aaa9a837e85b01106a54c39e42e97f62253" +dependencies = [ + "cfg-if", +] + [[package]] name = "dyn-clone" version = "1.0.20" @@ -476,6 +485,7 @@ dependencies = [ name = "xprobe-core" version = "0.3.3" dependencies = [ + "cpp_demangle", "nix", "object", "regex", diff --git a/Cargo.toml b/Cargo.toml index 6fe8c05..79b0711 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ rust-version = "1.85" [workspace.dependencies] clap = { version = "4.5", features = ["derive"] } +cpp_demangle = "0.4" libbpf-rs = { version = "0.26.2", features = ["vendored"] } nix = { version = "0.31.3", features = ["process", "ptrace", "signal", "uio"] } object = { version = "0.36", default-features = false, features = ["read", "std"] } diff --git a/docs/architecture.md b/docs/architecture.md index 7807356..b72a338 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -53,7 +53,10 @@ handling remain in the caller, and event correlation never crosses processes. Host selector resolution converts ELF virtual addresses through load segments to file offsets, then through `/proc//maps` to runtime addresses. This supports executables, PIE, and shared libraries without assuming that a mapping -base is a symbol address. +base is a symbol address. Raw ELF names are matched without demangling the +symbol table. A full C++ signature uses the explicit `symbol=` selector form; +the resolver checks exported dynamic symbols first, falls back to the complete +table, and returns both the attachable mangled name and readable signature. ## Validation diff --git a/docs/cli-contract.md b/docs/cli-contract.md index 2653303..06bf8b3 100644 --- a/docs/cli-contract.md +++ b/docs/cli-contract.md @@ -64,6 +64,8 @@ attachment. Host selector forms are: ```text uprobe:::entry uprobe:::return +uprobe::symbol=:entry +uprobe::symbol=:return uprobe::+0x:entry uprobe::+0x:return syscall::entry @@ -71,6 +73,12 @@ syscall::exit tracepoint:: ``` +The `symbol=` form allows `::` and other punctuation in a full C++ signature. +Resolution returns the attachable mangled ELF name in `symbol` and its readable +signature in `symbol_demangled`; captured host events retain both. It works for +mapped CPython, native extension, and framework libraries, but does not resolve +Python frames or `module.qualname` names. + Named syscall selectors are Linux x86_64 endpoints resolved by `validate`. Their eBPF path filters PID and syscall number before reserving an event, records the six scalar syscall ABI registers on entry, and records the scalar diff --git a/docs/development.md b/docs/development.md index 4eff9cd..220e71c 100644 --- a/docs/development.md +++ b/docs/development.md @@ -65,6 +65,20 @@ requires Docker daemon access and grants the container `BPF`, syscalls. It does not use `--privileged`, does not require GPU access, mounts the workspace read-only, and removes the container after the test. +Resolve real CPython, native extension, and libtorch C++ symbols with a Mamba +environment containing PyTorch: + +```bash +PYTORCH_PYTHON=/path/to/env/bin/python just test-pytorch-symbols +``` + +Run the corresponding live `torch.mm` entry/return measurement in the pinned +BPF container: + +```bash +PYTORCH_ENV=/path/to/env just test-pytorch-live +``` + ## GPU checks Run host diagnostics outside restricted sandboxes when GPU device access is diff --git a/justfile b/justfile index 2c59437..f2ec107 100644 --- a/justfile +++ b/justfile @@ -69,6 +69,12 @@ test-multisource-live: build test-multisource-live-cuda12: build python3 tests/integration/test_multisource.py "{{cuda12_devel_image}}" target/debug/xprobe +test-pytorch-symbols: build + python3 tests/integration/test_pytorch_symbols.py --python "${PYTORCH_PYTHON:?set PYTORCH_PYTHON to a Python with torch}" + +test-pytorch-live: build + python3 tests/integration/test_pytorch.py "{{cuda_smoke_image}}" "${PYTORCH_ENV:?set PYTORCH_ENV to a Mamba environment with torch}" + benchmark-gpu: python3 benchmarks/cuda-callback/run.py "{{cuda13_devel_image}}" diff --git a/schemas/event.schema.json b/schemas/event.schema.json index 6893296..3c6142e 100644 --- a/schemas/event.schema.json +++ b/schemas/event.schema.json @@ -374,6 +374,13 @@ "string", "null" ] + }, + "symbol_demangled": { + "type": [ + "string", + "null" + ], + "default": null } }, "additionalProperties": false, diff --git a/schemas/host-capture.schema.json b/schemas/host-capture.schema.json index 7857efd..2d4033c 100644 --- a/schemas/host-capture.schema.json +++ b/schemas/host-capture.schema.json @@ -431,6 +431,13 @@ "string", "null" ] + }, + "symbol_demangled": { + "type": [ + "string", + "null" + ], + "default": null } }, "additionalProperties": false, diff --git a/schemas/measurement-result.schema.json b/schemas/measurement-result.schema.json index 2dc9158..b41eec2 100644 --- a/schemas/measurement-result.schema.json +++ b/schemas/measurement-result.schema.json @@ -561,6 +561,13 @@ "string", "null" ] + }, + "symbol_demangled": { + "type": [ + "string", + "null" + ], + "default": null } }, "additionalProperties": false, diff --git a/schemas/resolve.schema.json b/schemas/resolve.schema.json index f9f7fca..bc7f286 100644 --- a/schemas/resolve.schema.json +++ b/schemas/resolve.schema.json @@ -46,6 +46,12 @@ "null" ] }, + "symbol_demangled": { + "type": [ + "string", + "null" + ] + }, "symbol_size": { "type": [ "integer", diff --git a/schemas/validate.schema.json b/schemas/validate.schema.json index b7a5b2c..b879380 100644 --- a/schemas/validate.schema.json +++ b/schemas/validate.schema.json @@ -338,6 +338,12 @@ "null" ] }, + "symbol_demangled": { + "type": [ + "string", + "null" + ] + }, "symbol_size": { "type": [ "integer", diff --git a/skills/xprobe-measure-latency/references/cli-contract.md b/skills/xprobe-measure-latency/references/cli-contract.md index f32aea5..0240ca5 100644 --- a/skills/xprobe-measure-latency/references/cli-contract.md +++ b/skills/xprobe-measure-latency/references/cli-contract.md @@ -16,10 +16,17 @@ Host function selectors use: ```text uprobe:::entry uprobe:::return +uprobe::symbol=:entry +uprobe::symbol=:return uprobe::+0x:entry uprobe::+0x:return ``` +Use the `symbol=` form for a C++ name containing `::`; pass the complete +demangled signature. Validation returns the actual mangled ELF symbol and the +readable signature. This resolves native code mapped by a Python process, not +Python frame or `module.qualname` names. + Linux selectors use `syscall::entry|exit` and `tracepoint::`. Syscall entry/exit for one name supports `exact` per-thread lifecycle matching. Entry evidence contains scalar ABI diff --git a/skills/xprobe-measure-latency/references/investigation.md b/skills/xprobe-measure-latency/references/investigation.md index cc0c74d..85e5cdf 100644 --- a/skills/xprobe-measure-latency/references/investigation.md +++ b/skills/xprobe-measure-latency/references/investigation.md @@ -84,14 +84,32 @@ readlink -f "/proc/$PID/exe" cat "/proc/$PID/maps" readelf -Ws /path/to/object nm -D --defined-only /path/to/object +nm -D --defined-only --demangle /path/to/cpp-object ``` Use `uprobe:::entry|return` when a symbol is available. For -stripped or local code, derive a file offset with `readelf`/`objdump` and use +an exact C++ signature containing `::`, use +`uprobe::symbol=:entry|return`; validation +returns both the mangled attach name and readable signature. For stripped or +local code, derive a file offset with `readelf`/`objdump` and use `uprobe::+0xOFFSET:entry|return`. Always pass the exact candidate to `validate`; do not infer a runtime virtual address from one process and reuse it as a file offset. +For eager PyTorch, inspect the mapped CPython executable, `torch._C`, and the +loaded libtorch objects. Prefer an exported dispatcher or native operator +signature observed in that exact installed build, such as an +`at::_ops::::call(...)` boundary, and validate both entry and return +before measuring with `stack-nested`. Treat `_PyEval_EvalFrameDefault` only as a +broad interpreter boundary; xprobe does not turn it into Python function names. + +After `torch.compile` or Triton warmup, do not assume an eager operator boundary +still encloses the fused work. Inventory CUDA kernels first and narrow using +the emitted kernel names. Generated CPU code without a stable file-backed ELF +symbol is not a valid uprobe target. Never reuse a C++ signature, mangled name, +file offset, or generated kernel name across PyTorch builds without resolving +and validating it again. + For kernel-facing latency, first use application logs, `/proc` state, or a bounded syscall summary to identify a candidate. Then validate `syscall:NAME:entry` to `syscall:NAME:exit` with `exact`. Use diff --git a/tests/fixtures/resolve_library.c b/tests/fixtures/resolve_library.c index 901a9b4..4750a94 100644 --- a/tests/fixtures/resolve_library.c +++ b/tests/fixtures/resolve_library.c @@ -1 +1,14 @@ -__attribute__((visibility("default"), noinline)) void xprobe_resolve_library_marker(void) {} +__attribute__((visibility("default"), noinline)) void +xprobe_cpp_native_operator(long value) + __asm__("_ZN14xprobe_fixture15native_operatorEl"); + +void xprobe_cpp_native_operator(long value) +{ + __asm__ volatile("" : : "r"(value) : "memory"); +} + +__attribute__((visibility("default"), noinline)) void +xprobe_resolve_library_marker(void) +{ + xprobe_cpp_native_operator(1); +} diff --git a/tests/integration/test_pytorch.py b/tests/integration/test_pytorch.py new file mode 100644 index 0000000..c52b8e6 --- /dev/null +++ b/tests/integration/test_pytorch.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +import pathlib +import subprocess +import sys + + +def main() -> None: + if len(sys.argv) != 3: + raise SystemExit("usage: test_pytorch.py ") + + workspace = pathlib.Path(__file__).resolve().parents[2] + pytorch_env = pathlib.Path(sys.argv[2]).resolve() + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--cap-add", + "BPF", + "--cap-add", + "PERFMON", + "--cap-add", + "SYS_ADMIN", + "--cap-add", + "SYS_RESOURCE", + "--security-opt", + "seccomp=unconfined", + "--volume", + f"{workspace}:/workspace:ro", + "--volume", + f"{pytorch_env}:/opt/xprobe-pytorch:ro", + "--workdir", + "/workspace", + sys.argv[1], + "/opt/xprobe-pytorch/bin/python", + "/workspace/tests/integration/test_pytorch_symbols.py", + "--python", + "/opt/xprobe-pytorch/bin/python", + "--xprobe", + "/workspace/target/debug/xprobe", + "--measure", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + raise SystemExit(completed.returncode) + sys.stdout.write(completed.stdout) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/test_pytorch_symbols.py b/tests/integration/test_pytorch_symbols.py new file mode 100755 index 0000000..9da2e29 --- /dev/null +++ b/tests/integration/test_pytorch_symbols.py @@ -0,0 +1,185 @@ +#!/usr/bin/env python3 +import argparse +import json +import pathlib +import subprocess +import sys + + +MM_SYMBOL = "at::_ops::mm::call(at::Tensor const&, at::Tensor const&)" + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--python", required=True) + parser.add_argument("--xprobe", default="target/debug/xprobe") + parser.add_argument("--measure", action="store_true") + args = parser.parse_args() + + workspace = pathlib.Path(__file__).resolve().parents[2] + workload = subprocess.Popen( + [args.python, "-u", "-c", WORKLOAD], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + try: + assert workload.stdout is not None + metadata = json.loads(workload.stdout.readline()) + pid = workload.pid + python_probe = resolve( + workspace, + args.xprobe, + pid, + metadata["cpython"], + "_PyEval_EvalFrameDefault", + ) + extension_probe = resolve( + workspace, + args.xprobe, + pid, + metadata["torch_extension"], + "PyInit__C", + ) + torch_probe = resolve( + workspace, + args.xprobe, + pid, + metadata["libtorch_cpu"], + MM_SYMBOL, + demangled=True, + ) + measurement = ( + measure(workspace, args.xprobe, pid, metadata["libtorch_cpu"]) + if args.measure + else None + ) + finally: + workload.terminate() + workload.wait(timeout=10) + + assert python_probe["symbol"] == "_PyEval_EvalFrameDefault" + assert extension_probe["symbol"] == "PyInit__C" + assert torch_probe["symbol"].startswith("_Z") + assert torch_probe["symbol_demangled"] == MM_SYMBOL + assert python_probe["object_kind"] in {"executable", "position_independent_executable"} + assert extension_probe["object_kind"] == "shared_library" + assert torch_probe["object_kind"] == "shared_library" + if measurement is not None: + assert measurement["measurement"]["samples"]["matched"] == 8 + assert len(measurement["evidence"]) == 8 + for evidence in measurement["evidence"]: + assert evidence["start"]["host"]["symbol_demangled"] == MM_SYMBOL + assert evidence["end"]["host"]["symbol_demangled"] == MM_SYMBOL + print( + json.dumps( + { + "ok": True, + "measured": measurement is not None, + "pytorch_version": metadata["torch_version"], + "resolved": ["cpython", "native_extension", "libtorch_cpp"], + }, + sort_keys=True, + ) + ) + + +def resolve( + workspace: pathlib.Path, + binary: str, + pid: int, + object_path: str, + symbol: str, + demangled: bool = False, +) -> dict: + target = f"symbol={symbol}" if demangled else symbol + selector = f"uprobe:{object_path}:{target}:entry" + completed = subprocess.run( + [ + binary, + "resolve", + "--pid", + str(pid), + "--selector", + selector, + "--json", + "--non-interactive", + "--no-color", + ], + cwd=workspace, + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise AssertionError( + f"resolve failed for {selector!r}:\n{completed.stdout}\n{completed.stderr}" + ) + return json.loads(completed.stdout) + + +def measure( + workspace: pathlib.Path, + binary: str, + pid: int, + object_path: str, +) -> dict: + selector = f"uprobe:{object_path}:symbol={MM_SYMBOL}" + completed = subprocess.run( + [ + binary, + "measure", + "--pid", + str(pid), + "--from", + f"{selector}:entry", + "--to", + f"{selector}:return", + "--match", + "stack-nested", + "--samples", + "8", + "--max-events", + "256", + "--timeout-ms", + "30000", + "--json", + "--non-interactive", + "--no-color", + ], + cwd=workspace, + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise AssertionError( + f"measure failed:\n{completed.stdout}\n{completed.stderr}" + ) + return json.loads(completed.stdout) + + +WORKLOAD = r""" +import json +import pathlib +import time + +import torch + +torch_root = pathlib.Path(torch.__file__).resolve().parent +print(json.dumps({ + "cpython": str(pathlib.Path("/proc/self/exe").resolve()), + "torch_extension": str(pathlib.Path(torch._C.__file__).resolve()), + "libtorch_cpu": str((torch_root / "lib" / "libtorch_cpu.so").resolve()), + "torch_version": torch.__version__, +}), flush=True) +a = torch.randn(64, 64) +b = torch.randn(64, 64) +while True: + torch.mm(a, b) + time.sleep(0.01) +""" + + +if __name__ == "__main__": + main() diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 7b3e85f..3b2c334 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -205,7 +205,7 @@ struct ResolveArgs { #[arg(long)] pid: u32, - /// Event selector: uprobe:::. + /// Event selector for a mapped ELF symbol, C++ signature, or file offset. #[arg(long, alias = "event")] selector: String, @@ -2182,6 +2182,7 @@ fn start_host_collectors( target: report.target.clone(), binary: PathBuf::from(&probe.binary_path), symbol: probe.symbol, + symbol_demangled: probe.symbol_demangled, offset, probe_kind: probe.probe_kind, probe_id: u32::try_from(index + 1).expect("two endpoints fit u32"), @@ -3168,6 +3169,7 @@ fn run_uprobe(args: UprobeArgs) -> ExitCode { target: report.target.clone(), binary, symbol: Some(symbol), + symbol_demangled: None, offset: 0, probe_kind: if return_probe { xprobe_protocol::HostProbeKind::Uretprobe @@ -3490,6 +3492,9 @@ fn print_resolved_probe(resolved: &ResolvedProbe) { if let Some(symbol) = &resolved.symbol { println!("Symbol: {symbol}"); } + if let Some(symbol) = &resolved.symbol_demangled { + println!("Demangled: {symbol}"); + } println!("File offset: {:#x}", resolved.file_offset); println!("Runtime address: {:#x}", resolved.runtime_address); println!( diff --git a/xprobe/cli/tests/measure.rs b/xprobe/cli/tests/measure.rs index c1d8e3b..709774a 100644 --- a/xprobe/cli/tests/measure.rs +++ b/xprobe/cli/tests/measure.rs @@ -96,6 +96,7 @@ fn write_host_capture(path: &PathBuf, timestamp_ns: u64) { binary_path: Some("/srv/libserver.so".to_owned()), build_id: None, symbol: Some("handle_request".to_owned()), + symbol_demangled: None, offset: None, return_value: None, arguments: Vec::new(), diff --git a/xprobe/cli/tests/resolve.rs b/xprobe/cli/tests/resolve.rs index ebe54dc..7717840 100644 --- a/xprobe/cli/tests/resolve.rs +++ b/xprobe/cli/tests/resolve.rs @@ -155,6 +155,22 @@ fn resolves_pie_shared_library_symbol_and_file_offset() { assert_eq!(library.probe_kind, HostProbeKind::Uretprobe); assert!(library.build_id.is_some()); + let cpp = resolve( + pid, + &format!( + "uprobe:{}:symbol=xprobe_fixture::native_operator(long):entry", + target.library.display() + ), + ); + assert_eq!( + cpp.symbol.as_deref(), + Some("_ZN14xprobe_fixture15native_operatorEl") + ); + assert_eq!( + cpp.symbol_demangled.as_deref(), + Some("xprobe_fixture::native_operator(long)") + ); + let by_offset = resolve( pid, &format!( diff --git a/xprobe/collector/src/linux.rs b/xprobe/collector/src/linux.rs index edce901..12ba494 100644 --- a/xprobe/collector/src/linux.rs +++ b/xprobe/collector/src/linux.rs @@ -522,6 +522,7 @@ fn normalize_event( binary_path: None, build_id: None, symbol: Some(probe.name.clone()), + symbol_demangled: None, offset: None, return_value: (probe.event_type == EventType::SyscallExit) .then_some(i64::from_ne_bytes(raw.values[0].to_ne_bytes())), diff --git a/xprobe/collector/src/uprobe.rs b/xprobe/collector/src/uprobe.rs index 73277e5..dde9c20 100644 --- a/xprobe/collector/src/uprobe.rs +++ b/xprobe/collector/src/uprobe.rs @@ -33,6 +33,7 @@ pub struct UprobeRequest { pub target: TargetIdentity, pub binary: PathBuf, pub symbol: Option, + pub symbol_demangled: Option, pub offset: u64, pub probe_kind: HostProbeKind, pub probe_id: u32, @@ -423,6 +424,7 @@ fn normalize_event( binary_path: Some(binary_path.to_owned()), build_id: None, symbol: request.symbol.clone(), + symbol_demangled: request.symbol_demangled.clone(), offset: request.symbol.is_none().then_some(request.offset), return_value: None, arguments: Vec::new(), @@ -492,6 +494,7 @@ mod tests { }, binary: PathBuf::from("/srv/server"), symbol: Some("handle_request".to_owned()), + symbol_demangled: None, offset: 0, probe_kind: HostProbeKind::Uretprobe, probe_id: 8, @@ -525,6 +528,7 @@ mod tests { }, binary: PathBuf::from("/srv/server"), symbol: Some("handle_request".to_owned()), + symbol_demangled: None, offset: 0, probe_kind: HostProbeKind::Kprobe, probe_id: 8, @@ -548,6 +552,7 @@ mod tests { }, binary: PathBuf::from("/srv/server"), symbol: None, + symbol_demangled: None, offset: 0x1234, probe_kind: HostProbeKind::Uprobe, probe_id: 8, diff --git a/xprobe/core/Cargo.toml b/xprobe/core/Cargo.toml index 6d1c256..6e95266 100644 --- a/xprobe/core/Cargo.toml +++ b/xprobe/core/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true rust-version.workspace = true [dependencies] +cpp_demangle.workspace = true object.workspace = true nix.workspace = true regex.workspace = true diff --git a/xprobe/core/src/resolve.rs b/xprobe/core/src/resolve.rs index aa9073d..c453a35 100644 --- a/xprobe/core/src/resolve.rs +++ b/xprobe/core/src/resolve.rs @@ -5,6 +5,7 @@ use std::{ path::{Path, PathBuf}, }; +use cpp_demangle::{DemangleOptions, Symbol as CppSymbol}; use object::{Object, ObjectKind, ObjectSegment, ObjectSymbol, SymbolKind}; use xprobe_protocol::{ ElfObjectKind, ErrorCode, HostProbeKind, ProcessMapping, ProcessReport, ResolvedProbe, @@ -216,27 +217,8 @@ pub fn run(report: &ProcessReport, selector_text: &str) -> Result { - let (address, size) = resolve_symbol(&file, name, &binary)?; - let offset = virtual_address_to_file_offset(&file, address).ok_or_else(|| { - ResolveError::OffsetNotLoadable { - offset: address, - path: binary.clone(), - } - })?; - (Some(name.clone()), Some(address), Some(size), offset) - } - ProbeTarget::FileOffset(offset) => { - if file_offset_to_virtual_address(&file, *offset).is_none() { - return Err(ResolveError::OffsetNotLoadable { - offset: *offset, - path: binary, - }); - } - (None, None, None, *offset) - } - }; + let (symbol, symbol_demangled, symbol_virtual_address, symbol_size, file_offset) = + resolve_probe_target(&file, &selector.target, &binary)?; let map = binary_maps .into_iter() .find(|region| region_contains_offset(region, file_offset)) @@ -261,6 +243,7 @@ pub fn run(report: &ProcessReport, selector_text: &str) -> Result Result { )); } }; - let (binary, target) = binary_and_target.rsplit_once(':').ok_or_else(|| { - ResolveError::InvalidSelector("expected binary path and symbol or offset".to_owned()) - })?; + let (binary, target) = if let Some(parts) = binary_and_target.split_once(":symbol=") { + parts + } else { + binary_and_target.rsplit_once(':').ok_or_else(|| { + ResolveError::InvalidSelector("expected binary path and symbol or offset".to_owned()) + })? + }; if binary.is_empty() || target.is_empty() { return Err(ResolveError::InvalidSelector( "binary path and probe target must not be empty".to_owned(), @@ -395,29 +382,135 @@ fn classify_object( } } +#[derive(Debug)] +struct ResolvedSymbol { + mangled: String, + demangled: Option, + address: u64, + size: u64, +} + +type ResolvedProbeTarget = ( + Option, + Option, + Option, + Option, + u64, +); + +fn resolve_probe_target( + file: &object::File<'_>, + target: &ProbeTarget, + path: &Path, +) -> Result { + match target { + ProbeTarget::Symbol(name) => { + let resolved = resolve_symbol(file, name, path)?; + let offset = + virtual_address_to_file_offset(file, resolved.address).ok_or_else(|| { + ResolveError::OffsetNotLoadable { + offset: resolved.address, + path: path.to_owned(), + } + })?; + Ok(( + Some(resolved.mangled), + resolved.demangled, + Some(resolved.address), + Some(resolved.size), + offset, + )) + } + ProbeTarget::FileOffset(offset) => { + if file_offset_to_virtual_address(file, *offset).is_none() { + return Err(ResolveError::OffsetNotLoadable { + offset: *offset, + path: path.to_owned(), + }); + } + Ok((None, None, None, None, *offset)) + } + } +} + fn resolve_symbol( file: &object::File<'_>, name: &str, path: &Path, -) -> Result<(u64, u64), ResolveError> { +) -> Result { + let mut matches = BTreeMap::new(); + collect_symbol_matches( + file.dynamic_symbols().chain(file.symbols()), + name, + false, + &mut matches, + ); + if let Some(resolved) = select_symbol_match(matches, name, path)? { + return Ok(resolved); + } + + let mut matches = BTreeMap::new(); + collect_symbol_matches(file.dynamic_symbols(), name, true, &mut matches); + if let Some(resolved) = select_symbol_match(matches, name, path)? { + return Ok(resolved); + } + let mut matches = BTreeMap::new(); - for symbol in file.symbols().chain(file.dynamic_symbols()) { - if symbol.is_definition() - && symbol.kind() == SymbolKind::Text - && symbol.name().ok() == Some(name) - { - matches - .entry(symbol.address()) - .and_modify(|size: &mut u64| *size = (*size).max(symbol.size())) - .or_insert(symbol.size()); + collect_symbol_matches(file.symbols(), name, true, &mut matches); + select_symbol_match(matches, name, path)?.ok_or_else(|| ResolveError::SymbolNotFound { + symbol: name.to_owned(), + path: path.to_owned(), + }) +} + +fn collect_symbol_matches<'data>( + symbols: impl Iterator>, + name: &str, + match_demangled: bool, + matches: &mut BTreeMap, +) { + for symbol in symbols { + if !symbol.is_definition() || symbol.kind() != SymbolKind::Text { + continue; + } + let Ok(mangled) = symbol.name() else { + continue; + }; + let demangled = if match_demangled { + demangle_cpp(mangled) + } else { + None + }; + let is_match = if match_demangled { + demangled.as_deref() == Some(name) + } else { + mangled == name + }; + if !is_match { + continue; } + matches + .entry(symbol.address()) + .and_modify(|resolved: &mut ResolvedSymbol| { + resolved.size = resolved.size.max(symbol.size()); + }) + .or_insert_with(|| ResolvedSymbol { + mangled: mangled.to_owned(), + demangled: demangled.or_else(|| demangle_cpp(mangled)), + address: symbol.address(), + size: symbol.size(), + }); } +} + +fn select_symbol_match( + matches: BTreeMap, + name: &str, + path: &Path, +) -> Result, ResolveError> { match matches.len() { - 0 => Err(ResolveError::SymbolNotFound { - symbol: name.to_owned(), - path: path.to_owned(), - }), - 1 => Ok(matches.into_iter().next().expect("one symbol match")), + 0 => Ok(None), + 1 => Ok(matches.into_values().next()), _ => Err(ResolveError::AmbiguousSymbol { symbol: name.to_owned(), path: path.to_owned(), @@ -425,6 +518,14 @@ fn resolve_symbol( } } +fn demangle_cpp(name: &str) -> Option { + if !name.starts_with("_Z") { + return None; + } + let symbol = CppSymbol::new(name).ok()?; + symbol.demangle(&DemangleOptions::default()).ok() +} + fn virtual_address_to_file_offset(file: &object::File<'_>, address: u64) -> Option { file.segments().find_map(|segment| { let delta = address.checked_sub(segment.address())?; @@ -475,6 +576,15 @@ mod tests { let offset = parse_selector("uprobe:/srv/app:+0x1234:entry").unwrap(); assert_eq!(offset.target, ProbeTarget::FileOffset(0x1234)); assert_eq!(offset.boundary, Boundary::Entry); + + let demangled = parse_selector( + "uprobe:/srv/libtorch_cpu.so:symbol=at::native::mm(at::Tensor const&):entry", + ) + .unwrap(); + assert_eq!( + demangled.target, + ProbeTarget::Symbol("at::native::mm(at::Tensor const&)".to_owned()) + ); } #[test] diff --git a/xprobe/core/src/validate.rs b/xprobe/core/src/validate.rs index ffaa55f..33ae4ad 100644 --- a/xprobe/core/src/validate.rs +++ b/xprobe/core/src/validate.rs @@ -1059,6 +1059,7 @@ mod tests { object_kind: ElfObjectKind::Executable, probe_kind, symbol: Some("work".to_owned()), + symbol_demangled: None, symbol_virtual_address: Some(0x1000), symbol_size: Some(16), file_offset: 0x1000, diff --git a/xprobe/correlator/src/lib.rs b/xprobe/correlator/src/lib.rs index 31674f1..4aee030 100644 --- a/xprobe/correlator/src/lib.rs +++ b/xprobe/correlator/src/lib.rs @@ -176,11 +176,15 @@ impl Selector { )); } }; - let (binary_path, target) = binary_and_target.rsplit_once(':').ok_or_else(|| { - MeasureError::InvalidSelector( - "uprobe selector requires a binary path and symbol or offset".to_owned(), - ) - })?; + let (binary_path, target) = if let Some(parts) = binary_and_target.split_once(":symbol=") { + parts + } else { + binary_and_target.rsplit_once(':').ok_or_else(|| { + MeasureError::InvalidSelector( + "uprobe selector requires a binary path and symbol or offset".to_owned(), + ) + })? + }; if binary_path.is_empty() || target.is_empty() { return Err(MeasureError::InvalidSelector( "uprobe binary path and target must not be empty".to_owned(), @@ -365,6 +369,7 @@ impl Selector { && match target { HostTarget::Symbol(symbol) => { host.symbol.as_deref() == Some(symbol.as_str()) + || host.symbol_demangled.as_deref() == Some(symbol.as_str()) } HostTarget::Offset(offset) => host.offset == Some(*offset), } @@ -1279,6 +1284,7 @@ mod tests { binary_path: Some("/srv/libserver.so".to_owned()), build_id: None, symbol: Some("handle_request".to_owned()), + symbol_demangled: None, offset: None, return_value: None, arguments: Vec::new(), @@ -1298,6 +1304,7 @@ mod tests { binary_path: None, build_id: None, symbol: Some(name.to_owned()), + symbol_demangled: None, offset: None, return_value: (event.event_type == EventType::SyscallExit).then_some(0), arguments: Vec::new(), @@ -1569,6 +1576,32 @@ mod tests { assert_eq!(result.correlation.confidence, CorrelationConfidence::High); } + #[test] + fn demangled_cpp_selectors_match_mangled_host_events() { + let mut entry = host_event(100); + entry.host.as_mut().unwrap().symbol = + Some("_ZN14xprobe_fixture15native_operatorEl".to_owned()); + entry.host.as_mut().unwrap().symbol_demangled = + Some("xprobe_fixture::native_operator(long)".to_owned()); + let mut exit = entry.clone(); + exit.timestamp_raw = 160; + exit.timestamp_ns = 160; + exit.event_type = EventType::HostFunctionExit; + exit.host.as_mut().unwrap().probe_kind = HostProbeKind::Uretprobe; + let mut options = options(MatchPolicy::StackNested); + options.start_selector = + "uprobe:/srv/libserver.so:symbol=xprobe_fixture::native_operator(long):entry" + .to_owned(); + options.end_selector = + "uprobe:/srv/libserver.so:symbol=xprobe_fixture::native_operator(long):return" + .to_owned(); + options.samples = Some(1); + + let result = measure(&[entry, exit], &options).unwrap(); + assert_eq!(result.measurement.samples.matched, 1); + assert_eq!(result.measurement.latency_ns.min, 60); + } + #[test] fn stream_order_never_pairs_across_cuda_streams() { let mut events = vec![ diff --git a/xprobe/protocol/src/event.rs b/xprobe/protocol/src/event.rs index 8561aba..b9cb5ff 100644 --- a/xprobe/protocol/src/event.rs +++ b/xprobe/protocol/src/event.rs @@ -86,6 +86,8 @@ pub struct HostEvent { pub binary_path: Option, pub build_id: Option, pub symbol: Option, + #[serde(default)] + pub symbol_demangled: Option, pub offset: Option, pub return_value: Option, #[serde(default)] diff --git a/xprobe/protocol/src/resolve.rs b/xprobe/protocol/src/resolve.rs index d112d34..ed15f87 100644 --- a/xprobe/protocol/src/resolve.rs +++ b/xprobe/protocol/src/resolve.rs @@ -31,6 +31,7 @@ pub struct ResolvedProbe { pub object_kind: ElfObjectKind, pub probe_kind: HostProbeKind, pub symbol: Option, + pub symbol_demangled: Option, pub symbol_virtual_address: Option, pub symbol_size: Option, pub file_offset: u64, diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index 5a1f305..474bbb0 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -39,6 +39,7 @@ fn event_contract_round_trips() { "binary_path": "/srv/app", "build_id": null, "symbol": "handle_request", + "symbol_demangled": null, "offset": 4096, "return_value": null, "arguments": [] @@ -80,6 +81,7 @@ fn host_capture_contract_round_trips() { "binary_path": "/srv/app", "build_id": null, "symbol": "handle_request", + "symbol_demangled": null, "offset": null, "return_value": null, "arguments": [] @@ -337,6 +339,7 @@ fn resolved_probe_contract_round_trips() { "object_kind": "shared_library", "probe_kind": "uprobe", "symbol": "handle_request", + "symbol_demangled": null, "symbol_virtual_address": 8192, "symbol_size": 32, "file_offset": 8192, From a8017f17a57d5ef4667caa4fa4b55adf0e2fb816 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 07:05:52 +0800 Subject: [PATCH 13/17] =?UTF-8?q?=E2=9A=A1=20perf:=20filter=20long=20frame?= =?UTF-8?q?work=20kernels?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cupti/src/cupti_agent.c | 9 +- docs/architecture.md | 5 +- docs/cli-contract.md | 4 + docs/development.md | 8 + justfile | 3 + .../aggregate-inventory-result.schema.json | 6 + .../references/cli-contract.md | 3 + .../references/investigation.md | 7 +- tests/integration/test_pytorch_cuda.py | 353 ++++++++++++++++++ xprobe/cli/src/main.rs | 40 +- xprobe/collector/src/cupti.rs | 31 +- xprobe/protocol/src/measurement.rs | 1 + xprobe/protocol/tests/contracts.rs | 1 + 13 files changed, 453 insertions(+), 18 deletions(-) create mode 100644 tests/integration/test_pytorch_cuda.py diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index ba9bad3..6ae86a0 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -374,12 +374,12 @@ static int name_matches(const struct xprobe_cupti_filter *filter, } filter_length = bounded_name_length(filter->name); record_length = bounded_name_length(record_name); - if (filter_length == XPROBE_CUPTI_NAME_LENGTH || - record_length == XPROBE_CUPTI_NAME_LENGTH) { + if (filter_length == XPROBE_CUPTI_NAME_LENGTH) { return 0; } if (filter->name_match == XPROBE_CUPTI_NAME_EXACT) { - return filter_length == record_length && + return record_length < XPROBE_CUPTI_NAME_LENGTH && + filter_length == record_length && memcmp(filter->name, record_name, filter_length) == 0; } if (filter->name_match == XPROBE_CUPTI_NAME_PREFIX) { @@ -387,7 +387,8 @@ static int name_matches(const struct xprobe_cupti_filter *filter, memcmp(filter->name, record_name, filter_length) == 0; } if (filter->name_match == XPROBE_CUPTI_NAME_SUFFIX) { - return filter_length <= record_length && + return record_length < XPROBE_CUPTI_NAME_LENGTH && + filter_length <= record_length && memcmp(filter->name, record_name + record_length - filter_length, filter_length) == 0; } diff --git a/docs/architecture.md b/docs/architecture.md index b72a338..0856f6a 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -103,7 +103,10 @@ Broad GPU inventory uses the same `measure` primitive with aggregate mode. The Agent updates a `--max-groups`-bounded table for matching kernel, memcpy, or memset activity and returns only final count/duration/byte summaries. Exact measurement remains the evidence path; aggregate output has a separate result -contract and cannot be exported as event JSONL. +contract and cannot be exported as event JSONL. Kernel names are retained as +fixed 127-byte observed prefixes. Complete names produce exact selector hints; +a full name buffer is marked incomplete and produces a prefix hint that remains +filterable before event reservation. CUPTI activity timestamps are normalized to `CLOCK_MONOTONIC` through its timestamp callback or an explicit CUDA 12 clock calibration. Activity that diff --git a/docs/cli-contract.md b/docs/cli-contract.md index 06bf8b3..bdbf2c5 100644 --- a/docs/cli-contract.md +++ b/docs/cli-contract.md @@ -181,6 +181,10 @@ inventory, not exact event evidence: it reports count, total/min/max/mean duration, optional transferred bytes, table occupancy, and drop completeness. Aggregate kernel regex must be reducible to an exact, prefix, suffix, or contains filter because this mode intentionally has no Rust-side event pass. +Each kernel group reports `name_complete`. CUPTI names that fill the fixed +127-byte observed prefix are marked `false`; their selector hints use that +prefix instead of claiming an exact full name, so the next capture can still +filter in the Agent hot path. Live host endpoints attach PID-scoped eBPF probes. Linux syscall endpoints use raw tracepoints so they do not depend on tracingfs event IDs; ordinary named diff --git a/docs/development.md b/docs/development.md index 220e71c..04343dd 100644 --- a/docs/development.md +++ b/docs/development.md @@ -79,6 +79,14 @@ BPF container: PYTORCH_ENV=/path/to/env just test-pytorch-live ``` +With a CUDA-enabled PyTorch environment, run eager matrix multiplication, +convolution, compiled Triton, bidirectional transfer, selected-kernel, and +stream-synchronization profiling on the local GPU: + +```bash +PYTORCH_ENV=/path/to/env just test-pytorch-cuda-live +``` + ## GPU checks Run host diagnostics outside restricted sandboxes when GPU device access is diff --git a/justfile b/justfile index f2ec107..1091eab 100644 --- a/justfile +++ b/justfile @@ -75,6 +75,9 @@ test-pytorch-symbols: build test-pytorch-live: build python3 tests/integration/test_pytorch.py "{{cuda_smoke_image}}" "${PYTORCH_ENV:?set PYTORCH_ENV to a Mamba environment with torch}" +test-pytorch-cuda-live: build + python3 tests/integration/test_pytorch_cuda.py --image "{{cuda12_devel_image}}" --pytorch-env "${PYTORCH_ENV:?set PYTORCH_ENV to a Mamba environment with torch}" + benchmark-gpu: python3 benchmarks/cuda-callback/run.py "{{cuda13_devel_image}}" diff --git a/schemas/aggregate-inventory-result.schema.json b/schemas/aggregate-inventory-result.schema.json index 5345973..ea4021b 100644 --- a/schemas/aggregate-inventory-result.schema.json +++ b/schemas/aggregate-inventory-result.schema.json @@ -172,6 +172,12 @@ "null" ] }, + "name_complete": { + "type": [ + "boolean", + "null" + ] + }, "start_selector_hint": { "type": "string" }, diff --git a/skills/xprobe-measure-latency/references/cli-contract.md b/skills/xprobe-measure-latency/references/cli-contract.md index 0240ca5..e7cfc7e 100644 --- a/skills/xprobe-measure-latency/references/cli-contract.md +++ b/skills/xprobe-measure-latency/references/cli-contract.md @@ -63,6 +63,9 @@ quality. It contains no event evidence, percentiles, or correlation confidence; `--samples`, `--input`, and `--events-out` are invalid in this mode. Kernel regex must be an exact, prefix, suffix, or contains shape that the Agent can apply before aggregation; other regex is rejected instead of widened. +Kernel groups expose `name_complete`. When it is false, the observed name is a +bounded prefix and the emitted selector hints deliberately remain prefix +selectors that the Agent can apply before reserving exact events. Kernel and other GPU activity durations require separate start and end records, so `max-events` is record capacity rather than sample capacity. Sample completion diff --git a/skills/xprobe-measure-latency/references/investigation.md b/skills/xprobe-measure-latency/references/investigation.md index 85e5cdf..b38443b 100644 --- a/skills/xprobe-measure-latency/references/investigation.md +++ b/skills/xprobe-measure-latency/references/investigation.md @@ -71,8 +71,11 @@ regex metacharacters and validate the final selector. For Triton and other JIT kernels, use captured names plus grid/block variants to identify a launch family, then correlate it with framework cache metadata, generated source, or application logs. xprobe does not read JIT cache contents -and cannot select by grid/block. Long mangled names should be narrowed to a -short, observed-unique literal instead of copied wholesale. +and cannot select by grid/block. Use each aggregate group's emitted selector +hint directly. A `name_complete: false` group contains a bounded observed prefix +and therefore emits a prefix hint; do not add an exact end anchor. When a hint +still covers several groups, derive a shorter observed-unique prefix or +contains literal and validate it before collection. ## Derive CPU selectors diff --git a/tests/integration/test_pytorch_cuda.py b/tests/integration/test_pytorch_cuda.py new file mode 100644 index 0000000..9f8ac0a --- /dev/null +++ b/tests/integration/test_pytorch_cuda.py @@ -0,0 +1,353 @@ +#!/usr/bin/env python3 +import argparse +import json +import pathlib +import subprocess +import sys +import tempfile +import time + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--image") + parser.add_argument("--pytorch-env") + parser.add_argument("--inner", action="store_true") + parser.add_argument("--xprobe", default="/workspace/target/debug/xprobe") + args = parser.parse_args() + if args.inner: + run_inner(pathlib.Path(args.xprobe)) + else: + run_container(args) + + +def run_container(args: argparse.Namespace) -> None: + if not args.image or not args.pytorch_env: + raise SystemExit("--image and --pytorch-env are required") + workspace = pathlib.Path(__file__).resolve().parents[2] + pytorch_env = pathlib.Path(args.pytorch_env).resolve() + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--gpus", + "all", + "--cap-add", + "SYS_PTRACE", + "--security-opt", + "seccomp=unconfined", + "--volume", + f"{workspace}:/workspace:ro", + "--volume", + f"{pytorch_env}:/opt/xprobe-pytorch:ro", + "--workdir", + "/workspace", + args.image, + "/opt/xprobe-pytorch/bin/python", + "/workspace/tests/integration/test_pytorch_cuda.py", + "--inner", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + raise SystemExit(completed.returncode) + sys.stdout.write(completed.stdout) + + +def run_inner(xprobe: pathlib.Path) -> None: + workspace = pathlib.Path("/workspace") + with tempfile.TemporaryDirectory(prefix="xprobe-pytorch-cuda-") as temp_dir: + temp = pathlib.Path(temp_dir) + agent = temp / "libxprobe-cupti.so" + mode = temp / "mode" + build_agent(workspace, agent) + set_mode(mode, "mm") + workload = subprocess.Popen( + [ + sys.executable, + "-u", + "-c", + WORKLOAD, + str(mode), + ], + stdout=subprocess.PIPE, + text=True, + ) + try: + assert workload.stdout is not None + ready = workload.stdout.readline() + if not ready: + raise AssertionError("PyTorch CUDA workload exited before readiness") + metadata = json.loads(ready) + pid = workload.pid + + mm = inventory(xprobe, agent, pid, mode, "mm", "kernel") + conv = inventory(xprobe, agent, pid, mode, "conv", "kernel") + compiled = inventory(xprobe, agent, pid, mode, "compiled", "kernel") + transfers = inventory(xprobe, agent, pid, mode, "transfer", "memcpy") + + mm_group = min(mm["inventory"]["groups"], key=lambda group: len(group["name"])) + exact_kernel = exact_measure( + xprobe, + agent, + pid, + mode, + "mm", + mm_group["start_selector_hint"], + mm_group["end_selector_hint"], + ) + synchronization = exact_measure( + xprobe, + agent, + pid, + mode, + "mm", + "cuda:runtime_api:cudaStreamSynchronize:entry", + "cuda:runtime_api:cudaStreamSynchronize:exit", + ) + finally: + workload.terminate() + workload.wait(timeout=10) + + assert_inventory(mm, "kernel") + assert_inventory(conv, "kernel") + assert_inventory(compiled, "kernel") + assert_inventory(transfers, "memcpy") + compiled_names = [ + group["name"].lower() for group in compiled["inventory"]["groups"] + ] + assert any("triton" in name for name in compiled_names), compiled_names + transfer_hints = { + group["start_selector_hint"] for group in transfers["inventory"]["groups"] + } + assert "cuda:memcpy_start:kind=HtoD" in transfer_hints + assert "cuda:memcpy_start:kind=DtoH" in transfer_hints + assert_exact(exact_kernel) + assert_exact(synchronization) + + print( + json.dumps( + { + "compiled_kernel_groups": len(compiled["inventory"]["groups"]), + "conv_kernel_groups": len(conv["inventory"]["groups"]), + "cuda": metadata["cuda"], + "device": metadata["device"], + "eager_mm_kernel_groups": len(mm["inventory"]["groups"]), + "exact_kernel_records": exact_kernel["collection"]["cupti"][ + "retained_records" + ], + "measured": ["kernel", "memcpy_htod", "memcpy_dtoh", "stream_sync"], + "ok": True, + "pytorch": metadata["pytorch"], + "stream_sync_records": synchronization["collection"]["cupti"][ + "retained_records" + ], + "triton": metadata["triton"], + }, + sort_keys=True, + ) + ) + + +def build_agent(workspace: pathlib.Path, output: pathlib.Path) -> None: + subprocess.run( + [ + "gcc", + "-std=c11", + "-D_GNU_SOURCE", + "-DXPROBE_HAS_CUPTI=1", + "-fPIC", + "-shared", + "-pthread", + "-O2", + "-Wall", + "-Wextra", + "-Wpedantic", + "-Werror", + f"-I{workspace / 'cupti/include'}", + "-isystem", + "/usr/local/cuda/include", + str(workspace / "cupti/src/cupti_agent.c"), + "-L/usr/local/cuda/lib64", + "-Wl,-rpath,/usr/local/cuda/lib64", + "-lcupti", + "-o", + str(output), + ], + check=True, + ) + + +def inventory( + xprobe: pathlib.Path, + agent: pathlib.Path, + pid: int, + mode_path: pathlib.Path, + mode: str, + activity: str, +) -> dict: + set_mode(mode_path, mode) + time.sleep(0.1) + return run_xprobe( + xprobe, + [ + "measure", + "--pid", + str(pid), + "--agent", + str(agent), + "--from", + f"cuda:{activity}_start", + "--to", + f"cuda:{activity}_end", + "--match", + "exact", + "--aggregate", + "--duration-ms", + "400", + "--max-groups", + "1024", + "--timeout-ms", + "30000", + ], + ) + + +def exact_measure( + xprobe: pathlib.Path, + agent: pathlib.Path, + pid: int, + mode_path: pathlib.Path, + mode: str, + start_selector: str, + end_selector: str, +) -> dict: + set_mode(mode_path, mode) + time.sleep(0.1) + return run_xprobe( + xprobe, + [ + "measure", + "--pid", + str(pid), + "--agent", + str(agent), + "--from", + start_selector, + "--to", + end_selector, + "--match", + "exact", + "--samples", + "8", + "--max-events", + "4096", + "--timeout-ms", + "30000", + ], + ) + + +def run_xprobe(xprobe: pathlib.Path, arguments: list[str]) -> dict: + completed = subprocess.run( + [ + str(xprobe), + *arguments, + "--json", + "--non-interactive", + "--no-color", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + raise AssertionError( + f"xprobe failed:\n{completed.stdout}\n{completed.stderr}" + ) + return json.loads(completed.stdout) + + +def set_mode(path: pathlib.Path, mode: str) -> None: + replacement = path.with_name(f"{path.name}.next") + replacement.write_text(mode) + replacement.replace(path) + + +def assert_inventory(result: dict, activity: str) -> None: + assert result["ok"] is True + assert result["status"] == "completed" + assert result["collection"]["dropped_activities"] == 0 + assert result["collection"]["observed_activities"] > 0 + assert result["inventory"]["groups"] + assert all(group["activity"] == activity for group in result["inventory"]["groups"]) + + +def assert_exact(result: dict) -> None: + assert result["ok"] is True + assert result["status"] == "completed" + assert result["measurement"]["samples"]["matched"] == 8 + assert result["collection"]["dropped_events"] == 0 + retained = result["collection"]["cupti"]["retained_records"] + assert retained <= 1024, retained + + +WORKLOAD = r""" +import json +import pathlib +import sys +import time + +import torch +import triton + +mode_path = pathlib.Path(sys.argv[1]) +torch.cuda.init() +device = torch.device("cuda") +stream = torch.cuda.current_stream() +a = torch.randn(256, 256, device=device) +b = torch.randn(256, 256, device=device) +image = torch.randn(8, 16, 64, 64, device=device) +conv = torch.nn.Conv2d(16, 32, 3, padding=1, device=device) +host_input = torch.randn(256, 256, pin_memory=True) +device_transfer = torch.empty_like(host_input, device=device) +host_output = torch.empty_like(host_input, pin_memory=True) + +@torch.compile +def compiled_step(left, right): + return torch.relu(torch.mm(left, right)) + +compiled_step(a, b) +stream.synchronize() +print(json.dumps({ + "cuda": torch.version.cuda, + "device": torch.cuda.get_device_name(), + "pytorch": torch.__version__, + "triton": triton.__version__, +}), flush=True) + +while True: + mode = mode_path.read_text().strip() + if mode == "mm": + torch.mm(a, b) + elif mode == "conv": + conv(image) + elif mode == "compiled": + compiled_step(a, b) + elif mode == "transfer": + device_transfer.copy_(host_input, non_blocking=True) + host_output.copy_(device_transfer, non_blocking=True) + else: + raise RuntimeError(f"unknown workload mode: {mode}") + stream.synchronize() + time.sleep(0.001) +""" + + +if __name__ == "__main__": + main() diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 3b2c334..3619ce7 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -1872,7 +1872,7 @@ fn bounded_kernel_name_filter(pattern: &str) -> cupti::CuptiNameFilter { let Some(literal) = regex_literal(inner) else { continue; }; - if literal.is_empty() || literal.len() >= 128 { + if literal.is_empty() || literal.len() >= cupti::CUPTI_NAME_CAPACITY { continue; } return match kind { @@ -2587,6 +2587,7 @@ fn protocol_aggregate_group(group: cupti::CuptiAggregateGroup) -> AggregateGroup cupti::CuptiAggregateActivity::Memset => AggregateActivity::Memset, }, name: group.name, + name_complete: group.name_complete, device_id: group.device_id, memcpy_kind: group.memcpy_kind, start_selector_hint, @@ -2607,9 +2608,14 @@ fn aggregate_selector_hints(group: &cupti::CuptiAggregateGroup) -> (String, Stri cupti::CuptiAggregateActivity::Kernel => { let name = group.name.as_deref().expect("kernel aggregate has a name"); let escaped = escape_selector_regex(name); + let suffix = if group.name_complete == Some(true) { + "$" + } else { + "" + }; ( - format!("cuda:kernel_start:name~^{escaped}$"), - format!("cuda:kernel_end:name~^{escaped}$"), + format!("cuda:kernel_start:name~^{escaped}{suffix}"), + format!("cuda:kernel_end:name~^{escaped}{suffix}"), ) } cupti::CuptiAggregateActivity::Memcpy => { @@ -3569,9 +3575,9 @@ const fn schema_version(version: SchemaVersion) -> &'static str { #[cfg(test)] mod tests { - use xprobe_collector::cupti::CuptiNameFilter; + use xprobe_collector::cupti::{CuptiAggregateActivity, CuptiAggregateGroup, CuptiNameFilter}; - use super::{bounded_kernel_name_filter, regex_literal}; + use super::{aggregate_selector_hints, bounded_kernel_name_filter, regex_literal}; #[test] fn lowers_only_equivalent_bounded_kernel_patterns() { @@ -3603,4 +3609,28 @@ mod tests { assert_eq!(regex_literal("kernel\\d"), None); assert_eq!(regex_literal("kernel+"), None); } + + #[test] + fn aggregate_hints_use_prefixes_for_incomplete_kernel_names() { + let group = CuptiAggregateGroup { + activity: CuptiAggregateActivity::Kernel, + name: Some("long_kernel_prefix".to_owned()), + name_complete: Some(false), + device_id: Some(0), + memcpy_kind: None, + count: 1, + total_duration_ns: 10, + min_duration_ns: 10, + max_duration_ns: 10, + total_bytes: None, + }; + + assert_eq!( + aggregate_selector_hints(&group), + ( + "cuda:kernel_start:name~^long_kernel_prefix".to_owned(), + "cuda:kernel_end:name~^long_kernel_prefix".to_owned(), + ) + ); + } } diff --git a/xprobe/collector/src/cupti.rs b/xprobe/collector/src/cupti.rs index 6af97d2..d27501f 100644 --- a/xprobe/collector/src/cupti.rs +++ b/xprobe/collector/src/cupti.rs @@ -26,7 +26,7 @@ const RECORD_SIZE_U32: u32 = 200; const CONTROL_SIZE: usize = 328; const FILTER_SIZE: usize = 144; const FILTER_COUNT: usize = 2; -const FILTER_NAME_SIZE: usize = 128; +pub const CUPTI_NAME_CAPACITY: usize = 128; const FEATURE_HOST_MONOTONIC_TIMESTAMPS: u32 = 1 << 0; const FEATURE_TRANSFER_RECORDS: u32 = 1 << 1; const FEATURE_AGGREGATE_RECORDS: u32 = 1 << 2; @@ -108,6 +108,7 @@ pub enum CuptiAggregateActivity { pub struct CuptiAggregateGroup { pub activity: CuptiAggregateActivity, pub name: Option, + pub name_complete: Option, pub device_id: Option, pub memcpy_kind: Option, pub count: u64, @@ -464,9 +465,9 @@ fn encode_control( CuptiNameFilter::Contains(name) => (4, name.as_str()), }; let name_bytes = name.as_bytes(); - if name_bytes.len() >= FILTER_NAME_SIZE || name_bytes.contains(&0) { + if name_bytes.len() >= CUPTI_NAME_CAPACITY || name_bytes.contains(&0) { return Err(CuptiSnapshotError::InvalidControl(format!( - "CUPTI filter name must be a NUL-free string shorter than {FILTER_NAME_SIZE} bytes" + "CUPTI filter name must be a NUL-free string shorter than {CUPTI_NAME_CAPACITY} bytes" ))); } control[offset + 12..offset + 16].copy_from_slice(&name_match.to_ne_bytes()); @@ -679,6 +680,7 @@ fn decode_aggregate_record( Ok(CuptiAggregateGroup { activity, name: has_name.then(|| name.to_owned()), + name_complete: has_name.then_some(name_length + 1 < CUPTI_NAME_CAPACITY), device_id: optional_unknown(read_u32(record, 44)), memcpy_kind, count, @@ -887,9 +889,9 @@ mod tests { use std::io::Cursor; use super::{ - CuptiApiDomain, CuptiArmConfig, CuptiCaptureMode, CuptiDecodeError, CuptiEventFilter, - CuptiMemcpyKind, CuptiNameFilter, CuptiRecordKind, HEADER_SIZE, OUTPUT_MAGIC, RECORD_SIZE, - decode_capture, encode_control, read_snapshot, + CUPTI_NAME_CAPACITY, CuptiApiDomain, CuptiArmConfig, CuptiCaptureMode, CuptiDecodeError, + CuptiEventFilter, CuptiMemcpyKind, CuptiNameFilter, CuptiRecordKind, HEADER_SIZE, + OUTPUT_MAGIC, RECORD_SIZE, decode_capture, encode_control, read_snapshot, }; use xprobe_protocol::{ClockDomain, EventSource, EventType, MemcpyKind}; @@ -1203,6 +1205,7 @@ mod tests { decoded.aggregate_groups[0].name.as_deref(), Some("vector_add") ); + assert_eq!(decoded.aggregate_groups[0].name_complete, Some(true)); assert_eq!(decoded.aggregate_groups[0].count, 4); assert_eq!( decoded.aggregate_groups[1].memcpy_kind, @@ -1211,6 +1214,22 @@ mod tests { assert_eq!(decoded.aggregate_groups[1].total_bytes, Some(4096)); } + #[test] + fn marks_a_full_name_buffer_as_possibly_truncated() { + let name = "x".repeat(CUPTI_NAME_CAPACITY - 1); + let decoded = decode_capture( + &aggregate_capture(&[aggregate_record(3, 1, 100, 0, 0, &name)], 1), + "xp_aggregate_name", + ) + .expect("aggregate capture must decode"); + + assert_eq!( + decoded.aggregate_groups[0].name.as_deref(), + Some(name.as_str()) + ); + assert_eq!(decoded.aggregate_groups[0].name_complete, Some(false)); + } + #[test] fn encodes_bounded_arm_filters() { let control = encode_control( diff --git a/xprobe/protocol/src/measurement.rs b/xprobe/protocol/src/measurement.rs index 3eaf4ea..e62ffc3 100644 --- a/xprobe/protocol/src/measurement.rs +++ b/xprobe/protocol/src/measurement.rs @@ -188,6 +188,7 @@ pub struct AggregateInventory { pub struct AggregateGroup { pub activity: AggregateActivity, pub name: Option, + pub name_complete: Option, pub device_id: Option, pub memcpy_kind: Option, pub start_selector_hint: String, diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index 474bbb0..761f0f9 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -195,6 +195,7 @@ fn aggregate_inventory_contract_round_trips() { "groups": [{ "activity": "kernel", "name": "vector_add", + "name_complete": true, "device_id": 0, "memcpy_kind": null, "start_selector_hint": "cuda:kernel_start:name~^vector_add$", From a510e64ae3bdc5b40fa8a1a547babe1021c2bb65 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 07:20:06 +0800 Subject: [PATCH 14/17] =?UTF-8?q?=E2=9C=A8=20feat:=20define=20bounded=20NV?= =?UTF-8?q?TX=20range=20events?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- schemas/event.schema.json | 52 +++++ schemas/host-capture.schema.json | 52 +++++ schemas/measurement-result.schema.json | 52 +++++ schemas/validate.schema.json | 16 +- xprobe/cli/tests/export.rs | 1 + xprobe/cli/tests/measure.rs | 1 + xprobe/collector/src/completed.rs | 1 + xprobe/collector/src/cupti.rs | 1 + xprobe/collector/src/linux.rs | 1 + xprobe/collector/src/uprobe.rs | 1 + xprobe/core/src/validate.rs | 160 ++++++++++++++- xprobe/correlator/src/lib.rs | 257 +++++++++++++++++++++++-- xprobe/exporter/src/lib.rs | 6 + xprobe/protocol/src/event.rs | 23 +++ xprobe/protocol/src/lib.rs | 2 +- xprobe/protocol/src/validate.rs | 5 + xprobe/protocol/tests/contracts.rs | 5 + 17 files changed, 616 insertions(+), 20 deletions(-) diff --git a/schemas/event.schema.json b/schemas/event.schema.json index 3c6142e..c19d114 100644 --- a/schemas/event.schema.json +++ b/schemas/event.schema.json @@ -45,6 +45,17 @@ } ] }, + "nvtx": { + "anyOf": [ + { + "$ref": "#/$defs/NvtxEvent" + }, + { + "type": "null" + } + ], + "default": null + }, "pid": { "type": "integer", "format": "uint32", @@ -326,6 +337,8 @@ "gpu_memcpy_end", "gpu_memset_start", "gpu_memset_end", + "nvtx_range_start", + "nvtx_range_end", "marker" ] }, @@ -411,6 +424,45 @@ "unknown" ] }, + "NvtxEvent": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "name_complete": { + "type": "boolean" + }, + "range_id": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "range_kind": { + "$ref": "#/$defs/NvtxRangeKind" + }, + "start_tid": { + "type": "integer", + "format": "uint32", + "minimum": 0 + } + }, + "additionalProperties": false, + "required": [ + "name", + "name_complete", + "range_kind", + "range_id", + "start_tid" + ] + }, + "NvtxRangeKind": { + "type": "string", + "enum": [ + "thread", + "process" + ] + }, "SchemaVersion": { "type": "string", "enum": [ diff --git a/schemas/host-capture.schema.json b/schemas/host-capture.schema.json index 2d4033c..0fd03bb 100644 --- a/schemas/host-capture.schema.json +++ b/schemas/host-capture.schema.json @@ -290,6 +290,17 @@ } ] }, + "nvtx": { + "anyOf": [ + { + "$ref": "#/$defs/NvtxEvent" + }, + { + "type": "null" + } + ], + "default": null + }, "pid": { "type": "integer", "format": "uint32", @@ -383,6 +394,8 @@ "gpu_memcpy_end", "gpu_memset_start", "gpu_memset_end", + "nvtx_range_start", + "nvtx_range_end", "marker" ] }, @@ -468,6 +481,45 @@ "unknown" ] }, + "NvtxEvent": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "name_complete": { + "type": "boolean" + }, + "range_id": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "range_kind": { + "$ref": "#/$defs/NvtxRangeKind" + }, + "start_tid": { + "type": "integer", + "format": "uint32", + "minimum": 0 + } + }, + "additionalProperties": false, + "required": [ + "name", + "name_complete", + "range_kind", + "range_id", + "start_tid" + ] + }, + "NvtxRangeKind": { + "type": "string", + "enum": [ + "thread", + "process" + ] + }, "SchemaVersion": { "type": "string", "enum": [ diff --git a/schemas/measurement-result.schema.json b/schemas/measurement-result.schema.json index b41eec2..82c3d32 100644 --- a/schemas/measurement-result.schema.json +++ b/schemas/measurement-result.schema.json @@ -420,6 +420,17 @@ } ] }, + "nvtx": { + "anyOf": [ + { + "$ref": "#/$defs/NvtxEvent" + }, + { + "type": "null" + } + ], + "default": null + }, "pid": { "type": "integer", "format": "uint32", @@ -513,6 +524,8 @@ "gpu_memcpy_end", "gpu_memset_start", "gpu_memset_end", + "nvtx_range_start", + "nvtx_range_end", "marker" ] }, @@ -691,6 +704,45 @@ "unknown" ] }, + "NvtxEvent": { + "type": "object", + "properties": { + "name": { + "type": "string" + }, + "name_complete": { + "type": "boolean" + }, + "range_id": { + "type": "integer", + "format": "uint64", + "minimum": 0 + }, + "range_kind": { + "$ref": "#/$defs/NvtxRangeKind" + }, + "start_tid": { + "type": "integer", + "format": "uint32", + "minimum": 0 + } + }, + "additionalProperties": false, + "required": [ + "name", + "name_complete", + "range_kind", + "range_id", + "start_tid" + ] + }, + "NvtxRangeKind": { + "type": "string", + "enum": [ + "thread", + "process" + ] + }, "SampleSummary": { "type": "object", "properties": { diff --git a/schemas/validate.schema.json b/schemas/validate.schema.json index b879380..60d0783 100644 --- a/schemas/validate.schema.json +++ b/schemas/validate.schema.json @@ -63,7 +63,8 @@ "enum": [ "not_required", "already_loaded", - "injection_required" + "injection_required", + "startup_required" ] }, "ElfObjectKind": { @@ -126,6 +127,8 @@ "gpu_memcpy_end", "gpu_memset_start", "gpu_memset_end", + "nvtx_range_start", + "nvtx_range_end", "marker" ] }, @@ -253,6 +256,13 @@ "type": "null" } ] + }, + "nvtx_name_regex": { + "type": [ + "string", + "null" + ], + "default": null } }, "additionalProperties": false, @@ -495,6 +505,10 @@ "needs_ebpf": { "type": "boolean" }, + "needs_nvtx": { + "type": "boolean", + "default": false + }, "target_mutation": { "type": "boolean" } diff --git a/xprobe/cli/tests/export.rs b/xprobe/cli/tests/export.rs index 3dae599..616f1e0 100644 --- a/xprobe/cli/tests/export.rs +++ b/xprobe/cli/tests/export.rs @@ -31,6 +31,7 @@ fn exports_event_jsonl_as_chrome_trace() { process_start_time: None, host: None, cuda: None, + nvtx: None, attributes: BTreeMap::new(), }; fs::write( diff --git a/xprobe/cli/tests/measure.rs b/xprobe/cli/tests/measure.rs index 709774a..511f807 100644 --- a/xprobe/cli/tests/measure.rs +++ b/xprobe/cli/tests/measure.rs @@ -102,6 +102,7 @@ fn write_host_capture(path: &PathBuf, timestamp_ns: u64) { arguments: Vec::new(), }), cuda: None, + nvtx: None, attributes: BTreeMap::new(), }; let capture = HostCaptureResult { diff --git a/xprobe/collector/src/completed.rs b/xprobe/collector/src/completed.rs index a147034..c347ff5 100644 --- a/xprobe/collector/src/completed.rs +++ b/xprobe/collector/src/completed.rs @@ -259,6 +259,7 @@ mod tests { process_start_time: Some(10), host: None, cuda: None, + nvtx: None, attributes: BTreeMap::new(), } } diff --git a/xprobe/collector/src/cupti.rs b/xprobe/collector/src/cupti.rs index d27501f..ba6f8dd 100644 --- a/xprobe/collector/src/cupti.rs +++ b/xprobe/collector/src/cupti.rs @@ -817,6 +817,7 @@ fn decode_record( bytes: is_transfer.then(|| read_split_u64(record, 44)), memcpy_kind: is_memcpy.then(|| decode_memcpy_kind(read_u32(record, 52))), }), + nvtx: None, attributes, }) } diff --git a/xprobe/collector/src/linux.rs b/xprobe/collector/src/linux.rs index 12ba494..5fc622c 100644 --- a/xprobe/collector/src/linux.rs +++ b/xprobe/collector/src/linux.rs @@ -529,6 +529,7 @@ fn normalize_event( arguments, }), cuda: None, + nvtx: None, attributes, }) } diff --git a/xprobe/collector/src/uprobe.rs b/xprobe/collector/src/uprobe.rs index dde9c20..9c8cbcf 100644 --- a/xprobe/collector/src/uprobe.rs +++ b/xprobe/collector/src/uprobe.rs @@ -430,6 +430,7 @@ fn normalize_event( arguments: Vec::new(), }), cuda: None, + nvtx: None, attributes: BTreeMap::new(), } } diff --git a/xprobe/core/src/validate.rs b/xprobe/core/src/validate.rs index 33ae4ad..79e7285 100644 --- a/xprobe/core/src/validate.rs +++ b/xprobe/core/src/validate.rs @@ -88,11 +88,12 @@ enum EndpointKind { Kernel, Memcpy, Memset, + NvtxRange { name_regex: String }, } impl EndpointKind { const fn needs_callback(&self) -> bool { - matches!(self, Self::CudaApi { .. }) + matches!(self, Self::CudaApi { .. } | Self::NvtxRange { .. }) } const fn needs_activity(&self) -> bool { @@ -119,8 +120,13 @@ impl EndpointKind { | Self::SyscallExit { .. } | Self::Tracepoint | Self::CudaApi { .. } + | Self::NvtxRange { .. } ) } + + const fn needs_nvtx(&self) -> bool { + matches!(self, Self::NvtxRange { .. }) + } } struct Endpoint { @@ -151,12 +157,15 @@ pub fn run( let needs_ebpf = start.kind.is_host() || end.kind.is_host(); let needs_cupti_callback = start.kind.needs_callback() || end.kind.needs_callback(); let needs_cupti_activity = start.kind.needs_activity() || end.kind.needs_activity(); + let needs_nvtx = start.kind.needs_nvtx() || end.kind.needs_nvtx(); let needs_cupti = needs_cupti_callback || needs_cupti_activity; let needs_clock_alignment = start.kind.uses_host_clock() != end.kind.uses_host_clock(); let agent_activation = if !needs_cupti { AgentActivation::NotRequired } else if report.cuda.xprobe_cupti_loaded { AgentActivation::AlreadyLoaded + } else if needs_nvtx { + AgentActivation::StartupRequired } else { AgentActivation::InjectionRequired }; @@ -165,6 +174,7 @@ pub fn run( needs_cupti, needs_cupti_callback, needs_cupti_activity, + needs_nvtx, needs_clock_alignment, agent_activation, target_mutation: agent_activation == AgentActivation::InjectionRequired, @@ -514,6 +524,7 @@ fn parse_cuda_selector( "kernel_start" | "kernel_end" => parse_kernel_selector(&fields), "memcpy_start" | "memcpy_end" => parse_memcpy_selector(&fields), "memset_start" | "memset_end" => parse_memset_selector(selector, &fields), + "nvtx_range_start" | "nvtx_range_end" => parse_nvtx_selector(&fields), "runtime_api" | "driver_api" => parse_api_selector(selector), event => Err(ValidateError::InvalidSelector(format!( "unsupported CUDA event {event:?}" @@ -556,6 +567,7 @@ fn parse_kernel_selector( api_domain: None, api_name: None, kernel_name_regex, + nvtx_name_regex: None, memcpy_kind: None, }, EndpointKind::Kernel, @@ -599,6 +611,7 @@ fn parse_memcpy_selector( api_domain: None, api_name: None, kernel_name_regex: None, + nvtx_name_regex: None, memcpy_kind, }, EndpointKind::Memcpy, @@ -626,6 +639,7 @@ fn parse_memset_selector( api_domain: None, api_name: None, kernel_name_regex: None, + nvtx_name_regex: None, memcpy_kind: None, }, EndpointKind::Memset, @@ -633,6 +647,89 @@ fn parse_memset_selector( )) } +fn parse_nvtx_selector( + fields: &[&str], +) -> Result<(ResolvedCudaSelector, EndpointKind, bool), ValidateError> { + let event_type = match fields[1] { + "nvtx_range_start" => EventType::NvtxRangeStart, + "nvtx_range_end" => EventType::NvtxRangeEnd, + _ => unreachable!("NVTX parser only receives NVTX range events"), + }; + let Some(filter) = fields.get(2) else { + return Err(ValidateError::InvalidSelector( + "NVTX range selector requires name~".to_owned(), + )); + }; + let pattern = filter.strip_prefix("name~").ok_or_else(|| { + ValidateError::InvalidSelector( + "NVTX range filter must use name~".to_owned(), + ) + })?; + Regex::new(pattern).map_err(|error| { + ValidateError::InvalidSelector(format!( + "invalid NVTX range name regular expression: {error}" + )) + })?; + if !is_bounded_name_regex(pattern) { + return Err(ValidateError::InvalidSelector( + "NVTX range name must reduce to one nonempty exact, prefix, suffix, or contains filter shorter than 128 bytes".to_owned(), + )); + } + Ok(( + ResolvedCudaSelector { + event_type, + api_domain: None, + api_name: None, + kernel_name_regex: None, + nvtx_name_regex: Some(pattern.to_owned()), + memcpy_kind: None, + }, + EndpointKind::NvtxRange { + name_regex: pattern.to_owned(), + }, + true, + )) +} + +fn is_bounded_name_regex(pattern: &str) -> bool { + const CANDIDATES: [(&str, &str); 8] = [ + ("^", "$"), + ("^", ".*"), + (".*", "$"), + (".*", ".*"), + ("^", ""), + ("", "$"), + (".*", ""), + ("", ""), + ]; + CANDIDATES.iter().any(|(prefix, suffix)| { + pattern + .strip_prefix(prefix) + .and_then(|value| value.strip_suffix(suffix)) + .and_then(regex_literal) + .is_some_and(|literal| !literal.is_empty() && literal.len() < 128) + }) +} + +fn regex_literal(pattern: &str) -> Option { + let mut literal = String::with_capacity(pattern.len()); + let mut characters = pattern.chars(); + while let Some(character) = characters.next() { + if character == '\\' { + let escaped = characters.next()?; + if escaped.is_ascii_alphanumeric() { + return None; + } + literal.push(escaped); + } else if ".+*?()|[]{}^$".contains(character) { + return None; + } else { + literal.push(character); + } + } + Some(literal) +} + fn parse_api_selector( selector: &str, ) -> Result<(ResolvedCudaSelector, EndpointKind, bool), ValidateError> { @@ -661,6 +758,7 @@ fn parse_api_selector( api_domain: Some(domain.clone()), api_name: Some(name.clone()), kernel_name_regex: None, + nvtx_name_regex: None, memcpy_kind: None, }, EndpointKind::CudaApi { domain, name }, @@ -725,12 +823,21 @@ fn check_capabilities( )); } } - if requirements.agent_activation == AgentActivation::InjectionRequired { - warnings.push(warning( + match requirements.agent_activation { + AgentActivation::InjectionRequired => warnings.push(warning( "TARGET_PROCESS_WILL_BE_MODIFIED", "measure must inject the xprobe CUPTI agent into the target process", - )); - } else { + )), + AgentActivation::StartupRequired => issues.push(issue( + ErrorCode::CuptiAgentNotLoaded, + "NVTX ranges cannot be enabled by online attach after NVTX initialization; restart the target with NVTX_INJECTION64_PATH set to the xprobe CUPTI agent".to_owned(), + )), + AgentActivation::NotRequired | AgentActivation::AlreadyLoaded => {} + } + if matches!( + requirements.agent_activation, + AgentActivation::NotRequired | AgentActivation::AlreadyLoaded + ) { if requirements.needs_cupti_callback && !report.capabilities.cuda_callback { issues.push(issue( ErrorCode::CuptiNotAvailable, @@ -798,6 +905,14 @@ fn supports_exact(start: &EndpointKind, end: &EndpointKind) -> bool { | (EndpointKind::Memcpy, EndpointKind::Memcpy) | (EndpointKind::Memset, EndpointKind::Memset) => true, ( + EndpointKind::NvtxRange { + name_regex: start_name, + }, + EndpointKind::NvtxRange { + name_regex: end_name, + }, + ) + | ( EndpointKind::SyscallEntry { name: start_name }, EndpointKind::SyscallExit { name: end_name }, ) => start_name == end_name, @@ -896,6 +1011,12 @@ mod tests { let (memset, _, collectable) = parse_cuda_selector("cuda:memset_start").unwrap(); assert_eq!(memset.event_type, EventType::GpuMemsetStart); assert!(collectable); + + let (nvtx, _, collectable) = + parse_cuda_selector("cuda:nvtx_range_start:name~^forward$").unwrap(); + assert_eq!(nvtx.event_type, EventType::NvtxRangeStart); + assert_eq!(nvtx.nvtx_name_regex.as_deref(), Some("^forward$")); + assert!(collectable); assert_eq!( parse_match_policy("first-after").unwrap(), MatchPolicy::FirstAfter @@ -906,6 +1027,8 @@ mod tests { fn rejects_invalid_filters_and_policies() { assert!(parse_cuda_selector("cuda:kernel_start:name~[").is_err()); assert!(parse_cuda_selector("cuda:memcpy_start:kind=sideways").is_err()); + assert!(parse_cuda_selector("cuda:nvtx_range_start").is_err()); + assert!(parse_cuda_selector("cuda:nvtx_range_start:name~(forward|backward)").is_err()); assert!(parse_match_policy("probably-near").is_err()); } @@ -1038,6 +1161,33 @@ mod tests { assert!(result.issues.is_empty()); } + #[test] + fn reports_nvtx_startup_requirement_before_agent_is_loaded() { + let report = inspect::run(std::process::id()).unwrap(); + let result = run( + &report, + "cuda:nvtx_range_start:name~^forward$", + "cuda:nvtx_range_end:name~^forward$", + "exact", + ) + .unwrap(); + + assert!(!result.valid); + assert!(result.requirements.needs_nvtx); + assert_eq!( + result.requirements.agent_activation, + AgentActivation::StartupRequired + ); + assert!(!result.requirements.target_mutation); + assert!( + result + .issues + .iter() + .any(|issue| issue.code == xprobe_protocol::ErrorCode::CuptiAgentNotLoaded) + ); + assert_eq!(result.policy_recommendation.policy, MatchPolicy::Exact); + } + #[test] fn recommends_stack_matching_for_host_function_spans() { let endpoint = |selector: &str, event_type, kind, probe_kind| Endpoint { diff --git a/xprobe/correlator/src/lib.rs b/xprobe/correlator/src/lib.rs index 4aee030..dab8297 100644 --- a/xprobe/correlator/src/lib.rs +++ b/xprobe/correlator/src/lib.rs @@ -6,8 +6,8 @@ use regex::Regex; use xprobe_protocol::{ CaptureCompleteness, ClockDomain, ClockQuality, CollectionSummary, CorrelationConfidence, CorrelationSummary, ErrorCode, Event, EventSource, EventType, HostProbeKind, LatencyStatistics, - MatchPolicy, MatchedEventPair, Measurement, MeasurementResult, MemcpyKind, SampleSummary, - SchemaVersion, SessionStatus, Warning, + MatchPolicy, MatchedEventPair, Measurement, MeasurementResult, MemcpyKind, NvtxRangeKind, + SampleSummary, SchemaVersion, SessionStatus, Warning, }; #[derive(Debug, Clone)] @@ -121,6 +121,11 @@ enum Selector { Memset { event_type: EventType, }, + Nvtx { + event_type: EventType, + name: Regex, + pattern: String, + }, } #[derive(Debug, PartialEq, Eq)] @@ -152,6 +157,7 @@ impl Selector { "kernel_start" | "kernel_end" => Self::parse_kernel(&fields), "memcpy_start" | "memcpy_end" => Self::parse_memcpy(&fields), "memset_start" | "memset_end" => Self::parse_memset(text, &fields), + "nvtx_range_start" | "nvtx_range_end" => Self::parse_nvtx(&fields), event => Err(MeasureError::InvalidSelector(format!( "event {event:?} is not present in completed CUPTI captures" ))), @@ -354,6 +360,39 @@ impl Selector { Ok(Self::Memset { event_type }) } + fn parse_nvtx(fields: &[&str]) -> Result { + let event_type = match fields[1] { + "nvtx_range_start" => EventType::NvtxRangeStart, + "nvtx_range_end" => EventType::NvtxRangeEnd, + _ => unreachable!("NVTX parser only receives NVTX range selectors"), + }; + let Some(filter) = fields.get(2) else { + return Err(MeasureError::InvalidSelector( + "NVTX range selector requires name~".to_owned(), + )); + }; + let pattern = filter.strip_prefix("name~").ok_or_else(|| { + MeasureError::InvalidSelector( + "NVTX range filter must use name~".to_owned(), + ) + })?; + if !is_bounded_name_regex(pattern) { + return Err(MeasureError::InvalidSelector( + "NVTX range name must reduce to one nonempty exact, prefix, suffix, or contains filter shorter than 128 bytes".to_owned(), + )); + } + let name = Regex::new(pattern).map_err(|error| { + MeasureError::InvalidSelector(format!( + "invalid NVTX range name regular expression: {error}" + )) + })?; + Ok(Self::Nvtx { + event_type, + name, + pattern: pattern.to_owned(), + }) + } + fn matches(&self, event: &Event) -> bool { match self { Self::Host { @@ -432,6 +471,15 @@ impl Selector { }) } Self::Memset { event_type } => event.event_type == *event_type, + Self::Nvtx { + event_type, name, .. + } => { + event.event_type == *event_type + && event + .nvtx + .as_ref() + .is_some_and(|nvtx| name.is_match(&nvtx.name)) + } } } @@ -447,8 +495,20 @@ impl Selector { name: end_name, }, ) => start_name == end_name, - (Self::Host { .. } | Self::Tracepoint { .. }, _) - | (_, Self::Host { .. } | Self::Tracepoint { .. }) => false, + ( + Self::Nvtx { + event_type: EventType::NvtxRangeStart, + pattern: start_pattern, + .. + }, + Self::Nvtx { + event_type: EventType::NvtxRangeEnd, + pattern: end_pattern, + .. + }, + ) => start_pattern == end_pattern, + (Self::Host { .. } | Self::Tracepoint { .. } | Self::Nvtx { .. }, _) + | (_, Self::Host { .. } | Self::Tracepoint { .. } | Self::Nvtx { .. }) => false, _ => true, } } @@ -469,6 +529,24 @@ impl Selector { ) } + fn is_nvtx_lifecycle(&self, end: &Self) -> bool { + matches!( + (self, end), + ( + Self::Nvtx { + event_type: EventType::NvtxRangeStart, + pattern: start_pattern, + .. + }, + Self::Nvtx { + event_type: EventType::NvtxRangeEnd, + pattern: end_pattern, + .. + }, + ) if start_pattern == end_pattern + ) + } + fn supports_stack_nested(&self, end: &Self) -> bool { match (self, end) { ( @@ -497,6 +575,45 @@ impl Selector { } } +fn is_bounded_name_regex(pattern: &str) -> bool { + const CANDIDATES: [(&str, &str); 8] = [ + ("^", "$"), + ("^", ".*"), + (".*", "$"), + (".*", ".*"), + ("^", ""), + ("", "$"), + (".*", ""), + ("", ""), + ]; + CANDIDATES.iter().any(|(prefix, suffix)| { + pattern + .strip_prefix(prefix) + .and_then(|value| value.strip_suffix(suffix)) + .and_then(regex_literal) + .is_some_and(|literal| !literal.is_empty() && literal.len() < 128) + }) +} + +fn regex_literal(pattern: &str) -> Option { + let mut literal = String::with_capacity(pattern.len()); + let mut characters = pattern.chars(); + while let Some(character) = characters.next() { + if character == '\\' { + let escaped = characters.next()?; + if escaped.is_ascii_alphanumeric() { + return None; + } + literal.push(escaped); + } else if ".+*?()|[]{}^$".contains(character) { + return None; + } else { + literal.push(character); + } + } + Some(literal) +} + fn valid_linux_component(value: &str) -> bool { !value.is_empty() && value @@ -558,6 +675,7 @@ pub fn measure( apply_duration_window(&mut starts, &mut ends, options.duration)?; let syscall_lifecycle = start_selector.is_syscall_lifecycle(&end_selector); + let nvtx_lifecycle = start_selector.is_nvtx_lifecycle(&end_selector); let outcome = correlate_selected( &starts, &ends, @@ -576,14 +694,11 @@ pub fn measure( .map(|pair| pair.latency_ns) .collect::>(); let denominator = matched + outcome.unmatched_start + outcome.unmatched_end + outcome.ambiguous; - let (method, confidence) = correlation_metadata(options.match_policy, syscall_lifecycle); + let (method, confidence) = + correlation_metadata(options.match_policy, syscall_lifecycle, nvtx_lifecycle); let warnings = measurement_warnings(options, &starts, &ends); - let host_events = events - .iter() - .filter(|event| event.source == EventSource::Ebpf) - .count() as u64; - let cuda_events = events.len() as u64 - host_events; + let (host_events, cuda_events) = collection_event_counts(events); Ok(MeasurementResult { schema_version: SchemaVersion::current(), ok: true, @@ -634,6 +749,14 @@ pub fn measure( }) } +fn collection_event_counts(events: &[Event]) -> (u64, u64) { + let host = events + .iter() + .filter(|event| event.source == EventSource::Ebpf) + .count() as u64; + (host, events.len() as u64 - host) +} + fn correlate_selected<'a>( starts: &[&'a Event], ends: &[&'a Event], @@ -677,11 +800,15 @@ fn validate_policy( const fn correlation_metadata( policy: MatchPolicy, syscall_lifecycle: bool, + nvtx_lifecycle: bool, ) -> (&'static str, CorrelationConfidence) { match policy { MatchPolicy::Exact if syscall_lifecycle => { ("exact_syscall_tid_lifecycle", CorrelationConfidence::Exact) } + MatchPolicy::Exact if nvtx_lifecycle => { + ("exact_nvtx_range_id", CorrelationConfidence::Exact) + } MatchPolicy::Exact => ("exact_cupti_correlation_id", CorrelationConfidence::Exact), MatchPolicy::FirstAfter => ("first_after", CorrelationConfidence::Heuristic), MatchPolicy::Nearest => ("nearest", CorrelationConfidence::Heuristic), @@ -1136,11 +1263,34 @@ fn group_by_stream<'a>( unmatched } -type CorrelationKey = (u32, u32, u32); +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum CorrelationKey { + Cuda { + pid: u32, + context_id: u32, + correlation_id: u32, + }, + Nvtx { + pid: u32, + range_kind: NvtxRangeKind, + range_id: u64, + }, +} fn correlation_key(event: &Event) -> Option { + if let Some(nvtx) = event.nvtx.as_ref() { + return Some(CorrelationKey::Nvtx { + pid: event.pid, + range_kind: nvtx.range_kind, + range_id: nvtx.range_id, + }); + } let cuda = event.cuda.as_ref()?; - Some((event.pid, cuda.context_id?, cuda.correlation_id?)) + Some(CorrelationKey::Cuda { + pid: event.pid, + context_id: cuda.context_id?, + correlation_id: cuda.correlation_id?, + }) } fn statistics(values: &[u64]) -> LatencyStatistics { @@ -1203,7 +1353,7 @@ mod tests { use xprobe_protocol::{ ClockDomain, CorrelationConfidence, CudaEvent, Event, EventSource, EventType, HostEvent, - HostProbeKind, MatchPolicy, MemcpyKind, SchemaVersion, + HostProbeKind, MatchPolicy, MemcpyKind, NvtxEvent, NvtxRangeKind, SchemaVersion, }; use super::{MeasureError, MeasureOptions, measure}; @@ -1242,10 +1392,34 @@ mod tests { bytes: None, memcpy_kind: None, }), + nvtx: None, attributes: BTreeMap::new(), } } + fn nvtx_event( + event_type: EventType, + timestamp: u64, + tid: u32, + range_kind: NvtxRangeKind, + range_id: u64, + start_tid: u32, + ) -> Event { + let mut event = event(event_type, timestamp, 0); + event.source = EventSource::CuptiCallback; + event.clock_domain = ClockDomain::HostMonotonic; + event.tid = tid; + event.cuda = None; + event.nvtx = Some(NvtxEvent { + name: "forward".to_owned(), + name_complete: true, + range_kind, + range_id, + start_tid, + }); + event + } + fn options(policy: MatchPolicy) -> MeasureOptions { MeasureOptions { session_id: "xp_test".to_owned(), @@ -1372,6 +1546,63 @@ mod tests { assert_eq!(result.evidence[0].end.tid, 10); } + #[test] + fn exact_nvtx_matching_uses_kind_and_range_id() { + let events = vec![ + nvtx_event( + EventType::NvtxRangeStart, + 100, + 10, + NvtxRangeKind::Thread, + 7, + 10, + ), + nvtx_event( + EventType::NvtxRangeStart, + 110, + 20, + NvtxRangeKind::Process, + 7, + 20, + ), + nvtx_event( + EventType::NvtxRangeEnd, + 160, + 30, + NvtxRangeKind::Process, + 7, + 20, + ), + nvtx_event( + EventType::NvtxRangeEnd, + 180, + 10, + NvtxRangeKind::Thread, + 7, + 10, + ), + ]; + let options = MeasureOptions { + session_id: "xp_nvtx".to_owned(), + name: Some("forward".to_owned()), + start_selector: "cuda:nvtx_range_start:name~^forward$".to_owned(), + end_selector: "cuda:nvtx_range_end:name~^forward$".to_owned(), + match_policy: MatchPolicy::Exact, + samples: Some(2), + duration: None, + max_events: 100, + dropped_events: 0, + }; + + let result = measure(&events, &options).unwrap(); + assert_eq!(result.measurement.samples.matched, 2); + assert_eq!(result.measurement.latency_ns.min, 50); + assert_eq!(result.measurement.latency_ns.max, 80); + assert_eq!(result.correlation.method, "exact_nvtx_range_id"); + assert_eq!(result.evidence[0].end.tid, 10); + assert_eq!(result.evidence[1].end.tid, 30); + } + #[test] fn exact_syscall_matching_does_not_cross_process_identity() { let mut exit = syscall_event(EventType::SyscallExit, 150, 10, "munmap"); diff --git a/xprobe/exporter/src/lib.rs b/xprobe/exporter/src/lib.rs index 0dff53c..69555da 100644 --- a/xprobe/exporter/src/lib.rs +++ b/xprobe/exporter/src/lib.rs @@ -52,6 +52,7 @@ fn chrome_event(event: &Event) -> Value { "correlation_id": cuda.and_then(|payload| payload.correlation_id), "context_id": cuda.and_then(|payload| payload.context_id), "stream_id": cuda.and_then(|payload| payload.stream_id), + "nvtx": event.nvtx, "attributes": event.attributes, } }) @@ -74,6 +75,7 @@ fn event_name(event: &Event) -> &str { .as_ref() .and_then(|cuda| cuda.kernel_name.as_deref()) }) + .or_else(|| event.nvtx.as_ref().map(|nvtx| nvtx.name.as_str())) .unwrap_or_else(|| event_type_name(&event.event_type)) } @@ -103,6 +105,8 @@ const fn event_type_name(event_type: &EventType) -> &'static str { EventType::GpuMemcpyEnd => "gpu_memcpy_end", EventType::GpuMemsetStart => "gpu_memset_start", EventType::GpuMemsetEnd => "gpu_memset_end", + EventType::NvtxRangeStart => "nvtx_range_start", + EventType::NvtxRangeEnd => "nvtx_range_end", EventType::Marker => "marker", } } @@ -134,6 +138,7 @@ mod tests { process_start_time: None, host: None, cuda: None, + nvtx: None, attributes: BTreeMap::new(), }; @@ -163,6 +168,7 @@ mod tests { process_start_time: None, host: None, cuda: None, + nvtx: None, attributes: BTreeMap::new(), }; diff --git a/xprobe/protocol/src/event.rs b/xprobe/protocol/src/event.rs index b9cb5ff..acdc9d9 100644 --- a/xprobe/protocol/src/event.rs +++ b/xprobe/protocol/src/event.rs @@ -33,6 +33,8 @@ pub enum EventType { GpuMemcpyEnd, GpuMemsetStart, GpuMemsetEnd, + NvtxRangeStart, + NvtxRangeEnd, Marker, } @@ -64,6 +66,8 @@ pub struct Event { pub host: Option, pub cuda: Option, #[serde(default)] + pub nvtx: Option, + #[serde(default)] pub attributes: BTreeMap, } @@ -123,6 +127,25 @@ pub struct CudaEvent { pub memcpy_kind: Option, } +#[derive( + Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, JsonSchema, +)] +#[serde(rename_all = "snake_case")] +pub enum NvtxRangeKind { + Thread, + Process, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] +#[serde(deny_unknown_fields)] +pub struct NvtxEvent { + pub name: String, + pub name_complete: bool, + pub range_kind: NvtxRangeKind, + pub range_id: u64, + pub start_tid: u32, +} + #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] #[serde(deny_unknown_fields)] pub struct Dim3 { diff --git a/xprobe/protocol/src/lib.rs b/xprobe/protocol/src/lib.rs index 99533f1..963155d 100644 --- a/xprobe/protocol/src/lib.rs +++ b/xprobe/protocol/src/lib.rs @@ -21,7 +21,7 @@ pub use discover::{CudaProcessCandidate, DiscoveryResult}; pub use error::{ErrorCode, ErrorResponse, XprobeError}; pub use event::{ ArgumentValue, ClockDomain, CudaEvent, Dim3, Event, EventSource, EventType, HostEvent, - HostProbeKind, MemcpyKind, + HostProbeKind, MemcpyKind, NvtxEvent, NvtxRangeKind, }; pub use export::{ExportFormat, TraceExportResult}; pub use measurement::{ diff --git a/xprobe/protocol/src/validate.rs b/xprobe/protocol/src/validate.rs index 5d49a36..e8c21d1 100644 --- a/xprobe/protocol/src/validate.rs +++ b/xprobe/protocol/src/validate.rs @@ -19,6 +19,7 @@ pub enum AgentActivation { NotRequired, AlreadyLoaded, InjectionRequired, + StartupRequired, } #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, JsonSchema)] @@ -28,6 +29,8 @@ pub struct ResolvedCudaSelector { pub api_domain: Option, pub api_name: Option, pub kernel_name_regex: Option, + #[serde(default)] + pub nvtx_name_regex: Option, pub memcpy_kind: Option, } @@ -61,6 +64,8 @@ pub struct ValidationRequirements { pub needs_cupti: bool, pub needs_cupti_callback: bool, pub needs_cupti_activity: bool, + #[serde(default)] + pub needs_nvtx: bool, pub needs_clock_alignment: bool, pub agent_activation: AgentActivation, pub target_mutation: bool, diff --git a/xprobe/protocol/tests/contracts.rs b/xprobe/protocol/tests/contracts.rs index 761f0f9..3b80afb 100644 --- a/xprobe/protocol/tests/contracts.rs +++ b/xprobe/protocol/tests/contracts.rs @@ -45,6 +45,7 @@ fn event_contract_round_trips() { "arguments": [] }, "cuda": null, + "nvtx": null, "attributes": {} })); } @@ -87,6 +88,7 @@ fn host_capture_contract_round_trips() { "arguments": [] }, "cuda": null, + "nvtx": null, "attributes": {} }] })); @@ -372,6 +374,7 @@ fn validation_result_contract_round_trips() { "api_domain": "runtime_api", "api_name": "cudaLaunchKernel", "kernel_name_regex": null, + "nvtx_name_regex": null, "memcpy_kind": null } }, @@ -387,6 +390,7 @@ fn validation_result_contract_round_trips() { "api_domain": null, "api_name": null, "kernel_name_regex": "flash.*", + "nvtx_name_regex": null, "memcpy_kind": null } }, @@ -401,6 +405,7 @@ fn validation_result_contract_round_trips() { "needs_cupti": true, "needs_cupti_callback": true, "needs_cupti_activity": true, + "needs_nvtx": false, "needs_clock_alignment": true, "agent_activation": "injection_required", "target_mutation": true From 36013b310fca1d7afa9f99d12d2f4af1fe2b3b17 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 08:16:17 +0800 Subject: [PATCH 15/17] =?UTF-8?q?=E2=9C=A8=20feat:=20collect=20bounded=20N?= =?UTF-8?q?VTX=20ranges?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- cupti/include/xprobe/cupti_agent.h | 14 +- cupti/src/cupti_agent.c | 594 ++++++++++++++++++++++++- cupti/tests/nvtx_range_fixture.cpp | 92 ++++ justfile | 6 + tests/integration/run-nvtx-live.sh | 91 ++++ tests/integration/test_nvtx.py | 104 +++++ tests/integration/test_pytorch_cuda.py | 36 +- xprobe/cli/src/main.rs | 21 +- xprobe/collector/src/cupti.rs | 205 +++++++-- xprobe/correlator/src/lib.rs | 69 +++ 10 files changed, 1185 insertions(+), 47 deletions(-) create mode 100644 cupti/tests/nvtx_range_fixture.cpp create mode 100755 tests/integration/run-nvtx-live.sh create mode 100755 tests/integration/test_nvtx.py diff --git a/cupti/include/xprobe/cupti_agent.h b/cupti/include/xprobe/cupti_agent.h index a5c5fc3..d48d031 100644 --- a/cupti/include/xprobe/cupti_agent.h +++ b/cupti/include/xprobe/cupti_agent.h @@ -18,7 +18,8 @@ extern "C" { enum xprobe_cupti_feature { XPROBE_CUPTI_FEATURE_HOST_MONOTONIC_TIMESTAMPS = 1U << 0, XPROBE_CUPTI_FEATURE_TRANSFER_RECORDS = 1U << 1, - XPROBE_CUPTI_FEATURE_AGGREGATE_RECORDS = 1U << 2 + XPROBE_CUPTI_FEATURE_AGGREGATE_RECORDS = 1U << 2, + XPROBE_CUPTI_FEATURE_NVTX_RECORDS = 1U << 3 }; enum xprobe_cupti_capture_mode { @@ -91,7 +92,9 @@ enum xprobe_cupti_record_kind { XPROBE_CUPTI_GPU_MEMCPY_START = 5, XPROBE_CUPTI_GPU_MEMCPY_END = 6, XPROBE_CUPTI_GPU_MEMSET_START = 7, - XPROBE_CUPTI_GPU_MEMSET_END = 8 + XPROBE_CUPTI_GPU_MEMSET_END = 8, + XPROBE_CUPTI_NVTX_RANGE_START = 9, + XPROBE_CUPTI_NVTX_RANGE_END = 10 }; struct xprobe_cupti_output_header { @@ -149,6 +152,10 @@ struct xprobe_cupti_aggregate_record { * For memcpy and memset records, grid_x/grid_y hold the low/high halves of the * byte count. grid_z holds the CUpti_ActivityMemcpyKind for memcpy records, * and block_x holds the assigned value for memset records. + * + * For NVTX records, grid_x/grid_y hold the low/high halves of the range ID, + * grid_z is 1 for a thread range or 2 for a process range, block_x is the + * thread that started the range, and block_y reports whether name is complete. */ unsigned int xprobe_cupti_agent_abi_version(void); @@ -160,6 +167,9 @@ int xprobe_cupti_agent_flush(void); /* CUDA calls this entry point when the library is loaded via CUDA_INJECTION64_PATH. */ int InitializeInjection(void); +/* NVTX calls this before its first API dispatch via NVTX_INJECTION64_PATH. */ +int InitializeInjectionNvtx(void *get_export_table); +int InitializeInjectionNvtx2(void *get_export_table); #ifdef __cplusplus } diff --git a/cupti/src/cupti_agent.c b/cupti/src/cupti_agent.c index 6ae86a0..abb3cf3 100644 --- a/cupti/src/cupti_agent.c +++ b/cupti/src/cupti_agent.c @@ -24,6 +24,18 @@ int InitializeInjection(void) return 0; } +int InitializeInjectionNvtx(void *get_export_table) +{ + (void)get_export_table; + return 0; +} + +int InitializeInjectionNvtx2(void *get_export_table) +{ + (void)get_export_table; + return 0; +} + int xprobe_cupti_agent_status(void) { return XPROBE_CUPTI_AGENT_UNAVAILABLE; @@ -46,6 +58,8 @@ int xprobe_cupti_agent_flush(void) #include #include #include +#include +#include #include #include @@ -65,6 +79,9 @@ int xprobe_cupti_agent_flush(void) #include #include +extern CUptiResult CUPTIAPI cuptiNvtxInitialize(void *get_export_table); +extern CUptiResult CUPTIAPI cuptiNvtxInitialize2(void *get_export_table); + #define XPROBE_CUPTI_DEFAULT_RECORD_CAPACITY 100000U #define XPROBE_CUPTI_ACTIVITY_BUFFER_SIZE (8U * 1024U * 1024U) #define XPROBE_CUPTI_CLOCK_MARGIN_NS 1000000000U @@ -94,9 +111,20 @@ struct xprobe_cupti_aggregate_slot { char name[XPROBE_CUPTI_NAME_LENGTH]; }; +struct xprobe_nvtx_range_slot { + _Atomic uint64_t state; + uint64_t range_id; + uint32_t range_kind; + uint32_t start_tid; + uint32_t name_complete; + char name[XPROBE_CUPTI_NAME_LENGTH]; +}; + static struct xprobe_cupti_record *records; static _Atomic unsigned char *record_ready; static struct xprobe_cupti_aggregate_slot *aggregate_slots; +static struct xprobe_nvtx_range_slot *nvtx_range_slots; +static uint64_t nvtx_range_capacity; static uint64_t record_capacity; static uint32_t capture_mode; static _Atomic uint64_t record_count; @@ -122,12 +150,16 @@ static uint32_t enabled_activity_mask; static char output_path[PATH_MAX]; static char snapshot_socket_path[sizeof(((struct sockaddr_un *)0)->sun_path)]; static pthread_mutex_t flush_mutex = PTHREAD_MUTEX_INITIALIZER; +static pthread_mutex_t subscriber_mutex = PTHREAD_MUTEX_INITIALIZER; static pthread_t snapshot_thread; static _Atomic int snapshot_thread_stop; static _Atomic int snapshot_listener = -1; static int snapshot_thread_started; static struct xprobe_cupti_filter capture_filters[XPROBE_CUPTI_FILTER_COUNT]; static _Atomic int capture_filter_enabled; +static _Atomic int nvtx_injection_active; +static pthread_once_t injection_agent_once = PTHREAD_ONCE_INIT; +static int injection_agent_status = XPROBE_CUPTI_AGENT_UNAVAILABLE; #define XPROBE_CUPTI_ACTIVITY_KERNEL (1U << 0) #define XPROBE_CUPTI_ACTIVITY_MEMCPY (1U << 1) #define XPROBE_CUPTI_ACTIVITY_MEMSET (1U << 2) @@ -135,11 +167,12 @@ static _Atomic int capture_filter_enabled; static int shutdown_agent(void); static int activate_capture(void); -static int allocate_capture(uint64_t capacity) +static int allocate_capture(uint64_t capacity, int needs_nvtx) { struct xprobe_cupti_record *new_records; _Atomic unsigned char *new_ready; struct xprobe_cupti_aggregate_slot *new_aggregate_slots; + struct xprobe_nvtx_range_slot *new_nvtx_range_slots; if (capacity == 0U || capacity > SIZE_MAX / sizeof(*records) || capacity > SIZE_MAX / sizeof(*record_ready)) { @@ -150,6 +183,7 @@ static int allocate_capture(uint64_t capacity) new_records = NULL; new_ready = NULL; new_aggregate_slots = NULL; + new_nvtx_range_slots = NULL; if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { if (capacity > SIZE_MAX / sizeof(*aggregate_slots)) { fprintf(stderr, "xprobe CUPTI: aggregate capacity exceeds size_t\n"); @@ -160,15 +194,28 @@ static int allocate_capture(uint64_t capacity) new_records = calloc((size_t)capacity, sizeof(*new_records)); new_ready = malloc((size_t)capacity * sizeof(*new_ready)); } + if (needs_nvtx != 0) { + if (capacity > SIZE_MAX / sizeof(*nvtx_range_slots)) { + fprintf(stderr, "xprobe CUPTI: NVTX range capacity exceeds size_t\n"); + free(new_records); + free(new_ready); + free(new_aggregate_slots); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } + new_nvtx_range_slots = + calloc((size_t)capacity, sizeof(*new_nvtx_range_slots)); + } if ((capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE && new_aggregate_slots == NULL) || (capture_mode == XPROBE_CUPTI_CAPTURE_EXACT && - (new_records == NULL || new_ready == NULL))) { + (new_records == NULL || new_ready == NULL)) || + (needs_nvtx != 0 && new_nvtx_range_slots == NULL)) { fprintf(stderr, "xprobe CUPTI: failed to allocate %llu capture records\n", (unsigned long long)capacity); free(new_records); free(new_ready); free(new_aggregate_slots); + free(new_nvtx_range_slots); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } if (new_ready != NULL) { @@ -179,9 +226,12 @@ static int allocate_capture(uint64_t capacity) free(records); free(record_ready); free(aggregate_slots); + free(nvtx_range_slots); records = new_records; record_ready = new_ready; aggregate_slots = new_aggregate_slots; + nvtx_range_slots = new_nvtx_range_slots; + nvtx_range_capacity = needs_nvtx != 0 ? capacity : 0U; record_capacity = capacity; return XPROBE_CUPTI_AGENT_READY; } @@ -200,6 +250,10 @@ static void reset_capture(void) atomic_store_explicit(&record_ready[index], 0U, memory_order_relaxed); } } + for (index = 0U; index < nvtx_range_capacity; ++index) { + atomic_store_explicit(&nvtx_range_slots[index].state, 0U, + memory_order_relaxed); + } atomic_store_explicit(&record_count, 0U, memory_order_relaxed); atomic_store_explicit(&committed_record_count, 0U, memory_order_relaxed); atomic_store_explicit(&agent_dropped_records, 0U, memory_order_relaxed); @@ -347,6 +401,49 @@ static void copy_name(char destination[XPROBE_CUPTI_NAME_LENGTH], destination[index] = '\0'; } +static uint32_t copy_nvtx_name(char destination[XPROBE_CUPTI_NAME_LENGTH], + const char *source) +{ + size_t index = 0U; + uint32_t complete; + + if (source == NULL) { + destination[0] = '\0'; + return 1U; + } + while (index + 1U < XPROBE_CUPTI_NAME_LENGTH && source[index] != '\0') { + destination[index] = source[index]; + ++index; + } + complete = source[index] == '\0' ? 1U : 0U; + if (complete == 0U && index > 0U) { + size_t lead = index; + unsigned char first; + size_t expected = 1U; + + while (lead > 0U && + ((unsigned char)destination[lead - 1U] & 0xc0U) == 0x80U) { + --lead; + } + if (lead > 0U) { + --lead; + first = (unsigned char)destination[lead]; + if ((first & 0xe0U) == 0xc0U) { + expected = 2U; + } else if ((first & 0xf0U) == 0xe0U) { + expected = 3U; + } else if ((first & 0xf8U) == 0xf0U) { + expected = 4U; + } + if (expected > index - lead) { + index = lead; + } + } + } + destination[index] = '\0'; + return complete; +} + static size_t bounded_name_length(const char *name) { size_t length = 0U; @@ -677,6 +774,169 @@ static void initialize_record(struct xprobe_cupti_record *record, uint32_t kind, record->stream_id = XPROBE_CUPTI_VALUE_UNKNOWN; } +#define XPROBE_NVTX_RANGE_THREAD 1U +#define XPROBE_NVTX_RANGE_PROCESS 2U +#define XPROBE_NVTX_SLOT_WRITING 1U +#define XPROBE_NVTX_SLOT_TOMBSTONE 2U + +static uint64_t nvtx_range_hash(uint32_t range_kind, uint64_t range_id) +{ + uint64_t hash = 1469598103934665603ULL; + uint64_t values[] = {(uint64_t)range_kind, range_id}; + + for (size_t index = 0U; index < sizeof(values); ++index) { + hash ^= ((const unsigned char *)values)[index]; + hash *= 1099511628211ULL; + } + if (hash <= XPROBE_NVTX_SLOT_TOMBSTONE) { + hash += XPROBE_NVTX_SLOT_TOMBSTONE + 1U; + } + return hash; +} + +static int nvtx_range_key_matches(const struct xprobe_nvtx_range_slot *slot, + uint32_t range_kind, uint64_t range_id) +{ + return slot->range_kind == range_kind && slot->range_id == range_id; +} + +static void nvtx_range_capacity_reached(void) +{ + atomic_fetch_add_explicit(&agent_dropped_records, 1U, memory_order_relaxed); + atomic_store_explicit(&stop_reason, XPROBE_CUPTI_STOP_RECORD_LIMIT, + memory_order_relaxed); + atomic_store_explicit(&capture_state, XPROBE_CUPTI_CAPTURE_LIMIT_REACHED, + memory_order_release); +} + +static int remember_nvtx_range(uint32_t range_kind, uint64_t range_id, + uint32_t start_tid, const char *name) +{ + uint64_t hash = nvtx_range_hash(range_kind, range_id); + uint64_t first = hash % nvtx_range_capacity; + uint64_t candidate = UINT64_MAX; + + for (uint64_t probe = 0U; probe < nvtx_range_capacity; ++probe) { + uint64_t index = (first + probe) % nvtx_range_capacity; + struct xprobe_nvtx_range_slot *slot = &nvtx_range_slots[index]; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + + if (state == hash && + nvtx_range_key_matches(slot, range_kind, range_id) != 0) { + atomic_fetch_add_explicit(&agent_dropped_records, 1U, + memory_order_relaxed); + return 0; + } + if (state == XPROBE_NVTX_SLOT_TOMBSTONE && candidate == UINT64_MAX) { + candidate = index; + continue; + } + if (state == 0U) { + if (candidate == UINT64_MAX) { + candidate = index; + } + break; + } + } + if (candidate != UINT64_MAX) { + struct xprobe_nvtx_range_slot *slot = &nvtx_range_slots[candidate]; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + + if ((state == 0U || state == XPROBE_NVTX_SLOT_TOMBSTONE) && + atomic_compare_exchange_strong_explicit( + &slot->state, &state, XPROBE_NVTX_SLOT_WRITING, + memory_order_acq_rel, memory_order_acquire)) { + slot->range_id = range_id; + slot->range_kind = range_kind; + slot->start_tid = start_tid; + slot->name_complete = copy_nvtx_name(slot->name, name); + atomic_store_explicit(&slot->state, hash, memory_order_release); + return 1; + } + atomic_fetch_add_explicit(&agent_dropped_records, 1U, + memory_order_relaxed); + return 0; + } + nvtx_range_capacity_reached(); + return 0; +} + +static int take_nvtx_range(uint32_t range_kind, uint64_t range_id, + struct xprobe_cupti_record *record) +{ + uint64_t hash = nvtx_range_hash(range_kind, range_id); + uint64_t first = hash % nvtx_range_capacity; + uint64_t timestamp_ns; + + for (uint64_t probe = 0U; probe < nvtx_range_capacity; ++probe) { + uint64_t index = (first + probe) % nvtx_range_capacity; + struct xprobe_nvtx_range_slot *slot = &nvtx_range_slots[index]; + uint64_t state = atomic_load_explicit(&slot->state, memory_order_acquire); + + if (state == 0U) { + return 0; + } + if (state != hash || + nvtx_range_key_matches(slot, range_kind, range_id) == 0) { + continue; + } + if (!atomic_compare_exchange_strong_explicit( + &slot->state, &state, XPROBE_NVTX_SLOT_WRITING, + memory_order_acq_rel, memory_order_acquire)) { + continue; + } + if (monotonic_timestamp_ns(×tamp_ns) != + XPROBE_CUPTI_AGENT_READY) { + atomic_store_explicit(&slot->state, XPROBE_NVTX_SLOT_TOMBSTONE, + memory_order_release); + return 0; + } + initialize_record(record, XPROBE_CUPTI_NVTX_RANGE_END, timestamp_ns); + record->grid_x = (uint32_t)range_id; + record->grid_y = (uint32_t)(range_id >> 32U); + record->grid_z = range_kind; + record->block_x = slot->start_tid; + record->block_y = slot->name_complete; + memcpy(record->name, slot->name, sizeof(record->name)); + atomic_store_explicit(&slot->state, XPROBE_NVTX_SLOT_TOMBSTONE, + memory_order_release); + return 1; + } + return 0; +} + +static void enqueue_nvtx_range_start(uint32_t range_kind, uint64_t range_id, + const char *name, uint64_t timestamp_ns) +{ + struct xprobe_cupti_record record; + uint32_t start_tid = current_tid(); + + if (atomic_load_explicit(&capture_state, memory_order_relaxed) != + XPROBE_CUPTI_CAPTURE_ACTIVE || + remember_nvtx_range(range_kind, range_id, start_tid, name) == 0) { + return; + } + initialize_record(&record, XPROBE_CUPTI_NVTX_RANGE_START, timestamp_ns); + record.grid_x = (uint32_t)range_id; + record.grid_y = (uint32_t)(range_id >> 32U); + record.grid_z = range_kind; + record.block_x = start_tid; + record.block_y = copy_nvtx_name(record.name, name); + enqueue_record(&record); +} + +static void enqueue_nvtx_range_end(uint32_t range_kind, uint64_t range_id) +{ + struct xprobe_cupti_record record; + + if (atomic_load_explicit(&capture_state, memory_order_relaxed) != + XPROBE_CUPTI_CAPTURE_ACTIVE || + take_nvtx_range(range_kind, range_id, &record) == 0) { + return; + } + enqueue_record(&record); +} + static void enqueue_api_record(const CUpti_CallbackData *data, CUpti_CallbackDomain domain, CUpti_CallbackId callback_id, uint32_t kind, @@ -694,19 +954,167 @@ static void enqueue_api_record(const CUpti_CallbackData *data, enqueue_record(&record); } +static int nvtx_ascii_message(CUpti_CallbackId callback_id, + const CUpti_NvtxData *data, const char **message) +{ + const nvtxEventAttributes_t *attributes; + + switch (callback_id) { + case CUPTI_CBID_NVTX_nvtxRangeStartA: + *message = + ((const nvtxRangeStartA_params *)data->functionParams)->message; + return *message == NULL ? -1 : 1; + case CUPTI_CBID_NVTX_nvtxRangePushA: + *message = + ((const nvtxRangePushA_params *)data->functionParams)->message; + return *message == NULL ? -1 : 1; + case CUPTI_CBID_NVTX_nvtxRangeStartEx: + attributes = + ((const nvtxRangeStartEx_params *)data->functionParams)->eventAttrib; + break; + case CUPTI_CBID_NVTX_nvtxRangePushEx: + attributes = + ((const nvtxRangePushEx_params *)data->functionParams)->eventAttrib; + break; + default: + return -1; + } + if (attributes == NULL || + attributes->size < + offsetof(nvtxEventAttributes_t, message) + + sizeof(attributes->message)) { + return -1; + } + if (attributes->messageType != NVTX_MESSAGE_TYPE_ASCII) { + return 0; + } + *message = attributes->message.ascii; + return *message == NULL ? -1 : 1; +} + +static void nvtx_callback(CUpti_CallbackId callback_id, + const CUpti_NvtxData *data) +{ + const char *name; + uint32_t tid; + uint64_t range_id; + uint64_t timestamp_ns; + int message_status; + int level; + + if (atomic_load_explicit(&capture_state, memory_order_relaxed) != + XPROBE_CUPTI_CAPTURE_ACTIVE) { + return; + } + if (data == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + switch (callback_id) { + case CUPTI_CBID_NVTX_nvtxRangeStartA: + case CUPTI_CBID_NVTX_nvtxRangeStartEx: + if (data->functionParams == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + message_status = nvtx_ascii_message(callback_id, data, &name); + if (message_status == 0) { + return; + } + if (message_status < 0 || data->functionReturnValue == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + if (values_match_capture(XPROBE_CUPTI_NVTX_RANGE_START, + XPROBE_CUPTI_NVTX_RANGE_END, 0U, 0U, + name) == 0 || + monotonic_timestamp_ns(×tamp_ns) != XPROBE_CUPTI_AGENT_READY) { + return; + } + range_id = *(const nvtxRangeId_t *)data->functionReturnValue; + if (range_id == 0U) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + enqueue_nvtx_range_start(XPROBE_NVTX_RANGE_PROCESS, range_id, name, + timestamp_ns); + return; + case CUPTI_CBID_NVTX_nvtxRangePushA: + case CUPTI_CBID_NVTX_nvtxRangePushEx: + if (data->functionParams == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + message_status = nvtx_ascii_message(callback_id, data, &name); + if (message_status == 0) { + return; + } + if (message_status < 0 || data->functionReturnValue == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + if (values_match_capture(XPROBE_CUPTI_NVTX_RANGE_START, + XPROBE_CUPTI_NVTX_RANGE_END, 0U, 0U, + name) == 0 || + monotonic_timestamp_ns(×tamp_ns) != XPROBE_CUPTI_AGENT_READY) { + return; + } + level = *(const int *)data->functionReturnValue; + if (level < 0) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + tid = current_tid(); + range_id = ((uint64_t)tid << 32U) | (uint32_t)level; + enqueue_nvtx_range_start(XPROBE_NVTX_RANGE_THREAD, range_id, name, + timestamp_ns); + return; + case CUPTI_CBID_NVTX_nvtxRangeEnd: + if (data->functionParams == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + range_id = + ((const nvtxRangeEnd_params *)data->functionParams)->id; + enqueue_nvtx_range_end(XPROBE_NVTX_RANGE_PROCESS, range_id); + return; + case CUPTI_CBID_NVTX_nvtxRangePop: + if (data->functionReturnValue == NULL) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + level = *(const int *)data->functionReturnValue; + if (level < 0) { + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + return; + } + tid = current_tid(); + range_id = ((uint64_t)tid << 32U) | (uint32_t)level; + enqueue_nvtx_range_end(XPROBE_NVTX_RANGE_THREAD, range_id); + return; + default: + atomic_fetch_add_explicit(&unknown_records, 1U, memory_order_relaxed); + } +} + static void CUPTIAPI api_callback(void *userdata, CUpti_CallbackDomain domain, CUpti_CallbackId callback_id, const void *callback_data) { - const CUpti_CallbackData *data = callback_data; + const CUpti_CallbackData *data; uint32_t kind; uint64_t timestamp_ns; (void)userdata; + if (domain == CUPTI_CB_DOMAIN_NVTX) { + nvtx_callback(callback_id, callback_data); + return; + } if (domain != CUPTI_CB_DOMAIN_RUNTIME_API && domain != CUPTI_CB_DOMAIN_DRIVER_API) { return; } + data = callback_data; kind = data->callbackSite == CUPTI_API_ENTER ? XPROBE_CUPTI_CUDA_API_ENTRY : XPROBE_CUPTI_CUDA_API_EXIT; if (values_match_capture(kind, 0U, (uint32_t)domain, 0U, @@ -1124,6 +1532,7 @@ static int activity_timestamps_are_host_monotonic(uint64_t available) static uint64_t aggregate_observed_records(void) { uint64_t observed = 0U; + uint64_t dropped; for (uint64_t index = 0U; index < record_capacity; ++index) { const struct xprobe_cupti_aggregate_slot *slot = &aggregate_slots[index]; @@ -1140,6 +1549,13 @@ static uint64_t aggregate_observed_records(void) } observed += count; } + dropped = + atomic_load_explicit(&agent_dropped_records, memory_order_relaxed); + if (UINT64_MAX - observed < dropped) { + remember_output_error(); + return UINT64_MAX; + } + observed += dropped; return observed; } @@ -1157,6 +1573,9 @@ static void initialize_output_header(struct xprobe_cupti_output_header *header, header->header_size = sizeof(*header); header->record_size = sizeof(records[0]); header->feature_flags = XPROBE_CUPTI_FEATURE_TRANSFER_RECORDS; + if (atomic_load_explicit(&nvtx_injection_active, memory_order_acquire) != 0) { + header->feature_flags |= XPROBE_CUPTI_FEATURE_NVTX_RECORDS; + } if (capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE) { header->feature_flags |= XPROBE_CUPTI_FEATURE_AGGREGATE_RECORDS; } else if (activity_timestamps_are_host_monotonic(available) != 0) { @@ -1300,7 +1719,7 @@ static int flush_activity_buffers(int allow_empty_flush) static int valid_filter(const struct xprobe_cupti_filter *filter) { - if (filter->record_kind > XPROBE_CUPTI_GPU_MEMSET_END || + if (filter->record_kind > XPROBE_CUPTI_NVTX_RANGE_END || filter->api_domain > CUPTI_CB_DOMAIN_RUNTIME_API || filter->memcpy_kind > 5U || filter->name_match > XPROBE_CUPTI_NAME_CONTAINS) { @@ -1334,6 +1753,7 @@ static int valid_aggregate_filters( static int arm_capture(const struct xprobe_cupti_control_request *request) { int status; + int needs_nvtx = 0; for (size_t index = 0U; index < XPROBE_CUPTI_FILTER_COUNT; ++index) { if (valid_filter(&request->filters[index]) == 0) { @@ -1341,6 +1761,12 @@ static int arm_capture(const struct xprobe_cupti_control_request *request) remember_output_error(); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } + if (request->filters[index].record_kind == + XPROBE_CUPTI_NVTX_RANGE_START || + request->filters[index].record_kind == + XPROBE_CUPTI_NVTX_RANGE_END) { + needs_nvtx = 1; + } } if (request->capture_mode == XPROBE_CUPTI_CAPTURE_AGGREGATE && valid_aggregate_filters(request->filters) == 0) { @@ -1357,7 +1783,8 @@ static int arm_capture(const struct xprobe_cupti_control_request *request) } } capture_mode = request->capture_mode; - if (allocate_capture(request->record_capacity) != XPROBE_CUPTI_AGENT_READY) { + if (allocate_capture(request->record_capacity, needs_nvtx) != + XPROBE_CUPTI_AGENT_READY) { remember_output_error(); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } @@ -1590,13 +2017,64 @@ static int shutdown_agent(void) CUptiResult result; int status = disable_activities(); + if (pthread_mutex_lock(&subscriber_mutex) != 0) { + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } if (subscriber_active != 0) { - result = cuptiUnsubscribe(subscriber); + if (atomic_load_explicit(&nvtx_injection_active, memory_order_acquire) != + 0) { + static const CUpti_CallbackDomain domains[] = { + CUPTI_CB_DOMAIN_DRIVER_API, + CUPTI_CB_DOMAIN_RUNTIME_API, + CUPTI_CB_DOMAIN_NVTX, + }; + + for (size_t index = 0U; + index < sizeof(domains) / sizeof(domains[0]); ++index) { + result = cuptiEnableDomain(0U, subscriber, domains[index]); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiEnableDomain(disable)", result); + status = XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + } + } else { + result = cuptiUnsubscribe(subscriber); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiUnsubscribe", result); + status = XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + subscriber_active = 0; + } + } + if (pthread_mutex_unlock(&subscriber_mutex) != 0) { + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } + return status; +} + +static int ensure_subscriber(void) +{ + CUptiResult result; + int status = XPROBE_CUPTI_AGENT_READY; + + if (pthread_mutex_lock(&subscriber_mutex) != 0) { + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; + } + if (subscriber_active == 0) { + result = cuptiSubscribe(&subscriber, api_callback, NULL); if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiUnsubscribe", result); + report_cupti_error("cuptiSubscribe", result); status = XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } else { + subscriber_active = 1; } - subscriber_active = 0; + } + if (pthread_mutex_unlock(&subscriber_mutex) != 0) { + remember_output_error(); + return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } return status; } @@ -1637,6 +2115,15 @@ static int capture_needs_api_domain(uint32_t domain) return 0; } +static int capture_needs_nvtx_records(void) +{ + if (atomic_load_explicit(&capture_filter_enabled, memory_order_relaxed) == 0) { + return 0; + } + return capture_needs_record_kind(XPROBE_CUPTI_NVTX_RANGE_START, + XPROBE_CUPTI_NVTX_RANGE_END); +} + static int enable_activity(uint32_t bit, CUpti_ActivityKind kind) { CUptiResult result = cuptiActivityEnable(kind); @@ -1761,7 +2248,22 @@ static int activate_capture(void) needs_memset != 0; int needs_runtime = capture_needs_api_domain(CUPTI_CB_DOMAIN_RUNTIME_API); int needs_driver = capture_needs_api_domain(CUPTI_CB_DOMAIN_DRIVER_API); + int needs_nvtx = capture_needs_nvtx_records(); + if (needs_nvtx != 0 && + atomic_load_explicit(&nvtx_injection_active, memory_order_acquire) == 0) { + fprintf(stderr, + "xprobe CUPTI: NVTX callback routing was not initialized; " + "restart with NVTX_INJECTION64_PATH set before the first NVTX " + "call\n"); + atomic_store_explicit(&agent_status, XPROBE_CUPTI_AGENT_UNAVAILABLE, + memory_order_relaxed); + atomic_store_explicit(&capture_state, XPROBE_CUPTI_CAPTURE_FAILED, + memory_order_relaxed); + atomic_store_explicit(&stop_reason, XPROBE_CUPTI_STOP_CUPTI_ERROR, + memory_order_relaxed); + return XPROBE_CUPTI_AGENT_UNAVAILABLE; + } result = cuptiGetVersion(&runtime_cupti_version); if (result != CUPTI_SUCCESS) { report_cupti_error("cuptiGetVersion", result); @@ -1771,13 +2273,10 @@ static int activate_capture(void) XPROBE_CUPTI_AGENT_READY) { return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } - if (needs_runtime != 0 || needs_driver != 0) { - result = cuptiSubscribe(&subscriber, api_callback, NULL); - if (result != CUPTI_SUCCESS) { - report_cupti_error("cuptiSubscribe", result); + if (needs_runtime != 0 || needs_driver != 0 || needs_nvtx != 0) { + if (ensure_subscriber() != XPROBE_CUPTI_AGENT_READY) { return XPROBE_CUPTI_AGENT_CUPTI_ERROR; } - subscriber_active = 1; } if (needs_activities != 0) { #if CUPTI_API_VERSION < 130000 @@ -1839,6 +2338,28 @@ static int activate_capture(void) return xprobe_cupti_agent_status(); } } + if (needs_nvtx != 0) { + static const CUpti_CallbackId nvtx_callback_ids[] = { + CUPTI_CBID_NVTX_nvtxRangeStartA, + CUPTI_CBID_NVTX_nvtxRangeStartEx, + CUPTI_CBID_NVTX_nvtxRangeEnd, + CUPTI_CBID_NVTX_nvtxRangePushA, + CUPTI_CBID_NVTX_nvtxRangePushEx, + CUPTI_CBID_NVTX_nvtxRangePop, + }; + + for (size_t index = 0U; + index < sizeof(nvtx_callback_ids) / sizeof(nvtx_callback_ids[0]); + ++index) { + result = cuptiEnableCallback(1U, subscriber, CUPTI_CB_DOMAIN_NVTX, + nvtx_callback_ids[index]); + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiEnableCallback(NVTX)", result); + shutdown_agent(); + return XPROBE_CUPTI_AGENT_CUPTI_ERROR; + } + } + } atomic_store_explicit(&agent_status, XPROBE_CUPTI_AGENT_READY, memory_order_relaxed); atomic_store_explicit(&capture_state, XPROBE_CUPTI_CAPTURE_ACTIVE, @@ -1870,7 +2391,7 @@ int xprobe_cupti_agent_start(const char *configured_socket, uint64_t capacity) atomic_store_explicit(&agent_status, XPROBE_CUPTI_AGENT_UNAVAILABLE, memory_order_release); snapshot_socket_path[0] = '\0'; - if (allocate_capture(capacity) != XPROBE_CUPTI_AGENT_READY) { + if (allocate_capture(capacity, 0) != XPROBE_CUPTI_AGENT_READY) { remember_output_error(); return XPROBE_CUPTI_AGENT_OUTPUT_ERROR; } @@ -1943,9 +2464,52 @@ int xprobe_cupti_agent_initialize(void) return xprobe_cupti_agent_start(configured_socket, capacity); } +static void initialize_injection_agent(void) +{ + if (getenv("XPROBE_CUPTI_OUTPUT") == NULL && + getenv("XPROBE_CUPTI_SOCKET") == NULL) { + injection_agent_status = XPROBE_CUPTI_AGENT_READY; + return; + } + injection_agent_status = xprobe_cupti_agent_initialize(); +} + int InitializeInjection(void) { - return xprobe_cupti_agent_initialize() == XPROBE_CUPTI_AGENT_READY; + if (pthread_once(&injection_agent_once, initialize_injection_agent) != 0) { + return 0; + } + return injection_agent_status == XPROBE_CUPTI_AGENT_READY; +} + +int InitializeInjectionNvtx(void *get_export_table) +{ + CUptiResult result = cuptiNvtxInitialize(get_export_table); + + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiNvtxInitialize", result); + return 0; + } + atomic_store_explicit(&nvtx_injection_active, 1, memory_order_release); + if (ensure_subscriber() != XPROBE_CUPTI_AGENT_READY) { + return 0; + } + return InitializeInjection(); +} + +int InitializeInjectionNvtx2(void *get_export_table) +{ + CUptiResult result = cuptiNvtxInitialize2(get_export_table); + + if (result != CUPTI_SUCCESS) { + report_cupti_error("cuptiNvtxInitialize2", result); + return 0; + } + atomic_store_explicit(&nvtx_injection_active, 1, memory_order_release); + if (ensure_subscriber() != XPROBE_CUPTI_AGENT_READY) { + return 0; + } + return InitializeInjection(); } int xprobe_cupti_agent_status(void) diff --git a/cupti/tests/nvtx_range_fixture.cpp b/cupti/tests/nvtx_range_fixture.cpp new file mode 100644 index 0000000..15728bd --- /dev/null +++ b/cupti/tests/nvtx_range_fixture.cpp @@ -0,0 +1,92 @@ +#include + +#include +#include +#include +#include +#include + +namespace { + +std::string read_mode(const char *path) +{ + std::ifstream input(path); + std::string mode; + input >> mode; + return mode; +} + +void nested_ranges() +{ + nvtxRangePushA("xprobe_outer"); + nvtxEventAttributes_t attributes = {}; + attributes.version = NVTX_VERSION; + attributes.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; + attributes.messageType = NVTX_MESSAGE_TYPE_ASCII; + attributes.message.ascii = "xprobe_inner_ex"; + nvtxRangePushEx(&attributes); + std::this_thread::sleep_for(std::chrono::microseconds(100)); + nvtxRangePop(); + nvtxRangePop(); +} + +void cross_thread_range() +{ + nvtxRangeId_t id = nvtxRangeStartA("xprobe_cross_thread"); + std::thread end_thread([id] { + std::this_thread::sleep_for(std::chrono::microseconds(100)); + nvtxRangeEnd(id); + }); + end_thread.join(); +} + +void long_range() +{ + static const std::string name = + "xprobe_long_range_" + std::string(180, 'x'); + nvtxRangePushA(name.c_str()); + std::this_thread::sleep_for(std::chrono::microseconds(100)); + nvtxRangePop(); +} + +void unsupported_registered_range() +{ + static const nvtxStringHandle_t name = + nvtxDomainRegisterStringA(nullptr, "xprobe_registered"); + nvtxEventAttributes_t attributes = {}; + attributes.version = NVTX_VERSION; + attributes.size = NVTX_EVENT_ATTRIB_STRUCT_SIZE; + attributes.messageType = NVTX_MESSAGE_TYPE_REGISTERED; + attributes.message.registered = name; + nvtxRangePushEx(&attributes); + nvtxRangePop(); +} + +} // namespace + +int main(int argc, char **argv) +{ + if (argc != 2) { + std::cerr << "usage: nvtx_range_fixture \n"; + return 2; + } + nvtxRangePushA("xprobe_init"); + nvtxRangePop(); + std::cout << "ready" << std::endl; + + for (;;) { + const std::string mode = read_mode(argv[1]); + unsupported_registered_range(); + if (mode == "nested") { + nested_ranges(); + } else if (mode == "cross") { + cross_thread_range(); + } else if (mode == "long") { + long_range(); + } else { + std::cerr << "unknown mode: " << mode << '\n'; + return 3; + } + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} diff --git a/justfile b/justfile index 1091eab..a06ca28 100644 --- a/justfile +++ b/justfile @@ -78,6 +78,12 @@ test-pytorch-live: build test-pytorch-cuda-live: build python3 tests/integration/test_pytorch_cuda.py --image "{{cuda12_devel_image}}" --pytorch-env "${PYTORCH_ENV:?set PYTORCH_ENV to a Mamba environment with torch}" +test-nvtx-live: build + python3 tests/integration/test_nvtx.py --image "{{cuda13_devel_image}}" + +test-nvtx-live-cuda12: build + python3 tests/integration/test_nvtx.py --image "{{cuda12_devel_image}}" + benchmark-gpu: python3 benchmarks/cuda-callback/run.py "{{cuda13_devel_image}}" diff --git a/tests/integration/run-nvtx-live.sh b/tests/integration/run-nvtx-live.sh new file mode 100755 index 0000000..09800bf --- /dev/null +++ b/tests/integration/run-nvtx-live.sh @@ -0,0 +1,91 @@ +#!/bin/sh +set -eu + +output_dir=$1 +agent="${output_dir}/libxprobe-cupti.so" +fixture="${output_dir}/nvtx-range-fixture" +mode="${output_dir}/mode" +xprobe=/workspace/target/debug/xprobe + +gcc -std=c11 -D_GNU_SOURCE -DXPROBE_HAS_CUPTI=1 -fPIC -shared -pthread \ + -O2 -Wall -Wextra -Wpedantic -Werror \ + -I/workspace/cupti/include -isystem /usr/local/cuda/include \ + /workspace/cupti/src/cupti_agent.c \ + -L/usr/local/cuda/lib64 -Wl,-rpath,/usr/local/cuda/lib64 -lcupti \ + -o "${agent}" + +g++ -std=c++17 -pthread -O2 -Wall -Wextra -Wpedantic -Werror \ + -isystem /usr/local/cuda/include \ + /workspace/cupti/tests/nvtx_range_fixture.cpp -ldl -o "${fixture}" + +printf '%s\n' nested >"${mode}" +NVTX_INJECTION64_PATH="${agent}" "${fixture}" "${mode}" \ + >"${output_dir}/fixture.log" 2>"${output_dir}/fixture.stderr" & +pid=$! +trap 'chmod 0644 "${output_dir}"/*.jsonl 2>/dev/null || true; kill "${pid}" 2>/dev/null || true; wait "${pid}" 2>/dev/null || true' EXIT + +attempt=0 +while ! grep -q '^ready$' "${output_dir}/fixture.log"; do + if ! kill -0 "${pid}" 2>/dev/null; then + set +e + wait "${pid}" + status=$? + set -e + echo "NVTX fixture exited before readiness with status ${status}" >&2 + exit 1 + fi + attempt=$((attempt + 1)) + if [ "${attempt}" -ge 100 ]; then + echo "timed out waiting for NVTX fixture" >&2 + exit 1 + fi + sleep 0.01 +done + +measure() +{ + result=$1 + workload_mode=$2 + pattern=$3 + printf '%s\n' "${workload_mode}" >"${mode}.next" + mv "${mode}.next" "${mode}" + sleep 0.05 + "${xprobe}" measure \ + --pid "${pid}" \ + --agent "${agent}" \ + --from "cuda:nvtx_range_start:name~${pattern}" \ + --to "cuda:nvtx_range_end:name~${pattern}" \ + --match exact \ + --samples 8 \ + --max-events 4096 \ + --timeout-ms 30000 \ + --events-out "${output_dir}/${result}.jsonl" \ + --json --non-interactive --no-color >"${output_dir}/${result}.json" +} + +measure nested nested '^xprobe_outer$' +measure extended nested '^xprobe_inner_ex$' +measure cross cross '^xprobe_cross_thread$' + +printf '%s\n' nested >"${mode}.next" +mv "${mode}.next" "${mode}" +sleep 0.05 +set +e +"${xprobe}" measure \ + --pid "${pid}" \ + --agent "${agent}" \ + --from 'cuda:nvtx_range_start:name~^xprobe_outer$' \ + --to 'cuda:nvtx_range_end:name~^xprobe_outer$' \ + --match exact \ + --samples 8 \ + --max-events 1 \ + --timeout-ms 30000 \ + --json --non-interactive --no-color >"${output_dir}/limit.json" +limit_status=$? +set -e +if [ "${limit_status}" -eq 0 ]; then + echo "NVTX record limit capture unexpectedly succeeded" >&2 + exit 1 +fi + +measure long long '^xprobe_long_range_.*' diff --git a/tests/integration/test_nvtx.py b/tests/integration/test_nvtx.py new file mode 100755 index 0000000..be5224d --- /dev/null +++ b/tests/integration/test_nvtx.py @@ -0,0 +1,104 @@ +#!/usr/bin/env python3 +import argparse +import json +import pathlib +import subprocess +import sys +import tempfile + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--image", required=True) + args = parser.parse_args() + workspace = pathlib.Path(__file__).resolve().parents[2] + + with tempfile.TemporaryDirectory(prefix="xprobe-nvtx-") as output_dir: + completed = subprocess.run( + [ + "docker", + "run", + "--rm", + "--gpus", + "all", + "--cap-add", + "SYS_PTRACE", + "--security-opt", + "seccomp=unconfined", + "--volume", + f"{workspace}:/workspace:ro", + "--volume", + f"{output_dir}:/output", + "--workdir", + "/workspace", + args.image, + "/workspace/tests/integration/run-nvtx-live.sh", + "/output", + ], + check=False, + capture_output=True, + text=True, + ) + if completed.returncode != 0: + sys.stdout.write(completed.stdout) + sys.stderr.write(completed.stderr) + for path in sorted(pathlib.Path(output_dir).glob("*")): + if path.suffix in {".json", ".jsonl", ".log", ".stderr"}: + sys.stderr.write( + f"\n--- {path.name} ---\n" + f"{path.read_text(errors='replace')}\n" + ) + raise SystemExit(completed.returncode) + + output = pathlib.Path(output_dir) + nested = json.loads((output / "nested.json").read_text()) + extended = json.loads((output / "extended.json").read_text()) + cross = json.loads((output / "cross.json").read_text()) + limit = json.loads((output / "limit.json").read_text()) + long_name = json.loads((output / "long.json").read_text()) + + assert_measurement(nested, "thread") + assert_measurement(extended, "thread") + assert_measurement(cross, "process") + assert_measurement(long_name, "thread") + assert limit["ok"] is False + assert limit["error"]["code"] == "EVENT_RATE_TOO_HIGH" + assert all( + pair["start"]["tid"] != pair["end"]["tid"] + and pair["start"]["nvtx"]["start_tid"] == pair["start"]["tid"] + and pair["end"]["nvtx"]["start_tid"] == pair["start"]["tid"] + for pair in cross["evidence"] + ) + assert all( + pair["start"]["nvtx"]["name_complete"] is False + and pair["end"]["nvtx"]["name_complete"] is False + for pair in long_name["evidence"] + ) + print( + json.dumps( + { + "cross_thread_samples": cross["measurement"]["samples"]["matched"], + "extended_samples": extended["measurement"]["samples"]["matched"], + "nested_samples": nested["measurement"]["samples"]["matched"], + "ok": True, + }, + sort_keys=True, + ) + ) + + +def assert_measurement(result: dict, range_kind: str) -> None: + assert result["ok"] is True + assert result["measurement"]["samples"]["matched"] == 8 + assert result["measurement"]["samples"]["dropped"] == 0 + assert result["correlation"]["method"] == "exact_nvtx_range_id" + assert all( + pair["start"]["nvtx"]["range_kind"] == range_kind + and pair["end"]["nvtx"]["range_kind"] == range_kind + and pair["start"]["nvtx"]["range_id"] == pair["end"]["nvtx"]["range_id"] + for pair in result["evidence"] + ) + + +if __name__ == "__main__": + main() diff --git a/tests/integration/test_pytorch_cuda.py b/tests/integration/test_pytorch_cuda.py index 9f8ac0a..d7f8c94 100644 --- a/tests/integration/test_pytorch_cuda.py +++ b/tests/integration/test_pytorch_cuda.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import argparse import json +import os import pathlib import subprocess import sys @@ -67,6 +68,8 @@ def run_inner(xprobe: pathlib.Path) -> None: mode = temp / "mode" build_agent(workspace, agent) set_mode(mode, "mm") + environment = os.environ.copy() + environment["NVTX_INJECTION64_PATH"] = str(agent) workload = subprocess.Popen( [ sys.executable, @@ -75,6 +78,7 @@ def run_inner(xprobe: pathlib.Path) -> None: WORKLOAD, str(mode), ], + env=environment, stdout=subprocess.PIPE, text=True, ) @@ -110,6 +114,15 @@ def run_inner(xprobe: pathlib.Path) -> None: "cuda:runtime_api:cudaStreamSynchronize:entry", "cuda:runtime_api:cudaStreamSynchronize:exit", ) + nvtx = exact_measure( + xprobe, + agent, + pid, + mode, + "nvtx", + "cuda:nvtx_range_start:name~^xprobe_pytorch_step$", + "cuda:nvtx_range_end:name~^xprobe_pytorch_step$", + ) finally: workload.terminate() workload.wait(timeout=10) @@ -129,6 +142,13 @@ def run_inner(xprobe: pathlib.Path) -> None: assert "cuda:memcpy_start:kind=DtoH" in transfer_hints assert_exact(exact_kernel) assert_exact(synchronization) + assert_exact(nvtx) + assert nvtx["correlation"]["method"] == "exact_nvtx_range_id" + assert all( + pair["start"]["nvtx"]["range_kind"] == "thread" + and pair["start"]["nvtx"]["range_id"] == pair["end"]["nvtx"]["range_id"] + for pair in nvtx["evidence"] + ) print( json.dumps( @@ -141,7 +161,14 @@ def run_inner(xprobe: pathlib.Path) -> None: "exact_kernel_records": exact_kernel["collection"]["cupti"][ "retained_records" ], - "measured": ["kernel", "memcpy_htod", "memcpy_dtoh", "stream_sync"], + "measured": [ + "kernel", + "memcpy_htod", + "memcpy_dtoh", + "stream_sync", + "nvtx_range", + ], + "nvtx_records": nvtx["collection"]["cupti"]["retained_records"], "ok": True, "pytorch": metadata["pytorch"], "stream_sync_records": synchronization["collection"]["cupti"][ @@ -324,6 +351,8 @@ def compiled_step(left, right): compiled_step(a, b) stream.synchronize() +torch.cuda.nvtx.range_push("xprobe_init") +torch.cuda.nvtx.range_pop() print(json.dumps({ "cuda": torch.version.cuda, "device": torch.cuda.get_device_name(), @@ -342,6 +371,11 @@ def compiled_step(left, right): elif mode == "transfer": device_transfer.copy_(host_input, non_blocking=True) host_output.copy_(device_transfer, non_blocking=True) + elif mode == "nvtx": + torch.cuda.nvtx.range_push("xprobe_pytorch_step") + torch.mm(a, b) + stream.synchronize() + torch.cuda.nvtx.range_pop() else: raise RuntimeError(f"unknown workload mode: {mode}") stream.synchronize() diff --git a/xprobe/cli/src/main.rs b/xprobe/cli/src/main.rs index 3619ce7..85a5cb6 100644 --- a/xprobe/cli/src/main.rs +++ b/xprobe/cli/src/main.rs @@ -1608,7 +1608,12 @@ fn prepare_live_collection(request: &LiveMeasureRequest) -> Result Result<(), Com fn arm_cupti_capture( activation: &CuptiActivation, config: Option<&cupti::CuptiArmConfig>, + needs_nvtx: bool, request: &LiveMeasureRequest, ) -> Result, CommandFailure> { let (Some(path), Some(config)) = (activation.socket.as_deref(), config) else { @@ -1721,6 +1727,15 @@ fn arm_cupti_capture( return cleanup_failed_arm(activation, request, failure); } }; + if needs_nvtx && !capture.nvtx_callbacks_active { + let failure = CommandFailure::new( + ErrorCode::CuptiNotAvailable, + "the mapped CUPTI agent was not initialized through NVTX; restart the target with NVTX_INJECTION64_PATH set to the same agent before the first NVTX call", + true, + ) + .with_hint("online CUDA injection cannot retrofit an initialized NVTX dispatch table"); + return cleanup_failed_arm(activation, request, failure); + } if capture.state != cupti::CuptiCaptureState::Active { let failure = CommandFailure::new( ErrorCode::CuptiNotAvailable, @@ -1805,6 +1820,8 @@ fn cupti_event_filter( EventType::GpuMemcpyEnd => cupti::CuptiRecordKind::GpuMemcpyEnd, EventType::GpuMemsetStart => cupti::CuptiRecordKind::GpuMemsetStart, EventType::GpuMemsetEnd => cupti::CuptiRecordKind::GpuMemsetEnd, + EventType::NvtxRangeStart => cupti::CuptiRecordKind::NvtxRangeStart, + EventType::NvtxRangeEnd => cupti::CuptiRecordKind::NvtxRangeEnd, _ => { return Err(CommandFailure::new( ErrorCode::InvalidEventSelector, @@ -1837,6 +1854,8 @@ fn cupti_event_filter( cupti::CuptiNameFilter::Exact(name.clone()) } else if let Some(pattern) = selector.kernel_name_regex.as_deref() { bounded_kernel_name_filter(pattern) + } else if let Some(pattern) = selector.nvtx_name_regex.as_deref() { + bounded_kernel_name_filter(pattern) } else { cupti::CuptiNameFilter::Any }; diff --git a/xprobe/collector/src/cupti.rs b/xprobe/collector/src/cupti.rs index ba6f8dd..2446a28 100644 --- a/xprobe/collector/src/cupti.rs +++ b/xprobe/collector/src/cupti.rs @@ -12,7 +12,8 @@ use std::{ use serde_json::Value; use xprobe_protocol::{ - ClockDomain, CudaEvent, Dim3, Event, EventSource, EventType, MemcpyKind, SchemaVersion, + ClockDomain, CudaEvent, Dim3, Event, EventSource, EventType, MemcpyKind, NvtxEvent, + NvtxRangeKind, SchemaVersion, }; const OUTPUT_MAGIC: &[u8; 8] = b"XPCUPTI\0"; @@ -30,8 +31,11 @@ pub const CUPTI_NAME_CAPACITY: usize = 128; const FEATURE_HOST_MONOTONIC_TIMESTAMPS: u32 = 1 << 0; const FEATURE_TRANSFER_RECORDS: u32 = 1 << 1; const FEATURE_AGGREGATE_RECORDS: u32 = 1 << 2; -const SUPPORTED_FEATURES: u32 = - FEATURE_HOST_MONOTONIC_TIMESTAMPS | FEATURE_TRANSFER_RECORDS | FEATURE_AGGREGATE_RECORDS; +const FEATURE_NVTX_RECORDS: u32 = 1 << 3; +const SUPPORTED_FEATURES: u32 = FEATURE_HOST_MONOTONIC_TIMESTAMPS + | FEATURE_TRANSFER_RECORDS + | FEATURE_AGGREGATE_RECORDS + | FEATURE_NVTX_RECORDS; const UNKNOWN_U32: u32 = u32::MAX; #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -45,6 +49,8 @@ pub enum CuptiRecordKind { GpuMemcpyEnd = 6, GpuMemsetStart = 7, GpuMemsetEnd = 8, + NvtxRangeStart = 9, + NvtxRangeEnd = 10, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -120,6 +126,7 @@ pub struct CuptiAggregateGroup { #[derive(Debug, Clone, PartialEq)] pub struct CuptiCapture { + pub nvtx_callbacks_active: bool, pub record_offset: u64, pub state: CuptiCaptureState, pub stop_reason: CuptiStopReason, @@ -181,6 +188,10 @@ pub enum CuptiDecodeError { index: usize, kind: u32, }, + InvalidNvtxRecord { + index: usize, + message: &'static str, + }, InvalidAggregateRecord { index: usize, message: &'static str, @@ -248,6 +259,9 @@ impl fmt::Display for CuptiDecodeError { "CUPTI record {index} of kind {kind} has a zero timestamp" ) } + Self::InvalidNvtxRecord { index, message } => { + write!(formatter, "CUPTI NVTX record {index} is invalid: {message}") + } Self::InvalidAggregateRecord { index, message } => { write!( formatter, @@ -582,6 +596,7 @@ pub fn decode_capture(bytes: &[u8], session_id: &str) -> Result Result { let kind = read_u32(record, 8); - let (source, event_type) = match kind { - 1 => (EventSource::CuptiCallback, EventType::CudaApiEntry), - 2 => (EventSource::CuptiCallback, EventType::CudaApiExit), - 3 => (EventSource::CuptiActivity, EventType::GpuKernelStart), - 4 => (EventSource::CuptiActivity, EventType::GpuKernelEnd), - 5 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { - (EventSource::CuptiActivity, EventType::GpuMemcpyStart) - } - 6 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { - (EventSource::CuptiActivity, EventType::GpuMemcpyEnd) - } - 7 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { - (EventSource::CuptiActivity, EventType::GpuMemsetStart) - } - 8 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { - (EventSource::CuptiActivity, EventType::GpuMemsetEnd) - } - _ => return Err(CuptiDecodeError::UnknownRecordKind { index, kind }), - }; + let (source, event_type) = decode_record_kind(kind, index, feature_flags)?; let timestamp_raw = read_u64(record, 0); if timestamp_raw == 0 { return Err(CuptiDecodeError::InvalidTimestamp { index, kind }); @@ -756,7 +753,8 @@ fn decode_record( let is_transfer = matches!(kind, 5..=8); let is_start = matches!(kind, 3 | 5 | 7); let is_end = matches!(kind, 4 | 6 | 8); - let (timestamp_ns, clock_domain, timestamp_error_ns) = if is_api { + let is_nvtx = matches!(kind, 9 | 10); + let (timestamp_ns, clock_domain, timestamp_error_ns) = if is_api || is_nvtx { (timestamp_raw, ClockDomain::HostMonotonic, None) } else if feature_flags & FEATURE_HOST_MONOTONIC_TIMESTAMPS != 0 { ( @@ -783,6 +781,7 @@ fn decode_record( if matches!(kind, 7 | 8) { attributes.insert("memset_value".to_owned(), Value::from(read_u32(record, 56))); } + let nvtx = decode_nvtx_event(record, index, is_nvtx, name)?; Ok(Event { schema_version: SchemaVersion::current(), @@ -800,7 +799,7 @@ fn decode_record( timestamp_error_ns, process_start_time: None, host: None, - cuda: Some(CudaEvent { + cuda: (!is_nvtx).then(|| CudaEvent { device_id: optional_unknown(read_u32(record, 20)), context_id: optional_unknown(read_u32(record, 24)), stream_id: optional_unknown(read_u32(record, 28)).map(u64::from), @@ -817,11 +816,92 @@ fn decode_record( bytes: is_transfer.then(|| read_split_u64(record, 44)), memcpy_kind: is_memcpy.then(|| decode_memcpy_kind(read_u32(record, 52))), }), - nvtx: None, + nvtx, attributes, }) } +fn decode_record_kind( + kind: u32, + index: usize, + feature_flags: u32, +) -> Result<(EventSource, EventType), CuptiDecodeError> { + Ok(match kind { + 1 => (EventSource::CuptiCallback, EventType::CudaApiEntry), + 2 => (EventSource::CuptiCallback, EventType::CudaApiExit), + 3 => (EventSource::CuptiActivity, EventType::GpuKernelStart), + 4 => (EventSource::CuptiActivity, EventType::GpuKernelEnd), + 5 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { + (EventSource::CuptiActivity, EventType::GpuMemcpyStart) + } + 6 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { + (EventSource::CuptiActivity, EventType::GpuMemcpyEnd) + } + 7 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { + (EventSource::CuptiActivity, EventType::GpuMemsetStart) + } + 8 if feature_flags & FEATURE_TRANSFER_RECORDS != 0 => { + (EventSource::CuptiActivity, EventType::GpuMemsetEnd) + } + 9 if feature_flags & FEATURE_NVTX_RECORDS != 0 => { + (EventSource::CuptiCallback, EventType::NvtxRangeStart) + } + 10 if feature_flags & FEATURE_NVTX_RECORDS != 0 => { + (EventSource::CuptiCallback, EventType::NvtxRangeEnd) + } + _ => return Err(CuptiDecodeError::UnknownRecordKind { index, kind }), + }) +} + +fn decode_nvtx_event( + record: &[u8], + index: usize, + is_nvtx: bool, + name: &str, +) -> Result, CuptiDecodeError> { + if !is_nvtx { + return Ok(None); + } + let range_kind = match read_u32(record, 52) { + 1 => NvtxRangeKind::Thread, + 2 => NvtxRangeKind::Process, + _ => { + return Err(CuptiDecodeError::InvalidNvtxRecord { + index, + message: "range kind is neither thread nor process", + }); + } + }; + let range_id = read_split_u64(record, 44); + if range_id == 0 { + return Err(CuptiDecodeError::InvalidNvtxRecord { + index, + message: "range ID is zero", + }); + } + let start_tid = read_u32(record, 56); + if start_tid == 0 { + return Err(CuptiDecodeError::InvalidNvtxRecord { + index, + message: "start thread ID is zero", + }); + } + let name_complete = read_u32(record, 60); + if name_complete > 1 { + return Err(CuptiDecodeError::InvalidNvtxRecord { + index, + message: "name completeness flag is not boolean", + }); + } + Ok(Some(NvtxEvent { + name: name.to_owned(), + name_complete: name_complete != 0, + range_kind, + range_id, + start_tid, + })) +} + fn read_split_u64(record: &[u8], offset: usize) -> u64 { u64::from(read_u32(record, offset)) | (u64::from(read_u32(record, offset + 4)) << 32) } @@ -894,7 +974,7 @@ mod tests { CuptiEventFilter, CuptiMemcpyKind, CuptiNameFilter, CuptiRecordKind, HEADER_SIZE, OUTPUT_MAGIC, RECORD_SIZE, decode_capture, encode_control, read_snapshot, }; - use xprobe_protocol::{ClockDomain, EventSource, EventType, MemcpyKind}; + use xprobe_protocol::{ClockDomain, EventSource, EventType, MemcpyKind, NvtxRangeKind}; fn capture(records: &[[u8; RECORD_SIZE]]) -> Vec { let mut bytes = vec![0_u8; HEADER_SIZE + records.len() * RECORD_SIZE]; @@ -936,6 +1016,12 @@ mod tests { bytes } + fn nvtx_capture(records: &[[u8; RECORD_SIZE]]) -> Vec { + let mut bytes = capture(records); + bytes[20..24].copy_from_slice(&super::FEATURE_NVTX_RECORDS.to_le_bytes()); + bytes + } + fn aggregate_record( kind: u32, count: u64, @@ -997,6 +1083,23 @@ mod tests { record } + fn nvtx_record( + kind: u32, + timestamp: u64, + range_id: u64, + range_kind: u32, + start_tid: u32, + name_complete: u32, + name: &str, + ) -> [u8; RECORD_SIZE] { + let mut record = record(kind, timestamp, 0, name); + record[44..52].copy_from_slice(&range_id.to_le_bytes()); + record[52..56].copy_from_slice(&range_kind.to_le_bytes()); + record[56..60].copy_from_slice(&start_tid.to_le_bytes()); + record[60..64].copy_from_slice(&name_complete.to_le_bytes()); + record + } + #[test] fn decodes_and_orders_callback_and_kernel_events() { let api = record(1, 200, 42, "cudaLaunchKernel"); @@ -1061,6 +1164,52 @@ mod tests { ); } + #[test] + fn decodes_nvtx_thread_and_process_range_boundaries() { + let thread_id = (1235_u64 << 32) | 2; + let process_id = 0x1122_3344_5566_7788; + let decoded = decode_capture( + &nvtx_capture(&[ + nvtx_record(9, 100, thread_id, 1, 1235, 1, "forward"), + nvtx_record(10, 180, thread_id, 1, 1235, 1, "forward"), + nvtx_record(9, 200, process_id, 2, 1235, 0, "long_range"), + nvtx_record(10, 290, process_id, 2, 1235, 0, "long_range"), + ]), + "xp_nvtx", + ) + .expect("NVTX capture must decode"); + + assert!(decoded.nvtx_callbacks_active); + assert_eq!(decoded.events[0].event_type, EventType::NvtxRangeStart); + assert_eq!(decoded.events[0].clock_domain, ClockDomain::HostMonotonic); + assert!(decoded.events[0].cuda.is_none()); + let thread = decoded.events[0].nvtx.as_ref().expect("NVTX payload"); + assert_eq!(thread.range_kind, NvtxRangeKind::Thread); + assert_eq!(thread.range_id, thread_id); + assert!(thread.name_complete); + let process = decoded.events[2].nvtx.as_ref().expect("NVTX payload"); + assert_eq!(process.range_kind, NvtxRangeKind::Process); + assert_eq!(process.range_id, process_id); + assert!(!process.name_complete); + } + + #[test] + fn rejects_invalid_nvtx_range_metadata() { + let error = decode_capture( + &nvtx_capture(&[nvtx_record(9, 100, 7, 3, 1235, 1, "bad")]), + "xp_nvtx", + ) + .expect_err("invalid NVTX range kind must fail"); + + assert_eq!( + error, + CuptiDecodeError::InvalidNvtxRecord { + index: 0, + message: "range kind is neither thread nor process", + } + ); + } + #[test] fn normalizes_flagged_gpu_timestamps_to_host_monotonic() { let api = record(1, 10_400, 42, "cudaLaunchKernel"); @@ -1114,10 +1263,10 @@ mod tests { #[test] fn rejects_unknown_feature_flags() { let mut bytes = capture(&[]); - bytes[20..24].copy_from_slice(&8_u32.to_le_bytes()); + bytes[20..24].copy_from_slice(&16_u32.to_le_bytes()); assert_eq!( decode_capture(&bytes, "xp_test"), - Err(CuptiDecodeError::UnsupportedFeatureFlags(8)) + Err(CuptiDecodeError::UnsupportedFeatureFlags(16)) ); } diff --git a/xprobe/correlator/src/lib.rs b/xprobe/correlator/src/lib.rs index dab8297..3b07fd9 100644 --- a/xprobe/correlator/src/lib.rs +++ b/xprobe/correlator/src/lib.rs @@ -682,6 +682,7 @@ pub fn measure( options.match_policy, options.samples, syscall_lifecycle, + nvtx_lifecycle, ); if outcome.pairs.is_empty() { return Err(MeasureError::NoMatchedSamples); @@ -763,9 +764,11 @@ fn correlate_selected<'a>( policy: MatchPolicy, limit: Option, syscall_lifecycle: bool, + nvtx_lifecycle: bool, ) -> Outcome<'a> { match policy { MatchPolicy::Exact if syscall_lifecycle => correlate_syscall_lifecycle(starts, ends, limit), + MatchPolicy::Exact if nvtx_lifecycle => correlate_nvtx_lifecycle(starts, ends, limit), MatchPolicy::Exact => correlate_exact(starts, ends, limit), MatchPolicy::FirstAfter => correlate_first_after(starts, ends, limit), MatchPolicy::Nearest => correlate_nearest(starts, ends, limit), @@ -865,6 +868,63 @@ fn correlate_syscall_lifecycle<'a>( } } +fn correlate_nvtx_lifecycle<'a>( + starts: &[&'a Event], + ends: &[&'a Event], + limit: Option, +) -> Outcome<'a> { + let mut timeline = starts + .iter() + .map(|event| (event.timestamp_ns, 0_u8, *event)) + .chain(ends.iter().map(|event| (event.timestamp_ns, 1_u8, *event))) + .collect::>(); + timeline.sort_by_key(|(timestamp, boundary, event)| (*timestamp, *boundary, event.sequence)); + + let mut pending = BTreeMap::::new(); + let mut pairs = Vec::new(); + let mut unmatched_start = 0_u64; + let mut unmatched_end = 0_u64; + let mut ambiguous = 0_u64; + for (_, boundary, event) in timeline { + let Some(key) = correlation_key(event) else { + if boundary == 0 { + unmatched_start += 1; + } else { + unmatched_end += 1; + } + continue; + }; + if boundary == 0 { + if pending.insert(key, event).is_some() { + ambiguous += 1; + } + } else if let Some(start) = pending.remove(&key) { + if let Some(latency_ns) = event.timestamp_ns.checked_sub(start.timestamp_ns) { + pairs.push(Pair { + start, + end: event, + latency_ns, + }); + } else { + ambiguous += 1; + } + } else { + unmatched_end += 1; + } + } + unmatched_start += pending.len() as u64; + pairs.sort_by_key(|pair| pair.start.timestamp_ns); + if let Some(limit) = limit { + pairs.truncate(limit); + } + Outcome { + pairs, + unmatched_start, + unmatched_end, + ambiguous, + } +} + fn validate_options(options: &MeasureOptions) -> Result<(), MeasureError> { if options.samples.is_none() && options.duration.is_none() { return Err(MeasureError::InvalidLimit( @@ -1581,6 +1641,14 @@ mod tests { 7, 10, ), + nvtx_event( + EventType::NvtxRangeStart, + 200, + 10, + NvtxRangeKind::Thread, + 7, + 10, + ), ]; let options = MeasureOptions { session_id: "xp_nvtx".to_owned(), @@ -1596,6 +1664,7 @@ mod tests { let result = measure(&events, &options).unwrap(); assert_eq!(result.measurement.samples.matched, 2); + assert_eq!(result.measurement.samples.unmatched_start, 1); assert_eq!(result.measurement.latency_ns.min, 50); assert_eq!(result.measurement.latency_ns.max, 80); assert_eq!(result.correlation.method, "exact_nvtx_range_id"); From 090adeddcb68f6a76a2c498a8e279180b1b9d6d8 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 08:16:27 +0800 Subject: [PATCH 16/17] =?UTF-8?q?=F0=9F=93=9D=20docs:=20guide=20NVTX=20sta?= =?UTF-8?q?rtup=20profiling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- AGENTS.md | 7 +++ README.md | 13 +++-- docs/architecture.md | 15 +++++- docs/cli-contract.md | 49 ++++++++++++++----- docs/cupti-agent.md | 28 ++++++++--- docs/development.md | 16 +++++- skills/xprobe-measure-latency/SKILL.md | 16 ++++-- .../references/cli-contract.md | 30 +++++++----- .../references/setup.md | 10 +++- 9 files changed, 144 insertions(+), 40 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 7c3ae4c..c408b15 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -33,6 +33,12 @@ - `validate` is read-only. `measure` may inject CUPTI only after validation reports `injection_required`; log the mutation and include a JSON warning. - Stop CUPTI logically after collection. Do not `dlclose` the injected agent. +- NVTX callback routing must be initialized before the target's first NVTX API + call. Once initialized, keep the CUPTI subscriber mapped and registered; + logical stop disables callback domains instead of unsubscribing it. +- Correlate NVTX IDs as ordered lifecycles. Thread range keys reuse push levels, + and a bounded capture may end with one range open; retain earlier closed + pairs and report only the trailing start as unmatched. - Never collect pointer-referenced payloads, environments, or GPU buffer data by default. Named tracepoints retain identity and timestamps unless a versioned scalar payload is explicitly designed. Never describe temporal correlation @@ -70,6 +76,7 @@ - Run `just fmt-check`, `just lint`, and `just test` for Rust changes. - Run `just test-bpf-live` for BPF attachment changes. - Run `just test-cupti-live` for CUPTI ABI or callback changes. +- Run `just test-nvtx-live` for NVTX callback, filtering, or lifecycle changes. - Run `just test-injection-live` for injection or agent lifecycle changes. - Run `just test-multisource-live` for host/GPU orchestration changes. - Run `just test-agent-contract` for CLI, schema, docs, or Skill changes. diff --git a/README.md b/README.md index d6f6115..50f263a 100644 --- a/README.md +++ b/README.md @@ -64,10 +64,10 @@ xprobe measure --pid 4242 \ Kernel launch latency is only one event pair. The same workflow measures host function spans, syscall latency, named Linux events, CUDA API calls, GPU -operation durations, transfers, and paths across CPU and GPU events after -selecting the correct process. Aggregate mode provides a bounded coarse -inventory of GPU operations before an exact evidence measurement narrows the -question. +operation durations, transfers, NVTX application ranges, and paths across CPU +and GPU events after selecting the correct process. Aggregate mode provides a +bounded coarse inventory of GPU operations before an exact evidence +measurement narrows the question. `measure` also accepts completed `--input` captures and versioned live `--spec` files. Evidence can be exported as `jsonl` or `chrome`. JSON results @@ -88,7 +88,9 @@ is also preserved when correlation or clock validation fails. `measure --pid` automatically loads the matching CUDA 12 or CUDA 13 CUPTI Agent when a selected endpoint requires it. It reports the target mutation on stderr and in JSON, disables collection afterward, and leaves the shared object mapped -for safe reactivation. +for safe reactivation. NVTX ranges are the exception: set +`NVTX_INJECTION64_PATH` before the target's first NVTX call because online +attach cannot retrofit an initialized NVTX dispatch. ## Support @@ -98,6 +100,7 @@ for safe reactivation. | Host events | PID-scoped ELF function, named syscall, and tracepoint boundaries | | CUDA callbacks | Runtime and Driver API entry/exit | | GPU activity | Kernel, memcpy, and memset start/end | +| Application ranges | Bounded ASCII NVTX thread and process ranges | | CUDA/CUPTI | 12.x and 13.x with automatic major selection | | Correlation | exact, first-after, nearest, stack-nested, stream-order | | Online injection | same mount namespace; ptrace permission required | diff --git a/docs/architecture.md b/docs/architecture.md index 0856f6a..cb63ad0 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -68,9 +68,13 @@ requirements. CUPTI activation is one of: - `not_required` - `already_loaded` - `injection_required` +- `startup_required` `injection_required` keeps validation valid when all semantic requirements are met, sets `target_mutation: true`, and emits `TARGET_PROCESS_WILL_BE_MODIFIED`. +`startup_required` is specific to an NVTX selector whose target did not load +the Agent before NVTX initialized. Online attach cannot rewrite that dispatch +table, so validation returns an explicit issue instead of promising injection. Malformed selectors, unresolved symbols, invalid policies, and unavailable required host collection remain explicit issues or errors. @@ -99,6 +103,15 @@ Driver callbacks plus kernel, memcpy, and memset activity are filtered before fixed records consume capacity. Callback hot paths do not perform blocking I/O or allocation. +NVTX startup injection installs CUPTI routing before the first NVTX call and +keeps one dormant subscriber alive for the process lifetime. Each ARM enables +only the six supported global ASCII/Ex range callbacks. A preallocated, +open-addressed table retains matching starts until a thread pop or process +range end recovers the same bounded name and start thread. Thread keys combine +TID with the returned nesting level; process keys use the returned 64-bit NVTX +range ID and permit a different end thread. STOP disables callback domains but +does not unsubscribe the routing required by CUDA 13. + Broad GPU inventory uses the same `measure` primitive with aggregate mode. The Agent updates a `--max-groups`-bounded table for matching kernel, memcpy, or memset activity and returns only final count/duration/byte summaries. Exact @@ -143,7 +156,7 @@ All sources normalize into the versioned `Event` type. `measure` supports: | Policy | Pairing | Confidence | | --- | --- | --- | -| `exact` | CUPTI correlation ID or same-thread syscall lifecycle | exact | +| `exact` | CUPTI correlation ID, NVTX range ID, or same-thread syscall lifecycle | exact | | `first-after` | First unused end at or after each start | heuristic | | `nearest` | Nearest unused end by timestamp | heuristic | | `stack-nested` | Per-thread LIFO host entry/return | high | diff --git a/docs/cli-contract.md b/docs/cli-contract.md index bdbf2c5..e5c95d0 100644 --- a/docs/cli-contract.md +++ b/docs/cli-contract.md @@ -86,9 +86,12 @@ return value on exit. It never dereferences pointer arguments. Named tracepoints record identity and timestamp only; they do not copy the tracepoint payload. Unsupported syscall names and unavailable tracepoints fail explicitly. -CUDA forms include Runtime/Driver API entry/exit and kernel, memcpy, or memset -activity start/end. Kernel selectors accept `name~REGEX`; memcpy selectors -accept `kind=`. +CUDA forms include Runtime/Driver API entry/exit, kernel/memcpy/memset activity +start/end, and NVTX range boundaries. Kernel selectors accept `name~REGEX`; +memcpy selectors accept `kind=`. NVTX selectors are +`cuda:nvtx_range_start:name~REGEX` and +`cuda:nvtx_range_end:name~REGEX`; their name expression must reduce to one +nonempty exact, prefix, suffix, or contains filter shorter than 128 bytes. ## `validate` @@ -104,18 +107,23 @@ parses CUDA filters, and checks collection and correlation requirements. Results conform to `schemas/validate.schema.json`. Supported policies are `exact`, `first-after`, `nearest`, `stack-nested`, and -`stream-order`. Exact uses deterministic CUPTI correlation IDs or one named -syscall's per-thread entry/exit lifecycle. Nested requires entry/return of the -same host function. Stream order requires GPU activity endpoints. Temporal -policies always warn that they are heuristic. +`stream-order`. Exact uses deterministic CUPTI correlation IDs, NVTX range +kind/ID, or one named syscall's per-thread entry/exit lifecycle. Nested +requires entry/return of the same host function. Stream order requires GPU +activity endpoints. Temporal policies always warn that they are heuristic. `policy_recommendation` reports the strongest compatible policy, a stable machine-readable reason, and all compatible alternatives. The caller still chooses the policy; xprobe does not silently replace or retry it. -`requirements.agent_activation` is `not_required`, `already_loaded`, or -`injection_required`. The last value sets `target_mutation: true` and emits -`TARGET_PROCESS_WILL_BE_MODIFIED`; it does not make an otherwise collectable -request invalid. Callers must still stop on `valid: false`. +`requirements.agent_activation` is `not_required`, `already_loaded`, +`injection_required`, or `startup_required`. Online injection handles ordinary +CUDA callbacks and activity; `injection_required` sets `target_mutation: true` +and emits `TARGET_PROCESS_WILL_BE_MODIFIED`. NVTX dispatch is fixed by its first +API initialization, so an unmapped Agent produces `startup_required` and an +invalid result. Restart the target with `NVTX_INJECTION64_PATH` set to the +matching xprobe Agent before its first NVTX call. A mapped Agent is verified +again at ARM through the capture feature flag. Callers must stop on +`valid: false`. Mapped CUDA/CUPTI majors that conflict or fall outside 12 and 13 produce `UNSUPPORTED_CUDA_VERSION`. CUDA major support does not imply host-clock alignment: `measure` accepts @@ -144,6 +152,16 @@ xprobe measure --pid 4242 \ --json --non-interactive --no-color ``` +NVTX application range: + +```bash +xprobe measure --pid 4242 \ + --from 'cuda:nvtx_range_start:name~^forward$' \ + --to 'cuda:nvtx_range_end:name~^forward$' \ + --match exact --samples 100 --max-events 1000 \ + --json --non-interactive --no-color +``` + Completed captures: ```bash @@ -203,6 +221,15 @@ a stderr warning and `CUPTI_AGENT_INJECTED`. Each live call arms a fresh automatic injection also removes its private socket while leaving the `.so` mapped. +NVTX thread ranges preserve nesting through the thread ID and zero-based push +level. Process ranges preserve the NVTX range ID and may end on a different +thread; evidence reports both boundary threads plus `start_tid`. Names are +copied into the fixed record bound with `name_complete`. ASCII +`nvtxRangePushA`/`nvtxRangeStartA` and ASCII messages in their `Ex` forms are +supported; wide and registered-string messages are not collected. This is +timeline boundary collection only. It does not enable CUPTI Range Profiling, +metric replay, or kernel performance-counter passes. + Completed inputs may be CUPTI ABI binary, bounded host-capture JSON, or Event JSONL. Repeated inputs are merged after target-PID checks. Unknown records, mixed targets, malformed captures, and excessive event counts fail explicitly. diff --git a/docs/cupti-agent.md b/docs/cupti-agent.md index 35edc28..27e6403 100644 --- a/docs/cupti-agent.md +++ b/docs/cupti-agent.md @@ -1,9 +1,9 @@ # CUPTI Agent and online injection -`libxprobe-cupti.so` collects CUDA Runtime/Driver callback boundaries and -kernel, memcpy, and memset activity inside the target process. Normal users -invoke it through `measure --pid`; it is a bounded collector, not a daemon or a -persistent profiling session. +`libxprobe-cupti.so` collects CUDA Runtime/Driver callback boundaries, +kernel/memcpy/memset activity, and selected NVTX timeline range boundaries +inside the target process. Normal users invoke it through `measure --pid`; it +is a bounded collector, not a daemon or a persistent profiling session. Release packages contain separate Agents linked to `libcupti.so.12` and `libcupti.so.13`. xprobe selects the major from the target's mapped CUDA @@ -25,6 +25,16 @@ directions, and simple kernel-name exact/prefix/suffix/contains patterns are filtered before they consume capture capacity. Complex kernel regular expressions use a wider Agent filter and retain exact Rust-side matching. +NVTX differs from ordinary online CUDA injection. Set +`NVTX_INJECTION64_PATH` to the matching Agent before the target's first NVTX +call; an initialized NVTX dispatch cannot be retrofitted by attach. Startup +installs routing and a dormant CUPTI subscriber. ARM enables only range +StartA/StartEx/End and PushA/PushEx/Pop callbacks, with bounded name matching +before timestamping or table insertion. STOP disables those callbacks but +retains the subscriber because CUDA 13 still references it. The Agent copies +only ASCII range names, IDs, kind, and thread identity. It does not copy NVTX +payload, color, category, registered strings, wide strings, or environments. + An aggregate ARM uses a fixed-capacity hash table instead of event records. Kernel activity is grouped by name and device, memcpy by direction and device, and memset by device. The activity callback updates count, total/min/max @@ -70,6 +80,10 @@ measurement to fail explicitly instead of returning partial success. Exact records preserve process/thread, device/context/stream, correlation IDs, callback domain/ID, dimensions or transfer metadata, and one bounded name. +NVTX records reuse the fixed layout for a 64-bit range ID, thread/process kind, +start thread, and explicit name completeness flag. Their exact correlation is +range identity, including a synthesized TID plus nesting level for thread +ranges; it is not a claim that arbitrary adjacent events are causally related. Aggregate records contain a bounded group key and exact integer counters, not event evidence or a latency distribution; they therefore do not report percentiles or correlation confidence. @@ -87,8 +101,10 @@ the validated endpoints. The API callback checks its fixed filter before reading time or constructing a record. Activity decoding checks event family, name, and transfer direction before copying fixed metadata, and stops converting records after the capture limit. These paths do not allocate, -perform file or socket I/O, symbolize names, or take blocking locks. Flushing -and incremental response I/O run on the control thread. +perform file or socket I/O, symbolize names, or take blocking locks. NVTX range +callbacks use only fixed records and a preallocated table; CUPTI Range +Profiling metrics and replay are outside this timeline collector. Flushing, +control, and incremental response I/O run on the control thread. ## Verification diff --git a/docs/development.md b/docs/development.md index 04343dd..b2aa794 100644 --- a/docs/development.md +++ b/docs/development.md @@ -81,7 +81,8 @@ PYTORCH_ENV=/path/to/env just test-pytorch-live With a CUDA-enabled PyTorch environment, run eager matrix multiplication, convolution, compiled Triton, bidirectional transfer, selected-kernel, and -stream-synchronization profiling on the local GPU: +stream-synchronization profiling on the local GPU. The test also wraps an eager +operation in an NVTX range and verifies exact range-ID matching: ```bash PYTORCH_ENV=/path/to/env just test-pytorch-cuda-live @@ -126,6 +127,19 @@ runs it against CUDA 12.0. GPU-only durations remain available when the runtime cannot prove host-clock alignment; cross CPU/GPU measurement then fails with `CLOCK_ALIGNMENT_FAILED` instead of reporting a shifted latency. +Run the bounded NVTX fixture against both supported CUPTI majors: + +```bash +just test-nvtx-live-cuda12 +just test-nvtx-live +``` + +It covers nested ASCII and extended ranges, cross-thread process ranges, +bounded-name truncation, capacity failure and reuse after logical stop. NVTX +initialization occurs before the fixture's first NVTX call; the test also +verifies that the subscriber remains mapped while callback domains are dormant +between captures. + Run online ptrace injection, stop, and reactivation twice against a target that does not preload the agent: diff --git a/skills/xprobe-measure-latency/SKILL.md b/skills/xprobe-measure-latency/SKILL.md index 91f2432..5019d04 100644 --- a/skills/xprobe-measure-latency/SKILL.md +++ b/skills/xprobe-measure-latency/SKILL.md @@ -53,14 +53,20 @@ For more than one selected process, follow `scripts/analyze_trace.py` and use launch variants, stream distribution, busy union, overlap factor, and adjacent gaps. Read [references/trace-analysis.md](references/trace-analysis.md) when interpreting - the report. For CPU-only work, use resolved host selectors or filtered - syscall/tracepoint evidence to form the hypothesis instead. + the report. When the application already marks the narrowed operation with a + bounded ASCII NVTX range, that range can provide exact application-level + start/end boundaries. For CPU-only work, use resolved host selectors or + filtered syscall/tracepoint evidence to form the hypothesis instead. 7. Run one read-only `xprobe validate` per selected worker. Compare every response target with the PID plus process start time retained from discovery before mutation, and stop that worker when `valid` is false. If `agent_activation` is `injection_required`, disclose that `measure` will ptrace the target and leave the CUPTI shared object mapped. Use - `policy_recommendation` explicitly; xprobe never changes policy for the caller. + `policy_recommendation` explicitly; xprobe never changes policy for the + caller. If it is `startup_required`, restart that worker with + `NVTX_INJECTION64_PATH` pointing to the matching xprobe CUPTI Agent before + its first NVTX call, reacquire PID plus start time, and validate again. + Online injection cannot retrofit NVTX dispatch into an initialized process. 8. Run one bounded `xprobe measure` for that hypothesis. Set samples or duration, timeout, and max-events; write `--events-out` when the capture may need audit or offline re-correlation. Use a versioned `--spec FILE` containing the stable @@ -91,7 +97,9 @@ hypothesis; orchestration remains the agent framework's responsibility. - Stop on target reuse, permission failure, invalid selectors, unavailable collectors, drops, incomplete capture, unknown clock alignment, or unexamined - ambiguity. Read structured `details`, `hints`, and any failed-capture artifact. + ambiguity. Treat an NVTX ARM-time feature mismatch as a startup failure rather + than retrying online injection. Read structured `details`, `hints`, and any + failed-capture artifact. - Do not claim request causality from `first-after` or `nearest`. Do not compare or sum events across streams as if they were serial. - Stop using xprobe once evidence isolates time inside one kernel. Kernel diff --git a/skills/xprobe-measure-latency/references/cli-contract.md b/skills/xprobe-measure-latency/references/cli-contract.md index e7cfc7e..577e26d 100644 --- a/skills/xprobe-measure-latency/references/cli-contract.md +++ b/skills/xprobe-measure-latency/references/cli-contract.md @@ -35,19 +35,23 @@ does not dereference pointers. Named tracepoints contain no payload fields. Always let `validate` resolve the named syscall or tracepoint on the current host before measurement. -CUDA selectors cover Runtime and Driver API entry/exit plus kernel, memcpy, and -memset activity start/end. Kernel activity accepts `name~REGEX`; memcpy accepts -`kind=`. `validate` must accept the complete selector -and policy before measurement. +CUDA selectors cover Runtime and Driver API entry/exit; kernel, memcpy, and +memset activity start/end; and bounded NVTX range start/end. Kernel activity +accepts `name~REGEX`; memcpy accepts `kind=`. NVTX +uses `cuda:nvtx_range_start:name~REGEX` and +`cuda:nvtx_range_end:name~REGEX`; the regex must reduce to an exact, prefix, +suffix, or contains match shorter than 128 bytes. `validate` must accept the +complete selector and policy before measurement. ## Correlation -Use `exact` for CUDA events with the same CUPTI correlation ID or entry/exit of -one named syscall, `stack-nested` for entry/return of the same host function, -and `stream-order` for activity events on one CUDA stream. `first-after` and -`nearest` are temporal heuristics and cannot establish causality. Read -`policy_recommendation.policy`, its machine-readable `reason`, and -`compatible_policies`; xprobe never silently changes the requested policy. +Use `exact` for CUDA events with the same CUPTI correlation ID, NVTX boundaries +with the same range kind and ID, or entry/exit of one named syscall. +Use `stack-nested` for entry/return of the same host function and `stream-order` +for activity events on one CUDA stream. `first-after` and `nearest` are temporal +heuristics and cannot establish causality. Read `policy_recommendation.policy`, +its machine-readable `reason`, and `compatible_policies`; xprobe never silently +changes the requested policy. ## Bounds and failures @@ -96,4 +100,8 @@ error code, message, details, and hints. Validation is read-only. When it reports `injection_required`, the following live `measure` may ptrace the target and load the matching CUDA 12 or CUDA 13 CUPTI Agent. Measurement disables the Agent logically after collection but -does not unload the shared object. +does not unload the shared object. When validation reports `startup_required` +for an NVTX selector, restart the worker with `NVTX_INJECTION64_PATH` pointing +to that Agent before its first NVTX call, reacquire its PID plus start time, and +validate again. Online injection cannot install NVTX dispatch after NVTX has +initialized. diff --git a/skills/xprobe-measure-latency/references/setup.md b/skills/xprobe-measure-latency/references/setup.md index d2bf61a..e4696c3 100644 --- a/skills/xprobe-measure-latency/references/setup.md +++ b/skills/xprobe-measure-latency/references/setup.md @@ -69,7 +69,15 @@ export XPROBE_CUPTI_AGENT_PATH="$PWD/build/cuda12/cupti/libxprobe-cupti.so" Replace `12` and the toolkit path with `13` for CUDA 13. CUDA/CUPTI majors other than 12 or 13 are not supported: do not force an Agent build or bypass the -version check. +version check. For an NVTX range measurement, start or restart the target with +the same local Agent configured before the target's first NVTX call: + +```bash +NVTX_INJECTION64_PATH="$XPROBE_CUPTI_AGENT_PATH" python serve.py +``` + +Record the new PID plus procfs start time and run `validate` again. Do not use +online injection as a fallback for an already initialized NVTX process. ## Repair the Skill only when needed From abe7319f4444297305c351495eb395f875aff883 Mon Sep 17 00:00:00 2001 From: itdevwu Date: Fri, 24 Jul 2026 14:59:35 +0800 Subject: [PATCH 17/17] =?UTF-8?q?=F0=9F=94=96=20chore:=20release=20v0.4.0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- CMakeLists.txt | 2 +- Cargo.lock | 12 ++++++------ Cargo.toml | 2 +- README.md | 2 +- docs/agent-integration.md | 2 +- docs/installation.md | 12 ++++++------ install.sh | 4 ++-- skills/xprobe-measure-latency/SKILL.md | 2 +- skills/xprobe-measure-latency/references/setup.md | 12 ++++++------ tests/agent-contract/test_contract.py | 2 +- 10 files changed, 26 insertions(+), 26 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f048aa7..3f38fd7 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,6 +1,6 @@ cmake_minimum_required(VERSION 3.20) -project(xprobe VERSION 0.3.3 LANGUAGES C) +project(xprobe VERSION 0.4.0 LANGUAGES C) include(CTest) diff --git a/Cargo.lock b/Cargo.lock index c178dee..5febe2a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -461,7 +461,7 @@ dependencies = [ [[package]] name = "xprobe-cli" -version = "0.3.3" +version = "0.4.0" dependencies = [ "clap", "serde_json", @@ -474,7 +474,7 @@ dependencies = [ [[package]] name = "xprobe-collector" -version = "0.3.3" +version = "0.4.0" dependencies = [ "libbpf-rs", "serde_json", @@ -483,7 +483,7 @@ dependencies = [ [[package]] name = "xprobe-core" -version = "0.3.3" +version = "0.4.0" dependencies = [ "cpp_demangle", "nix", @@ -494,7 +494,7 @@ dependencies = [ [[package]] name = "xprobe-correlator" -version = "0.3.3" +version = "0.4.0" dependencies = [ "regex", "serde_json", @@ -503,7 +503,7 @@ dependencies = [ [[package]] name = "xprobe-exporter" -version = "0.3.3" +version = "0.4.0" dependencies = [ "serde_json", "xprobe-protocol", @@ -511,7 +511,7 @@ dependencies = [ [[package]] name = "xprobe-protocol" -version = "0.3.3" +version = "0.4.0" dependencies = [ "schemars", "serde", diff --git a/Cargo.toml b/Cargo.toml index 79b0711..5a08b1f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,7 +10,7 @@ members = [ resolver = "2" [workspace.package] -version = "0.3.3" +version = "0.4.0" edition = "2024" license = "Apache-2.0" repository = "https://github.com/itdevwu/xprobe" diff --git a/README.md b/README.md index 50f263a..da2d179 100644 --- a/README.md +++ b/README.md @@ -22,7 +22,7 @@ only installation action required from the user: ```bash npx skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --global ``` diff --git a/docs/agent-integration.md b/docs/agent-integration.md index f7832e6..5c4333f 100644 --- a/docs/agent-integration.md +++ b/docs/agent-integration.md @@ -17,7 +17,7 @@ repairs the matching xprobe CLI itself: ```bash npx skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --global ``` diff --git a/docs/installation.md b/docs/installation.md index 3caf960..9e5657c 100644 --- a/docs/installation.md +++ b/docs/installation.md @@ -13,7 +13,7 @@ user needs to run: ```bash npx skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --global ``` @@ -28,7 +28,7 @@ The versioned bootstrap installs to `~/.local` without root access: ```bash curl --proto '=https' --tlsv1.2 -fsSL \ - https://raw.githubusercontent.com/itdevwu/xprobe/v0.3.3/install.sh | sh + https://raw.githubusercontent.com/itdevwu/xprobe/v0.4.0/install.sh | sh ``` The bootstrap downloads the release archive and its SHA256 file, verifies the @@ -47,7 +47,7 @@ prefix, download the script and pass `--prefix`: ```bash curl --proto '=https' --tlsv1.2 -fsSLO \ - https://raw.githubusercontent.com/itdevwu/xprobe/v0.3.3/install.sh + https://raw.githubusercontent.com/itdevwu/xprobe/v0.4.0/install.sh sh install.sh --prefix /opt/xprobe ``` @@ -59,7 +59,7 @@ script with `sudo`. The installer never elevates privileges itself. For a fully explicit archive workflow: ```bash -version=0.3.3 +version=0.4.0 base=https://github.com/itdevwu/xprobe/releases/download/v$version archive=xprobe-$version-linux-x86_64.tar.gz @@ -92,7 +92,7 @@ missing or damaged, manually install the complete version-matched directory: ```bash npx skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --global ``` @@ -101,7 +101,7 @@ installation names the target explicitly: ```bash npx --yes skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --agent codex --global --copy --yes ``` diff --git a/install.sh b/install.sh index c8b183f..70935f5 100755 --- a/install.sh +++ b/install.sh @@ -2,7 +2,7 @@ set -eu repository=${XPROBE_REPOSITORY:-itdevwu/xprobe} -version=${XPROBE_VERSION:-0.3.3} +version=${XPROBE_VERSION:-0.4.0} if [ -n "${XPROBE_PREFIX:-}" ]; then prefix=$XPROBE_PREFIX elif [ -n "${HOME:-}" ]; then @@ -19,7 +19,7 @@ Install a released xprobe binary and its CUDA Agents. Usage: install.sh [--version VERSION] [--prefix DIR] [--uninstall] Options: - --version VERSION Release to install (default: 0.3.3) + --version VERSION Release to install (default: 0.4.0) --prefix DIR Installation prefix (default: $HOME/.local) --uninstall Remove xprobe from the selected prefix -h, --help Show this help diff --git a/skills/xprobe-measure-latency/SKILL.md b/skills/xprobe-measure-latency/SKILL.md index 5019d04..a437d98 100644 --- a/skills/xprobe-measure-latency/SKILL.md +++ b/skills/xprobe-measure-latency/SKILL.md @@ -17,7 +17,7 @@ For more than one selected process, follow ## Workflow -1. Run `xprobe --version`. When the command is absent or not 0.3.3, read +1. Run `xprobe --version`. When the command is absent or not 0.4.0, read [references/setup.md](references/setup.md) and install or repair the CLI yourself; do not ask the user to perform a separate CLI installation. Confirm every JSON response has `schema_version: "2.0"`; do not assume a diff --git a/skills/xprobe-measure-latency/references/setup.md b/skills/xprobe-measure-latency/references/setup.md index e4696c3..5eb5e76 100644 --- a/skills/xprobe-measure-latency/references/setup.md +++ b/skills/xprobe-measure-latency/references/setup.md @@ -7,7 +7,7 @@ command unless their environment prevents the agent from writing a usable prefix ## Check and bootstrap xprobe -Check the executable first. Continue only when it reports 0.3.3; otherwise run +Check the executable first. Continue only when it reports 0.4.0; otherwise run the bootstrap below: ```bash @@ -21,7 +21,7 @@ installing under `~/.local`: ```bash curl --proto '=https' --tlsv1.2 -fsSL \ - https://raw.githubusercontent.com/itdevwu/xprobe/v0.3.3/install.sh \ + https://raw.githubusercontent.com/itdevwu/xprobe/v0.4.0/install.sh \ -o /tmp/xprobe-install.sh sh /tmp/xprobe-install.sh export PATH="$HOME/.local/bin:$PATH" @@ -44,7 +44,7 @@ host glibc instead. This is a local-use fallback, not permission to weaken the release package's `GLIBC_2.34` ceiling. ```bash -git clone --depth 1 --branch v0.3.3 https://github.com/itdevwu/xprobe.git +git clone --depth 1 --branch v0.4.0 https://github.com/itdevwu/xprobe.git cd xprobe mamba env create --file environment.yml mamba run -n xprobe-dev just build @@ -82,12 +82,12 @@ online injection as a fallback for an already initialized NVTX process. ## Repair the Skill only when needed The user normally installed this Skill before invoking the agent. When its files -are missing or the version is not 0.3.3, install the complete version-matched +are missing or the version is not 0.4.0, install the complete version-matched directory through the Agent Skills CLI: ```bash npx skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --global ``` @@ -95,7 +95,7 @@ For non-interactive automation, select the host explicitly: ```bash npx --yes skills@1 add \ - https://github.com/itdevwu/xprobe/tree/v0.3.3/skills/xprobe-measure-latency \ + https://github.com/itdevwu/xprobe/tree/v0.4.0/skills/xprobe-measure-latency \ --agent codex --global --copy --yes ``` diff --git a/tests/agent-contract/test_contract.py b/tests/agent-contract/test_contract.py index 392a7bc..9d5672c 100755 --- a/tests/agent-contract/test_contract.py +++ b/tests/agent-contract/test_contract.py @@ -165,7 +165,7 @@ def check_skill(workspace: pathlib.Path) -> None: ): assert required in normalized_trace_analysis for required in ( - "v0.3.3/install.sh", + "v0.4.0/install.sh", "npx skills@1 add", "xprobe --version", "xprobe doctor",