From 30dddd9bd903393384ffce662525867e123aed0b Mon Sep 17 00:00:00 2001 From: Mark Vrijlandt Date: Fri, 24 Jul 2026 21:39:35 +0200 Subject: [PATCH 1/5] prefect --- .github/workflows/ci.yml | 68 + .github/workflows/ci_linux.yml | 139 - .github/workflows/ci_win32.yml | 34 - .github/workflows/release.yml | 36 +- .gitignore | 226 +- .vscode/settings.json | 32 +- ci/linux/_load_dot_env.sh | 5 - ci/linux/build_python_package.sh | 8 - ci/linux/create_venv.sh | 8 - ci/linux/install_dependencies.sh | 9 - ci/linux/lint.sh | 8 - ci/linux/test_unit.sh | 8 - ci/linux/typecheck.sh | 8 - ci/win32/create_venv.cmd | 9 - ci/win32/install_dependencies.cmd | 6 - ci/win32/lint.cmd | 8 - ci/win32/test_unit.cmd | 9 - ci/win32/typecheck.cmd | 8 - ci/win32/update_dependencies.cmd | 6 - doc/dev_documentation/Makefile | 20 - doc/dev_documentation/Tools.rst | 60 - doc/dev_documentation/conf.py | 57 - doc/dev_documentation/index.rst | 24 - doc/dev_documentation/make.bat | 35 - doc/user_documentation/Makefile | 20 - doc/user_documentation/conf.py | 52 - doc/user_documentation/index.rst | 20 - doc/user_documentation/make.bat | 35 - justfile | 32 + pyproject.toml | 105 +- requirements.txt | 166 + src/omotes_sdk/__init__.py | 22 - src/omotes_sdk/config.py | 17 - src/omotes_sdk/esdl_messages.py | 25 + src/omotes_sdk/internal/__init__.py | 0 src/omotes_sdk/internal/common/__init__.py | 0 src/omotes_sdk/internal/common/app_logging.py | 96 - .../internal/common/broker_interface.py | 564 --- src/omotes_sdk/internal/common/config.py | 21 - src/omotes_sdk/internal/common/esdl_util.py | 13 - .../orchestrator_worker_events/__init__.py | 0 .../esdl_messages.py | 37 - .../orchestrator_worker_events/task_type.py | 20 - src/omotes_sdk/internal/worker/__init__.py | 1 - src/omotes_sdk/internal/worker/configs.py | 38 - src/omotes_sdk/internal/worker/worker.py | 400 -- src/omotes_sdk/job.py | 14 - src/omotes_sdk/job_status.py | 24 + src/omotes_sdk/log_forwarding.py | 241 ++ src/omotes_sdk/omotes_interface.py | 438 --- src/omotes_sdk/prefect_util.py | 663 ++++ src/omotes_sdk/py.typed | 0 src/omotes_sdk/queue_names.py | 84 - src/omotes_sdk/types.py | 9 - src/omotes_sdk/workflow_type.py | 1254 ------ tests/test_log_forwarding.py | 43 + tests/test_prefect_util.py | 151 + unit_test/internal/__init__.py | 0 unit_test/internal/common/__init__.py | 0 unit_test/internal/common/test_config.py | 9 - .../internal/common/test_queue_message_ttl.py | 38 - .../orchestrator_worker_events/__init__.py | 0 .../test_esdl_messages.py | 51 - unit_test/internal/worker/__init__.py | 0 unit_test/internal/worker/test_configs.py | 9 - ...rkflow_config_enum_option_missing_key.json | 48 - ...kflow_config_enum_options_not_as_list.json | 43 - .../test_config/workflow_config_happy.json | 50 - .../workflow_config_int_min_as_float.json | 49 - ...workflow_config_wrong_datetime_format.json | 50 - unit_test/test_omotes_interface.py | 51 - unit_test/test_workflow_type.py | 377 -- uv.lock | 3395 +++++++++++++++++ 73 files changed, 4904 insertions(+), 4702 deletions(-) create mode 100644 .github/workflows/ci.yml delete mode 100644 .github/workflows/ci_linux.yml delete mode 100644 .github/workflows/ci_win32.yml delete mode 100644 ci/linux/_load_dot_env.sh delete mode 100755 ci/linux/build_python_package.sh delete mode 100755 ci/linux/create_venv.sh delete mode 100755 ci/linux/install_dependencies.sh delete mode 100755 ci/linux/lint.sh delete mode 100755 ci/linux/test_unit.sh delete mode 100755 ci/linux/typecheck.sh delete mode 100644 ci/win32/create_venv.cmd delete mode 100644 ci/win32/install_dependencies.cmd delete mode 100644 ci/win32/lint.cmd delete mode 100644 ci/win32/test_unit.cmd delete mode 100644 ci/win32/typecheck.cmd delete mode 100644 ci/win32/update_dependencies.cmd delete mode 100644 doc/dev_documentation/Makefile delete mode 100644 doc/dev_documentation/Tools.rst delete mode 100644 doc/dev_documentation/conf.py delete mode 100644 doc/dev_documentation/index.rst delete mode 100644 doc/dev_documentation/make.bat delete mode 100644 doc/user_documentation/Makefile delete mode 100644 doc/user_documentation/conf.py delete mode 100644 doc/user_documentation/index.rst delete mode 100644 doc/user_documentation/make.bat create mode 100644 justfile create mode 100644 requirements.txt delete mode 100644 src/omotes_sdk/config.py create mode 100644 src/omotes_sdk/esdl_messages.py delete mode 100644 src/omotes_sdk/internal/__init__.py delete mode 100644 src/omotes_sdk/internal/common/__init__.py delete mode 100644 src/omotes_sdk/internal/common/app_logging.py delete mode 100644 src/omotes_sdk/internal/common/broker_interface.py delete mode 100644 src/omotes_sdk/internal/common/config.py delete mode 100644 src/omotes_sdk/internal/common/esdl_util.py delete mode 100644 src/omotes_sdk/internal/orchestrator_worker_events/__init__.py delete mode 100644 src/omotes_sdk/internal/orchestrator_worker_events/esdl_messages.py delete mode 100644 src/omotes_sdk/internal/orchestrator_worker_events/task_type.py delete mode 100644 src/omotes_sdk/internal/worker/__init__.py delete mode 100644 src/omotes_sdk/internal/worker/configs.py delete mode 100644 src/omotes_sdk/internal/worker/worker.py delete mode 100644 src/omotes_sdk/job.py create mode 100644 src/omotes_sdk/job_status.py create mode 100644 src/omotes_sdk/log_forwarding.py delete mode 100644 src/omotes_sdk/omotes_interface.py create mode 100644 src/omotes_sdk/prefect_util.py delete mode 100644 src/omotes_sdk/py.typed delete mode 100644 src/omotes_sdk/queue_names.py delete mode 100644 src/omotes_sdk/types.py delete mode 100644 src/omotes_sdk/workflow_type.py create mode 100644 tests/test_log_forwarding.py create mode 100644 tests/test_prefect_util.py delete mode 100644 unit_test/internal/__init__.py delete mode 100644 unit_test/internal/common/__init__.py delete mode 100644 unit_test/internal/common/test_config.py delete mode 100644 unit_test/internal/common/test_queue_message_ttl.py delete mode 100644 unit_test/internal/orchestrator_worker_events/__init__.py delete mode 100644 unit_test/internal/orchestrator_worker_events/test_esdl_messages.py delete mode 100644 unit_test/internal/worker/__init__.py delete mode 100644 unit_test/internal/worker/test_configs.py delete mode 100644 unit_test/test_config/workflow_config_enum_option_missing_key.json delete mode 100644 unit_test/test_config/workflow_config_enum_options_not_as_list.json delete mode 100644 unit_test/test_config/workflow_config_happy.json delete mode 100644 unit_test/test_config/workflow_config_int_min_as_float.json delete mode 100644 unit_test/test_config/workflow_config_wrong_datetime_format.json delete mode 100644 unit_test/test_omotes_interface.py delete mode 100644 unit_test/test_workflow_type.py create mode 100644 uv.lock diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..16c4167 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,68 @@ +name: Build-Test-Lint-Typecheck + +on: [push] + +jobs: + lint: + name: Lint + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install uv + uses: astral-sh/setup-uv@v2 + - name: Install just + uses: taiki-e/install-action@just + - name: Install dependencies + run: just install + - name: Run linter + run: just lint + - name: Check formatting + run: just format-check + + test: + name: Test + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install uv + uses: astral-sh/setup-uv@v2 + - name: Install just + uses: taiki-e/install-action@just + - name: Install dependencies + run: just install + - name: Run unit tests + run: just test + - name: Publish test results + if: always() + uses: pmeier/pytest-results-action@main + with: + path: test-results.xml + summary: true + display-options: fEX + fail-on-empty: true + + typecheck: + name: Type Check + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: "3.10" + - name: Install uv + uses: astral-sh/setup-uv@v2 + - name: Install just + uses: taiki-e/install-action@just + - name: Install dependencies + run: just install + - name: Run type checker + run: just typecheck \ No newline at end of file diff --git a/.github/workflows/ci_linux.yml b/.github/workflows/ci_linux.yml deleted file mode 100644 index dfa499b..0000000 --- a/.github/workflows/ci_linux.yml +++ /dev/null @@ -1,139 +0,0 @@ -name: Build-Test-Lint-etc (linux) - -on: [push] - -jobs: - lint: - name: Lint - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: [ "3.10", "3.11", "3.12", "3.13" ] - steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - cache: 'pip' - - name: Cache pip packages - uses: actions/cache@v3 - with: - path: ~/.cache/pip - key: pip-${{ matrix.python-version }}-${{ hashFiles('**/requirements.txt') }} - restore-keys: | - pip-${{ matrix.python-version }}- - - run: | - ./ci/linux/create_venv.sh - ./ci/linux/install_dependencies.sh - - name: lint - run: | - ./ci/linux/lint.sh - - test: - name: Test - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: [ "3.10", "3.11", "3.12", "3.13" ] - steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - cache: 'pip' - - name: Cache pip packages - uses: actions/cache@v3 - with: - path: ~/.cache/pip - key: pip-${{ matrix.python-version }}-${{ hashFiles('**/requirements.txt') }} - restore-keys: | - pip-${{ matrix.python-version }}- - - run: | - ./ci/linux/create_venv.sh - ./ci/linux/install_dependencies.sh - - name: run unit tests - run: | - ./ci/linux/test_unit.sh - - - name: Surface failing tests - if: always() - uses: pmeier/pytest-results-action@main - with: - # A list of JUnit XML files, directories containing the former, and wildcard - # patterns to process. - # See @actions/glob for supported patterns. - path: test-results.xml - - # Add a summary of the results at the top of the report - # Default: true - summary: true - - # Select which results should be included in the report. - # Follows the same syntax as - # `pytest -r` - # Default: fEX - display-options: fEX - - # Fail the workflow if no JUnit XML was found. - # Default: true - fail-on-empty: true - - typecheck: - name: Typecheck - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: [ "3.10", "3.11", "3.12", "3.13" ] - steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - cache: 'pip' - - name: Cache pip packages - uses: actions/cache@v3 - with: - path: ~/.cache/pip - key: pip-${{ matrix.python-version }}-${{ hashFiles('**/requirements.txt') }} - restore-keys: | - pip-${{ matrix.python-version }}- - - run: | - ./ci/linux/create_venv.sh - ./ci/linux/install_dependencies.sh - - name: run typechecker - run: | - ./ci/linux/typecheck.sh - - build: - name: Build - runs-on: ubuntu-latest - strategy: - fail-fast: false - matrix: - python-version: [ "3.10", "3.11", "3.12", "3.13" ] - steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v5 - with: - python-version: ${{ matrix.python-version }} - cache: 'pip' - - name: Cache pip packages - uses: actions/cache@v3 - with: - path: ~/.cache/pip - key: pip-${{ matrix.python-version }}-${{ hashFiles('**/requirements.txt') }} - restore-keys: | - pip-${{ matrix.python-version }}- - - run: | - ./ci/linux/create_venv.sh - ./ci/linux/install_dependencies.sh - - name: build - run: | - ./ci/linux/build_python_package.sh diff --git a/.github/workflows/ci_win32.yml b/.github/workflows/ci_win32.yml deleted file mode 100644 index 15b2bb5..0000000 --- a/.github/workflows/ci_win32.yml +++ /dev/null @@ -1,34 +0,0 @@ -name: Build-Test-Lint (win32) - -on: - pull_request: - types: [ opened, reopened, synchronize ] - -jobs: - build: - runs-on: windows-latest - strategy: - matrix: - python-version: [ "3.10", "3.11", "3.12", "3.13" ] - steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 - with: - python-version: ${{ matrix.python-version }} - cache: 'pip' - - run: | - .\ci\win32\create_venv.cmd - .\ci\win32\install_dependencies.cmd - - - name: Run lint - run: | - .\ci\win32\lint.cmd - - - name: Run unit tests - run: | - .\ci\win32\test_unit.cmd - - - name: Run typecheck - run: | - .\ci\win32\typecheck.cmd diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index e0ed9fd..a8ec65c 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,37 +1,35 @@ -name: PyPi release -run-name: Releasing next version 🚀 +name: Publish Python Package + on: push: tags: - - '*' + - "*" jobs: - pypi-publish: - name: upload release to PyPI + publish-python-package: runs-on: ubuntu-latest - strategy: - matrix: - python-version: [ "3.10" ] - # Specifying a GitHub environment is optional, but strongly encouraged environment: release permissions: - # IMPORTANT: this permission is mandatory for trusted publishing + contents: read id-token: write steps: - - uses: actions/checkout@v3 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v4 + - name: Checkout sources + uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 with: - fetch-depth: 0 - python-version: ${{ matrix.python-version }} - cache: 'pip' - - run: | + python-version: "3.10" + cache: "pip" + + - name: Install build dependencies + run: | ./ci/linux/create_venv.sh ./ci/linux/install_dependencies.sh - - name: build + - name: Build package run: | ./ci/linux/build_python_package.sh - - name: Publish package distributions to PyPI + - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.gitignore b/.gitignore index 765df51..a57d68f 100644 --- a/.gitignore +++ b/.gitignore @@ -1,225 +1,47 @@ -# Byte-compiled / optimized / DLL files +# Byte-compiled / optimized files __pycache__/ *.py[cod] *$py.class -# C extensions +# Native extensions *.so -# Distribution / packaging -.Python + # Build / packaging artifacts build/ -develop-eggs/ dist/ -downloads/ -eggs/ +*.egg-info/ .eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ pip-wheel-metadata/ -share/python-wheels/ -*.egg-info/ -.installed.cfg *.egg -MANIFEST -# PyInstaller -# Usually these files are written by a python script from a template -# before PyInstaller builds the exe, so as to inject date/other infos into it. -*.manifest -*.spec +# Virtual environments +.venv/ +venv/ +env/ +ENV/ -# Installer logs -pip-log.txt -pip-delete-this-directory.txt +# Local environment files +.env -# Unit test / coverage reports -htmlcov/ -.tox/ -.nox/ +# Test / coverage artifacts +.pytest_cache/ .coverage .coverage.* -.cache -nosetests.xml +htmlcov/ coverage.xml -*.cover -*.py,cover -.hypothesis/ -.pytest_cache/ - -# Translations -*.mo -*.pot - -# Django stuff: -*.log -local_settings.py -db.sqlite3 -db.sqlite3-journal - -# Flask stuff: -instance/ -.webassets-cache - -# Scrapy stuff: -.scrapy - -# Sphinx documentation -doc/*/_build/ -doc/*/_static/ -doc/*/_templates/ - -# PyBuilder -target/ - -# Jupyter Notebook -.ipynb_checkpoints - -# IPython -profile_default/ -ipython_config.py - -# pyenv -.python-version - -# pipenv -# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. -# However, in case of collaboration, if having platform-specific dependencies or dependencies -# having no cross-platform support, pipenv may install dependencies that don't work, or not -# install all needed dependencies. -#Pipfile.lock - -# PEP 582; used by e.g. github.com/David-OConnor/pyflow -__pypackages__/ - -# Celery stuff -celerybeat-schedule -celerybeat.pid - -# SageMath parsed files -*.sage.py - -# Environments -.env -.venv* -env/ -venv/ -ENV/ -env.bak/ -venv.bak/ - -# Spyder project settings -.spyderproject -.spyproject - -# Rope project settings -.ropeproject - -# mkdocs documentation -/site +test-results.xml +unit_test_coverage/ -# mypy +# Type checker / linter caches .mypy_cache/ -.dmypy.json -dmypy.json - -# Pyre type checker +.ruff_cache/ .pyre/ -### PyCharm+all ### -# Covers JetBrains IDEs: IntelliJ, RubyMine, PhpStorm, AppCode, PyCharm, CLion, Android Studio, WebStorm and Rider -# Reference: https://intellij-support.jetbrains.com/hc/en-us/articles/206544839 - -# User-specific stuff -.idea/**/workspace.xml -.idea/**/tasks.xml -.idea/**/usage.statistics.xml -.idea/**/dictionaries -.idea/**/shelf - -# AWS User-specific -.idea/**/aws.xml - -# Generated files -.idea/**/contentModel.xml - -# Sensitive or high-churn files -.idea/**/dataSources/ -.idea/**/dataSources.ids -.idea/**/dataSources.local.xml -.idea/**/sqlDataSources.xml -.idea/**/dynamic.xml -.idea/**/uiDesigner.xml -.idea/**/dbnavigator.xml - -# Gradle -.idea/**/gradle.xml -.idea/**/libraries - -# Gradle and Maven with auto-import -# When using Gradle or Maven with auto-import, you should exclude module files, -# since they will be recreated, and may cause churn. Uncomment if using -# auto-import. -# .idea/artifacts -# .idea/compiler.xml -# .idea/jarRepositories.xml -# .idea/modules.xml -# .idea/*.iml -# .idea/modules -# *.iml -# *.ipr - -# CMake -cmake-build-*/ - -# Mongo Explorer plugin -.idea/**/mongoSettings.xml +# Editor / IDE +.idea/ -# File-based project format -*.iws - -# IntelliJ -out/ - -# mpeltonen/sbt-idea plugin -.idea_modules/ - -# JIRA plugin -atlassian-ide-plugin.xml - -# Cursive Clojure plugin -.idea/replstate.xml - -# SonarLint plugin -.idea/sonarlint/ - -# Crashlytics plugin (for Android Studio and IntelliJ) -com_crashlytics_export_strings.xml -crashlytics.properties -crashlytics-build.properties -fabric.properties - -# Editor-based Rest Client -.idea/httpRequests - -# Android studio 3.1+ serialized cache file -.idea/caches/build_file_checksums.ser - -### PyCharm+all Patch ### -# Ignore everything but code style settings and run configurations -# that are supposed to be shared within teams. - -.idea/* - -!.idea/codeStyles -!.idea/runConfigurations - -unit_test_coverage/ -test-results.xml +# Notebook checkpoints +.ipynb_checkpoints/ -.env.* -requirements.txt +# Misc local output +temp/ \ No newline at end of file diff --git a/.vscode/settings.json b/.vscode/settings.json index 6728f09..ebbfc5d 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -1,23 +1,23 @@ { + "flake8.enabled": false, "[python]": { - "editor.formatOnSave": true, - "editor.defaultFormatter": "ms-python.black-formatter", + "editor.defaultFormatter": "charliermarsh.ruff", "editor.codeActionsOnSave": { - "source.organizeImports": true - }, + "source.fixAll.ruff": "always", + "source.organizeImports.ruff": "explicit" + } }, "python.analysis.typeCheckingMode": "basic", "python.testing.unittestArgs": [ "-v", "-s", - "./unit_test", + "./tests", "-p", "test_*.py" ], "python.testing.pytestEnabled": true, "python.testing.unittestEnabled": false, "python.testing.autoTestDiscoverOnSaveEnabled": true, - "mypy.runUsingActiveInterpreter": true, "isort.args": [ "--profile", "black" @@ -25,4 +25,24 @@ "terminal.integrated.env.windows": { "PYTHONPATH": "${workspaceFolder}/src" }, + "python.languageServer": "None", // disable Pylance (ruff/ty is used) + "ty.diagnosticMode": "workspace", + "ty.importStrategy": "fromEnvironment", + "ty.interpreter": "${workspaceFolder}/.venv/bin/python", + "ruff.organizeImports": true, + "ruff.lint.enable": true, + "[just]": { + "editor.formatOnSave": false, + "editor.defaultFormatter": null + }, + "files.associations": { + "justfile": "just", + "*.just": "just" + }, + "chat.tools.terminal.autoApprove": { + "/^python - <<'PY'\nimport inspect\nfrom prefect\\.client\\.orchestration import PrefectClient\nmethods = \\[name for name in dir\\(PrefectClient\\) if 'flow_run' in name\\.lower\\(\\) or 'delete' in name\\.lower\\(\\)\\]\nprint\\('\\\\n'\\.join\\(sorted\\(methods\\)\\)\\)\nPY$/": { + "approve": true, + "matchCommandLine": true + } + } } \ No newline at end of file diff --git a/ci/linux/_load_dot_env.sh b/ci/linux/_load_dot_env.sh deleted file mode 100644 index 434c067..0000000 --- a/ci/linux/_load_dot_env.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/usr/bin/env bash - -set -a -source $1 -set +a diff --git a/ci/linux/build_python_package.sh b/ci/linux/build_python_package.sh deleted file mode 100755 index 0d2d2af..0000000 --- a/ci/linux/build_python_package.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi - -python -m build diff --git a/ci/linux/create_venv.sh b/ci/linux/create_venv.sh deleted file mode 100755 index 71fd104..0000000 --- a/ci/linux/create_venv.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -python3 -m venv ./.venv -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi -pip3 install pip-tools diff --git a/ci/linux/install_dependencies.sh b/ci/linux/install_dependencies.sh deleted file mode 100755 index ebd6979..0000000 --- a/ci/linux/install_dependencies.sh +++ /dev/null @@ -1,9 +0,0 @@ -#!/usr/bin/env bash - -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi - -pip-compile -U --extra=dev --output-file=requirements.txt pyproject.toml -pip install -r ./requirements.txt diff --git a/ci/linux/lint.sh b/ci/linux/lint.sh deleted file mode 100755 index c2e91b5..0000000 --- a/ci/linux/lint.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi - -flake8 ./src/omotes_sdk ./unit_test/ diff --git a/ci/linux/test_unit.sh b/ci/linux/test_unit.sh deleted file mode 100755 index bff0c15..0000000 --- a/ci/linux/test_unit.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi - -PYTHONPATH='$PYTHONPATH:src/' pytest --junit-xml=test-results.xml unit_test/ diff --git a/ci/linux/typecheck.sh b/ci/linux/typecheck.sh deleted file mode 100755 index 7a64ecf..0000000 --- a/ci/linux/typecheck.sh +++ /dev/null @@ -1,8 +0,0 @@ -#!/usr/bin/env bash - -if [[ "$OSTYPE" != "win32" && "$OSTYPE" != "msys" ]]; then - echo "Activating .venv first." - . .venv/bin/activate -fi - -python -m mypy ./src/omotes_sdk ./unit_test/ diff --git a/ci/win32/create_venv.cmd b/ci/win32/create_venv.cmd deleted file mode 100644 index 3cd00af..0000000 --- a/ci/win32/create_venv.cmd +++ /dev/null @@ -1,9 +0,0 @@ -rem Short script to initialize virtual environment using venv and pip -rem @echo off - -pushd . -cd /D "%~dp0" -py -3.10 -m venv ..\..\venv -call ..\..\venv\Scripts\activate.bat -python -m pip install pip-tools -popd diff --git a/ci/win32/install_dependencies.cmd b/ci/win32/install_dependencies.cmd deleted file mode 100644 index e76c29d..0000000 --- a/ci/win32/install_dependencies.cmd +++ /dev/null @@ -1,6 +0,0 @@ - -pushd . -cd /D "%~dp0" -cd ..\..\ -pip-sync .\dev-requirements.txt .\requirements.txt -popd \ No newline at end of file diff --git a/ci/win32/lint.cmd b/ci/win32/lint.cmd deleted file mode 100644 index 098eb27..0000000 --- a/ci/win32/lint.cmd +++ /dev/null @@ -1,8 +0,0 @@ -rem Short script to run linting -rem @echo off - -pushd . -cd /D "%~dp0" -cd ..\..\ -flake8 .\src\omotes_sdk -popd diff --git a/ci/win32/test_unit.cmd b/ci/win32/test_unit.cmd deleted file mode 100644 index c2d51b8..0000000 --- a/ci/win32/test_unit.cmd +++ /dev/null @@ -1,9 +0,0 @@ - -pushd . -cd /D "%~dp0" - -cd ..\..\ -call .\venv\Scripts\activate -set PYTHONPATH=.\src\;%$PYTHONPATH% -pytest unit_test/ -popd \ No newline at end of file diff --git a/ci/win32/typecheck.cmd b/ci/win32/typecheck.cmd deleted file mode 100644 index 7783fff..0000000 --- a/ci/win32/typecheck.cmd +++ /dev/null @@ -1,8 +0,0 @@ -REM script to run mypy type checker on this source tree. -pushd . -cd /D "%~dp0" -cd ..\..\ -call .\venv\Scripts\activate -set PYTHONPATH=.\src\omotes_sdk;%$PYTHONPATH% -python -m mypy ./src/omotes_sdk ./unit_test/ -popd \ No newline at end of file diff --git a/ci/win32/update_dependencies.cmd b/ci/win32/update_dependencies.cmd deleted file mode 100644 index 2725179..0000000 --- a/ci/win32/update_dependencies.cmd +++ /dev/null @@ -1,6 +0,0 @@ - -pushd . -cd /D "%~dp0" -pip-compile --output-file=..\..\requirements.txt ..\..\pyproject.toml -pip-compile --extra=dev --output-file=..\..\dev-requirements.txt ..\..\pyproject.toml -popd \ No newline at end of file diff --git a/doc/dev_documentation/Makefile b/doc/dev_documentation/Makefile deleted file mode 100644 index d4bb2cb..0000000 --- a/doc/dev_documentation/Makefile +++ /dev/null @@ -1,20 +0,0 @@ -# Minimal makefile for Sphinx documentation -# - -# You can set these variables from the command line, and also -# from the environment for the first two. -SPHINXOPTS ?= -SPHINXBUILD ?= sphinx-build -SOURCEDIR = . -BUILDDIR = _build - -# Put it first so that "make" without argument is like "make help". -help: - @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) - -.PHONY: help Makefile - -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). -%: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/doc/dev_documentation/Tools.rst b/doc/dev_documentation/Tools.rst deleted file mode 100644 index 2f4653f..0000000 --- a/doc/dev_documentation/Tools.rst +++ /dev/null @@ -1,60 +0,0 @@ -Tools & IDE's -=================================================================== - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - -There are numerous IDE's for developing python applications, but we -will focus on two: Visual Studio Code and PyCharm - -General -------- -Python code can be writtin using any text editor. However, we've chosen the following tools to support us: - -- **Black** for code formatting -- **Flake8** for code linting -- **MyPy** for typehints/typeChecking -- **pytest** for unit testing -- - -PyCharm -------- - - - - - -VS Code -------- - -The following extensions should be installed: - -* https://marketplace.visualstudio.com/items?itemName=ms-python.python -* https://marketplace.visualstudio.com/items?itemName=ms-python.flake8 -* https://marketplace.visualstudio.com/items?itemName=ms-python.black-formatter -* https://marketplace.visualstudio.com/items?itemName=matangover.mypy - -settings.json for vscode: - -.. code-block:: JSON - - { - "python.analysis.typeCheckingMode": "basic", - "[python]": { - "editor.defaultFormatter": "ms-python.black-formatter" - }, - "python.formatting.provider": "black", - "mypy.runUsingActiveInterpreter": true, - } - - - - -Indices and tables ------------------- - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/doc/dev_documentation/conf.py b/doc/dev_documentation/conf.py deleted file mode 100644 index 2c567ef..0000000 --- a/doc/dev_documentation/conf.py +++ /dev/null @@ -1,57 +0,0 @@ -# Configuration file for the Sphinx documentation builder. -# -# This file only contains a selection of the most common options. For a full -# list see the documentation: -# https://www.sphinx-doc.org/en/master/usage/configuration.html - -# -- Path setup -------------------------------------------------------------- - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. -# -# import os -# import sys -# sys.path.insert(0, os.path.abspath('.')) - - -# -- Project information ----------------------------------------------------- - -project = 'Nieuwe Warmte Nu! python user documentation' -copyright = '2023, NieuweWarmteNu Design Toolkit' -author = 'Sebastiaan la Fleur' - - -# -- General configuration --------------------------------------------------- - -# Add any Sphinx extension module names here, as strings. They can be -# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom -# ones. -extensions = [ - "sphinx.ext.autodoc", - "sphinx.ext.intersphinx", - "sphinx.ext.autosummary", - "sphinx.ext.napoleon", - "sphinx_copybutton", -] - -# Add any paths that contain templates here, relative to this directory. -templates_path = ["_templates"] - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - - -# -- Options for HTML output ------------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "furo" - -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ["_static"] diff --git a/doc/dev_documentation/index.rst b/doc/dev_documentation/index.rst deleted file mode 100644 index 77be81e..0000000 --- a/doc/dev_documentation/index.rst +++ /dev/null @@ -1,24 +0,0 @@ -.. Deltares python Developer documentation documentation master file, created by - sphinx-quickstart on Fri Apr 21 15:13:59 2023. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. - -Welcome to Design Toolkit python Developer documentation's documentation! -=================================================================== - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - Tools - -This is the developer documentation for this project. This documentation provides -information on tools, configurations and howto's for developers. - - -Indices and tables -================== - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/doc/dev_documentation/make.bat b/doc/dev_documentation/make.bat deleted file mode 100644 index 954237b..0000000 --- a/doc/dev_documentation/make.bat +++ /dev/null @@ -1,35 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -) -set SOURCEDIR=. -set BUILDDIR=_build - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.https://www.sphinx-doc.org/ - exit /b 1 -) - -if "%1" == "" goto help - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/doc/user_documentation/Makefile b/doc/user_documentation/Makefile deleted file mode 100644 index d4bb2cb..0000000 --- a/doc/user_documentation/Makefile +++ /dev/null @@ -1,20 +0,0 @@ -# Minimal makefile for Sphinx documentation -# - -# You can set these variables from the command line, and also -# from the environment for the first two. -SPHINXOPTS ?= -SPHINXBUILD ?= sphinx-build -SOURCEDIR = . -BUILDDIR = _build - -# Put it first so that "make" without argument is like "make help". -help: - @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) - -.PHONY: help Makefile - -# Catch-all target: route all unknown targets to Sphinx using the new -# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). -%: Makefile - @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/doc/user_documentation/conf.py b/doc/user_documentation/conf.py deleted file mode 100644 index d31a147..0000000 --- a/doc/user_documentation/conf.py +++ /dev/null @@ -1,52 +0,0 @@ -# Configuration file for the Sphinx documentation builder. -# -# This file only contains a selection of the most common options. For a full -# list see the documentation: -# https://www.sphinx-doc.org/en/master/usage/configuration.html - -# -- Path setup -------------------------------------------------------------- - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. -# -# import os -# import sys -# sys.path.insert(0, os.path.abspath('.')) - - -# -- Project information ----------------------------------------------------- - -project = 'Nieuwe Warmte Nu! python user documentation' -copyright = '2023, NieuweWarmteNu Design Toolkit' -author = 'Sebastiaan la Fleur' - - -# -- General configuration --------------------------------------------------- - -# Add any Sphinx extension module names here, as strings. They can be -# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom -# ones. -extensions = [ -] - -# Add any paths that contain templates here, relative to this directory. -templates_path = ['_templates'] - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] - - -# -- Options for HTML output ------------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = 'alabaster' - -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ['_static'] \ No newline at end of file diff --git a/doc/user_documentation/index.rst b/doc/user_documentation/index.rst deleted file mode 100644 index 097493f..0000000 --- a/doc/user_documentation/index.rst +++ /dev/null @@ -1,20 +0,0 @@ -.. Deltares python user documentation documentation master file, created by - sphinx-quickstart on Fri Apr 21 15:13:29 2023. - You can adapt this file completely to your liking, but it should at least - contain the root `toctree` directive. - -Welcome to Design Toolkit python user documentation's documentation! -============================================================== - -.. toctree:: - :maxdepth: 2 - :caption: Contents: - - - -Indices and tables -================== - -* :ref:`genindex` -* :ref:`modindex` -* :ref:`search` diff --git a/doc/user_documentation/make.bat b/doc/user_documentation/make.bat deleted file mode 100644 index 954237b..0000000 --- a/doc/user_documentation/make.bat +++ /dev/null @@ -1,35 +0,0 @@ -@ECHO OFF - -pushd %~dp0 - -REM Command file for Sphinx documentation - -if "%SPHINXBUILD%" == "" ( - set SPHINXBUILD=sphinx-build -) -set SOURCEDIR=. -set BUILDDIR=_build - -%SPHINXBUILD% >NUL 2>NUL -if errorlevel 9009 ( - echo. - echo.The 'sphinx-build' command was not found. Make sure you have Sphinx - echo.installed, then set the SPHINXBUILD environment variable to point - echo.to the full path of the 'sphinx-build' executable. Alternatively you - echo.may add the Sphinx directory to PATH. - echo. - echo.If you don't have Sphinx installed, grab it from - echo.https://www.sphinx-doc.org/ - exit /b 1 -) - -if "%1" == "" goto help - -%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% -goto end - -:help -%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% - -:end -popd diff --git a/justfile b/justfile new file mode 100644 index 0000000..93ed975 --- /dev/null +++ b/justfile @@ -0,0 +1,32 @@ +# CI tasks + +# Install dependencies with dev group +install: + uv sync --locked --group dev + +# Run linter checks +lint: + uv run ruff check ./src ./tests + +# Fix linting issues +format: + uv run ruff format ./src ./tests + +# Check formatting without modifying files +format-check: + uv run ruff format --check ./src ./tests + +# Run type checker +typecheck: + uv run ty check ./src ./tests + +# Run tests +test: + uv run pytest --junit-xml=test-results.xml tests/ + +# Run all checks (install, lint, format-check, typecheck, test) +ci: install lint format-check typecheck test + +# Show this help message +help: + @just --list diff --git a/pyproject.toml b/pyproject.toml index 28622f5..514268c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "omotes-sdk-python" -requires-python = ">=3.10" +requires-python = ">=3.11" dynamic = ["version"] authors = [ { name = "Sebastiaan la Fleur", email = "sebastiaan.lafleur@tno.nl" }, @@ -24,33 +24,23 @@ classifiers = [ ] dependencies = [ - "aio-pika ~= 9.4, < 9.5", - "omotes-sdk-protocol ~= 1.2", - "pyesdl[profiles] ~= 26.7", - "pamqp ~= 3.3", - "celery ~= 5.3", - "typing-extensions ~= 4.11", + "prefect==3.7.3", + "s3fs", "streamcapture ~= 1.2.5", + "typing-extensions ~= 4.11", ] -[project.optional-dependencies] +[dependency-groups] dev = [ "setuptools ~= 75.6.0", "wheel ~= 0.45.1", "setuptools-git-versioning >= 2.0, < 3", - "black ~= 24.10.0", - "flake8 == 7.1.1", - "flake8-pyproject ~= 1.2.3", - "flake8-docstrings ~= 1.7.0", - "flake8-quotes ~= 3.4.0", - "flake8-bugbear ~= 24.10.31", - "flake8-tuple ~= 0.4.1", "pytest ~= 8.3.4", "pytest-cov ~= 6.0.0", - "mypy ~= 1.13.0", + "ruff ~= 0.15.22", + "ty ~= 0.0.61", "isort == 5.13.2", "build ~= 1.2.2", - "mypy-protobuf ~= 3.5.0", ] [project.urls] @@ -75,7 +65,7 @@ enabled = true starting_version = "0.0.1" [tool.pytest.ini_options] -addopts = "--cov=omotes_sdk --cov-report html --cov-report term-missing --cov-fail-under 60" +addopts = "--cov=omotes_sdk --cov-report html --cov-report term-missing --cov-fail-under 10" [tool.coverage.run] source = ["src"] @@ -84,65 +74,28 @@ omit = [ "*_pb2.py", ] -[tool.flake8] -exclude = [ - '.venv/*', - 'venv/*', - 'doc/*', - 'src/omotes_sdk/internal/orchestrator_worker_events/messages/*' +[tool.ruff] +line-length = 120 +preview = true + +[tool.ruff.lint] +select = [ + "E", # pycodestyle + "W", + "F", # Pyflakes + "UP", # pyupgrade + "B", # flake8-bugbear + "SIM", # flake8-simplify + "I", # isort + "D", # pydocstyle + "ANN", # flake8-annotations + "DOC", # flake8-docstrings ] ignore = [ - 'Q000', # Remove bad quotes - 'D401', # Docstring First line should be imperative - 'E203', # Space before colon (not PEP-8 compliant, and conflicts with black) - 'C408', # Suggestion to use dict() over {} - 'W503', # Starting lines with operators. - 'D104', # Missing docstring in public package - 'D100' # Missing docstring in public module -] -per-file-ignores = [ - '__init__.py:F401', - './unit_test/*:D100,D101,D102,D103' + "D100", # Missing docstring in public module + "D104", # Missing docstring in public package ] -max-line-length = 100 -count = true - -[tool.black] -line-length = 100 - -[tool.isort] -profile = "black" - -[tool.mypy] -python_version = "3.10" -warn_return_any = true -warn_unused_configs = true -disallow_untyped_defs = true -disallow_incomplete_defs = true -exclude = ['.venv/*', 'venv/*', 'doc/*', 'ci/*'] - -# mypy per-module options: -[[tool.mypy.overrides]] -module = "unit_test.*" -check_untyped_defs = true -ignore_missing_imports = true - -[[tool.mypy.overrides]] -module = "celery.*" -ignore_missing_imports = true - -[[tool.mypy.overrides]] -module = "kombu.*" -ignore_missing_imports = true - -[[tool.mypy.overrides]] -module = "billiard.*" -ignore_missing_imports = true - -[[tool.mypy.overrides]] -module = "streamcapture.*" -ignore_missing_imports = true +exclude = ["justfile", ".editorconfig"] -[[tool.mypy.overrides]] -module = "esdl.*" -ignore_missing_imports = true \ No newline at end of file +[tool.ruff.lint.pydocstyle] +convention = "google" diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e22bf32 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,166 @@ +# +# This file is autogenerated by pip-compile with Python 3.12 +# by the following command: +# +# pip-compile --cert=None --client-cert=None --extra=dev --index-url=None --output-file=requirements.txt --pip-args=None pyproject.toml +# +aio-pika==9.4.3 + # via omotes-sdk-python (pyproject.toml) +aiormq==6.8.1 + # via aio-pika +amqp==5.3.1 + # via kombu +attrs==25.4.0 + # via flake8-bugbear +billiard==4.2.4 + # via celery +black==24.10.0 + # via omotes-sdk-python (pyproject.toml) +build==1.2.2.post1 + # via omotes-sdk-python (pyproject.toml) +celery==5.6.2 + # via omotes-sdk-python (pyproject.toml) +click==8.3.1 + # via + # black + # celery + # click-didyoumean + # click-plugins + # click-repl +click-didyoumean==0.3.1 + # via celery +click-plugins==1.1.1.2 + # via celery +click-repl==0.3.0 + # via celery +coverage[toml]==7.13.4 + # via pytest-cov +flake8==7.1.1 + # via + # flake8-bugbear + # flake8-docstrings + # flake8-pyproject + # flake8-quotes + # flake8-tuple + # omotes-sdk-python (pyproject.toml) +flake8-bugbear==24.10.31 + # via omotes-sdk-python (pyproject.toml) +flake8-docstrings==1.7.0 + # via omotes-sdk-python (pyproject.toml) +flake8-pyproject==1.2.4 + # via omotes-sdk-python (pyproject.toml) +flake8-quotes==3.4.0 + # via omotes-sdk-python (pyproject.toml) +flake8-tuple==0.4.1 + # via omotes-sdk-python (pyproject.toml) +future-fstrings==1.2.0 + # via pyecore +idna==3.11 + # via yarl +iniconfig==2.3.0 + # via pytest +isort==5.13.2 + # via omotes-sdk-python (pyproject.toml) +kombu==5.6.2 + # via celery +lxml==6.0.2 + # via pyecore +mccabe==0.7.0 + # via flake8 +multidict==6.7.1 + # via yarl +mypy==1.13.0 + # via omotes-sdk-python (pyproject.toml) +mypy-extensions==1.1.0 + # via + # black + # mypy +mypy-protobuf==3.5.0 + # via omotes-sdk-python (pyproject.toml) +omotes-sdk-protocol==1.2.0 + # via omotes-sdk-python (pyproject.toml) +ordered-set==4.1.0 + # via pyecore +packaging==26.0 + # via + # black + # build + # kombu + # pytest + # setuptools-git-versioning +pamqp==3.3.0 + # via + # aiormq + # omotes-sdk-python (pyproject.toml) +pathspec==1.0.4 + # via black +platformdirs==4.9.4 + # via black +pluggy==1.6.0 + # via pytest +prompt-toolkit==3.0.52 + # via click-repl +propcache==0.4.1 + # via yarl +protobuf==5.29.6 + # via + # mypy-protobuf + # omotes-sdk-protocol +pycodestyle==2.12.1 + # via flake8 +pydocstyle==6.3.0 + # via flake8-docstrings +pyecore==0.13.2 + # via pyesdl +pyesdl==26.2 + # via omotes-sdk-python (pyproject.toml) +pyflakes==3.2.0 + # via flake8 +pyproject-hooks==1.2.0 + # via build +pytest==8.3.5 + # via + # omotes-sdk-python (pyproject.toml) + # pytest-cov +pytest-cov==6.0.0 + # via omotes-sdk-python (pyproject.toml) +python-dateutil==2.9.0.post0 + # via celery +restrictedpython==8.1 + # via pyecore +setuptools-git-versioning==2.1.0 + # via omotes-sdk-python (pyproject.toml) +six==1.17.0 + # via + # flake8-tuple + # python-dateutil +snowballstemmer==3.0.1 + # via pydocstyle +streamcapture==1.2.7 + # via omotes-sdk-python (pyproject.toml) +types-protobuf==6.32.1.20260221 + # via mypy-protobuf +typing-extensions==4.15.0 + # via + # mypy + # omotes-sdk-python (pyproject.toml) +tzdata==2025.3 + # via kombu +tzlocal==5.3.1 + # via celery +vine==5.1.0 + # via + # amqp + # celery + # kombu +wcwidth==0.6.0 + # via prompt-toolkit +wheel==0.45.1 + # via omotes-sdk-python (pyproject.toml) +yarl==1.23.0 + # via + # aio-pika + # aiormq + +# The following packages are considered to be unsafe in a requirements file: +# setuptools diff --git a/src/omotes_sdk/__init__.py b/src/omotes_sdk/__init__.py index a7cdab2..e69de29 100644 --- a/src/omotes_sdk/__init__.py +++ b/src/omotes_sdk/__init__.py @@ -1,22 +0,0 @@ -# Copyright (c) 2023. Deltares & TNO -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . -import os - -from omotes_sdk.internal.common.app_logging import setup_logging, LogLevel - -setup_logging(LogLevel.parse(os.environ.get("LOG_LEVEL", "INFO")), "omotes_sdk") -setup_logging(LogLevel.parse(os.environ.get("LOG_LEVEL", "INFO")), "omotes_sdk_internal") -setup_logging(LogLevel.parse(os.environ.get("LOG_LEVEL", "INFO")), "aio_pika") -setup_logging(LogLevel.parse(os.environ.get("LOG_LEVEL", "INFO")), "aiormq") diff --git a/src/omotes_sdk/config.py b/src/omotes_sdk/config.py deleted file mode 100644 index 62f3c98..0000000 --- a/src/omotes_sdk/config.py +++ /dev/null @@ -1,17 +0,0 @@ -from dataclasses import dataclass - - -@dataclass -class RabbitMQConfig: - """Configuration for connect to RabbitMQ.""" - - host: str = "localhost" - """RabbitMQ hostname""" - port: int = 5672 - """RabbitMQ port""" - username: str = "omotes" - """RabbitMQ username""" - password: str = "guest" - """RabbitMQ account password""" - virtual_host: str = "omotes" - """RabbitMQ virtual host""" diff --git a/src/omotes_sdk/esdl_messages.py b/src/omotes_sdk/esdl_messages.py new file mode 100644 index 0000000..bfea67a --- /dev/null +++ b/src/omotes_sdk/esdl_messages.py @@ -0,0 +1,25 @@ +"""ESDL feedback message models and severity definitions.""" + +from enum import Enum + +from pydantic import BaseModel + + +class MessageSeverity(Enum): + """Message severity options.""" + + DEBUG = "DEBUG" + INFO = "INFO" + WARNING = "WARNING" + ERROR = "ERROR" + + +class EsdlMessage(BaseModel): + """Esdl feedback message, optionally related to a specific object (asset).""" + + technical_message: str + """Technical message.""" + severity: MessageSeverity + """Message severity.""" + esdl_object_id: str | None = None + """Optional esdl object id, None implies a general energy system message.""" diff --git a/src/omotes_sdk/internal/__init__.py b/src/omotes_sdk/internal/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/omotes_sdk/internal/common/__init__.py b/src/omotes_sdk/internal/common/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/omotes_sdk/internal/common/app_logging.py b/src/omotes_sdk/internal/common/app_logging.py deleted file mode 100644 index c174b16..0000000 --- a/src/omotes_sdk/internal/common/app_logging.py +++ /dev/null @@ -1,96 +0,0 @@ -# Copyright (c) 2023. Deltares & TNO -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU General Public License as published by -# the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU General Public License for more details. -# -# You should have received a copy of the GNU General Public License -# along with this program. If not, see . - -"""Setup logging for this python application.""" - -import logging -import sys -from enum import Enum -from typing import Optional, Dict - -CONFIGURED_LOGGERS: Dict[str, logging.Logger] = {} - - -class LogLevel(Enum): - """Simple enum to cover log levels for logging library.""" - - DEBUG = logging.DEBUG - INFO = logging.INFO - WARNING = logging.WARNING - ERROR = logging.ERROR - - @staticmethod - def parse(value: str) -> "LogLevel": - """ - Parses a given string for LogLevel's. - - Parameters - ---------- - value : str - user provided string containing the requested log level - - Returns - ------- - LogLevel - Loglevel for this logger - """ - lowered = value.lower() - - if lowered == "debug": - result = LogLevel.DEBUG - elif lowered == "info": - result = LogLevel.INFO - elif lowered in ["warning", "warn"]: - result = LogLevel.WARNING - elif lowered in ["err", "error"]: - result = LogLevel.ERROR - else: - raise ValueError(f"Value {value} is not a valid log level.") - - return result - - -def setup_logging(log_level: LogLevel, logger_name: Optional[str]) -> logging.Logger: - """ - Initializes logging. - - Parameters - ---------- - log_level : LogLevel - The LogLevel for this logger. - logger_name : Optional[str] - Name for this logger. - """ - if "" not in CONFIGURED_LOGGERS: - root_logger = logging.getLogger() - root_logger.setLevel(LogLevel.DEBUG.value) - - log_handler = logging.StreamHandler(sys.stdout) - formatter = logging.Formatter( - fmt="%(asctime)s [%(name)s][%(threadName)s][%(filename)s:%(lineno)d]" - "[%(levelname)s]: %(message)s" - ) - log_handler.setFormatter(formatter) - root_logger.addHandler(log_handler) - CONFIGURED_LOGGERS[""] = root_logger - - logger = logging.getLogger(logger_name) - logger_name = logger_name if logger_name else "" - - if logger_name not in CONFIGURED_LOGGERS: - print(f"Will use log level {log_level} for logger '{logger_name}'") - logger.setLevel(log_level.value) - - return logger diff --git a/src/omotes_sdk/internal/common/broker_interface.py b/src/omotes_sdk/internal/common/broker_interface.py deleted file mode 100644 index 191dc3b..0000000 --- a/src/omotes_sdk/internal/common/broker_interface.py +++ /dev/null @@ -1,564 +0,0 @@ -import asyncio -import logging -from asyncio import Task -from concurrent.futures import Future -from dataclasses import dataclass -from enum import Enum -from functools import partial -import threading -from types import TracebackType -from typing import Callable, Optional, Dict, Type, TypedDict, cast -from datetime import timedelta - -from aio_pika import connect_robust, Message, DeliveryMode -from aio_pika.abc import ( - AbstractRobustConnection, - AbstractChannel, - AbstractQueue, - AbstractIncomingMessage, - AbstractExchange, -) -from aio_pika.exceptions import ChannelClosed -from pamqp.common import Arguments - -from omotes_sdk.config import RabbitMQConfig - -logger = logging.getLogger("omotes_sdk") - - -@dataclass -class QueueSubscriptionConsumer: - """Consumes a queue until stopped and forwards received messages using a callback.""" - - queue: AbstractQueue - """The queue to which is subscribed.""" - delete_after_messages: Optional[int] - """Delete the subscription & queue after successfully processing this amount of messages""" - callback_on_message: Callable[[bytes], None] - """Callback which is called on each message.""" - - async def run(self) -> None: - """Retrieve messages from the AMQP queue and run callback on each message. - - Requeue in case the handler generates an exception. Callback is a synchronous - function which is executed from the executor threadpool. - """ - number_of_messages_processed = 0 - delete_queue = False - async with self.queue.iterator() as queue_iter: - message: AbstractIncomingMessage - async for message in queue_iter: - async with message.process(requeue=True): - logger.debug( - "Received with queue subscription on queue %s: %s", - self.queue.name, - message.body, - ) - await asyncio.get_running_loop().run_in_executor( - None, partial(self.callback_on_message, message.body) - ) - number_of_messages_processed += 1 - - if ( - self.delete_after_messages is not None - and number_of_messages_processed >= self.delete_after_messages - ): - logger.debug( - "Successful message threshold %s reached for queue %s. " - "Stopping subscription and deleting queue.", - self.delete_after_messages, - self.queue.name, - ) - delete_queue = True - break - logger.debug( - "Stopping subscription on queue %s. Processed %s messages.", - self.queue.name, - number_of_messages_processed, - ) - - if delete_queue: - logger.debug("Deleting queue %s", self.queue.name) - await self.queue.delete() - - -class AioPikaQueueTypeArguments(TypedDict, total=False): - """The keyword arguments which may be used in aio-pika `AbstractChannel.declare_queue`.""" - - exclusive: bool - auto_delete: bool - durable: bool - - -class AMQPQueueType(Enum): - """The RabbitMQ AMQP queue types.""" - - EXCLUSIVE = "exclusive" - AUTO_DELETE = "auto_delete" - DURABLE = "durable" - - def to_argument(self) -> AioPikaQueueTypeArguments: - """Convert the RabbitMQ AMQP queue type to the aio-pika `declare_queue` keyword arguments. - - :return: The keyword arguments which may be used in `AbstractChannel.declare_queue` of the - aio-pika library. - """ - result = AioPikaQueueTypeArguments() - if self == AMQPQueueType.EXCLUSIVE: - result["exclusive"] = True - elif self == AMQPQueueType.AUTO_DELETE: - result["auto_delete"] = True - elif self == AMQPQueueType.DURABLE: - result["durable"] = True - else: - raise RuntimeError(f"Unknown AMQP queue type {self}") - - return result - - -@dataclass() -class QueueTTLArguments(): - """Construct additional time-to-live arguments when declaring a queue.""" - - queue_ttl: Optional[timedelta] = None - """Expires and deletes the queue after a period of time when it is unused. - The timedelta must be convertible into a positive integer. - Ref: https://www.rabbitmq.com/docs/ttl#queue-ttl""" - - def to_argument(self) -> Arguments: - """Convert the time-to-live variables to the aio-pika `declare_queue` keyword arguments. - - :return: The time-to-live keyword arguments in AMQP method arguments data type. - """ - arguments: Arguments = {} - # Ensure this is not None to avoid typecheck error. - arguments = cast(dict, arguments) - - if self.queue_ttl is not None: - if self.queue_ttl <= timedelta(0): - raise ValueError("queue_ttl must be a positive value, " - + f"{self.queue_ttl} received.") - arguments["x-expires"] = int(self.queue_ttl.total_seconds() * 1000) - - return arguments - - -class BrokerInterface(threading.Thread): - """Interface to RabbitMQ using aiopika.""" - - TIMEOUT_ON_STOP_SECONDS: int = 5 - """How long to wait till the broker connection has stopped before killing it non-gracefully.""" - - config: RabbitMQConfig - """The broker configuration.""" - - _loop: asyncio.AbstractEventLoop - """The eventloop in which all broker-related async traffic is run.""" - _connection: AbstractRobustConnection - """AMQP connection.""" - _channel: AbstractChannel - """AMQP channel.""" - _exchanges: Dict[str, AbstractExchange] - """Exchanges declared in this connection.""" - _queue_subscription_consumer_by_name: Dict[str, QueueSubscriptionConsumer] - """Task to consume messages when they are received ordered by queue name.""" - _queue_subscription_tasks: Dict[str, Task] - """Reference to the queue subscription task by queue name.""" - _ready_for_processing: threading.Event - """Thread-safe check which is set once the AMQP connection is up and running.""" - _stopping_lock: threading.Lock - """Lock to make sure only a single thread is stopping this interface.""" - _stopping: bool - """Value to check if a thread is stopping this interface.""" - _stopped: bool - """Value to check if this interface has successfully stopped.""" - - def __init__(self, config: RabbitMQConfig): - """Create the BrokerInterface. - - :param config: Configuration to connect to RabbitMQ. - """ - super().__init__() - self.config = config - - self._exchanges = {} - self._queue_subscription_consumer_by_name = {} - self._queue_subscription_tasks = {} - self._ready_for_processing = threading.Event() - self._stopping_lock = threading.Lock() - self._stopping = False - self._stopped = False - - def __enter__(self) -> "BrokerInterface": - """Start the interface when it is called as a context manager.""" - self.start() - return self - - def __exit__( - self, - exc_type: Optional[Type[BaseException]], - exc_val: Optional[BaseException], - exc_tb: Optional[TracebackType], - ) -> None: - """Stop the interface when it is called as a context manager.""" - self.stop() - - async def _declare_exchange(self, exchange_name: str) -> None: - """Declare an exchange on which messages may be published and routed to queues. - - :param exchange_name: Name of the exchange. - """ - new_exchange = await self._channel.declare_exchange(exchange_name) - self._exchanges[exchange_name] = new_exchange - - async def _send_message_to( - self, exchange_name: Optional[str], routing_key: str, message: bytes - ) -> None: - """Publish a message to a specific routing_key. - - :param exchange_name: The name of the exchange on which to publish the message. Must be - declared before publishing is possible. If None, the default exchange is used which is - already declared at startup - :param routing_key: The routing key to publish the message to. This may be a specific - routing key to which multiple queues are bound or the queue name (default routing key). - :param message: The message to publish. - """ - amqp_message = Message(body=message, delivery_mode=DeliveryMode.PERSISTENT) - - if exchange_name is None: - logger.debug( - "Sending a message to default exchange using routing key %s containing: %s", - routing_key, - message, - ) - await self._channel.default_exchange.publish(amqp_message, routing_key=routing_key) - elif exchange_name not in self._exchanges: - raise RuntimeError( - f"Unable to send message with routing key {routing_key} to exchange " - f"{exchange_name} as it hasn't been declared yet" - ) - else: - logger.debug( - "Sending a message to exchange %s using routing key %s containing: %s", - exchange_name, - routing_key, - message, - ) - await self._exchanges[exchange_name].publish(amqp_message, routing_key=routing_key) - - async def _declare_queue( - self, - queue_name: str, - queue_type: AMQPQueueType, - bind_to_routing_key: Optional[str] = None, - exchange_name: Optional[str] = None, - queue_ttl: Optional[QueueTTLArguments] = None - ) -> AbstractQueue: - """Declare an AMQP queue. - - :param queue_name: Name of the queue to declare. - :param queue_type: Declare the queue using one of the known queue types. - :param bind_to_routing_key: Bind the queue to this routing key next to the default routing - key of the queue name. If none, the queue is only bound to the name of the queue. - If not none, then the exchange_name must be set as well. - :param exchange_name: Name of the exchange on which the messages will be published. - :param queue_ttl: Additional queue TTL arguments. - """ - if bind_to_routing_key is not None and exchange_name is None: - raise RuntimeError( - f"Routing key for binding was set to {bind_to_routing_key} but no " - f"exchange name was provided." - ) - - if queue_ttl is not None: - ttl_arguments = queue_ttl.to_argument() - else: - ttl_arguments = None - - logger.info("Declaring queue %s as %s with arguments as %s", - queue_name, - queue_type, - ttl_arguments) - queue = await self._channel.declare_queue(queue_name, - **queue_type.to_argument(), - arguments=ttl_arguments) - - if exchange_name is not None: - if exchange_name not in self._exchanges: - raise RuntimeError( - f"Exchange {exchange_name} was not yet declared by this connection." - ) - exchange = self._exchanges[exchange_name] - logger.info("Binding queue %s to routing key %s", queue_name, bind_to_routing_key) - await queue.bind(exchange=exchange, routing_key=bind_to_routing_key) - - return queue - - async def _declare_queue_and_add_subscription( - self, - queue_name: str, - callback_on_message: Callable[[bytes], None], - queue_type: AMQPQueueType, - bind_to_routing_key: Optional[str] = None, - exchange_name: Optional[str] = None, - delete_after_messages: Optional[int] = None, - queue_ttl: Optional[QueueTTLArguments] = None - ) -> None: - """Declare an AMQP queue and subscribe to the messages. - - :param queue_name: Name of the queue to declare. - :param callback_on_message: Callback which is called from a separate thread to process the - message body. - :param queue_type: Declare the queue using one of the known queue types. - :param bind_to_routing_key: Bind the queue to this routing key next to the default routing - key of the queue name. If none, the queue is only bound to the name of the queue. - If not none, then the exchange_name must be set as well. - :param exchange_name: Name of the exchange on which the messages will be published. - :param delete_after_messages: Delete the subscription & queue after this limit of messages - have been successfully processed. - :param queue_ttl: Additional queue TTL arguments. - """ - if queue_name in self._queue_subscription_consumer_by_name: - logger.error( - "Attempting to declare a subscription on %s but a " - "subscription on this queue already exists." - ) - raise RuntimeError(f"Queue subscription for {queue_name} already exists.") - - queue = await self._declare_queue( - queue_name, queue_type, bind_to_routing_key, exchange_name, queue_ttl - ) - - queue_consumer = QueueSubscriptionConsumer( - queue, delete_after_messages, callback_on_message - ) - self._queue_subscription_consumer_by_name[queue_name] = queue_consumer - - queue_subscription_task = asyncio.create_task(queue_consumer.run()) - queue_subscription_task.add_done_callback( - partial(self._remove_queue_subscription_task, queue_name) - ) - self._queue_subscription_tasks[queue_name] = queue_subscription_task - - async def _queue_exists(self, queue_name: str) -> bool: - """Check if the queue exists. - - :param queue_name: Name of the queue to be checked. - """ - try: - await self._channel.get_queue(queue_name, ensure=True) - logger.info("The %s queue exists", queue_name) - return True - except ChannelClosed as err: - logger.warning(err) - return False - - async def _remove_queue_subscription(self, queue_name: str) -> None: - """Remove subscription from queue and delete the queue if one exists. - - :param queue_name: Name of the queue to unsubscribe from. - """ - if queue_name in self._queue_subscription_tasks: - logger.info("Stopping subscription to %s and remove queue", queue_name) - self._queue_subscription_tasks[queue_name].cancel() - await self._channel.queue_delete(queue_name) - - def _remove_queue_subscription_task(self, queue_name: str, future: Future) -> None: - """Remove the queue subscription from the internal cache. - - :param queue_name: Name of the queue to which is subscribed. - :param future: Required argument from Task.add_done_callback which also refers to the - task running the subscription but as a `Future`. - """ - if queue_name in self._queue_subscription_tasks: - logger.debug("Queue subscription %s is done. Calling termination callback", queue_name) - del self._queue_subscription_consumer_by_name[queue_name] - del self._queue_subscription_tasks[queue_name] - - async def _setup_broker_interface(self) -> None: - """Start the AMQP connection and channel.""" - logger.info( - "Broker interface connecting to %s:%s as %s at %s", - self.config.host, - self.config.port, - self.config.username, - self.config.virtual_host, - ) - - self._connection = await connect_robust( - host=self.config.host, - port=self.config.port, - login=self.config.username, - password=self.config.password, - virtualhost=self.config.virtual_host, - loop=self._loop, - fail_fast="false", # aiormq requires this to be str and not bool - ) - self._channel = await self._connection.channel() - await self._channel.set_qos(prefetch_count=1) - self._ready_for_processing.set() - - async def _stop_broker_interface(self) -> None: - """Cancel all subscriptions, close the channel and the connection.""" - logger.info("Stopping broker interface") - tasks_to_cancel = list(self._queue_subscription_tasks.values()) - for queue_task in tasks_to_cancel: - queue_task.cancel() - if hasattr(self, "_channel") and self._channel: - await self._channel.close() - if hasattr(self, "_connection") and self._connection: - await self._connection.close() - logger.info("Stopped broker interface") - - def start(self) -> None: - """Start the broker interface.""" - super().start() - self._ready_for_processing.wait() - - def run(self) -> None: - """Run the broker interface and start the AMQP connection. - - In a separate thread and starting a new, isolated eventloop. The AMQP connection and - channel are started as its first task. - """ - self._loop = asyncio.new_event_loop() - asyncio.set_event_loop(self._loop) - - setup_task = None - try: - setup_task = self._loop.create_task(self._setup_broker_interface()) - self._loop.run_forever() - finally: - # Setup task is destroyed if no reference to the task is kept. This is just to check - # if the task was successful but also to keep the reference. - if setup_task and not setup_task.done(): - logger.error("Setup task was not completed even though it was created.") - elif setup_task is None: - logger.error("Setup task for AMQP connection was not created.") - self._loop.close() - - def declare_exchange(self, exchange_name: str) -> None: - """Declare an exchange on which messages may be published and routed to queues. - - :param exchange_name: Name of the exchange. - """ - asyncio.run_coroutine_threadsafe(self._declare_exchange(exchange_name), self._loop).result() - - def declare_queue( - self, - queue_name: str, - queue_type: AMQPQueueType, - bind_to_routing_key: Optional[str] = None, - exchange_name: Optional[str] = None, - queue_ttl: Optional[QueueTTLArguments] = None - ) -> None: - """Declare an AMQP queue. - - :param queue_name: Name of the queue to declare. - :param queue_type: Declare the queue using one of the known queue types. - :param bind_to_routing_key: Bind the queue to this routing key next to the default routing - key of the queue name. If none, the queue is only bound to the name of the queue. - If not none, then the exchange_name must be set as well. - :param exchange_name: Name of the exchange on which the messages will be published. - :param queue_ttl: Additional queue TTL arguments. - """ - asyncio.run_coroutine_threadsafe( - self._declare_queue( - queue_name=queue_name, - queue_type=queue_type, - bind_to_routing_key=bind_to_routing_key, - exchange_name=exchange_name, - queue_ttl=queue_ttl, - ), - self._loop, - ).result() - - def declare_queue_and_add_subscription( - self, - queue_name: str, - callback_on_message: Callable[[bytes], None], - queue_type: AMQPQueueType, - bind_to_routing_key: Optional[str] = None, - exchange_name: Optional[str] = None, - delete_after_messages: Optional[int] = None, - queue_ttl: Optional[QueueTTLArguments] = None - ) -> None: - """Declare an AMQP queue and subscribe to the messages. - - :param queue_name: Name of the queue to declare. - :param callback_on_message: Callback which is called from a separate thread to process the - message body. - :param queue_type: Declare the queue using one of the known queue types. - :param bind_to_routing_key: Bind the queue to this routing key next to the default routing - key of the queue name. If none, the queue is only bound to the name of the queue. - :param exchange_name: Name of the exchange on which the messages will be published. - :param delete_after_messages: Delete the subscription & queue after this limit of messages - have been successfully processed. - :param queue_ttl: Additional queue TTL arguments. - """ - asyncio.run_coroutine_threadsafe( - self._declare_queue_and_add_subscription( - queue_name=queue_name, - callback_on_message=callback_on_message, - queue_type=queue_type, - bind_to_routing_key=bind_to_routing_key, - exchange_name=exchange_name, - delete_after_messages=delete_after_messages, - queue_ttl=queue_ttl, - ), - self._loop, - ).result() - - def queue_exists(self, queue_name: str) -> bool: - """Check if the queue exists. - - :param queue_name: Name of the queue to be checked. - """ - return asyncio.run_coroutine_threadsafe( - self._queue_exists(queue_name=queue_name), self._loop - ).result() - - def remove_queue_subscription(self, queue_name: str) -> None: - """Remove subscription from queue and delete the queue if one exists. - - :param queue_name: Name of the queue to unsubscribe from. - """ - asyncio.run_coroutine_threadsafe( - self._remove_queue_subscription(queue_name), self._loop - ).result() - - def send_message_to( - self, exchange_name: Optional[str], routing_key: str, message: bytes - ) -> None: - """Publish a single message to the queue. - - :param exchange_name: The name of the exchange on which to publish the message. Must be - declared before publishing is possible. If None, the default exchange is used which is - already declared at startup - :param routing_key: The routing key to publish the message to. This may be a specific - routing key to which multiple queues are bound or the queue name (default routing key). - :param message: The message to publish. - """ - asyncio.run_coroutine_threadsafe( - self._send_message_to(exchange_name, routing_key, message), self._loop - ).result() - - def stop(self) -> None: - """Stop the broker interface. - - By shutting down the AMQP connection and stopping the eventloop. - """ - will_stop = False - with self._stopping_lock: - if not self._stopping: - self._stopping = True - will_stop = True - - if will_stop: - future = asyncio.run_coroutine_threadsafe(self._stop_broker_interface(), self._loop) - try: - future.result(timeout=BrokerInterface.TIMEOUT_ON_STOP_SECONDS) - except Exception: - logger.exception("Could not stop the broker interface during shutdown.") - self._loop.call_soon_threadsafe(self._loop.stop) - self._stopped = True diff --git a/src/omotes_sdk/internal/common/config.py b/src/omotes_sdk/internal/common/config.py deleted file mode 100644 index ee7d312..0000000 --- a/src/omotes_sdk/internal/common/config.py +++ /dev/null @@ -1,21 +0,0 @@ -import os -from omotes_sdk.config import RabbitMQConfig as RabbitMQConfig - - -class EnvRabbitMQConfig(RabbitMQConfig): - """Retrieve RabbitMQ configuration from environment variables.""" - - def __init__(self, prefix: str = ""): - """Create the RabbitMQ configuration and retrieve values from env vars. - - :param prefix: Prefix to the name environment variables. - """ - super().__init__( - host=os.environ.get(f"{prefix}RABBITMQ_HOSTNAME", RabbitMQConfig.host), - port=int(os.environ.get(f"{prefix}RABBITMQ_PORT", RabbitMQConfig.port)), - username=os.environ.get(f"{prefix}RABBITMQ_USERNAME", RabbitMQConfig.username), - password=os.environ.get(f"{prefix}RABBITMQ_PASSWORD", RabbitMQConfig.password), - virtual_host=os.environ.get( - f"{prefix}RABBITMQ_VIRTUALHOST", RabbitMQConfig.virtual_host - ), - ) diff --git a/src/omotes_sdk/internal/common/esdl_util.py b/src/omotes_sdk/internal/common/esdl_util.py deleted file mode 100644 index 7b2b6ac..0000000 --- a/src/omotes_sdk/internal/common/esdl_util.py +++ /dev/null @@ -1,13 +0,0 @@ -from esdl.esdl_handler import EnergySystemHandler - - -def pyesdl_from_string(input_str: str) -> EnergySystemHandler: - """ - Loads esdl file from a string into memory. - - Please note that it is not checked if the contents of the string is a valid esdl. - :param input_str: string containing the contents of an esdl file. - """ - esh = EnergySystemHandler() - esh.load_from_string(input_str) - return esh diff --git a/src/omotes_sdk/internal/orchestrator_worker_events/__init__.py b/src/omotes_sdk/internal/orchestrator_worker_events/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/src/omotes_sdk/internal/orchestrator_worker_events/esdl_messages.py b/src/omotes_sdk/internal/orchestrator_worker_events/esdl_messages.py deleted file mode 100644 index f9507cf..0000000 --- a/src/omotes_sdk/internal/orchestrator_worker_events/esdl_messages.py +++ /dev/null @@ -1,37 +0,0 @@ -from dataclasses import dataclass -from enum import Enum -from typing import Optional - -from omotes_sdk_protocol.job_pb2 import EsdlMessage as EsdlMessagePb - - -class MessageSeverity(Enum): - """Message severity options.""" - - DEBUG = "DEBUG" - INFO = "INFO" - WARNING = "WARNING" - ERROR = "ERROR" - - -@dataclass -class EsdlMessage: - """Esdl feedback message, optionally related to a specific object (asset).""" - - technical_message: str - """Technical message.""" - severity: MessageSeverity - """Message severity.""" - esdl_object_id: Optional[str] = None - """Optional esdl object id, None implies a general energy system message.""" - - def to_protobuf_message(self) -> EsdlMessagePb: - """Generate a protobuf message from this class. - - :return: Protobuf message representation. - """ - return EsdlMessagePb( - technical_message=self.technical_message, - severity=EsdlMessagePb.Severity.Value(self.severity.value), - esdl_object_id=self.esdl_object_id, - ) diff --git a/src/omotes_sdk/internal/orchestrator_worker_events/task_type.py b/src/omotes_sdk/internal/orchestrator_worker_events/task_type.py deleted file mode 100644 index 58d391b..0000000 --- a/src/omotes_sdk/internal/orchestrator_worker_events/task_type.py +++ /dev/null @@ -1,20 +0,0 @@ -from dataclasses import dataclass -from typing import List - - -@dataclass -class TaskType: - """Celery task type which may belong to a workflow.""" - - task_type_name: str - """Technical name for the task.""" - task_type_description_name: str - """Human-readable name for the task.""" - - -@dataclass -class TaskTypeManager: - """Container for all possible Celery tasks.""" - - possible_tasks: List[TaskType] - """All possible Celery tasks.""" diff --git a/src/omotes_sdk/internal/worker/__init__.py b/src/omotes_sdk/internal/worker/__init__.py deleted file mode 100644 index e53ee6e..0000000 --- a/src/omotes_sdk/internal/worker/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from omotes_sdk.internal.worker.worker import Worker, WORKER diff --git a/src/omotes_sdk/internal/worker/configs.py b/src/omotes_sdk/internal/worker/configs.py deleted file mode 100644 index 73f6571..0000000 --- a/src/omotes_sdk/internal/worker/configs.py +++ /dev/null @@ -1,38 +0,0 @@ -import os -from dataclasses import dataclass - -from omotes_sdk.config import RabbitMQConfig -from omotes_sdk.internal.common.config import ( - EnvRabbitMQConfig, -) - - -@dataclass -class WorkerConfig: - """Base configuration for an OMOTES Worker. - - Values are retrieved from environment variables. - """ - - rabbitmq_config: RabbitMQConfig - """How to connect to the OMOTES Celery RabbitMQ.""" - task_result_queue_name: str - """Name of the queue to which the task result should be send.""" - task_progress_queue_name: str - """Name of the queue to which progress updates for the task should be send.""" - log_level: str - """Log level for any logging in the worker.""" - max_tasks_before_restart: int - """The number of tasks (runs) after which the celery worker will be restarted.""" - - def __init__(self) -> None: - """Create the worker config and retrieve values from environment variables.""" - self.rabbitmq_config = EnvRabbitMQConfig() - self.task_result_queue_name = os.environ.get( - "TASK_RESULT_QUEUE_NAME", "omotes_task_result_events" - ) - self.task_progress_queue_name = os.environ.get( - "TASK_PROGRESS_QUEUE_NAME", "omotes_task_progress_events" - ) - self.log_level = os.environ.get("LOG_LEVEL", "INFO") - self.max_tasks_before_restart = int(os.environ.get("MAX_TASKS_BEFORE_RESTART", 0)) diff --git a/src/omotes_sdk/internal/worker/worker.py b/src/omotes_sdk/internal/worker/worker.py deleted file mode 100644 index 6027fd1..0000000 --- a/src/omotes_sdk/internal/worker/worker.py +++ /dev/null @@ -1,400 +0,0 @@ -import io -import logging -import socket -import sys -from typing import Any, Callable, Dict, List, Optional, Tuple -from uuid import UUID - -import streamcapture -from billiard.einfo import ExceptionInfo -from celery import Celery -from celery import Task as CeleryTask -from celery.apps.worker import Worker as CeleryWorker -from esdl import EnergySystem -from kombu import Queue as KombuQueue -from omotes_sdk_protocol.internal.task_pb2 import TaskProgressUpdate, TaskResult - -from omotes_sdk.internal.common.broker_interface import BrokerInterface -from omotes_sdk.internal.common.esdl_util import pyesdl_from_string -from omotes_sdk.internal.orchestrator_worker_events.esdl_messages import EsdlMessage -from omotes_sdk.internal.worker.configs import WorkerConfig -from omotes_sdk.types import ProtobufDict - -logger = logging.getLogger("omotes_sdk_internal") - - -class TaskUtil: - """Utilities for a Celery task.""" - - job_id: UUID - """The id of the job this task belongs to.""" - task: CeleryTask - """Reference to the Celery task which is executed.""" - broker_if: BrokerInterface - """Interface to the OMOTES Celery RabbitMQ.""" - - def __init__(self, job_id: UUID, task: CeleryTask, broker_if: BrokerInterface): - """Create the task utilities. - - :param job_id: The id of the job this task belongs to. - :param task: Reference to the Celery task which is executed. - :param broker_if: Interface to the OMOTES Celery RabbitMQ. - """ - self.job_id = job_id - self.task = task - self.broker_if = broker_if - - def send_start(self) -> None: - """Send a START progress update to the orchestrator.""" - logger.debug( - "Sending START progress update for job %s (celery id %s)", - self.job_id, - self.task.request.id, - ) - self.broker_if.send_message_to( - None, - WORKER.config.task_progress_queue_name, - TaskProgressUpdate( - job_id=str(self.job_id), - celery_task_id=self.task.request.id, - celery_task_type=self.task.name, - status=TaskProgressUpdate.START, - message="Started job at worker.", - ).SerializeToString(), - ) - - def update_progress(self, fraction: float, message: str) -> None: - """Send a progress update to the orchestrator. - - :param fraction: Value between 0.0 and 1.0 to denote the progress. - :param message: Message which explains the current progress of the task. - """ - logger.debug( - "Sending progress update. Progress %s for job %s (celery id %s) with message %s", - fraction, - self.job_id, - self.task.request.id, - message, - ) - self.broker_if.send_message_to( - None, - WORKER.config.task_progress_queue_name, - TaskProgressUpdate( - job_id=str(self.job_id), - celery_task_id=self.task.request.id, - celery_task_type=self.task.name, - progress=float(fraction), - message=message, - ).SerializeToString(), - ) - - -class WorkerTask(CeleryTask): - """Wrapped CeleryTask to connect to a number of CeleryTask events.""" - - logs: io.BytesIO - stdout_capturer: streamcapture.StreamCapture - stderr_capturer: streamcapture.StreamCapture - broker_if: BrokerInterface - - output_esdl: Optional[str] - esdl_messages: List[EsdlMessage] - - def __init__(self) -> None: - """Create the worker task.""" - super().__init__() - self.output_esdl = None - self.esdl_messages = [] - - def before_start(self, task_id: str, args: List[Any], kwargs: Dict[str, Any]) -> None: - """Runs before task start. - - :param task_id: Celery task id of the task. - :param args: Unnamed arguments passed to the CeleryTask. - :param kwargs: Named arguments passed to the CeleryTask. - """ - self.logs = io.BytesIO() - self.stdout_capturer = streamcapture.StreamCapture(sys.stdout, self.logs) - self.stderr_capturer = streamcapture.StreamCapture(sys.stderr, self.logs) - - self.broker_if = BrokerInterface(config=WORKER.config.rabbitmq_config) - self.broker_if.start() - - def after_return( - self, - status: str, - retval: Any, - task_id: str, - args: List[Any], - kwargs: Dict[str, Any], - einfo: str, - ) -> None: - """Runs after task finished. - - :param status: Task status. - :param retval: Task return value. - :param task_id: Celery task id of the task. - :param args: Unnamed arguments passed to the CeleryTask. - :param kwargs: Named arguments passed to the CeleryTask. - :param einfo: Context information regarding the exception triggered. - """ - self.stdout_capturer.close() - self.stderr_capturer.close() - self.logs.flush() - logs = self.logs.getvalue().decode() - self.logs.close() - - # to protobuf esdl messages: - esdl_messages_pb = [ - esdl_message.to_protobuf_message() for esdl_message in self.esdl_messages - ] - - job_id: UUID = args[0] - job_reference: str = args[1] - - result_message = None - if status == "FAILURE" or not self.output_esdl: - logger.info( - "Job %s (celery task id %s) with reference %s failed.", - job_id, - self.request.id, - job_reference, - ) - result_message = TaskResult( - job_id=str(job_id), - celery_task_id=self.request.id, - celery_task_type=self.name, - result_type=TaskResult.ResultType.ERROR, - output_esdl="", - logs=logs, - esdl_messages=esdl_messages_pb, - ) - elif status == "SUCCESS": - logger.info( - "Job %s (celery task id %s) with reference %s was successful.", - job_id, - self.request.id, - job_reference, - ) - result_message = TaskResult( - job_id=str(job_id), - celery_task_id=self.request.id, - celery_task_type=self.name, - result_type=TaskResult.ResultType.SUCCEEDED, - output_esdl=self.output_esdl, - logs=logs, - esdl_messages=esdl_messages_pb, - ) - else: - logger.error( - "Job %s with reference %s led to unexpected celery task state %s. Please report or " - "implement how to handle this state", - job_id, - job_reference, - status, - ) - - if result_message: - logger.debug("Sending result for job %s with reference %s", job_id, job_reference) - self.broker_if.send_message_to( - None, - WORKER.config.task_result_queue_name, - result_message.SerializeToString(), - ) - else: - logger.error( - "Did not send a job result for job %s with reference %s. This should not " - "happen. Status: %s", - job_id, - job_reference, - status, - ) - - self.broker_if.stop() - - def on_failure( - self, - exc: Exception, - task_id: str, - args: List[Any], - kwargs: Dict[str, Any], - einfo: ExceptionInfo, - ) -> None: - """Runs when the CeleryTask fails on an exception. - - :param exc: The exception triggered by the task. - :param task_id: Celery task id of the task. - :param args: Unnamed arguments passed to the CeleryTask. - :param kwargs: Named arguments passed to the CeleryTask. - :param einfo: Context information regarding the exception triggered. - """ - super().on_failure(exc, task_id, args, kwargs, einfo) - job_id: UUID = args[0] - job_reference: str = args[1] - logger.error( - "Failure detected for job %s with reference %s celery task %s", - job_id, - job_reference, - task_id, - ) - - -def wrapped_worker_task( - task: WorkerTask, job_id: UUID, job_reference: str, input_esdl: str, params_dict: ProtobufDict -) -> None: - """Task performed by Celery. - - Note: Be careful! This spawns within a subprocess and gains a copy of memory from parent - process. You cannot open sockets and other resources in the main process and expect - it to be copied to subprocess. Any resources e.g. connections/sockets need to be opened - in this task by the subprocess. - - :param task: Reference to the CeleryTask. - :param job_id: Id of the job this task belongs to. - :param job_reference: A user-submitted reference to this job. - :param input_esdl: ESDL description which needs to be processed. - :param params_dict: job, non-ESDL, parameters. - """ - logger.info("Worker started new task %s with reference %s", job_id, job_reference) - task_util = TaskUtil( - job_id, - task, - task.broker_if, - ) - task_util.send_start() - output_esdl, esdl_messages = WORKER_TASK_FUNCTION( - input_esdl, params_dict, task_util.update_progress, task.name - ) - - if output_esdl: - input_esh = pyesdl_from_string(input_esdl) - input_energy_system: EnergySystem = input_esh.energy_system - if job_reference is None: - new_name = f"{input_energy_system.name}_{task.name}" - elif job_reference == "": - new_name = f"{input_energy_system.name}" - else: - new_name = f"{input_energy_system.name}_{job_reference}" - - output_esh = pyesdl_from_string(output_esdl) - output_energy_system: EnergySystem = output_esh.energy_system - output_energy_system.name = new_name - task.output_esdl = output_esh.to_string() - else: - task.output_esdl = None - - task.esdl_messages = esdl_messages - - task_util.update_progress(1.0, "Calculation finished.") - - -class Worker: - """Worker which works as an Celery Worker application and runs a single type of task.""" - - config: WorkerConfig - """Configuration for the worker.""" - - captured_logging_string = io.StringIO() - """Captured logging while performing the task.""" - - celery_app: Celery - """Contains the Celery app configuration.""" - celery_worker: CeleryWorker - """The Celery Worker app.""" - - def __init__(self) -> None: - """Instantiate Worker instance config attribute.""" - self.config = WorkerConfig() - - def start(self) -> None: - """Start the `Worker`. - - Creates the Celery application, Celery Worker application and connects to RabbitMQ. - """ - rabbitmq_config = self.config.rabbitmq_config - self.celery_app = Celery( - broker=f"amqp://{rabbitmq_config.username}:{rabbitmq_config.password}@" - f"{rabbitmq_config.host}:{rabbitmq_config.port}/{rabbitmq_config.virtual_host}", - ) - - # Config of celery app - queues = [] - for worker_task_type in WORKER_TASK_TYPES: - logger.info("Starting Worker to work on task %s", worker_task_type) - queues.append( - KombuQueue( - worker_task_type, - routing_key=worker_task_type, - queue_arguments={"x-max-priority": 10}, - ) - ) - self.celery_app.task( - wrapped_worker_task, base=WorkerTask, name=worker_task_type, bind=True - ) - - self.celery_app.conf.task_queues = queues - self.celery_app.conf.task_acks_late = True - self.celery_app.conf.task_reject_on_worker_lost = True - self.celery_app.conf.task_acks_on_failure_or_timeout = False - self.celery_app.conf.worker_prefetch_multiplier = 1 - self.celery_app.conf.broker_transport_options = { - "priority_step": 1 - } # Prioritize higher numbers - self.celery_app.conf.broker_connection_retry_on_startup = True - # app.conf.worker_send_task_events = True # Tell the worker to send task events. - self.celery_app.conf.worker_hijack_root_logger = False - self.celery_app.conf.worker_redirect_stdouts = False - # optionally restart celery worker after set number of tasks (runs) - if self.config.max_tasks_before_restart != 0: - self.celery_app.conf.worker_max_tasks_per_child = self.config.max_tasks_before_restart - - logger.info( - "Connected to broker rabbitmq (%s:%s/%s) as %s", - rabbitmq_config.host, - rabbitmq_config.port, - rabbitmq_config.virtual_host, - rabbitmq_config.username, - ) - - self.celery_worker = self.celery_app.Worker( - hostname=f"worker-{'_'.join(WORKER_TASK_TYPES)}@{socket.gethostname()}", - loglevel=logging.getLevelName(self.config.log_level), - autoscale=(1, 1), - ) - - self.celery_worker.start() - - -UpdateProgressHandler = Callable[[float, str], None] -WorkerTaskF = Callable[ - [str, ProtobufDict, UpdateProgressHandler, str], - Tuple[ - Optional[str], - List[EsdlMessage], - ], -] - -WORKER: Worker = None # type: ignore [assignment] # noqa -WORKER_TASK_FUNCTION: WorkerTaskF = None # type: ignore [assignment] # noqa -WORKER_TASK_TYPES: list[str] = None # type: ignore [assignment] # noqa - - -def initialize_worker( - task_types: list[str], - task_function: WorkerTaskF, -) -> None: - """Initialize and run the `Worker`. - - :param task_types: Technical name of the tasks. Needs to be equal to the name of the celery task - to which the orchestrator forwards the task. May connect to one or more tasks. - :param task_function: Function which performs the Celery task. - """ - global WORKER_TASK_FUNCTION, WORKER_TASK_TYPES, WORKER - WORKER_TASK_TYPES = task_types - if len(WORKER_TASK_TYPES) < 1: - raise RuntimeError( - f"Should connect to one or more worker task types. Only found {len(WORKER_TASK_TYPES)}" - ) - WORKER_TASK_FUNCTION = task_function - WORKER = Worker() - WORKER.start() diff --git a/src/omotes_sdk/job.py b/src/omotes_sdk/job.py deleted file mode 100644 index 562a3ee..0000000 --- a/src/omotes_sdk/job.py +++ /dev/null @@ -1,14 +0,0 @@ -import uuid -from dataclasses import dataclass - -from omotes_sdk.workflow_type import WorkflowType - - -@dataclass -class Job: - """Reference to a submitted job.""" - - id: uuid.UUID - """ID used to reference to job in OMOTES system.""" - workflow_type: WorkflowType - """Describes the type of job / workflow has been submitted.""" diff --git a/src/omotes_sdk/job_status.py b/src/omotes_sdk/job_status.py new file mode 100644 index 0000000..60e96cc --- /dev/null +++ b/src/omotes_sdk/job_status.py @@ -0,0 +1,24 @@ +from enum import Enum, auto + + +class JobStatus(str, Enum): # noqa: UP042 + """Possible job status.""" + + @staticmethod + def _generate_next_value_(name: str, start: int, count: int, last_values: list[object]) -> str: + return name + + REGISTERED = auto() + """Job is registered but not yet submitted .""" + ENQUEUED = auto() + """Job is submitted but not yet started.""" + RUNNING = auto() + """Job is started and waiting to complete.""" + SUCCEEDED = auto() + """Job is finished successfully.""" + CANCELLED = auto() + """Job was cancelled.""" + TIMEOUT = auto() + """Job ended due to a timeout.""" + ERROR = auto() + """Job ended due to an error.""" diff --git a/src/omotes_sdk/log_forwarding.py b/src/omotes_sdk/log_forwarding.py new file mode 100644 index 0000000..85f8595 --- /dev/null +++ b/src/omotes_sdk/log_forwarding.py @@ -0,0 +1,241 @@ +import asyncio +import io +import logging +import queue +import re +import sys +import threading +from collections.abc import Callable, Coroutine +from datetime import UTC, datetime +from types import TracebackType +from typing import Any, TextIO + +import streamcapture + +from omotes_sdk.prefect_util import get_client, get_run_context + +LogEntry = tuple[datetime, str] + +_FORWARDER_QUEUE_TIMEOUT_SECONDS = 0.05 +_FORWARDER_BATCH_SIZE = 10 +_PYTHON_LOG_LINE_RE = re.compile( + r"^(?:\d{2}:\d{2}:\d{2}\.\d{3} \| [A-Z]+\s+\||\d{4}-\d{2}-\d{2} \d{2}:\d{2}:\d{2},\d{3} \[)" +) + + +class SolverStdoutCapture(io.RawIOBase): + """Capture solver stdout lines and pass normalized entries to a callback.""" + + def __init__(self, log_fn: Callable[[LogEntry], None]) -> None: + """Initialize the line-capture sink. + + Args: + log_fn: Callback that receives normalized log entries. + + """ + self._log_fn = log_fn + self._buffer = b"" + + def write(self, data, /) -> int: # noqa: ANN001 + """Buffer incoming bytes and emit complete lines. + + Args: + data: Byte chunk written to the stream. + + Returns: + int: Number of bytes consumed. + + """ + chunk = bytes(data) + self._buffer += chunk + while b"\n" in self._buffer: + raw_line, self._buffer = self._buffer.split(b"\n", 1) + self._emit(raw_line) + return len(chunk) + + def flush_remaining(self) -> None: + """Emit any trailing buffered data as a final log line.""" + if self._buffer: + self._emit(self._buffer) + self._buffer = b"" + + def _emit(self, raw_line: bytes) -> None: + line = raw_line.decode(errors="replace").rstrip() + if not line: + return + if _PYTHON_LOG_LINE_RE.match(line): + return + self._log_fn((datetime.now(UTC), line)) + + +class BatchedLogForwarder: + """Forward captured log entries asynchronously in small batches.""" + + def __init__( + self, + send_batch: Callable[[list[LogEntry]], Coroutine[Any, Any, None]], + queue_timeout_seconds: float = _FORWARDER_QUEUE_TIMEOUT_SECONDS, + batch_size: int = _FORWARDER_BATCH_SIZE, + ) -> None: + """Initialize a background log forwarder. + + Args: + send_batch: Async callback that sends a batch of entries. + queue_timeout_seconds: Maximum wait before flushing pending entries. + batch_size: Maximum entries sent per batch. + + """ + self._send_batch = send_batch + self._queue_timeout_seconds = queue_timeout_seconds + self._batch_size = batch_size + self._queue: queue.Queue[LogEntry | None] = queue.Queue() + self._thread = threading.Thread(target=self._run, name="solver-log-forwarder") + + def start(self) -> None: + """Start the background forwarding thread.""" + self._thread.start() + + def enqueue(self, entry: LogEntry) -> None: + """Queue a single log entry for forwarding.""" + self._queue.put_nowait(entry) + + def stop(self) -> None: + """Stop the forwarding thread after flushing pending entries.""" + self._queue.put(None) + self._thread.join() + + def _run(self) -> None: + pending: list[LogEntry] = [] + + def _flush_pending() -> None: + nonlocal pending + if not pending: + return + asyncio.run(self._send_batch(pending)) + pending = [] + + while True: + try: + entry = self._queue.get(timeout=self._queue_timeout_seconds) + except queue.Empty: + self._safe_flush(_flush_pending) + continue + + if entry is None: + break + + pending.append(entry) + if len(pending) >= self._batch_size: + self._safe_flush(_flush_pending) + + self._safe_flush(_flush_pending) + + @staticmethod + def _safe_flush(flush_fn: Callable[[], None]) -> None: + try: + flush_fn() + except Exception as exc: + print(f"Solver log forwarding failed: {exc}", file=sys.__stderr__, flush=True) + + +def make_prefect_log_sender( + name: str = "solver", + level: int = logging.INFO, +) -> Callable[[list[LogEntry]], Coroutine[Any, Any, None]]: + """Build an async sender that writes batched log entries to Prefect. + + Returns: + Callable[[list[LogEntry]], Coroutine[Any, Any, None]]: Async batch sender. + + Raises: + RuntimeError: If no flow run id is available in the current Prefect context. + + """ + run_context = get_run_context() + flow_run_obj = getattr(run_context, "flow_run", None) + flow_run_id = str(flow_run_obj.id) if flow_run_obj is not None else None + if flow_run_id is None: + raise RuntimeError("Missing flow run id in Prefect run context") + + async def _send(entries: list[LogEntry]) -> None: + async with get_client() as client: + await client.create_logs([ + { + "name": name, + "level": level, + "message": line, + "timestamp": ts.isoformat(), + "flow_run_id": flow_run_id, + "task_run_id": None, + } + for ts, line in entries + ]) + + return _send + + +class StdCaptureToLogSession: + """Context manager that captures solver stdout and forwards it in batches.""" + + def __init__( + self, + capture_stderr: bool = True, + name: str = "solver", + level: int = logging.INFO, + ) -> None: + """Initialize a capture session. + + Args: + capture_stderr: Whether stderr should be captured alongside stdout. + name: Prefect log name. + level: Prefect log level. + + """ + send_batch = make_prefect_log_sender(name=name, level=level) + self._capture_stderr = capture_stderr + self._forwarder = BatchedLogForwarder(send_batch) + self._line_captures: list[SolverStdoutCapture] = [] + self._stream_targets: list[tuple[TextIO, SolverStdoutCapture]] = [] + self._active_stream_captures: list[streamcapture.StreamCapture] = [] + + stdout_capture = SolverStdoutCapture(self._forwarder.enqueue) + self._line_captures.append(stdout_capture) + self._stream_targets.append((sys.stdout, stdout_capture)) + + if self._capture_stderr: + stderr_capture = SolverStdoutCapture(self._forwarder.enqueue) + self._line_captures.append(stderr_capture) + self._stream_targets.append((sys.stderr, stderr_capture)) + + def __enter__(self) -> "StdCaptureToLogSession": + """Start stream capture and return the active context manager. + + Returns: + StdCaptureToLogSession: Active capture session. + + """ + self._forwarder.start() + self._active_stream_captures = [ + streamcapture.StreamCapture(source, target, echo=True) for source, target in self._stream_targets + ] + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_val: BaseException | None, + exc_tb: TracebackType | None, + ) -> None: + """Stop stream capture and flush pending log lines. + + Args: + exc_type: Exception type raised in the managed block, if any. + exc_val: Exception instance raised in the managed block, if any. + exc_tb: Exception traceback raised in the managed block, if any. + + """ + for stream_capture in self._active_stream_captures: + stream_capture.close() + for line_capture in self._line_captures: + line_capture.flush_remaining() + self._forwarder.stop() diff --git a/src/omotes_sdk/omotes_interface.py b/src/omotes_sdk/omotes_interface.py deleted file mode 100644 index 58d01c8..0000000 --- a/src/omotes_sdk/omotes_interface.py +++ /dev/null @@ -1,438 +0,0 @@ -from __future__ import annotations -import logging -import threading -import uuid -from dataclasses import dataclass -from datetime import timedelta -from typing import Callable, Optional, Union - -from omotes_sdk.internal.common.broker_interface import ( - BrokerInterface, - AMQPQueueType, - QueueTTLArguments, -) -from omotes_sdk.config import RabbitMQConfig -from omotes_sdk_protocol.job_pb2 import ( - JobResult, - JobProgressUpdate, - JobStatusUpdate, - JobSubmission, - JobDelete, -) -from omotes_sdk_protocol.workflow_pb2 import AvailableWorkflows, RequestAvailableWorkflows - -from omotes_sdk.job import Job -from omotes_sdk.queue_names import OmotesQueueNames -from omotes_sdk.types import ParamsDict -from omotes_sdk.workflow_type import ( - WorkflowType, - WorkflowTypeManager, - convert_params_dict_to_struct, -) - -logger = logging.getLogger("omotes_sdk") - - -@dataclass -class JobSubmissionCallbackHandler: - """Handler for updates on a submitted job.""" - - job: Job - """The job to listen to.""" - callback_on_finished: Callable[[Job, JobResult], None] - """Handler which is called when a job finishes.""" - callback_on_progress_update: Optional[Callable[[Job, JobProgressUpdate], None]] - """Handler which is called on a job progress update.""" - callback_on_status_update: Optional[Callable[[Job, JobStatusUpdate], None]] - """Handler which is called on a job status update.""" - auto_disconnect_on_result_handler: Optional[Callable[[Job], None]] - """Handler to remove/disconnect from all queues pertaining to this job once the result is - received and handled without exceptions through `callback_on_finished`.""" - - def callback_on_finished_wrapped(self, message: bytes) -> None: - """Parse a serialized JobResult message and call handler. - - :param message: Serialized message. - """ - job_result = JobResult() - job_result.ParseFromString(message) - self.callback_on_finished(self.job, job_result) - - if self.auto_disconnect_on_result_handler: - self.auto_disconnect_on_result_handler(self.job) - - def callback_on_progress_update_wrapped(self, message: bytes) -> None: - """Parse a serialized JobProgressUpdate message and call handler. - - :param message: Serialized message. - """ - progress_update = JobProgressUpdate() - progress_update.ParseFromString(message) - if self.callback_on_progress_update: - self.callback_on_progress_update(self.job, progress_update) - - def callback_on_status_update_wrapped(self, message: bytes) -> None: - """Parse a serialized JobStatusUpdate message and call handler. - - :param message: Serialized message. - """ - status_update = JobStatusUpdate() - status_update.ParseFromString(message) - if self.callback_on_status_update: - self.callback_on_status_update(self.job, status_update) - - -class UndefinedWorkflowsException(Exception): - """Thrown if the workflows are needed but not defined yet.""" - - ... - - -class UnknownWorkflowException(Exception): - """Thrown if a job is submitted using an unknown workflow type.""" - - ... - - -class OmotesInterface: - """SDK interface for other applications to communicate with OMOTES.""" - - broker_if: BrokerInterface - """Interface to RabbitMQ broker.""" - workflow_type_manager: Union[WorkflowTypeManager, None] - """All available workflow types.""" - _workflow_config_received: threading.Event - """Event triggered when workflow configuration is received.""" - client_id: str - """Identifier of this SDK instance. Should be unique from other SDKs.""" - timeout_on_initial_workflow_definitions: timedelta - """How long the SDK should wait for the first reply when requesting the current workflow - definitions from the orchestrator.""" - - JOB_QUEUES_TTL: timedelta = timedelta(hours=48) - """Default value of job result, progress, and status queue TTL.""" - - def __init__( - self, - rabbitmq_config: RabbitMQConfig, - client_id: str, - timeout_on_initial_workflow_definitions: timedelta = timedelta(minutes=1), # noqa: B008 - ): - """Create the OMOTES interface. - - NOTE: Needs to be started separately. - - :param rabbitmq_config: RabbitMQ configuration how to connect to OMOTES. - :param client_id: Identifier of this SDK instance. Should be unique from other SDKs. - :param timeout_on_initial_workflow_definitions: How long the SDK should wait for the first - reply when requesting the current workflow definitions from the orchestrator. - """ - self.broker_if = BrokerInterface(rabbitmq_config) - self.workflow_type_manager = None - self._workflow_config_received = threading.Event() - self.client_id = client_id - self.timeout_on_initial_workflow_definitions = timeout_on_initial_workflow_definitions - - def start(self) -> None: - """Start any other interfaces and request available workflows.""" - self.broker_if.start() - self.broker_if.declare_exchange(OmotesQueueNames.omotes_exchange_name()) - # Declare job submissions queue just in case the orchestrator isn't life yet so we do not - # lose work. - self.broker_if.declare_queue( - queue_name=OmotesQueueNames.job_submission_queue_name(), - queue_type=AMQPQueueType.DURABLE, - exchange_name=OmotesQueueNames.omotes_exchange_name(), - ) - self.connect_to_available_workflows_updates() - self.request_available_workflows() - - logger.info("Waiting for workflow definitions to be received from the orchestrator...") - if not self._workflow_config_received.wait( - timeout=self.timeout_on_initial_workflow_definitions.total_seconds() - ): - raise RuntimeError( - "The orchestrator did not send the available workflows within " - "the expected time. Is the orchestrator online and connected " - "to the broker?" - ) - - def stop(self) -> None: - """Stop any other interfaces.""" - self.broker_if.stop() - - def disconnect_from_submitted_job(self, job: Job) -> None: - """Disconnect from the submitted job and delete all queues on the broker. - - :param job: Job to disconnect from. - """ - self.broker_if.remove_queue_subscription(OmotesQueueNames.job_results_queue_name(job.id)) - self.broker_if.remove_queue_subscription(OmotesQueueNames.job_progress_queue_name(job.id)) - self.broker_if.remove_queue_subscription(OmotesQueueNames.job_status_queue_name(job.id)) - - def _autodelete_progres_status_queues_on_result(self, job: Job) -> None: - """Disconnect and delete the progress and status queues for some job.""" - self.broker_if.remove_queue_subscription(OmotesQueueNames.job_progress_queue_name(job.id)) - self.broker_if.remove_queue_subscription(OmotesQueueNames.job_status_queue_name(job.id)) - - def connect_to_submitted_job( - self, - job: Job, - callback_on_finished: Callable[[Job, JobResult], None], - callback_on_progress_update: Optional[Callable[[Job, JobProgressUpdate], None]], - callback_on_status_update: Optional[Callable[[Job, JobStatusUpdate], None]], - auto_disconnect_on_result: bool, - auto_cleanup_after_ttl: Optional[timedelta] = JOB_QUEUES_TTL, - reconnect: bool = True, - ) -> None: - """(Re)connect to the running job. - - Useful when the application using this SDK restarts and needs to reconnect to already - running jobs. Assumes that the job exists otherwise the callbacks will never be called. - - :param job: The job to reconnect to. - :param callback_on_finished: Called when the job has a result. - :param callback_on_progress_update: Called when there is a progress update for the job. - :param callback_on_status_update: Called when there is a status update for the job. - :param auto_disconnect_on_result: Remove/disconnect from all queues pertaining to this job - once the result is received and handled without exceptions through - `callback_on_finished`. - :param auto_cleanup_after_ttl: When erroneous situations occur (e.g. client is offline), - all queues pertaining to this job will be removed after the given TTL. - Default to 48 hours if unset. Set to `None` to turn off auto clean up, - but be aware this may lead to leftover messages and queues to be stored - in RabbitMQ indefinitely (which uses up memory & disk space). - :param reconnect: When True, first check the job queues status and raise an error if not - exist. Default to True. - """ - job_results_queue_name = OmotesQueueNames.job_results_queue_name(job.id) - job_progress_queue_name = OmotesQueueNames.job_progress_queue_name(job.id) - job_status_queue_name = OmotesQueueNames.job_status_queue_name(job.id) - - if reconnect: - logger.info( - "Reconnect to the submitted job %s is set to True. " - + "Checking job queues status...", - job.id, - ) - if not self.broker_if.queue_exists(job_results_queue_name): - raise RuntimeError( - f"The {job_results_queue_name} queue does not exist or is removed. " - "Abort reconnecting to the queue." - ) - if callback_on_progress_update and not self.broker_if.queue_exists( - job_progress_queue_name - ): - raise RuntimeError( - f"The {job_progress_queue_name} queue does not exist or is removed. " - "Abort reconnecting to the queue." - ) - if callback_on_status_update and not self.broker_if.queue_exists(job_status_queue_name): - raise RuntimeError( - f"The {job_status_queue_name} queue does not exist or is removed. " - "Abort reconnecting to the queue." - ) - - if auto_disconnect_on_result: - logger.info("Connecting to update for job %s with auto disconnect on result", job.id) - auto_disconnect_handler = self._autodelete_progres_status_queues_on_result - else: - logger.info("Connecting to update for job %s and expect manual disconnect", job.id) - auto_disconnect_handler = None - - if auto_cleanup_after_ttl is not None: - queue_ttl = auto_cleanup_after_ttl - logger.info( - "Auto job queues clean up on error after TTL is set. " - + "The leftover job messages will be dropped, " - + "and leftover job queues will be discarded after %s.", - queue_ttl, - ) - job_queue_ttl = QueueTTLArguments(queue_ttl=queue_ttl) - else: - logger.info( - "Auto job queues clean up on error after TTL is not set. " - + "Manual cleanup on leftover job queues and messages might be required." - ) - job_queue_ttl = None - - callback_handler = JobSubmissionCallbackHandler( - job, - callback_on_finished, - callback_on_progress_update, - callback_on_status_update, - auto_disconnect_handler, - ) - - if callback_on_progress_update: - self.broker_if.declare_queue_and_add_subscription( - queue_name=job_progress_queue_name, - callback_on_message=callback_handler.callback_on_progress_update_wrapped, - queue_type=AMQPQueueType.DURABLE, - exchange_name=OmotesQueueNames.omotes_exchange_name(), - queue_ttl=job_queue_ttl, - ) - if callback_on_status_update: - self.broker_if.declare_queue_and_add_subscription( - queue_name=job_status_queue_name, - callback_on_message=callback_handler.callback_on_status_update_wrapped, - queue_type=AMQPQueueType.DURABLE, - exchange_name=OmotesQueueNames.omotes_exchange_name(), - queue_ttl=job_queue_ttl, - ) - # Declares the job result queue last to ensure the job progress - # and status queues are already declared and exist. This order allows - # them to be properly removed by _autodelete_progres_status_queues_on_result() - self.broker_if.declare_queue_and_add_subscription( - queue_name=job_results_queue_name, - callback_on_message=callback_handler.callback_on_finished_wrapped, - queue_type=AMQPQueueType.DURABLE, - exchange_name=OmotesQueueNames.omotes_exchange_name(), - delete_after_messages=1, - queue_ttl=job_queue_ttl, - ) - - def submit_job( - self, - esdl: str, - params_dict: ParamsDict, - workflow_type: WorkflowType, - job_timeout: Optional[timedelta], - callback_on_finished: Callable[[Job, JobResult], None], - callback_on_progress_update: Optional[Callable[[Job, JobProgressUpdate], None]], - callback_on_status_update: Optional[Callable[[Job, JobStatusUpdate], None]], - auto_disconnect_on_result: bool, - job_reference: Optional[str] = None, - auto_cleanup_after_ttl: Optional[timedelta] = JOB_QUEUES_TTL, - job_priority: Union[JobSubmission.JobPriority, int] = JobSubmission.JobPriority.MEDIUM, - ) -> Job: - """Submit a new job and connect to progress and status updates and the job result. - - :param esdl: String containing the XML that make up the ESDL. - :param params_dict: Dictionary containing any job-specific, non-ESDL, configuration - parameters. Dictionary supports: - str, Union[Struct, ListValue, str, float, bool, None, Mapping[str, Any], Sequence] - :param workflow_type: Type of the workflow to start. - :param job_timeout: How long the job may take before it is considered to be timeout. - If None is given, the job is immune from the periodic timeout job cancellation task. - :param callback_on_finished: Callback which is called with the job result once the job is - done. - :param callback_on_progress_update: Callback which is called with any progress updates. - :param callback_on_status_update: Callback which is called with any status updates. - :param auto_disconnect_on_result: Remove/disconnect from all queues pertaining to this job - once the result is received and handled without exceptions through - `callback_on_finished`. - :param job_reference: An optional reference to the submitted job which is used in the - name of the output ESDL as well as in internal logging of OMOTES. - :param job_priority: An optional priority value for the job used in celery. - :param auto_cleanup_after_ttl: When erroneous situations occur (e.g. client is offline), - all queues pertaining to this job will be removed after the given TTL. - Default to 48 hours if unset. Set to `None` to turn off auto clean up, - but be aware this may lead to leftover messages and queues to be stored - in RabbitMQ indefinitely (which uses up memory & disk space). - :raises UnknownWorkflowException: If `workflow_type` is unknown as a possible workflow in - this interface. - :return: The job handle which is created. This object needs to be saved persistently by the - program using this SDK in order to resume listening to jobs in progress after a restart. - """ - if not self.workflow_type_manager or not self.workflow_type_manager.workflow_exists( - workflow_type - ): - raise UnknownWorkflowException() - - job = Job(id=uuid.uuid4(), workflow_type=workflow_type) - reconnect = False - logger.info("Submitting job %s with reference %s", job.id, job_reference) - self.connect_to_submitted_job( - job, - callback_on_finished, - callback_on_progress_update, - callback_on_status_update, - auto_disconnect_on_result, - auto_cleanup_after_ttl, - reconnect, - ) - - if job_timeout is not None: - timeout_ms = round(job_timeout.total_seconds() * 1000) - else: - timeout_ms = None - - job_submission_msg = JobSubmission( - uuid=str(job.id), - timeout_ms=timeout_ms, - workflow_type=workflow_type.workflow_type_name, - esdl=esdl, - params_dict=convert_params_dict_to_struct(workflow_type, params_dict), - job_reference=job_reference, - job_priority=job_priority, # type: ignore [arg-type] - ) - self.broker_if.send_message_to( - exchange_name=OmotesQueueNames.omotes_exchange_name(), - routing_key=OmotesQueueNames.job_submission_queue_name(), - message=job_submission_msg.SerializeToString(), - ) - logger.debug("Done submitting job %s with reference %s", job.id, job_reference) - - return job - - def delete_job(self, job: Job) -> None: - """Delete a job and all of its resources. - - This will delete the job regardless of its current state. If it is running, it will be - cancelled. If the job produced any timeseries data, it will be deleted eventually. - - Developers note: - If the jobs is successfully cancelled or not will be sent as a job status update through - the `callback_on_status_update` handler. This method will not disconnect from the submitted - job events. This will need to be done separately using `disconnect_from_submitted_job` - after receiving the job status update. - Deletion of the timeseries is done by the orchestrator. See: - https://github.com/Project-OMOTES/architecture-documentation/blob/main/Feature_Time_Series_DB_Cleanup/Feature_Time_Series_DB_Cleanup.md - - :param job: The job to delete. - """ - logger.info("Deleting job %s", job.id) - delete_msg = JobDelete(uuid=str(job.id)) - self.broker_if.send_message_to( - exchange_name=OmotesQueueNames.omotes_exchange_name(), - routing_key=OmotesQueueNames.job_delete_queue_name(), - message=delete_msg.SerializeToString(), - ) - - def connect_to_available_workflows_updates(self) -> None: - """Connect to updates of the available workflows.""" - self.broker_if.declare_queue_and_add_subscription( - queue_name=OmotesQueueNames.available_workflows_queue_name(self.client_id), - callback_on_message=self.callback_on_update_available_workflows, - queue_type=AMQPQueueType.EXCLUSIVE, - bind_to_routing_key=OmotesQueueNames.available_workflows_routing_key(), - exchange_name=OmotesQueueNames.omotes_exchange_name(), - ) - - def callback_on_update_available_workflows(self, message: bytes) -> None: - """Parse a serialized AvailableWorkflows message and update workflow type manager. - - :param message: Serialized message. - """ - available_workflows_pb = AvailableWorkflows() - available_workflows_pb.ParseFromString(message) - self.workflow_type_manager = WorkflowTypeManager.from_pb_message(available_workflows_pb) - self._workflow_config_received.set() - logger.info("Updated the available workflows") - - def request_available_workflows(self) -> None: - """Request the available workflows from the orchestrator.""" - request_available_workflows_pb = RequestAvailableWorkflows() - self.broker_if.send_message_to( - exchange_name=OmotesQueueNames.omotes_exchange_name(), - routing_key=OmotesQueueNames.request_available_workflows_queue_name(), - message=request_available_workflows_pb.SerializeToString(), - ) - - def get_workflow_type_manager(self) -> WorkflowTypeManager: - """Get the available workflows.""" - if self.workflow_type_manager: - return self.workflow_type_manager - else: - raise UndefinedWorkflowsException() diff --git a/src/omotes_sdk/prefect_util.py b/src/omotes_sdk/prefect_util.py new file mode 100644 index 0000000..45b1ec7 --- /dev/null +++ b/src/omotes_sdk/prefect_util.py @@ -0,0 +1,663 @@ +import json +import logging +import os +import re +from collections.abc import Callable +from typing import Any, cast +from uuid import UUID + +import httpx +from prefect.artifacts import ( + create_link_artifact, + create_progress_artifact, + update_progress_artifact, +) +from prefect.blocks.system import Secret +from prefect.client.orchestration import get_client +from prefect.client.schemas.filters import ( + ArtifactFilter, + ArtifactFilterFlowRunId, + DeploymentFilter, + DeploymentFilterName, + LogFilter, +) +from prefect.client.schemas.sorting import DeploymentSort +from prefect.context import ( + get_run_context, +) +from prefect.exceptions import ObjectNotFound +from prefect.filesystems import RemoteFileSystem +from prefect.runtime import flow_run +from prefect.states import StateType +from pydantic import BaseModel + +from omotes_sdk.job_status import JobStatus + + +def _build_minio_result_storage(minio_host: str, access_key: str, secret_key: str) -> RemoteFileSystem: + """Create MinIO-backed Prefect result storage block when env vars are available. + + Returns: + RemoteFileSystem: Configured Prefect result storage block. + + """ + bucket = "prefect-results" + prefix = "flow-results" + endpoint_url = f"http://{minio_host}" + + return RemoteFileSystem( + basepath=f"s3://{bucket}/{prefix}", + settings={ + "key": access_key, + "secret": secret_key, + "client_kwargs": {"endpoint_url": endpoint_url}, + }, + ) + + +for _noisy in ( + "httpcore", + "httpx", + "websockets", + "asyncio", + "hpack", + "botocore", + "aiobotocore", + "s3fs", + "fsspec", + "urllib3", +): + logging.getLogger(_noisy).setLevel(logging.WARNING) + + +def in_prefect_flow_context() -> bool: + """Check whether execution is inside a Prefect flow context. + + Returns: + bool: True when called in a Prefect flow context, otherwise False. + + """ + try: + get_run_context() + return True + except Exception: + return False + + +def get_flow_run_id_first_part() -> str: + """Get first part of flow run id. + + For local runs a default value is returned. + + Returns: + str: First 8 characters of the flow run id, or a default value for local runs. + + """ + return flow_run.id[:8] if flow_run and flow_run.id else "12345678" + + +def load_gurobi_license() -> None: + """Load Gurobi license content from Prefect Secret and write it to disk.""" + target_dir = "/app/gurobi" + target_file = os.path.join(target_dir, "gurobi.lic") + os.makedirs(target_dir, exist_ok=True) + + # get license content from Prefect Secret block + secret_block = cast(Secret[Any], Secret.load("gurobi-wls-secret")) + license_content = secret_block.get() + + with open(target_file, "w") as f: + f.write(license_content) + + logging.info(f"Successfully generated {target_file}") + + +def _sanitize_for_minio(value: str) -> str: + """Sanitize a value for Prefect artifact keys. + + Returns: + str: Lower-case alpha-numeric and dash-only string. + + """ + sanitized = re.sub(r"[^a-z0-9-]+", "-", value.lower().strip()) + sanitized = re.sub(r"-+", "-", sanitized).strip("-") + return sanitized or "field" + + +def _create_minio_presigned_url( + minio_block: RemoteFileSystem, object_path: str, expires_seconds: int = 7 * 24 * 60 * 60 +) -> str | None: + """Create a presigned GET URL for an object written via minio_block. + + Returns: + str | None: The generated URL or None when URL generation fails. + + """ + try: + resolved_path = minio_block._resolve_path(object_path) + return cast(str, minio_block.filesystem.sign(resolved_path, expiration=expires_seconds)) + except Exception: + logging.exception("Failed to create presigned URL for MinIO object") + return None + + +def _get_required_file_extension(result: BaseModel, field_name: str) -> str: + """Return the configured file extension for a BaseModel field. + + Returns: + str: Lower-cased file extension for the field. + + Raises: + ValueError: If the field has no valid file_extension entry. + + """ + field_info = type(result).model_fields.get(field_name) + extension = None + if field_info is not None and isinstance(field_info.json_schema_extra, dict): + extension = field_info.json_schema_extra.get("file_extension") + + if not isinstance(extension, str) or not extension.startswith("."): + raise ValueError( + f"Field '{field_name}' in model '{type(result).__name__}' must define " + "json_schema_extra={'file_extension': '.ext'}" + ) + + return extension.lower() + + +def write_flow_return_artifact_to_minio( + flow_result: BaseModel, minio_host: str, access_key: str, secret_key: str +) -> str | None: + """Persist flow return fields to MinIO and publish Prefect links to those objects. + + Returns: + str | None: Run folder path in MinIO, or None if not in flow context. + + """ + if not in_prefect_flow_context(): + return None + + minio_block = _build_minio_result_storage(minio_host, access_key, secret_key) + run_folder_path = _sanitize_for_minio(f"{flow_run.get_name()}-{get_flow_run_id_first_part()}") + + for field_name, field_value in flow_result: + if field_value is None: + continue + + artifact_key = _sanitize_for_minio(f"{field_name}-{get_flow_run_id_first_part()}") + field_extension = _get_required_file_extension(flow_result, field_name) + field_object_path = f"{run_folder_path}/{artifact_key}{field_extension}" + + field_bytes = ( + field_value.encode("utf-8") + if isinstance(field_value, str) + else json.dumps(field_value, indent=2, default=str).encode("utf-8") + ) + + try: + minio_block.write_path(path=field_object_path, content=field_bytes) + except Exception: + logging.exception("Failed to persist flow return field '%s' to MinIO", field_name) + continue + + presigned_url = _create_minio_presigned_url(minio_block, field_object_path) + if presigned_url is None: + continue + + create_link_artifact( + link=presigned_url, + link_text=f"Download '{field_name}' from MinIO", + key=artifact_key, + description=f"MinIO object: {field_object_path}", + ) + + return run_folder_path + + +def create_flow_progress_artifact( + key: str, + start_progress: float = 0.0, + start_description: str | None = None, +) -> UUID | None: + """Create a flow progress artifact and return its id. + + Returns: + UUID | None: Created artifact id, or None if unavailable. + + """ + if not in_prefect_flow_context(): + return None + + try: + return cast( + UUID, + create_progress_artifact( + key=key, + progress=start_progress, + description=start_description, + ), + ) + except Exception: + return None + + +def create_flow_progress_updater( + start_progress_fraction: float = 0.0, + start_description: str | None = None, +) -> Callable[[float, str | None], None]: + """Create a progress artifact and return a safe updater function. + + Returns: + Callable[[float, str | None], None]: Function that updates progress artifacts. + + """ + progress_key = f"progress-{get_flow_run_id_first_part()}" + artifact_id = create_flow_progress_artifact( + key=progress_key, + start_progress=start_progress_fraction * 100.0, + start_description=start_description, + ) + + def _update(progress_fraction: float, description: str | None = None) -> None: + if artifact_id is None: + logging.error(f"Error logging progress update for run '{flow_run.get_name()}': artifact creation failed") + else: + update_progress_artifact( + artifact_id=artifact_id, progress=progress_fraction * 100.0, description=description + ) + + return _update + + +def _memory_quantity_to_bytes(memory_limit: str) -> int: + """Convert a Kubernetes-style memory quantity into bytes for Docker. + + Returns: + int: Memory quantity in bytes. + + Raises: + ValueError: If the input memory quantity has an unsupported format. + + """ + normalized = memory_limit.strip() + match = re.fullmatch(r"(?i)(\d+(?:\.\d+)?)([kmgtpe]i?|)", normalized) + if not match: + raise ValueError(f"Unsupported memory quantity: {memory_limit!r}") + + value = float(match.group(1)) + suffix = match.group(2).lower() + binary_factors = { + "": 1, + "k": 10**3, + "m": 10**6, + "g": 10**9, + "t": 10**12, + "p": 10**15, + "e": 10**18, + "ki": 1024, + "mi": 1024**2, + "gi": 1024**3, + "ti": 1024**4, + "pi": 1024**5, + "ei": 1024**6, + } + return int(value * binary_factors[suffix]) + + +def build_universal_job_vars(memory_limit: str | None = None, base_vars: dict | None = None) -> dict | None: + """Build a job_variables payload for Kubernetes-style memory input. + + Kubernetes workers use memory_request and memory_limit directly. + Docker workers use mem_limit, so we convert the Kubernetes quantity to bytes in Python. + + Returns: + dict | None: Job variables suitable for worker deployment or None if unchanged. + + """ + job_vars = dict(base_vars or {}) + + if memory_limit: + job_vars["memory_limit"] = memory_limit + job_vars["memory_request"] = memory_limit + job_vars["mem_limit"] = _memory_quantity_to_bytes(memory_limit) + + return job_vars or None + + +async def deploy_flow( + flow_function: object, + deployment_name: str, + image_name: str, + job_variables: dict, + prefect_work_pool_name: str, +) -> None: + """Deploy prefect flow with variables. + + Will raise an error when trying to overwrite an existing semantic version (x.x.x). + Note: the 'flow' object cannot be annotated properly since @flow is dynamically typed. + + Args: + flow_function: Flow object exposing a deploy method. + deployment_name: Name of the deployment. + image_name: Name of the Docker image. + job_variables: Job variables. + prefect_work_pool_name: Name of the Prefect work pool. + + Raises: + RuntimeError: If a semantic-versioned deployment already exists. + + """ + version_name = image_name.rsplit(":", 1)[-1] + + async with get_client() as client: + deployments = await client.read_deployments( + deployment_filter=DeploymentFilter(name=DeploymentFilterName(any_=[deployment_name])), + sort=DeploymentSort.CREATED_DESC, + ) + if deployments and is_semantic_version(version_name): + raise RuntimeError(f"Prefect flow cannot be overwritten for semantic version '{deployment_name}'") + + deployable_flow = cast(Any, flow_function) + await deployable_flow.deploy( + name=deployment_name, + work_pool_name=prefect_work_pool_name, + image=image_name, + build=False, + job_variables=job_variables, + ) + + +# Support full semantic versions like 1.2.3-beta+exp.sha.5114f85 +SEMVER_REGEX = re.compile( + r""" + ^ + (0|[1-9]\d*)\. # major + (0|[1-9]\d*)\. # minor + (0|[1-9]\d*) # patch + (?:-([0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))? # prerelease (optional) + (?:\+([0-9A-Za-z-]+(?:\.[0-9A-Za-z-]+)*))? # build metadata (optional) + $ + """, + re.VERBOSE, +) + + +def is_semantic_version(version_name: str) -> bool: + """Return whether a version string follows Semantic Versioning. + + Returns: + bool: True for semantic version strings, otherwise False. + + """ + return bool(SEMVER_REGEX.match(version_name)) + + +async def get_flow_versions_by_name(flow_names: list[str]) -> dict[str, list[str]]: + """Return Prefect deployment versions keyed by flow name. + + Reads deployments once and groups version suffixes by the deployment + base name before the `:` separator. + """ + requested_flow_names = {flow_name for flow_name in flow_names if flow_name} + if not requested_flow_names: + return {} + + async with get_client() as client: + deployments = await client.read_deployments( + sort=DeploymentSort.CREATED_DESC, + ) + + result: dict[str, list[str]] = {flow_name: [] for flow_name in requested_flow_names} + for deployment in deployments: + if ":" not in deployment.name: + continue + + name_part, version_part = deployment.name.split(":", 1) + if name_part in requested_flow_names: + result[name_part].append(version_part) + + for flow_name, versions in result.items(): + result[flow_name] = sorted(versions, key=_version_sort_key, reverse=True) + + return result + + +def _version_sort_key(version_name: str) -> tuple[int, tuple[int, int, int], int, str]: + """Sort versions with semantic versions first, then other strings. + + Semantic versions are sorted by newest version. Non-semantic version + strings are kept sortable and come after semantic versions. + + Returns: + tuple[int, tuple[int, int, int], int, str]: Sort key that prefers + semantic versions and keeps non-semantic versions stable. + """ + match = SEMVER_REGEX.match(version_name) + if match is None: + return (0, (0, 0, 0), 0, version_name) + + major, minor, patch = (int(match.group(index)) for index in range(1, 4)) + prerelease = match.group(4) + return (1, (major, minor, patch), 1 if prerelease is None else 0, prerelease or "") + + +async def get_flows_with_versions(identifier: str | None = None) -> list[dict]: + """Return deployments grouped by flow name and version list. + + Returns: + list[dict]: Flow names with their available version names. + """ + async with get_client() as client: + deployments = await client.read_deployments( + sort=DeploymentSort.CREATED_DESC, + ) + result = {} + for deployment in deployments: + if (not identifier or identifier in deployment.name) and ":" in deployment.name: + name_part, version_part = deployment.name.split(":", 1) + result.setdefault(name_part, []).append(version_part) + + return [{"name": k, "version_names": v} for k, v in result.items()] + + +def from_prefect_state_type_to_job_status(prefect_state_type: StateType) -> JobStatus: + """Map a Prefect state type to the corresponding job status. + + Returns: + JobStatus: Mapped job status. + + Raises: + ValueError: If the Prefect state type is unexpected. + """ + if prefect_state_type in [StateType.PENDING]: + return JobStatus.ENQUEUED + elif prefect_state_type in [StateType.SCHEDULED, StateType.RUNNING, StateType.PAUSED]: + return JobStatus.RUNNING + elif prefect_state_type in [StateType.COMPLETED]: + return JobStatus.SUCCEEDED + elif prefect_state_type in [StateType.FAILED, StateType.CRASHED]: + return JobStatus.ERROR + elif prefect_state_type in [StateType.CANCELLED, StateType.CANCELLING]: + return JobStatus.CANCELLED + else: + raise ValueError(f"Unexpected prefect StateType: '{prefect_state_type}'.") + + +async def get_flow_run_status_and_results( + flow_run_id: str | UUID, +) -> tuple[str, StateType, dict[str, Any], dict[str, str], dict[str, dict[str, Any]], str]: + """Fetch detailed information for a specific Prefect flow run by ID. + + Returns: + tuple[str, StateType, dict[str, Any], dict[str, str], dict[str, dict[str, Any]], str]: + Flow run name, state, parameters, tags, artifacts, and logs. + + Raises: + RuntimeError: If the flow run has no state. + """ + flow_run_uuid = flow_run_id if isinstance(flow_run_id, UUID) else UUID(str(flow_run_id)) + async with get_client() as client: + flow_run = await client.read_flow_run(flow_run_uuid) + if flow_run.state is None: + raise RuntimeError(f"Prefect flow run '{flow_run_uuid}' does not have a state") + + run_state_type = flow_run.state.type + + artifacts = await client.read_artifacts( + artifact_filter=ArtifactFilter( + flow_run_id=ArtifactFilterFlowRunId(any_=[flow_run_uuid]), + # key=ArtifactFilterKey(any_=["metadata"]), + ) + ) + + tags_by_key: dict[str, str] = {} + for tag in flow_run.tags or []: + if ":" in tag: + tag_key, tag_value = tag.split(":", 1) + tags_by_key[tag_key] = tag_value + else: + tags_by_key[tag] = "" + + artifacts_by_key: dict[str, dict[str, Any]] = {} + for artifact in artifacts: + artifact_key = artifact.key or str(artifact.id) + artifact_data = await _resolve_artifact_data(artifact.data) + + artifacts_by_key[artifact_key] = { + "data": artifact_data, + "description": artifact.description, + } + + log_filter = LogFilter(flow_run_id={"any_": [flow_run.id]}) + logs = await client.read_logs(log_filter=log_filter) + log_lines: list[str] = [] + for log in logs: + level_value = getattr(log, "level", None) + if isinstance(level_value, int): + level_name = logging.getLevelName(level_value) + elif isinstance(level_value, str): + level_name = level_value + else: + level_name = "unknown" + + log_lines.append( + f"{getattr(log, 'timestamp', '')} " + f"[{str(level_name).lower()}]: " + f"{getattr(log, 'message', '')} " + f"[{getattr(log, 'name', 'unknown')}]" + ) + log_str = "\n".join(log_lines) + + logging.debug(f"Prefect run found with id '{flow_run.id}' in state '{run_state_type}'.") + + return flow_run.name, run_state_type, flow_run.parameters, tags_by_key, artifacts_by_key, log_str + + +_MINIO_MARKDOWN_LINK_RE = re.compile(r"^\[(?P