diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..490aa35 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,72 @@ +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.13" + - 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 (py${{ matrix.python-version }}) + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + python-version: ["3.11", "3.12", "3.13"] + steps: + - uses: actions/checkout@v4 + - name: Set up Python + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python-version }} + - 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.13" + - 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..6e10092 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -1,37 +1,36 @@ -name: PyPi release +name: Publish Python Package run-name: Releasing next version 🚀 + 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: | - ./ci/linux/create_venv.sh - ./ci/linux/install_dependencies.sh + python-version: "3.11" + cache: "pip" + + - name: Install uv + uses: astral-sh/setup-uv@v2 + + - name: Install build dependencies + run: uv sync --locked --group dev - - name: build - run: | - ./ci/linux/build_python_package.sh + - name: Build package + run: uv run python -m build - - 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/extensions.json b/.vscode/extensions.json deleted file mode 100644 index 742cb59..0000000 --- a/.vscode/extensions.json +++ /dev/null @@ -1,10 +0,0 @@ -{ - "recommendations": [ - "ms-python.python", - "ms-python.black-formatter", - "ms-python.flake8", - "ms-python.isort", - "matangover.mypy", - "lextudio.restructuredtext", - ] -} \ No newline at end of file diff --git a/.vscode/launch.json b/.vscode/launch.json new file mode 100644 index 0000000..e60a2f5 --- /dev/null +++ b/.vscode/launch.json @@ -0,0 +1,20 @@ +{ + "version": "0.2.0", + "configurations": [ + { + "name": "pytest", + "type": "debugpy", + "request": "launch", + "module": "pytest", + "args": [ + "tests", + "-q" + ], + "env": { + "PYTHONPATH": "${workspaceFolder}/src" + }, + "console": "integratedTerminal", + "justMyCode": false + } + ] +} \ 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/README.md b/README.md index 058a065..f363f5e 100644 --- a/README.md +++ b/README.md @@ -1,29 +1,38 @@ # OMOTES SDK Python -This repository is part of the 'Nieuwe Warmte Nu Design Toolkit' project. +This repository is part of the 'Nieuwe Warmte Nu Design Toolkit' project. Python implementation of the OMOTES SDK through jobs which may be submitted, receive status updates for submitted jobs or delete submitted jobs. -## Protobuf -Please install `protoc` on your machine and make sure it is available in your `PATH`. -Version 25.2 is used: https://github.com/protocolbuffers/protobuf/releases/tag/v25.2. +## Development -Installation example - Linux +### Tools -1. Manually download the release zip file corresponding to your operating system and computer architecture (`protoc---.zip`) from github, or fetch the file using the command below. +This project uses: - ``` - wget https://github.com/protocolbuffers/protobuf/releases/download/v25.2/protoc-25.2-linux-x86_64.zip - ``` +- **uv**: Fast Python package manager and resolver. Install via [https://docs.astral.sh/uv/](https://docs.astral.sh/uv/) +- **just**: Command runner for common tasks (similar to Make). Install via [https://github.com/casey/just](https://github.com/casey/just) -2. Unzip the file under `$HOME/.local` or a directory of your choice +### Setup - ``` - unzip protoc-25.2-linux-x86_64.zip -d $HOME/.local - ``` +Install dependencies: -3. Update your environment `PATH` variable to include the path to the protoc executable. +```bash +uv sync +``` - ``` - export PATH="$PATH:$HOME/.local/bin" - ``` \ No newline at end of file +### Lint/typecheck/test locally + +Run via just (also used in in github actions): + +```bash +just ci # run all CI checks (lint, format-check, typecheck, test) + +just lint # ruff checks +just format # ruff format +just format-check # verify formatting +just typecheck # ty type checking +just test # pytest +``` + +To debug test go to the debug view in vscode and run "pytest". 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/src/omotes_sdk/__init__.py b/src/omotes_sdk/__init__.py index a7cdab2..a6a93e6 100644 --- a/src/omotes_sdk/__init__.py +++ b/src/omotes_sdk/__init__.py @@ -1,22 +1,3 @@ -# 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.prefect_util import MemoryLimit, as_memory_limit -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") +__all__ = ["MemoryLimit", "as_memory_limit"] 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..c91f8b7 --- /dev/null +++ b/src/omotes_sdk/prefect_util.py @@ -0,0 +1,805 @@ +import json +import logging +import os +import re +from collections.abc import Callable +from typing import Any, NoReturn, 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.objects import FlowRun +from prefect.client.schemas.sorting import DeploymentSort +from prefect.context import get_run_context +from prefect.exceptions import ObjectNotFound, PrefectHTTPStatusError +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 + +MEMORY_LIMIT_RE = re.compile(r"(?i)(\d+(?:\.\d+)?)([kmgtpe]i?|)") + + +class MemoryLimit(str): + """Validated Kubernetes-style memory quantity. + + Examples: + MemoryLimit("512Mi") + MemoryLimit("2Gi") + MemoryLimit("750M") + + """ + + def __new__(cls, value: str) -> "MemoryLimit": + """Create a validated MemoryLimit value. + + Returns: + MemoryLimit: Normalized and validated memory quantity. + + Raises: + ValueError: If the value is not a supported memory quantity. + + """ + normalized = value.strip() + if not MEMORY_LIMIT_RE.fullmatch(normalized): + raise ValueError("Unsupported memory quantity format. Examples: '512Mi', '2Gi', '750M', '1000000'.") + return cast("MemoryLimit", str.__new__(cls, normalized)) + + +def as_memory_limit(value: str) -> MemoryLimit: + """Convert a plain string to a validated MemoryLimit value. + + Returns: + MemoryLimit: Normalized and validated memory quantity. + + """ + return MemoryLimit(value) + + +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 _raise_prefect_api_error( + exc: PrefectHTTPStatusError | httpx.RequestError, +) -> NoReturn: + """Raise a user-facing error for Prefect API authentication/connectivity failures. + + Raises: + RuntimeError: Describes unauthorized, upstream, or unavailable Prefect API errors. + + """ + prefect_api_url = os.getenv("PREFECT_API_URL", "") + + if isinstance(exc, PrefectHTTPStatusError): + if exc.response.status_code == 401: + raise RuntimeError( + "Unauthorized: Invalid or missing authentication for Prefect server at " + f"{prefect_api_url}. Check PREFECT_API_AUTH_STRING setting." + ) from exc + + raise RuntimeError( + f"Prefect server error (HTTP {exc.response.status_code}): {exc.response.reason_phrase}" + ) from exc + + raise RuntimeError( + f"Prefect server is unavailable at {prefect_api_url}. " + "Start Prefect server or update PREFECT_API_URL, then try again." + ) from exc + + +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: MemoryLimit | 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 = str(memory_limit).strip() + match = MEMORY_LIMIT_RE.fullmatch(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 _to_docker_mem_limit(value: MemoryLimit | str | int | float) -> str: + """Convert memory values into Docker-compatible mem_limit string. + + Docker expects values parseable by docker.utils.parse_bytes, e.g. "1024m" + or "8589934592b". We emit bytes with a trailing "b" for exact values. + + Returns: + str: Docker-compatible memory limit string. + + """ + if isinstance(value, (int, float)): + return f"{int(value)}b" + + normalized = str(value).strip() + + # Already a Docker parse_bytes value (supports b, k, m, g). + lower = normalized.lower() + docker_match = re.fullmatch(r"\d+(?:\.\d+)?[bkmg]", lower) + if docker_match: + return lower + + try: + as_memory = MemoryLimit(normalized) + except ValueError: + # Preserve caller-supplied Docker-style values (e.g. "8g", "1024m"). + return normalized + + # Preserve decimal SI units directly when Docker can parse them. + quantity_match = MEMORY_LIMIT_RE.fullmatch(str(as_memory).strip()) + if quantity_match: + number = quantity_match.group(1) + suffix = quantity_match.group(2).lower() + if suffix in {"k", "m", "g"}: + return f"{number}{suffix}" + + return f"{_memory_quantity_to_bytes(as_memory)}b" + + +def build_universal_job_vars( + memory_limit: MemoryLimit | 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 require mem_limit in Docker parse_bytes format. We emit a + byte-string value (e.g. "8589934592b") so it also satisfies schemas that + require a string. + + Returns: + dict | None: Job variables suitable for worker deployment or None if unchanged. + + """ + job_vars = dict(base_vars or {}) + + # Normalize provided mem_limit to Docker-compatible string. + if "mem_limit" in job_vars and job_vars["mem_limit"] is not None: + job_vars["mem_limit"] = _to_docker_mem_limit(job_vars["mem_limit"]) + + if memory_limit: + normalized_memory_limit = MemoryLimit(str(memory_limit)) + job_vars["memory_limit"] = str(normalized_memory_limit) + job_vars["memory_request"] = str(normalized_memory_limit) + job_vars["mem_limit"] = _to_docker_mem_limit(normalized_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, + max_concurrent_runs: int | None = None, +) -> 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. + max_concurrent_runs: Maximum number of concurrent flow runs for this deployment. + + Raises: + RuntimeError: If a semantic-versioned deployment already exists. + ValueError: If max_concurrent_runs is provided and is less than 1. + + """ + version_name = image_name.rsplit(":", 1)[-1] + if max_concurrent_runs is not None and max_concurrent_runs < 1: + raise ValueError("max_concurrent_runs must be at least 1") + + try: + 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}'") + except (PrefectHTTPStatusError, httpx.RequestError) as exc: + _raise_prefect_api_error(exc) + + 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, + concurrency_limit=max_concurrent_runs, + ) + + +# 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 {} + + try: + async with get_client() as client: + deployments = await client.read_deployments( + sort=DeploymentSort.CREATED_DESC, + ) + except (PrefectHTTPStatusError, httpx.RequestError) as exc: + _raise_prefect_api_error(exc) + + 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 "") + + +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)) + try: + 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, + ) + except (PrefectHTTPStatusError, httpx.RequestError) as exc: + _raise_prefect_api_error(exc) + + +_MINIO_MARKDOWN_LINK_RE = re.compile(r"^\[(?P