rarg-ray-patterns 0.0.1__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.
@@ -0,0 +1,6 @@
1
+ version: 2
2
+ updates:
3
+ - package-ecosystem: github-actions
4
+ directory: /
5
+ schedule:
6
+ interval: weekly
@@ -0,0 +1,30 @@
1
+ name: Continuous Integration
2
+
3
+ on:
4
+ push: { branches: [main], tags: ["*"] }
5
+ pull_request:
6
+ schedule: [{ cron: '30 2 * * 1,4' }] # Every Monday and Thursday @ 2h30am UTC
7
+
8
+ jobs:
9
+ test:
10
+ uses: ratt-ru/rarg-gh-workflows/.github/workflows/python-ci.yml@main
11
+ with:
12
+ os-matrix: '["ubuntu-24.04"]' # Add other OS (macos-15) if necessary
13
+ install-command: 'pip install --group test .'
14
+
15
+ deploy:
16
+ needs: ["test"]
17
+ runs-on: ubuntu-latest
18
+ environment:
19
+ name: pypi
20
+ url: https://pypi.org/p/rarg-ray-patterns
21
+ permissions:
22
+ id-token: write
23
+ steps:
24
+ - uses: actions/setup-python@v6
25
+ with: { python-version: '3.13' }
26
+ - uses: actions/checkout@v6
27
+ - run: pip install -U pip build && python -m build --sdist --wheel
28
+ - name: Publish distribution 📦 to PyPI
29
+ if: github.event_name == 'push' && startsWith(github.event.ref, 'refs/tags')
30
+ uses: pypa/gh-action-pypi-publish@release/v1
@@ -0,0 +1,220 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+
6
+ # C extensions
7
+ *.so
8
+
9
+ # Distribution / packaging
10
+ .Python
11
+ build/
12
+ develop-eggs/
13
+ dist/
14
+ downloads/
15
+ eggs/
16
+ .eggs/
17
+ lib/
18
+ lib64/
19
+ parts/
20
+ sdist/
21
+ var/
22
+ wheels/
23
+ share/python-wheels/
24
+ *.egg-info/
25
+ .installed.cfg
26
+ *.egg
27
+ MANIFEST
28
+
29
+ # PyInstaller
30
+ # Usually these files are written by a python script from a template
31
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
32
+ *.manifest
33
+ *.spec
34
+
35
+ # Installer logs
36
+ pip-log.txt
37
+ pip-delete-this-directory.txt
38
+
39
+ # Unit test / coverage reports
40
+ htmlcov/
41
+ .tox/
42
+ .nox/
43
+ .coverage
44
+ .coverage.*
45
+ .cache
46
+ nosetests.xml
47
+ coverage.xml
48
+ *.cover
49
+ *.py.cover
50
+ *.lcov
51
+ .hypothesis/
52
+ .pytest_cache/
53
+ cover/
54
+
55
+ # Translations
56
+ *.mo
57
+ *.pot
58
+
59
+ # Django stuff:
60
+ *.log
61
+ local_settings.py
62
+ db.sqlite3
63
+ db.sqlite3-journal
64
+
65
+ # Flask stuff:
66
+ instance/
67
+ .webassets-cache
68
+
69
+ # Scrapy stuff:
70
+ .scrapy
71
+
72
+ # Sphinx documentation
73
+ docs/_build/
74
+
75
+ # PyBuilder
76
+ .pybuilder/
77
+ target/
78
+
79
+ # Jupyter Notebook
80
+ .ipynb_checkpoints
81
+
82
+ # IPython
83
+ profile_default/
84
+ ipython_config.py
85
+
86
+ # pyenv
87
+ # For a library or package, you might want to ignore these files since the code is
88
+ # intended to run in multiple environments; otherwise, check them in:
89
+ # .python-version
90
+
91
+ # pipenv
92
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
93
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
94
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
95
+ # install all needed dependencies.
96
+ # Pipfile.lock
97
+
98
+ # UV
99
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
100
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
101
+ # commonly ignored for libraries.
102
+ # uv.lock
103
+
104
+ # poetry
105
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
106
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
107
+ # commonly ignored for libraries.
108
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
109
+ # poetry.lock
110
+ # poetry.toml
111
+
112
+ # pdm
113
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
114
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
115
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
116
+ # pdm.lock
117
+ # pdm.toml
118
+ .pdm-python
119
+ .pdm-build/
120
+
121
+ # pixi
122
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
123
+ # pixi.lock
124
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
125
+ # in the .venv directory. It is recommended not to include this directory in version control.
126
+ .pixi/*
127
+ !.pixi/config.toml
128
+
129
+ # PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
130
+ __pypackages__/
131
+
132
+ # Celery stuff
133
+ celerybeat-schedule*
134
+ celerybeat.pid
135
+
136
+ # Redis
137
+ *.rdb
138
+ *.aof
139
+ *.pid
140
+
141
+ # RabbitMQ
142
+ mnesia/
143
+ rabbitmq/
144
+ rabbitmq-data/
145
+
146
+ # ActiveMQ
147
+ activemq-data/
148
+
149
+ # SageMath parsed files
150
+ *.sage.py
151
+
152
+ # Environments
153
+ .env
154
+ .envrc
155
+ .venv
156
+ env/
157
+ venv/
158
+ ENV/
159
+ env.bak/
160
+ venv.bak/
161
+
162
+ # Spyder project settings
163
+ .spyderproject
164
+ .spyproject
165
+
166
+ # Rope project settings
167
+ .ropeproject
168
+
169
+ # mkdocs documentation
170
+ /site
171
+
172
+ # mypy
173
+ .mypy_cache/
174
+ .dmypy.json
175
+ dmypy.json
176
+
177
+ # Pyre type checker
178
+ .pyre/
179
+
180
+ # pytype static type analyzer
181
+ .pytype/
182
+
183
+ # Cython debug symbols
184
+ cython_debug/
185
+
186
+ # PyCharm
187
+ # JetBrains specific template is maintained in a separate JetBrains.gitignore that can
188
+ # be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
189
+ # and can be added to the global gitignore or merged into this file. For a more nuclear
190
+ # option (not recommended) you can uncomment the following to ignore the entire idea folder.
191
+ # .idea/
192
+
193
+ # Abstra
194
+ # Abstra is an AI-powered process automation framework.
195
+ # Ignore directories containing user credentials, local state, and settings.
196
+ # Learn more at https://abstra.io/docs
197
+ .abstra/
198
+
199
+ # Visual Studio Code
200
+ # Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
201
+ # that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
202
+ # and can be added to the global gitignore or merged into this file. However, if you prefer,
203
+ # you could uncomment the following to ignore the entire vscode folder
204
+ # .vscode/
205
+ # Temporary file for partial code execution
206
+ tempCodeRunnerFile.py
207
+
208
+ # Ruff stuff:
209
+ .ruff_cache/
210
+
211
+ # PyPI configuration file
212
+ .pypirc
213
+
214
+ # Marimo
215
+ marimo/_static/
216
+ marimo/_lsp/
217
+ __marimo__/
218
+
219
+ # Streamlit
220
+ .streamlit/secrets.toml
@@ -0,0 +1,27 @@
1
+ # See https://pre-commit.com for more information
2
+ # See https://pre-commit.com/hooks.html for more hooks
3
+ # Run `pre-commit autoupdate` after scaffolding to pin current hook revs.
4
+ repos:
5
+ - repo: https://github.com/pre-commit/pre-commit-hooks
6
+ rev: v5.0.0
7
+ hooks:
8
+ - id: trailing-whitespace
9
+ - id: end-of-file-fixer
10
+ - id: check-yaml
11
+ - id: check-added-large-files
12
+ - repo: https://github.com/astral-sh/ruff-pre-commit
13
+ rev: v0.12.3
14
+ hooks:
15
+ - id: ruff
16
+ # Handle imports, see
17
+ # https://github.com/astral-sh/ruff/issues/8232
18
+ # https://github.com/astral-sh/ruff/issues/10882
19
+ args: [ --fix, --extend-select, I ]
20
+ - id: ruff-format
21
+ name: ruff format
22
+ - repo: https://github.com/pre-commit/mirrors-mypy
23
+ rev: v2.1.0
24
+ hooks:
25
+ - id: mypy
26
+ name: mypy
27
+ additional_dependencies: [pytest, ray]
@@ -0,0 +1,28 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2026, Rhodes University Centre for Radio Astronomy Techniques & Technologies (RATT)
4
+
5
+ Redistribution and use in source and binary forms, with or without
6
+ modification, are permitted provided that the following conditions are met:
7
+
8
+ 1. Redistributions of source code must retain the above copyright notice, this
9
+ list of conditions and the following disclaimer.
10
+
11
+ 2. Redistributions in binary form must reproduce the above copyright notice,
12
+ this list of conditions and the following disclaimer in the documentation
13
+ and/or other materials provided with the distribution.
14
+
15
+ 3. Neither the name of the copyright holder nor the names of its
16
+ contributors may be used to endorse or promote products derived from
17
+ this software without specific prior written permission.
18
+
19
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
22
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
23
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
24
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
25
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
26
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
27
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
28
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,24 @@
1
+ Metadata-Version: 2.4
2
+ Name: rarg-ray-patterns
3
+ Version: 0.0.1
4
+ Summary: Common Ray distributed computing patterns for RARG codebases
5
+ Project-URL: Homepage, https://github.com/ratt-ru/rarg-ray-patterns
6
+ Project-URL: Repository, https://github.com/ratt-ru/rarg-ray-patterns
7
+ Author-email: Simon Perkins <simon.perkins@gmail.com>
8
+ License-Expression: BSD-3-Clause
9
+ License-File: LICENSE
10
+ Requires-Python: >=3.11
11
+ Requires-Dist: ray>=2.55.1
12
+ Description-Content-Type: text/markdown
13
+
14
+ # rarg-ray-patterns
15
+
16
+ Common Ray distributed computing patterns for RARG codebases
17
+
18
+ ## Development
19
+ This package is managed with [uv](https://docs.astral.sh/uv/).
20
+
21
+ - Install: `uv sync`
22
+ - Test: `uv run --group test pytest tests/`
23
+ - Pre-commit hooks gate commits: `uv run --group dev pre-commit run -a`
24
+ - Versioning is via `tbump` — never hand-edit version strings.
@@ -0,0 +1,11 @@
1
+ # rarg-ray-patterns
2
+
3
+ Common Ray distributed computing patterns for RARG codebases
4
+
5
+ ## Development
6
+ This package is managed with [uv](https://docs.astral.sh/uv/).
7
+
8
+ - Install: `uv sync`
9
+ - Test: `uv run --group test pytest tests/`
10
+ - Pre-commit hooks gate commits: `uv run --group dev pre-commit run -a`
11
+ - Versioning is via `tbump` — never hand-edit version strings.
@@ -0,0 +1,59 @@
1
+ [build-system]
2
+ requires = ["hatchling"]
3
+ build-backend = "hatchling.build"
4
+
5
+ [project]
6
+ name = "rarg-ray-patterns"
7
+ version = "0.0.1"
8
+ description = "Common Ray distributed computing patterns for RARG codebases"
9
+ readme = "README.md"
10
+ requires-python = ">=3.11"
11
+ license = "BSD-3-Clause"
12
+ authors = [{ name = "Simon Perkins", email = "simon.perkins@gmail.com" }]
13
+ dependencies = [
14
+ "ray>=2.55.1",
15
+ ]
16
+
17
+ [project.urls]
18
+ Homepage = "https://github.com/ratt-ru/rarg-ray-patterns"
19
+ Repository = "https://github.com/ratt-ru/rarg-ray-patterns"
20
+
21
+ [dependency-groups]
22
+ # ray[default] pulls in the dashboard / state-API server that the tests use to
23
+ # query scheduled actors via ray.util.state.
24
+ test = ["pytest", "ray[default]"]
25
+ dev = ["pre-commit", "mypy", "ruff", "tbump"]
26
+
27
+ [tool.hatch.build.targets.wheel]
28
+ packages = ["src/rarg_ray_patterns"]
29
+
30
+ [tool.ruff]
31
+ line-length = 88
32
+ indent-width = 2
33
+
34
+ [tool.ruff.lint]
35
+ extend-select = ["I"] # import sorting
36
+
37
+ [tool.mypy]
38
+ python_version = "3.13"
39
+ strict = true
40
+
41
+ [tool.pytest.ini_options]
42
+ testpaths = ["tests"]
43
+
44
+ [tool.tbump.version]
45
+ current = "0.0.1"
46
+ # https://semver.org/#spec-item-9
47
+ regex = '''(?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)-?(?P<ptype>(alpha|beta|rc)\.(?P<prerelease>\d+))?'''
48
+
49
+ [tool.tbump.git]
50
+ message_template = "Bump to {new_version}"
51
+ tag_template = "{new_version}"
52
+
53
+ [[tool.tbump.file]]
54
+ src = "pyproject.toml"
55
+ search = 'version = "{current_version}"'
56
+
57
+ [[tool.tbump.file]]
58
+ src = "src/rarg_ray_patterns/__init__.py"
59
+ search = '__version__ = "{current_version}"'
@@ -0,0 +1,3 @@
1
+ """Common Ray distributed computing patterns for RARG codebases"""
2
+
3
+ __version__ = "0.0.1"
@@ -0,0 +1,203 @@
1
+ import asyncio
2
+ import time
3
+ from typing import Any
4
+
5
+ import ray
6
+ from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy
7
+
8
+ from rarg_ray_patterns.utils import wrap_future
9
+
10
+ KEEPALIVE_DURATION = 60.0
11
+
12
+ # Sentinel distinguishing "head node id not yet looked up" from "no head found".
13
+ _UNSET = object()
14
+
15
+
16
+ @ray.remote
17
+ class MonitorActor:
18
+ async def heartbeat(self) -> None:
19
+ """Block forever; the awaiter sees RayActorError when the actor dies."""
20
+ await asyncio.Event().wait()
21
+
22
+
23
+ @ray.remote
24
+ def keep_alive(timeout: float) -> None:
25
+ """Sleep for timeout in order to keep the worker (and node) alive"""
26
+ time.sleep(timeout)
27
+
28
+
29
+ @ray.remote(num_cpus=0, num_returns=2)
30
+ def discover_node() -> tuple[Any, Any]:
31
+ """Return this node's id and a keep-alive ref pinned to the same node.
32
+
33
+ The keep-alive task takes 1 CPU on *this* node specifically (via
34
+ NodeAffinity), so the autoscaler will not reap the node before the caller
35
+ has a chance to schedule a long-lived actor on it.
36
+ """
37
+ node_id = ray.get_runtime_context().get_node_id()
38
+ keepalive_ref = keep_alive.options(
39
+ num_cpus=1,
40
+ scheduling_strategy=NodeAffinitySchedulingStrategy(node_id, soft=False),
41
+ ).remote(KEEPALIVE_DURATION)
42
+ return node_id, keepalive_ref
43
+
44
+
45
+ class ActorAutoscaler:
46
+ def __init__(
47
+ self, nworkers: int = 10, batch_size: int = 2, install_on_head: bool = False
48
+ ):
49
+ self._nworkers: int = nworkers
50
+ self._batch_size: int = batch_size
51
+ self._install_on_head: bool = install_on_head
52
+ # Resolved lazily on first autoscale(); cached because the head node id is
53
+ # stable for the cluster's lifetime. _UNSET until looked up.
54
+ self._head_node_id: Any = _UNSET
55
+ self._workers: dict[str, ray.actor.ActorHandle[Any]] = {}
56
+ # node-id future -> (node-id ObjectRef, keepalive ObjectRef)
57
+ self._pending: dict[
58
+ asyncio.Future[Any], tuple[ray.ObjectRef[Any], ray.ObjectRef[Any]]
59
+ ] = {}
60
+ # heartbeat futures held to keep them from being GC'd
61
+ self._heartbeats: set[asyncio.Future[Any]] = set()
62
+ self._closed: bool = False
63
+
64
+ async def close(self) -> None:
65
+ self._closed = True
66
+
67
+ # Cancel pending discovery tasks; collect their futures so we can drain
68
+ # them below. Without that drain, asyncio warns "Future exception was
69
+ # never retrieved" when the TaskCancelledError-bearing futures are GC'd.
70
+ pending_futures = list(self._pending.keys())
71
+ for node_ref, keepalive_ref in self._pending.values():
72
+ ray.cancel(node_ref, force=True)
73
+ ray.cancel(keepalive_ref, force=True)
74
+ self._pending.clear()
75
+
76
+ # Kill monitor actors. Their heartbeat futures will then resolve with
77
+ # RayActorError; same drain rule applies.
78
+ heartbeats = list(self._heartbeats)
79
+ self._heartbeats.clear()
80
+ for worker in self._workers.values():
81
+ ray.kill(worker)
82
+ self._workers.clear()
83
+
84
+ if pending_futures or heartbeats:
85
+ await asyncio.gather(*pending_futures, *heartbeats, return_exceptions=True)
86
+
87
+ def _resolve_head_node_id(self) -> str | None:
88
+ """Return the head node's id (cached), or None if no head node is found.
89
+
90
+ The head node is the live node carrying Ray's auto-resource
91
+ ``node:__internal_head__``.
92
+ """
93
+ if self._head_node_id is _UNSET:
94
+ self._head_node_id = next(
95
+ (
96
+ n["NodeID"]
97
+ for n in ray.nodes() # type: ignore[no-untyped-call]
98
+ if n.get("Alive") and "node:__internal_head__" in n.get("Resources", {})
99
+ ),
100
+ None,
101
+ )
102
+ return self._head_node_id # type: ignore[no-any-return]
103
+
104
+ async def autoscale(self, timeout: float = 1.0) -> None:
105
+ # Drain completed discoveries first so capacity decisions reflect them.
106
+ if self._pending:
107
+ ready, _ = await asyncio.wait(
108
+ self._pending.keys(),
109
+ return_when="ALL_COMPLETED",
110
+ timeout=timeout,
111
+ )
112
+ self._handle_ready(ready)
113
+
114
+ # If we hit capacity (e.g. nworkers shrunk, or a burst overshot),
115
+ # cancel surplus pending work instead of letting it leak extra nodes.
116
+ if len(self._workers) >= self._nworkers:
117
+ for node_ref, keepalive_ref in self._pending.values():
118
+ ray.cancel(node_ref, force=True)
119
+ ray.cancel(keepalive_ref, force=True)
120
+ self._pending.clear()
121
+ print(f"Autoscaler at capacity: {list(self._workers)}")
122
+ return
123
+
124
+ shortfall = self._nworkers - len(self._workers) - len(self._pending)
125
+ requested = min(shortfall, self._batch_size)
126
+ if requested <= 0:
127
+ return
128
+
129
+ # Keep discovery off already-monitored nodes and, unless explicitly enabled,
130
+ # off the head node too.
131
+ excluded = set(self._workers)
132
+ if not self._install_on_head:
133
+ if (head := self._resolve_head_node_id()) is not None:
134
+ excluded.add(head)
135
+
136
+ options: dict[str, Any] = {"num_cpus": 0, "scheduling_strategy": "SPREAD"}
137
+ if excluded:
138
+ options["label_selector"] = {"ray.io/node-id": f"!in({','.join(excluded)})"}
139
+
140
+ print(f"Requesting {requested} nodes")
141
+ for _ in range(requested):
142
+ # num_returns=2 makes .remote() yield a 2-tuple of ObjectRefs, but the
143
+ # type stubs only describe the single-return case.
144
+ node_ref, keepalive_ref = discover_node.options(**options).remote() # type: ignore[misc]
145
+ future = wrap_future(node_ref)
146
+ self._pending[future] = (node_ref, keepalive_ref)
147
+
148
+ def _handle_ready(self, ready: set[asyncio.Future[Any]]) -> None:
149
+ new_nodes = []
150
+ for future in ready:
151
+ _, keepalive_ref = self._pending.pop(future)
152
+
153
+ if future.cancelled():
154
+ ray.cancel(keepalive_ref, force=True)
155
+ continue
156
+
157
+ if exc := future.exception():
158
+ print(f"node discovery failed: {exc!r}")
159
+ ray.cancel(keepalive_ref, force=True)
160
+ continue
161
+
162
+ node_id = future.result()
163
+ if node_id in self._workers:
164
+ # Already monitored; release the keepalive so the node can return
165
+ # to its existing actor's exclusive use.
166
+ ray.cancel(keepalive_ref, force=True)
167
+ continue
168
+
169
+ if not self._install_on_head and node_id == self._resolve_head_node_id():
170
+ # Discovery raced ahead of head-node resolution and landed on the head
171
+ # node; drop it and release the keepalive.
172
+ ray.cancel(keepalive_ref, force=True)
173
+ continue
174
+
175
+ # Pin a long-lived monitor actor to this node. num_cpus=1 makes it the
176
+ # node's permanent CPU consumer so the autoscaler keeps the node alive.
177
+ monitor = MonitorActor.options( # type: ignore[attr-defined]
178
+ num_cpus=0,
179
+ label_selector={"ray.io/node-id": node_id},
180
+ ).remote()
181
+ self._workers[node_id] = monitor
182
+ new_nodes.append(node_id)
183
+
184
+ # Free the CPU so the queued actor can claim it. The autoscaler treats
185
+ # the queued actor as resource demand and will not reap the node.
186
+ ray.cancel(keepalive_ref, force=True)
187
+
188
+ # Watch for actor death so we'll reprovision next iteration.
189
+ heartbeat = wrap_future(monitor.heartbeat.remote())
190
+ self._heartbeats.add(heartbeat)
191
+ heartbeat.add_done_callback(self._heartbeats.discard)
192
+ heartbeat.add_done_callback(lambda f, nid=node_id: self._on_actor_dead(nid, f)) # type: ignore[misc]
193
+
194
+ if new_nodes:
195
+ print(f"Procured {len(new_nodes)} nodes: {new_nodes}")
196
+
197
+ def _on_actor_dead(self, node_id: str, future: asyncio.Future[Any]) -> None:
198
+ # Always retrieve the exception so asyncio doesn't warn about it.
199
+ exc = None if future.cancelled() else future.exception()
200
+ if self._closed:
201
+ return
202
+ print(f"Worker on {node_id} died ({exc!r}); will reprovision")
203
+ self._workers.pop(node_id, None)
@@ -0,0 +1,12 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ from typing import TYPE_CHECKING, Any
5
+
6
+ if TYPE_CHECKING:
7
+ import ray
8
+
9
+
10
+ def wrap_future(obj_ref: ray.ObjectRef[Any]) -> asyncio.Future[Any]:
11
+ """Wrap a ray ObjectReference as an asyncio Future"""
12
+ return asyncio.wrap_future(obj_ref.future())
File without changes
@@ -0,0 +1,27 @@
1
+ # Shared fixtures live here. Prefer fixtures over setup/teardown; keep tests
2
+ # deterministic (see the shared Python conventions).
3
+
4
+ import os
5
+ import os.path
6
+ from urllib.request import urlretrieve
7
+
8
+ import pytest
9
+ import ray
10
+
11
+ RAY_GITHUB_REPO = "https://raw.githubusercontent.com/ray-project/ray"
12
+ GITHUB_CONFIG_PATH = "python/ray/autoscaler/_private/fake_multi_node/example.yaml"
13
+ LOCAL_CONFIG_PATH = "autoscaler/_private/fake_multi_node/example.yaml"
14
+
15
+
16
+ @pytest.fixture(scope="session", autouse=True)
17
+ def install_example_yaml() -> None:
18
+ """Install the example config needed by AutoscalingCluster.
19
+
20
+ Pinned to the installed Ray version so the YAML schema can't drift away
21
+ from what the local Ray expects.
22
+ """
23
+ url = f"{RAY_GITHUB_REPO}/refs/tags/ray-{ray.__version__}/{GITHUB_CONFIG_PATH}"
24
+ file_path = os.path.join(os.path.dirname(ray.__file__), LOCAL_CONFIG_PATH)
25
+ if not os.path.exists(file_path):
26
+ os.makedirs(os.path.dirname(file_path), exist_ok=True)
27
+ urlretrieve(url, file_path)
@@ -0,0 +1,163 @@
1
+ import asyncio
2
+ import copy
3
+ import time
4
+ import typing
5
+
6
+ import pytest
7
+ import ray
8
+ from ray.cluster_utils import AutoscalingCluster
9
+ from ray.util.state import list_actors
10
+ from ray.util.state.common import ActorState
11
+
12
+ from rarg_ray_patterns.autoscaling import ActorAutoscaler
13
+
14
+ # Maximum number of cpu_node workers the cluster will provision. The autoscaler
15
+ # pins one MonitorActor per node, so this also caps how many actors can be
16
+ # scheduled regardless of how many workers are requested. Shared with the tests.
17
+ NODE_CAP = 20
18
+
19
+ # How long to wait for the cluster to provision nodes and converge. Bringing the
20
+ # fake multi-node cluster up to full capacity (NODE_CAP node processes) is slow,
21
+ # so give the convergence-to-capacity case generous headroom.
22
+ CONVERGE_TIMEOUT = 240.0
23
+ POLL_INTERVAL = 3.0
24
+ # A second sample window used to confirm the converged count is stable and
25
+ # doesn't overshoot the requested worker count.
26
+ STABILITY_WINDOW = 6.0
27
+
28
+ AUTOSCALING_CLUSTER_CONFIG = {
29
+ "head_resources": {"CPU": 0},
30
+ "worker_node_types": {
31
+ "cpu_node": {
32
+ "resources": {
33
+ "CPU": 1,
34
+ "object_store_memory": 100 * 1024 * 1024,
35
+ },
36
+ "node_config": {},
37
+ "min_workers": 0,
38
+ "max_workers": NODE_CAP,
39
+ },
40
+ },
41
+ "min_workers": 0,
42
+ "max_workers": NODE_CAP * 5,
43
+ "autoscaler_v2": True,
44
+ }
45
+
46
+
47
+ @ray.remote
48
+ class Scheduler:
49
+ def __init__(
50
+ self, nworkers: int = 10, batch_size: int = 2, install_on_head: bool = False
51
+ ):
52
+ self._autoscaler = ActorAutoscaler(nworkers, batch_size, install_on_head)
53
+ self._event = asyncio.Event()
54
+
55
+ async def run(self) -> None:
56
+ try:
57
+ while not self._event.is_set():
58
+ await asyncio.sleep(1.0)
59
+ await self._autoscaler.autoscale()
60
+ finally:
61
+ await self._autoscaler.close()
62
+
63
+ async def stop(self) -> None:
64
+ self._event.set()
65
+
66
+
67
+ def _alive_monitors() -> list[ActorState]:
68
+ """Return the live MonitorActor entries currently scheduled on the cluster."""
69
+ # list_actors is typed as returning Any; pin the element type here.
70
+ actors: list[ActorState] = list_actors(
71
+ filters=[
72
+ ("class_name", "=", "MonitorActor"),
73
+ ("state", "=", "ALIVE"),
74
+ ],
75
+ detail=True,
76
+ )
77
+ return actors
78
+
79
+
80
+ def _head_node_id() -> str:
81
+ """Return the cluster's head node id (the node with node:__internal_head__)."""
82
+ return typing.cast(
83
+ str,
84
+ next(
85
+ n["NodeID"]
86
+ for n in ray.nodes() # type: ignore[no-untyped-call]
87
+ if n.get("Alive") and "node:__internal_head__" in n.get("Resources", {})
88
+ ),
89
+ )
90
+
91
+
92
+ # The cheap nworkers=8 cases stay below the node ceiling and exercise the
93
+ # install_on_head flag in both states (head excluded vs. monitored). The single
94
+ # nworkers=100 case exercises convergence to the cluster's node capacity; it is
95
+ # slow, so we run it once with the default (head excluded) rather than for both
96
+ # flag values, since the flag behaviour is already covered by the nworkers=8 pair.
97
+ @pytest.mark.filterwarnings("ignore::FutureWarning")
98
+ @pytest.mark.parametrize(
99
+ "nworkers, batch_size, install_on_head", [(8, 2, False), (8, 2, True), (10, 2, False)]
100
+ )
101
+ def test_actor_autoscaling(
102
+ nworkers: int,
103
+ batch_size: int,
104
+ install_on_head: bool,
105
+ ) -> None:
106
+ schedulable = NODE_CAP + 1 if install_on_head else NODE_CAP
107
+ expected = min(nworkers, schedulable)
108
+
109
+ cfg = typing.cast(dict[str, typing.Any], copy.deepcopy(AUTOSCALING_CLUSTER_CONFIG))
110
+ cfg["worker_node_types"]["cpu_node"]["max_workers"] = schedulable
111
+ cfg["max_workers"] = schedulable * 5
112
+ with pytest.warns(ResourceWarning, match="unclosed file"):
113
+ cluster = AutoscalingCluster(**cfg)
114
+
115
+ try:
116
+ cluster.start()
117
+ ray.init("auto", runtime_env={"excludes": ".*"})
118
+ scheduler = Scheduler.options(num_cpus=0).remote( # type: ignore[attr-defined]
119
+ nworkers, batch_size, install_on_head
120
+ )
121
+ run_future = scheduler.run.remote()
122
+
123
+ # Poll until the cluster has scheduled the expected number of monitor
124
+ # actors, rather than waiting a fixed (and flaky) duration.
125
+ deadline = time.monotonic() + CONVERGE_TIMEOUT
126
+ monitors = _alive_monitors()
127
+ while len(monitors) < expected and time.monotonic() < deadline:
128
+ time.sleep(POLL_INTERVAL)
129
+ monitors = _alive_monitors()
130
+
131
+ assert len(monitors) == expected, (
132
+ f"expected {expected} MonitorActors scheduled, "
133
+ f"saw {len(monitors)} after {CONVERGE_TIMEOUT}s"
134
+ )
135
+
136
+ # Each actor must be pinned to a distinct node (one monitor per node).
137
+ node_ids = {m.node_id for m in monitors}
138
+ assert len(node_ids) == expected, (
139
+ f"expected actors on {expected} distinct nodes, saw {len(node_ids)}"
140
+ )
141
+
142
+ # The head node is monitored only when install_on_head is set.
143
+ head_node_id = _head_node_id()
144
+ if install_on_head:
145
+ assert head_node_id in node_ids, "expected a monitor on the head node"
146
+ else:
147
+ assert head_node_id not in node_ids, "head node should not be monitored"
148
+
149
+ # The converged count must be stable and must never overshoot the
150
+ # requested worker count.
151
+ time.sleep(STABILITY_WINDOW)
152
+ monitors = _alive_monitors()
153
+ assert len(monitors) == expected, (
154
+ f"actor count drifted to {len(monitors)} after convergence"
155
+ )
156
+ assert len(monitors) <= nworkers
157
+
158
+ ray.get(scheduler.stop.remote())
159
+ # Surface any error from the run loop.
160
+ ray.get(run_future)
161
+ finally:
162
+ ray.shutdown()
163
+ cluster.shutdown() # type: ignore[no-untyped-call]
@@ -0,0 +1,5 @@
1
+ import rarg_ray_patterns
2
+
3
+
4
+ def test_package_exposes_version() -> None:
5
+ assert rarg_ray_patterns.__version__