diff --git a/pkgs/clan-cli/clan_cli/generate_test_vars/cli.py b/pkgs/clan-cli/clan_cli/generate_test_vars/cli.py index 7ddf0856b..c980b7159 100755 --- a/pkgs/clan-cli/clan_cli/generate_test_vars/cli.py +++ b/pkgs/clan-cli/clan_cli/generate_test_vars/cli.py @@ -75,13 +75,14 @@ class TestFlake(Flake): def path(self) -> Path: return self.test_dir - def select_machine(self, machine_name: str, selector: str) -> Any: - """Select a nix attribute for a specific machine. + def machine_selector(self, machine_name: str, selector: str) -> str: + """Create a selector for a specific machine. Args: machine_name: The name of the machine selector: The attribute selector string relative to the machine config - apply: Optional function to apply to the result + Returns: + The full selector string for the machine """ config = nix_config() @@ -89,9 +90,7 @@ class TestFlake(Flake): test_system = system if system.endswith("-darwin"): test_system = system.rstrip("darwin") + "linux" - - full_selector = f'checks."{test_system}".{self.check_attr}.machinesCross.{system}."{machine_name}".{selector}' - return self.select(full_selector) + return f'checks."{test_system}".{self.check_attr}.machinesCross."{system}"."{machine_name}".{selector}' # we don't want to evaluate all machines of the flake. Only the ones defined in the test def set_machine_names(self, machine_names: list[str]) -> None: diff --git a/pkgs/clan-cli/clan_cli/tests/test_vars.py b/pkgs/clan-cli/clan_cli/tests/test_vars.py index 3fe679c2b..fb0f19547 100644 --- a/pkgs/clan-cli/clan_cli/tests/test_vars.py +++ b/pkgs/clan-cli/clan_cli/tests/test_vars.py @@ -1264,65 +1264,74 @@ def test_cache_misses_for_vars_operations( flake: ClanFlake, ) -> None: """Test that vars operations result in minimal cache misses.""" + # Set up first machine with two generators config = flake.machines["my_machine"] = create_test_machine_config() - # Set up a simple generator with a public value - my_generator = config["clan"]["core"]["vars"]["generators"]["my_generator"] - my_generator["files"]["my_value"]["secret"] = False - my_generator["script"] = 'echo -n "test_value" > "$out"/my_value' + # Set up two generators with public values + gen1 = config["clan"]["core"]["vars"]["generators"]["gen1"] + gen1["files"]["value1"]["secret"] = False + gen1["script"] = 'echo -n "test_value1" > "$out"/value1' + + gen2 = config["clan"]["core"]["vars"]["generators"]["gen2"] + gen2["files"]["value2"]["secret"] = False + gen2["script"] = 'echo -n "test_value2" > "$out"/value2' + + # Add a second machine with the same generator configuration + flake.machines["other_machine"] = config.copy() flake.refresh() monkeypatch.chdir(flake.path) - # Create a fresh machine object to ensure clean cache state - machine = Machine(name="my_machine", flake=Flake(str(flake.path))) + # Create fresh machine objects to ensure clean cache state + flake_obj = Flake(str(flake.path)) + machine1 = Machine(name="my_machine", flake=flake_obj) + machine2 = Machine(name="other_machine", flake=flake_obj) - # Test 1: Running vars generate with a fresh cache should result in exactly 3 cache misses - # Expected cache misses: - # 1. One for getting the list of generators - # 2. One for getting the final script of our test generator (my_generator) - # 3. One for getting the final script of the state version generator (added by default) - # TODO: The third cache miss is undesired in tests. disable state version module for tests + # Test 1: Running vars generate for BOTH machines simultaneously should still result in exactly 2 cache misses + # Even though we have: + # - 2 machines (my_machine and other_machine) + # - 2 generators per machine (gen1 and gen2) + # We still only get 2 cache misses when generating for both machines: + # 1. One for getting the list of generators for both machines + # 2. One batched evaluation for getting all generator scripts for both machines + # The key insight: the system should batch ALL evaluations across ALL machines into a single nix eval run_generators( - machines=[machine], + machines=[machine1, machine2], generators=None, # Generate all ) - # Print stack traces if we have more than 3 cache misses - if machine.flake._cache_misses != 3: - machine.flake.print_cache_miss_analysis( + # Print stack traces if we have more than 2 cache misses + if flake_obj._cache_misses != 2: + flake_obj.print_cache_miss_analysis( title="Cache miss analysis for vars generate" ) - assert machine.flake._cache_misses == 2, ( - f"Expected exactly 2 cache misses for vars generate, got {machine.flake._cache_misses}" + assert flake_obj._cache_misses == 2, ( + f"Expected exactly 2 cache misses for vars generate, got {flake_obj._cache_misses}" ) - # Verify the value was generated correctly - var_value = get_machine_var(machine, "my_generator/my_value") - assert var_value.printable_value == "test_value" - # Test 2: List all vars should result in exactly 1 cache miss # Force cache invalidation (this also resets cache miss tracking) invalidate_flake_cache(flake.path) - machine.flake.invalidate_cache() + flake_obj.invalidate_cache() - stringify_all_vars(machine) - assert machine.flake._cache_misses == 1, ( - f"Expected exactly 1 cache miss for vars list, got {machine.flake._cache_misses}" + stringify_all_vars(machine1) + assert flake_obj._cache_misses == 1, ( + f"Expected exactly 1 cache miss for vars list, got {flake_obj._cache_misses}" ) # Test 3: Getting a specific var with a fresh cache should result in exactly 1 cache miss # Force cache invalidation (this also resets cache miss tracking) invalidate_flake_cache(flake.path) - machine.flake.invalidate_cache() + flake_obj.invalidate_cache() - var_value = get_machine_var(machine, "my_generator/my_value") - assert var_value.printable_value == "test_value" + # Only test gen1 for the get operation + var_value = get_machine_var(machine1, "gen1/value1") + assert var_value.printable_value == "test_value1" - assert machine.flake._cache_misses == 1, ( - f"Expected exactly 1 cache miss for vars get with fresh cache, got {machine.flake._cache_misses}" + assert flake_obj._cache_misses == 1, ( + f"Expected exactly 1 cache miss for vars get with fresh cache, got {flake_obj._cache_misses}" ) diff --git a/pkgs/clan-cli/clan_cli/vars/generator.py b/pkgs/clan-cli/clan_cli/vars/generator.py index 1cf1bdda0..c1eae5775 100644 --- a/pkgs/clan-cli/clan_cli/vars/generator.py +++ b/pkgs/clan-cli/clan_cli/vars/generator.py @@ -245,15 +245,19 @@ class Generator: return sec_store.get(self, prompt.name).decode() return None + def final_script_selector(self, machine_name: str) -> str: + if self._flake is None: + msg = "Flake cannot be None" + raise ClanError(msg) + return self._flake.machine_selector( + machine_name, f'config.clan.core.vars.generators."{self.name}".finalScript' + ) + def final_script(self, machine: "Machine") -> Path: if self._flake is None: msg = "Flake cannot be None" raise ClanError(msg) - output = Path( - machine.select( - f'config.clan.core.vars.generators."{self.name}".finalScript', - ), - ) + output = Path(self._flake.select(self.final_script_selector(machine.name))) if tmp_store := nix_test_store(): output = tmp_store.joinpath(*output.parts[1:]) return output diff --git a/pkgs/clan-cli/clan_lib/flake/flake.py b/pkgs/clan-cli/clan_lib/flake/flake.py index e8346e508..792ea1151 100644 --- a/pkgs/clan-cli/clan_lib/flake/flake.py +++ b/pkgs/clan-cli/clan_lib/flake/flake.py @@ -1132,6 +1132,20 @@ class Flake: return self._cache.select(selector) + def machine_selector(self, machine_name: str, selector: str) -> str: + """Create a selector for a specific machine. + + Args: + machine_name: The name of the machine + selector: The attribute selector string relative to the machine config + Returns: + The full selector string for the machine + + """ + config = nix_config() + system = config["system"] + return f'clanInternals.machines."{system}"."{machine_name}".{selector}' + def select_machine(self, machine_name: str, selector: str) -> Any: """Select a nix attribute for a specific machine. @@ -1141,11 +1155,7 @@ class Flake: apply: Optional function to apply to the result """ - config = nix_config() - system = config["system"] - - full_selector = f'clanInternals.machines."{system}"."{machine_name}".{selector}' - return self.select(full_selector) + return self.select(self.machine_selector(machine_name, selector)) def list_machines( self, diff --git a/pkgs/clan-cli/clan_lib/vars/generate.py b/pkgs/clan-cli/clan_lib/vars/generate.py index c9b602199..67ab0e7f1 100644 --- a/pkgs/clan-cli/clan_lib/vars/generate.py +++ b/pkgs/clan-cli/clan_lib/vars/generate.py @@ -177,13 +177,25 @@ def run_generators( for machine in machines: _ensure_healthy(machine=machine) + # get the flake via any machine (they are all the same) + flake = machines[0].flake + + def get_generator_machine(generator: Generator) -> Machine: + if generator.machine is None: + # return first machine if generator is not tied to a specific one + return machines[0] + return Machine(name=generator.machine, flake=flake) + + # preheat the select cache, to reduce repeated calls during execution + selectors = [] + for generator in generator_objects: + machine = get_generator_machine(generator) + selectors.append(generator.final_script_selector(machine.name)) + flake.precache(selectors) + # execute generators for generator in generator_objects: - machine = ( - machines[0] - if generator.machine is None - else Machine(name=generator.machine, flake=machines[0].flake) - ) + machine = get_generator_machine(generator) if check_can_migrate(machine, generator): migrate_files(machine, generator) else: