python/run-with-viztracer

#!/usr/bin/env bash

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.
#

set -o pipefail
set -e

FWDIR="$(cd "`dirname $0`"; pwd)"

# Function to display help message
function usage {
  cat <<EOF
Usage:
  $(basename $0) your_original_commands

To view the profiling results, run:
  vizviewer pyspark_*.json

Environment:
  If SPARK_VIZTRACER_OUTPUT_DIR is set, the output will be saved to the directory.
  Otherwise, it will be saved to the current directory.

Requirements:
  - viztracer must be installed (pip install viztracer)
EOF
  exit 0
}

# Check for help flags or no arguments
if [ $# -eq 0 ] || [ "$1" = "-h" ] || [ "$1" = "--help" ]; then
  usage
fi

if ! hash viztracer 2>/dev/null; then
  echo "Error: viztracer is not installed. Please install it using 'pip install viztracer'."
  exit 1
fi

export PYTHONPATH="$FWDIR/conf_viztracer:$PYTHONPATH"
export SPARK_CONF_DIR="$FWDIR/conf_viztracer"
export SPARK_VIZTRACER_OUTPUT_DIR="${SPARK_VIZTRACER_OUTPUT_DIR:-"$(pwd)"}"

exec "$@"