diff --git a/.gitignore b/.gitignore index ddc5ec08..b48626fd 100644 --- a/.gitignore +++ b/.gitignore @@ -561,6 +561,9 @@ compile_commands.json # PRD drafts (local planning) PRD_*.md + +# setuptools-scm generated version file +morphling/_version.py *.bak *.swp diff --git a/Dockerfile b/Dockerfile index 2626f945..b069796b 100644 --- a/Dockerfile +++ b/Dockerfile @@ -77,6 +77,13 @@ ARG USE_CCACHE=1 ENV CCACHE_DIR=/ccache RUN if [ "$USE_CCACHE" = "1" ]; then ccache -M 5G; fi +# --no-build-isolation (below) does not install build-system.requires, and the +# build context excludes .git; install setuptools-scm and inject the git-derived +# version via a build arg (empty => pyproject [tool.setuptools_scm] fallback). +RUN pip install --no-cache "setuptools-scm>=8" +ARG MORPHLING_VERSION +ENV SETUPTOOLS_SCM_PRETEND_VERSION_FOR_MORPHLING=${MORPHLING_VERSION} + # 构建和安装项目（使用系统 python）with BuildKit cache mount for ccache RUN --mount=type=cache,target=/ccache \ if [ "$USE_CCACHE" = "1" ]; then export PATH="/usr/lib/ccache:$PATH"; fi && \ diff --git a/csrc/backend/operation_id.h b/csrc/backend/operation_id.h new file mode 100644 index 00000000..be4d4f2b --- /dev/null +++ b/csrc/backend/operation_id.h @@ -0,0 +1,32 @@ +#pragma once + +#include +#include +#include + +namespace morphling { +namespace backend { + +inline constexpr int kMaxLifetimeOperationCount = 65'536; + +inline void ValidateOperationId(int oid) { + if (oid < 0 || oid >= kMaxLifetimeOperationCount) { + throw std::out_of_range("matmul operation id " + std::to_string(oid) + + " is outside the valid range [0, " + + std::to_string(kMaxLifetimeOperationCount) + ")"); + } +} + +inline int ReserveOperationId(std::atomic_int& next_oid) noexcept { + int oid = next_oid.load(std::memory_order_relaxed); + while (oid >= 0 && oid < kMaxLifetimeOperationCount) { + if (next_oid.compare_exchange_weak(oid, oid + 1, + std::memory_order_relaxed)) { + return oid; + } + } + return -1; +} + +} // namespace backend +} // namespace morphling diff --git a/csrc/backend/partition_tracker.cpp b/csrc/backend/partition_tracker.cpp index 6c2d721a..b5534074 100644 --- a/csrc/backend/partition_tracker.cpp +++ b/csrc/backend/partition_tracker.cpp @@ -106,6 +106,47 @@ void PartitionTracker::RemovePartitionByKey(const std::string& partition_key) { } } +bool PartitionTracker::ReassignPartitionToDevice( + const std::string& partition_key, int64_t target_device_id) { + std::lock_guard lock(mutex_); + + auto partition_it = partition_map_.find(partition_key); + if (partition_it == partition_map_.end()) { + return false; + } + const auto& partition_info = partition_it->second; + + auto owner_it = partition_to_device_.find(partition_key); + if (owner_it != partition_to_device_.end()) { + auto device_it = device_partitions_.find(owner_it->second); + if (device_it != device_partitions_.end()) { + auto& partitions = device_it->second; + partitions.erase( + std::remove_if(partitions.begin(), partitions.end(), + [&partition_key](const PartitionInfoPtr& candidate) { + return candidate->key == partition_key; + }), + partitions.end()); + if (partitions.empty()) { + device_partitions_.erase(device_it); + } + } + } + + auto& target_partitions = device_partitions_[target_device_id]; + target_partitions.erase( + std::remove_if(target_partitions.begin(), target_partitions.end(), + [&partition_key](const PartitionInfoPtr& candidate) { + return candidate->key == partition_key; + }), + target_partitions.end()); + target_partitions.push_back(partition_info); + partition_to_device_[partition_key] = target_device_id; + partition_info->owner_device_id = target_device_id; + partition_info->partition->dev_id = target_device_id; + return true; +} + void PartitionTracker::MarkDevicePartitionsFailed(int64_t device_id) { std::lock_guard lock(mutex_); @@ -250,6 +291,19 @@ std::vector PartitionTracker::GetIdlePartitions() const { return idle_partitions; } +std::vector PartitionTracker::ClaimIdlePartitions() { + std::lock_guard lock(mutex_); + + std::vector claimed_partitions; + for (const auto& part_info : partitions_set_) { + if (part_info->state == PartitionState::IDLE) { + part_info->state = PartitionState::RUNNING; + claimed_partitions.push_back(part_info); + } + } + return claimed_partitions; +} + PartitionTracker::DeviceOidStats PartitionTracker::GetDeviceOidStats( int64_t device_id, int64_t oid) const { std::lock_guard lock(mutex_); diff --git a/csrc/backend/partition_tracker.h b/csrc/backend/partition_tracker.h index 59b6cd17..837a9b41 100644 --- a/csrc/backend/partition_tracker.h +++ b/csrc/backend/partition_tracker.h @@ -68,6 +68,8 @@ class PartitionTracker { int64_t oid, MatrixPartitionPtr partition); void RemovePartition(int64_t device_id, const std::string& partition_key); void RemovePartitionByKey(const std::string& partition_key); + bool ReassignPartitionToDevice(const std::string& partition_key, + int64_t target_device_id); // Mark all partitions owned by a device as failed (ownership removed) void MarkDevicePartitionsFailed(int64_t device_id); @@ -85,6 +87,7 @@ class PartitionTracker { size_t GetDevicePartitionCount(int64_t device_id) const; bool HasPendingPartitions(int64_t device_id) const; std::vector GetIdlePartitions() const; + std::vector ClaimIdlePartitions(); // Statistics for a specific OID on a device (for debugging) struct DeviceOidStats { diff --git a/csrc/backend/proxy_svr.cpp b/csrc/backend/proxy_svr.cpp index 87b8f64b..c270e5af 100644 --- a/csrc/backend/proxy_svr.cpp +++ b/csrc/backend/proxy_svr.cpp @@ -423,13 +423,9 @@ void ProxySvrHandle::HandleMatMul(const ConnectionUeventPtr& conn, start = std::chrono::high_resolution_clock::now(); auto output = torch::from_blob(o_ptr, {row_size, col_size}, FLOAT32_TENSOR_OPTIONS(torch::kCPU)); - { - // std::lock_guard lock(outputs_mutex_[partition.oid]); - auto& output_matrix = reinterpret_cast(ctx_.instance) - ->GetOutputMatrix(partition.oid); - IndexPutMatrixBlock(output_matrix, output, partition.row, partition.col, - partition.pivot, ctx_.block_size); - } + reinterpret_cast(ctx_.instance) + ->WriteResultBlock(partition.oid, output, partition.row, partition.col, + partition.pivot, ctx_.block_size); end = std::chrono::high_resolution_clock::now(); LOG_DEBUG << "UpdateMatrixBlock time: " << std::chrono::duration_cast(end - @@ -663,7 +659,9 @@ void ProxySvrHandle::ConnectionClosedCb(const ConnectionUeventPtr& conn) { /********************************ProxySvrImpl****************************************/ ProxySvrImpl::ProxySvrImpl(ProxyEnvCfg& ctx) - : ctx_(ctx), listener_(nullptr), rsp_cb_counts_(65536) { + : ctx_(ctx), + listener_(nullptr), + rsp_cb_counts_(kMaxLifetimeOperationCount) { // Initialize with greedy scheduling policy by default } @@ -714,7 +712,7 @@ void ProxySvrImpl::Initialize(UeventLoop* loop) { // Start(); // InitLogger(); - outputs_.resize(65536); + outputs_.resize(kMaxLifetimeOperationCount); for (size_t i = 0; i < outputs_.size(); i++) { outputs_[i] = torch::empty({0, 0}); rsp_cb_counts_[i] = 0; @@ -819,14 +817,14 @@ void ProxySvrImpl::RequestCb(const ConnectionUeventPtr& conn) { // handle->RequestCb(conn); // } -void ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, - torch::Tensor& mat_b) { +int ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, + torch::Tensor& mat_b) { auto* gate = DEVICE_TRACKER.GetDispatchGate(); if (gate != nullptr) { if (gate->GetMode() == DeviceMode::BARRIER) { if (!gate->WaitForReady()) { LOG_ERROR << "[DispatchMatMulAsync] DispatchGate WaitForReady timeout"; - return; + return -1; } } else if (gate->GetMode() == DeviceMode::DYNAMIC && DEVICE_TRACKER.GetConnectedDeviceCount() == 0) { @@ -839,21 +837,26 @@ void ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, }); LOG_INFO << "[DispatchMatMulAsync] No connected devices in DYNAMIC mode, " "work enqueued"; - return; + return -1; } } - LOG_INFO << "[DispatchMatMulAsync] Starting dispatch - mm_count=" - << mm_count_; + int oid = ReserveOperationId(mm_count_); + if (oid < 0) { + LOG_ERROR << "[DispatchMatMulAsync] Operation capacity exhausted at " + << kMaxLifetimeOperationCount << " lifetime dispatches"; + return -1; + } + LOG_INFO << "[DispatchMatMulAsync] Starting dispatch - oid=" << oid; - outputs_[mm_count_].set_data(CreateOutputMatrix(mat_a, mat_b)); + int gemm_id = gemm_id_count_.fetch_add(1); auto partitions = PartitionMatrices(mat_a, mat_b, ctx_.block_size); auto a_shape = mat_a.sizes().vec(); auto b_shape = mat_b.sizes().vec(); if (partitions.empty()) { LOG_ERROR << "[DispatchMatMulAsync] No partitions generated!"; - return; + return -1; } auto cur_ver = partitions[0]->version; @@ -874,7 +877,11 @@ void ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, // << partitions.size() << " partitions"; auto start = std::chrono::high_resolution_clock::now(); - DecRspCbCount(mm_count_, partitions.size()); + { + std::lock_guard lk(outputs_mutex_); + outputs_[oid].set_data(CreateOutputMatrix(mat_a, mat_b)); + DecRspCbCount(oid, partitions.size()); + } LOG_INFO << "[DispatchMatMulAsync] Creating " << partitions.size() << " partitions as IDLE"; @@ -882,21 +889,21 @@ void ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, // Add all partitions to tracker as IDLE - they will be dispatched by // SendIdlePartitions for (auto& partition : partitions) { - partition->oid = mm_count_; - partition->gemm_id = gemm_id_count_; // assign global gemm_id + partition->oid = oid; + partition->gemm_id = gemm_id; // assign global gemm_id partition->dev_id = -1; // Mark as unassigned, will be assigned by scheduling policy // Add partition to tracker with dev_id=-1 (unassigned, to be scheduled) // The tracker will use owner_device_id=-1 until scheduling assigns a real // device - PARTITION_TRACKER.AddPartition(-1, partition->GetPartitionKey(), mm_count_, + PARTITION_TRACKER.AddPartition(-1, partition->GetPartitionKey(), oid, partition); LOG_DEBUG << "[DispatchMatMulAsync] Created IDLE partition key=" << partition->GetPartitionKey() << ", dev_id=" << partition->dev_id << " (unassigned)" - << ", oid=" << mm_count_ << ", gemm_id=" << partition->gemm_id; + << ", oid=" << oid << ", gemm_id=" << partition->gemm_id; } auto end = std::chrono::high_resolution_clock::now(); LOG_INFO << "[DispatchMatMulAsync] Created " << partitions.size() @@ -904,141 +911,35 @@ void ProxySvrImpl::DispatchMatMulAsync(torch::Tensor& mat_a, << std::chrono::duration_cast(end - start) .count() << "us. Partitions will be sent by SendIdlePartitions timer. " - "gemm_id_count=" - << gemm_id_count_; - mm_count_++; - gemm_id_count_++; // increment global gemm_id for next operation + "gemm_id=" + << gemm_id; auto* handle = reinterpret_cast(loop_->GetLoopHandle()); loop_->QueueInLoop(bind(&ProxySvrHandle::SendIdlePartitions, handle)); + return oid; } torch::Tensor ProxySvrImpl::WaitMatMul(int oid) { - auto start = std::chrono::high_resolution_clock::now(); - LOG_INFO << "[WaitMatMul] Starting wait for oid=" << oid - << ", rsp_cb_counts_[oid]=" << rsp_cb_counts_[oid]; - - int poll_count = 0; - while (rsp_cb_counts_[oid] > 0) { - poll_count++; - if (poll_count % 50 == 0) { // Log every 5 seconds (50 * 100ms) - LOG_WARN << "[WaitMatMul] Still waiting for oid=" << oid - << ", rsp_cb_counts_[oid]=" << rsp_cb_counts_[oid] - << ", poll_count=" << poll_count * 100 << "ms"; - - // Diagnose partition states across all devices - auto connected_devices = DEVICE_TRACKER.GetConnectedDevices(); - LOG_WARN << "[WaitMatMul] Connected devices: " - << connected_devices.size(); - - size_t total_idle = 0, total_running = 0, total_finished = 0; - size_t devices_with_partitions = 0; - std::vector devices_with_oid; - - for (int64_t device_id : connected_devices) { - auto stats = PARTITION_TRACKER.GetDeviceOidStats(device_id, oid); - size_t device_total = - stats.idle_count + stats.running_count + stats.finished_count; - - if (device_total > 0) { - devices_with_partitions++; - devices_with_oid.push_back(device_id); - bool is_connected = DEVICE_TRACKER.IsDeviceConnected(device_id); - total_idle += stats.idle_count; - total_running += stats.running_count; - total_finished += stats.finished_count; - - LOG_WARN << "[WaitMatMul] Device " << device_id - << " (connected=" << (is_connected ? "YES" : "NO") << ")" - << ": IDLE=" << stats.idle_count - << ", RUNNING=" << stats.running_count - << ", FINISHED=" << stats.finished_count - << ", Total=" << device_total; - - // Show first few partition keys for debugging (only for first 3 - // devices) - if (devices_with_partitions <= 3) { - if (!stats.partition_keys.empty() && - stats.partition_keys.size() <= 5) { - std::string keys_str; - for (const auto& key : stats.partition_keys) { - if (!keys_str.empty()) keys_str += ", "; - keys_str += key; - } - LOG_WARN << "[WaitMatMul] Partition keys: " << keys_str; - } else if (stats.partition_keys.size() > 5) { - LOG_WARN << "[WaitMatMul] First partition key: " - << stats.partition_keys[0] << " (+ " - << (stats.partition_keys.size() - 1) << " more)"; - } - } - } - } - - // Summary with device distribution info - LOG_WARN << "[WaitMatMul] Summary for oid=" << oid - << ": Devices with partitions=" << devices_with_partitions << "/" - << connected_devices.size() << ", Total IDLE=" << total_idle - << ", RUNNING=" << total_running - << ", FINISHED=" << total_finished - << ", Expected remaining=" << rsp_cb_counts_[oid]; - - // Critical: if only 1 device has all partitions, this is a scheduling - // problem! - if (devices_with_partitions == 1 && connected_devices.size() > 1) { - LOG_ERROR << "[WaitMatMul] ⚠️ SCHEDULING ISSUE: All " - << (total_idle + total_running + total_finished) - << " partitions assigned to single device " - << devices_with_oid[0] << " while " - << (connected_devices.size() - 1) - << " other devices are idle!"; - } else if (devices_with_partitions > 0 && devices_with_partitions <= 10) { - std::string device_list; - for (auto dev_id : devices_with_oid) { - if (!device_list.empty()) device_list += ", "; - device_list += std::to_string(dev_id); - } - LOG_WARN << "[WaitMatMul] Devices with oid=" << oid << ": [" - << device_list << "]"; - } - - // Check if partitions are stuck in RUNNING state - if (total_running > 0 && total_running == rsp_cb_counts_[oid] && - total_idle == 0) { - LOG_ERROR << "[WaitMatMul] ⚠️ STUCK PARTITIONS: All " << total_running - << " partitions stuck in RUNNING state for " - << poll_count * 100 << "ms"; - LOG_ERROR << "[WaitMatMul] Possible causes: 1) Devices not responding " - "2) Network issues 3) Devices processing too slowly"; - - // Sample a few devices to check connection quality - size_t check_count = std::min(size_t(5), devices_with_oid.size()); - for (size_t i = 0; i < check_count; ++i) { - int64_t dev_id = devices_with_oid[i]; - auto conn = DEVICE_TRACKER.GetDeviceConnection(dev_id); - bool has_conn = (conn != nullptr); - bool conn_closed = has_conn ? conn->IsClosed() : true; - LOG_ERROR << "[WaitMatMul] Sample device " << dev_id - << ": has_connection=" << (has_conn ? "YES" : "NO") - << ", connection_closed=" << (conn_closed ? "YES" : "NO"); - } - } - } - std::this_thread::sleep_for(std::chrono::milliseconds(100)); + ValidateOperationId(oid); + std::unique_lock lk(outputs_mutex_); + while (!outputs_cv_.wait_for(lk, std::chrono::seconds(5), + [&] { return rsp_cb_counts_[oid] == 0; })) { + LOG_WARN << "[WaitMatMul] Still waiting for oid=" << oid + << ", remaining=" << rsp_cb_counts_[oid]; } - auto end = std::chrono::high_resolution_clock::now(); - auto shape = outputs_[oid].sizes().vec(); - auto wait_time = - std::chrono::duration_cast(end - start) - .count(); - LOG_INFO << "[WaitMatMul] Completed for oid=" << oid - << ", wait_time=" << wait_time << "us, shape=" << shape; - - return outputs_[oid]; + return outputs_[oid].clone(); +} + +void ProxySvrImpl::WriteResultBlock(int oid, torch::Tensor& block, int64_t row, + int64_t col, int64_t pivot, + int block_size) { + std::lock_guard lock(outputs_mutex_); + IndexPutMatrixBlock(outputs_[oid], block, row, col, pivot, block_size); } void ProxySvrImpl::IncRspCbCount(int oid, size_t count) { - unsigned long long prev = rsp_cb_counts_[oid].load(); + std::lock_guard lock(outputs_mutex_); + uint64_t prev = rsp_cb_counts_[oid]; if (prev < count) { LOG_WARN << "[IncRspCbCount] Clamping underflow for oid=" << oid << ", current=" << prev << ", decrement=" << count; @@ -1046,6 +947,9 @@ void ProxySvrImpl::IncRspCbCount(int oid, size_t count) { } else { rsp_cb_counts_[oid] -= count; } + if (rsp_cb_counts_[oid] == 0) { + outputs_cv_.notify_all(); + } LOG_DEBUG << "[IncRspCbCount] oid=" << oid << ", count=" << count << ", prev=" << prev << ", now=" << rsp_cb_counts_[oid]; } @@ -1146,7 +1050,7 @@ void ProxySvrImpl::HandleDeviceFailure(int64_t failed_device_id) { } void ProxySvrHandle::SendIdlePartitions() { - auto idle_partitions = PARTITION_TRACKER.GetIdlePartitions(); + auto idle_partitions = PARTITION_TRACKER.ClaimIdlePartitions(); if (idle_partitions.empty()) { LOG_DEBUG << "[SendIdlePartitions] No IDLE partitions to send"; @@ -1161,42 +1065,35 @@ void ProxySvrHandle::SendIdlePartitions() { if (redistributed.empty()) { LOG_DEBUG << "[SendIdlePartitions] No available devices for redistribution"; + for (const auto& part_info : idle_partitions) { + PARTITION_TRACKER.MarkPartitionIdle(part_info->key); + } return; } LOG_INFO << "[SendIdlePartitions] Scheduling complete, moving partitions to " "assigned devices"; - // IMPORTANT: After scheduling, partitions' owner_device_id has been updated - // We need to move them from device -1 to their assigned devices in tracker - // Do this BEFORE marking as RUNNING to avoid race conditions + std::vector ready_partitions; + ready_partitions.reserve(idle_partitions.size()); for (const auto& part_info : idle_partitions) { - int64_t old_device = - part_info - ->owner_device_id; // This might be wrong due to scheduling update - // Find old device from partition_to_device_ map - - LOG_DEBUG << "[SendIdlePartitions] Moving partition " << part_info->key - << " (oid=" << part_info->oid << ") to device " - << part_info->owner_device_id; - - // Remove from old location and add to new location - PARTITION_TRACKER.RemovePartitionByKey(part_info->key); - PARTITION_TRACKER.AddPartition(part_info->owner_device_id, part_info->key, - part_info->oid, part_info->partition); - - // Update partition's dev_id to match the assigned device - part_info->partition->dev_id = part_info->owner_device_id; - - // Mark partition as RUNNING after it's in the correct device list - PARTITION_TRACKER.MarkPartitionRunning(part_info->key); + auto assignment = redistributed.find(part_info->key); + if (assignment == redistributed.end() || + !PARTITION_TRACKER.ReassignPartitionToDevice(part_info->key, + assignment->second)) { + LOG_ERROR << "[SendIdlePartitions] Failed to reassign partition " + << part_info->key; + PARTITION_TRACKER.MarkPartitionIdle(part_info->key); + continue; + } + ready_partitions.push_back(part_info); } - LOG_INFO << "[SendIdlePartitions] Sending " << idle_partitions.size() + LOG_INFO << "[SendIdlePartitions] Sending " << ready_partitions.size() << " partitions to devices"; // Send each partition - for (const auto& part_info : idle_partitions) { + for (const auto& part_info : ready_partitions) { LOG_DEBUG << "[SendIdlePartitions] Sending partition " << part_info->key << " to device " << part_info->owner_device_id; @@ -1216,7 +1113,7 @@ void ProxySvrHandle::SendIdlePartitions() { } LOG_INFO << "[SendIdlePartitions] Completed sending " - << idle_partitions.size() << " partitions to devices"; + << ready_partitions.size() << " partitions to devices"; } void ProxySvrImpl::CheckFailedPartitions() { diff --git a/csrc/backend/proxy_svr.h b/csrc/backend/proxy_svr.h index 842b455c..566d4d59 100644 --- a/csrc/backend/proxy_svr.h +++ b/csrc/backend/proxy_svr.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -16,6 +17,7 @@ #include "morphling.pb.h" #include "network/uevent.h" #include "network/ueventloop_thread.h" +#include "operation_id.h" #include "sched_policy.h" #include "server_base.h" @@ -121,9 +123,11 @@ class ProxySvrImpl : public std::enable_shared_from_this { ~ProxySvrImpl(); void Initialize(uevent::UeventLoop* loop); - void DispatchMatMulAsync(torch::Tensor& mat_a, torch::Tensor& mat_b); + int DispatchMatMulAsync(torch::Tensor& mat_a, torch::Tensor& mat_b); torch::Tensor WaitMatMul(int oid); torch::Tensor& GetOutputMatrix(int oid) { return outputs_[oid]; } + void WriteResultBlock(int oid, torch::Tensor& block, int64_t row, int64_t col, + int64_t pivot, int block_size); void IncRspCbCount(int oid, size_t count); void DecRspCbCount(int oid, size_t count) { rsp_cb_counts_[oid] += count; } @@ -175,8 +179,10 @@ class ProxySvrImpl : public std::enable_shared_from_this { std::atomic_int mm_count_{0}; std::atomic_int gemm_id_count_{0}; // global gemm operation id counter std::vector outputs_; - std::vector rsp_cb_counts_; + std::vector rsp_cb_counts_; std::vector> device_tensors_; + std::mutex outputs_mutex_; + std::condition_variable outputs_cv_; // Note: Partition tracking is now handled by DevicePartitionTracker singleton // Access via DEVICE_TRACKER macro @@ -224,8 +230,8 @@ class ProxySvr { void SetCacheEnabled(bool enabled) { context_.enable_cli_cache = enabled ? 1 : 0; } - void DispatchMatMulAsync(torch::Tensor& mat_a, torch::Tensor& mat_b) { - svr_->DispatchMatMulAsync(mat_a, mat_b); + int DispatchMatMulAsync(torch::Tensor& mat_a, torch::Tensor& mat_b) { + return svr_->DispatchMatMulAsync(mat_a, mat_b); } torch::Tensor WaitMatMul(int oid) { return svr_->WaitMatMul(oid); } size_t GetConnectionCount() const { return svr_->GetConnectionCount(); } diff --git a/docs/GEMM_ID_ISSUES.md b/docs/GEMM_ID_ISSUES.md index 631c5407..ca190e77 100644 --- a/docs/GEMM_ID_ISSUES.md +++ b/docs/GEMM_ID_ISSUES.md @@ -25,10 +25,44 @@ describing each record format. time. In [`csrc/backend/proxy_svr.cpp`](../csrc/backend/proxy_svr.cpp) it is sourced from the atomic `gemm_id_count_` (`proxy_svr.h`), stamped onto every `MatrixPartition` in `DispatchMatMulAsync` (`partition->gemm_id = -gemm_id_count_`), and incremented after each dispatch. It appears in both +gemm_id`), and reserved with `fetch_add` for each dispatch. It appears in both the VTIME and Throughput records so events can be grouped by the GEMM that produced them; it starts at 0 and increases by one per `DispatchMatMulAsync`. +## Live-dispatch synchronization invariants + +Live training can have several GEMMs and dispatch callbacks in flight. The +proxy runtime preserves operation and partition ownership through three +invariants. + +1. `DispatchMatMulAsync` returns a distinct operation ID. Callers pass that ID + to `WaitMatMul`, so a later GEMM cannot consume an earlier output. +2. Output block writes, response-count decrements, and completion waits share + one mutex and condition variable. `WaitMatMul` returns a cloned completed + tensor rather than storage that a later dispatch can reuse. +3. `PartitionTracker::ClaimIdlePartitions` changes `IDLE` partitions to + `RUNNING` while holding the tracker lock. Assignment then uses + `ReassignPartitionToDevice`, which moves the same `PartitionInfo` between + device indexes atomically without exposing an intermediate `IDLE` state. + +The third invariant prevents concurrent `SendIdlePartitions` callbacks from +sending the same shard twice and prematurely completing an operation. The +claim and reassignment lifecycle is covered by +[`tests/cpp/unit/backend/test_partition_tracker_claim.cpp`](../tests/cpp/unit/backend/test_partition_tracker_claim.cpp). + +Operation IDs index fixed storage and are therefore bounded by +`kMaxLifetimeOperationCount`. The current lifetime limit is 65,536 dispatches +per proxy instance. The final valid ID is reserved normally, but subsequent +dispatches fail with `-1` before any output storage is accessed. Storage is not +resized because callbacks retain operation IDs for the lifetime of an in-flight +GEMM. + +`WaitMatMul` accepts only IDs in the valid storage range. Negative IDs and IDs +at or above the lifetime limit raise `std::out_of_range`, which pybind exposes +as a Python exception. The autograd hook checks dispatch results before calling +`wait_matmul` and raises `RuntimeError` with the failed GEMM phase when dispatch +returns the `-1` sentinel. + ## VTIME format ``` diff --git a/morphling/__init__.py b/morphling/__init__.py index 01c859dc..a2d66125 100644 --- a/morphling/__init__.py +++ b/morphling/__init__.py @@ -4,6 +4,16 @@ inference device emulation. """ +try: + from morphling._version import __version__ +except Exception: + try: + from importlib.metadata import version as _pkg_version + + __version__ = _pkg_version("morphling") + except Exception: + __version__ = "0.0.0+unknown" + import morphling import morphling.hooks as hooks @@ -16,4 +26,4 @@ def set_backend(backend): autograd.set_backend(backend) -__all__ = ["set_backend"] +__all__ = ["set_backend", "__version__"] diff --git a/morphling/hooks/autograd.py b/morphling/hooks/autograd.py index 71a566bc..388fc34b 100644 --- a/morphling/hooks/autograd.py +++ b/morphling/hooks/autograd.py @@ -16,6 +16,20 @@ _backend: Any = None +def _reference_grad_input( + grad_output: torch.Tensor, weight: torch.Tensor +) -> torch.Tensor: + return torch.matmul(grad_output, weight) + + +def _wait_for_dispatched_matmul(oid: int, phase: str) -> torch.Tensor: + if oid < 0: + raise RuntimeError( + f"{phase} matmul dispatch failed with operation id {oid}" + ) + return _backend.wait_matmul(oid) + + def set_backend(backend: Any) -> None: """Set the active compute backend for morphling hooks. @@ -181,7 +195,6 @@ class LinearFunction(torch.autograd.Function): def forward(ctx, input, weight, bias=None): global _gemm_idx - print("LinearFunction forward", input.shape, weight.shape) ctx.save_for_backward(input, weight, bias) # output = input.mm(weight.t()) # logger.debug(f"input shape: {input.shape}") @@ -191,13 +204,15 @@ def forward(ctx, input, weight, bias=None): start_us: Optional[float] = None sm_count: Optional[int] = None gemm_idx: Optional[int] = None + _fwd_oid = -1 + _fwd_weight = weight.transpose(-2, -1).contiguous() if _greenctx is not None: start_us = _elapsed_us() activated_sm_count = None try: sm_count, _ = _greenctx.activate_for_time(int(start_us)) activated_sm_count = sm_count - _backend.async_dispatch_matmul(input, weight.transpose(-2, -1)) + _fwd_oid = _backend.async_dispatch_matmul(input, _fwd_weight) finally: if activated_sm_count is not None: _greenctx.deactivate(activated_sm_count) @@ -205,8 +220,8 @@ def forward(ctx, input, weight, bias=None): gemm_idx = _gemm_idx _gemm_idx += 1 else: - _backend.async_dispatch_matmul(input, weight.transpose(-2, -1)) - output = _backend.wait_matmul(0) + _fwd_oid = _backend.async_dispatch_matmul(input, _fwd_weight) + output = _wait_for_dispatched_matmul(_fwd_oid, "forward") if _gemm_log_enabled and start_us is not None and gemm_idx is not None: m = int(input.shape[0]) k = int(input.shape[1]) @@ -243,13 +258,13 @@ def backward(ctx, grad_output): global _gemm_idx input, weight, bias = ctx.saved_tensors - print( - "LinearFunction backward", - grad_output.shape, - weight.shape, - input.shape, - ) grad_input = grad_weight = grad_bias = None + _gi_oid = -1 + _gw_oid = -1 + _dg = grad_output.contiguous() + _dw = weight.contiguous() + _dg_t = grad_output.transpose(-2, -1).contiguous() + _din_t = input.transpose(-2, -1).contiguous() grad_input_start_us: Optional[float] = None grad_input_sm_count: Optional[int] = None grad_input_gemm_idx: Optional[int] = None @@ -272,7 +287,7 @@ def backward(ctx, grad_output): int(grad_input_start_us) ) activated_sm_count = grad_input_sm_count - _backend.async_dispatch_matmul(grad_output, weight) + _gi_oid = _backend.async_dispatch_matmul(_dg, _dw) finally: if activated_sm_count is not None: _greenctx.deactivate(activated_sm_count) @@ -280,7 +295,7 @@ def backward(ctx, grad_output): grad_input_gemm_idx = _gemm_idx _gemm_idx += 1 else: - _backend.async_dispatch_matmul(grad_output, weight) + _gi_oid = _backend.async_dispatch_matmul(_dg, _dw) if ctx.needs_input_grad[1]: # grad_weight = grad_output.t().mm(input) # grad_weight = torch.as_tensor( @@ -297,10 +312,7 @@ def backward(ctx, grad_output): int(grad_weight_start_us) ) activated_sm_count = grad_weight_sm_count - _backend.async_dispatch_matmul( - grad_output.transpose(-2, -1), - input.transpose(-2, -1), - ) + _gw_oid = _backend.async_dispatch_matmul(_dg_t, _din_t) finally: if activated_sm_count is not None: _greenctx.deactivate(activated_sm_count) @@ -308,15 +320,12 @@ def backward(ctx, grad_output): grad_weight_gemm_idx = _gemm_idx _gemm_idx += 1 else: - _backend.async_dispatch_matmul( - grad_output.transpose(-2, -1), - input.transpose(-2, -1), - ) + _gw_oid = _backend.async_dispatch_matmul(_dg_t, _din_t) - dispatch_count = 0 if ctx.needs_input_grad[0]: - grad_input = _backend.wait_matmul(dispatch_count) - dispatch_count += 1 + grad_input = _wait_for_dispatched_matmul( + _gi_oid, "backward_grad_input" + ) if ( grad_input_start_us is not None and grad_input_gemm_idx is not None @@ -336,8 +345,9 @@ def backward(ctx, grad_output): ) if ctx.needs_input_grad[1]: - grad_weight = _backend.wait_matmul(dispatch_count).transpose(-2, -1) - dispatch_count += 1 + grad_weight = _wait_for_dispatched_matmul( + _gw_oid, "backward_grad_weight" + ).transpose(-2, -1) if ( grad_weight_start_us is not None and grad_weight_gemm_idx is not None @@ -361,7 +371,7 @@ def backward(ctx, grad_output): if _enable_verification: if ctx.needs_input_grad[0]: - ref_grad_input = torch.matmul(grad_output, weight, atol=1e-5) + ref_grad_input = _reference_grad_input(grad_output, weight) assert torch.allclose(grad_input, ref_grad_input), ( f"grad_input is not close! max diff: {torch.max(torch.abs(grad_input - ref_grad_input))}" ) diff --git a/morphling/runtime/coordinator_metrics.py b/morphling/runtime/coordinator_metrics.py new file mode 100644 index 00000000..7ddd15f2 --- /dev/null +++ b/morphling/runtime/coordinator_metrics.py @@ -0,0 +1,252 @@ +"""Low-overhead sampled metrics for coordinator training processes.""" + +from __future__ import annotations + +import json +import sys +import threading +import time +from contextlib import contextmanager +from typing import Iterator, Protocol, TextIO + +import psutil + +from morphling.runtime.coordinator_metrics_cli import ( + MetricsConfig, + add_metrics_arguments, + metrics_config_from_args, + metrics_output_path, +) +from morphling.runtime.coordinator_metrics_models import ( + MetricReading, + MetricsSample, + NicCounters, + PhaseSnapshot, + SelectedNicCounters, + aggregate_nic_counters, + build_sample, +) + + +class PhaseRecorder: + """Accumulate phase events in memory for sampled publication.""" + + def __init__(self) -> None: + self._lock = threading.Lock() + self._counts: dict[str, int] = {} + self._durations: dict[str, float] = {} + self._timestamps: dict[str, float] = {} + + def record( + self, name: str, *, duration_seconds: float, timestamp_unix_s: float + ) -> None: + with self._lock: + self._counts[name] = self._counts.get(name, 0) + 1 + self._durations[name] = ( + self._durations.get(name, 0.0) + duration_seconds + ) + self._timestamps[name] = timestamp_unix_s + + @contextmanager + def track(self, name: str) -> Iterator[None]: + started = time.perf_counter() + try: + yield + finally: + self.record( + name, + duration_seconds=time.perf_counter() - started, + timestamp_unix_s=time.time(), + ) + + def snapshot(self) -> PhaseSnapshot: + with self._lock: + return PhaseSnapshot( + counts=dict(self._counts), + durations_seconds=dict(self._durations), + last_timestamps_unix_s=dict(self._timestamps), + ) + + +@contextmanager +def track_phase(recorder: PhaseRecorder | None, name: str) -> Iterator[None]: + """Track a phase only when metrics collection is enabled.""" + if recorder is None: + yield + else: + with recorder.track(name): + yield + + +class MetricsCollectionError(RuntimeError): + """Represent a terminal collector failure without stopping training.""" + + def __init__(self, cause: Exception) -> None: + self.error_type = type(cause).__name__ + super().__init__(str(cause)) + + +class _MetricsReader(Protocol): + def prime(self) -> None: ... + + def read(self, nic_names: tuple[str, ...] | None) -> MetricReading: ... + + +class _PsutilMetricsReader: + def __init__(self) -> None: + self._process = psutil.Process() + + def prime(self) -> None: + self._process.cpu_percent(interval=None) + + def read(self, nic_names: tuple[str, ...] | None) -> MetricReading: + per_nic = psutil.net_io_counters(pernic=True) + selected = aggregate_nic_counters( + { + name: NicCounters(value.bytes_sent, value.bytes_recv) + for name, value in per_nic.items() + }, + nic_names, + ) + with self._process.oneshot(): + return MetricReading( + timestamp_unix_s=time.time(), + monotonic_s=time.monotonic(), + process_cpu_percent=self._process.cpu_percent(interval=None), + rss_bytes=self._process.memory_info().rss, + nic=selected.counters, + nic_names=selected.nic_names, + ) + + +class CoordinatorMetricsCollector: + """Sample psutil metrics on a background thread and append JSONL.""" + + def __init__( + self, config: MetricsConfig, phase_recorder: PhaseRecorder | None = None + ) -> None: + self._config = config + self._phases = phase_recorder or PhaseRecorder() + self._reader: _MetricsReader = _PsutilMetricsReader() + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + self._output: TextIO | None = None + self._previous: MetricReading | None = None + self._failure: MetricsCollectionError | None = None + + @property + def failure(self) -> MetricsCollectionError | None: + return self._failure + + def start(self) -> None: + self._config.output_path.parent.mkdir(parents=True, exist_ok=True) + self._output = self._config.output_path.open("w", encoding="utf-8") + try: + self._reader.prime() + self._sample_once() + except Exception as error: # noqa: BLE001 - thread/process boundary + failure = self._record_failure(error) + self._close_output() + raise failure from error + self._thread = threading.Thread( + target=self._run, + name="coordinator-metrics", + daemon=True, + ) + self._thread.start() + + def stop(self) -> None: + if self._output is None: + return + self._stop_event.set() + if self._thread is not None: + self._thread.join() + if self._failure is None: + try: + self._sample_once() + except Exception as error: # noqa: BLE001 - shutdown boundary + self._record_failure(error) + self._close_output() + self._thread = None + + def _run(self) -> None: + while not self._stop_event.wait(self._config.interval_seconds): + try: + self._sample_once() + except Exception as error: # noqa: BLE001 - daemon boundary + self._record_failure(error) + self._stop_event.set() + return + + def _sample_once(self) -> None: + reading = self._reader.read(self._config.nic_names) + sample = build_sample(reading, self._previous, self._phases.snapshot()) + self._previous = reading + if self._output is not None: + self._output.write(sample.to_json_line()) + self._output.flush() + + def _record_failure(self, error: Exception) -> MetricsCollectionError: + if self._failure is not None: + return self._failure + failure = MetricsCollectionError(error) + self._failure = failure + terminal = { + "record_type": "terminal_error", + "timestamp_unix_s": time.time(), + "error_type": failure.error_type, + "error_message": str(failure), + } + if self._output is not None: + try: + self._output.write( + json.dumps(terminal, sort_keys=True, separators=(",", ":")) + + "\n" + ) + self._output.flush() + except (OSError, ValueError) as output_error: + print( + f"Coordinator metrics terminal write failed: {output_error}", + file=sys.stderr, + flush=True, + ) + print( + f"Coordinator metrics stopped: {failure}", + file=sys.stderr, + flush=True, + ) + return failure + + def _close_output(self) -> None: + output = self._output + self._output = None + if output is None: + return + try: + output.flush() + except (OSError, ValueError) as output_error: + print( + f"Coordinator metrics flush failed: {output_error}", + file=sys.stderr, + flush=True, + ) + try: + output.close() + except (OSError, ValueError) as output_error: + print( + f"Coordinator metrics close failed: {output_error}", + file=sys.stderr, + flush=True, + ) + + +def start_metrics_collector( + config: MetricsConfig, phase_recorder: PhaseRecorder +) -> CoordinatorMetricsCollector | None: + """Start metrics without allowing setup failure to abort training.""" + collector = CoordinatorMetricsCollector(config, phase_recorder) + try: + collector.start() + except MetricsCollectionError: + return None + return collector diff --git a/morphling/runtime/coordinator_metrics_cli.py b/morphling/runtime/coordinator_metrics_cli.py new file mode 100644 index 00000000..2f7aa56a --- /dev/null +++ b/morphling/runtime/coordinator_metrics_cli.py @@ -0,0 +1,98 @@ +"""CLI and output-path configuration for coordinator metrics.""" + +from __future__ import annotations + +import argparse +from dataclasses import dataclass +from pathlib import Path +from typing import Protocol, Sequence + + +class MetricsConfigurationError(ValueError): + """Report invalid coordinator metrics configuration.""" + + +@dataclass(frozen=True, slots=True) +class MetricsConfig: + output_path: Path + interval_seconds: float = 1.0 + nic_names: tuple[str, ...] | None = None + + def __post_init__(self) -> None: + if self.interval_seconds <= 0: + raise MetricsConfigurationError( + "metrics interval must be greater than zero" + ) + + +class MetricsCliArgs(Protocol): + metrics_output: Path | None + metrics_interval: float + metrics_nics: Sequence[str] | None + + +def _positive_interval(raw: str) -> float: + try: + interval = float(raw) + except ValueError as error: + raise argparse.ArgumentTypeError( + "metrics interval must be a number" + ) from error + if interval <= 0: + raise argparse.ArgumentTypeError( + "metrics interval must be greater than zero" + ) + return interval + + +def add_metrics_arguments(parser: argparse.ArgumentParser) -> None: + """Add shared opt-in coordinator metrics flags.""" + parser.add_argument( + "--metrics_output", + type=Path, + default=None, + help=( + "Metrics destination: a path ending in .jsonl is a file; " + "every other path is treated as a directory." + ), + ) + parser.add_argument( + "--metrics_interval", + type=_positive_interval, + default=1.0, + help="Sampling interval in seconds; must be greater than zero.", + ) + parser.add_argument( + "--metrics_nics", + nargs="+", + default=None, + help="Optional network interface names to aggregate.", + ) + + +def metrics_output_path(destination: Path, rank: int | None = None) -> Path: + """Resolve .jsonl files directly and all other paths as directories.""" + if rank is None: + if destination.suffix == ".jsonl": + return destination + return destination / "coordinator.jsonl" + if destination.suffix == ".jsonl": + return destination.with_name( + f"{destination.stem}-rank-{rank}{destination.suffix}" + ) + return destination / f"coordinator-rank-{rank}.jsonl" + + +def metrics_config_from_args( + args: MetricsCliArgs, rank: int | None = None +) -> MetricsConfig | None: + """Build validated metrics configuration when collection is enabled.""" + if args.metrics_output is None: + return None + return MetricsConfig( + output_path=metrics_output_path(args.metrics_output, rank=rank), + interval_seconds=args.metrics_interval, + nic_names=( + tuple(args.metrics_nics) if args.metrics_nics is not None else None + ), + ) diff --git a/morphling/runtime/coordinator_metrics_models.py b/morphling/runtime/coordinator_metrics_models.py new file mode 100644 index 00000000..88176443 --- /dev/null +++ b/morphling/runtime/coordinator_metrics_models.py @@ -0,0 +1,125 @@ +"""Value types and calculations for sampled coordinator metrics.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Mapping, Sequence + + +@dataclass(frozen=True, slots=True) +class NicCounters: + bytes_sent: int + bytes_recv: int + + +@dataclass(frozen=True, slots=True) +class SelectedNicCounters: + nic_names: tuple[str, ...] + counters: NicCounters + + +@dataclass(frozen=True, slots=True) +class MetricReading: + timestamp_unix_s: float + monotonic_s: float + process_cpu_percent: float + rss_bytes: int + nic: NicCounters + nic_names: tuple[str, ...] + + +@dataclass(frozen=True, slots=True) +class PhaseSnapshot: + counts: Mapping[str, int] + durations_seconds: Mapping[str, float] + last_timestamps_unix_s: Mapping[str, float] + + @classmethod + def empty(cls) -> PhaseSnapshot: + return cls(counts={}, durations_seconds={}, last_timestamps_unix_s={}) + + +@dataclass(frozen=True, slots=True) +class MetricsSample: + reading: MetricReading + nic_tx_bytes_per_sec: float + nic_rx_bytes_per_sec: float + phases: PhaseSnapshot + + def to_json_line(self) -> str: + record = { + "record_type": "sample", + "timestamp_unix_s": self.reading.timestamp_unix_s, + "process_cpu_percent": self.reading.process_cpu_percent, + "rss_bytes": self.reading.rss_bytes, + "nic_names": self.reading.nic_names, + "nic_tx_bytes": self.reading.nic.bytes_sent, + "nic_rx_bytes": self.reading.nic.bytes_recv, + "nic_tx_bytes_per_sec": self.nic_tx_bytes_per_sec, + "nic_rx_bytes_per_sec": self.nic_rx_bytes_per_sec, + "phase_counts": self.phases.counts, + "phase_durations_seconds": self.phases.durations_seconds, + "phase_last_timestamps_unix_s": ( + self.phases.last_timestamps_unix_s + ), + } + return json.dumps(record, sort_keys=True, separators=(",", ":")) + "\n" + + +def aggregate_nic_counters( + counters: Mapping[str, NicCounters], + nic_names: Sequence[str] | None, +) -> SelectedNicCounters: + """Aggregate selected interfaces in stable lexical order.""" + selected_names = tuple( + sorted(counters if nic_names is None else set(nic_names)) + ) + missing_names = tuple( + name for name in selected_names if name not in counters + ) + if missing_names: + missing = ", ".join(missing_names) + raise KeyError(f"network interface(s) unavailable: {missing}") + return SelectedNicCounters( + nic_names=selected_names, + counters=NicCounters( + bytes_sent=sum( + counters[name].bytes_sent for name in selected_names + ), + bytes_recv=sum( + counters[name].bytes_recv for name in selected_names + ), + ), + ) + + +def build_sample( + current: MetricReading, + previous: MetricReading | None, + phases: PhaseSnapshot, +) -> MetricsSample: + """Compute byte rates from cumulative NIC readings.""" + if previous is None: + tx_rate = 0.0 + rx_rate = 0.0 + else: + elapsed = current.monotonic_s - previous.monotonic_s + if elapsed <= 0: + tx_rate = 0.0 + rx_rate = 0.0 + else: + tx_rate = ( + max(0, current.nic.bytes_sent - previous.nic.bytes_sent) + / elapsed + ) + rx_rate = ( + max(0, current.nic.bytes_recv - previous.nic.bytes_recv) + / elapsed + ) + return MetricsSample( + reading=current, + nic_tx_bytes_per_sec=tx_rate, + nic_rx_bytes_per_sec=rx_rate, + phases=phases, + ) diff --git a/pyproject.toml b/pyproject.toml index a1670952..e382920e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,7 @@ [build-system] requires = [ "setuptools", + "setuptools-scm>=8", "wheel", "cmake>=4.3.4", "ninja", @@ -10,7 +11,6 @@ build-backend = "setuptools.build_meta" [project] name = "morphling" -version = "0.0.1" description = "Morphling: an emulator for distributed machine-learning training on heterogeneous edge devices (EdgeSys '26 companion code)." readme = {file = "README.md", content-type = "text/markdown"} requires-python = ">=3.10" @@ -38,7 +38,7 @@ classifiers = [ "Topic :: System :: Distributed Computing", "Intended Audience :: Science/Research", ] -dynamic = ["dependencies"] +dynamic = ["version", "dependencies"] [project.urls] Homepage = "https://github.com/drunkcoding/DeviceEmulator" @@ -83,6 +83,15 @@ install = "_setup_helpers.CustomInstall" sdist = "_setup_helpers.CustomBuild" bdist_wheel = "_setup_helpers.CustomBdistWheel" +[tool.setuptools_scm] +# The Docker build context excludes .git, so the image injects the version via +# SETUPTOOLS_SCM_PRETEND_VERSION_FOR_MORPHLING (see Dockerfile); fallback_version +# applies only to a bare sdist build with neither git nor that env var. +version_file = "morphling/_version.py" +version_scheme = "guess-next-dev" +local_scheme = "node-and-date" +fallback_version = "0.0.0+unknown" + [tool.ruff] line-length = 80 exclude = [] diff --git a/scripts/_multi_coordinator_fleet.py b/scripts/_multi_coordinator_fleet.py new file mode 100644 index 00000000..d6ce5bfa --- /dev/null +++ b/scripts/_multi_coordinator_fleet.py @@ -0,0 +1,82 @@ +"""Rank-owned fake-device process lifecycle.""" + +from __future__ import annotations + +import os +import subprocess +from pathlib import Path +from typing import Protocol, Sequence + +import torch + + +class FleetProcess(Protocol): + def poll(self) -> int | None: ... + + def terminate(self) -> None: ... + + def kill(self) -> None: ... + + def wait(self, timeout: float | None = None) -> int: ... + + +def spawn_fleet( + rank: int, count: int, device_config: Path +) -> tuple[subprocess.Popen[bytes], ...]: + gpu_count = max(1, torch.cuda.device_count()) + log_directory = os.environ.get("MORPHLING_DEV_LOG_DIR") + processes: list[subprocess.Popen[bytes]] = [] + for index in range(count): + environment = os.environ.copy() + environment["CUDA_VISIBLE_DEVICES"] = str( + (rank * count + index) % gpu_count + ) + command = ( + "morphling_device", + "--id", + str(rank * count + index), + "--flops", + "100T", + "--memory", + "8G", + "--ul_bw", + "10G", + "--dl_bw", + "10G", + "--ul_lat", + "0", + "--dl_lat", + "0", + "--backend", + "proxy", + "--cfg", + str(device_config), + ) + output_path = ( + Path(log_directory) / f"dev_c{rank}_{index}.log" + if log_directory + else Path(os.devnull) + ) + with output_path.open("wb") as output: + processes.append( + subprocess.Popen( + command, + env=environment, + stdout=output, + stderr=output, + start_new_session=False, + ) + ) + return tuple(processes) + + +def terminate_fleet(processes: Sequence[FleetProcess]) -> None: + for process in processes: + if process.poll() is None: + process.terminate() + for process in processes: + try: + process.wait(timeout=10.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait() diff --git a/scripts/_multi_coordinator_training.py b/scripts/_multi_coordinator_training.py new file mode 100644 index 00000000..2d6096d4 --- /dev/null +++ b/scripts/_multi_coordinator_training.py @@ -0,0 +1,201 @@ +"""Rank-local training workload for the multi-coordinator launcher.""" + +from __future__ import annotations + +import argparse +import json +import os +import time +import traceback +from datetime import timedelta +from pathlib import Path +from typing import Final + +import torch.distributed as dist + +from morphling.runtime.coordinator_metrics import ( + CoordinatorMetricsCollector, + PhaseRecorder, + metrics_config_from_args, + start_metrics_collector, + track_phase, +) +from scripts._multi_coordinator_fleet import spawn_fleet as _spawn_fleet +from scripts._multi_coordinator_fleet import terminate_fleet as _terminate_fleet +from scripts._multi_coordinator_workload import run_workload +from scripts.multi_coordinator_scaling_config import ScalingConfig +from scripts.multi_coordinator_scaling_results import ( + RankMeasurement, + build_rank_result, + rank_result_payload, +) + + +def _write_rank_config( + template_path: Path, rank: int, listen_port: int, device_count: int +) -> Path: + replacements = { + "listen_port": f"listen_port = {listen_port}\n", + "listen_ip": "listen_ip = 127.0.0.1\n", + "num_device": f"num_device = {device_count}\n", + } + lines = template_path.read_text(encoding="utf-8").splitlines(keepends=True) + configured = [ + replacements.get(line.lstrip().split("=", maxsplit=1)[0].strip(), line) + for line in lines + ] + path = Path("/tmp") / f"svr_c{rank}.ini" + path.write_text("".join(configured), encoding="utf-8") + return path + + +def _scaling_config(args: argparse.Namespace) -> ScalingConfig: + return ScalingConfig( + mode=args.scaling_mode, + coordinators=args.coords, + total_devices=args.coords * args.devices_per_coord, + global_batch=args.coords * args.local_batch, + devices_per_coordinator=args.devices_per_coord, + local_batch=args.local_batch, + warmup_iterations=args.warmup_steps, + measured_iterations=args.steps, + relative_tolerance=args.rtol, + tiny=args.tiny, + model_name=args.model_name, + sequence_length=args.sequence_length, + block_size=args.block_size, + learning_rate=args.lr, + distributed_timeout_seconds=args.distributed_timeout_seconds, + ) + + +def _write_result( + args: argparse.Namespace, + rank: int, + recorder: PhaseRecorder, + losses: tuple[float, ...], + golden: tuple[float, ...], +) -> None: + if args.result_output is None: + return + config = _scaling_config(args) + result = build_rank_result( + config, + RankMeasurement( + rank, recorder.snapshot().durations_seconds, losses, golden + ), + ) + path = Path(args.result_output) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + rank_result_payload(config, result), indent=2, sort_keys=True + ) + + "\n", + encoding="utf-8", + ) + + +def _write_failure( + args: argparse.Namespace, rank: int, error: Exception +) -> None: + if args.result_output is None: + return + path = Path(args.result_output) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "schema_version": 1, + "rank": rank, + "status": "failed", + "error_type": type(error).__name__, + "error_message": str(error), + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + + +def run_coordinator(rank: int, args: argparse.Namespace) -> int: + proxy_port = args.base_proxy_port + rank + os.environ.update( + MORPHLING_PROXY_HOST="127.0.0.1", + MORPHLING_PROXY_PORT=str(proxy_port), + NUM_DEVICES=str(args.devices_per_coord), + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(args.master_port), + ) + recorder = PhaseRecorder() + collector: CoordinatorMetricsCollector | None = None + backend = None + fleet = () + group_initialized = False + metrics_config = metrics_config_from_args(args, rank=rank) + try: + if metrics_config is not None: + collector = start_metrics_collector(metrics_config, recorder) + with track_phase(recorder, "idle_wait"): + dist.init_process_group( + "gloo", + rank=rank, + world_size=args.coords, + timeout=timedelta(seconds=args.distributed_timeout_seconds), + ) + group_initialized = True + from scripts._runtime_common import start_backend, wait_for_connections + + root = Path(__file__).resolve().parents[1] + rank_config = _write_rank_config( + root / "config/proxy/svr.ini", + rank, + proxy_port, + args.devices_per_coord, + ) + backend = start_backend("proxy", args.block_size, str(rank_config)) + import morphling.hooks.autograd as hooks_autograd + + hooks_autograd._backend = backend + with track_phase(recorder, "idle_wait"): + time.sleep(3) + fleet = _spawn_fleet( + rank, args.devices_per_coord, root / "config/proxy/cli.ini" + ) + if hasattr(backend, "get_connection_count"): + with track_phase(recorder, "idle_wait"): + wait_for_connections(backend, args.devices_per_coord, 120) + with track_phase(recorder, "idle_wait"): + time.sleep(3) + config = _scaling_config(args) + losses, golden = run_workload(rank, config, recorder) + _write_result(args, rank, recorder, losses, golden) + passed = build_rank_result( + config, + RankMeasurement( + rank, recorder.snapshot().durations_seconds, losses, golden + ), + ).loss_correctness.passed + with track_phase(recorder, "idle_wait"): + dist.barrier() + return 0 if passed else 1 + finally: + if collector is not None: + collector.stop() + if group_initialized and dist.is_initialized(): + dist.destroy_process_group() + _terminate_fleet(fleet) + if backend is not None and hasattr(backend, "stop"): + backend.stop() + + +def worker_boundary(rank: int, args: argparse.Namespace) -> None: + try: + status = run_coordinator(rank, args) + except Exception as error: # noqa: BLE001 - top-level worker boundary + _write_failure(args, rank, error) + traceback.print_exc() + status = 1 + os._exit(status) diff --git a/scripts/_multi_coordinator_workload.py b/scripts/_multi_coordinator_workload.py new file mode 100644 index 00000000..73e29ab1 --- /dev/null +++ b/scripts/_multi_coordinator_workload.py @@ -0,0 +1,193 @@ +"""Tiny and OPT rank workloads sharing the single-coordinator model helpers.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Callable + +import torch +import torch.distributed as dist +import torch.nn.functional as F + +from morphling.runtime.coordinator_metrics import PhaseRecorder, track_phase +from scripts._runtime_common import load_model_and_tokenizer +from scripts._single_coordinator_training_workload import ( + TokenBatchSpec, + apply_restricted_hook, + extract_loss, + infer_vocab_size, + make_token_batch, + remove_restricted_hook, +) +from scripts.multi_coordinator_scaling_config import ScalingConfig + +LossFunction = Callable[[torch.nn.Module], torch.Tensor] + + +@dataclass(frozen=True, slots=True) +class ModelStep: + model: torch.nn.Module + loss_function: LossFunction + + +def _optimizer( + config: ScalingConfig, model: torch.nn.Module +) -> torch.optim.Optimizer: + return torch.optim.AdamW(model.parameters(), lr=config.learning_rate) + + +def _native_losses( + config: ScalingConfig, + step_function: ModelStep, +) -> tuple[float, ...]: + optimizer = _optimizer(config, step_function.model) + losses: list[float] = [] + for step in range(config.warmup_iterations + config.measured_iterations): + optimizer.zero_grad() + loss = step_function.loss_function(step_function.model) + loss.backward() + if step >= config.warmup_iterations: + losses.append(float(loss.detach())) + optimizer.step() + return tuple(losses) + + +def _distributed_losses( + config: ScalingConfig, + step_function: ModelStep, + recorder: PhaseRecorder, +) -> tuple[float, ...]: + optimizer = _optimizer(config, step_function.model) + losses: list[float] = [] + for step in range(config.warmup_iterations + config.measured_iterations): + warmup = step < config.warmup_iterations + measured_recorder = None if warmup else recorder + phase_name = "warmup_iteration" if warmup else "iteration_total" + with track_phase(recorder, phase_name): + optimizer.zero_grad() + with track_phase( + measured_recorder, "forward_device_dispatch_aggregation" + ): + loss = step_function.loss_function(step_function.model) + with track_phase( + measured_recorder, "backward_device_dispatch_aggregation" + ): + loss.backward() + with track_phase(measured_recorder, "gradient_sync"): + for parameter in step_function.model.parameters(): + if parameter.grad is not None: + dist.all_reduce(parameter.grad, op=dist.ReduceOp.SUM) + parameter.grad /= config.coordinators + global_loss = loss.detach().clone() + dist.all_reduce(global_loss, op=dist.ReduceOp.SUM) + global_loss /= config.coordinators + if not warmup: + losses.append(float(global_loss)) + with track_phase(measured_recorder, "optimizer"): + optimizer.step() + return tuple(losses) + + +def _tiny_model(seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + return torch.nn.Sequential( + torch.nn.Linear(8, 16), + torch.nn.ReLU(), + torch.nn.Linear(16, 4), + ) + + +def _tiny_inputs( + seed: int, batch_size: int +) -> tuple[torch.Tensor, torch.Tensor]: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed) + return ( + torch.randn(batch_size, 8, generator=generator), + torch.randn(batch_size, 4, generator=generator), + ) + + +def _run_tiny( + rank: int, config: ScalingConfig, recorder: PhaseRecorder +) -> tuple[tuple[float, ...], tuple[float, ...]]: + inputs, targets = _tiny_inputs(42, config.global_batch) + golden = _native_losses( + config, + ModelStep( + _tiny_model(42), + lambda model: F.mse_loss(model(inputs), targets), + ), + ) + lower = rank * config.local_batch + local_inputs = inputs[lower : lower + config.local_batch] + local_targets = targets[lower : lower + config.local_batch] + dispatched = _distributed_losses( + config, + ModelStep( + _tiny_model(42), + lambda model: F.mse_loss(model(local_inputs), local_targets), + ), + recorder, + ) + return dispatched, golden + + +def _load_full_model(config: ScalingConfig) -> tuple[torch.nn.Module, int]: + torch.manual_seed(42) + model, tokenizer = load_model_and_tokenizer( + config.model_name, dtype=torch.float32 + ) + return model.to(torch.device("cpu")), infer_vocab_size(model, tokenizer) + + +def _run_full( + rank: int, config: ScalingConfig, recorder: PhaseRecorder +) -> tuple[tuple[float, ...], tuple[float, ...]]: + golden_model, vocab_size = _load_full_model(config) + tokens = make_token_batch( + TokenBatchSpec( + seed=42, + batch_size=config.global_batch, + sequence_length=config.sequence_length, + vocab_size=vocab_size, + ), + device=torch.device("cpu"), + ) + golden = _native_losses( + config, + ModelStep( + golden_model, + lambda model: extract_loss( + model(input_ids=tokens, labels=tokens), tokens + ), + ), + ) + del golden_model + dispatched_model, _ = _load_full_model(config) + lower = rank * config.local_batch + local_tokens = tokens[lower : lower + config.local_batch] + apply_restricted_hook() + try: + dispatched = _distributed_losses( + config, + ModelStep( + dispatched_model, + lambda model: extract_loss( + model(input_ids=local_tokens, labels=local_tokens), + local_tokens, + ), + ), + recorder, + ) + finally: + remove_restricted_hook() + return dispatched, golden + + +def run_workload( + rank: int, config: ScalingConfig, recorder: PhaseRecorder +) -> tuple[tuple[float, ...], tuple[float, ...]]: + if config.tiny: + return _run_tiny(rank, config, recorder) + return _run_full(rank, config, recorder) diff --git a/scripts/_single_coordinator_fleet.py b/scripts/_single_coordinator_fleet.py new file mode 100644 index 00000000..371e0358 --- /dev/null +++ b/scripts/_single_coordinator_fleet.py @@ -0,0 +1,77 @@ +"""Local emulated-device fleet lifecycle for coordinator experiments.""" + +from __future__ import annotations + +import os +import subprocess + +import torch + +_device_processes: list[subprocess.Popen[bytes]] = [] + + +def cleanup_fake_fleet() -> None: + for process in _device_processes: + if process.poll() is not None: + continue + process.terminate() + try: + process.wait(timeout=5.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=5.0) + _device_processes.clear() + + +def spawn_fake_fleet( + num_devices: int, backend_name: str, proxy_host: str +) -> None: + cleanup_fake_fleet() + num_gpus = max(1, torch.cuda.device_count()) + root = os.path.dirname(os.path.dirname(os.path.realpath(__file__))) + device_cfg = os.environ.get("CFG_PATH") or os.path.join( + root, "config", "proxy", "cli.ini" + ) + log_dir = os.environ.get("MORPHLING_DEV_LOG_DIR", "") + for index in range(num_devices): + environment = os.environ.copy() + environment["CUDA_VISIBLE_DEVICES"] = str(index % num_gpus) + command = [ + "morphling_device", + "--id", + str(index), + "--flops", + "100T", + "--memory", + "8G", + "--ul_bw", + "10G", + "--dl_bw", + "10G", + "--ul_lat", + "0", + "--dl_lat", + "0", + "--backend", + backend_name, + "--cfg", + device_cfg, + ] + if proxy_host: + command += ["--proxy_host", proxy_host] + output = ( + open(os.path.join(log_dir, f"dev_{index}.log"), "wb") + if log_dir + else open(os.devnull, "wb") + ) + try: + process = subprocess.Popen( + command, + env=environment, + stdout=output, + stderr=output, + start_new_session=True, + ) + _device_processes.append(process) + finally: + output.close() diff --git a/scripts/_single_coordinator_training_workload.py b/scripts/_single_coordinator_training_workload.py new file mode 100644 index 00000000..7b7b2b8e --- /dev/null +++ b/scripts/_single_coordinator_training_workload.py @@ -0,0 +1,163 @@ +"""Training workloads shared by the single-coordinator evaluation launcher.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, List + +import torch +import torch.nn.functional as F + +from morphling.hooks.autograd import LinearFunction +from morphling.runtime.coordinator_metrics import PhaseRecorder, track_phase + +_LARGE_DIM = 10000 +_orig_linear_forward = torch.nn.Linear.forward + + +@dataclass(frozen=True, slots=True) +class TokenBatchSpec: + seed: int + batch_size: int + sequence_length: int + vocab_size: int + + +def make_token_batch( + spec: TokenBatchSpec, + device: torch.device, +) -> torch.Tensor: + generator = torch.Generator(device="cpu") + generator.manual_seed(spec.seed) + return torch.randint( + low=0, + high=spec.vocab_size, + size=(spec.batch_size, spec.sequence_length), + dtype=torch.long, + generator=generator, + ).to(device) + + +def extract_loss(outputs: Any, input_ids: torch.Tensor) -> torch.Tensor: + loss = getattr(outputs, "loss", None) + if isinstance(loss, torch.Tensor): + return loss + logits = outputs["logits"] if isinstance(outputs, dict) else outputs.logits + return F.cross_entropy(logits.view(-1, logits.size(-1)), input_ids.view(-1)) + + +def restricted_linear_forward(self: torch.nn.Linear, inp: torch.Tensor): + if self.out_features >= _LARGE_DIM or self.in_features >= _LARGE_DIM: + return _orig_linear_forward(self, inp) + return LinearFunction.apply(inp, self.weight.t(), self.bias) + + +def apply_restricted_hook() -> None: + """Dispatch transformer Linear GEMMs while keeping wide operations native.""" + torch.nn.Linear.forward = restricted_linear_forward + print( + f"Restricted Linear dispatch hook applied (skip dims >= {_LARGE_DIM}); " + "attention bmm/matmul and lm_head stay native.", + flush=True, + ) + + +def remove_restricted_hook() -> None: + torch.nn.Linear.forward = _orig_linear_forward + + +def infer_vocab_size(model: torch.nn.Module, tokenizer: Any) -> int: + vs = getattr(getattr(model, "config", object()), "vocab_size", None) + if isinstance(vs, int): + return vs + tvs = getattr(tokenizer, "vocab_size", None) + if isinstance(tvs, int): + return tvs + raise ValueError("Unable to infer vocab_size") + + +def train_loop( + model: torch.nn.Module, + *, + steps: int, + warmup_steps: int = 0, + seed: int, + batch_size: int, + seq_length: int, + device: torch.device, + vocab_size: int, + phase_recorder: PhaseRecorder | None = None, +) -> List[float]: + model.eval() + optimizer = torch.optim.AdamW(model.parameters(), lr=5e-5) + losses: List[float] = [] + for step_index in range(warmup_steps + steps): + is_warmup = step_index < warmup_steps + measured_recorder = None if is_warmup else phase_recorder + phase_name = "warmup_iteration" if is_warmup else "iteration_total" + with track_phase(phase_recorder, phase_name): + input_ids = make_token_batch( + TokenBatchSpec( + seed=seed + step_index, + batch_size=batch_size, + sequence_length=seq_length, + vocab_size=vocab_size, + ), + device=device, + ) + optimizer.zero_grad() + with track_phase( + measured_recorder, "forward_device_dispatch_aggregation" + ): + outputs = model(input_ids=input_ids, labels=input_ids) + loss = extract_loss(outputs, input_ids) + with track_phase( + measured_recorder, "backward_device_dispatch_aggregation" + ): + loss.backward() + losses.append(float(loss.detach().item())) + with track_phase(measured_recorder, "optimizer"): + optimizer.step() + return losses + + +def make_tiny_model(seed: int) -> torch.nn.Module: + torch.manual_seed(seed) + return torch.nn.Sequential(torch.nn.Linear(8, 4)) + + +def train_loop_tiny( + model: torch.nn.Module, + *, + steps: int, + warmup_steps: int = 0, + seed: int, + device: torch.device, + phase_recorder: PhaseRecorder | None = None, +) -> List[float]: + model.train() + optimizer = torch.optim.AdamW(model.parameters(), lr=0.01) + gen = torch.Generator(device="cpu") + gen.manual_seed(seed) + x = torch.randn(16, 8, generator=gen).to(device) + y = torch.randn(16, 4, generator=gen).to(device) + losses: List[float] = [] + for step_index in range(warmup_steps + steps): + is_warmup = step_index < warmup_steps + measured_recorder = None if is_warmup else phase_recorder + phase_name = "warmup_iteration" if is_warmup else "iteration_total" + with track_phase(phase_recorder, phase_name): + optimizer.zero_grad() + with track_phase( + measured_recorder, "forward_device_dispatch_aggregation" + ): + out = model(x) + loss = F.mse_loss(out, y) + with track_phase( + measured_recorder, "backward_device_dispatch_aggregation" + ): + loss.backward() + losses.append(float(loss.detach().item())) + with track_phase(measured_recorder, "optimizer"): + optimizer.step() + return losses diff --git a/scripts/compute_batch_runtime.py b/scripts/compute_batch_runtime.py index f8c8fc13..8f29f51c 100644 --- a/scripts/compute_batch_runtime.py +++ b/scripts/compute_batch_runtime.py @@ -446,6 +446,46 @@ def _baseline_uses_pipeline_overlay( return pp_size > 1 +def _within_batch_multiplier( + trace: dict[str, Any] | None, + device_id: int, + level: int, + num_levels: int, +) -> tuple[float, float, float]: + """(flops_mult, ul_bw_mult, dl_bw_mult) for an entry from a within-batch + trace, indexed by device and batch progress = level / num_levels. Returns + (1, 1, 1) when no trace or the device is absent.""" + if not trace: + return 1.0, 1.0, 1.0 + dev = trace.get("devices", {}).get(str(device_id)) + if not dev: + return 1.0, 1.0, 1.0 + num_bins = int(trace.get("num_bins", 1)) or 1 + frac = 0.0 if num_levels <= 1 else min(0.999999, max(0.0, level / num_levels)) + b = min(num_bins - 1, int(frac * num_bins)) + + def _at(key: str) -> float: + arr = dev.get(key) + if not arr: + return 1.0 + v = float(arr[b]) if b < len(arr) else float(arr[-1]) + return v if v > 1e-6 else 1e-6 + + return _at("flops_mult"), _at("ul_bw_mult"), _at("dl_bw_mult") + + +def _scale_phase( + phase: str, dur_us: float, fm: float, um: float, dm: float +) -> float: + if phase == "COMPUTE": + return dur_us / fm + if phase == "DOWNLOAD": + return dur_us / dm + if phase == "UPLOAD": + return dur_us / um + return dur_us + + def compute_batch_runtime( vtime_events: list[VTimeEvent], manifest: dict[str, Any] | list[dict[str, Any]], @@ -457,6 +497,7 @@ def compute_batch_runtime( device_profiles: dict[int, dict[str, float]] | None = None, bytes_per_element: float = 2.0, allreduce_bandwidth_bps: float | None = None, + within_batch_trace: dict[str, Any] | None = None, ) -> BatchRuntimeResult: device_profiles = device_profiles or {} baseline = baseline_type.lower() @@ -466,6 +507,8 @@ def compute_batch_runtime( if not entries: raise ValueError("Manifest has no dispatch entries") + num_levels = max((e["level"] for e in entries), default=0) + 1 + by_device_gemm, by_gemm = _build_entry_lookup(entries) collective_groups = _build_collective_groups(entries) @@ -499,8 +542,13 @@ def compute_batch_runtime( entry["stage_id"], entry["is_local"], ) + fm, um, dm = _within_batch_multiplier( + within_batch_trace, entry["device_id"], entry["level"], num_levels + ) for phase, dur_us in phases.items(): - phase_us_by_key[key][phase] += dur_us + phase_us_by_key[key][phase] += _scale_phase( + phase, dur_us, fm, um, dm + ) else: for entry in entries: key = ( @@ -516,8 +564,13 @@ def compute_batch_runtime( device_profiles, bytes_per_element, ) + fm, um, dm = _within_batch_multiplier( + within_batch_trace, entry["device_id"], entry["level"], num_levels + ) for phase, dur_us in synthetic.items(): - phase_us_by_key[key][phase] += max(dur_us, 0.0) + phase_us_by_key[key][phase] += _scale_phase( + phase, max(dur_us, 0.0), fm, um, dm + ) per_device_breakdown: dict[int, dict[str, float]] = defaultdict( lambda: {"compute_ms": 0.0, "network_ms": 0.0, "total_ms": 0.0} diff --git a/scripts/coord_scaling_plot/__init__.py b/scripts/coord_scaling_plot/__init__.py new file mode 100644 index 00000000..3f46cc62 --- /dev/null +++ b/scripts/coord_scaling_plot/__init__.py @@ -0,0 +1,12 @@ +"""Coordinator-scaling figure pipeline (paper figure D8). + +Typed, boundary-validated data-prep (:mod:`schema`) plus a three-panel +double-column figure builder (:mod:`figure`) and a CLI (:mod:`cli`) that +consumes the real ``summary.json`` / ``strong.json`` / ``weak.json`` / +``breakdown.json`` emitted by the coordinator-scaling experiments. + +The data-prep layer is intentionally free of any matplotlib dependency so it +can be unit-tested in isolation; import :mod:`figure` only when rendering. +""" + +from __future__ import annotations diff --git a/scripts/coord_scaling_plot/__main__.py b/scripts/coord_scaling_plot/__main__.py new file mode 100644 index 00000000..bc158152 --- /dev/null +++ b/scripts/coord_scaling_plot/__main__.py @@ -0,0 +1,8 @@ +"""Module entry point: ``python3 -m scripts.coord_scaling_plot ...``.""" + +from __future__ import annotations + +from .cli import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/coord_scaling_plot/_plot_utils.py b/scripts/coord_scaling_plot/_plot_utils.py new file mode 100644 index 00000000..af2ea92f --- /dev/null +++ b/scripts/coord_scaling_plot/_plot_utils.py @@ -0,0 +1,74 @@ +"""Locate and adapt the conference-plot skill's ``plot_utils`` helpers. + +The paper style, Wong colorblind-safe palette and dual-output saver live in the +``conference-plot`` skill outside this repository. This shim finds that module +(via the ``CONFERENCE_PLOT_SCRIPTS`` environment variable, else the known skill +location), forces a headless ``Agg`` backend before matplotlib initializes, and +re-exports the pieces the figure builder needs. + +It also adds :func:`save_paper_figure`, which writes a vector PDF through the +skill's ``save_dual_output`` and a 300-dpi PNG directly -- the skill helper +emits PNGs at 150 dpi, so the raster path is written here to meet the paper's +300-dpi requirement. +""" + +from __future__ import annotations + +import importlib.util +import os +from pathlib import Path +from typing import Tuple + +os.environ.setdefault("MPLBACKEND", "Agg") +import matplotlib # noqa: E402 (backend must be set before pyplot import) + +matplotlib.use("Agg", force=True) + +_DEFAULT_SCRIPTS = Path( + "/home/xly/.opencode/skills/conference-plot/scripts/plot_utils.py" +) + + +def _load_plot_utils(): + override = os.environ.get("CONFERENCE_PLOT_SCRIPTS") + candidate = ( + Path(override) / "plot_utils.py" if override else _DEFAULT_SCRIPTS + ) + if not candidate.is_file(): + raise ImportError( + "conference-plot plot_utils.py not found at " + f"{candidate}; set CONFERENCE_PLOT_SCRIPTS to its scripts directory" + ) + spec = importlib.util.spec_from_file_location( + "conference_plot_plot_utils", candidate + ) + if spec is None or spec.loader is None: # pragma: no cover - defensive + raise ImportError(f"cannot load a module spec from {candidate}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_plot_utils = _load_plot_utils() + +paper_style = _plot_utils.paper_style +save_dual_output = _plot_utils.save_dual_output +WONG_PALETTE = _plot_utils.WONG_PALETTE +HATCHES = _plot_utils.HATCHES + + +def save_paper_figure(fig, out_base: Path) -> Tuple[Path, Path]: + """Write ``out_base`` as a vector PDF and a 300-dpi PNG; return both paths.""" + base = Path(out_base) + pdf_path = base.with_suffix(".pdf") + png_path = base.with_suffix(".png") + save_dual_output(fig, pdf_path, None, save_both=False) + png_path.parent.mkdir(parents=True, exist_ok=True) + fig.savefig( + png_path, + dpi=300, + bbox_inches="tight", + pad_inches=0.05, + facecolor="white", + ) + return pdf_path, png_path diff --git a/scripts/coord_scaling_plot/cli.py b/scripts/coord_scaling_plot/cli.py new file mode 100644 index 00000000..6c857146 --- /dev/null +++ b/scripts/coord_scaling_plot/cli.py @@ -0,0 +1,99 @@ +"""Command-line entry point for the coordinator-scaling figure. + +Every input path is required and must point at an existing file; a missing or +malformed input fails clearly with a non-zero exit code and a message naming +the offending file rather than emitting a partial figure. The figure is never +generated from placeholder data -- the caller supplies the real +``summary.json`` / ``strong.json`` / ``weak.json`` / ``breakdown.json``. + +Usage:: + + python3 -m scripts.coord_scaling_plot \\ + --summary results/coord_scaling/single/summary.json \\ + --strong results/coord_scaling/strong.json \\ + --weak results/coord_scaling/weak.json \\ + --breakdown results/coord_scaling/breakdown.json +""" + +from __future__ import annotations + +import argparse +import sys +from pathlib import Path +from typing import Optional, Sequence, Tuple + +from . import schema + +_DEFAULT_OUT = "results/coord_scaling/figures/coordinator_scaling" + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="coord_scaling_plot", + description="Render the three-panel coordinator-scaling paper figure.", + ) + parser.add_argument( + "--summary", + required=True, + type=Path, + help="single-coordinator device-scaling summary.json", + ) + parser.add_argument( + "--strong", required=True, type=Path, help="strong-scaling strong.json" + ) + parser.add_argument( + "--weak", required=True, type=Path, help="weak-scaling weak.json" + ) + parser.add_argument( + "--breakdown", + required=True, + type=Path, + help="per-iteration breakdown.json", + ) + parser.add_argument( + "--out", + type=Path, + default=Path(_DEFAULT_OUT), + help="output path stem (a .pdf and .png are written); " + f"default: {_DEFAULT_OUT}", + ) + return parser + + +def _load_inputs(args: argparse.Namespace): + device = schema.parse_device_scaling(schema.load_json(args.summary)) + strong = schema.parse_scaling_efficiency( + schema.load_json(args.strong), "strong" + ) + weak = schema.parse_scaling_efficiency(schema.load_json(args.weak), "weak") + breakdown = schema.parse_breakdown(schema.load_json(args.breakdown)) + return device, strong, weak, breakdown + + +def main(argv: Optional[Sequence[str]] = None) -> int: + args = build_parser().parse_args(argv) + try: + device, strong, weak, breakdown = _load_inputs(args) + except schema.SchemaError as exc: + print(f"error: {exc}", file=sys.stderr) + return 2 + + from . import ( + figure, # imported lazily so schema-only errors avoid matplotlib + ) + + pdf_path, png_path = figure.render( + device, strong, weak, breakdown, args.out + ) + _report(pdf_path, png_path) + return 0 + + +def _report(pdf_path: Path, png_path: Path) -> Tuple[Path, Path]: + print(f"wrote {pdf_path}") + print(f"wrote {png_path}") + return pdf_path, png_path + + +if __name__ == "__main__": # pragma: no cover + raise SystemExit(main()) diff --git a/scripts/coord_scaling_plot/figure.py b/scripts/coord_scaling_plot/figure.py new file mode 100644 index 00000000..452b7ddb --- /dev/null +++ b/scripts/coord_scaling_plot/figure.py @@ -0,0 +1,246 @@ +"""Three-panel double-column coordinator-scaling figure (paper figure D8). + +Panels, left to right: + +* (a) Measured iteration runtime and sustained normalized CPU versus emulated + device count, with the loopback/emulated substrate and the measured device + range marked and no exact value labels. +* (b) Strong and weak multi-coordinator scaling efficiency for coordinator + counts 1 and 2, as separate marked lines against an ideal reference; only + the available points are drawn. +* (c) Per-iteration time composition as stacked bars for each available + (mode, coordinator-count) run, using the canonical component order that + reconciles to the iteration total. + +Rendering uses the conference-plot skill's paper style, Wong colorblind-safe +palette and hatches (via :mod:`._plot_utils`). The figure carries no +figure-level title; each panel gets a concise ``(a)/(b)/(c)`` label. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Tuple + +import matplotlib.pyplot as plt +import numpy as np + +from . import schema +from ._plot_utils import HATCHES, WONG_PALETTE, paper_style, save_paper_figure + +COMPONENT_LABELS: Tuple[str, ...] = ( + "Dispatch+agg", + "Gradient sync", + "Optimizer", + "Idle/other", +) + +_FIG_WIDTH_INCHES = 7.0 +_FIG_HEIGHT_INCHES = 2.4 +_SMALL = 6 + +_RUNTIME_COLOR = WONG_PALETTE[5] +_CPU_COLOR = WONG_PALETTE[6] +_STRONG_COLOR = WONG_PALETTE[1] +_WEAK_COLOR = WONG_PALETTE[3] +_COMPONENT_COLORS = ( + WONG_PALETTE[2], + WONG_PALETTE[3], + WONG_PALETTE[6], + WONG_PALETTE[7], +) +_COMPONENT_HATCHES = (HATCHES[1], HATCHES[2], HATCHES[3], HATCHES[4]) + + +def _panel_device(ax, device: schema.DeviceScaling) -> None: + counts = [point.device_count for point in device.points] + runtimes = [point.iteration_runtime_seconds for point in device.points] + cpu = [point.sustained_cpu_percent for point in device.points] + + line_rt = ax.plot( + counts, + runtimes, + color=_RUNTIME_COLOR, + marker="o", + markersize=3.5, + label="Iteration runtime", + ) + twin = ax.twinx() + line_cpu = twin.plot( + counts, + cpu, + color=_CPU_COLOR, + marker="s", + ls="--", + markersize=3.5, + label="Sustained CPU", + ) + + ax.set_xscale("log", base=2) + ax.set_xticks(counts) + ax.set_xticklabels([str(count) for count in counts]) + ax.set_xlim(min(counts) * 2**-0.4, max(counts) * 2**0.4) + ax.axvspan(min(counts), max(counts), color="0.85", alpha=0.35, zorder=0) + ax.set_xlabel("Emulated devices") + ax.set_ylabel("Iteration runtime (s)", color=_RUNTIME_COLOR) + ax.tick_params(axis="y", colors=_RUNTIME_COLOR) + ax.set_ylim(bottom=0) + twin.set_ylabel("Sustained CPU (% host)", color=_CPU_COLOR) + twin.tick_params(axis="y", colors=_CPU_COLOR) + twin.set_ylim(bottom=0) + twin.grid(False) + + ax.text( + 0.5, + 0.12, + "loopback/emulated substrate", + transform=ax.transAxes, + ha="center", + va="bottom", + fontsize=_SMALL, + style="italic", + color="0.35", + ) + ax.text( + 0.5, + 0.03, + "measured range", + transform=ax.transAxes, + ha="center", + va="bottom", + fontsize=_SMALL, + color="0.4", + ) + handles = line_rt + line_cpu + ax.legend( + handles, + [handle.get_label() for handle in handles], + fontsize=_SMALL, + loc="center right", + framealpha=0.9, + ) + + +def _plot_efficiency( + ax, + efficiency: schema.ScalingEfficiency, + color: str, + marker: str, + style: str, +) -> None: + coordinators = [point.coordinators for point in efficiency.points] + values = [point.efficiency for point in efficiency.points] + ax.plot( + coordinators, + values, + color=color, + marker=marker, + ls=style, + markersize=4, + label=efficiency.mode.capitalize(), + ) + + +def _panel_efficiency( + ax, strong: schema.ScalingEfficiency, weak: schema.ScalingEfficiency +) -> None: + _plot_efficiency(ax, strong, _STRONG_COLOR, "o", "-") + _plot_efficiency(ax, weak, _WEAK_COLOR, "s", "--") + ax.axhline(1.0, color="0.5", ls=":", lw=0.8, label="Ideal") + + ticks = sorted( + {point.coordinators for point in (*strong.points, *weak.points)} + ) + ax.set_xticks(ticks) + ax.set_xticklabels([str(tick) for tick in ticks]) + ax.set_xlim(min(ticks) - 0.2, max(ticks) + 0.2) + ax.set_ylim(0.0, 1.15) + ax.set_xlabel("Coordinators") + ax.set_ylabel("Scaling efficiency") + ax.legend(fontsize=_SMALL, loc="lower left", framealpha=0.9) + + +def _panel_breakdown(ax, breakdown: schema.Breakdown) -> None: + rows = breakdown.rows + positions = np.arange(len(rows)) + labels = [ + ("S" if row.mode == "strong" else "W") + str(row.coordinators) + for row in rows + ] + bottoms = np.zeros(len(rows)) + for index, component_label in enumerate(COMPONENT_LABELS): + heights = np.array([row.components[index] for row in rows]) + ax.bar( + positions, + heights, + bottom=bottoms, + width=0.7, + color=_COMPONENT_COLORS[index], + hatch=_COMPONENT_HATCHES[index], + edgecolor="black", + linewidth=0.4, + label=component_label, + ) + bottoms += heights + + ax.set_xticks(positions) + ax.set_xticklabels(labels) + ax.set_xlabel("Mode-coordinators") + ax.set_ylabel("Per-iteration time (s)") + top = max((row.iteration_total for row in rows), default=1.0) + ax.set_ylim(0, top * 1.35) + ax.legend( + fontsize=_SMALL, + loc="upper right", + framealpha=0.9, + handlelength=1.4, + handleheight=1.2, + ) + + +def _assemble( + device: schema.DeviceScaling, + strong: schema.ScalingEfficiency, + weak: schema.ScalingEfficiency, + breakdown: schema.Breakdown, +): + fig, axes = plt.subplots(1, 3) + _panel_device(axes[0], device) + _panel_efficiency(axes[1], strong, weak) + _panel_breakdown(axes[2], breakdown) + axes[0].set_title("(a) Coordinator device scaling") + axes[1].set_title("(b) Multi-coordinator efficiency") + axes[2].set_title("(c) Per-iteration breakdown") + fig.tight_layout(pad=0.4, w_pad=0.8) + return fig + + +def build_figure( + device: schema.DeviceScaling, + strong: schema.ScalingEfficiency, + weak: schema.ScalingEfficiency, + breakdown: schema.Breakdown, +): + """Build and return the three-panel figure (caller owns ``plt.close``).""" + with paper_style( + width=_FIG_WIDTH_INCHES, height=_FIG_HEIGHT_INCHES, font_size=7 + ): + return _assemble(device, strong, weak, breakdown) + + +def render( + device: schema.DeviceScaling, + strong: schema.ScalingEfficiency, + weak: schema.ScalingEfficiency, + breakdown: schema.Breakdown, + out_base: Path, +) -> Tuple[Path, Path]: + """Build the figure and save a vector PDF + 300-dpi PNG under ``out_base``.""" + with paper_style( + width=_FIG_WIDTH_INCHES, height=_FIG_HEIGHT_INCHES, font_size=7 + ): + fig = _assemble(device, strong, weak, breakdown) + try: + return save_paper_figure(fig, out_base) + finally: + plt.close(fig) diff --git a/scripts/coord_scaling_plot/schema.py b/scripts/coord_scaling_plot/schema.py new file mode 100644 index 00000000..c0783532 --- /dev/null +++ b/scripts/coord_scaling_plot/schema.py @@ -0,0 +1,298 @@ +"""Typed JSON boundary parsing for the coordinator-scaling figure. + +Untrusted result JSON crosses into typed, immutable dataclasses exactly once, +here. Everything downstream (the figure builder, the CLI) receives validated +values and never re-checks them. This module deliberately depends only on the +standard library so it can be unit-tested without matplotlib. + +Consumed inputs and the fields this parser actually reads: + +* ``summary.json`` (single-coordinator device scaling, panel a) -- + ``environment.logical_cpu_count``, ``environment.host_nic_substrate`` and, + per ``rows[]`` entry, ``device_count``, ``iteration_runtime_seconds`` and + ``median_cpu_percent``. Sustained CPU is normalized to a percentage of the + full host by dividing process CPU percent by the logical CPU count. +* ``strong.json`` / ``weak.json`` (multi-coordinator scaling, panel b) -- + per ``points[]`` entry, ``configuration.coordinators`` and + ``result.iteration_total_seconds``. Efficiency is derived per point against + the single-coordinator baseline: strong keeps global work fixed + (``base / (coordinators * seconds)``); weak keeps local work fixed + (``base / seconds``). +* ``breakdown.json`` (per-iteration composition, panel c) -- per ``rows[]`` + entry, ``mode``, ``coordinators``, ``iteration_total`` and the four + components in :data:`BREAKDOWN_COMPONENTS`, which must reconcile to the + iteration total. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Optional, Tuple + +# Canonical bottom-to-top stack order for the per-iteration breakdown. +BREAKDOWN_COMPONENTS: Tuple[str, ...] = ( + "device_dispatch_aggregation", + "gradient_sync", + "optimizer", + "idle_other", +) + +SCALING_MODES: Tuple[str, ...] = ("strong", "weak") + +_DEFAULT_SUBSTRATE = "single-host loopback/emulated devices" +_RECON_RTOL = 0.02 +_RECON_ATOL = 1e-6 + + +class SchemaError(ValueError): + """Raised when result JSON is missing required fields or is malformed.""" + + +class MissingDataError(SchemaError): + """Raised when required JSON is well-formed but has no usable points.""" + + +# ── Typed, immutable views ────────────────────────────────────────────────── + + +@dataclass(frozen=True) +class DeviceScalingPoint: + device_count: int + iteration_runtime_seconds: float + sustained_cpu_percent: float + peak_cpu_percent: Optional[float] + + +@dataclass(frozen=True) +class DeviceScaling: + substrate: str + logical_cpu_count: int + points: Tuple[DeviceScalingPoint, ...] + + +@dataclass(frozen=True) +class EfficiencyPoint: + coordinators: int + iteration_total_seconds: float + efficiency: float + throughput_samples_per_second: Optional[float] + + +@dataclass(frozen=True) +class ScalingEfficiency: + mode: str + points: Tuple[EfficiencyPoint, ...] + + +@dataclass(frozen=True) +class BreakdownRow: + mode: str + coordinators: int + iteration_total: float + components: Tuple[float, ...] + + +@dataclass(frozen=True) +class Breakdown: + rows: Tuple[BreakdownRow, ...] + component_semantics: Tuple[Tuple[str, str], ...] + + +# ── Small validation helpers ──────────────────────────────────────────────── + + +def _mapping(value: Any, ctx: str) -> dict: + if not isinstance(value, dict): + raise SchemaError( + f"{ctx}: expected an object, got {type(value).__name__}" + ) + return value + + +def _sequence(value: Any, ctx: str) -> list: + if not isinstance(value, (list, tuple)): + raise SchemaError(f"{ctx}: expected a list, got {type(value).__name__}") + return list(value) + + +def _require(mapping: dict, key: str, ctx: str) -> Any: + if key not in mapping: + raise SchemaError(f"{ctx}: missing required field '{key}'") + return mapping[key] + + +def _as_float(value: Any, ctx: str) -> float: + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise SchemaError(f"{ctx}: expected a number, got {value!r}") + return float(value) + + +def _as_int(value: Any, ctx: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise SchemaError(f"{ctx}: expected an integer, got {value!r}") + return value + + +def load_json(path: Any) -> Any: + """Read and decode a JSON file, failing clearly on missing/invalid input.""" + resolved = Path(path) + if not resolved.is_file(): + raise SchemaError(f"input file not found: {resolved}") + try: + return json.loads(resolved.read_text(encoding="utf-8")) + except json.JSONDecodeError as exc: + raise SchemaError(f"invalid JSON in {resolved}: {exc}") from exc + + +# ── Parsers ───────────────────────────────────────────────────────────────── + + +def parse_device_scaling(data: Any) -> DeviceScaling: + """Parse single-coordinator device-scaling ``summary.json`` (panel a).""" + root = _mapping(data, "summary") + environment = _mapping( + _require(root, "environment", "summary"), "summary.environment" + ) + logical = _as_int( + _require(environment, "logical_cpu_count", "summary.environment"), + "logical_cpu_count", + ) + if logical <= 0: + raise SchemaError( + "summary.environment.logical_cpu_count must be positive" + ) + rows = _sequence(_require(root, "rows", "summary"), "summary.rows") + + points = [] + substrate = environment.get("host_nic_substrate") + for index, entry in enumerate(rows): + row = _mapping(entry, f"summary.rows[{index}]") + ctx = f"summary.rows[{index}]" + runtime = _require(row, "iteration_runtime_seconds", ctx) + if runtime is None: # failed repetition -- keep only measured points + continue + substrate = substrate or row.get("substrate") + peak = row.get("peak_cpu_percent") + points.append( + DeviceScalingPoint( + device_count=_as_int(_require(row, "device_count", ctx), ctx), + iteration_runtime_seconds=_as_float(runtime, ctx), + sustained_cpu_percent=_as_float( + _require(row, "median_cpu_percent", ctx), ctx + ) + / logical, + peak_cpu_percent=None if peak is None else _as_float(peak, ctx), + ) + ) + + if not points: + raise MissingDataError("summary.rows has no measured device points") + points.sort(key=lambda point: point.device_count) + return DeviceScaling( + substrate=substrate or _DEFAULT_SUBSTRATE, + logical_cpu_count=logical, + points=tuple(points), + ) + + +def parse_scaling_efficiency(data: Any, mode: str) -> ScalingEfficiency: + """Parse ``strong.json``/``weak.json`` and derive per-point efficiency.""" + if mode not in SCALING_MODES: + raise SchemaError( + f"unknown scaling mode {mode!r}; expected {SCALING_MODES}" + ) + root = _mapping(data, mode) + entries = _sequence(_require(root, "points", mode), f"{mode}.points") + + seconds_by_coord = {} + raw = [] + for index, entry in enumerate(entries): + point = _mapping(entry, f"{mode}.points[{index}]") + config = _mapping( + _require(point, "configuration", f"{mode}.points[{index}]"), + "configuration", + ) + result = _mapping( + _require(point, "result", f"{mode}.points[{index}]"), "result" + ) + coordinators = _as_int( + _require(config, "coordinators", "configuration"), "coordinators" + ) + seconds = _as_float( + _require(result, "iteration_total_seconds", "result"), + "iteration_total_seconds", + ) + throughput = result.get("throughput_samples_per_second") + seconds_by_coord[coordinators] = seconds + raw.append( + ( + coordinators, + seconds, + None if throughput is None else _as_float(throughput, "result"), + ) + ) + + if not raw: + raise MissingDataError(f"{mode}.points is empty") + baseline = seconds_by_coord.get(1) + if baseline is None: + raise MissingDataError( + f"{mode} scaling lacks the single-coordinator baseline point" + ) + + points = [] + for coordinators, seconds, throughput in sorted( + raw, key=lambda item: item[0] + ): + if mode == "strong": + efficiency = baseline / (coordinators * seconds) + else: + efficiency = baseline / seconds + points.append( + EfficiencyPoint(coordinators, seconds, efficiency, throughput) + ) + return ScalingEfficiency(mode=mode, points=tuple(points)) + + +def parse_breakdown(data: Any) -> Breakdown: + """Parse ``breakdown.json`` stacked-bar components (panel c).""" + root = _mapping(data, "breakdown") + entries = _sequence(_require(root, "rows", "breakdown"), "breakdown.rows") + if not entries: + raise MissingDataError("breakdown.rows is empty") + + rows = [] + for index, entry in enumerate(entries): + row = _mapping(entry, f"breakdown.rows[{index}]") + ctx = f"breakdown.rows[{index}]" + iteration_total = _as_float(_require(row, "iteration_total", ctx), ctx) + components = tuple( + _as_float(_require(row, name, ctx), ctx) + for name in BREAKDOWN_COMPONENTS + ) + if abs(sum(components) - iteration_total) > ( + _RECON_ATOL + _RECON_RTOL * abs(iteration_total) + ): + raise SchemaError( + f"{ctx}: components {sum(components):.4f}s do not reconcile with " + f"iteration_total {iteration_total:.4f}s" + ) + rows.append( + BreakdownRow( + mode=str(_require(row, "mode", ctx)), + coordinators=_as_int(_require(row, "coordinators", ctx), ctx), + iteration_total=iteration_total, + components=components, + ) + ) + + semantics = root.get("component_semantics", {}) + semantics_pairs = tuple( + (str(key), str(value)) + for key, value in _mapping( + semantics, "breakdown.component_semantics" + ).items() + ) + return Breakdown(rows=tuple(rows), component_semantics=semantics_pairs) diff --git a/scripts/coordinator_scaling_analysis.py b/scripts/coordinator_scaling_analysis.py new file mode 100644 index 00000000..d94a4876 --- /dev/null +++ b/scripts/coordinator_scaling_analysis.py @@ -0,0 +1,268 @@ +"""Pure analysis functions for coordinator scaling measurements.""" + +from __future__ import annotations + +import json +import statistics +from dataclasses import dataclass +from pathlib import Path +from typing import Mapping, Sequence + +from scripts.coordinator_scaling_conclusions import ( + ScalingPoint, + assess_saturation, +) +from scripts.multi_coordinator_scaling_analysis import ( + BreakdownReconciliationError, + PhaseBreakdown, + reconcile_breakdown, + response_counters_are_balanced, + scaling_efficiency, +) + +__all__ = ( + "BreakdownReconciliationError", + "PhaseBreakdown", + "ScalingPoint", + "assess_saturation", + "reconcile_breakdown", + "scaling_efficiency", +) + + +@dataclass(frozen=True, slots=True) +class RunMeasurement: + iteration_runtime_seconds: float | None + peak_cpu_percent: float | None + median_cpu_percent: float | None + peak_rss_bytes: int | None + peak_nic_tx_bytes_per_sec: float | None + median_nic_tx_bytes_per_sec: float | None + peak_nic_rx_bytes_per_sec: float | None + median_nic_rx_bytes_per_sec: float | None + phase_counts: Mapping[str, int] + phase_durations_seconds: Mapping[str, float] + error: str | None + warnings: tuple[str, ...] = () + + +@dataclass(frozen=True, slots=True) +class ScalingAggregate: + device_count: int + repetitions: int + successful_repetitions: int + iteration_runtime_seconds: float | None + peak_cpu_percent: float | None + median_cpu_percent: float | None + peak_rss_bytes: int | None + peak_nic_tx_bytes_per_sec: float | None + median_nic_tx_bytes_per_sec: float | None + peak_nic_rx_bytes_per_sec: float | None + median_nic_rx_bytes_per_sec: float | None + phase_counts: Mapping[str, float] + phase_durations_seconds: Mapping[str, float] + + +def _empty_measurement( + error: str, warnings: tuple[str, ...] = () +) -> RunMeasurement: + return RunMeasurement( + iteration_runtime_seconds=None, + peak_cpu_percent=None, + median_cpu_percent=None, + peak_rss_bytes=None, + peak_nic_tx_bytes_per_sec=None, + median_nic_tx_bytes_per_sec=None, + peak_nic_rx_bytes_per_sec=None, + median_nic_rx_bytes_per_sec=None, + phase_counts={}, + phase_durations_seconds={}, + error=error, + warnings=warnings, + ) + + +def parse_metrics_jsonl(path: Path) -> RunMeasurement: + numbered_lines = [ + (line_number, line) + for line_number, line in enumerate( + path.read_text(encoding="utf-8").splitlines(), start=1 + ) + if line.strip() + ] + records = [] + warnings: list[str] = [] + for position, (line_number, line) in enumerate(numbered_lines): + try: + records.append(json.loads(line)) + except json.JSONDecodeError as error: + message = ( + f"malformed metrics JSONL at line {line_number}: {error.msg}" + ) + if position != len(numbered_lines) - 1: + return _empty_measurement(message) + warnings.append( + f"trailing partial metrics JSONL record at line {line_number}: " + f"{error.msg}" + ) + terminal_errors = [ + record + for record in records + if record.get("record_type") == "terminal_error" + ] + samples = [ + record for record in records if record.get("record_type") == "sample" + ] + measured = [ + sample + for sample in samples + if sample.get("phase_counts", {}).get("iteration_total", 0) > 0 + ] + if not measured: + if terminal_errors: + terminal = terminal_errors[-1] + return _empty_measurement( + f"{terminal['error_type']}: {terminal['error_message']}", + tuple(warnings), + ) + return _empty_measurement( + "no measured iteration samples", tuple(warnings) + ) + + final = measured[-1] + phase_counts = { + str(name): int(count) for name, count in final["phase_counts"].items() + } + phase_durations = { + str(name): float(duration) + for name, duration in final["phase_durations_seconds"].items() + } + iteration_count = phase_counts.get("iteration_total", 0) + iteration_total = phase_durations.get("iteration_total", 0.0) + error = None + if terminal_errors: + terminal = terminal_errors[-1] + error = f"{terminal['error_type']}: {terminal['error_message']}" + return RunMeasurement( + iteration_runtime_seconds=( + iteration_total / iteration_count if iteration_count > 0 else None + ), + peak_cpu_percent=max( + float(sample["process_cpu_percent"]) for sample in measured + ), + median_cpu_percent=statistics.median( + float(sample["process_cpu_percent"]) for sample in measured + ), + peak_rss_bytes=max(int(sample["rss_bytes"]) for sample in measured), + peak_nic_tx_bytes_per_sec=max( + float(sample["nic_tx_bytes_per_sec"]) for sample in measured + ), + median_nic_tx_bytes_per_sec=statistics.median( + float(sample["nic_tx_bytes_per_sec"]) for sample in measured + ), + peak_nic_rx_bytes_per_sec=max( + float(sample["nic_rx_bytes_per_sec"]) for sample in measured + ), + median_nic_rx_bytes_per_sec=statistics.median( + float(sample["nic_rx_bytes_per_sec"]) for sample in measured + ), + phase_counts=phase_counts, + phase_durations_seconds=phase_durations, + error=error, + warnings=tuple(warnings), + ) + + +def _mean_present(values: Sequence[float | int | None]) -> float | None: + present = [float(value) for value in values if value is not None] + return statistics.mean(present) if present else None + + +def aggregate_repetitions( + device_count: int, + measurements: Sequence[RunMeasurement], + exit_statuses: Sequence[int], +) -> ScalingAggregate: + successful = [ + measurement + for measurement, exit_status in zip(measurements, exit_statuses) + if exit_status == 0 + and measurement.error is None + and measurement.iteration_runtime_seconds is not None + ] + phase_names = { + name for measurement in successful for name in measurement.phase_counts + } + duration_names = { + name + for measurement in successful + for name in measurement.phase_durations_seconds + } + peak_cpu_values = [ + measurement.peak_cpu_percent for measurement in successful + ] + peak_rss_values = [measurement.peak_rss_bytes for measurement in successful] + return ScalingAggregate( + device_count=device_count, + repetitions=len(measurements), + successful_repetitions=len(successful), + iteration_runtime_seconds=_mean_present( + [ + measurement.iteration_runtime_seconds + for measurement in successful + ] + ), + peak_cpu_percent=max(peak_cpu_values) if peak_cpu_values else None, + median_cpu_percent=_mean_present( + [measurement.median_cpu_percent for measurement in successful] + ), + peak_rss_bytes=( + max(value for value in peak_rss_values if value is not None) + if any(value is not None for value in peak_rss_values) + else None + ), + peak_nic_tx_bytes_per_sec=max( + ( + measurement.peak_nic_tx_bytes_per_sec or 0.0 + for measurement in successful + ), + default=0.0, + ) + if successful + else None, + median_nic_tx_bytes_per_sec=_mean_present( + [ + measurement.median_nic_tx_bytes_per_sec + for measurement in successful + ] + ), + peak_nic_rx_bytes_per_sec=max( + ( + measurement.peak_nic_rx_bytes_per_sec or 0.0 + for measurement in successful + ), + default=0.0, + ) + if successful + else None, + median_nic_rx_bytes_per_sec=_mean_present( + [ + measurement.median_nic_rx_bytes_per_sec + for measurement in successful + ] + ), + phase_counts={ + name: statistics.mean( + measurement.phase_counts.get(name, 0) + for measurement in successful + ) + for name in sorted(phase_names) + }, + phase_durations_seconds={ + name: statistics.mean( + measurement.phase_durations_seconds.get(name, 0.0) + for measurement in successful + ) + for name in sorted(duration_names) + }, + ) diff --git a/scripts/coordinator_scaling_conclusions.py b/scripts/coordinator_scaling_conclusions.py new file mode 100644 index 00000000..99db7e14 --- /dev/null +++ b/scripts/coordinator_scaling_conclusions.py @@ -0,0 +1,68 @@ +"""Conclusion rules for measured coordinator scaling points.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Sequence + + +@dataclass(frozen=True, slots=True) +class ScalingPoint: + device_count: int + iteration_runtime_seconds: float + median_cpu_percent: float + peak_cpu_percent: float + + +@dataclass(frozen=True, slots=True) +class SaturationAssessment: + runtime_plateau_observed: bool + resource_saturation_observed: bool + median_cpu_normalized_percent: float + peak_cpu_normalized_percent: float + bottleneck: str + runtime_plateau_criterion: str + resource_saturation_criterion: str + + +def assess_saturation( + points: Sequence[ScalingPoint], logical_cpu_count: int +) -> SaturationAssessment: + runtime_criterion = "largest-device iteration runtime is at least 95% of the preceding point" + resource_criterion = "largest-device median process CPU is at least 90% of logical host CPU capacity" + if len(points) < 2 or logical_cpu_count <= 0: + return SaturationAssessment( + False, + False, + 0.0, + 0.0, + "none within measured range through local emulated devices", + runtime_criterion, + resource_criterion, + ) + previous, current = points[-2:] + median_normalized = current.median_cpu_percent / logical_cpu_count + peak_normalized = current.peak_cpu_percent / logical_cpu_count + runtime_plateau = ( + current.iteration_runtime_seconds / previous.iteration_runtime_seconds + >= 0.95 + ) + resource_saturation = median_normalized >= 90.0 + if runtime_plateau and resource_saturation: + bottleneck = "coordinator CPU" + elif runtime_plateau: + bottleneck = "coordinator-side host work / specific resource unisolated" + else: + bottleneck = ( + "none within measured range through " + f"{current.device_count} local emulated devices" + ) + return SaturationAssessment( + runtime_plateau, + resource_saturation, + median_normalized, + peak_normalized, + bottleneck, + runtime_criterion, + resource_criterion, + ) diff --git a/scripts/coordinator_scaling_execution.py b/scripts/coordinator_scaling_execution.py new file mode 100644 index 00000000..5e5f8110 --- /dev/null +++ b/scripts/coordinator_scaling_execution.py @@ -0,0 +1,190 @@ +"""Subprocess execution for coordinator scaling measurements.""" + +from __future__ import annotations + +import json +import os +import re +import subprocess +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +IMAGE = "device-emulator:latest" +SUBSTRATE = "single-host loopback/emulated devices" + + +@dataclass(frozen=True, slots=True) +class RunSpec: + device_count: int + repetition: int + port: int + run_directory: Path + container_directory: Path + + +def render_proxy_config(template: str, device_count: int, port: int) -> str: + loopback = re.sub( + r"(?m)^listen_ip\s*=\s*\S+\s*$", + "listen_ip = 127.0.0.1", + template, + count=1, + ) + configured = re.sub( + r"(?m)^listen_port\s*=\s*\d+\s*$", + f"listen_port = {port}", + loopback, + count=1, + ) + return re.sub( + r"(?m)^num_device\s*=\s*\d+\s*$", + f"num_device = {device_count}", + configured, + count=1, + ) + + +def build_inner_command(run: RunSpec, metrics_interval: float) -> list[str]: + return [ + "python3", + "scripts/run_single_coordinator_training.py", + "--num_devices", + str(run.device_count), + "--model_name", + "facebook/opt-125m", + "--steps", + "3", + "--warmup_steps", + "1", + "--batch_size", + "1", + "--seq_length", + "32", + "--block_size", + "256", + "--cfg", + str(run.container_directory / "proxy.ini"), + "--metrics_output", + str(run.container_directory / "metrics.jsonl"), + "--metrics_interval", + str(metrics_interval), + "--metrics_nics", + "lo", + ] + + +def docker_container_name(run: RunSpec) -> str: + return f"morphling-d2-d{run.device_count}-r{run.repetition}-p{run.port}" + + +def build_docker_command( + run: RunSpec, output_root: Path, image: str, metrics_interval: float +) -> list[str]: + return [ + "docker", + "run", + "--rm", + "--name", + docker_container_name(run), + "--gpus", + "all", + "--ulimit", + "memlock=-1", + "--ipc", + "host", + "-v", + f"{output_root.resolve()}:/scaling-output", + "-e", + f"CFG_PATH={run.container_directory / 'proxy.ini'}", + image, + *build_inner_command(run, metrics_interval), + ] + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _write_json(path: Path, value: dict[str, object]) -> None: + path.write_text( + json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + +def run_command( + run: RunSpec, + output_root: Path, + image: str, + metrics_interval: float, + timeout_seconds: int, +) -> None: + run.run_directory.mkdir(parents=True, exist_ok=True) + template = Path("config/proxy/svr.ini").read_text(encoding="utf-8") + config_text = render_proxy_config(template, run.device_count, run.port) + (run.run_directory / "proxy.ini").write_text(config_text, encoding="utf-8") + inner = build_inner_command(run, metrics_interval) + inside_docker = Path("/.dockerenv").exists() + if inside_docker: + command = inner + environment = os.environ.copy() + environment["CFG_PATH"] = str(run.container_directory / "proxy.ini") + else: + command = build_docker_command( + run, output_root, image, metrics_interval + ) + environment = None + started = _utc_now() + timed_out = False + try: + completed = subprocess.run( + command, + capture_output=True, + text=True, + env=environment, + timeout=timeout_seconds, + check=False, + ) + exit_status = completed.returncode + stdout = completed.stdout + stderr = completed.stderr + except subprocess.TimeoutExpired as error: + timed_out = True + exit_status = 124 + stdout = error.stdout or "" + stderr = error.stderr or "" + if not inside_docker: + subprocess.run( + ["docker", "rm", "-f", docker_container_name(run)], + capture_output=True, + text=True, + check=False, + ) + (run.run_directory / "stdout.log").write_text(stdout, encoding="utf-8") + (run.run_directory / "stderr.log").write_text(stderr, encoding="utf-8") + _write_json( + run.run_directory / "result.json", + { + "command": command, + "config": { + "device_count": run.device_count, + "repetition": run.repetition, + "port": run.port, + "model": "facebook/opt-125m", + "batch_size": 1, + "sequence_length": 32, + "block_size": 256, + "warmup_optimizer_steps": 1, + "measured_optimizer_steps": 3, + "worker_thread_count": 56, + "metrics_interval_seconds": metrics_interval, + "metrics_nics": ["lo"], + "proxy_config": config_text, + }, + "started_at": started, + "finished_at": _utc_now(), + "exit_status": exit_status, + "timed_out": timed_out, + "substrate": SUBSTRATE, + "real_vs_emulated": "real coordinator process; emulated devices", + }, + ) diff --git a/scripts/export_selected_layouts.py b/scripts/export_selected_layouts.py new file mode 100644 index 00000000..e2fd6de7 --- /dev/null +++ b/scripts/export_selected_layouts.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +"""A3: derive and export the selected DP/PP/TP (or pipeline-stage) layout that +each baseline's planner chose for each principal model/device configuration. + +The layout is derived from the dispatch manifests (the planners' outputs): +pipeline degree = number of distinct stage_ids; the within-stage device count +splits into TP (if a tp_allreduce collective is present, e.g. Alpa), +DP (if a dp_allreduce collective is present, e.g. DTFM/Asteroid), or pure +pipeline (Confidant). WASP dispatches sub-GEMM shards under selective TP across +the whole fleet. Emits results/analytical_scaling/selected_layouts.json and a +LaTeX table body. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +BASELINES = ["cleave", "dtfm", "asteroid", "confident", "alpa"] +LABELS = { + "cleave": "WASP", + "dtfm": "DTFM", + "asteroid": "Asteroid", + "confident": "Confidant", + "alpa": "Alpa", +} + + +def derive_layout(manifest: dict[str, Any]) -> dict[str, Any]: + ents = manifest.get("entries", []) + meta = manifest.get("metadata", {}) + nd = int(meta.get("num_devices", 0)) or len( + {int(e.get("device_id", -1)) for e in ents if int(e.get("device_id", -1)) >= 0} + ) + ptypes = {str(e.get("parallelism_type", e.get("type", ""))) for e in ents} + + if "cleave_tp" in ptypes: + return { + "num_devices": nd, + "granularity": "sub-GEMM shard", + "dp": 1, + "pp": 1, + "tp": nd, + "summary": f"selective sub-GEMM TP across {nd}", + } + + stages = sorted({int(e.get("stage_id", -1)) for e in ents}) + pp = max(1, len(stages)) + dev_per_stage: dict[int, set[int]] = {} + for e in ents: + dev_per_stage.setdefault(int(e.get("stage_id", -1)), set()).add( + int(e.get("device_id", -1)) + ) + per_stage = max((len(v) for v in dev_per_stage.values()), default=nd) + + has_dp = any("dp_allreduce" in p for p in ptypes) + has_tp = any("tp_allreduce" in p for p in ptypes) + + if has_tp: + tp, dp = per_stage, 1 + elif has_dp: + dp, tp = per_stage, 1 + else: + dp, tp = (per_stage, 1) if per_stage > 1 else (1, 1) + + parts = [] + if dp > 1: + parts.append(f"DP{dp}") + if pp > 1: + parts.append(f"PP{pp}") + if tp > 1: + parts.append(f"TP{tp}") + summary = " x ".join(parts) if parts else "PP1" + return { + "num_devices": nd, + "granularity": "pipeline/layer", + "dp": dp, + "pp": pp, + "tp": tp, + "summary": summary, + } + + +def _manifest_path(root: Path, cfg: dict[str, str], baseline: str) -> Path | None: + for pat in cfg["patterns"]: + p = root / pat.format(b=baseline) + if p.exists(): + return p + return None + + +def build_configs(root: Path) -> dict[str, dict[str, str]]: + configs: dict[str, dict[str, str]] = {} + for n, pt in [(64, "000_64"), (128, "001_128"), (256, "002_256"), (512, "003_512"), (1024, "004_1024")]: + configs[f"opt-13b/{n}"] = { + "patterns": [ + f"results/vtime_scaling/num_devices/points/{pt}/planning/manifests/{{b}}_manifest.json" + ] + } + configs["llama2-13b"] = { + "patterns": ["results/vtime_models/llama2-13b/{b}_results/manifests/{b}_manifest.json"] + } + configs["llama2-70b"] = { + "patterns": ["results/vtime_models/llama2-70b/{b}_results/manifests/{b}_manifest.json"] + } + return configs + + +def main() -> int: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--root", type=str, default=".") + ap.add_argument("--output", type=str, default="results/analytical_scaling/selected_layouts.json") + args = ap.parse_args() + root = Path(args.root) + + configs = build_configs(root) + out: dict[str, Any] = {"description": "Selected planner layouts per (config, baseline), derived from dispatch manifests.", "configs": {}} + for cfg_name, cfg in configs.items(): + cfg_out: dict[str, Any] = {} + for b in BASELINES: + mp = _manifest_path(root, cfg, b) + if mp is None: + continue + cfg_out[b] = derive_layout(json.loads(mp.read_text())) + if cfg_out: + out["configs"][cfg_name] = cfg_out + + outp = Path(args.output) + outp.parent.mkdir(parents=True, exist_ok=True) + outp.write_text(json.dumps(out, indent=2)) + print(f"Wrote {outp}\n") + for cfg_name, cfg_out in out["configs"].items(): + row = " ".join(f"{LABELS[b]}={v['summary']}({v['num_devices']}d)" for b, v in cfg_out.items()) + print(f"{cfg_name:16s}: {row}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/generate_trace.py b/scripts/generate_trace.py index a02257a0..4d6bd5f7 100644 --- a/scripts/generate_trace.py +++ b/scripts/generate_trace.py @@ -5,7 +5,9 @@ import argparse import csv import importlib +import json import subprocess +import sys from pathlib import Path import numpy as np @@ -274,7 +276,7 @@ def parse_args() -> argparse.Namespace: return parser.parse_args() -def main() -> None: +def main_ran() -> None: args = parse_args() if args.num_slots <= 0: @@ -355,5 +357,192 @@ def main() -> None: print(f"Wrote: {out_path}") +def _drift_series( + n: int, hurst: float, std: float, rng: np.random.Generator +) -> NDArray[np.float64]: + noise = generate_fgn(n, hurst, rng) + s = float(np.std(noise)) + norm = noise / s if s > 1e-12 else np.zeros_like(noise) + return np.clip(1.0 + std * norm, 0.2, 1.2) + + +def generate_within_batch_traces( + num_devices: int, + num_bins: int, + magnitude: float, + event_rate: float, + duration_frac: float, + network_drift_std: float, + hurst: float, + seed: int, +) -> dict: + """Per-device within-batch time-varying compute and network multipliers. + + The batch is divided into ``num_bins`` equal virtual-time bins. For each + device, ``flops_mult`` models thermal-throttle / foreground-contention + events (Poisson-arrival count ``event_rate``, each of depth ``magnitude`` + and length ``duration_frac`` of the batch); ``ul_bw_mult`` / ``dl_bw_mult`` + model network-rate drift as fractional Gaussian noise (std + ``network_drift_std``) with an added dip while a device is throttled + (contention couples compute and link). Multipliers lie in (0, ~1.2]; the + consumer divides a phase's nominal duration by the multiplier. + """ + rng = np.random.default_rng(seed) + dur_bins = max(1, int(round(duration_frac * num_bins))) + devices: dict[str, dict[str, list[float]]] = {} + + for dev in range(num_devices): + flops = np.ones(num_bins, dtype=np.float64) + n_events = int(rng.poisson(event_rate)) + for _ in range(n_events): + start = int(rng.integers(0, num_bins)) + depth = magnitude * float(rng.uniform(0.7, 1.0)) + mult = max(0.1, 1.0 - depth) + end = min(num_bins, start + dur_bins) + flops[start:end] = np.minimum(flops[start:end], mult) + + ul = _drift_series(num_bins, hurst, network_drift_std, rng) + dl = _drift_series(num_bins, hurst, network_drift_std, rng) + contention = flops < 1.0 + ul[contention] *= 0.7 + dl[contention] *= 0.7 + + devices[str(dev)] = { + "flops_mult": [round(float(x), 6) for x in flops], + "ul_bw_mult": [round(float(x), 6) for x in ul], + "dl_bw_mult": [round(float(x), 6) for x in dl], + } + + return { + "schema": "within_batch_trace_v1", + "num_bins": num_bins, + "meta": { + "num_devices": num_devices, + "magnitude": magnitude, + "event_rate": event_rate, + "duration_frac": duration_frac, + "network_drift_std": network_drift_std, + "hurst": hurst, + "seed": seed, + }, + "devices": devices, + } + + +def validate_within_batch(trace: dict) -> dict: + meta = trace["meta"] + devices = trace["devices"] + flops = np.array([devices[k]["flops_mult"] for k in devices]) + ul = np.array([devices[k]["ul_bw_mult"] for k in devices]) + + throttled_fraction = float(np.mean(flops < 1.0)) + expected_fraction = float( + 1.0 - np.exp(-float(meta["event_rate"]) * float(meta["duration_frac"])) + ) + min_flops_mult = float(flops.min()) + network_std = float(np.std(ul)) + + stats = { + "throttled_bin_fraction": throttled_fraction, + "expected_throttled_fraction": expected_fraction, + "min_flops_mult": min_flops_mult, + "network_mult_std": network_std, + "requested_network_std": float(meta["network_drift_std"]), + } + + if float(meta["magnitude"]) > 0 and float(meta["event_rate"]) > 0: + assert min_flops_mult <= 1.0 - 0.5 * float(meta["magnitude"]), ( + f"throttle depth too shallow: min_mult={min_flops_mult}" + ) + tol = max(0.15, 0.6 * expected_fraction) + assert abs(throttled_fraction - expected_fraction) <= tol, ( + f"event coverage off: got {throttled_fraction}, " + f"expected ~{expected_fraction}" + ) + if float(meta["network_drift_std"]) > 0: + assert network_std >= 0.3 * float(meta["network_drift_std"]), ( + f"network drift too small: std={network_std}" + ) + return stats + + +def parse_within_batch_args(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser( + prog="generate_trace.py within-batch", + description="Generate per-device within-batch time-varying " + "compute/network multiplier traces for the C2 experiment", + ) + parser.add_argument("--num-devices", type=int, default=256) + parser.add_argument("--num-bins", type=int, default=100) + parser.add_argument("--magnitude", type=float, default=0.5) + parser.add_argument("--event-rate", type=float, default=3.0) + parser.add_argument("--duration-frac", type=float, default=0.15) + parser.add_argument("--network-drift-std", type=float, default=0.15) + parser.add_argument("--hurst", type=float, default=0.85) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "-o", + "--output", + default=None, + help="Output JSON path (default under results/within_batch/traces/)", + ) + return parser.parse_args(argv) + + +def main_within_batch(argv: list[str]) -> None: + args = parse_within_batch_args(argv) + + if args.num_devices <= 0 or args.num_bins <= 0: + raise ValueError("--num-devices and --num-bins must be > 0") + if not (0.0 <= args.magnitude <= 1.0): + raise ValueError("--magnitude must be in [0, 1]") + if not (0.0 < args.hurst < 1.0): + raise ValueError("--hurst must be in (0, 1)") + if not (0.0 <= args.duration_frac <= 1.0): + raise ValueError("--duration-frac must be in [0, 1]") + + trace = generate_within_batch_traces( + num_devices=args.num_devices, + num_bins=args.num_bins, + magnitude=args.magnitude, + event_rate=args.event_rate, + duration_frac=args.duration_frac, + network_drift_std=args.network_drift_std, + hurst=args.hurst, + seed=args.seed, + ) + stats = validate_within_batch(trace) + + if args.output: + out_path = Path(args.output) + else: + out_path = Path( + "results/within_batch/traces/" + f"within_batch_seed{args.seed}_mag{args.magnitude}.json" + ) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w") as f: + json.dump(trace, f) + + print(f"Devices: {args.num_devices}, bins: {args.num_bins}") + print( + f"Throttled bin fraction: {stats['throttled_bin_fraction']:.3f} " + f"(expected ~{stats['expected_throttled_fraction']:.3f})" + ) + print(f"Min flops mult: {stats['min_flops_mult']:.3f}") + print( + f"Network mult std: {stats['network_mult_std']:.3f} " + f"(requested {stats['requested_network_std']:.3f})" + ) + print(f"Wrote: {out_path}") + + +def main() -> None: + if len(sys.argv) > 1 and sys.argv[1] == "within-batch": + main_within_batch(sys.argv[2:]) + else: + main_ran() + + if __name__ == "__main__": main() diff --git a/scripts/multi_coordinator_rank_process.py b/scripts/multi_coordinator_rank_process.py new file mode 100644 index 00000000..ac7dcfad --- /dev/null +++ b/scripts/multi_coordinator_rank_process.py @@ -0,0 +1,57 @@ +"""Rank process-group lifecycle for the scaling launcher.""" + +from __future__ import annotations + +import os +import signal +import subprocess +from typing import Sequence + + +class RankProcess: + def __init__(self, process: subprocess.Popen[bytes]) -> None: + self._process = process + + def wait(self, timeout: float | None = None) -> int: + return self._process.wait(timeout=timeout) + + def terminate_group(self) -> None: + try: + os.killpg(self._process.pid, signal.SIGTERM) + except ProcessLookupError: + return + + def kill_group(self) -> None: + try: + os.killpg(self._process.pid, signal.SIGKILL) + except ProcessLookupError: + return + + +def cleanup_rank_processes(processes: Sequence[RankProcess]) -> None: + for process in processes: + process.terminate_group() + for process in processes: + try: + process.wait(timeout=5.0) + except subprocess.TimeoutExpired: + pass + for process in processes: + process.kill_group() + process.wait() + + +def wait_for_rank_processes( + processes: Sequence[RankProcess], timeout_seconds: int +) -> tuple[int, ...]: + try: + statuses = tuple( + process.wait(timeout=timeout_seconds) for process in processes + ) + except subprocess.TimeoutExpired: + cleanup_rank_processes(processes) + raise + if any(status != 0 for status in statuses): + cleanup_rank_processes(processes) + raise subprocess.CalledProcessError(max(statuses), "rank workers") + return statuses diff --git a/scripts/multi_coordinator_scaling_analysis.py b/scripts/multi_coordinator_scaling_analysis.py new file mode 100644 index 00000000..4876076c --- /dev/null +++ b/scripts/multi_coordinator_scaling_analysis.py @@ -0,0 +1,94 @@ +"""Pure D5-D7 scaling and phase-breakdown calculations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Sequence + +try: # Python 3.11+ + from typing import assert_never +except ImportError: # Python 3.10 + from typing_extensions import assert_never + +ScalingMode = Literal["strong", "weak"] +ResponseCounterRecord = tuple[int, int, int, int] + + +@dataclass(frozen=True, slots=True) +class BreakdownReconciliationError(ValueError): + residual_seconds: float + + def __str__(self) -> str: + return f"breakdown components exceed iteration total by {-self.residual_seconds:.9f}s" + + +@dataclass(frozen=True, slots=True) +class PhaseBreakdown: + iteration_total_seconds: float + device_dispatch_aggregation_seconds: float + gradient_sync_seconds: float + optimizer_seconds: float + idle_other_seconds: float + reconciles: bool + + +def scaling_efficiency( + *, + mode: ScalingMode, + baseline_seconds: float, + measured_seconds: float, + coordinators: int, +) -> float: + """Compute conventional strong- or weak-scaling efficiency.""" + match mode: + case "strong": + return baseline_seconds / (coordinators * measured_seconds) + case "weak": + return baseline_seconds / measured_seconds + case unreachable: + assert_never(unreachable) + + +def response_counters_are_balanced( + records: Sequence[ResponseCounterRecord], +) -> bool: + grouped: dict[int, list[tuple[int, int, int]]] = {} + for operation_id, count, previous, current in records: + if count != 1 or previous != current + 1 or current < 0: + return False + grouped.setdefault(operation_id, []).append((count, previous, current)) + if not grouped: + return False + return all( + sorted(previous for _, previous, _ in operation_records) + == list(range(1, len(operation_records) + 1)) + for operation_records in grouped.values() + ) + + +def reconcile_breakdown( + *, + iteration_total_seconds: float, + forward_device_seconds: float, + backward_device_seconds: float, + gradient_sync_seconds: float, + optimizer_seconds: float, + tolerance_seconds: float = 1e-9, +) -> PhaseBreakdown: + """Combine device phases and assign uninstrumented time to residual idle.""" + device_seconds = forward_device_seconds + backward_device_seconds + accounted = device_seconds + gradient_sync_seconds + optimizer_seconds + residual = iteration_total_seconds - accounted + if residual < -tolerance_seconds: + raise BreakdownReconciliationError(residual) + idle_other = max(0.0, residual) + reconciled_total = accounted + idle_other + return PhaseBreakdown( + iteration_total_seconds=iteration_total_seconds, + device_dispatch_aggregation_seconds=device_seconds, + gradient_sync_seconds=gradient_sync_seconds, + optimizer_seconds=optimizer_seconds, + idle_other_seconds=idle_other, + reconciles=abs(reconciled_total - iteration_total_seconds) + <= tolerance_seconds, + ) diff --git a/scripts/multi_coordinator_scaling_cli.py b/scripts/multi_coordinator_scaling_cli.py new file mode 100644 index 00000000..320d004b --- /dev/null +++ b/scripts/multi_coordinator_scaling_cli.py @@ -0,0 +1,63 @@ +"""CLI boundary for the multi-coordinator scaling experiment.""" + +from __future__ import annotations + +import argparse +import dataclasses +from dataclasses import dataclass +from pathlib import Path + +from scripts.multi_coordinator_scaling_config import ( + OPT_125M_WORKLOAD, + TINY_WORKLOAD, + ScalingMode, + WorkloadConfig, +) + + +@dataclass(frozen=True, slots=True) +class ScalingCliConfig: + modes: tuple[ScalingMode, ...] + output_directory: Path + timeout_seconds: int + dry_run: bool + workload: WorkloadConfig + + +def parse_scaling_cli() -> ScalingCliConfig: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--modes", + nargs="+", + choices=("strong", "weak"), + default=("strong", "weak"), + ) + parser.add_argument( + "--output-dir", type=Path, default=Path("results/coord_scaling") + ) + parser.add_argument("--timeout-seconds", type=int, default=1800) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--tiny", action="store_true") + parser.add_argument("--model-name", default="facebook/opt-125m") + parser.add_argument("--sequence-length", type=int, default=32) + parser.add_argument("--measured-iterations", type=int, default=2) + parser.add_argument("--block-size", type=int, default=256) + args = parser.parse_args() + workload = ( + TINY_WORKLOAD + if args.tiny + else dataclasses.replace( + OPT_125M_WORKLOAD, + model_name=args.model_name, + sequence_length=args.sequence_length, + measured_iterations=args.measured_iterations, + block_size=args.block_size, + ) + ) + return ScalingCliConfig( + modes=tuple(args.modes), + output_directory=args.output_dir, + timeout_seconds=args.timeout_seconds, + dry_run=args.dry_run, + workload=workload, + ) diff --git a/scripts/multi_coordinator_scaling_conclusions.py b/scripts/multi_coordinator_scaling_conclusions.py new file mode 100644 index 00000000..7a4fd884 --- /dev/null +++ b/scripts/multi_coordinator_scaling_conclusions.py @@ -0,0 +1,46 @@ +"""Evidence-derived conclusions for D5 and D6 scaling points.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from scripts.multi_coordinator_scaling_analysis import scaling_efficiency +from scripts.multi_coordinator_scaling_config import ScalingConfig +from scripts.multi_coordinator_scaling_results import GlobalResult + + +@dataclass(frozen=True, slots=True) +class ScalingConclusion: + mode: str + speedup: float + efficiency: float + throughput_ratio: float + correctness_passed: bool + + +def build_scaling_conclusion( + baseline_config: ScalingConfig, + scaled_config: ScalingConfig, + baseline: GlobalResult, + scaled: GlobalResult, +) -> ScalingConclusion: + """Compare two measured points without inferring beyond the observations.""" + return ScalingConclusion( + mode=baseline_config.mode, + speedup=( + baseline.iteration_total_seconds / scaled.iteration_total_seconds + ), + efficiency=scaling_efficiency( + mode=baseline_config.mode, + baseline_seconds=baseline.iteration_total_seconds, + measured_seconds=scaled.iteration_total_seconds, + coordinators=scaled_config.coordinators, + ), + throughput_ratio=( + scaled.throughput_samples_per_second + / baseline.throughput_samples_per_second + ), + correctness_passed=( + baseline.loss_correctness.passed and scaled.loss_correctness.passed + ), + ) diff --git a/scripts/multi_coordinator_scaling_config.py b/scripts/multi_coordinator_scaling_config.py new file mode 100644 index 00000000..d71f8bf8 --- /dev/null +++ b/scripts/multi_coordinator_scaling_config.py @@ -0,0 +1,207 @@ +"""Configuration and rank launch commands for multi-coordinator scaling.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Final, Literal, Mapping + +try: # Python 3.11+ + from typing import assert_never +except ImportError: # Python 3.10 + from typing_extensions import assert_never + +ScalingMode = Literal["strong", "weak"] +AffinityMode = Literal["numactl", "taskset"] +IterationKind = Literal["warmup", "measured"] + +NUMA_CPU_RANGES: Final = ("0-27", "28-55") +WARMUP_ITERATIONS: Final = 1 +MEASURED_ITERATIONS: Final = 3 + + +@dataclass(frozen=True, slots=True) +class WorkloadConfig: + tiny: bool + model_name: str + sequence_length: int + measured_iterations: int + block_size: int + learning_rate: float + + +TINY_WORKLOAD: Final = WorkloadConfig( + tiny=True, + model_name="tiny-mlp", + sequence_length=1, + measured_iterations=3, + block_size=128, + learning_rate=0.01, +) +OPT_125M_WORKLOAD: Final = WorkloadConfig( + tiny=False, + model_name="facebook/opt-125m", + sequence_length=32, + measured_iterations=2, + block_size=256, + learning_rate=5e-5, +) + + +@dataclass(frozen=True, slots=True) +class ScalingConfig: + mode: ScalingMode + coordinators: int + total_devices: int + global_batch: int + devices_per_coordinator: int + local_batch: int + warmup_iterations: int = WARMUP_ITERATIONS + measured_iterations: int = MEASURED_ITERATIONS + affinity_mode: AffinityMode = "numactl" + relative_tolerance: float = 0.05 + tiny: bool = True + model_name: str = "tiny-mlp" + sequence_length: int = 1 + block_size: int = 128 + learning_rate: float = 0.01 + distributed_timeout_seconds: int = 120 + + +@dataclass(frozen=True, slots=True) +class RankLaunch: + rank: int + numa_node: int + cpu_range: str + command: tuple[str, ...] + environment: Mapping[str, str] + substrate: str + + +def build_scaling_configs( + mode: ScalingMode, + workload: WorkloadConfig = TINY_WORKLOAD, +) -> tuple[ScalingConfig, ...]: + """Return the fixed D5 strong or D6 weak scaling matrix.""" + if workload.tiny: + strong_points = ((1, 8, 16, 8, 16), (2, 8, 16, 4, 8)) + weak_points = ((1, 4, 8, 4, 8), (2, 8, 16, 4, 8)) + else: + strong_points = ((1, 8, 2, 8, 2), (2, 8, 2, 4, 1)) + weak_points = ((1, 4, 1, 4, 1), (2, 8, 2, 4, 1)) + match mode: + case "strong": + points = strong_points + case "weak": + points = weak_points + case unreachable: + assert_never(unreachable) + return tuple( + ScalingConfig( + mode=mode, + coordinators=coordinators, + total_devices=total_devices, + global_batch=global_batch, + devices_per_coordinator=devices_per_coordinator, + local_batch=local_batch, + measured_iterations=workload.measured_iterations, + tiny=workload.tiny, + model_name=workload.model_name, + sequence_length=workload.sequence_length, + block_size=workload.block_size, + learning_rate=workload.learning_rate, + ) + for ( + coordinators, + total_devices, + global_batch, + devices_per_coordinator, + local_batch, + ) in points + ) + + +def select_affinity_mode( + *, numactl_available: bool, memory_binding_available: bool +) -> AffinityMode: + """Prefer NUMA memory binding and otherwise use CPU first-touch placement.""" + if numactl_available and memory_binding_available: + return "numactl" + return "taskset" + + +def build_iteration_plan(config: ScalingConfig) -> tuple[IterationKind, ...]: + return ("warmup",) * config.warmup_iterations + ( + "measured", + ) * config.measured_iterations + + +def build_rank_launch( + config: ScalingConfig, + rank: int, + output_directory: Path, +) -> RankLaunch: + node = rank + cpu_range = NUMA_CPU_RANGES[node] + common = ( + "python3", + "scripts/run_multi_coordinator.py", + "--worker", + "--coords", + str(config.coordinators), + "--devices_per_coord", + str(config.devices_per_coordinator), + "--local_batch", + str(config.local_batch), + "--warmup_steps", + str(config.warmup_iterations), + "--steps", + str(config.measured_iterations), + "--rtol", + str(config.relative_tolerance), + "--model-name", + config.model_name, + "--sequence-length", + str(config.sequence_length), + "--block_size", + str(config.block_size), + "--lr", + str(config.learning_rate), + "--distributed-timeout-seconds", + str(config.distributed_timeout_seconds), + "--metrics_output", + str(output_directory / "metrics.jsonl"), + "--result_output", + str(output_directory / f"rank-{rank}.json"), + "--metrics_nics", + "lo", + ) + match config.affinity_mode: + case "numactl": + prefix = ( + "numactl", + f"--cpunodebind={node}", + f"--membind={node}", + ) + substrate = "same-host Gloo loopback lower bound; NUMA CPU and memory binding" + case "taskset": + prefix = ("taskset", "--cpu-list", cpu_range) + substrate = ( + "same-host Gloo loopback lower bound; CPU affinity with " + "first-touch memory" + ) + case unreachable: + assert_never(unreachable) + tiny_flag = ("--tiny",) if config.tiny else () + return RankLaunch( + rank=rank, + numa_node=node, + cpu_range=cpu_range, + command=prefix + common + tiny_flag, + environment={ + "RANK": str(rank), + "WORLD_SIZE": str(config.coordinators), + "MASTER_ADDR": "127.0.0.1", + }, + substrate=substrate, + ) diff --git a/scripts/multi_coordinator_scaling_execution.py b/scripts/multi_coordinator_scaling_execution.py new file mode 100644 index 00000000..f365687c --- /dev/null +++ b/scripts/multi_coordinator_scaling_execution.py @@ -0,0 +1,63 @@ +"""Execution plan values for D5-D7 multi-coordinator runs.""" + +from __future__ import annotations + +import dataclasses +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence + +from scripts.multi_coordinator_scaling_config import ( + TINY_WORKLOAD, + AffinityMode, + ScalingConfig, + ScalingMode, + WorkloadConfig, + build_scaling_configs, +) + + +@dataclass(frozen=True, slots=True) +class ExperimentRun: + config: ScalingConfig + run_directory: Path + master_port: int + base_proxy_port: int + + +def build_experiment_plan( + *, + modes: Sequence[ScalingMode], + output_root: Path, + workload: WorkloadConfig = TINY_WORKLOAD, +) -> tuple[ExperimentRun, ...]: + runs: list[ExperimentRun] = [] + index = 0 + for mode in modes: + for config in build_scaling_configs(mode, workload): + runs.append( + ExperimentRun( + config=config, + run_directory=( + output_root + / mode + / f"coordinators-{config.coordinators}" + ), + master_port=29600 + index, + base_proxy_port=39200 + index * 10, + ) + ) + index += 1 + return tuple(runs) + + +def apply_affinity_mode( + runs: Sequence[ExperimentRun], affinity_mode: AffinityMode +) -> tuple[ExperimentRun, ...]: + return tuple( + dataclasses.replace( + run, + config=dataclasses.replace(run.config, affinity_mode=affinity_mode), + ) + for run in runs + ) diff --git a/scripts/multi_coordinator_scaling_results.py b/scripts/multi_coordinator_scaling_results.py new file mode 100644 index 00000000..25c45073 --- /dev/null +++ b/scripts/multi_coordinator_scaling_results.py @@ -0,0 +1,231 @@ +"""Typed rank and global results for multi-coordinator measurements.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Mapping, TypedDict + +try: # Python 3.11+ + from typing import assert_never +except ImportError: # Python 3.10 + from typing_extensions import assert_never + +from scripts.multi_coordinator_scaling_analysis import reconcile_breakdown +from scripts.multi_coordinator_scaling_config import ScalingConfig + + +@dataclass(frozen=True, slots=True) +class RankMeasurement: + rank: int + phase_durations: Mapping[str, float] + measured_losses: tuple[float, ...] + golden_losses: tuple[float, ...] + + +@dataclass(frozen=True, slots=True) +class LossCorrectness: + passed: bool + decreasing: bool + tracks_golden: bool + max_relative_error: float + + +@dataclass(frozen=True, slots=True) +class RankResult: + rank: int + warmup_iteration_seconds: float + iteration_total_seconds: float + device_dispatch_aggregation_seconds: float + gradient_sync_seconds: float + optimizer_seconds: float + idle_other_seconds: float + throughput_samples_per_second: float + measured_losses: tuple[float, ...] + golden_losses: tuple[float, ...] + loss_correctness: LossCorrectness + + +@dataclass(frozen=True, slots=True) +class GlobalResult: + coordinators: int + iteration_total_seconds: float + throughput_samples_per_second: float + loss_correctness: LossCorrectness + substrate: str + ranks: tuple[RankResult, ...] + + +class LossCorrectnessPayload(TypedDict): + passed: bool + decreasing: bool + tracks_golden: bool + max_relative_error: float + + +class PhasePayload(TypedDict): + iteration_total: float + device_dispatch_aggregation: float + gradient_sync: float + optimizer: float + idle_other: float + + +class RankResultPayload(TypedDict): + schema_version: int + rank: int + coordinators: int + warmup_iterations: int + measured_iterations: int + local_batch: int + devices_per_coordinator: int + relative_tolerance: float + model_name: str + sequence_length: int + warmup_iteration_seconds: float + throughput_samples_per_second: float + measured_losses: tuple[float, ...] + golden_losses: tuple[float, ...] + loss_correctness: LossCorrectnessPayload + phases_seconds: PhasePayload + + +def _loss_correctness( + measured: tuple[float, ...], + golden: tuple[float, ...], + relative_tolerance: float, +) -> LossCorrectness: + finite = all(math.isfinite(loss) for loss in measured) + decreasing = finite and len(measured) > 1 and measured[-1] < measured[0] + relative_errors = tuple( + abs(actual - expected) / max(abs(expected), 1e-6) + for actual, expected in zip(measured, golden) + ) + max_relative_error = max(relative_errors, default=math.inf) + tracks = ( + finite + and len(measured) == len(golden) + and max_relative_error <= relative_tolerance + ) + return LossCorrectness( + passed=tracks, + decreasing=decreasing, + tracks_golden=tracks, + max_relative_error=max_relative_error, + ) + + +def build_rank_result( + config: ScalingConfig, measurement: RankMeasurement +) -> RankResult: + phases = measurement.phase_durations + iteration_total = phases.get("iteration_total", 0.0) + breakdown = reconcile_breakdown( + iteration_total_seconds=iteration_total, + forward_device_seconds=phases.get( + "forward_device_dispatch_aggregation", 0.0 + ), + backward_device_seconds=phases.get( + "backward_device_dispatch_aggregation", 0.0 + ), + gradient_sync_seconds=phases.get("gradient_sync", 0.0), + optimizer_seconds=phases.get("optimizer", 0.0), + tolerance_seconds=max(1e-9, iteration_total * 1e-6), + ) + throughput = ( + config.local_batch * config.measured_iterations / iteration_total + if iteration_total > 0 + else 0.0 + ) + return RankResult( + rank=measurement.rank, + warmup_iteration_seconds=phases.get("warmup_iteration", 0.0), + iteration_total_seconds=iteration_total, + device_dispatch_aggregation_seconds=( + breakdown.device_dispatch_aggregation_seconds + ), + gradient_sync_seconds=breakdown.gradient_sync_seconds, + optimizer_seconds=breakdown.optimizer_seconds, + idle_other_seconds=breakdown.idle_other_seconds, + throughput_samples_per_second=throughput, + measured_losses=measurement.measured_losses, + golden_losses=measurement.golden_losses, + loss_correctness=_loss_correctness( + measurement.measured_losses, + measurement.golden_losses, + config.relative_tolerance, + ), + ) + + +def build_global_result( + config: ScalingConfig, ranks: tuple[RankResult, ...] +) -> GlobalResult: + makespan = max(rank.iteration_total_seconds for rank in ranks) + throughput = ( + config.global_batch * config.measured_iterations / makespan + if makespan > 0 + else 0.0 + ) + correctness = LossCorrectness( + passed=all(rank.loss_correctness.passed for rank in ranks), + decreasing=all(rank.loss_correctness.decreasing for rank in ranks), + tracks_golden=all( + rank.loss_correctness.tracks_golden for rank in ranks + ), + max_relative_error=max( + rank.loss_correctness.max_relative_error for rank in ranks + ), + ) + match config.affinity_mode: + case "numactl": + placement = "NUMA CPU and memory binding" + case "taskset": + placement = "CPU affinity with first-touch memory" + case unreachable: + assert_never(unreachable) + return GlobalResult( + coordinators=config.coordinators, + iteration_total_seconds=makespan, + throughput_samples_per_second=throughput, + loss_correctness=correctness, + substrate=f"same-host Gloo loopback lower bound; {placement}; " + "no physical NIC-to-NIC routing", + ranks=ranks, + ) + + +def rank_result_payload( + config: ScalingConfig, result: RankResult +) -> RankResultPayload: + return { + "schema_version": 1, + "rank": result.rank, + "coordinators": config.coordinators, + "warmup_iterations": config.warmup_iterations, + "measured_iterations": config.measured_iterations, + "local_batch": config.local_batch, + "devices_per_coordinator": config.devices_per_coordinator, + "relative_tolerance": config.relative_tolerance, + "model_name": config.model_name, + "sequence_length": config.sequence_length, + "warmup_iteration_seconds": result.warmup_iteration_seconds, + "throughput_samples_per_second": result.throughput_samples_per_second, + "measured_losses": result.measured_losses, + "golden_losses": result.golden_losses, + "loss_correctness": { + "passed": result.loss_correctness.passed, + "decreasing": result.loss_correctness.decreasing, + "tracks_golden": result.loss_correctness.tracks_golden, + "max_relative_error": result.loss_correctness.max_relative_error, + }, + "phases_seconds": { + "iteration_total": result.iteration_total_seconds, + "device_dispatch_aggregation": ( + result.device_dispatch_aggregation_seconds + ), + "gradient_sync": result.gradient_sync_seconds, + "optimizer": result.optimizer_seconds, + "idle_other": result.idle_other_seconds, + }, + } diff --git a/scripts/run_coordinator_scaling.py b/scripts/run_coordinator_scaling.py new file mode 100644 index 00000000..9931020a --- /dev/null +++ b/scripts/run_coordinator_scaling.py @@ -0,0 +1,239 @@ +#!/usr/bin/env python3 +"""Run and analyze the D2 single-coordinator scaling sweep.""" + +from __future__ import annotations + +import argparse +import dataclasses +import json +import os +import sys +from datetime import datetime, timezone +from pathlib import Path +from typing import Sequence + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +BASE_PORT = 39100 +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from scripts.coordinator_scaling_analysis import ( + RunMeasurement, + aggregate_repetitions, + parse_metrics_jsonl, +) +from scripts.coordinator_scaling_conclusions import ( + ScalingPoint, + assess_saturation, +) +from scripts.coordinator_scaling_execution import ( + IMAGE, + SUBSTRATE, + RunSpec, + build_inner_command, + run_command, +) + + +class PortAllocationError(ValueError): + pass + + +def port_for_run(run_index: int) -> int: + port = BASE_PORT + run_index + if run_index < 0 or port > 65535: + raise PortAllocationError(f"run index {run_index} exceeds port range") + return port + + +def _failed_measurement(error: str) -> RunMeasurement: + return RunMeasurement( + None, None, None, None, None, None, None, None, {}, {}, error + ) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _write_json(path: Path, value: dict[str, object]) -> None: + path.write_text( + json.dumps(value, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + +def _load_run( + run_directory: Path, +) -> tuple[RunMeasurement, int, dict[str, object]]: + result_path = run_directory / "result.json" + if not result_path.exists(): + error = "run result is unavailable" + return _failed_measurement(error), 125, {"error": error} + result = json.loads(result_path.read_text(encoding="utf-8")) + metrics_path = run_directory / "metrics.jsonl" + measurement = ( + parse_metrics_jsonl(metrics_path) + if metrics_path.exists() + else _failed_measurement("metrics JSONL is unavailable") + ) + return measurement, int(result["exit_status"]), result + + +def _write_summary( + output_root: Path, + device_counts: Sequence[int], + repetitions: int, + image: str, + invocation: Sequence[str], +) -> Path: + logical_cpu_count = os.cpu_count() or 1 + rows: list[dict[str, object]] = [] + scaling_points: list[ScalingPoint] = [] + for device_count in device_counts: + loaded = [ + _load_run( + output_root / f"devices-{device_count}" / f"rep-{repetition}" + ) + for repetition in range(1, repetitions + 1) + ] + measurements = tuple(item[0] for item in loaded) + statuses = tuple(item[1] for item in loaded) + aggregate = aggregate_repetitions(device_count, measurements, statuses) + if ( + aggregate.iteration_runtime_seconds is not None + and aggregate.median_cpu_percent is not None + and aggregate.peak_cpu_percent is not None + ): + scaling_points.append( + ScalingPoint( + device_count, + aggregate.iteration_runtime_seconds, + aggregate.median_cpu_percent, + aggregate.peak_cpu_percent, + ) + ) + rows.append( + { + **dataclasses.asdict(aggregate), + "runs": [ + { + "exit_status": status, + "measurement": dataclasses.asdict(measurement), + "evidence": evidence, + } + for measurement, status, evidence in loaded + ], + "substrate": SUBSTRATE, + } + ) + saturation = assess_saturation(scaling_points, logical_cpu_count) + summary = { + "schema_version": 1, + "experiment": "D2 measured single-coordinator operating envelope", + "timestamp": _utc_now(), + "invocation": list(invocation), + "configuration": { + "image": image, + "device_counts": list(device_counts), + "repetitions": repetitions, + "model": "facebook/opt-125m", + "batch_size": 1, + "sequence_length": 32, + "block_size": 256, + "warmup_optimizer_steps": 1, + "measured_optimizer_steps": 3, + "worker_thread_count": 56, + "metrics_nic": "lo", + "aggregate_phase_semantics": ( + "per-repetition mean across successful repetitions; " + "raw per-run totals retained in rows[].runs" + ), + }, + "environment": { + "real_vs_emulated": "real coordinator process; emulated devices", + "host_nic_substrate": SUBSTRATE, + "logical_cpu_count": logical_cpu_count, + "worker_thread_count": 56, + "coordinator_cpu_scope": ( + "coordinator process CPU includes local non-GEMM model work " + "and optimizer execution" + ), + }, + "saturation": dataclasses.asdict(saturation), + "observed_bottleneck": saturation.bottleneck, + "limitations": [ + "queue occupancy is unavailable in the current backend", + "worker thread count is fixed at 56", + "coordinator process CPU includes local non-GEMM model work and optimizer execution", + "single host and loopback networking do not represent a physical edge network", + "results cover the measured range through 8 local emulated devices and do not establish a host real-connection limit", + "one or few repetitions bound repeatability and do not establish tail distributions", + ], + "rows": rows, + } + summary_path = output_root / "summary.json" + _write_json(summary_path, summary) + return summary_path + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--device-counts", nargs="+", type=int, default=[1, 2, 4, 8] + ) + parser.add_argument("--repetitions", type=int, default=1) + parser.add_argument( + "--output-dir", type=Path, default=Path("results/coord_scaling/single") + ) + parser.add_argument("--image", default=IMAGE) + parser.add_argument("--metrics-interval", type=float, default=0.2) + parser.add_argument("--timeout-seconds", type=int, default=1800) + parser.add_argument("--dry-run", action="store_true") + parser.add_argument("--analyze-only", action="store_true") + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + inside_docker = Path("/.dockerenv").exists() + if not args.analyze_only: + run_index = 0 + for device_count in args.device_counts: + for repetition in range(1, args.repetitions + 1): + relative = Path(f"devices-{device_count}") / f"rep-{repetition}" + run = RunSpec( + device_count, + repetition, + port_for_run(run_index), + args.output_dir / relative, + (args.output_dir / relative) + if inside_docker + else Path("/scaling-output") / relative, + ) + if args.dry_run: + print( + " ".join( + build_inner_command(run, args.metrics_interval) + ) + ) + else: + run_command( + run, + args.output_dir, + args.image, + args.metrics_interval, + args.timeout_seconds, + ) + run_index += 1 + if args.dry_run: + return 0 + summary_path = _write_summary( + args.output_dir, + args.device_counts, + args.repetitions, + args.image, + sys.argv, + ) + print(summary_path) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_multi_coordinator.py b/scripts/run_multi_coordinator.py new file mode 100644 index 00000000..0264b956 --- /dev/null +++ b/scripts/run_multi_coordinator.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +"""Run correct data-parallel training across one or two coordinators.""" + +from __future__ import annotations + +import argparse +import os + +import torch.multiprocessing as mp + +from morphling.runtime.coordinator_metrics import add_metrics_arguments +from scripts._multi_coordinator_training import worker_boundary +from scripts.multi_coordinator_scaling_conclusions import ( + ScalingConclusion, + build_scaling_conclusion, +) +from scripts.multi_coordinator_scaling_config import ( + OPT_125M_WORKLOAD, + TINY_WORKLOAD, + RankLaunch, + ScalingConfig, + WorkloadConfig, + build_iteration_plan, + build_rank_launch, + build_scaling_configs, + select_affinity_mode, +) +from scripts.multi_coordinator_scaling_execution import ( + ExperimentRun, + apply_affinity_mode, + build_experiment_plan, +) +from scripts.multi_coordinator_scaling_results import ( + GlobalResult, + LossCorrectness, + RankMeasurement, + RankResult, + build_global_result, + build_rank_result, + rank_result_payload, +) + + +def _parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--worker", action="store_true") + parser.add_argument( + "--scaling_mode", choices=("strong", "weak"), default="strong" + ) + parser.add_argument("--coords", type=int, choices=(1, 2), default=2) + parser.add_argument("--devices_per_coord", type=int, default=1) + parser.add_argument("--steps", type=int, default=3) + parser.add_argument("--warmup_steps", type=int, default=1) + parser.add_argument("--local_batch", type=int, default=8) + parser.add_argument("--tiny", action="store_true") + parser.add_argument("--model-name", default="facebook/opt-125m") + parser.add_argument("--sequence-length", type=int, default=32) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--lr", type=float, default=0.01) + parser.add_argument("--block_size", type=int, default=128) + parser.add_argument("--cfg", type=str, default=None) + parser.add_argument("--base_proxy_port", type=int, default=39000) + parser.add_argument("--master_port", type=int, default=29500) + parser.add_argument("--distributed-timeout-seconds", type=int, default=120) + parser.add_argument("--rtol", type=float, default=0.05) + parser.add_argument("--result_output", type=str, default=None) + add_metrics_arguments(parser) + return parser + + +def _spawned_worker(rank: int, args: argparse.Namespace) -> None: + worker_boundary(rank, args) + + +def main() -> int: + args = _parser().parse_args() + if args.worker: + rank = int(os.environ["RANK"]) + worker_boundary(rank, args) + return 0 + + mp.spawn(_spawned_worker, args=(args,), nprocs=args.coords, join=True) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_multi_coordinator_scaling.py b/scripts/run_multi_coordinator_scaling.py new file mode 100644 index 00000000..c9559984 --- /dev/null +++ b/scripts/run_multi_coordinator_scaling.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""Execute D5-D7 strong/weak scaling and emit provenance-rich summaries.""" + +from __future__ import annotations + +import dataclasses +import json +import os +import shutil +import subprocess +import sys +from contextlib import ExitStack +from datetime import datetime, timezone +from pathlib import Path +from typing import Sequence + +from scripts.multi_coordinator_rank_process import RankProcess +from scripts.multi_coordinator_rank_process import ( + wait_for_rank_processes as _wait_for_rank_processes, +) +from scripts.multi_coordinator_scaling_cli import parse_scaling_cli +from scripts.multi_coordinator_scaling_conclusions import ( + build_scaling_conclusion, +) +from scripts.multi_coordinator_scaling_config import ( + AffinityMode, + build_rank_launch, + select_affinity_mode, +) +from scripts.multi_coordinator_scaling_execution import ( + ExperimentRun, + apply_affinity_mode, + build_experiment_plan, +) +from scripts.multi_coordinator_scaling_results import ( + GlobalResult, + LossCorrectness, + RankResult, + build_global_result, +) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _memory_binding_available() -> bool: + if shutil.which("numactl") is None: + return False + probe = subprocess.run( + ("numactl", "--cpunodebind=0", "--membind=0", "true"), + capture_output=True, + check=False, + ) + return probe.returncode == 0 + + +def _load_rank_result(path: Path) -> RankResult: + payload = json.loads(path.read_text(encoding="utf-8")) + correctness = payload["loss_correctness"] + phases = payload["phases_seconds"] + return RankResult( + rank=int(payload["rank"]), + warmup_iteration_seconds=float(payload["warmup_iteration_seconds"]), + iteration_total_seconds=float(phases["iteration_total"]), + device_dispatch_aggregation_seconds=float( + phases["device_dispatch_aggregation"] + ), + gradient_sync_seconds=float(phases["gradient_sync"]), + optimizer_seconds=float(phases["optimizer"]), + idle_other_seconds=float(phases["idle_other"]), + throughput_samples_per_second=float( + payload["throughput_samples_per_second"] + ), + measured_losses=tuple( + float(value) for value in payload["measured_losses"] + ), + golden_losses=tuple(float(value) for value in payload["golden_losses"]), + loss_correctness=LossCorrectness( + passed=bool(correctness["passed"]), + decreasing=bool(correctness["decreasing"]), + tracks_golden=bool(correctness["tracks_golden"]), + max_relative_error=float(correctness["max_relative_error"]), + ), + ) + + +def _rank_commands( + run: ExperimentRun, affinity_mode: AffinityMode +) -> tuple[tuple[str, ...], ...]: + config = dataclasses.replace(run.config, affinity_mode=affinity_mode) + commands: list[tuple[str, ...]] = [] + for rank in range(config.coordinators): + launch = build_rank_launch(config, rank, run.run_directory) + commands.append( + launch.command + + ( + "--scaling_mode", + config.mode, + "--master_port", + str(run.master_port), + "--base_proxy_port", + str(run.base_proxy_port), + ) + ) + return tuple(commands) + + +def _execute_run( + run: ExperimentRun, affinity_mode: AffinityMode, timeout_seconds: int +) -> GlobalResult: + run.run_directory.mkdir(parents=True, exist_ok=True) + config = dataclasses.replace(run.config, affinity_mode=affinity_mode) + commands = _rank_commands(run, affinity_mode) + processes: list[RankProcess] = [] + with ExitStack() as stack: + for rank, command in enumerate(commands): + launch = build_rank_launch(config, rank, run.run_directory) + environment = os.environ.copy() + environment.update(launch.environment) + environment["GLOO_SOCKET_IFNAME"] = "lo" + stdout = stack.enter_context( + (run.run_directory / f"rank-{rank}.stdout.log").open("wb") + ) + stderr = stack.enter_context( + (run.run_directory / f"rank-{rank}.stderr.log").open("wb") + ) + processes.append( + RankProcess( + subprocess.Popen( + command, + env=environment, + stdout=stdout, + stderr=stderr, + start_new_session=True, + ) + ) + ) + _wait_for_rank_processes(processes, timeout_seconds) + ranks = tuple( + _load_rank_result(run.run_directory / f"rank-{rank}.json") + for rank in range(config.coordinators) + ) + result = build_global_result(config, ranks) + payload = { + "schema_version": 1, + "experiment": f"D5-D7 {config.mode} multi-coordinator scaling", + "timestamp": _utc_now(), + "configuration": dataclasses.asdict(config), + "commands": commands, + "transport": "Gloo over loopback", + "substrate": result.substrate, + "physical_nic_routing": False, + "result": dataclasses.asdict(result), + } + (run.run_directory / "global.json").write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return result + + +def _write_summaries( + output_root: Path, + runs: Sequence[ExperimentRun], + results: Sequence[GlobalResult], + invocation: Sequence[str], +) -> None: + grouped: dict[str, list[tuple[ExperimentRun, GlobalResult]]] = {} + for run, result in zip(runs, results): + grouped.setdefault(run.config.mode, []).append((run, result)) + breakdown_rows: list[dict[str, object]] = [] + for mode, points in grouped.items(): + baseline, scaled = points + conclusion = build_scaling_conclusion( + baseline[0].config, + scaled[0].config, + baseline[1], + scaled[1], + ) + summary = { + "schema_version": 1, + "experiment": f"{mode} multi-coordinator scaling", + "timestamp": _utc_now(), + "invocation": list(invocation), + "transport": "Gloo over loopback", + "substrate": "same-host loopback lower bound", + "physical_nic_routing": False, + "points": [ + { + "configuration": dataclasses.asdict(run.config), + "result": dataclasses.asdict(result), + } + for run, result in points + ], + "conclusion": dataclasses.asdict(conclusion), + } + (output_root / f"{mode}.json").write_text( + json.dumps(summary, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + for run, result in points: + slowest = max( + result.ranks, key=lambda rank: rank.iteration_total_seconds + ) + breakdown_rows.append( + { + "mode": mode, + "coordinators": run.config.coordinators, + "iteration_total": slowest.iteration_total_seconds, + "device_dispatch_aggregation": slowest.device_dispatch_aggregation_seconds, + "gradient_sync": slowest.gradient_sync_seconds, + "optimizer": slowest.optimizer_seconds, + "idle_other": slowest.idle_other_seconds, + } + ) + breakdown = { + "schema_version": 1, + "timestamp": _utc_now(), + "component_semantics": { + "device_dispatch_aggregation": "combined forward and backward device-facing dispatch and aggregation", + "gradient_sync": "inter-coordinator Gloo AllReduce over loopback", + "optimizer": "rank-local optimizer step", + "idle_other": "residual uninstrumented work and waiting", + }, + "rows": breakdown_rows, + } + (output_root / "breakdown.json").write_text( + json.dumps(breakdown, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def main() -> int: + cli = parse_scaling_cli() + runs = build_experiment_plan( + modes=cli.modes, + output_root=cli.output_directory, + workload=cli.workload, + ) + memory_binding = _memory_binding_available() + affinity_mode = select_affinity_mode( + numactl_available=shutil.which("numactl") is not None, + memory_binding_available=memory_binding, + ) + runs = apply_affinity_mode(runs, affinity_mode) + if cli.dry_run: + for run in runs: + for command in _rank_commands(run, affinity_mode): + print(" ".join(command)) + return 0 + results = tuple( + _execute_run(run, affinity_mode, cli.timeout_seconds) for run in runs + ) + _write_summaries(cli.output_directory, runs, results, sys.argv) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_single_coordinator_training.py b/scripts/run_single_coordinator_training.py new file mode 100644 index 00000000..ba9d3f51 --- /dev/null +++ b/scripts/run_single_coordinator_training.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 +# pyright: reportAny=false, reportExplicitAny=false, reportUnknownMemberType=false +# pyright: reportUnknownVariableType=false, reportUnknownArgumentType=false +# pyright: reportMissingParameterType=false, reportUnusedCallResult=false +"""D1: real single-coordinator training loop through the proxy backend + fleet. + +Runs a multi-step AdamW training loop where the transformer's ``nn.Linear`` +GEMMs are dispatched to a real ``ProxySvr`` backend + fake-device fleet (the +same path the paper uses), while attention ``bmm`` and the wide ``lm_head`` stay +native so a full forward is tractable. It first trains the model natively +(golden) and then dispatched, from the same seed/inputs, and reports whether the +dispatched loss trajectory tracks golden. + +Usage (inside the Docker image, with scripts/ mounted): + python3 scripts/run_single_coordinator_training.py \ + --num_devices 2 --model_name facebook/opt-125m \ + --steps 6 --seq_length 128 --batch_size 1 +""" + +from __future__ import annotations + +import argparse +import os +import time + +import torch + +import morphling.hooks.autograd as hooks_autograd +from morphling.runtime.coordinator_metrics import ( + CoordinatorMetricsCollector, + PhaseRecorder, + add_metrics_arguments, + metrics_config_from_args, + start_metrics_collector, + track_phase, +) +from scripts._runtime_common import ( + load_model_and_tokenizer, + start_backend, + wait_for_connections, +) +from scripts._single_coordinator_fleet import ( + cleanup_fake_fleet, +) +from scripts._single_coordinator_fleet import ( + spawn_fake_fleet as _spawn_fake_fleet, +) +from scripts._single_coordinator_training_workload import ( + apply_restricted_hook as _apply_restricted_hook, +) +from scripts._single_coordinator_training_workload import ( + infer_vocab_size as _infer_vocab_size, +) +from scripts._single_coordinator_training_workload import ( + make_tiny_model as _make_tiny_model, +) +from scripts._single_coordinator_training_workload import ( + remove_restricted_hook as _remove_restricted_hook, +) +from scripts._single_coordinator_training_workload import ( + train_loop as _train_loop, +) +from scripts._single_coordinator_training_workload import ( + train_loop_tiny as _train_loop_tiny, +) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num_devices", type=int, default=2) + parser.add_argument("--model_name", type=str, default="facebook/opt-125m") + parser.add_argument("--steps", type=int, default=3) + parser.add_argument("--warmup_steps", type=int, default=1) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--batch_size", type=int, default=1) + parser.add_argument("--seq_length", type=int, default=128) + parser.add_argument("--backend", type=str, default="proxy") + parser.add_argument("--block_size", type=int, default=128) + parser.add_argument("--cfg", type=str, default=None) + parser.add_argument("--proxy_host", type=str, default="") + parser.add_argument("--rtol", type=float, default=0.05) + parser.add_argument( + "--tiny", + action="store_true", + help="Use a small MLP + regression loop so a full dispatched run " + "completes quickly (correctness check without the opt-125m dispatch " + "overhead).", + ) + add_metrics_arguments(parser) + args = parser.parse_args() + + device = torch.device("cpu") + vocab_size = 0 + os.environ["NUM_DEVICES"] = str(args.num_devices) + + phase_recorder: PhaseRecorder | None = None + metrics_collector: CoordinatorMetricsCollector | None = None + metrics_config = metrics_config_from_args(args) + if metrics_config is not None: + phase_recorder = PhaseRecorder() + metrics_collector = start_metrics_collector( + metrics_config, phase_recorder + ) + if metrics_collector is None: + phase_recorder = None + + # ---- 1. Golden: native (no dispatch), same seed/inputs ----------------- + _remove_restricted_hook() + if args.tiny: + model_g = _make_tiny_model(args.seed).to(device) + golden = _train_loop_tiny( + model_g, + steps=args.steps, + warmup_steps=args.warmup_steps, + seed=args.seed, + device=device, + ) + else: + model_g, tokenizer = load_model_and_tokenizer( + args.model_name, dtype=torch.float32 + ) + model_g = model_g.to(device) + vocab_size = _infer_vocab_size(model_g, tokenizer) + golden = _train_loop( + model_g, + steps=args.steps, + warmup_steps=args.warmup_steps, + seed=args.seed, + batch_size=args.batch_size, + seq_length=args.seq_length, + device=device, + vocab_size=vocab_size, + ) + del model_g + if torch.cuda.is_available(): + torch.cuda.empty_cache() + print("GOLDEN (native) losses:", [round(x, 4) for x in golden], flush=True) + + # ---- 2. Start real proxy backend + fake fleet -------------------------- + backend = start_backend( + backend_name=args.backend, + block_size=args.block_size, + cfg_path=args.cfg, + ) + hooks_autograd._backend = backend + with track_phase(phase_recorder, "idle_wait"): + time.sleep(5) + print(f"Spawning {args.num_devices} fast fake devices...", flush=True) + _spawn_fake_fleet(args.num_devices, args.backend, args.proxy_host) + if args.backend == "proxy": + with track_phase(phase_recorder, "idle_wait"): + connected = wait_for_connections( + backend, min_devices=args.num_devices, timeout=120 + ) + if connected < args.num_devices: + print( + f"WARNING: only {connected}/{args.num_devices} devices connected", + flush=True, + ) + with track_phase(phase_recorder, "idle_wait"): + time.sleep(5) + + # ---- 3. Dispatched: Linear GEMMs through the fleet --------------------- + if args.tiny: + model_d = _make_tiny_model(args.seed).to(device) + _apply_restricted_hook() + t0 = time.time() + dispatched = _train_loop_tiny( + model_d, + steps=args.steps, + warmup_steps=args.warmup_steps, + seed=args.seed, + device=device, + phase_recorder=phase_recorder, + ) + elapsed = time.time() - t0 + else: + model_d, _ = load_model_and_tokenizer( + args.model_name, dtype=torch.float32 + ) + model_d = model_d.to(device) + _apply_restricted_hook() + t0 = time.time() + dispatched = _train_loop( + model_d, + steps=args.steps, + warmup_steps=args.warmup_steps, + seed=args.seed, + batch_size=args.batch_size, + seq_length=args.seq_length, + device=device, + vocab_size=vocab_size, + phase_recorder=phase_recorder, + ) + elapsed = time.time() - t0 + _remove_restricted_hook() + print( + "DISPATCHED losses:", + [round(x, 4) if x == x else x for x in dispatched], + flush=True, + ) + print( + f"dispatched wall time: {elapsed:.1f}s for {args.warmup_steps} warmup " + f"and {args.steps} measured steps", + flush=True, + ) + + # ---- 4. Verdict -------------------------------------------------------- + any_nan = any(x != x for x in dispatched) + decreasing = (not any_nan) and dispatched[-1] < dispatched[0] + tracks = (not any_nan) and all( + abs(d - g) <= args.rtol * max(abs(g), 1e-6) + for d, g in zip(dispatched, golden) + ) + max_rel = ( + max(abs(d - g) / max(abs(g), 1e-6) for d, g in zip(dispatched, golden)) + if not any_nan + else float("nan") + ) + print(f"any NaN: {any_nan}", flush=True) + print(f"loss decreasing: {decreasing}", flush=True) + print( + f"full trajectory tracks golden (rtol={args.rtol}): {tracks} " + f"(max rel err {max_rel})", + flush=True, + ) + + try: + if hasattr(backend, "stop"): + backend.stop() + except Exception: + pass + cleanup_fake_fleet() + + ok = (not any_nan) and decreasing and tracks + print("D1 RESULT:", "PASS" if ok else "FAIL", flush=True) + if metrics_collector is not None: + metrics_collector.stop() + # ProxySvr's libevent loop thread is non-daemon with no stop binding, so a + # normal return leaves the process alive; force-exit after the verdict. + os._exit(0 if ok else 1) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/run_within_batch_experiment.py b/scripts/run_within_batch_experiment.py new file mode 100644 index 00000000..fe91e9ac --- /dev/null +++ b/scripts/run_within_batch_experiment.py @@ -0,0 +1,336 @@ +#!/usr/bin/env python3 +"""C2: within-batch dynamic-performance experiment (Wave C). + +Takes the real 256-device OPT-13B VTIME logs (measured base per-GEMM +durations) for every baseline, overlays a C1 within-batch time-varying trace +(per-device compute-throttle + network-rate drift), and reports per-batch +runtime and recovered useful FLOPs versus the fixed-straggler setup across a +variation-magnitude sweep. Post-processing only; reuses compute_batch_runtime. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +from compute_batch_runtime import ( + _load_device_profiles, + _load_json, + compute_batch_runtime, + parse_vtime_log, +) +from generate_trace import generate_within_batch_traces + +BASELINES = ["cleave", "dtfm", "asteroid", "confident", "alpa"] + + +def _inputs(root: Path, baseline: str) -> tuple[Path, Path, Path]: + vt_candidates = [ + root / baseline / "vtime.log", + root / "vtime" / baseline / "vtime.log", + ] + vt = next((p for p in vt_candidates if p.exists()), vt_candidates[0]) + mf_candidates = [ + root / f"{baseline}_results" / "manifests" / f"{baseline}_manifest.json", + root / "planning" / "manifests" / f"{baseline}_manifest.json", + ] + mf = next((p for p in mf_candidates if p.exists()), mf_candidates[0]) + dc_candidates = [ + root / f"{baseline}_results" / "generated_device_fleet.json", + root / "generated_device_fleet.json", + ] + dc = next((p for p in dc_candidates if p.exists()), dc_candidates[0]) + return vt, mf, dc + + +def _useful_flops(manifest: Any) -> float: + ents = manifest.get("entries", []) if isinstance(manifest, dict) else manifest + seen: dict[tuple[int, int], float] = {} + for e in ents: + key = (int(e.get("level", 0)), int(e.get("gemm_id", 0))) + if key not in seen: + m = float(e.get("m_total") or e.get("alpha") or 0.0) + n = float(e.get("n") or 0.0) + q = float(e.get("q_total") or e.get("beta") or 0.0) + seen[key] = 2.0 * m * n * q + return float(sum(seen.values())) + + +def _num_devices(manifest: Any) -> int: + if isinstance(manifest, dict): + meta = manifest.get("metadata", {}) + nd = int(meta.get("num_devices", 0) or 0) + if nd > 0: + return nd + ents = manifest.get("entries", []) if isinstance(manifest, dict) else manifest + return len({int(e.get("device_id", -1)) for e in ents if int(e.get("device_id", -1)) >= 0}) + + +def _idle_fraction(result: Any, num_devices: int) -> float: + provisioned = num_devices * result.total_runtime_ms + if provisioned <= 0: + return 0.0 + busy = sum(v["total_ms"] for v in result.per_device_breakdown.values()) + return max(0.0, 1.0 - busy / provisioned) + + +def _level_devices(manifest: Any) -> dict[int, list[int]]: + ents = manifest.get("entries", []) if isinstance(manifest, dict) else manifest + out: dict[int, set[int]] = {} + for e in ents: + lvl = int(e.get("level", 0)) + dev = int(e.get("device_id", -1)) + if dev >= 0: + out.setdefault(lvl, set()).add(dev) + return {k: sorted(v) for k, v in out.items()} + + +def _mean_rate( + trace: dict[str, Any], + devices: list[int], + level: int, + num_levels: int, + num_bins: int, +) -> float: + if not devices: + return 1.0 + frac = 0.0 if num_levels <= 1 else min(0.999999, level / num_levels) + b = min(num_bins - 1, int(frac * num_bins)) + tdev = trace.get("devices", {}) + rates: list[float] = [] + for d in devices: + dd = tdev.get(str(d)) + if not dd: + rates.append(1.0) + continue + r = min(dd["flops_mult"][b], dd["ul_bw_mult"][b], dd["dl_bw_mult"][b]) + rates.append(max(1e-6, float(r))) + return sum(rates) / len(rates) + + +def _cleave_reweighted_runtime( + nominal: Any, + trace: dict[str, Any], + level_devices: dict[int, list[int]], + num_bins: int, + num_levels: int, +) -> float: + """CLEAVE work-conserving reweighting: per level, the nominal bottleneck + time is inflated by 1/mean_rate (work redistributed proportional to current + device rates), so CLEAVE is limited by the AVERAGE device slowdown rather + than the slowest participant (paper mechanism, evaluation.tex L402).""" + total = 0.0 + for lvl in nominal.per_level_breakdown: + l = int(lvl["level"]) + t = float(lvl["runtime_ms"]) + mr = _mean_rate(trace, level_devices.get(l, []), l, num_levels, num_bins) + total += t / mr if mr > 0 else t + return total + float(nominal.optimizer_tail_ms) + + +def run( + root: Path, + magnitudes: list[float], + num_bins: int, + event_rate: float, + duration_frac: float, + seed: int, +) -> dict[str, Any]: + cache: dict[str, dict[str, Any]] = {} + for b in BASELINES: + vt, mf, dc = _inputs(root, b) + if not (vt.exists() and mf.exists()): + print(f"skip {b}: missing inputs ({vt.exists()=}, {mf.exists()=})") + continue + manifest = _load_json(str(mf)) + ld = _level_devices(manifest) + cache[b] = { + "vtime": parse_vtime_log(str(vt)), + "manifest": manifest, + "profiles": _load_device_profiles(str(dc)) if dc.exists() else {}, + "useful_flops": _useful_flops(manifest), + "num_devices": _num_devices(manifest), + "level_devices": ld, + "num_levels": max(ld.keys(), default=0) + 1, + "nominal": compute_batch_runtime( + vtime_events=parse_vtime_log(str(vt)), + manifest=manifest, + baseline_type=b, + device_profiles=_load_device_profiles(str(dc)) if dc.exists() else {}, + ), + } + + per_baseline: dict[str, dict[str, Any]] = {b: {"points": []} for b in cache} + base_runtime: dict[str, float] = { + b: cache[b]["nominal"].total_runtime_ms for b in cache + } + + for m in magnitudes: + nd = max((cache[b]["num_devices"] for b in cache), default=256) + trace = ( + None + if m <= 0.0 + else generate_within_batch_traces( + num_devices=nd, + num_bins=num_bins, + magnitude=m, + event_rate=event_rate, + duration_frac=duration_frac, + network_drift_std=0.5 * m, + hurst=0.85, + seed=seed, + ) + ) + for b, data in cache.items(): + idle: float | None + if m <= 0.0: + rt = base_runtime[b] + idle = _idle_fraction(data["nominal"], data["num_devices"]) + elif b == "cleave": + rt = _cleave_reweighted_runtime( + data["nominal"], trace, data["level_devices"], + num_bins, data["num_levels"], + ) + idle = None + else: + result = compute_batch_runtime( + vtime_events=data["vtime"], + manifest=data["manifest"], + baseline_type=b, + device_profiles=data["profiles"], + within_batch_trace=trace, + ) + rt = result.total_runtime_ms + idle = _idle_fraction(result, data["num_devices"]) + per_baseline[b]["points"].append( + { + "magnitude": m, + "runtime_ms": rt, + "normalized_runtime": rt / base_runtime[b] + if base_runtime.get(b) + else 1.0, + "idle_fraction": idle, + "useful_flop_throughput": data["useful_flops"] / rt + if rt > 0 + else 0.0, + } + ) + + recovered = _recovered_useful_flops(cache, per_baseline, base_runtime, magnitudes) + + return { + "experiment": "within_batch_dynamic_performance", + "model": "opt-13b", + "num_devices": max((cache[b]["num_devices"] for b in cache), default=0), + "magnitudes": magnitudes, + "variation_model": { + "num_bins": num_bins, + "event_rate": event_rate, + "duration_frac": duration_frac, + "network_drift_std": "0.5 * magnitude", + "seed": seed, + }, + "base_runtime_ms": base_runtime, + "per_baseline": per_baseline, + "recovered_useful_flops": recovered, + } + + +def _recovered_useful_flops( + cache: dict[str, Any], + per_baseline: dict[str, Any], + base_runtime: dict[str, float], + magnitudes: list[float], +) -> dict[str, Any]: + if "cleave" not in cache: + return {} + baselines = [b for b in per_baseline if b != "cleave"] + out: list[dict[str, Any]] = [] + for i, m in enumerate(magnitudes): + if m <= 0.0: + continue + cl = per_baseline["cleave"]["points"][i] + useful = cache["cleave"]["useful_flops"] + cl_lost = useful * max( + 0.0, 1.0 - base_runtime["cleave"] / cl["runtime_ms"] + ) + base_losses = [] + for b in baselines: + pt = per_baseline[b]["points"][i] + base_losses.append( + cache[b]["useful_flops"] + * max(0.0, 1.0 - base_runtime[b] / pt["runtime_ms"]) + ) + mean_base_lost = sum(base_losses) / len(base_losses) if base_losses else 0.0 + recovered = max(0.0, mean_base_lost - cl_lost) + out.append( + { + "magnitude": m, + "cleave_lost_flops": cl_lost, + "mean_baseline_lost_flops": mean_base_lost, + "recovered_flops": recovered, + "recovered_fraction_of_useful": recovered / useful if useful else 0.0, + } + ) + return {"per_magnitude": out} + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--root", + type=str, + default="results/vtime_models/opt-13b", + help="Directory with {baseline}/vtime.log and {baseline}_results/", + ) + parser.add_argument( + "--magnitudes", + type=str, + default="0.0,0.2,0.4,0.6", + help="Comma-separated variation magnitudes (throttle depth)", + ) + parser.add_argument("--num-bins", type=int, default=100) + parser.add_argument("--event-rate", type=float, default=3.0) + parser.add_argument("--duration-frac", type=float, default=0.15) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument( + "--output", + type=str, + default="results/within_batch/summary.json", + ) + args = parser.parse_args() + + magnitudes = [float(x) for x in args.magnitudes.split(",")] + summary = run( + root=Path(args.root), + magnitudes=magnitudes, + num_bins=args.num_bins, + event_rate=args.event_rate, + duration_frac=args.duration_frac, + seed=args.seed, + ) + + out_path = Path(args.output) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w") as f: + json.dump(summary, f, indent=2) + + print(f"\nWrote: {out_path}") + for b, data in summary["per_baseline"].items(): + pts = data["points"] + norm = ", ".join(f"{p['normalized_runtime']:.2f}" for p in pts) + print(f" {b:10s} normalized runtime: [{norm}]") + rec = summary.get("recovered_useful_flops", {}).get("per_magnitude", []) + for r in rec: + print( + f" m={r['magnitude']}: CLEAVE recovers " + f"{100 * r['recovered_fraction_of_useful']:.1f}% of useful FLOPs " + f"vs mean baseline" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/verify_absolute_times.py b/scripts/verify_absolute_times.py new file mode 100644 index 00000000..7b46d8b0 --- /dev/null +++ b/scripts/verify_absolute_times.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +"""B1 verifier: recompute + validate the absolute per-batch times in +tab:absolute-times (results/analytical_scaling/absolute_times.json). + +Reproduces the analytical cloud-A100 baseline for OPT-13B, Llama2-13B, and +Llama2-70B, recomputes the OPT-13B/256 sys (CLEAVE) time from the planner +manifest with a uniform median-edge fleet, and asserts the matched envelope +(batch=128, seq=1024) so every plotted point is checked against a common model. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent)) +from compute_batch_runtime import compute_batch_runtime, _load_json + +BATCH = 128 +SEQ = 1024 +TOKENS = BATCH * SEQ +A100_BF16_PEAK = 3.12e14 +A100_UTIL = 0.975 +MEDIAN_EDGE = {"flops": 6e12, "ul_bw": 7.5e6, "dl_bw": 55e6, "ul_lat": 0.0, "dl_lat": 0.0} + +CLOUD_PARAMS = {"opt-13b": 13e9, "llama2-13b": 13e9, "llama2-70b": 70e9} +CLOUD_EXPECTED = {"opt-13b": 33.6, "llama2-13b": 33.6, "llama2-70b": 180.8} + +OPT13B_256_MANIFESTS = ( + "results/vtime_scaling/model_size/points/002_opt-13b/planning/manifests" +) + + +def cloud_seconds(num_params: float) -> float: + return 6.0 * num_params * TOKENS / (A100_BF16_PEAK * A100_UTIL) + + +def main() -> int: + ok = True + + print("== Cloud A100 analytical baseline (6*P*tokens / (312 TFLOPS * 0.975)) ==") + for model, params in CLOUD_PARAMS.items(): + got = cloud_seconds(params) + exp = CLOUD_EXPECTED[model] + close = abs(got - exp) <= 0.6 + ok = ok and close + print(f" {model:11s}: {got:6.1f} s (paper {exp}) {'OK' if close else 'MISMATCH'}") + + print("\n== sys (CLEAVE) OPT-13B/256, uniform median-edge fleet, bytes/elem=1 ==") + manifest = _load_json(f"{OPT13B_256_MANIFESTS}/cleave_manifest.json") + meta = manifest.get("metadata", {}).get("model", {}) + profiles = {i: dict(MEDIAN_EDGE) for i in range(256)} + res = compute_batch_runtime( + [], manifest, "cleave", device_profiles=profiles, bytes_per_element=1.0 + ) + cleave_s = res.total_runtime_ms / 1000.0 + cleave_close = abs(cleave_s - 37.3) <= 2.0 + ok = ok and cleave_close + print(f" cleave: {cleave_s:.1f} s (paper 37.3) {'OK' if cleave_close else 'MISMATCH'}") + + print("\n== matched envelope ==") + seq_ok = int(meta.get("seq_length", 0)) == SEQ + ok = ok and seq_ok + print(f" OPT-13B/256 manifest: model={meta.get('model_name')} " + f"seq={meta.get('seq_length')} (expected {SEQ}) {'OK' if seq_ok else 'MISMATCH'}") + print(f" batch={BATCH}, tokens={TOKENS} used for the cloud FLOP count") + + print("\nRESULT:", "PASS" if ok else "FAIL") + return 0 if ok else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt index 685efe90..b35b89c1 100644 --- a/tests/cpp/CMakeLists.txt +++ b/tests/cpp/CMakeLists.txt @@ -246,6 +246,18 @@ target_link_libraries(test_dispatch_gate PRIVATE ) add_test(NAME test_dispatch_gate COMMAND test_dispatch_gate) +add_executable(test_operation_id + unit/backend/test_operation_id.cpp +) +target_include_directories(test_operation_id PRIVATE + ${PROJECT_ROOT_DG}/csrc +) +target_link_libraries(test_operation_id PRIVATE + GTest::gtest_main + GTest::gtest +) +add_test(NAME test_operation_id COMMAND test_operation_id) + # ============================================================================ # GemmArgs layout tests (unit/intercept/) — Issue #48 collapse # ============================================================================ diff --git a/tests/cpp/unit/backend/test_operation_id.cpp b/tests/cpp/unit/backend/test_operation_id.cpp new file mode 100644 index 00000000..076e3e14 --- /dev/null +++ b/tests/cpp/unit/backend/test_operation_id.cpp @@ -0,0 +1,31 @@ +#include + +#include +#include + +#include "backend/operation_id.h" + +namespace { + +using morphling::backend::kMaxLifetimeOperationCount; +using morphling::backend::ReserveOperationId; +using morphling::backend::ValidateOperationId; + +TEST(OperationIdTest, WaitMatMulRejectsNegativeOid) { + EXPECT_THROW(ValidateOperationId(-1), std::out_of_range); +} + +TEST(OperationIdTest, WaitMatMulRejectsOidAtCapacity) { + EXPECT_THROW(ValidateOperationId(kMaxLifetimeOperationCount), + std::out_of_range); +} + +TEST(OperationIdTest, ReservesLastAvailableOidWithoutExceedingCapacity) { + std::atomic_int next_oid{kMaxLifetimeOperationCount - 1}; + + EXPECT_EQ(ReserveOperationId(next_oid), kMaxLifetimeOperationCount - 1); + EXPECT_EQ(ReserveOperationId(next_oid), -1); + EXPECT_EQ(next_oid.load(), kMaxLifetimeOperationCount); +} + +} // namespace diff --git a/tests/cpp/unit/backend/test_partition_tracker_claim.cpp b/tests/cpp/unit/backend/test_partition_tracker_claim.cpp new file mode 100644 index 00000000..12faee9f --- /dev/null +++ b/tests/cpp/unit/backend/test_partition_tracker_claim.cpp @@ -0,0 +1,129 @@ +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "backend/partition_tracker.h" + +namespace morphling::backend { +namespace { + +TEST(PartitionTrackerClaimTest, ConcurrentClaimersClaimEachPartitionOnce) { + // Given + auto& tracker = PARTITION_TRACKER; + tracker.Reset(); + constexpr int kPartitionCount = 100; + constexpr int kClaimerCount = 8; + for (int index = 0; index < kPartitionCount; ++index) { + auto partition = std::make_shared(); + tracker.AddPartition(-1, "partition-" + std::to_string(index), 7, + partition); + } + std::atomic start{false}; + std::mutex claimed_mutex; + std::vector claimed_keys; + std::vector claimers; + + // When + for (int index = 0; index < kClaimerCount; ++index) { + claimers.emplace_back([&, index] { + while (!start.load(std::memory_order_acquire)) { + std::this_thread::yield(); + } + auto claimed = tracker.ClaimIdlePartitions(); + std::lock_guard lock(claimed_mutex); + for (const auto& partition : claimed) { + EXPECT_TRUE(tracker.ReassignPartitionToDevice(partition->key, index)); + claimed_keys.push_back(partition->key); + } + }); + } + start.store(true, std::memory_order_release); + for (auto& claimer : claimers) { + claimer.join(); + } + + // Then + const std::unordered_set unique_keys(claimed_keys.begin(), + claimed_keys.end()); + EXPECT_EQ(claimed_keys.size(), kPartitionCount); + EXPECT_EQ(unique_keys.size(), kPartitionCount); + EXPECT_TRUE(tracker.GetIdlePartitions().empty()); + size_t relocated_count = 0; + for (int index = 0; index < kClaimerCount; ++index) { + relocated_count += tracker.GetDevicePartitionCount(index); + } + EXPECT_EQ(relocated_count, kPartitionCount); + tracker.Reset(); +} + +TEST(PartitionTrackerClaimTest, ReassignmentPreservesClaimedPartitionIdentity) { + // Given + auto& tracker = PARTITION_TRACKER; + tracker.Reset(); + auto matrix_partition = std::make_shared(); + tracker.AddPartition(-1, "partition", 9, matrix_partition); + auto claimed = tracker.ClaimIdlePartitions(); + ASSERT_EQ(claimed.size(), 1); + const auto claimed_info = claimed.front(); + + // When + ASSERT_TRUE(tracker.ReassignPartitionToDevice("partition", 4)); + ASSERT_TRUE(tracker.ReassignPartitionToDevice("partition", 2)); + + // Then + EXPECT_TRUE(tracker.GetDevicePartitions(-1).empty()); + EXPECT_TRUE(tracker.GetDevicePartitions(4).empty()); + const auto target_partitions = tracker.GetDevicePartitions(2); + ASSERT_EQ(target_partitions.size(), 1); + EXPECT_EQ(target_partitions.front(), claimed_info); + EXPECT_EQ(target_partitions.front()->partition, matrix_partition); + EXPECT_EQ(target_partitions.front()->owner_device_id, 2); + EXPECT_EQ(target_partitions.front()->partition->dev_id, 2); + EXPECT_EQ(target_partitions.front()->state, PartitionState::RUNNING); + EXPECT_TRUE(tracker.GetIdlePartitions().empty()); + tracker.Reset(); +} + +TEST(PartitionTrackerClaimTest, RevertedClaimCanBeClaimedOnceAgain) { + // Given + auto& tracker = PARTITION_TRACKER; + tracker.Reset(); + auto matrix_partition = std::make_shared(); + tracker.AddPartition(-1, "partition", 11, matrix_partition); + auto first_claim = tracker.ClaimIdlePartitions(); + ASSERT_EQ(first_claim.size(), 1); + ASSERT_TRUE(tracker.ReassignPartitionToDevice("partition", 3)); + + // When + tracker.MarkPartitionIdle("partition"); + auto second_claim = tracker.ClaimIdlePartitions(); + auto third_claim = tracker.ClaimIdlePartitions(); + + // Then + ASSERT_EQ(second_claim.size(), 1); + EXPECT_EQ(second_claim.front(), first_claim.front()); + EXPECT_TRUE(third_claim.empty()); + EXPECT_TRUE(tracker.GetIdlePartitions().empty()); + tracker.Reset(); +} + +TEST(PartitionTrackerClaimTest, ReassignmentRejectsUnknownPartition) { + // Given + auto& tracker = PARTITION_TRACKER; + tracker.Reset(); + + // When + const bool reassigned = tracker.ReassignPartitionToDevice("missing", 1); + + // Then + EXPECT_FALSE(reassigned); +} + +} // namespace +} // namespace morphling::backend diff --git a/tests/cpp/unit/zerocopy/CMakeLists.txt b/tests/cpp/unit/zerocopy/CMakeLists.txt index 3e552fda..53969227 100644 --- a/tests/cpp/unit/zerocopy/CMakeLists.txt +++ b/tests/cpp/unit/zerocopy/CMakeLists.txt @@ -69,6 +69,25 @@ add_zerocopy_test(test_profile_delta_log) add_zerocopy_test(test_device_measurement_service) add_zerocopy_test(test_device_measurement_session) +add_executable(test_partition_tracker_claim + ../backend/test_partition_tracker_claim.cpp + ${PROJECT_ROOT}/csrc/backend/partition_tracker.cpp + ${ZEROCOPY_TEST_DEPS} + ${BASE_LOGGING_SRC} + ${PROTO_SRCS} +) +target_compile_options(test_partition_tracker_claim PRIVATE -ffunction-sections + -fdata-sections) +target_link_options(test_partition_tracker_claim PRIVATE -Wl,--gc-sections) +target_include_directories(test_partition_tracker_claim PRIVATE + ${ZEROCOPY_INCLUDE_DIRS}) +target_link_libraries(test_partition_tracker_claim PRIVATE + ${ZEROCOPY_COMMON_LIBS}) +if(TARGET generate_proto_files) + add_dependencies(test_partition_tracker_claim generate_proto_files) +endif() +add_test(NAME test_partition_tracker_claim COMMAND test_partition_tracker_claim) + # ============================================================================ # CUDA Pinned Pool Test (uses CUDAToolkit from parent project) # ============================================================================ diff --git a/tests/python/unit/conftest.py b/tests/python/unit/conftest.py new file mode 100644 index 00000000..9da4b11a --- /dev/null +++ b/tests/python/unit/conftest.py @@ -0,0 +1,31 @@ +"""Collection guards for coordinator-evaluation unit tests. + +The coordinator-evaluation tests import heavy, optional dependencies at module +scope: ``psutil`` (coordinator metrics), ``matplotlib`` and the external +conference-plot helper (scaling figure), and the C-extension-backed +``morphling`` runtime. pytest imports every module during collection before it +can read markers, so these imports would abort the CPU smoke job -- which +installs only ``pytest``/``torch``/``numpy`` -- even though the tests are not +smoke-marked and are never selected there. + +``psutil`` is present in every environment that can actually run these tests +(the full ``requirements.txt`` install used by the Docker test image), and +absent from the minimal smoke environment, so it is a reliable signal. When it +is missing we skip collecting these files; full environments collect and run +them normally. +""" + +from __future__ import annotations + +import importlib.util + +collect_ignore_glob: list[str] = [] + +if importlib.util.find_spec("psutil") is None: + collect_ignore_glob += [ + "test_coord_scaling_*.py", + "test_coordinator_*.py", + "test_multi_coordinator_*.py", + "test_single_coordinator_*.py", + "test_run_coordinator_scaling.py", + ] diff --git a/tests/python/unit/hooks/test_autograd_dispatch_safety.py b/tests/python/unit/hooks/test_autograd_dispatch_safety.py new file mode 100644 index 00000000..47ab1e6f --- /dev/null +++ b/tests/python/unit/hooks/test_autograd_dispatch_safety.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +import torch + +from morphling.hooks import autograd as ag + + +class _SequenceBackend: + def __init__(self, operation_ids: list[int]) -> None: + self._operation_ids = iter(operation_ids) + self._outputs: dict[int, torch.Tensor] = {} + self.dispatch_operands: list[tuple[torch.Tensor, torch.Tensor]] = [] + self.waited_ids: list[int] = [] + + def async_dispatch_matmul( + self, mat_a: torch.Tensor, mat_b: torch.Tensor + ) -> int: + oid = next(self._operation_ids) + self.dispatch_operands.append((mat_a, mat_b)) + if oid >= 0: + self._outputs[oid] = torch.matmul(mat_a, mat_b.transpose(-2, -1)) + return oid + + def wait_matmul(self, oid: int) -> torch.Tensor: + self.waited_ids.append(oid) + return self._outputs.pop(oid) + + +@pytest.fixture(autouse=True) +def _isolate_autograd_state() -> Iterator[None]: + previous_backend = ag._backend + previous_greenctx = ag._greenctx + ag.set_greenctx(None) + yield + ag._backend = previous_backend + ag._greenctx = previous_greenctx + + +def test_forward_rejects_failed_dispatch_before_wait() -> None: + backend = _SequenceBackend([-1]) + ag.set_backend(backend) + input_tensor = torch.randn(3, 4) + weight = torch.randn(2, 4) + + with pytest.raises(RuntimeError, match="forward.*-1"): + ag.LinearFunction.apply(input_tensor, weight.transpose(-2, -1), None) + + assert backend.waited_ids == [] + + +def test_backward_rejects_failed_dispatch_before_wait() -> None: + backend = _SequenceBackend([0, -1, 1]) + ag.set_backend(backend) + input_tensor = torch.randn(3, 4, requires_grad=True) + weight = torch.randn(2, 4, requires_grad=True) + output = ag.LinearFunction.apply( + input_tensor, weight.transpose(-2, -1), None + ) + + with pytest.raises(RuntimeError, match="backward_grad_input.*-1"): + output.sum().backward() + + assert backend.waited_ids == [0] + + +def test_forward_dispatches_contiguous_operands_with_native_linear_parity() -> ( + None +): + backend = _SequenceBackend([7]) + ag.set_backend(backend) + input_tensor = torch.randn(3, 4) + weight = torch.randn(4, 2) + bias = torch.randn(2) + expected = torch.nn.functional.linear( + input_tensor, weight.transpose(-2, -1), bias + ) + + actual = ag.LinearFunction.apply(input_tensor, weight, bias) + + mat_a, mat_b = backend.dispatch_operands[0] + assert mat_a.is_contiguous() + assert mat_b.is_contiguous() + assert torch.allclose(actual, expected) diff --git a/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py b/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py index b9ac292d..b050a0b4 100644 --- a/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py +++ b/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py @@ -22,17 +22,19 @@ class _StubBackend: def __init__(self) -> None: self.dispatch_calls = 0 self.wait_calls = 0 - self._queue: list[torch.Tensor] = [] + self._outputs: dict[int, torch.Tensor] = {} def async_dispatch_matmul( self, mat_a: torch.Tensor, mat_b: torch.Tensor - ) -> None: + ) -> int: + oid = self.dispatch_calls self.dispatch_calls += 1 - self._queue.append(torch.matmul(mat_a, mat_b.transpose(-2, -1))) + self._outputs[oid] = torch.matmul(mat_a, mat_b.transpose(-2, -1)) + return oid - def wait_matmul(self, _idx: int) -> torch.Tensor: + def wait_matmul(self, oid: int) -> torch.Tensor: self.wait_calls += 1 - return self._queue.pop(0) + return self._outputs.pop(oid) class _StubGreenCtx: diff --git a/tests/python/unit/hooks/test_per_gemm_greenctx.py b/tests/python/unit/hooks/test_per_gemm_greenctx.py index aebe9345..4bc2e2ae 100644 --- a/tests/python/unit/hooks/test_per_gemm_greenctx.py +++ b/tests/python/unit/hooks/test_per_gemm_greenctx.py @@ -16,15 +16,19 @@ class _LocalMatmulBackend: def __init__(self) -> None: - self._queue = [] + self._next_oid = 0 + self._outputs = {} - def async_dispatch_matmul(self, mat_a, mat_b) -> None: - self._queue.append(torch.matmul(mat_a, mat_b.transpose(-2, -1))) + def async_dispatch_matmul(self, mat_a, mat_b) -> int: + oid = self._next_oid + self._next_oid += 1 + self._outputs[oid] = torch.matmul(mat_a, mat_b.transpose(-2, -1)) + return oid - def wait_matmul(self, _idx: int): - if not self._queue: + def wait_matmul(self, oid: int): + if oid not in self._outputs: raise RuntimeError("wait_matmul called without pending result") - return self._queue.pop(0) + return self._outputs.pop(oid) class _DeactivateRecorder: diff --git a/tests/python/unit/test_coord_scaling_figure_cli.py b/tests/python/unit/test_coord_scaling_figure_cli.py new file mode 100644 index 00000000..579050bd --- /dev/null +++ b/tests/python/unit/test_coord_scaling_figure_cli.py @@ -0,0 +1,194 @@ +"""Figure-builder and CLI tests for the coordinator-scaling pipeline. + +These exercise the matplotlib rendering path with the ``Agg`` backend: the +three-panel figure structure (no figure-level title, concise panel labels, +canonical breakdown component order) and the CLI contract (clear failure on a +missing input path, successful dual PDF + PNG output from fixture JSON). +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from scripts.coord_scaling_plot import cli, figure, schema + +# ── Fixture builders (mirror the real emitter field names) ────────────────── + + +def _summary() -> dict: + def _row(device_count: int, seconds: float, cpu: float) -> dict: + return { + "device_count": device_count, + "iteration_runtime_seconds": seconds, + "median_cpu_percent": cpu, + "peak_cpu_percent": 5500.0, + "substrate": "single-host loopback/emulated devices", + } + + return { + "environment": { + "logical_cpu_count": 56, + "host_nic_substrate": "single-host loopback/emulated devices", + }, + "rows": [ + _row(1, 12.0, 2200.0), + _row(2, 9.6, 2600.0), + _row(4, 9.7, 2950.0), + _row(8, 9.9, 2860.0), + ], + } + + +def _scaling(mode: str, base: float, second: float) -> dict: + def _point(coordinators: int, seconds: float) -> dict: + return { + "configuration": {"mode": mode, "coordinators": coordinators}, + "result": { + "coordinators": coordinators, + "iteration_total_seconds": seconds, + "throughput_samples_per_second": 16.0 / seconds, + }, + } + + return { + "points": [_point(1, base), _point(2, second)], + "conclusion": {"mode": mode, "efficiency": base / second}, + } + + +def _breakdown() -> dict: + def _row(mode: str, coordinators: int, sync: float) -> dict: + idle = 10.0 - 6.0 - sync - 1.0 + return { + "mode": mode, + "coordinators": coordinators, + "iteration_total": 10.0, + "device_dispatch_aggregation": 6.0, + "gradient_sync": sync, + "optimizer": 1.0, + "idle_other": idle, + } + + return { + "component_semantics": {"gradient_sync": "AllReduce over loopback"}, + "rows": [ + _row("strong", 1, 0.5), + _row("strong", 2, 1.5), + _row("weak", 1, 0.5), + _row("weak", 2, 1.5), + ], + } + + +def _parsed(): + device = schema.parse_device_scaling(_summary()) + strong = schema.parse_scaling_efficiency( + _scaling("strong", 10.0, 6.0), "strong" + ) + weak = schema.parse_scaling_efficiency(_scaling("weak", 8.0, 9.0), "weak") + breakdown = schema.parse_breakdown(_breakdown()) + return device, strong, weak, breakdown + + +def _write_inputs(tmp_path: Path) -> dict: + paths = { + "summary": tmp_path / "summary.json", + "strong": tmp_path / "strong.json", + "weak": tmp_path / "weak.json", + "breakdown": tmp_path / "breakdown.json", + } + paths["summary"].write_text(json.dumps(_summary()), encoding="utf-8") + paths["strong"].write_text( + json.dumps(_scaling("strong", 10.0, 6.0)), "utf-8" + ) + paths["weak"].write_text(json.dumps(_scaling("weak", 8.0, 9.0)), "utf-8") + paths["breakdown"].write_text(json.dumps(_breakdown()), encoding="utf-8") + return paths + + +# ── Figure structure ──────────────────────────────────────────────────────── + + +def test_build_figure_has_three_panels_and_no_figure_title() -> None: + import matplotlib.pyplot as plt + + fig = figure.build_figure(*_parsed()) + try: + titles = [ax.get_title() for ax in fig.axes] + assert fig._suptitle is None + assert any(title.startswith("(a)") for title in titles) + assert any(title.startswith("(b)") for title in titles) + assert any(title.startswith("(c)") for title in titles) + # Figure is the 7.0-inch double-column width. + assert fig.get_figwidth() == pytest.approx(7.0) + finally: + plt.close(fig) + + +def test_breakdown_panel_legend_uses_canonical_component_order() -> None: + import matplotlib.pyplot as plt + + fig = figure.build_figure(*_parsed()) + try: + assert len(figure.COMPONENT_LABELS) == len(schema.BREAKDOWN_COMPONENTS) + panel_c = next( + ax for ax in fig.axes if ax.get_title().startswith("(c)") + ) + legend = panel_c.get_legend() + labels = [text.get_text() for text in legend.get_texts()] + assert labels == list(figure.COMPONENT_LABELS) + finally: + plt.close(fig) + + +# ── CLI contract ──────────────────────────────────────────────────────────── + + +def test_cli_missing_input_path_fails_clearly(tmp_path: Path, capsys) -> None: + missing = tmp_path / "absent_summary.json" + exit_code = cli.main( + [ + "--summary", + str(missing), + "--strong", + str(tmp_path / "strong.json"), + "--weak", + str(tmp_path / "weak.json"), + "--breakdown", + str(tmp_path / "breakdown.json"), + "--out", + str(tmp_path / "fig"), + ] + ) + + assert exit_code == 2 + assert "absent_summary.json" in capsys.readouterr().err + + +def test_cli_generates_vector_pdf_and_png(tmp_path: Path) -> None: + inputs = _write_inputs(tmp_path) + out_base = tmp_path / "coordinator_scaling" + + exit_code = cli.main( + [ + "--summary", + str(inputs["summary"]), + "--strong", + str(inputs["strong"]), + "--weak", + str(inputs["weak"]), + "--breakdown", + str(inputs["breakdown"]), + "--out", + str(out_base), + ] + ) + + assert exit_code == 0 + pdf = out_base.with_suffix(".pdf") + png = out_base.with_suffix(".png") + assert pdf.is_file() and pdf.stat().st_size > 0 + assert png.is_file() and png.stat().st_size > 0 diff --git a/tests/python/unit/test_coord_scaling_schema.py b/tests/python/unit/test_coord_scaling_schema.py new file mode 100644 index 00000000..c37cf075 --- /dev/null +++ b/tests/python/unit/test_coord_scaling_schema.py @@ -0,0 +1,272 @@ +"""Unit tests for the coordinator-scaling plot data-prep boundary parser. + +These tests exercise ``scripts.coord_scaling_plot.schema`` in isolation +(no matplotlib): typed dataclass construction, JSON boundary validation, +per-point scaling-efficiency computation, and stacked-breakdown component +ordering / reconciliation. Every fixture is a temp JSON built in-memory so +the tests never touch real result files. +""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from scripts.coord_scaling_plot import schema + +# ── Fixture builders ──────────────────────────────────────────────────────── + + +def _summary(logical_cpu_count: int = 50) -> dict: + return { + "schema_version": 1, + "environment": { + "logical_cpu_count": logical_cpu_count, + "host_nic_substrate": "single-host loopback/emulated devices", + }, + "configuration": {"device_counts": [1, 2]}, + "rows": [ + { + "device_count": 2, + "iteration_runtime_seconds": 9.6, + "median_cpu_percent": 2600.0, + "peak_cpu_percent": 5500.0, + "substrate": "single-host loopback/emulated devices", + }, + { + "device_count": 1, + "iteration_runtime_seconds": 12.0, + "median_cpu_percent": 2200.0, + "peak_cpu_percent": 5560.0, + "substrate": "single-host loopback/emulated devices", + }, + ], + } + + +def _scaling(mode: str, base_seconds: float, second_seconds: float) -> dict: + def _point(coordinators: int, total_devices: int, seconds: float) -> dict: + return { + "configuration": { + "mode": mode, + "coordinators": coordinators, + "total_devices": total_devices, + }, + "result": { + "coordinators": coordinators, + "iteration_total_seconds": seconds, + "throughput_samples_per_second": 16.0 / seconds, + }, + } + + return { + "schema_version": 1, + "experiment": f"{mode} multi-coordinator scaling", + "points": [ + _point(1, 8 if mode == "strong" else 4, base_seconds), + _point(2, 8, second_seconds), + ], + "conclusion": {"mode": mode, "efficiency": 0.5, "speedup": 1.0}, + } + + +def _breakdown() -> dict: + def _row(mode: str, coordinators: int) -> dict: + # Deliberately scrambled key order to prove canonical reordering. + return { + "idle_other": 1.5, + "coordinators": coordinators, + "optimizer": 1.0, + "mode": mode, + "gradient_sync": 1.5, + "iteration_total": 10.0, + "device_dispatch_aggregation": 6.0, + } + + return { + "schema_version": 1, + "component_semantics": {"gradient_sync": "AllReduce over loopback"}, + "rows": [ + _row("strong", 1), + _row("strong", 2), + _row("weak", 1), + _row("weak", 2), + ], + } + + +def _write(tmp_path: Path, name: str, payload: dict) -> Path: + path = tmp_path / name + path.write_text(json.dumps(payload), encoding="utf-8") + return path + + +# ── Device-scaling parsing (panel a) ──────────────────────────────────────── + + +def test_parse_device_scaling_sorts_points_and_normalizes_cpu() -> None: + # Given a summary with rows out of device-count order + device = schema.parse_device_scaling(_summary(logical_cpu_count=50)) + + # Then points are sorted ascending and CPU is normalized by logical CPUs + assert [p.device_count for p in device.points] == [1, 2] + assert device.points[0].sustained_cpu_percent == pytest.approx(44.0) + assert device.points[1].sustained_cpu_percent == pytest.approx(52.0) + assert device.points[0].iteration_runtime_seconds == pytest.approx(12.0) + assert "loopback" in device.substrate + assert device.logical_cpu_count == 50 + + +def test_parse_device_scaling_missing_runtime_field_raises_schema_error() -> ( + None +): + payload = _summary() + del payload["rows"][0]["iteration_runtime_seconds"] + + with pytest.raises(schema.SchemaError): + schema.parse_device_scaling(payload) + + +def test_parse_device_scaling_missing_logical_cpu_raises_schema_error() -> None: + payload = _summary() + del payload["environment"]["logical_cpu_count"] + + with pytest.raises(schema.SchemaError): + schema.parse_device_scaling(payload) + + +def test_parse_device_scaling_empty_rows_raises_missing_data() -> None: + payload = _summary() + payload["rows"] = [] + + with pytest.raises(schema.MissingDataError): + schema.parse_device_scaling(payload) + + +def test_parse_device_scaling_skips_null_runtime_points() -> None: + payload = _summary() + payload["rows"][0]["iteration_runtime_seconds"] = None + + device = schema.parse_device_scaling(payload) + + assert [p.device_count for p in device.points] == [1] + + +# ── Scaling efficiency parsing (panel b) ──────────────────────────────────── + + +def test_strong_efficiency_computed_per_point_with_unit_baseline() -> None: + # Given strong points: base=10s @1 coord, 6s @2 coords + strong = schema.parse_scaling_efficiency( + _scaling("strong", base_seconds=10.0, second_seconds=6.0), "strong" + ) + + # Then efficiency(1)=1.0 and efficiency(2)=base/(2*meas)=10/12 + assert strong.mode == "strong" + assert [p.coordinators for p in strong.points] == [1, 2] + assert strong.points[0].efficiency == pytest.approx(1.0) + assert strong.points[1].efficiency == pytest.approx(10.0 / 12.0) + + +def test_weak_efficiency_uses_local_baseline() -> None: + weak = schema.parse_scaling_efficiency( + _scaling("weak", base_seconds=8.0, second_seconds=9.0), "weak" + ) + + assert weak.points[0].efficiency == pytest.approx(1.0) + assert weak.points[1].efficiency == pytest.approx(8.0 / 9.0) + + +def test_scaling_efficiency_missing_baseline_raises_missing_data() -> None: + payload = _scaling("strong", 10.0, 6.0) + payload["points"] = [payload["points"][1]] # only coordinators=2 + + with pytest.raises(schema.MissingDataError): + schema.parse_scaling_efficiency(payload, "strong") + + +def test_scaling_efficiency_rejects_unknown_mode() -> None: + with pytest.raises(schema.SchemaError): + schema.parse_scaling_efficiency(_scaling("strong", 10.0, 6.0), "bogus") + + +# ── Breakdown parsing (panel c) ───────────────────────────────────────────── + + +def test_breakdown_orders_components_canonically() -> None: + breakdown = schema.parse_breakdown(_breakdown()) + + assert schema.BREAKDOWN_COMPONENTS == ( + "device_dispatch_aggregation", + "gradient_sync", + "optimizer", + "idle_other", + ) + row = breakdown.rows[0] + # components tuple is aligned to BREAKDOWN_COMPONENTS regardless of key order + assert row.components == pytest.approx((6.0, 1.5, 1.0, 1.5)) + assert [(r.mode, r.coordinators) for r in breakdown.rows] == [ + ("strong", 1), + ("strong", 2), + ("weak", 1), + ("weak", 2), + ] + + +def test_breakdown_components_reconcile_to_iteration_total() -> None: + breakdown = schema.parse_breakdown(_breakdown()) + + row = breakdown.rows[0] + assert sum(row.components) == pytest.approx(row.iteration_total) + + +def test_breakdown_non_reconciling_row_raises_schema_error() -> None: + payload = _breakdown() + payload["rows"][0]["idle_other"] = 99.0 # sum no longer == iteration_total + + with pytest.raises(schema.SchemaError): + schema.parse_breakdown(payload) + + +def test_breakdown_missing_component_raises_schema_error() -> None: + payload = _breakdown() + del payload["rows"][0]["gradient_sync"] + + with pytest.raises(schema.SchemaError): + schema.parse_breakdown(payload) + + +def test_breakdown_empty_rows_raises_missing_data() -> None: + payload = _breakdown() + payload["rows"] = [] + + with pytest.raises(schema.MissingDataError): + schema.parse_breakdown(payload) + + +# ── JSON loading boundary ─────────────────────────────────────────────────── + + +def test_load_json_missing_file_raises_clear_error(tmp_path: Path) -> None: + missing = tmp_path / "nope.json" + + with pytest.raises(schema.SchemaError) as excinfo: + schema.load_json(missing) + + assert "nope.json" in str(excinfo.value) + + +def test_load_json_invalid_json_raises_clear_error(tmp_path: Path) -> None: + bad = tmp_path / "bad.json" + bad.write_text("{not json", encoding="utf-8") + + with pytest.raises(schema.SchemaError): + schema.load_json(bad) + + +def test_load_json_roundtrip(tmp_path: Path) -> None: + path = _write(tmp_path, "summary.json", _summary()) + + assert schema.load_json(path)["schema_version"] == 1 diff --git a/tests/python/unit/test_coordinator_metrics.py b/tests/python/unit/test_coordinator_metrics.py new file mode 100644 index 00000000..62c69c12 --- /dev/null +++ b/tests/python/unit/test_coordinator_metrics.py @@ -0,0 +1,447 @@ +from __future__ import annotations + +import argparse +import importlib.util +import json +import sys +import threading +from pathlib import Path + +import pytest + +MODULE_PATH = ( + Path(__file__).resolve().parents[3] + / "morphling" + / "runtime" + / "coordinator_metrics.py" +) +CLI_MODULE_PATH = MODULE_PATH.with_name("coordinator_metrics_cli.py") +CLI_SPEC = importlib.util.spec_from_file_location( + "morphling.runtime.coordinator_metrics_cli", CLI_MODULE_PATH +) +assert CLI_SPEC is not None and CLI_SPEC.loader is not None +coordinator_metrics_cli = importlib.util.module_from_spec(CLI_SPEC) +sys.modules[CLI_SPEC.name] = coordinator_metrics_cli +CLI_SPEC.loader.exec_module(coordinator_metrics_cli) +MODELS_MODULE_PATH = MODULE_PATH.with_name("coordinator_metrics_models.py") +MODELS_SPEC = importlib.util.spec_from_file_location( + "morphling.runtime.coordinator_metrics_models", MODELS_MODULE_PATH +) +assert MODELS_SPEC is not None and MODELS_SPEC.loader is not None +coordinator_metrics_models = importlib.util.module_from_spec(MODELS_SPEC) +sys.modules[MODELS_SPEC.name] = coordinator_metrics_models +MODELS_SPEC.loader.exec_module(coordinator_metrics_models) +SPEC = importlib.util.spec_from_file_location( + "coordinator_metrics", MODULE_PATH +) +assert SPEC is not None and SPEC.loader is not None +coordinator_metrics = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = coordinator_metrics +SPEC.loader.exec_module(coordinator_metrics) + +MetricReading = coordinator_metrics.MetricReading +NicCounters = coordinator_metrics.NicCounters +PhaseRecorder = coordinator_metrics.PhaseRecorder +PhaseSnapshot = coordinator_metrics.PhaseSnapshot +MetricsConfig = coordinator_metrics.MetricsConfig +CoordinatorMetricsCollector = coordinator_metrics.CoordinatorMetricsCollector +aggregate_nic_counters = coordinator_metrics.aggregate_nic_counters +build_sample = coordinator_metrics.build_sample + + +class FakeSamplingError(RuntimeError): + pass + + +class FakeMetricsReader: + def __init__( + self, + readings: list[MetricReading | Exception], + *, + prime_error: Exception | None = None, + read_attempted: threading.Event | None = None, + ) -> None: + self._readings = iter(readings) + self._prime_error = prime_error + self._read_attempted = read_attempted + + def prime(self) -> None: + if self._prime_error is not None: + raise self._prime_error + + def read(self, _nic_names: tuple[str, ...] | None) -> MetricReading: + reading = next(self._readings) + if isinstance(reading, Exception): + if self._read_attempted is not None: + self._read_attempted.set() + raise reading + return reading + + +def _reading(monotonic_s: float, sent: int, received: int) -> MetricReading: + return MetricReading( + timestamp_unix_s=100.0 + monotonic_s, + monotonic_s=monotonic_s, + process_cpu_percent=25.0, + rss_bytes=1_000, + nic=NicCounters(bytes_sent=sent, bytes_recv=received), + nic_names=("eth0",), + ) + + +def test_metrics_cli_defaults_preserve_disabled_behavior() -> None: + # Given + parser = argparse.ArgumentParser() + coordinator_metrics.add_metrics_arguments(parser) + + # When + args = parser.parse_args([]) + + # Then + assert args.metrics_output is None + assert args.metrics_interval == 1.0 + assert args.metrics_nics is None + + +def test_metrics_cli_parses_output_interval_and_nics(tmp_path: Path) -> None: + # Given + parser = argparse.ArgumentParser() + coordinator_metrics.add_metrics_arguments(parser) + output_path = tmp_path / "coordinator.jsonl" + + # When + args = parser.parse_args( + [ + "--metrics_output", + str(output_path), + "--metrics_interval", + "0.5", + "--metrics_nics", + "ib0", + "eth0", + ] + ) + + # Then + assert args.metrics_output == output_path + assert args.metrics_interval == 0.5 + assert args.metrics_nics == ["ib0", "eth0"] + + +def test_multi_coordinator_output_paths_do_not_collide(tmp_path: Path) -> None: + # Given + output_directory = tmp_path / "metrics" + + # When + rank_zero = coordinator_metrics.metrics_output_path( + output_directory, rank=0 + ) + rank_one = coordinator_metrics.metrics_output_path(output_directory, rank=1) + + # Then + assert rank_zero == output_directory / "coordinator-rank-0.jsonl" + assert rank_one == output_directory / "coordinator-rank-1.jsonl" + assert rank_zero != rank_one + + +def test_jsonl_sample_contains_resource_totals_and_computed_rates() -> None: + # Given + previous = MetricReading( + timestamp_unix_s=100.0, + monotonic_s=10.0, + process_cpu_percent=25.0, + rss_bytes=1_000, + nic=NicCounters(bytes_sent=2_000, bytes_recv=4_000), + nic_names=("eth0",), + ) + current = MetricReading( + timestamp_unix_s=102.0, + monotonic_s=12.0, + process_cpu_percent=50.0, + rss_bytes=1_500, + nic=NicCounters(bytes_sent=2_600, bytes_recv=5_000), + nic_names=("eth0",), + ) + + # When + line = build_sample(current, previous, PhaseSnapshot.empty()).to_json_line() + sample = json.loads(line) + + # Then + assert sample["timestamp_unix_s"] == 102.0 + assert sample["process_cpu_percent"] == 50.0 + assert sample["rss_bytes"] == 1_500 + assert sample["nic_tx_bytes"] == 2_600 + assert sample["nic_rx_bytes"] == 5_000 + assert sample["nic_tx_bytes_per_sec"] == 300.0 + assert sample["nic_rx_bytes_per_sec"] == 500.0 + + +@pytest.mark.parametrize("current_monotonic_s", [10.0, 9.0]) +def test_sample_uses_zero_rates_when_elapsed_is_nonpositive( + current_monotonic_s: float, +) -> None: + # Given + previous = MetricReading( + timestamp_unix_s=100.0, + monotonic_s=10.0, + process_cpu_percent=25.0, + rss_bytes=1_000, + nic=NicCounters(bytes_sent=2_000, bytes_recv=4_000), + nic_names=("eth0",), + ) + current = MetricReading( + timestamp_unix_s=101.0, + monotonic_s=current_monotonic_s, + process_cpu_percent=50.0, + rss_bytes=1_500, + nic=NicCounters(bytes_sent=2_600, bytes_recv=5_000), + nic_names=("eth0",), + ) + + # When + sample = build_sample(current, previous, PhaseSnapshot.empty()) + + # Then + assert sample.nic_tx_bytes_per_sec == 0.0 + assert sample.nic_rx_bytes_per_sec == 0.0 + + +def test_sample_clamps_rates_when_nic_counters_reset() -> None: + # Given + previous = _reading(monotonic_s=10.0, sent=2_000, received=4_000) + current = _reading(monotonic_s=12.0, sent=100, received=200) + + # When + sample = build_sample(current, previous, PhaseSnapshot.empty()) + + # Then + assert sample.nic_tx_bytes_per_sec == 0.0 + assert sample.nic_rx_bytes_per_sec == 0.0 + + +def test_selected_nic_aggregation_is_deterministic() -> None: + # Given + counters = { + "wlan0": NicCounters(bytes_sent=50, bytes_recv=60), + "eth1": NicCounters(bytes_sent=20, bytes_recv=30), + "eth0": NicCounters(bytes_sent=10, bytes_recv=40), + } + + # When + first = aggregate_nic_counters(counters, ("eth1", "eth0")) + second = aggregate_nic_counters(counters, ("eth0", "eth1")) + + # Then + assert first == second + assert first.nic_names == ("eth0", "eth1") + assert first.counters == NicCounters(bytes_sent=30, bytes_recv=70) + + +def test_phase_recorder_accumulates_without_per_event_file_io( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given + recorder = PhaseRecorder() + + def fail_on_file_io(*_args: str, **_kwargs: str) -> None: + raise AssertionError("phase recording must not perform file I/O") + + monkeypatch.setattr(Path, "open", fail_on_file_io) + + # When + recorder.record("dispatch", duration_seconds=1.25, timestamp_unix_s=10.0) + recorder.record("dispatch", duration_seconds=0.75, timestamp_unix_s=12.0) + snapshot = recorder.snapshot() + + # Then + assert snapshot.counts == {"dispatch": 2} + assert snapshot.durations_seconds == {"dispatch": 2.0} + assert snapshot.last_timestamps_unix_s == {"dispatch": 12.0} + + +def test_phase_track_records_when_body_raises() -> None: + # Given + recorder = PhaseRecorder() + + # When + with pytest.raises(RuntimeError, match="training failed"): + with recorder.track("iteration_total"): + raise RuntimeError("training failed") + + # Then + snapshot = recorder.snapshot() + assert snapshot.counts == {"iteration_total": 1} + assert snapshot.durations_seconds["iteration_total"] >= 0.0 + + +def test_named_training_phases_accumulate_without_file_io( + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given + recorder = PhaseRecorder() + phase_names = ( + "iteration_total", + "forward_device_dispatch_aggregation", + "backward_device_dispatch_aggregation", + "gradient_sync", + "optimizer", + "idle_wait", + ) + + def fail_on_file_io(*_args: str, **_kwargs: str) -> None: + raise AssertionError("phase tracking must not perform file I/O") + + monkeypatch.setattr(Path, "open", fail_on_file_io) + + # When + for phase_name in phase_names: + with coordinator_metrics.track_phase(recorder, phase_name): + pass + with coordinator_metrics.track_phase(recorder, "iteration_total"): + pass + snapshot = recorder.snapshot() + + # Then + assert snapshot.counts == { + "iteration_total": 2, + "forward_device_dispatch_aggregation": 1, + "backward_device_dispatch_aggregation": 1, + "gradient_sync": 1, + "optimizer": 1, + "idle_wait": 1, + } + assert set(snapshot.durations_seconds) == set(phase_names) + assert set(snapshot.last_timestamps_unix_s) == set(phase_names) + + +def test_collector_start_stop_writes_jsonl_and_stop_is_idempotent( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given + output_path = tmp_path / "metrics.jsonl" + reader = FakeMetricsReader( + [ + _reading(monotonic_s=10.0, sent=1_000, received=2_000), + _reading(monotonic_s=12.0, sent=1_600, received=3_000), + ] + ) + monkeypatch.setattr( + coordinator_metrics, "_PsutilMetricsReader", lambda: reader + ) + collector = CoordinatorMetricsCollector( + MetricsConfig(output_path=output_path, interval_seconds=60.0) + ) + + # When + collector.start() + collector.stop() + collector.stop() + samples = [ + json.loads(line) for line in output_path.read_text().splitlines() + ] + + # Then + assert [sample["record_type"] for sample in samples] == ["sample", "sample"] + assert samples[1]["nic_tx_bytes_per_sec"] == 300.0 + assert samples[1]["nic_rx_bytes_per_sec"] == 500.0 + + +def test_background_sampling_failure_writes_terminal_error_and_stop_returns( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given + output_path = tmp_path / "metrics.jsonl" + read_attempted = threading.Event() + reader = FakeMetricsReader( + [ + _reading(monotonic_s=10.0, sent=1_000, received=2_000), + FakeSamplingError("NIC sampling failed"), + ], + read_attempted=read_attempted, + ) + monkeypatch.setattr( + coordinator_metrics, "_PsutilMetricsReader", lambda: reader + ) + collector = CoordinatorMetricsCollector( + MetricsConfig(output_path=output_path, interval_seconds=0.001) + ) + + # When + collector.start() + assert read_attempted.wait(timeout=1.0) + collector.stop() + records = [ + json.loads(line) for line in output_path.read_text().splitlines() + ] + + # Then + assert set(records[-1]) == { + "error_message", + "error_type", + "record_type", + "timestamp_unix_s", + } + assert records[-1]["record_type"] == "terminal_error" + assert records[-1]["error_type"] == "FakeSamplingError" + assert records[-1]["error_message"] == "NIC sampling failed" + assert collector.failure is not None + + +def test_start_closes_output_when_priming_fails( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + # Given + output_path = tmp_path / "metrics.jsonl" + reader = FakeMetricsReader( + [], prime_error=FakeSamplingError("CPU priming failed") + ) + monkeypatch.setattr( + coordinator_metrics, "_PsutilMetricsReader", lambda: reader + ) + collector = CoordinatorMetricsCollector( + MetricsConfig(output_path=output_path, interval_seconds=1.0) + ) + + # When + with pytest.raises( + coordinator_metrics.MetricsCollectionError, + match="CPU priming failed", + ): + collector.start() + with output_path.open("a", encoding="utf-8") as output: + output.write("") + + # Then + assert collector._output is None + terminal = json.loads(output_path.read_text().splitlines()[-1]) + assert terminal["record_type"] == "terminal_error" + + +def test_start_helper_reports_priming_failure_without_raising( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], +) -> None: + # Given + reader = FakeMetricsReader( + [], prime_error=FakeSamplingError("CPU priming failed") + ) + monkeypatch.setattr( + coordinator_metrics, "_PsutilMetricsReader", lambda: reader + ) + config = MetricsConfig( + output_path=tmp_path / "metrics.jsonl", interval_seconds=1.0 + ) + + # When + collector = coordinator_metrics.start_metrics_collector( + config, PhaseRecorder() + ) + + # Then + assert collector is None + assert "CPU priming failed" in capsys.readouterr().err diff --git a/tests/python/unit/test_coordinator_metrics_cli.py b/tests/python/unit/test_coordinator_metrics_cli.py new file mode 100644 index 00000000..c087f037 --- /dev/null +++ b/tests/python/unit/test_coordinator_metrics_cli.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +import argparse +import importlib.util +import sys +from pathlib import Path + +import pytest + +MODULE_PATH = ( + Path(__file__).resolve().parents[3] + / "morphling" + / "runtime" + / "coordinator_metrics_cli.py" +) +SPEC = importlib.util.spec_from_file_location( + "coordinator_metrics_cli", MODULE_PATH +) +assert SPEC is not None and SPEC.loader is not None +coordinator_metrics_cli = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = coordinator_metrics_cli +SPEC.loader.exec_module(coordinator_metrics_cli) + + +@pytest.mark.parametrize("interval", ["0", "-0.5"]) +def test_cli_rejects_nonpositive_interval(interval: str) -> None: + # Given + parser = argparse.ArgumentParser() + coordinator_metrics_cli.add_metrics_arguments(parser) + + # When / Then + with pytest.raises(SystemExit): + parser.parse_args(["--metrics_interval", interval]) + + +@pytest.mark.parametrize("interval", [0.0, -0.5]) +def test_config_rejects_nonpositive_interval( + tmp_path: Path, interval: float +) -> None: + # Given / When / Then + with pytest.raises( + coordinator_metrics_cli.MetricsConfigurationError, + match="greater than zero", + ): + coordinator_metrics_cli.MetricsConfig( + output_path=tmp_path / "metrics.jsonl", + interval_seconds=interval, + ) + + +def test_output_path_preserves_explicit_jsonl_file(tmp_path: Path) -> None: + # Given + destination = tmp_path / "metrics.jsonl" + + # When + output_path = coordinator_metrics_cli.metrics_output_path(destination) + + # Then + assert output_path == destination + + +def test_output_path_uses_default_file_for_directory(tmp_path: Path) -> None: + # Given + destination = tmp_path / "metrics" + + # When + output_path = coordinator_metrics_cli.metrics_output_path(destination) + + # Then + assert output_path == destination / "coordinator.jsonl" + + +def test_ranked_output_path_suffixes_explicit_jsonl_file( + tmp_path: Path, +) -> None: + # Given + destination = tmp_path / "metrics.jsonl" + + # When + output_path = coordinator_metrics_cli.metrics_output_path( + destination, rank=2 + ) + + # Then + assert output_path == tmp_path / "metrics-rank-2.jsonl" + + +def test_ranked_output_path_uses_rank_file_for_directory( + tmp_path: Path, +) -> None: + # Given + destination = tmp_path / "metrics.results" + + # When + output_path = coordinator_metrics_cli.metrics_output_path( + destination, rank=2 + ) + + # Then + assert output_path == destination / "coordinator-rank-2.jsonl" + + +def test_config_from_args_keeps_metrics_disabled() -> None: + # Given + args = argparse.Namespace( + metrics_output=None, + metrics_interval=1.0, + metrics_nics=None, + ) + + # When + config = coordinator_metrics_cli.metrics_config_from_args(args) + + # Then + assert config is None + + +def test_config_from_args_resolves_ranked_output_and_nics( + tmp_path: Path, +) -> None: + # Given + args = argparse.Namespace( + metrics_output=tmp_path / "metrics.jsonl", + metrics_interval=0.5, + metrics_nics=["ib0", "eth0"], + ) + + # When + config = coordinator_metrics_cli.metrics_config_from_args(args, rank=3) + + # Then + assert config == coordinator_metrics_cli.MetricsConfig( + output_path=tmp_path / "metrics-rank-3.jsonl", + interval_seconds=0.5, + nic_names=("ib0", "eth0"), + ) diff --git a/tests/python/unit/test_coordinator_scaling_analysis.py b/tests/python/unit/test_coordinator_scaling_analysis.py new file mode 100644 index 00000000..67e54a8f --- /dev/null +++ b/tests/python/unit/test_coordinator_scaling_analysis.py @@ -0,0 +1,308 @@ +from __future__ import annotations + +import json +from pathlib import Path + +from scripts import coordinator_scaling_analysis + + +def _write_jsonl(path: Path, records: list[dict[str, object]]) -> None: + path.write_text( + "".join(f"{json.dumps(record)}\n" for record in records), + encoding="utf-8", + ) + + +def test_parse_metrics_uses_measured_iteration_window(tmp_path: Path) -> None: + # Given + metrics_path = tmp_path / "metrics.jsonl" + _write_jsonl( + metrics_path, + [ + { + "record_type": "sample", + "process_cpu_percent": 10.0, + "rss_bytes": 100, + "nic_tx_bytes_per_sec": 10.0, + "nic_rx_bytes_per_sec": 20.0, + "phase_counts": {}, + "phase_durations_seconds": {}, + }, + { + "record_type": "sample", + "process_cpu_percent": 120.0, + "rss_bytes": 300, + "nic_tx_bytes_per_sec": 300.0, + "nic_rx_bytes_per_sec": 500.0, + "phase_counts": {"iteration_total": 1, "optimizer": 1}, + "phase_durations_seconds": { + "iteration_total": 3.0, + "optimizer": 0.25, + }, + }, + { + "record_type": "sample", + "process_cpu_percent": 80.0, + "rss_bytes": 250, + "nic_tx_bytes_per_sec": 500.0, + "nic_rx_bytes_per_sec": 700.0, + "phase_counts": {"iteration_total": 2, "optimizer": 2}, + "phase_durations_seconds": { + "iteration_total": 8.0, + "optimizer": 0.5, + }, + }, + ], + ) + + # When + metrics = coordinator_scaling_analysis.parse_metrics_jsonl(metrics_path) + + # Then + assert metrics.iteration_runtime_seconds == 4.0 + assert metrics.peak_cpu_percent == 120.0 + assert metrics.median_cpu_percent == 100.0 + assert metrics.peak_rss_bytes == 300 + assert metrics.peak_nic_tx_bytes_per_sec == 500.0 + assert metrics.median_nic_tx_bytes_per_sec == 400.0 + assert metrics.peak_nic_rx_bytes_per_sec == 700.0 + assert metrics.median_nic_rx_bytes_per_sec == 600.0 + assert metrics.phase_counts == {"iteration_total": 2, "optimizer": 2} + assert metrics.phase_durations_seconds == { + "iteration_total": 8.0, + "optimizer": 0.5, + } + + +def test_parse_metrics_retains_terminal_error(tmp_path: Path) -> None: + # Given + metrics_path = tmp_path / "metrics.jsonl" + _write_jsonl( + metrics_path, + [ + { + "record_type": "terminal_error", + "error_type": "OSError", + "error_message": "sampling failed", + "timestamp_unix_s": 1.0, + } + ], + ) + + # When + metrics = coordinator_scaling_analysis.parse_metrics_jsonl(metrics_path) + + # Then + assert metrics.error == "OSError: sampling failed" + assert metrics.iteration_runtime_seconds is None + + +def test_parse_metrics_skips_blank_lines(tmp_path: Path) -> None: + # Given + metrics_path = tmp_path / "metrics.jsonl" + metrics_path.write_text( + "\n \n" + + json.dumps( + { + "record_type": "sample", + "process_cpu_percent": 10.0, + "rss_bytes": 100, + "nic_tx_bytes_per_sec": 1.0, + "nic_rx_bytes_per_sec": 2.0, + "phase_counts": {"iteration_total": 1}, + "phase_durations_seconds": {"iteration_total": 3.0}, + } + ) + + "\n\n", + encoding="utf-8", + ) + + # When + metrics = coordinator_scaling_analysis.parse_metrics_jsonl(metrics_path) + + # Then + assert metrics.error is None + assert metrics.iteration_runtime_seconds == 3.0 + + +def test_parse_metrics_warns_for_trailing_partial_record( + tmp_path: Path, +) -> None: + # Given + metrics_path = tmp_path / "metrics.jsonl" + _write_jsonl( + metrics_path, + [ + { + "record_type": "sample", + "process_cpu_percent": 10.0, + "rss_bytes": 100, + "nic_tx_bytes_per_sec": 1.0, + "nic_rx_bytes_per_sec": 2.0, + "phase_counts": {"iteration_total": 1}, + "phase_durations_seconds": {"iteration_total": 3.0}, + } + ], + ) + with metrics_path.open("a", encoding="utf-8") as output: + output.write('{"record_type":') + + # When + metrics = coordinator_scaling_analysis.parse_metrics_jsonl(metrics_path) + + # Then + assert metrics.error is None + assert metrics.iteration_runtime_seconds == 3.0 + assert len(metrics.warnings) == 1 + assert "trailing partial" in metrics.warnings[0] + + +def test_parse_metrics_rejects_malformed_non_trailing_record( + tmp_path: Path, +) -> None: + # Given + metrics_path = tmp_path / "metrics.jsonl" + metrics_path.write_text( + '{"record_type":\n' + '{"record_type":"sample","process_cpu_percent":10.0}\n', + encoding="utf-8", + ) + + # When + metrics = coordinator_scaling_analysis.parse_metrics_jsonl(metrics_path) + + # Then + assert metrics.iteration_runtime_seconds is None + assert metrics.error is not None + assert "malformed metrics JSONL" in metrics.error + + +def test_aggregate_repetitions_reports_means_and_success_count() -> None: + # Given + first = coordinator_scaling_analysis.RunMeasurement( + iteration_runtime_seconds=4.0, + peak_cpu_percent=120.0, + median_cpu_percent=100.0, + peak_rss_bytes=300, + peak_nic_tx_bytes_per_sec=500.0, + median_nic_tx_bytes_per_sec=400.0, + peak_nic_rx_bytes_per_sec=700.0, + median_nic_rx_bytes_per_sec=600.0, + phase_counts={"iteration_total": 2}, + phase_durations_seconds={"iteration_total": 8.0}, + error=None, + ) + second = coordinator_scaling_analysis.RunMeasurement( + iteration_runtime_seconds=6.0, + peak_cpu_percent=140.0, + median_cpu_percent=120.0, + peak_rss_bytes=500, + peak_nic_tx_bytes_per_sec=700.0, + median_nic_tx_bytes_per_sec=600.0, + peak_nic_rx_bytes_per_sec=900.0, + median_nic_rx_bytes_per_sec=800.0, + phase_counts={"iteration_total": 2}, + phase_durations_seconds={"iteration_total": 12.0}, + error=None, + ) + + # When + aggregate = coordinator_scaling_analysis.aggregate_repetitions( + device_count=4, + measurements=(first, second), + exit_statuses=(0, 0), + ) + + # Then + assert aggregate.device_count == 4 + assert aggregate.repetitions == 2 + assert aggregate.successful_repetitions == 2 + assert aggregate.iteration_runtime_seconds == 5.0 + assert aggregate.peak_cpu_percent == 140.0 + assert aggregate.median_cpu_percent == 110.0 + assert aggregate.phase_counts == {"iteration_total": 2.0} + assert aggregate.phase_durations_seconds == {"iteration_total": 10.0} + + +def test_runtime_plateau_without_sustained_cpu_does_not_name_cpu_bottleneck() -> ( + None +): + # Given + rows = ( + coordinator_scaling_analysis.ScalingPoint(1, 8.0, 100.0, 200.0), + coordinator_scaling_analysis.ScalingPoint(2, 5.0, 150.0, 300.0), + coordinator_scaling_analysis.ScalingPoint(4, 4.8, 200.0, 760.0), + coordinator_scaling_analysis.ScalingPoint(8, 4.7, 220.0, 780.0), + ) + + # When + assessment = coordinator_scaling_analysis.assess_saturation( + rows, logical_cpu_count=8 + ) + + # Then + assert assessment.runtime_plateau_observed is True + assert assessment.resource_saturation_observed is False + assert assessment.bottleneck == ( + "coordinator-side host work / specific resource unisolated" + ) + assert assessment.median_cpu_normalized_percent == 27.5 + assert assessment.peak_cpu_normalized_percent == 97.5 + + +def test_sustained_cpu_and_runtime_plateau_are_reported_separately() -> None: + # Given + rows = ( + coordinator_scaling_analysis.ScalingPoint(1, 8.0, 100.0, 200.0), + coordinator_scaling_analysis.ScalingPoint(2, 5.0, 200.0, 300.0), + coordinator_scaling_analysis.ScalingPoint(4, 4.8, 720.0, 760.0), + coordinator_scaling_analysis.ScalingPoint(8, 4.7, 736.0, 780.0), + ) + + # When + assessment = coordinator_scaling_analysis.assess_saturation( + rows, logical_cpu_count=8 + ) + + # Then + assert assessment.runtime_plateau_observed is True + assert assessment.resource_saturation_observed is True + assert assessment.bottleneck == "coordinator CPU" + + +def test_scaling_without_runtime_plateau_reports_only_measured_range() -> None: + # Given + rows = ( + coordinator_scaling_analysis.ScalingPoint(1, 8.0, 100.0, 200.0), + coordinator_scaling_analysis.ScalingPoint(2, 5.0, 150.0, 300.0), + coordinator_scaling_analysis.ScalingPoint(4, 3.0, 200.0, 400.0), + coordinator_scaling_analysis.ScalingPoint(8, 2.0, 250.0, 500.0), + ) + + # When + assessment = coordinator_scaling_analysis.assess_saturation( + rows, logical_cpu_count=8 + ) + + # Then + assert assessment.runtime_plateau_observed is False + assert assessment.bottleneck == ( + "none within measured range through 8 local emulated devices" + ) + + +def test_response_counter_validation_accepts_partitioned_decrements() -> None: + # Given + records = ( + (0, 1, 3, 2), + (0, 1, 2, 1), + (0, 1, 1, 0), + (1, 1, 2, 1), + (1, 1, 1, 0), + ) + + # When + valid = coordinator_scaling_analysis.response_counters_are_balanced(records) + + # Then + assert valid is True diff --git a/tests/python/unit/test_multi_coordinator_full_model.py b/tests/python/unit/test_multi_coordinator_full_model.py new file mode 100644 index 00000000..14400323 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_full_model.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +from pathlib import Path + +import torch + +from scripts import _single_coordinator_training_workload as shared_workload +from scripts import run_multi_coordinator + + +def test_opt_strong_scaling_keeps_full_model_global_work_fixed() -> None: + # Given / When + configs = run_multi_coordinator.build_scaling_configs( + "strong", run_multi_coordinator.OPT_125M_WORKLOAD + ) + + # Then + assert [ + (config.coordinators, config.local_batch) for config in configs + ] == [ + (1, 2), + (2, 1), + ] + assert {config.total_devices for config in configs} == {8} + assert {config.global_batch for config in configs} == {2} + assert [(config.devices_per_coordinator) for config in configs] == [8, 4] + assert all(config.model_name == "facebook/opt-125m" for config in configs) + assert all(config.sequence_length == 32 for config in configs) + assert all(config.measured_iterations == 2 for config in configs) + assert all(config.block_size == 256 for config in configs) + + +def test_opt_weak_scaling_keeps_full_model_local_work_fixed() -> None: + # Given / When + configs = run_multi_coordinator.build_scaling_configs( + "weak", run_multi_coordinator.OPT_125M_WORKLOAD + ) + + # Then + assert [ + (config.coordinators, config.global_batch) for config in configs + ] == [ + (1, 1), + (2, 2), + ] + assert {config.local_batch for config in configs} == {1} + assert {config.devices_per_coordinator for config in configs} == {4} + assert [(config.total_devices) for config in configs] == [4, 8] + + +def test_full_model_rank_command_carries_model_sequence_and_block_size( + tmp_path: Path, +) -> None: + # Given + config = run_multi_coordinator.build_scaling_configs( + "strong", run_multi_coordinator.OPT_125M_WORKLOAD + )[0] + + # When + launch = run_multi_coordinator.build_rank_launch(config, 0, tmp_path) + + # Then + assert launch.command[launch.command.index("--model-name") + 1] == ( + "facebook/opt-125m" + ) + assert launch.command[launch.command.index("--sequence-length") + 1] == "32" + assert launch.command[launch.command.index("--block_size") + 1] == "256" + assert "--tiny" not in launch.command + + +def test_tiny_rank_command_preserves_quick_workload_mode( + tmp_path: Path, +) -> None: + # Given + config = run_multi_coordinator.build_scaling_configs( + "strong", run_multi_coordinator.TINY_WORKLOAD + )[0] + + # When + launch = run_multi_coordinator.build_rank_launch(config, 0, tmp_path) + + # Then + assert "--tiny" in launch.command + assert config.measured_iterations == 3 + assert config.global_batch == 16 + + +def test_result_provenance_records_model_and_sequence_length() -> None: + # Given + config = run_multi_coordinator.build_scaling_configs( + "weak", run_multi_coordinator.OPT_125M_WORKLOAD + )[0] + result = run_multi_coordinator.build_rank_result( + config, + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations={"iteration_total": 4.0}, + measured_losses=(2.0, 1.0), + golden_losses=(2.0, 1.0), + ), + ) + + # When + payload = run_multi_coordinator.rank_result_payload(config, result) + + # Then + assert payload["model_name"] == "facebook/opt-125m" + assert payload["sequence_length"] == 32 + assert result.throughput_samples_per_second == 0.5 + + +def test_deterministic_token_batch_supports_global_rank_split() -> None: + # Given / When + first = shared_workload.make_token_batch( + shared_workload.TokenBatchSpec(42, 2, 4, 32), + device=torch.device("cpu"), + ) + second = shared_workload.make_token_batch( + shared_workload.TokenBatchSpec(42, 2, 4, 32), + device=torch.device("cpu"), + ) + + # Then + assert torch.equal(first, second) + assert first.shape == (2, 4) + assert torch.equal(torch.cat((first[:1], first[1:])), first) + + +def test_worker_parser_accepts_full_model_flags() -> None: + # Given / When + args = run_multi_coordinator._parser().parse_args( + [ + "--model-name", + "facebook/opt-125m", + "--sequence-length", + "32", + "--block_size", + "256", + ] + ) + + # Then + assert args.model_name == "facebook/opt-125m" + assert args.sequence_length == 32 + assert args.block_size == 256 diff --git a/tests/python/unit/test_multi_coordinator_lifecycle.py b/tests/python/unit/test_multi_coordinator_lifecycle.py new file mode 100644 index 00000000..8cbf7e29 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_lifecycle.py @@ -0,0 +1,155 @@ +from __future__ import annotations + +import dataclasses +import subprocess +from pathlib import Path + +import pytest +import torch + +from morphling.hooks import autograd +from scripts import _single_coordinator_training_workload as workload +from scripts import multi_coordinator_scaling_results as results +from scripts import run_multi_coordinator +from scripts import run_multi_coordinator_scaling as runner + + +class FakeRankProcess: + def __init__(self, *, status: int = 0, timeout: bool = False) -> None: + self.status = status + self.timeout = timeout + self.terminated = False + self.killed = False + self.wait_count = 0 + + def wait(self, timeout: float | None = None) -> int: + self.wait_count += 1 + if self.timeout and self.wait_count == 1: + raise subprocess.TimeoutExpired("rank", timeout) + return self.status + + def terminate_group(self) -> None: + self.terminated = True + + def kill_group(self) -> None: + self.killed = True + + +def test_timeout_terminates_waits_and_kills_all_rank_groups() -> None: + # Given + ranks = (FakeRankProcess(timeout=True), FakeRankProcess()) + + # When / Then + with pytest.raises(subprocess.TimeoutExpired): + runner._wait_for_rank_processes(ranks, timeout_seconds=1) + assert all(rank.terminated and rank.killed for rank in ranks) + assert all(rank.wait_count >= 2 for rank in ranks) + + +def test_rank_error_terminates_all_other_rank_groups() -> None: + # Given + ranks = (FakeRankProcess(status=7), FakeRankProcess()) + + # When / Then + with pytest.raises(subprocess.CalledProcessError): + runner._wait_for_rank_processes(ranks, timeout_seconds=1) + assert all(rank.terminated and rank.killed for rank in ranks) + + +def test_distributed_timeout_flows_to_command_and_config( + tmp_path: Path, +) -> None: + # Given + config = dataclasses.replace( + run_multi_coordinator.build_scaling_configs("strong")[0], + distributed_timeout_seconds=37, + ) + + # When + launch = run_multi_coordinator.build_rank_launch(config, 0, tmp_path) + + # Then + assert ( + launch.command[ + launch.command.index("--distributed-timeout-seconds") + 1 + ] + == "37" + ) + assert dataclasses.asdict(config)["distributed_timeout_seconds"] == 37 + + +class EvalTrackingModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(1)) + self.eval_called = False + + def eval(self) -> EvalTrackingModel: + self.eval_called = True + return self + + def forward( + self, *, input_ids: torch.Tensor, labels: torch.Tensor + ) -> dict[str, torch.Tensor]: + return { + "logits": torch.nn.functional.one_hot(input_ids, 4).float() + * self.weight + } + + +def test_single_coordinator_full_workload_uses_eval_mode() -> None: + # Given + model = EvalTrackingModel() + + # When + workload.train_loop( + model, + steps=1, + seed=1, + batch_size=1, + seq_length=2, + device=torch.device("cpu"), + vocab_size=4, + ) + + # Then + assert model.eval_called is True + + +def test_matching_rising_trajectory_still_passes_equivalence() -> None: + # Given + config = run_multi_coordinator.build_scaling_configs("strong")[0] + + # When + result = results.build_rank_result( + config, + results.RankMeasurement( + 0, {"iteration_total": 1.0}, (1.0, 1.01), (1.0, 1.01) + ), + ) + + # Then + assert result.loss_correctness.passed is True + assert result.loss_correctness.decreasing is False + + +@pytest.mark.parametrize("measured", [(float("nan"), 1.0), (1.0, 2.0)]) +def test_nan_or_mismatch_fails_equivalence( + measured: tuple[float, float], +) -> None: + config = run_multi_coordinator.build_scaling_configs("strong")[0] + result = results.build_rank_result( + config, + results.RankMeasurement( + 0, {"iteration_total": 1.0}, measured, (1.0, 1.0) + ), + ) + assert result.loss_correctness.passed is False + + +def test_reference_grad_input_uses_supported_matmul_signature() -> None: + grad = torch.ones((2, 3)) + weight = torch.ones((3, 4)) + assert torch.equal( + autograd._reference_grad_input(grad, weight), torch.matmul(grad, weight) + ) diff --git a/tests/python/unit/test_multi_coordinator_review_handbacks.py b/tests/python/unit/test_multi_coordinator_review_handbacks.py new file mode 100644 index 00000000..da6e5bd3 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_review_handbacks.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +import dataclasses +import json +from pathlib import Path + +from scripts import _multi_coordinator_training as training +from scripts import run_multi_coordinator +from scripts import run_multi_coordinator_scaling as scaling_runner + + +class FakeFleetProcess: + def __init__(self, *, running: bool = True) -> None: + self.running = running + self.terminated = False + self.killed = False + self.waited = False + + def poll(self) -> int | None: + return None if self.running else 0 + + def terminate(self) -> None: + self.terminated = True + self.running = False + + def kill(self) -> None: + self.killed = True + self.running = False + + def wait(self, timeout: float | None = None) -> int: + self.waited = True + return 0 + + +def _rank_result( + rank: int, *, passed: bool = True +) -> run_multi_coordinator.RankResult: + return run_multi_coordinator.RankResult( + rank=rank, + warmup_iteration_seconds=1.0, + iteration_total_seconds=3.0 + rank, + device_dispatch_aggregation_seconds=1.0, + gradient_sync_seconds=0.5, + optimizer_seconds=0.25, + idle_other_seconds=1.25 + rank, + throughput_samples_per_second=8.0, + measured_losses=(3.0, 2.0, 1.0), + golden_losses=(3.0, 2.0, 1.0), + loss_correctness=run_multi_coordinator.LossCorrectness( + passed=passed, + decreasing=passed, + tracks_golden=passed, + max_relative_error=0.0 if passed else 0.2, + ), + ) + + +def test_summary_serializes_executed_affinity_config(tmp_path: Path) -> None: + # Given + planned = run_multi_coordinator.build_experiment_plan( + modes=("strong",), output_root=tmp_path + ) + executed = run_multi_coordinator.apply_affinity_mode(planned, "taskset") + results = tuple( + run_multi_coordinator.build_global_result( + run.config, (_rank_result(0),) + ) + for run in executed + ) + + # When + scaling_runner._write_summaries(tmp_path, executed, results, ("test",)) + + # Then + summary = json.loads((tmp_path / "strong.json").read_text(encoding="utf-8")) + assert { + point["configuration"]["affinity_mode"] for point in summary["points"] + } == {"taskset"} + + +def test_relative_tolerance_controls_loss_correctness() -> None: + # Given + config = dataclasses.replace( + run_multi_coordinator.build_scaling_configs("strong")[0], + relative_tolerance=0.01, + ) + + # When + result = run_multi_coordinator.build_rank_result( + config, + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations={"iteration_total": 3.0}, + measured_losses=(3.12, 2.08, 1.04), + golden_losses=(3.0, 2.0, 1.0), + ), + ) + + # Then + assert result.loss_correctness.tracks_golden is False + assert result.loss_correctness.passed is False + + +def test_rank_command_and_result_provenance_include_relative_tolerance( + tmp_path: Path, +) -> None: + # Given + config = dataclasses.replace( + run_multi_coordinator.build_scaling_configs("strong")[0], + relative_tolerance=0.0125, + ) + result = run_multi_coordinator.build_rank_result( + config, + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations={"iteration_total": 3.0}, + measured_losses=(3.0, 2.0, 1.0), + golden_losses=(3.0, 2.0, 1.0), + ), + ) + + # When + launch = run_multi_coordinator.build_rank_launch(config, 0, tmp_path) + payload = run_multi_coordinator.rank_result_payload(config, result) + + # Then + assert launch.command[launch.command.index("--rtol") + 1] == "0.0125" + assert payload["relative_tolerance"] == 0.0125 + + +def test_rank_cleanup_terminates_only_owned_fleet_processes() -> None: + # Given + owned = (FakeFleetProcess(), FakeFleetProcess()) + unrelated = FakeFleetProcess() + + # When + training._terminate_fleet(owned) + + # Then + assert all(process.terminated and process.waited for process in owned) + assert unrelated.terminated is False + assert unrelated.waited is False + + +def test_global_correctness_aggregates_every_rank() -> None: + # Given + config = run_multi_coordinator.build_scaling_configs("strong")[1] + + # When + result = run_multi_coordinator.build_global_result( + config, + (_rank_result(0, passed=True), _rank_result(1, passed=False)), + ) + + # Then + assert result.loss_correctness.passed is False + assert result.loss_correctness.decreasing is False + assert result.loss_correctness.tracks_golden is False + assert result.loss_correctness.max_relative_error == 0.2 diff --git a/tests/python/unit/test_multi_coordinator_scaling.py b/tests/python/unit/test_multi_coordinator_scaling.py new file mode 100644 index 00000000..c6ca12e3 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_scaling.py @@ -0,0 +1,338 @@ +from __future__ import annotations + +import dataclasses +from pathlib import Path + +import pytest + +from scripts import coordinator_scaling_analysis, run_multi_coordinator + + +def test_strong_scaling_keeps_global_work_fixed() -> None: + # Given / When + configs = run_multi_coordinator.build_scaling_configs("strong") + + # Then + assert [config.coordinators for config in configs] == [1, 2] + assert {config.total_devices for config in configs} == {8} + assert {config.global_batch for config in configs} == {16} + assert [ + (config.devices_per_coordinator, config.local_batch) + for config in configs + ] == [ + (8, 16), + (4, 8), + ] + assert all(config.warmup_iterations == 1 for config in configs) + assert all(config.measured_iterations == 3 for config in configs) + + +def test_weak_scaling_keeps_local_work_fixed() -> None: + # Given / When + configs = run_multi_coordinator.build_scaling_configs("weak") + + # Then + assert [config.coordinators for config in configs] == [1, 2] + assert {config.devices_per_coordinator for config in configs} == {4} + assert {config.local_batch for config in configs} == {8} + assert [ + (config.total_devices, config.global_batch) for config in configs + ] == [ + (4, 8), + (8, 16), + ] + + +@pytest.mark.parametrize( + ("rank", "expected_node", "expected_range"), + [(0, 0, "0-27"), (1, 1, "28-55")], +) +def test_rank_command_pins_each_rank_to_its_numa_node( + rank: int, + expected_node: int, + expected_range: str, + tmp_path: Path, +) -> None: + # Given + config = run_multi_coordinator.build_scaling_configs("strong")[1] + + # When + launch = run_multi_coordinator.build_rank_launch( + config=config, + rank=rank, + output_directory=tmp_path, + ) + + # Then + assert launch.numa_node == expected_node + assert launch.cpu_range == expected_range + assert launch.command[:4] == ( + "numactl", + f"--cpunodebind={expected_node}", + f"--membind={expected_node}", + "python3", + ) + assert launch.environment["RANK"] == str(rank) + assert launch.environment["WORLD_SIZE"] == "2" + assert launch.substrate == ( + "same-host Gloo loopback lower bound; NUMA CPU and memory binding" + ) + + +def test_affinity_fallback_labels_cpu_first_touch_substrate( + tmp_path: Path, +) -> None: + # Given + config = dataclasses.replace( + run_multi_coordinator.build_scaling_configs("weak")[1], + affinity_mode="taskset", + ) + + # When + launch = run_multi_coordinator.build_rank_launch( + config=config, + rank=1, + output_directory=tmp_path, + ) + + # Then + assert launch.command[:3] == ("taskset", "--cpu-list", "28-55") + assert launch.substrate == ( + "same-host Gloo loopback lower bound; CPU affinity with first-touch memory" + ) + + +def test_affinity_selection_falls_back_when_memory_binding_is_unavailable() -> ( + None +): + # Given / When + privileged = run_multi_coordinator.select_affinity_mode( + numactl_available=True, memory_binding_available=True + ) + restricted = run_multi_coordinator.select_affinity_mode( + numactl_available=True, memory_binding_available=False + ) + + # Then + assert privileged == "numactl" + assert restricted == "taskset" + + +def test_scaling_efficiency_uses_strong_and_weak_definitions() -> None: + # Given / When + strong = coordinator_scaling_analysis.scaling_efficiency( + mode="strong", + baseline_seconds=12.0, + measured_seconds=7.5, + coordinators=2, + ) + weak = coordinator_scaling_analysis.scaling_efficiency( + mode="weak", + baseline_seconds=12.0, + measured_seconds=15.0, + coordinators=2, + ) + + # Then + assert strong == pytest.approx(0.8) + assert weak == pytest.approx(0.8) + + +def test_breakdown_reconciles_device_communication_optimizer_and_residual() -> ( + None +): + # Given / When + breakdown = coordinator_scaling_analysis.reconcile_breakdown( + iteration_total_seconds=10.0, + forward_device_seconds=2.0, + backward_device_seconds=3.0, + gradient_sync_seconds=1.5, + optimizer_seconds=0.5, + ) + + # Then + assert breakdown.device_dispatch_aggregation_seconds == 5.0 + assert breakdown.gradient_sync_seconds == 1.5 + assert breakdown.optimizer_seconds == 0.5 + assert breakdown.idle_other_seconds == 3.0 + assert breakdown.reconciles is True + + +def test_breakdown_rejects_components_exceeding_iteration_total() -> None: + # Given / When / Then + with pytest.raises( + coordinator_scaling_analysis.BreakdownReconciliationError + ): + coordinator_scaling_analysis.reconcile_breakdown( + iteration_total_seconds=1.0, + forward_device_seconds=0.5, + backward_device_seconds=0.4, + gradient_sync_seconds=0.3, + optimizer_seconds=0.2, + ) + + +def test_rank_result_records_warmup_separately_and_measured_throughput() -> ( + None +): + # Given + phases = { + "warmup_iteration": 4.0, + "iteration_total": 6.0, + "forward_device_dispatch_aggregation": 1.0, + "backward_device_dispatch_aggregation": 2.0, + "gradient_sync": 1.0, + "optimizer": 0.5, + } + + # When + result = run_multi_coordinator.build_rank_result( + run_multi_coordinator.build_scaling_configs("strong")[1], + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations=phases, + measured_losses=(2.0, 1.5, 1.0), + golden_losses=(2.0, 1.5, 1.0), + ), + ) + + # Then + assert result.warmup_iteration_seconds == 4.0 + assert result.iteration_total_seconds == 6.0 + assert result.device_dispatch_aggregation_seconds == 3.0 + assert result.gradient_sync_seconds == 1.0 + assert result.optimizer_seconds == 0.5 + assert result.idle_other_seconds == 1.5 + assert result.throughput_samples_per_second == 4.0 + assert result.loss_correctness.passed is True + + +def test_global_result_uses_slowest_rank_as_iteration_makespan() -> None: + # Given + config = run_multi_coordinator.build_scaling_configs("weak")[1] + first = run_multi_coordinator.build_rank_result( + config, + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations={"iteration_total": 6.0}, + measured_losses=(2.0, 1.5, 1.0), + golden_losses=(2.0, 1.5, 1.0), + ), + ) + second = dataclasses.replace( + first, + rank=1, + iteration_total_seconds=8.0, + throughput_samples_per_second=3.0, + ) + + # When + result = run_multi_coordinator.build_global_result(config, (first, second)) + + # Then + assert result.iteration_total_seconds == 8.0 + assert result.throughput_samples_per_second == 6.0 + assert result.loss_correctness.passed is True + assert result.substrate.startswith("same-host Gloo loopback lower bound") + + +def test_iteration_plan_has_one_warmup_before_three_measured_iterations() -> ( + None +): + # Given + config = run_multi_coordinator.build_scaling_configs("strong")[0] + + # When + plan = run_multi_coordinator.build_iteration_plan(config) + + # Then + assert plan == ("warmup", "measured", "measured", "measured") + + +def test_rank_result_payload_contains_required_phase_and_loss_fields() -> None: + # Given + config = run_multi_coordinator.build_scaling_configs("strong")[0] + result = run_multi_coordinator.build_rank_result( + config, + run_multi_coordinator.RankMeasurement( + rank=0, + phase_durations={ + "warmup_iteration": 1.0, + "iteration_total": 3.0, + "forward_device_dispatch_aggregation": 0.5, + "backward_device_dispatch_aggregation": 1.0, + "gradient_sync": 0.5, + "optimizer": 0.25, + }, + measured_losses=(3.0, 2.0, 1.0), + golden_losses=(3.0, 2.0, 1.0), + ), + ) + + # When + payload = run_multi_coordinator.rank_result_payload(config, result) + + # Then + assert payload["warmup_iterations"] == 1 + assert payload["measured_iterations"] == 3 + assert payload["loss_correctness"]["passed"] is True + assert payload["phases_seconds"] == { + "iteration_total": 3.0, + "device_dispatch_aggregation": 1.5, + "gradient_sync": 0.5, + "optimizer": 0.25, + "idle_other": 0.75, + } + + +def test_experiment_plan_covers_strong_and_weak_without_writing_results( + tmp_path: Path, +) -> None: + # Given / When + plan = run_multi_coordinator.build_experiment_plan( + modes=("strong", "weak"), output_root=tmp_path + ) + + # Then + assert [(run.config.mode, run.config.coordinators) for run in plan] == [ + ("strong", 1), + ("strong", 2), + ("weak", 1), + ("weak", 2), + ] + assert len({run.master_port for run in plan}) == 4 + assert plan[0].run_directory == tmp_path / "strong" / "coordinators-1" + assert (tmp_path / "strong").exists() is False + + +def test_scaling_conclusion_reports_speedup_and_efficiency() -> None: + # Given + config_one, config_two = run_multi_coordinator.build_scaling_configs( + "strong" + ) + baseline = run_multi_coordinator.GlobalResult( + coordinators=1, + iteration_total_seconds=12.0, + throughput_samples_per_second=4.0, + loss_correctness=run_multi_coordinator.LossCorrectness( + True, True, True, 0.0 + ), + substrate="same-host Gloo loopback lower bound", + ranks=(), + ) + scaled = dataclasses.replace( + baseline, + coordinators=2, + iteration_total_seconds=7.5, + throughput_samples_per_second=6.4, + ) + + # When + conclusion = run_multi_coordinator.build_scaling_conclusion( + config_one, config_two, baseline, scaled + ) + + # Then + assert conclusion.speedup == pytest.approx(1.6) + assert conclusion.efficiency == pytest.approx(0.8) + assert conclusion.throughput_ratio == pytest.approx(1.6) diff --git a/tests/python/unit/test_run_coordinator_scaling.py b/tests/python/unit/test_run_coordinator_scaling.py new file mode 100644 index 00000000..45e6df74 --- /dev/null +++ b/tests/python/unit/test_run_coordinator_scaling.py @@ -0,0 +1,219 @@ +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +import pytest + +from scripts import coordinator_scaling_execution, run_coordinator_scaling + + +def test_render_proxy_config_updates_device_count_and_port() -> None: + # Given + template = """[network]\nlisten_ip = 0.0.0.0\nlisten_port = 39000\n\n[worker]\nnum_device = 3\n""" + + # When + rendered = coordinator_scaling_execution.render_proxy_config( + template, device_count=8, port=39108 + ) + + # Then + assert "listen_ip = 127.0.0.1" in rendered + assert "listen_port = 39108" in rendered + assert "num_device = 8" in rendered + assert "listen_port = 39000" not in rendered + assert "num_device = 3" not in rendered + + +def test_build_inner_command_uses_proven_d1_settings(tmp_path: Path) -> None: + # Given + run = run_coordinator_scaling.RunSpec( + device_count=4, + repetition=1, + port=39104, + run_directory=tmp_path, + container_directory=Path("/scaling-output/devices-4/rep-1"), + ) + + # When + command = run_coordinator_scaling.build_inner_command( + run, metrics_interval=0.2 + ) + + # Then + assert command[:2] == [ + "python3", + "scripts/run_single_coordinator_training.py", + ] + assert command[command.index("--num_devices") + 1] == "4" + assert command[command.index("--model_name") + 1] == "facebook/opt-125m" + assert command[command.index("--warmup_steps") + 1] == "1" + assert command[command.index("--steps") + 1] == "3" + assert command[command.index("--batch_size") + 1] == "1" + assert command[command.index("--seq_length") + 1] == "32" + assert command[command.index("--block_size") + 1] == "256" + assert command[command.index("--metrics_nics") + 1] == "lo" + assert command[command.index("--cfg") + 1].endswith("proxy.ini") + assert command[command.index("--metrics_output") + 1].endswith( + "metrics.jsonl" + ) + + +def test_ports_are_unique_across_device_counts_and_repetitions() -> None: + # Given + run_count = 48 + + # When + ports = [ + run_coordinator_scaling.port_for_run(index) + for index in range(run_count) + ] + + # Then + assert len(set(ports)) == run_count + assert ports == sorted(ports) + + +def test_docker_command_has_deterministic_unique_name_and_auto_remove( + tmp_path: Path, +) -> None: + # Given + run = run_coordinator_scaling.RunSpec( + device_count=4, + repetition=2, + port=39107, + run_directory=tmp_path / "run", + container_directory=Path("/scaling-output/devices-4/rep-2"), + ) + + # When + command = coordinator_scaling_execution.build_docker_command( + run, tmp_path, "device-emulator:latest", 0.2 + ) + + # Then + assert command[command.index("--name") + 1] == "morphling-d2-d4-r2-p39107" + assert "--rm" in command + + +def test_timeout_forcibly_removes_only_named_docker_container( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Given + run = run_coordinator_scaling.RunSpec( + device_count=4, + repetition=1, + port=39103, + run_directory=tmp_path / "run", + container_directory=Path("/scaling-output/devices-4/rep-1"), + ) + calls: list[list[str]] = [] + original_exists = Path.exists + + def fake_exists(path: Path) -> bool: + if path == Path("/.dockerenv"): + return False + return original_exists(path) + + def fake_run(command: list[str], **_kwargs): + calls.append(command) + if len(calls) == 1: + raise subprocess.TimeoutExpired( + command, 1, output="partial", stderr="late" + ) + return subprocess.CompletedProcess(command, 0, "", "") + + monkeypatch.setattr(Path, "exists", fake_exists) + monkeypatch.setattr( + coordinator_scaling_execution.subprocess, "run", fake_run + ) + + # When + coordinator_scaling_execution.run_command( + run, tmp_path, "device-emulator:latest", 0.2, 1 + ) + + # Then + assert calls[1] == ["docker", "rm", "-f", "morphling-d2-d4-r1-p39103"] + result = json.loads((run.run_directory / "result.json").read_text()) + assert result["timed_out"] is True + + +def test_normal_docker_completion_relies_only_on_auto_remove( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Given + run = run_coordinator_scaling.RunSpec( + device_count=2, + repetition=1, + port=39101, + run_directory=tmp_path / "run", + container_directory=Path("/scaling-output/devices-2/rep-1"), + ) + calls: list[list[str]] = [] + original_exists = Path.exists + + def fake_exists(path: Path) -> bool: + if path == Path("/.dockerenv"): + return False + return original_exists(path) + + def fake_run(command: list[str], **_kwargs): + calls.append(command) + return subprocess.CompletedProcess(command, 0, "ok", "") + + monkeypatch.setattr(Path, "exists", fake_exists) + monkeypatch.setattr( + coordinator_scaling_execution.subprocess, "run", fake_run + ) + + # When + coordinator_scaling_execution.run_command( + run, tmp_path, "device-emulator:latest", 0.2, 1 + ) + + # Then + assert len(calls) == 1 + assert "--rm" in calls[0] + assert calls[0][0:3] != ["docker", "rm", "-f"] + + +def test_summary_keeps_good_repetition_when_other_metrics_are_invalid( + tmp_path: Path, +) -> None: + # Given + for repetition in (1, 2): + run_directory = tmp_path / "devices-1" / f"rep-{repetition}" + run_directory.mkdir(parents=True) + (run_directory / "result.json").write_text( + json.dumps({"exit_status": 0}), encoding="utf-8" + ) + valid = tmp_path / "devices-1" / "rep-1" / "metrics.jsonl" + valid.write_text( + json.dumps( + { + "record_type": "sample", + "process_cpu_percent": 10.0, + "rss_bytes": 100, + "nic_tx_bytes_per_sec": 1.0, + "nic_rx_bytes_per_sec": 2.0, + "phase_counts": {"iteration_total": 1}, + "phase_durations_seconds": {"iteration_total": 3.0}, + } + ) + + "\n", + encoding="utf-8", + ) + invalid = tmp_path / "devices-1" / "rep-2" / "metrics.jsonl" + invalid.write_text("{bad\n{}\n", encoding="utf-8") + + # When + summary_path = run_coordinator_scaling._write_summary( + tmp_path, (1,), 2, "image", ("command",) + ) + summary = json.loads(summary_path.read_text()) + + # Then + assert summary["rows"][0]["successful_repetitions"] == 1 + assert summary["rows"][0]["runs"][1]["measurement"]["error"] is not None diff --git a/tests/python/unit/test_single_coordinator_fleet.py b/tests/python/unit/test_single_coordinator_fleet.py new file mode 100644 index 00000000..43d38bd6 --- /dev/null +++ b/tests/python/unit/test_single_coordinator_fleet.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +import subprocess +from pathlib import Path + +import pytest + +from scripts import _single_coordinator_fleet as fleet + + +class FakeProcess: + def __init__(self, *, requires_kill: bool = False) -> None: + self.requires_kill = requires_kill + self.terminate_calls = 0 + self.kill_calls = 0 + self.wait_calls = 0 + + def poll(self) -> int | None: + return None + + def terminate(self) -> None: + self.terminate_calls += 1 + + def kill(self) -> None: + self.kill_calls += 1 + + def wait(self, timeout: float | None = None) -> int: + self.wait_calls += 1 + if self.requires_kill and self.kill_calls == 0: + raise subprocess.TimeoutExpired("fake", timeout) + return 0 + + +def test_spawn_cleans_only_owned_processes_and_closes_parent_logs( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + # Given + stale = FakeProcess() + created: list[FakeProcess] = [] + parent_logs = [] + fleet._device_processes[:] = [stale] + monkeypatch.setenv("CFG_PATH", str(tmp_path / "proxy.ini")) + monkeypatch.setenv("MORPHLING_DEV_LOG_DIR", str(tmp_path)) + monkeypatch.setattr(fleet.torch.cuda, "device_count", lambda: 2) + + def fail_global_command(*_args, **_kwargs) -> None: + raise AssertionError("fleet lifecycle must not run global commands") + + def fake_popen(*_args, **kwargs): + process = FakeProcess() + created.append(process) + parent_logs.append(kwargs["stdout"]) + return process + + monkeypatch.setattr(fleet.subprocess, "run", fail_global_command) + monkeypatch.setattr(fleet.subprocess, "Popen", fake_popen) + + # When + fleet.spawn_fake_fleet(2, "proxy", "") + + # Then + assert stale.terminate_calls == 1 + assert stale.wait_calls == 1 + assert fleet._device_processes == created + assert all(log.closed for log in parent_logs) + fleet.cleanup_fake_fleet() + assert all(process.terminate_calls == 1 for process in created) + assert fleet._device_processes == [] + + +def test_cleanup_kills_owned_process_that_does_not_terminate() -> None: + # Given + process = FakeProcess(requires_kill=True) + fleet._device_processes[:] = [process] + + # When + fleet.cleanup_fake_fleet() + + # Then + assert process.terminate_calls == 1 + assert process.kill_calls == 1 + assert process.wait_calls == 2 + assert fleet._device_processes == [] diff --git a/tests/python/unit/test_single_coordinator_workload.py b/tests/python/unit/test_single_coordinator_workload.py new file mode 100644 index 00000000..2d634858 --- /dev/null +++ b/tests/python/unit/test_single_coordinator_workload.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import torch + +from morphling.runtime.coordinator_metrics import PhaseRecorder +from scripts._single_coordinator_training_workload import ( + make_tiny_model, + train_loop_tiny, +) + + +def test_tiny_training_excludes_warmup_from_measured_phases() -> None: + # Given + recorder = PhaseRecorder() + model = make_tiny_model(seed=42) + + # When + losses = train_loop_tiny( + model, + steps=3, + warmup_steps=1, + seed=42, + device=torch.device("cpu"), + phase_recorder=recorder, + ) + + # Then + snapshot = recorder.snapshot() + assert len(losses) == 4 + assert snapshot.counts["warmup_iteration"] == 1 + assert snapshot.counts["iteration_total"] == 3 + assert snapshot.counts["forward_device_dispatch_aggregation"] == 3 + assert snapshot.counts["backward_device_dispatch_aggregation"] == 3 + assert snapshot.counts["optimizer"] == 3