@@ -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+ // Creates an Interpreter which is managed by `manager` and using the environment `env`
8793 Interpreter (InterpreterManager* manager, std::shared_ptr<Environment> env);
94+
95+ // Creates an Interpreter manager using environment `env` which is not tied to an Interpreter Manager.
8896 explicit Interpreter (std::shared_ptr<Environment> env)
8997 : Interpreter(nullptr , env) {}
9098
99+ // Gets 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+ // Creates 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+ // Changes 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+ // Allocates an subinterpreter, and return its ID which is used to free it.
126143 int acquire ();
144+
145+ // Frees the subinterpreter with ID `where`. This ID is returned by `LoadBalancer::acquire()`
127146 void free (int where);
128147
129148 private:
@@ -134,13 +153,21 @@ 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 `Environment`.
162+ // The defualt uses the local python env.
138163 explicit InterpreterManager (
139164 size_t nInterp = 2 ,
140165 std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141166
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.
167+ // get a free InterpreterSession, guarenteed that no other user of acquireOne has the same
168+ // InterpreterSession. It _is_ possible that other users will be using the interpreter if there are
169+ // no free InterpreterSessions. Unless you are very careful to only use free interpreters, then do not assume
170+ // that the `Obj`s are isolated from each other.
144171 InterpreterSession acquireOne () {
145172 int where = resources_.acquire ();
146173 InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +181,18 @@ struct TORCH_API InterpreterManager {
154181 at::ArrayRef<Interpreter> allInstances () {
155182 return instances_;
156183 }
184+
185+ // debugging tool to control the size of the loadBalancer
186+ // and change the number of interpreters on the fly
157187 void debugLimitInterpreters (size_t N) {
158188 AT_ASSERT (N <= instances_.size ());
159189 resources_.setResourceLimit (N);
160190 }
191+
192+ // loads a package from a file with name `uri`
161193 Package loadPackage (const std::string& uri);
194+
195+ // loads a package from a `PyTorchStreamReader` or any class other which uses `ReadAdapterInterface`
162196 Package loadPackage (
163197 std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164198
@@ -171,10 +205,12 @@ struct TORCH_API InterpreterManager {
171205 registeredModuleSource_[std::move (name)] = std::move (src);
172206 }
173207
174- // Util function for debugging.
208+ // Util function for debugging which outputs the number of registered modules .
175209 size_t countRegisteredModuleSources () {
176210 return registeredModuleSource_.size ();
177211 }
212+
213+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178214 ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179215 InterpreterManager (const InterpreterManager&) = delete ;
180216 InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -204,33 +240,51 @@ struct TORCH_API ReplicatedObjImpl {
204240 InterpreterManager* manager_;
205241};
206242
243+ // A python object which is Replicated from an `Obj` such that it is able to move around to different `InterpreterSessions`
244+ // by using `InterpreterSession::fromMovable(ReplicatedObj)`
207245struct TORCH_API ReplicatedObj {
246+
247+ // Default constructor for `ReplicatedObj`
208248 ReplicatedObj () : pImpl_(nullptr ) {}
249+
250+ // Creates an `InterpreterSession` using `onThisInterpreter`. If `onThisInterpreter` is
251+ // a `nullptr', then the associated `InterpreterManager` allocates it.
209252 InterpreterSession acquireSession (
210253 const Interpreter* onThisInterpreter = nullptr ) const ;
211254 at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212255 auto I = acquireSession ();
213256 return I.self (args).toIValue ();
214257 }
215258
259+ // Calls an `ReplicatedObj` callable, with arguments given by the tuple args and named arguments given by the dictionary
260+ // kwargs. This is done on an arbitrary `InterpreterSession` which belongs to the `ReplicatedObj`'s manager.
216261 [[nodiscard]] at::IValue callKwargs (
217262 std::vector<at::IValue> args,
218263 std::unordered_map<std::string, c10::IValue> kwargs) const {
219264 auto I = acquireSession ();
220265 return I.self .callKwargs (std::move (args), std::move (kwargs)).toIValue ();
221266 }
222267
268+ // Calls an `ReplicatedObj` callable, with named arguments given by the dictionary kwargs.
269+ // This is done on an arbitrary `InterpreterSession` which belongs to the `ReplicatedObj`'s manager.
223270 [[nodiscard]] at::IValue callKwargs (
224271 std::unordered_map<std::string, c10::IValue> kwargs) const {
225272 auto I = acquireSession ();
226273 return I.self .callKwargs (std::move (kwargs)).toIValue ();
227274 }
228275
229- [[nodiscard]] bool hasattr (const char * name) const {
276+ // Returns true if `ReplicatedObj` has attribute with name `attr` and false otherwise.
277+ // This is done on an arbitrary `InterpreterSession` which belongs to the `ReplicatedObj`'s manager.
278+ [[nodiscard]] bool hasattr (const char * attr) const {
230279 auto I = acquireSession ();
231- return I.self .hasattr (name );
280+ return I.self .hasattr (attr );
232281 }
282+
283+ // Deletes `ReplicatedObj` from onThisInterpreter, if onThisInterpreter is `nullptr`,
284+ // unload is called on all interpreters belonging to the ReplicatedObject's InterpreterManager
233285 void unload (const Interpreter* onThisInterpreter = nullptr );
286+
287+ // Converts `ReplicatedObj` to `Obj` on `InterpreterSession` `I`
234288 Obj toObj (InterpreterSession* I);
235289
236290 private:
@@ -242,21 +296,24 @@ struct TORCH_API ReplicatedObj {
242296 friend struct InterpreterManager ;
243297};
244298
299+ // PythonMethodWrapper is a more specific instance of a
300+ // ReplicatedObj which represents a python method, and
301+ // is therefore callable and has argument names accessible.
245302class 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.
249303 public:
250304 // TODO(whc) make bound method pickleable, then directly construct from that
305+
251306 PythonMethodWrapper (
252307 torch::deploy::ReplicatedObj model,
253308 std::string methodName)
254309 : model_(std::move(model)), methodName_(std::move(methodName)) {}
255310
311+ // return the name of the python method.
256312 const std::string& name () const override {
257313 return methodName_;
258314 }
259315
316+ // overrides the `()` operater to call the underlying python method.
260317 c10::IValue operator ()(
261318 std::vector<c10::IValue> args,
262319 const IValueMap& kwargs = IValueMap()) const override {
@@ -274,6 +331,7 @@ class PythonMethodWrapper : public torch::IMethod {
274331 std::string methodName_;
275332};
276333
334+ // An object to encapsulate a `torch::package` which can act as part (or entire) environment for subinterpreters.
277335struct TORCH_API Package {
278336 // shorthand for getting the object as a pickle resource in the package
279337 ReplicatedObj loadPickle (const std::string& module , const std::string& file) {
@@ -308,12 +366,15 @@ struct TORCH_API Package {
308366 }
309367#endif
310368
369+ // Allocates an `InterpreterSession` and load the appropriate torch.package with it.
311370 InterpreterSession acquireSession () {
312371 auto I = manager_->acquireOne ();
313372 I.self =
314373 I.impl_ ->createOrGetPackageImporterFromContainerFile (containerFile_);
315374 return I;
316375 }
376+
377+ // Converts an `Obj` from `InterpreterSession` `I` into a `ReplicatedObj`.
317378 ReplicatedObj createMovable (Obj obj, InterpreterSession* I) {
318379 return manager_->createMovable (obj, I);
319380 }
0 commit comments