mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-03-16 21:27:39 +00:00
06ad66b3e4
[CK Tools] Auto-enable unbuffered output for Python commands (#4265) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ck-docker exec and ck-exec now automatically detect Python commands and set PYTHONUNBUFFERED=1 to enable live output streaming. This eliminates the need to manually set the environment variable when running Python scripts that print progress updates. The detection matches python, python3, or any .py file argument. This helps in watching live terminal output when a python script is running inside the container.
121 lines
3.1 KiB
Bash
Executable File
121 lines
3.1 KiB
Bash
Executable File
#!/bin/bash
|
|
# Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
# CK Exec - Execute arbitrary commands in Docker container
|
|
|
|
set -e
|
|
set -o pipefail
|
|
|
|
# Find script directory and load common utilities
|
|
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
|
source "${SCRIPT_DIR}/common.sh"
|
|
|
|
# Initialize configuration
|
|
PROJECT_ROOT=$(get_project_root "${SCRIPT_DIR}")
|
|
CONTAINER_NAME=$(get_container_name "${PROJECT_ROOT}")
|
|
|
|
# Help message
|
|
show_help() {
|
|
cat << EOF
|
|
CK Exec - Execute arbitrary commands in Docker container
|
|
|
|
Usage: ck-exec [options] <command> [args...]
|
|
|
|
Options:
|
|
-h, --help Show this help message
|
|
--name <name> Specify container name
|
|
-w <dir> Working directory (default: /workspace)
|
|
-i, --interactive Interactive mode (allocate TTY)
|
|
|
|
Arguments:
|
|
command Command to execute (required)
|
|
args Arguments to the command
|
|
|
|
Environment:
|
|
CK_CONTAINER_NAME - Override default container name
|
|
|
|
Examples:
|
|
ck-exec rocm-smi # Run rocm-smi
|
|
ck-exec rocminfo # Run rocminfo
|
|
ck-exec ls -la build/bin # List build binaries
|
|
ck-exec -w /workspace/build ninja -t commands # Run ninja commands
|
|
ck-exec --interactive python3 # Interactive Python session
|
|
|
|
Common Commands:
|
|
ck-exec rocm-smi # Check GPU status
|
|
ck-exec rocminfo \| grep gfx # Check GPU architecture
|
|
ck-exec hipcc --version # Check HIP compiler version
|
|
ck-exec cmake --version # Check CMake version
|
|
ck-exec ninja -C build -t targets # List all build targets
|
|
|
|
EOF
|
|
}
|
|
|
|
# Parse arguments
|
|
workdir="/workspace"
|
|
interactive=false
|
|
command_args=()
|
|
|
|
while [[ $# -gt 0 ]]; do
|
|
case $1 in
|
|
-h|--help)
|
|
show_help
|
|
exit 0
|
|
;;
|
|
--name)
|
|
CONTAINER_NAME="$2"
|
|
shift 2
|
|
;;
|
|
-w)
|
|
workdir="$2"
|
|
shift 2
|
|
;;
|
|
-i|--interactive)
|
|
interactive=true
|
|
shift
|
|
;;
|
|
*)
|
|
command_args+=("$1")
|
|
shift
|
|
;;
|
|
esac
|
|
done
|
|
|
|
# Validate command
|
|
if [ ${#command_args[@]} -eq 0 ]; then
|
|
echo "Error: command required"
|
|
echo ""
|
|
show_help
|
|
exit 1
|
|
fi
|
|
|
|
# Ensure container is running
|
|
if ! container_is_running "${CONTAINER_NAME}"; then
|
|
echo "Container '${CONTAINER_NAME}' not running. Starting..."
|
|
"${SCRIPT_DIR}/ck-start" "${CONTAINER_NAME}"
|
|
echo ""
|
|
fi
|
|
|
|
# Build command string
|
|
cmd_string=""
|
|
for arg in "${command_args[@]}"; do
|
|
cmd_string="${cmd_string} $(printf '%q' "$arg")"
|
|
done
|
|
|
|
# Auto-detect Python commands and enable unbuffered output for live streaming
|
|
env_flags=""
|
|
for arg in "${command_args[@]}"; do
|
|
if [[ "$arg" == "python" || "$arg" == "python3" || "$arg" == *.py ]]; then
|
|
env_flags="-e PYTHONUNBUFFERED=1"
|
|
break
|
|
fi
|
|
done
|
|
|
|
# Execute command
|
|
if [ "$interactive" = true ]; then
|
|
docker exec -it ${env_flags} -w "${workdir}" "${CONTAINER_NAME}" bash -c "${cmd_string}"
|
|
else
|
|
docker exec ${env_flags} -w "${workdir}" "${CONTAINER_NAME}" bash -c "${cmd_string}"
|
|
fi
|