horus-agent 0.1.2__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.
Files changed (38) hide show
  1. horus_agent-0.1.2/.gitignore +54 -0
  2. horus_agent-0.1.2/LICENSE +21 -0
  3. horus_agent-0.1.2/PKG-INFO +130 -0
  4. horus_agent-0.1.2/README.md +118 -0
  5. horus_agent-0.1.2/pyproject.toml +124 -0
  6. horus_agent-0.1.2/src/horus_agent/__init__.py +6 -0
  7. horus_agent-0.1.2/src/horus_agent/__main__.py +13 -0
  8. horus_agent-0.1.2/src/horus_agent/blobs.py +89 -0
  9. horus_agent-0.1.2/src/horus_agent/broker.py +223 -0
  10. horus_agent-0.1.2/src/horus_agent/cli.py +153 -0
  11. horus_agent-0.1.2/src/horus_agent/config.py +62 -0
  12. horus_agent-0.1.2/src/horus_agent/daemon.py +281 -0
  13. horus_agent-0.1.2/src/horus_agent/logs.py +139 -0
  14. horus_agent-0.1.2/src/horus_agent/ops.py +297 -0
  15. horus_agent-0.1.2/src/horus_agent/protocol.py +156 -0
  16. horus_agent-0.1.2/src/horus_agent/py.typed +0 -0
  17. horus_agent-0.1.2/src/horus_agent/target.py +423 -0
  18. horus_agent-0.1.2/tests/__init__.py +6 -0
  19. horus_agent-0.1.2/tests/conftest.py +93 -0
  20. horus_agent-0.1.2/tests/fake_broker.py +272 -0
  21. horus_agent-0.1.2/tests/fake_slurm/sacct +16 -0
  22. horus_agent-0.1.2/tests/fake_slurm/sbatch +19 -0
  23. horus_agent-0.1.2/tests/fake_slurm/scancel +11 -0
  24. horus_agent-0.1.2/tests/fake_slurm/sinfo +6 -0
  25. horus_agent-0.1.2/tests/fake_slurm/squeue +14 -0
  26. horus_agent-0.1.2/tests/integration/__init__.py +6 -0
  27. horus_agent-0.1.2/tests/integration/test_debug_logs.py +150 -0
  28. horus_agent-0.1.2/tests/integration/test_primitives.py +234 -0
  29. horus_agent-0.1.2/tests/integration/test_slurm.py +213 -0
  30. horus_agent-0.1.2/tests/integration/test_slurm_cluster.py +371 -0
  31. horus_agent-0.1.2/tests/integration/test_work_dir.py +75 -0
  32. horus_agent-0.1.2/tests/integration/test_workflow.py +86 -0
  33. horus_agent-0.1.2/tests/slurm_cluster/install_agent.sh +29 -0
  34. horus_agent-0.1.2/tests/unit/__init__.py +6 -0
  35. horus_agent-0.1.2/tests/unit/test_broker.py +142 -0
  36. horus_agent-0.1.2/tests/unit/test_cli.py +168 -0
  37. horus_agent-0.1.2/tests/unit/test_ops.py +112 -0
  38. horus_agent-0.1.2/tests/unit/test_target.py +80 -0
@@ -0,0 +1,54 @@
1
+ # Python
2
+ __pycache__/
3
+ *.py[cod]
4
+ *$py.class
5
+ *.so
6
+ .Python
7
+ build/
8
+ develop-eggs/
9
+ dist/
10
+ downloads/
11
+ eggs/
12
+ .eggs/
13
+ lib/
14
+ lib64/
15
+ parts/
16
+ sdist/
17
+ var/
18
+ wheels/
19
+ *.egg-info/
20
+ .installed.cfg
21
+ *.egg
22
+ *.ipynb
23
+
24
+ # Virtual environments
25
+ venv/
26
+ .venv/
27
+ ENV/
28
+ env/
29
+
30
+ # IDEs
31
+ .vscode/
32
+ .idea/
33
+ *.swp
34
+ *.swo
35
+
36
+ # Testing
37
+ .pytest_cache/
38
+ .coverage
39
+ htmlcov/
40
+ .tox/
41
+
42
+ # OS
43
+ .DS_Store
44
+ Thumbs.db
45
+
46
+ *.xml
47
+
48
+ # i18n
49
+ *.mo
50
+
51
+ # logs
52
+ logs/
53
+
54
+ no-commit/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Temple Compute
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,130 @@
1
+ Metadata-Version: 2.5
2
+ Name: horus_agent
3
+ Version: 0.1.2
4
+ Summary: Horus target for machines that only allow outbound connections: a small agent polls a broker for target operations and runs them locally.
5
+ License-File: LICENSE
6
+ Classifier: License :: OSI Approved :: MIT License
7
+ Classifier: Programming Language :: Python :: 3
8
+ Requires-Python: >=3.13
9
+ Requires-Dist: horus-runtime>=0.8.4
10
+ Requires-Dist: httpx>=0.28
11
+ Description-Content-Type: text/markdown
12
+
13
+ # horus-agent
14
+
15
+ A Horus target for machines the orchestrator cannot connect to, for example an HPC login node or a lab server behind a VPN.
16
+
17
+ The machine runs `horus-agent`. It connects **out** to a broker over HTTPS, polls it for work and runs it locally. Nothing listens on the machine, so no inbound port or VPN access is needed.
18
+
19
+ The target, kind `agent`, works like the SSH target. Each channel primitive (`launch`, `poll`, `put_file`, `get_file`, `mkdir`, ...) becomes an op that the broker queues for the agent, and the agent runs it with Horus's own `LocalTarget`. These work on agent targets exactly as they do on SSH:
20
+
21
+ - executors (shell, containers, Python environments);
22
+ - Slurm with `inner: {kind: agent, ...}`;
23
+ - the resource monitor;
24
+ - transfers.
25
+
26
+ ## Install and register
27
+
28
+ The machine needs Python 3.13 or newer. `uv` downloads it if it is missing.
29
+
30
+ ```sh
31
+ uv tool install horus-agent # or: pip install horus-agent
32
+ horus-agent register --url https://api.example.com/api/v1 --token <one-time token>
33
+ horus-agent run
34
+ ```
35
+
36
+ `register` stores the agent's credential in `~/.config/horus-agent/agent.json` (mode 0600). To use another location, set `HORUS_AGENT_CONFIG`. The agent's name comes from the token when the admin set one in the UI, else from `--name`, else the hostname.
37
+
38
+ `--work-dir PATH` sets where tasks run on this machine; the default is `~/horus-agent`. It is created if missing, stored in the config, and reported to the broker, so workflows do not set it. `horus-agent run --work-dir PATH` overrides it for one run.
39
+
40
+ ### Keep it running
41
+
42
+ As a systemd user service (see `examples/horus-agent.service`):
43
+
44
+ ```sh
45
+ cp examples/horus-agent.service ~/.config/systemd/user/
46
+ systemctl --user enable --now horus-agent
47
+ loginctl enable-linger "$USER" # keep it running after logout
48
+ ```
49
+
50
+ Without systemd (for example on a login node), run it in `tmux` or with `nohup horus-agent run &`.
51
+
52
+ Jobs run detached, in their own session, with their output in log files. Stopping, restarting or upgrading the agent does not stop them; the orchestrator picks them up again once the agent is back.
53
+
54
+ Proxies: the agent honours `HTTPS_PROXY`, `NO_PROXY` and `SSL_CERT_FILE`.
55
+
56
+ ### Logs
57
+
58
+ By default the agent logs each op it receives at info level: the command it runs, or the path it reads or writes, with the run it belongs to. It also logs exit codes, the pids of launched jobs, and when each job exits. The status checks it repeats every few seconds for running jobs (`poll`, `read_output`) appear only in debug logs.
59
+
60
+ ### Debug logs
61
+
62
+ `horus-agent run --debug` (or `HORUS_AGENT_DEBUG=1`) logs everything the agent does:
63
+ - every TCP and TLS connection and HTTP request;
64
+ - each poll, with how long it waited and how many ops it got;
65
+ - each op with its arguments, its duration and its result;
66
+ - every command it runs, with a preview of the output;
67
+ - file transfers, job pids, cancels, and results delivered.
68
+
69
+ Secrets stay out of the log, so it can be shared. Environment variables appear by name only, because they carry task secrets. Presigned URLs lose their query strings, which hold the signatures. Tokens are never logged.
70
+
71
+ ## Use it in a workflow
72
+
73
+ ```yaml
74
+ target:
75
+ kind: agent
76
+ agent: lab-server # the agent's name
77
+ ```
78
+
79
+ Slurm on the agent's machine:
80
+
81
+ ```yaml
82
+ target:
83
+ kind: slurm
84
+ inner:
85
+ kind: agent
86
+ agent: hpc-login
87
+ ```
88
+
89
+ | Field | Default | Meaning |
90
+ |---|---|---|
91
+ | `agent` | required | Agent name |
92
+ | `working_directory` | the agent's `--work-dir` | Usually unset. Absolute: used instead of the agent's directory. Relative: taken under it |
93
+ | `offline_timeout` | `900` | Seconds without a poll before the agent's tasks count as lost |
94
+
95
+ ## Orchestrator configuration
96
+
97
+ The broker is set on the orchestrator, not in workflows:
98
+
99
+ | Variable | Meaning |
100
+ |---|---|
101
+ | `HORUS_AGENT_BROKER_URL` | Broker base URL |
102
+ | `HORUS_AGENT_BROKER_TOKEN` | Token for the broker |
103
+ | `HORUS_AGENT_BROKER_TOKEN_CONTEXT_KEY` | Instead of a fixed token: the `HorusContext.data` key holding a per-run token |
104
+
105
+ The broker HTTP API is described in `src/horus_agent/protocol.py`. Temple Compute OS implements it.
106
+
107
+ ## Limits
108
+
109
+ - Files over 4 MiB, in either direction, go through presigned object-storage URLs that the broker issues.
110
+ - Output from `run_command(detach=False)` is buffered and capped at 4 MiB. Long-running commands go through `launch`, the default for this target.
111
+ - The `python_function_external` executor is refused: it would load bytes written on the agent's machine as code in the orchestrator.
112
+ - Task commands run as the agent's Unix user and can read its config file. Run the agent as a dedicated user, the same advice as for GitHub self-hosted runners.
113
+
114
+ ## Development
115
+
116
+ ```sh
117
+ uv sync --group dev
118
+ uv run make lint
119
+ uv run make test
120
+ ```
121
+
122
+ `tests/integration/test_slurm.py` runs Slurm through an agent against a fake Slurm CLI (`tests/fake_slurm`) and is part of `make test`. To check it against a real cluster, start horus-slurm's docker cluster and run the opt-in test, which installs the agent in the login container:
123
+
124
+ ```sh
125
+ # from a horus-slurm checkout
126
+ ITEST_DIR=/tmp/horus-agent-slurm-itest-$(id -u) \
127
+ ITEST_PROJECT=horus-agent-slurm-itest tests/integration/cluster/up.sh
128
+ # from this repo
129
+ HORUS_AGENT_SLURM_ITEST=1 uv run pytest --no-cov tests/integration/test_slurm_cluster.py
130
+ ```
@@ -0,0 +1,118 @@
1
+ # horus-agent
2
+
3
+ A Horus target for machines the orchestrator cannot connect to, for example an HPC login node or a lab server behind a VPN.
4
+
5
+ The machine runs `horus-agent`. It connects **out** to a broker over HTTPS, polls it for work and runs it locally. Nothing listens on the machine, so no inbound port or VPN access is needed.
6
+
7
+ The target, kind `agent`, works like the SSH target. Each channel primitive (`launch`, `poll`, `put_file`, `get_file`, `mkdir`, ...) becomes an op that the broker queues for the agent, and the agent runs it with Horus's own `LocalTarget`. These work on agent targets exactly as they do on SSH:
8
+
9
+ - executors (shell, containers, Python environments);
10
+ - Slurm with `inner: {kind: agent, ...}`;
11
+ - the resource monitor;
12
+ - transfers.
13
+
14
+ ## Install and register
15
+
16
+ The machine needs Python 3.13 or newer. `uv` downloads it if it is missing.
17
+
18
+ ```sh
19
+ uv tool install horus-agent # or: pip install horus-agent
20
+ horus-agent register --url https://api.example.com/api/v1 --token <one-time token>
21
+ horus-agent run
22
+ ```
23
+
24
+ `register` stores the agent's credential in `~/.config/horus-agent/agent.json` (mode 0600). To use another location, set `HORUS_AGENT_CONFIG`. The agent's name comes from the token when the admin set one in the UI, else from `--name`, else the hostname.
25
+
26
+ `--work-dir PATH` sets where tasks run on this machine; the default is `~/horus-agent`. It is created if missing, stored in the config, and reported to the broker, so workflows do not set it. `horus-agent run --work-dir PATH` overrides it for one run.
27
+
28
+ ### Keep it running
29
+
30
+ As a systemd user service (see `examples/horus-agent.service`):
31
+
32
+ ```sh
33
+ cp examples/horus-agent.service ~/.config/systemd/user/
34
+ systemctl --user enable --now horus-agent
35
+ loginctl enable-linger "$USER" # keep it running after logout
36
+ ```
37
+
38
+ Without systemd (for example on a login node), run it in `tmux` or with `nohup horus-agent run &`.
39
+
40
+ Jobs run detached, in their own session, with their output in log files. Stopping, restarting or upgrading the agent does not stop them; the orchestrator picks them up again once the agent is back.
41
+
42
+ Proxies: the agent honours `HTTPS_PROXY`, `NO_PROXY` and `SSL_CERT_FILE`.
43
+
44
+ ### Logs
45
+
46
+ By default the agent logs each op it receives at info level: the command it runs, or the path it reads or writes, with the run it belongs to. It also logs exit codes, the pids of launched jobs, and when each job exits. The status checks it repeats every few seconds for running jobs (`poll`, `read_output`) appear only in debug logs.
47
+
48
+ ### Debug logs
49
+
50
+ `horus-agent run --debug` (or `HORUS_AGENT_DEBUG=1`) logs everything the agent does:
51
+ - every TCP and TLS connection and HTTP request;
52
+ - each poll, with how long it waited and how many ops it got;
53
+ - each op with its arguments, its duration and its result;
54
+ - every command it runs, with a preview of the output;
55
+ - file transfers, job pids, cancels, and results delivered.
56
+
57
+ Secrets stay out of the log, so it can be shared. Environment variables appear by name only, because they carry task secrets. Presigned URLs lose their query strings, which hold the signatures. Tokens are never logged.
58
+
59
+ ## Use it in a workflow
60
+
61
+ ```yaml
62
+ target:
63
+ kind: agent
64
+ agent: lab-server # the agent's name
65
+ ```
66
+
67
+ Slurm on the agent's machine:
68
+
69
+ ```yaml
70
+ target:
71
+ kind: slurm
72
+ inner:
73
+ kind: agent
74
+ agent: hpc-login
75
+ ```
76
+
77
+ | Field | Default | Meaning |
78
+ |---|---|---|
79
+ | `agent` | required | Agent name |
80
+ | `working_directory` | the agent's `--work-dir` | Usually unset. Absolute: used instead of the agent's directory. Relative: taken under it |
81
+ | `offline_timeout` | `900` | Seconds without a poll before the agent's tasks count as lost |
82
+
83
+ ## Orchestrator configuration
84
+
85
+ The broker is set on the orchestrator, not in workflows:
86
+
87
+ | Variable | Meaning |
88
+ |---|---|
89
+ | `HORUS_AGENT_BROKER_URL` | Broker base URL |
90
+ | `HORUS_AGENT_BROKER_TOKEN` | Token for the broker |
91
+ | `HORUS_AGENT_BROKER_TOKEN_CONTEXT_KEY` | Instead of a fixed token: the `HorusContext.data` key holding a per-run token |
92
+
93
+ The broker HTTP API is described in `src/horus_agent/protocol.py`. Temple Compute OS implements it.
94
+
95
+ ## Limits
96
+
97
+ - Files over 4 MiB, in either direction, go through presigned object-storage URLs that the broker issues.
98
+ - Output from `run_command(detach=False)` is buffered and capped at 4 MiB. Long-running commands go through `launch`, the default for this target.
99
+ - The `python_function_external` executor is refused: it would load bytes written on the agent's machine as code in the orchestrator.
100
+ - Task commands run as the agent's Unix user and can read its config file. Run the agent as a dedicated user, the same advice as for GitHub self-hosted runners.
101
+
102
+ ## Development
103
+
104
+ ```sh
105
+ uv sync --group dev
106
+ uv run make lint
107
+ uv run make test
108
+ ```
109
+
110
+ `tests/integration/test_slurm.py` runs Slurm through an agent against a fake Slurm CLI (`tests/fake_slurm`) and is part of `make test`. To check it against a real cluster, start horus-slurm's docker cluster and run the opt-in test, which installs the agent in the login container:
111
+
112
+ ```sh
113
+ # from a horus-slurm checkout
114
+ ITEST_DIR=/tmp/horus-agent-slurm-itest-$(id -u) \
115
+ ITEST_PROJECT=horus-agent-slurm-itest tests/integration/cluster/up.sh
116
+ # from this repo
117
+ HORUS_AGENT_SLURM_ITEST=1 uv run pytest --no-cov tests/integration/test_slurm_cluster.py
118
+ ```
@@ -0,0 +1,124 @@
1
+ [project]
2
+ name = "horus_agent"
3
+ description = "Horus target for machines that only allow outbound connections: a small agent polls a broker for target operations and runs them locally."
4
+ readme = "README.md"
5
+ requires-python = ">=3.13"
6
+ classifiers = [
7
+ "License :: OSI Approved :: MIT License",
8
+ "Programming Language :: Python :: 3",
9
+ ]
10
+ dependencies = ["horus-runtime>=0.8.4", "httpx>=0.28"]
11
+ dynamic = ["version"]
12
+
13
+ [project.scripts]
14
+ horus-agent = "horus_agent.cli:main"
15
+
16
+ # ================
17
+ # Plugin entry-points
18
+ # ================
19
+ [project.entry-points."horus.target"]
20
+ agent = "horus_agent.target"
21
+
22
+ # ================
23
+
24
+ [dependency-groups]
25
+ dev = [
26
+ "licenseheaders~=0.8",
27
+ "mypy~=1.19",
28
+ "pre_commit~=4.0",
29
+ "pytest~=9.0",
30
+ "pytest-asyncio~=1.3.0",
31
+ "pytest-cov~=7.0",
32
+ "ruff~=0.15",
33
+ "horus-slurm==0.6.2",
34
+ ]
35
+
36
+ [build-system]
37
+ requires = ["hatchling>=1.29,<2.0", "uv-dynamic-versioning>=0.14.0,<1.0"]
38
+ build-backend = "hatchling.build"
39
+
40
+ [tool.coverage.run]
41
+ source = ["horus_agent"]
42
+
43
+ [tool.coverage.report]
44
+ omit = ["tests/*"]
45
+
46
+ # Other tool configurations
47
+ [tool.hatch.version]
48
+ source = "uv-dynamic-versioning"
49
+
50
+ [tool.hatch.build.targets.wheel]
51
+ packages = ["src/horus_agent"]
52
+
53
+ [tool.hatch.build.targets.sdist]
54
+ include = ["src/horus_agent", "tests"]
55
+
56
+ [tool.mypy]
57
+ python_version = "3.13"
58
+ strict = true
59
+ exclude = ["build", "dist"]
60
+ plugins = ["pydantic.mypy"]
61
+
62
+ [tool.pydantic-mypy]
63
+ init_typed = true
64
+ init_forbid_extra = true
65
+ warn_required_dynamic_aliases = true
66
+
67
+ [tool.pytest.ini_options]
68
+ addopts = """
69
+ --verbose
70
+ --cov=src/horus_agent
71
+ --cov-report=html
72
+ --cov-report=term-missing
73
+ --cov-fail-under=90
74
+ """
75
+ minversion = "7.0"
76
+ python_classes = ["Test*"]
77
+ python_files = ["test_*.py", "*_test.py"]
78
+ python_functions = ["test_*"]
79
+ testpaths = ["tests"]
80
+ asyncio_mode = "auto"
81
+
82
+ [tool.ruff]
83
+ line-length = 79
84
+ target-version = "py313"
85
+ exclude = [".git", ".venv", "build", "dist", "__pycache__"]
86
+
87
+ [tool.ruff.lint]
88
+ select = [
89
+ "D", # pydocstyle
90
+ "E", # pycodestyle errors
91
+ "W", # pycodestyle warnings
92
+ "F", # pyflakes
93
+ "I", # isort
94
+ "N", # pep8-naming
95
+ "UP", # pyupgrade
96
+ "B", # flake8-bugbear
97
+ "PL", # pylint
98
+ "RUF", # ruff-specific
99
+ "ARG", # ruff-arglint
100
+ "SLF", # ruff-slf
101
+ ]
102
+ extend-select = ["T201"] # Prevent use of print()
103
+
104
+ ignore = [
105
+ "E203", # Whitespace before ':', conflicts with formatter
106
+ "D104", # Missing docstring in __init__ (often redundant)
107
+ "D205", # 1 blank line required between summary line and description
108
+ "D212", # Multi-line docstring summary should start at the first line
109
+ "D200", # One-line docstring should fit on one line
110
+ "PLR2004", # Magic numbers are often more readable than constants for things like HTTP status codes
111
+ ]
112
+
113
+ [tool.ruff.lint.pydocstyle]
114
+ convention = "google"
115
+
116
+ [tool.ruff.lint.per-file-ignores]
117
+ "tests/**/*.py" = ["SLF001", "PLC0415", "ARG002", "D"]
118
+
119
+ [tool.ruff.format]
120
+ quote-style = "double"
121
+ indent-style = "space"
122
+
123
+ [tool.uv-dynamic-versioning]
124
+ fallback-version = "0.0.1"
@@ -0,0 +1,6 @@
1
+ #
2
+ # horus_agent
3
+ # Copyright (c) 2026 Temple Compute
4
+ #
5
+ # MIT License
6
+ #
@@ -0,0 +1,13 @@
1
+ #
2
+ # horus_agent
3
+ # Copyright (c) 2026 Temple Compute
4
+ #
5
+ # MIT License
6
+ #
7
+ """``python -m horus_agent``."""
8
+
9
+ import sys
10
+
11
+ from horus_agent.cli import main
12
+
13
+ sys.exit(main())
@@ -0,0 +1,89 @@
1
+ #
2
+ # horus_agent
3
+ # Copyright (c) 2026 Temple Compute
4
+ #
5
+ # MIT License
6
+ #
7
+ """
8
+ Streamed transfers to and from presigned staging URLs.
9
+
10
+ Used on both sides: the orchestrator and the agent each move large files
11
+ through object storage instead of through the broker.
12
+ """
13
+
14
+ import time
15
+ from collections.abc import AsyncIterator
16
+ from pathlib import Path
17
+
18
+ import httpx
19
+ from horus_runtime.logging import horus_logger
20
+
21
+ from horus_agent.logs import safe_url
22
+
23
+ _CHUNK = 1024 * 1024
24
+ MAX_SINGLE_PUT = 5 * 1024**3
25
+ """S3's limit for one PUT."""
26
+
27
+
28
+ async def _read_chunks(path: Path) -> AsyncIterator[bytes]:
29
+ with path.open("rb") as fh:
30
+ while chunk := fh.read(_CHUNK):
31
+ yield chunk
32
+
33
+
34
+ async def upload(url: str, source: Path | bytes) -> None:
35
+ """
36
+ PUT *source* to the presigned *url*, streaming a file from disk.
37
+
38
+ ``Content-Length`` is always sent: S3 refuses a chunked presigned PUT.
39
+ """
40
+ if isinstance(source, Path):
41
+ size = source.stat().st_size
42
+ if size > MAX_SINGLE_PUT:
43
+ # shortcut: one presigned PUT, S3's limit; add multipart if
44
+ # agents need to move larger single files.
45
+ raise ValueError(
46
+ f"{source} is {size} bytes; files over 5 GiB cannot be "
47
+ "moved to or from an agent yet"
48
+ )
49
+ content: bytes | AsyncIterator[bytes] = _read_chunks(source)
50
+ else:
51
+ size = len(source)
52
+ content = source
53
+ log = horus_logger.log
54
+ log.debug(f"Upload {size} bytes to {safe_url(url)}")
55
+ started = time.monotonic()
56
+ async with httpx.AsyncClient(timeout=None) as client:
57
+ response = await client.put(
58
+ url, content=content, headers={"Content-Length": str(size)}
59
+ )
60
+ response.raise_for_status()
61
+ log.debug(f"Uploaded {size} bytes in {time.monotonic() - started:.2f}s")
62
+
63
+
64
+ async def download_to(url: str, dest: Path) -> None:
65
+ """GET the presigned *url* into *dest*, streamed, creating parents."""
66
+ dest.parent.mkdir(parents=True, exist_ok=True)
67
+ log = horus_logger.log
68
+ log.debug(f"Download {safe_url(url)} to {dest}")
69
+ started = time.monotonic()
70
+ async with (
71
+ httpx.AsyncClient(timeout=None) as client,
72
+ client.stream("GET", url) as response,
73
+ ):
74
+ response.raise_for_status()
75
+ with dest.open("wb") as fh:
76
+ async for chunk in response.aiter_bytes(_CHUNK):
77
+ fh.write(chunk)
78
+ log.debug(
79
+ f"Downloaded {dest.stat().st_size} bytes to {dest} in "
80
+ f"{time.monotonic() - started:.2f}s"
81
+ )
82
+
83
+
84
+ async def download(url: str) -> bytes:
85
+ """GET the presigned *url* into memory."""
86
+ async with httpx.AsyncClient(timeout=None) as client:
87
+ response = await client.get(url)
88
+ response.raise_for_status()
89
+ return response.content