diff --git a/README.md b/README.md index dc5ee06d..937bdd28 100644 --- a/README.md +++ b/README.md @@ -257,6 +257,7 @@ tpums benchmark run-config configs/sample_configs/parameter_sweep.yaml \ - **Report Output & Compiler Flags**: - **`--output_dir `**: Directory to save `summary.csv` and `detailed.json` (defaults to `results/`). - **`--xla_flags_file_path `**: Optional path to a custom YAML file overriding default per-benchmark XLA / compiler runtime flags. +- **Standalone Single-Host Init (`--single_host_tpu_init` / `single_host_tpu_init: true`, optional)**: Initializes this host as a standalone 1-host topology inside a multi-host slice (skips SliceBuilder, disables the cross-host launch barrier) so that a single-host benchmark can be used for node-local diagnostics without the other hosts of the slice participating. The slice's ICI routing is left untouched, so device-to-device transfers hang in this mode; only `gemm`, `gemm_throttling` and `hbm` are known to avoid them (a warning is printed for other benchmarks). Off by default; has no effect on normal single- or multi-host runs. --- @@ -273,6 +274,7 @@ Configuration files define a top-level `benchmark:` mapping containing: - **`name:` (``)** — Target benchmark name (e.g., `gemm`, `hbm`, `all_reduce`; see [Benchmark Catalog](#benchmark-catalog)). - **`xprof_timing:`** *(optional)* — Boolean (`true` / `false`) to enable XProf hardware trace collection and device timing analysis. - **`xprof_device_timing_aggregation:`** *(optional)* — Device timing aggregation mode when `xprof_timing: true` (`measured_device_or_first` [default], `local_slowest_device`, or `global_slowest_per_step`). +- **`single_host_tpu_init:`** *(optional)* — Boolean (`true` / `false`, default `false`) to initialize this host standalone inside a multi-host slice for node-local diagnostics (see [Common Execution Flags](#4-common-execution-flags)). - **`params:`** *(optional)* — Baseline execution parameters shared across all generated runs. - **`cases:` / `cases_from_csv:` / `sweep:`** *(optional)* — Case override and parameter sweep generators. diff --git a/src/accelerator_microbenchmarks/benchmarks/matmul.py b/src/accelerator_microbenchmarks/benchmarks/matmul.py index f4fa05b6..2da26549 100644 --- a/src/accelerator_microbenchmarks/benchmarks/matmul.py +++ b/src/accelerator_microbenchmarks/benchmarks/matmul.py @@ -82,6 +82,11 @@ class GemmParams(base.BaseBenchmarkParams): TGemmConfig = TypeVar("TGemmConfig", bound=base.BaseBenchmarkParams) +def _sample_normal(key: jax.Array, shape: tuple[int, ...], dtype: Any): + """Module-level so `jax.jit` reuses its compile cache across calls.""" + return jax.random.normal(key, shape).astype(dtype) + + class BaseGemmBenchmark(base.BaseBenchmark[TGemmConfig]): """Base class for GEMM benchmarks. @@ -198,6 +203,27 @@ def get_run_identifier(self) -> str: run_identifier += "_sf" return run_identifier + def _replicated_normal( + self, key: jax.Array, shape: tuple[int, ...], dtype: Any + ) -> jax.Array: + """Returns standard-normal samples replicated across the mesh. + + Generated on every mesh device under `jit` with a replicated + `out_shardings` (as in `hbm.py`), so no device-to-device broadcast is + needed; required for `single_host_tpu_init`, where such transfers hang. + """ + if self.mesh is None: + raise ValueError("Mesh not initialized.") + replicated_sharding = jax.sharding.NamedSharding( + self.mesh, jax.sharding.PartitionSpec(*([None] * len(shape))) + ) + generate = jax.jit( + _sample_normal, + static_argnums=(1, 2), + out_shardings=replicated_sharding, + ) + return generate(key, shape, dtype) + def generate_inputs(self) -> tuple[Any, ...]: ( m, @@ -219,57 +245,34 @@ def generate_inputs(self) -> tuple[Any, ...]: key = jax.random.PRNGKey(self.config.seed) k1, k2, k3, k4, k5 = jax.random.split(key, 5) - # Data generation in HBM + # Data generation in HBM, replicated across the mesh to avoid + # discrepancies between devices. # Note: JAX might require intermediate conversion for random.normal if # dtypes aren't supported a_shape = (k, m) if transpose_a else (m, k) b_shape = (n, k) if transpose_b else (k, n) - a = jax.random.normal(k1, a_shape).astype(in_dtype) - b = jax.random.normal(k2, b_shape).astype(in_dtype) - - # Optional scaling factors - use_sf = self.config.use_scaling_factors - sf0 = jax.random.normal(k3, (m, 1)).astype(jnp.float32) if use_sf else None - sf1 = jax.random.normal(k4, (1, n)).astype(jnp.float32) if use_sf else None - - use_accumulator_matrix = True if self.config.beta != 0.0 else False - c = ( - jax.random.normal(k5, (m, n)).astype(in_dtype) - if use_accumulator_matrix - else None - ) - - # Replicated sharding for computation ops to avoid discrepancies - if self.mesh is None: - raise ValueError("Mesh not initialized.") - replicated_sharding = jax.sharding.NamedSharding( - self.mesh, jax.sharding.PartitionSpec(None, None) - ) + a = self._replicated_normal(k1, a_shape, in_dtype) + b = self._replicated_normal(k2, b_shape, in_dtype) # Row and column scaling factor vectors (sf0: m x 1, sf1: 1 x n) # Unpacked as sf0, sf1 args into gemm_fn(a, b, sf0, sf1, c) if enabled, # else (None, None). - if use_sf: + if self.config.use_scaling_factors: sf_args = [ - jax.device_put(sf0, replicated_sharding), - jax.device_put(sf1, replicated_sharding), + self._replicated_normal(k3, (m, 1), jnp.float32), + self._replicated_normal(k4, (1, n), jnp.float32), ] else: sf_args = [None, None] # Accumulator matrix C (M x N) for fused GEMM addition when beta != 0.0. # Unpacked as c arg into gemm_fn(a, b, sf0, sf1, c) if enabled, else []. - if use_accumulator_matrix: - acc_args = [jax.device_put(c, replicated_sharding)] + if self.config.beta != 0.0: + acc_args = [self._replicated_normal(k5, (m, n), in_dtype)] else: acc_args = [] - return ( - jax.device_put(a, replicated_sharding), - jax.device_put(b, replicated_sharding), - *sf_args, - *acc_args, - ) + return (a, b, *sf_args, *acc_args) def run_op(self, *args) -> jnp.ndarray: if self._jit_fn is None: diff --git a/src/accelerator_microbenchmarks/cli.py b/src/accelerator_microbenchmarks/cli.py index bdf177ab..e081f0d2 100644 --- a/src/accelerator_microbenchmarks/cli.py +++ b/src/accelerator_microbenchmarks/cli.py @@ -39,6 +39,17 @@ def _add_common_execution_args(parser: argparse.ArgumentParser) -> None: default=None, help="Optional path to custom op_flags YAML file.", ) + parser.add_argument( + "--single_host_tpu_init", + action="store_true", + default=False, + help=( + "Initialize this host standalone inside a multi-host slice for" + " node-local diagnostics (forces a 1-host topology and skips" + " SliceBuilder). Device-to-device transfers hang in this mode;" + " supported benchmarks: gemm, gemm_throttling, hbm." + ), + ) def _str_to_bool(value: str) -> bool: @@ -305,12 +316,18 @@ def run(argv: Sequence[str]) -> None: "xprof_device_timing_aggregation", constants.XprofDeviceTimingAggregation.MEASURED_DEVICE_OR_FIRST, ) + effective_single_host_tpu_init = args.single_host_tpu_init or ( + bool(raw_configs[0].get("single_host_tpu_init", False)) + if raw_configs + else False + ) tasks = [] for raw in raw_configs: name = raw.pop("name") raw.pop("xprof_timing", None) raw.pop("xprof_device_timing_aggregation", None) + raw.pop("single_host_tpu_init", None) bench_cls = registry.benchmark_registry.get_benchmark(name) tasks.append((name, bench_cls.Config(**raw))) @@ -324,6 +341,7 @@ def run(argv: Sequence[str]) -> None: xprof_dir=args.xprof_dir, config_path=args.config_path, xla_flags_file_path=args.xla_flags_file_path, + single_host_tpu_init=effective_single_host_tpu_init, ) return @@ -338,6 +356,7 @@ def run(argv: Sequence[str]) -> None: xprof_device_timing_aggregation=args.xprof_device_timing_aggregation, xprof_dir=args.xprof_dir, xla_flags_file_path=args.xla_flags_file_path, + single_host_tpu_init=args.single_host_tpu_init, ) return diff --git a/src/accelerator_microbenchmarks/core/config.py b/src/accelerator_microbenchmarks/core/config.py index a18caf39..ae13c42e 100644 --- a/src/accelerator_microbenchmarks/core/config.py +++ b/src/accelerator_microbenchmarks/core/config.py @@ -20,6 +20,7 @@ class BenchmarkKey(str, enum.Enum): SWEEP = "sweep" XPROF_TIMING = "xprof_timing" XPROF_DEVICE_TIMING_AGGREGATION = "xprof_device_timing_aggregation" + SINGLE_HOST_TPU_INIT = "single_host_tpu_init" _ALLOWED_BENCHMARK_KEYS = frozenset(BenchmarkKey) @@ -304,6 +305,16 @@ def load_config(path: str) -> list[dict[str, Any]]: xprof_device_timing_aggregation ) + single_host_tpu_init = benchmark_spec.get(BenchmarkKey.SINGLE_HOST_TPU_INIT) + if single_host_tpu_init is not None: + if not isinstance(single_host_tpu_init, bool): + raise ValueError( + f"Expected '{BenchmarkKey.SINGLE_HOST_TPU_INIT.value}' in benchmark" + f" '{benchmark_name}' to be a bool, but got" + f" {type(single_host_tpu_init).__name__}." + ) + base_params["single_host_tpu_init"] = single_host_tpu_init + # Stage 3: Cases (per-case parameter overrides over base_params) case_dicts = ( [{**base_params, **case} for case in cases_list] diff --git a/src/accelerator_microbenchmarks/core/runner.py b/src/accelerator_microbenchmarks/core/runner.py index 6473ca85..47bdf8e2 100644 --- a/src/accelerator_microbenchmarks/core/runner.py +++ b/src/accelerator_microbenchmarks/core/runner.py @@ -14,6 +14,7 @@ from accelerator_microbenchmarks.core import platform from accelerator_microbenchmarks.core import registry from accelerator_microbenchmarks.core import report +from accelerator_microbenchmarks.core import single_host from accelerator_microbenchmarks.core import system import jax import yaml @@ -147,6 +148,7 @@ def run_benchmarks( xprof_device_timing_aggregation: str = ( constants.XprofDeviceTimingAggregation.MEASURED_DEVICE_OR_FIRST ), + single_host_tpu_init: bool = False, ) -> List[base.BenchmarkResult]: """Core execution engine for typed benchmark task configurations.""" xprof_config = base.XprofConfig( @@ -162,6 +164,20 @@ def run_benchmarks( [{"name": task_name} for task_name, _ in tasks], xla_flags_file_path=xla_flags_file_path, ) + # Optional standalone single-host init. Must follow set_xla_flags (which + # replaces LIBTPU_INIT_ARGS) and precede any JAX device call. + if single_host_tpu_init: + unsupported = sorted( + {name for name, _ in tasks} - single_host.SUPPORTED_BENCHMARKS + ) + if unsupported: + print( + f"Warning: single_host_tpu_init is not supported for {unsupported}:" + " device-to-device transfers hang in this mode, so collectives and" + " device_to_device cannot work and other benchmarks may hang." + f" Supported: {sorted(single_host.SUPPORTED_BENCHMARKS)}." + ) + single_host.configure_single_host_tpu() # 2. Ensure JAX is initialized init_jax_distributed() diff --git a/src/accelerator_microbenchmarks/core/single_host.py b/src/accelerator_microbenchmarks/core/single_host.py new file mode 100644 index 00000000..9b9939a3 --- /dev/null +++ b/src/accelerator_microbenchmarks/core/single_host.py @@ -0,0 +1,68 @@ +"""Standalone single-host TPU initialization inside a multi-host slice. + +Opt-in (`single_host_tpu_init` config key / CLI flag) for node-local +diagnostics. `configure_single_host_tpu()` forces a 1-host topology, skips +SliceBuilder/tpunetd and disables the cross-host XLA launch barrier. It is +process-wide and must run before the first JAX device call. The slice's ICI +routing is left untouched, so device-to-device transfers hang in this mode; +only benchmarks listed in `SUPPORTED_BENCHMARKS` are known to avoid them. +""" + +import os + +# Benchmarks known to perform no device-to-device transfers in this mode +# (gemm_throttling shares GeneralizedGemmBenchmark.generate_inputs). Lists +# primary names and registered aliases; keep in sync with +# `@benchmark_registry.register(..., aliases=...)`. +SUPPORTED_BENCHMARKS = frozenset({ + "gemm", + "gemm_generalized", + "gemm_throttling", + "hbm", + "hbm_bandwidth", +}) + +_SINGLE_HOST_ENV = { + "TPU_HOST_BOUNDS": "1,1,1", + "TPU_SKIP_MDS_QUERY": "true", + "TPU_WORKER_ID": "0", + "TPU_WORKER_HOSTNAMES": "localhost", +} + +_SINGLE_HOST_LIBTPU_FLAGS = { + "--deepsea_hal_test_skip_slicebuilder": "true", + "--deepsea_hal_test_allow_multichip_skip_slicebuilder": "true", + "--xla_tpu_use_enhanced_launch_barrier": "false", +} + + +def _set_env(name: str, value: str) -> None: + """Sets `os.environ[name]`, logging the previous value when it changes.""" + previous = os.environ.get(name) + if previous == value: + return + os.environ[name] = value + if previous is None: + print(f"single_host_tpu_init: set {name}={value}") + else: + print( + f"single_host_tpu_init: override {name}={value} (was {previous!r})" + ) + + +def configure_single_host_tpu() -> None: + """Configures libtpu to initialize this host standalone. Call before JAX init.""" + for name, value in _SINGLE_HOST_ENV.items(): + _set_env(name, value) + + # Merge into LIBTPU_INIT_ARGS by flag name so pre-existing values (e.g. from + # op_flags.yaml) are overridden rather than duplicated or left in place. + flags: dict[str, str | None] = {} + for token in os.environ.get("LIBTPU_INIT_ARGS", "").split(): + name, sep, value = token.partition("=") + flags[name] = value if sep else None + flags.update(_SINGLE_HOST_LIBTPU_FLAGS) + _set_env( + "LIBTPU_INIT_ARGS", + " ".join(k if v is None else f"{k}={v}" for k, v in flags.items()), + ) diff --git a/tests/test_cli.py b/tests/test_cli.py index 2f7b6dc5..de189f97 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -551,6 +551,57 @@ def test_benchmark_run_without_xprof_timing(self, mock_run_benchmarks): self.assertFalse(kwargs["xprof_timing"]) self.assertEqual(kwargs["xprof_dir"], "/tmp/tensorboard") + @parameterized.named_parameters( + ("flag_set", ["--single_host_tpu_init"], True), + ("flag_omitted", [], False), + ) + @mock.patch.object(runner, "run_benchmarks") + def test_benchmark_run_single_host_tpu_init( + self, extra_flags, expected, mock_run_benchmarks + ): + """Verifies that `benchmark run` forwards --single_host_tpu_init.""" + cli.run( + ["benchmark", "run", "gemm", "-m", "128", "-n", "128", "-k", "128"] + + extra_flags + ) + mock_run_benchmarks.assert_called_once() + _, kwargs = mock_run_benchmarks.call_args + self.assertIs(kwargs["single_host_tpu_init"], expected) + + @parameterized.named_parameters( + ("yaml_true", " single_host_tpu_init: true\n", [], True), + ("yaml_false", " single_host_tpu_init: false\n", [], False), + ("omitted", "", [], False), + ("cli_flag", "", ["--single_host_tpu_init"], True), + ( + "cli_flag_overrides_yaml_false", + " single_host_tpu_init: false\n", + ["--single_host_tpu_init"], + True, + ), + ) + @mock.patch.object(runner, "run_benchmarks") + def test_benchmark_run_config_single_host_tpu_init( + self, yaml_snippet, extra_flags, expected, mock_run_benchmarks + ): + """Verifies single_host_tpu_init resolution for benchmark run-config.""" + config_yaml = ( + "benchmark:\n" + " name: gemm\n" + f"{yaml_snippet}" + " params:\n" + " warmup_tries: 1\n" + " num_runs: 1\n" + ) + fake_config = self.create_tempfile(content=config_yaml) + cli.run(["benchmark", "run-config", fake_config.full_path] + extra_flags) + mock_run_benchmarks.assert_called_once() + _, kwargs = mock_run_benchmarks.call_args + self.assertIs(kwargs["single_host_tpu_init"], expected) + # The root-level key must not leak into the benchmark Config dataclass. + _, task_config = kwargs["tasks"][0] + self.assertFalse(hasattr(task_config, "single_host_tpu_init")) + @parameterized.named_parameters( ( "yaml_xprof_true", diff --git a/tests/test_config.py b/tests/test_config.py index 752bc0f6..7436970f 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -500,6 +500,9 @@ def test_benchmark_key_enum(self): config.BenchmarkKey.XPROF_DEVICE_TIMING_AGGREGATION, "xprof_device_timing_aggregation", ) + self.assertEqual( + config.BenchmarkKey.SINGLE_HOST_TPU_INIT, "single_host_tpu_init" + ) self.assertEqual( {k.value for k in config.BenchmarkKey}, { @@ -511,6 +514,7 @@ def test_benchmark_key_enum(self): "sweep", "xprof_timing", "xprof_device_timing_aggregation", + "single_host_tpu_init", }, ) @@ -654,6 +658,76 @@ def test_load_config_xprof_timing_non_bool_raises_error(self, value): self.assertIn("Expected 'xprof_timing' in benchmark", str(ctx.exception)) self.assertIn("to be a bool", str(ctx.exception)) + @parameterized.named_parameters( + ("true", "benchmark:\n name: gemm\n single_host_tpu_init: true\n", True), + ( + "false", + "benchmark:\n name: gemm\n single_host_tpu_init: false\n", + False, + ), + ("omitted", "benchmark:\n name: gemm\n", None), + ) + def test_load_config_root_single_host_tpu_init_extracted( + self, yaml_content, expected + ): + """Verifies that root single_host_tpu_init is extracted into base params.""" + config_path = os.path.join(self.test_dir.name, "config_single_host.yaml") + with open(config_path, "w", encoding="utf-8") as f: + f.write(yaml_content) + expanded = config.load_config(config_path) + self.assertNotEmpty(expanded) + self.assertEqual(expanded[0].get("single_host_tpu_init"), expected) + + @parameterized.named_parameters( + ("in_params", "params:\n single_host_tpu_init: true\n"), + ("in_sweep", "sweep:\n single_host_tpu_init: [true, false]\n"), + ) + def test_load_config_single_host_tpu_init_reserved_key_raises_error( + self, snippet + ): + """Verifies reserved key single_host_tpu_init in config sections raises.""" + yaml_content = "benchmark:\n name: gemm\n" + "".join( + f" {line}\n" for line in snippet.splitlines() + ) + config_path = os.path.join( + self.test_dir.name, "config_reserved_single_host.yaml" + ) + with open(config_path, "w", encoding="utf-8") as f: + f.write(yaml_content) + + with self.assertRaises(ValueError) as ctx: + config.load_config(config_path) + self.assertIn("reserved key(s)", str(ctx.exception)) + self.assertIn("single_host_tpu_init", str(ctx.exception)) + + @parameterized.named_parameters( + ("string_true", '"true"'), + ("integer", "1"), + ) + def test_load_config_single_host_tpu_init_non_bool_raises_error( + self, value + ): + """Verifies non-bool single_host_tpu_init at benchmark root raises.""" + yaml_content = f""" +benchmark: + name: gemm + single_host_tpu_init: {value} + params: + m: 1024 +""" + config_path = os.path.join( + self.test_dir.name, "config_invalid_single_host.yaml" + ) + with open(config_path, "w", encoding="utf-8") as f: + f.write(yaml_content) + + with self.assertRaises(ValueError) as ctx: + config.load_config(config_path) + self.assertIn( + "Expected 'single_host_tpu_init' in benchmark", str(ctx.exception) + ) + self.assertIn("to be a bool", str(ctx.exception)) + @parameterized.named_parameters(_find_all_benchmark_configs()) def test_all_configs_load_successfully(self, config_path): """Verifies that every YAML config loads and expands successfully.""" diff --git a/tests/test_matmul.py b/tests/test_matmul.py index 76b30569..4e4f2f12 100644 --- a/tests/test_matmul.py +++ b/tests/test_matmul.py @@ -80,6 +80,40 @@ def test_generate_inputs_with_scaling_factors(self): self.assertEqual(sf0.dtype, jnp.float32) self.assertEqual(sf1.dtype, jnp.float32) + def test_generate_inputs_replicated_without_device_put(self): + """Inputs are replicated, match eager sampling, and use no `device_put`. + + `single_host_tpu_init` relies on operands being generated on every device + rather than broadcast from one; a `jax.device_put` broadcast would hang. + """ + self._setup_benchmark(seed=7, use_scaling_factors=True, beta=1.0) + with mock.patch.object( + jax, "device_put", wraps=jax.device_put + ) as device_put: + inputs = self.bm.generate_inputs() + device_put.assert_not_called() + self.assertLen(inputs, 5) + for arr in inputs: + self.assertTrue(arr.sharding.is_fully_replicated) + self.assertEqual( + set(arr.sharding.device_set), set(self.mock_mesh.devices.flat) + ) + self.assertLen(arr.addressable_shards, len(self.mock_mesh.devices.flat)) + first = np.asarray(arr.addressable_shards[0].data) + for shard in arr.addressable_shards[1:]: + np.testing.assert_array_equal(np.asarray(shard.data), first) + + key = jax.random.PRNGKey(7) + k1, _, _, _, k5 = jax.random.split(key, 5) + np.testing.assert_array_equal( + np.asarray(inputs[0]), + np.asarray(jax.random.normal(k1, (64, 64)).astype(jnp.bfloat16)), + ) + np.testing.assert_array_equal( + np.asarray(inputs[4]), + np.asarray(jax.random.normal(k5, (64, 64)).astype(jnp.bfloat16)), + ) + def test_run_op(self): """Test run op without scaling factors.""" self._setup_benchmark() diff --git a/tests/test_runner.py b/tests/test_runner.py index 52846994..777667e4 100644 --- a/tests/test_runner.py +++ b/tests/test_runner.py @@ -1,6 +1,8 @@ """Unit tests for the benchmark execution runner (runner.py).""" +import io import os +import sys import tempfile from unittest import mock @@ -206,6 +208,66 @@ def test_init_jax_distributed(self): runner.init_jax_distributed() mock_init.assert_called_once() + def _run_benchmarks_with_stages_mocked(self, task_names, **kwargs): + """Runs run_benchmarks for `task_names` with TPU-touching stages mocked. + + Returns the ordered list of init stages invoked and the captured stdout. + """ + calls = [] + params = base.BaseBenchmarkParams(warmup_tries=1, num_runs=1) + with mock.patch.object( + runner, "set_xla_flags", side_effect=lambda *a, **k: calls.append("xla") + ), mock.patch.object( + runner.single_host, + "configure_single_host_tpu", + side_effect=lambda: calls.append("single_host"), + ), mock.patch.object( + runner, "init_jax_distributed", side_effect=lambda: calls.append("jax") + ), mock.patch.object( + runner.platform, "get_platform_info" + ), mock.patch.object( + runner.system, "get_hardware_spec" + ), mock.patch.object( + runner.benchmark_loader, "load_all_benchmarks" + ), mock.patch.object( + runner.registry.benchmark_registry, "get_benchmark" + ), mock.patch.object( + runner.report, "report_results" + ), mock.patch.object( + sys, "stdout", new=io.StringIO() + ) as fake_out: + runner.run_benchmarks( + tasks=[(name, params) for name in task_names], + output_dir=self.temp_dir.name, + **kwargs, + ) + return calls, fake_out.getvalue() + + def test_run_benchmarks_single_host_tpu_init_off_by_default(self): + """Without the flag, single-host configuration is never applied.""" + calls, _ = self._run_benchmarks_with_stages_mocked(["gemm"]) + self.assertEqual(calls, ["xla", "jax"]) + + def test_run_benchmarks_single_host_tpu_init_order(self): + """Single-host env is applied after op_flags and before JAX init.""" + calls, out = self._run_benchmarks_with_stages_mocked( + ["gemm", "hbm"], single_host_tpu_init=True + ) + self.assertEqual(calls, ["xla", "single_host", "jax"]) + self.assertNotIn("Warning: single_host_tpu_init", out) + + def test_run_benchmarks_single_host_tpu_init_warns_unsupported(self): + """Benchmarks outside the allow-list produce a warning but still run.""" + calls, out = self._run_benchmarks_with_stages_mocked( + ["all_reduce"], single_host_tpu_init=True + ) + self.assertEqual(calls, ["xla", "single_host", "jax"]) + self.assertIn( + "Warning: single_host_tpu_init is not supported for ['all_reduce']", + out, + ) + self.assertIn("hang", out) + def test_run_benchmarks(self): """Verifies that run_benchmarks executes typed benchmark tasks.""" dummy_params = base.BaseBenchmarkParams(warmup_tries=1, num_runs=1) diff --git a/tests/test_single_host.py b/tests/test_single_host.py new file mode 100644 index 00000000..7aed86bf --- /dev/null +++ b/tests/test_single_host.py @@ -0,0 +1,129 @@ +"""Unit tests for single_host.py.""" + +import io +import os +import sys +from unittest import mock + +from absl.testing import absltest +from absl.testing import parameterized +from accelerator_microbenchmarks.benchmarks import benchmark_loader +from accelerator_microbenchmarks.core import registry +from accelerator_microbenchmarks.core import single_host + + +def _flags(args: str) -> dict[str, str | None]: + """Parses a LIBTPU_INIT_ARGS string into a name -> value mapping.""" + parsed = {} + for token in args.split(): + name, sep, value = token.partition("=") + parsed[name] = value if sep else None + return parsed + + +class SingleHostTest(parameterized.TestCase): + """Unit tests for configure_single_host_tpu().""" + + _EXPECTED_FLAGS = { + "--deepsea_hal_test_skip_slicebuilder": "true", + "--deepsea_hal_test_allow_multichip_skip_slicebuilder": "true", + "--xla_tpu_use_enhanced_launch_barrier": "false", + } + + def _configure(self, env: dict[str, str]) -> tuple[dict[str, str], str]: + """Runs configure_single_host_tpu() in an isolated env; returns env, log.""" + with mock.patch.dict(os.environ, env, clear=True): + with mock.patch.object(sys, "stdout", new=io.StringIO()) as fake_out: + single_host.configure_single_host_tpu() + return dict(os.environ), fake_out.getvalue() + + def test_sets_single_host_topology_env(self): + """Forces a 1-host topology and skips the metadata server query.""" + env, _ = self._configure({}) + self.assertEqual(env["TPU_HOST_BOUNDS"], "1,1,1") + self.assertEqual(env["TPU_WORKER_ID"], "0") + self.assertEqual(env["TPU_WORKER_HOSTNAMES"], "localhost") + self.assertEqual(env["TPU_SKIP_MDS_QUERY"], "true") + self.assertEqual(_flags(env["LIBTPU_INIT_ARGS"]), self._EXPECTED_FLAGS) + + def test_overrides_injected_multi_host_env_and_logs_previous_value(self): + """GKE-injected multi-host values are replaced and the change is logged.""" + env, log = self._configure({ + "TPU_WORKER_ID": "2", + "TPU_WORKER_HOSTNAMES": "host-0,host-1,host-2,host-3", + "TPU_HOST_BOUNDS": "2,2,1", + }) + self.assertEqual(env["TPU_WORKER_ID"], "0") + self.assertEqual(env["TPU_WORKER_HOSTNAMES"], "localhost") + self.assertEqual(env["TPU_HOST_BOUNDS"], "1,1,1") + self.assertIn("override TPU_WORKER_ID=0 (was '2')", log) + self.assertIn("override TPU_HOST_BOUNDS=1,1,1 (was '2,2,1')", log) + + def test_preserves_existing_libtpu_flags(self): + """Flags from op_flags.yaml are kept, in order, ahead of the additions.""" + env, _ = self._configure( + {"LIBTPU_INIT_ARGS": "--xla_tpu_foo=1 --xla_tpu_bar --xla_tpu_baz=x"} + ) + args = env["LIBTPU_INIT_ARGS"] + self.assertTrue( + args.startswith("--xla_tpu_foo=1 --xla_tpu_bar --xla_tpu_baz=x ") + ) + self.assertEqual( + _flags(args), + { + "--xla_tpu_foo": "1", + "--xla_tpu_bar": None, + "--xla_tpu_baz": "x", + **self._EXPECTED_FLAGS, + }, + ) + + @parameterized.named_parameters( + ("lowercase_true", "--xla_tpu_use_enhanced_launch_barrier=true"), + ("capitalised_true", "--xla_tpu_use_enhanced_launch_barrier=True"), + ("numeric_true", "--xla_tpu_use_enhanced_launch_barrier=1"), + ("bare_flag", "--xla_tpu_use_enhanced_launch_barrier"), + ) + def test_overrides_launch_barrier_flag_regardless_of_spelling(self, flag): + """Any pre-set launch-barrier value is overridden in place, not appended.""" + env, log = self._configure( + {"LIBTPU_INIT_ARGS": f"--xla_tpu_foo=1 {flag} --xla_tpu_bar=2"} + ) + args = env["LIBTPU_INIT_ARGS"] + self.assertEqual( + args.split()[:3], + [ + "--xla_tpu_foo=1", + "--xla_tpu_use_enhanced_launch_barrier=false", + "--xla_tpu_bar=2", + ], + ) + self.assertEqual(args.count("--xla_tpu_use_enhanced_launch_barrier"), 1) + self.assertIn("override LIBTPU_INIT_ARGS=", log) + + def test_idempotent(self): + """A second call changes nothing and does not duplicate flags.""" + with mock.patch.dict( + os.environ, {"LIBTPU_INIT_ARGS": "--xla_tpu_foo=1"}, clear=True + ): + with mock.patch.object(sys, "stdout", new=io.StringIO()): + single_host.configure_single_host_tpu() + first = dict(os.environ) + with mock.patch.object(sys, "stdout", new=io.StringIO()) as fake_out: + single_host.configure_single_host_tpu() + self.assertEqual(dict(os.environ), first) + self.assertEqual(fake_out.getvalue(), "") + + def test_supported_benchmarks_are_registered_names(self): + """The allow-list only names benchmarks (or aliases) that exist.""" + benchmark_loader.load_all_benchmarks() + registered = set( + registry.benchmark_registry.list_benchmark_names( + include_experimental=True, include_aliases=True + ) + ) + self.assertContainsSubset(single_host.SUPPORTED_BENCHMARKS, registered) + + +if __name__ == "__main__": + absltest.main()