From abba53d4427e4d65cb10de398f429acdc4c78670 Mon Sep 17 00:00:00 2001 From: Darshan Sanghani Date: Mon, 2 Mar 2026 18:17:15 -0800 Subject: [PATCH] Add PortConfigLoader for on-demand GPU tracing via TCP (#1280) Summary: This implements Phase 1.1 of the lightweight on-demand GPU tracing system for Kubernetes environments. It provides an alternative to DaemonConfigLoader that receives trace configs via a TCP port instead of IPC Fabric, enabling GPU tracing in K8s pods without requiring a full Dynolog daemon. The PortConfigLoader implements the IDaemonConfigLoader interface with a simple JSON protocol (PING/TRACE commands) and includes a socket abstraction layer (ISocket) for testability. Unit tests use MockSocket to enable deterministic testing without binding to real ports. Differential Revision: D91278956 --- libkineto/src/ISocket.h | 51 ++++ libkineto/src/IpcFabricConfigClient.cpp | 4 +- libkineto/src/Logger.cpp | 1 - libkineto/src/PortConfigLoader.cpp | 264 ++++++++++++++++++ libkineto/src/PortConfigLoader.h | 99 +++++++ libkineto/src/TraceProtocol.cpp | 156 +++++++++++ libkineto/src/TraceProtocol.h | 51 ++++ libkineto/test/MockSocket.h | 132 +++++++++ libkineto/test/PortConfigLoaderTest.cpp | 352 ++++++++++++++++++++++++ libkineto/test/TraceProtocolTest.cpp | 229 +++++++++++++++ 10 files changed, 1336 insertions(+), 3 deletions(-) create mode 100644 libkineto/src/ISocket.h create mode 100644 libkineto/src/PortConfigLoader.cpp create mode 100644 libkineto/src/PortConfigLoader.h create mode 100644 libkineto/src/TraceProtocol.cpp create mode 100644 libkineto/src/TraceProtocol.h create mode 100644 libkineto/test/MockSocket.h create mode 100644 libkineto/test/PortConfigLoaderTest.cpp create mode 100644 libkineto/test/TraceProtocolTest.cpp diff --git a/libkineto/src/ISocket.h b/libkineto/src/ISocket.h new file mode 100644 index 000000000..d6f98c28f --- /dev/null +++ b/libkineto/src/ISocket.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include + +namespace KINETO_NAMESPACE { + +// Abstract socket interface for dependency injection and testability. +// This allows PortConfigLoader to be tested without binding to real ports. +class ISocket { + public: + virtual ~ISocket() = default; + + // Create a TCP server socket bound to the given port. + // Returns the server file descriptor, or -1 on error. + virtual int createServer(uint16_t port) = 0; + + // Accept a connection on the server socket. + // Returns the client file descriptor, or -1 on error. + virtual int + accept(int server_fd, struct sockaddr* addr, socklen_t* addrlen) = 0; + + // Read from a file descriptor. + // Returns the number of bytes read, 0 on EOF, or -1 on error. + virtual ssize_t read(int fd, void* buf, size_t count) = 0; + + // Write to a file descriptor. + // Returns the number of bytes written, or -1 on error. + virtual ssize_t write(int fd, const void* buf, size_t count) = 0; + + // Close a file descriptor. + // Returns 0 on success, or -1 on error. + virtual int close(int fd) = 0; +}; + +// Factory function type for creating socket implementations +using SocketFactory = std::unique_ptr (*)(); + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/src/IpcFabricConfigClient.cpp b/libkineto/src/IpcFabricConfigClient.cpp index 36a179ff4..c52edccd8 100644 --- a/libkineto/src/IpcFabricConfigClient.cpp +++ b/libkineto/src/IpcFabricConfigClient.cpp @@ -176,8 +176,8 @@ std::string IpcFabricConfigClient::getLibkinetoOndemandConfig(int32_t type) { try { if (!fabricManager_->sync_send(*msg, std::string(kDynoIpcName))) { - LOG(ERROR) << "Failed to send config type=" << type - << " to dyno: IPC sync_send fail"; + VLOG(1) << "Failed to send config type=" << type + << " to dyno: IPC sync_send fail"; free(req); req = nullptr; return ""; diff --git a/libkineto/src/Logger.cpp b/libkineto/src/Logger.cpp index a81914920..2cb292402 100644 --- a/libkineto/src/Logger.cpp +++ b/libkineto/src/Logger.cpp @@ -16,7 +16,6 @@ #include #include #include -#include #include #include diff --git a/libkineto/src/PortConfigLoader.cpp b/libkineto/src/PortConfigLoader.cpp new file mode 100644 index 000000000..e1669fa0d --- /dev/null +++ b/libkineto/src/PortConfigLoader.cpp @@ -0,0 +1,264 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "PortConfigLoader.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include "ILoggerObserver.h" +#include "Logger.h" +#include "TraceProtocol.h" + +namespace KINETO_NAMESPACE { + +namespace { + +// Default socket implementation using POSIX sockets. +class PosixSocket : public ISocket { + public: + int createServer(uint16_t port) override { + int const fd = ::socket(AF_INET, SOCK_STREAM, 0); + if (fd < 0) { + return -1; + } + + int opt = 1; + if (::setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt)) < 0) { + ::close(fd); + return -1; + } + + struct sockaddr_in addr{}; + addr.sin_family = AF_INET; + addr.sin_addr.s_addr = INADDR_ANY; + addr.sin_port = htons(port); + + if (::bind(fd, reinterpret_cast(&addr), sizeof(addr)) < + 0) { + ::close(fd); + return -1; + } + + if (::listen(fd, 5) < 0) { + ::close(fd); + return -1; + } + + return fd; + } + + int accept(int serverFd, struct sockaddr* addr, socklen_t* addrlen) override { + return ::accept(serverFd, addr, addrlen); + } + + ssize_t read(int fd, void* buf, size_t count) override { + return ::read(fd, buf, count); + } + + ssize_t write(int fd, const void* buf, size_t count) override { + return ::write(fd, buf, count); + } + + int close(int fd) override { + return ::close(fd); + } +}; + +} // namespace + +PortConfigLoader::PortConfigLoader() + : PortConfigLoader(std::make_unique()) {} + +PortConfigLoader::PortConfigLoader(std::unique_ptr socket) + : socket_(std::move(socket)) { + // Read port from environment variable + const char* portEnv = std::getenv("KINETO_TRACE_PORT"); + if (portEnv != nullptr) { + port_ = static_cast(std::atoi(portEnv)); + } + startServer(); +} + +PortConfigLoader::PortConfigLoader( + uint16_t port, + std::unique_ptr socket) + : socket_(std::move(socket)), port_(port) { + // Don't start server automatically in test mode + // Tests will call testHandleOneConnection() directly +} + +PortConfigLoader::~PortConfigLoader() { + running_ = false; + if (serverFd_ >= 0) { + socket_->close(serverFd_); + } + if (serverThread_ && serverThread_->joinable()) { + serverThread_->join(); + } +} + +std::string PortConfigLoader::readOnDemandConfig( + bool /* events */, + bool activities) { + std::lock_guard const lock(configMutex_); + if (!activities || !configPending_) { + return ""; + } + + std::string config = std::move(pendingConfig_); + configPending_ = false; + pendingConfig_.clear(); + LOG(INFO) << "Received on-demand config from TCP port " << port_ << ":\n" + << config; + return config; +} + +bool PortConfigLoader::hasConfigPending() const { + std::lock_guard const lock(configMutex_); + return configPending_; +} + +void PortConfigLoader::testHandleOneConnection() { + // Initialize server if not already done + if (serverFd_ < 0) { + serverFd_ = socket_->createServer(port_); + } + + struct sockaddr_in clientAddr{}; + socklen_t addrLen = sizeof(clientAddr); + int const clientFd = socket_->accept( + serverFd_, reinterpret_cast(&clientAddr), &addrLen); + + if (clientFd >= 0) { + handleClient(clientFd); + socket_->close(clientFd); + } +} + +void PortConfigLoader::startServer() { + serverFd_ = socket_->createServer(port_); + if (serverFd_ < 0) { + LOG(WARNING) << "PortConfigLoader: Failed to create server on port " + << port_; + return; + } + + running_ = true; + serverThread_ = std::make_unique([this]() { serverLoop(); }); +} + +void PortConfigLoader::serverLoop() { + while (running_) { + struct sockaddr_in clientAddr{}; + socklen_t addrLen = sizeof(clientAddr); + int const clientFd = socket_->accept( + serverFd_, reinterpret_cast(&clientAddr), &addrLen); + + if (clientFd < 0) { + if (running_) { + LOG(WARNING) << "PortConfigLoader: Accept failed"; + } + continue; + } + + handleClient(clientFd); + socket_->close(clientFd); + } +} + +void PortConfigLoader::handleClient(int clientFd) { + char buffer[4096]; + ssize_t const bytesRead = socket_->read(clientFd, buffer, sizeof(buffer) - 1); + if (bytesRead <= 0) { + return; + } + buffer[bytesRead] = '\0'; + + std::string const request(buffer, bytesRead); + std::string response; + + // Parse JSON to determine message type + std::string const msgType = extractJsonString(request, "type", ""); + + if (msgType == "PING") { + response = handlePing(); + } else if (msgType == "TRACE") { + response = handleTrace(request); + } else { + response = R"({"type":"ERROR","message":"unknown_command"})" + "\n"; + } + + socket_->write(clientFd, response.c_str(), response.size()); +} + +std::string PortConfigLoader::handlePing() const { + std::lock_guard const lock(configMutex_); + if (configPending_) { + return R"({"type":"PONG","status":"BUSY"})" + "\n"; + } + return R"({"type":"PONG","status":"READY"})" + "\n"; +} + +std::string PortConfigLoader::handleTrace(std::string_view jsonPayload) { + std::lock_guard const lock(configMutex_); + + // Extract trace_id for response + std::string const traceId = + extractJsonString(jsonPayload, "trace_id", "unknown"); + + // Config Pending Guard: reject if config not yet consumed + if (configPending_) { + return R"({"type":"TRACE_ACK","status":"REJECTED","trace_id":")" + traceId + + R"(","error_code":"CONFIG_PENDING","error_message":"A trace config is pending"})" + "\n"; + } + + // Parse JSON and build config string + int const durationMs = extractJsonInt(jsonPayload, "duration_ms", 5000); + std::string const activities = + extractJsonString(jsonPayload, "activities", "CUDA,CPU"); + bool const recordShapes = + extractJsonBool(jsonPayload, "record_shapes", false); + bool const profileMemory = + extractJsonBool(jsonPayload, "profile_memory", false); + bool const withStack = extractJsonBool(jsonPayload, "with_stack", false); + bool const withFlops = extractJsonBool(jsonPayload, "with_flops", false); + bool const withModules = extractJsonBool(jsonPayload, "with_modules", false); + std::string const outputDir = + extractJsonString(jsonPayload, "output_dir", ""); + + pendingConfig_ = buildConfigString( + durationMs, + activities, + recordShapes, + profileMemory, + withStack, + withFlops, + withModules, + outputDir, + traceId); + configPending_ = true; + + return R"({"type":"TRACE_ACK","status":"ACCEPTED","trace_id":")" + traceId + + R"("})" + "\n"; +} + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/src/PortConfigLoader.h b/libkineto/src/PortConfigLoader.h new file mode 100644 index 000000000..38643db9d --- /dev/null +++ b/libkineto/src/PortConfigLoader.h @@ -0,0 +1,99 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "DaemonConfigLoader.h" +#include "ISocket.h" + +namespace KINETO_NAMESPACE { + +// PortConfigLoader implements IDaemonConfigLoader to receive on-demand trace +// requests via a TCP port. This provides an alternative to DaemonConfigLoader +// for Kubernetes environments where IPC Fabric is not available. +// +// Protocol: +// PING -> READY | BUSY (health check) +// TRACE -> TRACE_ACK | TRACE_BUSY | TRACE_ERROR (trigger trace) +// +// Port is configured via KINETO_TRACE_PORT environment variable (default: +// 20599). The loader runs a TCP server thread that accepts connections and +// processes commands. +class PortConfigLoader : public IDaemonConfigLoader { + public: + // Create with default PosixSocket. + PortConfigLoader(); + + // Create with injected socket (for testing). + explicit PortConfigLoader(std::unique_ptr socket); + + // Create with port and injected socket (for testing). + PortConfigLoader(uint16_t port, std::unique_ptr socket); + + ~PortConfigLoader() override; + + // IDaemonConfigLoader interface + std::string readBaseConfig() override { + return ""; // Not supported for port-based loader + } + + std::string readOnDemandConfig(bool events, bool activities) override; + + int gpuContextCount(uint32_t /*device*/) override { + return 0; // Not applicable for port-based loader + } + + void setCommunicationFabric(bool enabled) override { + // No-op for port-based loader + } + + // Check if a trace request is pending (Config Pending Guard). + bool hasConfigPending() const; + + // Test hook: handle one connection manually (for unit tests). + void testHandleOneConnection(); + + private: + // Start the TCP server thread. + void startServer(); + + // Server thread main loop. + void serverLoop(); + + // Handle a single client connection. + void handleClient(int clientFd); + + // Handle PING command. + std::string handlePing() const; + + // Handle TRACE command. + std::string handleTrace(std::string_view jsonPayload); + + std::unique_ptr socket_; + std::unique_ptr serverThread_; + std::atomic running_{false}; + + // Config Pending Guard: prevents overwriting pending config. + mutable std::mutex configMutex_; + bool configPending_{false}; + std::string pendingConfig_; + + // Port configuration. + uint16_t port_{20599}; + int serverFd_{-1}; +}; + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/src/TraceProtocol.cpp b/libkineto/src/TraceProtocol.cpp new file mode 100644 index 000000000..3bc6baf0c --- /dev/null +++ b/libkineto/src/TraceProtocol.cpp @@ -0,0 +1,156 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include "TraceProtocol.h" + +#include +#include + +namespace KINETO_NAMESPACE { + +namespace { + +// Find the position of a key in JSON, returning npos if not found. +size_t findKey(std::string_view json, std::string_view key) { + // Look for "key" pattern + std::string const pattern = "\"" + std::string(key) + "\""; + return json.find(pattern); +} + +// Skip whitespace and colon after key. +size_t skipToValue(std::string_view json, size_t pos) { + while (pos < json.size() && + ((std::isspace(json[pos]) != 0) || json[pos] == ':')) { + ++pos; + } + return pos; +} + +} // namespace + +int extractJsonInt( + std::string_view json, + std::string_view key, + int defaultVal) { + size_t const keyPos = findKey(json, key); + if (keyPos == std::string_view::npos) { + return defaultVal; + } + + size_t pos = skipToValue(json, keyPos + key.size() + 2); + if (pos >= json.size()) { + return defaultVal; + } + + // Parse integer + int result = 0; + bool negative = false; + if (json[pos] == '-') { + negative = true; + ++pos; + } + while (pos < json.size() && (std::isdigit(json[pos]) != 0)) { + result = result * 10 + (json[pos] - '0'); + ++pos; + } + return negative ? -result : result; +} + +bool extractJsonBool( + std::string_view json, + std::string_view key, + bool defaultVal) { + size_t const keyPos = findKey(json, key); + if (keyPos == std::string_view::npos) { + return defaultVal; + } + + size_t const pos = skipToValue(json, keyPos + key.size() + 2); + if (pos >= json.size()) { + return defaultVal; + } + + if (json.substr(pos, 4) == "true") { + return true; + } else if (json.substr(pos, 5) == "false") { + return false; + } + return defaultVal; +} + +std::string extractJsonString( + std::string_view json, + std::string_view key, + std::string_view defaultVal) { + size_t const keyPos = findKey(json, key); + if (keyPos == std::string_view::npos) { + return std::string(defaultVal); + } + + size_t pos = skipToValue(json, keyPos + key.size() + 2); + if (pos >= json.size() || json[pos] != '"') { + return std::string(defaultVal); + } + + // Skip opening quote + ++pos; + size_t const endPos = json.find('"', pos); + if (endPos == std::string_view::npos) { + return std::string(defaultVal); + } + + return std::string(json.substr(pos, endPos - pos)); +} + +std::string buildConfigString( + int durationMs, + const std::string& activities, + bool recordShapes, + bool profileMemory, + bool withStack, + bool withFlops, + bool withModules, + const std::string& outputDir, + const std::string& traceId) { + std::string config; + + config += "ACTIVITIES_DURATION_MSECS=" + std::to_string(durationMs) + "\n"; + + if (!activities.empty()) { + config += "ACTIVITIES=" + activities + "\n"; + } + + if (recordShapes) { + config += "PROFILE_REPORT_INPUT_SHAPES=true\n"; + } + + if (profileMemory) { + config += "PROFILE_PROFILE_MEMORY=true\n"; + } + + if (withStack) { + config += "PROFILE_WITH_STACK=true\n"; + } + + if (withFlops) { + config += "PROFILE_WITH_FLOPS=true\n"; + } + + if (withModules) { + config += "PROFILE_WITH_MODULES=true\n"; + } + + // Set the output file path if outputDir and traceId are provided + if (!outputDir.empty() && !traceId.empty()) { + config += "ACTIVITIES_LOG_FILE=" + outputDir + "/" + traceId + ".json\n"; + } + + return config; +} + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/src/TraceProtocol.h b/libkineto/src/TraceProtocol.h new file mode 100644 index 000000000..cf5f3fefa --- /dev/null +++ b/libkineto/src/TraceProtocol.h @@ -0,0 +1,51 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +namespace KINETO_NAMESPACE { + +// JSON field extraction for TRACE request parsing. +// These functions use simple hand-rolled parsing to avoid adding JSON library +// dependencies. + +// Extract an integer value from JSON for the given key. +// Returns defaultVal if the key is not found or parsing fails. +int extractJsonInt(std::string_view json, std::string_view key, int defaultVal); + +// Extract a boolean value from JSON for the given key. +// Returns defaultVal if the key is not found or parsing fails. +bool extractJsonBool( + std::string_view json, + std::string_view key, + bool defaultVal); + +// Extract a string value from JSON for the given key. +// Returns defaultVal if the key is not found or parsing fails. +std::string extractJsonString( + std::string_view json, + std::string_view key, + std::string_view defaultVal); + +// Build a Kineto config string from parsed JSON fields. +// Uses Kineto's KEY=VALUE\n format. +std::string buildConfigString( + int durationMs, + const std::string& activities, + bool recordShapes, + bool profileMemory, + bool withStack, + bool withFlops, + bool withModules, + const std::string& outputDir = "", + const std::string& traceId = ""); + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/test/MockSocket.h b/libkineto/test/MockSocket.h new file mode 100644 index 000000000..89d77a865 --- /dev/null +++ b/libkineto/test/MockSocket.h @@ -0,0 +1,132 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include "src/ISocket.h" + +namespace KINETO_NAMESPACE { + +// Mock socket implementation for unit testing. +// Uses in-memory buffers instead of real sockets, enabling deterministic +// tests without binding to real ports. +class MockSocket : public ISocket { + public: + MockSocket() = default; + + // Pre-configure an expected connection with request data. + // When accept() is called, it will return a fake client FD and + // subsequent read() calls will return this request data. + void expectConnection(const std::string& requestData) { + pendingConnections_.push(requestData); + } + + // Check if there are pending connections + [[nodiscard]] bool hasPendingConnections() const { + return !pendingConnections_.empty(); + } + + // Get the data that was written to the socket (for verification) + [[nodiscard]] const std::string& getWrittenData() const { + return writeBuffer_; + } + + // Clear the write buffer for the next test + void clearWrittenData() { + writeBuffer_.clear(); + } + + // Get the number of times accept() was called + [[nodiscard]] int getAcceptCount() const { + return acceptCount_; + } + + // Get the number of times close() was called + [[nodiscard]] int getCloseCount() const { + return closeCount_; + } + + // ISocket interface implementation + + int createServer(uint16_t port) override { + serverPort_ = port; + return kFakeServerFd; + } + + int accept(int serverFd, struct sockaddr* /*addr*/, socklen_t* /*addrlen*/) + override { + if (serverFd != kFakeServerFd) { + return -1; + } + + if (pendingConnections_.empty()) { + // No more connections expected - simulate blocking or error + return -1; + } + + // Set up the read buffer with the expected request data + currentReadBuffer_ = pendingConnections_.front(); + pendingConnections_.pop(); + readPos_ = 0; + + acceptCount_++; + return kFakeClientFd; + } + + ssize_t read(int fd, void* buf, size_t count) override { + if (fd != kFakeClientFd) { + return -1; + } + + if (readPos_ >= currentReadBuffer_.size()) { + return 0; // EOF + } + + size_t const toRead = std::min(count, currentReadBuffer_.size() - readPos_); + std::memcpy(buf, currentReadBuffer_.data() + readPos_, toRead); + readPos_ += toRead; + return static_cast(toRead); + } + + ssize_t write(int fd, const void* buf, size_t count) override { + if (fd != kFakeClientFd) { + return -1; + } + + writeBuffer_.append(static_cast(buf), count); + return static_cast(count); + } + + int close(int fd) override { + if (fd == kFakeClientFd || fd == kFakeServerFd) { + closeCount_++; + return 0; + } + return -1; + } + + private: + static constexpr int kFakeServerFd = 100; + static constexpr int kFakeClientFd = 101; + + uint16_t serverPort_{0}; + std::queue pendingConnections_; + std::string currentReadBuffer_; + size_t readPos_{0}; + std::string writeBuffer_; + int acceptCount_{0}; + int closeCount_{0}; +}; + +} // namespace KINETO_NAMESPACE diff --git a/libkineto/test/PortConfigLoaderTest.cpp b/libkineto/test/PortConfigLoaderTest.cpp new file mode 100644 index 000000000..ad4c2c8b5 --- /dev/null +++ b/libkineto/test/PortConfigLoaderTest.cpp @@ -0,0 +1,352 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include +#include + +#include "MockSocket.h" +#include "PortConfigLoader.h" + +using namespace KINETO_NAMESPACE; + +// ============================================================================ +// MockSocket Tests (verify mock works correctly) +// ============================================================================ + +TEST(MockSocketTest, CreateServer) { + MockSocket socket; + int const fd = socket.createServer(20599); + EXPECT_GE(fd, 0); +} + +TEST(MockSocketTest, AcceptWithPendingConnection) { + MockSocket socket; + socket.createServer(20599); + socket.expectConnection(R"({"type":"PING"})"); + + int const clientFd = socket.accept(100, nullptr, nullptr); + EXPECT_GE(clientFd, 0); + EXPECT_EQ(socket.getAcceptCount(), 1); +} + +TEST(MockSocketTest, AcceptWithNoConnections) { + MockSocket socket; + socket.createServer(20599); + + int const clientFd = socket.accept(100, nullptr, nullptr); + EXPECT_EQ(clientFd, -1); +} + +TEST(MockSocketTest, ReadFromConnection) { + MockSocket socket; + socket.createServer(20599); + socket.expectConnection(R"({"type":"PING"})"); + + int const clientFd = socket.accept(100, nullptr, nullptr); + ASSERT_GE(clientFd, 0); + + char buffer[256]; + ssize_t const bytesRead = socket.read(clientFd, buffer, sizeof(buffer)); + EXPECT_EQ(bytesRead, 15); // Length of {"type":"PING"} + EXPECT_EQ(std::string(buffer, bytesRead), R"({"type":"PING"})"); +} + +TEST(MockSocketTest, WriteToConnection) { + MockSocket socket; + socket.createServer(20599); + socket.expectConnection(""); + + int const clientFd = socket.accept(100, nullptr, nullptr); + ASSERT_GE(clientFd, 0); + + std::string response = R"({"type":"PONG","status":"READY"})"; + ssize_t const bytesWritten = + socket.write(clientFd, response.data(), response.size()); + EXPECT_EQ(bytesWritten, static_cast(response.size())); + EXPECT_EQ(socket.getWrittenData(), response); +} + +TEST(MockSocketTest, CloseConnection) { + MockSocket socket; + socket.createServer(20599); + socket.expectConnection(""); + + int const clientFd = socket.accept(100, nullptr, nullptr); + ASSERT_GE(clientFd, 0); + + EXPECT_EQ(socket.close(clientFd), 0); + EXPECT_EQ(socket.getCloseCount(), 1); +} + +// ============================================================================ +// PortConfigLoader - PING Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, HandlePingReturnsReady) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Setup: Client sends PING request + mockSocketPtr->expectConnection( + R"({"type":"PING"})" + "\n"); + + // Create loader with mock socket + PortConfigLoader loader(20599, std::move(mockSocket)); + + // Handle one connection + loader.testHandleOneConnection(); + + // Verify response contains PONG and READY status + std::string const response = mockSocketPtr->getWrittenData(); + EXPECT_NE(response.find("\"type\":\"PONG\""), std::string::npos); + EXPECT_NE(response.find("\"status\":\"READY\""), std::string::npos); +} + +TEST(PortConfigLoaderTest, HandlePingReturnsBusyWhenConfigPending) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // First, send a TRACE request to make config pending + mockSocketPtr->expectConnection( + R"({"type":"TRACE","trace_id":"first","config":{"duration_ms":500}})" + "\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // Now check that the config is pending + EXPECT_TRUE(loader.hasConfigPending()); + + // Setup second connection for PING + auto mockSocket2 = std::make_unique(); + auto* mockSocket2Ptr = mockSocket2.get(); + mockSocket2Ptr->expectConnection( + R"({"type":"PING"})" + "\n"); + + // Note: In real implementation, we'd inject the new socket + // For this test, we're just verifying the concept +} + +// ============================================================================ +// PortConfigLoader - TRACE Accept Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, HandleTraceAcceptsWhenIdle) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Setup: Client sends TRACE request + std::string const traceRequest = R"({ + "type": "TRACE", + "trace_id": "test-trace-123", + "config": { + "duration_ms": 1000, + "record_shapes": true, + "profile_memory": false + } + })" + "\n"; + mockSocketPtr->expectConnection(traceRequest); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // Verify response shows ACCEPTED + std::string const response = mockSocketPtr->getWrittenData(); + EXPECT_NE(response.find("\"type\":\"TRACE_ACK\""), std::string::npos); + EXPECT_NE(response.find("\"status\":\"ACCEPTED\""), std::string::npos); + EXPECT_NE( + response.find("\"trace_id\":\"test-trace-123\""), std::string::npos); + + // Verify config is now pending + EXPECT_TRUE(loader.hasConfigPending()); +} + +TEST(PortConfigLoaderTest, HandleTraceAcceptsWithMinimalConfig) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Minimal config with just duration + std::string const traceRequest = + R"({"type":"TRACE","trace_id":"minimal","config":{"duration_ms":500}})" + "\n"; + mockSocketPtr->expectConnection(traceRequest); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + std::string const response = mockSocketPtr->getWrittenData(); + EXPECT_NE(response.find("\"status\":\"ACCEPTED\""), std::string::npos); +} + +// ============================================================================ +// PortConfigLoader - TRACE Reject Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, HandleTraceRejectsWhenConfigPending) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // First TRACE request - should be accepted + mockSocketPtr->expectConnection( + R"({"type":"TRACE","trace_id":"first","config":{"duration_ms":500}})" + "\n"); + // Second TRACE request - should be rejected + mockSocketPtr->expectConnection( + R"({"type":"TRACE","trace_id":"second","config":{"duration_ms":500}})" + "\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + + // First request - accepted + loader.testHandleOneConnection(); + std::string const response1 = mockSocketPtr->getWrittenData(); + EXPECT_NE(response1.find("\"status\":\"ACCEPTED\""), std::string::npos); + + // Clear write buffer and handle second request + mockSocketPtr->clearWrittenData(); + loader.testHandleOneConnection(); + + // Second request should be rejected with CONFIG_PENDING + std::string const response2 = mockSocketPtr->getWrittenData(); + EXPECT_NE(response2.find("\"status\":\"REJECTED\""), std::string::npos); + EXPECT_NE( + response2.find("\"error_code\":\"CONFIG_PENDING\""), std::string::npos); +} + +// ============================================================================ +// PortConfigLoader - readOnDemandConfig Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, ReadOnDemandConfigConsumesPendingConfig) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Send a TRACE request + mockSocketPtr->expectConnection( + R"({"type":"TRACE","trace_id":"test","config":{"duration_ms":1000}})" + "\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // Config should be pending + EXPECT_TRUE(loader.hasConfigPending()); + + // First call to readOnDemandConfig should return the config + std::string const config = loader.readOnDemandConfig(false, true); + EXPECT_FALSE(config.empty()); + EXPECT_NE(config.find("ACTIVITIES_DURATION_MSECS=1000"), std::string::npos); + + // Config should no longer be pending + EXPECT_FALSE(loader.hasConfigPending()); + + // Second call should return empty + std::string const config2 = loader.readOnDemandConfig(false, true); + EXPECT_TRUE(config2.empty()); +} + +TEST(PortConfigLoaderTest, ReadOnDemandConfigReturnsEmptyWhenNoPending) { + auto mockSocket = std::make_unique(); + + PortConfigLoader loader(20599, std::move(mockSocket)); + + // No pending config - should return empty + std::string const config = loader.readOnDemandConfig(false, true); + EXPECT_TRUE(config.empty()); +} + +TEST(PortConfigLoaderTest, ReadOnDemandConfigReturnsEmptyWhenActivitiesFalse) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Send a TRACE request + mockSocketPtr->expectConnection( + R"({"type":"TRACE","trace_id":"test","config":{"duration_ms":500}})" + "\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // activities=false should not return the config + std::string const config = loader.readOnDemandConfig(false, false); + EXPECT_TRUE(config.empty()); + + // Config should still be pending + EXPECT_TRUE(loader.hasConfigPending()); +} + +// ============================================================================ +// PortConfigLoader - readBaseConfig Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, ReadBaseConfigReturnsEmpty) { + auto mockSocket = std::make_unique(); + + PortConfigLoader loader(20599, std::move(mockSocket)); + + // PortConfigLoader doesn't support base config - should return empty + std::string const config = loader.readBaseConfig(); + EXPECT_TRUE(config.empty()); +} + +// ============================================================================ +// PortConfigLoader - Error Handling Tests +// ============================================================================ + +TEST(PortConfigLoaderTest, HandleUnknownMessageType) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Unknown message type + mockSocketPtr->expectConnection( + R"({"type":"UNKNOWN"})" + "\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // Should respond with an error or ignore + // (implementation will define exact behavior) + std::string const response = mockSocketPtr->getWrittenData(); + // At minimum, it should not crash and should write something + EXPECT_FALSE(response.empty()); +} + +TEST(PortConfigLoaderTest, HandleMalformedJson) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Malformed JSON + mockSocketPtr->expectConnection("not valid json\n"); + + PortConfigLoader loader(20599, std::move(mockSocket)); + loader.testHandleOneConnection(); + + // Should handle gracefully - not crash + std::string const response = mockSocketPtr->getWrittenData(); + // Response should indicate error + EXPECT_FALSE(response.empty()); +} + +TEST(PortConfigLoaderTest, HandleEmptyRequest) { + auto mockSocket = std::make_unique(); + auto* mockSocketPtr = mockSocket.get(); + + // Empty request + mockSocketPtr->expectConnection(""); + + PortConfigLoader loader(20599, std::move(mockSocket)); + + // Should handle gracefully - connection closed immediately + loader.testHandleOneConnection(); +} diff --git a/libkineto/test/TraceProtocolTest.cpp b/libkineto/test/TraceProtocolTest.cpp new file mode 100644 index 000000000..1c1a1eff4 --- /dev/null +++ b/libkineto/test/TraceProtocolTest.cpp @@ -0,0 +1,229 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include "TraceProtocol.h" + +using namespace KINETO_NAMESPACE; + +// ============================================================================ +// extractJsonInt Tests +// ============================================================================ + +TEST(TraceProtocolTest, ExtractJsonIntBasic) { + std::string const json = R"({"duration_ms": 500})"; + + int const val = extractJsonInt(json, "duration_ms", 0); + EXPECT_EQ(val, 500); +} + +TEST(TraceProtocolTest, ExtractJsonIntMultipleKeys) { + std::string const json = + R"({"duration_ms": 500, "warmup_ms": 100, "count": 42})"; + + int const duration = extractJsonInt(json, "duration_ms", 0); + EXPECT_EQ(duration, 500); + + int const warmup = extractJsonInt(json, "warmup_ms", 0); + EXPECT_EQ(warmup, 100); + + int const count = extractJsonInt(json, "count", 0); + EXPECT_EQ(count, 42); +} + +TEST(TraceProtocolTest, ExtractJsonIntMissingKey) { + std::string const json = R"({"duration_ms": 500})"; + + int const val = extractJsonInt(json, "missing_key", -1); + EXPECT_EQ(val, -1); +} + +TEST(TraceProtocolTest, ExtractJsonIntNegativeValue) { + std::string const json = R"({"offset": -100})"; + + int const val = extractJsonInt(json, "offset", 0); + EXPECT_EQ(val, -100); +} + +TEST(TraceProtocolTest, ExtractJsonIntZero) { + std::string const json = R"({"value": 0})"; + + int const val = extractJsonInt(json, "value", 999); + EXPECT_EQ(val, 0); +} + +TEST(TraceProtocolTest, ExtractJsonIntWhitespace) { + std::string const json = R"({"duration_ms" : 1000})"; + + int const val = extractJsonInt(json, "duration_ms", 0); + EXPECT_EQ(val, 1000); +} + +// ============================================================================ +// extractJsonBool Tests +// ============================================================================ + +TEST(TraceProtocolTest, ExtractJsonBoolTrue) { + std::string const json = R"({"enabled": true})"; + + bool const val = extractJsonBool(json, "enabled", false); + EXPECT_TRUE(val); +} + +TEST(TraceProtocolTest, ExtractJsonBoolFalse) { + std::string const json = R"({"disabled": false})"; + + bool const val = extractJsonBool(json, "disabled", true); + EXPECT_FALSE(val); +} + +TEST(TraceProtocolTest, ExtractJsonBoolMultipleBools) { + std::string const json = + R"({"record_shapes": true, "profile_memory": false, "with_stack": true})"; + + bool const recordShapes = extractJsonBool(json, "record_shapes", false); + EXPECT_TRUE(recordShapes); + + bool const profileMemory = extractJsonBool(json, "profile_memory", true); + EXPECT_FALSE(profileMemory); + + bool const withStack = extractJsonBool(json, "with_stack", false); + EXPECT_TRUE(withStack); +} + +TEST(TraceProtocolTest, ExtractJsonBoolMissingKey) { + std::string const json = R"({"enabled": true})"; + + bool val = extractJsonBool(json, "missing_key", true); + EXPECT_TRUE(val); // Returns default + + val = extractJsonBool(json, "missing_key", false); + EXPECT_FALSE(val); // Returns default +} + +TEST(TraceProtocolTest, ExtractJsonBoolWhitespace) { + std::string const json = R"({"enabled" : true})"; + + bool const val = extractJsonBool(json, "enabled", false); + EXPECT_TRUE(val); +} + +// ============================================================================ +// extractJsonString Tests +// ============================================================================ + +TEST(TraceProtocolTest, ExtractJsonStringBasic) { + std::string const json = R"({"trace_id": "abc-123"})"; + + std::string const val = extractJsonString(json, "trace_id", ""); + EXPECT_EQ(val, "abc-123"); +} + +TEST(TraceProtocolTest, ExtractJsonStringPath) { + std::string const json = R"({"output_dir": "/tmp/kineto_traces"})"; + + std::string const val = extractJsonString(json, "output_dir", ""); + EXPECT_EQ(val, "/tmp/kineto_traces"); +} + +TEST(TraceProtocolTest, ExtractJsonStringEmptyString) { + std::string const json = R"({"trace_id": ""})"; + + std::string const val = extractJsonString(json, "trace_id", "default"); + EXPECT_EQ(val, ""); +} + +TEST(TraceProtocolTest, ExtractJsonStringMissingKey) { + std::string const json = R"({"trace_id": "abc"})"; + + std::string const val = extractJsonString(json, "missing_key", "default"); + EXPECT_EQ(val, "default"); +} + +TEST(TraceProtocolTest, ExtractJsonStringWhitespace) { + std::string const json = R"({"trace_id" : "xyz-789"})"; + + std::string const val = extractJsonString(json, "trace_id", ""); + EXPECT_EQ(val, "xyz-789"); +} + +TEST(TraceProtocolTest, ExtractJsonStringUUID) { + std::string const json = + R"({"trace_id": "a1b2c3d4-e5f6-7890-abcd-ef1234567890"})"; + + std::string const val = extractJsonString(json, "trace_id", ""); + EXPECT_EQ(val, "a1b2c3d4-e5f6-7890-abcd-ef1234567890"); +} + +// ============================================================================ +// buildConfigString Tests +// ============================================================================ + +TEST(TraceProtocolTest, BuildConfigStringFullConfig) { + std::string const config = buildConfigString( + 1000, // durationMs + "CUDA,CPU", // activities + true, // recordShapes + false, // profileMemory + true, // withStack + false, // withFlops + true // withModules + ); + + // Verify key config values are present + EXPECT_NE(config.find("ACTIVITIES_DURATION_MSECS=1000"), std::string::npos); + EXPECT_NE(config.find("ACTIVITIES=CUDA,CPU"), std::string::npos); + EXPECT_NE(config.find("PROFILE_REPORT_INPUT_SHAPES=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_WITH_STACK=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_WITH_MODULES=true"), std::string::npos); + + // These should NOT be present since they are false + EXPECT_EQ(config.find("PROFILE_PROFILE_MEMORY=true"), std::string::npos); + EXPECT_EQ(config.find("PROFILE_WITH_FLOPS=true"), std::string::npos); +} + +TEST(TraceProtocolTest, BuildConfigStringMinimalConfig) { + std::string const config = buildConfigString( + 500, // durationMs + "", // activities (empty) + false, // recordShapes + false, // profileMemory + false, // withStack + false, // withFlops + false // withModules + ); + + // Duration should be present + EXPECT_NE(config.find("ACTIVITIES_DURATION_MSECS=500"), std::string::npos); + + // Empty activities should not add ACTIVITIES line + EXPECT_EQ(config.find("ACTIVITIES="), std::string::npos); + + // No boolean flags should be set + EXPECT_EQ(config.find("PROFILE_REPORT_INPUT_SHAPES"), std::string::npos); +} + +TEST(TraceProtocolTest, BuildConfigStringAllBooleansTrue) { + std::string const config = buildConfigString( + 100, // durationMs + "CUDA", // activities + true, // recordShapes + true, // profileMemory + true, // withStack + true, // withFlops + true // withModules + ); + + EXPECT_NE(config.find("PROFILE_REPORT_INPUT_SHAPES=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_PROFILE_MEMORY=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_WITH_STACK=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_WITH_FLOPS=true"), std::string::npos); + EXPECT_NE(config.find("PROFILE_WITH_MODULES=true"), std::string::npos); +}