utils.py

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Utils module."""

from __future__ import absolute_import

import os
import json
import subprocess
import tempfile
from pathlib import Path

from datetime import datetime
from typing import Literal, Any

from sagemaker.core.git_utils import _redact_credentials_from_url
from sagemaker.core.shapes import Unassigned
from sagemaker.core.modules import logger


def _is_valid_s3_uri(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
    """Check if the path is a valid S3 URI.

    This method checks if the path is a valid S3 URI. If the path_type is specified,
    it will also check if the path is a file or a directory.
    This method does not check if the S3 bucket or object exists.

    Args:
        path (str): S3 URI to validate
        path_type (Optional(Literal["File", "Directory", "Any"])): The type of the path to validate.
            Defaults to "Any".

    Returns:
        bool: True if the path is a valid S3 URI, False otherwise
    """
    # Check if the path is a valid S3 URI
    if not path.startswith("s3://"):
        return False

    if path_type == "File":
        # If it's a file, it should not end with a slash
        return not path.endswith("/")
    if path_type == "Directory":
        # If it's a directory, it should end with a slash
        return path.endswith("/")

    return path_type == "Any"


def _is_valid_path(path: str, path_type: Literal["File", "Directory", "Any"] = "Any") -> bool:
    """Check if the path is a valid local path.

    Args:
        path (str): Local path to validate
        path_type (Optional(Literal["File", "Directory", "Any"])): The type of the path to validate.
            Defaults to "Any".

    Returns:
        bool: True if the path is a valid local path, False otherwise
    """
    if not os.path.exists(path):
        return False

    if path_type == "File":
        return os.path.isfile(path)
    if path_type == "Directory":
        return os.path.isdir(path)

    return path_type == "Any"


def _get_unique_name(base, max_length=63):
    """Generate a unique name based on the base name.

    This method generates a unique name based on the base name.
    The unique name is generated by appending the current timestamp
    to the base name.

    Args:
        base (str): The base name to use
        max_length (int): The maximum length of the unique name. Defaults to 63.

    Returns:
        str: The unique name
    """
    current_time = datetime.now().strftime("%Y%m%d%H%M%S")
    base = base.replace("_", "-")
    unique_name = f"{base}-{current_time}"
    unique_name = unique_name[:max_length]  # Truncate to max_length
    return unique_name


def _get_repo_name_from_image(image: str) -> str:
    """Get the repository name from the image URI.

    Example:
    ``` python
    _get_repo_name_from_image("123456789012.dkr.ecr.us-west-2.amazonaws.com/my-repo:latest")
    # Returns "my-repo"
    ```

    Args:
        image (str): The image URI

    Returns:
        str: The repository name
    """
    return image.split("/")[-1].split(":")[0]


def convert_unassigned_to_none(instance) -> Any:
    """Convert Unassigned values to None for any instance."""
    for name, value in instance.__dict__.items():
        if isinstance(value, Unassigned):
            setattr(instance, name, None)
    return instance


def safe_serialize(data):
    """Serialize the data without wrapping strings in quotes.

    This function handles the following cases:
    1. If `data` is a string, it returns the string as-is without wrapping in quotes.
    2. If `data` is serializable (e.g., a dictionary, list, int, float), it returns
       the JSON-encoded string using `json.dumps()`.
    3. If `data` cannot be serialized (e.g., a custom object), it returns the string
       representation of the data using `str(data)`.

    Args:
        data (Any): The data to serialize.

    Returns:
        str: The serialized JSON-compatible string or the string representation of the input.
    """
    if isinstance(data, str):
        return data
    try:
        return json.dumps(data)
    except TypeError:
        return str(data)


def _run_clone_command_silent(repo_url, dest_dir):
    """Run the 'git clone' command with the repo url and the directory to clone the repo into.

    Args:
        repo_url (str): Git repo url to be cloned.
        dest_dir: (str): Local path where the repo should be cloned into.

    Raises:
        CalledProcessError: If failed to clone git repo.
    """
    my_env = os.environ.copy()
    if repo_url.startswith("https://"):
        try:
            my_env["GIT_TERMINAL_PROMPT"] = "0"
            subprocess.check_call(
                ["git", "clone", repo_url, dest_dir],
                env=my_env,
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
            )
        except subprocess.CalledProcessError as e:
            safe_url = _redact_credentials_from_url(repo_url)
            error = subprocess.CalledProcessError(
                e.returncode,
                ["git", "clone", safe_url, dest_dir],
                output=e.output,
                stderr=e.stderr,
            )
            logger.error(f"Failed to clone repository: {safe_url}")
            logger.error(f"Error output:\n{error}")
            raise error from None
    elif repo_url.startswith("git@") or repo_url.startswith("ssh://"):
        try:
            with tempfile.TemporaryDirectory() as tmp_dir:
                custom_ssh_executable = Path(tmp_dir) / "ssh_batch"
                with open(custom_ssh_executable, "w") as pipe:
                    print("#!/bin/sh", file=pipe)
                    print("ssh -oBatchMode=yes $@", file=pipe)
                os.chmod(custom_ssh_executable, 0o511)
                my_env["GIT_SSH"] = str(custom_ssh_executable)
                subprocess.check_call(
                    ["git", "clone", repo_url, dest_dir],
                    env=my_env,
                    stdout=subprocess.DEVNULL,
                    stderr=subprocess.DEVNULL,
                )
        except subprocess.CalledProcessError as e:
            del my_env["GIT_SSH"]
            logger.error(f"Failed to clone repository: {repo_url}")
            logger.error(f"Error output:\n{e}")
            raise


def validate_instance_preferences(compute) -> None:
    """Client-side validation for Compute.instance_preferences (server remains the source of truth).

    - instance_preferences is mutually exclusive with the classic
      single-cluster fields instance_type / instance_groups /
      instance_placement_config, and with managed spot training.
    - Instance types must not repeat across preferences.
    - Count mode: exactly one of the top-level instance_count (applies to
      whichever preference wins) with no per-preference instance_count, or an
      instance_count on EVERY element with the top-level unset. Both-set,
      partial, and neither are rejected.
    - Training plans: the top-level (whole-job) training_plan_arn is mutually
      exclusive with per-preference training_plan_arns.

    List-size limits (max preferences, max plans per preference) are
    deliberately NOT enforced client-side: they are server-side configurable,
    so raising them must not require a new SDK release.

    No-op when instance_preferences is not set.
    """
    preferences = getattr(compute, "instance_preferences", None)
    if not preferences or isinstance(preferences, Unassigned):
        return

    def _value(obj, field):
        value = getattr(obj, field, None)
        if value is None or isinstance(value, Unassigned):
            return None
        return value

    for field in ("instance_type", "instance_groups", "instance_placement_config"):
        if _value(compute, field) is not None:
            raise ValueError(
                f"instance_preferences is mutually exclusive with {field}; "
                "specify either a single fixed cluster or instance_preferences, not both."
            )
    # Spot is a bool: only an explicit True conflicts (False/None is the default).
    if _value(compute, "enable_managed_spot_training") is True:
        raise ValueError(
            "instance_preferences is mutually exclusive with managed spot training "
            "(enable_managed_spot_training=True)."
        )

    instance_types = [_value(p, "instance_type") for p in preferences]
    duplicates = sorted(
        {t for t in instance_types if t is not None and instance_types.count(t) > 1}
    )
    if duplicates:
        raise ValueError(
            "instance_preferences must not contain duplicate instance types: " f"{duplicates}."
        )

    per_preference_counts = [_value(p, "instance_count") is not None for p in preferences]
    if _value(compute, "instance_count") is not None:
        if any(per_preference_counts):
            raise ValueError(
                "The top-level instance_count and per-preference instance_count are "
                "mutually exclusive; set the top-level instance_count (applies to "
                "whichever preference wins) or an instance_count on every element of "
                "instance_preferences, not both."
            )
    elif not all(per_preference_counts):
        raise ValueError(
            "When the top-level instance_count is not set, every element of "
            "instance_preferences must set its own instance_count."
        )

    per_preference_plans = [_value(p, "training_plan_arns") for p in preferences]
    if _value(compute, "training_plan_arn") is not None and any(per_preference_plans):
        raise ValueError(
            "The top-level (whole-job) training_plan_arn and per-preference "
            "training_plan_arns are mutually exclusive; set one or the other, not both."
        )