Feat/mut resource eagerly

oneflow/api/python/session/session.cpp

@@ -33,6 +33,7 @@ ONEFLOW_API_PYBIND11_MODULE("", m) {
   // multi-client lazy global session context
   m.def("CreateMultiClientSessionContext", &CreateMultiClientSessionContext);
   m.def("InitMultiClientSessionContext", &InitMultiClientSessionContext);
+  m.def("MultiClientSessionContextUpdateResource", &MultiClientSessionContextUpdateResource);
   m.def("MultiClientSessionContextAddCGraph", &MultiClientSessionContextAddCGraph);
   m.def("TryDestroyMultiClientSessionContext", &TryDestroyMultiClientSessionContext);
 

oneflow/api/python/session/session.h

@@ -18,6 +18,7 @@ limitations under the License.
 
 #include <string>
 #include <google/protobuf/text_format.h>
+#include "oneflow/api/python/session/session_api.h"
 #include "oneflow/core/common/protobuf.h"
 #include "oneflow/core/control/ctrl_client.h"
 #include "oneflow/core/control/global_process_ctx.h"
@@ -128,8 +129,18 @@ inline Maybe<void> InitMultiClientSessionContext(const std::string& config_proto
   return Maybe<void>::Ok();
 }
 
+inline Maybe<void> MultiClientSessionContextUpdateResource(const std::string& resource_proto_str) {
+  CHECK_NOTNULL_OR_RETURN(Global<MultiClientSessionContext>::Get());
+  Resource reso_proto;
+  CHECK_OR_RETURN(TxtString2PbMessage(resource_proto_str, &reso_proto))
+      << "failed to parse config_proto: " << resource_proto_str;
+  JUST(Global<MultiClientSessionContext>::Get()->UpdateResource(reso_proto));
+  return Maybe<void>::Ok();
+}
+
 inline Maybe<void> MultiClientSessionContextAddCGraph(
     const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
+  CHECK_NOTNULL_OR_RETURN(Global<MultiClientSessionContext>::Get());
   JUST(Global<MultiClientSessionContext>::Get()->AddCGraph(c_graph_ptr));
   return Maybe<void>::Ok();
 }

oneflow/api/python/session/session_api.h

@@ -43,6 +43,10 @@ inline void InitMultiClientSessionContext(const std::string& config_proto_str) {
   return oneflow::InitMultiClientSessionContext(config_proto_str).GetOrThrow();
 }
 
+inline void MultiClientSessionContextUpdateResource(const std::string& resource_proto_str) {
+  return oneflow::MultiClientSessionContextUpdateResource(resource_proto_str).GetOrThrow();
+}
+
 inline void MultiClientSessionContextAddCGraph(
     const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
   return oneflow::MultiClientSessionContextAddCGraph(c_graph_ptr).GetOrThrow();

oneflow/core/framework/multi_client_session_context.cpp

@@ -15,9 +15,11 @@ limitations under the License.
 */
 
 #include "oneflow/core/common/buffer_manager.h"
+#include "oneflow/core/common/maybe.h"
 #include "oneflow/core/common/multi_client.h"
 #include "oneflow/core/framework/multi_client_session_context.h"
 #include "oneflow/core/framework/load_library.h"
+#include "oneflow/core/job/resource.pb.h"
 #include "oneflow/core/job/version.h"
 #include "oneflow/core/job/global_for.h"
 #include "oneflow/core/job/id_manager.h"
@@ -117,6 +119,12 @@ Maybe<void> MultiClientSessionContext::TryInit(const ConfigProto& config_proto)
   return Maybe<void>::Ok();
 }
 
+Maybe<void> MultiClientSessionContext::UpdateResource(const Resource& reso_proto) {
+  CHECK_NOTNULL_OR_RETURN((Global<ResourceDesc, ForSession>::Get()));
+  Global<ResourceDesc, ForSession>::Get()->Update(reso_proto);
+  return Maybe<void>::Ok();
+}
+
 Maybe<void> MultiClientSessionContext::AddCGraph(
     const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr) {
   graphs_.emplace_back(c_graph_ptr);

oneflow/core/framework/multi_client_session_context.h

@@ -31,6 +31,7 @@ class MultiClientSessionContext {
   ~MultiClientSessionContext() {}
 
   Maybe<void> TryInit(const ConfigProto& config_proto);
+  Maybe<void> UpdateResource(const Resource& reso_proto);
   Maybe<void> AddCGraph(const std::shared_ptr<oneflow::NNGraph>& c_graph_ptr);
   Maybe<void> TryClose();
 

oneflow/core/job/resource_desc.cpp

@@ -14,6 +14,7 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 #include <algorithm>
+#include "oneflow/core/job/resource.pb.h"
 #include "oneflow/core/job/resource_desc.h"
 #include "oneflow/core/common/util.h"
 #include "oneflow/core/control/global_process_ctx.h"
@@ -110,4 +111,14 @@ void ResourceDesc::DumpCudnnConf(const JobConfigProto& job_conf) {
   }
 }
 
+void ResourceDesc::Update(const Resource& reso_conf) {
+  if (reso_conf.has_nccl_use_compute_stream()) {
+    resource_.set_nccl_use_compute_stream(reso_conf.nccl_use_compute_stream());
+  }
+  if (reso_conf.has_disable_group_boxing_by_dst_parallel()) {
+    resource_.set_disable_group_boxing_by_dst_parallel(
+        reso_conf.disable_group_boxing_by_dst_parallel());
+  }
+}
+
 }  // namespace oneflow

oneflow/core/job/resource_desc.h

@@ -56,6 +56,7 @@ class ResourceDesc final {
   bool enable_tensor_float_32_compute() const { return resource_.enable_tensor_float_32_compute(); }
   const Resource& resource() const { return resource_; }
   void DumpCudnnConf(const JobConfigProto& job_conf);
+  void Update(const Resource& reso_conf);
 
  private:
   Resource resource_;

python/oneflow/framework/config_util.py

@@ -17,11 +17,22 @@
 import traceback
 
 import oneflow._oneflow_internal
+import oneflow.core.job.resource_pb2 as resource_util
 import oneflow.framework.hob as hob
 import oneflow.framework.session_context as session_ctx
 import oneflow.support.enable_if as enable_if
 
 
+def _set_attr_to_resource(attr_name, attr_value):
+    sess = session_ctx.GetDefaultSession()
+    if sess.status_ == sess.Status.INITED:
+        reso_config = resource_util.Resource()
+        setattr(reso_config, attr_name, attr_value)
+        sess.update_resource_eagerly(reso_config)
+    else:
+        setattr(sess.config_proto.resource, attr_name, attr_value)
+
+
 def api_load_library(val: str) -> None:
     """Load necessary library for job
 
@@ -347,14 +358,8 @@ def api_nccl_use_compute_stream(val: bool = False) -> None:
     Args:
         val (bool, optional): True or False. Defaults to False.
     """
-    return enable_if.unique([nccl_use_compute_stream, do_nothing])(val=val)
-
-
-@enable_if.condition(hob.in_normal_mode & ~hob.session_initialized)
-def nccl_use_compute_stream(val=False):
-    sess = session_ctx.GetDefaultSession()
     assert type(val) is bool
-    sess.config_proto.resource.nccl_use_compute_stream = val
+    _set_attr_to_resource("nccl_use_compute_stream", val)
 
 
 def api_disable_group_boxing_by_dst_parallel(val: bool = False) -> None:
@@ -363,14 +368,8 @@ def api_disable_group_boxing_by_dst_parallel(val: bool = False) -> None:
     Args:
         val (bool, optional): True or False. Defaults to False.
     """
-    return enable_if.unique([disable_group_boxing_by_dst_parallel, do_nothing])(val=val)
-
-
-@enable_if.condition(hob.in_normal_mode & ~hob.session_initialized)
-def disable_group_boxing_by_dst_parallel(val=False):
-    sess = session_ctx.GetDefaultSession()
     assert type(val) is bool
-    sess.config_proto.resource.disable_group_boxing_by_dst_parallel = val
+    _set_attr_to_resource("disable_group_boxing_by_dst_parallel", val)
 
 
 def api_nccl_num_streams(val: int) -> None:
@@ -553,5 +552,6 @@ def nccl_enable_mixed_fusion(val):
 
 @enable_if.condition(hob.in_normal_mode & hob.session_initialized)
 def do_nothing(*args, **kwargs):
-    print("Nothing happened because the session is running")
-    return False
+    raise NotImplementedError(
+        "This action donot working because session is initialized."
+    )

python/oneflow/framework/multi_client_session.py

@@ -119,3 +119,10 @@ def _update_function_flag_name2defaultVal(self):
     def _update_scope_attr_name2defaultVal(self):
         items = c_api_util.GetScopeConfigDef().attr_name2attr_def.items()
         self.scope_attr_name2default_val_ = {k: v.default_val for (k, v) in items}
+
+    def update_resource_eagerly(self, resource_config):
+        self._check_status(self.Status.INITED)
+        config_proto_str = text_format.MessageToString(resource_config)
+        oneflow._oneflow_internal.MultiClientSessionContextUpdateResource(
+            config_proto_str
+        )

python/oneflow/test/graph/test_optimization_conf.py

@@ -75,9 +75,19 @@ def build(self, x):
         g = CustomGraphSysConf()
 
         print("optimization conf: \n", g._optimization_conf_proto)
+        test_case.assertTrue(g._optimization_conf_proto.nccl_use_compute_stream)
         g._generate_config_proto()
         print("graph conf: \n", g._config_proto)
 
+        flow.boxing.nccl.enable_use_compute_stream(False)
+        test_case.assertTrue(not g._optimization_conf_proto.nccl_use_compute_stream)
+        flow.boxing.nccl.disable_group_boxing_by_dst_parallel(False)
+        test_case.assertTrue(
+            not g._optimization_conf_proto.disable_group_boxing_by_dst_parallel
+        )
+
+        print("optimization conf after session init: \n", g._optimization_conf_proto)
+
 
 if __name__ == "__main__":
     unittest.main()