Skip to content

Adds a closeable session result #411

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 9 commits into from
Feb 8, 2022

Conversation

Craigacp
Copy link
Collaborator

This PR changes the output of Session.Runner.run() from List<Tensor> to a Session.Result class which is autocloseable and has both string and index accessors. This is an intermediate step before we have proper resource scopes (#354) which we don't currently have time to work on.

One open question is if we should change the return type of TensorFunction.call(Map<String,Tensor>) so it returns the Result class instead of Map<String,Tensor>. If so, then it probably should be promoted into a top level class rather than a static inner class on Session.

I'll run it through the formatter after review.

@rnett
Copy link
Contributor

rnett commented Jan 30, 2022

One open question is if we should change the return type of TensorFunction.call(Map<String,Tensor>) so it returns the Result class instead of Map<String,Tensor>. If so, then it probably should be promoted into a top level class rather than a static inner class on Session.

IMO we should.

I'm not sure what happens if you double close a tensor but I think it's worthy checking that its nothing bad, or doing something to avoid it. Using the Result class shouldn't prevent closing individual tensors from it, and then the whole result.

@Craigacp
Copy link
Collaborator Author

Craigacp commented Jan 30, 2022

I think I'd prefer that it returned Result rather than Map<String,Tensor> as well though the asymmetry with the method argument is a bit annoying. @karllessard thoughts?

@karllessard
Copy link
Collaborator

(sorry for the delay @Craigacp)

I definitely think that we need to support both TensorFunction.call and Session.run with the same Result class, for simplicity.

In addition, if that's possible, the same structure could be used also to pass the input tensors to a function call, meaning that a user could instantiate it by passing a list of tensors (so instead of Result should be called something like TensorMap?). The danger of doing this though is to make sure that tensors won't leak if one failed to be initialize. For example in here:

try (Tensor.mapOf(
    "x", TFloat64.tensorOf(...),
    "y", TFloat32.vectorOf(...),
    "z", TInt32.vectorOf(...)
)) {
   ...
}

What will happen with x and y if z fails to allocate? It will throw before the autocloseable map is created. This could be handled I guess using lambdas but I'm just throwing ideas here.

@Craigacp
Copy link
Collaborator Author

Craigacp commented Feb 3, 2022

Yeah, so the separate lifetime issue is why it isn't used as a way of supplying arguments to a call. My justification is that users are in control of Tensor creation for arguments, so they can manage the lifetimes, however for the return values of Session.run and TensorFunction.call we create the tensors and so should provide a lifetime management mechanism. I can definitely see the other side of the argument though, as passing in Map<String,Tensor> and getting back Result which is near enough an AutoCloseable Map<String,Tensor> feels wrong. It feels wrong when I do it in ONNX Runtime too. But it does avoid resource leaks of the kind you mention.

@Craigacp
Copy link
Collaborator Author

Craigacp commented Feb 5, 2022

I've refactored TensorFunction.call so it returns Result and moved Result so it's a top level class rather than an inner class on Session.

*
* <p>When this is closed it closes all the {@link Tensor}s inside it. If you maintain a
* reference to a value after this object has been closed it will throw an {@link
* IllegalStateException} upon access.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it can still be useful in some cases for a user to retain the lifetime of a specific tensor beyond the scope of the Result, and it shouldn't be hard to do using reference counts internally.

What about adding a getAndRetain(...) endpoint to do this? If not being added as part of this PR, I can play around it after.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Rather than do reference counting it could have a BitSet denoting if the tensor should be closed? It would be a one way operation, you could getAndRetain and then the lifetime of the tensor is up to you with no way of putting it back under the Result management.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I already use reference counting for managing sparse tensors, I think if we can stick internally to a single method for extending the life of a tensor, it would be better.

Copy link
Collaborator Author

@Craigacp Craigacp Feb 7, 2022

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After some investigation we couldn't use the retainReference idiom to extend the lifespan of the Tensor as then the RawTensor's PointerScope loses track of it. So we're going to leave the getAndRetain method as a todo for when we have a look at managing Tensor lifespans better.

karllessard
karllessard previously approved these changes Feb 7, 2022
Copy link
Collaborator

@karllessard karllessard left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you want to run spotless now @Craigacp ?

@Craigacp
Copy link
Collaborator Author

Craigacp commented Feb 7, 2022

Done.

@karllessard karllessard merged commit 8fd9362 into tensorflow:master Feb 8, 2022
@Craigacp Craigacp deleted the session-result branch February 8, 2022 15:12
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants