From f7e1b676a433406b9845b3a92a0136949911e886 Mon Sep 17 00:00:00 2001 From: Victor Morand Date: Mon, 19 Oct 2026 17:34:48 +0200 Subject: [PATCH 1/3] feat: proposition to improve launchers installation and debug --- src/experimaestro/launcherfinder/__init__.py | 2 +- src/experimaestro/launcherfinder/registry.py | 41 ++- src/experimaestro/launchers/direct.py | 178 ++++++++--- src/experimaestro/launchers/slurm/cli.py | 276 ++++++++++++++++++ .../launchers/slurm/templates/launchers.py.j2 | 3 +- src/experimaestro/launchers/slurm/tui.py | 15 +- .../tests/test_direct_generate.py | 49 ++++ src/experimaestro/tests/test_findlauncher.py | 47 +++ .../tests/test_slurm_generate.py | 86 ++++++ 9 files changed, 638 insertions(+), 59 deletions(-) diff --git a/src/experimaestro/launcherfinder/__init__.py b/src/experimaestro/launcherfinder/__init__.py index c3f698ee..62e7dc09 100644 --- a/src/experimaestro/launcherfinder/__init__.py +++ b/src/experimaestro/launcherfinder/__init__.py @@ -15,5 +15,5 @@ HostRequirement, MatchRequirement, ) -from .registry import find_launcher, LauncherRegistry +from .registry import find_launcher, LauncherRegistry, LauncherNotFoundError from .parser import parse diff --git a/src/experimaestro/launcherfinder/registry.py b/src/experimaestro/launcherfinder/registry.py index 81e58aa2..a31c2b1a 100644 --- a/src/experimaestro/launcherfinder/registry.py +++ b/src/experimaestro/launcherfinder/registry.py @@ -77,6 +77,8 @@ def set_config_dir(config_dir: Path): LauncherRegistry.CURRENT_CONFIG_DIR = config_dir def __init__(self, basepath: Path): + self.basepath = basepath + self.launchers_py = basepath / "launchers.py" self.connectors_schema: Dict[str, Type] = {} self.tokens_schema: Dict[str, Type] = {} self.find_launcher_fn = None @@ -89,13 +91,12 @@ def __init__(self, basepath: Path): entry_point.load().init_registry(self) # Register the find launcher function if it exists - launchers_py = basepath / "launchers.py" - if launchers_py.is_file(): - logger.info("Loading %s", launchers_py) + if self.launchers_py.is_file(): + logger.info("Loading %s", self.launchers_py) from importlib import util - with ensure_enter(launchers_py.__fspath__()) as fp: + with ensure_enter(self.launchers_py.__fspath__()) as fp: spec = util.spec_from_file_location("xpm_launchers_conf", fp) module = util.module_from_spec(spec) spec.loader.exec_module(module) @@ -103,7 +104,7 @@ def __init__(self, basepath: Path): self.find_launcher_fn = getattr(module, "find_launcher", None) if self.find_launcher_fn is None: logger.warning( - "No find_launcher() function was found in %s", launchers_py + "No find_launcher() function was found in %s", self.launchers_py ) # Read the configuration file @@ -147,7 +148,14 @@ def find( """ if self.find_launcher_fn is None: - logger.info("No launchers.yaml file: using local host ") + logger.warning( + "No launcher configuration found at %s: falling back to local DirectLauncher.\n" + "Tip: To configure launchers with resource management:\n" + " - Local machine / laptop: experimaestro launchers direct generate -y\n" + " - SLURM cluster: experimaestro launchers slurm generate -y\n" + " - Documentation: https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html", + self.launchers_py, + ) from experimaestro.launchers.direct import DirectLauncher from experimaestro.connectors.local import LocalConnector @@ -170,7 +178,7 @@ def find( for spec in specs.requirements: if launcher := self.find_launcher_fn(spec, tags): assert isinstance(launcher, Launcher), ( - "f{self.find_launcher_fn} did not return a Launcher but {type(launcher)}" + f"{self.find_launcher_fn} did not return a Launcher but {type(launcher)}" ) return launcher @@ -181,11 +189,22 @@ def find_launcher( *specs: Union[HostRequirement, str], tags: Set[str] = set() ) -> "Launcher": """Find a launcher matching a given specification""" - launcher = LauncherRegistry.instance().find(*specs, tags=tags) + registry = LauncherRegistry.instance() + launcher = registry.find(*specs, tags=tags) if not launcher: + specs_str = ", ".join(str(s) for s in specs) + tags_info = f" (tags={tags})" if tags else "" + sample_spec = str(specs[0]) if specs else "cpu(cores=2)" raise LauncherNotFoundError( - f"No launcher with specification: {specs}." - "Please refer to the documentation at the following URL: " - "https://experimaestro-python.readthedocs.io/en/latest/launchers/" + f"No launcher matched specification: {specs_str}{tags_info} " + f"(using {registry.launchers_py}).\n" + "How to fix:\n" + " 1. Adjust the requested resources or add a fallback with '|' " + "(e.g. 'gpu(mem=4G) | cpu(cores=2)') if they exceed your host capabilities.\n" + f" 2. Regenerate or edit {registry.launchers_py}:\n" + " - Local machine / laptop: experimaestro launchers direct generate -y\n" + " - SLURM cluster: experimaestro launchers slurm generate -y\n" + f" 3. Test your specification with: experimaestro find-launchers \"{sample_spec}\"\n" + " 4. Documentation: https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html" ) return launcher diff --git a/src/experimaestro/launchers/direct.py b/src/experimaestro/launchers/direct.py index dc4a1a9c..7dfd70d3 100644 --- a/src/experimaestro/launchers/direct.py +++ b/src/experimaestro/launchers/direct.py @@ -189,11 +189,62 @@ def direct_cli(): is_flag=True, help="Overwrite existing file without prompting", ) -def generate(output: Path | None, force: bool): +@click.option( + "--yes", + "-y", + is_flag=True, + help="Non-interactive mode: accept defaults without prompting", +) +@click.option( + "--memory-tokens/--no-memory-tokens", + default=None, + help="Enable or disable CPU memory tokens (default: enabled)", +) +@click.option( + "--reserved-memory", + type=float, + default=None, + help="Reserved memory for OS in GiB (default: 2.0)", +) +@click.option( + "--memory-token-unit", + type=int, + default=None, + help=f"Memory token unit size in MiB (default: {DEFAULT_TOKEN_UNIT_MIB})", +) +@click.option( + "--gpu-tokens/--no-gpu-tokens", + default=None, + help="Enable or disable GPU tokens when GPUs are detected (default: enabled)", +) +@click.option( + "--gpu-mode", + type=click.Choice(["exclusive", "memory"]), + default=None, + help="GPU token mode: 'exclusive' (1 token/GPU) or 'memory' (default: exclusive for CUDA, memory for MPS)", +) +@click.option( + "--gpu-token-unit", + type=int, + default=None, + help="GPU memory token unit size in MiB when using memory mode (default: 1024)", +) +def generate( + output: Path | None, + force: bool, + yes: bool, + memory_tokens: bool | None, + reserved_memory: float | None, + memory_token_unit: int | None, + gpu_tokens: bool | None, + gpu_mode: str | None, + gpu_token_unit: int | None, +): """Generate a launchers.py file for direct (local) execution. - This command detects system resources and interactively configures - memory-based tokens for CPU and GPU resource management. + This command detects system resources and configures memory-based tokens + for CPU and GPU resource management. Use `-y / --yes` for non-interactive + generation with sensible defaults. Token system: - Memory tokens: 1 token = configurable unit (default 256 MiB) @@ -205,8 +256,9 @@ def generate(output: Path | None, force: bool): - GPU memory requests reduce available CPU memory tokens Examples: + experimaestro launchers direct generate -y experimaestro launchers direct generate - experimaestro launchers direct generate -o ./launchers.py + experimaestro launchers direct generate -o ./launchers.py -y """ from rich.console import Console from rich.panel import Panel @@ -214,6 +266,7 @@ def generate(output: Path | None, force: bool): from rich.table import Table console = Console() + non_interactive = yes or not sys.stdin.isatty() # Default output path if output is None: @@ -233,8 +286,12 @@ def generate(output: Path | None, force: bool): ) # Check if file exists - if output.exists() and not force: + if output.exists() and not (force or yes): console.print(f"\n[yellow]File already exists:[/yellow] {output}") + if non_interactive: + raise click.ClickException( + f"File already exists: {output}. Pass --force (-f) or --yes (-y) to overwrite." + ) if not Confirm.ask("Overwrite?", default=False): console.print("[yellow]Aborted.[/yellow]") return @@ -272,31 +329,47 @@ def generate(output: Path | None, force: bool): config.is_mps = sys_info.has_mps # === Memory tokens configuration === - console.print( - Panel( - "[bold]Memory Tokens[/bold]\n\n" - "Memory tokens prevent over-committing system RAM.\n" - "Each token represents a unit of memory (default: 256 MiB).\n" - "Tasks request tokens based on their memory needs.\n\n" - "[dim]Example: 1 GiB requirement = 4 tokens (with 256 MiB unit)[/dim]", - border_style="blue", + if not non_interactive: + console.print( + Panel( + "[bold]Memory Tokens[/bold]\n\n" + "Memory tokens prevent over-committing system RAM.\n" + "Each token represents a unit of memory (default: 256 MiB).\n" + "Tasks request tokens based on their memory needs.\n\n" + "[dim]Example: 1 GiB requirement = 4 tokens (with 256 MiB unit)[/dim]", + border_style="blue", + ) ) - ) - config.use_memory_tokens = Confirm.ask("\nEnable memory tokens?", default=True) + if memory_tokens is not None: + config.use_memory_tokens = memory_tokens + elif non_interactive: + config.use_memory_tokens = True + else: + config.use_memory_tokens = Confirm.ask("\nEnable memory tokens?", default=True) if config.use_memory_tokens: # Reserved memory for OS - config.reserved_memory_gib = FloatPrompt.ask( - " Reserved memory for OS (GiB)", - default=2.0, - ) + if reserved_memory is not None: + config.reserved_memory_gib = reserved_memory + elif non_interactive: + config.reserved_memory_gib = 2.0 + else: + config.reserved_memory_gib = FloatPrompt.ask( + " Reserved memory for OS (GiB)", + default=2.0, + ) # Token unit size - config.memory_token_unit_mib = IntPrompt.ask( - " Token unit size (MiB)", - default=DEFAULT_TOKEN_UNIT_MIB, - ) + if memory_token_unit is not None: + config.memory_token_unit_mib = memory_token_unit + elif non_interactive: + config.memory_token_unit_mib = DEFAULT_TOKEN_UNIT_MIB + else: + config.memory_token_unit_mib = IntPrompt.ask( + " Token unit size (MiB)", + default=DEFAULT_TOKEN_UNIT_MIB, + ) # Calculate available tokens available_bytes = sys_info.total_memory_bytes - int( @@ -312,24 +385,30 @@ def generate(output: Path | None, force: bool): # === GPU tokens configuration === if sys_info.gpus: - console.print() - console.print( - Panel( - "[bold]GPU Tokens[/bold]\n\n" - "GPU tokens control concurrent GPU task execution.\n\n" - "[cyan]Exclusive mode:[/cyan] 1 token = 1 GPU (simple, recommended)\n" - "[cyan]Memory mode:[/cyan] Tokens based on GPU memory units\n" - + ( - "\n[yellow]Note: Apple MPS uses unified memory - GPU memory\n" - "requests will reduce available CPU memory tokens.[/yellow]" - if sys_info.has_mps - else "" - ), - border_style="blue", + if not non_interactive: + console.print() + console.print( + Panel( + "[bold]GPU Tokens[/bold]\n\n" + "GPU tokens control concurrent GPU task execution.\n\n" + "[cyan]Exclusive mode:[/cyan] 1 token = 1 GPU (simple, recommended)\n" + "[cyan]Memory mode:[/cyan] Tokens based on GPU memory units\n" + + ( + "\n[yellow]Note: Apple MPS uses unified memory - GPU memory\n" + "requests will reduce available CPU memory tokens.[/yellow]" + if sys_info.has_mps + else "" + ), + border_style="blue", + ) ) - ) - config.use_gpu_tokens = Confirm.ask("\nEnable GPU tokens?", default=True) + if gpu_tokens is not None: + config.use_gpu_tokens = gpu_tokens + elif non_interactive: + config.use_gpu_tokens = True + else: + config.use_gpu_tokens = Confirm.ask("\nEnable GPU tokens?", default=True) if config.use_gpu_tokens: config.gpu_count = len(sys_info.gpus) @@ -342,6 +421,10 @@ def generate(output: Path | None, force: bool): ) config.gpu_mode = "memory" config.mps_shared_memory = True + elif gpu_mode is not None: + config.gpu_mode = gpu_mode + elif non_interactive: + config.gpu_mode = "exclusive" else: config.gpu_mode = Prompt.ask( " GPU token mode", @@ -356,10 +439,15 @@ def generate(output: Path | None, force: bool): f"(1 per GPU)[/green]" ) else: - config.gpu_token_unit_mib = IntPrompt.ask( - " GPU token unit size (MiB)", - default=1024, - ) + if gpu_token_unit is not None: + config.gpu_token_unit_mib = gpu_token_unit + elif non_interactive: + config.gpu_token_unit_mib = 1024 + else: + config.gpu_token_unit_mib = IntPrompt.ask( + " GPU token unit size (MiB)", + default=1024, + ) # Calculate GPU tokens from first GPU (or total for MPS) if sys_info.has_mps: # For MPS, GPU memory is same as system memory @@ -410,7 +498,7 @@ def generate(output: Path | None, force: bool): console.print() # Confirm generation - if not Confirm.ask("Generate configuration?", default=True): + if not non_interactive and not Confirm.ask("Generate configuration?", default=True): console.print("[yellow]Aborted.[/yellow]") return @@ -425,7 +513,7 @@ def generate(output: Path | None, force: bool): config.use_gpu_tokens and config.total_gpu_tokens > 0 ): tokens_yaml_path = output.parent / "tokens.yaml" - write_tokens = force or not tokens_yaml_path.exists() + write_tokens = force or yes or non_interactive or not tokens_yaml_path.exists() if not write_tokens: write_tokens = Confirm.ask(f"Generate {tokens_yaml_path}?", default=True) diff --git a/src/experimaestro/launchers/slurm/cli.py b/src/experimaestro/launchers/slurm/cli.py index a4603aec..bfb8f044 100644 --- a/src/experimaestro/launchers/slurm/cli.py +++ b/src/experimaestro/launchers/slurm/cli.py @@ -920,3 +920,279 @@ def configure(cache_file: Path | None, use_cache: bool, config_path: Path | None use_cache=use_cache, config_path=config_path, ) + + +_KNOWN_GPU_MEMORY_GB: list[tuple[str, int]] = [ + ("h200", 141), + ("b200", 192), + ("h100", 80), + ("a100_80", 80), + ("a100-80", 80), + ("a100", 40), + ("l40", 48), + ("a40", 48), + ("a6000", 48), + ("rtx6000", 48), + ("v100_32", 32), + ("v100-32", 32), + ("v100", 16), + ("a10g", 24), + ("a10", 24), + ("4090", 24), + ("3090", 24), + ("rtx8000", 48), + ("t4", 16), + ("p100", 16), + ("a30", 24), +] + + +def _infer_gpu_memory_gb( + gpu_type: str | None, default_gb: int = 16 +) -> tuple[int, bool]: + """Infer GPU memory in GB from GPU type string when not explicitly configured. + + Returns: + Tuple of (memory_gb, matched_known_type) + """ + if gpu_type: + lower = gpu_type.lower() + # Check for explicit size suffix like "80g" or "40gb" + m = re.search(r"(\d+)\s*gb?\b", lower) + if m: + return int(m.group(1)), True + for prefix, mem_gb in _KNOWN_GPU_MEMORY_GB: + if prefix in lower: + return mem_gb, True + return default_gb, False + + +def _generate_slurm_launchers_py( + config: "SlurmConfig", + config_path: Path, + default_gpu_memory_gb: int = 16, +) -> str: + """Render SLURM launchers.py from a SlurmConfig using Jinja2 template.""" + import logging + import jinja2 + + logger = logging.getLogger("xpm.slurm") + template_env = jinja2.Environment( + loader=jinja2.PackageLoader("experimaestro.launchers.slurm", "templates"), + keep_trailing_newline=True, + ) + template = template_env.get_template("launchers.py.j2") + + # Sort enabled partitions by priority descending, then name + enabled = sorted( + config.enabled_partitions.values(), + key=lambda p: (-p.priority.value, p.name), + ) + + partitions_ctx = [] + for p in enabled: + mem_bytes = ( + p.effective_memory_gb * (1024**3) + if (p.memory_gb.is_overridden or p.memory_is_inferred) + else p.memory_bytes + ) + gpu_mem_gb = p.effective_gpu_memory_gb + if p.gpus_per_node > 0 and gpu_mem_gb <= 0: + gpu_mem_gb, matched = _infer_gpu_memory_gb( + p.effective_gpu_type, default_gb=default_gpu_memory_gb + ) + source_desc = ( + f"inferred from GPU type '{p.effective_gpu_type}'" + if matched + else "fallback default" + ) + logger.warning( + "Partition '%s' has %d GPU(s) without explicit GPU memory in %s; " + "using %d GB (%s). Set 'gpu_memory_gb' in %s to override.", + p.name, + p.gpus_per_node, + config_path, + gpu_mem_gb, + source_desc, + config_path, + ) + + allowed_qos = p.effective_qos + qos = allowed_qos[0] if allowed_qos else None + + partitions_ctx.append( + { + "name": p.name, + "priority": p.priority.value, + "cpus": p.effective_cores, + "memory_bytes": mem_bytes, + "gpus": p.gpus_per_node, + "gpu_type": p.effective_gpu_type, + "gpu_memory_bytes": gpu_mem_gb * (1024**3) if p.gpus_per_node > 0 else 0, + "time_limit_seconds": p.time_limit_seconds, + "qos": qos, + } + ) + + return template.render( + cluster_name=config.cluster_name, + config_file=str(config_path), + default_account=config.default_account.value, + consider_priority=config.consider_priority.value, + partitions=partitions_ctx, + ) + + +@cli.command("generate") +@click.option( + "--output", + "-o", + type=click.Path(path_type=Path), + default=None, + help="Output path for launchers.py (default: ~/.config/experimaestro/launchers.py)", +) +@click.option( + "--config", + "-c", + "config_path", + type=click.Path(path_type=Path), + default=None, + help="Path to SLURM configuration YAML (default: ~/.config/experimaestro/slurm.yaml)", +) +@click.option( + "--cache-file", + type=click.Path(path_type=Path), + default=None, + help="JSON file for cached SLURM command outputs", +) +@click.option( + "--use-cache", + is_flag=True, + help="Use cached SLURM outputs instead of running commands", +) +@click.option( + "--default-gpu-memory", + type=int, + default=16, + help="Fallback GPU memory in GB when not reported by SLURM or slurm.yaml (default: 16)", +) +@click.option( + "--force", + "-f", + is_flag=True, + help="Overwrite existing launchers.py without prompting", +) +@click.option( + "--yes", + "-y", + is_flag=True, + help="Non-interactive mode: accept defaults without prompting", +) +def generate( + output: Path | None, + config_path: Path | None, + cache_file: Path | None, + use_cache: bool, + default_gpu_memory: int, + force: bool, + yes: bool, +): + """Generate a launchers.py file from SLURM cluster information and slurm.yaml. + + Detects SLURM partitions (or reads from cache), merges any overrides from + slurm.yaml, saves slurm.yaml if not present, and writes launchers.py. + Use `-y / --yes` for non-interactive generation. + + Examples: + experimaestro launchers slurm generate -y + experimaestro launchers slurm generate -o ./launchers.py -y + """ + from rich.console import Console + from rich.prompt import Confirm + from rich.table import Table + from .models import SlurmConfig + + console = Console() + non_interactive = yes or not sys.stdin.isatty() + + if config_path is None: + config_path = Path("~/.config/experimaestro/slurm.yaml").expanduser() + else: + config_path = config_path.expanduser() + + if output is None: + output = config_path.parent / "launchers.py" + else: + output = output.expanduser() + + if cache_file is not None: + cache_file = cache_file.expanduser() + + if output.exists() and not (force or yes): + console.print(f"\n[yellow]File already exists:[/yellow] {output}") + if non_interactive: + raise click.ClickException( + f"File already exists: {output}. Pass --force (-f) or --yes (-y) to overwrite." + ) + if not Confirm.ask("Overwrite?", default=False): + console.print("[yellow]Aborted.[/yellow]") + return + + console.print("[cyan]Detecting SLURM cluster configuration...[/cyan]") + cache = SlurmCommandCache( + cache_file=cache_file or get_default_cache_file(), + use_cache=use_cache, + ) + cluster_data = detect_cluster_info(cache) + config = SlurmConfig.from_cluster_and_file(cluster_data, config_path) + + if not config.enabled_partitions: + raise click.ClickException( + f"No enabled SLURM partitions found for cluster '{config.cluster_name}'. " + "Make sure SLURM commands (sinfo, scontrol) are available or provide --cache-file / --use-cache." + ) + + table = Table( + title=f"SLURM Partitions ({config.cluster_name})", border_style="blue" + ) + table.add_column("Partition", style="cyan") + table.add_column("Priority", style="green") + table.add_column("CPUs", style="green") + table.add_column("Memory (GB)", style="green") + table.add_column("GPUs", style="green") + + for p in sorted( + config.enabled_partitions.values(), key=lambda x: (-x.priority.value, x.name) + ): + gpu_str = ( + f"{p.gpus_per_node}x {p.effective_gpu_type or 'GPU'}" + if p.gpus_per_node > 0 + else "-" + ) + table.add_row( + p.name, + str(p.priority.value), + str(p.effective_cores), + str(p.effective_memory_gb), + gpu_str, + ) + + console.print() + console.print(table) + console.print() + + if not non_interactive and not Confirm.ask("Generate configuration?", default=True): + console.print("[yellow]Aborted.[/yellow]") + return + + if not config_path.exists(): + config.save_to_yaml(config_path) + console.print(f"[green]✓[/green] Saved cluster config to [bold]{config_path}[/bold]") + + content = _generate_slurm_launchers_py( + config, config_path, default_gpu_memory_gb=default_gpu_memory + ) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text(content) + console.print(f"[green]✓[/green] Generated [bold]{output}[/bold]") + diff --git a/src/experimaestro/launchers/slurm/templates/launchers.py.j2 b/src/experimaestro/launchers/slurm/templates/launchers.py.j2 index 75a53b49..18aa268a 100644 --- a/src/experimaestro/launchers/slurm/templates/launchers.py.j2 +++ b/src/experimaestro/launchers/slurm/templates/launchers.py.j2 @@ -57,6 +57,7 @@ def _get_partition_spec(partition: dict) -> HostSpecification: ) return HostSpecification( + max_duration=partition["time_limit_seconds"] or 0, cpu=CPUSpecification( memory=partition["memory_bytes"], cores=partition["cpus"], @@ -116,7 +117,7 @@ def find_launcher(requirements: HostRequirement, tags: Set[str] = set()): # Get duration for time limit duration_seconds = None if hasattr(req, "duration") and req.duration: - duration_seconds = req.duration + duration_seconds = int(req.duration) # Build options options = SlurmOptions( diff --git a/src/experimaestro/launchers/slurm/tui.py b/src/experimaestro/launchers/slurm/tui.py index 9294457d..a5348637 100644 --- a/src/experimaestro/launchers/slurm/tui.py +++ b/src/experimaestro/launchers/slurm/tui.py @@ -2238,7 +2238,20 @@ def _auto_save(self) -> None: def action_generate(self) -> None: """Generate launchers.py.""" - self.set_status("Generate not yet implemented") + if not self.config: + self.set_status("No configuration to generate from") + return + try: + from .cli import _generate_slurm_launchers_py + + self.config.save_to_yaml(self.config_path) + output_path = self.config_path.parent / "launchers.py" + content = _generate_slurm_launchers_py(self.config, self.config_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text(content) + self.set_status(f"Generated {output_path}") + except Exception as e: + self.set_status(f"Generate error: {e}") def action_save(self) -> None: """Save configuration to YAML.""" diff --git a/src/experimaestro/tests/test_direct_generate.py b/src/experimaestro/tests/test_direct_generate.py index c5233412..019d2c5c 100644 --- a/src/experimaestro/tests/test_direct_generate.py +++ b/src/experimaestro/tests/test_direct_generate.py @@ -419,3 +419,52 @@ def test_no_gpu_section_when_zero_tokens(self): yaml = _generate_tokens_yaml(config) assert "gpu:" not in yaml + + +class TestDirectGenerateCLI: + """Test non-interactive CLI execution of `experimaestro launchers direct generate`.""" + + def test_direct_generate_yes(self, tmp_path): + """`direct generate -y` writes launchers.py and tokens.yaml without prompts.""" + from click.testing import CliRunner + from experimaestro.launchers.direct import direct_cli + + out_file = tmp_path / "launchers.py" + runner = CliRunner() + result = runner.invoke( + direct_cli, + ["generate", "-o", str(out_file), "-y"], + ) + + assert result.exit_code == 0, result.output + assert out_file.is_file() + assert (tmp_path / "tokens.yaml").is_file() + assert "def find_launcher(" in out_file.read_text() + + def test_direct_generate_custom_flags(self, tmp_path): + """`direct generate` respects explicit token CLI flags in non-interactive mode.""" + from click.testing import CliRunner + from experimaestro.launchers.direct import direct_cli + + out_file = tmp_path / "launchers.py" + runner = CliRunner() + result = runner.invoke( + direct_cli, + [ + "generate", + "-o", + str(out_file), + "-y", + "--reserved-memory", + "1.0", + "--memory-token-unit", + "512", + "--no-gpu-tokens", + ], + ) + + assert result.exit_code == 0, result.output + assert out_file.is_file() + content = out_file.read_text() + assert "512 * (1024 ** 2)" in content + diff --git a/src/experimaestro/tests/test_findlauncher.py b/src/experimaestro/tests/test_findlauncher.py index cfb04c5a..80f58eba 100644 --- a/src/experimaestro/tests/test_findlauncher.py +++ b/src/experimaestro/tests/test_findlauncher.py @@ -466,3 +466,50 @@ def test_unified_memory_not_applied_to_cuda(self): # Request 30G CPU + 20G GPU - should pass (separate memory pools) req = cuda_gpu(mem="20G") & cpu(mem="30G") assert req.match(cuda_host) is not None + + +class TestLauncherFallbackAndErrors: + """Test missing launchers.py warning fallback and LauncherNotFoundError message""" + + def test_missing_launchers_py_warns_and_falls_back(self, tmp_path, caplog): + """Missing launchers.py should log a warning with setup tips and return DirectLauncher""" + import logging + from experimaestro.launchers.direct import DirectLauncher + + registry = LauncherRegistry(tmp_path) + with caplog.at_level(logging.WARNING): + launcher = registry.find("duration=1h & gpu(mem=4G)*1") + + assert isinstance(launcher, DirectLauncher) + assert str(tmp_path / "launchers.py") in caplog.text + assert "experimaestro launchers direct generate" in caplog.text + assert ( + "https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html" + in caplog.text + ) + + def test_unmatched_launcher_raises_helpful_error(self, tmp_path, monkeypatch): + """When launchers.py returns None, find_launcher raises LauncherNotFoundError with tips""" + from experimaestro.launcherfinder import LauncherNotFoundError, find_launcher + + launchers_py = tmp_path / "launchers.py" + launchers_py.write_text( + "def find_launcher(requirements, tags=set()):\n return None\n" + ) + + monkeypatch.setattr(LauncherRegistry, "CURRENT_CONFIG_DIR", tmp_path) + monkeypatch.setattr(LauncherRegistry, "INSTANCES", {}) + + with pytest.raises(LauncherNotFoundError) as exc_info: + find_launcher("duration=1h & cuda(mem=80G)*8") + + msg = str(exc_info.value) + assert str(launchers_py) in msg + assert "experimaestro launchers direct generate" in msg + assert "experimaestro launchers slurm generate" in msg + assert "experimaestro find-launchers" in msg + assert ( + "https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html" + in msg + ) + diff --git a/src/experimaestro/tests/test_slurm_generate.py b/src/experimaestro/tests/test_slurm_generate.py index a146c1d9..a3f58a2f 100644 --- a/src/experimaestro/tests/test_slurm_generate.py +++ b/src/experimaestro/tests/test_slurm_generate.py @@ -59,3 +59,89 @@ def test_empty(self): count, gpu_type = parse_gres("") assert count == 0 assert gpu_type is None + + +class TestSlurmGenerateCLI: + """Test non-interactive SLURM launcher generation from cached cluster data.""" + + def test_slurm_generate_yes_from_cache(self, tmp_path): + import json + from click.testing import CliRunner + from experimaestro.launchers.slurm.cli import cli as slurm_cli + from experimaestro.launchers.slurm import SlurmLauncher + from experimaestro.launcherfinder.parser import parse + + cache_data = { + "sinfo": { + "meta": {"slurm": {"cluster": "test_cluster"}}, + "nodes": [], + "sinfo": [], + }, + "partitions": { + "partitions": [ + { + "name": "gpu_a100", + "partition": {"state": ["UP"]}, + "cpus": {"total": 64}, + "nodes": {"configured": "node[01-02]", "total": 2}, + "tres": { + "configured": "cpu=64,mem=256G,node=1,gres/gpu:a100=4" + }, + "maximums": {"time": {"number": 120, "infinite": False}}, + }, + { + "name": "cpu_std", + "partition": {"state": ["UP"]}, + "cpus": {"total": 32}, + "nodes": {"configured": "node[03-04]", "total": 2}, + "tres": {"configured": "cpu=32,mem=128G,node=1"}, + "maximums": {"time": {"number": 60, "infinite": False}}, + }, + ] + }, + "qos": {"qos": []}, + "associations": { + "associations": [ + {"account": "myproj", "is_default": True, "qos": ["normal"]} + ] + }, + "config": {}, + } + cache_file = tmp_path / "slurm_cache.json" + cache_file.write_text(json.dumps(cache_data)) + config_file = tmp_path / "slurm.yaml" + out_file = tmp_path / "launchers.py" + + runner = CliRunner() + result = runner.invoke( + slurm_cli, + [ + "generate", + "--cache-file", + str(cache_file), + "--use-cache", + "-c", + str(config_file), + "-o", + str(out_file), + "-y", + ], + ) + + assert result.exit_code == 0, result.output + assert config_file.is_file() + assert out_file.is_file() + + # Execute the generated launchers.py and verify find_launcher works + ns = {} + exec(out_file.read_text(), ns) + find_launcher_fn = ns["find_launcher"] + + req = parse("duration=1h & gpu(mem=20G)*2 & cpu(cores=8)").requirements[0] + launcher = find_launcher_fn(req) + assert isinstance(launcher, SlurmLauncher) + assert launcher.options.partition == "gpu_a100" + assert launcher.options.gpus == 2 + assert launcher.options.account == "myproj" + assert launcher.options.time == "1:0:0" + From 8622fe7c6e57dab5a496608dd0e4ec943e6735fd Mon Sep 17 00:00:00 2001 From: Victor Morand Date: Tue, 20 Oct 2026 17:22:53 +0200 Subject: [PATCH 2/3] chore: ruff and linting --- src/experimaestro/launcherfinder/registry.py | 2 +- src/experimaestro/launchers/slurm/cli.py | 12 ++++++++---- src/experimaestro/tests/test_direct_generate.py | 1 - src/experimaestro/tests/test_findlauncher.py | 1 - src/experimaestro/tests/test_slurm_generate.py | 1 - 5 files changed, 9 insertions(+), 8 deletions(-) diff --git a/src/experimaestro/launcherfinder/registry.py b/src/experimaestro/launcherfinder/registry.py index a31c2b1a..697ab1cd 100644 --- a/src/experimaestro/launcherfinder/registry.py +++ b/src/experimaestro/launcherfinder/registry.py @@ -204,7 +204,7 @@ def find_launcher( f" 2. Regenerate or edit {registry.launchers_py}:\n" " - Local machine / laptop: experimaestro launchers direct generate -y\n" " - SLURM cluster: experimaestro launchers slurm generate -y\n" - f" 3. Test your specification with: experimaestro find-launchers \"{sample_spec}\"\n" + f' 3. Test your specification with: experimaestro find-launchers "{sample_spec}"\n' " 4. Documentation: https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html" ) return launcher diff --git a/src/experimaestro/launchers/slurm/cli.py b/src/experimaestro/launchers/slurm/cli.py index bfb8f044..3074cc13 100644 --- a/src/experimaestro/launchers/slurm/cli.py +++ b/src/experimaestro/launchers/slurm/cli.py @@ -11,6 +11,7 @@ # Import data models from .models import ( + SlurmConfig, ClusterPartition, ClusterQoS, ClusterAccount, @@ -968,7 +969,7 @@ def _infer_gpu_memory_gb( def _generate_slurm_launchers_py( - config: "SlurmConfig", + config: SlurmConfig, config_path: Path, default_gpu_memory_gb: int = 16, ) -> str: @@ -1028,7 +1029,9 @@ def _generate_slurm_launchers_py( "memory_bytes": mem_bytes, "gpus": p.gpus_per_node, "gpu_type": p.effective_gpu_type, - "gpu_memory_bytes": gpu_mem_gb * (1024**3) if p.gpus_per_node > 0 else 0, + "gpu_memory_bytes": gpu_mem_gb * (1024**3) + if p.gpus_per_node > 0 + else 0, "time_limit_seconds": p.time_limit_seconds, "qos": qos, } @@ -1187,7 +1190,9 @@ def generate( if not config_path.exists(): config.save_to_yaml(config_path) - console.print(f"[green]✓[/green] Saved cluster config to [bold]{config_path}[/bold]") + console.print( + f"[green]✓[/green] Saved cluster config to [bold]{config_path}[/bold]" + ) content = _generate_slurm_launchers_py( config, config_path, default_gpu_memory_gb=default_gpu_memory @@ -1195,4 +1200,3 @@ def generate( output.parent.mkdir(parents=True, exist_ok=True) output.write_text(content) console.print(f"[green]✓[/green] Generated [bold]{output}[/bold]") - diff --git a/src/experimaestro/tests/test_direct_generate.py b/src/experimaestro/tests/test_direct_generate.py index 019d2c5c..0a6cebe6 100644 --- a/src/experimaestro/tests/test_direct_generate.py +++ b/src/experimaestro/tests/test_direct_generate.py @@ -467,4 +467,3 @@ def test_direct_generate_custom_flags(self, tmp_path): assert out_file.is_file() content = out_file.read_text() assert "512 * (1024 ** 2)" in content - diff --git a/src/experimaestro/tests/test_findlauncher.py b/src/experimaestro/tests/test_findlauncher.py index 80f58eba..66133dd9 100644 --- a/src/experimaestro/tests/test_findlauncher.py +++ b/src/experimaestro/tests/test_findlauncher.py @@ -512,4 +512,3 @@ def test_unmatched_launcher_raises_helpful_error(self, tmp_path, monkeypatch): "https://experimaestro-python.readthedocs.io/en/latest/launchers/index.html" in msg ) - diff --git a/src/experimaestro/tests/test_slurm_generate.py b/src/experimaestro/tests/test_slurm_generate.py index a3f58a2f..12b5de74 100644 --- a/src/experimaestro/tests/test_slurm_generate.py +++ b/src/experimaestro/tests/test_slurm_generate.py @@ -144,4 +144,3 @@ def test_slurm_generate_yes_from_cache(self, tmp_path): assert launcher.options.gpus == 2 assert launcher.options.account == "myproj" assert launcher.options.time == "1:0:0" - From 0a384b3458855c3dd691f1e3fa83b5652d757a88 Mon Sep 17 00:00:00 2001 From: Victor Morand Date: Wed, 21 Oct 2026 10:11:43 +0200 Subject: [PATCH 3/3] fix: ruff --- src/experimaestro/launchers/direct.py | 488 ++++++++++++++++---------- 1 file changed, 295 insertions(+), 193 deletions(-) diff --git a/src/experimaestro/launchers/direct.py b/src/experimaestro/launchers/direct.py index 7dfd70d3..963b6a1e 100644 --- a/src/experimaestro/launchers/direct.py +++ b/src/experimaestro/launchers/direct.py @@ -2,6 +2,7 @@ from pathlib import Path import platform import sys +from typing import TYPE_CHECKING import click import jinja2 @@ -9,6 +10,9 @@ from experimaestro.scriptbuilder import PythonScriptBuilder from . import Launcher +if TYPE_CHECKING: + from rich.console import Console + class DirectLauncher(Launcher): """Launcher that runs tasks directly as local processes. @@ -175,117 +179,16 @@ def direct_cli(): pass -@direct_cli.command("generate") -@click.option( - "--output", - "-o", - type=click.Path(path_type=Path), - default=None, - help="Output path for launchers.py (default: ~/.config/experimaestro/launchers.py)", -) -@click.option( - "--force", - "-f", - is_flag=True, - help="Overwrite existing file without prompting", -) -@click.option( - "--yes", - "-y", - is_flag=True, - help="Non-interactive mode: accept defaults without prompting", -) -@click.option( - "--memory-tokens/--no-memory-tokens", - default=None, - help="Enable or disable CPU memory tokens (default: enabled)", -) -@click.option( - "--reserved-memory", - type=float, - default=None, - help="Reserved memory for OS in GiB (default: 2.0)", -) -@click.option( - "--memory-token-unit", - type=int, - default=None, - help=f"Memory token unit size in MiB (default: {DEFAULT_TOKEN_UNIT_MIB})", -) -@click.option( - "--gpu-tokens/--no-gpu-tokens", - default=None, - help="Enable or disable GPU tokens when GPUs are detected (default: enabled)", -) -@click.option( - "--gpu-mode", - type=click.Choice(["exclusive", "memory"]), - default=None, - help="GPU token mode: 'exclusive' (1 token/GPU) or 'memory' (default: exclusive for CUDA, memory for MPS)", -) -@click.option( - "--gpu-token-unit", - type=int, - default=None, - help="GPU memory token unit size in MiB when using memory mode (default: 1024)", -) -def generate( - output: Path | None, +def _check_output_file( + output: Path, force: bool, yes: bool, - memory_tokens: bool | None, - reserved_memory: float | None, - memory_token_unit: int | None, - gpu_tokens: bool | None, - gpu_mode: str | None, - gpu_token_unit: int | None, -): - """Generate a launchers.py file for direct (local) execution. - - This command detects system resources and configures memory-based tokens - for CPU and GPU resource management. Use `-y / --yes` for non-interactive - generation with sensible defaults. - - Token system: - - Memory tokens: 1 token = configurable unit (default 256 MiB) - - Tasks request tokens based on their memory requirements - - Example: 1 GiB requirement = 4 tokens (with 256 MiB unit) + non_interactive: bool, + console: "Console", +) -> bool: + """Check if output file exists and handle overwrite confirmation.""" + from rich.prompt import Confirm - Apple Silicon (MPS): - - GPU uses unified memory (shared with CPU) - - GPU memory requests reduce available CPU memory tokens - - Examples: - experimaestro launchers direct generate -y - experimaestro launchers direct generate - experimaestro launchers direct generate -o ./launchers.py -y - """ - from rich.console import Console - from rich.panel import Panel - from rich.prompt import Confirm, FloatPrompt, IntPrompt, Prompt - from rich.table import Table - - console = Console() - non_interactive = yes or not sys.stdin.isatty() - - # Default output path - if output is None: - output = Path("~/.config/experimaestro/launchers.py").expanduser() - else: - output = output.expanduser() - - # Header - console.print() - console.print( - Panel.fit( - "[bold cyan]Direct Launcher Configuration[/bold cyan]\n\n" - "This wizard configures resource tokens for local task execution.\n" - "Tokens ensure tasks don't exceed available memory.", - border_style="cyan", - ) - ) - - # Check if file exists if output.exists() and not (force or yes): console.print(f"\n[yellow]File already exists:[/yellow] {output}") if non_interactive: @@ -294,13 +197,14 @@ def generate( ) if not Confirm.ask("Overwrite?", default=False): console.print("[yellow]Aborted.[/yellow]") - return + return False + return True - # Detect system resources - console.print("\n[cyan]Detecting system resources...[/cyan]") - sys_info = detect_system_info() - # Display system info +def _display_system_info(console: "Console", sys_info: SystemInfo) -> None: + """Display detected system resources table.""" + from rich.table import Table + table = Table(title="System Resources", border_style="blue") table.add_column("Resource", style="cyan") table.add_column("Value", style="green") @@ -324,11 +228,20 @@ def generate( console.print(table) console.print() - # Initialize config - config = LauncherConfig() - config.is_mps = sys_info.has_mps - # === Memory tokens configuration === +def _configure_memory_tokens( + console: "Console", + config: LauncherConfig, + sys_info: SystemInfo, + non_interactive: bool, + memory_tokens: bool | None, + reserved_memory: float | None, + memory_token_unit: int | None, +) -> None: + """Prompt and configure memory tokens on LauncherConfig.""" + from rich.panel import Panel + from rich.prompt import Confirm, FloatPrompt, IntPrompt + if not non_interactive: console.print( Panel( @@ -383,89 +296,107 @@ def generate( f" [green]→ {config.total_memory_tokens} memory tokens available[/green]" ) - # === GPU tokens configuration === - if sys_info.gpus: - if not non_interactive: - console.print() - console.print( - Panel( - "[bold]GPU Tokens[/bold]\n\n" - "GPU tokens control concurrent GPU task execution.\n\n" - "[cyan]Exclusive mode:[/cyan] 1 token = 1 GPU (simple, recommended)\n" - "[cyan]Memory mode:[/cyan] Tokens based on GPU memory units\n" - + ( - "\n[yellow]Note: Apple MPS uses unified memory - GPU memory\n" - "requests will reduce available CPU memory tokens.[/yellow]" - if sys_info.has_mps - else "" - ), - border_style="blue", - ) + +def _configure_gpu_tokens( + console: "Console", + config: LauncherConfig, + sys_info: SystemInfo, + non_interactive: bool, + gpu_tokens: bool | None, + gpu_mode: str | None, + gpu_token_unit: int | None, +) -> None: + """Prompt and configure GPU tokens on LauncherConfig.""" + from rich.panel import Panel + from rich.prompt import Confirm, IntPrompt, Prompt + + if not sys_info.gpus: + return + + if not non_interactive: + console.print() + console.print( + Panel( + "[bold]GPU Tokens[/bold]\n\n" + "GPU tokens control concurrent GPU task execution.\n\n" + "[cyan]Exclusive mode:[/cyan] 1 token = 1 GPU (simple, recommended)\n" + "[cyan]Memory mode:[/cyan] Tokens based on GPU memory units\n" + + ( + "\n[yellow]Note: Apple MPS uses unified memory - GPU memory\n" + "requests will reduce available CPU memory tokens.[/yellow]" + if sys_info.has_mps + else "" + ), + border_style="blue", ) + ) + + if gpu_tokens is not None: + config.use_gpu_tokens = gpu_tokens + elif non_interactive: + config.use_gpu_tokens = True + else: + config.use_gpu_tokens = Confirm.ask("\nEnable GPU tokens?", default=True) + + if not config.use_gpu_tokens: + return - if gpu_tokens is not None: - config.use_gpu_tokens = gpu_tokens + config.gpu_count = len(sys_info.gpus) + + # GPU mode selection + if sys_info.has_mps: + console.print( + " [yellow]MPS detected: using memory-based tokens " + "(shared with CPU memory)[/yellow]" + ) + config.gpu_mode = "memory" + config.mps_shared_memory = True + elif gpu_mode is not None: + config.gpu_mode = gpu_mode + elif non_interactive: + config.gpu_mode = "exclusive" + else: + config.gpu_mode = Prompt.ask( + " GPU token mode", + choices=["exclusive", "memory"], + default="exclusive", + ) + + if config.gpu_mode == "exclusive": + config.total_gpu_tokens = config.gpu_count + console.print( + f" [green]→ {config.total_gpu_tokens} GPU tokens (1 per GPU)[/green]" + ) + else: + if gpu_token_unit is not None: + config.gpu_token_unit_mib = gpu_token_unit elif non_interactive: - config.use_gpu_tokens = True + config.gpu_token_unit_mib = 1024 + else: + config.gpu_token_unit_mib = IntPrompt.ask( + " GPU token unit size (MiB)", + default=1024, + ) + # Calculate GPU tokens from first GPU (or total for MPS) + if sys_info.has_mps: + # For MPS, GPU memory is same as system memory + # Don't double-count - just note that GPU uses memory tokens + console.print(" [yellow]→ MPS: GPU tasks will use memory tokens[/yellow]") + config.total_gpu_tokens = 0 # Use memory tokens instead else: - config.use_gpu_tokens = Confirm.ask("\nEnable GPU tokens?", default=True) + total_gpu_mem = sum(g.memory_bytes for g in sys_info.gpus) + config.total_gpu_tokens = max( + 1, total_gpu_mem // (config.gpu_token_unit_mib * (1024**2)) + ) + console.print( + f" [green]→ {config.total_gpu_tokens} GPU memory tokens[/green]" + ) - if config.use_gpu_tokens: - config.gpu_count = len(sys_info.gpus) - # GPU mode selection - if sys_info.has_mps: - console.print( - " [yellow]MPS detected: using memory-based tokens " - "(shared with CPU memory)[/yellow]" - ) - config.gpu_mode = "memory" - config.mps_shared_memory = True - elif gpu_mode is not None: - config.gpu_mode = gpu_mode - elif non_interactive: - config.gpu_mode = "exclusive" - else: - config.gpu_mode = Prompt.ask( - " GPU token mode", - choices=["exclusive", "memory"], - default="exclusive", - ) - - if config.gpu_mode == "exclusive": - config.total_gpu_tokens = config.gpu_count - console.print( - f" [green]→ {config.total_gpu_tokens} GPU tokens " - f"(1 per GPU)[/green]" - ) - else: - if gpu_token_unit is not None: - config.gpu_token_unit_mib = gpu_token_unit - elif non_interactive: - config.gpu_token_unit_mib = 1024 - else: - config.gpu_token_unit_mib = IntPrompt.ask( - " GPU token unit size (MiB)", - default=1024, - ) - # Calculate GPU tokens from first GPU (or total for MPS) - if sys_info.has_mps: - # For MPS, GPU memory is same as system memory - # Don't double-count - just note that GPU uses memory tokens - console.print( - " [yellow]→ MPS: GPU tasks will use memory tokens[/yellow]" - ) - config.total_gpu_tokens = 0 # Use memory tokens instead - else: - total_gpu_mem = sum(g.memory_bytes for g in sys_info.gpus) - config.total_gpu_tokens = max( - 1, total_gpu_mem // (config.gpu_token_unit_mib * (1024**2)) - ) - console.print( - f" [green]→ {config.total_gpu_tokens} GPU memory tokens[/green]" - ) +def _display_summary(console: "Console", output: Path, config: LauncherConfig) -> None: + """Display configuration summary table.""" + from rich.table import Table - # === Summary === console.print() summary = Table(title="Configuration Summary", border_style="green") summary.add_column("Setting", style="cyan") @@ -497,6 +428,19 @@ def generate( console.print(summary) console.print() + +def _write_outputs( + output: Path, + config: LauncherConfig, + sys_info: SystemInfo, + force: bool, + yes: bool, + non_interactive: bool, + console: "Console", +) -> None: + """Prompt confirmation (if interactive) and write output files.""" + from rich.prompt import Confirm + # Confirm generation if not non_interactive and not Confirm.ask("Generate configuration?", default=True): console.print("[yellow]Aborted.[/yellow]") @@ -527,6 +471,164 @@ def generate( ) +@direct_cli.command("generate") +@click.option( + "--output", + "-o", + type=click.Path(path_type=Path), + default=None, + help="Output path for launchers.py (default: ~/.config/experimaestro/launchers.py)", +) +@click.option( + "--force", + "-f", + is_flag=True, + help="Overwrite existing file without prompting", +) +@click.option( + "--yes", + "-y", + is_flag=True, + help="Non-interactive mode: accept defaults without prompting", +) +@click.option( + "--memory-tokens/--no-memory-tokens", + default=None, + help="Enable or disable CPU memory tokens (default: enabled)", +) +@click.option( + "--reserved-memory", + type=float, + default=None, + help="Reserved memory for OS in GiB (default: 2.0)", +) +@click.option( + "--memory-token-unit", + type=int, + default=None, + help=f"Memory token unit size in MiB (default: {DEFAULT_TOKEN_UNIT_MIB})", +) +@click.option( + "--gpu-tokens/--no-gpu-tokens", + default=None, + help="Enable or disable GPU tokens when GPUs are detected (default: enabled)", +) +@click.option( + "--gpu-mode", + type=click.Choice(["exclusive", "memory"]), + default=None, + help="GPU token mode: 'exclusive' (1 token/GPU) or 'memory' (default: exclusive for CUDA, memory for MPS)", +) +@click.option( + "--gpu-token-unit", + type=int, + default=None, + help="GPU memory token unit size in MiB when using memory mode (default: 1024)", +) +def generate( + output: Path | None, + force: bool, + yes: bool, + memory_tokens: bool | None, + reserved_memory: float | None, + memory_token_unit: int | None, + gpu_tokens: bool | None, + gpu_mode: str | None, + gpu_token_unit: int | None, +): + """Generate a launchers.py file for direct (local) execution. + + This command detects system resources and configures memory-based tokens + for CPU and GPU resource management. Use `-y / --yes` for non-interactive + generation with sensible defaults. + + Token system: + - Memory tokens: 1 token = configurable unit (default 256 MiB) + - Tasks request tokens based on their memory requirements + - Example: 1 GiB requirement = 4 tokens (with 256 MiB unit) + + Apple Silicon (MPS): + - GPU uses unified memory (shared with CPU) + - GPU memory requests reduce available CPU memory tokens + + Examples: + experimaestro launchers direct generate -y + experimaestro launchers direct generate + experimaestro launchers direct generate -o ./launchers.py -y + """ + from rich.console import Console + from rich.panel import Panel + + console = Console() + non_interactive = yes or not sys.stdin.isatty() + + # Default output path + if output is None: + output = Path("~/.config/experimaestro/launchers.py").expanduser() + else: + output = output.expanduser() + + # Header + console.print() + console.print( + Panel.fit( + "[bold cyan]Direct Launcher Configuration[/bold cyan]\n\n" + "This wizard configures resource tokens for local task execution.\n" + "Tokens ensure tasks don't exceed available memory.", + border_style="cyan", + ) + ) + + # Check if file exists + if not _check_output_file(output, force, yes, non_interactive, console): + return + + # Detect system resources + console.print("\n[cyan]Detecting system resources...[/cyan]") + sys_info = detect_system_info() + _display_system_info(console, sys_info) + + # Initialize config + config = LauncherConfig() + config.is_mps = sys_info.has_mps + + # Configure memory tokens + _configure_memory_tokens( + console, + config, + sys_info, + non_interactive, + memory_tokens, + reserved_memory, + memory_token_unit, + ) + + # Configure GPU tokens + _configure_gpu_tokens( + console, + config, + sys_info, + non_interactive, + gpu_tokens, + gpu_mode, + gpu_token_unit, + ) + + # Display configuration summary + _display_summary(console, output, config) + + # Write generated files + _write_outputs( + output, + config, + sys_info, + force, + yes, + non_interactive, + console, + ) + + # Template environment for code generation _template_env = jinja2.Environment( loader=jinja2.PackageLoader("experimaestro.launchers", "templates"),