@@ -57,6 +57,8 @@ struct TORCH_API InterpreterSession {
5757 // InterpreterSession* I)' instead. We will have no backwards compatibility
5858 // guarentees for this function.
5959 ReplicatedObj createMovable (Obj obj);
60+
61+ // Converts a `ReplicatedObj` to an `Obj` on this InterpreterSession.
6062 Obj fromMovable (const ReplicatedObj& obj);
6163
6264 protected:
@@ -73,6 +75,9 @@ struct TORCH_API InterpreterSession {
7375 std::function<void ()> deconstruction_callback_ = nullptr ;
7476};
7577
78+ // An `Interpreter` represents an invidual subinterpreter created by `torch::deploy`.
79+ // It allows for the creation of `InterpreterSession` objects which allow users to interact with
80+ // python objects.
7681class TORCH_API Interpreter {
7782 private:
7883 void * handle_;
@@ -84,17 +89,22 @@ class TORCH_API Interpreter {
8489 multipy::optional<EmbeddedFile> torchPluginFile_;
8590
8691 public:
92+ // Create an Interpreter which is managed by `manager` and using the enviroment `env`
8793 Interpreter (InterpreterManager* manager, std::shared_ptr<Environment> env);
94+
95+ // Create an Interpreter manager using enviroment `env` which is not tied to an Interpreter Manager.
8896 explicit Interpreter (std::shared_ptr<Environment> env)
8997 : Interpreter(nullptr , env) {}
9098
99+ // Get a new `InterpreterSession` from this Interpreter.
91100 InterpreterSession acquireSession () const {
92101 if (manager_) {
93102 return InterpreterSession (pImpl_->acquireSession (), manager_);
94103 } else {
95104 return InterpreterSession (pImpl_->acquireSession ());
96105 }
97106 }
107+
98108 ~Interpreter ();
99109 Interpreter (Interpreter&& rhs) noexcept
100110 : handle_(rhs.handle_),
@@ -113,17 +123,26 @@ class TORCH_API Interpreter {
113123
114124struct Package ;
115125
126+ // The default LoadBalancer for torch::deploy which handles allocating and freeing subinterpreters.
116127struct TORCH_API LoadBalancer {
128+
129+ // create a Loadbalancer which handles `n` interpreters.
117130 explicit LoadBalancer (size_t n)
118131 : uses_(new uint64_t [8 * n]), allocated_(n), n_(n) {
119132 // 8*... to avoid false sharing of atomics on the same cache line
120133 memset (uses_.get (), 0 , 8 * n_ * sizeof (uint64_t ));
121134 }
135+
136+ // change the amount of subinterpreters which is handled by the load balancer.
122137 void setResourceLimit (size_t n) {
123138 MULTIPY_INTERNAL_ASSERT (n <= allocated_);
124139 n_ = n;
125140 }
141+
142+ // allocate an subinterpreter, and return its ID which is used to free it.
126143 int acquire ();
144+
145+ // free the subinterpreter with ID `where`. This ID is returned by `LoadBalancer::acquire()`
127146 void free (int where);
128147
129148 private:
@@ -134,13 +153,18 @@ struct TORCH_API LoadBalancer {
134153 size_t n_;
135154};
136155
156+ // An `InterpreterManager` handles the interaction of multiple subinterpreters such as allocating
157+ // subinterpreters, or load balancing the subinterpreters.
137158struct TORCH_API InterpreterManager {
159+
160+ // constructor for `InterpreterManager` which takes the number of interpreters
161+ // (usually correlates to number of cores on your cpu), and a pointer to an `Enviroment`.
138162 explicit InterpreterManager (
139163 size_t nInterp = 2 ,
140164 std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141165
142- // get a free model , guarenteed that no other user of acquireOne has the same
143- // model . It _is_ possible that other users will be using the interpreter.
166+ // get a free InterpreterSession , guarenteed that no other user of acquireOne has the same
167+ // InterpreterSession . It _is_ possible that other users will be using the interpreter.
144168 InterpreterSession acquireOne () {
145169 int where = resources_.acquire ();
146170 InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +178,18 @@ struct TORCH_API InterpreterManager {
154178 at::ArrayRef<Interpreter> allInstances () {
155179 return instances_;
156180 }
181+
182+ // debugging tool to control the size of the loadBalancer
183+ // and change the number of interpreters on the fly
157184 void debugLimitInterpreters (size_t N) {
158185 AT_ASSERT (N <= instances_.size ());
159186 resources_.setResourceLimit (N);
160187 }
188+
189+ // loads a package from a file with name `uri`
161190 Package loadPackage (const std::string& uri);
191+
192+ // loads a package from a `PyTorchStreamReader` or any class other which uses `ReadAdapterInterface`
162193 Package loadPackage (
163194 std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164195
@@ -171,10 +202,12 @@ struct TORCH_API InterpreterManager {
171202 registeredModuleSource_[std::move (name)] = std::move (src);
172203 }
173204
174- // Util function for debugging.
205+ // Util function for debugging which outputs the number of registered modules .
175206 size_t countRegisteredModuleSources () {
176207 return registeredModuleSource_.size ();
177208 }
209+
210+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178211 ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179212 InterpreterManager (const InterpreterManager&) = delete ;
180213 InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -189,6 +222,7 @@ struct TORCH_API InterpreterManager {
189222 std::unordered_map<std::string, std::string> registeredModuleSource_;
190223};
191224
225+ // Underlying implementation for ReplicatedObj.
192226struct TORCH_API ReplicatedObjImpl {
193227 ReplicatedObjImpl (
194228 size_t object_id,
@@ -204,32 +238,44 @@ struct TORCH_API ReplicatedObjImpl {
204238 InterpreterManager* manager_;
205239};
206240
241+ // A python object which is Replicated from an `Obj` such that it is able to move around to different `InterpreterSessions`
242+ // by using `InterpreterSession::fromMovable(ReplicatedObj)`
207243struct TORCH_API ReplicatedObj {
244+
245+ // default constructor for `ReplicatedObj`
208246 ReplicatedObj () : pImpl_(nullptr ) {}
247+
248+ // Create an `InterpreterSession` using `onThisInterpreter`. If `onThisInterpreter` is
249+ // a `nullptr', then the associated `InterpreterManager` allocates it.
209250 InterpreterSession acquireSession (
210251 const Interpreter* onThisInterpreter = nullptr ) const ;
211252 at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212253 auto I = acquireSession ();
213254 return I.self (args).toIValue ();
214255 }
215256
257+ // invokes `callKwargs` using the underlying python object, and returns the output as an `IValue`.
216258 [[nodiscard]] at::IValue callKwargs (
217259 std::vector<at::IValue> args,
218260 std::unordered_map<std::string, c10::IValue> kwargs) const {
219261 auto I = acquireSession ();
220262 return I.self .callKwargs (std::move (args), std::move (kwargs)).toIValue ();
221263 }
222264
265+ // invokes `callKwargs` using the underlying python object, and returns the output as an `IValue`.
223266 [[nodiscard]] at::IValue callKwargs (
224267 std::unordered_map<std::string, c10::IValue> kwargs) const {
225268 auto I = acquireSession ();
226269 return I.self .callKwargs (std::move (kwargs)).toIValue ();
227270 }
228271
272+ // invokes `hasattr` using the underlying python object, and returns the output as an `IValue`.
229273 [[nodiscard]] bool hasattr (const char * name) const {
230274 auto I = acquireSession ();
231275 return I.self .hasattr (name);
232276 }
277+
278+
233279 void unload (const Interpreter* onThisInterpreter = nullptr );
234280 Obj toObj (InterpreterSession* I);
235281
@@ -242,21 +288,24 @@ struct TORCH_API ReplicatedObj {
242288 friend struct InterpreterManager ;
243289};
244290
291+ // PythonMethodWrapper is a more specific instance of a
292+ // ReplicatedObj which represents a python method, and
293+ // is therefore callable and has argument names accessible.
245294class PythonMethodWrapper : public torch ::IMethod {
246- // PythonMethodWrapper is a more specific instance of a
247- // ReplicatedObj which represents a python method, and
248- // is therefore callable and has argument names accessible.
249295 public:
250296 // TODO(whc) make bound method pickleable, then directly construct from that
297+
251298 PythonMethodWrapper (
252299 torch::deploy::ReplicatedObj model,
253300 std::string methodName)
254301 : model_(std::move(model)), methodName_(std::move(methodName)) {}
255302
303+ // return the name of the python method.
256304 const std::string& name () const override {
257305 return methodName_;
258306 }
259307
308+ // overrides the `()` operater to call the underlying python method.
260309 c10::IValue operator ()(
261310 std::vector<c10::IValue> args,
262311 const IValueMap& kwargs = IValueMap()) const override {
@@ -274,6 +323,7 @@ class PythonMethodWrapper : public torch::IMethod {
274323 std::string methodName_;
275324};
276325
326+ // An object to encapsulate a `torch::package` which can act as part (or entire) enviroment for subinterpreters.
277327struct TORCH_API Package {
278328 // shorthand for getting the object as a pickle resource in the package
279329 ReplicatedObj loadPickle (const std::string& module , const std::string& file) {
@@ -308,12 +358,15 @@ struct TORCH_API Package {
308358 }
309359#endif
310360
361+ // allocate an `InterpreterSession` and load the appropriate torch.package with it.
311362 InterpreterSession acquireSession () {
312363 auto I = manager_->acquireOne ();
313364 I.self =
314365 I.impl_ ->createOrGetPackageImporterFromContainerFile (containerFile_);
315366 return I;
316367 }
368+
369+ // Converts an `Obj` from `InterpreterSession` `I` into a `ReplicatedObj`.
317370 ReplicatedObj createMovable (Obj obj, InterpreterSession* I) {
318371 return manager_->createMovable (obj, I);
319372 }
0 commit comments