| 69 | } |
| 70 | |
| 71 | Status SingleMachine::Provision() { |
| 72 | // This is really ugly: to avoid leaking variables, we need to reset the tf |
| 73 | // session every time we're done processing a grappler item. However, |
| 74 | // variables are global, and therefore we can't have more than 1 session alive |
| 75 | // at a time. This check detects when more that one cluster is provisioned. |
| 76 | if (already_provisioned) { |
| 77 | return errors::Unavailable( |
| 78 | "Can't provision more than one single cluster at a time"); |
| 79 | } |
| 80 | |
| 81 | TF_RETURN_IF_ERROR(ResetSession()); |
| 82 | |
| 83 | std::vector<DeviceAttributes> devices; |
| 84 | TF_RETURN_IF_ERROR(session_->ListDevices(&devices)); |
| 85 | for (const auto& dev : devices) { |
| 86 | DeviceProperties attr; |
| 87 | if (dev.device_type() == "CPU") { |
| 88 | attr = GetLocalCPUInfo(); |
| 89 | } else if (dev.device_type() == "GPU") { |
| 90 | DeviceNameUtils::ParsedName parsed; |
| 91 | if (!DeviceNameUtils::ParseFullName(dev.name(), &parsed)) { |
| 92 | return errors::InvalidArgument( |
| 93 | strings::StrCat("Not able to parse GPU device name: ", dev.name())); |
| 94 | } |
| 95 | TfGpuId tf_gpu_id(parsed.id); |
| 96 | PlatformGpuId platform_gpu_id; |
| 97 | Status s = GpuIdManager::TfToPlatformGpuId(tf_gpu_id, &platform_gpu_id); |
| 98 | if (!s.ok()) { |
| 99 | return errors::Unavailable("Unknown TF GPU device with id ", |
| 100 | tf_gpu_id.value(), ": ", s.ToString()); |
| 101 | } |
| 102 | attr = GetLocalGPUInfo(platform_gpu_id); |
| 103 | } else if (dev.device_type().find("XLA") == string::npos) { |
| 104 | // Filter out the fake XLA devices to avoid double counting the actual |
| 105 | // hardware resources that are available. |
| 106 | attr.set_type(dev.device_type()); |
| 107 | } |
| 108 | // Overwrite the memory size since users might have requested to use only a |
| 109 | // fraction of the available device memory. |
| 110 | attr.set_memory_size(dev.memory_limit()); |
| 111 | devices_[dev.name()] = attr; |
| 112 | } |
| 113 | already_provisioned = true; |
| 114 | |
| 115 | // Clear highmark stats of all local allocators. |
| 116 | if (cpu_allocator_stats_enabled_) { |
| 117 | TF_RETURN_IF_ERROR(ClearAllocatorStats()); |
| 118 | } |
| 119 | return Status::OK(); |
| 120 | } |
| 121 | |
| 122 | Status SingleMachine::Initialize(const GrapplerItem& item) { |
| 123 | mutex_lock l(this->last_graph_mu_); |