!1017 add profiling_mode and profiling_options to context

Merge pull request !1017 from jinyaohui/context_opt

mindspore/ccsrc/pipeline/init.cc

@@ -145,7 +145,11 @@ PYBIND11_MODULE(_c_expression, m) {
     .def("set_save_dump_path", &mindspore::MsContext::set_save_dump_path, "Set path to dump.")
     .def("set_graph_memory_max_size", &mindspore::MsContext::set_graph_memory_max_size, "set graph memory max size.")
     .def("set_variable_memory_max_size", &mindspore::MsContext::set_variable_memory_max_size,
-         "set variable memory max size");
+         "set variable memory max size")
+    .def("get_enable_profiling", &mindspore::MsContext::enable_profiling, "Get whether to open profiling.")
+    .def("set_enable_profiling", &mindspore::MsContext::set_enable_profiling, "Set whether to open profiling.")
+    .def("get_profiling_options", &mindspore::MsContext::profiling_options, "Get options to profiling.")
+    .def("set_profiling_options", &mindspore::MsContext::set_profiling_options, "Set options to profiling.");
 
   (void)py::class_<ParallelContext, std::shared_ptr<ParallelContext>>(m, "AutoParallelContext")
     .def_static("get_instance", &ParallelContext::GetInstance, "Get auto parallel context instance.")

mindspore/ccsrc/utils/context/ms_context.cc

@@ -78,6 +78,8 @@ MsContext::MsContext(const std::string &policy, const std::string &target) {
   graph_memory_max_size_ = "0";
   variable_memory_max_size_ = "0";
   enable_loop_sink_ = target == kAscendDevice || target == kDavinciDevice;
+  profiling_mode_ = false;
+  profiling_options_ = "training_trace";
 }
 
 std::shared_ptr<MsContext> MsContext::GetInstance() {
@@ -279,6 +281,11 @@ void MsContext::GetGeOptions(std::map<std::string, std::string> *ge_options) con
   (*ge_options)["device_id"] = "0";
   (*ge_options)["ge.exec.enableDump"] = std::to_string(enable_dump_);
   (*ge_options)["ge.exec.dumpPath"] = save_dump_path_;
+  (*ge_options)["ge.exec.profilingMode"] = std::to_string(profiling_mode_);
+  if (profiling_mode_) {
+    (*ge_options)["ge.exec.profilingOptions"] = profiling_options_;
+  }
+
   // only not supported in ge
   auto tbe_plugin_path = common::GetEnv("ME_TBE_PLUGIN_PATH");
   if (!tbe_plugin_path.empty()) {

mindspore/ccsrc/utils/context/ms_context.h

@@ -138,6 +138,12 @@ class MsContext {
     variable_memory_max_size_ = variable_memory_max_size;
   }
 
+  void set_enable_profiling(bool flag) { profiling_mode_ = flag; }
+  bool enable_profiling() const { return profiling_mode_; }
+
+  void set_profiling_options(const std::string &options) { profiling_options_ = options; }
+  std::string profiling_options() const { return profiling_options_; }
+
  private:
   MsContext(const std::string &backend_policy, const std::string &target);
   void GetGeOptions(std::map<std::string, std::string> *ge_options) const;
@@ -174,6 +180,8 @@ class MsContext {
   std::string graph_memory_max_size_;
   std::string variable_memory_max_size_;
   std::thread tdt_print_;
+  bool profiling_mode_;
+  std::string profiling_options_;
 };
 
 }  // namespace mindspore

mindspore/context.py

@@ -26,7 +26,6 @@ from mindspore._checkparam import args_type_check
 from mindspore.parallel._auto_parallel_context import _set_auto_parallel_context, _get_auto_parallel_context, \
     _reset_auto_parallel_context
 
-
 __all__ = ['GRAPH_MODE', 'PYNATIVE_MODE', 'set_context', 'get_context', 'set_auto_parallel_context',
            'get_auto_parallel_context', 'reset_auto_parallel_context']
 
@@ -297,6 +296,26 @@ class _Context:
     def save_dump_path(self, save_dump_path):
         self._context_handle.set_save_dump_path(save_dump_path)
 
+    @property
+    def enable_profiling(self):
+        return self._context_handle.get_enable_profiling()
+
+    @enable_profiling.setter
+    def enable_profiling(self, flag):
+        self._context_handle.set_enable_profiling(flag)
+
+    @property
+    def profiling_options(self):
+        return self._context_handle.get_profiling_options()
+
+    @profiling_options.setter
+    def profiling_options(self, option):
+        options = ["training_trace", "task_trace", "task_trace:training_trace", "training_trace:task_trace", "op_trace"]
+        if option not in options:
+            raise ValueError("Profiling options must be in 'training_trace' 'task_trace' "
+                             "'task_trace:training_trace' 'training_trace:task_trace' or 'op_trace'.")
+        self._context_handle.set_profiling_options(option)
+
     @property
     def reserve_class_name_in_scope(self):
         """Gets whether to save the network class name in the scope."""
@@ -472,7 +491,7 @@ def reset_auto_parallel_context():
                  enable_mem_reuse=bool, save_ms_model=bool, save_ms_model_path=str,
                  enable_auto_mixed_precision=bool, enable_dump=bool, save_dump_path=str,
                  enable_reduce_precision=bool, graph_memory_max_size=str,
-                 variable_memory_max_size=str)
+                 variable_memory_max_size=str, enable_profiling=bool, profiling_options=str)
 def set_context(**kwargs):
     """
     Sets context for running environment.
@@ -515,6 +534,21 @@ def set_context(**kwargs):
             So the real dump path is "{configured root dump path}/{`save_dump_path`}". Default: ".".
         graph_memory_max_size (str): Sets graph memory max size. Default: "26GB".
         variable_memory_max_size (str): Sets variable memory max size. Default: "5GB".
+        enable_profiling (bool): Whether to open profiling. Default: False.
+        profiling_options (str): Sets profiling collection options, operators can profiling data here.
+            Profiling collection options, the values are as follows, supporting the collection of multiple data.
+
+            - training_trace: collect iterative trajectory data, that is, the training task and software information of
+              the AI software stack, to achieve performance analysis of the training task, focusing on data
+              enhancement, forward and backward calculation, gradient aggregation update and other related data.
+
+            - task_trace: collect task trajectory data, that is, the hardware information of the HWTS/AICore of
+              the Ascend 910 processor, and analyze the information of start and end of the task.
+
+            - op_trace: collect single operator performance data.
+            The profiling can choose training_trace, task_trace, training_trace and task_trace combination and
+            separated by colons; single operator can choose op_trace, op_trace cannot be combined with
+            training_trace and task_trace. Default: "training_trace".
 
     Raises:
         ValueError: If input key is not an attribute in context.
@@ -536,6 +570,7 @@ def set_context(**kwargs):
         >>> context.set_context(mode=context.GRAPH_MODE,
         >>>                     device_target="Ascend",device_id=0, save_graphs=True,
         >>>                     save_graphs_path="/mindspore")
+        >>> context.set_context(enable_profiling=True, profiling_options="training_trace")
     """
     for key, value in kwargs.items():
         if not hasattr(_context(), key):

tests/ut/python/pynative_mode/test_context.py

@@ -16,6 +16,8 @@
 import os
 import pytest
 from mindspore import context
+
+
 # pylint: disable=W0212
 # W0212: protected-access
 
@@ -72,6 +74,34 @@ def test_dump_target():
     assert context.get_context("save_dump_path") == "."
 
 
+def test_enable_profiling():
+    """ test_profiling_mode """
+    with pytest.raises(TypeError):
+        context.set_context(enable_profiling=1)
+    with pytest.raises(TypeError):
+        context.set_context(enable_profiling="1")
+    context.set_context(enable_profiling=True)
+    assert context.get_context("enable_profiling") is True
+    context.set_context(enable_profiling=False)
+    assert context.get_context("enable_profiling") is False
+
+
+def test_profiling_options():
+    """ test_profiling_options """
+    with pytest.raises(TypeError):
+        context.set_context(profiling_options=True)
+    with pytest.raises(TypeError):
+        context.set_context(profiling_options=1)
+    with pytest.raises(ValueError):
+        context.set_context(profiling_options="training_")
+    with pytest.raises(ValueError):
+        context.set_context(profiling_options="training_trace:op_trace")
+    context.set_context(profiling_options="training_trace")
+    assert context.get_context("profiling_options") == "training_trace"
+    context.set_context(profiling_options="training_trace:task_trace")
+    assert context.get_context("profiling_options") == "training_trace:task_trace"
+
+
 def test_set_context():
     """ test_set_context """
     context.set_context(mode=context.GRAPH_MODE, device_target="Ascend",
@@ -101,4 +131,3 @@ def teardown_module():
             os.rmdir(item_name)
         elif os.path.isfile(item_name):
             os.remove(item_name)
-