truss 0.11.6rc102__py3-none-any.whl → 0.11.24rc2__py3-none-any.whl

truss/cli/train/deploy_checkpoints/deploy_checkpoints.py

@@ -1,7 +1,6 @@
+import json
 import re
-import tempfile
 from collections import OrderedDict
-from pathlib import Path
 from typing import List, Optional, Union
 
 import rich_click as click
@@ -11,7 +10,8 @@ from truss.base import truss_config
 from truss.cli.train import common
 from truss.cli.train.types import (
     DeployCheckpointsConfigComplete,
-    PrepareCheckpointResult,
+    DeploySuccessModelVersion,
+    DeploySuccessResult,
 )
 from truss.cli.utils.output import console
 from truss.remote.baseten.remote import BasetenRemote
@@ -25,18 +25,9 @@ from truss_train.definitions import (
     SecretReference,
 )
 
-from .deploy_full_checkpoints import (
-    hydrate_full_checkpoint,
-    render_vllm_full_truss_config,
-)
-from .deploy_lora_checkpoints import (
-    hydrate_lora_checkpoint,
-    render_vllm_lora_truss_config,
-)
-from .deploy_whisper_checkpoints import (
-    hydrate_whisper_checkpoint,
-    render_vllm_whisper_truss_config,
-)
+from .deploy_full_checkpoints import hydrate_full_checkpoint
+from .deploy_lora_checkpoints import hydrate_lora_checkpoint
+from .deploy_whisper_checkpoints import hydrate_whisper_checkpoint
 
 HF_TOKEN_ENVVAR_NAME = "HF_TOKEN"
 # If we change this, make sure to update the logic in backend codebase
@@ -44,28 +35,184 @@ CHECKPOINT_PATTERN = re.compile(r".*checkpoint-\d+(?:-\d+)?$")
 ALLOWED_DEPLOYMENT_NAMES = re.compile(r"^[0-9a-zA-Z_\-\.]*$")
 
 
-def prepare_checkpoint_deploy(
+def create_model_version_from_inference_template(
     remote_provider: BasetenRemote,
     checkpoint_deploy_config: DeployCheckpointsConfig,
     project_id: Optional[str],
     job_id: Optional[str],
-) -> PrepareCheckpointResult:
+    dry_run: bool,
+) -> DeploySuccessResult:
     checkpoint_deploy_config = _hydrate_deploy_config(
         checkpoint_deploy_config, remote_provider, project_id, job_id
     )
-    rendered_truss = _render_truss_config_for_checkpoint_deployment(
-        checkpoint_deploy_config
-    )
-    truss_directory = Path(tempfile.mkdtemp())
-    truss_config_path = truss_directory / "config.yaml"
-    rendered_truss.write_to_yaml_file(truss_config_path)
-    console.print(rendered_truss, style="green")
-    console.print(f"Writing truss config to {truss_config_path}", style="yellow")
-    return PrepareCheckpointResult(
-        truss_directory=truss_directory,
+
+    request_data = _build_inference_template_request(
         checkpoint_deploy_config=checkpoint_deploy_config,
+        remote_provider=remote_provider,
+        dry_run=dry_run,
     )
 
+    # Call the GraphQL mutation to create model version from inference template
+    try:
+        result = remote_provider.api.create_model_version_from_inference_template(
+            request_data
+        )
+        truss_config_result = _get_truss_config_from_result(result)
+
+        model_version = None
+        if result and result.get("model_version"):
+            console.print(
+                f"Successfully created model version: {result['model_version']['name']}",
+                style="green",
+            )
+            console.print(
+                f"Model version ID: {result['model_version']['id']}", style="yellow"
+            )
+            model_version = DeploySuccessModelVersion.model_validate(
+                result["model_version"]
+            )
+        elif not dry_run:
+            console.print(
+                "Warning: Unexpected response format from server", style="yellow"
+            )
+            console.print(f"Response: {result}", style="yellow")
+
+    except Exception as e:
+        console.print(f"Error creating model version: {e}", style="red")
+        raise
+
+    return DeploySuccessResult(
+        deploy_config=checkpoint_deploy_config,
+        truss_config=truss_config_result,
+        model_version=model_version,
+    )
+
+
+def _get_truss_config_from_result(result: dict) -> Optional[truss_config.TrussConfig]:
+    if result and result.get("truss_config"):
+        truss_config_dict = json.loads(result["truss_config"])
+        return truss_config.TrussConfig.from_dict(truss_config_dict)
+    # Although this should never happen, we defensively allow ourselves to return None
+    # because we need a failure to handle the truss config doesn't necessarily mean we failed to deploy
+    # the model version.
+    console.print(
+        "No truss config returned. Reach out to Baseten for support if this persists.",
+        style="red",
+    )
+    return None
+
+
+def _build_inference_template_request(
+    checkpoint_deploy_config: DeployCheckpointsConfigComplete,
+    remote_provider: BasetenRemote,
+    dry_run: bool,
+) -> dict:
+    """
+    Build the GraphQL request data structure for createModelVersionFromInferenceTemplate mutation.
+    """
+
+    # Build weights sources
+    weights_sources = []
+    for checkpoint in checkpoint_deploy_config.checkpoint_details.checkpoints:
+        # Extract checkpoint name from the first path
+        weights_source = {
+            "weight_source_type": "B10_CHECKPOINTING",
+            "b10_training_checkpoint_weights_source": {
+                "checkpoint": {
+                    "training_job_id": checkpoint.training_job_id,
+                    "checkpoint_name": checkpoint.checkpoint_name,
+                }
+            },
+        }
+        weights_sources.append(weights_source)
+
+    # Build environment variables
+    environment_variables = []
+    for name, value in checkpoint_deploy_config.runtime.environment_variables.items():
+        if isinstance(value, SecretReference):
+            env_var = {"name": name, "value": value.name, "is_secret_reference": True}
+        else:
+            env_var = {"name": name, "value": str(value), "is_secret_reference": False}
+        environment_variables.append(env_var)
+
+    # Build inference stack
+    inference_stack = {
+        "stack_type": "VLLM",
+        "environment_variables": environment_variables,
+    }
+
+    # Get instance type ID from compute spec
+    instance_type_id = _get_instance_type_id(
+        checkpoint_deploy_config.compute, remote_provider
+    )
+
+    # Build the complete request
+    request_data = {
+        "metadata": {"oracle_name": checkpoint_deploy_config.model_name},
+        "weights_sources": weights_sources,
+        "inference_stack": inference_stack,
+        "instance_type_id": instance_type_id,
+        "dry_run": dry_run,
+    }
+
+    return request_data
+
+
+def _get_instance_type_id(compute: Compute, remote_provider: BasetenRemote) -> str:
+    """
+    Get the instance type ID based on the compute specification.
+    Fetches available instance types from the API and maps compute specs to instance type IDs.
+    Only considers single-node instances (node_count == 1).
+    """
+    # step 1: fetch the instance types from the API
+    instance_types = remote_provider.api.get_instance_types()
+    # step 2: sort them into two different dictionaries, excluding multi-node instances:
+    cpu_instance_types = {
+        it.id: it for it in instance_types if it.gpu_count == 0 and it.node_count == 1
+    }
+    gpu_instance_types = {
+        it.id: it for it in instance_types if it.gpu_count > 0 and it.node_count == 1
+    }
+    # step 3: if compute is cpu, find the smallest such cpu that matches the compute request
+    if not compute.accelerator or compute.accelerator.accelerator is None:
+        compute_as_truss_config = compute.to_truss_config()
+        smallest_cpu_instance_type = None
+        for it in cpu_instance_types.values():
+            if (
+                it.millicpu_limit / 1000 >= compute.cpu_count
+                and it.memory_limit >= compute_as_truss_config.memory_in_bytes
+            ):
+                if (
+                    smallest_cpu_instance_type is None
+                    or it.millicpu_limit < smallest_cpu_instance_type.millicpu_limit
+                ):
+                    smallest_cpu_instance_type = it
+        if not smallest_cpu_instance_type:
+            raise ValueError(
+                f"Unable to find single-node instance type for {compute.cpu_count} CPU and {compute.memory} memory. Reach out to Baseten for support if this persists."
+            )
+        return smallest_cpu_instance_type.id
+    # step 4: if compute is gpu, find the smallest such gpu by instance type
+    else:
+        assert compute.accelerator.accelerator is not None
+        compute_as_truss_config = compute.to_truss_config()
+        smallest_gpu_instance_type = None
+        for it in gpu_instance_types.values():
+            if (
+                it.gpu_type == compute.accelerator.accelerator.value
+                and it.gpu_count >= compute.accelerator.count
+            ):
+                if (
+                    smallest_gpu_instance_type is None
+                    or it.gpu_count < smallest_gpu_instance_type.gpu_count
+                ):
+                    smallest_gpu_instance_type = it
+        if not smallest_gpu_instance_type:
+            raise ValueError(
+                f"Unable to find single-node instance type for {compute.accelerator}:{compute.accelerator.count}. Reach out to Baseten for support if this persists."
+            )
+        return smallest_gpu_instance_type.id
+
 
 def _validate_base_model_id(
     base_model_id: Optional[str], model_weight_format: ModelWeightsFormat
@@ -93,18 +240,12 @@ def _get_model_name(
         else ""
     )
 
-    model_name = inquirer.text(
+    return inquirer.text(
         message=f"Enter the model name for your {model_weight_format.value} model.",
         validate=lambda s: s and s.strip(),
         default=default,
     ).execute()
 
-    if model_weight_format == ModelWeightsFormat.FULL:
-        model_name += "-vLLM-Full"
-    elif model_weight_format == ModelWeightsFormat.LORA:
-        model_name += "-vLLM-LORA"
-    return model_name
-
 
 def _hydrate_deploy_config(
     deploy_config: DeployCheckpointsConfig,
@@ -123,53 +264,18 @@ def _hydrate_deploy_config(
     else:
         model_name = _get_model_name(model_weight_format, base_model_id)
 
-    compute = _ensure_compute_spec(deploy_config.compute)
+    compute = _ensure_compute_spec(deploy_config.compute, remote_provider)
 
     runtime = _ensure_runtime_config(deploy_config.runtime)
-    deployment_name = _ensure_deployment_name(
-        deploy_config.deployment_name, checkpoint_details.checkpoints
-    )
 
     return DeployCheckpointsConfigComplete(
         checkpoint_details=checkpoint_details,
         model_name=model_name,
-        deployment_name=deployment_name,
         runtime=runtime,
         compute=compute,
-        model_weight_format=model_weight_format.to_truss_config(),  # type: ignore[attr-defined]
     )
 
 
-def _ensure_deployment_name(
-    deploy_config_deployment_name: Optional[str], checkpoints: List[Checkpoint]
-) -> str:
-    if deploy_config_deployment_name:
-        return deploy_config_deployment_name
-
-    default_deployment_name = "checkpoint"
-
-    if checkpoints and checkpoints[0].paths:
-        first_checkpoint_name = checkpoints[0].paths[0].strip("/").split("/")[-1]
-
-        if ALLOWED_DEPLOYMENT_NAMES.match(first_checkpoint_name):
-            # Allow autoincrementing if the checkpoint matches both regexes
-            if (
-                CHECKPOINT_PATTERN.match(first_checkpoint_name)
-                and len(checkpoints) == 1
-            ):
-                return first_checkpoint_name
-
-    # If no valid autoincrementing checkpoint name is found, prompt the user
-    deployment_name = inquirer.text(
-        message="Enter the deployment name.", default=default_deployment_name
-    ).execute()
-
-    if not deployment_name:
-        raise click.UsageError("Deployment name is required.")
-
-    return deployment_name
-
-
 def hydrate_checkpoint(
     job_id: str, checkpoint_id: str, checkpoint: dict, checkpoint_type: str
 ) -> Checkpoint:
@@ -190,26 +296,6 @@ def hydrate_checkpoint(
         )
 
 
-def _render_truss_config_for_checkpoint_deployment(
-    checkpoint_deploy: DeployCheckpointsConfigComplete,
-) -> truss_config.TrussConfig:
-    """
-    Render truss config for checkpoint deployment.
-    Currently supports LoRA checkpoints via vLLM, but can be extended for other formats.
-    """
-    # Delegate to specific rendering function based on model weight format
-    if checkpoint_deploy.model_weight_format == ModelWeightsFormat.LORA:
-        return render_vllm_lora_truss_config(checkpoint_deploy)
-    elif checkpoint_deploy.model_weight_format == ModelWeightsFormat.FULL:
-        return render_vllm_full_truss_config(checkpoint_deploy)
-    elif checkpoint_deploy.model_weight_format == ModelWeightsFormat.WHISPER:
-        return render_vllm_whisper_truss_config(checkpoint_deploy)
-    else:
-        raise ValueError(
-            f"Unsupported model weight format: {checkpoint_deploy.model_weight_format}. Please upgrade to the latest Truss version to access the latest supported formats. Contact Baseten if you would like us to support additional formats."
-        )
-
-
 def _ensure_checkpoint_details(
     remote_provider: BasetenRemote,
     checkpoint_details: Optional[CheckpointList],
@@ -217,6 +303,7 @@ def _ensure_checkpoint_details(
     job_id: Optional[str],
 ) -> CheckpointList:
     if checkpoint_details and checkpoint_details.checkpoints:
+        # TODO: check here
         return _process_user_provided_checkpoints(checkpoint_details, remote_provider)
     else:
         return _prompt_user_for_checkpoint_details(
@@ -309,31 +396,77 @@ def _select_multiple_checkpoints(checkpoint_id_options: List[str]) -> List[str]:
     return checkpoint_ids
 
 
-def _ensure_compute_spec(compute: Optional[Compute]) -> Compute:
+def _ensure_compute_spec(
+    compute: Optional[Compute], remote_provider: BasetenRemote
+) -> Compute:
     if not compute:
         compute = Compute(cpu_count=0, memory="0Mi")
-    compute.accelerator = _get_accelerator_if_specified(compute.accelerator)
+    compute = _get_accelerator_if_specified(compute, remote_provider)
     return compute
 
 
 def _get_accelerator_if_specified(
-    user_input: Optional[truss_config.AcceleratorSpec],
-) -> Optional[truss_config.AcceleratorSpec]:
-    if user_input:
+    user_input: Optional[Compute], remote_provider: BasetenRemote
+) -> Compute:
+    if user_input and user_input.accelerator:
         return user_input
+
+    # Fetch available instance types to get valid GPU options
+    instance_types = remote_provider.api.get_instance_types()
+
+    # Extract unique accelerator types from instance types
+    accelerator_options = set()
+    for it in instance_types:
+        if it.gpu_type and it.gpu_count > 0:
+            accelerator_options.add(it.gpu_type)
+
+    # Convert to sorted list and add CPU option
+    choices = sorted(list(accelerator_options)) + [None]
+
+    if not choices or choices == [None]:
+        console.print("No GPU instance types available, using CPU", style="yellow")
+        return Compute(cpu_count=0, memory="0Mi", accelerator=None)
+
     # prompt user for accelerator
     gpu_type = inquirer.select(
         message="Select the GPU type to use for deployment. Select None for CPU.",
-        choices=[x.value for x in truss_config.Accelerator] + [None],
+        choices=choices,
     ).execute()
+
     if gpu_type is None:
-        return None
-    count = inquirer.text(
-        message="Enter the number of accelerators to use for deployment.",
-        default="1",
-        validate=lambda x: x.isdigit() and int(x) > 0 and int(x) <= 8,
-    ).execute()
-    return truss_config.AcceleratorSpec(accelerator=gpu_type, count=int(count))
+        return Compute(cpu_count=0, memory="0Mi", accelerator=None)
+
+    # Get available counts for the selected GPU type
+    available_counts = set()
+    for it in instance_types:
+        if it.gpu_type == gpu_type and it.gpu_count > 0:
+            available_counts.add(it.gpu_count)
+    if not available_counts:
+        raise ValueError(
+            f"No available counts for {gpu_type}. Reach out to Baseten for support if this persists."
+        )
+
+    if available_counts:
+        count_choices = sorted(list(available_counts))
+        count = inquirer.select(
+            message=f"Select the number of {gpu_type} GPUs to use for deployment.",
+            choices=count_choices,
+            default=str(count_choices[0]),
+        ).execute()
+    else:
+        count = inquirer.text(
+            message=f"Enter the number of {gpu_type} accelerators to use for deployment.",
+            default="1",
+            validate=lambda x: x.isdigit() and int(x) > 0 and int(x) <= 8,
+        ).execute()
+
+    return Compute(
+        cpu_count=0,
+        memory="0Mi",
+        accelerator=truss_config.AcceleratorSpec(
+            accelerator=gpu_type.replace("-", "_"), count=int(count)
+        ),
+    )
 
 
 def _get_base_model_id(user_input: Optional[str], checkpoint: dict) -> Optional[str]:

truss/cli/train/deploy_checkpoints/deploy_checkpoints_helpers.py

@@ -1,52 +1 @@
-import os
-from pathlib import Path
-
-from truss.base import truss_config
-from truss.cli.train.types import DeployCheckpointsConfigComplete
-from truss_train.definitions import ModelWeightsFormat, SecretReference
-
-START_COMMAND_ENVVAR_NAME = "BT_DOCKER_SERVER_START_CMD"
-
-
-def setup_base_truss_config(
-    checkpoint_deploy: DeployCheckpointsConfigComplete,
-) -> truss_config.TrussConfig:
-    """Set up the base truss config with common properties."""
-    truss_deploy_config = None
-    truss_base_file = (
-        "deploy_from_checkpoint_config_whisper.yml"
-        if checkpoint_deploy.model_weight_format == ModelWeightsFormat.WHISPER
-        else "deploy_from_checkpoint_config.yml"
-    )
-    truss_deploy_config = truss_config.TrussConfig.from_yaml(
-        Path(os.path.dirname(__file__), "..", truss_base_file)
-    )
-    if not truss_deploy_config.docker_server:
-        raise ValueError(
-            "Unexpected checkpoint deployment config: missing docker_server"
-        )
-
-    truss_deploy_config.model_name = checkpoint_deploy.model_name
-    truss_deploy_config.training_checkpoints = (
-        checkpoint_deploy.checkpoint_details.to_truss_config()
-    )
-    truss_deploy_config.resources = checkpoint_deploy.compute.to_truss_config()
-
-    return truss_deploy_config
-
-
-def setup_environment_variables_and_secrets(
-    truss_deploy_config: truss_config.TrussConfig,
-    checkpoint_deploy: DeployCheckpointsConfigComplete,
-) -> str:
-    """Set up environment variables and secrets, return start command envvars string."""
-    start_command_envvars = ""
-
-    for key, value in checkpoint_deploy.runtime.environment_variables.items():
-        if isinstance(value, SecretReference):
-            truss_deploy_config.secrets[value.name] = "set token in baseten workspace"
-            start_command_envvars = f"{key}=$(cat /secrets/{value.name})"
-        else:
-            truss_deploy_config.environment_variables[key] = value
-
-    return start_command_envvars
+# This file is kept for potential future use but currently contains no active code

truss/cli/train/deploy_checkpoints/deploy_full_checkpoints.py

@@ -1,94 +1,9 @@
-from pathlib import Path
-
-from jinja2 import Template
-
-from truss.base import truss_config
-from truss.cli.train.deploy_checkpoints.deploy_checkpoints_helpers import (
-    START_COMMAND_ENVVAR_NAME,
-)
-from truss.cli.train.types import DeployCheckpointsConfigComplete
 from truss_train.definitions import FullCheckpoint
 
-from .deploy_checkpoints_helpers import (
-    setup_base_truss_config,
-    setup_environment_variables_and_secrets,
-)
-
-# NB(aghilan): Transformers was recently changed to save a chat_template.jinja file instead of inside the tokenizer_config.json file.
-# Old Models will not have this file, so we check for it and use it if it exists.
-# vLLM will not automatically resolve the chat_template.jinja file, so we need to pass it to the start command.
-# This logic is needed for any models trained using Transformers v4.51.3 or later
-VLLM_FULL_START_COMMAND = Template(
-    "sh -c '{% if envvars %}{{ envvars }} {% endif %}"
-    'HF_TOKEN="$$(cat /secrets/hf_access_token)" && export HF_TOKEN && '
-    "if [ -f {{ model_path }}/chat_template.jinja ]; then "
-    "  vllm serve {{ model_path }} --chat-template {{ model_path }}/chat_template.jinja "
-    "  --port 8000 --tensor-parallel-size {{ specify_tensor_parallelism }} --dtype bfloat16; "
-    "else "
-    "  vllm serve {{ model_path }} --port 8000 --tensor-parallel-size {{ specify_tensor_parallelism }} --dtype bfloat16; "
-    "fi'"
-)
-
-
-def render_vllm_full_truss_config(
-    checkpoint_deploy: DeployCheckpointsConfigComplete,
-) -> truss_config.TrussConfig:
-    """Render truss config specifically for full checkpoints using vLLM."""
-    truss_deploy_config = setup_base_truss_config(checkpoint_deploy)
-
-    start_command_envvars = setup_environment_variables_and_secrets(
-        truss_deploy_config, checkpoint_deploy
-    )
-
-    checkpoint_str = build_full_checkpoint_string(truss_deploy_config)
-
-    accelerator = checkpoint_deploy.compute.accelerator
-
-    start_command_args = {
-        "model_path": checkpoint_str,
-        "envvars": start_command_envvars,
-        "specify_tensor_parallelism": accelerator.count if accelerator else 1,
-    }
-    # Note: we set the start command as an environment variable in supervisord config.
-    # This is so that we don't have to change the supervisord config when the start command changes.
-    # Our goal is to reduce the number of times we need to rebuild the image, and allow us to deploy faster.
-    start_command = VLLM_FULL_START_COMMAND.render(**start_command_args)
-    truss_deploy_config.environment_variables[START_COMMAND_ENVVAR_NAME] = start_command
-    # Note: supervisord uses the convention %(ENV_VAR_NAME)s to access environment variable VAR_NAME
-    truss_deploy_config.docker_server.start_command = (  # type: ignore[union-attr]
-        f"%(ENV_{START_COMMAND_ENVVAR_NAME})s"
-    )
-
-    return truss_deploy_config
-
 
 def hydrate_full_checkpoint(
     job_id: str, checkpoint_id: str, checkpoint: dict
 ) -> FullCheckpoint:
     """Create a Checkpoint object for full model weights."""
     # NOTE: Slash at the end is important since it means the checkpoint is a directory
-    paths = [f"rank-0/{checkpoint_id}/"]
-    return FullCheckpoint(training_job_id=job_id, paths=paths)
-
-
-def build_full_checkpoint_string(truss_deploy_config) -> str:
-    """Build checkpoint string from artifact references for full checkpoints.
-
-    Args:
-        truss_deploy_config: The truss deploy configuration containing training checkpoints.
-
-    Returns:
-        A space-separated string of checkpoint paths.
-    """
-    checkpoint_parts = []
-    for (
-        truss_checkpoint
-    ) in truss_deploy_config.training_checkpoints.artifact_references:  # type: ignore
-        ckpt_path = Path(
-            truss_deploy_config.training_checkpoints.download_folder,  # type: ignore
-            truss_checkpoint.training_job_id,
-            truss_checkpoint.paths[0],
-        )
-        checkpoint_parts.append(str(ckpt_path))
-
-    return " ".join(checkpoint_parts)
+    return FullCheckpoint(training_job_id=job_id, checkpoint_name=checkpoint_id)