Delete Python ScopedAnnotation
PiperOrigin-RevId: 312316009 Change-Id: I6d5d9afc55b33bb832ba5c7a867c65efdea4a5c2
This commit is contained in:
parent
0fd60fe3a3
commit
273617ad91
@ -139,12 +139,6 @@ cc_library(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
tf_pybind_cc_library_wrapper(
|
|
||||||
name = "scoped_annotation_headers",
|
|
||||||
visibility = ["//tensorflow/python/profiler/internal:__pkg__"],
|
|
||||||
deps = [":scoped_annotation"],
|
|
||||||
)
|
|
||||||
|
|
||||||
cc_library(
|
cc_library(
|
||||||
name = "scoped_annotation",
|
name = "scoped_annotation",
|
||||||
hdrs = ["scoped_annotation.h"],
|
hdrs = ["scoped_annotation.h"],
|
||||||
|
@ -239,13 +239,3 @@ py_library(
|
|||||||
":trace",
|
":trace",
|
||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
py_library(
|
|
||||||
name = "scoped_annotation",
|
|
||||||
srcs = ["scoped_annotation.py"],
|
|
||||||
srcs_version = "PY2AND3",
|
|
||||||
deps = [
|
|
||||||
"//tensorflow/python/profiler/internal:_pywrap_scoped_annotation",
|
|
||||||
"@six_archive//:six",
|
|
||||||
],
|
|
||||||
)
|
|
||||||
|
@ -95,19 +95,6 @@ tf_python_pybind_extension(
|
|||||||
],
|
],
|
||||||
)
|
)
|
||||||
|
|
||||||
tf_python_pybind_extension(
|
|
||||||
name = "_pywrap_scoped_annotation",
|
|
||||||
srcs = ["scoped_annotation_wrapper.cc"],
|
|
||||||
features = ["-layering_check"],
|
|
||||||
module_name = "_pywrap_scoped_annotation",
|
|
||||||
deps = [
|
|
||||||
"//tensorflow/core:lib",
|
|
||||||
"//tensorflow/core/profiler/lib:scoped_annotation_headers",
|
|
||||||
"@com_google_absl//absl/types:optional",
|
|
||||||
"@pybind11",
|
|
||||||
],
|
|
||||||
)
|
|
||||||
|
|
||||||
tf_python_pybind_extension(
|
tf_python_pybind_extension(
|
||||||
name = "_pywrap_profiler",
|
name = "_pywrap_profiler",
|
||||||
srcs = ["profiler_wrapper.cc"],
|
srcs = ["profiler_wrapper.cc"],
|
||||||
|
@ -1,55 +0,0 @@
|
|||||||
/* 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.
|
|
||||||
==============================================================================*/
|
|
||||||
|
|
||||||
#include <utility>
|
|
||||||
|
|
||||||
#include "absl/types/optional.h"
|
|
||||||
#include "pybind11/pybind11.h"
|
|
||||||
#include "tensorflow/core/platform/types.h"
|
|
||||||
#include "tensorflow/core/profiler/lib/scoped_annotation.h"
|
|
||||||
|
|
||||||
namespace py = pybind11;
|
|
||||||
|
|
||||||
namespace {
|
|
||||||
|
|
||||||
// Helper to implement ScopedAnnotation as a context manager in Python.
|
|
||||||
class ScopedAnnotationWrapper {
|
|
||||||
public:
|
|
||||||
explicit ScopedAnnotationWrapper(const tensorflow::string& name)
|
|
||||||
: name_(name) {}
|
|
||||||
|
|
||||||
void Enter() { annotation_.emplace(std::move(name_)); }
|
|
||||||
|
|
||||||
void Exit() { annotation_.reset(); }
|
|
||||||
|
|
||||||
static bool IsEnabled() {
|
|
||||||
return tensorflow::profiler::ScopedAnnotation::IsEnabled();
|
|
||||||
}
|
|
||||||
|
|
||||||
private:
|
|
||||||
tensorflow::string name_;
|
|
||||||
absl::optional<tensorflow::profiler::ScopedAnnotation> annotation_;
|
|
||||||
};
|
|
||||||
|
|
||||||
} // namespace
|
|
||||||
|
|
||||||
PYBIND11_MODULE(_pywrap_scoped_annotation, m) {
|
|
||||||
py::class_<ScopedAnnotationWrapper> scoped_annotation_class(
|
|
||||||
m, "ScopedAnnotation");
|
|
||||||
scoped_annotation_class.def(py::init<const tensorflow::string&>())
|
|
||||||
.def("Enter", &ScopedAnnotationWrapper::Enter)
|
|
||||||
.def("Exit", &ScopedAnnotationWrapper::Exit)
|
|
||||||
.def_static("IsEnabled", &ScopedAnnotationWrapper::IsEnabled);
|
|
||||||
};
|
|
@ -1,49 +0,0 @@
|
|||||||
# 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.
|
|
||||||
# ==============================================================================
|
|
||||||
"""ScopedAnnotation allows the profiler to annotate device (e.g., GPU) events.
|
|
||||||
|
|
||||||
Usage:
|
|
||||||
with scoped_annotation.ScopedAnnotation('name'):
|
|
||||||
...
|
|
||||||
"""
|
|
||||||
|
|
||||||
from __future__ import absolute_import
|
|
||||||
from __future__ import division
|
|
||||||
from __future__ import print_function
|
|
||||||
|
|
||||||
import six
|
|
||||||
|
|
||||||
from tensorflow.python.profiler.internal import _pywrap_scoped_annotation
|
|
||||||
|
|
||||||
|
|
||||||
class ScopedAnnotation(object):
|
|
||||||
"""Context manager that generates an annotation for the profiler."""
|
|
||||||
|
|
||||||
def __init__(self, name, **kwargs):
|
|
||||||
if _pywrap_scoped_annotation.ScopedAnnotation.IsEnabled():
|
|
||||||
if kwargs:
|
|
||||||
name += '#' + ','.join(key + '=' + str(value)
|
|
||||||
for key, value in six.iteritems(kwargs)) + '#'
|
|
||||||
self._scoped_annotation = _pywrap_scoped_annotation.ScopedAnnotation(name)
|
|
||||||
else:
|
|
||||||
self._scoped_annotation = None
|
|
||||||
|
|
||||||
def __enter__(self):
|
|
||||||
if self._scoped_annotation:
|
|
||||||
self._scoped_annotation.Enter()
|
|
||||||
|
|
||||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
||||||
if self._scoped_annotation:
|
|
||||||
self._scoped_annotation.Exit()
|
|
Loading…
Reference in New Issue
Block a user