Add ability to set ON_DEMAND using flags

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

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

Maniphest Tasks: T69999

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

Author: hmellor

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

@@ -1562,7 +1562,8 @@ Status PoplarExecutor::AttachToPoplarDevice() {
   }
 
   const bool wait_for_device =
-      ConnectionType() == IpuDeviceConnectionType::ON_DEMAND;
+      ConnectionType() == IpuDeviceConnectionType::ON_DEMAND ||
+      PoplarXlaFlags::Get().use_on_demand;
   const bool use_ipu_model = PoplarXlaFlags::Get().use_ipu_model;
 
   try {

tensorflow/compiler/plugin/poplar/driver/tools/flags.cc

@@ -112,6 +112,8 @@ absl::flat_hash_map<std::string, std::string> GetFlagUsage() {
       {"show_progress_bar",
        "Whether to show the compilation progress bar. Either \"true\", "
        "\"false\" or \"auto\". Defaults to \"auto\"."},
+      {"use_on_demand",
+       "When true, sets connection type to \"ON_DEMAND\", useful for CI."},
       {"on_demand_device_poll_time",
        "When using \"ON_DEMAND\" connection type, configure how often to poll "
        "for the device (in miliseconds) when a device is not available - "
@@ -166,6 +168,7 @@ PoplarXlaFlags::PoplarXlaFlags() {
     ADD_FLAG(allow_nans)
     ADD_FLAG(null_data_feed)
     ADD_FLAG(show_progress_bar)
+    ADD_FLAG(use_on_demand)
     ADD_FLAG(on_demand_device_poll_time)
     ADD_FLAG(on_demand_device_timeout)
     ADD_FLAG(ipu_model_tiles)

tensorflow/compiler/plugin/poplar/driver/tools/flags.h

@@ -100,6 +100,9 @@ class PoplarXlaFlags {
   // Whether to show the compilation progress bar.
   std::string show_progress_bar = "auto";
 
+  // When true, sets connection type to 'ON_DEMAND', useful for CI.
+  bool use_on_demand = false;
+
   // When using 'ON_DEMAND' connection type, configure how often to poll for the
   // device (in milliseconds) when a device is not available - defaults to
   // 1000ms. Minimum is 100ms.