Fix --use_on_demand flag and add test

Summary:
--use_on_demand's previous impl was not equivalent to ON_DEMAND connection type: it would try to connect to IPUs before compilation rather than after.
This is because it wasn't overriding the config's device connection type, so the program still thought we were using ALWAYS in some places.
This diff changes it so that the flag makes us override the connection type in our config to ON_DEMAND at configuration time. The rest is free.
Done at C++ level so users of the kernel directly can also use the flag.
This means the flag will now make TF connect to IPUs after compilation.

Ref T69999

Test Plan: Added test to check our internal config is correct when not using and using the flag.

Reviewers: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, krzysztofk, gauthamg, mateuszk

Reviewed By: #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved, krzysztofk, gauthamg, mateuszk

Subscribers: harrym

Maniphest Tasks: T69999

Differential Revision: https://phabricator.sourcevertex.net/D80874

Author: Milly Fowden

tensorflow/compiler/plugin/poplar/driver/poplar_executor.cc

@@ -1562,8 +1562,7 @@ Status PoplarExecutor::AttachToPoplarDevice() {
   }
 
   const bool wait_for_device =
-      ConnectionType() == IpuDeviceConnectionType::ON_DEMAND ||
-      PoplarXlaFlags::Get().use_on_demand;
+      ConnectionType() == IpuDeviceConnectionType::ON_DEMAND;
   const bool use_ipu_model = PoplarXlaFlags::Get().use_ipu_model;
 
   try {
@@ -1802,10 +1801,21 @@ Status PoplarExecutor::CreatePoplarTarget() {
 
 Status PoplarExecutor::ConfigurePoplarDevice(const IpuOptions& cfg) {
   TENSORFLOW_TRACEPOINT();
+
+  // Override config options with global flags.
+  IpuOptions overridden_cfg = cfg;
+  if (PoplarXlaFlags::Get().use_on_demand) {
+    VLOG(1) << "Overriding device connection type to ON_DEMAND due to "
+               "--use_on_demand flag";
+    overridden_cfg.set_device_connection_type(
+        IpuDeviceConnectionType::ON_DEMAND);
+  }
+
   bool has_user_config = (current_config_.device_config_size() > 0);
-  if (!DeviceConfigurationsEqual(cfg, current_config_) && has_user_config) {
+  if (!DeviceConfigurationsEqual(overridden_cfg, current_config_) &&
+      has_user_config) {
     XLA_VLOG_LINES(1, "Current config: " + current_config_.DebugString() +
-                          "\nNew config: " + cfg.DebugString());
+                          "\nNew config: " + overridden_cfg.DebugString());
     return FailedPrecondition(
         "IPU system configuration has already been set in this process, "
         "but it should have been reset automatically by the call to "
@@ -1822,12 +1832,13 @@ Status PoplarExecutor::ConfigurePoplarDevice(const IpuOptions& cfg) {
               << ordinal_ << " is already configured: staying attached to it.";
     }
   }
-  current_config_ = cfg;
+  current_config_ = overridden_cfg;
   configured_ = true;
 
   if (!ipu_.DeviceAttached()) {
     TF_RETURN_IF_ERROR(CreatePoplarTarget());
-    if (cfg.device_connection_type() == IpuDeviceConnectionType::ALWAYS) {
+    if (overridden_cfg.device_connection_type() ==
+        IpuDeviceConnectionType::ALWAYS) {
       TF_RETURN_IF_ERROR(AttachToPoplarDevice());
     }
   }

tensorflow/python/ipu/BUILD

@@ -1038,6 +1038,24 @@ tf_py_test(
     ],
 )
 
+tf_py_test(
+    name = "utils_use_on_demand_test",
+    size = "medium",
+    srcs = ["tests/utils_use_on_demand_test.py"],
+    # Shard count set equal to the number of tests as one process per test is
+    # required
+    shard_count = 2,
+    deps = [
+        "//tensorflow:tensorflow_py",
+        "//tensorflow/compiler/tests:xla_test",
+        "//tensorflow/python:framework_test_lib",
+        "//tensorflow/python/compiler/xla:compiler_py",
+        "//tensorflow/python/ipu:ipu_lib",
+        "//tensorflow/python/ipu/test_utils",
+        "//tensorflow/python/keras",
+    ],
+)
+
 tf_py_test(
     name = "utils_use_synthetic_data_for_test",
     size = "medium",

tensorflow/python/ipu/tests/utils_use_on_demand_test.py

@@ -0,0 +1,69 @@
+# Copyright 2023 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.
+# =============================================================================
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+
+from tensorflow.compiler.plugin.poplar.driver.config_pb2 import IpuOptions
+from tensorflow.compiler.plugin.poplar.ops import gen_ipu_ops
+from tensorflow.python.client import session as session_lib
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ipu import config
+from tensorflow.python.ipu.utils import DeviceConnectionType
+from tensorflow.python.platform import googletest
+from tensorflow.python.platform import test
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+
+
+class UseOnDemandTest(test_util.TensorFlowTestCase):
+  def _configure(self, configured_type, expected_type):
+    cfg = config.IPUConfig()
+    cfg.device_connection.type = configured_type
+    cfg.configure_ipu_system()
+
+    # Get the current config.
+    g = ops.Graph()
+    with g.as_default():
+      with ops.device("CPU"):
+        with session_lib.Session(graph=g) as s:
+          configurations = s.run(gen_ipu_ops.ipu_get_configuration())
+
+    self.assertEqual(len(configurations), 1)
+    actual_cfg = IpuOptions()
+    actual_cfg.ParseFromString(configurations[0])
+
+    self.assertEqual(actual_cfg.device_connection_type, expected_type)
+
+  @test_util.deprecated_graph_mode_only
+  def testNoFlag(self):
+    self._configure(DeviceConnectionType.ALWAYS,
+                    DeviceConnectionType.ALWAYS.value)
+
+  @test_util.deprecated_graph_mode_only
+  def testFlag(self):
+    flags = os.environ.get("TF_POPLAR_FLAGS", "") + ' --use_on_demand'
+    with test.mock.patch.dict("os.environ", {"TF_POPLAR_FLAGS": flags}):
+      self._configure(DeviceConnectionType.ALWAYS,
+                      DeviceConnectionType.ON_DEMAND.value)
+
+
+if __name__ == "__main__":
+  os.environ['TF_XLA_FLAGS'] = ('--tf_xla_min_cluster_size=1' +
+                                os.environ.get('TF_XLA_FLAGS', ''))
+  googletest.main()