| 39 | static std::atomic<bool> already_provisioned(false); |
| 40 | |
| 41 | SingleMachine::SingleMachine(int timeout_s, int num_cpu_cores, int num_gpus) |
| 42 | : Cluster(timeout_s), expected_init_time_s_(0), closing_(false) { |
| 43 | VLOG(1) << "Number of CPU cores: " << num_cpu_cores |
| 44 | << " Number of GPUs: " << num_gpus; |
| 45 | thread_pool_.reset(new thread::ThreadPool( |
| 46 | Env::Default(), SanitizeThreadSuffix("single_machine"), 2)); |
| 47 | |
| 48 | (*options_.config.mutable_device_count())["CPU"] = 1; |
| 49 | if (num_gpus > 0) { |
| 50 | (*options_.config.mutable_device_count())["GPU"] = num_gpus; |
| 51 | } |
| 52 | CHECK_GE(num_cpu_cores, 1); |
| 53 | options_.config.set_intra_op_parallelism_threads(num_cpu_cores); |
| 54 | // Create a session specific thread pool to ensure the threads are reset when |
| 55 | // the session is reset. |
| 56 | options_.config.add_session_inter_op_thread_pool()->set_num_threads( |
| 57 | num_cpu_cores); |
| 58 | if (timeout_s > 0) { |
| 59 | options_.config.set_operation_timeout_in_ms(timeout_s * 1000); |
| 60 | } |
| 61 | } |
| 62 | |
| 63 | SingleMachine::~SingleMachine() { |
| 64 | CloseSession(false /*use_timeout*/).IgnoreError(); |
nothing calls this directly
no test coverage detected