callserver 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- callserver-0.1.0/LICENSE +21 -0
- callserver-0.1.0/PKG-INFO +134 -0
- callserver-0.1.0/README.md +118 -0
- callserver-0.1.0/callserver/__init__.py +54 -0
- callserver-0.1.0/callserver/__main__.py +4 -0
- callserver-0.1.0/callserver/allocationmanager.py +176 -0
- callserver-0.1.0/callserver/allocationset.py +381 -0
- callserver-0.1.0/callserver/callserver.md +315 -0
- callserver-0.1.0/callserver/callserver.py +511 -0
- callserver-0.1.0/callserver.egg-info/PKG-INFO +134 -0
- callserver-0.1.0/callserver.egg-info/SOURCES.txt +16 -0
- callserver-0.1.0/callserver.egg-info/dependency_links.txt +1 -0
- callserver-0.1.0/callserver.egg-info/entry_points.txt +2 -0
- callserver-0.1.0/callserver.egg-info/requires.txt +4 -0
- callserver-0.1.0/callserver.egg-info/top_level.txt +1 -0
- callserver-0.1.0/pyproject.toml +34 -0
- callserver-0.1.0/setup.cfg +4 -0
- callserver-0.1.0/test/test_callserver.py +551 -0
callserver-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 dexman-ai
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,134 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: callserver
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Shared-memory RPC for fast inter-process function calls with numpy arrays.
|
|
5
|
+
Author: dexman-ai
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
Classifier: Programming Language :: Python :: 3
|
|
8
|
+
Classifier: Operating System :: OS Independent
|
|
9
|
+
Requires-Python: >=3.8
|
|
10
|
+
Description-Content-Type: text/markdown
|
|
11
|
+
License-File: LICENSE
|
|
12
|
+
Requires-Dist: numpy
|
|
13
|
+
Provides-Extra: test
|
|
14
|
+
Requires-Dist: pytest; extra == "test"
|
|
15
|
+
Dynamic: license-file
|
|
16
|
+
|
|
17
|
+
# callserver
|
|
18
|
+
|
|
19
|
+
Fast inter-process function calls for Python, with numpy arrays passed through shared memory.
|
|
20
|
+
|
|
21
|
+
## Overview
|
|
22
|
+
|
|
23
|
+
`callserver` runs Python functions in another process on the same machine. Large numpy arrays in arguments and return values are moved to shared memory automatically, so they are not pickled or pushed through a pipe. Shared memory is allocated on the first call and reused on every call after that.
|
|
24
|
+
|
|
25
|
+
It is built for low-latency, high-throughput workloads that need a separate process, such as:
|
|
26
|
+
|
|
27
|
+
- Real-time robot control
|
|
28
|
+
- Vision processing pipelines
|
|
29
|
+
- Compute and data-intensive workloads that benefit from zero-copy data sharing
|
|
30
|
+
|
|
31
|
+
**Features**
|
|
32
|
+
|
|
33
|
+
- Automatic shared memory transfer of large numpy arrays, in both directions
|
|
34
|
+
- Zero-copy transfer with pre-allocated shared memory buffers
|
|
35
|
+
- Synchronous and asynchronous calls, with double-buffering to overlap work
|
|
36
|
+
- Persistent server-side state
|
|
37
|
+
- Runs the server as a child process, or connects to a standalone server over TCP
|
|
38
|
+
|
|
39
|
+
`callserver` works on a single machine only. It is not meant for distributed execution.
|
|
40
|
+
|
|
41
|
+
## Installation
|
|
42
|
+
|
|
43
|
+
```bash
|
|
44
|
+
pip install callserver
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Requires Python 3.8+ and numpy.
|
|
48
|
+
|
|
49
|
+
## Usage
|
|
50
|
+
### Simple Example
|
|
51
|
+
Arrays are transferred via shared memory both ways. No serialization!
|
|
52
|
+
Shared memory is only allocated the first time the RPC is made. Subsequent calls reuse the same buffers.
|
|
53
|
+
```python
|
|
54
|
+
import numpy as np
|
|
55
|
+
from callserver import CallServer
|
|
56
|
+
|
|
57
|
+
if __name__ == "__main__":
|
|
58
|
+
# Spawns a child process. np.add runs there; the arrays travel through shared memory.
|
|
59
|
+
with CallServer("adder") as server:
|
|
60
|
+
result = np.zeros((1024, 1024))
|
|
61
|
+
ones = np.ones((1024, 1024))
|
|
62
|
+
for _ in range(100):
|
|
63
|
+
result = server.call(np.add, result, ones)
|
|
64
|
+
print(result[0, 0]) # 100.0
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
### Server-Side State
|
|
68
|
+
A context passed to the server (instance or callable) holds state for subsequent calls:
|
|
69
|
+
|
|
70
|
+
```python
|
|
71
|
+
import numpy as np
|
|
72
|
+
from callserver import CallServer
|
|
73
|
+
|
|
74
|
+
class Accumulator:
|
|
75
|
+
def __init__(self, shape):
|
|
76
|
+
self.total = np.zeros(shape)
|
|
77
|
+
|
|
78
|
+
def accumulate(acc, values):
|
|
79
|
+
acc.total += values
|
|
80
|
+
return acc.total
|
|
81
|
+
|
|
82
|
+
if __name__ == "__main__":
|
|
83
|
+
# Arguments after the initializer are passed to it on the server.
|
|
84
|
+
with CallServer("accumulator", Accumulator, (512, 512)) as server:
|
|
85
|
+
for _ in range(10):
|
|
86
|
+
total = server.call(accumulate, np.ones((512, 512)))
|
|
87
|
+
print(total[0, 0]) # 10.0
|
|
88
|
+
```
|
|
89
|
+
|
|
90
|
+
### Standalone Server
|
|
91
|
+
A server can also run as a separate, long-lived process that clients connect to over TCP. Start it from the command line (the port is required):
|
|
92
|
+
|
|
93
|
+
```bash
|
|
94
|
+
callserver --port 5000
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
Then connect by passing a `(name, (address, port))` tuple instead of a name:
|
|
98
|
+
|
|
99
|
+
```python
|
|
100
|
+
import numpy as np
|
|
101
|
+
from callserver import CallServer
|
|
102
|
+
|
|
103
|
+
with CallServer(("adder", ("127.0.0.1", 5000))) as server:
|
|
104
|
+
result = server.call(np.add, np.ones((1024, 1024)), np.ones((1024, 1024)))
|
|
105
|
+
print(result[0, 0]) # 2.0
|
|
106
|
+
```
|
|
107
|
+
|
|
108
|
+
The server handles one client at a time and keeps running after a client disconnects. Functions sent to it must be importable by the server process.
|
|
109
|
+
|
|
110
|
+
> Server-side functions must be defined at module level so they can be pickled. The `if __name__ == "__main__":` guard is required on Windows and macOS, because the child process re-imports the main module.
|
|
111
|
+
|
|
112
|
+
## Documentation
|
|
113
|
+
|
|
114
|
+
- [Usage guide](callserver/callserver.md): async calls, double-buffering, pre-allocated buffers, standalone servers, the API reference, and how it works.
|
|
115
|
+
- [Examples](examples): runnable scripts.
|
|
116
|
+
|
|
117
|
+
| Example | Shows |
|
|
118
|
+
| --- | --- |
|
|
119
|
+
| [simple_call.py](examples/simple_call.py) | Basic synchronous call |
|
|
120
|
+
| [server_state.py](examples/server_state.py) | Server-side state with a context initializer |
|
|
121
|
+
| [async_calls.py](examples/async_calls.py) | Async calls with double-buffering |
|
|
122
|
+
| [preallocated_buffers.py](examples/preallocated_buffers.py) | Zero-copy with pre-allocated shared memory |
|
|
123
|
+
| [standalone_client.py](examples/standalone_client.py) | Connecting to a standalone server |
|
|
124
|
+
|
|
125
|
+
## Development
|
|
126
|
+
|
|
127
|
+
```bash
|
|
128
|
+
pip install -e .[test]
|
|
129
|
+
python -m pytest
|
|
130
|
+
```
|
|
131
|
+
|
|
132
|
+
## License
|
|
133
|
+
|
|
134
|
+
MIT. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,118 @@
|
|
|
1
|
+
# callserver
|
|
2
|
+
|
|
3
|
+
Fast inter-process function calls for Python, with numpy arrays passed through shared memory.
|
|
4
|
+
|
|
5
|
+
## Overview
|
|
6
|
+
|
|
7
|
+
`callserver` runs Python functions in another process on the same machine. Large numpy arrays in arguments and return values are moved to shared memory automatically, so they are not pickled or pushed through a pipe. Shared memory is allocated on the first call and reused on every call after that.
|
|
8
|
+
|
|
9
|
+
It is built for low-latency, high-throughput workloads that need a separate process, such as:
|
|
10
|
+
|
|
11
|
+
- Real-time robot control
|
|
12
|
+
- Vision processing pipelines
|
|
13
|
+
- Compute and data-intensive workloads that benefit from zero-copy data sharing
|
|
14
|
+
|
|
15
|
+
**Features**
|
|
16
|
+
|
|
17
|
+
- Automatic shared memory transfer of large numpy arrays, in both directions
|
|
18
|
+
- Zero-copy transfer with pre-allocated shared memory buffers
|
|
19
|
+
- Synchronous and asynchronous calls, with double-buffering to overlap work
|
|
20
|
+
- Persistent server-side state
|
|
21
|
+
- Runs the server as a child process, or connects to a standalone server over TCP
|
|
22
|
+
|
|
23
|
+
`callserver` works on a single machine only. It is not meant for distributed execution.
|
|
24
|
+
|
|
25
|
+
## Installation
|
|
26
|
+
|
|
27
|
+
```bash
|
|
28
|
+
pip install callserver
|
|
29
|
+
```
|
|
30
|
+
|
|
31
|
+
Requires Python 3.8+ and numpy.
|
|
32
|
+
|
|
33
|
+
## Usage
|
|
34
|
+
### Simple Example
|
|
35
|
+
Arrays are transferred via shared memory both ways. No serialization!
|
|
36
|
+
Shared memory is only allocated the first time the RPC is made. Subsequent calls reuse the same buffers.
|
|
37
|
+
```python
|
|
38
|
+
import numpy as np
|
|
39
|
+
from callserver import CallServer
|
|
40
|
+
|
|
41
|
+
if __name__ == "__main__":
|
|
42
|
+
# Spawns a child process. np.add runs there; the arrays travel through shared memory.
|
|
43
|
+
with CallServer("adder") as server:
|
|
44
|
+
result = np.zeros((1024, 1024))
|
|
45
|
+
ones = np.ones((1024, 1024))
|
|
46
|
+
for _ in range(100):
|
|
47
|
+
result = server.call(np.add, result, ones)
|
|
48
|
+
print(result[0, 0]) # 100.0
|
|
49
|
+
```
|
|
50
|
+
|
|
51
|
+
### Server-Side State
|
|
52
|
+
A context passed to the server (instance or callable) holds state for subsequent calls:
|
|
53
|
+
|
|
54
|
+
```python
|
|
55
|
+
import numpy as np
|
|
56
|
+
from callserver import CallServer
|
|
57
|
+
|
|
58
|
+
class Accumulator:
|
|
59
|
+
def __init__(self, shape):
|
|
60
|
+
self.total = np.zeros(shape)
|
|
61
|
+
|
|
62
|
+
def accumulate(acc, values):
|
|
63
|
+
acc.total += values
|
|
64
|
+
return acc.total
|
|
65
|
+
|
|
66
|
+
if __name__ == "__main__":
|
|
67
|
+
# Arguments after the initializer are passed to it on the server.
|
|
68
|
+
with CallServer("accumulator", Accumulator, (512, 512)) as server:
|
|
69
|
+
for _ in range(10):
|
|
70
|
+
total = server.call(accumulate, np.ones((512, 512)))
|
|
71
|
+
print(total[0, 0]) # 10.0
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
### Standalone Server
|
|
75
|
+
A server can also run as a separate, long-lived process that clients connect to over TCP. Start it from the command line (the port is required):
|
|
76
|
+
|
|
77
|
+
```bash
|
|
78
|
+
callserver --port 5000
|
|
79
|
+
```
|
|
80
|
+
|
|
81
|
+
Then connect by passing a `(name, (address, port))` tuple instead of a name:
|
|
82
|
+
|
|
83
|
+
```python
|
|
84
|
+
import numpy as np
|
|
85
|
+
from callserver import CallServer
|
|
86
|
+
|
|
87
|
+
with CallServer(("adder", ("127.0.0.1", 5000))) as server:
|
|
88
|
+
result = server.call(np.add, np.ones((1024, 1024)), np.ones((1024, 1024)))
|
|
89
|
+
print(result[0, 0]) # 2.0
|
|
90
|
+
```
|
|
91
|
+
|
|
92
|
+
The server handles one client at a time and keeps running after a client disconnects. Functions sent to it must be importable by the server process.
|
|
93
|
+
|
|
94
|
+
> Server-side functions must be defined at module level so they can be pickled. The `if __name__ == "__main__":` guard is required on Windows and macOS, because the child process re-imports the main module.
|
|
95
|
+
|
|
96
|
+
## Documentation
|
|
97
|
+
|
|
98
|
+
- [Usage guide](callserver/callserver.md): async calls, double-buffering, pre-allocated buffers, standalone servers, the API reference, and how it works.
|
|
99
|
+
- [Examples](examples): runnable scripts.
|
|
100
|
+
|
|
101
|
+
| Example | Shows |
|
|
102
|
+
| --- | --- |
|
|
103
|
+
| [simple_call.py](examples/simple_call.py) | Basic synchronous call |
|
|
104
|
+
| [server_state.py](examples/server_state.py) | Server-side state with a context initializer |
|
|
105
|
+
| [async_calls.py](examples/async_calls.py) | Async calls with double-buffering |
|
|
106
|
+
| [preallocated_buffers.py](examples/preallocated_buffers.py) | Zero-copy with pre-allocated shared memory |
|
|
107
|
+
| [standalone_client.py](examples/standalone_client.py) | Connecting to a standalone server |
|
|
108
|
+
|
|
109
|
+
## Development
|
|
110
|
+
|
|
111
|
+
```bash
|
|
112
|
+
pip install -e .[test]
|
|
113
|
+
python -m pytest
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
## License
|
|
117
|
+
|
|
118
|
+
MIT. See [LICENSE](LICENSE).
|
|
@@ -0,0 +1,54 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Shared-memory RPC module for high-performance inter-process communication.
|
|
3
|
+
|
|
4
|
+
This module provides an efficient Remote Procedure Call (RPC) mechanism that uses
|
|
5
|
+
shared memory for transferring large numpy arrays between processes on the same machine.
|
|
6
|
+
It is designed for use cases where low-latency, high-throughput data exchange is required,
|
|
7
|
+
such as real-time robot control, vision processing pipelines, or any compute-intensive
|
|
8
|
+
workload that benefits from zero-copy data sharing.
|
|
9
|
+
|
|
10
|
+
Key Features:
|
|
11
|
+
- Automatic promotion of large numpy arrays to shared memory for zero-copy transfer
|
|
12
|
+
- Support for both synchronous and asynchronous calls
|
|
13
|
+
- Front/back buffer pattern for overlapping computation and communication
|
|
14
|
+
- Connection to child processes or existing server processes via TCP
|
|
15
|
+
|
|
16
|
+
Limitations:
|
|
17
|
+
- Only supports inter-process communication on the same machine (shared memory is local)
|
|
18
|
+
- Cannot be used for distributed or remote execution across machines
|
|
19
|
+
|
|
20
|
+
Classes:
|
|
21
|
+
CallServer: Main client interface for making RPC calls to a server process.
|
|
22
|
+
AllocationSet: Manages a set of shared memory allocations for a single function.
|
|
23
|
+
AllocationDescriptor: Descriptor for a shared memory buffer (name, shape, dtype).
|
|
24
|
+
AllocationSetManager: Manages multiple AllocationSet instances for resource allocation.
|
|
25
|
+
|
|
26
|
+
Functions:
|
|
27
|
+
run: Run the server loop, either on a child-process connection or as a
|
|
28
|
+
standalone TCP server on a given (address, port, authkey) endpoint.
|
|
29
|
+
run_console: Command-line entry point for running a standalone server.
|
|
30
|
+
|
|
31
|
+
Example:
|
|
32
|
+
>>> import numpy as np
|
|
33
|
+
>>> from callserver import CallServer
|
|
34
|
+
>>> with CallServer('adder') as server:
|
|
35
|
+
... result = server.call(np.add, np.ones((1024, 1024)), np.ones((1024, 1024)))
|
|
36
|
+
>>> print(result[0, 0])
|
|
37
|
+
2.0
|
|
38
|
+
|
|
39
|
+
For more details, see the callserver.md documentation file.
|
|
40
|
+
"""
|
|
41
|
+
|
|
42
|
+
from .allocationmanager import AllocationSetManager
|
|
43
|
+
from .allocationset import AllocationDescriptor, AllocationSet
|
|
44
|
+
from .callserver import DEFAULT_AUTHKEY, CallServer, run, run_console
|
|
45
|
+
|
|
46
|
+
__all__ = [
|
|
47
|
+
"CallServer",
|
|
48
|
+
"run",
|
|
49
|
+
"run_console",
|
|
50
|
+
"DEFAULT_AUTHKEY",
|
|
51
|
+
"AllocationSet",
|
|
52
|
+
"AllocationDescriptor",
|
|
53
|
+
"AllocationSetManager",
|
|
54
|
+
]
|
|
@@ -0,0 +1,176 @@
|
|
|
1
|
+
"""
|
|
2
|
+
AllocationSetManager for managing multiple function allocation sets.
|
|
3
|
+
|
|
4
|
+
This module provides the AllocationSetManager class which coordinates multiple
|
|
5
|
+
AllocationSet instances for different functions in an RPC workflow.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
from typing import Callable, Dict, List, Optional, Union
|
|
11
|
+
|
|
12
|
+
from .allocationset import AllocationSet
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class AllocationSetManager:
|
|
16
|
+
"""
|
|
17
|
+
Manages multiple AllocationSet instances for resource allocation in RPC workflows.
|
|
18
|
+
|
|
19
|
+
This class implements a singleton pattern per name, ensuring that all references
|
|
20
|
+
to the same named manager share state. It supports front- and back-buffering
|
|
21
|
+
for function arguments, tracks active sets, and handles cleanup.
|
|
22
|
+
|
|
23
|
+
The manager provides multiple allocation sets per function to support asynchronous
|
|
24
|
+
call patterns where one set is being used for a call in progress while another
|
|
25
|
+
is being prepared for the next call.
|
|
26
|
+
|
|
27
|
+
Attributes:
|
|
28
|
+
_managers: Class-level dict of all manager instances by name.
|
|
29
|
+
|
|
30
|
+
Example:
|
|
31
|
+
>>> manager = AllocationSetManager("my_server")
|
|
32
|
+
>>> alloc_set = manager.add_allocation_set("my_function")
|
|
33
|
+
>>> # Add a second set for async double-buffering
|
|
34
|
+
>>> alloc_set2 = manager.add_allocation_set("my_function")
|
|
35
|
+
"""
|
|
36
|
+
|
|
37
|
+
_managers: Dict[str, "AllocationSetManager"] = {}
|
|
38
|
+
|
|
39
|
+
def __new__(cls, name: Optional[str] = None) -> "AllocationSetManager":
|
|
40
|
+
"""
|
|
41
|
+
Return existing manager for the given name, or create a new one.
|
|
42
|
+
|
|
43
|
+
Args:
|
|
44
|
+
name: The unique name for this manager.
|
|
45
|
+
|
|
46
|
+
Returns:
|
|
47
|
+
An AllocationSetManager instance (existing or new).
|
|
48
|
+
"""
|
|
49
|
+
if name is not None and name in cls._managers:
|
|
50
|
+
return cls._managers[name]
|
|
51
|
+
return super().__new__(cls)
|
|
52
|
+
|
|
53
|
+
def __init__(self, name: Optional[str] = None) -> None:
|
|
54
|
+
"""
|
|
55
|
+
Initialize the AllocationSetManager.
|
|
56
|
+
|
|
57
|
+
If a manager with the given name already exists, this is a no-op.
|
|
58
|
+
|
|
59
|
+
Args:
|
|
60
|
+
name: The unique name for this manager.
|
|
61
|
+
"""
|
|
62
|
+
if name is not None and name in AllocationSetManager._managers:
|
|
63
|
+
return
|
|
64
|
+
self._name = name
|
|
65
|
+
self._arg_sets: Dict[str, List[AllocationSet]] = {}
|
|
66
|
+
self._active_set_ids: Dict[str, int] = {}
|
|
67
|
+
if self._name is not None:
|
|
68
|
+
assert (
|
|
69
|
+
self._name not in AllocationSetManager._managers
|
|
70
|
+
), f"An AllocationSetManager with name {self._name} already exists."
|
|
71
|
+
AllocationSetManager._managers[self._name] = self
|
|
72
|
+
|
|
73
|
+
def _make_name(self, fn: Union[Callable, str]) -> str:
|
|
74
|
+
"""
|
|
75
|
+
Convert a function or string to a string name.
|
|
76
|
+
|
|
77
|
+
Args:
|
|
78
|
+
fn: A callable or string identifier.
|
|
79
|
+
|
|
80
|
+
Returns:
|
|
81
|
+
The function's __name__ attribute, or the string itself.
|
|
82
|
+
"""
|
|
83
|
+
return fn.__name__ if callable(fn) else fn
|
|
84
|
+
|
|
85
|
+
def add_allocation_set(self, fn: Union[Callable, str]) -> AllocationSet:
|
|
86
|
+
"""
|
|
87
|
+
Add a new AllocationSet for the given function.
|
|
88
|
+
|
|
89
|
+
Creates a new allocation set and adds it to the list of sets for the
|
|
90
|
+
specified function. Multiple sets can be created for double-buffering.
|
|
91
|
+
|
|
92
|
+
Args:
|
|
93
|
+
fn: A callable or string identifier for the function.
|
|
94
|
+
|
|
95
|
+
Returns:
|
|
96
|
+
The newly created AllocationSet.
|
|
97
|
+
"""
|
|
98
|
+
fn_name = self._make_name(fn)
|
|
99
|
+
if fn_name not in self._arg_sets:
|
|
100
|
+
self._arg_sets[fn_name] = []
|
|
101
|
+
self._active_set_ids[fn_name] = -1
|
|
102
|
+
|
|
103
|
+
arg_set = AllocationSet(
|
|
104
|
+
len(self._arg_sets[fn_name]), f"{self._name}.{fn_name}"
|
|
105
|
+
)
|
|
106
|
+
self._arg_sets[fn_name].append(arg_set)
|
|
107
|
+
return arg_set
|
|
108
|
+
|
|
109
|
+
def get_allocation_set(
|
|
110
|
+
self, fn: Union[Callable, str], idx: Optional[int] = None
|
|
111
|
+
) -> AllocationSet:
|
|
112
|
+
"""
|
|
113
|
+
Get an AllocationSet for the given function.
|
|
114
|
+
|
|
115
|
+
If idx is None, rotates to the next allocation set (for double-buffering).
|
|
116
|
+
If idx is specified, returns or creates the set at that index.
|
|
117
|
+
|
|
118
|
+
Args:
|
|
119
|
+
fn: A callable or string identifier for the function.
|
|
120
|
+
idx: Optional specific index to retrieve.
|
|
121
|
+
|
|
122
|
+
Returns:
|
|
123
|
+
The requested AllocationSet.
|
|
124
|
+
"""
|
|
125
|
+
fn_name = self._make_name(fn)
|
|
126
|
+
if fn_name not in self._arg_sets:
|
|
127
|
+
self.add_allocation_set(fn_name)
|
|
128
|
+
if idx is None:
|
|
129
|
+
idx = self._active_set_ids[fn_name]
|
|
130
|
+
idx = (1 + idx) % len(self._arg_sets[fn_name])
|
|
131
|
+
self._active_set_ids[fn_name] = idx
|
|
132
|
+
else:
|
|
133
|
+
for i in range(len(self._arg_sets[fn_name]), idx + 1):
|
|
134
|
+
self.add_allocation_set(fn_name)
|
|
135
|
+
active_set = self._arg_sets[fn_name][idx]
|
|
136
|
+
return active_set
|
|
137
|
+
|
|
138
|
+
def get_active_id(self, fn: Union[Callable, str]) -> int:
|
|
139
|
+
"""
|
|
140
|
+
Get the currently active allocation set index for a function.
|
|
141
|
+
|
|
142
|
+
Args:
|
|
143
|
+
fn: A callable or string identifier for the function.
|
|
144
|
+
|
|
145
|
+
Returns:
|
|
146
|
+
The index of the currently active AllocationSet.
|
|
147
|
+
"""
|
|
148
|
+
fn_name = self._make_name(fn)
|
|
149
|
+
return self._active_set_ids[fn_name]
|
|
150
|
+
|
|
151
|
+
def close(self) -> None:
|
|
152
|
+
"""
|
|
153
|
+
Release all resources held by all allocation sets.
|
|
154
|
+
|
|
155
|
+
Closes all allocation sets for all functions managed by this manager.
|
|
156
|
+
After calling this method, the manager should not be used.
|
|
157
|
+
"""
|
|
158
|
+
for arg_set_group in self._arg_sets.values():
|
|
159
|
+
for arg_set in arg_set_group:
|
|
160
|
+
arg_set.close()
|
|
161
|
+
self._arg_sets = {}
|
|
162
|
+
|
|
163
|
+
def __del__(self) -> None:
|
|
164
|
+
"""Clean up resources when the manager is garbage collected."""
|
|
165
|
+
self.close()
|
|
166
|
+
|
|
167
|
+
|
|
168
|
+
if __name__ == "__main__":
|
|
169
|
+
manager = AllocationSetManager("test")
|
|
170
|
+
manager.add_allocation_set("func1")
|
|
171
|
+
manager2 = AllocationSetManager("test")
|
|
172
|
+
assert manager is manager2
|
|
173
|
+
assert manager2.get_allocation_set("func1") != None
|
|
174
|
+
|
|
175
|
+
print("Test passed.")
|
|
176
|
+
|