Add truss train exec to run a local directory as a Baseten training job

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",
+]

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)

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

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.)"
+    )