dragon-torch 1.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.
Files changed (103) hide show
  1. dragon_torch-1.1.0/.github/workflows/release.yml +107 -0
  2. dragon_torch-1.1.0/.github/workflows/tests.yml +138 -0
  3. dragon_torch-1.1.0/.gitignore +223 -0
  4. dragon_torch-1.1.0/.python-version +1 -0
  5. dragon_torch-1.1.0/LICENSE +21 -0
  6. dragon_torch-1.1.0/PKG-INFO +118 -0
  7. dragon_torch-1.1.0/README.md +78 -0
  8. dragon_torch-1.1.0/docs/rl_trainers.md +145 -0
  9. dragon_torch-1.1.0/docs/tuning.md +214 -0
  10. dragon_torch-1.1.0/dragon_torch/__init__.py +7 -0
  11. dragon_torch-1.1.0/dragon_torch/common/__init__.py +0 -0
  12. dragon_torch-1.1.0/dragon_torch/common/defaults.py +6 -0
  13. dragon_torch-1.1.0/dragon_torch/common/torch/__init__.py +0 -0
  14. dragon_torch-1.1.0/dragon_torch/common/torch/tree.py +35 -0
  15. dragon_torch-1.1.0/dragon_torch/common/torch/types.py +7 -0
  16. dragon_torch-1.1.0/dragon_torch/common/types.py +22 -0
  17. dragon_torch-1.1.0/dragon_torch/common/utils.py +6 -0
  18. dragon_torch-1.1.0/dragon_torch/data/RL/__init__.py +3 -0
  19. dragon_torch-1.1.0/dragon_torch/data/RL/collector.py +169 -0
  20. dragon_torch-1.1.0/dragon_torch/data/RL/episodes.py +47 -0
  21. dragon_torch-1.1.0/dragon_torch/data/RL/rollout.py +122 -0
  22. dragon_torch-1.1.0/dragon_torch/data/__init__.py +0 -0
  23. dragon_torch-1.1.0/dragon_torch/data/dataloaders.py +1 -0
  24. dragon_torch-1.1.0/dragon_torch/data/datasets/__init__.py +0 -0
  25. dragon_torch-1.1.0/dragon_torch/data/datasets/aggregator.py +120 -0
  26. dragon_torch-1.1.0/dragon_torch/data/datasets/dragon.py +74 -0
  27. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/__init__.py +0 -0
  28. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/noise_zoo.py +0 -0
  29. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/semasia/__init__.py +0 -0
  30. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/semasia/config.py +23 -0
  31. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/semasia/semasia_celeba.py +44 -0
  32. dragon_torch-1.1.0/dragon_torch/data/datasets/hugging_face/semasia/semasia_imagenet_1k.py +42 -0
  33. dragon_torch-1.1.0/dragon_torch/data/envs/__init__.py +5 -0
  34. dragon_torch-1.1.0/dragon_torch/data/envs/config.py +77 -0
  35. dragon_torch-1.1.0/dragon_torch/data/envs/normalize.py +77 -0
  36. dragon_torch-1.1.0/dragon_torch/data/envs/spaces.py +39 -0
  37. dragon_torch-1.1.0/dragon_torch/data/envs/wrappers.py +33 -0
  38. dragon_torch-1.1.0/dragon_torch/data/sources/__init__.py +7 -0
  39. dragon_torch-1.1.0/dragon_torch/data/sources/base.py +55 -0
  40. dragon_torch-1.1.0/dragon_torch/data/sources/disk.py +66 -0
  41. dragon_torch-1.1.0/dragon_torch/data/sources/hf.py +130 -0
  42. dragon_torch-1.1.0/dragon_torch/loss/RL.py +80 -0
  43. dragon_torch-1.1.0/dragon_torch/loss/__init__.py +1 -0
  44. dragon_torch-1.1.0/dragon_torch/loss/base.py +29 -0
  45. dragon_torch-1.1.0/dragon_torch/models/RL/__init__.py +1 -0
  46. dragon_torch-1.1.0/dragon_torch/models/RL/base.py +112 -0
  47. dragon_torch-1.1.0/dragon_torch/models/__init__.py +0 -0
  48. dragon_torch-1.1.0/dragon_torch/models/base.py +395 -0
  49. dragon_torch-1.1.0/dragon_torch/models/external/__init__.py +0 -0
  50. dragon_torch-1.1.0/dragon_torch/models/external/hf.py +26 -0
  51. dragon_torch-1.1.0/dragon_torch/modules/RL/__init__.py +3 -0
  52. dragon_torch-1.1.0/dragon_torch/modules/RL/heads.py +110 -0
  53. dragon_torch-1.1.0/dragon_torch/modules/RL/policy.py +91 -0
  54. dragon_torch-1.1.0/dragon_torch/modules/__init__.py +1 -0
  55. dragon_torch-1.1.0/dragon_torch/modules/base.py +20 -0
  56. dragon_torch-1.1.0/dragon_torch/modules/external/__init__.py +0 -0
  57. dragon_torch-1.1.0/dragon_torch/modules/external/sb3.py +31 -0
  58. dragon_torch-1.1.0/dragon_torch/modules/hugging_face/__init__.py +0 -0
  59. dragon_torch-1.1.0/dragon_torch/modules/hugging_face/image_transformer.py +25 -0
  60. dragon_torch-1.1.0/dragon_torch/trainers/RL/A2C/__init__.py +2 -0
  61. dragon_torch-1.1.0/dragon_torch/trainers/RL/A2C/config.py +22 -0
  62. dragon_torch-1.1.0/dragon_torch/trainers/RL/A2C/trainer.py +18 -0
  63. dragon_torch-1.1.0/dragon_torch/trainers/RL/PPO/__init__.py +2 -0
  64. dragon_torch-1.1.0/dragon_torch/trainers/RL/PPO/config.py +19 -0
  65. dragon_torch-1.1.0/dragon_torch/trainers/RL/PPO/trainer.py +18 -0
  66. dragon_torch-1.1.0/dragon_torch/trainers/RL/__init__.py +6 -0
  67. dragon_torch-1.1.0/dragon_torch/trainers/RL/callbacks.py +116 -0
  68. dragon_torch-1.1.0/dragon_torch/trainers/RL/config.py +186 -0
  69. dragon_torch-1.1.0/dragon_torch/trainers/RL/trainer.py +137 -0
  70. dragon_torch-1.1.0/dragon_torch/trainers/__init__.py +3 -0
  71. dragon_torch-1.1.0/dragon_torch/trainers/base.py +112 -0
  72. dragon_torch-1.1.0/dragon_torch/utils/__init__.py +2 -0
  73. dragon_torch-1.1.0/dragon_torch/utils/random_sources.py +64 -0
  74. dragon_torch-1.1.0/dragon_torch/utils/tuning/RL.py +49 -0
  75. dragon_torch-1.1.0/dragon_torch/utils/tuning/__init__.py +7 -0
  76. dragon_torch-1.1.0/dragon_torch/utils/tuning/config.py +122 -0
  77. dragon_torch-1.1.0/dragon_torch/utils/tuning/monitor.py +130 -0
  78. dragon_torch-1.1.0/dragon_torch/utils/tuning/objective.py +157 -0
  79. dragon_torch-1.1.0/dragon_torch/utils/tuning/presets/__init__.py +2 -0
  80. dragon_torch-1.1.0/dragon_torch/utils/tuning/presets/a2c.py +19 -0
  81. dragon_torch-1.1.0/dragon_torch/utils/tuning/presets/ppo.py +31 -0
  82. dragon_torch-1.1.0/dragon_torch/utils/tuning/reporting.py +71 -0
  83. dragon_torch-1.1.0/dragon_torch/utils/tuning/space.py +90 -0
  84. dragon_torch-1.1.0/dragon_torch/utils/tuning/study.py +51 -0
  85. dragon_torch-1.1.0/main.py +6 -0
  86. dragon_torch-1.1.0/pyproject.toml +102 -0
  87. dragon_torch-1.1.0/scripts/release.py +176 -0
  88. dragon_torch-1.1.0/scripts/tune_rl.py +86 -0
  89. dragon_torch-1.1.0/tests/conftest.py +20 -0
  90. dragon_torch-1.1.0/tests/test_RL/test_rl_collector.py +147 -0
  91. dragon_torch-1.1.0/tests/test_RL/test_rl_envs.py +132 -0
  92. dragon_torch-1.1.0/tests/test_RL/test_rl_training.py +137 -0
  93. dragon_torch-1.1.0/tests/test_datasets/test_aggregator.py +230 -0
  94. dragon_torch-1.1.0/tests/test_datasets/test_semasia_celeba.py +154 -0
  95. dragon_torch-1.1.0/tests/test_models/test_base_model.py +423 -0
  96. dragon_torch-1.1.0/tests/test_release.py +95 -0
  97. dragon_torch-1.1.0/tests/test_sources/test_directory_source.py +171 -0
  98. dragon_torch-1.1.0/tests/test_sources/test_file_source.py +156 -0
  99. dragon_torch-1.1.0/tests/test_sources/test_hf_source.py +165 -0
  100. dragon_torch-1.1.0/tests/test_trainers/test_rl_trainer.py +275 -0
  101. dragon_torch-1.1.0/tests/test_trainers/test_trainer.py +221 -0
  102. dragon_torch-1.1.0/tests/test_tuning/test_tuning.py +221 -0
  103. dragon_torch-1.1.0/tests/test_tuning/test_tuning_rl.py +145 -0
@@ -0,0 +1,107 @@
1
+ name: Release
2
+
3
+ # Builds dragon-torch and publishes it to PyPI with the next version number
4
+ # (scripts/release.py: patch bump, each part capped at 10.10.20).
5
+ #
6
+ # Started by tests.yml after the tests pass on main. It's a separate workflow
7
+ # because tests.yml cancels in-progress runs when a new push arrives, and an
8
+ # upload must never be cancelled halfway. It can also be run by hand for a commit
9
+ # on main; either way it re-checks that the commit's tests passed on main.
10
+ #
11
+ # One-time setup: on PyPI, add a trusted publisher for this repo with workflow
12
+ # `release.yml` and environment `pypi` (a "pending publisher" before the first
13
+ # upload), and create the `pypi` environment under the repo's settings.
14
+
15
+ on:
16
+ workflow_dispatch:
17
+ inputs:
18
+ sha:
19
+ description: Commit on main to release (its Tests run on main must have passed)
20
+ required: true
21
+
22
+ # releases run one at a time and are never cancelled: each one reads the last
23
+ # released version, so two at once could pick the same number
24
+ concurrency:
25
+ group: release
26
+ cancel-in-progress: false
27
+
28
+ jobs:
29
+ release:
30
+ runs-on: ubuntu-latest
31
+ timeout-minutes: 20
32
+ environment: pypi
33
+ permissions:
34
+ contents: write # push the v* release tag
35
+ id-token: write # PyPI trusted publishing
36
+ actions: read # look up the commit's test runs
37
+ env:
38
+ SHA: ${{ inputs.sha }}
39
+ steps:
40
+ - uses: actions/checkout@v4
41
+ with:
42
+ ref: ${{ inputs.sha }}
43
+ fetch-depth: 0 # tags: released versions + "already released" check
44
+
45
+ - name: Check the commit is releasable
46
+ id: gate
47
+ env:
48
+ GH_TOKEN: ${{ github.token }}
49
+ run: |
50
+ if ! git merge-base --is-ancestor "$SHA" origin/main; then
51
+ echo "::error::${SHA::7} is not on main"; exit 1
52
+ fi
53
+ # check the `test` job, not the whole run: the run that dispatched this
54
+ # release is still in progress (its release job is what started us)
55
+ passed=false
56
+ for run in $(gh api "repos/$GITHUB_REPOSITORY/actions/workflows/tests.yml/runs?head_sha=$SHA&branch=main" \
57
+ --jq '.workflow_runs[].id'); do
58
+ conclusion=$(gh api "repos/$GITHUB_REPOSITORY/actions/runs/$run/jobs" \
59
+ --jq '.jobs[] | select(.name == "test") | .conclusion')
60
+ if [ "$conclusion" = "success" ]; then passed=true; break; fi
61
+ done
62
+ if [ "$passed" != "true" ]; then
63
+ echo "::error::no Tests run on main passed its test job for ${SHA::7}"; exit 1
64
+ fi
65
+ released=$(git tag --points-at "$SHA" --list 'v*')
66
+ if [ -n "$released" ]; then
67
+ echo "${SHA::7} is already released as $released, nothing to do"
68
+ echo "skip=true" >> "$GITHUB_OUTPUT"
69
+ else
70
+ echo "skip=false" >> "$GITHUB_OUTPUT"
71
+ fi
72
+
73
+ - uses: astral-sh/setup-uv@v6
74
+ if: steps.gate.outputs.skip == 'false'
75
+ with:
76
+ python-version: "3.13"
77
+
78
+ - name: Pick the next version
79
+ id: version
80
+ if: steps.gate.outputs.skip == 'false'
81
+ run: uv run --no-project python scripts/release.py next-version --write
82
+
83
+ - name: Build
84
+ if: steps.gate.outputs.skip == 'false'
85
+ run: uv build --out-dir dist
86
+
87
+ - name: Check the package
88
+ if: steps.gate.outputs.skip == 'false'
89
+ run: |
90
+ # refuse wheels that would install stray top-level packages (utils, tests, ...)
91
+ uv run --no-project python scripts/release.py check-wheel dist/*.whl
92
+ uvx twine check --strict dist/*
93
+
94
+ - name: Publish to PyPI
95
+ if: steps.gate.outputs.skip == 'false'
96
+ uses: pypa/gh-action-pypi-publish@release/v1
97
+
98
+ - name: Tag the release
99
+ if: steps.gate.outputs.skip == 'false'
100
+ env:
101
+ VERSION: ${{ steps.version.outputs.version }}
102
+ run: |
103
+ git config user.name "github-actions[bot]"
104
+ git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
105
+ git tag -a "v$VERSION" "$SHA" -m "dragon-torch $VERSION"
106
+ git push origin "v$VERSION"
107
+ echo "Released dragon-torch $VERSION (${SHA::7}): https://pypi.org/project/dragon-torch/$VERSION/" >> "$GITHUB_STEP_SUMMARY"
@@ -0,0 +1,138 @@
1
+ name: Tests
2
+
3
+ on:
4
+ push:
5
+ branches: [dev, main]
6
+ workflow_dispatch:
7
+
8
+ # a newer push supersedes an in-flight run on the same branch, so main never gets a stale merge
9
+ concurrency:
10
+ group: tests-${{ github.ref }}
11
+ cancel-in-progress: true
12
+
13
+ jobs:
14
+ test:
15
+ runs-on: ubuntu-latest
16
+ timeout-minutes: 30
17
+ permissions:
18
+ contents: write # publish the test-count badge to the `badges` branch
19
+ steps:
20
+ - uses: actions/checkout@v4
21
+
22
+ - uses: astral-sh/setup-uv@v6
23
+ with:
24
+ python-version: "3.13"
25
+ enable-cache: true
26
+ cache-dependency-glob: pyproject.toml
27
+
28
+ - name: Install dependencies
29
+ run: uv sync --extra test --extra torch-cpu --extra torch-ext --extra RL --extra tune
30
+
31
+ - name: Cache Hugging Face downloads
32
+ uses: actions/cache@v4
33
+ with:
34
+ path: ~/.cache/huggingface
35
+ key: hf-${{ runner.os }}-${{ hashFiles('tests/conftest.py') }}
36
+
37
+ - name: Run tests
38
+ run: uv run pytest -v --junitxml=report.xml
39
+
40
+ # the README badge tracks main only, and is updated whether the tests passed or failed
41
+ - name: Build test-count badge
42
+ if: always() && github.ref == 'refs/heads/main'
43
+ run: |
44
+ python3 - <<'PY'
45
+ import json, os
46
+ import xml.etree.ElementTree as ET
47
+
48
+ if os.path.exists('report.xml'):
49
+ root = ET.parse('report.xml').getroot()
50
+ suites = [root] if root.tag == 'testsuite' else root.findall('testsuite')
51
+ count = lambda key: sum(int(s.get(key, 0)) for s in suites)
52
+ total, skipped = count('tests'), count('skipped')
53
+ failed = count('failures') + count('errors')
54
+ passed = total - skipped - failed
55
+ ran = total - skipped
56
+ message = f'{passed} passed' if failed == 0 else f'{passed}/{ran} passed'
57
+ color = 'brightgreen' if failed == 0 else 'red'
58
+ else: # pytest never ran (e.g. dependency install failed)
59
+ message, color = 'error', 'lightgrey'
60
+
61
+ with open('tests.json', 'w') as f:
62
+ json.dump({'schemaVersion': 1, 'label': 'tests on main',
63
+ 'message': message, 'color': color}, f)
64
+ print(message)
65
+ PY
66
+
67
+ - name: Publish test-count badge
68
+ if: always() && github.ref == 'refs/heads/main' && hashFiles('tests.json') != ''
69
+ run: |
70
+ cp tests.json "$RUNNER_TEMP/tests.json"
71
+ git config user.name "github-actions[bot]"
72
+ git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
73
+ git fetch origin badges || true
74
+ if git rev-parse --verify -q origin/badges; then
75
+ git switch -f -C badges origin/badges
76
+ else
77
+ git switch -f --orphan badges
78
+ git rm -rfq --cached . || true
79
+ fi
80
+ git clean -fdxq
81
+ cp "$RUNNER_TEMP/tests.json" tests.json
82
+ git add tests.json
83
+ git diff --cached --quiet || git commit -m "badge: $(jq -r .message tests.json) @ ${GITHUB_SHA::7}"
84
+ git push origin badges
85
+
86
+ merge-to-main:
87
+ needs: test
88
+ # only commits pushed to dev are merged; pushes to main just run the tests
89
+ if: github.ref == 'refs/heads/dev'
90
+ runs-on: ubuntu-latest
91
+ permissions:
92
+ contents: write # push to main
93
+ actions: write # trigger the main test run
94
+ steps:
95
+ - uses: actions/checkout@v4
96
+ with:
97
+ ref: main
98
+ fetch-depth: 0
99
+
100
+ - name: Merge dev into main
101
+ id: merge
102
+ run: |
103
+ # skip commits main already has (e.g. dev was fast-forwarded to main)
104
+ if git merge-base --is-ancestor "$GITHUB_SHA" HEAD; then
105
+ echo "Commit ${GITHUB_SHA::7} is already on main, nothing to merge"
106
+ echo "merged=false" >> "$GITHUB_OUTPUT"
107
+ exit 0
108
+ fi
109
+ git config user.name "github-actions[bot]"
110
+ git config user.email "41898282+github-actions[bot]@users.noreply.github.com"
111
+ # merge the exact commit that was tested, not whatever dev points at now
112
+ git merge --ff-only "$GITHUB_SHA" || git merge --no-edit "$GITHUB_SHA"
113
+ git push origin main
114
+ echo "merged=true" >> "$GITHUB_OUTPUT"
115
+
116
+ # pushes made with GITHUB_TOKEN don't trigger workflows, so start the
117
+ # main run (which refreshes the badge) explicitly
118
+ - name: Run tests on main
119
+ if: steps.merge.outputs.merged == 'true'
120
+ env:
121
+ GH_TOKEN: ${{ github.token }}
122
+ run: gh workflow run tests.yml --ref main --repo "$GITHUB_REPOSITORY"
123
+
124
+ release:
125
+ needs: test
126
+ # every commit whose tests pass on main is released (release.yml skips commits
127
+ # that already have a v* tag, e.g. a re-run)
128
+ if: github.ref == 'refs/heads/main'
129
+ runs-on: ubuntu-latest
130
+ permissions:
131
+ actions: write # trigger the release workflow
132
+ steps:
133
+ # release.yml runs outside this workflow's cancel-in-progress group, so a
134
+ # newer push to main can't cancel an upload halfway
135
+ - name: Start the release
136
+ env:
137
+ GH_TOKEN: ${{ github.token }}
138
+ run: gh workflow run release.yml --ref main --repo "$GITHUB_REPOSITORY" -f sha="$GITHUB_SHA"
@@ -0,0 +1,223 @@
1
+ # Byte-compiled / optimized / DLL files
2
+ __pycache__/
3
+ *.py[codz]
4
+ *$py.class
5
+ uv.lock
6
+
7
+ # C extensions
8
+ *.so
9
+
10
+ # Distribution / packaging
11
+ .Python
12
+ build/
13
+ develop-eggs/
14
+ dist/
15
+ downloads/
16
+ eggs/
17
+ .eggs/
18
+ lib/
19
+ lib64/
20
+ parts/
21
+ sdist/
22
+ var/
23
+ wheels/
24
+ share/python-wheels/
25
+ *.egg-info/
26
+ .installed.cfg
27
+ *.egg
28
+ MANIFEST
29
+ *.ipynb
30
+
31
+ # PyInstaller
32
+ # Usually these files are written by a python script from a template
33
+ # before PyInstaller builds the exe, so as to inject date/other infos into it.
34
+ *.manifest
35
+ *.spec
36
+
37
+ # Installer logs
38
+ pip-log.txt
39
+ pip-delete-this-directory.txt
40
+
41
+ # Unit test / coverage reports
42
+ htmlcov/
43
+ .tox/
44
+ .nox/
45
+ .coverage
46
+ .coverage.*
47
+ .cache
48
+ nosetests.xml
49
+ coverage.xml
50
+ *.cover
51
+ *.py.cover
52
+ .hypothesis/
53
+ .pytest_cache/
54
+ cover/
55
+
56
+ # Translations
57
+ *.mo
58
+ *.pot
59
+
60
+ # Django stuff:
61
+ *.log
62
+ local_settings.py
63
+ db.sqlite3
64
+ db.sqlite3-journal
65
+
66
+ # Flask stuff:
67
+ instance/
68
+ .webassets-cache
69
+
70
+ # Scrapy stuff:
71
+ .scrapy
72
+
73
+ # Sphinx documentation
74
+ docs/_build/
75
+
76
+ # PyBuilder
77
+ .pybuilder/
78
+ target/
79
+
80
+ # Jupyter Notebook
81
+ .ipynb_checkpoints
82
+
83
+ # IPython
84
+ profile_default/
85
+ ipython_config.py
86
+
87
+ # pyenv
88
+ # For a library or package, you might want to ignore these files since the code is
89
+ # intended to run in multiple environments; otherwise, check them in:
90
+ # .python-version
91
+
92
+ # pipenv
93
+ # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
94
+ # However, in case of collaboration, if having platform-specific dependencies or dependencies
95
+ # having no cross-platform support, pipenv may install dependencies that don't work, or not
96
+ # install all needed dependencies.
97
+ # Pipfile.lock
98
+
99
+ # UV
100
+ # Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
101
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
102
+ # commonly ignored for libraries.
103
+ # uv.lock
104
+
105
+ # poetry
106
+ # Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
107
+ # This is especially recommended for binary packages to ensure reproducibility, and is more
108
+ # commonly ignored for libraries.
109
+ # https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
110
+ # poetry.lock
111
+ # poetry.toml
112
+
113
+ # pdm
114
+ # Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
115
+ # pdm recommends including project-wide configuration in pdm.toml, but excluding .pdm-python.
116
+ # https://pdm-project.org/en/latest/usage/project/#working-with-version-control
117
+ # pdm.lock
118
+ # pdm.toml
119
+ .pdm-python
120
+ .pdm-build/
121
+
122
+ # pixi
123
+ # Similar to Pipfile.lock, it is generally recommended to include pixi.lock in version control.
124
+ # pixi.lock
125
+ # Pixi creates a virtual environment in the .pixi directory, just like venv module creates one
126
+ # in the .venv directory. It is recommended not to include this directory in version control.
127
+ .pixi
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
221
+
222
+ # optuna studies / trial dirs
223
+ runs/
@@ -0,0 +1 @@
1
+ 3.13
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Campbell Rankine
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,118 @@
1
+ Metadata-Version: 2.5
2
+ Name: dragon-torch
3
+ Version: 1.1.0
4
+ Summary: PyTorch Deep Learning Research framework. Focused on RL but support for extending the framework to other areas of research
5
+ License-File: LICENSE
6
+ Requires-Python: >=3.13
7
+ Requires-Dist: jax
8
+ Requires-Dist: optuna>=5.0.0
9
+ Requires-Dist: pydantic>=2.13.5
10
+ Requires-Dist: python-dotenv>=1.2.4
11
+ Requires-Dist: requests>=2.34.2
12
+ Provides-Extra: rl
13
+ Requires-Dist: gymnasium>=1.1; extra == 'rl'
14
+ Requires-Dist: stable-baselines3; extra == 'rl'
15
+ Requires-Dist: stable-retro[extra]; extra == 'rl'
16
+ Provides-Extra: test
17
+ Requires-Dist: pytest>=9.1.1; extra == 'test'
18
+ Provides-Extra: torch-cpu
19
+ Requires-Dist: jax>=0.11.2; extra == 'torch-cpu'
20
+ Requires-Dist: torch>=2.14.0; extra == 'torch-cpu'
21
+ Requires-Dist: torchvision>=0.29.0; extra == 'torch-cpu'
22
+ Requires-Dist: transformers>=5.20.2; extra == 'torch-cpu'
23
+ Provides-Extra: torch-cuda12
24
+ Requires-Dist: jax[cuda12]>=0.11.2; extra == 'torch-cuda12'
25
+ Requires-Dist: torch>=2.11.0; extra == 'torch-cuda12'
26
+ Requires-Dist: torchvision>=0.26.0; extra == 'torch-cuda12'
27
+ Requires-Dist: transformers>=5.20.2; extra == 'torch-cuda12'
28
+ Provides-Extra: torch-cuda13
29
+ Requires-Dist: jax[cuda13]>=0.11.2; extra == 'torch-cuda13'
30
+ Requires-Dist: torch>=2.14.0; extra == 'torch-cuda13'
31
+ Requires-Dist: torchvision>=0.29.0; extra == 'torch-cuda13'
32
+ Requires-Dist: transformers>=5.20.2; extra == 'torch-cuda13'
33
+ Provides-Extra: torch-ext
34
+ Requires-Dist: accelerate>=1.15.0; extra == 'torch-ext'
35
+ Requires-Dist: datasets>=5.1.0; extra == 'torch-ext'
36
+ Requires-Dist: transformers[torch]>=5.20.2; extra == 'torch-ext'
37
+ Provides-Extra: tune
38
+ Requires-Dist: optuna-dashboard; extra == 'tune'
39
+ Description-Content-Type: text/markdown
40
+
41
+ # dragon-torch
42
+
43
+ [![Tests on main](https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2FCampbell-Rankine%2Fdragon-torch%2Fbadges%2Ftests.json)](https://github.com/Campbell-Rankine/dragon-torch/actions/workflows/tests.yml)
44
+
45
+ PyTorch Deep Learning Dev Modules
46
+
47
+ ### Sources:
48
+
49
+ Define source for data in a dataset. Source.__getitem__() is cached and returns the full dataset. Supported types:
50
+
51
+ - disk
52
+ - hugging-face
53
+ - hugging-face-polars
54
+
55
+ ### Datasets:
56
+
57
+ take a source as input. Dataset.__getitem__() returns an iterable of batches from the input source
58
+
59
+ **NOTE**: Dataset aggregators define batching rules / filters for combining two datasets
60
+
61
+ ### Tests:
62
+
63
+ Pushes to `dev` run the full test suite in GitHub Actions; when every test passes, `dev` is merged into `main` automatically. Pushes to `main` run the tests without merging anything. The badge above shows how many tests pass on `main`, and is refreshed after every `main` run. Run them locally with:
64
+
65
+ ```bash
66
+ uv run --extra test --extra torch-cpu --extra torch-ext pytest # everything
67
+ uv run --extra test --extra torch-cpu --extra torch-ext pytest -m "not network" # skip Hugging Face Hub tests
68
+ ```
69
+
70
+ ### Releases:
71
+
72
+ Install with `pip install dragon-torch`. Add extras as needed: `dragon-torch[RL,torch-cpu]` for RL training, `[tune]` for the dashboard. Everything imports from `dragon_torch`, for example `from dragon_torch.trainers.RL import PPOTrainer`.
73
+
74
+ When the tests pass on `main`, `tests.yml` starts `release.yml`. That workflow:
75
+
76
+ 1. Re-checks that the commit is on `main` and that its `test` job passed.
77
+ 2. Picks the next version with `scripts/release.py`.
78
+ 3. Builds the package and uploads it to PyPI.
79
+ 4. Tags the commit `vX.Y.Z`. A commit that already has a tag is never released twice.
80
+
81
+ Each release bumps the patch number of the latest release, carrying over when a part reaches its cap. The caps are 10, 10 and 20, so `1.1.20` is followed by `1.2.0`, `1.10.20` by `2.0.0`, and `10.10.20` is the last possible version. The version in `pyproject.toml` acts as a floor: set it higher, for example to `2.0.0`, to jump straight to that version.
82
+
83
+ Uploads use PyPI trusted publishing. Add a publisher for this repo on PyPI with workflow `release.yml` and environment `pypi`.
84
+
85
+ ### Tuning / monitoring:
86
+
87
+ `dragon_torch/utils/tuning/` runs optuna studies over any `_DragonModel` (pruning, multi-objective, shared storage for parallel workers) and serves them in `optuna-dashboard`. `dragon_torch/utils/tuning/RL.py` and `dragon_torch/utils/tuning/presets/` tune any RL trainer with PPO / A2C search spaces; `scripts/tune_rl.py` runs a study worker. See [docs/tuning.md](docs/tuning.md).
88
+
89
+ ```bash
90
+ uv sync --extra tune
91
+ dragon-tune dashboard --storage sqlite:///runs/optuna.db
92
+ ```
93
+
94
+ ### TODO:
95
+
96
+ - [ ] : Setup Lazy Loading for hugging face sources.
97
+ - [x] : Set up monitoring / experimentation with optuna
98
+ - [ ] : Implement off policy algorithms / dependencies.
99
+ - [ ] : Test Accelerate.prepare works with loaders
100
+ - [ ] : Test Accelerate.prepare works with gyms
101
+ - [ ] : OUActionNoise for RL
102
+ - [ ] : Research better state memory solutions
103
+ - [ ] : Trading environment as a test case in self implemented problem
104
+
105
+
106
+ ### RL
107
+
108
+ | Layer | Files | Contents |
109
+ |---|---|---|
110
+ | Env utils | `dragon_torch/data/envs/` | `DragonEnvConfig` and `make_vec_env`, which always uses `SAME_STEP` autoreset. Also `validate_env`, `register_dragon_env`, space conversions (`obs_to_tensor`, `action_to_env`), running normalisation, and `Discretizer` for stable-retro's button actions |
111
+ | Sources | `dragon_torch/data/RL/` | `EnvCollector` (env stepping, shared by all sources), the `Actor` protocol, `Transition`, `RolloutSource` with GAE, and `EpisodeSource` for evaluation |
112
+ | Networks | `dragon_torch/modules/RL/` | `ActorCritic`, `PolicyOutput`, and distribution heads for Discrete, MultiDiscrete, MultiBinary and Box actions |
113
+ | Losses | `dragon_torch/loss/RL.py` | `PPOLoss` and `A2CLoss`; each records its components (KL, clip fraction, entropy) |
114
+ | Model | `dragon_torch/models/RL/base.py` | `DragonRLModel` adds gradient clipping and returns episode return from `evaluate`. Its epoch count comes from `total_timesteps`, minus the steps already taken |
115
+ | Trainers | `dragon_torch/trainers/RL/` | `DragonRLTrainer` builds a whole run from one config: rollout source, evaluator, network, loss, optimizer, LR annealing, step-scheduled eval and a divergence guard. `PPOTrainer` and `A2CTrainer` add their loss. See [docs/rl_trainers.md](docs/rl_trainers.md) |
116
+ | SB3 | `dragon_torch/modules/external/sb3.py` | `DragonFeaturesExtractor` and `sb3_policy_kwargs` |
117
+
118
+
@@ -0,0 +1,78 @@
1
+ # dragon-torch
2
+
3
+ [![Tests on main](https://img.shields.io/endpoint?url=https%3A%2F%2Fraw.githubusercontent.com%2FCampbell-Rankine%2Fdragon-torch%2Fbadges%2Ftests.json)](https://github.com/Campbell-Rankine/dragon-torch/actions/workflows/tests.yml)
4
+
5
+ PyTorch Deep Learning Dev Modules
6
+
7
+ ### Sources:
8
+
9
+ Define source for data in a dataset. Source.__getitem__() is cached and returns the full dataset. Supported types:
10
+
11
+ - disk
12
+ - hugging-face
13
+ - hugging-face-polars
14
+
15
+ ### Datasets:
16
+
17
+ take a source as input. Dataset.__getitem__() returns an iterable of batches from the input source
18
+
19
+ **NOTE**: Dataset aggregators define batching rules / filters for combining two datasets
20
+
21
+ ### Tests:
22
+
23
+ Pushes to `dev` run the full test suite in GitHub Actions; when every test passes, `dev` is merged into `main` automatically. Pushes to `main` run the tests without merging anything. The badge above shows how many tests pass on `main`, and is refreshed after every `main` run. Run them locally with:
24
+
25
+ ```bash
26
+ uv run --extra test --extra torch-cpu --extra torch-ext pytest # everything
27
+ uv run --extra test --extra torch-cpu --extra torch-ext pytest -m "not network" # skip Hugging Face Hub tests
28
+ ```
29
+
30
+ ### Releases:
31
+
32
+ Install with `pip install dragon-torch`. Add extras as needed: `dragon-torch[RL,torch-cpu]` for RL training, `[tune]` for the dashboard. Everything imports from `dragon_torch`, for example `from dragon_torch.trainers.RL import PPOTrainer`.
33
+
34
+ When the tests pass on `main`, `tests.yml` starts `release.yml`. That workflow:
35
+
36
+ 1. Re-checks that the commit is on `main` and that its `test` job passed.
37
+ 2. Picks the next version with `scripts/release.py`.
38
+ 3. Builds the package and uploads it to PyPI.
39
+ 4. Tags the commit `vX.Y.Z`. A commit that already has a tag is never released twice.
40
+
41
+ Each release bumps the patch number of the latest release, carrying over when a part reaches its cap. The caps are 10, 10 and 20, so `1.1.20` is followed by `1.2.0`, `1.10.20` by `2.0.0`, and `10.10.20` is the last possible version. The version in `pyproject.toml` acts as a floor: set it higher, for example to `2.0.0`, to jump straight to that version.
42
+
43
+ Uploads use PyPI trusted publishing. Add a publisher for this repo on PyPI with workflow `release.yml` and environment `pypi`.
44
+
45
+ ### Tuning / monitoring:
46
+
47
+ `dragon_torch/utils/tuning/` runs optuna studies over any `_DragonModel` (pruning, multi-objective, shared storage for parallel workers) and serves them in `optuna-dashboard`. `dragon_torch/utils/tuning/RL.py` and `dragon_torch/utils/tuning/presets/` tune any RL trainer with PPO / A2C search spaces; `scripts/tune_rl.py` runs a study worker. See [docs/tuning.md](docs/tuning.md).
48
+
49
+ ```bash
50
+ uv sync --extra tune
51
+ dragon-tune dashboard --storage sqlite:///runs/optuna.db
52
+ ```
53
+
54
+ ### TODO:
55
+
56
+ - [ ] : Setup Lazy Loading for hugging face sources.
57
+ - [x] : Set up monitoring / experimentation with optuna
58
+ - [ ] : Implement off policy algorithms / dependencies.
59
+ - [ ] : Test Accelerate.prepare works with loaders
60
+ - [ ] : Test Accelerate.prepare works with gyms
61
+ - [ ] : OUActionNoise for RL
62
+ - [ ] : Research better state memory solutions
63
+ - [ ] : Trading environment as a test case in self implemented problem
64
+
65
+
66
+ ### RL
67
+
68
+ | Layer | Files | Contents |
69
+ |---|---|---|
70
+ | Env utils | `dragon_torch/data/envs/` | `DragonEnvConfig` and `make_vec_env`, which always uses `SAME_STEP` autoreset. Also `validate_env`, `register_dragon_env`, space conversions (`obs_to_tensor`, `action_to_env`), running normalisation, and `Discretizer` for stable-retro's button actions |
71
+ | Sources | `dragon_torch/data/RL/` | `EnvCollector` (env stepping, shared by all sources), the `Actor` protocol, `Transition`, `RolloutSource` with GAE, and `EpisodeSource` for evaluation |
72
+ | Networks | `dragon_torch/modules/RL/` | `ActorCritic`, `PolicyOutput`, and distribution heads for Discrete, MultiDiscrete, MultiBinary and Box actions |
73
+ | Losses | `dragon_torch/loss/RL.py` | `PPOLoss` and `A2CLoss`; each records its components (KL, clip fraction, entropy) |
74
+ | Model | `dragon_torch/models/RL/base.py` | `DragonRLModel` adds gradient clipping and returns episode return from `evaluate`. Its epoch count comes from `total_timesteps`, minus the steps already taken |
75
+ | Trainers | `dragon_torch/trainers/RL/` | `DragonRLTrainer` builds a whole run from one config: rollout source, evaluator, network, loss, optimizer, LR annealing, step-scheduled eval and a divergence guard. `PPOTrainer` and `A2CTrainer` add their loss. See [docs/rl_trainers.md](docs/rl_trainers.md) |
76
+ | SB3 | `dragon_torch/modules/external/sb3.py` | `DragonFeaturesExtractor` and `sb3_policy_kwargs` |
77
+
78
+