From 353a35ddb78e6fdfb76f028ca4a16f6cab8c9c27 Mon Sep 17 00:00:00 2001 From: xly Date: Fri, 10 Jul 2026 10:00:47 +0100 Subject: [PATCH 1/7] build(version): derive package version from git tags via setuptools-scm Replace the hardcoded version = "0.0.1" with a git-tag-driven version. - pyproject.toml: add setuptools-scm>=8 to build requires, make version dynamic, configure [tool.setuptools_scm] with a fallback_version for non-git builds. - morphling/__init__.py: expose __version__ from the generated _version.py, falling back to importlib.metadata. - Dockerfile: the build context excludes .git and installs with --no-build-isolation, so install setuptools-scm and inject the version via SETUPTOOLS_SCM_PRETEND_VERSION_FOR_MORPHLING (empty => fallback_version). - gitignore the generated morphling/_version.py. Verified in-container: docker build --build-arg MORPHLING_VERSION=0.1.0 yields morphling.__version__ == 0.1.0 with _C/_Msg/_GreenCtx importing. --- .gitignore | 3 +++ Dockerfile | 7 +++++++ morphling/__init__.py | 12 +++++++++++- pyproject.toml | 13 +++++++++++-- 4 files changed, 32 insertions(+), 3 deletions(-) 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/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/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 = [] From 9dec2c1b6c0a47d6876e62e290df3f30bcd91f57 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 10:02:44 +0100 Subject: [PATCH 2/7] feat(eval): coordinator scaling evaluation with race-safe live dispatch Add single- and multi-coordinator evaluation over the real dispatch path and fix the concurrency defects that blocked it. Runtime (csrc/backend, morphling/hooks/autograd.py): - DispatchMatMulAsync returns a distinct operation id and callers wait on it, so concurrent GEMMs no longer alias each other's output buffers. - Guard the result hand-off with a mutex and condition variable; WaitMatMul clones the completed tensor and rejects out-of-range operation ids. - Add PartitionTracker::ClaimIdlePartitions and ReassignPartitionToDevice so concurrent SendIdlePartitions callbacks cannot double-send a shard. - Bound the operation id by kMaxLifetimeOperationCount and fail fast instead of writing out of bounds; materialize contiguous forward operands. Evaluation (morphling/runtime, scripts): - Sampled coordinator CPU, RSS, NIC, and phase metrics. - Single-coordinator device-count sweep and one/two-coordinator strong and weak scaling with per-iteration breakdowns. - Three-panel coordinator scaling figure pipeline. Tests: C++ operation-id and partition-claim concurrency tests; Python coordinator metrics, scaling, workload, dispatch-safety, and figure tests. Docs: document the live-dispatch synchronization invariants and the operation-id bound in docs/GEMM_ID_ISSUES.md. --- csrc/backend/operation_id.h | 33 ++ csrc/backend/partition_tracker.cpp | 54 +++ csrc/backend/partition_tracker.h | 3 + csrc/backend/proxy_svr.cpp | 240 +++------- csrc/backend/proxy_svr.h | 14 +- docs/GEMM_ID_ISSUES.md | 36 +- morphling/hooks/autograd.py | 62 +-- morphling/runtime/coordinator_metrics.py | 250 ++++++++++ morphling/runtime/coordinator_metrics_cli.py | 98 ++++ .../runtime/coordinator_metrics_models.py | 117 +++++ scripts/_multi_coordinator_fleet.py | 67 +++ scripts/_multi_coordinator_training.py | 192 ++++++++ scripts/_multi_coordinator_workload.py | 189 ++++++++ scripts/_single_coordinator_fleet.py | 68 +++ .../_single_coordinator_training_workload.py | 165 +++++++ scripts/coord_scaling_plot/__init__.py | 12 + scripts/coord_scaling_plot/__main__.py | 8 + scripts/coord_scaling_plot/_plot_utils.py | 64 +++ scripts/coord_scaling_plot/cli.py | 80 ++++ scripts/coord_scaling_plot/figure.py | 161 +++++++ scripts/coord_scaling_plot/schema.py | 257 ++++++++++ scripts/coordinator_scaling_analysis.py | 244 ++++++++++ scripts/coordinator_scaling_conclusions.py | 71 +++ scripts/coordinator_scaling_execution.py | 178 +++++++ scripts/multi_coordinator_rank_process.py | 57 +++ scripts/multi_coordinator_scaling_analysis.py | 91 ++++ scripts/multi_coordinator_scaling_cli.py | 58 +++ .../multi_coordinator_scaling_conclusions.py | 47 ++ scripts/multi_coordinator_scaling_config.py | 204 ++++++++ .../multi_coordinator_scaling_execution.py | 65 +++ scripts/multi_coordinator_scaling_results.py | 224 +++++++++ scripts/run_coordinator_scaling.py | 223 +++++++++ scripts/run_multi_coordinator.py | 85 ++++ scripts/run_multi_coordinator_scaling.py | 255 ++++++++++ scripts/run_single_coordinator_training.py | 245 ++++++++++ tests/cpp/CMakeLists.txt | 12 + tests/cpp/unit/backend/test_operation_id.cpp | 31 ++ .../backend/test_partition_tracker_claim.cpp | 129 +++++ tests/cpp/unit/zerocopy/CMakeLists.txt | 19 + .../hooks/test_autograd_dispatch_safety.py | 86 ++++ .../test_autograd_greenctx_decoupling.py | 14 +- .../unit/hooks/test_per_gemm_greenctx.py | 20 +- .../unit/test_coord_scaling_figure_cli.py | 157 +++++++ .../python/unit/test_coord_scaling_schema.py | 249 ++++++++++ tests/python/unit/test_coordinator_metrics.py | 439 ++++++++++++++++++ .../unit/test_coordinator_metrics_cli.py | 126 +++++ .../unit/test_coordinator_scaling_analysis.py | 304 ++++++++++++ .../unit/test_multi_coordinator_full_model.py | 139 ++++++ .../unit/test_multi_coordinator_lifecycle.py | 135 ++++++ ...test_multi_coordinator_review_handbacks.py | 157 +++++++ .../unit/test_multi_coordinator_scaling.py | 311 +++++++++++++ .../unit/test_run_coordinator_scaling.py | 203 ++++++++ .../unit/test_single_coordinator_fleet.py | 83 ++++ .../unit/test_single_coordinator_workload.py | 34 ++ 54 files changed, 6649 insertions(+), 216 deletions(-) create mode 100644 csrc/backend/operation_id.h create mode 100644 morphling/runtime/coordinator_metrics.py create mode 100644 morphling/runtime/coordinator_metrics_cli.py create mode 100644 morphling/runtime/coordinator_metrics_models.py create mode 100644 scripts/_multi_coordinator_fleet.py create mode 100644 scripts/_multi_coordinator_training.py create mode 100644 scripts/_multi_coordinator_workload.py create mode 100644 scripts/_single_coordinator_fleet.py create mode 100644 scripts/_single_coordinator_training_workload.py create mode 100644 scripts/coord_scaling_plot/__init__.py create mode 100644 scripts/coord_scaling_plot/__main__.py create mode 100644 scripts/coord_scaling_plot/_plot_utils.py create mode 100644 scripts/coord_scaling_plot/cli.py create mode 100644 scripts/coord_scaling_plot/figure.py create mode 100644 scripts/coord_scaling_plot/schema.py create mode 100644 scripts/coordinator_scaling_analysis.py create mode 100644 scripts/coordinator_scaling_conclusions.py create mode 100644 scripts/coordinator_scaling_execution.py create mode 100644 scripts/multi_coordinator_rank_process.py create mode 100644 scripts/multi_coordinator_scaling_analysis.py create mode 100644 scripts/multi_coordinator_scaling_cli.py create mode 100644 scripts/multi_coordinator_scaling_conclusions.py create mode 100644 scripts/multi_coordinator_scaling_config.py create mode 100644 scripts/multi_coordinator_scaling_execution.py create mode 100644 scripts/multi_coordinator_scaling_results.py create mode 100644 scripts/run_coordinator_scaling.py create mode 100644 scripts/run_multi_coordinator.py create mode 100644 scripts/run_multi_coordinator_scaling.py create mode 100644 scripts/run_single_coordinator_training.py create mode 100644 tests/cpp/unit/backend/test_operation_id.cpp create mode 100644 tests/cpp/unit/backend/test_partition_tracker_claim.cpp create mode 100644 tests/python/unit/hooks/test_autograd_dispatch_safety.py create mode 100644 tests/python/unit/test_coord_scaling_figure_cli.py create mode 100644 tests/python/unit/test_coord_scaling_schema.py create mode 100644 tests/python/unit/test_coordinator_metrics.py create mode 100644 tests/python/unit/test_coordinator_metrics_cli.py create mode 100644 tests/python/unit/test_coordinator_scaling_analysis.py create mode 100644 tests/python/unit/test_multi_coordinator_full_model.py create mode 100644 tests/python/unit/test_multi_coordinator_lifecycle.py create mode 100644 tests/python/unit/test_multi_coordinator_review_handbacks.py create mode 100644 tests/python/unit/test_multi_coordinator_scaling.py create mode 100644 tests/python/unit/test_run_coordinator_scaling.py create mode 100644 tests/python/unit/test_single_coordinator_fleet.py create mode 100644 tests/python/unit/test_single_coordinator_workload.py diff --git a/csrc/backend/operation_id.h b/csrc/backend/operation_id.h new file mode 100644 index 00000000..4af7f321 --- /dev/null +++ b/csrc/backend/operation_id.h @@ -0,0 +1,33 @@ +#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; +} + +} +} 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..739d3bef 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,7 @@ 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 +710,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 +815,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 +835,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 +875,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 +887,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 +909,34 @@ 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 +944,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 +1047,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 +1062,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 +1110,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/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..0dc09d76 --- /dev/null +++ b/morphling/runtime/coordinator_metrics.py @@ -0,0 +1,250 @@ +"""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..e21e0b09 --- /dev/null +++ b/morphling/runtime/coordinator_metrics_models.py @@ -0,0 +1,117 @@ +"""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/scripts/_multi_coordinator_fleet.py b/scripts/_multi_coordinator_fleet.py new file mode 100644 index 00000000..904cc8a2 --- /dev/null +++ b/scripts/_multi_coordinator_fleet.py @@ -0,0 +1,67 @@ +"""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..6d93529e --- /dev/null +++ b/scripts/_multi_coordinator_training.py @@ -0,0 +1,192 @@ +"""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..69c1a452 --- /dev/null +++ b/scripts/_multi_coordinator_workload.py @@ -0,0 +1,189 @@ +"""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..4067349d --- /dev/null +++ b/scripts/_single_coordinator_fleet.py @@ -0,0 +1,68 @@ +"""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..3f7ca3da --- /dev/null +++ b/scripts/_single_coordinator_training_workload.py @@ -0,0 +1,165 @@ +"""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/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..b426688d --- /dev/null +++ b/scripts/coord_scaling_plot/_plot_utils.py @@ -0,0 +1,64 @@ +"""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..2b93dbc4 --- /dev/null +++ b/scripts/coord_scaling_plot/cli.py @@ -0,0 +1,80 @@ +"""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..e18ae47b --- /dev/null +++ b/scripts/coord_scaling_plot/figure.py @@ -0,0 +1,161 @@ +"""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..fbc59829 --- /dev/null +++ b/scripts/coord_scaling_plot/schema.py @@ -0,0 +1,257 @@ +"""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..39f48d37 --- /dev/null +++ b/scripts/coordinator_scaling_analysis.py @@ -0,0 +1,244 @@ +"""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..6a30c155 --- /dev/null +++ b/scripts/coordinator_scaling_conclusions.py @@ -0,0 +1,71 @@ +"""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..c141e9c0 --- /dev/null +++ b/scripts/coordinator_scaling_execution.py @@ -0,0 +1,178 @@ +"""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/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..385e7402 --- /dev/null +++ b/scripts/multi_coordinator_scaling_analysis.py @@ -0,0 +1,91 @@ +"""Pure D5-D7 scaling and phase-breakdown calculations.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal, Sequence, 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..c41ce3e2 --- /dev/null +++ b/scripts/multi_coordinator_scaling_cli.py @@ -0,0 +1,58 @@ +"""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..ad033311 --- /dev/null +++ b/scripts/multi_coordinator_scaling_conclusions.py @@ -0,0 +1,47 @@ +"""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..b1ea9a60 --- /dev/null +++ b/scripts/multi_coordinator_scaling_config.py @@ -0,0 +1,204 @@ +"""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, 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..f376abff --- /dev/null +++ b/scripts/multi_coordinator_scaling_execution.py @@ -0,0 +1,65 @@ +"""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..75c3de54 --- /dev/null +++ b/scripts/multi_coordinator_scaling_results.py @@ -0,0 +1,224 @@ +"""Typed rank and global results for multi-coordinator measurements.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Mapping, TypedDict, 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..8c9f9c11 --- /dev/null +++ b/scripts/run_coordinator_scaling.py @@ -0,0 +1,223 @@ +#!/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..a80a29dc --- /dev/null +++ b/scripts/run_multi_coordinator.py @@ -0,0 +1,85 @@ +#!/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..fb45fe93 --- /dev/null +++ b/scripts/run_multi_coordinator_scaling.py @@ -0,0 +1,255 @@ +#!/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/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..7fedc7f7 --- /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); +} + +} 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..51ff2401 --- /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); +} + +} +} 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/hooks/test_autograd_dispatch_safety.py b/tests/python/unit/hooks/test_autograd_dispatch_safety.py new file mode 100644 index 00000000..a824e6e4 --- /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..73112e4e 100644 --- a/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py +++ b/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py @@ -22,17 +22,21 @@ 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..98863d00 100644 --- a/tests/python/unit/hooks/test_per_gemm_greenctx.py +++ b/tests/python/unit/hooks/test_per_gemm_greenctx.py @@ -16,15 +16,21 @@ class _LocalMatmulBackend: def __init__(self) -> None: - self._queue = [] - - def async_dispatch_matmul(self, mat_a, mat_b) -> None: - self._queue.append(torch.matmul(mat_a, mat_b.transpose(-2, -1))) + self._next_oid = 0 + self._outputs = {} + + 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..ba7d87a4 --- /dev/null +++ b/tests/python/unit/test_coord_scaling_figure_cli.py @@ -0,0 +1,157 @@ +"""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..470a4123 --- /dev/null +++ b/tests/python/unit/test_coord_scaling_schema.py @@ -0,0 +1,249 @@ +"""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..6cae6326 --- /dev/null +++ b/tests/python/unit/test_coordinator_metrics.py @@ -0,0 +1,439 @@ +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..fbb678c4 --- /dev/null +++ b/tests/python/unit/test_coordinator_metrics_cli.py @@ -0,0 +1,126 @@ +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..9ba1117d --- /dev/null +++ b/tests/python/unit/test_coordinator_scaling_analysis.py @@ -0,0 +1,304 @@ +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..93f23096 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_full_model.py @@ -0,0 +1,139 @@ +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..6fc65aa4 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_lifecycle.py @@ -0,0 +1,135 @@ +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..2aa7d937 --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_review_handbacks.py @@ -0,0 +1,157 @@ +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..44c6721b --- /dev/null +++ b/tests/python/unit/test_multi_coordinator_scaling.py @@ -0,0 +1,311 @@ +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..87e2ca65 --- /dev/null +++ b/tests/python/unit/test_run_coordinator_scaling.py @@ -0,0 +1,203 @@ +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 From add15c6ac39be816a48f5e28aa75b226176f1780 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 10:33:04 +0100 Subject: [PATCH 3/7] ci: fix formatting, py3.10 compat, and smoke-test collection - Apply ruff-format (0.9.4) and clang-format to the coordinator sources and tests so the Formatting pre-commit job passes. - Import assert_never compatibly on Python 3.10 (fall back to typing_extensions) in the multi-coordinator scaling modules; typing.assert_never is 3.11+ and the project supports 3.10. - Add tests/python/unit/conftest.py to skip coordinator-evaluation tests when psutil is absent, so the CPU smoke job (pytest/torch/numpy only) no longer aborts collection on module-scope heavy-dependency imports. Full/Docker environments still collect and run them. --- csrc/backend/operation_id.h | 11 +- csrc/backend/proxy_svr.cpp | 7 +- morphling/runtime/coordinator_metrics.py | 10 +- .../runtime/coordinator_metrics_models.py | 16 +- scripts/_multi_coordinator_fleet.py | 23 ++- scripts/_multi_coordinator_training.py | 17 +- scripts/_multi_coordinator_workload.py | 8 +- scripts/_single_coordinator_fleet.py | 27 ++- .../_single_coordinator_training_workload.py | 4 +- scripts/coord_scaling_plot/_plot_utils.py | 22 ++- scripts/coord_scaling_plot/cli.py | 47 +++-- scripts/coord_scaling_plot/figure.py | 177 +++++++++++++----- scripts/coord_scaling_plot/schema.py | 117 ++++++++---- scripts/coordinator_scaling_analysis.py | 46 +++-- scripts/coordinator_scaling_conclusions.py | 11 +- scripts/coordinator_scaling_execution.py | 30 ++- scripts/multi_coordinator_scaling_analysis.py | 11 +- scripts/multi_coordinator_scaling_cli.py | 9 +- .../multi_coordinator_scaling_conclusions.py | 3 +- scripts/multi_coordinator_scaling_config.py | 11 +- .../multi_coordinator_scaling_execution.py | 4 +- scripts/multi_coordinator_scaling_results.py | 11 +- scripts/run_coordinator_scaling.py | 30 ++- scripts/run_multi_coordinator.py | 4 +- scripts/run_multi_coordinator_scaling.py | 13 +- tests/cpp/unit/backend/test_operation_id.cpp | 2 +- .../backend/test_partition_tracker_claim.cpp | 6 +- tests/python/unit/conftest.py | 31 +++ .../hooks/test_autograd_dispatch_safety.py | 8 +- .../test_autograd_greenctx_decoupling.py | 4 +- .../unit/hooks/test_per_gemm_greenctx.py | 4 +- .../unit/test_coord_scaling_figure_cli.py | 89 ++++++--- .../python/unit/test_coord_scaling_schema.py | 43 ++++- tests/python/unit/test_coordinator_metrics.py | 16 +- .../unit/test_coordinator_metrics_cli.py | 20 +- .../unit/test_coordinator_scaling_analysis.py | 8 +- .../unit/test_multi_coordinator_full_model.py | 12 +- .../unit/test_multi_coordinator_lifecycle.py | 36 +++- ...test_multi_coordinator_review_handbacks.py | 4 +- .../unit/test_multi_coordinator_scaling.py | 51 +++-- .../unit/test_run_coordinator_scaling.py | 30 ++- 41 files changed, 743 insertions(+), 290 deletions(-) create mode 100644 tests/python/unit/conftest.py diff --git a/csrc/backend/operation_id.h b/csrc/backend/operation_id.h index 4af7f321..be4d4f2b 100644 --- a/csrc/backend/operation_id.h +++ b/csrc/backend/operation_id.h @@ -11,10 +11,9 @@ 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) + ")"); + throw std::out_of_range("matmul operation id " + std::to_string(oid) + + " is outside the valid range [0, " + + std::to_string(kMaxLifetimeOperationCount) + ")"); } } @@ -29,5 +28,5 @@ inline int ReserveOperationId(std::atomic_int& next_oid) noexcept { return -1; } -} -} +} // namespace backend +} // namespace morphling diff --git a/csrc/backend/proxy_svr.cpp b/csrc/backend/proxy_svr.cpp index 739d3bef..c270e5af 100644 --- a/csrc/backend/proxy_svr.cpp +++ b/csrc/backend/proxy_svr.cpp @@ -659,7 +659,9 @@ void ProxySvrHandle::ConnectionClosedCb(const ConnectionUeventPtr& conn) { /********************************ProxySvrImpl****************************************/ ProxySvrImpl::ProxySvrImpl(ProxyEnvCfg& ctx) - : ctx_(ctx), listener_(nullptr), rsp_cb_counts_(kMaxLifetimeOperationCount) { + : ctx_(ctx), + listener_(nullptr), + rsp_cb_counts_(kMaxLifetimeOperationCount) { // Initialize with greedy scheduling policy by default } @@ -929,7 +931,8 @@ torch::Tensor ProxySvrImpl::WaitMatMul(int oid) { } void ProxySvrImpl::WriteResultBlock(int oid, torch::Tensor& block, int64_t row, - int64_t col, int64_t pivot, int block_size) { + int64_t col, int64_t pivot, + int block_size) { std::lock_guard lock(outputs_mutex_); IndexPutMatrixBlock(outputs_[oid], block, row, col, pivot, block_size); } diff --git a/morphling/runtime/coordinator_metrics.py b/morphling/runtime/coordinator_metrics.py index 0dc09d76..7ddd15f2 100644 --- a/morphling/runtime/coordinator_metrics.py +++ b/morphling/runtime/coordinator_metrics.py @@ -69,9 +69,7 @@ def snapshot(self) -> PhaseSnapshot: @contextmanager -def track_phase( - recorder: PhaseRecorder | None, name: str -) -> Iterator[None]: +def track_phase(recorder: PhaseRecorder | None, name: str) -> Iterator[None]: """Track a phase only when metrics collection is enabled.""" if recorder is None: yield @@ -212,7 +210,11 @@ def _record_failure(self, error: Exception) -> MetricsCollectionError: file=sys.stderr, flush=True, ) - print(f"Coordinator metrics stopped: {failure}", file=sys.stderr, flush=True) + print( + f"Coordinator metrics stopped: {failure}", + file=sys.stderr, + flush=True, + ) return failure def _close_output(self) -> None: diff --git a/morphling/runtime/coordinator_metrics_models.py b/morphling/runtime/coordinator_metrics_models.py index e21e0b09..88176443 100644 --- a/morphling/runtime/coordinator_metrics_models.py +++ b/morphling/runtime/coordinator_metrics_models.py @@ -72,16 +72,24 @@ def aggregate_nic_counters( 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) + 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), + bytes_sent=sum( + counters[name].bytes_sent for name in selected_names + ), + bytes_recv=sum( + counters[name].bytes_recv for name in selected_names + ), ), ) diff --git a/scripts/_multi_coordinator_fleet.py b/scripts/_multi_coordinator_fleet.py index 904cc8a2..d6ce5bfa 100644 --- a/scripts/_multi_coordinator_fleet.py +++ b/scripts/_multi_coordinator_fleet.py @@ -32,10 +32,25 @@ def spawn_fleet( (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), + "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" diff --git a/scripts/_multi_coordinator_training.py b/scripts/_multi_coordinator_training.py index 6d93529e..2d6096d4 100644 --- a/scripts/_multi_coordinator_training.py +++ b/scripts/_multi_coordinator_training.py @@ -81,18 +81,24 @@ def _write_result( config = _scaling_config(args) result = build_rank_result( config, - RankMeasurement(rank, recorder.snapshot().durations_seconds, losses, golden), + 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) + 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: +def _write_failure( + args: argparse.Namespace, rank: int, error: Exception +) -> None: if args.result_output is None: return path = Path(args.result_output) @@ -144,7 +150,10 @@ def run_coordinator(rank: int, args: argparse.Namespace) -> int: root = Path(__file__).resolve().parents[1] rank_config = _write_rank_config( - root / "config/proxy/svr.ini", rank, proxy_port, args.devices_per_coord + 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 diff --git a/scripts/_multi_coordinator_workload.py b/scripts/_multi_coordinator_workload.py index 69c1a452..73e29ab1 100644 --- a/scripts/_multi_coordinator_workload.py +++ b/scripts/_multi_coordinator_workload.py @@ -30,7 +30,9 @@ class ModelStep: loss_function: LossFunction -def _optimizer(config: ScalingConfig, model: torch.nn.Module) -> torch.optim.Optimizer: +def _optimizer( + config: ScalingConfig, model: torch.nn.Module +) -> torch.optim.Optimizer: return torch.optim.AdamW(model.parameters(), lr=config.learning_rate) @@ -95,7 +97,9 @@ def _tiny_model(seed: int) -> torch.nn.Module: ) -def _tiny_inputs(seed: int, batch_size: int) -> tuple[torch.Tensor, torch.Tensor]: +def _tiny_inputs( + seed: int, batch_size: int +) -> tuple[torch.Tensor, torch.Tensor]: generator = torch.Generator(device="cpu") generator.manual_seed(seed) return ( diff --git a/scripts/_single_coordinator_fleet.py b/scripts/_single_coordinator_fleet.py index 4067349d..371e0358 100644 --- a/scripts/_single_coordinator_fleet.py +++ b/scripts/_single_coordinator_fleet.py @@ -38,15 +38,24 @@ def spawn_fake_fleet( 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, + "--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] diff --git a/scripts/_single_coordinator_training_workload.py b/scripts/_single_coordinator_training_workload.py index 3f7ca3da..7b7b2b8e 100644 --- a/scripts/_single_coordinator_training_workload.py +++ b/scripts/_single_coordinator_training_workload.py @@ -43,9 +43,7 @@ def extract_loss(outputs: Any, input_ids: torch.Tensor) -> torch.Tensor: 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) - ) + 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): diff --git a/scripts/coord_scaling_plot/_plot_utils.py b/scripts/coord_scaling_plot/_plot_utils.py index b426688d..af2ea92f 100644 --- a/scripts/coord_scaling_plot/_plot_utils.py +++ b/scripts/coord_scaling_plot/_plot_utils.py @@ -25,18 +25,23 @@ matplotlib.use("Agg", force=True) _DEFAULT_SCRIPTS = Path( - "/home/xly/.opencode/skills/conference-plot/scripts/plot_utils.py") + "/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 + 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") + f"{candidate}; set CONFERENCE_PLOT_SCRIPTS to its scripts directory" + ) spec = importlib.util.spec_from_file_location( - "conference_plot_plot_utils", candidate) + "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) @@ -59,6 +64,11 @@ def save_paper_figure(fig, out_base: Path) -> Tuple[Path, Path]: 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") + 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 index 2b93dbc4..6c857146 100644 --- a/scripts/coord_scaling_plot/cli.py +++ b/scripts/coord_scaling_plot/cli.py @@ -30,24 +30,41 @@ 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}") + 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") + 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 @@ -65,7 +82,9 @@ def main(argv: Optional[Sequence[str]] = None) -> int: figure, # imported lazily so schema-only errors avoid matplotlib ) - pdf_path, png_path = figure.render(device, strong, weak, breakdown, args.out) + pdf_path, png_path = figure.render( + device, strong, weak, breakdown, args.out + ) _report(pdf_path, png_path) return 0 diff --git a/scripts/coord_scaling_plot/figure.py b/scripts/coord_scaling_plot/figure.py index e18ae47b..452b7ddb 100644 --- a/scripts/coord_scaling_plot/figure.py +++ b/scripts/coord_scaling_plot/figure.py @@ -29,7 +29,11 @@ from ._plot_utils import HATCHES, WONG_PALETTE, paper_style, save_paper_figure COMPONENT_LABELS: Tuple[str, ...] = ( - "Dispatch+agg", "Gradient sync", "Optimizer", "Idle/other") + "Dispatch+agg", + "Gradient sync", + "Optimizer", + "Idle/other", +) _FIG_WIDTH_INCHES = 7.0 _FIG_HEIGHT_INCHES = 2.4 @@ -40,7 +44,11 @@ _STRONG_COLOR = WONG_PALETTE[1] _WEAK_COLOR = WONG_PALETTE[3] _COMPONENT_COLORS = ( - WONG_PALETTE[2], WONG_PALETTE[3], WONG_PALETTE[6], WONG_PALETTE[7]) + WONG_PALETTE[2], + WONG_PALETTE[3], + WONG_PALETTE[6], + WONG_PALETTE[7], +) _COMPONENT_HATCHES = (HATCHES[1], HATCHES[2], HATCHES[3], HATCHES[4]) @@ -49,16 +57,29 @@ def _panel_device(ax, device: schema.DeviceScaling) -> None: 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") + 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") + 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.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) @@ -69,30 +90,67 @@ def _panel_device(ax, device: schema.DeviceScaling) -> None: 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") + 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: + 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: + 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)}) + 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) @@ -105,14 +163,24 @@ def _panel_efficiency(ax, strong: schema.ScalingEfficiency, 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] + 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) + 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) @@ -121,13 +189,21 @@ def _panel_breakdown(ax, breakdown: schema.Breakdown) -> None: 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): + 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) @@ -139,21 +215,30 @@ def _assemble(device: schema.DeviceScaling, strong: schema.ScalingEfficiency, return fig -def build_figure(device: schema.DeviceScaling, strong: schema.ScalingEfficiency, - weak: schema.ScalingEfficiency, - breakdown: schema.Breakdown): +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): + 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]: +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): + 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) diff --git a/scripts/coord_scaling_plot/schema.py b/scripts/coord_scaling_plot/schema.py index fbc59829..c0783532 100644 --- a/scripts/coord_scaling_plot/schema.py +++ b/scripts/coord_scaling_plot/schema.py @@ -56,6 +56,7 @@ class MissingDataError(SchemaError): # ── Typed, immutable views ────────────────────────────────────────────────── + @dataclass(frozen=True) class DeviceScalingPoint: device_count: int @@ -101,9 +102,12 @@ class Breakdown: # ── 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__}") + raise SchemaError( + f"{ctx}: expected an object, got {type(value).__name__}" + ) return value @@ -144,15 +148,21 @@ def load_json(path: Any) -> Any: # ── 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") + 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") + raise SchemaError( + "summary.environment.logical_cpu_count must be positive" + ) rows = _sequence(_require(root, "rows", "summary"), "summary.rows") points = [] @@ -165,13 +175,17 @@ def parse_device_scaling(data: Any) -> DeviceScaling: 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), - )) + 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") @@ -186,7 +200,9 @@ def parse_device_scaling(data: Any) -> DeviceScaling: 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}") + raise SchemaError( + f"unknown scaling mode {mode!r}; expected {SCALING_MODES}" + ) root = _mapping(data, mode) entries = _sequence(_require(root, "points", mode), f"{mode}.points") @@ -194,33 +210,49 @@ def parse_scaling_efficiency(data: Any, mode: str) -> ScalingEfficiency: 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") + 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"))) + 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") + f"{mode} scaling lacks the single-coordinator baseline point" + ) points = [] - for coordinators, seconds, throughput in sorted(raw, key=lambda item: item[0]): + 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)) + points.append( + EfficiencyPoint(coordinators, seconds, efficiency, throughput) + ) return ScalingEfficiency(mode=mode, points=tuple(points)) @@ -237,21 +269,30 @@ def parse_breakdown(data: Any) -> Breakdown: 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) + _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)): + _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, - )) + 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()) + (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 index 39f48d37..d94a4876 100644 --- a/scripts/coordinator_scaling_analysis.py +++ b/scripts/coordinator_scaling_analysis.py @@ -96,7 +96,9 @@ def parse_metrics_jsonl(path: Path) -> RunMeasurement: try: records.append(json.loads(line)) except json.JSONDecodeError as error: - message = f"malformed metrics JSONL at line {line_number}: {error.msg}" + message = ( + f"malformed metrics JSONL at line {line_number}: {error.msg}" + ) if position != len(numbered_lines) - 1: return _empty_measurement(message) warnings.append( @@ -104,9 +106,13 @@ def parse_metrics_jsonl(path: Path) -> RunMeasurement: f"{error.msg}" ) terminal_errors = [ - record for record in records if record.get("record_type") == "terminal_error" + record + for record in records + if record.get("record_type") == "terminal_error" + ] + samples = [ + record for record in records if record.get("record_type") == "sample" ] - samples = [record for record in records if record.get("record_type") == "sample"] measured = [ sample for sample in samples @@ -141,7 +147,9 @@ def parse_metrics_jsonl(path: Path) -> 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), + 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 ), @@ -190,14 +198,19 @@ def aggregate_repetitions( for measurement in successful for name in measurement.phase_durations_seconds } - peak_cpu_values = [measurement.peak_cpu_percent for measurement in successful] + 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] + [ + 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( @@ -214,9 +227,14 @@ def aggregate_repetitions( for measurement in successful ), default=0.0, - ) if successful else None, + ) + if successful + else None, median_nic_tx_bytes_per_sec=_mean_present( - [measurement.median_nic_tx_bytes_per_sec for measurement in successful] + [ + measurement.median_nic_tx_bytes_per_sec + for measurement in successful + ] ), peak_nic_rx_bytes_per_sec=max( ( @@ -224,13 +242,19 @@ def aggregate_repetitions( for measurement in successful ), default=0.0, - ) if successful else None, + ) + if successful + else None, median_nic_rx_bytes_per_sec=_mean_present( - [measurement.median_nic_rx_bytes_per_sec for measurement in successful] + [ + 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 + measurement.phase_counts.get(name, 0) + for measurement in successful ) for name in sorted(phase_names) }, diff --git a/scripts/coordinator_scaling_conclusions.py b/scripts/coordinator_scaling_conclusions.py index 6a30c155..99db7e14 100644 --- a/scripts/coordinator_scaling_conclusions.py +++ b/scripts/coordinator_scaling_conclusions.py @@ -28,12 +28,8 @@ class SaturationAssessment: 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" - ) + 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, @@ -48,7 +44,8 @@ def assess_saturation( 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 + current.iteration_runtime_seconds / previous.iteration_runtime_seconds + >= 0.95 ) resource_saturation = median_normalized >= 90.0 if runtime_plateau and resource_saturation: diff --git a/scripts/coordinator_scaling_execution.py b/scripts/coordinator_scaling_execution.py index c141e9c0..5e5f8110 100644 --- a/scripts/coordinator_scaling_execution.py +++ b/scripts/coordinator_scaling_execution.py @@ -74,20 +74,30 @@ def build_inner_command(run: RunSpec, metrics_interval: float) -> list[str]: def docker_container_name(run: RunSpec) -> str: - return ( - f"morphling-d2-d{run.device_count}-r{run.repetition}-p{run.port}" - ) + 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), + "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), ] @@ -119,7 +129,9 @@ def run_command( environment = os.environ.copy() environment["CFG_PATH"] = str(run.container_directory / "proxy.ini") else: - command = build_docker_command(run, output_root, image, metrics_interval) + command = build_docker_command( + run, output_root, image, metrics_interval + ) environment = None started = _utc_now() timed_out = False diff --git a/scripts/multi_coordinator_scaling_analysis.py b/scripts/multi_coordinator_scaling_analysis.py index 385e7402..4876076c 100644 --- a/scripts/multi_coordinator_scaling_analysis.py +++ b/scripts/multi_coordinator_scaling_analysis.py @@ -3,7 +3,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Literal, Sequence, assert_never +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] @@ -51,9 +56,7 @@ def response_counters_are_balanced( 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) - ) + grouped.setdefault(operation_id, []).append((count, previous, current)) if not grouped: return False return all( diff --git a/scripts/multi_coordinator_scaling_cli.py b/scripts/multi_coordinator_scaling_cli.py index c41ce3e2..320d004b 100644 --- a/scripts/multi_coordinator_scaling_cli.py +++ b/scripts/multi_coordinator_scaling_cli.py @@ -27,9 +27,14 @@ class ScalingCliConfig: def parse_scaling_cli() -> ScalingCliConfig: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( - "--modes", nargs="+", choices=("strong", "weak"), default=("strong", "weak") + "--modes", + nargs="+", + choices=("strong", "weak"), + default=("strong", "weak"), + ) + parser.add_argument( + "--output-dir", type=Path, default=Path("results/coord_scaling") ) - 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") diff --git a/scripts/multi_coordinator_scaling_conclusions.py b/scripts/multi_coordinator_scaling_conclusions.py index ad033311..7a4fd884 100644 --- a/scripts/multi_coordinator_scaling_conclusions.py +++ b/scripts/multi_coordinator_scaling_conclusions.py @@ -41,7 +41,6 @@ def build_scaling_conclusion( / baseline.throughput_samples_per_second ), correctness_passed=( - baseline.loss_correctness.passed - and scaled.loss_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 index b1ea9a60..d71f8bf8 100644 --- a/scripts/multi_coordinator_scaling_config.py +++ b/scripts/multi_coordinator_scaling_config.py @@ -4,7 +4,12 @@ from dataclasses import dataclass from pathlib import Path -from typing import Final, Literal, Mapping, assert_never +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"] @@ -178,9 +183,7 @@ def build_rank_launch( f"--cpunodebind={node}", f"--membind={node}", ) - substrate = ( - "same-host Gloo loopback lower bound; NUMA CPU and memory binding" - ) + substrate = "same-host Gloo loopback lower bound; NUMA CPU and memory binding" case "taskset": prefix = ("taskset", "--cpu-list", cpu_range) substrate = ( diff --git a/scripts/multi_coordinator_scaling_execution.py b/scripts/multi_coordinator_scaling_execution.py index f376abff..f365687c 100644 --- a/scripts/multi_coordinator_scaling_execution.py +++ b/scripts/multi_coordinator_scaling_execution.py @@ -57,9 +57,7 @@ def apply_affinity_mode( return tuple( dataclasses.replace( run, - config=dataclasses.replace( - run.config, affinity_mode=affinity_mode - ), + 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 index 75c3de54..25c45073 100644 --- a/scripts/multi_coordinator_scaling_results.py +++ b/scripts/multi_coordinator_scaling_results.py @@ -4,7 +4,12 @@ import math from dataclasses import dataclass -from typing import Mapping, TypedDict, assert_never +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 @@ -165,7 +170,9 @@ def build_global_result( 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), + 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 ), diff --git a/scripts/run_coordinator_scaling.py b/scripts/run_coordinator_scaling.py index 8c9f9c11..9931020a 100644 --- a/scripts/run_coordinator_scaling.py +++ b/scripts/run_coordinator_scaling.py @@ -57,10 +57,14 @@ def _utc_now() -> str: 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") + 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]]: +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" @@ -87,7 +91,9 @@ def _write_summary( scaling_points: list[ScalingPoint] = [] for device_count in device_counts: loaded = [ - _load_run(output_root / f"devices-{device_count}" / f"rep-{repetition}") + _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) @@ -172,9 +178,13 @@ def _write_summary( def main() -> int: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--device-counts", nargs="+", type=int, default=[1, 2, 4, 8]) + 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( + "--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) @@ -193,10 +203,16 @@ def main() -> int: repetition, port_for_run(run_index), args.output_dir / relative, - (args.output_dir / relative) if inside_docker else Path("/scaling-output") / 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))) + print( + " ".join( + build_inner_command(run, args.metrics_interval) + ) + ) else: run_command( run, diff --git a/scripts/run_multi_coordinator.py b/scripts/run_multi_coordinator.py index a80a29dc..0264b956 100644 --- a/scripts/run_multi_coordinator.py +++ b/scripts/run_multi_coordinator.py @@ -44,7 +44,9 @@ 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( + "--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) diff --git a/scripts/run_multi_coordinator_scaling.py b/scripts/run_multi_coordinator_scaling.py index fb45fe93..c9559984 100644 --- a/scripts/run_multi_coordinator_scaling.py +++ b/scripts/run_multi_coordinator_scaling.py @@ -72,7 +72,9 @@ def _load_rank_result(path: Path) -> RankResult: throughput_samples_per_second=float( payload["throughput_samples_per_second"] ), - measured_losses=tuple(float(value) for value in payload["measured_losses"]), + 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"]), @@ -93,9 +95,12 @@ def _rank_commands( commands.append( launch.command + ( - "--scaling_mode", config.mode, - "--master_port", str(run.master_port), - "--base_proxy_port", str(run.base_proxy_port), + "--scaling_mode", + config.mode, + "--master_port", + str(run.master_port), + "--base_proxy_port", + str(run.base_proxy_port), ) ) return tuple(commands) diff --git a/tests/cpp/unit/backend/test_operation_id.cpp b/tests/cpp/unit/backend/test_operation_id.cpp index 7fedc7f7..076e3e14 100644 --- a/tests/cpp/unit/backend/test_operation_id.cpp +++ b/tests/cpp/unit/backend/test_operation_id.cpp @@ -28,4 +28,4 @@ TEST(OperationIdTest, ReservesLastAvailableOidWithoutExceedingCapacity) { 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 index 51ff2401..12faee9f 100644 --- a/tests/cpp/unit/backend/test_partition_tracker_claim.cpp +++ b/tests/cpp/unit/backend/test_partition_tracker_claim.cpp @@ -50,7 +50,7 @@ TEST(PartitionTrackerClaimTest, ConcurrentClaimersClaimEachPartitionOnce) { // Then const std::unordered_set unique_keys(claimed_keys.begin(), - claimed_keys.end()); + claimed_keys.end()); EXPECT_EQ(claimed_keys.size(), kPartitionCount); EXPECT_EQ(unique_keys.size(), kPartitionCount); EXPECT_TRUE(tracker.GetIdlePartitions().empty()); @@ -125,5 +125,5 @@ TEST(PartitionTrackerClaimTest, ReassignmentRejectsUnknownPartition) { EXPECT_FALSE(reassigned); } -} -} +} // namespace +} // namespace morphling::backend 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 index a824e6e4..47ab1e6f 100644 --- a/tests/python/unit/hooks/test_autograd_dispatch_safety.py +++ b/tests/python/unit/hooks/test_autograd_dispatch_safety.py @@ -21,9 +21,7 @@ def async_dispatch_matmul( 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) - ) + self._outputs[oid] = torch.matmul(mat_a, mat_b.transpose(-2, -1)) return oid def wait_matmul(self, oid: int) -> torch.Tensor: @@ -68,7 +66,9 @@ def test_backward_rejects_failed_dispatch_before_wait() -> None: assert backend.waited_ids == [0] -def test_forward_dispatches_contiguous_operands_with_native_linear_parity() -> None: +def test_forward_dispatches_contiguous_operands_with_native_linear_parity() -> ( + None +): backend = _SequenceBackend([7]) ag.set_backend(backend) input_tensor = torch.randn(3, 4) diff --git a/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py b/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py index 73112e4e..b050a0b4 100644 --- a/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py +++ b/tests/python/unit/hooks/test_autograd_greenctx_decoupling.py @@ -29,9 +29,7 @@ def async_dispatch_matmul( ) -> int: oid = self.dispatch_calls self.dispatch_calls += 1 - self._outputs[oid] = 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, oid: int) -> torch.Tensor: diff --git a/tests/python/unit/hooks/test_per_gemm_greenctx.py b/tests/python/unit/hooks/test_per_gemm_greenctx.py index 98863d00..4bc2e2ae 100644 --- a/tests/python/unit/hooks/test_per_gemm_greenctx.py +++ b/tests/python/unit/hooks/test_per_gemm_greenctx.py @@ -22,9 +22,7 @@ def __init__(self) -> None: 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) - ) + self._outputs[oid] = torch.matmul(mat_a, mat_b.transpose(-2, -1)) return oid def wait_matmul(self, oid: int): diff --git a/tests/python/unit/test_coord_scaling_figure_cli.py b/tests/python/unit/test_coord_scaling_figure_cli.py index ba7d87a4..579050bd 100644 --- a/tests/python/unit/test_coord_scaling_figure_cli.py +++ b/tests/python/unit/test_coord_scaling_figure_cli.py @@ -17,6 +17,7 @@ # ── Fixture builders (mirror the real emitter field names) ────────────────── + def _summary() -> dict: def _row(device_count: int, seconds: float, cpu: float) -> dict: return { @@ -32,8 +33,12 @@ def _row(device_count: int, seconds: float, cpu: float) -> dict: "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)], + "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), + ], } @@ -48,29 +53,41 @@ def _point(coordinators: int, seconds: float) -> dict: }, } - return {"points": [_point(1, base), _point(2, second)], - "conclusion": {"mode": mode, "efficiency": base / second}} + 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, + "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)], + "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") + 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 @@ -84,7 +101,9 @@ def _write_inputs(tmp_path: Path) -> dict: "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["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 @@ -92,6 +111,7 @@ def _write_inputs(tmp_path: Path) -> dict: # ── Figure structure ──────────────────────────────────────────────────────── + def test_build_figure_has_three_panels_and_no_figure_title() -> None: import matplotlib.pyplot as plt @@ -114,7 +134,9 @@ def test_breakdown_panel_legend_uses_canonical_component_order() -> None: 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)")) + 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) @@ -124,15 +146,23 @@ def test_breakdown_panel_legend_uses_canonical_component_order() -> None: # ── 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"), - ]) + 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 @@ -142,13 +172,20 @@ 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), - ]) + 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") diff --git a/tests/python/unit/test_coord_scaling_schema.py b/tests/python/unit/test_coord_scaling_schema.py index 470a4123..c37cf075 100644 --- a/tests/python/unit/test_coord_scaling_schema.py +++ b/tests/python/unit/test_coord_scaling_schema.py @@ -18,6 +18,7 @@ # ── Fixture builders ──────────────────────────────────────────────────────── + def _summary(logical_cpu_count: int = 50) -> dict: return { "schema_version": 1, @@ -48,8 +49,11 @@ def _summary(logical_cpu_count: int = 50) -> dict: 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}, + "configuration": { + "mode": mode, + "coordinators": coordinators, + "total_devices": total_devices, + }, "result": { "coordinators": coordinators, "iteration_total_seconds": seconds, @@ -84,8 +88,12 @@ def _row(mode: str, coordinators: int) -> dict: return { "schema_version": 1, "component_semantics": {"gradient_sync": "AllReduce over loopback"}, - "rows": [_row("strong", 1), _row("strong", 2), - _row("weak", 1), _row("weak", 2)], + "rows": [ + _row("strong", 1), + _row("strong", 2), + _row("weak", 1), + _row("weak", 2), + ], } @@ -97,6 +105,7 @@ def _write(tmp_path: Path, name: str, payload: dict) -> 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)) @@ -110,7 +119,9 @@ def test_parse_device_scaling_sorts_points_and_normalizes_cpu() -> None: assert device.logical_cpu_count == 50 -def test_parse_device_scaling_missing_runtime_field_raises_schema_error() -> None: +def test_parse_device_scaling_missing_runtime_field_raises_schema_error() -> ( + None +): payload = _summary() del payload["rows"][0]["iteration_runtime_seconds"] @@ -145,10 +156,12 @@ def test_parse_device_scaling_skips_null_runtime_points() -> None: # ── 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") + _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" @@ -159,7 +172,8 @@ def test_strong_efficiency_computed_per_point_with_unit_baseline() -> None: def test_weak_efficiency_uses_local_baseline() -> None: weak = schema.parse_scaling_efficiency( - _scaling("weak", base_seconds=8.0, second_seconds=9.0), "weak") + _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) @@ -180,17 +194,25 @@ def test_scaling_efficiency_rejects_unknown_mode() -> None: # ── 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") + "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)] + ("strong", 1), + ("strong", 2), + ("weak", 1), + ("weak", 2), + ] def test_breakdown_components_reconcile_to_iteration_total() -> None: @@ -226,6 +248,7 @@ def test_breakdown_empty_rows_raises_missing_data() -> None: # ── JSON loading boundary ─────────────────────────────────────────────────── + def test_load_json_missing_file_raises_clear_error(tmp_path: Path) -> None: missing = tmp_path / "nope.json" diff --git a/tests/python/unit/test_coordinator_metrics.py b/tests/python/unit/test_coordinator_metrics.py index 6cae6326..62c69c12 100644 --- a/tests/python/unit/test_coordinator_metrics.py +++ b/tests/python/unit/test_coordinator_metrics.py @@ -31,7 +31,9 @@ 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) +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 @@ -131,7 +133,9 @@ def test_multi_coordinator_output_paths_do_not_collide(tmp_path: Path) -> None: output_directory = tmp_path / "metrics" # When - rank_zero = coordinator_metrics.metrics_output_path(output_directory, rank=0) + rank_zero = coordinator_metrics.metrics_output_path( + output_directory, rank=0 + ) rank_one = coordinator_metrics.metrics_output_path(output_directory, rank=1) # Then @@ -334,7 +338,9 @@ def test_collector_start_stop_writes_jsonl_and_stop_is_idempotent( collector.start() collector.stop() collector.stop() - samples = [json.loads(line) for line in output_path.read_text().splitlines()] + samples = [ + json.loads(line) for line in output_path.read_text().splitlines() + ] # Then assert [sample["record_type"] for sample in samples] == ["sample", "sample"] @@ -367,7 +373,9 @@ def test_background_sampling_failure_writes_terminal_error_and_stop_returns( collector.start() assert read_attempted.wait(timeout=1.0) collector.stop() - records = [json.loads(line) for line in output_path.read_text().splitlines()] + records = [ + json.loads(line) for line in output_path.read_text().splitlines() + ] # Then assert set(records[-1]) == { diff --git a/tests/python/unit/test_coordinator_metrics_cli.py b/tests/python/unit/test_coordinator_metrics_cli.py index fbb678c4..c087f037 100644 --- a/tests/python/unit/test_coordinator_metrics_cli.py +++ b/tests/python/unit/test_coordinator_metrics_cli.py @@ -70,23 +70,31 @@ def test_output_path_uses_default_file_for_directory(tmp_path: Path) -> None: assert output_path == destination / "coordinator.jsonl" -def test_ranked_output_path_suffixes_explicit_jsonl_file(tmp_path: Path) -> None: +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) + 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: +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) + output_path = coordinator_metrics_cli.metrics_output_path( + destination, rank=2 + ) # Then assert output_path == destination / "coordinator-rank-2.jsonl" @@ -107,7 +115,9 @@ def test_config_from_args_keeps_metrics_disabled() -> None: assert config is None -def test_config_from_args_resolves_ranked_output_and_nics(tmp_path: Path) -> 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", diff --git a/tests/python/unit/test_coordinator_scaling_analysis.py b/tests/python/unit/test_coordinator_scaling_analysis.py index 9ba1117d..67e54a8f 100644 --- a/tests/python/unit/test_coordinator_scaling_analysis.py +++ b/tests/python/unit/test_coordinator_scaling_analysis.py @@ -125,7 +125,9 @@ def test_parse_metrics_skips_blank_lines(tmp_path: Path) -> None: assert metrics.iteration_runtime_seconds == 3.0 -def test_parse_metrics_warns_for_trailing_partial_record(tmp_path: Path) -> None: +def test_parse_metrics_warns_for_trailing_partial_record( + tmp_path: Path, +) -> None: # Given metrics_path = tmp_path / "metrics.jsonl" _write_jsonl( @@ -222,7 +224,9 @@ def test_aggregate_repetitions_reports_means_and_success_count() -> None: assert aggregate.phase_durations_seconds == {"iteration_total": 10.0} -def test_runtime_plateau_without_sustained_cpu_does_not_name_cpu_bottleneck() -> None: +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), diff --git a/tests/python/unit/test_multi_coordinator_full_model.py b/tests/python/unit/test_multi_coordinator_full_model.py index 93f23096..14400323 100644 --- a/tests/python/unit/test_multi_coordinator_full_model.py +++ b/tests/python/unit/test_multi_coordinator_full_model.py @@ -15,7 +15,9 @@ def test_opt_strong_scaling_keeps_full_model_global_work_fixed() -> None: ) # Then - assert [(config.coordinators, config.local_batch) for config in configs] == [ + assert [ + (config.coordinators, config.local_batch) for config in configs + ] == [ (1, 2), (2, 1), ] @@ -35,7 +37,9 @@ def test_opt_weak_scaling_keeps_full_model_local_work_fixed() -> None: ) # Then - assert [(config.coordinators, config.global_batch) for config in configs] == [ + assert [ + (config.coordinators, config.global_batch) for config in configs + ] == [ (1, 1), (2, 2), ] @@ -64,7 +68,9 @@ def test_full_model_rank_command_carries_model_sequence_and_block_size( assert "--tiny" not in launch.command -def test_tiny_rank_command_preserves_quick_workload_mode(tmp_path: Path) -> None: +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 diff --git a/tests/python/unit/test_multi_coordinator_lifecycle.py b/tests/python/unit/test_multi_coordinator_lifecycle.py index 6fc65aa4..8cbf7e29 100644 --- a/tests/python/unit/test_multi_coordinator_lifecycle.py +++ b/tests/python/unit/test_multi_coordinator_lifecycle.py @@ -56,7 +56,9 @@ def test_rank_error_terminates_all_other_rank_groups() -> None: assert all(rank.terminated and rank.killed for rank in ranks) -def test_distributed_timeout_flows_to_command_and_config(tmp_path: Path) -> None: +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], @@ -67,7 +69,12 @@ def test_distributed_timeout_flows_to_command_and_config(tmp_path: Path) -> None launch = run_multi_coordinator.build_rank_launch(config, 0, tmp_path) # Then - assert launch.command[launch.command.index("--distributed-timeout-seconds") + 1] == "37" + assert ( + launch.command[ + launch.command.index("--distributed-timeout-seconds") + 1 + ] + == "37" + ) assert dataclasses.asdict(config)["distributed_timeout_seconds"] == 37 @@ -81,8 +88,13 @@ 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 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: @@ -111,7 +123,9 @@ def test_matching_rising_trajectory_still_passes_equivalence() -> None: # When result = results.build_rank_result( config, - results.RankMeasurement(0, {"iteration_total": 1.0}, (1.0, 1.01), (1.0, 1.01)), + results.RankMeasurement( + 0, {"iteration_total": 1.0}, (1.0, 1.01), (1.0, 1.01) + ), ) # Then @@ -120,11 +134,15 @@ def test_matching_rising_trajectory_still_passes_equivalence() -> None: @pytest.mark.parametrize("measured", [(float("nan"), 1.0), (1.0, 2.0)]) -def test_nan_or_mismatch_fails_equivalence(measured: tuple[float, float]) -> None: +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)), + results.RankMeasurement( + 0, {"iteration_total": 1.0}, measured, (1.0, 1.0) + ), ) assert result.loss_correctness.passed is False @@ -132,4 +150,6 @@ def test_nan_or_mismatch_fails_equivalence(measured: tuple[float, float]) -> Non 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)) + 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 index 2aa7d937..da6e5bd3 100644 --- a/tests/python/unit/test_multi_coordinator_review_handbacks.py +++ b/tests/python/unit/test_multi_coordinator_review_handbacks.py @@ -32,7 +32,9 @@ def wait(self, timeout: float | None = None) -> int: return 0 -def _rank_result(rank: int, *, passed: bool = True) -> run_multi_coordinator.RankResult: +def _rank_result( + rank: int, *, passed: bool = True +) -> run_multi_coordinator.RankResult: return run_multi_coordinator.RankResult( rank=rank, warmup_iteration_seconds=1.0, diff --git a/tests/python/unit/test_multi_coordinator_scaling.py b/tests/python/unit/test_multi_coordinator_scaling.py index 44c6721b..c6ca12e3 100644 --- a/tests/python/unit/test_multi_coordinator_scaling.py +++ b/tests/python/unit/test_multi_coordinator_scaling.py @@ -16,7 +16,10 @@ def test_strong_scaling_keeps_global_work_fixed() -> None: 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] == [ + assert [ + (config.devices_per_coordinator, config.local_batch) + for config in configs + ] == [ (8, 16), (4, 8), ] @@ -32,7 +35,9 @@ def test_weak_scaling_keeps_local_work_fixed() -> None: 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] == [ + assert [ + (config.total_devices, config.global_batch) for config in configs + ] == [ (4, 8), (8, 16), ] @@ -74,7 +79,9 @@ def test_rank_command_pins_each_rank_to_its_numa_node( ) -def test_affinity_fallback_labels_cpu_first_touch_substrate(tmp_path: Path) -> None: +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], @@ -95,7 +102,9 @@ def test_affinity_fallback_labels_cpu_first_touch_substrate(tmp_path: Path) -> N ) -def test_affinity_selection_falls_back_when_memory_binding_is_unavailable() -> None: +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 @@ -112,10 +121,16 @@ def test_affinity_selection_falls_back_when_memory_binding_is_unavailable() -> N 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 + 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 + mode="weak", + baseline_seconds=12.0, + measured_seconds=15.0, + coordinators=2, ) # Then @@ -123,7 +138,9 @@ def test_scaling_efficiency_uses_strong_and_weak_definitions() -> None: assert weak == pytest.approx(0.8) -def test_breakdown_reconciles_device_communication_optimizer_and_residual() -> None: +def test_breakdown_reconciles_device_communication_optimizer_and_residual() -> ( + None +): # Given / When breakdown = coordinator_scaling_analysis.reconcile_breakdown( iteration_total_seconds=10.0, @@ -143,7 +160,9 @@ def test_breakdown_reconciles_device_communication_optimizer_and_residual() -> N def test_breakdown_rejects_components_exceeding_iteration_total() -> None: # Given / When / Then - with pytest.raises(coordinator_scaling_analysis.BreakdownReconciliationError): + with pytest.raises( + coordinator_scaling_analysis.BreakdownReconciliationError + ): coordinator_scaling_analysis.reconcile_breakdown( iteration_total_seconds=1.0, forward_device_seconds=0.5, @@ -153,7 +172,9 @@ def test_breakdown_rejects_components_exceeding_iteration_total() -> None: ) -def test_rank_result_records_warmup_separately_and_measured_throughput() -> None: +def test_rank_result_records_warmup_separately_and_measured_throughput() -> ( + None +): # Given phases = { "warmup_iteration": 4.0, @@ -215,7 +236,9 @@ def test_global_result_uses_slowest_rank_as_iteration_makespan() -> None: assert result.substrate.startswith("same-host Gloo loopback lower bound") -def test_iteration_plan_has_one_warmup_before_three_measured_iterations() -> None: +def test_iteration_plan_has_one_warmup_before_three_measured_iterations() -> ( + None +): # Given config = run_multi_coordinator.build_scaling_configs("strong")[0] @@ -284,12 +307,16 @@ def test_experiment_plan_covers_strong_and_weak_without_writing_results( def test_scaling_conclusion_reports_speedup_and_efficiency() -> None: # Given - config_one, config_two = run_multi_coordinator.build_scaling_configs("strong") + 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), + loss_correctness=run_multi_coordinator.LossCorrectness( + True, True, True, 0.0 + ), substrate="same-host Gloo loopback lower bound", ranks=(), ) diff --git a/tests/python/unit/test_run_coordinator_scaling.py b/tests/python/unit/test_run_coordinator_scaling.py index 87e2ca65..45e6df74 100644 --- a/tests/python/unit/test_run_coordinator_scaling.py +++ b/tests/python/unit/test_run_coordinator_scaling.py @@ -37,10 +37,15 @@ def test_build_inner_command_uses_proven_d1_settings(tmp_path: Path) -> None: ) # When - command = run_coordinator_scaling.build_inner_command(run, metrics_interval=0.2) + 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[: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" @@ -50,7 +55,9 @@ def test_build_inner_command_uses_proven_d1_settings(tmp_path: Path) -> None: 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") + assert command[command.index("--metrics_output") + 1].endswith( + "metrics.jsonl" + ) def test_ports_are_unique_across_device_counts_and_repetitions() -> None: @@ -58,7 +65,10 @@ def test_ports_are_unique_across_device_counts_and_repetitions() -> None: run_count = 48 # When - ports = [run_coordinator_scaling.port_for_run(index) for index in range(run_count)] + ports = [ + run_coordinator_scaling.port_for_run(index) + for index in range(run_count) + ] # Then assert len(set(ports)) == run_count @@ -109,11 +119,15 @@ def fake_exists(path: Path) -> bool: def fake_run(command: list[str], **_kwargs): calls.append(command) if len(calls) == 1: - raise subprocess.TimeoutExpired(command, 1, output="partial", stderr="late") + 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) + monkeypatch.setattr( + coordinator_scaling_execution.subprocess, "run", fake_run + ) # When coordinator_scaling_execution.run_command( @@ -150,7 +164,9 @@ def fake_run(command: list[str], **_kwargs): return subprocess.CompletedProcess(command, 0, "ok", "") monkeypatch.setattr(Path, "exists", fake_exists) - monkeypatch.setattr(coordinator_scaling_execution.subprocess, "run", fake_run) + monkeypatch.setattr( + coordinator_scaling_execution.subprocess, "run", fake_run + ) # When coordinator_scaling_execution.run_command( From d3f5d022dd9a2402b7869732fc1b80be0e54b088 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 15:26:29 +0100 Subject: [PATCH 4/7] feat(scripts): time-varying within-batch trace generator Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- scripts/generate_trace.py | 191 +++++++++++++++++++++++++++++++++++++- 1 file changed, 190 insertions(+), 1 deletion(-) 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() From e4d0a948b45e202f822d362f46203697dc1c9a72 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 15:26:29 +0100 Subject: [PATCH 5/7] feat(eval): within-batch dynamic-performance replay and runtime overlay Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- scripts/compute_batch_runtime.py | 57 ++++- scripts/run_within_batch_experiment.py | 336 +++++++++++++++++++++++++ 2 files changed, 391 insertions(+), 2 deletions(-) create mode 100644 scripts/run_within_batch_experiment.py 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/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()) From 080fb14428d4830460ffada65577523514c368c7 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 15:26:30 +0100 Subject: [PATCH 6/7] feat(eval): export selected baseline layouts Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- scripts/export_selected_layouts.py | 141 +++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 scripts/export_selected_layouts.py 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()) From 37fc6190202d20ee4963b25d752779a997acd8c5 Mon Sep 17 00:00:00 2001 From: xly Date: Sun, 9 Aug 2026 15:26:30 +0100 Subject: [PATCH 7/7] feat(eval): verifier for absolute per-batch times Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus --- scripts/verify_absolute_times.py | 73 ++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 scripts/verify_absolute_times.py 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())