diff --git a/truss/cli/train/exec/__init__.py b/truss/cli/train/exec/__init__.py new file mode 100644 index 000000000..2de238d58 --- /dev/null +++ b/truss/cli/train/exec/__init__.py @@ -0,0 +1,40 @@ +"""`truss train exec`: run a local directory as a Baseten training job.""" + +from .builder import ( + DEFAULT_CPU_COUNT, + DEFAULT_EXEC_PROJECT_NAME, + DEFAULT_MEMORY, + PYTHON_BASE_IMAGE, + SUPPORTED_EXEC_ACCELERATORS, + build_exec_project, + build_start_commands, + default_base_image, + resolve_workspace_root, + validate_workspace_root, +) +from .project import Project, get_project_type +from .secrets import ( + SECRETS_SETTINGS_URL, + parse_environment_variables, + validate_secret_references, +) +from .uv import UvProject + +__all__ = [ + "DEFAULT_CPU_COUNT", + "DEFAULT_EXEC_PROJECT_NAME", + "DEFAULT_MEMORY", + "PYTHON_BASE_IMAGE", + "SECRETS_SETTINGS_URL", + "SUPPORTED_EXEC_ACCELERATORS", + "Project", + "UvProject", + "build_exec_project", + "build_start_commands", + "default_base_image", + "get_project_type", + "parse_environment_variables", + "resolve_workspace_root", + "validate_secret_references", + "validate_workspace_root", +] diff --git a/truss/cli/train/exec/builder.py b/truss/cli/train/exec/builder.py new file mode 100644 index 000000000..5b84b8855 --- /dev/null +++ b/truss/cli/train/exec/builder.py @@ -0,0 +1,174 @@ +"""Assembles a `TrainingProject` from `truss train exec` CLI input.""" + +import shlex +from pathlib import Path +from typing import List, Mapping, Optional, Sequence, Union + +import rich_click as click + +from truss.base import truss_config +from truss.cli.train import workstation +from truss_train.definitions import ( + Compute, + Image, + InteractiveSession, + InteractiveSessionProvider, + InteractiveSessionTrigger, + Runtime, + SecretReference, + TrainingJob, + TrainingProject, + Workspace, +) + +from .project import Project + +# A CPU-only job doesn't need a CUDA image. +PYTHON_BASE_IMAGE = "python:3.12-slim" + +# An empty project name fails server-side validation, which the filesystem root +# would otherwise produce. +DEFAULT_EXEC_PROJECT_NAME = "truss-train-exec" + +SUPPORTED_EXEC_ACCELERATORS = workstation.SUPPORTED_WORKSTATION_ACCELERATORS + +# Read from the model so the CLI defaults cannot drift from it. +DEFAULT_CPU_COUNT: int = Compute.model_fields["cpu_count"].default +DEFAULT_MEMORY: str = Compute.model_fields["memory"].default + + +def default_base_image(accelerator: Optional[str], project: Optional[Project]) -> str: + """The base image to use when the user didn't pass --image. + + An accelerator wins over the project's preference: a GPU job needs the CUDA image + regardless of how the project installs its dependencies, and the project's setup + steps cover the difference. + """ + if accelerator is not None: + return workstation.default_base_image(accelerator) + if project is not None: + return project.base_image() + return PYTHON_BASE_IMAGE + + +def resolve_workspace_root(source_dir: Path, workspace_root: Optional[str]) -> Path: + """The directory that actually gets archived and becomes the job's root.""" + if not workspace_root: + return source_dir + root = Path(workspace_root) + if not root.is_absolute(): + root = source_dir / root + return root.resolve() + + +def validate_workspace_root(source_dir: Path, workspace_root: Optional[str]) -> Path: + """Resolve and check `--workspace-root`, returning the effective job root. + + `truss_train` runs the same containment check inside `push`, but only after the + training project has been created, so a bad value there leaves a stray empty + project behind. Checking here keeps that from happening. + """ + root = resolve_workspace_root(source_dir, workspace_root) + if not workspace_root: + return root + + if not root.is_dir(): + raise click.UsageError( + f"--workspace-root '{workspace_root}' resolves to {root}, " + "which is not a directory." + ) + try: + source_dir.resolve().relative_to(root) + except ValueError: + raise click.UsageError( + f"--workspace-root '{workspace_root}' resolves to {root}, which does not " + f"contain the current directory ({source_dir}); it must be a parent of it." + ) + return root + + +def build_start_commands( + start_command: Sequence[str], setup_steps: Sequence[str] = () +) -> List[str]: + """Build `Runtime.start_commands` for `start_command`. + + The user's command always runs last and verbatim. `setup_steps` (from the + detected project type) run first, chained into a single `/bin/sh -c` entry to + match the pattern in `truss/templates/train/config.py`, because whether the + platform runs more than the first entry can't be established from this repo. + """ + command = shlex.join(start_command) + if not setup_steps: + return [command] + return [f"/bin/sh -c {shlex.quote(' && '.join([*setup_steps, command]))}"] + + +def build_exec_project( + *, + start_command: Sequence[str], + project_name: str, + accelerator: Optional[str], + gpu_count: int, + cpu_count: int, + memory: str, + base_image: Optional[str], + project: Optional[Project], + workspace_root: Optional[str], + exclude_dirs: Sequence[str], + external_dirs: Sequence[str], + environment_variables: Mapping[str, Union[str, SecretReference]], +) -> TrainingProject: + """Build the training project for `truss train exec`. + + Every parameter is required and keyword-only, so a caller cannot build a + partially-specified project and the CLI stays the single source of defaults. The + `Optional` types are values, not omissions: `accelerator` None means CPU-only, + `base_image` None means derive it, `project` None means run against a plain image + with no setup steps, and `workspace_root` None means archive the invocation + directory. + """ + accelerator_spec = None + if accelerator is not None: + accelerator_spec = truss_config.AcceleratorSpec( + accelerator=truss_config.Accelerator(accelerator), count=gpu_count + ) + + compute = Compute(cpu_count=cpu_count, memory=memory, accelerator=accelerator_spec) + + resolved_base_image = base_image or default_base_image(accelerator, project) + + # A one-off command needs no persistent storage, hence no cache or + # checkpointing config. + runtime = Runtime( + start_commands=build_start_commands( + start_command=start_command, + setup_steps=project.setup(resolved_base_image) if project else (), + ), + environment_variables=dict(environment_variables), + ) + + # SSH available on demand, rather than a session live from job startup: the + # session timeout applies once the job ends, so it is not a concern for a + # long-running job. No timeout is set here; the model default still applies. + interactive_session = InteractiveSession( + trigger=InteractiveSessionTrigger.ON_DEMAND, + session_provider=InteractiveSessionProvider.SSH, + ) + + workspace_config = None + if workspace_root or exclude_dirs or external_dirs: + workspace_config = Workspace( + workspace_root=workspace_root, + exclude_dirs=list(exclude_dirs), + external_dirs=list(external_dirs), + ) + + job = TrainingJob( + image=Image(base_image=resolved_base_image), + compute=compute, + runtime=runtime, + interactive_session=interactive_session, + workspace=workspace_config, + ) + + return TrainingProject(name=project_name, job=job) diff --git a/truss/cli/train/exec/project.py b/truss/cli/train/exec/project.py new file mode 100644 index 000000000..8d0694a0d --- /dev/null +++ b/truss/cli/train/exec/project.py @@ -0,0 +1,37 @@ +"""Project-type detection for `truss train exec`. + +A project type answers two questions about the directory being pushed: which base +image suits it, and what has to happen in the job before the user's command runs. +`get_project_type` is the single place that maps a directory to one, so supporting +pip or poetry later means adding a branch here plus a sibling module. +""" + +from pathlib import Path +from typing import List, Optional, Protocol + +from . import uv + + +class Project(Protocol): + """A recognised project type in the directory `truss train exec` pushes.""" + + #: Short label used in CLI messages, e.g. "uv". + label: str + + def base_image(self) -> str: + """The base image this project type wants when the user didn't pass --image.""" + ... + + def setup(self, base_image: str) -> List[str]: + """Shell steps to run before the user's command, given the resolved image. + + Empty when the image already provides everything the project needs. + """ + ... + + +def get_project_type(source_dir: Path) -> Optional[Project]: + """The project type detected in `source_dir`, or None if none is recognised.""" + if uv.is_uv_project(source_dir): + return uv.UvProject() + return None diff --git a/truss/cli/train/exec/secrets.py b/truss/cli/train/exec/secrets.py new file mode 100644 index 000000000..f88f5fd40 --- /dev/null +++ b/truss/cli/train/exec/secrets.py @@ -0,0 +1,133 @@ +"""`--env` / `--secret` handling for `truss train exec`.""" + +import logging +from typing import Any, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Union + +import rich_click as click + +from truss.remote.baseten.api import BasetenApi +from truss_train.definitions import SecretReference + +logger = logging.getLogger(__name__) + +# There is no CLI command to create a workspace secret, so the settings page is the +# only actionable next step we can point at. +SECRETS_SETTINGS_URL = "https://app.baseten.co/settings/secrets" + + +def _parse_key_value_flag( + flag: str, expected: str, entry: str, require_value: bool = False +) -> Tuple[str, str]: + # partition, not split: a value may itself contain `=`. + key, separator, value = entry.partition("=") + # An empty --env value is legitimate; an empty secret *name* is not. + if not separator or not key or (require_value and not value): + raise click.UsageError(f"Invalid {flag} value '{entry}'. Expected {expected}.") + return key, value + + +def parse_environment_variables( + env: Sequence[str] = (), secrets: Sequence[str] = () +) -> Dict[str, Union[str, SecretReference]]: + """Turn `--env KEY=VALUE` and `--secret KEY=SECRET_NAME` flags into the + `Runtime.environment_variables` mapping.""" + entries: List[Tuple[str, Union[str, SecretReference]]] = [] + for entry in env: + key, value = _parse_key_value_flag("--env", "KEY=VALUE", entry) + entries.append((key, value)) + for entry in secrets: + key, secret_name = _parse_key_value_flag( + "--secret", "KEY=SECRET_NAME", entry, require_value=True + ) + entries.append((key, SecretReference(name=secret_name))) + + environment_variables: Dict[str, Union[str, SecretReference]] = {} + for key, resolved in entries: + if key in environment_variables: + raise click.UsageError( + f"Environment variable '{key}' is set more than once by " + "--env / --secret." + ) + environment_variables[key] = resolved + return environment_variables + + +def _known_secret_names(response: Any) -> Optional[Set[str]]: + """Secret names from a `GET v1/secrets` payload, or None if it is unrecognized. + + `get_all_secrets` had no callers before this, and the response shape is not + pinned down by any test or doc in this repo, so accept the plausible shapes and + give up rather than guess -- returning None means "don't check", which is very + different from returning an empty set. + """ + if isinstance(response, dict): + entries = response.get("secrets") + elif isinstance(response, list): + entries = response + else: + return None + if not isinstance(entries, list): + return None + + names: Set[str] = set() + for entry in entries: + if isinstance(entry, str): + names.add(entry) + elif isinstance(entry, dict) and isinstance(entry.get("name"), str): + names.add(entry["name"]) + else: + # An unfamiliar entry shape would mean guessing, and a wrong guess now + # fails the command rather than just warning. + return None + return names + + +def validate_secret_references( + api: BasetenApi, environment_variables: Mapping[str, Union[str, SecretReference]] +) -> None: + """Fail before pushing if a `--secret` names a secret the workspace doesn't have. + + Two cases, deliberately treated differently: + + * The listing came back and the name isn't in it -> hard error. The job would + fail to start, so failing here is faster and clearer. + * The listing call failed, or returned something we can't parse -> continue. That + is an API problem, not evidence the secret is missing, and a convenience check + must not break the command over an API blip or a permissions quirk. + """ + referenced = sorted( + { + value.name + for value in environment_variables.values() + if isinstance(value, SecretReference) + } + ) + if not referenced: + # No --secret flags, so don't spend a round trip on the common path. + return + + try: + response = api.get_all_secrets() + except Exception: + logger.debug("Could not list workspace secrets; skipping check.", exc_info=True) + return + + # Outside the try: a bug in the parser should surface, not be mistaken for an + # unreachable API. + known = _known_secret_names(response) + if known is None: + logger.debug("Unrecognized v1/secrets payload; skipping check.") + return + + missing = [name for name in referenced if name not in known] + if not missing: + return + + plural = len(missing) > 1 + raise click.UsageError( + f"{'Secrets' if plural else 'Secret'} {', '.join(missing)} " + f"{'were' if plural else 'was'} not found in this workspace's secrets. " + f"Create {'them' if plural else 'it'} at {SECRETS_SETTINGS_URL}. " + "(The listing this checks against is not team-scoped, so if the secret does " + "exist for the team this job runs in, please report the mismatch.)" + ) diff --git a/truss/cli/train/exec/uv.py b/truss/cli/train/exec/uv.py new file mode 100644 index 000000000..a11d5c6fd --- /dev/null +++ b/truss/cli/train/exec/uv.py @@ -0,0 +1,80 @@ +"""uv-specific pieces of `truss train exec`. + +Everything that knows what uv is lives here, so adding another project type (pip, +poetry, ...) means adding a sibling module rather than editing the exec builder. +""" + +import logging +from pathlib import Path +from typing import List + +import tomlkit + +logger = logging.getLogger(__name__) + +# The official uv image is the same slim Python as the non-uv default, with uv +# preinstalled. (uv stopped publishing bookworm-slim variants for current versions, +# hence trixie-slim.) +UV_BASE_IMAGE = "ghcr.io/astral-sh/uv:0.12.6-python3.12-trixie-slim" + +UV_LOCK_FILE = "uv.lock" +PYPROJECT_FILE = "pyproject.toml" + +# Skip-if-present, so this is safe on an image that already has uv. Needs curl, +# which the CUDA base image has but a custom --image may not. +# +# The download is separate from running it because in `curl ... | sh` the pipeline's +# status is `sh`'s: a failed download would feed an empty script to a shell that +# exits 0, so the `&&` chain would continue and the job would fail later with an +# opaque "uv: not found". +UV_INSTALL_SCRIPT_PATH = "/tmp/uv-install.sh" +UV_INSTALL_STEPS = [ + "{ command -v uv >/dev/null 2>&1 || { " + f"curl -LsSf https://astral.sh/uv/install.sh -o {UV_INSTALL_SCRIPT_PATH} && " + f"sh {UV_INSTALL_SCRIPT_PATH}" + " ; } ; }", + 'export PATH="$HOME/.local/bin:$PATH"', +] + + +def is_uv_project(source_dir: Path) -> bool: + """Whether `source_dir` carries uv project metadata. + + A `pyproject.toml` on its own is not enough -- poetry, hatch, PDM and setuptools + all ship one -- so require a uv lockfile or an explicit `[tool.uv]` section. + """ + if (source_dir / UV_LOCK_FILE).exists(): + return True + + pyproject = source_dir / PYPROJECT_FILE + if not pyproject.is_file(): + return False + try: + return "uv" in (tomlkit.parse(pyproject.read_text()).get("tool") or {}) + except Exception: + logger.debug("Could not read %s for uv detection.", pyproject, exc_info=True) + return False + + +class UvProject: + """A project whose dependencies are managed by uv. + + Satisfies the `Project` protocol structurally; there is no `build()` because + nothing is built locally -- the user's command does whatever building is needed. + """ + + label = "uv" + + def base_image(self) -> str: + return UV_BASE_IMAGE + + def setup(self, base_image: str) -> List[str]: + """Steps to run before the user's command, given the resolved base image. + + Empty on the uv image, which already ships uv. Anything else -- the CUDA + image used for GPU jobs, or a custom --image we know nothing about -- gets + the skip-if-present install. + """ + if base_image == UV_BASE_IMAGE: + return [] + return list(UV_INSTALL_STEPS) diff --git a/truss/cli/train/poller.py b/truss/cli/train/poller.py index 062e04fa1..8a16699b9 100644 --- a/truss/cli/train/poller.py +++ b/truss/cli/train/poller.py @@ -18,6 +18,7 @@ ] JOB_LOGGING_STATES = ["TRAINING_JOB_DEPLOYING", "TRAINING_JOB_RUNNING"] STATES_WITH_ERROR_MESSAGES = ["TRAINING_JOB_DEPLOY_FAILED"] +JOB_FAILED_STATES = ["TRAINING_JOB_FAILED", "TRAINING_JOB_DEPLOY_FAILED"] @dataclass @@ -92,6 +93,15 @@ def after_polling(self) -> None: elif self._current_status.status == "TRAINING_JOB_DEPLOY_FAILED": console.print("Training job failed during deployment.", style="red") + @property + def failed(self) -> bool: + """Whether the last status seen while polling was a failure. + + Lets a caller pick an exit code once `watch()` returns, instead of reaching + for the private status. + """ + return self._current_status.status in JOB_FAILED_STATES + def _update_from_current_status(self) -> None: current_job = self.api.get_training_job(self.project_id, self.job_id) self._current_status = Status( diff --git a/truss/cli/train_commands.py b/truss/cli/train_commands.py index 34a9fc393..b7c543842 100644 --- a/truss/cli/train_commands.py +++ b/truss/cli/train_commands.py @@ -7,6 +7,7 @@ import rich.table import rich_click as click +from rich.markup import escape import truss.cli.train.core as train_cli from truss.base.constants import TRAINING_TEMPLATE_DIR @@ -32,6 +33,18 @@ SORT_ORDER_ASC, SORT_ORDER_DESC, ) +from truss.cli.train.exec import ( + DEFAULT_CPU_COUNT, + DEFAULT_EXEC_PROJECT_NAME, + DEFAULT_MEMORY, + SUPPORTED_EXEC_ACCELERATORS, + UvProject, + build_exec_project, + get_project_type, + parse_environment_variables, + validate_secret_references, + validate_workspace_root, +) from truss.cli.train.workstation import ( SUPPORTED_WORKSTATION_ACCELERATORS, build_workstation_project, @@ -45,6 +58,7 @@ from truss.remote.remote_factory import RemoteFactory from truss.util.path import copy_tree_path from truss_train import TrainingJob +from truss_train import public_api as train_public_api @click.group() @@ -1366,3 +1380,244 @@ def workstation( watcher = TrainingLogWatcher(remote_provider.api, project_resp_id, job_id) for log in watcher.watch(): cli_log_utils.output_log(log) + + +@train.command(name="exec", context_settings={"ignore_unknown_options": True}) +@click.argument("start_command", nargs=-1, type=click.UNPROCESSED) +@click.option( + "--accelerator", + type=click.Choice(SUPPORTED_EXEC_ACCELERATORS, case_sensitive=False), + default=None, + help="GPU accelerator type. Omit for a CPU-only job (the default).", +) +@click.option( + "--gpu-count", + type=click.IntRange(1, 8), + default=None, + help="Number of GPUs (1-8, default: 1). Requires --accelerator.", +) +@click.option( + "--cpu-count", + type=click.IntRange(min=1), + default=DEFAULT_CPU_COUNT, + show_default=True, + help="Number of CPUs to request.", +) +@click.option( + "--memory", + type=str, + default=DEFAULT_MEMORY, + show_default=True, + help="Memory to request (e.g. 8Gi).", +) +@click.option( + "--project-name", + type=str, + required=False, + help="Training project name (default: the name of the current directory).", +) +@click.option("--image", type=str, required=False, help="Custom Docker base image.") +@click.option( + "--workspace-root", + type=str, + required=False, + help=( + "Directory to upload instead of just the current directory. Must be a " + "parent of the current directory." + ), +) +@click.option( + "--exclude-dir", + "exclude_dirs", + type=str, + multiple=True, + help=( + "Top-level directory of the workspace root to leave out of the upload. " + "Repeatable." + ), +) +@click.option( + "--external-dir", + "external_dirs", + type=str, + multiple=True, + help="Directory outside the workspace root to include in the upload. Repeatable.", +) +@click.option( + "--env", + type=str, + multiple=True, + help="Environment variable for the job as KEY=VALUE. Repeatable.", +) +@click.option( + "--secret", + "secrets", + type=str, + multiple=True, + help=( + "Environment variable sourced from a Baseten workspace secret, as " + "KEY=SECRET_NAME. Create secrets at https://app.baseten.co/settings/secrets. " + "Repeatable." + ), +) +@click.option( + "--with-uv", + is_flag=True, + default=False, + help=( + "Make uv available in the job image. Your command should invoke uv itself, " + "e.g. `truss train exec --with-uv -- uv run python my_script.py`." + ), +) +@click.option("--remote", type=str, required=False, help="Remote to use.") +@click.option( + "--team", + "provided_team_name", + type=str, + required=False, + help="Team name for the training project", +) +@click.option( + "--tail/--no-tail", + default=False, + show_default=True, + help=( + "Stream status + logs after push instead of returning immediately. With " + "--tail, exec exits non-zero if the job fails." + ), +) +@common.common_options() +def exec_training_job( + start_command: tuple[str, ...], + accelerator: Optional[str], + gpu_count: Optional[int], + cpu_count: int, + memory: str, + project_name: Optional[str], + image: Optional[str], + workspace_root: Optional[str], + exclude_dirs: tuple[str, ...], + external_dirs: tuple[str, ...], + env: tuple[str, ...], + secrets: tuple[str, ...], + with_uv: bool, + remote: Optional[str], + provided_team_name: Optional[str], + tail: bool, +): + """Run a command from the current directory as a training job. + + Archives the directory the command is invoked from, ships it to a training job, + and runs START_COMMAND there. Pass the command after `--`: + + truss train exec -- python my_script.py --steps 100 + + START_COMMAND always runs last and verbatim. Pass --with-uv to get uv in the + job image, then invoke `uv run` yourself. SSH into the job is available on + demand. The job gets no persistent storage and no checkpointing. + """ + if not start_command: + raise click.UsageError( + "No start command given. Pass the command to run after `--`, " + "e.g. `truss train exec -- python my_script.py`." + ) + if gpu_count is not None and accelerator is None: + raise click.UsageError("--gpu-count requires --accelerator.") + + if accelerator: + accelerator = accelerator.upper() + gpu_count = gpu_count or 1 + + environment_variables = parse_environment_variables(env=env, secrets=secrets) + + source_dir = Path.cwd() + # Validate before any API call: truss_train's own check runs after the training + # project has been created, which would leave a stray empty project behind. + workspace_dir = validate_workspace_root(source_dir, workspace_root) + if not project_name: + # Repeated runs from the same checkout should group into one project. + project_name = source_dir.name or DEFAULT_EXEC_PROJECT_NAME + + if not remote: + remote = remote_cli.inquire_remote_name() + + remote_provider: BasetenRemote = cast( + BasetenRemote, RemoteFactory.create(remote=remote) + ) + effective_team_name = provided_team_name or RemoteFactory.get_remote_team(remote) + _, team_id = _resolve_team_name( + remote_provider, effective_team_name, existing_project_name=project_name + ) + validate_secret_references(remote_provider.api, environment_variables) + + # --with-uv names uv explicitly, so it selects UvProject directly. Detection only + # drives the warning below, and is the hook a future --project-type would use. + detected_project = get_project_type(workspace_dir) + training_project = build_exec_project( + start_command=start_command, + project_name=project_name, + accelerator=accelerator, + gpu_count=gpu_count, + cpu_count=cpu_count, + memory=memory, + base_image=image, + project=UvProject() if with_uv else None, + workspace_root=workspace_root, + exclude_dirs=exclude_dirs, + external_dirs=external_dirs, + environment_variables=environment_variables, + ) + + compute_str = ( + f"{gpu_count}x {accelerator}" if accelerator else f"{cpu_count} CPU / {memory}" + ) + if not with_uv and detected_project is not None: + console.print( + f"Warning: this looks like a {detected_project.label} project, but " + "--with-uv was not passed, so uv will not be present in the job image.", + style="yellow", + ) + + # Escaped: `myproj[v2]` would otherwise be read as console markup, reporting a + # different name than the one being pushed. + console.print( + f"Launching [cyan]{escape(project_name)}[/cyan] from " + f"[cyan]{escape(str(source_dir))}[/cyan] on [cyan]{escape(compute_str)}[/cyan]..." + ) + + job_resp = train_public_api.push( + config=training_project, remote=remote, source_dir=source_dir, team_id=team_id + ) + + job_id = job_resp["id"] + console.print( + f"\n[green]Job created![/green]\n" + f"\n" + f"SSH is available on demand. Check the interactive session with:\n" + f" [cyan]truss train isession --job-id {job_id}[/cyan]\n" + f"\n" + f"Then SSH in with:\n" + f" [cyan]ssh training-job-{job_id}-0.ssh.baseten.co[/cyan]\n" + f"\n" + f"If you haven't set up SSH yet, run:\n" + f" [cyan]truss ssh setup[/cyan]\n" + f"\n" + f"View logs:\n" + f" [cyan]truss train logs --job-id {job_id} --tail[/cyan]\n" + f"\n" + f"Stop the job:\n" + f" [cyan]truss train stop --job-id {job_id}[/cyan]" + ) + + if tail: + project_resp_id = job_resp["training_project"]["id"] + watcher = TrainingLogWatcher(remote_provider.api, project_resp_id, job_id) + for log in watcher.watch(): + cli_log_utils.output_log(log) + + if watcher.failed: + # Without this, `truss train exec --tail -- pytest` is green in CI no + # matter what the job did. sys.exit rather than click's Exit, which + # subclasses RuntimeError and would be caught by `common_options`' error + # handler and reported as "ERROR Exit: 1". + sys.exit(1) diff --git a/truss/tests/cli/train/test_exec.py b/truss/tests/cli/train/test_exec.py new file mode 100644 index 000000000..b83dfbf06 --- /dev/null +++ b/truss/tests/cli/train/test_exec.py @@ -0,0 +1,1121 @@ +import os +import re +import subprocess +from contextlib import contextmanager +from pathlib import Path +from unittest.mock import Mock, patch + +import pytest +import rich_click as click +from click.testing import CliRunner + +from truss.cli.cli import truss_cli +from truss.cli.train.exec import ( + DEFAULT_CPU_COUNT, + DEFAULT_MEMORY, + PYTHON_BASE_IMAGE, + SECRETS_SETTINGS_URL, + SUPPORTED_EXEC_ACCELERATORS, + UvProject, + build_exec_project, + build_start_commands, + default_base_image, + get_project_type, + parse_environment_variables, + resolve_workspace_root, + validate_secret_references, + validate_workspace_root, +) +from truss.cli.train.exec.uv import ( + PYPROJECT_FILE, + UV_BASE_IMAGE, + UV_INSTALL_STEPS, + UV_LOCK_FILE, + is_uv_project, +) +from truss.cli.train.workstation import DEFAULT_BASE_IMAGE +from truss.remote.baseten.api import BasetenApi +from truss.remote.baseten.custom_types import TeamType +from truss.remote.baseten.remote import BasetenRemote +from truss_train.definitions import ( + InteractiveSessionProvider, + InteractiveSessionTrigger, + SecretReference, +) + +USER_COMMAND = ["uv", "run", "python", "my_script.py"] +USER_COMMAND_STR = "uv run python my_script.py" + +_ANSI_RE = re.compile(r"\x1b\[[0-9;]*m") +_BOX_DRAWING_RE = re.compile(r"[\u2500-\u257f]") + + +def _plain(text: str) -> str: + """`text` as plain, single-line text. + + rich renders errors and warnings as colorized, width-wrapped panels, and turns + color on under GITHUB_ACTIONS (and FORCE_COLOR) -- which splits tokens like + `--flag` with ANSI codes. Strip those and the panel borders, and collapse + whitespace, so assertions match the message regardless of terminal width or + color support. + """ + return " ".join(_BOX_DRAWING_RE.sub(" ", _ANSI_RE.sub("", text)).split()) + + +def _message_text(result) -> str: + """The CLI result's output as plain, single-line text.""" + return _plain(result.output) + + +def _uv_project(tmp_path: Path, lock: bool = True) -> Path: + (tmp_path / "pyproject.toml").write_text("[project]\nname = 'x'\n") + if lock: + (tmp_path / "uv.lock").write_text("") + return tmp_path + + +def _build(**overrides): + """Build with a complete argument set, so each test overrides only what it + exercises. `build_exec_project` takes no defaults of its own -- the CLI is the + single source of those -- so a baseline has to live somewhere, and a test helper + is the right place for it.""" + kwargs = dict( + start_command=["python", "my_script.py"], + project_name="my-project", + accelerator=None, + gpu_count=1, + cpu_count=DEFAULT_CPU_COUNT, + memory=DEFAULT_MEMORY, + base_image=None, + project=None, + workspace_root=None, + exclude_dirs=(), + external_dirs=(), + environment_variables={}, + ) + kwargs.update(overrides) + return build_exec_project(**kwargs) + + +def test_build_exec_project_requires_every_argument(): + """The looseness this guards against: a caller half-specifying a project.""" + with pytest.raises(TypeError): + build_exec_project(start_command=["python", "x.py"], project_name="p") + + +def test_build_exec_project_is_keyword_only(): + with pytest.raises(TypeError): + build_exec_project(["python", "x.py"], "p") # type: ignore[misc] + + +# --- project types ----------------------------------------------------------- + + +def test_get_project_type_detects_uv(tmp_path): + (tmp_path / UV_LOCK_FILE).write_text("") + project = get_project_type(tmp_path) + assert isinstance(project, UvProject) + assert project.label == "uv" + + +def test_get_project_type_returns_none_for_an_unrecognised_directory(tmp_path): + assert get_project_type(tmp_path) is None + + +def test_uv_project_wants_the_uv_image(): + assert UvProject().base_image() == UV_BASE_IMAGE + + +def test_uv_project_setup_is_empty_on_the_uv_image(): + """The image already ships uv, so there is nothing to prepend.""" + assert UvProject().setup(UV_BASE_IMAGE) == [] + + +@pytest.mark.parametrize("image", ["nvidia/cuda:12.8.1-devel-ubuntu24.04", "custom:1"]) +def test_uv_project_setup_installs_uv_on_any_other_image(image): + steps = UvProject().setup(image) + assert steps == UV_INSTALL_STEPS + assert "astral.sh/uv/install.sh" in steps[0] + + +def test_default_base_image_prefers_the_accelerator_over_the_project(): + """A GPU job needs the CUDA image; the project's setup steps cover the gap.""" + assert default_base_image("H100", UvProject()) == DEFAULT_BASE_IMAGE + assert default_base_image(None, UvProject()) == UV_BASE_IMAGE + assert default_base_image(None, None) == PYTHON_BASE_IMAGE + + +# --- builder: compute, image, session, workspace ----------------------------- + + +def test_build_exec_project_cpu_defaults(tmp_path): + project = _build( + start_command=["python", "my_script.py"], project_name="my-project" + ) + assert project.name == "my-project" + + job = project.job + assert job.compute.accelerator is None + assert job.compute.cpu_count == DEFAULT_CPU_COUNT + assert job.compute.memory == DEFAULT_MEMORY + assert job.compute.node_count == 1 + assert job.image.base_image == PYTHON_BASE_IMAGE + assert job.workspace is None + + +def test_build_exec_project_gpu(tmp_path): + job = _build( + start_command=["python", "my_script.py"], + project_name="my-project", + accelerator="H100", + gpu_count=4, + ).job + assert job.compute.accelerator is not None + assert job.compute.accelerator.accelerator.value == "H100" + assert job.compute.accelerator.count == 4 + assert job.image.base_image == DEFAULT_BASE_IMAGE + + +@pytest.mark.parametrize("accelerator", SUPPORTED_EXEC_ACCELERATORS) +def test_build_exec_project_supported_accelerators(accelerator, tmp_path): + job = _build( + start_command=["python", "my_script.py"], + project_name="my-project", + accelerator=accelerator, + ).job + assert job.compute.accelerator.accelerator.value == accelerator + + +def test_build_exec_project_invalid_accelerator(tmp_path): + with pytest.raises(ValueError): + _build( + start_command=["python", "my_script.py"], + project_name="my-project", + accelerator="INVALID", + ) + + +def test_build_exec_project_enables_ssh_on_demand(tmp_path): + job = _build( + start_command=["python", "my_script.py"], project_name="my-project" + ).job + assert job.interactive_session is not None + assert job.interactive_session.trigger == InteractiveSessionTrigger.ON_DEMAND + assert job.interactive_session.session_provider == InteractiveSessionProvider.SSH + + +def test_build_exec_project_leaves_storage_disabled(tmp_path): + runtime = _build( + start_command=["python", "my_script.py"], project_name="my-project" + ).job.runtime + assert runtime.cache_config is None + assert runtime.load_checkpoint_config is None + assert runtime.checkpointing_config.enabled is False + assert runtime.checkpointing_config.checkpoint_path is None + + +def test_build_exec_project_custom_image_wins(tmp_path): + job = _build( + start_command=["python", "my_script.py"], + project_name="my-project", + base_image="my-registry/my-image:latest", + ).job + assert job.image.base_image == "my-registry/my-image:latest" + + +def test_build_exec_project_workspace_from_dir_flags(tmp_path): + job = _build( + start_command=["python", "my_script.py"], + project_name="my-project", + workspace_root="..", + exclude_dirs=["data", "checkpoints"], + external_dirs=["../shared"], + ).job + assert job.workspace is not None + assert job.workspace.workspace_root == ".." + assert job.workspace.exclude_dirs == ["data", "checkpoints"] + assert job.workspace.external_dirs == ["../shared"] + + +def test_build_exec_project_environment_variables(tmp_path): + job = _build( + start_command=["python", "my_script.py"], + project_name="my-project", + environment_variables={ + "PLAIN": "1", + "BASETEN_API_KEY": SecretReference(name="my_api_key"), + }, + ).job + assert job.runtime.environment_variables == { + "PLAIN": "1", + "BASETEN_API_KEY": SecretReference(name="my_api_key"), + } + + +def test_build_exec_project_no_environment_variables_by_default(tmp_path): + job = _build( + start_command=["python", "my_script.py"], project_name="my-project" + ).job + assert job.runtime.environment_variables == {} + + +# --- builder: --env / --secret parsing --------------------------------------- + + +def test_parse_environment_variables_splits_on_first_equals_only(): + parsed = parse_environment_variables(env=["TOKEN=abc=def==", "PLAIN=1"]) + assert parsed == {"TOKEN": "abc=def==", "PLAIN": "1"} + + +def test_parse_environment_variables_builds_secret_references(): + parsed = parse_environment_variables(secrets=["BASETEN_API_KEY=my_api_key"]) + assert parsed == {"BASETEN_API_KEY": SecretReference(name="my_api_key")} + assert isinstance(parsed["BASETEN_API_KEY"], SecretReference) + + +@pytest.mark.parametrize( + "kwargs, expected", + [ + ({"env": ["NO_EQUALS"]}, "Invalid --env value 'NO_EQUALS'"), + ({"env": ["=novalue"]}, "Invalid --env value '=novalue'"), + ({"secrets": ["NO_EQUALS"]}, "Invalid --secret value 'NO_EQUALS'"), + ], +) +def test_parse_environment_variables_rejects_entries_without_a_key(kwargs, expected): + with pytest.raises(click.UsageError, match=re.escape(expected)): + parse_environment_variables(**kwargs) + + +def test_parse_environment_variables_rejects_duplicate_keys(): + with pytest.raises(click.UsageError, match="set more than once"): + parse_environment_variables(env=["KEY=literal"], secrets=["KEY=secret_name"]) + + +# --- builder: uv detection and start commands -------------------------------- + + +def test_build_start_commands_runs_the_command_verbatim(): + assert build_start_commands( + start_command=["python", "my script.py", "--steps", "100"] + ) == ["python 'my script.py' --steps 100"] + + +def test_build_start_commands_prepends_the_idempotent_uv_install(): + assert build_start_commands( + start_command=USER_COMMAND, setup_steps=UV_INSTALL_STEPS + ) == [ + "/bin/sh -c '{ command -v uv >/dev/null 2>&1 || { " + "curl -LsSf https://astral.sh/uv/install.sh -o /tmp/uv-install.sh && " + "sh /tmp/uv-install.sh ; } ; } && " + 'export PATH="$HOME/.local/bin:$PATH" && ' + f"{USER_COMMAND_STR}'" + ] + + +def test_build_exec_project_with_uv_selects_the_uv_image_on_cpu(): + job = _build( + start_command=USER_COMMAND, project_name="my-project", project=UvProject() + ).job + assert job.image.base_image == UV_BASE_IMAGE + # The uv image already ships uv, so the command is the only start command. + assert job.runtime.start_commands == [USER_COMMAND_STR] + + +def test_build_exec_project_with_uv_installs_uv_on_the_gpu_image(): + job = _build( + start_command=USER_COMMAND, + project_name="my-project", + accelerator="H100", + project=UvProject(), + ).job + assert job.image.base_image == DEFAULT_BASE_IMAGE + assert "astral.sh/uv/install.sh" in job.runtime.start_commands[0] + assert job.runtime.start_commands[0].endswith(f"{USER_COMMAND_STR}'") + + +def test_build_exec_project_with_uv_installs_uv_on_a_custom_image(): + job = _build( + start_command=USER_COMMAND, + project_name="my-project", + base_image="my-registry/my-image:latest", + project=UvProject(), + ).job + assert job.image.base_image == "my-registry/my-image:latest" + assert "astral.sh/uv/install.sh" in job.runtime.start_commands[0] + assert job.runtime.start_commands[0].endswith(f"{USER_COMMAND_STR}'") + + +@pytest.mark.parametrize("accelerator", [None, "H100"]) +def test_build_exec_project_without_with_uv_injects_nothing(accelerator): + job = _build( + start_command=USER_COMMAND, project_name="my-project", accelerator=accelerator + ).job + assert job.runtime.start_commands == [USER_COMMAND_STR] + + +# --- uv install failure mode, secret/env edge cases, workspace root ---------- + + +def test_uv_install_step_fails_when_the_download_fails(): + """`curl | sh` would exit 0 on a failed download, hiding the error until the + job died with an opaque `uv: not found`.""" + command = build_start_commands( + start_command=["uv", "--version"], setup_steps=UV_INSTALL_STEPS + )[0] + script = command[len("/bin/sh -c ") :] + # The install group must be its own command list, not a pipeline into sh. + assert "install.sh | sh" not in script + assert "-o /tmp/uv-install.sh && sh /tmp/uv-install.sh" in script + + # Syntactically valid, and the group reports failure when the download fails. + assert subprocess.run(["sh", "-n", "-c", script.strip("'")]).returncode == 0 + failing_group = ( + "{ command -v definitely_not_a_real_binary >/dev/null 2>&1 || { " + "curl -LsSf https://astral.sh/NOPE-404-xyz/install.sh -o /tmp/uv-probe.sh " + "&& sh /tmp/uv-probe.sh ; } ; }" + ) + assert ( + subprocess.run(["sh", "-c", failing_group], capture_output=True).returncode != 0 + ) + + +@pytest.mark.parametrize("entry", ["FOO=", "FOO"]) +def test_parse_environment_variables_rejects_an_empty_secret_name(entry): + with pytest.raises(click.UsageError, match="Invalid --secret value"): + parse_environment_variables(secrets=[entry]) + + +def test_parse_environment_variables_still_allows_an_empty_env_value(): + """Unlike a secret name, an empty --env value is legitimate.""" + assert parse_environment_variables(env=["EMPTY="]) == {"EMPTY": ""} + + +def test_is_uv_project_accepts_a_lockfile(tmp_path): + (tmp_path / UV_LOCK_FILE).write_text("") + assert is_uv_project(tmp_path) + + +def test_is_uv_project_accepts_a_tool_uv_section(tmp_path): + (tmp_path / PYPROJECT_FILE).write_text( + "[project]\nname = 'x'\n\n[tool.uv]\ndev-dependencies = []\n" + ) + assert is_uv_project(tmp_path) + + +def test_is_uv_project_rejects_a_poetry_project(tmp_path): + """A pyproject.toml alone is not a uv project -- poetry, hatch, PDM and + setuptools all ship one, and warning on those is noise.""" + (tmp_path / PYPROJECT_FILE).write_text( + "[tool.poetry]\nname = 'x'\nversion = '0.1.0'\n" + ) + assert not is_uv_project(tmp_path) + + +def test_is_uv_project_rejects_malformed_toml(tmp_path): + (tmp_path / PYPROJECT_FILE).write_text("this is [not valid toml") + assert not is_uv_project(tmp_path) + + +def test_is_uv_project_rejects_a_plain_directory(tmp_path): + assert not is_uv_project(tmp_path) + + +def test_resolve_workspace_root_defaults_to_the_source_dir(tmp_path): + assert resolve_workspace_root(tmp_path, None) == tmp_path + + +def test_resolve_workspace_root_handles_relative_and_absolute(tmp_path): + child = tmp_path / "child" + child.mkdir() + assert resolve_workspace_root(child, "..") == tmp_path.resolve() + assert resolve_workspace_root(child, str(tmp_path)) == tmp_path.resolve() + + +def test_validate_workspace_root_accepts_a_parent(tmp_path): + child = tmp_path / "child" + child.mkdir() + assert validate_workspace_root(child, "..") == tmp_path.resolve() + + +def test_validate_workspace_root_rejects_a_non_parent(tmp_path): + sibling = tmp_path / "sibling" + sibling.mkdir() + here = tmp_path / "here" + here.mkdir() + with pytest.raises(click.UsageError, match="does not contain the current"): + validate_workspace_root(here, str(sibling)) + + +def test_validate_workspace_root_rejects_a_missing_directory(tmp_path): + with pytest.raises(click.UsageError, match="not a directory"): + validate_workspace_root(tmp_path, "nope-does-not-exist") + + +# --- CLI --------------------------------------------------------------------- + + +def _mock_remote(secrets=("my_api_key",)): + mock_remote = Mock(spec=BasetenRemote) + mock_remote.api = Mock(spec=BasetenApi) + mock_remote.api.get_teams.return_value = { + "team-a": TeamType(id="team1", name="team-a", default=True) + } + mock_remote.api.list_training_projects.return_value = [] + mock_remote.api.get_all_secrets.return_value = { + "secrets": [{"name": name} for name in secrets] + } + return mock_remote + + +@contextmanager +def _chdir(directory: Path): + original_cwd = Path.cwd() + os.chdir(directory) + try: + yield + finally: + os.chdir(original_cwd) + + +def _invoke_exec(args, cwd: Path, tail: bool = False, remote=None): + """Invoke `truss train exec` from `cwd`, returning (result, mock_push).""" + base_args = ["train", "exec", "--remote", "test_remote"] + if tail: + base_args.append("--tail") + remote = remote if remote is not None else _mock_remote() + + with ( + _chdir(cwd), + patch("truss_train.public_api.push") as mock_push, + patch( + "truss.cli.train_commands.RemoteFactory.get_remote_team", return_value=None + ), + patch("truss.cli.train_commands.RemoteFactory.create", return_value=remote), + ): + mock_push.return_value = { + "id": "job123", + "training_project": {"id": "proj123", "name": cwd.name}, + } + result = CliRunner().invoke(truss_cli, base_args + list(args)) + + return result, mock_push + + +def test_exec_passes_start_command_through_after_double_dash(tmp_path): + result, mock_push = _invoke_exec( + ["--", "python", "my_script.py", "--steps", "100", "--verbose"], tmp_path + ) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["config"].job.runtime.start_commands == [ + "python my_script.py --steps 100 --verbose" + ] + + +def test_exec_start_command_may_reuse_our_own_flag_names(tmp_path): + """Everything after `--` belongs to the command, even `--memory`/`--tail`.""" + result, mock_push = _invoke_exec( + [ + "--memory", + "16Gi", + "--", + "python", + "my_script.py", + "--memory", + "4Gi", + "--tail", + ], + tmp_path, + ) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.runtime.start_commands == ["python my_script.py --memory 4Gi --tail"] + assert job.compute.memory == "16Gi" + + +def test_exec_defaults_to_cpu_only_job(tmp_path): + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.compute.accelerator is None + assert job.compute.cpu_count == DEFAULT_CPU_COUNT + assert job.compute.memory == DEFAULT_MEMORY + + +def test_exec_pushes_current_directory_as_source_dir(tmp_path): + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["source_dir"] == tmp_path.resolve() + + +def test_exec_defaults_project_name_to_directory_name(tmp_path): + work_dir = tmp_path / "my-checkout" + work_dir.mkdir() + + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], work_dir) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["config"].name == "my-checkout" + + +def test_exec_project_name_overrides_directory_name(tmp_path): + result, mock_push = _invoke_exec( + ["--project-name", "explicit-name", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["config"].name == "explicit-name" + + +def test_exec_requires_a_start_command(tmp_path): + result, mock_push = _invoke_exec([], tmp_path) + + assert result.exit_code != 0 + assert "No start command given" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_gpu_count_requires_accelerator(tmp_path): + result, mock_push = _invoke_exec( + ["--gpu-count", "2", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code != 0 + assert "--gpu-count requires --accelerator" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_accelerator_is_normalized(tmp_path): + result, mock_push = _invoke_exec( + ["--accelerator", "h100", "--gpu-count", "2", "--", "python", "my_script.py"], + tmp_path, + ) + + assert result.exit_code == 0, result.output + accelerator = mock_push.call_args[1]["config"].job.compute.accelerator + assert accelerator.accelerator.value == "H100" + assert accelerator.count == 2 + + +def test_exec_cpu_and_memory_flags(tmp_path): + result, mock_push = _invoke_exec( + ["--cpu-count", "8", "--memory", "16Gi", "--", "python", "my_script.py"], + tmp_path, + ) + + assert result.exit_code == 0, result.output + compute = mock_push.call_args[1]["config"].job.compute + assert compute.cpu_count == 8 + assert compute.memory == "16Gi" + + +def test_exec_directory_flags_build_a_workspace(tmp_path): + result, mock_push = _invoke_exec( + [ + "--workspace-root", + "..", + "--exclude-dir", + "data", + "--exclude-dir", + "logs", + "--", + "python", + "my_script.py", + ], + tmp_path, + ) + + assert result.exit_code == 0, result.output + workspace = mock_push.call_args[1]["config"].job.workspace + assert workspace is not None + assert workspace.workspace_root == ".." + assert workspace.exclude_dirs == ["data", "logs"] + + +def test_exec_passes_team_id_to_push(tmp_path): + result, mock_push = _invoke_exec( + ["--team", "team-a", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["team_id"] == "team1" + + +def test_exec_env_and_secret_flags(tmp_path): + result, mock_push = _invoke_exec( + [ + "--env", + "MY_URL=https://example.com/?a=1&b=2", + "--secret", + "BASETEN_API_KEY=my_api_key", + "--", + "python", + "my_script.py", + ], + tmp_path, + ) + + assert result.exit_code == 0, result.output + assert mock_push.call_args[1]["config"].job.runtime.environment_variables == { + "MY_URL": "https://example.com/?a=1&b=2", + "BASETEN_API_KEY": SecretReference(name="my_api_key"), + } + + +def test_exec_rejects_env_without_equals(tmp_path): + result, mock_push = _invoke_exec( + ["--env", "BROKEN", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code != 0 + assert "Invalid --env value" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_rejects_key_set_by_both_env_and_secret(tmp_path): + result, mock_push = _invoke_exec( + [ + "--env", + "KEY=literal", + "--secret", + "KEY=secret_name", + "--", + "python", + "my_script.py", + ], + tmp_path, + ) + + assert result.exit_code != 0 + assert "set more than once" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_with_uv_selects_the_uv_image_and_keeps_the_command(tmp_path): + result, mock_push = _invoke_exec( + ["--with-uv", "--"] + USER_COMMAND, _uv_project(tmp_path) + ) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.image.base_image == UV_BASE_IMAGE + assert job.runtime.start_commands == [USER_COMMAND_STR] + + +def test_exec_without_with_uv_uses_the_plain_python_image(tmp_path): + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.image.base_image == PYTHON_BASE_IMAGE + assert job.runtime.start_commands == ["python my_script.py"] + + +@pytest.mark.parametrize("with_uv", [True, False]) +def test_exec_image_flag_overrides_the_base_image(with_uv, tmp_path): + args = ["--image", "my-registry/my-image:latest", "--"] + USER_COMMAND + if with_uv: + args.insert(0, "--with-uv") + + result, mock_push = _invoke_exec(args, _uv_project(tmp_path)) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.image.base_image == "my-registry/my-image:latest" + if with_uv: + # An unknown image may not ship uv, so the install step is prepended. + assert "astral.sh/uv/install.sh" in job.runtime.start_commands[0] + assert job.runtime.start_commands[0].endswith(f"{USER_COMMAND_STR}'") + else: + assert job.runtime.start_commands == [USER_COMMAND_STR] + + +def test_exec_warns_about_a_uv_project_without_with_uv(tmp_path): + result, _ = _invoke_exec(["--", "python", "my_script.py"], _uv_project(tmp_path)) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" in _message_text(result) + + +def test_exec_does_not_warn_for_a_poetry_project(tmp_path): + """The warning is about uv specifically; a bare pyproject.toml is not a signal.""" + (tmp_path / "pyproject.toml").write_text("[tool.poetry]\nname = 'x'\n") + + result, _ = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" not in _message_text(result) + + +def test_exec_does_not_warn_when_with_uv_is_passed(tmp_path): + result, _ = _invoke_exec(["--with-uv", "--"] + USER_COMMAND, _uv_project(tmp_path)) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" not in _message_text(result) + + +def test_exec_does_not_warn_for_a_plain_directory(tmp_path): + result, _ = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" not in _message_text(result) + + +# --- secret existence validation --------------------------------------------- + + +def _api(secrets_response): + # spec'd: renaming get_all_secrets should fail loudly, not silently disable this. + api = Mock(spec=BasetenApi) + api.get_all_secrets.return_value = secrets_response + return api + + +def test_validate_secret_references_skips_the_call_without_any_secret_refs(capsys): + api = _api({"secrets": []}) + + validate_secret_references(api, {"PLAIN": "literal"}) + + api.get_all_secrets.assert_not_called() + assert _plain(capsys.readouterr().out) == "" + + +def test_validate_secret_references_is_quiet_when_the_secret_exists(capsys): + api = _api({"secrets": [{"name": "my_api_key"}]}) + + validate_secret_references(api, {"K": SecretReference(name="my_api_key")}) + + api.get_all_secrets.assert_called_once_with() + assert _plain(capsys.readouterr().out) == "" + + +def test_validate_secret_references_errors_when_the_secret_is_absent(): + api = _api({"secrets": [{"name": "other"}]}) + + with pytest.raises(click.UsageError) as excinfo: + validate_secret_references(api, {"K": SecretReference(name="my_api_key")}) + + message = str(excinfo.value) + assert "Secret my_api_key was not found in this workspace's secrets" in message + assert "Create it at" in message + assert SECRETS_SETTINGS_URL in message + assert "not team-scoped" in message + + +def test_validate_secret_references_errors_for_an_empty_workspace(): + """An empty listing is a real answer, unlike an unreadable one.""" + with pytest.raises(click.UsageError, match="my_api_key"): + validate_secret_references( + _api({"secrets": []}), {"K": SecretReference(name="my_api_key")} + ) + + +def test_validate_secret_references_names_every_missing_secret(): + with pytest.raises(click.UsageError) as excinfo: + validate_secret_references( + _api({"secrets": [{"name": "present"}]}), + { + "A": SecretReference(name="missing_a"), + "B": SecretReference(name="present"), + "C": SecretReference(name="missing_b"), + "D": "literal", + }, + ) + + message = str(excinfo.value) + assert "missing_a" in message and "missing_b" in message + assert "present" not in message + assert "Secrets missing_a, missing_b were not found" in message + assert "Create them at" in message + + +@pytest.mark.parametrize( + "response", + [ + ["my_api_key"], + [{"name": "my_api_key"}], + {"secrets": ["my_api_key"]}, + {"secrets": [{"name": "my_api_key"}]}, + ], +) +def test_validate_secret_references_accepts_the_plausible_payload_shapes( + response, capsys +): + validate_secret_references( + _api(response), {"K": SecretReference(name="my_api_key")} + ) + + assert _plain(capsys.readouterr().out) == "" + + +@pytest.mark.parametrize( + "response", ["a string", 42, None, {"data": []}, {"secrets": "nope"}, [123]] +) +def test_validate_secret_references_stays_silent_on_an_unreadable_payload( + response, capsys +): + """Rather than guess a shape and warn about a secret that is really there.""" + validate_secret_references( + _api(response), {"K": SecretReference(name="my_api_key")} + ) + + assert _plain(capsys.readouterr().out) == "" + + +def test_validate_secret_references_swallows_api_errors(capsys): + api = Mock(spec=BasetenApi) + api.get_all_secrets.side_effect = RuntimeError("403 Forbidden") + + validate_secret_references(api, {"K": SecretReference(name="my_api_key")}) + + assert _plain(capsys.readouterr().out) == "" + + +def test_validate_secret_references_reports_the_name_verbatim(): + """click.UsageError messages are not markup-parsed, so escaping them would leak + a backslash into the name the user sees.""" + with pytest.raises(click.UsageError) as excinfo: + validate_secret_references( + _api({"secrets": []}), {"K": SecretReference(name="[bold]weird")} + ) + + message = str(excinfo.value) + assert "Secret [bold]weird was not found" in message + assert "\\" not in message + + +def test_exec_errors_and_does_not_push_when_a_secret_is_missing(tmp_path): + """Absent from a listing we could read means the job would fail to start.""" + remote = _mock_remote(secrets=("some_other_secret",)) + + result, mock_push = _invoke_exec( + ["--secret", "BASETEN_API_KEY=my_api_key", "--", "python", "my_script.py"], + tmp_path, + remote=remote, + ) + + assert result.exit_code != 0 + assert "my_api_key was not found in this workspace" in _message_text(result) + assert SECRETS_SETTINGS_URL in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_pushes_when_the_secret_exists(tmp_path): + result, mock_push = _invoke_exec( + ["--secret", "BASETEN_API_KEY=my_api_key", "--", "python", "my_script.py"], + tmp_path, + remote=_mock_remote(secrets=("my_api_key",)), + ) + + assert result.exit_code == 0, result.output + assert "was not found in this workspace" not in _message_text(result) + mock_push.assert_called_once() + + +def test_exec_still_pushes_when_listing_secrets_fails(tmp_path): + remote = _mock_remote() + remote.api.get_all_secrets.side_effect = RuntimeError("403 Forbidden") + + result, mock_push = _invoke_exec( + ["--secret", "BASETEN_API_KEY=my_api_key", "--", "python", "my_script.py"], + tmp_path, + remote=remote, + ) + + assert result.exit_code == 0, result.output + assert "Traceback" not in result.output + assert "403 Forbidden" not in result.output + mock_push.assert_called_once() + + +def test_exec_still_pushes_when_the_secrets_payload_is_unreadable(tmp_path): + """An unparseable listing is an API problem, not proof the secret is missing.""" + remote = _mock_remote() + remote.api.get_all_secrets.return_value = {"unexpected": "shape"} + + result, mock_push = _invoke_exec( + ["--secret", "BASETEN_API_KEY=my_api_key", "--", "python", "my_script.py"], + tmp_path, + remote=remote, + ) + + assert result.exit_code == 0, result.output + assert "was not found in this workspace" not in _message_text(result) + mock_push.assert_called_once() + + +def test_exec_with_uv_uses_the_uv_image_even_without_uv_metadata(tmp_path): + """--with-uv is about the image, not about detection: it names uv explicitly, so + it applies whether or not the directory carries uv metadata.""" + result, mock_push = _invoke_exec(["--with-uv", "--"] + USER_COMMAND, tmp_path) + + assert result.exit_code == 0, result.output + job = mock_push.call_args[1]["config"].job + assert job.image.base_image == UV_BASE_IMAGE + assert job.runtime.start_commands == [USER_COMMAND_STR] + + +def test_exec_skips_the_secrets_call_without_secret_flags(tmp_path): + remote = _mock_remote() + + result, mock_push = _invoke_exec( + ["--env", "PLAIN=1", "--", "python", "my_script.py"], tmp_path, remote=remote + ) + + assert result.exit_code == 0, result.output + remote.api.get_all_secrets.assert_not_called() + mock_push.assert_called_once() + + +def test_exec_does_not_tail_by_default(tmp_path): + """The motivating use case is a client running for hours, so blocking the + terminal is the wrong default; matches `push` and `workstation`.""" + with patch("truss.cli.train_commands.TrainingLogWatcher") as mock_watcher: + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], tmp_path) + + assert result.exit_code == 0, result.output + mock_push.assert_called_once() + mock_watcher.assert_not_called() + + +def test_exec_tails_when_asked(tmp_path): + with patch("truss.cli.train_commands.TrainingLogWatcher") as mock_watcher: + mock_watcher.return_value.watch.return_value = [] + mock_watcher.return_value.failed = False + result, _ = _invoke_exec(["--", "python", "my_script.py"], tmp_path, tail=True) + + assert result.exit_code == 0, result.output + assert mock_watcher.call_args[0][1:] == ("proj123", "job123") + + +def test_exec_exits_nonzero_when_the_job_fails(tmp_path): + """Otherwise `truss train exec --tail -- pytest` is green in CI regardless of + outcome. --tail is passed explicitly: it is opt-in, so this cannot rely on a + default.""" + with patch("truss.cli.train_commands.TrainingLogWatcher") as mock_watcher: + mock_watcher.return_value.watch.return_value = [] + mock_watcher.return_value.failed = True + result, mock_push = _invoke_exec( + ["--", "python", "my_script.py"], tmp_path, tail=True + ) + + assert result.exit_code == 1, result.output + mock_push.assert_called_once() + # A clean exit, not an error surfaced by the common_options handler. + assert "ERROR" not in _message_text(result) + assert "Traceback" not in result.output + + +def test_exec_exits_zero_when_the_job_succeeds(tmp_path): + with patch("truss.cli.train_commands.TrainingLogWatcher") as mock_watcher: + mock_watcher.return_value.watch.return_value = [] + mock_watcher.return_value.failed = False + result, _ = _invoke_exec(["--", "python", "my_script.py"], tmp_path, tail=True) + + assert result.exit_code == 0, result.output + + +def test_exec_escapes_brackets_in_the_launch_line(tmp_path): + """A directory named `myproj[v2]` must not be reported as `myproj`.""" + work_dir = tmp_path / "myproj[v2]" + work_dir.mkdir() + + result, mock_push = _invoke_exec(["--", "python", "my_script.py"], work_dir) + + assert result.exit_code == 0, result.output + assert "myproj[v2]" in _message_text(result) + assert mock_push.call_args[1]["config"].name == "myproj[v2]" + + +def test_exec_survives_a_closing_tag_in_the_project_name(tmp_path): + """`[/cyan]` in interpolated text would otherwise raise a rich MarkupError. + + A directory name can't contain `/`, but --project-name can. + """ + result, mock_push = _invoke_exec( + ["--project-name", "weird[/cyan]name", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code == 0, result.output + assert "MarkupError" not in result.output + assert mock_push.call_args[1]["config"].name == "weird[/cyan]name" + + +@pytest.mark.parametrize("value", ["0", "-4"]) +def test_exec_rejects_a_nonpositive_cpu_count(value, tmp_path): + result, mock_push = _invoke_exec( + ["--cpu-count", value, "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code != 0 + mock_push.assert_not_called() + + +def test_exec_rejects_an_empty_secret_name(tmp_path): + result, mock_push = _invoke_exec( + ["--secret", "KEY=", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code != 0 + assert "Invalid --secret value" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_rejects_a_workspace_root_that_is_not_a_parent(tmp_path): + """Client-side, so a bad value can't leave a stray empty training project.""" + sibling = tmp_path / "sibling" + sibling.mkdir() + here = tmp_path / "here" + here.mkdir() + + result, mock_push = _invoke_exec( + ["--workspace-root", str(sibling), "--", "python", "my_script.py"], here + ) + + assert result.exit_code != 0 + assert "does not contain the current" in _message_text(result) + mock_push.assert_not_called() + + +def test_exec_uv_warning_follows_the_workspace_root(tmp_path): + """With --workspace-root the parent is what gets archived and executed, so that + is the directory whose uv metadata matters.""" + (tmp_path / "uv.lock").write_text("") + child = tmp_path / "child" + child.mkdir() + + result, _ = _invoke_exec( + ["--workspace-root", "..", "--", "python", "my_script.py"], child + ) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" in _message_text(result) + + +def test_exec_uv_warning_ignores_the_cwd_when_workspace_root_is_set(tmp_path): + """The mirror of the above: uv metadata in the cwd is irrelevant when the + archived root is elsewhere.""" + child = tmp_path / "child" + child.mkdir() + (child / "uv.lock").write_text("") + + result, _ = _invoke_exec( + ["--workspace-root", "..", "--", "python", "my_script.py"], child + ) + + assert result.exit_code == 0, result.output + assert "--with-uv was not passed" not in _message_text(result) + + +def test_exec_no_tail_flag_is_still_accepted(tmp_path): + """Now the same as the default, but the paired form must keep working.""" + with patch("truss.cli.train_commands.TrainingLogWatcher") as mock_watcher: + result, _ = _invoke_exec( + ["--no-tail", "--", "python", "my_script.py"], tmp_path + ) + + assert result.exit_code == 0, result.output + mock_watcher.assert_not_called()