chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
License Check / license-check (push) Failing after 2s

This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
+62
View File
@@ -0,0 +1,62 @@
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:__subpackages__",
],
licenses = ["notice"],
)
py_library(
name = "profiler",
srcs = ["__init__.py"],
strict_deps = True,
deps = [
"//tensorflow/core/profiler:profiler_analysis_proto_py",
"//tensorflow/python/util:all_util",
"@tsl//tsl/profiler/protobuf:trace_events_proto_py",
],
)
py_library(
name = "profiler_analysis_pb2_grpc",
srcs = ["profiler_analysis_pb2_grpc.py"],
strict_deps = True,
visibility = ["//visibility:public"],
deps = ["//tensorflow/core/profiler:profiler_analysis_proto_py"],
)
py_library(
name = "capture_tpu_profile_lib",
srcs = [
"capture_tpu_profile.py",
"version.py",
],
strict_deps = True,
deps = [
"//tensorflow/python/distribute/cluster_resolver:tpu_cluster_resolver_py",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:versions",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/profiler:profiler_client",
"//tensorflow/python/profiler:profiler_v2",
"@absl_py//absl:app",
"@absl_py//absl/flags",
"@pypi//packaging",
],
)
py_binary(
name = "capture_tpu_profile_bin",
srcs = ["capture_tpu_profile.py"],
main = "capture_tpu_profile.py",
strict_deps = True,
deps = [
":capture_tpu_profile_lib",
"@absl_py//absl/flags",
"@pypi//packaging",
],
)
@@ -0,0 +1,27 @@
# Copyright 2017 The TensorFlow Authors. 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.
# 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.
# =============================================================================
"""Classes for TPU trace events."""
# pylint: disable=wildcard-import,unused-import
from tensorflow.core.profiler.profiler_analysis_pb2 import *
from tensorflow.python.util.all_util import remove_undocumented
from tsl.profiler.protobuf.trace_events_pb2 import *
# pylint: enable=wildcard-import,unused-import
_allowed_symbols = ['Trace', 'Resource', 'Device', 'TraceEvent']
remove_undocumented(__name__, _allowed_symbols)
@@ -0,0 +1,206 @@
# Copyright 2019 The TensorFlow Authors. 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.
# 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.
# =============================================================================
"""Cloud TPU profiler client."""
import os
import sys
from absl import app
from absl import flags
from packaging.version import Version
from tensorflow.python.distribute.cluster_resolver import tpu_cluster_resolver as resolver
from tensorflow.python.framework import errors
from tensorflow.python.framework import versions
from tensorflow.python.platform import gfile
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.profiler import profiler_client
from tensorflow.python.profiler import profiler_v2 as profiler
from tensorflow.python.tpu.profiler import version as profiler_version
FLAGS = flags.FLAGS
# Cloud TPU Cluster Resolvers
flags.DEFINE_string(
'gcp_project', None,
'Project name for the Cloud TPU-enabled project. If not specified, we '
'will attempt to automatically detect the GCE project from metadata.')
flags.DEFINE_string(
'tpu_zone',
None,
help='GCE zone where the Cloud TPU is located in. If not specified, we '
'will attempt to automatically detect the GCE project from metadata.')
flags.DEFINE_string(
'tpu', None, 'Name of the Cloud TPU for Cluster Resolvers. You must '
'specify either this flag or --service_addr.')
# Tool specific parameters
flags.DEFINE_string(
'service_addr', None, 'Address of TPU profiler service e.g. '
'localhost:8466, you must specify either this flag or --tpu.')
flags.DEFINE_string(
'workers_list', None, 'The list of worker TPUs that we are about to profile'
' e.g. 10.0.1.2:8466, 10.0.1.3:8466. You can specify this flag with --tpu '
'or --service_addr to profile a subset of tpu nodes. You can also use only'
'--tpu and leave this flag unspecified to profile all the tpus.')
flags.DEFINE_string(
'logdir', None, 'Path of TensorBoard log directory e.g. /tmp/tb_log, '
'gs://tb_bucket')
flags.DEFINE_integer('duration_ms', 0,
'Duration of tracing or monitoring in ms.')
flags.DEFINE_integer(
'num_tracing_attempts', 3, 'Automatically retry N times when no trace '
'event is collected.')
flags.DEFINE_boolean('include_dataset_ops', True, 'Deprecated.')
flags.DEFINE_integer(
'host_tracer_level', 2, 'Adjust host tracer level to control the verbosity '
' of the TraceMe event being collected.')
# Monitoring parameters
flags.DEFINE_integer(
'monitoring_level', 0, 'Choose a monitoring level between '
'1 and 2 to monitor your TPU job continuously. Level 2 is more verbose than'
' level 1 and shows more metrics.')
flags.DEFINE_integer(
'num_queries', 100,
'This script will run monitoring for num_queries before it stops.')
flags.DEFINE_boolean('display_timestamp', True, 'Deprecated.')
def get_workers_list(cluster_resolver):
"""Returns a comma separated list of TPU worker host:port pairs.
Gets cluster_spec from cluster_resolver. Use the worker's task indices to
obtain and return a list of host:port pairs.
Args:
cluster_resolver: TensorFlow TPUClusterResolver instance.
Returns:
A string of comma separated list of host:port pairs. For example:
'10.2.0.1:8466,10.2.0.2:8466,10.2.0.3:8466,10.2.0.4:8466'
Raises:
UnavailableError: cluster_resolver doesn't contain a valid cluster_spec.
"""
worker_job_name = 'worker'
cluster_spec = cluster_resolver.cluster_spec()
if not cluster_spec:
raise errors.UnavailableError(
'None', 'None',
'Cluster spec not found, your client must run in GCE environment.')
task_indices = cluster_spec.task_indices(worker_job_name)
workers_list = [
cluster_spec.task_address(worker_job_name, i).replace(':8470', ':8466')
for i in task_indices
]
return ','.join(workers_list)
def monitoring_helper(service_addr, duration_ms, monitoring_level, num_queries):
"""Helper function to print monitoring results.
Helper function to print monitoring results for num_queries times.
Args:
service_addr: Address of the TPU profiler service.
duration_ms: Duration of one monitoring sample in milliseconds.
monitoring_level: An integer between 1 and 2. Level 2 is more verbose than
level 1 and shows more metrics.
num_queries: Number of monitoring samples to collect.
"""
if monitoring_level <= 0 or monitoring_level > 2:
sys.exit('Please choose a monitoring level between 1 and 2.')
for query in range(0, num_queries):
res = profiler_client.monitor(service_addr, duration_ms, monitoring_level)
print('Cloud TPU Monitoring Results (Sample ', query, '):\n\n', res)
def run_main():
app.run(main)
def main(unused_argv=None):
logging.set_verbosity(logging.INFO)
tf_version = versions.__version__
print('TensorFlow version %s detected' % tf_version)
print('Welcome to the Cloud TPU Profiler v%s' % profiler_version.__version__)
if Version(tf_version) < Version('2.2.0'):
sys.exit('You must install tensorflow >= 2.2.0 to use this plugin.')
if not FLAGS.service_addr and not FLAGS.tpu:
sys.exit('You must specify either --service_addr or --tpu.')
tpu_cluster_resolver = None
if FLAGS.service_addr:
if FLAGS.tpu:
logging.warn('Both --service_addr and --tpu are set. Ignoring '
'--tpu and using --service_addr.')
service_addr = FLAGS.service_addr
else:
try:
tpu_cluster_resolver = (
resolver.TPUClusterResolver([FLAGS.tpu],
zone=FLAGS.tpu_zone,
project=FLAGS.gcp_project))
service_addr = tpu_cluster_resolver.get_master()
except (ValueError, TypeError):
sys.exit('Failed to find TPU %s in zone %s project %s. You may use '
'--tpu_zone and --gcp_project to specify the zone and project of'
' your TPU.' % (FLAGS.tpu, FLAGS.tpu_zone, FLAGS.gcp_project))
service_addr = service_addr.replace('grpc://', '').replace(':8470', ':8466')
workers_list = ''
if FLAGS.workers_list is not None:
workers_list = FLAGS.workers_list
elif tpu_cluster_resolver is not None:
workers_list = get_workers_list(tpu_cluster_resolver)
# If profiling duration was not set by user or set to a non-positive value,
# we set it to a default value of 1000ms.
duration_ms = FLAGS.duration_ms if FLAGS.duration_ms > 0 else 1000
if FLAGS.monitoring_level > 0:
print('Since monitoring level is provided, profile', service_addr, ' for ',
FLAGS.duration_ms, ' ms and show metrics for ', FLAGS.num_queries,
' time(s).')
monitoring_helper(service_addr, duration_ms, FLAGS.monitoring_level,
FLAGS.num_queries)
else:
if not FLAGS.logdir:
sys.exit('You must specify either --logdir or --monitoring_level.')
if not gfile.Exists(FLAGS.logdir):
gfile.MakeDirs(FLAGS.logdir)
try:
if Version(tf_version) < Version('2.3.0'):
profiler_client.trace(service_addr, os.path.expanduser(FLAGS.logdir),
duration_ms, workers_list,
FLAGS.num_tracing_attempts)
else:
options = profiler.ProfilerOptions(
host_tracer_level=FLAGS.host_tracer_level)
profiler_client.trace(service_addr, os.path.expanduser(FLAGS.logdir),
duration_ms, workers_list,
FLAGS.num_tracing_attempts, options)
except errors.UnavailableError:
sys.exit(0)
if __name__ == '__main__':
run_main()
@@ -0,0 +1,15 @@
# Description:
# Tools for building the Cloud TPU profiler pip package.
package(default_visibility = ["//visibility:private"])
licenses(["notice"]) # Apache 2.0
sh_binary(
name = "build_pip_package",
srcs = ["build_pip_package.sh"],
data = [
"setup.py",
"//tensorflow/python/tpu/profiler:capture_tpu_profile_bin",
],
)
@@ -0,0 +1,2 @@
This package contains a Python wrapper around Tensorflow Python APIs that are
used to profile Cloud TPU.
@@ -0,0 +1,65 @@
#!/usr/bin/env bash
# Copyright 2017 The TensorFlow Authors. 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.
# 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 -e
if [ "$(uname)" = "Darwin" ]; then
sedi="sed -i ''"
else
sedi="sed -i"
fi
PACKAGE_NAME="cloud_tpu_profiler"
PIP_PACKAGE="tensorflow/python/tpu/profiler/pip_package"
RUNFILES="bazel-bin/tensorflow/python/tpu/profiler/pip_package/build_pip_package.runfiles/org_tensorflow/tensorflow/python/tpu/profiler"
function main() {
if [ $# -lt 1 ] ; then
echo "No destination dir provided"
exit 1
fi
DEST=$1
TMPDIR=$(mktemp -d -t tmp.XXXXXXXXXX)
echo $(date) : "=== Using tmpdir: ${TMPDIR}"
cp ${PIP_PACKAGE}/README ${TMPDIR}
cp ${PIP_PACKAGE}/setup.py ${TMPDIR}
mkdir ${TMPDIR}/${PACKAGE_NAME}
cp -a ${RUNFILES}/. ${TMPDIR}/${PACKAGE_NAME}/
# Fix the import statements to reflect the copied over path.
find ${TMPDIR}/${PACKAGE_NAME} -name \*.py |
xargs $sedi -e '
s/^from tensorflow.python.tpu.profiler/from '${PACKAGE_NAME}'/
'
echo $(ls $TMPDIR)
pushd ${TMPDIR}
echo $(date) : "=== Building wheel"
echo $(pwd)
python setup.py bdist_wheel >/dev/null
python3 setup.py bdist_wheel >/dev/null
mkdir -p ${DEST}
cp dist/* ${DEST}
popd
rm -rf ${TMPDIR}
echo $(date) : "=== Output wheel file is in: ${DEST}"
}
main "$@"
@@ -0,0 +1,58 @@
# Copyright 2017 The TensorFlow Authors. 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.
# 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.
# =============================================================================
"""Cloud TPU profiler package."""
from setuptools import setup
from cloud_tpu_profiler.version import __version__
CONSOLE_SCRIPTS = [
'capture_tpu_profile=cloud_tpu_profiler.capture_tpu_profile:run_main',
]
setup(
name='cloud_tpu_profiler',
version=__version__.replace('-', ''),
description='Trace and profile Cloud TPU performance',
long_description='Tools for capture TPU profile',
url='https://www.tensorflow.org/tfrc/',
author='Google Inc.',
author_email='packages@tensorflow.org',
packages=['cloud_tpu_profiler'],
entry_points={
'console_scripts': CONSOLE_SCRIPTS,
},
classifiers=[
'Development Status :: 5 - Production/Stable',
'Intended Audience :: Developers',
'Intended Audience :: Education',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: Apache Software License',
'Programming Language :: Python :: 2',
'Programming Language :: Python :: 2.7',
'Programming Language :: Python :: 3',
'Programming Language :: Python :: 3.4',
'Programming Language :: Python :: 3.5',
'Programming Language :: Python :: 3.6',
'Topic :: Scientific/Engineering',
'Topic :: Scientific/Engineering :: Mathematics',
'Topic :: Scientific/Engineering :: Artificial Intelligence',
'Topic :: Software Development',
'Topic :: Software Development :: Libraries',
'Topic :: Software Development :: Libraries :: Python Modules',
],
license='Apache 2.0',
keywords='tensorflow performance tpu',
)
@@ -0,0 +1,121 @@
# Copyright 2019 The TensorFlow Authors. 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.
# 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.
# ==============================================================================
# Generated by the gRPC Python protocol compiler plugin. DO NOT EDIT!
#
# Do not use pylint on generated code.
# pylint: disable=missing-docstring,g-short-docstring-punctuation,g-no-space-after-docstring-summary,invalid-name,line-too-long,unused-argument,g-doc-args
import grpc
from tensorflow.core.profiler import profiler_analysis_pb2 as third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
class ProfileAnalysisStub(object):
"""//////////////////////////////////////////////////////////////////////////////
ProfileAnalysis service provide entry point for profiling TPU and for
serving profiled data to Tensorboard through GRPC
//////////////////////////////////////////////////////////////////////////////
"""
def __init__(self, channel):
"""Constructor.
Args:
channel: A grpc.Channel.
"""
self.NewSession = channel.unary_unary(
'/tensorflow.ProfileAnalysis/NewSession',
request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.NewProfileSessionRequest.SerializeToString,
response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.NewProfileSessionResponse.FromString,
)
self.EnumSessions = channel.unary_unary(
'/tensorflow.ProfileAnalysis/EnumSessions',
request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.EnumProfileSessionsAndToolsRequest.SerializeToString,
response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.EnumProfileSessionsAndToolsResponse.FromString,
)
self.GetSessionToolData = channel.unary_unary(
'/tensorflow.ProfileAnalysis/GetSessionToolData',
request_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.ProfileSessionDataRequest.SerializeToString,
response_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.ProfileSessionDataResponse.FromString,
)
class ProfileAnalysisServicer(object):
"""//////////////////////////////////////////////////////////////////////////////
ProfileAnalysis service provide entry point for profiling TPU and for
serving profiled data to Tensorboard through GRPC
//////////////////////////////////////////////////////////////////////////////
"""
def NewSession(self, request, context):
"""Starts a profiling session, blocks until it completes.
TPUProfileAnalysis service delegate this to TPUProfiler service.
Populate the profiled data in repository, then return status to caller.
"""
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
def EnumSessions(self, request, context):
"""Enumerate existing sessions and return available profile tools."""
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
def GetSessionToolData(self, request, context):
"""Retrieve specific tool's data for specific session."""
context.set_code(grpc.StatusCode.UNIMPLEMENTED)
context.set_details('Method not implemented!')
raise NotImplementedError('Method not implemented!')
def add_ProfileAnalysisServicer_to_server(servicer, server):
rpc_method_handlers = {
'NewSession':
grpc.unary_unary_rpc_method_handler(
servicer.NewSession,
request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.NewProfileSessionRequest.FromString,
response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.NewProfileSessionResponse.SerializeToString,
),
'EnumSessions':
grpc.unary_unary_rpc_method_handler(
servicer.EnumSessions,
request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.EnumProfileSessionsAndToolsRequest.FromString,
response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.EnumProfileSessionsAndToolsResponse.SerializeToString,
),
'GetSessionToolData':
grpc.unary_unary_rpc_method_handler(
servicer.GetSessionToolData,
request_deserializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.ProfileSessionDataRequest.FromString,
response_serializer=third__party_dot_tensorflow_dot_core_dot_profiler_dot_profiler__analysis__pb2
.ProfileSessionDataResponse.SerializeToString,
),
}
generic_handler = grpc.method_handlers_generic_handler(
'tensorflow.ProfileAnalysis', rpc_method_handlers)
server.add_generic_rpc_handlers((generic_handler,))
+20
View File
@@ -0,0 +1,20 @@
# Copyright 2019 The TensorFlow Authors. 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.
# 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.
# =============================================================================
"""Cloud TPU profiler version information."""
# Cloud TPU profiler uses semantic versioning, see http://semver.org/.
# A version string consists of
# major_version.minor_version.patch_version-build_metadata.
__version__ = "1.14.1-a0"