From dbdacc46c70e2becf655cbae64a73d265ae1227c Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Fri, 11 Jul 2025 18:06:34 +0800
Subject: [PATCH 1/9] allow extra fields in runtime data file; update some type
 hints

---
 testplan/base.py               | 6 +++---
 testplan/common/entity/base.py | 2 +-
 testplan/parser.py             | 1 +
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/testplan/base.py b/testplan/base.py
index bf8b0ce4b..6009ba552 100644
--- a/testplan/base.py
+++ b/testplan/base.py
@@ -174,7 +174,7 @@ def __init__(
         parse_cmdline: bool = True,
         parser: Type[TestplanParser] = TestplanParser,
         interactive_port: Optional[int] = None,
-        abort_signals: Optional[List[int]] = None,
+        abort_signals: Optional[List[signal.Signals]] = None,
         logger_level: int = logger.USER_INFO,
         file_log_level: int = logger.DEBUG,
         runpath: Union[str, Callable] = path.default_runpath,
@@ -197,8 +197,8 @@ def __init__(
         browse: bool = False,
         ui_port: Optional[int] = None,
         web_server_startup_timeout: int = defaults.WEB_SERVER_TIMEOUT,
-        test_filter: Type[BaseFilter] = filtering.Filter(),
-        test_sorter: Type[BaseSorter] = ordering.NoopSorter(),
+        test_filter: BaseFilter = filtering.Filter(),
+        test_sorter: BaseSorter = ordering.NoopSorter(),
         test_lister: Optional[MetadataBasedLister] = None,
         test_lister_output: Optional[os.PathLike] = None,
         verbose: bool = False,
diff --git a/testplan/common/entity/base.py b/testplan/common/entity/base.py
index 50e5cf8ab..45c1e0673 100644
--- a/testplan/common/entity/base.py
+++ b/testplan/common/entity/base.py
@@ -906,7 +906,7 @@ class RunnableResult:
 
     def __init__(self):
         self.step_results = OrderedDict()
-        self.run = False
+        self.run: Union[bool, Exception] = False
 
     def __repr__(self):
         return f"{self.__class__.__name__}[{vars(self)}]"
diff --git a/testplan/parser.py b/testplan/parser.py
index 6cab50770..d16a3f1ac 100644
--- a/testplan/parser.py
+++ b/testplan/parser.py
@@ -600,6 +600,7 @@ def _read_text_file(file: str) -> List[str]:
             "execution_time": schema.Or(int, float),
             "setup_time": schema.Or(int, float),
             schema.Optional("teardown_time"): schema.Or(int, float),
+            schema.Optional(str): object,  # allow other keys
         }
     }
 )

From c81fc046619675d8b04d4bc23f24e10fa2f2916c Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Tue, 22 Jul 2025 16:29:10 +0800
Subject: [PATCH 2/9] adjust estimated test execution time based on case count
 change

---
 testplan/parser.py        |  5 ++--
 testplan/runnable/base.py | 53 +++++++++++++++++++++++++++++----------
 2 files changed, 43 insertions(+), 15 deletions(-)

diff --git a/testplan/parser.py b/testplan/parser.py
index d16a3f1ac..5ec617380 100644
--- a/testplan/parser.py
+++ b/testplan/parser.py
@@ -202,6 +202,7 @@ def generate_parser(self) -> HelpParser:
         "execution_time": 199.99,
         "setup_time": 39.99,
         "teardown_time": 39.99,
+        "testcase_count": 10
     },
     ......
 }""",
@@ -599,8 +600,8 @@ def _read_text_file(file: str) -> List[str]:
         str: {
             "execution_time": schema.Or(int, float),
             "setup_time": schema.Or(int, float),
-            schema.Optional("teardown_time"): schema.Or(int, float),
-            schema.Optional(str): object,  # allow other keys
+            "teardown_time": schema.Or(int, float),
+            schema.Optional("testcase_count"): int,
         }
     }
 )
diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
index 709e680c3..809fd677e 100644
--- a/testplan/runnable/base.py
+++ b/testplan/runnable/base.py
@@ -991,10 +991,29 @@ def _calculate_parts_and_weights(
         )  # @task_target(multitest_parts=...)
         uid = task_info.uid
         runtime_data: dict = self.cfg.runtime_data or {}
-        time_info = runtime_data.get(uid, None)
+        time_info: Optional[dict] = runtime_data.get(uid, None)
 
         partitioned: List[TaskInformation] = []
 
+        adjusted_exec_time = prev_case_count = curr_case_count = 0
+        if time_info:
+            adjusted_exec_time = time_info["execution_time"]
+            if prev_case_count := time_info.get("testcase_count", 0):
+                # XXX: cache dry_run result to task_info?
+                if (
+                    curr_case_count
+                    := task_info.materialized_test.dry_run().report.counter[
+                        "total"
+                    ]
+                ):
+                    # testcase count factor caps at 2
+                    # XXX: declare const?
+                    # XXX: only apply on mt?
+                    adjusted_exec_time *= min(
+                        curr_case_count / prev_case_count, 2
+                    )
+                # XXX: shoutout if curr_case_count is 0?
+
         if num_of_parts:
             if not isinstance(task_info.materialized_test, MultiTest):
                 raise TypeError(
@@ -1010,25 +1029,33 @@ def _calculate_parts_and_weights(
                     )
                     num_of_parts = 1
                 else:
-                    # the setup time shall take no more than 50% of runtime
-                    cap = math.ceil(
-                        time_info["execution_time"]
-                        / auto_part_runtime_limit
-                        * 2
-                    )
+                    if prev_case_count and curr_case_count:
+                        adjust_formula_part = f"""
+                * min(
+                    curr_case_count {curr_case_count}
+                    / prev_case_count {prev_case_count},
+                    2
+                )"""
+                    else:
+                        adjust_formula_part = ""
                     formula = f"""
             num_of_parts = math.ceil(
-                time_info["execution_time"] {time_info["execution_time"]}
+                time_info["execution_time"] {time_info["execution_time"]}{adjust_formula_part}
                 / (
                     self.cfg.auto_part_runtime_limit {auto_part_runtime_limit}
                     - time_info["setup_time"] {time_info["setup_time"]}
                     - time_info["teardown_time"] {time_info["teardown_time"]}
                 )
             )
-        """
+"""
+
+                    # the setup time shall take no more than 50% of runtime
+                    cap = math.ceil(
+                        adjusted_exec_time / auto_part_runtime_limit * 2
+                    )
                     try:
                         num_of_parts = math.ceil(
-                            time_info["execution_time"]
+                            adjusted_exec_time
                             / (
                                 auto_part_runtime_limit
                                 - time_info["setup_time"]
@@ -1058,12 +1085,12 @@ def _calculate_parts_and_weights(
             if "weight" not in task_arguments:
                 task_arguments["weight"] = (
                     math.ceil(
-                        (time_info["execution_time"] / num_of_parts)
+                        (adjusted_exec_time / num_of_parts)
                         + time_info["setup_time"]
                         + time_info["teardown_time"]
                     )
                     if time_info
-                    else auto_part_runtime_limit
+                    else int(auto_part_runtime_limit)
                 )
             self.logger.user_info(
                 "%s: parts=%d, weight=%d",
@@ -1085,7 +1112,7 @@ def _calculate_parts_and_weights(
         else:
             if time_info and not task_info.target.weight:
                 task_info.target.weight = math.ceil(
-                    time_info["execution_time"]
+                    adjusted_exec_time
                     + time_info["setup_time"]
                     + time_info["teardown_time"]
                 )

From 96490b01b473e67f22e16f0044df0120f1691c2e Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Thu, 24 Jul 2025 18:30:05 +0800
Subject: [PATCH 3/9] place adjusted time values in execution_time fields of
 cfg.runtime_data

---
 testplan/common/utils/logger.py    |  2 +-
 testplan/runnable/base.py          | 68 ++++++++++++++++++------------
 testplan/testing/multitest/base.py |  2 +-
 3 files changed, 43 insertions(+), 29 deletions(-)

diff --git a/testplan/common/utils/logger.py b/testplan/common/utils/logger.py
index cb678119b..2895cc7a5 100644
--- a/testplan/common/utils/logger.py
+++ b/testplan/common/utils/logger.py
@@ -106,7 +106,7 @@ def _custom_log(self, level, msg, *args, **kwargs):
             self._log(level, msg, args, **kwargs)
 
 
-def _initial_setup():
+def _initial_setup() -> tuple[TestplanLogger, logging.StreamHandler]:
     """
     Perform initial setup for the logger. Creates and adds a handler to log
     to stdout with default level USER_INFO.
diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
index 809fd677e..302a7d5a0 100644
--- a/testplan/runnable/base.py
+++ b/testplan/runnable/base.py
@@ -913,22 +913,29 @@ def auto_part(self, tasks: List[Task]) -> List[Task]:
         discovered: List[TaskInformation] = [
             _detach_task_info(task) for task in tasks
         ]
-        auto_part_runtime_limit = self._calculate_part_runtime(discovered)
-
+        runtime_data = self.cfg.runtime_data or {}
+        auto_part_runtime_limit = self._calculate_part_runtime(
+            discovered, runtime_data
+        )
         for task_info in discovered:
             partitioned.extend(
                 self._calculate_parts_and_weights(
-                    task_info, auto_part_runtime_limit
+                    task_info, auto_part_runtime_limit, runtime_data
                 )
             )
 
+        # here we replace the original overall runtime data with "partitioned" values
+        # XXX: testcase_count are still sum-up value from previous run, what to do?
+        self.cfg.set_local("runtime_data", runtime_data)
+
         return [_attach_task_info(task_info) for task_info in partitioned]
 
-    def _calculate_part_runtime(self, discovered: List[TaskInformation]):
+    def _calculate_part_runtime(
+        self, discovered: List[TaskInformation], runtime_data: dict
+    ) -> float:
         if self.cfg.auto_part_runtime_limit != "auto":
             return self.cfg.auto_part_runtime_limit
 
-        runtime_data = self.cfg.runtime_data or {}
         if not runtime_data:
             self.logger.warning(
                 "Cannot derive auto_part_runtime_limit without runtime data, "
@@ -984,13 +991,15 @@ def _calculate_part_runtime(self, discovered: List[TaskInformation]):
         return auto_part_runtime_limit
 
     def _calculate_parts_and_weights(
-        self, task_info: TaskInformation, auto_part_runtime_limit: float
+        self,
+        task_info: TaskInformation,
+        auto_part_runtime_limit: float,
+        runtime_data: dict,
     ):
         num_of_parts = (
             task_info.num_of_parts
         )  # @task_target(multitest_parts=...)
         uid = task_info.uid
-        runtime_data: dict = self.cfg.runtime_data or {}
         time_info: Optional[dict] = runtime_data.get(uid, None)
 
         partitioned: List[TaskInformation] = []
@@ -998,21 +1007,29 @@ def _calculate_parts_and_weights(
         adjusted_exec_time = prev_case_count = curr_case_count = 0
         if time_info:
             adjusted_exec_time = time_info["execution_time"]
-            if prev_case_count := time_info.get("testcase_count", 0):
-                # XXX: cache dry_run result to task_info?
-                if (
-                    curr_case_count
-                    := task_info.materialized_test.dry_run().report.counter[
-                        "total"
-                    ]
-                ):
-                    # testcase count factor caps at 2
-                    # XXX: declare const?
-                    # XXX: only apply on mt?
-                    adjusted_exec_time *= min(
-                        curr_case_count / prev_case_count, 2
-                    )
-                # XXX: shoutout if curr_case_count is 0?
+            if isinstance(task_info.materialized_test, MultiTest):
+                if prev_case_count := time_info.get("testcase_count", 0):
+                    # XXX: cache dry_run result somewhere?
+                    # NOTE: get_metadata won't work here since filters not applied
+                    if (
+                        curr_case_count
+                        := task_info.materialized_test.dry_run().report.counter[
+                            "total"
+                        ]
+                    ):
+                        # XXX: define lb & ub of testcase-count factor?
+                        adjusted_exec_time *= curr_case_count / prev_case_count
+                        self.logger.user_info(
+                            "%s: estimated total execution time %f -> %f "
+                            "(prev total tc: %d, curr total tc: %d)",
+                            uid,
+                            time_info["execution_time"],
+                            adjusted_exec_time,
+                            prev_case_count,
+                            curr_case_count,
+                        )
+                        time_info["execution_time"] = adjusted_exec_time
+                    # XXX: shoutout if curr_case_count is 0?
 
         if num_of_parts:
             if not isinstance(task_info.materialized_test, MultiTest):
@@ -1031,11 +1048,8 @@ def _calculate_parts_and_weights(
                 else:
                     if prev_case_count and curr_case_count:
                         adjust_formula_part = f"""
-                * min(
-                    curr_case_count {curr_case_count}
-                    / prev_case_count {prev_case_count},
-                    2
-                )"""
+                * curr_case_count {curr_case_count}
+                / prev_case_count {prev_case_count}"""
                     else:
                         adjust_formula_part = ""
                     formula = f"""
diff --git a/testplan/testing/multitest/base.py b/testplan/testing/multitest/base.py
index 9e7398f32..45c389038 100644
--- a/testplan/testing/multitest/base.py
+++ b/testplan/testing/multitest/base.py
@@ -1,7 +1,7 @@
 """MultiTest test execution framework."""
 
 import collections.abc
-import concurrent
+import concurrent.futures
 import functools
 import itertools
 import warnings

From 81346852cd84441174851a28840cd8fe62805381 Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Fri, 25 Jul 2025 14:39:46 +0800
Subject: [PATCH 4/9] add tests & newsfrag

---
 ...hanged.testcase_count_adjusted_runtime.rst |  2 +
 .../testplan/runners/pools/test_auto_part.py  | 63 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst

diff --git a/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst b/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst
new file mode 100644
index 000000000..d1f5ce049
--- /dev/null
+++ b/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst
@@ -0,0 +1,2 @@
+Runtime data JSON file specified through ``--runtime-data`` now accepts an optional ``testcase_count`` field, which should hold the total number of testcases from the previous run. If ``testcase_count`` is specified, estimated MultiTest execution time will be adjusted by the ratio of current run's testcase count to that of previous run, which should bring more accurate auto-partitioning and improved testing resource utilization.
+
diff --git a/tests/functional/testplan/runners/pools/test_auto_part.py b/tests/functional/testplan/runners/pools/test_auto_part.py
index b5df23429..87e1f61be 100755
--- a/tests/functional/testplan/runners/pools/test_auto_part.py
+++ b/tests/functional/testplan/runners/pools/test_auto_part.py
@@ -1,5 +1,7 @@
 import os
+import math
 import tempfile
+from itertools import count
 from pathlib import Path
 
 import pytest
@@ -379,3 +381,64 @@ def test_auto_plan_runtime_target():
             assert task.weight == 900
         mockplan.run()
         assert pool.size == 6
+
+
+@pytest.mark.parametrize(
+    "tc, adjusted_exec_t, exp_weight, exp_parts",
+    [
+        (10, 60, 35, 2),
+        (5, 120, 35, 4),
+        (20, 30, 35, 1),
+        (8, 75, 30, 3),
+    ],
+    ids=count(0),
+)
+def test_multitest_weight_adjusted_by_relative_testcase_count(
+    tc, adjusted_exec_t, exp_weight, exp_parts
+):
+    with tempfile.TemporaryDirectory() as runpath:
+        mockplan = TestplanMock(
+            "plan",
+            runpath=runpath,
+            merge_scheduled_parts=True,
+            auto_part_runtime_limit=40,
+            plan_runtime_target=80,
+            runtime_data={
+                "Proj1-suite": {
+                    "execution_time": 60,
+                    "setup_time": 5,
+                    "teardown_time": 0,
+                    "testcase_count": tc,
+                }
+            },
+        )
+        pool = ProcessPool(name="MyPool", size="auto")
+        mockplan.add_resource(pool)
+        current_folder = Path(__file__).resolve().parent
+        assert pool.cfg.runtime_data == {
+            "Proj1-suite": {
+                "execution_time": 60,
+                "setup_time": 5,
+                "teardown_time": 0,
+                "testcase_count": tc,
+            }
+        }
+        # curr tc: 10
+        mockplan.schedule_all(
+            path=current_folder / "discover_tasks",
+            name_pattern=r".*auto_parts_tasks\.py$",
+            resource="MyPool",
+        )
+        assert pool.cfg.runtime_data == {
+            "Proj1-suite": {
+                "execution_time": adjusted_exec_t,
+                "setup_time": 5,
+                "teardown_time": 0,
+                "testcase_count": tc,
+            }
+        }
+        assert len(pool.added_items) == exp_parts
+        for task in pool.added_items.values():
+            assert task.weight == exp_weight
+        mockplan.run()
+        assert pool.size == math.ceil(exp_weight * exp_parts / 80)

From 74c85055a291e2973e16449c9bbe4d460ab80089 Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Fri, 25 Jul 2025 16:02:02 +0800
Subject: [PATCH 5/9] update doc; hitchhike ns package debug msg

---
 doc/en/introduction.rst                               |  3 ++-
 doc/en/pools.rst                                      |  3 ++-
 .../3294_changed.testcase_count_adjusted_runtime.rst  |  3 +--
 examples/ExecutionPools/AutoPart/runtime_data.txt     |  3 ++-
 testplan/runnable/interactive/reloader.py             | 11 ++++++++---
 5 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/doc/en/introduction.rst b/doc/en/introduction.rst
index 5cc7fe8c1..42f55b0b3 100644
--- a/doc/en/introduction.rst
+++ b/doc/en/introduction.rst
@@ -346,7 +346,8 @@ Command line
                                 "<Multitest>": {
                                     | "execution_time": 199.99,
                                     | "setup_time": 39.99,
-                                    | "teardown_time": 39.99, // optional
+                                    | "teardown_time": 39.99,
+                                    | "testcase_count": 10  // optional
                                 },
                             }
 
diff --git a/doc/en/pools.rst b/doc/en/pools.rst
index 643357c53..d5455930d 100644
--- a/doc/en/pools.rst
+++ b/doc/en/pools.rst
@@ -390,7 +390,8 @@ enabled by providing runtime data like following via ``--runtime-data`` command
         "<Multitest>": {
             "execution_time": 199.99,
             "setup_time": 39.99,
-            "teardown_time": 0
+            "teardown_time": 0,
+            "testcase_count": 10
         },
         ......
     }
diff --git a/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst b/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst
index d1f5ce049..9a4200e18 100644
--- a/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst
+++ b/doc/newsfragments/3294_changed.testcase_count_adjusted_runtime.rst
@@ -1,2 +1 @@
-Runtime data JSON file specified through ``--runtime-data`` now accepts an optional ``testcase_count`` field, which should hold the total number of testcases from the previous run. If ``testcase_count`` is specified, estimated MultiTest execution time will be adjusted by the ratio of current run's testcase count to that of previous run, which should bring more accurate auto-partitioning and improved testing resource utilization.
-
+Runtime data JSON file specified through ``--runtime-data`` now accepts an optional ``testcase_count`` field, which should hold the total number of testcases from the previous run. If ``testcase_count`` is specified, estimated MultiTest execution time will be adjusted by the ratio of current run's testcase count to that of previous run, which should bring more accurate auto-partitioning and improved testing resource utilization.
\ No newline at end of file
diff --git a/examples/ExecutionPools/AutoPart/runtime_data.txt b/examples/ExecutionPools/AutoPart/runtime_data.txt
index 4d3703844..8b795bbf0 100755
--- a/examples/ExecutionPools/AutoPart/runtime_data.txt
+++ b/examples/ExecutionPools/AutoPart/runtime_data.txt
@@ -2,6 +2,7 @@
   "Proj1-suite": {
     "execution_time": 199.99,
     "setup_time": 5,
-    "teardown_time": 0
+    "teardown_time": 0,
+    "testcase_count": 8
   }
 }
diff --git a/testplan/runnable/interactive/reloader.py b/testplan/runnable/interactive/reloader.py
index 769d2c8c6..fcb12af9b 100644
--- a/testplan/runnable/interactive/reloader.py
+++ b/testplan/runnable/interactive/reloader.py
@@ -14,15 +14,19 @@
 from testplan.common.utils import path as path_utils
 from testplan.common.utils import logger
 from testplan.common.utils import strings
+from testplan.common.utils.logger import TESTPLAN_LOGGER
 from testplan.common.utils.package import import_tmp_module
 from testplan.testing.multitest import suite, MultiTest
 
 
-def _patched_find_module(name, path=None):
+def _patched_find_module(name, path=None, logger=None):
     """
     modified from <3.13 branches, to handle namespace packages
     """
 
+    if logger is None:
+        logger = TESTPLAN_LOGGER
+
     # Old imp constants:
     _SEARCH_ERROR = 0
     _PY_SOURCE = 1
@@ -64,6 +68,7 @@ def _patched_find_module(name, path=None):
         ):
             # ModuleFinder.find_module is designed to only return one package dir,
             # while namespace packages can have multiple
+            logger.debug("Skipping namespace package %s under %s", name, path)
             return None, None, ("", "", _NAMESPACE_IGNORED)
 
     if spec.loader is importlib.machinery.BuiltinImporter:
@@ -150,7 +155,7 @@ def __init__(self, extra_deps=None, scheduled_modules=None):
         ) = self._build_dependencies()
 
         # Last recorded reload time for watched modules.
-        self._last_reload_time = {}  # type: Dict[str, float]
+        self._last_reload_time: dict[str, float] = {}
         self._init_time = time.time()
 
     def reload(self, tests, rebuild_dependencies=False):
@@ -663,7 +668,7 @@ def find_module(self, name, path, parent=None):
 
             path = self.path
 
-        return _patched_find_module(name, path)
+        return _patched_find_module(name, path, logger=self.logger)
 
 
 class _ModuleNode:

From fab123571004ed98eeac084f931f385ed25d3b3c Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Tue, 29 Jul 2025 11:51:35 +0800
Subject: [PATCH 6/9] address review comments

---
 testplan/runnable/base.py                     | 85 ++++++++++---------
 .../testplan/runners/pools/test_auto_part.py  |  1 +
 2 files changed, 48 insertions(+), 38 deletions(-)

diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
index 302a7d5a0..0ea01f1b6 100644
--- a/testplan/runnable/base.py
+++ b/testplan/runnable/base.py
@@ -914,6 +914,7 @@ def auto_part(self, tasks: List[Task]) -> List[Task]:
             _detach_task_info(task) for task in tasks
         ]
         runtime_data = self.cfg.runtime_data or {}
+        self._adjust_runtime_data(discovered, runtime_data)
         auto_part_runtime_limit = self._calculate_part_runtime(
             discovered, runtime_data
         )
@@ -930,6 +931,45 @@ def auto_part(self, tasks: List[Task]) -> List[Task]:
 
         return [_attach_task_info(task_info) for task_info in partitioned]
 
+    def _adjust_runtime_data(
+        self, discovered: List[TaskInformation], runtime_data: dict
+    ):
+        """
+        Adjust the runtime data to ensure that all discovered tasks have their
+        runtime data available. If a task's UID is not found in the runtime data,
+        it will be added with default values.
+        """
+        for task_info in discovered:
+            uid = task_info.uid
+            time_info = runtime_data.get(uid, None)
+            if time_info and isinstance(
+                task_info.materialized_test, MultiTest
+            ):
+                if prev_case_count := time_info.get("testcase_count", 0):
+                    # XXX: cache dry_run result somewhere?
+                    # NOTE: get_metadata won't work here since filters not applied
+                    if (
+                        curr_case_count
+                        := task_info.materialized_test.dry_run().report.counter[
+                            "total"
+                        ]
+                    ):
+                        # XXX: define lb & ub of testcase-count factor?
+                        adjusted_exec_time = time_info["execution_time"] * max(
+                            curr_case_count / prev_case_count, 0.25
+                        )
+                        self.logger.user_info(
+                            "%s: estimated total execution time %f -> %f "
+                            "(prev total testcase number: %d, curr total testcase number: %d)",
+                            uid,
+                            time_info["execution_time"],
+                            adjusted_exec_time,
+                            prev_case_count,
+                            curr_case_count,
+                        )
+                        time_info["execution_time"] = adjusted_exec_time
+                    # XXX: shoutout if curr_case_count is 0?
+
     def _calculate_part_runtime(
         self, discovered: List[TaskInformation], runtime_data: dict
     ) -> float:
@@ -1004,33 +1044,6 @@ def _calculate_parts_and_weights(
 
         partitioned: List[TaskInformation] = []
 
-        adjusted_exec_time = prev_case_count = curr_case_count = 0
-        if time_info:
-            adjusted_exec_time = time_info["execution_time"]
-            if isinstance(task_info.materialized_test, MultiTest):
-                if prev_case_count := time_info.get("testcase_count", 0):
-                    # XXX: cache dry_run result somewhere?
-                    # NOTE: get_metadata won't work here since filters not applied
-                    if (
-                        curr_case_count
-                        := task_info.materialized_test.dry_run().report.counter[
-                            "total"
-                        ]
-                    ):
-                        # XXX: define lb & ub of testcase-count factor?
-                        adjusted_exec_time *= curr_case_count / prev_case_count
-                        self.logger.user_info(
-                            "%s: estimated total execution time %f -> %f "
-                            "(prev total tc: %d, curr total tc: %d)",
-                            uid,
-                            time_info["execution_time"],
-                            adjusted_exec_time,
-                            prev_case_count,
-                            curr_case_count,
-                        )
-                        time_info["execution_time"] = adjusted_exec_time
-                    # XXX: shoutout if curr_case_count is 0?
-
         if num_of_parts:
             if not isinstance(task_info.materialized_test, MultiTest):
                 raise TypeError(
@@ -1046,15 +1059,9 @@ def _calculate_parts_and_weights(
                     )
                     num_of_parts = 1
                 else:
-                    if prev_case_count and curr_case_count:
-                        adjust_formula_part = f"""
-                * curr_case_count {curr_case_count}
-                / prev_case_count {prev_case_count}"""
-                    else:
-                        adjust_formula_part = ""
                     formula = f"""
             num_of_parts = math.ceil(
-                time_info["execution_time"] {time_info["execution_time"]}{adjust_formula_part}
+                time_info["execution_time"] {time_info["execution_time"]}
                 / (
                     self.cfg.auto_part_runtime_limit {auto_part_runtime_limit}
                     - time_info["setup_time"] {time_info["setup_time"]}
@@ -1065,11 +1072,13 @@ def _calculate_parts_and_weights(
 
                     # the setup time shall take no more than 50% of runtime
                     cap = math.ceil(
-                        adjusted_exec_time / auto_part_runtime_limit * 2
+                        time_info["execution_time"]
+                        / auto_part_runtime_limit
+                        * 2
                     )
                     try:
                         num_of_parts = math.ceil(
-                            adjusted_exec_time
+                            time_info["execution_time"]
                             / (
                                 auto_part_runtime_limit
                                 - time_info["setup_time"]
@@ -1099,7 +1108,7 @@ def _calculate_parts_and_weights(
             if "weight" not in task_arguments:
                 task_arguments["weight"] = (
                     math.ceil(
-                        (adjusted_exec_time / num_of_parts)
+                        (time_info["execution_time"] / num_of_parts)
                         + time_info["setup_time"]
                         + time_info["teardown_time"]
                     )
@@ -1126,7 +1135,7 @@ def _calculate_parts_and_weights(
         else:
             if time_info and not task_info.target.weight:
                 task_info.target.weight = math.ceil(
-                    adjusted_exec_time
+                    time_info["execution_time"]
                     + time_info["setup_time"]
                     + time_info["teardown_time"]
                 )
diff --git a/tests/functional/testplan/runners/pools/test_auto_part.py b/tests/functional/testplan/runners/pools/test_auto_part.py
index 87e1f61be..cc38870d6 100755
--- a/tests/functional/testplan/runners/pools/test_auto_part.py
+++ b/tests/functional/testplan/runners/pools/test_auto_part.py
@@ -390,6 +390,7 @@ def test_auto_plan_runtime_target():
         (5, 120, 35, 4),
         (20, 30, 35, 1),
         (8, 75, 30, 3),
+        (100, 15, 20, 1),
     ],
     ids=count(0),
 )

From df5e7b5e807a7317b1fd17f49c0addc1827beeb8 Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Tue, 29 Jul 2025 14:23:20 +0800
Subject: [PATCH 7/9] tweak log msg; revert some unnecessary changes

---
 testplan/runnable/base.py | 38 +++++++++++++++++---------------------
 1 file changed, 17 insertions(+), 21 deletions(-)

diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
index 0ea01f1b6..d8fa325c5 100644
--- a/testplan/runnable/base.py
+++ b/testplan/runnable/base.py
@@ -915,20 +915,18 @@ def auto_part(self, tasks: List[Task]) -> List[Task]:
         ]
         runtime_data = self.cfg.runtime_data or {}
         self._adjust_runtime_data(discovered, runtime_data)
-        auto_part_runtime_limit = self._calculate_part_runtime(
-            discovered, runtime_data
-        )
+        # here we replace the original runtime data with adjusted values
+        # XXX: testcase_count are still sum-up value from previous run, what to do?
+        self.cfg.set_local("runtime_data", runtime_data)
+
+        auto_part_runtime_limit = self._calculate_part_runtime(discovered)
         for task_info in discovered:
             partitioned.extend(
                 self._calculate_parts_and_weights(
-                    task_info, auto_part_runtime_limit, runtime_data
+                    task_info, auto_part_runtime_limit
                 )
             )
 
-        # here we replace the original overall runtime data with "partitioned" values
-        # XXX: testcase_count are still sum-up value from previous run, what to do?
-        self.cfg.set_local("runtime_data", runtime_data)
-
         return [_attach_task_info(task_info) for task_info in partitioned]
 
     def _adjust_runtime_data(
@@ -959,7 +957,7 @@ def _adjust_runtime_data(
                             curr_case_count / prev_case_count, 0.25
                         )
                         self.logger.user_info(
-                            "%s: estimated total execution time %f -> %f "
+                            "%s: adjust estimated total execution time %.2f -> %.2f "
                             "(prev total testcase number: %d, curr total testcase number: %d)",
                             uid,
                             time_info["execution_time"],
@@ -971,11 +969,12 @@ def _adjust_runtime_data(
                     # XXX: shoutout if curr_case_count is 0?
 
     def _calculate_part_runtime(
-        self, discovered: List[TaskInformation], runtime_data: dict
+        self, discovered: List[TaskInformation]
     ) -> float:
         if self.cfg.auto_part_runtime_limit != "auto":
             return self.cfg.auto_part_runtime_limit
 
+        runtime_data = self.cfg.runtime_data or {}
         if not runtime_data:
             self.logger.warning(
                 "Cannot derive auto_part_runtime_limit without runtime data, "
@@ -1031,15 +1030,13 @@ def _calculate_part_runtime(
         return auto_part_runtime_limit
 
     def _calculate_parts_and_weights(
-        self,
-        task_info: TaskInformation,
-        auto_part_runtime_limit: float,
-        runtime_data: dict,
+        self, task_info: TaskInformation, auto_part_runtime_limit: float
     ):
         num_of_parts = (
             task_info.num_of_parts
         )  # @task_target(multitest_parts=...)
         uid = task_info.uid
+        runtime_data: dict = self.cfg.runtime_data or {}
         time_info: Optional[dict] = runtime_data.get(uid, None)
 
         partitioned: List[TaskInformation] = []
@@ -1059,6 +1056,12 @@ def _calculate_parts_and_weights(
                     )
                     num_of_parts = 1
                 else:
+                    # the setup time shall take no more than 50% of runtime
+                    cap = math.ceil(
+                        time_info["execution_time"]
+                        / auto_part_runtime_limit
+                        * 2
+                    )
                     formula = f"""
             num_of_parts = math.ceil(
                 time_info["execution_time"] {time_info["execution_time"]}
@@ -1069,13 +1072,6 @@ def _calculate_parts_and_weights(
                 )
             )
 """
-
-                    # the setup time shall take no more than 50% of runtime
-                    cap = math.ceil(
-                        time_info["execution_time"]
-                        / auto_part_runtime_limit
-                        * 2
-                    )
                     try:
                         num_of_parts = math.ceil(
                             time_info["execution_time"]

From eb75ffdada83f6145683dead456f14e10f39f346 Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Tue, 29 Jul 2025 15:01:53 +0800
Subject: [PATCH 8/9] address more review comments

---
 testplan/runnable/base.py | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/testplan/runnable/base.py b/testplan/runnable/base.py
index d8fa325c5..1db8b6f30 100644
--- a/testplan/runnable/base.py
+++ b/testplan/runnable/base.py
@@ -86,6 +86,9 @@
 TestTask = Union[Test, Task, Callable]
 
 
+MULTITEST_EXEC_TIME_ADJUST_FACTOR_LB = 0.25
+
+
 @dataclass
 class TaskInformation:
     target: TestTask
@@ -916,7 +919,7 @@ def auto_part(self, tasks: List[Task]) -> List[Task]:
         runtime_data = self.cfg.runtime_data or {}
         self._adjust_runtime_data(discovered, runtime_data)
         # here we replace the original runtime data with adjusted values
-        # XXX: testcase_count are still sum-up value from previous run, what to do?
+        # and "previous" testcase count with current run count
         self.cfg.set_local("runtime_data", runtime_data)
 
         auto_part_runtime_limit = self._calculate_part_runtime(discovered)
@@ -952,9 +955,10 @@ def _adjust_runtime_data(
                             "total"
                         ]
                     ):
-                        # XXX: define lb & ub of testcase-count factor?
+                        # XXX: lb defined, ub?
                         adjusted_exec_time = time_info["execution_time"] * max(
-                            curr_case_count / prev_case_count, 0.25
+                            curr_case_count / prev_case_count,
+                            MULTITEST_EXEC_TIME_ADJUST_FACTOR_LB,
                         )
                         self.logger.user_info(
                             "%s: adjust estimated total execution time %.2f -> %.2f "
@@ -966,7 +970,7 @@ def _adjust_runtime_data(
                             curr_case_count,
                         )
                         time_info["execution_time"] = adjusted_exec_time
-                    # XXX: shoutout if curr_case_count is 0?
+                        time_info["testcase_count"] = curr_case_count
 
     def _calculate_part_runtime(
         self, discovered: List[TaskInformation]

From af4aef8603a90b9603ee275bb7948a402452e88d Mon Sep 17 00:00:00 2001
From: zhenyu-ms <111329301+zhenyu-ms@users.noreply.github.com>
Date: Tue, 29 Jul 2025 15:30:30 +0800
Subject: [PATCH 9/9] fix tests

---
 tests/functional/testplan/runners/pools/test_auto_part.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/functional/testplan/runners/pools/test_auto_part.py b/tests/functional/testplan/runners/pools/test_auto_part.py
index cc38870d6..ab1adac95 100755
--- a/tests/functional/testplan/runners/pools/test_auto_part.py
+++ b/tests/functional/testplan/runners/pools/test_auto_part.py
@@ -435,7 +435,7 @@ def test_multitest_weight_adjusted_by_relative_testcase_count(
                 "execution_time": adjusted_exec_t,
                 "setup_time": 5,
                 "teardown_time": 0,
-                "testcase_count": tc,
+                "testcase_count": 10,
             }
         }
         assert len(pool.added_items) == exp_parts