# 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.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]
[docs]
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
[docs]
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:
logger.error(f"Failed to clone repository: {repo_url}")
logger.error(f"Error output:\n{e}")
raise
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