@@ -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
79+ // `torch::deploy`. It allows for the creation of `InterpreterSession` objects
80+ // which allow users to interact with python objects.
7681class TORCH_API Interpreter {
7782 private:
7883 void * handle_;
@@ -84,17 +89,24 @@ class TORCH_API Interpreter {
8489 multipy::optional<EmbeddedFile> torchPluginFile_;
8590
8691 public:
92+ // Creates an Interpreter which is managed by `manager` and using the
93+ // environment `env`
8794 Interpreter (InterpreterManager* manager, std::shared_ptr<Environment> env);
95+
96+ // Creates an Interpreter manager using environment `env` which is not tied to
97+ // an Interpreter Manager.
8898 explicit Interpreter (std::shared_ptr<Environment> env)
8999 : Interpreter(nullptr , env) {}
90100
101+ // Gets a new `InterpreterSession` from this Interpreter.
91102 InterpreterSession acquireSession () const {
92103 if (manager_) {
93104 return InterpreterSession (pImpl_->acquireSession (), manager_);
94105 } else {
95106 return InterpreterSession (pImpl_->acquireSession ());
96107 }
97108 }
109+
98110 ~Interpreter ();
99111 Interpreter (Interpreter&& rhs) noexcept
100112 : handle_(rhs.handle_),
@@ -113,17 +125,28 @@ class TORCH_API Interpreter {
113125
114126struct Package ;
115127
128+ // The default LoadBalancer for torch::deploy which handles allocating and
129+ // freeing subinterpreters.
116130struct TORCH_API LoadBalancer {
131+ // Creates a Loadbalancer which handles `n` interpreters.
117132 explicit LoadBalancer (size_t n)
118133 : uses_(new uint64_t [8 * n]), allocated_(n), n_(n) {
119134 // 8*... to avoid false sharing of atomics on the same cache line
120135 memset (uses_.get (), 0 , 8 * n_ * sizeof (uint64_t ));
121136 }
137+
138+ // Changes the amount of subinterpreters which is handled by the load
139+ // balancer.
122140 void setResourceLimit (size_t n) {
123141 MULTIPY_INTERNAL_ASSERT (n <= allocated_);
124142 n_ = n;
125143 }
144+
145+ // Allocates an subinterpreter, and return its ID which is used to free it.
126146 int acquire ();
147+
148+ // Frees the subinterpreter with ID `where`. This ID is returned by
149+ // `LoadBalancer::acquire()`
127150 void free (int where);
128151
129152 private:
@@ -134,13 +157,21 @@ struct TORCH_API LoadBalancer {
134157 size_t n_;
135158};
136159
160+ // An `InterpreterManager` handles the interaction of multiple subinterpreters
161+ // such as allocating subinterpreters, or load balancing the subinterpreters.
137162struct TORCH_API InterpreterManager {
163+ // constructor for `InterpreterManager` which takes the number of interpreters
164+ // (usually correlates to number of cores on your cpu), and a pointer to an
165+ // `Environment`. The default uses the local python env.
138166 explicit InterpreterManager (
139167 size_t nInterp = 2 ,
140168 std::shared_ptr<Environment> env = std::make_shared<NoopEnvironment>());
141169
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.
170+ // get a free InterpreterSession, guarenteed that no other user of acquireOne
171+ // has the same InterpreterSession. It _is_ possible that other users will be
172+ // using the interpreter if there are no free InterpreterSessions. Unless you
173+ // are very careful to only use free interpreters, then do not assume that the
174+ // `Obj`s are isolated from each other.
144175 InterpreterSession acquireOne () {
145176 int where = resources_.acquire ();
146177 InterpreterSession I = instances_[where].acquireSession ();
@@ -154,11 +185,19 @@ struct TORCH_API InterpreterManager {
154185 at::ArrayRef<Interpreter> allInstances () {
155186 return instances_;
156187 }
188+
189+ // debugging tool to control the size of the loadBalancer
190+ // and change the number of interpreters on the fly
157191 void debugLimitInterpreters (size_t N) {
158192 AT_ASSERT (N <= instances_.size ());
159193 resources_.setResourceLimit (N);
160194 }
195+
196+ // loads a package from a file with name `uri`
161197 Package loadPackage (const std::string& uri);
198+
199+ // loads a package from a `PyTorchStreamReader` or any class other which uses
200+ // `ReadAdapterInterface`
162201 Package loadPackage (
163202 std::shared_ptr<caffe2::serialize::ReadAdapterInterface> reader);
164203
@@ -171,10 +210,12 @@ struct TORCH_API InterpreterManager {
171210 registeredModuleSource_[std::move (name)] = std::move (src);
172211 }
173212
174- // Util function for debugging.
213+ // Util function for debugging which outputs the number of registered modules .
175214 size_t countRegisteredModuleSources () {
176215 return registeredModuleSource_.size ();
177216 }
217+
218+ // Converts `obj` from on `InterpreterSession` I into a `ReplicatedObj`.
178219 ReplicatedObj createMovable (Obj obj, InterpreterSession* I);
179220 InterpreterManager (const InterpreterManager&) = delete ;
180221 InterpreterManager& operator =(const InterpreterManager&) = delete ;
@@ -204,33 +245,57 @@ struct TORCH_API ReplicatedObjImpl {
204245 InterpreterManager* manager_;
205246};
206247
248+ // A python object which is Replicated from an `Obj` such that it is able to
249+ // move around to different `InterpreterSessions` by using
250+ // `InterpreterSession::fromMovable(ReplicatedObj)`
207251struct TORCH_API ReplicatedObj {
252+ // Default constructor for `ReplicatedObj`
208253 ReplicatedObj () : pImpl_(nullptr ) {}
254+
255+ // Creates an `InterpreterSession` using `onThisInterpreter`. If
256+ // `onThisInterpreter` is a `nullptr', then the associated
257+ // `InterpreterManager` allocates it.
209258 InterpreterSession acquireSession (
210259 const Interpreter* onThisInterpreter = nullptr ) const ;
211260 at::IValue operator ()(at::ArrayRef<at::IValue> args) const {
212261 auto I = acquireSession ();
213262 return I.self (args).toIValue ();
214263 }
215264
265+ // Calls an `ReplicatedObj` callable, with arguments given by the tuple args
266+ // and named arguments given by the dictionary kwargs. This is done on an
267+ // arbitrary `InterpreterSession` which belongs to the `ReplicatedObj`'s
268+ // manager.
216269 [[nodiscard]] at::IValue callKwargs (
217270 std::vector<at::IValue> args,
218271 std::unordered_map<std::string, c10::IValue> kwargs) const {
219272 auto I = acquireSession ();
220273 return I.self .callKwargs (std::move (args), std::move (kwargs)).toIValue ();
221274 }
222275
276+ // Calls an `ReplicatedObj` callable, with named arguments given by the
277+ // dictionary kwargs. This is done on an arbitrary `InterpreterSession` which
278+ // belongs to the `ReplicatedObj`'s manager.
223279 [[nodiscard]] at::IValue callKwargs (
224280 std::unordered_map<std::string, c10::IValue> kwargs) const {
225281 auto I = acquireSession ();
226282 return I.self .callKwargs (std::move (kwargs)).toIValue ();
227283 }
228284
229- [[nodiscard]] bool hasattr (const char * name) const {
285+ // Returns true if `ReplicatedObj` has attribute with name `attr` and false
286+ // otherwise. This is done on an arbitrary `InterpreterSession` which belongs
287+ // to the `ReplicatedObj`'s manager.
288+ [[nodiscard]] bool hasattr (const char * attr) const {
230289 auto I = acquireSession ();
231- return I.self .hasattr (name );
290+ return I.self .hasattr (attr );
232291 }
292+
293+ // Deletes `ReplicatedObj` from onThisInterpreter, if onThisInterpreter is
294+ // `nullptr`, unload is called on all interpreters belonging to the
295+ // ReplicatedObject's InterpreterManager
233296 void unload (const Interpreter* onThisInterpreter = nullptr );
297+
298+ // Converts `ReplicatedObj` to `Obj` on `InterpreterSession` `I`
234299 Obj toObj (InterpreterSession* I);
235300
236301 private:
@@ -242,21 +307,24 @@ struct TORCH_API ReplicatedObj {
242307 friend struct InterpreterManager ;
243308};
244309
310+ // PythonMethodWrapper is a more specific instance of a
311+ // ReplicatedObj which represents a python method, and
312+ // is therefore callable and has argument names accessible.
245313class 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.
249314 public:
250315 // TODO(whc) make bound method pickleable, then directly construct from that
316+
251317 PythonMethodWrapper (
252318 torch::deploy::ReplicatedObj model,
253319 std::string methodName)
254320 : model_(std::move(model)), methodName_(std::move(methodName)) {}
255321
322+ // return the name of the python method.
256323 const std::string& name () const override {
257324 return methodName_;
258325 }
259326
327+ // overrides the `()` operater to call the underlying python method.
260328 c10::IValue operator ()(
261329 std::vector<c10::IValue> args,
262330 const IValueMap& kwargs = IValueMap()) const override {
@@ -274,6 +342,8 @@ class PythonMethodWrapper : public torch::IMethod {
274342 std::string methodName_;
275343};
276344
345+ // An object to encapsulate a `torch::package` which can act as part (or entire)
346+ // environment for subinterpreters.
277347struct TORCH_API Package {
278348 // shorthand for getting the object as a pickle resource in the package
279349 ReplicatedObj loadPickle (const std::string& module , const std::string& file) {
@@ -308,12 +378,16 @@ struct TORCH_API Package {
308378 }
309379#endif
310380
381+ // Allocates an `InterpreterSession` and load the appropriate torch.package
382+ // with it.
311383 InterpreterSession acquireSession () {
312384 auto I = manager_->acquireOne ();
313385 I.self =
314386 I.impl_ ->createOrGetPackageImporterFromContainerFile (containerFile_);
315387 return I;
316388 }
389+
390+ // Converts an `Obj` from `InterpreterSession` `I` into a `ReplicatedObj`.
317391 ReplicatedObj createMovable (Obj obj, InterpreterSession* I) {
318392 return manager_->createMovable (obj, I);
319393 }
0 commit comments