diff --git a/.github/actions/setup-python/action.yml b/.github/actions/setup-python/action.yml new file mode 100644 index 000000000..d412562bf --- /dev/null +++ b/.github/actions/setup-python/action.yml @@ -0,0 +1,45 @@ +name: Setup Python Environment +description: Setup Python, reinstall pip, and install dependencies +inputs: + python-version: + required: true + +runs: + using: "composite" + steps: + - name: Cache pip dependencies + uses: actions/cache@v4 + with: + path: ~/.cache/pip + key: ${{ hashFiles('requirements/release.txt') }}-${{ hashFiles('requirements/dev.txt') }} + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: ${{ inputs.python-version }} + + - name: Reinstall pip + shell: bash + run: | + PY_VER=$(python -c "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}')") + PY_MAJOR=$(echo $PY_VER | cut -d. -f1) + PY_MINOR=$(echo $PY_VER | cut -d. -f2) + + if [ "$PY_MAJOR" -eq 3 ] && [ "$PY_MINOR" -le 8 ]; then + URL="https://bootstrap.pypa.io/pip/${PY_VER}/get-pip.py" + else + URL="https://bootstrap.pypa.io/get-pip.py" + fi + + curl -sS "$URL" -o /tmp/get-pip.py + python /tmp/get-pip.py --force-reinstall + + pip --version + pip3 --version + + - name: Install dependencies + shell: bash + run: | + python3 -m pip install --upgrade pip + if [ -f requirements.txt ]; then pip install -r requirements.txt; fi + pip3 install -r requirements/dev.txt diff --git a/.github/actions/setup-singularity/action.yml b/.github/actions/setup-singularity/action.yml new file mode 100644 index 000000000..14e376e7d --- /dev/null +++ b/.github/actions/setup-singularity/action.yml @@ -0,0 +1,45 @@ +name: Install Singularity +description: Install Go and Singularity +inputs: + go-version: + required: true + singularity-version: + required: true + os: + default: linux + arch: + default: amd64 + +runs: + using: "composite" + steps: + - name: Install system dependencies + shell: bash + run: | + sudo apt-get update && sudo apt-get install -y \ + build-essential \ + libssl-dev \ + uuid-dev \ + libgpgme11-dev \ + squashfs-tools \ + libseccomp-dev \ + pkg-config + + - name: Download and install Go + shell: bash + run: | + wget https://go.dev/dl/go${{ inputs.go-version }}.${{ inputs.os }}-${{ inputs.arch }}.tar.gz + sudo tar -C /usr/local -xzf go${{ inputs.go-version }}.${{ inputs.os }}-${{ inputs.arch }}.tar.gz + rm go${{ inputs.go-version }}.${{ inputs.os }}-${{ inputs.arch }}.tar.gz + echo "/usr/local/go/bin" >> $GITHUB_PATH + + - name: Download and install Singularity + shell: bash + run: | + wget https://github.com/sylabs/singularity/releases/download/v${{ inputs.singularity-version }}/singularity-ce-${{ inputs.singularity-version }}.tar.gz + tar -xzf singularity-ce-${{ inputs.singularity-version }}.tar.gz + cd singularity-ce-${{ inputs.singularity-version }} + ./mconfig + make -C ./builddir + sudo make -C ./builddir install + diff --git a/.github/workflows/push-pr_workflow.yml b/.github/workflows/push-pr_workflow.yml index 1d3c9d958..1e49543c5 100644 --- a/.github/workflows/push-pr_workflow.yml +++ b/.github/workflows/push-pr_workflow.yml @@ -9,29 +9,28 @@ jobs: if: github.event_name == 'pull_request' steps: - - name: Checkout code - uses: actions/checkout@v2 - with: - fetch-depth: 0 # Checkout the whole history, in case the target is way far behind - - - name: Check if target branch has been merged - run: | - if git merge-base --is-ancestor ${{ github.event.pull_request.base.sha }} ${{ github.sha }}; then - echo "Target branch has been merged into the source branch." - else - echo "Target branch has not been merged into the source branch. Please merge in target first." - exit 1 - fi - - - name: Check that CHANGELOG has been updated - run: | - # If this step fails, this means you haven't updated the CHANGELOG.md file with notes on your contribution. - if git diff --name-only ${{ github.event.pull_request.base.sha }} ${{ github.sha }} | grep -q '^CHANGELOG.md$'; then - echo "Thanks for helping keep our CHANGELOG up-to-date!" - else - echo "Please update the CHANGELOG.md file with notes on your contribution." - exit 1 - fi + - uses: actions/checkout@v4 + with: + fetch-depth: 0 # Checkout the whole history, in case the target is way far behind + + - name: Check if target branch has been merged + run: | + if git merge-base --is-ancestor ${{ github.event.pull_request.base.sha }} ${{ github.sha }}; then + echo "Target branch has been merged into the source branch." + else + echo "Target branch has not been merged into the source branch. Please merge in target first." + exit 1 + fi + + - name: Check that CHANGELOG has been updated + run: | + # If this step fails, this means you haven't updated the CHANGELOG.md file with notes on your contribution. + if git diff --name-only ${{ github.event.pull_request.base.sha }} ${{ github.sha }} | grep -q '^CHANGELOG.md$'; then + echo "Thanks for helping keep our CHANGELOG up-to-date!" + else + echo "Please update the CHANGELOG.md file with notes on your contribution." + exit 1 + fi Lint: runs-on: ubuntu-latest @@ -40,51 +39,74 @@ jobs: MAX_COMPLEXITY: 15 steps: - - uses: actions/checkout@v2 - - name: Set up Python - uses: actions/setup-python@v2 - with: - python-version: '3.x' - - - name: Check cache - uses: actions/cache@v2 - with: - path: ~/.cache/pip - key: ${{ hashFiles('requirements/release.txt') }}-${{ hashFiles('requirements/dev.txt') }} - - - name: Install dependencies - run: | - python3 -m pip install --upgrade pip - if [ -f requirements.txt ]; then pip install --upgrade -r requirements.txt; fi - pip3 install --upgrade -r requirements/dev.txt - - - name: Lint with flake8 - run: | - # stop the build if there are Python syntax errors or undefined names - flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics - # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide - flake8 . --count --max-complexity=$MAX_COMPLEXITY --statistics --max-line-length=$MAX_LINE_LENGTH - - - name: Lint with isort - run: | - python3 -m isort --check --line-length $MAX_LINE_LENGTH merlin - python3 -m isort --check --line-length $MAX_LINE_LENGTH tests - python3 -m isort --check --line-length $MAX_LINE_LENGTH *.py - - - name: Lint with Black - run: | - python3 -m black --check --line-length $MAX_LINE_LENGTH --target-version py38 merlin - python3 -m black --check --line-length $MAX_LINE_LENGTH --target-version py38 tests - python3 -m black --check --line-length $MAX_LINE_LENGTH --target-version py38 *.py - - - name: Lint with PyLint - run: | - python3 -m pylint merlin --rcfile=setup.cfg --exit-zero - python3 -m pylint tests --rcfile=setup.cfg --exit-zero + - uses: actions/checkout@v4 + + - uses: ./.github/actions/setup-python + with: + python-version: '3.x' + + - name: Lint with flake8 + run: | + # stop the build if there are Python syntax errors or undefined names + flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics + # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide + flake8 . --count --max-complexity=$MAX_COMPLEXITY --statistics --max-line-length=$MAX_LINE_LENGTH + + - name: Lint with isort + run: | + isort --check --line-length $MAX_LINE_LENGTH merlin tests *.py + + - name: Lint with Black + run: | + black --check --line-length $MAX_LINE_LENGTH --target-version py38 merlin tests *.py + + - name: Lint with PyLint + run: | + pylint merlin tests --rcfile=setup.cfg --exit-zero Local-test-suite: runs-on: ubuntu-latest - env: + env: + GO_VERSION: 1.18.1 + SINGULARITY_VERSION: 3.9.9 + OS: linux + ARCH: amd64 + + strategy: + matrix: + python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13'] + + steps: + - uses: actions/checkout@v4 + + - uses: ./.github/actions/setup-python + with: + python-version: ${{ matrix.python-version }} + + - uses: ./.github/actions/setup-singularity + with: + go-version: ${{ env.GO_VERSION }} + singularity-version: ${{ env.SINGULARITY_VERSION }} + os: ${{ env.OS }} + arch: ${{ env.ARCH }} + + - name: Install merlin and setup + run: | + pip3 install -e . + merlin config + + - name: Install CLI task dependencies from 'feature_demo' workflow + run: | + merlin example feature_demo + pip3 install -r feature_demo/requirements.txt + + - name: Run integration test suite for local tests + run: | + python3 tests/integration/run_tests.py --verbose --local + + Unit-tests: + runs-on: ubuntu-latest + env: GO_VERSION: 1.18.1 SINGULARITY_VERSION: 3.9.9 OS: linux @@ -92,138 +114,75 @@ jobs: strategy: matrix: - python-version: ['3.7', '3.8', '3.9', '3.10', '3.11'] + python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13'] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - - name: Check cache - uses: actions/cache@v2 - with: - path: ${{ env.pythonLocation }} - key: ${{ env.pythonLocation }}-${{ hashFiles('requirements/release.txt') }}-${{ hashFiles('requirements/dev.txt') }} - - - name: Install dependencies - run: | - python3 -m pip install --upgrade pip - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - pip3 install -r requirements/dev.txt - pip freeze - - - name: Install singularity - run: | - sudo apt-get update && sudo apt-get install -y \ - build-essential \ - libssl-dev \ - uuid-dev \ - libgpgme11-dev \ - squashfs-tools \ - libseccomp-dev \ - pkg-config - wget https://go.dev/dl/go$GO_VERSION.$OS-$ARCH.tar.gz - sudo tar -C /usr/local -xzf go$GO_VERSION.$OS-$ARCH.tar.gz - rm go$GO_VERSION.$OS-$ARCH.tar.gz - export PATH=$PATH:/usr/local/go/bin - wget https://github.com/sylabs/singularity/releases/download/v$SINGULARITY_VERSION/singularity-ce-$SINGULARITY_VERSION.tar.gz - tar -xzf singularity-ce-$SINGULARITY_VERSION.tar.gz - cd singularity-ce-$SINGULARITY_VERSION - ./mconfig && \ - make -C ./builddir && \ - sudo make -C ./builddir install - - - name: Install merlin to run unit tests - run: | - pip3 install -e . - merlin config - - - name: Install CLI task dependencies generated from the 'feature demo' workflow - run: | - merlin example feature_demo - pip3 install -r feature_demo/requirements.txt - - - name: Run pytest over unit test suite - run: | - python3 -m pytest -v --order-scope=module tests/unit/ - - - name: Run integration test suite for local tests - run: | - python3 tests/integration/run_tests.py --verbose --local - - Distributed-test-suite: + - uses: actions/checkout@v4 + + - uses: ./.github/actions/setup-python + with: + python-version: ${{ matrix.python-version }} + + - uses: ./.github/actions/setup-singularity + with: + go-version: ${{ env.GO_VERSION }} + singularity-version: ${{ env.SINGULARITY_VERSION }} + os: ${{ env.OS }} + arch: ${{ env.ARCH }} + + - name: Install merlin and setup + run: | + pip3 install -e . + merlin config + + - name: Install CLI task dependencies from 'feature_demo' workflow + run: | + merlin example feature_demo + pip3 install -r feature_demo/requirements.txt + + - name: Run pytest over unit test suite + run: | + python3 -m pytest -v --order-scope=module tests/unit/ + + Integration-tests: runs-on: ubuntu-latest - services: - # rabbitmq: - # image: rabbitmq:latest - # ports: - # - 5672:5672 - # options: --health-cmd "rabbitmqctl node_health_check" --health-interval 10s --health-timeout 5s --health-retries 5 - # Label used to access the service container - redis: - # Docker Hub image - image: redis - # Set health checks to wait until redis has started - options: >- - --health-cmd "redis-cli ping" - --health-interval 10s - --health-timeout 5s - --health-retries 5 - ports: - - 6379:6379 + env: + GO_VERSION: 1.18.1 + SINGULARITY_VERSION: 3.9.9 + OS: linux + ARCH: amd64 strategy: matrix: - python-version: ['3.7', '3.8', '3.9', '3.10', '3.11'] + python-version: ['3.8', '3.9', '3.10', '3.11', '3.12', '3.13'] steps: - - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} - uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - - name: Check cache - uses: actions/cache@v2 - with: - path: ${{ env.pythonLocation }} - key: ${{ env.pythonLocation }}-${{ hashFiles('requirements/release.txt') }}-${{ hashFiles('requirements/dev.txt') }} - - - name: Install dependencies - run: | - python3 -m pip install --upgrade pip - if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - pip3 install -r requirements/dev.txt - - - name: Install merlin and setup redis as the broker - run: | - pip3 install -e . - merlin config --broker redis - - - name: Install CLI task dependencies generated from the 'feature demo' workflow - run: | - merlin example feature_demo - pip3 install -r feature_demo/requirements.txt - - - name: Run integration test suite for distributed tests - env: - REDIS_HOST: redis - REDIS_PORT: 6379 - run: | - python3 tests/integration/run_tests.py --verbose --distributed - - # - name: Setup rabbitmq config - # run: | - # merlin config --test rabbitmq - - # - name: Run integration test suite for rabbitmq - # env: - # AMQP_URL: amqp://localhost:${{ job.services.rabbitmq.ports[5672] }} - # RABBITMQ_USER: Jimmy_Space - # RABBITMQ_PASS: Alexander_Rules - # ports: - # - ${{ job.services.rabbitmq.ports['5672'] }} - # run: | - # python3 tests/integration/run_tests.py --verbose --ids 31 32 + - uses: actions/checkout@v4 + + - uses: ./.github/actions/setup-python + with: + python-version: ${{ matrix.python-version }} + + - uses: ./.github/actions/setup-singularity + with: + go-version: ${{ env.GO_VERSION }} + singularity-version: ${{ env.SINGULARITY_VERSION }} + os: ${{ env.OS }} + arch: ${{ env.ARCH }} + + - name: Install merlin + run: | + pip --version + pip3 --version + pip3 install -e . + merlin config + + - name: Install CLI task dependencies from 'feature_demo' workflow + run: | + merlin example feature_demo + pip3 install -r feature_demo/requirements.txt + + # TODO remove the --ignore statement once those tests are fixed + - name: Run integration test suite for distributed tests + run: | + python3 -m pytest -v --ignore tests/integration/test_celeryadapter.py tests/integration/ diff --git a/.gitignore b/.gitignore index c22521934..cec577a85 100644 --- a/.gitignore +++ b/.gitignore @@ -39,8 +39,9 @@ flux.out slurm*.out docs/build/ -# Tox files +# Test files .tox/* +.coverage # Jupyter jupyter/.ipynb_checkpoints diff --git a/.readthedocs.yaml b/.readthedocs.yaml index c3cfbbe07..7c5ed2c04 100644 --- a/.readthedocs.yaml +++ b/.readthedocs.yaml @@ -3,7 +3,7 @@ version: 2 build: os: "ubuntu-20.04" tools: - python: "3.8" + python: "3.13" python: install: @@ -11,5 +11,6 @@ python: mkdocs: fail_on_warning: false + configuration: mkdocs.yml formats: [pdf] \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 21b4427b1..a73b05ab7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,28 @@ All notable changes to Merlin will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). -## [1.12.2b1] +## [Unreleased] +### Added +- API documentation for Merlin's core codebase +- Added support for Python 3.12 and 3.13 +- Added additional tests for the `merlin run` and `merlin purge` commands +- Aliased types to represent different types of pytest fixtures +- New test condition `StepFinishedFilesCount` to help search for `MERLIN_FINISHED` files in output workspaces +- Added "Unit-tests" GitHub action to run the unit test suite +- Added `CeleryTaskManager` context manager to the test suite to ensure tasks are safely purged from queues if tests fail +- Added `command-tests`, `workflow-tests`, and `integration-tests` to the Makefile +- Python 3.8 now requires `orderly-set==5.3.0` to avoid a bug with the deepdiff library +- New GitHub actions to reduce common code in CI + +### Changed +- Dropped support for Python 3.7 +- Ported all distributed tests of the integration test suite to pytest + - There is now a `commands/` directory and a `workflows/` directory under the integration suite to house these tests + - Removed the "Distributed-tests" GitHub action as these tests will now be run under "Integration-tests" +- Removed `e2e-distributed*` definitions from the Makefile +- CI to use new actions + +## [1.12.2] ### Added - Conflict handler option to the `dict_deep_merge` function in `utils.py` - Ability to add module-specific pytest fixtures @@ -19,6 +40,16 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - New step return code `$(MERLIN_RAISE_ERROR)` to force an error to be raised by a task (mainly for testing) - Added description of this to docs - New test to ensure a single failed task won't break a workflow +- Several new unit tests for the following subdirectories: + - `merlin/common/` + - `merlin/config/` + - `merlin/examples/` + - `merlin/server/` +- Context managers for the `conftest.py` file to ensure safe spin up and shutdown of fixtures + - `RedisServerManager`: context to help with starting/stopping a redis server for tests + - `CeleryWorkersManager`: context to help with starting/stopping workers for tests +- Ability to copy and print the `Config` object from `merlin/config/__init__.py` +- Equality method to the `ContainerFormatConfig` and `ContainerConfig` objects from `merlin/server/server_util.py` ### Changed - `merlin info` is cleaner and gives python package info @@ -28,6 +59,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - Added fixtures for `merlin status` tests that copy the workspace to a temporary directory so you can see exactly what's run in a test - Batch block and workers now allow for variables to be used in node settings - Task id is now the path to the directory +- Split the `start_server` and `config_server` functions of `merlin/server/server_commands.py` into multiple functions to make testing easier +- Split the `create_server_config` function of `merlin/server/server_config.py` into two functions to make testing easier +- Combined `set_snapshot_seconds` and `set_snapshot_changes` methods of `RedisConfig` into one method `set_snapshot` ### Fixed - Bugfix for output of `merlin example openfoam_wf_singularity` @@ -95,8 +129,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - this required adding a decent amount of test files to help with the tests; these can be found under the tests/unit/study/status_test_files directory - Pytest fixtures in the `conftest.py` file of the integration test suite - NOTE: an export command `export LC_ALL='C'` had to be added to fix a bug in the WEAVE CI. This can be removed when we resolve this issue for the `merlin server` command -- Tests for the `celeryadapter.py` module -- New CeleryTestWorkersManager context to help with starting/stopping workers for tests +- Coverage to the test suite. This includes adding tests for: + - `merlin/common/` + - `merlin/config/` + - `merlin/examples/` + - `celeryadapter.py` +- Context managers for the `conftest.py` file to ensure safe spin up and shutdown of fixtures + - `RedisServerManager`: context to help with starting/stopping a redis server for tests + - `CeleryWorkersManager`: context to help with starting/stopping workers for tests +- Ability to copy and print the `Config` object from `merlin/config/__init__.py` ### Changed - Reformatted the entire `merlin status` command @@ -132,7 +173,6 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 - The `merlin monitor` command will now keep an allocation up if the queues are empty and workers are still processing tasks - Add the restart keyword to the specification docs - Cyclical imports and config imports that could easily cause ci issues - ## [1.11.1] ### Fixed - Typo in `batch.py` that caused lsf launches to fail (`ALL_SGPUS` changed to `ALL_GPUS`) diff --git a/Makefile b/Makefile index 08fb7d0f8..f04f6ba40 100644 --- a/Makefile +++ b/Makefile @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -34,12 +34,13 @@ include config.mk .PHONY : install-workflow-deps .PHONY : install-dev .PHONY : unit-tests +.PHONY : command-tests +.PHONY : workflow-tests +.PHONY : integration-tests .PHONY : e2e-tests .PHONY : e2e-tests-diagnostic .PHONY : e2e-tests-local .PHONY : e2e-tests-local-diagnostic -.PHONY : e2e-tests-distributed -.PHONY : e2e-tests-distributed-diagnostic .PHONY : tests .PHONY : check-flake8 .PHONY : check-black @@ -89,6 +90,18 @@ unit-tests: . $(VENV)/bin/activate; \ $(PYTHON) -m pytest -v --order-scope=module $(UNIT); \ +command-tests: + . $(VENV)/bin/activate; \ + $(PYTHON) -m pytest -v $(TEST)/integration/commands/; \ + + +workflow-tests: + . $(VENV)/bin/activate; \ + $(PYTHON) -m pytest -v $(TEST)/integration/workflows/; \ + + +integration-tests: command-tests workflow-tests + # run CLI tests - these require an active install of merlin in a venv e2e-tests: @@ -111,18 +124,8 @@ e2e-tests-local-diagnostic: $(PYTHON) $(TEST)/integration/run_tests.py --local --verbose -e2e-tests-distributed: - . $(VENV)/bin/activate; \ - $(PYTHON) $(TEST)/integration/run_tests.py --distributed; \ - - -e2e-tests-distributed-diagnostic: - . $(VENV)/bin/activate; \ - $(PYTHON) $(TEST)/integration/run_tests.py --distributed --verbose - - # run unit and CLI tests -tests: unit-tests e2e-tests +tests: unit-tests integration-tests e2e-tests check-flake8: @@ -135,9 +138,9 @@ check-flake8: check-black: . $(VENV)/bin/activate; \ - $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version py38 $(MRLN); \ - $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version py38 $(TEST); \ - $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version py38 *.py; \ + $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version $(PY_TARGET_VER) $(MRLN); \ + $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version $(PY_TARGET_VER) $(TEST); \ + $(PYTHON) -m black --check --line-length $(MAX_LINE_LENGTH) --target-version $(PY_TARGET_VER) *.py; \ check-isort: @@ -179,9 +182,9 @@ fix-style: $(PYTHON) -m isort -w $(MAX_LINE_LENGTH) $(MRLN); \ $(PYTHON) -m isort -w $(MAX_LINE_LENGTH) $(TEST); \ $(PYTHON) -m isort -w $(MAX_LINE_LENGTH) *.py; \ - $(PYTHON) -m black --target-version py38 -l $(MAX_LINE_LENGTH) $(MRLN); \ - $(PYTHON) -m black --target-version py38 -l $(MAX_LINE_LENGTH) $(TEST); \ - $(PYTHON) -m black --target-version py38 -l $(MAX_LINE_LENGTH) *.py; \ + $(PYTHON) -m black --target-version $(PY_TARGET_VER) -l $(MAX_LINE_LENGTH) $(MRLN); \ + $(PYTHON) -m black --target-version $(PY_TARGET_VER) -l $(MAX_LINE_LENGTH) $(TEST); \ + $(PYTHON) -m black --target-version $(PY_TARGET_VER) -l $(MAX_LINE_LENGTH) *.py; \ # Increment the Merlin version. USE ONLY ON DEVELOP BEFORE MERGING TO MASTER. diff --git a/README.md b/README.md index e47f7744b..d2b179e09 100644 --- a/README.md +++ b/README.md @@ -86,7 +86,7 @@ Need help? ## Quick Start -Note: Merlin supports Python 3.6+. +Note: Merlin supports Python 3.8+. To install Merlin and its dependencies, run: diff --git a/config.mk b/config.mk index f1cfbcea3..b7c2ea3d8 100644 --- a/config.mk +++ b/config.mk @@ -1,4 +1,5 @@ PYTHON?=python3 +PY_TARGET_VER?=py311 PYV=$(shell $(PYTHON) -c "import sys;t='{v[0]}_{v[1]}'.format(v=list(sys.version_info[:2]));sys.stdout.write(t)") PYVD=$(shell $(PYTHON) -c "import sys;t='{v[0]}.{v[1]}'.format(v=list(sys.version_info[:2]));sys.stdout.write(t)") VENV?=venv_merlin_py_$(PYV) @@ -18,7 +19,7 @@ else endif VER?=1.0.0 -VSTRING=[0-9]\+\.[0-9]\+\.[0-9]\+ +VSTRING=[0-9]\+\.[0-9]\+\.[0-9]\+\(b[0-9]\+\)\? YEAR=20[0-9][0-9] NEW_YEAR?=2023 CHANGELOG_VSTRING="## \[$(VSTRING)\]" diff --git a/docs/README.md b/docs/README.md index cbddc54a4..e427b7d0e 100644 --- a/docs/README.md +++ b/docs/README.md @@ -41,4 +41,84 @@ MkDocs relies on an `mkdocs.yml` file for almost everything to do with configura ## How Do API Docs Work? -Coming soon... +The API documentation in this project is automatically generated using a combination of MkDocs plugins and a custom Python script. This ensures that the documentation stays up-to-date with your codebase and provides a structured reference for all Python modules, classes, and functions in the merlin directory. + +This section will discuss: + +- [Code Reference Generation](#code-reference-generation) + - [How the Script Works](#how-the-script-works) +- [Viewing the API Docs](#viewing-the-api-docs) +- [Keeping API Docs Up-to-Date](#keeping-api-docs-up-to-date) +- [Example Docstring](#example-docstring) +- [Plugins Involved](#plugins-involved) + +### Code Reference Generation + +The script `docs/gen_ref_pages.py` is responsible for generating the API reference pages. It scans the `merlin` directory for Python files and creates Markdown files for each module. These Markdown files are then included in the `api_reference` section of the documentation. + +#### How the Script Works + +1. File Scanning: + + The script recursively scans all Python files in the `merlin` directory using `Path.rglob("*.py")`. + +2. Ignore Patterns: + + Certain files and directories are excluded from the API docs based on the `IGNORE_PATTERNS` list. For example: + + - `merlin/examples/workflows` + - `merlin/examples/dev_workflows` + - `merlin/data` + - Files like `ascii_art.py` + + The `should_ignore()` function checks each file against these patterns and skips them if they match. + +3. Markdown File Creation: + + For each valid Python file: + + - The script determines the module path (e.g., merlin.module_name) and the corresponding Markdown file path. + - Special cases like `__init__.py` are handled by renaming the generated file to index.md for better navigation. + - Files like `__main__.py` are ignored entirely. + + The script then writes the mkdocstrings syntax (::: module_name) into the Markdown file, which tells the mkdocstrings plugin to generate the documentation for that module. + +4. Navigation File: + + The script builds a navigation structure using the `mkdocs_gen_files.Nav` class. This structure is saved into a `SUMMARY.md` file, which is used by the `literate-nav` plugin to define the navigation for the API reference section. + +### Viewing the API Docs + +Once the script generates the Markdown files, they are included in the documentation site under the `api_reference` section. You can explore the API docs in the navigation bar under `Reference Guide` -> `API Reference`, with the navigation organized based on the module hierarchy. + +### Keeping API Docs Up-to-Date + +To ensure the API documentation remains accurate: + +Update the docstrings in Merlin's code whenever changes are made to functions, classes, or modules. The `docs/gen_ref_pages.py` file will run automatically when docs are created (i.e. when you run `mkdocs serve`). + +### Example Docstring + +The API documentation relies on properly formatted docstrings. Here’s an example using the Google-style docstring format: + +```python +def add_numbers(a: int, b: int) -> int: + """ + Adds two numbers. + + Args: + a (int): The first number. + b (int): The second number. + + Returns: + The sum of the two numbers. + """ +``` + +### Plugins Involved + +Several MkDocs plugins work together to generate and display the API documentation: + +- `mkdocstrings`: Parses Python docstrings and generates the actual API content. +- `mkdocs-gen-files`: Handles the creation of Markdown files and navigation structure. +- `literate-nav`: Uses the `SUMMARY.md` file to organize the API reference section in the documentation sidebar. diff --git a/docs/api_reference/index.md b/docs/api_reference/index.md index 6707457de..072a7b679 100644 --- a/docs/api_reference/index.md +++ b/docs/api_reference/index.md @@ -1,5 +1,3 @@ # Merlin API Reference -Coming soon! - - +Welcome to the Application Program Interface (API) Reference Guide for Merlin! This comprehensive guide is designed to provide developers with a detailed understanding of the various modules, classes, and functions available within the Merlin API. diff --git a/docs/faq.md b/docs/faq.md index 0b0397a45..a934df5c2 100644 --- a/docs/faq.md +++ b/docs/faq.md @@ -245,7 +245,7 @@ See the docs on all [Merlin Commands](./user_guide/command_line.md) that are ava ### How do I set up a workspace without executing step scripts? -Use [Merlin's Dry Run](./user_guide/command_line.md#dry-run) capability: +Use [Merlin's Dry Run](./user_guide/running_studies.md#dry-runs) capability: === "Locally" @@ -438,7 +438,7 @@ run: procs: 3 ``` -See [The `LAUNCHER` and `VLAUNCHER` Variables](./user_guide/variables.md#the-launcher-and-vlauncher-variables) and the [Scheduler Specific Properties](./user_guide/specification.md#scheduler-specific-properties) sections for more information. +See [The `LAUNCHER` and `VLAUNCHER` Variables](./user_guide/variables.md#the-launcher-and-vlauncher-variables) and [The Run Property](./user_guide/specification.md#the-run-property) sections for more information. ### What is `level_max_dirs`? diff --git a/docs/gen_ref_pages.py b/docs/gen_ref_pages.py index e46c0eb93..717c47d13 100644 --- a/docs/gen_ref_pages.py +++ b/docs/gen_ref_pages.py @@ -6,10 +6,27 @@ nav = mkdocs_gen_files.Nav() -# print(sorted(Path("merlin").rglob("*.py"))) +IGNORE_PATTERNS = [ + Path("merlin/examples/workflows"), + Path("merlin/examples/dev_workflows"), + Path("merlin/data"), + "*/ascii_art.py", +] + + +def should_ignore(path): + """Check if the given path matches any ignore patterns.""" + for pattern in IGNORE_PATTERNS: + # if Path(pattern).is_relative_to(path): + if path.is_relative_to(Path(pattern)): + return True + if path.match(pattern): + return True + return False + for path in sorted(Path("merlin").rglob("*.py")): - if "merlin/examples" in str(path): + if should_ignore(path): continue module_path = path.relative_to("merlin").with_suffix("") doc_path = path.relative_to("merlin").with_suffix(".md") diff --git a/docs/index.md b/docs/index.md index 46466494f..1abc4936a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -71,7 +71,7 @@ First, let's create a folder to store our server files and our examples. We'll a mkdir merlin_examples ; cd merlin_examples/ ``` -Now let's set up a [containerized server](./user_guide/configuration/merlin_server.md) that Merlin can connect to. +Now let's set up a [containerized server](./user_guide/configuration/containerized_server.md) that Merlin can connect to. 1. Initialize the server files: diff --git a/docs/user_guide/command_line.md b/docs/user_guide/command_line.md index 30faee9ce..096891a47 100644 --- a/docs/user_guide/command_line.md +++ b/docs/user_guide/command_line.md @@ -95,7 +95,7 @@ Create a local containerized server for Merlin to connect to. Merlin server crea Merlin server has a list of commands for interacting with the broker and results server. These commands allow the user to manage and monitor the exisiting server and create instances of servers if needed. -More information on configuring with Merlin server can be found at the [Merlin Server Configuration](./configuration/merlin_server.md) page. +More information on configuring with Merlin server can be found at the [Containerized Server Configuration](./configuration/containerized_server.md) page. **Usage:** diff --git a/docs/user_guide/configuration/containerized_server.md b/docs/user_guide/configuration/containerized_server.md index 218e97f23..93dea140f 100644 --- a/docs/user_guide/configuration/containerized_server.md +++ b/docs/user_guide/configuration/containerized_server.md @@ -616,6 +616,6 @@ You can check that everything ran properly with: merlin status hello_samples.yaml ``` -Or, if you're using a version of Merlin prior to v1.12.0, you can ensure that the `hello_samples_/` output workspace was created. More info on the expected output can be found in [the Hello World Examples page](../../examples/hello.md#expected-output-1). +Or, if you're using a version of Merlin prior to v1.12.0, you can ensure that the `hello_samples_/` output workspace was created. More info on the expected output can be found in [the Hello World Examples page](../../examples/hello.md#expected-output_1). Congratulations, you just ran a cross-node workflow with a containerized server! diff --git a/docs/user_guide/configuration/merlin_server.md b/docs/user_guide/configuration/merlin_server.md deleted file mode 100644 index 39cf1327d..000000000 --- a/docs/user_guide/configuration/merlin_server.md +++ /dev/null @@ -1,177 +0,0 @@ -# Merlin Server Configuration - -!!! warning - - It's recommended that you read through the [Configuration Overview](./index.md) page before proceeding with this module. - -The merlin server command allows users easy access to containerized broker and results servers for Merlin workflows. This allows users to run Merlin without a dedicated external server. - -The main configuration will be stored in the subdirectory called `server/` by default in the main Merlin configuration directory `~/.merlin`. However, different server images can be created for different use cases or studies by simplying creating a new directory to store local configuration files for Merlin server instances. - -This module will walk through how to initalize the server, start it, and ensure it's linked to Merlin. - -## Initializing the Server - -First create and navigate into a directory to store your local Merlin configuration for a specific use case or study: - -```bash -mkdir study1/ ; cd study1/ -``` - -Afterwards you can instantiate Merlin server in this directory by running: - -```bash -merlin server init -``` - -A main server configuration will be created in the `~/.merlin/server/` directory. This will have the following files: - -- docker.yaml -- merlin_server.yaml -- podman.yaml -- singularity.yaml - -The main configuration in `~/.merlin/server/` deals with defaults and technical commands that might be used for setting up the Merlin server local configuration and its containers. Each container has their own configuration file to allow users to be able to switch between different containerized services freely. - -In addition to the main server configuration, a local server configuration will be created in your current working directory in a folder called `merlin_server/`. This directory will contain: - -- `redis.conf`: The Redis configuration file that contains all of the settings to be used for our Redis server -- `redis.pass`: A password for the Redis server that we'll start up next -- `redis.users`: A file defining the users that are allowed to access the Redis server and their permissions -- `redis_latest.sif`: A singularity file that contains the latest Redis Docker image that was pulled behind the scenes by Merlin - -The local configuration `merlin_server/` folder contains configuration files specific to a certain use case or run. In the case above you can see that we have a Redis singularity container called `redis_latest.sif` with the Redis configuration file called `redis.conf`. This Redis configuration will allow the user to configure Redis to their specified needs without have to manage or edit the Redis container. When the server is run this configuration will be dynamically read, so settings can be changed between runs if needed. - -Once the Merlin server has been initialized in the local directory the user will be allowed to run other Merlin server commands such as `start`, `status`, and `stop` to interact with the Merlin server. A detailed list of commands can be found in the [Merlin Server](../command_line.md#server-merlin-server) section of the [Command Line](../command_line.md) page. - -!!! note - - Running `merlin server init` again will *not* override any exisiting configuration that the users might have set or edited. By running this command again any missing files will be created for the users with exisiting defaults. *However,* it is highly advised that users back up their configuration in case an error occurs where configuration files are overriden. - -## Starting the Server and Linking it to Merlin - -!!! bug - - For LC users, servers cannot be started outside your home (`~/`) directory. - -!!! warning - - Newer versions of Redis have started requiring a global variable `LC_ALL` to be set in order for this to work. To set this properly, run: - - ```bash - export LC_ALL="C" - ``` - - If this is not set, the `merlin server start` command may seem to run forever until you manually terminate it. - -After initializing the server, starting the server is as simple as running: - -```bash -merlin server start -``` - -You can check that the server was started properly with: - -```bash -merlin server status -``` - -The `merlin server start` command will add new files to the local configuration `merlin_server/` folder: - -- `merlin_server.pf`: A process file containing information regarding the Redis process -- `app.yaml`: A new app.yaml file configured specifically for the containerized Redis server that we just started - -To have Merlin read this server configuration: - -=== "Copy Configuration to CWD" - - ```bash - cp merlin_server/app.yaml . - ``` - -=== "Make This Server Configuration Your Main Configuration" - - If you're going to use the server configuration as your main configuration, it's a good idea to make a backup of your current server configuration (if you have one): - - ```bash - mv ~/.merlin/app.yaml ~/.merlin/app.yaml.bak - ``` - - From here, simply copy the server configuration to your `~/.merlin/` folder: - - ```bash - cp merlin_server/app.yaml ~/.merlin/app.yaml - ``` - -You can check that Merlin recognizes the containerized server connection with: - -```bash -merlin info -``` - -If your servers are running and set up properly, this should output something similar to this: - -???+ success - - ```bash - * - *~~~~~ - *~~*~~~* __ __ _ _ - / ~~~~~ | \/ | | (_) - ~~~~~ | \ / | ___ _ __| |_ _ __ - ~~~~~* | |\/| |/ _ \ '__| | | '_ \ - *~~~~~~~ | | | | __/ | | | | | | | - ~~~~~~~~~~ |_| |_|\___|_| |_|_|_| |_| - *~~~~~~~~~~~ - ~~~*~~~* Machine Learning for HPC Workflows - - - - Merlin Configuration - ------------------------- - - config_file | /path/to/app.yaml - is_debug | False - merlin_home | /path/to/.merlin - merlin_home_exists | True - broker server | redis://default:******@127.0.0.1:6379/0 - broker ssl | False - results server | redis://default:******@127.0.0.1:6379/0 - results ssl | False - - Checking server connections: - ---------------------------- - broker server connection: OK - results server connection: OK - - Python Configuration - ------------------------- - - $ which python3 - /path/to/python3 - - $ python3 --version - Python x.y.z - - $ which pip3 - /path/to/pip3 - - $ pip3 --version - pip x.y.x from /path/to/pip (python x.y) - - "echo $PYTHONPATH" - ``` - -## Stopping the Server - -Once you're done using your containerized server, it can be stopped with: - -```bash -merlin server stop -``` - -You can check that it's no longer running with: - -```bash -merlin server status -``` diff --git a/docs/user_guide/installation.md b/docs/user_guide/installation.md index d73d51db5..29f4df04b 100644 --- a/docs/user_guide/installation.md +++ b/docs/user_guide/installation.md @@ -1,6 +1,6 @@ # Installation -The Merlin library can be installed by using [virtual environments and pip](#installing-with-pip) or [spack](#installing-with-spack). +The Merlin library can be installed by using [virtual environments and pip](#installing-with-virtual-environments-pip) or [spack](#installing-with-spack). Contributors to Merlin should follow the [Developer Setup](#developer-setup) below. diff --git a/docs/user_guide/interpreting_output.md b/docs/user_guide/interpreting_output.md index a60e5e490..96c81862b 100644 --- a/docs/user_guide/interpreting_output.md +++ b/docs/user_guide/interpreting_output.md @@ -677,7 +677,7 @@ A visual representation of the `merlin_info/` subdirectory with sample generatio f.write(result) ``` - Since this script uses some third party libraries ([`names`](https://pypi.org/project/names/) and [`numpy`](https://numpy.org/)), you'll need to install them to your current environment in order to run this example. If you're using a [virtual environment](./installation.md#installing-with-virtual-environments--pip), these can be installed with: + Since this script uses some third party libraries ([`names`](https://pypi.org/project/names/) and [`numpy`](https://numpy.org/)), you'll need to install them to your current environment in order to run this example. If you're using a [virtual environment](./installation.md#installing-with-virtual-environments-pip), these can be installed with: ```bash pip install names numpy diff --git a/docs/user_guide/running_studies.md b/docs/user_guide/running_studies.md index 91215c19e..5ff9dbfbe 100644 --- a/docs/user_guide/running_studies.md +++ b/docs/user_guide/running_studies.md @@ -152,7 +152,7 @@ merlin stop-workers --spec my_specification.yaml !!! note - If you wish to execute a workflow after dry-running it, simply use [`merlin restart`](#restart-merlin-restart) (to understand why this works, see the section below on [Restarting Workflows](#restarting-workflows)). + If you wish to execute a workflow after dry-running it, simply use [`merlin restart`](./command_line.md#restart-merlin-restart) (to understand why this works, see the section below on [Restarting Workflows](#restarting-workflows)). 'Dry run' means telling workers to create a study's workspace and all of its necessary subdirectories and scripts (with variables expanded) without actually executing the scripts. diff --git a/docs/user_guide/specification.md b/docs/user_guide/specification.md index 53c682875..cabf16b1c 100644 --- a/docs/user_guide/specification.md +++ b/docs/user_guide/specification.md @@ -82,7 +82,7 @@ The `batch` block is an optional block that enables specification of HPC schedul | Property Name | Required? | Type | Description | | ------------- | --------- | ---- | ----------- | | `bank` | Yes | str | Account to charge computing time to | -| `dry_run` | No | bool | Execute a [dry run](./command_line.md#dry-run) of the study | +| `dry_run` | No | bool | Execute a [dry run](./running_studies.md#dry-runs) of the study | | `launch_args` | No | str | Extra arguments for the parallel launch command | | `launch_pre` | No | str | Any configuration needed before the scheduler launch command (`srun`, `jsrun`, etc.) | | `nodes` | No | int | The number of nodes to use for all workers. This can be overridden in [the `resources` property of the `merlin` block](#resources). If this is unset the number of nodes will be queried from the environment, failing that, the number of nodes will be set to 1. | diff --git a/merlin/__init__.py b/merlin/__init__.py index a2d173a8f..b6c065fd7 100644 --- a/merlin/__init__.py +++ b/merlin/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -38,7 +38,7 @@ import sys -__version__ = "1.12.2b1" +__version__ = "1.12.2" VERSION = __version__ PATH_TO_PROJ = os.path.join(os.path.dirname(__file__), "") diff --git a/merlin/ascii_art.py b/merlin/ascii_art.py index 5c90a4b12..bcada70af 100644 --- a/merlin/ascii_art.py +++ b/merlin/ascii_art.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/merlin/celery.py b/merlin/celery.py index eb10f1a12..431d7928f 100644 --- a/merlin/celery.py +++ b/merlin/celery.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -33,7 +33,7 @@ import logging import os -from typing import Dict, Optional, Union +from typing import Any, Dict, List, Optional, Union import billiard import celery @@ -58,7 +58,8 @@ def patch_celery(): Celery has error callbacks but they do not work properly on chords that are nested within chains. - Credit to this function goes to: https://danidee10.github.io/2019/07/09/celery-chords.html + Credit to this function goes to + [the following post](https://danidee10.github.io/2019/07/09/celery-chords.html). """ def _unpack_chord_result( @@ -84,9 +85,41 @@ def _unpack_chord_result( # This function has to have specific args/return values for celery so ignore pylint -def route_for_task(name, args, kwargs, options, task=None, **kw): # pylint: disable=W0613,R1710 +def route_for_task( + name: str, + args: List[Any], + kwargs: Dict[Any, Any], + options: Dict[Any, Any], + task: celery.Task = None, + **kw: Dict[Any, Any], +) -> Dict[Any, Any]: # pylint: disable=W0613,R1710 """ - Custom task router for queues + Custom task router for Celery queues. + + This function routes tasks to specific queues based on the task name. + If the task name contains a colon, it splits the name to determine the queue. + + Args: + name: The name of the task being routed. + args: The positional arguments passed to the task. + kwargs: The keyword arguments passed to the task. + options: Additional options for the task. + task: The task instance (default is None). + **kw: Additional keyword arguments for THIS function (not the task). + + Returns: + A dictionary specifying the queue to route the task to. + If the task name contains a colon, it returns a dictionary with + the key "queue" set to the queue name. Otherwise, it returns + an empty dictionary. + + Example: + Using a colon in the name will return the string before the colon as the queue: + + ```python + >>> route_for_task("my_queue:my_task") + {"queue": "my_queue"} + ``` """ if ":" in name: queue, _ = name.split(":") @@ -114,9 +147,11 @@ def route_for_task(name, args, kwargs, options, task=None, **kw): # pylint: dis BROKER_URI = None RESULTS_BACKEND_URI = None +app_name = "merlin_test_app" if os.getenv("CELERY_ENV") == "test" else "merlin" + # initialize app with essential properties app: Celery = patch_celery().Celery( - "merlin", + app_name, broker=BROKER_URI, backend=RESULTS_BACKEND_URI, broker_use_ssl=BROKER_SSL, @@ -169,11 +204,12 @@ def route_for_task(name, args, kwargs, options, task=None, **kw): # pylint: dis # Pylint believes the args are unused, I believe they're used after decoration @worker_process_init.connect() -def setup(**kwargs): # pylint: disable=W0613 +def setup(**kwargs: Dict[Any, Any]): # pylint: disable=W0613 """ - Set affinity for the worker on startup (works on toss3 nodes) + Set affinity for the worker on startup (works on toss3 nodes). - :param `**kwargs`: keyword arguments + Args: + **kwargs: Keyword arguments. """ if "CELERY_AFFINITY" in os.environ and int(os.environ["CELERY_AFFINITY"]) > 1: # Number of cpus between workers. diff --git a/merlin/common/__init__.py b/merlin/common/__init__.py index 57477ea1f..4f0b58ae6 100644 --- a/merlin/common/__init__.py +++ b/merlin/common/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,3 +27,24 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `common` package provides shared utilities, classes, and logic used across Merlin. +It includes functionality for managing encryption, handling data sampling, working with +enumerations, and defining Celery tasks. + +Subpackages: + - `security/`: Contains functionality for managing encryption and ensuring secure + communication. Includes modules for general encryption logic and encrypting backend traffic. + +Modules: + dumper.py: Provides functionality for dumping information to files. + enums.py: Defines enumerations for interfaces. + sample_index_factory.py: Houses factory methods for creating + [`SampleIndex`][common.sample_index.SampleIndex] objects. + sample_index.py: Implements the logic for managing the sample hierarchy, including + the [`SampleIndex`][common.sample_index.SampleIndex] class. + tasks.py: Defines Celery tasks, breaking down the Directed Acyclic Graph ([`DAG`][study.dag.DAG]) + into smaller tasks that Celery can manage. + util_sampling.py: Contains utility functions for data sampling. +""" diff --git a/merlin/common/abstracts/__init__.py b/merlin/common/abstracts/__init__.py deleted file mode 100644 index 57477ea1f..000000000 --- a/merlin/common/abstracts/__init__.py +++ /dev/null @@ -1,29 +0,0 @@ -############################################################################### -# Copyright (c) 2023, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory -# Written by the Merlin dev team, listed in the CONTRIBUTORS file. -# -# -# LLNL-CODE-797170 -# All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. -# -# For details, see https://github.com/LLNL/merlin. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -############################################################################### diff --git a/merlin/common/dumper.py b/merlin/common/dumper.py index 96a940357..c687c4bfc 100644 --- a/merlin/common/dumper.py +++ b/merlin/common/dumper.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1 +# This file is part of Merlin, Version: 1.12.2 # # For details, see https://github.com/LLNL/merlin. # @@ -43,34 +43,53 @@ class Dumper: # pylint: disable=R0903 """ The dumper class is intended to help write information to files. - Currently, the supported file types to dump to are csv and json. - - Example csv usage: - dumper = Dumper("populations.csv") - # Eugene, OR has a population of 175096 - # Livermore, CA has a population of 86803 - population_data = { - "City": ["Eugene", "Livermore"], - "State": ["OR", "CA"], - "Population": [175096, 86803] - } - dumper.write(population_data, "w") - |---> Output will be written to populations.csv - - Example json usage: - dumper = Dumper("populations.json") - population_data = { - "OR": {"Eugene": 175096, "Portland": 641162}, - "CA": {"Livermore": 86803, "San Francisco": 815201} - } - dumper.write(population_data, "w") - |---> Output will be written to populations.json + Currently, the supported dump file types are: csv or json. + + Attributes: + file_name (str): The name of the file to write data to. + file_type (str): The type of the file (either "csv" or "json") determined + from the file name. + + Methods: + write: Writes information to the specified output file based on the file type. + _csv_write: Writes information to a CSV file. + _json_write: Writes information to a JSON file. + + Example: + CSV usage: + ```python + dumper = Dumper("populations.csv") + # Eugene, OR has a population of 175096 + # Livermore, CA has a population of 86803 + population_data = { + "City": ["Eugene", "Livermore"], + "State": ["OR", "CA"], + "Population": [175096, 86803] + } + dumper.write(population_data, "w") # Output will be written to populations.csv + ``` + + Example: + JSON usage: + ```python + dumper = Dumper("populations.json") + population_data = { + "OR": {"Eugene": 175096, "Portland": 641162}, + "CA": {"Livermore": 86803, "San Francisco": 815201} + } + dumper.write(population_data, "w") # Output will be written to populations.json + ``` """ - def __init__(self, file_name): + def __init__(self, file_name: str): """ - Initialize the class and ensure the file is of a supported type. - :param `file_name`: The name of the file to dump to eventually + Initializes the Dumper class and validates the file type. + + Args: + file_name: The name of the file to write data to. + + Raises: + ValueError: If the file type is not supported. Supported types are CSV and JSON. """ supported_types = ["csv", "json"] @@ -87,9 +106,14 @@ def __init__(self, file_name): def write(self, info_to_write: Dict, fmode: str): """ - Write information to an outfile. - :param `info_to_write`: The information you want to write to the output file - :param `fmode`: The file write mode ("w", "a", etc.) + Writes information to the specified output file. + + This method determines the file type and calls the appropriate + method to write the data. + + Args: + info_to_write: The information to write to the output file. + fmode: The file write mode ("w" for write, "a" for append, etc.). """ if self.file_type == "csv": self._csv_write(info_to_write, fmode) @@ -98,10 +122,12 @@ def write(self, info_to_write: Dict, fmode: str): def _csv_write(self, csv_to_dump: Dict[str, List], fmode: str): """ - Write information to a csv file. - :param `csv_to_dump`: The information to write to the csv file. - Dict keys will be the column headers and values will be the column values. - :param `fmode`: The file write mode ("w", "a", etc.) + Writes information to a CSV file. + + Args: + csv_to_dump: The data to write to the CSV file. Keys are column + headers and values are column values. + fmode: The file write mode ("w" for write, "a" for append, etc.). """ # If we have statuses to write, create a csv writer object and write to the csv file with open(self.file_name, fmode) as outfile: @@ -112,9 +138,11 @@ def _csv_write(self, csv_to_dump: Dict[str, List], fmode: str): def _json_write(self, json_to_dump: Dict[str, Dict], fmode: str): """ - Write information to a json file. - :param `json_to_dump`: The information to write to the json file. - :param `fmode`: The file write mode ("w", "a", etc.) + Writes information to a JSON file. + + Args: + json_to_dump: The data to write to the JSON file. + fmode: The file write mode ("w" for write, "a" for append, etc.). """ # Appending to json requires file mode to be r+ for json.load if fmode == "a": @@ -132,11 +160,18 @@ def _json_write(self, json_to_dump: Dict[str, Dict], fmode: str): def dump_handler(dump_file: str, dump_info: Dict): """ - Help handle the process of creating a Dumper object and writing + Handles the process of creating a Dumper object and writing data to an output file. - :param `dump_file`: A filepath to the file we're dumping to - :param `dump_info`: A dict of information that we'll be dumping to `dump_file` + This function checks if the specified dump file exists to determine + the appropriate file write mode (append or write). It then creates + a Dumper object and writes the provided information to the file, + logging the process. + + Args: + dump_file: The filepath to the file where data will be dumped. + dump_info: A dictionary containing the information to be written + to the `dump_file`. """ # Create a dumper object to help us write to dump_file dumper = Dumper(dump_file) diff --git a/merlin/common/abstracts/enums/__init__.py b/merlin/common/enums.py similarity index 54% rename from merlin/common/abstracts/enums/__init__.py rename to merlin/common/enums.py index a95ba0872..0faa6f4c7 100644 --- a/merlin/common/abstracts/enums/__init__.py +++ b/merlin/common/enums.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -28,7 +28,7 @@ # SOFTWARE. ############################################################################### -"""Package for providing enumerations for interfaces""" +"""This module provides enumerations for interfaces.""" from enum import IntEnum @@ -37,15 +37,30 @@ class ReturnCode(IntEnum): """ - Merlin return codes. + Enum for Merlin return codes. + + This class defines various return codes used in the Merlin system to indicate + the status of operations. Each return code corresponds to a specific outcome + of a process. + + Attributes: + OK (int): Indicates a successful operation. Numeric value: 0. + ERROR (int): Indicates a general error occurred. Numeric value: 1. + RESTART (int): Indicates that the process should be restarted. Numeric value: 100. + SOFT_FAIL (int): Indicates a non-critical failure that allows for recovery. Numeric value: 101. + HARD_FAIL (int): Indicates a critical failure that cannot be recovered from. Numeric value: 102. + DRY_OK (int): Indicates a successful operation in a dry run (no changes made). Numeric value: 103. + RETRY (int): Indicates that the operation should be retried. Numeric value: 104. + STOP_WORKERS (int): Indicates that worker processes should be stopped. Numeric value: 105. + RAISE_ERROR (int): Indicates that an error should be raised. Numeric value: 106. """ - OK = 0 - ERROR = 1 - RESTART = 100 - SOFT_FAIL = 101 - HARD_FAIL = 102 - DRY_OK = 103 - RETRY = 104 - STOP_WORKERS = 105 - RAISE_ERROR = 106 + OK: int = 0 + ERROR: int = 1 + RESTART: int = 100 + SOFT_FAIL: int = 101 + HARD_FAIL: int = 102 + DRY_OK: int = 103 + RETRY: int = 104 + STOP_WORKERS: int = 105 + RAISE_ERROR: int = 106 diff --git a/merlin/common/openfilelist.py b/merlin/common/openfilelist.py deleted file mode 100644 index 16aa2f87c..000000000 --- a/merlin/common/openfilelist.py +++ /dev/null @@ -1,198 +0,0 @@ -#!/usr/bin/env python - -############################################################################### -# Copyright (c) 2023, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory -# Written by the Merlin dev team, listed in the CONTRIBUTORS file. -# -# -# LLNL-CODE-797170 -# All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. -# -# For details, see https://github.com/LLNL/merlin. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -############################################################################### - -""" - OpenFileList - - A synthetic file class that opens a list of files and reads them as if they - were a single file - - SYNOPSIS: - - with OpenFileList(["file1.txt","file2.txt",...]) as f : - print f.read(); - - reads the concatenation of file1.txt, file2.txt, etc. - - file methods supported : - - f.read([bytes]) - f.readlines([bytes]) - f.readline([bytes]) - f.tell() - f.close() - f.__iter__() - - TODO: - implement a seek method - -""" - -# This file is not currently used so we don't care what pylint has to say -# pylint: skip-file - -import copy - - -class OpenFileList: - openwas = open - - def __new__(cls, files, *v, **kw): - if isinstance(files, str): - return open(files, *v, **kw) - return super(OpenFileList, cls).__new__(cls) - - def __init__(self, files, *v, **kw): - self.files = copy.copy(files) - self.argv, self.argkw = (v, kw) - if self.files: - self.fnnow = self.files.pop(0) - self.fnow = open(self.fnnow, *v, **kw) if files else None # noqa - self.atend = False - else: - self.fnnow = self.fnow = None - self.atend = True - self._tell = 0 - - def _errclosed(self): - raise ValueError("I/O operation on closed file") - - def _tonext(self): - if self.fnow is not None: - self._tell += self.fnow.tell() - self.fnow.close() - if self.files: - self.fnnow = self.files.pop(0) - self.fnow = open(self.fnnow, *self.argv, **self.argkw) # noqa - else: - self.fnnow = self.fnow = None - self.atend = True - - def tell(self): - if self.fnow is None: - return self._tell - return self._tell + self.fnow.tell() - - def read(self, n=None): - if self.atend: - return "" - if self.fnow is None: - self._errclosed() - if n is None: - n = 1 << 32 - s = "" - while n and (self.files or self.fnow is not None): - ns = self.fnow.read(n) - if ns: - n -= len(ns) - s += ns - else: - self._tonext() - return s - - def readlines(self, b=None): - if self.atend: - return [] - if self.fnow is None: - self._errclosed() - s = self.read(b) - if not s: - return [] - if not s.endswith("\n"): - ch = "" - while ch != "\n": - ch = self.read(1) - if ch == "" or ch == "\n": - break - s += ch - return s.split("\n") - else: - return s[:-1].split("\n") - - def readline(self, b=None): - if self.atend: - return [] - if self.fnow is None: - return self._errclosed() - if b is None: - s = self.fnow.readline() - else: - s = self.fnow.readline(b) - if not s: - self._tonext() - return s - - def __iter__(self): - while not self.atend: - yield self.readline() - - def close(self): - if self.fnow is not None: - self.fnow.close() - self.files = [] - self.atend = True - self.fnow = self.fnnow = None - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, tback): - self.close() - - -if __name__ == "__main__": - import os - import unittest - import uuid - - import numpy - - class TestOpenFileList(unittest.TestCase): - def test_opener(self): - stride = 5 - fn = [str(uuid.uuid1()) for i in range(3)] - e = numpy.diag(numpy.arange(float(stride) * len(fn))) - - for n, ff in enumerate(fn): # Create files. - with open(ff, "w") as f: - for i in range(stride * n, stride * (n + 1)): - print(" ".join(map(str, e[i])), file=f) - - with OpenFileList(fn) as f: # Load using loadtxt method. - ep = numpy.loadtxt(f) - - for i in fn: - os.unlink(i) # Delete files. - - self.assertTrue(numpy.all(ep == e)) - - unittest.main() diff --git a/merlin/common/opennpylib.py b/merlin/common/opennpylib.py deleted file mode 100644 index 872699ae1..000000000 --- a/merlin/common/opennpylib.py +++ /dev/null @@ -1,391 +0,0 @@ -#!/usr/bin/env python - -############################################################################### -# Copyright (c) 2023, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory -# Written by the Merlin dev team, listed in the CONTRIBUTORS file. -# -# -# LLNL-CODE-797170 -# All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. -# -# For details, see https://github.com/LLNL/merlin. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -############################################################################### - -""" -smallnpylib - -A simple library to read the .npy file header and return a dict -containing the header from the .npy file as well as a few other -keys. Also provides the OpenNPY class which is a way of seeking -into .npy files without using the memory mapping functionality -in the load function. Finally, there is the OpenNPYList class -which opens a list of OpenNPY files and allows for random -access among all of them. - - 'shape' shape for the array - 'fortran_order' is it in fortran order (normally false) - 'descr' the dtype description of the array - 'size' total size of the data - 'offset' position in file of start of data - 'rowsize' the total size in bytes of a single row in the file (e.g. sample) - 'itemsize' the size of a single element in the array - - SYNOPSIS : - - d = get_npy_info("myfile.txt"); - print d['offset'] # offset to data in file - print d['rowsize'] # number of bytes per row - - with open("myfile.npy") as f : - f.seek(d['offset']+d['rowsize]*N); - - OpenNPY - - with OpenNPY("myfile.npy") as a : - print a[5] # print row number 5 - print a[1:4] # print rows from 1,2,3 - my_array = a.to_array(); - print len(a) # number of rows in file - for i in a : - print i # print all the rows in a - print a.shape # shape of array - print a.dtype # dtype of array - - with OpenNPYList(["myfile1.npy","myfile2.npy",...]) as a : - print a[5] # print row number 5 - print a[1:4] # print rows from 1,2,3 - my_array = a.to_array(); - print len(a) # number of rows in file - for i in a : - print i # print all the rows in a - print a.shape # shape of array - print a.dtype # dtype of array - -""" -# This file is not currently used so we don't care what pylint has to say -# pylint: skip-file - -from typing import List, Tuple - -import numpy as np - - -try: - unistr = (unicode, str) - npy_magic = "\x93NUMPY" -except NameError: - unistr = str - npy_magic = b"\x93NUMPY" - - -def _get_npy_info2(f): - if isinstance(f, unistr): - f = open(f, "rb") # noqa - magic = f.read(6) # must be .npy file - assert magic == npy_magic # must be .npy file or ELSE - major, _ = list(map(ord, f.read(2))) - if major == 1: - hlen_char = list(map(ord, f.read(2))) - hlen = hlen_char[0] + 256 * hlen_char[1] - elif major == 2: - hlen_char = list(map(ord, f.read(4))) - # fmt: off - hlen = ( - hlen_char[0] - + 256 * hlen_char[1] - + 65536 * hlen_char[2] - + (1 << 24) * hlen_char[3] - ) - # fmt: on - else: - raise Exception("unknown .npy format, e.g. not 1 or 2") - hdr = eval(f.read(hlen)) # TODO remove eval - hdr["dtype"] = np.dtype(hdr["descr"]) - hdr["offset"] = f.tell() # location of data start - hdr["itemsize"] = np.dtype(hdr["descr"]).itemsize - hdr["rowsize"] = hdr["itemsize"] * np.product(hdr["shape"][1:]) - hdr["items"] = np.product(hdr["shape"]) - hdr["size"] = hdr["itemsize"] * hdr["items"] - return f, hdr - - -def _get_npy_info3(f): - if isinstance(f, unistr): - f = open(f, "rb") # noqa - magic = f.read(6) # must be .npy file - assert magic == npy_magic # must be .npy file or ELSE - major, _ = list(f.read(2)) - if major == 1: - hlen_char = list(f.read(2)) - hlen = hlen_char[0] + 256 * hlen_char[1] - elif major == 2: - hlen_char = list(f.read(4)) - # fmt: off - hlen = ( - hlen_char[0] - + 256 * hlen_char[1] - + 65536 * hlen_char[2] - + (1 << 24) * hlen_char[3] - ) - # fmt: on - else: - raise Exception("unknown .npy format, e.g. not 1 or 2") - hdr = eval(f.read(hlen)) # TODO remove eval - hdr["dtype"] = np.dtype(hdr["descr"]) - hdr["offset"] = f.tell() # location of data start - hdr["itemsize"] = np.dtype(hdr["descr"]).itemsize - hdr["rowsize"] = hdr["itemsize"] * np.product(hdr["shape"][1:]) - hdr["items"] = np.product(hdr["shape"]) - hdr["size"] = hdr["itemsize"] * hdr["items"] - return f, hdr - - -def _get_npy_info(f): - d = None - try: - d = _get_npy_info2(f) - except TypeError: - d = _get_npy_info3(f) - - return d - - -def get_npy_info(f): - try: - d = _get_npy_info2(f) - except TypeError: - d = _get_npy_info3(f) - d[0].close() - return d[1] - - -def read_items(f, hdr, idx, n=-1, sep=""): - f.seek(hdr["offset"] + idx * hdr["itemsize"]) - if n < 0: - n = hdr["items"] - idx - n = min(hdr["items"] - idx, n) - print(f"n is {n}") - return np.fromfile(f, dtype=hdr["dtype"], count=n, sep=sep) - - -def read_rows(f, hdr, idx, n=-1, sep=""): - f.seek(hdr["offset"] + idx * hdr["rowsize"]) - if n < 0: - n = hdr["shape"][0] - idx - n = min(hdr["shape"][0] - idx, n) - a = np.fromfile(f, dtype=hdr["dtype"], count=n * hdr["rowsize"] // hdr["itemsize"], sep=sep) - return np.reshape(a, (n,) + hdr["shape"][1:]) - - -def verify_open(func): # A wrapper function used by FileSamples. - """ - :param func: (function) a class instance method that needs to call - _verify_open before doing anything else. - """ - - def wrapper(self, *v, **kw): - self._verify_open() - return func(self, *v, **kw) - - return wrapper - - -class OpenNPY: - def __init__(self, f): - self.hdr = self.f = None - if isinstance(f, unistr): - self.fname = f - else: - self.f = f - self._verify_open() - - def _verify_open(self): - if self.f is None: - self.f, self.hdr = _get_npy_info(self.f if self.f is not None else self.fname) - - @verify_open - def load_header(self, close=True): - if close: - self.close() - return self.hdr - - def close(self): - if self.f is not None: - self.f.close() - self.f = None - - def __del__(self): - self.close() - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, tback): - self.close() - - @verify_open - def _shape(self): - return self.hdr["shape"] - - shape = property(fget=_shape) - - @verify_open - def _dtype(self): - return self.hdr["dtype"] - - dtype = property(fget=_dtype) - - @verify_open - def __getitem__(self, k): - if isinstance(k, int): - return read_rows(self.f, self.hdr, k, 1)[0] - elif isinstance(k, slice): - if k.step == 1 or k.step is None: - return read_rows(self.f, self.hdr, k.start, k.stop - k.start) - else: - return np.asarray( - [read_rows(self.f, self.hdr, _, 1)[0] for _ in range(k.start, k.stop, 1 if k.step is None else k.step)] - ) - - @verify_open - def __len__(self): - return self.hdr["items"] - - @verify_open - def __iter__(self): - for k in range(self.hdr["shape"][0]): - yield self[k] - - @verify_open - def to_array(self): - return read_rows(self.f, self.hdr, 0) - - -class OpenNPYList: - def __init__(self, filename_strs: List[str]): - self.filenames: List[str] = filename_strs - self.files: List[OpenNPY] = [OpenNPY(file_str) for file_str in self.filenames] - i: OpenNPY - for i in self.files: - i.load_header() - self.shapes: List[Tuple[int]] = [openNPY_obj.hdr["shape"] for openNPY_obj in self.files] - k: Tuple[int] - for k in self.shapes[1:]: - # Match subsequent axes shapes. - if k[1:] != self.shapes[0][1:]: - raise AttributeError(f"Mismatch in subsequent axes shapes: {k[1:]} != {self.shapes[0][1:]}") - self.tells: np.ndarray = np.cumsum([arr_shape[0] for arr_shape in self.shapes]) # Tell locations. - self.tells = np.hstack(([0], self.tells)) - - def close(self): - for i in self.files: - i.close() - - def __del__(self): - self.close() - - def __iter__(self): - for i in self.files: - for j in i: - yield j - - def __getitem__(self, k): - if isinstance(k, int): - if k < 0: - k = self.tells[-1] + k # Negative indexing. - if k >= self.tells[-1]: - raise IndexError("index %d is out of bounds" % k) - fno = (self.tells > k).argmax() - return self.files[fno - 1][k - self.tells[fno - 1]] - else: # Slice indexing. - # TODO : Implement a faster version. - return np.asarray([self[_] for _ in np.arange(k.start, k.stop, k.step if k.step is not None else 1)]) - - def to_array(self): - return np.vstack([_.to_array() for _ in self.files]) - - def __len__(self): - return self.tells[-1] - - def __enter__(self): - return self - - def __exit__(self, exc_type, exc_value, tback): - self.close() - - -__all__ = [ - "OpenNPYList", - "OpenNPY", - "unistr", - "read_items", - "read_rows", - "get_npy_info", -] - -if __name__ == "__main__": - import os - import sys - import unittest - import uuid - - if "-h" in sys.argv: - print(__doc__) - sys.exit(1) - - class TestOpenNPY(unittest.TestCase): - def test_seeknpy(self): - e = np.diag(np.arange(5.0)) - try: - fn = unicode(str(uuid.uuid1()) + ".npy") - except NameError: - fn = str(str(uuid.uuid1()) + ".npy") - np.save(fn, e) - with OpenNPY(fn) as a: - ep = np.asarray([_ for _ in a]) - os.unlink(fn) - self.assertTrue((ep == e).all()) - - def test_seeknpylist(self): - e = np.diag(np.arange(5.0)) - fn = str(uuid.uuid1()) + ".npy" - np.save(fn, e) - with OpenNPYList([fn, fn, fn]) as a: - ep = np.asarray([_ for _ in a]) - en = np.asarray(a[5:10]) - en2 = np.asarray(a[11:14]) - en3 = np.asarray(a[1:14]) - en4 = a.to_array() - # test __len__ method - self.assertEqual(len(a), 3 * len(e)) - os.unlink(fn) - # test read slice of whole file - self.assertTrue((en == e).all()) - self.assertTrue((en2 == e[1:4]).all()) # test slice - # test read all - self.assertTrue((ep == np.vstack((e, e, e))).all()) - # test to_array method - self.assertTrue((en4 == np.vstack((e, e, e))).all()) - # test slice read across files - self.assertTrue((en3 == np.vstack((e, e, e))[1:14]).all()) - - unittest.main() diff --git a/merlin/common/sample_index.py b/merlin/common/sample_index.py index c7808bd3b..bfdaec4b2 100644 --- a/merlin/common/sample_index.py +++ b/merlin/common/sample_index.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,12 +29,14 @@ ############################################################################### """ -The merlin sample_index module, which contains the SampleIndex class. +This module contains the logic for managing the sample hierarchy, including +the implementation of the [`SampleIndex`][common.sample_index.SampleIndex] class. """ import logging import os from contextlib import suppress +from typing import Callable, Dict, Generator, List, Tuple LOG = logging.getLogger(__name__) @@ -42,14 +44,36 @@ MAX_SAMPLE = 99999999999999999999999999999999999 -def new_dir(path): - """Create a new directory at the given path if it does not exist.""" +def new_dir(path: str): + """ + Create a new directory at the specified path if it does not already exist. + + This function attempts to create the directory and suppresses any + OSError that may occur if the directory already exists. + + Args: + path: The path where the new directory should be created. + """ with suppress(OSError): os.makedirs(path) -def uniform_directories(num_samples=MAX_SAMPLE, bundle_size=1, level_max_dirs=100): - """Create a directory hierarchy uniformly stepping up directory sizes.""" +def uniform_directories(num_samples: int = MAX_SAMPLE, bundle_size: int = 1, level_max_dirs: int = 100) -> List[int]: + """ + Create a directory hierarchy with uniformly increasing directory sizes. + + This function generates a list of directory sizes, starting from + the specified `bundle_size` and increasing by a factor of + `level_max_dirs` until the total number of samples is reached. + + Args: + num_samples: The total number of samples to consider. + bundle_size: The initial size of each bundle. + level_max_dirs: The factor by which to increase directory sizes at each level. + + Returns: + A list of integers representing the sizes of directories in the hierarchy. + """ directory_sizes = [bundle_size] while directory_sizes[0] < num_samples: directory_sizes.insert(0, directory_sizes[0] * level_max_dirs) @@ -60,17 +84,80 @@ def uniform_directories(num_samples=MAX_SAMPLE, bundle_size=1, level_max_dirs=10 class SampleIndex: """ - A SampleIndex is the insitu representation of a directory hierarchy where - bundles of samples result files will be stored. Factory methods to produce - full hierarchies are provided, as well as to read an index from one - previously stored on disk. + Represents a hierarchical structure for managing bundles of sample files. + + A [`SampleIndex`][common.sample_index.SampleIndex] serves as an in-situ + representation of a directory hierarchy where bundles of sample result files + are stored. This class provides factory methods to create complete hierarchies + and to read an index from a previously stored file on disk. + + Attributes: + address (str): The full address of this node. + children (Dict[str, SampleIndex]): A dictionary containing the direct children of this node, which are also + of type SampleIndex. + depth (int): (class attribute) A class variable indicating the current depth in the + hierarchy, primarily used for pretty printing. + is_leaf (bool): Returns whether this node is a leaf in the hierarchy. + is_directory (bool): Returns whether this node is a directory (not a leaf). + is_parent_of_leaf (bool): Returns whether this node is the direct parent of a leaf. + is_grandparent_of_leaf (bool): Returns whether this node is the parent of a parent of a leaf. + is_great_grandparent_of_leaf (bool): Returns whether this node is the parent of a grandparent of a leaf. + leafid (int): The unique leaf ID of this node. + max (int): The maximum global sample ID of any child node. + min (int): The minimum global sample ID of any child node. + name (str): The name of this node in the hierarchy. + num_bundles (int): The total number of bundles in this index. + + Methods: + traverse: Yields the full path and associated node for each node meeting a specified condition. + traverse_all: Returns a generator that traverses all nodes in the + [`SampleIndex`][common.sample_index.SampleIndex]. + traverse_bundles: Returns a generator that traverses all bundles (leaves) in the + [`SampleIndex`][common.sample_index.SampleIndex]. + traverse_directories: Returns a generator that traverses all directories in the + [`SampleIndex`][common.sample_index.SampleIndex]. + check_valid_addresses_for_insertion: Validates addresses for insertion into the hierarchy. + __getitem__: Retrieves a child node by its full address. + __setitem__: Sets a child node at a specified full address. + write_directory: Creates a new directory associated with this node. + write_directories: Recursively writes directories for this node and its children. + get_path_to_sample: Retrieves the file path to the bundle file with a specified sample ID. + write_single_sample_index_file: Writes a single sample index file for this node. + write_multiple_sample_index_files: Writes multiple sample index files for this node and its children. + make_directory_string: Creates a delimited string representation of the directories in the index. + __str__: Returns a string representation of the [`SampleIndex`][common.sample_index.SampleIndex], + including its children. """ # Class variable to indicate depth (mostly used for pretty printing). - depth = -1 - - def __init__(self, minid, maxid, children, name, leafid=-1, num_bundles=0, address=""): # pylint: disable=R0913 - """The constructor.""" + depth: int = -1 + + def __init__( + self, + minid: int, + maxid: int, + children: Dict[str, "SampleIndex"], + name: str, + leafid: int = -1, + num_bundles: int = 0, + address: str = "", + ): # pylint: disable=R0913 + """ + Initializes a new instance of the `SampleIndex` class. + + Args: + minid: The minimum global sample ID of any child node. + maxid: The maximum global sample ID of any child node. + children: A dictionary containing the direct children of this node, + where the keys are the children's full addresses and the values are `SampleIndex` instances. + name: The name of this node in the hierarchy. + leafid: The unique leaf ID of this node. + num_bundles: The total number of bundles in this index. + address: The full address of this node. + + Raises: + TypeError: If `children` is not of type `dict`. + """ # The direct children of this node, generally also of type SampleIndex. # A dictionary keyed by the childrens' full addresses. @@ -78,39 +165,56 @@ def __init__(self, minid, maxid, children, name, leafid=-1, num_bundles=0, addre LOG.error("SampleIndex children must be of type dict") raise TypeError - self.children = children - self.name = name # The name of this node. - self.address = str(address) # The address of this node + self.children: Dict[str, "SampleIndex"] = children + self.name: str = name # The name of this node. + self.address: str = str(address) # The address of this node # @NOTE: The following are only valid if no insertions,splits, or # joins have taken place since the last renumber() call, or this index # was constructed with apriori global knowledge of its contents. # The minimum global sample ID of any of the children of this node. - self.min = minid + self.min: int = minid # The maximum global sample ID of any of the children of this node. - self.max = maxid + self.max: int = maxid # The total number of bundles in this index. - self.num_bundles = num_bundles + self.num_bundles: int = num_bundles # The unique leaf ID of this node - self.leafid = leafid + self.leafid: int = leafid @property - def is_leaf(self): - """Returns whether this is a leaf in the graph""" + def is_leaf(self) -> bool: + """ + Indicates whether this node is a leaf in the hierarchy. + + A leaf node is defined as a node that has no children. This property + returns True if the node has no children, and False otherwise. + """ return len(self.children) == 0 @property - def is_directory(self): - """Returns whether this is a directory (not a leaf) in the graph""" + def is_directory(self) -> bool: + """ + Indicates whether this node is a directory (not a leaf). + + A directory node is defined as a node that has one or more children. + This property returns True if the node has children, and False if it + is a leaf node. + """ return len(self.children) > 0 @property - def is_parent_of_leaf(self): - """Returns whether this is the direct parent of a leaf in the graph""" + def is_parent_of_leaf(self) -> bool: + """ + Indicates whether this node is the direct parent of a leaf. + + This property checks if the current node is a directory and if any of + its children are leaf nodes. It returns True if at least one child is + a leaf, and False otherwise. + """ if not self.is_directory: return False for child_val in self.children.values(): @@ -119,8 +223,14 @@ def is_parent_of_leaf(self): return False @property - def is_grandparent_of_leaf(self): - """Returns whether this is the parent of a parent of a leaf in the graph""" + def is_grandparent_of_leaf(self) -> bool: + """ + Indicates whether this node is the parent of a parent of a leaf. + + This property checks if the current node is a directory and if any of + its children are parents of leaf nodes. It returns True if at least + one child is a parent of a leaf, and False otherwise. + """ if not self.is_directory: return False for child_val in self.children.values(): @@ -129,8 +239,14 @@ def is_grandparent_of_leaf(self): return False @property - def is_great_grandparent_of_leaf(self): - """Returns whether this is the parent of a parent of a leaf in the graph""" + def is_great_grandparent_of_leaf(self) -> bool: + """ + Indicates whether this node is the parent of a grandparent of a leaf. + + This property checks if the current node is a directory and if any of + its children are grandparents of leaf nodes. It returns True if at + least one child is a grandparent of a leaf, and False otherwise. + """ if not self.is_directory: return False for child_val in self.children.values(): @@ -138,19 +254,44 @@ def is_great_grandparent_of_leaf(self): return True return False - def traverse(self, path=None, conditional=lambda c: True, bottom_up=True, top_level=True): + def traverse( + self, + path: str = None, + conditional: Callable = lambda c: True, + bottom_up: bool = True, + top_level: bool = True, + ) -> Generator[Tuple[str, "SampleIndex"], None, None]: """ - Yield the full path and associated node for each node that meets the - conditional - param:path: The path to this node. - param:conditional: A lambda that returns a boolean, takes a - SampleIndex as its only argument. - - param:bottom_up: If True, yield leaves of the tree first. Otherwise, - yield top level nodes first. - param:top_level: used to allow filtering of yielded values based off - the conditional. Should only be set to False internally for the - recursive calls. + Traverse the tree structure and yield the full path and associated node + for each node that meets the specified conditional criteria. + + This method allows for flexible traversal of the tree, either from + the top down or bottom up, depending on the `bottom_up` parameter. + Nodes are yielded based on whether they satisfy the provided + conditional function. + + Notes: + - The method uses a "SKIP ME" placeholder to manage the + recursion flow, ensuring that the traversal can skip + non-qualifying nodes without breaking the iteration. + + Args: + path: The current path to this node. If None, it defaults to + the name of the node. + conditional: A function that takes a + [`SampleIndex`][common.sample_index.SampleIndex] instance as its + only argument and returns a boolean. It determines whether a + node should be yielded. + bottom_up: If True, yields leaf nodes first, moving upwards through + the tree. If False, yields top-level nodes first. + top_level: A flag used internally to control filtering of yielded + values based on the conditional. Should only be set to False + for recursive calls. + + Yields: + A tuple containing the full path to the node and the + associated [`SampleIndex`][common.sample_index.SampleIndex] node + that meets the conditional criteria. """ if path is None: path = self.name @@ -171,36 +312,118 @@ def traverse(self, path=None, conditional=lambda c: True, bottom_up=True, top_le if not top_level: yield "SKIP ME" - def traverse_all(self, bottom_up=True): + def traverse_all(self, bottom_up: bool = True) -> Generator[Tuple[str, "SampleIndex"], None, None]: """ - Returns a generator that will traverse all nodes in the SampleIndex. + Traverse all nodes in the [`SampleIndex`][common.sample_index.SampleIndex]. + + This method returns a generator that yields all nodes in the + [`SampleIndex`][common.sample_index.SampleIndex], regardless of their type + (leaf or directory). The traversal order can be controlled by the `bottom_up` + parameter. + + Notes: + This method calls the [`traverse`][common.sample_index.SampleIndex.traverse] + method with a conditional that always returns True, ensuring all nodes are + included. + + Args: + bottom_up: If True, yields leaf nodes first, moving upwards + through the tree. If False, yields top-level nodes first. + + Returns: + A tuple containing the full path to each node and the associated + [`SampleIndex`][common.sample_index.SampleIndex] node. """ return self.traverse(path=self.name, conditional=lambda c: True, bottom_up=bottom_up) - def traverse_bundles(self): + def traverse_bundles(self) -> Generator[Tuple[str, "SampleIndex"], None, None]: """ - Returns a generator that will traverse all Bundles (leaves) in the - SampleIndex. + Traverse all Bundles (leaf nodes) in the [`SampleIndex`][common.sample_index.SampleIndex]. + + This method returns a generator that yields only the leaf nodes + (Bundles) in the [`SampleIndex`][common.sample_index.SampleIndex]. + It filters the nodes based on their type, ensuring that only leaves are returned. + + Notes: + This method calls the [`traverse`][common.sample_index.SampleIndex.traverse] + method with a conditional that checks if a node is a leaf, ensuring only Bundles + are yielded. + + Returns: + A tuple containing the full path to each Bundle and the associated + [`SampleIndex`][common.sample_index.SampleIndex] node. """ return self.traverse(path=self.name, conditional=lambda c: c.is_leaf) - def traverse_directories(self, bottom_up=False): + def traverse_directories(self, bottom_up: bool = False) -> Generator[Tuple[str, "SampleIndex"], None, None]: """ - Returns a generator that will traverse all Directories in the - SampleIndex. + Traverse all Directories in the [`SampleIndex`][common.sample_index.SampleIndex]. + + This method returns a generator that yields all directory nodes + in the [`SampleIndex`][common.sample_index.SampleIndex]. The + traversal order can be controlled by the `bottom_up` parameter. + + Notes: + This method calls the [`traverse`][common.sample_index.SampleIndex.traverse] + method with a conditional that checks if a node is a directory, ensuring only + directories are yielded. + + Args: + bottom_up: If True, yields leaf directories first, moving + upwards through the tree. If False, yields top-level + directories first. + + Yields: + A tuple containing the full path to each Directory and the + associated [`SampleIndex`][common.sample_index.SampleIndex] + node. """ return self.traverse(path=self.name, conditional=lambda c: c.is_directory, bottom_up=bottom_up) @staticmethod - def check_valid_addresses_for_insertion(full_address, sub_tree): + def check_valid_addresses_for_insertion(full_address: str, sub_tree: "SampleIndex"): """ - TODO + Check if the provided address is valid for insertion into the subtree. + + This method traverses all nodes in the given subtree and verifies + that no existing node's address conflicts with the specified + `full_address`. If any node's address starts with the `full_address`, + a TypeError is raised, indicating that the insertion would create + an invalid state. + + Args: + full_address: The full address to be checked for validity before + insertion. + sub_tree: The subtree in which the address will be checked. + + Raises: + TypeError: If any node in the subtree has an address that + conflicts with the `full_address`. """ for _, node in sub_tree.traverse_all(): if node.address[0 : len(full_address)] != full_address: raise TypeError - def __getitem__(self, full_address): + def __getitem__(self, full_address: str) -> "SampleIndex": + """ + Retrieve the subtree associated with the given full address. + + This method allows for accessing nodes in the + [`SampleIndex`][common.sample_index.SampleIndex] using + the full address. If the full address matches the current node's + address, the node itself is returned. Otherwise, it recursively + searches through the children to find the corresponding node. + + Args: + full_address: The full address of the node to retrieve. + + Returns: + The node associated with the specified full address. + + Raises: + KeyError: If no node with the specified full address exists + in the [`SampleIndex`][common.sample_index.SampleIndex]. + """ if full_address == self.address: return self for child_val in list(self.children.values()): @@ -208,7 +431,28 @@ def __getitem__(self, full_address): return child_val[full_address] raise KeyError - def __setitem__(self, full_address, sub_tree): + def __setitem__(self, full_address: str, sub_tree: "SampleIndex"): + """ + Set or replace the subtree associated with the given full address. + + This method allows for inserting or updating a node in the + [`SampleIndex`][common.sample_index.SampleIndex]. If the full + address matches the current node's address, a KeyError is raised + to prevent self-assignment. The method searches through the children + to find the appropriate location for insertion. If a node already + exists at the specified address, it will be replaced after validating + the insertion. + + Args: + full_address: The full address of the node to set or replace. + sub_tree: The subtree to insert or replace at the specified + address. + + Raises: + KeyError: If the full address matches the current node's + address or if no matching child node is found for + insertion. + """ if full_address == self.address: # This should never happen. raise KeyError @@ -225,29 +469,74 @@ def __setitem__(self, full_address, sub_tree): # Replace if we already have something at this address. if delete_me is not None: - self.children.__delitem__(full_address) SampleIndex.check_valid_addresses_for_insertion(full_address, sub_tree) + self.children.__delitem__(full_address) self.children[full_address] = sub_tree return raise KeyError - def write_directory(self, path): - """Creates a new directory associated with this node in the graph.""" + def write_directory(self, path: str): + """ + Create a new directory associated with this node in the graph. + + This method checks if the current node is a directory and, if so, + creates a new directory at the specified path using the node's + name. The directory is created using the + [`new_dir`][common.sample_index.new_dir] function. + + Args: + path: The base path where the new directory will be created. + """ if self.is_directory: new_dir(os.path.join(path, self.name)) - def write_directories(self, path="."): + def write_directories(self, path: str = "."): """ - Creates the directory tree associated with this node and its children. + Create the directory tree associated with this node and its children. + + This method initiates the creation of the directory structure + starting from the current node. It first creates the directory for + the current node and then recursively creates directories for all + child nodes. The base path can be specified, and the directories + will be created relative to this path. + + Args: + path: The base path where the directory tree will be created. + Defaults to the current directory ("."), meaning the + directories will be created in the current working + directory. """ self.write_directory(path) for child_val in list(self.children.values()): child_val.write_directories(os.path.join(path, self.name)) - def get_path_to_sample(self, sample_id): + def get_path_to_sample(self, sample_id: int) -> str: """ - Retrieves the file path to the bundle file with the sample_id of - interest. Note this only works when the global numbering is known. + Retrieve the file path to the bundle file associated with the specified + sample ID. + + This method constructs the path to the bundle file by traversing the + directory structure represented by the current node and its children. + It checks each child node to determine if the provided `sample_id` + falls within the range defined by the child's `min` and `max` attributes. + If a matching child is found, the method recursively calls itself to + build the complete path. + + Notes: + - This method assumes that the global numbering system is known + and that the `min` and `max` attributes of child nodes are + correctly defined. + - If no matching child is found, the method will return the + current node's name as the base path. + + Args: + sample_id: The identifier of the sample for which the file path is + to be retrieved. This ID must correspond to a sample within the + known global numbering system. + + Returns: + The constructed file path to the bundle file associated + with the specified `sample_id`. """ path = self.name for child_val in self.children.values(): @@ -255,8 +544,32 @@ def get_path_to_sample(self, sample_id): path = os.path.join(path, child_val.get_path_to_sample(sample_id)) return path - def write_single_sample_index_file(self, path): - """Writes the index file associated with this node.""" + def write_single_sample_index_file(self, path: str) -> str: + """ + Write the index file associated with this node. + + This method creates an index file named "sample_index.txt" in the + specified directory path. The index file contains information about + the child nodes of the current node. For each child, it records + whether the child is a leaf or a directory, along with its address, + name, and the range of sample IDs it covers. + + Notes: + - This method will only execute if the current node is identified + as a directory + ([`self.is_directory`][common.sample_index.SampleIndex.is_directory]). + - The index file format includes lines for each child in the + following format: + - For leaf nodes: `BUNDLE:
\tname:\tSAMPLES:[, )` + - For directory nodes: `DIR:
\tname:\tSAMPLES:[, )` + + Args: + path: The base path where the index file will be created. + + Returns: + The full path to the created index file if the current node is + a directory; otherwise, returns None. + """ if not self.is_directory: return None @@ -273,10 +586,23 @@ def write_single_sample_index_file(self, path): ) return fname - def write_multiple_sample_index_files(self, path="."): + def write_multiple_sample_index_files(self, path: str = ".") -> List[str]: """ - Write index files that couple with location in directory hierarchy, - contain necessary info to create a new index. + Write index files that correspond to the location in the directory hierarchy. + + This method generates index files for the current node and all its + children, allowing for a structured representation of the sample + indices in the directory hierarchy. It first writes a single index + file for the current node and then recursively writes index files + for each child node. + + Args: + path: The base path where the index files will be created. + Defaults to the current directory ("."), meaning the index + files will be created in the current working directory. + + Returns: + A list of file paths to the created index files. """ filepath = self.write_single_sample_index_file(path) filepaths = [] @@ -286,18 +612,29 @@ def write_multiple_sample_index_files(self, path="."): filepaths += child_val.write_multiple_sample_index_files(os.path.join(path, self.name)) return filepaths - def make_directory_string(self, delimiter=" ", just_leaf_directories=True): + def make_directory_string(self, delimiter: str = " ", just_leaf_directories: bool = True) -> str: """ - Make a string that is a delimited list of the directories in the - index. - - :param delimiter: the characters used to separate the directories - :param just_leaf_directories: A boolean on whether just to return the - leaf (bottom) directories - :returns: A string representation of the directories - e.g. - "0/0 0/1 0/2 1/0 1/1 1/2" - + Create a delimited string representation of the directories in the index. + + This method generates a string that lists the directories in the + current index, separated by the specified delimiter. The user can + choose to include only the leaf directories or all directories. + + Notes: + - The method utilizes the + [`traverse_directories`][common.sample_index.SampleIndex.traverse_directories] + function to retrieve the directory paths and their corresponding nodes. + + Args: + delimiter: The characters used to separate the directories in the + resulting string. + just_leaf_directories: If True, only leaf directories (the bottom-level + directories) will be included in the output. If False, all directories + will be included. + + Returns: + A string representation of the directories, formatted as a delimited + list. For example: "0/0 0/1 0/2 1/0 1/1 1/2". """ # fmt: off if just_leaf_directories: @@ -309,8 +646,28 @@ def make_directory_string(self, delimiter=" ", just_leaf_directories=True): # fmt: on return delimiter.join([path for path, _ in self.traverse_directories()]) - def __str__(self): - """String representation.""" + def __str__(self) -> str: + """ + Return a string representation of the + [`SampleIndex`][common.sample_index.SampleIndex] object. + + This method provides a formatted string that represents the current + node in the sample index, including its address, type (BUNDLE or + DIRECTORY), and relevant attributes such as minimum and maximum + sample IDs. If the node is a directory, it recursively includes + the string representations of its child nodes. + + Notes: + - The method uses a class variable `depth` to manage indentation + levels for nested directories, enhancing readability. + - The output format varies depending on whether the node is a + leaf or a directory. + + Returns: + A formatted string representation of the + [`SampleIndex`][common.sample_index.SampleIndex] object, + including its children if applicable. + """ SampleIndex.depth = SampleIndex.depth + 1 if self.is_leaf: result = (" " * SampleIndex.depth) + f"{self.address}: BUNDLE {self.leafid} MIN {self.min} MAX {self.max}\n" diff --git a/merlin/common/sample_index_factory.py b/merlin/common/sample_index_factory.py index eb4fbcc61..211e9864c 100644 --- a/merlin/common/sample_index_factory.py +++ b/merlin/common/sample_index_factory.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,8 +29,10 @@ ############################################################################### """ -SampleIndex factory methods +This module houses [`SampleIndex`][common.sample_index.SampleIndex] factory methods. """ +from typing import List + from parse import parse from merlin.common.sample_index import MAX_SAMPLE, SampleIndex @@ -43,27 +45,38 @@ def create_hierarchy( - num_samples, - bundle_size, - directory_sizes=None, - root=".", - start_sample_id=0, - start_bundle_id=0, - address="", - n_digits=1, -): + num_samples: int, + bundle_size: int, + directory_sizes: List[int] = None, + root: str = ".", + start_sample_id: int = 0, + start_bundle_id: int = 0, + address: str = "", + n_digits: int = 1, +) -> SampleIndex: """ - SampleIndex Hierarchy Factory method. Wraps - create_hierarchy_from_max_sample, which is a max_sample-based API, not a - numSample-based API like this method. - - :param num_samples: The total number of samples. - :bundle_size: The max number of samples a bundle file is responsible for. - :directory_sizes: The number of samples each directory is responsible - for - a list, one value for each level in the directory hierarchy. - :root: The root path of this index. Defaults to ".". - :start_sample_id: The start of the sample count. Defaults to 0. - :n_digits: The number of digits to pad the directories with + Factory method to create a [`SampleIndex`][common.sample_index.SampleIndex] + hierarchy based on the number of samples. + + This method wraps the + [`create_hierarchy_from_max_sample`][common.sample_index_factory.create_hierarchy_from_max_sample] + function, which operates on a maximum sample basis rather than a total + sample count. + + Args: + num_samples (int): The total number of samples. + bundle_size (int): The maximum number of samples a bundle file can handle. + directory_sizes (List[int]): A list specifying the number of samples each directory + is responsible for. + root (str): The root path of the index. + start_sample_id (int): The starting sample ID. + start_bundle_id (int): The starting bundle ID. + address (str): An optional address prefix for the hierarchy. + n_digits (int): The number of digits to pad the directory names. + + Returns: + (common.sample_index.SampleIndex): The root [`SampleIndex`][common.sample_index.SampleIndex] + object representing the hierarchy. """ if directory_sizes is None: directory_sizes = [] @@ -80,29 +93,37 @@ def create_hierarchy( def create_hierarchy_from_max_sample( - max_sample, - bundle_size, - directory_sizes=None, - root=".", - start_bundle_id=0, - min_sample=0, - address="", - n_digits=1, -): - """ " - Construct the SampleIndex based off the total number of samples and the - chunking size at each depth in the hierarchy. - - This method will add new SampleIndex objects as this SampleIndex's - children if directory_sizes is not the empty set. - - :param max_sample: The max Sample ID this hierarchy is responsible for. - :bundle_size: The max number of samples a bundle file is responsible for. - :directory_sizes: The number of samples each directory is responsible - for - a list, one value for each level in the directory hierarchy. - :bundle_id: The current bundle_id count. - :min_sample: The start of the sample count. - :n_digits: The number of digits to pad the directories with + max_sample: int, + bundle_size: int, + directory_sizes: List[int] = None, + root: str = ".", + start_bundle_id: int = 0, + min_sample: int = 0, + address: str = "", + n_digits: int = 1, +) -> SampleIndex: + """ + Constructs a [`SampleIndex`][common.sample_index.SampleIndex] hierarchy based on + the maximum sample ID and chunking size at each depth. + + This method adds new [`SampleIndex`][common.sample_index.SampleIndex] objects as + children if `directory_sizes` is provided. + + Args: + max_sample: The maximum Sample ID this hierarchy is responsible for. + bundle_size: The maximum number of samples a bundle file can handle. + directory_sizes: A list specifying the number of samples each directory + is responsible for. + root: The root path of this index. + start_bundle_id: The starting bundle ID. + min_sample: The starting sample ID. + address: An optional address prefix for the hierarchy. + n_digits: The number of digits to pad the directory names. + + Returns: + (common.sample_index.SampleIndex): The root + [`SampleIndex`][common.sample_index.SampleIndex] object representing + the constructed hierarchy. """ if directory_sizes is None: directory_sizes = [] @@ -158,9 +179,23 @@ def create_hierarchy_from_max_sample( return SampleIndex(min_sample, max_sample, children, root, num_bundles=num_bundles, address=address) -def read_hierarchy(path): +def read_hierarchy(path: str) -> SampleIndex: """ - TODO + Reads a hierarchy from a specified path and constructs a + [`SampleIndex`][common.sample_index.SampleIndex]. + + This function reads a file named "sample_index.txt" in the given path, + parsing its contents to create a hierarchical structure of + [`SampleIndex`][common.sample_index.SampleIndex] objects based on the + information found in the file. + + Args: + path: The directory path where the sample index file is located. + + Returns: + (common.sample_index.SampleIndex): The root + [`SampleIndex`][common.sample_index.SampleIndex] object representing + the hierarchy read from the file. """ children = {} min_sample = MAX_SAMPLE diff --git a/merlin/common/security/__init__.py b/merlin/common/security/__init__.py index 57477ea1f..759d8bf2d 100644 --- a/merlin/common/security/__init__.py +++ b/merlin/common/security/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,3 +27,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `security` package contains functionality for managing encryption within Merlin, +ensuring secure communication and data protection. + +Modules: + encrypt.py: Handles general encryption logic. + encrypt_backend_traffic.py: Provides functions for encrypting backend traffic. +""" diff --git a/merlin/common/security/encrypt.py b/merlin/common/security/encrypt.py index b1932cd28..5777d20c2 100644 --- a/merlin/common/security/encrypt.py +++ b/merlin/common/security/encrypt.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -45,22 +45,37 @@ LOG = logging.getLogger(__name__) -def _get_key_path(): - """Loads the redis encryption key path from the file described in config.""" +def _get_key_path() -> str: + """ + Loads the path to the Redis encryption key from the configuration. + + If the path is not specified in the configuration, it defaults to + "~/.merlin/encrypt_data_key". + + Returns: + A string representing the absolute path to the encryption key file. + + Raises: + ValueError: If there is an issue retrieving the key path from the configuration. + """ try: key_filepath = CONFIG.results_backend.encryption_key except AttributeError: key_filepath = "~/.merlin/encrypt_data_key" - try: - key_filepath = os.path.abspath(os.path.expanduser(key_filepath)) - except KeyError as e: - raise ValueError("Error! No password provided for RabbitMQ") from e - return key_filepath + if key_filepath is None: + raise ValueError("Error! No password provided for RabbitMQ") + + return os.path.abspath(os.path.expanduser(key_filepath)) -def _gen_key(key_path): - """generates an encryption key and writes it to the given key_path""" +def _gen_key(key_path: str): + """ + Generates a new encryption key and writes it to the specified key path. + + Args: + key_path: The path where the encryption key will be stored. + """ key = Fernet.generate_key() parent_dir = os.path.dirname(os.path.normpath(key_path)) if not os.path.isdir(parent_dir): @@ -69,9 +84,20 @@ def _gen_key(key_path): f.write(key) -def _get_key(): - """get a valid encryption key. Loads from CONFIG.results_backend.encryption_key if possible, - initializes that key if it does not yet exist.""" +def _get_key() -> bytes: + """ + Retrieves a valid encryption key. + + This function attempts to load the key from the path specified in the + configuration. If the key does not exist, it generates a new key and + saves it to the specified path. + + Returns: + The encryption key in bytes format. + + Raises: + IOError: If there is an issue reading the key file or generating a new key. + """ key_path = _get_key_path() try: with open(key_path, "rb") as f: @@ -88,9 +114,15 @@ def _get_key(): return key -def encrypt(payload): +def encrypt(payload: bytes) -> bytes: """ - TODO + Encrypts the given payload using a Fernet key. + + Args: + payload: The data to be encrypted. Must be in bytes format. + + Returns: + The encrypted data in bytes format. """ key = _get_key() f = Fernet(key) @@ -98,9 +130,15 @@ def encrypt(payload): return f.encrypt(payload) -def decrypt(payload): +def decrypt(payload: bytes) -> bytes: """ - TODO + Decrypts the given payload using a Fernet key. + + Args: + payload: The encrypted data to be decrypted. Must be in bytes format. + + Returns: + The decrypted data in bytes format. """ key = _get_key() f = Fernet(key) @@ -110,7 +148,9 @@ def decrypt(payload): def init_key(): """ - Initialize the key to disk on import to prevent race conditions later on, or at least drastically reduce - the number of corner cases where they could appear. + Initializes the Fernet key and stores it on disk. + + This function is called on import to prevent race conditions later on, + or at least drastically reduce the number of corner cases where they could appear. """ Fernet(_get_key()) diff --git a/merlin/common/security/encrypt_backend_traffic.py b/merlin/common/security/encrypt_backend_traffic.py index d597f084b..7542c17a8 100644 --- a/merlin/common/security/encrypt_backend_traffic.py +++ b/merlin/common/security/encrypt_backend_traffic.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -31,7 +31,10 @@ """ Functions for encrypting backend traffic. """ +from typing import Any + import celery.backends.base +from celery.backends.base import Backend from merlin.common.security import encrypt @@ -44,25 +47,51 @@ old_decode = celery.backends.base.Backend.decode -def _encrypt_encode(*args, **kwargs): +def _encrypt_encode(*args, **kwargs) -> bytes: """ - Intercept all celery.backends.Backend.encode calls and encrypt them after - encoding + Intercepts calls to the encode method of the Celery backend and encrypts + the encoded data. + + This function wraps the original encode method, encrypting the result + after encoding. + + Args: + *args: Positional arguments passed to the original encode method. + **kwargs: Keyword arguments passed to the original encode method. + + Returns: + The encrypted encoded data in bytes format. """ return encrypt.encrypt(old_encode(*args, **kwargs)) -def _decrypt_decode(self, payload): +def _decrypt_decode(self: Backend, payload: bytes) -> Any: """ - Intercept all celery.backends.Backend.decode calls and decrypt them before - decoding. + Intercepts calls to the decode method of the Celery backend and decrypts + the payload before decoding. + + This function wraps the original decode method, decrypting the payload + prior to decoding. + + Args: + self: The instance of the backend from which the decode method is called. + payload: The encrypted data to be decrypted. + + Returns: + The decoded data after decryption. Can be any format. """ return old_decode(self, encrypt.decrypt(payload)) def set_backend_funcs(): """ - Set the encode / decode to our own encrypt_encode / encrypt_decode. + Sets the encode and decode methods of the Celery backend to custom + implementations that handle encryption and decryption. + + This function replaces the default encode and decode methods with + `_encrypt_encode` and `_decrypt_decode`, respectively, ensuring that + all data processed by the Celery backend is encrypted and decrypted + appropriately. """ celery.backends.base.Backend.encode = _encrypt_encode celery.backends.base.Backend.decode = _decrypt_decode diff --git a/merlin/common/tasks.py b/merlin/common/tasks.py index 143d3bf12..56b28165a 100644 --- a/merlin/common/tasks.py +++ b/merlin/common/tasks.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -28,29 +28,39 @@ # SOFTWARE. ############################################################################### -"""Test tasks.""" +""" +This module contains Celery task definitions. + +The purpose of this module is to convert the Directed Acyclic Graph +([`DAG`][study.dag.DAG]) provided by Maestro into smaller tasks that +Celery can manage. +""" from __future__ import absolute_import, unicode_literals import json import logging import os -from typing import Any, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional # Need to disable an overwrite warning here since celery has an exception that we need that directly # overwrites a python built-in exception -from celery import chain, chord, group, shared_task, signature +from celery import Signature, Task, chain, chord, group, shared_task, signature from celery.exceptions import MaxRetriesExceededError, OperationalError, TimeoutError # pylint: disable=W0622 +from celery.result import AsyncResult from filelock import FileLock, Timeout from redis.exceptions import TimeoutError as RedisTimeoutError -from merlin.common.abstracts.enums import ReturnCode -from merlin.common.sample_index import uniform_directories +from merlin.common.enums import ReturnCode +from merlin.common.sample_index import SampleIndex, uniform_directories from merlin.common.sample_index_factory import create_hierarchy from merlin.config.utils import Priority, get_priority from merlin.exceptions import HardFailException, InvalidChainException, RestartException, RetryException from merlin.router import stop_workers from merlin.spec.expansion import parameter_substitutions_for_cmd, parameter_substitutions_for_sample +from merlin.study.dag import DAG from merlin.study.status import read_status, status_conflict_handler +from merlin.study.step import Step +from merlin.study.study import MerlinStudy from merlin.utils import dict_deep_merge @@ -84,20 +94,42 @@ retry_backoff=True, priority=get_priority(Priority.HIGH), ) -def merlin_step(self, *args: Any, **kwargs: Any) -> Optional[ReturnCode]: # noqa: C901 pylint: disable=R0912,R0915 +def merlin_step(self: Task, *args: Any, **kwargs: Any) -> ReturnCode: # noqa: C901 pylint: disable=R0912,R0915 """ - Executes a Merlin Step - :param args: The arguments, one of which should be an instance of Step - :param kwargs: The optional keyword arguments that describe adapter_config and - the next step in the chain, if there is one. - - Example kwargs dict: - {"adapter_config": {'type':'local'}, - "next_in_chain": } # merlin_step will be added to the current chord - # with next_in_chain as an argument + Executes a Merlin step. + + This task executes a step in the Merlin workflow, handling various + outcomes such as success, retries, and failures. It can also manage + chaining to the next step in the workflow. + + Notes: + - If the step has already been completed, it will be skipped. + + Args: + self: The current task instance. + *args: Positional arguments, one of which should be an instance + of [`Step`][study.step.Step]. + **kwargs: Optional keyword arguments that include:\n + - adapter_config (`Dict`): Configuration for the adapter, + defaulting to `{'type': 'local'}`. + - next_in_chain ([`Step`][study.step.Step]): The next step in + the workflow chain, if applicable.\n + Example kwargs dict where `merlin_step` will be added to the + current chord with `next_in_chain` as an argument:\n + ``` + { + "adapter_config": { + 'type': 'local' + }, + "next_in_chain": + } + ``` + + Returns: + (common.enums.ReturnCode): The result of the step + execution, which can indicate success, various failure modes, + or a request to retry. """ - from merlin.study.step import Step # pylint: disable=C0415 - step: Optional[Step] = None LOG.debug(f"args is {len(args)} long") @@ -204,15 +236,29 @@ def merlin_step(self, *args: Any, **kwargs: Any) -> Optional[ReturnCode]: # noq return None -def is_chain_expandable(chain_, labels): +def is_chain_expandable(chain_: List[Step], labels: List[str]) -> bool: """ - Returns whether to expand the steps in the given chain. - A chain_ is expandable if all the steps are expandable. - It is not expandable if none of the steps are expandable. - If neither expandable nor not expandable, we raise an InvalidChainException. - :param chain_: A list of Step objects representing chain of dependent steps. - :param labels: The labels - + Determine if the steps in the given chain are expandable. + + A chain is considered expandable if all steps within the chain require + expansion. Conversely, if none of the steps require expansion, the chain + is not expandable. If there is a mix of steps that require expansion and + those that do not, an `InvalidChainException` is raised, indicating that + the chain is incompatible. + + Args: + chain_ (List[study.step.Step]): A list of [`Step`][study.step.Step] + objects representing a chain of dependent steps. + labels: The labels associated with the steps in the chain, used to + determine if expansion is needed. + + Returns: + True if all steps in the chain are expandable, False if none are + expandable. + + Raises: + InvalidChainException: If there is a mix of steps that require + expansion and those that do not, indicating an incompatible chain. """ array_of_bools = [step.check_if_expansion_needed(labels) for step in chain_] @@ -235,11 +281,19 @@ def is_chain_expandable(chain_, labels): return needs_expansion -def prepare_chain_workspace(sample_index, chain_): +def prepare_chain_workspace(sample_index: SampleIndex, chain_: List[Step]): """ - Prepares a user's workspace for each step in the given chain. - :param chain_: A list of Step objects representing chain of dependent steps. - :param labels: The labels + Prepares a user's workspace for each step in the given chain of dependent steps. + + This function iterates through a list of [`Step`][study.step.Step] objects and + prepares the necessary workspace for each step by creating directories and writing + sample index files. + + Args: + sample_index (common.sample_index.SampleIndex): An object that manages sample + indexing and workspace preparation. + chain_ (List[study.step.Step]): A list of [`Step`][study.step.Step] objects + representing a chain of dependent steps. Each step's workspace will be prepared. """ # TODO: figure out faster way to create these directories (probably using # yet another task) @@ -261,25 +315,36 @@ def prepare_chain_workspace(sample_index, chain_): priority=get_priority(Priority.LOW), ) def add_merlin_expanded_chain_to_chord( # pylint: disable=R0913,R0914 - self, - task_type, - chain_, - samples, - labels, - sample_index, - adapter_config, - min_sample_id, + self: Task, + task_type: Signature, + chain_: List[Step], + samples: List[Any], + labels: List[str], + sample_index: SampleIndex, + adapter_config: Dict, + min_sample_id: int, ): """ - Expands tasks in a chain, then adds the expanded tasks to the current chord. - :param self: The current task. - :param task_type: The celery task signature type the new tasks should be. - :param chain_: The list of tasks to expand. - :param samples: The sample values to use for each new task. - :param labels: The sample labels. - :param sample_index: The sample index that contains the directory structure for tasks. - :param adapter_config: The adapter config. - :param min_sample_id: offset to use for the sample_index. + Expand tasks in a chain and add the expanded tasks to the current chord. + + This Celery task recursively expands a chain of tasks based on provided + sample values and their corresponding labels. The expanded tasks are + configured with specific parameters and added to the current chord for + execution. The function handles both the expansion of tasks and the + management of task dependencies. + + Args: + self: The current task instance. + task_type: The Celery task signature type for the new tasks to be + created. + chain_ (List[study.step.Step]): A list of tasks to expand into a chain. + samples: The sample values to use for each new task. + labels: The sample labels corresponding to the samples. + sample_index (common.sample_index.SampleIndex): The sample index that + contains the directory structure for tasks. + adapter_config: Configuration settings for the adapter used in task + execution. + min_sample_id: An offset to use for the sample index. """ num_samples = len(samples) # Use the index to get a path to each sample @@ -368,13 +433,27 @@ def add_merlin_expanded_chain_to_chord( # pylint: disable=R0913,R0914 return ReturnCode.OK -def add_simple_chain_to_chord(self, task_type, chain_, adapter_config): +def add_simple_chain_to_chord(self: Task, task_type: Signature, chain_: List[Step], adapter_config: Dict): """ - Adds a chain of tasks to the current chord. - :param self: The current task. - :param task_type: The celery task signature type the new tasks should be. - :param chain_: The list of tasks to expand. - :param adapter_config: The adapter config. + Add a chain of tasks to the current chord for execution. + + This function takes a list of tasks, modifies their signatures based on + provided parameters, and adds them to the current chord. Each task in the + chain is transformed into a new task signature with specific configurations + such as queue and task ID. + + This function takes a list of steps and creates signatures based on the + parameters they provide, such as queue and workspace. It then adds these + signatures to the current chord for later execution. + + Args: + self: The current task instance invoking this method. + task_type: The Celery task signature type that the new tasks should be + based on. + chain_ (List[study.step.Step]): A list of tasks to expand into a chain. + Each task should provide necessary parameters for signature creation. + adapter_config: Configuration settings for the adapter used in task + execution. """ LOG.debug(f"simple chain with {chain_}") all_chains = [] @@ -394,18 +473,20 @@ def add_simple_chain_to_chord(self, task_type, chain_, adapter_config): launch_chain(self, chain_1d) -def launch_chain(self: "Task", chain_1d: List["Signature"], condense_sig: "Signature" = None): # noqa: F821 +def launch_chain(self: Task, chain_1d: List[Signature], condense_sig: Signature = None): """ - Given a 1D chain, appropriately launch the signatures it contains. - If this is a local run, launch the signatures instantly. - Otherwise, there's two cases: - a. The chain is dealing with samples (i.e. we'll need to condense status files) - so create a new chord and add it to the current chord - b. The chain is NOT dealing with samples so we can just add the signatures to the current chord - - :param `self`: The current task - :param `chain_1d`: A 1-dimensional list of signatures to launch - :param `condense_sig`: A signature for condensing the status files. None if condensing isn't needed. + Launch a 1D chain of task signatures appropriately based on the execution context. + + This function handles the launching of a list of task signatures in a + one-dimensional chain. The behavior varies depending on whether the + execution is local or remote, and whether the tasks involve sample + processing that requires condensing status files. + + Args: + self: The current task instance invoking this method. + chain_1d: A one-dimensional list of task signatures to be launched. + condense_sig: A signature for condensing the status files after task execution. + If None, condensing is not required. """ # If there's nothing in the chain then we won't have to launch anything so check that first if chain_1d: @@ -426,11 +507,28 @@ def launch_chain(self: "Task", chain_1d: List["Signature"], condense_sig: "Signa self.add_to_chord(sig, lazy=False) -def get_1d_chain(all_chains: List[List["Signature"]]) -> List["Signature"]: # noqa: F821 +def get_1d_chain(all_chains: List[List[Signature]]) -> List[Signature]: """ - Convert a 2D list of chains into a 1D list. - :param all_chains: Two-dimensional list of chains [chain_length][number_of_chains] - :returns: A one-dimensional list representing a chain of tasks + Convert a 2D list of task chains into a 1D list of task signatures. + + This function takes a two-dimensional list of task signatures, where each + inner list represents a parallel group of tasks. It transforms this structure + into a one-dimensional list suitable for creating a linear chain of tasks. + If there is only one chain, it returns that chain directly. If there are + multiple chains, it sets up dependencies between tasks to ensure proper + execution order. + + Notes: + - The function processes the chains in reverse order to correctly + set up the dependencies before adding them to the final list. + + Args: + all_chains: A two-dimensional list of task signatures, where each inner + list represents a group of tasks that can be executed in parallel. + + Returns: + A one-dimensional list of task signatures representing a chain of tasks, + with dependencies set up for proper execution order. """ chain_steps = [] if len(all_chains) == 1: @@ -465,17 +563,37 @@ def get_1d_chain(all_chains: List[List["Signature"]]) -> List["Signature"]: # n return chain_steps -def gather_statuses( - sample_index: "SampleIndex", workspace: str, condensed_workspace: str, files_to_remove: List[str] # noqa: F821 -) -> Dict: +def gather_statuses(sample_index: SampleIndex, workspace: str, condensed_workspace: str, files_to_remove: List[str]) -> Dict: """ - Traverse the sample index and gather all of the statuses into one. - - :param `sample_index`: A SampleIndex object to track this specific sample hierarchy - :param `workspace`: The full workspace path to the step we're condensing for - :param `condensed_workspace`: A shortened version of `workspace` that's saved in the status files - :param `files_to_remove`: An empty list that we'll add filepaths to that need removed - :returns: A dict of condensed statuses + Traverse the sample index and gather all statuses into a single dictionary. + + This function iterates through the provided + [`SampleIndex`][common.sample_index.SampleIndex] object, + reading status files from each sample's workspace. It condenses + the statuses into a single dictionary while tracking which files + need to be removed after condensing. The function ensures that + only completed statuses are included in the condensed output. + + Args: + sample_index (common.sample_index.SampleIndex): A + [`SampleIndex`][common.sample_index.SampleIndex] object + representing the specific sample hierarchy to traverse. + workspace: The full path to the workspace for the step being + condensed. + condensed_workspace: A shortened version of the workspace + path that will be used in the status files. + files_to_remove: A list that will be populated with file paths + of status files that need to be removed after condensing. + + Returns: + A dictionary containing the condensed statuses gathered + from the status files. + + Raises: + TimeoutError: If a timeout occurs while reading a status file, + triggering a restart of the task. + FileNotFoundError: If a status file is not found during the + condensing process. """ LOG.info(f"Gathering statuses to condense for '{condensed_workspace}'") condensed_statuses = {} @@ -520,15 +638,38 @@ def gather_statuses( retry_backoff=True, priority=get_priority(Priority.LOW), ) -def condense_status_files(self, *args: Any, **kwargs: Any) -> ReturnCode: # pylint: disable=R0914,W0613 +def condense_status_files(self: Task, *args: Any, **kwargs: Any) -> ReturnCode: # pylint: disable=R0914,W0613 """ - After a section of the sample tree has finished, condense the status files. - - kwargs should look like so: - kwargs = { - "sample_index": SampleIndex Object, - "workspace": str representing the step's workspace - } + Condenses status files after a section of the sample tree has completed processing. + + This task gathers status information from a specified + [`SampleIndex`][common.sample_index.SampleIndex] and condenses it into a single + JSON file. It handles potential race conditions by using a file lock during + the write operation. If the condensed status file already exists, it merges + the new statuses with the existing ones. + + Notes: + - The task will remove the original status files after condensing them + into the JSON file. + + Args: + self: The current task instance. + *args: Additional positional arguments (not used in this task). + **kwargs: Keyword arguments containing:\n + - `sample_index` ([`SampleIndex`][common.sample_index.SampleIndex]): + The [`SampleIndex`][common.sample_index.SampleIndex] object used + for gathering statuses. + - `workspace` (str): The workspace path for the step. + - `condensed_workspace` (str): The workspace path for the + condensed status. + + Returns: + (common.enums.ReturnCode): A [`ReturnCode.OK`][common.enums.ReturnCode] + message if the operation was successful. None, otherwise. + + Raises: + TimeoutError: If the file lock cannot be acquired within the + specified timeout period, which triggers a task restart. """ # Get the sample index object that we'll use for condensing sample_index = kwargs.pop("sample_index", None) @@ -592,26 +733,40 @@ def condense_status_files(self, *args: Any, **kwargs: Any) -> ReturnCode: # pyl priority=get_priority(Priority.LOW), ) def expand_tasks_with_samples( # pylint: disable=R0913,R0914 - self, - dag, - chain_, - samples, - labels, - task_type, - adapter_config, - level_max_dirs, + self: Task, + dag: DAG, + chain_: List[str], + samples: List[List[str]], + labels: List[str], + task_type: Callable, + adapter_config: Dict, + level_max_dirs: int, ): """ - Generate a group of celery chains of tasks from a chain of task names, using merlin - samples and labels to do variable substitution. - - :param dag : A Merlin DAG. - :param chain_ : The list of task names to expand into a celery group of celery chains. - :param samples : The list of lists of merlin sample values to do substitution for. - :labels : A list of strings containing the label associated with each column in the samples. - :task_type : The celery task type to create. Currently always merlin_step. - :adapter_config : A dictionary used for configuring maestro script adapters. - :level_max_dirs : The max number of directories per level in the sample hierarchy. + Expands a chain of task names into a group of Celery chains, using samples + and labels for variable substitution. + + This task determines whether the provided chain of tasks requires + expansion based on the structure of the Directed Acyclic Graph ([`DAG`][study.dag.DAG]), + samples, and labels. If expansion is needed, it generates and queues new tasks + for each range of samples. Otherwise, it queues a simple chain task. + + Args: + self: The current task instance. + dag (study.dag.DAG): A Merlin Directed Acyclic Graph + ([`DAG`][study.dag.DAG]) representing the workflow. + chain_: A list of task names to be expanded into a + Celery group of chains. + samples: A list of lists containing Merlin sample values for + variable substitution. + labels: A list of strings representing the labels associated + with each column in the samples. + task_type: The Celery task type to create, currently expected + to be [`merlin_step`][common.tasks.merlin_step]. + adapter_config: A configuration dictionary for Maestro + script adapters. + level_max_dirs: The maximum number of directories allowed per + level in the sample hierarchy. """ LOG.debug(f"expand_tasks_with_samples called with chain,{chain_}\n") # Figure out how many directories there are, make a glob string @@ -703,15 +858,19 @@ def expand_tasks_with_samples( # pylint: disable=R0913,R0914 name="merlin:shutdown_workers", priority=get_priority(Priority.HIGH), ) -def shutdown_workers(self, shutdown_queues): # pylint: disable=W0613 +def shutdown_workers(self: Task, shutdown_queues: List[str]): # pylint: disable=W0613 """ - This task issues a call to shutdown workers. + Initiates the shutdown of Celery workers. - It wraps the stop_celery_workers call as a task. - It is acknolwedged right away, so that it will not be requeued when - executed by a worker. + This task wraps the [`stop_celery_workers`][study.celeryadapter.stop_celery_workers] + function, allowing for the graceful shutdown of specified Celery worker queues. It is + acknowledged immediately upon execution, ensuring that it will not be requeued, even + if executed by a worker. - :param: shutdown_queues: The specific queues to shutdown (list) + Args: + self: The current task instance. + shutdown_queues: A list of specific queues to shut down. If None, all queues will + be shut down. """ if shutdown_queues is not None: LOG.warning(f"Shutting down workers in queues {shutdown_queues}!") @@ -728,13 +887,27 @@ def shutdown_workers(self, shutdown_queues): # pylint: disable=W0613 name="merlin:chordfinisher", priority=get_priority(Priority.LOW), ) -def chordfinisher(*args, **kwargs): # pylint: disable=W0613 - """. - It turns out that chain(group,group) in celery does not execute one group - after another, but executes the groups as if they were independent from - one another. To get a sync point between groups, we use this method as a - callback to enforce sync points for chords so we can declare chains of groups - dynamically. +def chordfinisher(*args: List, **kwargs: Dict) -> str: # pylint: disable=W0613 + """ + Synchronization callback for Celery chords. + + This function serves as a synchronization point between groups of tasks + in a Celery workflow. In Celery, using `chain(group, group)` does not + guarantee that the second group will execute only after the first group + has completed. Instead, both groups are executed independently. + + To enforce a synchronization point between these groups, this function + is used as a callback in a chord. It allows for the declaration of chains + of groups dynamically, ensuring that subsequent tasks wait for the + completion of all tasks in the preceding groups. + + Args: + *args: Variable length argument list. Needed by Celery. + **kwargs: Arbitrary keyword arguments. Needed by Celery. + + Returns: + A constant string "SYNC" indicating the synchronization point + has been reached. """ return "SYNC" @@ -745,9 +918,26 @@ def chordfinisher(*args, **kwargs): # pylint: disable=W0613 name="merlin:queue_merlin_study", priority=get_priority(Priority.LOW), ) -def queue_merlin_study(study, adapter): +def queue_merlin_study(study: MerlinStudy, adapter: Dict) -> AsyncResult: """ - Launch a chain of tasks based off of a MerlinStudy. + Launch a chain of tasks based on a MerlinStudy. + + This Celery task initiates a series of tasks derived from a + [`MerlinStudy`][study.study.MerlinStudy] object. It processes + the study's Directed Acyclic Graph ([`DAG`][study.dag.DAG]) + to group tasks and convert them into a chain of Celery tasks + for execution. + + Args: + study: The study object containing samples, sample labels, + and the Directed Acyclic Graph ([`DAG`][study.dag.DAG]) + structure that defines the task dependencies. + adapter: An adapter object used to facilitate interactions with + the study's data or processing logic. + + Returns: + An instance representing the asynchronous result of the task chain, + allowing for tracking and management of the task's execution. """ samples = study.samples sample_labels = study.sample_labels diff --git a/merlin/common/util_sampling.py b/merlin/common/util_sampling.py index 0a6c585cf..c8c273b72 100644 --- a/merlin/common/util_sampling.py +++ b/merlin/common/util_sampling.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -31,39 +31,53 @@ """ Utility functions for sampling. """ +from typing import List, Tuple, Union import numpy as np -def scale_samples(samples_norm, limits, limits_norm=(0, 1), do_log=False): - """Scale samples to new limits, either log10 or linearly. +# TODO should we move this to merlin-spellbook? +def scale_samples( + samples_norm: np.ndarray, + limits: List[Tuple[int, int]], + limits_norm: Tuple[int, int] = (0, 1), + do_log: Union[bool, List[bool]] = False, +) -> np.ndarray: + """ + Scale samples to new limits, either logarithmically or linearly. + + This function transforms normalized samples to specified limits, + allowing for both linear and logarithmic scaling based on the + provided parameters. Args: - samples_norm (ndarray): The normalized samples to scale, - with dimensions (nsamples,ndims). - limits (list of tuples): A list of (min, max) for the various - dimensions. Length of list is ndims. - limits_norm (tuple of floats, optional): The (min, max) from which - samples_norm were drawn. Defaults to (0,1). - do_log (boolean or list of booleans, optional): Whether - to log10 scale each dimension. Either a single boolean or - a list of length ndims, for each dimension. - Defaults to ndims*[False]. + samples_norm: The normalized samples to scale, with dimensions + (nsamples, ndims). + limits: A list of (min, max) tuples for the various dimensions. + The length of the list must match the number of dimensions (ndims). + limits_norm: The (min, max) values from which `samples_norm` were + derived. Defaults to (0, 1). + do_log: Indicates whether to apply log10 scaling to each dimension. + This can be a single boolean or a list of length ndims. Defaults + to a list of `ndims` containing `False`. Returns: - ndarray: The scaled samples. + The scaled samples, with the same shape as `samples_norm`. - Note: - We follow the sklearn convention of requiring samples to be - given as an (nsamples, ndims) array. + Raises: + ValueError: If `samples_norm` does not have two dimensions. - To transform 1-D arrays: - - >>> samples = samples.reshape((-1,1)) # ndims = 1 - >>> samples = samples.reshape((1,-1)) # nsamples = 1 + Notes: + - The function follows the sklearn convention, requiring + samples to be provided as an (nsamples, ndims) array. + - To transform 1-D arrays, reshape them accordingly: + ```python + >>> samples = samples.reshape((-1, 1)) # ndims = 1 + >>> samples = samples.reshape((1, -1)) # nsamples = 1 + ``` Example: - + ```python >>> # Turn 0:1 samples into -1:1 >>> import numpy as np >>> norm_values = np.linspace(0,1,5).reshape((-1,1)) @@ -82,6 +96,7 @@ def scale_samples(samples_norm, limits, limits_norm=(0, 1), do_log=False): [ 1.00000000e+02] [ 1.00000000e+03] [ 1.00000000e+04]] + ``` """ norms = np.asarray(samples_norm) if len(norms.shape) != 2: diff --git a/merlin/config/__init__.py b/merlin/config/__init__.py index 41645e249..98e03bccb 100644 --- a/merlin/config/__init__.py +++ b/merlin/config/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -30,8 +30,21 @@ """ Used to store the application configuration. -""" +The `config` package provides functionality for managing and configuring various aspects +of the Merlin application, including broker settings, results backends, Celery configurations, +and application-level settings. It serves as the central hub for loading, processing, and +utilizing configuration data defined in the `app.yaml` file and other related resources. + +Modules: + broker.py: Manages broker configurations and connection strings for messaging systems. + celeryconfig.py: Contains default Celery configuration settings for Merlin. + configfile.py: Handles the loading and processing of application configuration files + and SSL-related settings. + results_backend.py: Configures connection strings and SSL settings for results backends. + utils.py: Provides utilities for broker priority handling and validation. +""" +from copy import copy from types import SimpleNamespace from typing import Dict, List, Optional @@ -46,9 +59,29 @@ class Config: # pylint: disable=R0903 The Config class, meant to store all Merlin config settings in one place. Regardless of the config data loading method, this class is meant to standardize config data retrieval throughout all parts of Merlin. + + Attributes: + celery (Optional[SimpleNamespace]): A namespace containing Celery configuration settings. + broker (Optional[SimpleNamespace]): A namespace containing broker configuration settings. + results_backend (Optional[SimpleNamespace]): A namespace containing results backend configuration settings. + + Methods: + __copy__: Creates a shallow copy of the Config instance. + __str__: Returns a formatted string representation of the Config instance. + load_app_into_namespaces: Converts the provided configuration dictionary into namespaces + and assigns them to the Config instance's attributes. """ - def __init__(self, app_dict): + def __init__(self, app_dict: Dict): + """ + Initializes the Config instance with configuration data from a dictionary. + + Args: + app_dict: A dictionary containing configuration data for the application. + The dictionary may include keys such as "celery", "broker", and "results_backend", + each of which is converted into a `SimpleNamespace` and assigned to the corresponding + attribute of the Config instance. + """ # I think this ends up a SimpleNamespace from load_app_into_namespaces, but it seems like it should be typed as # the app var in celery.py, as celery.app.base.Celery self.celery: Optional[SimpleNamespace] @@ -56,9 +89,48 @@ def __init__(self, app_dict): self.results_backend: Optional[SimpleNamespace] self.load_app_into_namespaces(app_dict) - def load_app_into_namespaces(self, app_dict: Dict) -> None: + def __copy__(self) -> "Config": """ - Makes the application dictionary into a namespace, sets the attributes of the Config from the namespace values. + Creates a shallow copy of the Config instance. + + Returns: + A new Config instance with copied `celery`, `broker`, and `results_backend` attributes. + """ + cls = self.__class__ + result = cls.__new__(cls) + copied_attrs = { + "celery": copy(self.__dict__["celery"]), + "broker": copy(self.__dict__["broker"]), + "results_backend": copy(self.__dict__["results_backend"]), + } + result.__dict__.update(copied_attrs) + return result + + def __str__(self) -> str: + """ + Returns a formatted string representation of the Config instance. + + Returns: + str: A string containing the values of the `celery`, `broker`, and `results_backend` attributes. + """ + formatted_str = "config:" + attrs = {"celery": self.celery, "broker": self.broker, "results_backend": self.results_backend} + for name, attr in attrs.items(): + if attr is not None: + items = (f" {k}: {v!r}" for k, v in attr.__dict__.items()) + joined_items = "\n".join(items) + formatted_str += f"\n {name}:\n{joined_items}" + else: + formatted_str += f"\n {name}:\n None" + return formatted_str + + def load_app_into_namespaces(self, app_dict: Dict): + """ + Converts the provided application dictionary into namespaces and assigns them + to the Config instance's attributes. + + Args: + app_dict: A dictionary containing configuration data for the application. """ fields: List[str] = ["celery", "broker", "results_backend"] for field in fields: diff --git a/merlin/config/broker.py b/merlin/config/broker.py index dc8131c28..065ea415f 100644 --- a/merlin/config/broker.py +++ b/merlin/config/broker.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -28,7 +28,15 @@ # SOFTWARE. ############################################################################### -"""Logic for configuring the celery broker.""" +""" +This module provides utility functions and constants to manage broker configurations and connection strings +for various messaging systems, including RabbitMQ and Redis. It supports multiple connection protocols +and configurations, such as SSL, Unix sockets, and password inclusion. + +The module defines constants for supported brokers and connection string templates, along with functions +to construct and retrieve connection strings and SSL configurations based on settings defined in the +`app.yaml` configuration file. +""" from __future__ import print_function import getpass @@ -37,16 +45,11 @@ import ssl from os.path import expanduser from typing import Dict, List, Optional, Union +from urllib.parse import quote from merlin.config.configfile import CONFIG, get_ssl_entries -try: - from urllib import quote -except ImportError: - from urllib.parse import quote - - LOG: logging.Logger = logging.getLogger(__name__) BROKERS: List[str] = ["rabbitmq", "redis", "rediss", "redis+socket", "amqps", "amqp"] @@ -56,19 +59,41 @@ USER = getpass.getuser() -def read_file(filepath): - "Safe file read from filepath" +def read_file(filepath: str) -> str: + """ + Safely reads the first line from a file and returns it with special characters URL-encoded. + + Args: + filepath (str): The path to the file to be read. + + Returns: + The first line of the file, stripped of leading and trailing whitespace, + with special characters URL-encoded. + """ with open(filepath, "r") as f: # pylint: disable=C0103 line = f.readline().strip() return quote(line, safe="") -def get_rabbit_connection(include_password, conn="amqps"): +def get_rabbit_connection(include_password: bool, conn: str = "amqps") -> str: """ - Given the path to the directory where the broker configurations are stored - setup and return the RabbitMQ connection string. + Constructs and returns a RabbitMQ connection string based on broker configurations. - :param include_password : Format the connection for ouput by setting this True + This function reads broker configurations (such as server, port, username, password, and vhost) + and formats them into a RabbitMQ connection string. Optionally, the password can be included + in the connection string if `include_password` is set to `True`. + + Args: + include_password (bool): Whether to include the password in the connection string. + conn (str, optional): The connection protocol to use. Defaults to "amqps". + Supported values are "amqp" and "amqps". + + Returns: + A formatted RabbitMQ connection string. + + Raises: + ValueError: If the password file path is not provided in the broker configuration, or if + the password file does not exist or cannot be read. """ LOG.debug(f"Broker: connection = {conn}") @@ -85,13 +110,13 @@ def get_rabbit_connection(include_password, conn="amqps"): password_filepath = CONFIG.broker.password LOG.debug(f"Broker: password filepath = {password_filepath}") password_filepath = os.path.abspath(expanduser(password_filepath)) - except KeyError as e: # pylint: disable=C0103 - raise ValueError("Broker: No password provided for RabbitMQ") from e + except (AttributeError, KeyError) as exc: + raise ValueError("Broker: No password provided for RabbitMQ") from exc try: password = read_file(password_filepath) - except IOError as e: # pylint: disable=C0103 - raise ValueError(f"Broker: RabbitMQ password file {password_filepath} does not exist") from e + except IOError as exc: + raise ValueError(f"Broker: RabbitMQ password file {password_filepath} does not exist") from exc try: port = CONFIG.broker.port @@ -119,10 +144,17 @@ def get_rabbit_connection(include_password, conn="amqps"): return RABBITMQ_CONNECTION.format(**rabbitmq_config) -def get_redissock_connection(): +def get_redissock_connection() -> str: """ - Given the path to the directory where the broker configurations are stored - setup and return the redis+socket connection string. + Constructs and returns a Redis connection string using a Unix socket. + + This function retrieves broker configurations, such as the database number (`db_num`) and + the Unix socket file path (`path`), and formats them into a Redis connection string. + + If the database number is not specified in the configuration, it defaults to `0`. + + Returns: + A formatted Redis connection string using a Unix socket. """ try: db_num = CONFIG.broker.db_num @@ -137,12 +169,20 @@ def get_redissock_connection(): # flake8 complains this function is too complex, we don't gain much nesting any of this as a separate function, # however, cyclomatic complexity examination is off to get around this -def get_redis_connection(include_password, use_ssl=False): # noqa C901 +def get_redis_connection(include_password: bool, use_ssl: bool = False) -> str: # noqa C901 """ - Return the redis or rediss specific connection + Constructs and returns a Redis connection string, optionally using SSL and including a password. - :param include_password : Format the connection for ouput by setting this True - :param use_ssl : Flag to use rediss output + This function retrieves broker configurations (such as server, port, username, password, and database number) + and formats them into a Redis connection string. The connection can be configured to use SSL (`rediss` protocol) + and optionally include the password in the connection string. + + Args: + include_password (bool): Whether to include the password in the connection string. + use_ssl (bool, optional): Whether to use the `rediss` protocol (SSL). + + Returns: + A formatted Redis connection string. """ server = CONFIG.broker.server LOG.debug(f"Broker: server = {server}") @@ -184,15 +224,22 @@ def get_redis_connection(include_password, use_ssl=False): # noqa C901 return f"{urlbase}://{spass}{server}:{port}/{db_num}" -def get_connection_string(include_password=True): +def get_connection_string(include_password: bool = True) -> str: """ - Return the connection string based on the configuration specified in the - `app.yaml` config file. + Constructs and returns a connection string based on the broker configuration. + + This function retrieves the connection string from the `CONFIG.broker.url` if available. + Otherwise, it determines the connection string based on the broker name specified in the + configuration file (`app.yaml`). If the broker name is not supported, a `ValueError` is raised. + + Args: + include_password (bool): Whether to include the password in the connection string. - If the url variable is present, return that as the connection string. + Returns: + A formatted connection string based on the broker configuration. - :param include_password : The connection can be formatted for output by - setting this to True + Raises: + ValueError: If the broker name is not supported. """ try: return CONFIG.broker.url @@ -205,18 +252,27 @@ def get_connection_string(include_password=True): except AttributeError: broker = "" - try: - config_path = CONFIG.celery.certs - config_path = os.path.abspath(os.path.expanduser(config_path)) - except AttributeError: - config_path = None - if broker not in BROKERS: raise ValueError(f"Error: {broker} is not a supported broker.") return _sort_valid_broker(broker, include_password) -def _sort_valid_broker(broker, include_password): +def _sort_valid_broker(broker: str, include_password: bool) -> str: + """ + Determines and returns the appropriate connection string for a given broker. + + This function selects the connection string generation method based on the broker type + provided as input. Supported brokers include RabbitMQ (`amqp` or `amqps`), Redis (`redis`), + Redis over SSL (`rediss`), and Redis over a socket (`redis+socket`). + + Args: + broker (str): The name of the broker. Must be one of the supported broker types: + `rabbitmq`, `amqps`, `amqp`, `redis+socket`, `redis`, or `rediss`. + include_password (bool): Whether to include the password in the connection string. + + Returns: + A formatted connection string for the specified broker. + """ if broker in ("rabbitmq", "amqps"): return get_rabbit_connection(include_password, conn="amqps") @@ -235,11 +291,17 @@ def _sort_valid_broker(broker, include_password): def get_ssl_config() -> Union[bool, Dict[str, Union[str, ssl.VerifyMode]]]: """ - Return the ssl config based on the configuration specified in the - `app.yaml` config file. + Retrieves the SSL configuration for the broker based on the settings in the `app.yaml` configuration file. + + This function determines whether SSL should be used for the broker connection and, if applicable, + returns the SSL configuration details. If the broker does not require SSL or is unsupported, + the function returns `False`. - :return: Returns either False if no ssl - :rtype: Union[bool, Dict[str, Union[str, ssl.VerifyMode]]] + Returns: + This returns either:\n + - `False` if SSL is not required or the broker is unsupported. + - A dictionary containing SSL configuration details if SSL is required. + The dictionary may include keys such as certificate paths and verification modes. """ broker: Union[bool, str] = "" try: diff --git a/merlin/config/celeryconfig.py b/merlin/config/celeryconfig.py index 5794599dc..ece91e0c3 100644 --- a/merlin/config/celeryconfig.py +++ b/merlin/config/celeryconfig.py @@ -1,7 +1,3 @@ -""" -Default celery configuration for merlin -""" - ############################################################################### # Copyright (c) 2023, Lawrence Livermore National Security, LLC. # Produced at the Lawrence Livermore National Laboratory @@ -10,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -32,6 +28,8 @@ # SOFTWARE. ############################################################################### +"""This module houses the default Celery configuration settings for Merlin.""" + from merlin.log_formatter import FORMATS diff --git a/merlin/config/configfile.py b/merlin/config/configfile.py index 1634b833f..deff360b7 100644 --- a/merlin/config/configfile.py +++ b/merlin/config/configfile.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,8 +29,12 @@ ############################################################################### """ -This module handles the logic for the Merlin config files for setting up all -configurations. +This module provides functionality for managing and loading application configuration files, +default settings, and SSL-related configurations. It includes utilities for locating, +reading, and processing configuration files, as well as handling SSL certificates and protocols +for various server types. + +It houses the `CONFIG` object that's used throughout Merlin's codebase. """ import getpass import logging @@ -51,12 +55,15 @@ MERLIN_HOME: str = os.path.join(USER_HOME, ".merlin") -def load_config(filepath): +def load_config(filepath: str) -> Dict: """ - Given the path to the merlin YAML config file, read the file and return - a dictionary of the contents. + Reads a Merlin YAML configuration file and returns its contents as a dictionary. + + Args: + filepath (str): The path to the YAML configuration file. - :param filepath : Read a yaml file given by filepath + Returns: + A dictionary containing the contents of the YAML file. """ if not os.path.isfile(filepath): LOG.info(f"No app config file at {filepath}") @@ -65,12 +72,19 @@ def load_config(filepath): return load_yaml(filepath) -def find_config_file(path=None): +def find_config_file(path: str = None) -> str: """ - Given a dir path, find and return the path to the merlin application - config file. + Finds and returns the path to the Merlin application configuration file (`app.yaml`). + + If no `path` is provided, the function searches in the current working directory + and the `MERLIN_HOME` directory for the configuration file. If `path` is provided, + it checks for the configuration file in the specified directory. + + Args: + path (str, optional): The directory path to search for the `app.yaml` file. - :param path : The path to search for the app.yaml file + Returns: + The full path to the `app.yaml` file if found. """ if path is None: local_app = os.path.join(os.getcwd(), APP_FILENAME) @@ -89,14 +103,18 @@ def find_config_file(path=None): return None -def load_default_user_names(config): +def set_username_and_vhost(config: Dict): """ - Load broker.username and broker.vhost defaults if they are not present in - the current configuration. Doing this here prevents other areas that rely - on config from needing to know that those fields could not be defined by - the user. + Ensures that `broker.username` and `broker.vhost` default values are set in the + configuration if they are not already defined. - :param config : The namespace config object + This function checks the `config` object for the presence of `broker.username` + and `broker.vhost`. If either is missing, it sets their default values using + the current system username. This prevents other parts of the code from having + to handle missing values for these fields. + + Args: + config (Dict): The configuration object containing the `broker` namespace. """ try: config["broker"]["username"] @@ -112,12 +130,22 @@ def load_default_user_names(config): def get_config(path: Optional[str]) -> Dict: """ - Load a merlin configuration file and return a dictionary of the - configurations. + Loads a Merlin configuration file and returns a dictionary containing the configuration data. + + This function locates the configuration file using the provided `path` or default search locations, + loads the configuration data, and applies default values where necessary. If the configuration file + cannot be found, it raises a `ValueError`. + + Args: + path (str, optional): The directory path to search for the configuration file. + If `None`, default search paths are used. - :param [Optional[str]] path : The path to search for the config file. - :return: the config file to coordinate brokers/results backend/task manager." - :rtype: A Dict with all the config data. + Returns: + A dictionary containing all the configuration data, including broker, + results backend, and task manager settings. + + Raises: + ValueError: If the configuration file cannot be found. """ filepath: Optional[str] = find_config_file(path) @@ -131,8 +159,17 @@ def get_config(path: Optional[str]) -> Dict: return config -def load_default_celery(config): - """Creates the celery default configuration""" +def load_default_celery(config: Dict): + """ + Initializes the default Celery configuration within the provided configuration object. + + This function ensures that the `celery` section of the configuration exists and sets + default values for specific Celery-related settings if they are not already defined. + These defaults include `omit_queue_tag`, `queue_tag`, and `override`. + + Args: + config (Dict): The configuration object where the Celery settings will be initialized. + """ try: config["celery"] except KeyError: @@ -151,23 +188,52 @@ def load_default_celery(config): config["celery"]["override"] = None -def load_defaults(config): - """Loads default configuration values""" - load_default_user_names(config) +def load_defaults(config: Dict): + """ + Loads default configuration values into the provided configuration dictionary. + + This function initializes default values for various configuration sections, + including user-related settings and Celery-specific settings, by calling + `set_username_and_vhost` and `load_default_celery`. + + Args: + config (Dict): The configuration dictionary to be updated with default values. + """ + set_username_and_vhost(config) load_default_celery(config) -def is_debug(): +def is_debug() -> bool: """ - Check for MERLIN_DEBUG in environment to set a debugging flag + Determines whether the application is running in debug mode. + + This function checks the environment variable `MERLIN_DEBUG`. If the variable + exists and its value is set to `1`, debug mode is enabled. + + Returns: + True if `MERLIN_DEBUG` is set to `1` in the environment, otherwise False. """ if "MERLIN_DEBUG" in os.environ and int(os.environ["MERLIN_DEBUG"]) == 1: return True return False -def default_config_info(): - """Return information about Merlin's default configurations.""" +def default_config_info() -> Dict: + """ + Returns information about Merlin's default configurations. + + This function gathers and returns key details about the current configuration + of the Merlin application, including the location of the configuration file, + debug mode status, the Merlin home directory, and whether the Merlin home + directory exists. + + Returns: + A dictionary containing the following keys:\n + - `config_file` (str): Path to the Merlin configuration file. + - `is_debug` (bool): Whether debug mode is enabled. + - `merlin_home` (str): Path to the Merlin home directory. + - `merlin_home_exists` (bool): True if the Merlin home directory exists, otherwise False. + """ return { "config_file": find_config_file(), "is_debug": is_debug(), @@ -176,14 +242,22 @@ def default_config_info(): } -def get_cert_file(server_type, config, cert_name, cert_path): +def get_cert_file(server_type: str, config: Config, cert_name: str, cert_path: str) -> str: """ - Check if a ssl certificate file is present in the config + Determines the SSL certificate file for a given server configuration. - :param server_type : The server type for output (Broker, Results Backend) - :param config : The server config - :param cert_name : The argument in cert argument name - :param cert_path : The optional cert path + This function checks if an SSL certificate file is specified in the server configuration. + If the file does not exist, it attempts to locate the certificate in an optional + certificate path. If the certificate file cannot be found, an error is logged. + + Args: + server_type (str): The type of server (e.g., Broker, Results Backend) for logging purposes. + config (config.Config): The server configuration object containing certificate details. + cert_name (str): The name of the certificate attribute in the configuration. + cert_path (str): An optional directory path to search for the certificate file. + + Returns: + The absolute path to the certificate file if found, otherwise `None`. """ cert_file = None try: @@ -208,14 +282,25 @@ def get_ssl_entries( server_type: str, server_name: str, server_config: Config, cert_path: str ) -> Dict[str, Union[str, ssl.VerifyMode]]: """ - Check if a ssl certificate file is present in the config - - :param [str] server_type : The server type - :param [str] server_name : The server name for output - :param [Config] server_config : The server config - :param [str] cert_path : The optional cert path - :return : The data needed to manage an ssl certification. - :rtype : A Dict. + Retrieves SSL configuration entries for a given server. + + This function checks for SSL certificate files and other SSL-related settings + in the server configuration. It builds and returns a dictionary containing + the necessary data to manage SSL certificates and protocols for the server. + + Args: + server_type (str): The type of server (e.g., Broker, Results Backend) for logging purposes. + server_name (str): The name of the server, used for output and mapping SSL configurations. + server_config (config.Config): The server configuration object containing SSL settings. + cert_path (str): An optional directory path to search for certificate files. + + Returns: + A dictionary containing SSL configuration entries, including:\n + - `keyfile` (str): Path to the SSL key file, if present. + - `certfile` (str): Path to the SSL certificate file, if present. + - `ca_certs` (str): Path to the CA certificates file, if present. + - `cert_reqs` (ssl.VerifyMode): SSL certificate requirements (e.g., `CERT_REQUIRED`, `CERT_OPTIONAL`, `CERT_NONE`). + - `ssl_protocol` (str): SSL protocol used, if specified. """ server_ssl: Dict[str, Union[str, ssl.VerifyMode]] = {} @@ -259,11 +344,22 @@ def get_ssl_entries( return server_ssl -def process_ssl_map(server_name: str) -> Optional[Dict[str, str]]: +def process_ssl_map(server_name: str) -> Dict[str, str]: """ - Process a special map for rediss and mysql. + Processes and returns a mapping of SSL-related configuration keys + specific to certain server types (e.g., Redis and MySQL). + + This function generates a dictionary mapping standard SSL configuration keys + (e.g., `keyfile`, `certfile`, `ca_certs`, `cert_reqs`) to server-specific key names + required by Redis (`rediss`) or MySQL server configurations. - :param server_name : The server name for output + Args: + server_name (str): The name of the server (e.g., "rediss", "mysql") used to determine + the appropriate SSL key mappings. + + Returns: + A dictionary containing the SSL key mappings for the given server type. Returns an empty + dictionary if the server type is not `rediss` or `mysql`. """ ssl_map: Dict[str, str] = {} # The redis server requires key names with ssl_ @@ -284,11 +380,20 @@ def process_ssl_map(server_name: str) -> Optional[Dict[str, str]]: def merge_sslmap(server_ssl: Dict[str, Union[str, ssl.VerifyMode]], ssl_map: Dict[str, str]) -> Dict: """ - The different servers have different key var expectations, this updates the keys of the ssl_server dict with keys from - the ssl_map if using rediss or mysql. + Updates the keys of the `server_ssl` dictionary based on the `ssl_map` for specific server types. + + This function modifies the `server_ssl` dictionary by replacing its keys with the corresponding + keys from the `ssl_map` when the server type requires specialized key names (e.g., `rediss` or `mysql`). + If a key in `server_ssl` is not found in `ssl_map`, it remains unchanged. + + Args: + server_ssl (Dict[str, Union[str, ssl.VerifyMode]]): The dictionary constructed in `get_ssl_entries`, + containing SSL configuration entries such as `keyfile`, `certfile`, `ca_certs`, and `cert_reqs`. + ssl_map (Dict[str, str]): A dictionary mapping standard SSL keys to server-specific keys. - : param server_ssl : the dict constructed in get_ssl_entries, here updated with keys from ssl_map - : param ssl_map : the dict holding special key:value pairs for rediss and mysql + Returns: + A new dictionary with updated keys based on the `ssl_map`. Keys not present in `ssl_map` + remain unchanged. """ new_server_ssl: Dict[str, Union[str, ssl.VerifyMode]] = {} diff --git a/merlin/config/results_backend.py b/merlin/config/results_backend.py index 259e249a6..1f8f56ef5 100644 --- a/merlin/config/results_backend.py +++ b/merlin/config/results_backend.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,22 +29,21 @@ ############################################################################### """ -This module contains the logic for configuring the Celery results backend. +This module provides functionality for managing and configuring connection strings +and SSL settings for various results backends, including MySQL, Redis, Rediss, and SQLite. +The module relies on the application's configuration file (`app.yaml`) to determine backend +settings and certificate paths. """ from __future__ import print_function import logging import os +from typing import Dict +from urllib.parse import quote from merlin.config.configfile import CONFIG, get_ssl_entries -try: - from urllib import quote -except ImportError: - from urllib.parse import quote - - LOG = logging.getLogger(__name__) BACKENDS = ["sqlite", "mysql", "redis", "rediss", "none"] @@ -72,16 +71,27 @@ SQLITE_CONNECTION_STRING = "db+sqlite:///results.db" -def get_backend_password(password_file, certs_path=None): +def get_backend_password(password_file: str, certs_path: str = None) -> str: """ - Check for password in file. - If the password is not found in the given password_file, - then the certs_path will be searched for the file, - if this file cannot be found, the password value will - be returned. - - :param password_file : The file path for the password - :param certs_path : The path for ssl certificates and passwords + Retrieves the backend password from a specified file or returns the provided password value. + + This function attempts to locate the password file in several locations: + + 1. The default Merlin directory (`~/.merlin`). + 2. The path specified by `password_file`. + 3. A directory specified by `certs_path` (if provided). + + If the password file is found, the password is read from the file. If the file cannot be + found, the value of `password_file` is treated as the password itself and returned. + + Args: + password_file (str): The file path or value for the password. If this is not a valid + file path, it is treated as the password itself. + certs_path (str, optional): An optional directory path where SSL certificates and + password files may be located. + + Returns: + The backend password, either retrieved from the file or the provided value. """ password = None @@ -113,13 +123,19 @@ def get_backend_password(password_file, certs_path=None): # flake8 complains about cyclomatic complexity because of all the try-excepts, # this isn't so complicated it can't be followed and tucking things in functions # would make it less readable, so complexity evaluation is off -def get_redis(certs_path=None, include_password=True, ssl=False): # noqa C901 +def get_redis(certs_path: str = None, include_password: bool = True, ssl: bool = False) -> str: # noqa C901 """ - Return the redis or rediss specific connection + Constructs and returns a Redis or Rediss connection URL based on the provided parameters and configuration. + + Args: + certs_path (str, optional): The path to SSL certificates and password files. + include_password (bool, optional): Whether to include the password in the connection URL. + If True, the password will be included; otherwise, it will be masked. + ssl (bool, optional): Flag indicating whether to use SSL for the connection (Rediss). + If True, the connection URL will use the "rediss" protocol; otherwise, it will use "redis". - :param certs_path : The path for ssl certificates and passwords - :param include_password : Format the connection for ouput by setting this True - :param ssl : Flag to use rediss output + Returns: + A Redis or Rediss connection URL formatted based on the provided parameters and configuration. """ server = CONFIG.results_backend.server password_file = "" @@ -165,13 +181,18 @@ def get_redis(certs_path=None, include_password=True, ssl=False): # noqa C901 return f"{urlbase}://{spass}{server}:{port}/{db_num}" -def get_mysql_config(certs_path, mysql_certs): +def get_mysql_config(certs_path: str, mysql_certs: Dict) -> Dict: """ - Determine if all the information for connecting MySQL as the Celery - results backend exists. + Determines whether all required information for connecting to MySQL as the Celery + results backend is available, and returns the MySQL SSL configuration or certificate paths. - :param certs_path : The path for ssl certificates and passwords - :param mysql_certs : The dict of mysql certificates + Args: + certs_path (str): The path to the directory containing SSL certificates and password files. + mysql_certs (Dict): A dictionary mapping certificate keys (e.g., 'cert', 'key', 'ca') + to their expected filenames. + + Returns: + A dictionary containing the paths to the required MySQL certificates if they exist. """ mysql_ssl = get_ssl_config(celery_check=False) if mysql_ssl: @@ -192,13 +213,25 @@ def get_mysql_config(certs_path, mysql_certs): return certs -def get_mysql(certs_path=None, mysql_certs=None, include_password=True): +def get_mysql(certs_path: str = None, mysql_certs: Dict = None, include_password: bool = True) -> str: """ - Returns the formatted MySQL connection string. - - :param certs_path : The path for ssl certificates and passwords - :param mysql_certs : The dict of mysql certificates - :param include_password : Format the connection for ouput by setting this True + Constructs and returns a formatted MySQL connection string based on the provided parameters + and application configuration. + + Args: + certs_path (str, optional): The path to the directory containing SSL certificates and password files. + mysql_certs (dict, optional): A dictionary mapping MySQL certificate keys (e.g., 'ssl_key', 'ssl_cert', 'ssl_ca') + to their expected filenames. If this is None, it uses the default `MYSQL_CONFIG_FILENAMES`. + include_password (bool, optional): Whether to include the password in the connection string. + If True, the password will be included; otherwise, it will be masked. + + Returns: + A formatted MySQL connection string. + + Raises: + TypeError: \n + - If the `server` configuration is missing or invalid. + - If the MySQL connection information cannot be set due to missing certificates or configuration. """ dbname = CONFIG.results_backend.dbname password_file = CONFIG.results_backend.password @@ -236,18 +269,32 @@ def get_mysql(certs_path=None, mysql_certs=None, include_password=True): mysql_config["password"] = "******" mysql_config["server"] = server + # Ensure the ssl_key, ssl_ca, and ssl_cert keys are all set + if mysql_certs == MYSQL_CONFIG_FILENAMES: + for key, cert_file in mysql_certs.items(): + if key not in mysql_config: + mysql_config[key] = os.path.join(certs_path, cert_file) + return MYSQL_CONNECTION_STRING.format(**mysql_config) -def get_connection_string(include_password=True): +def get_connection_string(include_password: bool = True) -> str: """ - Given the package configuration determine what results backend to use and - return the connection string. + Determines the appropriate results backend to use based on the package configuration + and returns the corresponding connection string. + + If a URL is explicitly defined in the configuration (`CONFIG.results_backend.url`), + it is returned as the connection string. + + Args: + include_password (bool, optional): Whether to include the password in the connection string. + If True, the password will be included; otherwise, it will be masked. - If the url variable is present, return that as the connection string. + Returns: + The connection string for the configured results backend. - :param config_path : The path for ssl certificates and passwords - :param include_password : Format the connection for ouput by setting this True + Raises: + ValueError: If the specified results backend in the configuration is not supported. """ try: return CONFIG.results_backend.url @@ -272,7 +319,23 @@ def get_connection_string(include_password=True): return _resolve_backend_string(backend, certs_path, include_password) -def _resolve_backend_string(backend, certs_path, include_password): +def _resolve_backend_string(backend: str, certs_path: str, include_password: bool) -> str: + """ + Resolves and returns the connection string for the specified results backend. + + Based on the backend type provided, this function delegates the connection string + generation to the appropriate helper function or returns a predefined connection string. + + Args: + backend (str): The name of the results backend (e.g., "mysql", "sqlite", "redis", "rediss"). + certs_path (str): The path to SSL certificates and password files, used for certain backends + (e.g., MySQL and Redis). + include_password (bool): Whether to include the password in the connection string. + If True, the password will be included; otherwise, it will be masked. + + Returns: + The connection string for the specified backend, or `None` if the backend is unsupported. + """ if "mysql" in backend: return get_mysql(certs_path=certs_path, include_password=include_password) @@ -288,12 +351,22 @@ def _resolve_backend_string(backend, certs_path, include_password): return None -def get_ssl_config(celery_check=False): +def get_ssl_config(celery_check: bool = False) -> bool: """ - Return the ssl config based on the configuration specified in the - `app.yaml` config file. + Retrieves the SSL configuration for the results backend based on the settings + specified in the `app.yaml` configuration file. + + This function determines whether SSL should be enabled for the results backend + and returns the appropriate configuration. It supports various backend types + such as MySQL, Redis, and Rediss. + + Args: + celery_check (bool, optional): If True, the function returns the SSL settings + specifically for configuring Celery. - :param celery_check : Return the proper results ssl setting when configuring celery + Returns: + The SSL configuration for the results backend. Returns `True` if SSL is enabled, + `False` otherwise. """ results_backend = "" try: diff --git a/merlin/config/utils.py b/merlin/config/utils.py index 46672ba1f..bace9fd52 100644 --- a/merlin/config/utils.py +++ b/merlin/config/utils.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,7 +27,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module contains priority handling""" +""" +This module provides utility functions and classes for handling broker priorities +and determining configurations for supported brokers such as RabbitMQ and Redis. +It includes functionality for mapping priority levels to integer values based on +the broker type and validating broker configurations. +""" import enum from typing import Dict @@ -36,7 +41,18 @@ class Priority(enum.Enum): - """Enumerated Priorities""" + """ + Enumerated Priorities. + + This enumeration defines the different priority levels that can be used + for message handling with brokers. + + Attributes: + HIGH (int): Represents the highest priority level. Numeric value: 1. + MID (int): Represents the medium priority level. Numeric value: 2. + LOW (int): Represents the lowest priority level. Numeric value: 3. + RETRY (int): Represents the priority level for retrying messages. Numeric value: 4. + """ HIGH = 1 MID = 2 @@ -44,22 +60,54 @@ class Priority(enum.Enum): RETRY = 4 -def is_rabbit_broker(broker: str) -> bool: - """Check if the broker is a rabbit server""" - return broker in ["rabbitmq", "amqps", "amqp"] +def is_rabbit_broker(broker_name: str) -> bool: + """ + Check if the given broker is a RabbitMQ server. + + This function checks whether the provided broker name matches any of the + RabbitMQ-related broker types. + + Args: + broker_name: The name of the broker to check. + + Returns: + True if the broker is a RabbitMQ server, False otherwise. + """ + return broker_name in ["rabbitmq", "amqps", "amqp"] + + +def is_redis_broker(broker_name: str) -> bool: + """ + Check if the given broker is a Redis server. + This function checks whether the provided broker name matches any of the + Redis-related broker types. -def is_redis_broker(broker: str) -> bool: - """Check if the broker is a redis server""" - return broker in ["redis", "rediss", "redis+socket"] + Args: + broker_name: The name of the broker to check. + + Returns: + True if the broker is a Redis server, False otherwise. + """ + return broker_name in ["redis", "rediss", "redis+socket"] def determine_priority_map(broker_name: str) -> Dict[Priority, int]: """ - Returns the priority mapping for the given broker name. + Determine the priority mapping for the given broker name. + + This function returns a mapping of [`Priority`][config.utils.Priority] + enum values to integer priority levels based on the type of broker provided. - :param broker_name: The name of the broker that we need the priority map for - :returns: The priority map associated with `broker_name` + Args: + broker_name: The name of the broker for which to determine the priority map. + + Returns: + (Dict[config.utils.Priority, int]): A dictionary mapping + [`Priority`][config.utils.Priority] enum values to integer levels. + + Raises: + ValueError: If the broker name is not supported. """ if is_rabbit_broker(broker_name): return {Priority.LOW: 1, Priority.MID: 5, Priority.HIGH: 9, Priority.RETRY: 10} @@ -71,14 +119,33 @@ def determine_priority_map(broker_name: str) -> Dict[Priority, int]: def get_priority(priority: Priority) -> int: """ - Gets the priority level as an integer based on the broker. - For a rabbit broker a low priority is 1 and high is 10. For redis it's the opposite. + Get the integer priority level for a given [`Priority`][config.utils.Priority] + enum value. + + This function determines the priority level as an integer based on the + broker configuration. For RabbitMQ brokers, lower numbers represent lower + priorities, while for Redis brokers, higher numbers represent lower + priorities. + + Args: + priority (config.utils.Priority): The [`Priority`][config.utils.Priority] + enum value for which to get the integer level. + + Returns: + The integer priority level corresponding to the given [`Priority`][config.utils.Priority]. - :param priority: The priority value that we want - :returns: The priority value as an integer + Raises: + ValueError: If the provided `priority` is invalid or not part of the + [`Priority`][config.utils.Priority] enum. """ - if priority not in Priority: - raise ValueError(f"Invalid priority: {priority}") + priority_err_msg = f"Invalid priority: {priority}" + try: + # In python 3.12+ if something is not in the enum it will just return False + if priority not in Priority: + raise ValueError(priority_err_msg) + # In python 3.11 and below, a TypeError is raised when looking for something in an enum that is not there + except TypeError: + raise ValueError(priority_err_msg) priority_map = determine_priority_map(CONFIG.broker.name.lower()) return priority_map.get(priority, priority_map[Priority.MID]) # Default to MID priority for unknown priorities diff --git a/merlin/data/celery/__init__.py b/merlin/data/celery/__init__.py index 57477ea1f..37cabcad1 100644 --- a/merlin/data/celery/__init__.py +++ b/merlin/data/celery/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/merlin/display.py b/merlin/display.py index a1af0ac28..cbf88fc65 100644 --- a/merlin/display.py +++ b/merlin/display.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -37,11 +37,13 @@ import shutil import time import traceback +from argparse import Namespace from datetime import datetime from multiprocessing import Pipe, Process -from typing import Dict +from multiprocessing.connection import Connection +from typing import Any, Dict, List, Union -from kombu import Connection +from kombu import Connection as KombuConnection from tabulate import tabulate from merlin.ascii_art import banner_small @@ -72,16 +74,39 @@ class ConnProcess(Process): """ - An extension of Multiprocessing's Process class in order - to overwrite the run and exception defintions. + An extension of the multiprocessing's Process class that allows for + custom handling of exceptions and inter-process communication. + + This class overrides the `run` method to capture exceptions that occur + during the execution of the process and sends them back to the parent + process via a pipe. It also provides a property to retrieve any + exceptions that were raised during execution. + + Attributes: + _pconn: The parent connection for inter-process communication. + _cconn: The child connection for inter-process communication. + exception: Stores the exception raised during the process run. + + Methods: + run: Executes the process's main logic. """ def __init__(self, *args, **kwargs): Process.__init__(self, *args, **kwargs) + self._pconn: Connection + self._cconn: Connection self._pconn, self._cconn = Pipe() self._exception = None def run(self): + """ + Executes the process's main logic. + + This method overrides the default run method of the Process class. + It attempts to run the process and captures any exceptions that occur. + If an exception is raised, it sends the exception and its traceback + back to the parent process via the child connection. + """ try: Process.run(self) self._cconn.send(None) @@ -91,17 +116,35 @@ def run(self): # raise e # You can still rise this exception if you need to @property - def exception(self): - """Create custom exception""" + def exception(self) -> Union[Exception, None]: + """ + Retrieves the exception raised during the process execution. + + This property checks if there is an exception available from the + parent connection. If an exception was raised, it is received and + stored for later access. + + Returns: + The exception raised during the process run, or None if no exception occurred. + """ if self._pconn.poll(): self._exception = self._pconn.recv() return self._exception -def check_server_access(sconf): +def check_server_access(sconf: Dict[str, Any]): """ Check if there are any issues connecting to the servers. If there are, output the errors. + + This function iterates through a predefined list of servers and checks + their connectivity based on the provided server configuration. If any + connection issues are detected, the exceptions are collected and printed. + + Args: + sconf: A dictionary containing server configurations, where keys + represent server names and values contain connection details. + The function expects keys corresponding to the servers being checked. """ servers = ["broker server", "results server"] @@ -120,7 +163,24 @@ def check_server_access(sconf): print(f"{key}: {val}") -def _examine_connection(server, sconf, excpts): +def _examine_connection(server: str, sconf: Dict[str, Any], excpts: Dict[str, Exception]): + """ + Examine the connection to a specified server and handle any exceptions. + + This function attempts to establish a connection to the given server using + the configuration provided in `sconf`. It utilizes a separate process to + manage the connection attempt and checks for timeouts. If the connection + fails or times out, the error is recorded in the `excpts` dictionary. + + Args: + server: A string representing the name of the server to connect to. + This should correspond to a key in the `sconf` dictionary. + sconf: A dictionary containing server configurations, where keys + represent server names and values contain connection details. + excpts: A dictionary to store exceptions encountered during the + connection attempt, with server names as keys and exceptions + as values. + """ from merlin.config import broker, results_backend # pylint: disable=C0415 connect_timeout = 60 @@ -130,7 +190,7 @@ def _examine_connection(server, sconf, excpts): ssl_conf = broker.get_ssl_config() if "results" in server: ssl_conf = results_backend.get_ssl_config() - conn = Connection(sconf[server], ssl=ssl_conf) + conn = KombuConnection(sconf[server], ssl=ssl_conf) conn_check = ConnProcess(target=conn.connect) conn_check.start() counter = 0 @@ -153,7 +213,11 @@ def _examine_connection(server, sconf, excpts): def display_config_info(): """ - Prints useful configuration information to the console. + Prints useful configuration information for the Merlin application to the console. + + This function retrieves and displays the connection strings and SSL configurations + for the broker and results servers. It handles any exceptions that may occur during + the retrieval process, providing error messages for any issues encountered. """ from merlin.config import broker, results_backend # pylint: disable=C0415 from merlin.config.configfile import default_config_info # pylint: disable=C0415 @@ -191,12 +255,13 @@ def display_config_info(): check_server_access(sconf) -def display_multiple_configs(files, configs): +def display_multiple_configs(files: List[str], configs: List[Dict]): """ Logic for displaying multiple Merlin config files. - :param `files`: List of merlin config files - :param `configs`: List of merlin configurations + Args: + files: List of merlin config files + configs: List of merlin configurations """ print("=" * 50) print(" MERLIN CONFIG ") @@ -211,13 +276,19 @@ def display_multiple_configs(files, configs): # Might use args here in the future so we'll disable the pylint warning for now -def print_info(args): # pylint: disable=W0613 +def print_info(args: Namespace): # pylint: disable=W0613 """ Provide version and location information about python and packages to facilitate user troubleshooting. Also provides info about server connections and configurations. - :param `args`: parsed CLI arguments + Note: + The `args` parameter is currently unused but is included for + compatibility with the command-line interface (CLI) in case we decide to use + args here in the future. + + Args: + args: parsed CLI arguments (currently unused). """ print(banner_small) display_config_info() @@ -235,16 +306,30 @@ def print_info(args): # pylint: disable=W0613 def display_status_task_by_task(status_obj: "DetailedStatus", test_mode: bool = False): # noqa: F821 """ - Displays a low level overview of the status of a study. This is a task-by-task - status display where each task will show: - step name, worker name, task queue, cmd & restart parameters, - step workspace, step status, return code, elapsed time, run time, and num restarts. - If too many tasks are found and the pager is disabled, prompts will appear for the user to decide - what to do that way we don't overload the terminal (unless the no-prompts flag is provided). - - :param `status_obj`: A DetailedStatus object - :param `test_mode`: If true, run this in testing mode and don't print any output. This will also - decrease the limit on the number of tasks allowed before a prompt is displayed. + Displays a low-level overview of the status of a study in a task-by-task format. + + Each task will display the following details: + - Step name + - Worker name + - Task queue + - Command and restart parameters + - Step workspace + - Step status + - Return code + - Elapsed time + - Run time + - Number of restarts + + If the number of tasks exceeds a certain limit and the pager is disabled, the user + will be prompted to apply additional filters to avoid overwhelming the terminal output, + unless the prompts are disabled through the no-prompts flag. + + Args: + status_obj (study.status.DetailedStatus): An instance of + [`DetailedStatus`][study.status.DetailedStatus] containing information about + the current state of tasks. + test_mode: If True, runs the function in testing mode, suppressing output and + reducing the task limit for prompts. Defaults to False. """ args = status_obj.args try: @@ -299,10 +384,18 @@ def display_status_task_by_task(status_obj: "DetailedStatus", test_mode: bool = def _display_summary(state_info: Dict[str, str], cb_help: bool): """ - Given a dict of state info for a step, print a summary of the task states. - - :param `state_info`: A dictionary of information related to task states for a step - :param `cb_help`: True if colorblind assistance (using symbols) is needed. False otherwise. + Prints a summary of task states based on the provided state information. + + This function takes a dictionary of state information for a step and + prints a formatted summary, including optional colorblind assistance using + symbols if specified. + + Args: + state_info: A dictionary containing information related to task states + for a step. Each entry should correspond to a specific task state + with its associated properties (e.g., count, total, name). + cb_help: If True, provides colorblind assistance by using symbols in the + display. Defaults to False for standard output. """ # Build a summary list of task info print("\nSUMMARY:") @@ -338,18 +431,27 @@ def _display_summary(state_info: Dict[str, str], cb_help: bool): def display_status_summary( # pylint: disable=R0912 - status_obj: "Status", non_workspace_keys: set, test_mode=False # noqa: F821 + status_obj: "Status", non_workspace_keys: set, test_mode: bool = False # noqa: F821 ) -> Dict: """ - Displays a high level overview of the status of a study. This includes - progress bars for each step and a summary of the number of initialized, - running, finished, cancelled, dry ran, failed, and unknown tasks. - - :param `status_obj`: A Status object - :param `non_workspace_keys`: A set of keys in requested_statuses that are not workspace keys. - This will be set("parameters", "task_queue", "workers") - :param `test_mode`: If True, don't print anything and just return a dict of all the state info for each step - :returns: A dict that's empty usually. If ran in test_mode it will be a dict of state_info for every step. + Displays a high-level overview of the status of a study, including progress bars for each step + and a summary of the number of initialized, running, finished, cancelled, dry ran, failed, and + unknown tasks. + + The function prints a summary for each step and collects state information. In test mode, + it suppresses output and returns a dictionary of state information instead. + + Args: + status_obj (study.status.Status): An instance of [`Status`][study.status.Status] containing + information about task states and associated data for the study. + non_workspace_keys: A set of keys in requested_statuses that are not workspace keys. + Typically includes keys like "parameters", "task_queue", and "workers". + test_mode: If True, runs in test mode; suppresses printing and returns a dictionary + of state information for each step. Defaults to False. + + Returns: + An empty dictionary in regular mode. In test mode, returns a dictionary containing + the state information for each step. """ all_state_info = {} if not test_mode: @@ -429,32 +531,40 @@ def display_status_summary( # pylint: disable=R0912 # Credit to this stack overflow post: https://stackoverflow.com/a/34325723 def display_progress_bar( # pylint: disable=R0913,R0914 - current, - total, - state_info=None, - prefix="", - suffix="", - decimals=1, - length=80, - fill="█", - print_end="\n", - color=None, - cb_help=False, + current: int, + total: int, + state_info: Dict[str, Any] = None, + prefix: str = "", + suffix: str = "", + decimals: int = 1, + length: int = 80, + fill: str = "█", + print_end: str = "\n", + color: str = None, + cb_help: bool = False, ): """ - Prints a progress bar based on current and total. - - :param `current`: current number (Int) - :param `total`: total number (Int) - :param `state_info`: information about the state of tasks (Dict) (overrides color) - :param `prefix`: prefix string (Str) - :param `suffix`: suffix string (Str) - :param `decimals`: positive number of decimals in percent complete (Int) - :param `length`: character length of bar (Int) - :param `fill`: bar fill character (Str) - :param `print_end`: end character (e.g. "\r", "\r\n") (Str) - :param `color`: color of the progress bar (ANSI Str) (overridden by state_info) - :param `cb_help`: true if color blind help is needed; false otherwise (Bool) + Prints a customizable progress bar that visually represents the completion percentage + relative to a given total. + + The function can display additional state information for detailed tracking, including + support for color customization and adaptation for color-blind users. It updates the + display based on current progress and optionally accepts state information to adjust the + appearance of the progress bar. + + Args: + current: Current progress value. + total: Total value representing 100% completion. + state_info: Dictionary containing state information about tasks. This can override + color settings and modifies how the progress bar is displayed. + prefix: Optional prefix string to display before the progress bar. + suffix: Optional suffix string to display after the progress bar. + decimals: Number of decimal places to display in the percentage (default is 1). + length: Character length of the progress bar (default is 80). + fill: Character used to fill the progress bar (default is "█"). + print_end: Character(s) to print at the end of the line (e.g., '\\r', '\\n'). + color: ANSI color string for the progress bar. Overrides state_info colors. + cb_help: If True, provides color-blind assistance by adapting the fill characters. """ # Set the color of the bar if color and color in ANSI_COLORS: diff --git a/merlin/examples/__init__.py b/merlin/examples/__init__.py index 57477ea1f..4db966018 100644 --- a/merlin/examples/__init__.py +++ b/merlin/examples/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,3 +27,12 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `examples` package provides resources for learning and setting up Merlin workflows. + +Modules: + examples.py: Contains example specification files for Merlin workflows, along with + detailed explanations of each block in the specification. + generator.py: Provides utilities for managing and generating example workflows. +""" diff --git a/merlin/examples/dev_workflows/multiple_workers.yaml b/merlin/examples/dev_workflows/multiple_workers.yaml index 8785d9e9a..967582a53 100644 --- a/merlin/examples/dev_workflows/multiple_workers.yaml +++ b/merlin/examples/dev_workflows/multiple_workers.yaml @@ -46,11 +46,11 @@ merlin: resources: workers: step_1_merlin_test_worker: - args: -l INFO + args: -l INFO --concurrency 1 steps: [step_1] step_2_merlin_test_worker: - args: -l INFO + args: -l INFO --concurrency 1 steps: [step_2] other_merlin_test_worker: - args: -l INFO + args: -l INFO --concurrency 1 steps: [step_3, step_4] diff --git a/merlin/examples/examples.py b/merlin/examples/examples.py index 371471550..2dd8dad1b 100644 --- a/merlin/examples/examples.py +++ b/merlin/examples/examples.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,7 +27,13 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module contains example spec files with explanations of each block""" +""" +This module contains example specification files for Merlin workflows, +along with detailed explanations of each block in the specification. + +The examples and templates are useful for understanding how to structure +Merlin workflows, define tasks, manage parameters, and configure resources. +""" # Taken from https://lc.llnl.gov/mlsi/docs/merlin/merlin_config.html TEMPLATE_FILE_CONTENTS = """ diff --git a/merlin/examples/generator.py b/merlin/examples/generator.py index a553d703b..3b0896754 100644 --- a/merlin/examples/generator.py +++ b/merlin/examples/generator.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -33,12 +33,13 @@ Merlin, or for setting up new workflows. Examples are packaged in directories, with the directory name denoting -the example name. This must match the name of the merlin specification inside. +the example name. This must match the name of the Merlin specification inside. """ import glob import logging import os import shutil +from typing import Dict, List, Union import tabulate import yaml @@ -48,26 +49,46 @@ EXAMPLES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "workflows") +# TODO modify the example command to eliminate redundancy +# - e.g. running `merlin example flux_local` will produce the same output +# as running `merlin example flux_par` or `merlin example flux_par_restart`. +# This should just be `merlin example flux`. +# - restart and restart delay should be one example +# - feature demo and remote feature demo should be one example +# - all openfoam examples should just be under one openfoam label -def gather_example_dirs(): - """Get all the example directories""" + +def gather_example_dirs() -> Dict[str, str]: + """ + Get all the example directories. + + Returns: + A dictionary where the keys and values are the names of example directories. + """ result = {} - for directory in os.listdir(EXAMPLES_DIR): + for directory in sorted(os.listdir(EXAMPLES_DIR)): result[directory] = directory return result -def gather_all_examples(): - """Get all the example yaml files""" +def gather_all_examples() -> List[str]: + """ + Get all the example YAML files. + + Returns: + A list of file paths to all YAML files in the example directories. + """ path = os.path.join(os.path.join(EXAMPLES_DIR, ""), os.path.join("*", "*.yaml")) return glob.glob(path) -def write_example(src_path, dst_path): +def write_example(src_path: str, dst_path: str): """ - Write out the example workflow to a file. - :param src_path: The path to copy from. - :param content: The formatted content to write the file to. + Write out the example workflow to a file or directory. + + Args: + src_path: The path to copy the example from. + dst_path: The destination path to copy the example to. """ if os.path.isdir(src_path): shutil.copytree(src_path, dst_path) @@ -75,14 +96,19 @@ def write_example(src_path, dst_path): shutil.copy(src_path, dst_path) -def list_examples(): - """List all available examples.""" +def list_examples() -> str: + """ + List all available examples with their descriptions. + + Returns: + A formatted string table of example names and their descriptions. + """ headers = ["name", "description"] rows = [] for example_dir in gather_example_dirs(): directory = os.path.join(os.path.join(EXAMPLES_DIR, example_dir), "") specs = glob.glob(directory + "*.yaml") - for spec in specs: + for spec in sorted(specs): if "template" in spec: continue with open(spec) as f: # pylint: disable=C0103 @@ -103,8 +129,17 @@ def list_examples(): return "\n" + tabulate.tabulate(rows, headers) + "\n" -def setup_example(name, outdir): - """Setup the given example.""" +def setup_example(name: str, outdir: str) -> Union[str, None]: + """ + Set up the given example by copying it to the specified output directory. + + Args: + name: The name of the example to set up. + outdir: The output directory where the example will be copied. + + Returns: + The name of the example if successful, or None if the example was not found or an error occurred. + """ example = None spec_paths = gather_all_examples() spec_path = None diff --git a/merlin/examples/workflows/feature_demo/requirements.txt b/merlin/examples/workflows/feature_demo/requirements.txt index e308e1895..3eee4d90c 100644 --- a/merlin/examples/workflows/feature_demo/requirements.txt +++ b/merlin/examples/workflows/feature_demo/requirements.txt @@ -1,2 +1,3 @@ scikit-learn -merlin-spellbook +merlin-spellbook; python_version < "3.12" +merlin-spellbook>=0.9.0; python_version >= "3.12" diff --git a/merlin/examples/workflows/optimization/requirements.txt b/merlin/examples/workflows/optimization/requirements.txt index a487c9f39..378678b1e 100644 --- a/merlin/examples/workflows/optimization/requirements.txt +++ b/merlin/examples/workflows/optimization/requirements.txt @@ -1,4 +1,5 @@ -merlin-spellbook +merlin-spellbook; python_version < "3.12" +merlin-spellbook>=0.9.0; python_version >= "3.12" numpy scikit-learn matplotlib diff --git a/merlin/examples/workflows/remote_feature_demo/requirements.txt b/merlin/examples/workflows/remote_feature_demo/requirements.txt index e308e1895..3eee4d90c 100644 --- a/merlin/examples/workflows/remote_feature_demo/requirements.txt +++ b/merlin/examples/workflows/remote_feature_demo/requirements.txt @@ -1,2 +1,3 @@ scikit-learn -merlin-spellbook +merlin-spellbook; python_version < "3.12" +merlin-spellbook>=0.9.0; python_version >= "3.12" diff --git a/merlin/exceptions/__init__.py b/merlin/exceptions/__init__.py index 572135aec..ade4862d6 100644 --- a/merlin/exceptions/__init__.py +++ b/merlin/exceptions/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -78,7 +78,7 @@ def __init__(self): class InvalidChainException(Exception): """ - Exception for invalid Merlin step DAGs. + Exception for invalid Merlin step Directed Acyclic Graphs (DAGs). """ def __init__(self): diff --git a/merlin/log_formatter.py b/merlin/log_formatter.py index 6cd6a745a..b67bf96dc 100644 --- a/merlin/log_formatter.py +++ b/merlin/log_formatter.py @@ -8,7 +8,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -43,12 +43,14 @@ } -def setup_logging(logger, log_level="INFO", colors=True): +def setup_logging(logger: logging.Logger, log_level: str = "INFO", colors: bool = True): """ Setup and configure Python logging. - :param `logger`: a logging.Logger object - :param `log_level`: logger level + Args: + logger: A logging.Logger object. + log_level: Logger level. + colors: If True use colored logs. """ formatter = logging.Formatter() handler = logging.StreamHandler(sys.stdout) diff --git a/merlin/main.py b/merlin/main.py index 4bb005985..e8fdaa241 100644 --- a/merlin/main.py +++ b/merlin/main.py @@ -8,7 +8,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -46,7 +46,7 @@ RawTextHelpFormatter, ) from contextlib import suppress -from typing import Dict, List, Optional, Union +from typing import Dict, List, Optional, Tuple, Union from tabulate import tabulate @@ -69,10 +69,21 @@ class HelpParser(ArgumentParser): - """This class overrides the error message of the argument parser to - print the help message when an error happens.""" + """ + This class overrides the error message of the argument parser to + print the help message when an error happens. + + Methods: + error: Override the error message of the `ArgumentParser` class. + """ + + def error(self, message: str): + """ + Override the error message of the `ArgumentParser` class. - def error(self, message): + Args: + message: The error message to log. + """ sys.stderr.write(f"error: {message}\n") self.print_help() sys.exit(2) @@ -82,13 +93,27 @@ def parse_override_vars( variables_list: Optional[List[str]], ) -> Optional[Dict[str, Union[str, int]]]: """ - Parse a list of variables from command line syntax - into a valid dictionary of variable keys and values. - - :param [List[str]] `variables_list`: an optional list of strings, e.g. ["KEY=val",...] - - :return: returns either None or a Dict keyed with strs, linked to strs and ints. - :rtype: Dict + Parse a list of command-line variables into a dictionary of key-value pairs. + + This function takes an optional list of strings following the syntax + "KEY=val" and converts them into a dictionary. It validates the format + of the variables and ensures that keys are valid according to specified rules. + + Args: + variables_list: An optional list of strings, where each string should be in the + format "KEY=val", e.g., ["KEY1=value1", "KEY2=42"]. + + Returns: + A dictionary where the keys are variable names (str) and the + values are either strings or integers. If `variables_list` is + None or empty, returns None. + + Raises: + ValueError: If the input format is incorrect, including:\n + - Missing '=' operator. + - Excess '=' operators in a variable assignment. + - Invalid variable names (must be alphanumeric and underscores). + - Attempting to override reserved variable names. """ if variables_list is None: return None @@ -121,11 +146,26 @@ def parse_override_vars( return result -def get_merlin_spec_with_override(args): +def get_merlin_spec_with_override(args: Namespace) -> Tuple[MerlinSpec, str]: """ - Shared command to return the spec object. - - :param 'args': parsed CLI arguments + Shared command to retrieve a [`MerlinSpec`][spec.specification.MerlinSpec] object + and an expanded filepath. + + This function processes parsed command-line interface (CLI) arguments to validate + and expand the specified filepath and any associated variables. It then constructs + and returns a [`MerlinSpec`][spec.specification.MerlinSpec] object based on the + provided specification. + + Args: + args: Parsed CLI arguments containing:\n + - `specification`: the path to the specification file + - `variables`: optional variable overrides to customize the spec. + + Returns: + spec (spec.specification.MerlinSpec): An instance of the + [`MerlinSpec`][spec.specification.MerlinSpec] class with the expanded + configuration based on the provided filepath and variables. + filepath: The expanded filepath derived from the specification. """ filepath = verify_filepath(args.specification) variables_dict = parse_override_vars(args.variables) @@ -133,11 +173,27 @@ def get_merlin_spec_with_override(args): return spec, filepath -def process_run(args: Namespace) -> None: +def process_run(args: Namespace): """ CLI command for running a study. - :param [Namespace] `args`: parsed CLI arguments + This function initializes and runs a study using the specified parameters. + It handles file verification, variable parsing, and checks for required + arguments related to the study configuration and execution. + + Args: + args: Parsed CLI arguments containing:\n + - `specification`: Path to the specification file for the study. + - `variables`: Optional variable overrides for the study. + - `samples_file`: Optional path to a samples file. + - `dry`: If True, runs the study in dry-run mode (without actual execution). + - `no_errors`: If True, suppresses error reporting. + - `pgen_file`: Optional path to the pgen file, required if `pargs` is specified. + - `pargs`: Additional arguments for parallel processing. + + Raises: + ValueError: + If the `pargs` parameter is used without specifying a `pgen_file`. """ print(banner_small) filepath: str = verify_filepath(args.specification) @@ -164,11 +220,22 @@ def process_run(args: Namespace) -> None: router.run_task_server(study, args.run_mode) -def process_restart(args: Namespace) -> None: +def process_restart(args: Namespace): """ CLI command for restarting a study. - :param [Namespace] `args`: parsed CLI arguments + This function handles the restart process by verifying the specified restart + directory, locating a valid provenance specification file, and initiating + the study from that point. + + Args: + args: Parsed CLI arguments containing:\n + - `restart_dir`: Path to the directory where the restart specifications are located. + - `run_mode`: The mode for running the study (e.g., normal, dry-run). + + Raises: + ValueError: If the `restart_dir` does not contain a valid provenance spec file or + if multiple files match the specified pattern. """ print(banner_small) restart_dir: str = verify_dirpath(args.restart_dir) @@ -184,11 +251,20 @@ def process_restart(args: Namespace) -> None: router.run_task_server(study, args.run_mode) -def launch_workers(args): +def launch_workers(args: Namespace): """ CLI command for launching workers. - :param `args`: parsed CLI arguments + This function initializes worker processes for executing tasks as defined + in the Merlin specification. + + Args: + args: Parsed CLI arguments containing:\n + - `worker_echo_only`: If True, don't start the workers and just echo the launch command + - Additional worker-related parameters such as: + - `worker_steps`: Only start workers for these steps. + - `worker_args`: Arguments to pass to the worker processes. + - `disable_logs`: If True, disables logging for the worker processes. """ if not args.worker_echo_only: print(banner_small) @@ -204,11 +280,17 @@ def launch_workers(args): LOG.debug(f"celery command: {launch_worker_status}") -def purge_tasks(args): +def purge_tasks(args: Namespace): """ - CLI command for purging tasks. + CLI command for purging tasks from the task server. + + This function removes specified tasks from the task server based on the provided + Merlin specification. It allows for targeted purging or forced removal of tasks. - :param `args`: parsed CLI arguments + Args: + args: Parsed CLI arguments containing:\n + - `purge_force`: If True, forces the purge operation without confirmation. + - `purge_steps`: Steps or criteria based on which tasks will be purged. """ print(banner_small) spec, _ = get_merlin_spec_with_override(args) @@ -222,15 +304,28 @@ def purge_tasks(args): LOG.info(f"Purge return = {ret} .") -def query_status(args): +def query_status(args: Namespace): """ - CLI command for querying status of studies. - Based on the parsed CLI args, construct either a Status object or a DetailedStatus object - and display the appropriate output. - Object mapping is as follows: - merlin status -> Status object ; merlin detailed-status -> DetailedStatus object - - :param `args`: parsed CLI arguments + CLI command for querying the status of studies. + + This function processes the given command-line arguments to determine the + status of a study. It constructs either a [`Status`][study.status.Status] object + or a [`DetailedStatus`][study.status.DetailedStatus] object based on the specified + command and the arguments provided. The function handles validations for the task + server input and the output format specified for status dumping. + + Object mapping: + - `merlin status` -> [`Status`][study.status.Status] object + - `merlin detailed-status` -> [`DetailedStatus`][study.status.DetailedStatus] + object + + Args: + args: Parsed CLI arguments containing user inputs for the status query. + + Raises: + ValueError: + - If the task server specified is not supported (only "celery" is valid). + - If the --dump filename provided does not end with ".csv" or ".json". """ print(banner_small) @@ -274,11 +369,24 @@ def query_status(args): return None -def query_queues(args): +def query_queues(args: Namespace): """ - CLI command for finding all workers. - - :param args: parsed CLI arguments + CLI command for finding all workers and their associated queues. + + This function processes the command-line arguments to retrieve and display + information about the available workers and their queues within the task server. + It validates the necessary parameters, handles potential file dumping, and + formats the output for easy readability. + + Args: + args: Parsed CLI arguments containing user inputs related to the query. + + Raises: + ValueError: + - If a specification is not provided when steps are specified and the + steps do not include "all". + - If variables are included without a corresponding specification. + - If the specified dump filename does not end with '.json' or '.csv'. """ print(banner_small) @@ -318,11 +426,19 @@ def query_queues(args): router.dump_queue_info(args.task_server, queue_information, args.dump) -def query_workers(args): +def query_workers(args: Namespace): """ CLI command for finding all workers. - :param `args`: parsed CLI arguments + This function retrieves and queries the names of any active workers. + If the `--spec` argument is included, only query the workers defined in the spec file. + + Args: + args: Parsed command-line arguments, which may include:\n + - `spec`: Path to the specification file. + - `task_server`: Address of the task server to query. + - `queues`: List of queue names to filter workers. + - `workers`: List of specific worker names to query. """ print(banner_small) @@ -340,11 +456,20 @@ def query_workers(args): router.query_workers(args.task_server, worker_names, args.queues, args.workers) -def stop_workers(args): +def stop_workers(args: Namespace): """ CLI command for stopping all workers. - :param `args`: parsed CLI arguments + This function stops any active workers connected to a user's task server. + If the `--spec` argument is provided, this function retrieves the names of + workers from a the spec file and then issues a command to stop them. + + Args: + args: Parsed command-line arguments, which may include:\n + - `spec`: Path to the specification file to load worker names. + - `task_server`: Address of the task server to send the stop command to. + - `queues`: List of queue names to filter the workers. + - `workers`: List of specific worker names to stop. """ print(banner_small) worker_names = [] @@ -362,11 +487,12 @@ def stop_workers(args): router.stop_workers(args.task_server, worker_names, args.queues, args.workers) -def print_info(args): +def print_info(args: Namespace): """ - CLI command to print merlin config info. + CLI command to print merlin configuration info. - :param `args`: parsed CLI arguments + Args: + args: Parsed CLI arguments. """ # if this is moved to the toplevel per standard style, merlin is unable to generate the (needed) default config file from merlin import display # pylint: disable=import-outside-toplevel @@ -374,11 +500,24 @@ def print_info(args): display.print_info(args) -def config_merlin(args: Namespace) -> None: +def config_merlin(args: Namespace): """ - CLI command to setup default merlin config. - - :param [Namespace] `args`: parsed CLI arguments + CLI command to set up the default Merlin configuration. + + This function initializes the configuration app.yaml file that's + necessary to connect Merlin to a central server. If the output + directory is not specified via the command-line arguments, it + defaults to the user's home directory under `.merlin`. + + Args: + args: Parsed command-line arguments, which may include:\n + - `output_dir`: Path to the output directory for + configuration files. If not provided, defaults to + `~/.merlin`. + - `task_server`: Address of the task server for the + configuration. + - `broker`: Address of the broker service to use. + - `test`: Flag indicating whether to run in test mode. """ output_dir: Optional[str] = args.output_dir if output_dir is None: @@ -389,9 +528,20 @@ def config_merlin(args: Namespace) -> None: def process_example(args: Namespace) -> None: - """Either lists all example workflows, or sets up an example as a workflow to be run at root dir. - - :param [Namespace] `args`: parsed CLI arguments + """ + CLI command to set up or list Merlin example workflows. + + This function either lists all available example workflows or sets + up a specified example workflow to be run in the root directory. The + behavior is determined by the `workflow` argument. + + Args: + args: Parsed command-line arguments, which may include:\n + - `workflow`: The action to perform; should be "list" + to display all examples or the name of a specific example + workflow to set up. + - `path`: The directory where the example workflow + should be set up. Only applicable when `workflow` is not "list". """ if args.workflow == "list": print(list_examples()) @@ -400,12 +550,20 @@ def process_example(args: Namespace) -> None: setup_example(args.workflow, args.path) -def process_monitor(args): +def process_monitor(args: Namespace): """ - CLI command to monitor merlin workers and queues to keep - the allocation alive - - :param `args`: parsed CLI arguments + CLI command to monitor Merlin workers and queues to maintain + allocation status. + + This function periodically checks the status of Merlin workers and + the associated queues to ensure that the allocation remains active. + It includes a sleep interval to wait before each check, including + the initial one. + + Args: + args: Parsed command-line arguments, which may include:\n + - `sleep`: The duration (in seconds) to wait before + checking the queue status again. """ LOG.info("Monitor: checking queues ...") spec, _ = get_merlin_spec_with_override(args) @@ -422,9 +580,33 @@ def process_monitor(args): def process_server(args: Namespace): """ - Route to the correct function based on the command - given via the CLI + Route to the appropriate server function based on the command + specified via the CLI. + + This function processes commands related to server management, + directing the flow to the corresponding function for actions such + as initializing, starting, stopping, checking status, restarting, + or configuring the server. + + Args: + args: Parsed command-line arguments, which includes:\n + - `commands`: The server management command to execute. + Possible values are: + - `init`: Initialize the server. + - `start`: Start the server. + - `stop`: Stop the server. + - `status`: Check the server status. + - `restart`: Restart the server. + - `config`: Configure the server. """ + try: + lc_all_val = os.environ["LC_ALL"] + if lc_all_val != "C": + raise ValueError(f"The 'LC_ALL' environment variable is currently set to {lc_all_val} but it must be set to 'C'.") + except KeyError: + LOG.debug("The 'LC_ALL' environment variable was not set. Setting this to 'C'.") + os.environ["LC_ALL"] = "C" # Necessary for Redis to configure LOCALE + if args.commands == "init": init_server() elif args.commands == "start": @@ -443,7 +625,12 @@ def process_server(args: Namespace): # to split the function up but that wouldn't make much sense so we ignore it def setup_argparse() -> None: # pylint: disable=R0915 """ - Setup argparse and any CLI options we want available via the package. + Set up the command-line argument parser for the Merlin package. + + This function configures the ArgumentParser for the Merlin CLI, allowing users + to interact with various commands related to workflow management and task handling. + It includes options for running a workflow, restarting tasks, purging task queues, + generating configuration files, and managing/configuring the server. """ parser: HelpParser = HelpParser( prog="merlin", @@ -791,10 +978,17 @@ def setup_argparse() -> None: # pylint: disable=R0915 def generate_worker_touching_parsers(subparsers: ArgumentParser) -> None: - """All CLI arg parsers directly controlling or invoking workers are generated here. + """ + Generate command-line argument parsers for managing worker operations. + + This function sets up subparsers for CLI commands that directly control or invoke + workers in the context of the Merlin framework. It provides options for running, + querying, stopping, and monitoring workers associated with a Merlin YAML study + specification. - :param [ArgumentParser] `subparsers`: the subparsers needed for every CLI command that directly controls or invokes - workers. + Args: + subparsers: An instance of ArgumentParser for adding command-line subcommands + related to worker management. """ # merlin run-workers run_workers: ArgumentParser = subparsers.add_parser( @@ -939,11 +1133,18 @@ def generate_worker_touching_parsers(subparsers: ArgumentParser) -> None: monitor.set_defaults(func=process_monitor) -def generate_diagnostic_parsers(subparsers: ArgumentParser) -> None: - """All CLI arg parsers generally used diagnostically are generated here. +def generate_diagnostic_parsers(subparsers: ArgumentParser): + """ + Generate command-line argument parsers for diagnostic operations in the Merlin framework. + + This function sets up subparsers for CLI commands that handle diagnostics related + to Merlin jobs. It provides options to check the status of studies, gather queue + statistics, and retrieve configuration information, making it easier for users to + diagnose issues with their workflows. - :param [ArgumentParser] `subparsers`: the subparsers needed for every CLI command that handles diagnostics for a - Merlin job. + Args: + subparsers: An instance of ArgumentParser that will be used to add command-line + subcommands for various diagnostic activities. """ # merlin status status_cmd: ArgumentParser = subparsers.add_parser( @@ -1133,7 +1334,13 @@ def generate_diagnostic_parsers(subparsers: ArgumentParser) -> None: def main(): """ - High-level CLI operations. + Entry point for the Merlin command-line interface (CLI) operations. + + This function sets up the argument parser, handles command-line arguments, + initializes logging, and executes the appropriate function based on the + provided command. It ensures that the user receives help information if + no arguments are provided and performs error handling for any exceptions + that may occur during command execution. """ parser = setup_argparse() if len(sys.argv) == 1: diff --git a/merlin/merlin_templates.py b/merlin/merlin_templates.py deleted file mode 100644 index 5253c79e4..000000000 --- a/merlin/merlin_templates.py +++ /dev/null @@ -1,75 +0,0 @@ -############################################################################### -# Copyright (c) 2023, Lawrence Livermore National Security, LLC. -# Produced at the Lawrence Livermore National Laboratory -# Written by the Merlin dev team, listed in the CONTRIBUTORS file. -# -# -# LLNL-CODE-797170 -# All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. -# -# For details, see https://github.com/LLNL/merlin. -# -# Permission is hereby granted, free of charge, to any person obtaining a copy -# of this software and associated documentation files (the "Software"), to deal -# in the Software without restriction, including without limitation the rights -# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -# copies of the Software, and to permit persons to whom the Software is -# furnished to do so, subject to the following conditions: -# The above copyright notice and this permission notice shall be included in -# all copies or substantial portions of the Software. -# -# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -# SOFTWARE. -############################################################################### - -""" -This module handles the CLI for the deprecated `merlin-templates` command. -""" -import argparse -import logging -import sys - -from merlin.ascii_art import banner_small -from merlin.log_formatter import setup_logging - - -LOG = logging.getLogger("merlin-templates") -DEFAULT_LOG_LEVEL = "ERROR" - -# We disable all pylint errors in this file since this is deprecated anyways - - -def process_templates(args): # pylint: disable=W0613,C0116 - LOG.error("The command `merlin-templates` has been deprecated in favor of `merlin example`.") - - -def setup_argparse(): # pylint: disable=C0116 - parser = argparse.ArgumentParser( - prog="Merlin Examples", - description=banner_small, - formatter_class=argparse.RawDescriptionHelpFormatter, - ) - parser.set_defaults(func=process_templates) - return parser - - -def main(): # pylint: disable=C0116 - try: - parser = setup_argparse() - args = parser.parse_args() - setup_logging(logger=LOG, log_level=DEFAULT_LOG_LEVEL, colors=True) - args.func(args) - sys.exit() - except Exception as ex: # pylint: disable=W0718 - print(ex) - sys.exit(1) - - -if __name__ == "__main__": - main() diff --git a/merlin/router.py b/merlin/router.py index d9114bbcd..4ba186ec7 100644 --- a/merlin/router.py +++ b/merlin/router.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -38,9 +38,12 @@ import logging import os import time +from argparse import Namespace +from importlib import resources from typing import Dict, List, Tuple from merlin.exceptions import NoWorkersException +from merlin.spec.specification import MerlinSpec from merlin.study.celeryadapter import ( build_set_of_queues, check_celery_workers_processing, @@ -55,12 +58,7 @@ start_celery_workers, stop_celery_workers, ) - - -try: - from importlib import resources -except ImportError: - import importlib_resources as resources +from merlin.study.study import MerlinStudy LOG = logging.getLogger(__name__) @@ -70,12 +68,20 @@ # and try to resolve them -def run_task_server(study, run_mode=None): +def run_task_server(study: MerlinStudy, run_mode: str = None): """ - Creates the task server interface for communicating the tasks. - - :param `study`: The MerlinStudy object - :param `run_mode`: The type of run mode, e.g. local, batch + Creates the task server interface for managing task communications. + + This function determines which server to send tasks to. It checks if + Celery is set as the task server; if not, it logs an error message. + The run mode can be specified to determine how tasks should be executed. + + Args: + study (study.study.MerlinStudy): The study object representing the + current experiment setup, containing configuration details for + the task server. + run_mode: The type of run mode to use for task execution. This can + include options such as 'local' or 'batch'. """ if study.expanded_spec.merlin["resources"]["task_server"] == "celery": run_celery(study, run_mode) @@ -83,14 +89,36 @@ def run_task_server(study, run_mode=None): LOG.error("Celery is not specified as the task server!") -def launch_workers(spec, steps, worker_args="", disable_logs=False, just_return_command=False): +def launch_workers( + spec: MerlinSpec, + steps: List[str], + worker_args: str = "", + disable_logs: bool = False, + just_return_command: bool = False, +) -> str: """ - Launches workers for the specified study. - - :param `specs`: Tuple of (YAMLSpecification, MerlinSpec) - :param `steps`: The steps in the spec to tie the workers to - :param `worker_args`: Optional arguments for the workers - :param `just_return_command`: Don't execute, just return the command + Launches workers for the specified study based on the provided + specification and steps. + + This function checks if Celery is configured as the task server + and initiates the specified workers accordingly. It provides options + for additional worker arguments, logging control, and command-only + execution without launching the workers. + + Args: + spec (spec.specification.MerlinSpec): Specification details + necessary for launching the workers. + steps: The specific steps in the specification that the workers + will be associated with. + worker_args: Additional arguments to be passed to the workers. + Defaults to an empty string. + disable_logs: Flag to disable logging during worker execution. + Defaults to False. + just_return_command: If True, the function will not execute the + command but will return it instead. Defaults to False. + + Returns: + A string of the worker launch command(s). """ if spec.merlin["resources"]["task_server"] == "celery": # pylint: disable=R1705 # Start workers @@ -101,15 +129,30 @@ def launch_workers(spec, steps, worker_args="", disable_logs=False, just_return_ return "No workers started" -def purge_tasks(task_server, spec, force, steps): +def purge_tasks(task_server: str, spec: MerlinSpec, force: bool, steps: List[str]) -> int: """ - Purges all tasks. - - :param `task_server`: The task server from which to purge tasks. - :param `spec`: A MerlinSpec object - :param `force`: Purge without asking for confirmation - :param `steps`: Space-separated list of stepnames defining queues to purge, - default is all steps + Purges all tasks from the specified task server. + + This function removes tasks from the designated queues associated + with the specified steps. It operates without confirmation if + the `force` parameter is set to True. The function logs the + steps being purged and checks if Celery is the configured task + server before proceeding. + + Args: + task_server: The task server from which to purge tasks. + spec (spec.specification.MerlinSpec): A + [`MerlinSpec`][spec.specification.MerlinSpec] object + containing the configuration needed to generate queue + specifications. + force: If True, purge the tasks without any confirmation prompt. + steps: A space-separated list of step names that define + which queues to purge. If not specified, defaults to purging + all steps. + + Returns: + The result of the purge operation; -1 if the task server is not + supported (i.e., not Celery). """ LOG.info(f"Purging queues for steps = {steps}") @@ -124,12 +167,18 @@ def purge_tasks(task_server, spec, force, steps): def dump_queue_info(task_server: str, query_return: List[Tuple[str, int, int]], dump_file: str): """ - Format the information we're going to dump in a way that the Dumper class can - understand and add a timestamp to the info. - - :param task_server: The task server from which to query queues - :param query_return: The output of `query_queues` - :param dump_file: The filepath of the file we'll dump queue info to + Formats and dumps queue information for the specified task server. + + This function prepares the queue data returned from the queue + query and formats it in a way that the [`Dumper`][common.dumper.Dumper] + class can process. It also adds a timestamp to the information before + dumping it to the specified file. + + Args: + task_server: The task server from which to query queues. + query_return: The output from the [`query_queues`][router.query_queues] + function, containing tuples of queue information. + dump_file: The filepath where the queue information will be dumped. """ if task_server == "celery": dump_celery_queue_info(query_return, dump_file) @@ -139,19 +188,35 @@ def dump_queue_info(task_server: str, query_return: List[Tuple[str, int, int]], def query_queues( task_server: str, - spec: "MerlinSpec", # noqa: F821 + spec: MerlinSpec, steps: List[str], specific_queues: List[str], verbose: bool = True, -): +) -> Dict[str, Dict[str, int]]: """ - Queries status of queues. - - :param task_server: The task server from which to query queues - :param spec: A MerlinSpec object or None - :param steps: Spaced-separated list of stepnames to query. Default is all - :param specific_queues: A list of queue names to query or None - :param verbose: A bool to determine whether to output log statements or not + Queries the status of queues from the specified task server. + + This function checks the status of queues tied to a given task server, + building a list of queues based on the provided steps and specific queue + names. It supports querying Celery task servers and returns the results + in a structured format. Logging behavior can be controlled with the verbose + parameter. + + Args: + task_server: The task server from which to query queues. + spec (spec.specification.MerlinSpec): A + [`MerlinSpec`][spec.specification.MerlinSpec] object used to define + the configuration of queues. Can also be None. + steps: A space-separated list of step names to query. Default is to query + all available steps if this is empty. + specific_queues: A list of specific queue names to query. Can be empty or + None to query all relevant queues. + verbose: If True, enables logging of query operations. Defaults to True. + + Returns: + A dictionary where the keys are queue names and the values are dictionaries + containing the number of workers (consumers) and tasks (jobs) attached + to each queue. """ if task_server == "celery": # pylint: disable=R1705 # Build a set of queues to query and query them @@ -159,14 +224,18 @@ def query_queues( return query_celery_queues(queues) else: LOG.error("Celery is not specified as the task server!") - return [] + return {} -def query_workers(task_server, spec_worker_names, queues, workers_regex): +def query_workers(task_server: str, spec_worker_names: List[str], queues: List[str], workers_regex: str): """ - Gets info from workers. + Retrieves information from workers associated with the specified task server. - :param `task_server`: The task server to query. + Args: + task_server: The task server to query. + spec_worker_names: A list of specific worker names to query. + queues: A list of queues to search for associated workers. + workers_regex: A regex pattern used to filter worker names during the query. """ LOG.info("Searching for workers...") @@ -176,12 +245,17 @@ def query_workers(task_server, spec_worker_names, queues, workers_regex): LOG.error("Celery is not specified as the task server!") -def get_workers(task_server): - """Get all workers. +def get_workers(task_server: str) -> List[str]: + """ + This function queries the designated task server to obtain a list of all + workers that are currently connected. - :param `task_server`: The task server to query. - :return: A list of all connected workers - :rtype: list + Args: + task_server: The task server to query. + + Returns: + A list of all connected workers. If the task server is not supported, + an empty list is returned. """ if task_server == "celery": # pylint: disable=R1705 return get_workers_from_app() @@ -190,14 +264,17 @@ def get_workers(task_server): return [] -def stop_workers(task_server, spec_worker_names, queues, workers_regex): +def stop_workers(task_server: str, spec_worker_names: List[str], queues: List[str], workers_regex: str): """ - Stops workers. - - :param `task_server`: The task server from which to stop workers. - :param `spec_worker_names`: Worker names to stop, drawn from a spec. - :param `queues` : The queues to stop - :param `workers_regex` : Regex for workers to stop + This function sends a command to stop workers that match the specified + criteria from the designated task server. + + Args: + task_server: The task server from which to stop workers. + spec_worker_names: A list of worker names to stop, as defined + in a specification. + queues: A list of queues from which to stop associated workers. + workers_regex: A regex pattern used to filter the workers to stop. """ LOG.info("Stopping workers...") @@ -208,14 +285,23 @@ def stop_workers(task_server, spec_worker_names, queues, workers_regex): LOG.error("Celery is not specified as the task server!") -def create_config(task_server: str, config_dir: str, broker: str, test: str) -> None: +def create_config(task_server: str, config_dir: str, broker: str, test: str): """ - Create a config for the given task server. - - :param [str] `task_server`: The task server from which to stop workers. - :param [str] `config_dir`: Optional directory to install the config. - :param [str] `broker`: string indicated the broker, used to check for redis. - :param [str] `test`: string indicating if the app.yaml is used for testing. + Create a configuration app.yaml that Merlin will use to connect to the + specified task server. + + This function generates a configuration file for the given task server. + It creates the necessary directories if they do not exist and determines + the appropriate configuration file based on the provided broker and testing + parameters. + + Args: + task_server: The task server for which to create the configuration. + config_dir: The directory where the configuration files will be installed. + If the directory does not exist, it will be created. + broker: A string indicating the broker type. + test: A string that indicates whether the application should use a test + configuration file. If set, a test configuration is created. """ if test: LOG.info("Creating test config ...") @@ -240,10 +326,20 @@ def create_config(task_server: str, config_dir: str, broker: str, test: str) -> def get_active_queues(task_server: str) -> Dict[str, List[str]]: """ - Get a dictionary of active queues and the workers attached to these queues. + Retrieve a dictionary of active queues and their associated workers for the specified task server. + + This function queries the given task server for its active queues and gathers + information about which workers are currently monitoring these queues. It supports + the 'celery' task server and returns a structured dictionary containing the queue + names as keys and lists of worker names as values. + + Args: + task_server: The task server to query for active queues. - :param `task_server`: The task server to query for active queues - :returns: A dict where keys are queue names and values are a list of workers watching them + Returns: + A dictionary where:\n + - The keys are the names of the active queues. + - The values are lists of worker names that are currently attached to those queues. """ active_queues = {} @@ -257,15 +353,25 @@ def get_active_queues(task_server: str) -> Dict[str, List[str]]: return active_queues -def wait_for_workers(sleep: int, task_server: str, spec: "MerlinSpec"): # noqa +def wait_for_workers(sleep: int, task_server: str, spec: MerlinSpec): # noqa """ - Wait on workers to start up. Check on worker start 10 times with `sleep` seconds between - each check. If no workers are started in time, raise an error to kill the monitor (there - was likely an issue with the task server that caused worker launch to fail). - - :param `sleep`: An integer representing the amount of seconds to sleep between each check - :param `task_server`: The task server from which to look for workers - :param `spec`: A MerlinSpec object representing the spec we're monitoring + Wait for workers to start up by checking their status at regular intervals. + + This function monitors the specified task server for the startup of worker processes. + It checks for the existence of the expected workers up to 10 times, sleeping for a + specified number of seconds between each check. If no workers are detected after + the maximum number of attempts, it raises an error to terminate the monitoring + process, indicating a potential issue with the task server. + + Args: + sleep: The number of seconds to pause between each check for worker status. + task_server: The task server from which to query for worker status. + spec (spec.specification.MerlinSpec): An instance of the + [`MerlinSpec`][spec.specification.MerlinSpec] class that contains the + specification for the workers being monitored. + + Raises: + NoWorkersException: If no workers are detected after the maximum number of checks. """ # Get the names of the workers that we're looking for worker_names = spec.get_worker_names() @@ -297,9 +403,12 @@ def check_workers_processing(queues_in_spec: List[str], task_server: str) -> boo """ Check if any workers are still processing tasks by querying the task server. - :param `queues_in_spec`: A list of queues to check if tasks are still active in - :param `task_server`: The task server from which to query - :returns: True if workers are still processing tasks, False otherwise + Args: + queues_in_spec: A list of queue names to check for active tasks. + task_server: The task server from which to query the processing status. + + Returns: + True if workers are still processing tasks, False otherwise. """ result = False @@ -313,13 +422,23 @@ def check_workers_processing(queues_in_spec: List[str], task_server: str) -> boo return result -def check_merlin_status(args: "Namespace", spec: "MerlinSpec") -> bool: # noqa +def check_merlin_status(args: Namespace, spec: MerlinSpec) -> bool: """ - Function to check merlin workers and queues to keep the allocation alive + Function to check Merlin workers and queues to keep the allocation alive. + + This function monitors the status of workers and jobs within the specified task server + and the provided Merlin specification. It checks for active tasks and workers, ensuring + that the allocation remains valid. + + Args: + args: Parsed command-line interface arguments, including task server + specifications and sleep duration. + spec (spec.specification.MerlinSpec): The parsed spec.yaml as a + [`MerlinSpec`][spec.specification.MerlinSpec] object, containing queue + and worker definitions. - :param `args`: parsed CLI arguments - :param `spec`: the parsed spec.yaml as a MerlinSpec object - :returns: True if there are still tasks being processed, False otherwise + Returns: + True if there are still tasks being processed, False otherwise. """ # Initialize the variable to track if there are still active tasks active_tasks = False diff --git a/merlin/server/__init__.py b/merlin/server/__init__.py index 522d67d1d..2953df4b7 100644 --- a/merlin/server/__init__.py +++ b/merlin/server/__init__.py @@ -6,8 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. - +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -28,3 +27,13 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `server` package defines the functionality for managing a containerized server +through Merlin. + +Modules: + server_commands.py: Main functions to interact with the server. + server_config.py: Server configuration functions. + server_util.py: Defines the structure of our server configurations. +""" diff --git a/merlin/server/server_commands.py b/merlin/server/server_commands.py index be2b944a0..bb505b6e9 100644 --- a/merlin/server/server_commands.py +++ b/merlin/server/server_commands.py @@ -8,7 +8,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -48,15 +48,15 @@ pull_server_config, pull_server_image, ) -from merlin.server.server_util import AppYaml, RedisConfig, RedisUsers +from merlin.server.server_util import AppYaml, RedisConfig, RedisUsers, ServerConfig LOG = logging.getLogger("merlin") -def init_server() -> None: +def init_server(): """ - Initialize merlin server by checking and initializing main configuration directory + Initializes the Merlin server by setting up the main configuration directory and local server configuration. """ @@ -70,17 +70,21 @@ def init_server() -> None: LOG.info("Merlin server initialization successful.") -# Pylint complains that there's too many branches in this function but -# it looks clean to me so we'll ignore it -def config_server(args: Namespace) -> None: # pylint: disable=R0912 +def apply_config_changes(server_config: ServerConfig, args: Namespace): """ - Process the merlin server config flags to make changes and edits to appropriate configurations - based on the input passed in by the user. + Applies configuration changes to the Merlin server based on user-provided arguments. + + This function modifies the Redis configuration and user settings as specified by + the user through the provided arguments. It updates various Redis settings, such as + IP address, port, password, directory, snapshot settings, append mode, and file paths. + If any changes are made, the updated configuration is written to the appropriate files. + + Args: + server_config (server.server_util.ServerConfig): An instance of `ServerConfig` containing + all the necessary configuration values for the server. + args (Namespace): An argparse `Namespace` object containing user-provided arguments + from the argument parser. """ - server_config = pull_server_config() - if not server_config: - LOG.error('Try to run "merlin server init" again to reinitialize values.') - return False redis_config = RedisConfig(server_config.container.get_config_path()) redis_config.set_ip_address(args.ipaddress) @@ -98,9 +102,7 @@ def config_server(args: Namespace) -> None: # pylint: disable=R0912 redis_config.set_directory(args.directory) - redis_config.set_snapshot_seconds(args.snapshot_seconds) - - redis_config.set_snapshot_changes(args.snapshot_changes) + redis_config.set_snapshot(seconds=args.snapshot_seconds, changes=args.snapshot_changes) redis_config.set_snapshot_file(args.snapshot_file) @@ -116,14 +118,33 @@ def config_server(args: Namespace) -> None: # pylint: disable=R0912 else: LOG.info("Add changes to config file and exisiting containers.") - server_config = pull_server_config() - if not server_config: + +# Pylint complains that there's too many branches in this function but +# it looks clean to me so we'll ignore it +def config_server(args: Namespace): # pylint: disable=R0912 + """ + Processes the Merlin server configuration flags to make changes and edits + to appropriate configurations based on user input. + + Args: + args (Namespace): An argparse `Namespace` object containing user-provided arguments + from the argument parser. + """ + server_config_before_changes = pull_server_config() + if not server_config_before_changes: + LOG.error('Try to run "merlin server init" again to reinitialize values.') + return False + + apply_config_changes(server_config_before_changes, args) + + server_config_after_changes = pull_server_config() + if not server_config_after_changes: LOG.error('Try to run "merlin server init" again to reinitialize values.') return False # Read the user from the list of avaliable users - redis_users = RedisUsers(server_config.container.get_user_file_path()) - redis_config = RedisConfig(server_config.container.get_config_path()) + redis_users = RedisUsers(server_config_after_changes.container.get_user_file_path()) + redis_config = RedisConfig(server_config_after_changes.container.get_config_path()) if args.add_user is not None: # Log the user in a file @@ -152,12 +173,15 @@ def config_server(args: Namespace) -> None: # pylint: disable=R0912 return None -def status_server() -> None: +def status_server(): """ - Get the server status of the any current running containers for merlin server + Retrieves and displays the current status of the Merlin server. + + This function checks the status of any running containers for the Merlin server + and logs appropriate messages based on the server's state. """ current_status = get_server_status() - if current_status == ServerStatus.NOT_INITALIZED: + if current_status == ServerStatus.NOT_INITIALIZED: LOG.info("Merlin server has not been initialized.") LOG.info("Please initalize server by running 'merlin server init'") elif current_status == ServerStatus.MISSING_CONTAINER: @@ -169,15 +193,20 @@ def status_server() -> None: LOG.info("Merlin server is running.") -def start_server() -> bool: # pylint: disable=R0911 +def check_for_not_running_server() -> bool: """ - Start a merlin server container using singularity. - :return:: True if server was successful started and False if failed. + Checks if the Merlin server status is `NOT_RUNNING` before starting a new server. + + If the server status is anything other than `NOT_RUNNING`, logs an appropriate + error message to inform the user. + + Returns: + True if the server status is `NOT_RUNNING`, False otherwise. """ current_status = get_server_status() uninitialized_err = "Merlin server has not been intitialized. Please run 'merlin server init' first." status_errors = { - ServerStatus.NOT_INITALIZED: uninitialized_err, + ServerStatus.NOT_INITIALIZED: uninitialized_err, ServerStatus.MISSING_CONTAINER: uninitialized_err, ServerStatus.RUNNING: """Merlin server already running. Stop current server with 'merlin server stop' before attempting to start a new server.""", @@ -187,11 +216,26 @@ def start_server() -> bool: # pylint: disable=R0911 LOG.info(status_errors[current_status]) return False - server_config = pull_server_config() - if not server_config: - LOG.error('Try to run "merlin server init" again to reinitialize values.') - return False + return True + + +def start_container(server_config: ServerConfig) -> subprocess.Popen: + """ + Starts a container based on the provided server configuration. + This function uses the server configuration to locate the necessary image and + configuration files, validates their existence, and starts the container using + a subprocess. + + Args: + server_config (server.server_util.ServerConfig): An instance of `ServerConfig` + containing information about the server to start, including paths to the image + and configuration files. + + Returns: + A subprocess object representing the running container process, or `None` if + required files are missing. + """ image_path = server_config.container.get_image_path() config_path = server_config.container.get_config_path() path_errors = { @@ -202,7 +246,7 @@ def start_server() -> bool: # pylint: disable=R0911 for path in (image_path, config_path): if not os.path.exists(path): LOG.error(f"Unable to find {path_errors[path]} at {path}") - return False + return None # Pylint wants us to use with here but we don't need that process = subprocess.Popen( # pylint: disable=R1732 @@ -222,6 +266,27 @@ def start_server() -> bool: # pylint: disable=R0911 time.sleep(1) + return process + + +def server_started(process: subprocess.Popen, server_config: ServerConfig) -> bool: + """ + Verifies that the server started by [`start_container`][server.server_commands.start_container] + is running properly. + + This function checks the Redis output to ensure the server started successfully, + creates a process file for the container, and validates that the server status + is `RUNNING`. + + Args: + process (subprocess.Popen): The subprocess object representing the container + process started by `start_container`. + server_config (server.server_util.ServerConfig): An instance of `ServerConfig` + containing information about the Redis server configuration. + + Returns: + True if the server started successfully, False otherwise. + """ redis_start, redis_out = parse_redis_output(process.stdout) if not redis_start: @@ -243,6 +308,36 @@ def start_server() -> bool: # pylint: disable=R0911 LOG.info(f"Server started with PID {str(process.pid)}.") LOG.info(f'Merlin server operating on "{redis_out["hostname"]}" and port "{redis_out["port"]}".') + return True + + +def start_server() -> bool: # pylint: disable=R0911 + """ + Starts a Merlin server container. + + This function performs several steps to start the server, including checking + for an existing non-running server, pulling the server configuration, starting + the container, verifying the server startup, and applying Redis user and + configuration settings. It also generates a new `app.yaml` file for the server + configuration. + + Returns: + True if the server was successfully started, False otherwise. + """ + if not check_for_not_running_server(): + return False + + server_config = pull_server_config() + if not server_config: + LOG.error('Try to run "merlin server init" again to reinitialize values.') + return False + + process = start_container(server_config) + if process is None: + return False + + if not server_started(process, server_config): + return False redis_users = RedisUsers(server_config.container.get_user_file_path()) redis_config = RedisConfig(server_config.container.get_config_path()) @@ -259,10 +354,17 @@ def start_server() -> bool: # pylint: disable=R0911 return True -def stop_server(): +def stop_server() -> bool: """ - Stop running merlin server containers. - :return:: True if server was stopped successfully and False if failed. + Stops a running Merlin server container. + + This function checks the current server status, retrieves the server configuration, + and attempts to terminate the running server process. If successful, the server + process is stopped, and the function returns True. Otherwise, it logs errors + and returns False. + + Returns: + True if the server was successfully stopped, False otherwise. """ if get_server_status() != ServerStatus.RUNNING: LOG.info("There is no instance of merlin server running.") @@ -307,8 +409,13 @@ def stop_server(): def restart_server() -> bool: """ - Restart a running merlin server instance. - :return:: True if server was restarted successfully and False if failed. + Restarts a running Merlin server instance. + + This function stops the currently running Merlin server and then starts it again. + If the server is not running, it logs a message and returns False. + + Returns: + True if the server was successfully restarted, False otherwise. """ if get_server_status() != ServerStatus.RUNNING: LOG.info("Merlin server is not currently running.") diff --git a/merlin/server/server_config.py b/merlin/server/server_config.py index f58c7567a..d27874de9 100644 --- a/merlin/server/server_config.py +++ b/merlin/server/server_config.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -35,8 +35,9 @@ import random import string import subprocess +from importlib import resources from io import BufferedReader -from typing import Tuple +from typing import Dict, Tuple import yaml @@ -52,12 +53,6 @@ ) -try: - from importlib import resources -except ImportError: - import importlib_resources as resources - - LOG = logging.getLogger("merlin") # Default values for configuration @@ -73,11 +68,18 @@ class ServerStatus(enum.Enum): """ - Different states in which the server can be in. + Represents different states that a server can be in. + + Attributes: + RUNNING (int): Indicates the server is running and operational. Numeric value: 0. + NOT_INITIALIZED (int): Indicates the server has not been initialized yet. Numeric value: 1. + MISSING_CONTAINER (int): Indicates the server is missing a required container. Numeric value: 2. + NOT_RUNNING (int): Indicates the server is not currently running. Numeric value: 3. + ERROR (int): Indicates the server encountered an error. Numeric value: 4. """ RUNNING = 0 - NOT_INITALIZED = 1 + NOT_INITIALIZED = 1 MISSING_CONTAINER = 2 NOT_RUNNING = 3 ERROR = 4 @@ -85,15 +87,23 @@ class ServerStatus(enum.Enum): def generate_password(length, pass_command: str = None) -> str: """ - Function for generating passwords for redis container. If a specified command is given - then a password would be generated with the given command. If not a password will be - created by combining a string a characters based on the given length. + Generates a password for a Redis container. + + If a specific command is provided, the password will be generated using the output + of the given command. Otherwise, a random password will be created by combining + characters (letters, digits, and special symbols) based on the specified length. + + Args: + length (int): The desired length of the password. + pass_command (str, optional): A shell command to generate the password. + If provided, the command's output will be used as the password. - :return:: string value with given length + Returns: + The generated password. """ if pass_command: - process = subprocess.run(pass_command.split(), shell=True, stdout=subprocess.PIPE) - return process.stdout + process = subprocess.run(pass_command, shell=True, capture_output=True, text=True) + return process.stdout.strip() characters = list(string.ascii_letters + string.digits + "!@#$%^&*()") @@ -109,17 +119,26 @@ def generate_password(length, pass_command: str = None) -> str: def parse_redis_output(redis_stdout: BufferedReader) -> Tuple[bool, str]: """ - Parse the redis output for a the redis container. It will get all the necessary information - from the output and returns a dictionary of those values. + Parses the Redis output from a Redis container. - :return:: two values is_successful, dictionary of values from redis output + This function processes the Redis container's output to extract necessary information, + such as configuration details and server state. It determines whether the server was + successfully initialized and ready to accept connections, or if an error occurred. + + Args: + redis_stdout (BufferedReader): A buffered reader object containing the Redis container's output. + + Returns: + A tuple containing:\n + - A boolean indicating whether the server was successfully initialized and ready. + - A dictionary containing parsed configuration values if successful, or an error message otherwise. """ if redis_stdout is None: return False, "None passed as redis output" server_init = False redis_config = {} line = redis_stdout.readline() - while line != "" or line is not None: + while line != b"" and line is not None: if not server_init: values = [ln for ln in line.split() if b"=" in ln] for val in values: @@ -136,13 +155,44 @@ def parse_redis_output(redis_stdout: BufferedReader) -> Tuple[bool, str]: return False, "Reached end of redis output without seeing 'Ready to accept connections'" +def copy_container_command_files(config_dir: str) -> bool: + """ + Copies YAML files containing command instructions for container types to the specified configuration directory. + + Args: + config_dir (str): The path to the configuration directory where the YAML files will be copied. + + Returns: + True if all files are successfully copied or already exist. False otherwise. + """ + files = [i + ".yaml" for i in CONTAINER_TYPES] + for file in files: + file_path = os.path.join(config_dir, file) + if os.path.exists(file_path): + LOG.info(f"{file} already exists.") + continue + LOG.info(f"Copying file {file} to configuration directory.") + try: + with resources.path("merlin.server", file) as config_file: + with open(file_path, "w") as outfile, open(config_file, "r") as infile: + outfile.write(infile.read()) + except OSError: + LOG.error(f"Destination location {config_dir} is not writable.") + return False + return True + + def create_server_config() -> bool: """ - Create main configuration file for merlin server in the - merlin configuration directory. If a configuration already - exists it will not replace the current configuration and exit. + Creates the main configuration file for the Merlin server in the Merlin configuration directory. + + This function checks for the existence of the Merlin configuration directory and creates a default + server configuration if none exists. It also copies necessary container command files, applies the + server configuration to `app.yaml`, and initializes the server configuration directory. If the + configuration already exists, it will not overwrite it. - :return:: True if success and False if fail + Returns: + True if the configuration is successfully created and applied. False otherwise. """ if not os.path.exists(MERLIN_CONFIG_DIR): LOG.error(f"Unable to find main merlin configuration directory at {MERLIN_CONFIG_DIR}") @@ -158,20 +208,8 @@ def create_server_config() -> bool: LOG.error(err) return False - files = [i + ".yaml" for i in CONTAINER_TYPES] - for file in files: - file_path = os.path.join(config_dir, file) - if os.path.exists(file_path): - LOG.info(f"{file} already exists.") - continue - LOG.info(f"Copying file {file} to configuration directory.") - try: - with resources.path("merlin.server", file) as config_file: - with open(file_path, "w") as outfile, open(config_file, "r") as infile: - outfile.write(infile.read()) - except OSError: - LOG.error(f"Destination location {config_dir} is not writable.") - return False + if not copy_container_command_files(config_dir): + return False # Load Merlin Server Configuration and apply it to app.yaml with resources.path("merlin.server", MERLIN_SERVER_CONFIG) as merlin_server_config: @@ -197,7 +235,11 @@ def create_server_config() -> bool: def config_merlin_server(): """ - Configurate the merlin server with configurations such as username password and etc. + Configures the Merlin server with necessary settings, including username and password. + + This function sets up the Merlin server by generating and storing a password file, creating a user file, + and configuring Redis settings. If the password or user files already exist, it skips the respective + setup steps. The function ensures that default and environment-specific users are added to the user file. """ server_config = pull_server_config() @@ -209,9 +251,6 @@ def config_merlin_server(): if os.path.exists(pass_file): LOG.info("Password file already exists. Skipping password generation step.") else: - # if "pass_command" in server_config["container"]: - # password = generate_password(PASSWORD_LENGTH, server_config["container"]["pass_command"]) - # else: password = generate_password(PASSWORD_LENGTH) with open(pass_file, "w+") as f: # pylint: disable=C0103 @@ -238,10 +277,14 @@ def config_merlin_server(): def pull_server_config() -> ServerConfig: """ - Pull the main configuration file and corresponding format configuration file - as well. Returns the values as a dictionary. + Retrieves the main configuration file and its corresponding format configuration file for the Merlin server. - :return: A instance of ServerConfig containing all the necessary configuration values. + This function reads the `app.yaml` configuration file and additional format-specific configuration files + to construct a complete configuration dictionary. It validates the presence of required keys in the format + and process configurations. If any required configuration is missing, an error is logged and `None` is returned. + + Returns: + An instance of [`ServerConfig`][server.server_util.ServerConfig] containing all necessary configuration values. """ return_data = {} format_needed_keys = ["command", "run_command", "stop_command", "pull_command"] @@ -285,9 +328,14 @@ def pull_server_config() -> ServerConfig: def pull_server_image() -> bool: """ - Fetch the server image using singularity. + Fetches the server image and ensures the necessary configuration files are in place. + + This function retrieves the server image from a specified URL and saves it locally if it does not already exist. + Additionally, it copies the default Redis configuration file to the appropriate location if it is missing. + The function relies on the server configuration to determine the image URL, image path, and configuration file details. - :return:: True if success and False if fail + Returns: + True if the server image and configuration file are successfully set up, False if an error occurs. """ server_config = pull_server_config() if not server_config: @@ -318,8 +366,8 @@ def pull_server_image() -> bool: with resources.path("merlin.server", config_file) as file: with open(os.path.join(config_dir, config_file), "w") as outfile, open(file, "r") as infile: outfile.write(infile.read()) - except OSError: - LOG.error(f"Destination location {config_dir} is not writable.") + except OSError as exc: + LOG.error(f"Destination location {config_dir} is not writable. Raised from:\n{exc}") return False else: LOG.info("Redis configuration file already exist.") @@ -327,22 +375,27 @@ def pull_server_image() -> bool: return True -def get_server_status(): +def get_server_status() -> ServerStatus: """ - Determine the status of the current server. - This function can be used to check if the servers - have been initalized, started, or stopped. - - :param `server_dir`: location of all server related files. - :param `image_name`: name of the image when fetched. - :return:: A enum value of ServerStatus describing its current state. + Determines the current status of the server. + + This function checks the server's state by verifying the existence of necessary files, + including configuration files, the container image, and the process file. It also checks + if the server process is actively running. + + Returns: + An enum value representing the server's current state:\n + - `ServerStatus.NOT_INITIALIZED`: The server has not been initialized. + - `ServerStatus.MISSING_CONTAINER`: The server container image is missing. + - `ServerStatus.NOT_RUNNING`: The server process is not running. + - `ServerStatus.RUNNING`: The server is actively running. """ server_config = pull_server_config() if not server_config: - return ServerStatus.NOT_INITALIZED + return ServerStatus.NOT_INITIALIZED if not os.path.exists(server_config.container.get_config_dir()): - return ServerStatus.NOT_INITALIZED + return ServerStatus.NOT_INITIALIZED if not os.path.exists(server_config.container.get_image_path()): return ServerStatus.MISSING_CONTAINER @@ -365,10 +418,18 @@ def get_server_status(): return ServerStatus.RUNNING -def check_process_file_format(data: dict) -> bool: +def check_process_file_format(data: Dict) -> bool: """ - Check to see if the process file has the correct format and contains the expected key values. - :return:: True if success and False if fail + Validates the format of a process file. + + This function checks if the given process file data (in dictionary format) contains all the + required keys: "parent_pid", "image_pid", "port", and "hostname". + + Args: + data (Dict): The process file data to validate. + + Returns: + True if the process file contains all required keys, False otherwise. """ required_keys = ["parent_pid", "image_pid", "port", "hostname"] for key in required_keys: @@ -377,11 +438,19 @@ def check_process_file_format(data: dict) -> bool: return True -def pull_process_file(file_path: str) -> dict: +def pull_process_file(file_path: str) -> Dict: """ - Pull the data from the process file. If one is found returns the data in a dictionary - if not returns None - :return:: Data containing in process file. + Reads and parses data from a process file. + + This function attempts to load the contents of a process file located at the specified + file path. If the file exists and its format is valid, the data is returned as a dictionary. + If the format is invalid or the file cannot be processed, `None` is returned. + + Args: + file_path (str): The path to the process file. + + Returns: + A dictionary containing the data from the process file if the format is valid. """ with open(file_path, "r") as f: # pylint: disable=C0103 data = yaml.load(f, yaml.Loader) @@ -390,10 +459,22 @@ def pull_process_file(file_path: str) -> dict: return None -def dump_process_file(data: dict, file_path: str): +def dump_process_file(data: Dict, file_path: str) -> bool: """ - Dump the process data from the dictionary to the specified file path. - :return:: True if success and False if fail + Writes process data to a specified file. + + This function takes a dictionary containing process data and writes it to the specified + file path in YAML format. Before writing, the function validates the format of the data. + If the data format is invalid, the function returns `False`. If the operation is successful, + it returns `True`. + + Args: + data (Dict): The process data to be written to the file. + file_path (str): The path to the file where the data will be written. + + Returns: + True if the data is successfully written to the file, False if the data + format is invalid or the operation fails. """ if not check_process_file_format(data): return False diff --git a/merlin/server/server_util.py b/merlin/server/server_util.py index aa7c2765b..24ac4d993 100644 --- a/merlin/server/server_util.py +++ b/merlin/server/server_util.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -32,6 +32,7 @@ import hashlib import logging import os +from typing import Dict, List import redis import yaml @@ -50,7 +51,17 @@ def valid_ipv4(ip: str) -> bool: # pylint: disable=C0103 """ - Checks valid ip address + Validates whether a given string is a valid IPv4 address. + + An IPv4 address consists of four octets separated by dots, where each octet + is a number between 0 and 255 (inclusive). This function checks if the input + string meets these criteria. + + Args: + ip: The string to validate as an IPv4 address. + + Returns: + True if the input string is a valid IPv4 address, False otherwise. """ if not ip: return False @@ -60,7 +71,7 @@ def valid_ipv4(ip: str) -> bool: # pylint: disable=C0103 return False for i in arr: - if int(i) < 0 and int(i) > 255: + if int(i) < 0 or int(i) > 255: return False return True @@ -68,7 +79,16 @@ def valid_ipv4(ip: str) -> bool: # pylint: disable=C0103 def valid_port(port: int) -> bool: """ - Checks valid network port + Validates whether a given integer is a valid network port number. + + A valid network port number is an integer in the range 1 to 65535 (inclusive). + This function checks if the provided port falls within this range. + + Args: + port: The port number to validate. + + Returns: + True if the port is valid, False otherwise. """ if 0 < port < 65536: return True @@ -78,107 +98,605 @@ def valid_port(port: int) -> bool: # Pylint complains about too many instance variables but it's necessary here so ignore class ContainerConfig: # pylint: disable=R0902 """ - ContainerConfig provides interface for parsing and interacting with the container value specified within - the merlin_server.yaml configuration file. Dictionary of the config values should be passed when initialized - to parse values. This can be done after parsing yaml to data dictionary. - If there are missing values within the configuration it will be populated with default values for - singularity container. - - Configuration contains values for setting up containers and storing values specific to each container. - Values that are stored consist of things within the local configuration directory as different runs - can have differnt configuration values. + A class for parsing and interacting with container configuration values. + + The `ContainerConfig` class provides an interface for handling container-related + configuration values specified in the `merlin_server.yaml` file. It initializes + with a dictionary of configuration values, allowing for default values to be + populated for a Singularity container if any values are missing. + + The configuration contains values for setting up containers and storing alues specific + to each container. These values are used to manage container setup and store configuration + details specific to each container run. + + Attributes: + FORMAT (str): Default container format (e.g., "singularity"). + IMAGE_TYPE (str): Default image type (e.g., "redis"). + IMAGE_NAME (str): Default image name (e.g., "redis_latest.sif"). + REDIS_URL (str): Default URL for the container image (e.g., "docker://redis"). + CONFIG_FILE (str): Default name of the configuration file (e.g., "redis.conf"). + CONFIG_DIR (str): Default path to the configuration directory. + PROCESS_FILE (str): Default name of the process file (e.g., "merlin_server.pf"). + PASSWORD_FILE (str): Default name of the password file (e.g., "redis.pass"). + USERS_FILE (str): Default name of the users file (e.g., "redis.users"). + + format (str): Container format, initialized with the provided data or default. + image_type (str): Image type, initialized with the provided data or default. + image (str): Image name, initialized with the provided data or default. + url (str): Image URL, initialized with the provided data or default. + config (str): Configuration file name, initialized with the provided data or default. + config_dir (str): Configuration directory path, initialized with the provided data or default. + pfile (str): Process file name, initialized with the provided data or default. + pass_file (str): Password file name, initialized with the provided data or default. + user_file (str): Users file name, initialized with the provided data or default. + + Methods: + __eq__: Determines equality between two `ContainerConfig` instances. + get_format: Returns the container format. + get_image_type: Returns the image type. + get_image_name: Returns the image name. + get_image_url: Returns the image URL. + get_image_path: Returns the full path to the image file. + get_config_name: Returns the name of the configuration file. + get_config_path: Returns the full path to the configuration file. + get_config_dir: Returns the configuration directory path. + get_pfile_name: Returns the name of the process file. + get_pfile_path: Returns the full path to the process file. + get_pass_file_name: Returns the name of the password file. + get_pass_file_path: Returns the full path to the password file. + get_user_file_name: Returns the name of the users file. + get_user_file_path: Returns the full path to the users file. + get_container_password: Reads and returns the container password from the password file. """ # Default values for configuration - FORMAT = "singularity" - IMAGE_TYPE = "redis" - IMAGE_NAME = "redis_latest.sif" - REDIS_URL = "docker://redis" - CONFIG_FILE = "redis.conf" - CONFIG_DIR = os.path.abspath("./merlin_server/") - PROCESS_FILE = "merlin_server.pf" - PASSWORD_FILE = "redis.pass" - USERS_FILE = "redis.users" - - format = FORMAT - image_type = IMAGE_TYPE - image = IMAGE_NAME - url = REDIS_URL - config = CONFIG_FILE - config_dir = CONFIG_DIR - pfile = PROCESS_FILE - pass_file = PASSWORD_FILE - user_file = USERS_FILE - - def __init__(self, data: dict) -> None: - self.format = data["format"] if "format" in data else self.FORMAT - self.image_type = data["image_type"] if "image_type" in data else self.IMAGE_TYPE - self.image = data["image"] if "image" in data else self.IMAGE_NAME - self.url = data["url"] if "url" in data else self.REDIS_URL - self.config = data["config"] if "config" in data else self.CONFIG_FILE - self.config_dir = os.path.abspath(data["config_dir"]) if "config_dir" in data else self.CONFIG_DIR - self.pfile = data["pfile"] if "pfile" in data else self.PROCESS_FILE - self.pass_file = data["pass_file"] if "pass_file" in data else self.PASSWORD_FILE - self.user_file = data["user_file"] if "user_file" in data else self.USERS_FILE + FORMAT: str = "singularity" + IMAGE_TYPE: str = "redis" + IMAGE_NAME: str = "redis_latest.sif" + REDIS_URL: str = "docker://redis" + CONFIG_FILE: str = "redis.conf" + CONFIG_DIR: str = os.path.abspath("./merlin_server/") + PROCESS_FILE: str = "merlin_server.pf" + PASSWORD_FILE: str = "redis.pass" + USERS_FILE: str = "redis.users" + + format: str = FORMAT + image_type: str = IMAGE_TYPE + image: str = IMAGE_NAME + url: str = REDIS_URL + config: str = CONFIG_FILE + config_dir: str = CONFIG_DIR + pfile: str = PROCESS_FILE + pass_file: str = PASSWORD_FILE + user_file: str = USERS_FILE + + def __init__(self, data: Dict): + """ + Initializes a `ContainerConfig` instance with configuration values. + + Take in a dictionary of configuration values and set up the attributes of the + `ContainerConfig` instance. If any values are missing from the provided dictionary, + default values for a Singularity container are used. + + Args: + data: A dictionary containing configuration values. Keys can include:\n + - `format` (str): Container format (e.g., "singularity"). + - `image_type` (str): Image type (e.g., "redis"). + - `image` (str): Image name (e.g., "redis_latest.sif"). + - `url` (str): URL for the container image (e.g., "docker://redis"). + - `config` (str): Name of the configuration file (e.g., "redis.conf"). + - `config_dir` (str): Path to the configuration directory. + - `pfile` (str): Name of the process file (e.g., "merlin_server.pf"). + - `pass_file` (str): Name of the password file (e.g., "redis.pass"). + - `user_file` (str): Name of the users file (e.g., "redis.users"). + """ + self.format: str = data["format"] if "format" in data else self.FORMAT + self.image_type: str = data["image_type"] if "image_type" in data else self.IMAGE_TYPE + self.image: str = data["image"] if "image" in data else self.IMAGE_NAME + self.url: str = data["url"] if "url" in data else self.REDIS_URL + self.config: str = data["config"] if "config" in data else self.CONFIG_FILE + self.config_dir: str = os.path.abspath(data["config_dir"]) if "config_dir" in data else self.CONFIG_DIR + self.pfile: str = data["pfile"] if "pfile" in data else self.PROCESS_FILE + self.pass_file: str = data["pass_file"] if "pass_file" in data else self.PASSWORD_FILE + self.user_file: str = data["user_file"] if "user_file" in data else self.USERS_FILE + + def __eq__(self, other: "ContainerConfig") -> bool: + """ + Checks equality between two `ContainerConfig` instances. + + This magic method overrides the equality operator (`==`) to compare two + `ContainerConfig` objects. It checks if all relevant attributes + of the two objects are equal. + + Args: + other: Another instance of `ContainerConfig` to compare against. + + Returns: + True if all attributes are equal between the two objects. False otherwise. + + Example: + ```python + >>> config1_data = { + ... "format": "singularity", + ... "image": "redis_latest.sif", + ... "config_dir": "/configs", + ... } + >>> config2_data = { + ... "format": "singularity", + ... "image": "redis_latest.sif", + ... "config_dir": "/configs", + ... } + >>> config3_data = { + ... "format": "singularity", + ... "image": "redis_latest.sif", + ... "config_dir": "/other_configs", + ... } + >>> config1 = ContainerConfig(config1_data) + >>> config2 = ContainerConfig(config2_data) + >>> config3 = ContainerConfig(config3_data) + >>> config1 == config2 + True + >>> config1 == config3 + False + ``` + """ + variables = ("format", "image_type", "image", "url", "config", "config_dir", "pfile", "pass_file", "user_file") + return all(getattr(self, attr) == getattr(other, attr) for attr in variables) def get_format(self) -> str: - """Getter method to get the container format""" + """ + Retrieves the container format. + + This method returns the format of the container, which specifies the type + of container being used (e.g., "singularity", "docker"). + + Returns: + The container format. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_format() + 'singularity' + ``` + """ return self.format def get_image_type(self) -> str: - """Getter method to get the image type""" + """ + Retrieves the image type. + + This method returns the type of the container image, which typically + describes the application or service associated with the image + (e.g., "redis", "mysql"). + + Returns: + The image type. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_image_type() + 'redis' + ``` + """ return self.image_type def get_image_name(self) -> str: - """Getter method to get the image name""" + """ + Retrieves the image name. + + This method returns the name of the container image, which may include + the version or tag (e.g., "redis_latest.sif", "mysql:8.0"). + + Returns: + The image name. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_image_name() + 'redis_latest.sif' + ``` + """ return self.image def get_image_url(self) -> str: - """Getter method to get the image url""" + """ + Retrieves the URL of the image. + + This method returns the URL where the container image is hosted or can + be downloaded from (e.g., a public or private registry URL). + + Returns: + The URL of the container image. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_image_url() + 'docker://redis' + ``` + """ return self.url def get_image_path(self) -> str: - """Getter method to get the path to the image""" + """ + Retrieves the full path to the image file. + + This method constructs and returns the absolute path to the container + image by combining the configuration directory and the image name. + + Returns: + The full path to the container image file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_image_path() + '/configs/redis_latest.sif' + ``` + """ return os.path.join(self.config_dir, self.image) def get_config_name(self) -> str: - """Getter method to get the configuration file name""" + """ + Retrieves the name of the configuration file. + + This method returns the name of the configuration file associated with + the container or application (e.g., "redis.conf", "my.cnf"). + + Returns: + The name of the configuration file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_config_name() + 'redis.conf' + ``` + """ return self.config def get_config_path(self) -> str: - """Getter method to get the configuration file path""" + """ + Retrieves the full path to the configuration file. + + This method constructs and returns the absolute path to the configuration + file by combining the configuration directory and the configuration file name. + + Returns: + The full path to the configuration file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_config_path() + '/configs/redis.conf' + ``` + """ return os.path.join(self.config_dir, self.config) def get_config_dir(self) -> str: - """Getter method to get the configuration directory""" + """ + Retrieves the path to the configuration directory. + + This method returns the directory where configuration files are stored, + which can be used as a base path for accessing specific configuration files. + + Returns: + The path to the configuration directory. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_config_dir() + '/configs' + ``` + """ return self.config_dir def get_pfile_name(self) -> str: - """Getter method to get the process file name""" + """ + Retrieves the name of the process file. + + This method returns the name of the process file, which may represent + a file used to store process-related information (e.g., PID files or + other runtime data files). + + Returns: + The name of the process file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_pfile_name() + 'merlin_server.pf' + ``` + """ return self.pfile def get_pfile_path(self) -> str: - """Getter method to get the process file path""" + """ + Retrieves the full path to the process file. + + This method constructs and returns the absolute path to the process file + by combining the configuration directory and the process file name. + + Returns: + The full path to the process file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_pfile_path() + '/configs/merlin_server.pf' + ``` + """ return os.path.join(self.config_dir, self.pfile) def get_pass_file_name(self) -> str: - """Getter method to get the password file name""" + """ + Retrieves the name of the password file. + + This method returns the name of the password file, which is typically used + to store sensitive information such as user credentials or authentication keys. + + Returns: + The name of the password file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_pass_file_name() + 'redis.pass' + ``` + """ return self.pass_file def get_pass_file_path(self) -> str: - """Getter method to get the password file path""" + """ + Retrieves the full path to the password file. + + This method constructs and returns the absolute path to the password file + by combining the configuration directory and the password file name. + + Returns: + The full path to the password file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_pass_file_path() + '/configs/redis.pass' + ``` + """ return os.path.join(self.config_dir, self.pass_file) def get_user_file_name(self) -> str: - """Getter method to get the user file name""" + """ + Retrieves the name of the user file. + + This method returns the name of the user file, which may be used to store + information related to users, such as user configurations or metadata. + + Returns: + The name of the user file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_user_file_name() + 'redis.users' + ``` + """ return self.user_file def get_user_file_path(self) -> str: - """Getter method to get the user file path""" + """ + Retrieves the full path to the user file. + + This method constructs and returns the absolute path to the user file + by combining the configuration directory and the user file name. + + Returns: + The full path to the user file. + + Example: + ```python + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_user_file_path() + '/configs/redis.users' + ``` + """ return os.path.join(self.config_dir, self.user_file) def get_container_password(self) -> str: - """Getter method to get the container password""" + """ + Retrieves the password for the container. + + This method reads the container password from the password file, which is + located at the path returned by + [`get_pass_file_path`][server.server_util.ContainerConfig.get_pass_file_path]. + The password is read as plain text from the file. + + Returns: + The container password. + + Example: + ```python + >>> with open("redis.pass", "w") as passfile: + ... passfile.write("redis_password") + >>> config_data = { + ... "format": "singularity", + ... "image_type": "redis", + ... "image": "redis_latest.sif", + ... "url": "docker://redis", + ... "config": "redis.conf", + ... "config_dir": "/configs", + ... "pfile": "merlin_server.pf", + ... "pass_file": "redis.pass", + ... "user_file": "redis.users", + ... } + >>> container_config = ContainerConfig(config_data) + >>> container_config.get_container_password() + 'redis_password' + ``` + """ password = None with open(self.get_pass_file_path(), "r") as f: # pylint: disable=C0103 password = f.read() @@ -187,67 +705,295 @@ def get_container_password(self) -> str: class ContainerFormatConfig: """ - ContainerFormatConfig provides an interface for parsing and interacting with container specific - configuration files .yaml. These configuration files contain container specific - commands to run containerizers such as singularity, docker, and podman. + `ContainerFormatConfig` provides an interface for parsing and interacting with container-specific + configuration files, such as .yaml. These configuration files define + container-specific commands for containerizers like Singularity, Docker, and Podman. + + This class allows you to customize and retrieve commands for running, stopping, and pulling + container images. + + Attributes: + COMMAND (str): Default command for running the container (default is "singularity"). + RUN_COMMAND (str): Default template for the run command. + STOP_COMMAND (str): Default command for stopping the container (default is "kill"). + PULL_COMMAND (str): Default template for the pull command. + + command (str): The container command, initialized from the configuration data or defaulting to `COMMAND`. + run_command (str): The run command, initialized from the configuration data or defaulting to `RUN_COMMAND`. + stop_command (str): The stop command, initialized from the configuration data or defaulting to `STOP_COMMAND`. + pull_command (str): The pull command, initialized from the configuration data or defaulting to `PULL_COMMAND`. + + Methods: + __eq__: Compares two `ContainerFormatConfig` objects for equality. + get_command: Retrieves the container command. + get_run_command: Retrieves the run command. + get_stop_command: Retrieves the stop command. + get_pull_command: Retrieves the pull command. """ - COMMAND = "singularity" - RUN_COMMAND = "{command} run {image} {config}" - STOP_COMMAND = "kill" - PULL_COMMAND = "{command} pull {image} {url}" + COMMAND: str = "singularity" + RUN_COMMAND: str = "{command} run {image} {config}" + STOP_COMMAND: str = "kill" + PULL_COMMAND: str = "{command} pull {image} {url}" + + command: str = COMMAND + run_command: str = RUN_COMMAND + stop_command: str = STOP_COMMAND + pull_command: str = PULL_COMMAND + + def __init__(self, data: Dict): + """ + Initializes a `ContainerFormatConfig` object with container-specific configuration data. + + This constructor takes a dictionary of configuration data and initializes the container + command attributes. If any of the keys are missing in the provided data, default values + are used instead. + + Args: + data: A dictionary containing configuration data. Expected keys are:\n + - `command` (str): The container command (e.g., "singularity", "docker"). + - `run_command` (str): The template for the run command. + - `stop_command` (str): The command to stop the container. + - `pull_command` (str): The template for the pull command. + """ + self.command: str = data["command"] if "command" in data else self.COMMAND + self.run_command: str = data["run_command"] if "run_command" in data else self.RUN_COMMAND + self.stop_command: str = data["stop_command"] if "stop_command" in data else self.STOP_COMMAND + self.pull_command: str = data["pull_command"] if "pull_command" in data else self.PULL_COMMAND + + def __eq__(self, other: "ContainerFormatConfig") -> bool: + """ + Checks equality between two `ContainerFormatConfig` objects. + + This method compares the attributes of the current object with those of another + `ContainerFormatConfig` object to determine if they are equal. - command = COMMAND - run_command = RUN_COMMAND - stop_command = STOP_COMMAND - pull_command = PULL_COMMAND + Args: + other: Another instance of the `ContainerFormatConfig` class to compare against. - def __init__(self, data: dict) -> None: - self.command = data["command"] if "command" in data else self.COMMAND - self.run_command = data["run_command"] if "run_command" in data else self.RUN_COMMAND - self.stop_command = data["stop_command"] if "stop_command" in data else self.STOP_COMMAND - self.pull_command = data["pull_command"] if "pull_command" in data else self.PULL_COMMAND + Returns: + True if all corresponding attributes are equal between the two objects, otherwise False. + + Example: + ```python + >>> config1 = ContainerFormatConfig({"command": "singularity"}) + >>> config2 = ContainerFormatConfig({"command": "singularity"}) + >>> config1 == config2 + True + ``` + """ + variables = ("command", "run_command", "stop_command", "pull_command") + return all(getattr(self, attr) == getattr(other, attr) for attr in variables) def get_command(self) -> str: - """Getter method to get the container command""" + """ + Retrieves the container command. + + This method returns the value of the `command` attribute, + which specifies the container command (e.g., "singularity", "docker"). + + Returns: + The container command. + + Example: + ```python + >>> config = ContainerFormatConfig( + ... command="docker", + ... run_command="docker run --name my_container", + ... stop_command="docker stop my_container", + ... pull_command="docker pull my_image" + ... ) + >>> config.get_command() + 'docker' + ``` + """ return self.command def get_run_command(self) -> str: - """Getter method to get the run command""" + """ + Retrieves the run command. + + This method returns the value of the `run_command` attribute, + which specifies the template or command used to run the container. + + Returns: + The run command. + + Example: + ```python + >>> config = ContainerFormatConfig( + ... command="docker", + ... run_command="docker run --name my_container", + ... stop_command="docker stop my_container", + ... pull_command="docker pull my_image" + ... ) + >>> config.get_run_command() + 'docker run --name my_container' + ``` + """ return self.run_command def get_stop_command(self) -> str: - """Getter method to get the stop command""" + """ + Retrieves the stop command. + + This method returns the value of the `stop_command` attribute, + which specifies the command used to stop the container. + + Returns: + The stop command. + + Example: + ```python + >>> config = ContainerFormatConfig( + ... command="docker", + ... run_command="docker run --name my_container", + ... stop_command="docker stop my_container", + ... pull_command="docker pull my_image" + ... ) + >>> config.get_stop_command() + 'docker stop my_container' + ``` + """ return self.stop_command def get_pull_command(self) -> str: - """Getter method to get the pull command""" + """ + Retrieves the pull command. + + This method returns the value of the `pull_command` attribute, + which specifies the template or command used to pull the container image. + + Returns: + The pull command. + + Example: + ```python + >>> config = ContainerFormatConfig( + ... command="docker", + ... run_command="docker run --name my_container", + ... stop_command="docker stop my_container", + ... pull_command="docker pull my_image" + ... ) + >>> config.get_pull_command() + 'docker pull my_image' + ``` + """ return self.pull_command class ProcessConfig: """ - ProcessConfig provides an interface for parsing and interacting with process config specified - in merlin_server.yaml configuration. This configuration provide commands for interfacing with - host machine while the containers are running. + `ProcessConfig` provides an interface for parsing and interacting with process configuration + specified in the `merlin_server.yaml` configuration file. This configuration defines commands + for interacting with the host machine while containers are running, such as checking the status + of processes or terminating them. + + Attributes: + STATUS_COMMAND (str): Default template for the status command, which checks if a process + is running using its parent process ID (PID). Default is "pgrep -P {pid}". + KILL_COMMAND (str): Default template for the kill command, which terminates a process + using its PID. Default is "kill {pid}". + + status (str): The status command template to check the status of a process. This is + initialized from the provided configuration or defaults to `STATUS_COMMAND`. + kill (str): The kill command template to terminate a process. This is initialized from + the provided configuration or defaults to `KILL_COMMAND`. + + Methods: + __eq__: Compares two ProcessConfig objects for equality based on their `status` and + `kill` attributes. + get_status_command: Retrieves the status command template. + get_kill_command: Retrieves the kill command template. """ - STATUS_COMMAND = "pgrep -P {pid}" - KILL_COMMAND = "kill {pid}" + STATUS_COMMAND: str = "pgrep -P {pid}" + KILL_COMMAND: str = "kill {pid}" - status = STATUS_COMMAND - kill = KILL_COMMAND + status: str = STATUS_COMMAND + kill: str = KILL_COMMAND - def __init__(self, data: dict) -> None: - self.status = data["status"] if "status" in data else self.STATUS_COMMAND - self.kill = data["kill"] if "kill" in data else self.KILL_COMMAND + def __init__(self, data: Dict): + """ + Initializes the ProcessConfig object with custom or default process commands. + + This constructor takes a dictionary containing configuration data and initializes + the `status` and `kill` attributes. If the keys `status` or `kill` are not present + in the provided dictionary, their values default to `STATUS_COMMAND` and `KILL_COMMAND`, + respectively. + + Args: + data: A dictionary containing process configuration. Expected keys are:\n + - "status": A string representing the status command template. + - "kill": A string representing the kill command template. + """ + self.status: str = data["status"] if "status" in data else self.STATUS_COMMAND + self.kill: str = data["kill"] if "kill" in data else self.KILL_COMMAND + + def __eq__(self, other: "ProcessConfig") -> bool: + """ + Checks equality between two `ProcessConfig` objects. + + This method compares the attributes of the current object with those of another + `ProcessConfig` object to determine if they are equal. + + Args: + other: Another instance of the `ProcessConfig` class to compare with. + + Returns: + `True` if the `status` and `kill` attributes of both instances are equal, + otherwise `False`. + + Example: + ```python + >>> config1 = ProcessConfig({"status": "check_status", "kill": "terminate_process"}) + >>> config2 = ProcessConfig({"status": "check_status", "kill": "terminate_process"}) + >>> config3 = ProcessConfig({"status": "check_status", "kill": "stop_process"}) + >>> config1 == config2 + True + >>> config1 == config3 + False + ``` + """ + variables = ("status", "kill") + return all(getattr(self, attr) == getattr(other, attr) for attr in variables) def get_status_command(self) -> str: - """Getter method to get the status command""" + """ + Retrieves the status command for the process. + + This method returns the command used to check the status of the process + managed by the `ProcessConfig` instance. + + Returns: + The status command as a string. + + Example: + ```python + >>> process_config = ProcessConfig(status="ps aux | grep process_name", kill="kill -9 process_id") + >>> process_config.get_status_command() + 'ps aux | grep process_name' + ``` + """ return self.status def get_kill_command(self) -> str: - """Getter method to get the kill command""" + """ + Retrieves the kill command for the process. + + This method returns the command used to terminate the process + managed by the `ProcessConfig` instance. + + Returns: + The kill command as a string. + + Example: + ```python + >>> process_config = ProcessConfig(status="ps aux | grep process_name", kill="kill -9 process_id") + >>> process_config.get_kill_command() + 'kill -9 process_id' + ``` + """ return self.kill @@ -255,21 +1001,38 @@ def get_kill_command(self) -> str: # classes so we can ignore it class ServerConfig: # pylint: disable=R0903 """ - ServerConfig is an interface for storing all the necessary configuration for merlin server. - These configuration container things such as ContainerConfig, ProcessConfig, and ContainerFormatConfig. + `ServerConfig` is an interface for storing all the necessary configuration for the Merlin server. + + This class encapsulates configurations related to containers, processes, and container formats, + making it easier to manage and access these settings in a structured way. + + Attributes: + container (ContainerConfig): Configuration related to the container. + process (ProcessConfig): Configuration related to the process. + container_format (ContainerFormatConfig): Configuration for the container format. """ container: ContainerConfig = None process: ProcessConfig = None container_format: ContainerFormatConfig = None - def __init__(self, data: dict) -> None: - if "container" in data: - self.container = ContainerConfig(data["container"]) - if "process" in data: - self.process = ProcessConfig(data["process"]) - if self.container.get_format() in data: - self.container_format = ContainerFormatConfig(data[self.container.get_format()]) + def __init__(self, data: Dict): + """ + Initializes a ServerConfig instance with the provided configuration data. + + Args: + data: A dictionary containing configuration data. Expected keys include:\n + - `container` ([`ContainerConfig`][server.server_util.ContainerConfig]): Configuration data for the container. + - `process` ([`ProcessConfig`][server.server_util.ProcessConfig]): Configuration data for the process. + - `container_format` ([`ContainerFormatConfig`][server.server_util.ContainerFormatConfig]): Configuration + data for the container format. + """ + self.container: ContainerConfig = ContainerConfig(data["container"]) if "container" in data else None + self.process: ProcessConfig = ProcessConfig(data["process"]) if "process" in data else None + container_format_data: str = data.get(self.container.get_format() if self.container else None) + self.container_format: ContainerFormatConfig = ( + ContainerFormatConfig(container_format_data) if container_format_data else None + ) class RedisConfig: @@ -277,22 +1040,110 @@ class RedisConfig: RedisConfig is an interface for parsing and interacing with redis.conf file that is provided by redis. This allows users to parse the given redis configuration and make edits and allow users to write those changes into a redis readable config file. + + `RedisConfig` is an interface for parsing and interacting with a Redis configuration file (`redis.conf`). + + This class allows users to: + - Parse an existing Redis configuration file. + - Modify configuration settings such as IP address, port, password, snapshot settings, directories, etc. + - Write the updated configuration back to a file in a Redis-readable format. + + Attributes: + filename (str): The path to the Redis configuration file. + changed (bool): A flag indicating whether any changes have been made to the configuration. + entry_order (List): A list maintaining the order of configuration entries as they appear in the file. + entries (Dict): A dictionary storing configuration keys and their corresponding values. + comments (Dict): A dictionary storing comments associated with each configuration entry. + trailing_comments (str): Any comments that appear at the end of the configuration file. + + Methods: + parse: Parses the Redis configuration file and populates the `entries`, `comments`, and `entry_order` + attributes. + write: Writes the current configuration (including comments) back to the file. + set_filename: Updates the filename of the configuration file. + set_config_value: Updates the value of a given configuration key. + get_config_value: Retrieves the value of a given configuration key. + changes_made: Returns whether any changes have been made to the configuration. + get_ip_address: Retrieves the IP address (`bind`) setting from the configuration. + set_ip_address: Validates and sets the IP address (`bind`) in the configuration. + get_port: Retrieves the port setting from the configuration. + set_port: Validates and sets the port in the configuration. + set_password: Sets the Redis password (`requirepass`) in the configuration. + get_password: Retrieves the Redis password (`requirepass`) from the configuration. + set_directory: Sets the save directory (`dir`) in the configuration. Creates the directory if it + does not exist. + set_snapshot: Updates the snapshot settings (`save`) for the configuration. + set_snapshot_file: Sets the snapshot file name (`dbfilename`) in the configuration. + set_append_mode: Sets the append mode (`appendfsync`) in the configuration. + set_append_file: Sets the append file name (`appendfilename`) in the configuration. """ - filename = "" - entry_order = [] - entries = {} - comments = {} - trailing_comments = "" - changed = False + def __init__(self, filename: str): + """ + Initializes a `RedisConfig` instance and parses the given Redis configuration file. - def __init__(self, filename) -> None: - self.filename = filename - self.changed = False + Args: + filename: The path to the Redis configuration file (`redis.conf`). + + Notes: + - Parses the configuration file immediately after initialization by calling the `parse()` method. + - Populates the `entries`, `comments`, and `entry_order` attributes based on the file's contents. + """ + self.filename: str = filename + self.changed: bool = False + self.entry_order: List[str] = [] + self.entries: Dict[str, str] = {} + self.comments: Dict[str, str] = {} + self.trailing_comments: str = "" + self.changed: bool = False self.parse() - def parse(self) -> None: - """Parses the redis configuration file""" + def parse(self): + """ + Parses the Redis configuration file and populates the configuration data. + + This method reads the Redis configuration file specified by `self.filename` and extracts: + + - Configuration entries (key-value pairs). + - Associated comments for each entry. + - The order of entries as they appear in the file. + - Any trailing comments at the end of the file. + + Behavior: + - Lines starting with `#` are treated as comments and stored in the `comments` dictionary or `trailing_comments`. + - Non-comment lines are split into key-value pairs and stored in the `entries` dictionary. + - The order of configuration keys is preserved in the `entry_order` list. + - Handles duplicate keys by appending additional parts to the key (e.g., for special cases like `save`). + + Attributes Updated: + - `self.entries`: Stores configuration key-value pairs. + - `self.comments`: Stores comments associated with each configuration key. + - `self.entry_order`: Preserves the order of configuration keys as they appear in the file. + - `self.trailing_comments`: Stores any comments that appear at the end of the file. + + Example: + Assume we have a Redis configuration file named "redis.conf" with the following content: + ``` + # Redis configuration file + maxmemory 256mb + save 900 1 + # End of configuration + ``` + + Then we'd see the following: + ```python + >>> config = RedisConfig("redis.conf") + >>> config.parse() + >>> print(config.entries) + {'maxmemory': '256mb', 'save 900': 1} + >>> print(config.comments) + {'maxmemory': '# Redis configuration file\\n', 'save 900': ''} + >>> print(config.entry_order) + ['maxmemory', 'save 900'] + >>> print(config.trailing_comments) + '# End of configuration' + ``` + """ self.entries = {} self.comments = {} with open(self.filename, "r+") as f: # pylint: disable=C0103 @@ -314,20 +1165,88 @@ def parse(self) -> None: comments += line + "\n" self.trailing_comments = comments[:-1] - def write(self) -> None: - """Writes to the redis configuration file""" + def write(self): + """ + Writes the current configuration and comments back to the Redis configuration file. + + This method writes the configuration entries, their associated comments, and any + trailing comments to the file specified by `self.filename`. The order of entries is + preserved as per the `self.entry_order` list. + + Example: + Assume we start with a Redis configuration file named "redis.conf" with the following content: + ``` + # Redis configuration file + maxmemory 256mb + save 900 1 + # End of configuration + ``` + + We then update it with the write method: + ```python + >>> config = RedisConfig("redis.conf") + >>> config.set_config_value("maxmemory", "512mb") + >>> config.write() + ``` + + Which updates our "redis.conf" file to be: + ``` + # Redis configuration file + maxmemory 512mb + save 900 1 + # End of configuration + ``` + """ with open(self.filename, "w") as f: # pylint: disable=C0103 for entry in self.entry_order: f.write(self.comments[entry]) f.write(f"{entry} {self.entries[entry]}\n") f.write(self.trailing_comments) - def set_filename(self, filename: str) -> None: - """Setter method to set the filename""" + def set_filename(self, filename: str): + """ + Sets a new filename for the Redis configuration file. + + Args: + filename: The new file path to be set as the Redis configuration file. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> print(config.filename) + 'redis.conf' + >>> config.set_filename("/path/to/new/redis.conf") + >>> print(config.filename) + '/path/to/new/redis.conf' + ``` + """ self.filename = filename def set_config_value(self, key: str, value: str) -> bool: - """Changes a configuration value""" + """ + Updates the value of a specific configuration key. + + This method changes the value of an existing configuration key in the `entries` dictionary. + If the key does not exist, the method returns `False`. If the key is updated successfully, + the `changed` attribute is set to `True`. + + Args: + key: The configuration key to update. + value: The new value to set for the specified key. + + Returns: + True if the key exists and the value is successfully updated. False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_config_value("maxmemory", "512mb") + >>> if success: + ... print("Configuration updated successfully!") + ... else: + ... print("Key not found in the configuration.") + ``` + """ if key not in self.entries: return False self.entries[key] = value @@ -335,21 +1254,102 @@ def set_config_value(self, key: str, value: str) -> bool: return True def get_config_value(self, key: str) -> str: - """Given an entry in the config, get the value""" + """ + Retrieves the value of a specific configuration key. + + This method looks up the value of the specified key in the `entries` dictionary + and returns it. If the key does not exist, the method returns `None`. + + Args: + key: The configuration key to retrieve the value for. + + Returns: + The value associated with the specified key as a string, or `None` if the key is not found. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> value = config.get_config_value("maxmemory") + >>> print(value) + '256mb' + >>> value = config.get_config_value("nonexistent_key") + >>> print(value) + None + ``` + """ if key in self.entries: return self.entries[key] return None def changes_made(self) -> bool: - """Getter method to get the changes made""" + """ + Checks if any changes have been made to the configuration. + + This method returns the value of the `self.changed` attribute, which indicates + whether any configuration values have been modified since the last parse or write. + + Returns: + True if changes have been made to the configuration, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> print(config.changes_made()) + False + >>> config.set_config_value("maxmemory", "512mb") + >>> print(config.changes_made()) + True + ``` + """ return self.changed def get_ip_address(self) -> str: - """Getter method to get the ip from the redis config""" + """ + Retrieves the IP address bound in the Redis configuration. + + This method uses the [`get_config_value`][server.server_util.RedisConfig.get_config_value] + method to fetch the value of the `bind` key from the configuration. The `bind` key typically + specifies the IP address that Redis binds to. If the `bind` key is not present in the configuration, + the method returns `None`. + + Returns: + The IP address as a string if the `bind` key exists, or `None` if it does not. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> ip_address = config.get_ip_address() + >>> print(ip_address) + '127.0.0.1' + ``` + """ return self.get_config_value("bind") def set_ip_address(self, ipaddress: str) -> bool: - """Validates and sets a given ip address""" + """ + Validates and sets the given IP address in the Redis configuration. + + This method checks if the provided IP address is a valid IPv4 address. + If valid, it updates the `bind` key in the Redis configuration with the new IP address. + If the IP address is invalid or the update fails, the method logs an error and returns `False`. + + Args: + ipaddress: The IP address to set in the Redis configuration. + + Returns: + True if the IP address is successfully validated and set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_ip_address("192.168.1.1") + >>> print(success) + True + >>> success = config.set_ip_address("invalid_ip") + >>> print(success) + False + ``` + """ if ipaddress is None: return False # Check if ipaddress is valid @@ -365,11 +1365,52 @@ def set_ip_address(self, ipaddress: str) -> bool: return True def get_port(self) -> str: - """Getter method to get the port from the redis config""" + """ + Retrieves the port number from the Redis configuration. + + This method fetches the value of the `port` key from the configuration + using the [`get_config_value`][server.server_util.RedisConfig.get_config_value] + method. If the `port` key is not present, the method returns `None`. + + Returns: + The port number as a string if the `port` key exists, or `None` if it does not. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> port = config.get_port() + >>> print(port) + '6379' + ``` + """ return self.get_config_value("port") - def set_port(self, port: str) -> bool: - """Validates and sets a given port""" + def set_port(self, port: int) -> bool: + """ + Validates and sets the given port number in the Redis configuration. + + This method checks if the provided port number is valid. If valid, it updates + the `port` key in the Redis configuration with the new port number. If the port + is invalid or the update fails, the method logs an error and returns `False`. + + Args: + port: The port number to set in the Redis configuration. + + Returns: + True if the port number is successfully validated and set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_port(6379) + >>> print(success) + True + >>> success = config.set_port(99999) + ERROR: Invalid port given + >>> print(success) + False + ``` + """ if port is None: return False # Check if port is valid @@ -385,7 +1426,26 @@ def set_port(self, port: str) -> bool: return True def set_password(self, password: str) -> bool: - """Changes the password""" + """ + Sets a new password in the Redis configuration. + + This method updates the `requirepass` key in the Redis configuration with the provided password. + If the password is `None`, the method returns `False` without making any changes. + + Args: + password: The new password to set in the Redis configuration. + + Returns: + True if the password is successfully set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_password("my_secure_password") + >>> print(success) + True + ``` + """ if password is None: return False self.set_config_value("requirepass", password) @@ -393,73 +1453,146 @@ def set_password(self, password: str) -> bool: return True def get_password(self) -> str: - """Getter method to get the config password""" + """ + Retrieves the password from the Redis configuration. + + This method fetches the value of the `requirepass` key from the configuration + using the [`get_config_value`][server.server_util.RedisConfig.get_config_value] + method. If the `requirepass` key is not present, the method returns `None`. + + Returns: + The password as a string if the `requirepass` key exists, or `None` if it does not. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> password = config.get_password() + >>> print(password) + 'my_secure_password' + ``` + """ return self.get_config_value("requirepass") def set_directory(self, directory: str) -> bool: """ - Sets the save directory in the redis config file. - Creates the directory if necessary. + Sets the save directory in the Redis configuration file. + + This method updates the `dir` key in the Redis configuration with the provided directory path. + If the directory does not exist, it is created. If the directory is `None` or the update fails, + the method logs an error and returns `False`. + + Args: + directory: The directory path to set as the save directory in the Redis configuration. + + Returns: + True if the directory is successfully validated, created (if necessary), and set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_directory("/var/lib/redis") + >>> print(success) + True + ``` """ if directory is None: return False + # Create the directory if it doesn't exist if not os.path.exists(directory): os.mkdir(directory) LOG.info(f"Created directory {directory}") - # Validate the directory input - if os.path.exists(directory): - # Set the save directory to the redis config - if not self.set_config_value("dir", directory): - LOG.error("Unable to set directory for redis config") - return False - else: - LOG.error(f"Directory {directory} given does not exist and could not be created.") + # Set the save directory to the redis config + if not self.set_config_value("dir", directory): + LOG.error("Unable to set directory for redis config") return False LOG.info(f"Directory is set to {directory}") return True - def set_snapshot_seconds(self, seconds: int) -> bool: - """Sets the snapshot wait time""" - if seconds is None: - return False - # Set the snapshot second in the redis config - value = self.get_config_value("save") - if value is None: - LOG.error("Unable to get exisiting parameter values for snapshot") - return False + def set_snapshot(self, seconds: int = None, changes: int = None) -> bool: + """ + Updates the snapshot configuration in the Redis settings. + + This method allows you to set the snapshot parameters, which determine + when Redis creates a snapshot of the dataset. The snapshot is triggered + based on a combination of time (`seconds`) and the number of changes + (`changes`) made to the dataset. If either parameter is `None`, it will + remain unchanged. + + Args: + seconds: The time interval (in seconds) after which a snapshot should be created. + If `None`, the existing value for `seconds` remains unchanged. + changes: The number of changes to the dataset that trigger a snapshot. + If `None`, the existing value for `changes` remains unchanged. + + Returns: + True if the snapshot configuration is successfully updated, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_snapshot(seconds=300, changes=10) + >>> print(success) + True + >>> success = config.set_snapshot(seconds=600) + >>> print(success) + True + >>> success = config.set_snapshot() # No changes + >>> print(success) + False + ``` + """ - value = value.split() - value[0] = str(seconds) - value = " ".join(value) - if not self.set_config_value("save", value): - LOG.error("Unable to set snapshot value seconds") + # If both values are None, this method is doing nothing + if seconds is None and changes is None: return False - LOG.info(f"Snapshot wait time is set to {seconds} seconds") - return True - - def set_snapshot_changes(self, changes: int) -> bool: - """Sets the snapshot threshold""" - if changes is None: - return False - # Set the snapshot changes into the redis config + # Grab the snapshot value from the redis config value = self.get_config_value("save") if value is None: LOG.error("Unable to get exisiting parameter values for snapshot") return False + # Update the snapshot value value = value.split() - value[1] = str(changes) + log_msg = "" + if seconds is not None: + value[0] = str(seconds) + log_msg += f"Snapshot wait time is set to {seconds} seconds. " + if changes is not None: + value[1] = str(changes) + log_msg += f"Snapshot threshold is set to {changes} changes." value = " ".join(value) + + # Set the new snapshot value if not self.set_config_value("save", value): - LOG.error("Unable to set snapshot value seconds") + LOG.error("Unable to set snapshot value") return False - LOG.info(f"Snapshot threshold is set to {changes} changes") + LOG.info(log_msg) return True def set_snapshot_file(self, file: str) -> bool: - """Sets the snapshot file""" + """ + Sets the name of the snapshot file in the Redis configuration. + + This method updates the `dbfilename` parameter in the Redis configuration + with the provided file name. The snapshot file is where Redis saves the + dataset during a snapshot operation. + + Args: + file: The name of the snapshot file to set in the Redis configuration. + + Returns: + True if the snapshot file name is successfully set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_snapshot_file("dump.rdb") + >>> print(success) + True + ``` + """ if file is None: return False # Set the snapshot file in the redis config @@ -471,7 +1604,30 @@ def set_snapshot_file(self, file: str) -> bool: return True def set_append_mode(self, mode: str) -> bool: - """Sets the append mode""" + """ + Sets the append mode in the Redis configuration. + + The append mode determines how Redis handles data persistence to the append-only file (AOF). + Valid modes are: + + - "always": Redis appends data to the AOF after every write operation. + - "everysec": Redis appends data to the AOF every second (default and recommended). + - "no": Disables AOF persistence. + + Args: + mode: The append mode to set. Must be one of "always", "everysec", or "no". + + Returns: + True if the append mode is successfully set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_append_mode("everysec") + >>> print(success) + True + ``` + """ if mode is None: return False valid_modes = ["always", "everysec", "no"] @@ -483,14 +1639,34 @@ def set_append_mode(self, mode: str) -> bool: LOG.error("Unable to set append_mode in redis config") return False else: - LOG.error("Not a valid append_mode(Only valid modes are always, everysec, no)") + LOG.error("Not a valid append_mode (Only valid modes are always, everysec, no)") return False LOG.info(f"Append mode is set to {mode}") return True def set_append_file(self, file: str) -> bool: - """Sets the append file""" + """ + Sets the name of the append-only file (AOF) in the Redis configuration. + + The append-only file is used by Redis for data persistence in append-only mode. + This method updates the `appendfilename` parameter in the Redis configuration + with the provided file name. + + Args: + file: The name of the append-only file to set in the Redis configuration. + + Returns: + True if the append file name is successfully set, False otherwise. + + Example: + ```python + >>> config = RedisConfig("redis.conf") + >>> success = config.set_append_file("appendonly.aof") + >>> print(success) + True + ``` + """ if file is None: return False # Set the append file in the redis config @@ -503,35 +1679,107 @@ def set_append_file(self, file: str) -> bool: class RedisUsers: """ - RedisUsers provides an interface for parsing and interacting with redis.users configuration - file. Allow users and merlin server to create, remove, and edit users within the redis files. - Changes can be sync and push to an exisiting redis server if one is available. + `RedisUsers` provides an interface for parsing and interacting with a Redis `users` configuration + file. This class allows users and the Merlin server to create, remove, and edit user entries + within the Redis configuration files. Changes can be synchronized and pushed to an existing + Redis server if one is available. + + Attributes: + filename (str): The path to the Redis user configuration file. + users (Dict[str, User]): A dictionary + containing user data, where keys are usernames and values are instances of the + `User` class. + + Methods: + parse: Parses the Redis user configuration file and populates the `users` dictionary with + [`User`][server.server_util.RedisUsers.User] objects. + write: Writes the current users' data back to the Redis user configuration file. + add_user: Adds a new user to the `users` dictionary. + set_password: Sets a new password for a specific user. + remove_user: Removes a user from the `users` dictionary. + apply_to_redis: Applies the current user configuration to a running Redis server by synchronizing + the local user data with the Redis server's ACL configuration. """ class User: - """Embedded class to store user specific information""" + """ + An embedded class that represents an individual Redis user with attributes and methods for managing + user-specific data. + + Attributes: + status (str): The status of the user, either "on" (enabled) or "off" (disabled). + hash_password (str): The hashed password of the user. + keys (str): The keys the user has access to (e.g., "*" for all keys). + channels (str): The channels the user can access (e.g., "*" for all channels). + commands (str): The commands the user is allowed to execute (e.g., "@all" for all commands). + + Methods: + parse_dict: Parses a dictionary of user data and updates the `User` object's attributes. + get_user_dict: Returns a dictionary representation of the `User` object. + __repr__(): Returns a string representation of the `User` object. + __str__(): Returns a string representation of the `User` object (same as `__repr__`). + set_password: Sets the user's hashed password based on the provided plaintext password. + """ - status = "on" - hash_password = hashlib.sha256(b"password").hexdigest() - keys = "*" - channels = "*" - commands = "@all" + status: str = "on" + hash_password: str = hashlib.sha256(b"password").hexdigest() + keys: str = "*" + channels: str = "*" + commands: str = "@all" def __init__( # pylint: disable=R0913 - self, status="on", keys="*", channels="*", commands="@all", password=None - ) -> None: - self.status = status - self.keys = keys - self.channels = channels - self.commands = commands + self, + status: str = "on", + keys: str = "*", + channels: str = "*", + commands: str = "@all", + password: str = None, + ): + """ + Initializes a `User` object with the provided attributes. + + Args: + status: The status of the user, either "on" (enabled) or "off" (disabled). + keys: The keys the user has access to (e.g., "*" for all keys). + channels: The channels the user can access (e.g., "*" for all channels). + commands: The commands the user is allowed to execute (e.g., "@all" for all commands). + password: The plaintext password for the user. If provided, it will be hashed and stored. + """ + self.status: str = status + self.keys: str = keys + self.channels: str = channels + self.commands: str = commands if password is not None: self.set_password(password) - def parse_dict(self, dictionary: dict) -> None: + def parse_dict(self, dictionary: Dict[str, str]): """ - Given a dict of user info, parse the dict and store - the values as class attributes. - :param `dictionary`: The dict to parse + Parses a dictionary containing user information and updates the attributes of the `User` object. + + Args: + dictionary: A dictionary containing user data. Expected keys are:\n + - `status` (str): The user's status ("on" or "off"). + - `keys` (str): The keys the user can access. + - `channels` (str): The channels the user can access. + - `commands` (str): The commands the user can execute. + - `hash_password` (str): The hashed password of the user. + + Example: + ```python + >>> user_data = { + ... "status": "on", + ... "keys": "*", + ... "channels": "*", + ... "commands": "@all", + ... "hash_password": "hashed_password_value" + ... } + >>> user = User() + >>> user.parse_dict(user_data) + >>> print(user.status) + 'on' + >>> print(user.hash_password) + 'hashed_password_value' + ``` """ self.status = dictionary["status"] self.keys = dictionary["keys"] @@ -539,8 +1787,27 @@ def parse_dict(self, dictionary: dict) -> None: self.commands = dictionary["commands"] self.hash_password = dictionary["hash_password"] - def get_user_dict(self) -> dict: - """Getter method to get the user info""" + def get_user_dict(self) -> Dict: + """ + Returns a dictionary representation of the `User` object. + + Returns: + A dictionary containing the user's attributes. + + Example: + ```python + >>> user = User(status="on", keys="*", channels="*", commands="@all", password="secure_password") + >>> user_dict = user.get_user_dict() + >>> print(user_dict) + { + "status": "on", + "hash_password": "hashed_password_value", + "keys": "*", + "channels": "*", + "commands": "@all" + } + ``` + """ self.status = "on" return { "status": self.status, @@ -551,27 +1818,89 @@ def get_user_dict(self) -> dict: } def __repr__(self) -> str: - """Repr magic method for User class""" + """ + Returns a string representation of the `User` object for debugging purposes. + + Returns: + A string representation of the user's attributes in dictionary format. + """ return str(self.get_user_dict()) def __str__(self) -> str: - """Str magic method for User class""" + """ + Returns a string representation of the `User` object. + + Returns: + A string representation of the user's attributes in dictionary format. + """ return self.__repr__() - def set_password(self, password: str) -> None: - """Setter method to set the user's hash password""" + def set_password(self, password: str): + """ + Sets the user's hashed password based on the provided plaintext password. + + Args: + password: The plaintext password to hash and store. + + Example: + ```python + >>> user = User() + >>> user.set_password("secure_password") + >>> print(user.hash_password) + # Output: The hashed value of "secure_password" + ``` + """ self.hash_password = hashlib.sha256(bytes(password, "utf-8")).hexdigest() - filename = "" - users = {} + filename: str = "" + users: Dict[str, User] = {} - def __init__(self, filename) -> None: - self.filename = filename + def __init__(self, filename: str): + """ + Initializes a `RedisUsers` object and parses the Redis user configuration file if it exists. + + Args: + filename: The path to the Redis user configuration file. + """ + self.filename: str = filename if os.path.exists(self.filename): self.parse() - def parse(self) -> None: - """Parses the redis user configuration file""" + def parse(self): + """ + Parses the Redis user configuration file and populates the `users` dictionary. + + This method reads the YAML configuration file specified by `self.filename` and converts + the user data into a dictionary. Each user entry is converted into an instance of the + [`User`][server.server_util.RedisUsers.User] class, which stores the user's attributes. + + Example: + Assume `users.yaml` contains: + ``` + user1: + status: "on" + hash_password: "hashed_password_1" + keys: "*" + channels: "*" + commands: "@all" + user2: + status: "off" + hash_password: "hashed_password_2" + keys: "key1" + channels: "channel1" + commands: "command1" + ``` + + This would then be parsed like so: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> redis_users.parse() + >>> print(redis_users.users["user1"].status) + 'on' + >>> print(redis_users.users["user2"].commands) + 'command1' + ``` + """ with open(self.filename, "r") as f: # pylint: disable=C0103 self.users = yaml.load(f, yaml.Loader) for user in self.users: @@ -579,8 +1908,30 @@ def parse(self) -> None: new_user.parse_dict(self.users[user]) self.users[user] = new_user - def write(self) -> None: - """Writes to the redis user configuration file""" + def write(self): + """ + Writes the current `users` dictionary back to the Redis user configuration file. + + This method converts the `users` dictionary, which contains [`User`][server.server_util.RedisUsers.User] + objects, into a format suitable for storage in the YAML configuration file. The file + specified by `self.filename` is then updated with the current user data. + + Example: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> redis_users.add_user( + ... user="new_user", + ... status="on", + ... keys="*", + ... channels="*", + ... commands="@all", + ... password="secure_password" + ... ) + >>> redis_users.write() + ``` + + After calling `write`, the `users.yaml` file will be updated with the new user's data. + """ data = self.users.copy() for key in data: data[key] = self.users[key].get_user_dict() @@ -588,30 +1939,121 @@ def write(self) -> None: yaml.dump(data, f, yaml.Dumper) def add_user( # pylint: disable=R0913 - self, user, status="on", keys="*", channels="*", commands="@all", password=None + self, + user: str, + status: str = "on", + keys: str = "*", + channels: str = "*", + commands: str = "@all", + password: str = None, ) -> bool: - """Add a user to the dict of Redis users""" + """ + Adds a new user to the dictionary of Redis users. + + Args: + user: The username of the new user. + status: The status of the user, either "on" (enabled) or "off" (disabled). + keys: The keys the user has access to (e.g., "*" for all keys). + channels: The channels the user can access (e.g., "*" for all channels). + commands: The commands the user is allowed to execute (e.g., "@all" for all commands). + password: The plaintext password for the user. If provided, it will be hashed and stored. + + Returns: + True if the user was successfully added, False if the user already exists. + + Example: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> success = redis_users.add_user( + ... user="new_user", + ... status="on", + ... keys="*", + ... channels="*", + ... commands="@all", + ... password="secure_password" + ... ) + >>> print(success) + True + ``` + """ if user in self.users: return False self.users[user] = self.User(status, keys, channels, commands, password) return True - def set_password(self, user: str, password: str): - """Set the password for a specific user""" + def set_password(self, user: str, password: str) -> bool: + """ + Sets the password for an existing user. + + Args: + user: The username of the user whose password is to be updated. + password: The plaintext password to hash and store for the user. + + Returns: + True if the password was successfully updated, False if the user does not exist. + + Example: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> redis_users.add_user("existing_user", password="old_password") + >>> success = redis_users.set_password("existing_user", "new_password") + >>> print(success) + True + ``` + """ if user not in self.users: return False self.users[user].set_password(password) return True - def remove_user(self, user) -> bool: - """Remove a user from the dict of users""" + def remove_user(self, user: str) -> bool: + """ + Removes a user from the dictionary of Redis users. + + Args: + user: The username of the user to be removed. + + Returns: + True if the user was successfully removed, False if the user does not exist. + + Example: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> redis_users.add_user("user_to_remove", password="password") + >>> success = redis_users.remove_user("user_to_remove") + >>> print(success) + True + ``` + """ if user in self.users: del self.users[user] return True return False - def apply_to_redis(self, host: str, port: int, password: str) -> None: - """Apply the changes to users to redis""" + def apply_to_redis(self, host: str, port: int, password: str): + """ + Applies the user configuration changes to a Redis instance. + + This method synchronizes the current user configuration stored in `self.users` + with the Redis instance specified by the provided connection details. It performs + the following actions: + + - Adds or updates users in Redis based on the `self.users` dictionary. + - Removes users from Redis that are not present in `self.users`. + + Args: + host: The hostname or IP address of the Redis server. + port: The port number of the Redis server. + password: The password for authenticating with the Redis server. + + Example: + ```python + >>> redis_users = RedisUsers("users.yaml") + >>> redis_users.add_user("user1", password="password1") + >>> redis_users.add_user("user2", password="password2") + >>> redis_users.apply_to_redis(host="127.0.0.1", port=6379, password="redis_password") + ``` + """ database = redis.Redis(host=host, port=port, password=password) current_users = database.acl_users() for user in self.users: @@ -633,23 +2075,57 @@ def apply_to_redis(self, host: str, port: int, password: str) -> None: class AppYaml: """ - AppYaml allows for an structured way to interact with any app.yaml main merlin configuration file. - It helps to parse each component of the app.yaml and allow users to edit, configure and write the - file. + `AppYaml` provides a structured way to interact with the main `app.yaml` configuration file for Merlin. + This class allows users to parse, edit, configure, and write the `app.yaml` file, which contains + the application's main configuration settings. + + Attributes: + default_filename (str): The default file path for the `app.yaml` configuration file. + data (Dict): A dictionary that holds the parsed configuration data from the `app.yaml` file. + broker_name (str): The key name in the configuration file representing the broker settings. + results_name (str): The key name in the configuration file representing the results backend settings. + + Methods: + apply_server_config: Updates the `data` dictionary with Redis server configuration details based on + the provided `ServerConfig` object. + update_data: Updates the `data` dictionary with new entries. + get_data: Returns the current configuration data stored in the `data` attribute. + read: Reads a YAML file and populates the `data` attribute with its contents. + write: Writes the current `data` dictionary to the specified YAML file. """ - default_filename = os.path.join(MERLIN_CONFIG_DIR, "app.yaml") - data = {} - broker_name = "broker" - results_name = "results_backend" + default_filename: str = os.path.join(MERLIN_CONFIG_DIR, "app.yaml") + data: Dict = {} + broker_name: str = "broker" + results_name: str = "results_backend" + + def __init__(self, filename: str = default_filename): + """ + Initializes the `AppYaml` object and loads the configuration file. - def __init__(self, filename: str = default_filename) -> None: + Args: + filename: The path to the `app.yaml` file. If the file does not exist, + the default file path (`default_filename`) is used. + """ if not os.path.exists(filename): filename = self.default_filename self.read(filename) def apply_server_config(self, server_config: ServerConfig): - """Store the redis configuration""" + """ + Updates the `data` dictionary with Redis server configuration details + based on the provided `ServerConfig` object. + + Args: + server_config: An object containing the server configuration. + + Example: + ```python + >>> server_config = ServerConfig(...) + >>> app_yaml = AppYaml() + >>> app_yaml.apply_server_config(server_config) + ``` + """ redis_config = RedisConfig(server_config.container.get_config_path()) self.data[self.broker_name]["name"] = server_config.container.get_image_type() @@ -664,19 +2140,67 @@ def apply_server_config(self, server_config: ServerConfig): self.data[self.results_name]["server"] = redis_config.get_ip_address() self.data[self.results_name]["port"] = redis_config.get_port() - def update_data(self, new_data: dict): - """Update the data dict with new entries""" + def update_data(self, new_data: Dict): + """ + Updates the `data` dictionary with new entries. + + Args: + new_data: A dictionary containing the new data to be merged + into the existing `data` dictionary. + + Example: + ```python + >>> new_data = {"custom_key": {"sub_key": "value"}} + >>> app_yaml = AppYaml() + >>> app_yaml.update_data(new_data) + ``` + """ self.data.update(new_data) - def get_data(self): - """Getter method to obtain the data""" + def get_data(self) -> Dict: + """ + Retrieves the current configuration data stored in the `data` attribute. + + Returns: + The current configuration data. + + Example: + ```python + >>> app_yaml = AppYaml() + >>> current_data = app_yaml.get_data() + ``` + """ return self.data def read(self, filename: str = default_filename): - """Load in a yaml file and save it to the data attribute""" + """ + Reads a YAML file and populates the `data` attribute with its contents. + + Args: + filename: The path to the YAML file to be read. If not provided, + the default file path (`default_filename`) is used. + + Example: + ```python + >>> app_yaml = AppYaml() + >>> app_yaml.read("custom_app.yaml") + ``` + """ self.data = merlin.utils.load_yaml(filename) def write(self, filename: str = default_filename): - """Given a filename, dump the data to the file""" + """ + Writes the current `data` dictionary to the specified YAML file. + + Args: + filename: The path to the YAML file where the data should be written. + If not provided, the default file path (`default_filename`) is used. + + Example: + ```python + >>> app_yaml = AppYaml() + >>> app_yaml.write("output_app.yaml") + ``` + """ with open(filename, "w+") as f: # pylint: disable=C0103 yaml.dump(self.data, f, yaml.Dumper) diff --git a/merlin/spec/__init__.py b/merlin/spec/__init__.py index 57477ea1f..198202cd7 100644 --- a/merlin/spec/__init__.py +++ b/merlin/spec/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,3 +27,20 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `spec` package defines the structure, defaults, and functionality for working with +Merlin specification files. + +Modules: + all_keys.py: Defines all the valid keys for each block in a Merlin specification file, + ensuring consistency and validation. + defaults.py: Provides the default values for each block in a spec file, enabling workflows + to execute even when fields are omitted. + expansion.py: Handles the expansion of variables within a spec file, including user-defined, + environment, and reserved variables, as well as parameter substitutions. + override.py: Supports overriding variables in a spec file via the command-line interface, + with functions for validation and replacement. + specification.py: Contains the `MerlinSpec` class, which represents the raw data from a + Merlin specification file and provides methods for interacting with it. +""" diff --git a/merlin/spec/all_keys.py b/merlin/spec/all_keys.py index fbb70f8d7..0e288bf8d 100644 --- a/merlin/spec/all_keys.py +++ b/merlin/spec/all_keys.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,7 +27,16 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module defines all the keys possible in each block of a merlin spec file""" +""" +This module defines all the possible keys for each block in a Merlin specification file. + +Merlin specification files are used to configure and manage studies, workflows, and tasks +in Merlin. Each block in the spec file corresponds to a specific aspect of the workflow, +such as batch settings, environment variables, study steps, parameters, and resources. +This module provides sets of valid keys for each block to ensure consistency and validation. + +This module serves as a reference for the structure and valid keys of Merlin specification files. +""" DESCRIPTION = {"description", "name"} diff --git a/merlin/spec/defaults.py b/merlin/spec/defaults.py index 6c9bd9c09..558eeb9ee 100644 --- a/merlin/spec/defaults.py +++ b/merlin/spec/defaults.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,7 +27,16 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module defines the default values of every block in the merlin spec""" +""" +This module defines the default values for every block in a Merlin specification file. + +Merlin specification files are used to configure workflows, tasks, and resources in +Merlin. This module provides the default values for each block in the spec file, ensuring +that workflows can execute even if certain fields are omitted by the user. These defaults +serve as a fallback mechanism to maintain consistency and simplify the configuration process. + +This module serves as a reference for the default configuration of Merlin specification files. +""" DESCRIPTION = {"description": {}} diff --git a/merlin/spec/expansion.py b/merlin/spec/expansion.py index ac514a369..1f62f74a2 100644 --- a/merlin/spec/expansion.py +++ b/merlin/spec/expansion.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,14 +27,22 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module handles expanding variables in the merlin spec""" +""" +This module handles the expansion of variables in a Merlin spec file. + +It provides functionality to expand user-defined variables, environment variables, and reserved variables +within a spec file. The module also supports variable substitution for specific use cases, such as parameter +substitutions for samples and commands, and allows for the processing of override variables provided via +the command-line interface. +""" import logging from collections import ChainMap from copy import deepcopy from os.path import expanduser, expandvars +from typing import Dict, List, Tuple -from merlin.common.abstracts.enums import ReturnCode +from merlin.common.enums import ReturnCode from merlin.spec.override import error_override_vars, replace_override_vars from merlin.spec.specification import MerlinSpec from merlin.utils import contains_shell_ref, contains_token, verify_filepath @@ -66,10 +74,27 @@ LOG = logging.getLogger(__name__) -def var_ref(string): +def var_ref(string: str) -> str: """ - Given a string , return that string surrounded - by $(). + Format a string as a variable reference. + + This function takes a string, converts it to uppercase, and returns it + wrapped in the format `$()`. If the string already contains + a token (e.g., it is already formatted as a variable reference), a warning + is logged and the original string is returned unchanged. + + Args: + string: The input string to format as a variable reference. + + Returns: + The formatted variable reference, or the original string if it + already contains a token. + + Example: + ```python + >>> var_ref("example") + '$(EXAMPLE)' + ``` """ string = string.upper() if contains_token(string): @@ -78,10 +103,31 @@ def var_ref(string): return f"$({string})" -def expand_line(line, var_dict, env_vars=False): +def expand_line(line: str, var_dict: Dict[str, str], env_vars: bool = False) -> str: """ - Expand one line of text by substituting user variables, - optionally environment variables, as well as variables in 'var_dict'. + Expand a single line of text by substituting variables. + + This function replaces variable references in a given line of text with + their corresponding values from a provided dictionary. Optionally, it can + also expand environment variables and user home directory shortcuts (e.g., `~`). + + Args: + line: The input line of text to expand. + var_dict: A dictionary of variable names and their corresponding values + to substitute in the line. + env_vars: If True, environment variables and home directory shortcuts + will also be expanded. + + Returns: + The expanded line of text with all applicable substitutions applied. + + Example: + ```python + >>> line = "Path: $(VAR1) and $(VAR2)" + >>> var_dict = {"VAR1": "/path/to/dir", "VAR2": "/another/path"} + >>> expand_line(line, var_dict) + 'Path: /path/to/dir and /another/path' + ``` """ # fmt: off if ( @@ -99,10 +145,30 @@ def expand_line(line, var_dict, env_vars=False): return line -def expand_by_line(text, var_dict): +def expand_by_line(text: str, var_dict: Dict[str, str]) -> str: """ - Given a text (yaml spec), and a dictionary of variable names - and values, expand variables in the text line by line. + Expand variables in a text line by line. + + This function processes a multi-line text (e.g., a YAML specification) and + replaces variable references in each line using a provided dictionary of + variable names and their corresponding values. + + Args: + text: The input multi-line text to process. + var_dict: A dictionary of variable names and their corresponding values + to substitute in the text. + + Returns: + The text with all applicable variable substitutions applied, processed + line by line. + + Example: + ```python + >>> text = "Path is $(VAR1) and stores $(VAR2)" + >>> var_dict = {"VAR1": "/path/to/dir", "VAR2": "value"} + >>> expand_by_line(text, var_dict) + 'Path is /path/to/dir and stores value' + ``` """ text = text.splitlines() result = "" @@ -112,11 +178,23 @@ def expand_by_line(text, var_dict): return result -def expand_env_vars(spec): +def expand_env_vars(spec: MerlinSpec) -> MerlinSpec: """ - Expand environment variables for all sections of a spec, except - for values with the key 'cmd' or 'restart' (these are executable - shell scripts, so environment variable expansion would be redundant). + Expand environment variables in all sections of a spec. + + This function processes all sections of a given spec object and expands + environment variables (e.g., `$HOME` or `~`) in string values. It skips + expansion for values associated with the keys 'cmd' or 'restart', as these + are typically shell scripts where environment variable expansion would + already occur during execution. + + Args: + spec (spec.specification.MerlinSpec): The spec object + containing sections to process. + + Returns: + (spec.specification.MerlinSpec): The updated spec object with environment variables expanded in all + applicable sections. """ def recurse(section): @@ -139,25 +217,35 @@ def recurse(section): return spec -def determine_user_variables(*user_var_dicts): +def determine_user_variables(*user_var_dicts: List[Dict]) -> Dict: """ - Given an arbitrary number of dictionaries, determine them - in order. - - param `user_var_dicts`: A list of dictionaries of user variables. - For example: - [variables, labels] - - A single user var dict may look like: - {'OUTPUT_PATH':'./studies', 'N_SAMPLES':10} - - This user var dict: - {'TARGET': 'target_dir', - 'PATH': '$(SPECROOT)/$(TARGET)'} - - ...would be determined as: - {'TARGET': 'target_dir', - 'PATH': '$(SPECROOT)/target_dir'} + Determine user-defined variables from multiple dictionaries. + + This function takes an arbitrary number of dictionaries containing user-defined + variables and resolves them in order, handling variable references and expansions + (e.g., environment variables, user home directory shortcuts). Variable names are + converted to uppercase, and reserved words cannot be reassigned. + + Args: + user_var_dicts: One or more dictionaries of user variables. Each dictionary + contains key-value pairs where the key is the variable name, and the value + is the variable's definition. + + Returns: + A dictionary of resolved user variables, with variable names in uppercase + and all references expanded. + + Raises: + ValueError: If a reserved word is attempted to be reassigned. + + Example: + ```python + >>> user_vars_1 = {'OUTPUT_PATH': './studies', 'N_SAMPLES': 10} + >>> user_vars_2 = {'TARGET': 'target_dir', 'PATH': '$(SPECROOT)/$(TARGET)'} + >>> determine_user_variables(user_vars_1, user_vars_2) + {'OUTPUT_PATH': './studies', 'N_SAMPLES': '10', + 'TARGET': 'target_dir', 'PATH': '$(SPECROOT)/target_dir'} + ``` """ all_var_dicts = dict(ChainMap(*user_var_dicts)) determined_results = {} @@ -175,15 +263,41 @@ def determine_user_variables(*user_var_dicts): return determined_results -def parameter_substitutions_for_sample(sample, labels, sample_id, relative_path_to_sample): +def parameter_substitutions_for_sample( + sample: List[str], labels: List[str], sample_id: int, relative_path_to_sample: str +) -> List[Tuple[str, str]]: """ - :param sample : The sample to do substitution for. - :param labels : The column labels of the sample. - :param sample_id : The merlin sample id for this sample. - :param relative_path_to_sample : The relative path to this sample. - - :return : list of pairs indicating what needs to be substituted for a - merlin sample + Generate parameter substitutions for a specific sample. + + This function creates a list of substitution pairs for a given sample, + mapping variable references (e.g., `$(LABEL)`) to their corresponding + values in the sample. It also includes metadata substitutions such as + the sample ID and the relative path to the sample. + + Args: + sample: A list of values representing the sample. + labels: A list of column labels corresponding to the sample values. + sample_id: The unique integer ID of the sample. + relative_path_to_sample: The relative path to the sample. + + Returns: + A list of tuples, where each tuple contains a variable reference + (e.g., `$(LABEL)`) and its corresponding value. + + Example: + ```python + >>> sample = [10, 20] + >>> labels = ["X", "Y"] + >>> sample_id = 0 + >>> relative_path_to_sample = "/0/3/4/8/9/" + >>> parameter_substitutions_for_sample(sample, labels, sample_id, relative_path_to_sample) + [ + ("$(X)", "10"), + ("$(Y)", "20"), + ("$(MERLIN_SAMPLE_ID)", "0"), + ("$(MERLIN_SAMPLE_PATH)", "/0/3/4/8/9/") + ] + ``` """ substitutions = [] for label, axis in zip(labels, sample): @@ -198,13 +312,50 @@ def parameter_substitutions_for_sample(sample, labels, sample_id, relative_path_ return substitutions -def parameter_substitutions_for_cmd(glob_path, sample_paths): +def parameter_substitutions_for_cmd(glob_path: str, sample_paths: str) -> List[Tuple[str, str]]: """ - :param glob_path: a glob that should yield the paths to all merlin samples - :param sample_paths: a delimited list of all of the samples - - :return : list of pairs indicating what needs to be substituted for a - merlin cmd + Generate parameter substitutions for a Merlin command. + + This function creates a list of substitution pairs for a Merlin command, + mapping variable references to their corresponding values. It also includes + predefined return codes for various Merlin states. + + Substitutions: + - `$(MERLIN_GLOB_PATH)`: The provided `glob_path`. + - `$(MERLIN_PATHS_ALL)`: The provided `sample_paths`. + - `$(MERLIN_SUCCESS)`: The return code for a successful operation. + - `$(MERLIN_RESTART)`: The return code for a restart operation. + - `$(MERLIN_SOFT_FAIL)`: The return code for a soft failure. + - `$(MERLIN_HARD_FAIL)`: The return code for a hard failure. + - `$(MERLIN_RETRY)`: The return code for a retry operation. + - `$(MERLIN_STOP_WORKERS)`: The return code for stopping workers. + - `$(MERLIN_RAISE_ERROR)`: The return code for raising an error. + + Args: + glob_path: A glob pattern that yields the paths to all Merlin samples. + sample_paths: A delimited string containing the paths to all samples. + + Returns: + A list of tuples, where each tuple contains a variable reference and its + corresponding value. + + Example: + ```python + >>> glob_path = "/path/to/samples/*" + >>> sample_paths = "/path/to/sample1:/path/to/sample2" + >>> parameter_substitutions_for_cmd(glob_path, sample_paths) + [ + ("$(MERLIN_GLOB_PATH)", "/path/to/samples/*"), + ("$(MERLIN_PATHS_ALL)", "/path/to/sample1:/path/to/sample2"), + ("$(MERLIN_SUCCESS)", "0"), + ("$(MERLIN_RESTART)", "100"), + ("$(MERLIN_SOFT_FAIL)", "101"), + ("$(MERLIN_HARD_FAIL)", "102"), + ("$(MERLIN_RETRY)", "104"), + ("$(MERLIN_STOP_WORKERS)", "105"), + ("$(MERLIN_RAISE_ERROR)", "106") + ] + ``` """ substitutions = [] substitutions.append(("$(MERLIN_GLOB_PATH)", glob_path)) @@ -223,12 +374,26 @@ def parameter_substitutions_for_cmd(glob_path, sample_paths): # There's similar code inside study.py but the whole point of this function is to not use # the MerlinStudy object so we disable this pylint error # pylint: disable=duplicate-code -def expand_spec_no_study(filepath, override_vars=None): +def expand_spec_no_study(filepath: str, override_vars: Dict[str, str] = None) -> str: """ Get the expanded text of a spec without creating a MerlinStudy. Expansion is limited to user variables (the ones defined inside the yaml spec or at the command line). + + Expand a spec without creating a [`MerlinStudy`][study.study.MerlinStudy]. + + This function processes a spec file to expand user-defined variables (those defined + in the YAML spec or provided via `override_vars`) without creating a `MerlinStudy` + object. It returns the expanded text of the specification. + + Args: + filepath: The path to the YAML specification file. + override_vars: A dictionary of variable overrides to apply during the expansion. + These overrides replace or supplement the variables defined in the spec. + + Returns: + The expanded YAML specification as a string, with user-defined variables resolved. """ error_override_vars(override_vars, filepath) spec = MerlinSpec.load_specification(filepath) @@ -248,10 +413,22 @@ def expand_spec_no_study(filepath, override_vars=None): # pylint: enable=duplicate-code -def get_spec_with_expansion(filepath, override_vars=None): +def get_spec_with_expansion(filepath: str, override_vars: Dict[str, str] = None) -> MerlinSpec: """ - Return a MerlinSpec with overrides and expansion, without - creating a MerlinStudy. + Load and expand a Merlin YAML specification with overrides, without creating a + [`MerlinStudy`][study.study.MerlinStudy] object. + + This function returns a [`MerlinSpec`][spec.specification.MerlinSpec] object with + variables expanded and overrides applied. It processes the YAML specification file + and resolves user-defined variables without creating a `MerlinStudy` object. + + Args: + filepath: The path to the YAML specification file. + override_vars: A dictionary of variable overrides to apply during the expansion. + These overrides replace or supplement the variables defined in the YAML spec. + + Returns: + (spec.specification.MerlinSpec): A `MerlinSpec` object with expanded variables and applied overrides. """ filepath = verify_filepath(filepath) expanded_spec_text = expand_spec_no_study(filepath, override_vars) diff --git a/merlin/spec/override.py b/merlin/spec/override.py index 316e76f86..cf5242d88 100644 --- a/merlin/spec/override.py +++ b/merlin/spec/override.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,18 +27,34 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module handles overriding variables in a spec file via the CLI""" +""" +This module provides functionality to handle overriding variables in a spec file +via the command-line interface. It includes functions to validate and replace variables +in the spec file or environment block based on user-provided overrides. +""" import logging from copy import deepcopy +from typing import Dict LOG = logging.getLogger(__name__) -def error_override_vars(override_vars, spec_filepath): +def error_override_vars(override_vars: Dict[str, str], spec_filepath: str): """ - Warn user if any given variable name isn't found in the original spec file. + Warns the user if any given variable name in the override list is not found + in the original spec file. + + Args: + override_vars: A dictionary of variables to override, where keys are + variable names and values are their corresponding new values. Can + be `None` if no overrides are provided. + spec_filepath: The file path to the original spec file. + + Raises: + ValueError: If any variable name in `override_vars` is not found in the + spec file. """ if override_vars is None: return @@ -49,8 +65,21 @@ def error_override_vars(override_vars, spec_filepath): raise ValueError(f"Command line override variable '{variable}' not found in spec file '{spec_filepath}'.") -def replace_override_vars(env, override_vars): - """Replace override variables in the environment block""" +def replace_override_vars(env: Dict[str, str], override_vars: Dict[str, str]) -> Dict[str, str]: + """ + Replaces variables in the given environment block with the provided override + values. + + Args: + env: The environment block, represented as a dictionary where keys are + environment variable names and values are their corresponding values. + override_vars: A dictionary of variables to override, where keys are + variable names and values are their corresponding new values. Can be + `None` if no overrides are provided. + + Returns: + A new environment block with the override variables replaced. + """ if override_vars is None: return env result = deepcopy(env) diff --git a/merlin/spec/specification.py b/merlin/spec/specification.py index 5ffffd959..92da49498 100644 --- a/merlin/spec/specification.py +++ b/merlin/spec/specification.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,8 +29,9 @@ ############################################################################### """ -This module contains a class, MerlinSpec, which holds the unchanged +This module contains a class, `MerlinSpec`, which holds the unchanged data from the Merlin specification file. + To see examples of yaml specifications, run `merlin example`. """ import json @@ -40,9 +41,10 @@ from copy import deepcopy from datetime import timedelta from io import StringIO -from typing import Dict, List +from typing import Any, Dict, List, Set, TextIO, Union import yaml +from maestrowf.datastructures.core.parameters import ParameterGenerator from maestrowf.specification import YAMLSpecification from merlin.spec import all_keys, defaults @@ -55,32 +57,62 @@ # Pylint complains we have too many instance attributes but it's fine class MerlinSpec(YAMLSpecification): # pylint: disable=R0902 """ - This class represents the logic for parsing the Merlin yaml - specification. - - Example spec_file contents: - - --spec_file.yaml-- - ... - merlin: - resources: - task_server: celery - samples: - generate: - cmd: python make_samples.py -outfile=$(OUTPUT_PATH)/merlin_info/samples.npy - file: $(OUTPUT_PATH)/merlin_info/samples.npy - column_labels: [X0, X1] + A class to represent and manage the specifications for a Merlin workflow. + + This class provides methods to verify, load, and process various sections of a + workflow specification file, including the merlin block, batch block, and user block. + It also handles default values and parameter mapping. + + Attributes: + batch (Dict): A dictionary representing the batch section of the spec file. + description (Dict): A dictionary representing the description section of the spec file. + environment (Dict): A dictionary representing the environment section of the spec file. + globals (Dict): A dictionary representing global parameters in the spec file. + merlin (Dict): A dictionary representing the merlin section of the spec file. + sections (Dict): A dictionary of all sections in the spec file. + study (Dict): A dictionary representing the study section of the spec file. + user (Dict): A dictionary representing the user section of the spec file. + yaml_sections (Dict): A dictionary for YAML representation of the sections. + + Methods: + check_section: Checks sections of the spec file for unrecognized keys. + dump: Dumps the current spec to a pretty YAML string. + fill_missing_defaults: Merges default values into an object. + get_queue_list: Returns a sorted set of queues for specified steps. + get_queue_step_relationship: Maps task queues to their associated steps. + get_step_param_map: Creates a mapping of parameters used for each step. + get_step_worker_map: Maps step names to associated workers. + get_study_step_names: Returns a list of the names of the steps in the spec file. + get_task_queues: Maps steps to their corresponding task queues. + get_tasks_per_step: Returns the number of tasks needed for each step. + get_worker_names: Returns a list of worker names. + get_worker_step_map: Maps worker names to associated steps. + load_merlin_block: Loads the merlin block from a YAML stream. + load_spec_from_string: Creates a `MerlinSpec` object from a string (or stream) representing + a spec file. + load_specification: Creates a `MerlinSpec` object based on the contents of a spec file. + load_user_block: Loads the user block from a YAML stream. + make_queue_string: Returns a unique queue string for specified steps. + process_spec_defaults: Fills in default values for missing sections. + verify: Verify the spec against a valid schema. + verify_batch_block: Validates the batch block against a predefined schema. + verify_merlin_block: Validates the merlin block against a predefined schema. + warn_unrecognized_keys: Checks for unrecognized keys in the spec file. """ # Pylint says this call to super is useless but we'll leave it in case we want to add to __init__ in the future def __init__(self): # pylint: disable=W0246 + """Initializes a MerlinSpec object.""" super().__init__() @property - def yaml_sections(self): + def yaml_sections(self) -> Dict: """ - Returns a nested dictionary of all sections of the specification - as used in a yaml spec. + Returns a nested dictionary of all sections of the specification as used in a YAML + specification. The structure is tailored for YAML representation. + + Returns: + A dictionary containing the sections of the specification formatted for YAML. """ return { "description": self.description, @@ -93,10 +125,14 @@ def yaml_sections(self): } @property - def sections(self): + def sections(self) -> Dict: """ - Returns a nested dictionary of all sections of the specification - as referenced by Maestro's YAMLSpecification class. + Returns a nested dictionary of all sections of the specification as referenced by + [Maestro's `YAMLSpecification` class](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/specification/yamlspecification.html). + The structure is aligned with the expectations of Maestro's `YAMLSpecification` class. + + Returns: + A dictionary containing the sections of the specification formatted for Maestro. """ return { "description": self.description, @@ -109,7 +145,6 @@ def sections(self): } def __str__(self): - """Magic method to print an instance of our MerlinSpec class.""" env = "" globs = "" merlin = "" @@ -129,14 +164,23 @@ def __str__(self): return result @classmethod - def load_specification(cls, path, suppress_warning=True): + def load_specification(cls, path: str, suppress_warning: bool = True) -> "MerlinSpec": """ - Load in a spec file and create a MerlinSpec object based on its' contents. + Load a specification file and create a `MerlinSpec` object based on its contents. + + This method reads a YAML specification file from the provided path, + processes its contents, and returns a `MerlinSpec` object. It can also + suppress warnings about unrecognized keys in the specification. - :param `cls`: The class reference (like self) - :param `path`: A path to the spec file we're loading in - :param `suppress_warning`: A bool representing whether to warn the user about unrecognized keys - :returns: A MerlinSpec object + Args: + path: The path to the specification file to be loaded. + suppress_warning: Whether to suppress warnings about unrecognized keys. + + Returns: + A `MerlinSpec` object created from the contents of the specification file. + + Raises: + Exception: If there is an error loading the specification file. """ LOG.info("Loading specification from path: %s", path) try: @@ -157,16 +201,23 @@ def load_specification(cls, path, suppress_warning=True): return spec @classmethod - def load_spec_from_string(cls, string, needs_IO=True, needs_verification=False): # pylint: disable=C0103 + def load_spec_from_string( + cls, string: Union[str, TextIO], needs_IO: bool = True, needs_verification: bool = False + ) -> "MerlinSpec": # pylint: disable=C0103 """ - Read in a spec file from a string (or stream) and create a MerlinSpec object from it. + Read a specification from a string (or stream) and create a `MerlinSpec` object from it. + + This method processes a string or stream containing the specification + and returns a `MerlinSpec` object. It can also verify the specification + if required. + + Args: + string: A string or stream of the specification content. + needs_IO: Whether to treat the string as a file object. + needs_verification: Whether to verify the specification after loading. - :param `cls`: The class reference (like self) - :param `string`: A string or stream of the file we're reading in - :param `needs_IO`: A bool representing whether we need to turn the string into a file - object or not - :param `needs_verification`: A bool representing whether we need to verify the spec - :returns: A MerlinSpec object + Returns: + A `MerlinSpec` object created from the provided specification content. """ LOG.debug("Creating Merlin spec object...") # Create and populate the MerlinSpec object @@ -197,19 +248,28 @@ def load_spec_from_string(cls, string, needs_IO=True, needs_verification=False): return spec @classmethod - def _populate_spec(cls, data): + def _populate_spec(cls, data: TextIO) -> "MerlinSpec": """ - Helper method to load a study spec and populate it's fields. - - NOTE: This is basically a direct copy of YAMLSpecification's - load_specification method from Maestro just without the call to verify. - The verify method was breaking our code since we have no way of modifying - Maestro's schema that they use to verify yaml files. The work around - is to load the yaml file ourselves and create our own schema to verify - against. - - :param data: Raw text stream to study YAML spec data - :returns: A MerlinSpec object containing information from the path + Helper method to load a study specification and populate its fields. + + This method reads a YAML specification from a raw text stream and + populates the fields of a `MerlinSpec` object. It is a modified version + of the `load_specification` method from Maestro's YAMLSpecification class, + excluding the verification step due to compatibility issues with Maestro's schema. + + Note: + This is basically a direct copy of YAMLSpecification's + load_specification method from Maestro just without the call to verify. + The verify method was breaking our code since we have no way of modifying + Maestro's schema that they use to verify yaml files. The work around + is to load the yaml file ourselves and create our own schema to verify + against. + + Args: + data: A raw text stream containing the study YAML specification data. + + Returns: + A `MerlinSpec` object populated with information extracted from the YAML specification. """ # Read in the spec file try: @@ -244,12 +304,22 @@ def _populate_spec(cls, data): def verify(self): """ - Verify the spec against a valid schema. Similar to YAMLSpecification's verify - method from Maestro but specific for Merlin yaml specs. - - NOTE: Maestro v2.0 may add the ability to customize the schema files it - compares against. If that's the case then we can convert this file back to - using Maestro's verification. + Verify the specification against a valid schema. + + This method checks the current `MerlinSpec` object against a predefined + schema to ensure that it adheres to the expected structure and + constraints. It is similar to the verify method from Maestro's + YAMLSpecification class but is tailored specifically for Merlin YAML + specifications. + + Note: + Maestro v2.0 may introduce the ability to customize the schema files + used for verification. If that feature becomes available, then we can + convert this file back to using Maestro's verification. + + Raises: + Exception: If the specification does not conform to the schema, + appropriate exceptions will be raised during the verification process. """ # Load the MerlinSpec schema file dir_path = os.path.dirname(os.path.abspath(__file__)) @@ -267,11 +337,16 @@ def verify(self): self.verify_merlin_block(schema["MERLIN"]) self.verify_batch_block(schema["BATCH"]) - def get_study_step_names(self): + def get_study_step_names(self) -> List[str]: """ - Get a list of the names of steps in our study. + Retrieve the names of steps in the study. + + This method iterates through the study steps and collects their names + into a list. The returned list is unsorted. - :returns: an unsorted list of study step names + Returns: + An unsorted list of strings representing the names of the + study steps. """ names = [] for step in self.study: @@ -280,8 +355,16 @@ def get_study_step_names(self): def _verify_workers(self): """ - Helper method to verify the workers section located within the Merlin block - of our spec file. + Verify the workers section in the Merlin block of the specification. + + This helper method checks that the steps referenced in the workers + section of the Merlin block exist in the study steps. It raises a + ValueError if any step specified for a worker does not match the + defined study steps. + + Raises: + ValueError: If a step specified in the workers section does not + exist in the list of study step names. """ # Retrieve the names of the steps in our study actual_steps = self.get_study_step_names() @@ -301,24 +384,35 @@ def _verify_workers(self): except Exception: # pylint: disable=W0706 raise - def verify_merlin_block(self, schema): + def verify_merlin_block(self, schema: Dict): """ - Method to verify the merlin section of our spec file. - :param schema: The section of the predefined schema (merlinspec.json) to check - our spec file against. + Verify the Merlin section of the specification file against a schema. + + This method validates the Merlin block of the specification file + against a predefined JSON schema and verifies the workers section + to ensure that all specified steps are defined in the study. + + Args: + schema: The section of the predefined schema (merlinspec.json) to + check the Merlin block against. """ # Validate merlin block against the json schema YAMLSpecification.validate_schema("merlin", self.merlin, schema) # Verify the workers section within merlin block self._verify_workers() - def verify_batch_block(self, schema): + def verify_batch_block(self, schema: Dict): """ - Method to verify the batch section of our spec file. + Verify the batch section of the specification file against a schema. - :param schema: The section of the predefined schema (merlinspec.json) to check - our spec file against. + This method validates the batch block of the specification file + against a predefined JSON schema and performs additional checks + related to the walltime parameter for the LSF batch type. + + Args: + schema: The section of the predefined schema (merlinspec.json) to + check the batch block against. """ # Validate batch block against the json schema YAMLSpecification.validate_schema("batch", self.batch, schema) @@ -328,8 +422,23 @@ def verify_batch_block(self, schema): LOG.warning("The walltime argument is not available in lsf.") @staticmethod - def load_merlin_block(stream): - """Loads in the merlin block of the spec file""" + def load_merlin_block(stream: TextIO) -> Dict: + """ + Load the Merlin block from a specification file stream. + + This static method reads a YAML stream and attempts to extract + the 'merlin' section. If the 'merlin' section is missing, it + logs a warning and returns an empty dictionary, indicating that + the default configuration will be used without sampling. + + Args: + stream: A file-like object or string stream containing the + YAML specification. + + Returns: + The Merlin block extracted from the YAML stream. If the 'merlin' + section is not found, an empty dictionary is returned. + """ try: merlin_block = yaml.safe_load(stream)["merlin"] except KeyError: @@ -343,8 +452,22 @@ def load_merlin_block(stream): return merlin_block @staticmethod - def load_user_block(stream): - """Loads in the user block of the spec file""" + def load_user_block(stream: TextIO) -> Dict: + """ + Load the user block from a specification file stream. + + This static method reads a YAML stream and attempts to extract + the 'user' section. If the 'user' section is not present, it + returns an empty dictionary. + + Args: + stream: A file-like object or string stream containing the + YAML specification. + + Returns: + The user block extracted from the YAML stream. If the 'user' + section is not found, an empty dictionary is returned. + """ try: user_block = yaml.safe_load(stream)["user"] except KeyError: @@ -352,7 +475,23 @@ def load_user_block(stream): return user_block def process_spec_defaults(self): - """Fills in the default values if they aren't there already""" + """ + Fill in default values for specification sections if they are missing. + + This method iterates through the sections of the specification and + populates any that are `None` with empty dictionaries. It then fills + in default values for various sections, including batch, environment, + global parameters, and step sections within the study. + + The method also handles specific cases for the VLAUNCHER variables + in the command of each step, ensuring that default values are set + if they are not defined by the user. Additionally, it ensures that + workers are assigned to steps appropriately, filling in defaults + where necessary. + + The method modifies the instance's attributes directly, ensuring that + the specification is complete and ready for further processing. + """ for name, section in self.sections.items(): if section is None: setattr(self, name, {}) @@ -415,15 +554,52 @@ def process_spec_defaults(self): # no defaults for user block @staticmethod - def fill_missing_defaults(object_to_update, default_dict): + def fill_missing_defaults(object_to_update: Dict, default_dict: Dict): """ - Merge keys and values from a dictionary of defaults - into a parallel object that may be missing attributes. - Only adds missing attributes to object; does not overwrite - existing ones. + Merge default values into an object, filling in missing attributes. + + This static method takes an object and a dictionary of default values, + and merges the defaults into the object. It only adds missing attributes + to the object and does not overwrite any existing attributes. If an + attribute is present in the object but its value is `None`, it will be + updated with the corresponding value from the defaults. + + The method works recursively, allowing for nested dictionaries. + + The method modifies the `object_to_update` in place. + + Args: + object_to_update: The object (as a dictionary) that needs to be + updated with default values. + default_dict: A dictionary containing default values to merge into + the object. + + Example: + ```python + >>> obj = {'a': 1, 'b': None} + >>> defaults = {'a': 2, 'b': 3, 'c': 4} + >>> fill_missing_defaults(obj, defaults) + >>> print(obj) + {'a': 1, 'b': 3, 'c': 4} + ``` """ - def recurse(result, recurse_defaults): + def recurse(result: Dict, recurse_defaults: Dict): + """ + Recursively merge default values into the result object. + + This helper function checks if the current level of the `recurse_defaults` + dictionary is a dictionary itself. If it is, it iterates through each key-value + pair. If a key is not present in the `result` or its value is `None`, it + assigns the value from `recurse_defaults`. If the key exists and has a value, + it recursively calls itself to handle nested dictionaries. + + The function modifies the `result` in place. + + Args: + result: The current state of the object being updated. + recurse_defaults: The current level of defaults to merge. + """ if not isinstance(recurse_defaults, dict): return for key, val in recurse_defaults.items(): @@ -440,7 +616,16 @@ def recurse(result, recurse_defaults): # ***Unsure if this method is still needed after adding json schema verification*** def warn_unrecognized_keys(self): - """Checks if there are any unrecognized keys in the spec file""" + """ + Check for unrecognized keys in the specification file. + + This method verifies that all keys present in the specification file + conform to the expected structure defined by the `MerlinSpec` class. + It checks various sections of the specification, including "description", + "batch", "env", "global parameters", "steps", and "merlin". For each + section, it calls the `check_section` method to ensure that the keys + are recognized and valid according to predefined criteria. + """ # check description MerlinSpec.check_section("description", self.description, all_keys.DESCRIPTION) @@ -470,8 +655,21 @@ def warn_unrecognized_keys(self): # user block is not checked @staticmethod - def check_section(section_name, section, known_keys): - """Checks a section of the spec file to see if there are any unrecognized keys""" + def check_section(section_name: str, section: Dict, known_keys: Set[str]): + """ + Check a section of the specification file for unrecognized keys. + + This static method compares the keys present in a specified section + of the specification file against a set of known keys. If any keys + are found that are not recognized, a warning is logged indicating + the unrecognized key and the section in which it was found. + + Args: + section_name: The name of the section being checked. + section: The section of the specification file to validate. + known_keys: A set of keys that are recognized as valid for + the specified section. + """ diff = set(section.keys()).difference(known_keys) # TODO: Maybe add a check here for required keys @@ -479,9 +677,23 @@ def check_section(section_name, section, known_keys): for extra in diff: LOG.warning(f"Unrecognized key '{extra}' found in spec section '{section_name}'.") - def dump(self): + def dump(self) -> str: """ - Dump this MerlinSpec to a pretty yaml string. + Dump the `MerlinSpec` instance to a formatted YAML string. + + This method converts the current state of the `MerlinSpec` instance + into a YAML formatted string. It utilizes the `_dict_to_yaml` + method to handle the conversion and prettification of the data. + Additionally, it ensures that the resulting YAML string is valid + by attempting to parse it with `yaml.safe_load`. If parsing fails, + a ValueError is raised with details about the error. + + Returns: + A pretty formatted YAML string representation of the + `MerlinSpec` instance. + + Raises: + ValueError: If there is an error while parsing the YAML string. """ tab = 3 * " " result = self._dict_to_yaml(self.yaml_sections, "", [], tab) @@ -493,9 +705,26 @@ def dump(self): raise ValueError(f"Error parsing provenance spec:\n{e}") from e return result - def _dict_to_yaml(self, obj, string, key_stack, tab): + def _dict_to_yaml(self, obj: Any, string: str, key_stack: List[str], tab: int) -> str: """ - The if-else ladder for sorting the yaml string prettification of dump(). + Convert a Python object to a formatted YAML string. + + This private method handles the conversion of various Python data + types (strings, booleans, lists, and dictionaries) into a + formatted YAML string. It uses an if-else structure to determine + the type of the input object and calls the appropriate processing + methods for each type. The method also manages indentation based + on the current level of nesting. + + Args: + obj: The object to convert to YAML format. + string: The current string representation being built. + key_stack: A stack of keys representing the current level of + nesting in the YAML structure. + tab: The number of spaces to use for indentation. + + Returns: + A formatted YAML string representation of the input object. """ if obj is None: return "" @@ -512,18 +741,57 @@ def _dict_to_yaml(self, obj, string, key_stack, tab): return self._process_dict(obj, string, key_stack, lvl, tab) return obj - def _process_string(self, obj, lvl, tab): + def _process_string(self, obj: str, lvl: int, tab: int) -> str: """ - Processes strings for _dict_to_yaml() in the dump() method. + Process a string for YAML formatting in the dump method. + + This private method takes a string and formats it for inclusion + in a YAML output. If the string contains multiple lines, it + transforms the string into a block scalar format using the pipe + (`|`) character, which is suitable for YAML representation. + The indentation is adjusted based on the current level of + nesting. + + Args: + obj: The string to be processed. + lvl: The current level of indentation for the YAML output. + tab: The number of spaces to use for indentation. + + Returns: + The formatted string ready for YAML output. """ split = obj.splitlines() if len(split) > 1: obj = "|\n" + tab * (lvl + 1) + ("\n" + tab * (lvl + 1)).join(split) return obj - def _process_list(self, obj, string, key_stack, lvl, tab): # pylint: disable=R0913 + def _process_list( + self, + obj: List[Any], + string: str, + key_stack: List[str], + lvl: int, + tab: int, + ) -> str: """ - Processes lists for _dict_to_yaml() in the dump() method. + Process a list for YAML formatting in the dump method. + + This private method handles the conversion of a list into a + YAML formatted string. It determines whether to use hyphens + for list items based on the context provided by the key stack. + The method recursively processes each element in the list and + manages indentation based on the current level of nesting. + + Args: + obj: The list to be processed. + string: The current string representation being built. + key_stack: A stack of keys representing the current + level of nesting in the YAML structure. + lvl: The current level of indentation for the YAML output. + tab: The number of spaces to use for indentation. + + Returns: + A formatted YAML string representation of the input list. """ num_entries = len(obj) use_hyphens = key_stack[-1] in ["paths", "sources", "git", "study"] or key_stack[0] in ["user"] @@ -545,9 +813,33 @@ def _process_list(self, obj, string, key_stack, lvl, tab): # pylint: disable=R0 string += "]" return string - def _process_dict(self, obj, string, key_stack, lvl, tab): # pylint: disable=R0913 + def _process_dict( + self, + obj: Dict, + string: str, + key_stack: List[str], + lvl: int, + tab: int, + ) -> str: # pylint: disable=R0913 """ - Processes dicts for _dict_to_yaml() in the dump() method + Process a dictionary for YAML formatting in the dump method. + + This private method converts a dictionary into a YAML formatted + string. It iterates over the dictionary's key-value pairs, + formatting each pair according to YAML syntax. The method + handles indentation and manages the key stack to maintain the + correct nesting level in the output. + + Args: + obj: The dictionary to be processed. + string: The current string representation being built. + key_stack: A stack of keys representing the current + level of nesting in the YAML structure. + lvl: The current level of indentation for the YAML output. + tab: The number of spaces to use for indentation. + + Returns: + A formatted YAML string representation of the input dictionary. """ list_offset = 2 * " " if len(key_stack) > 0 and key_stack[-1] != "elem": @@ -570,10 +862,18 @@ def _process_dict(self, obj, string, key_stack, lvl, tab): # pylint: disable=R0 def get_step_worker_map(self) -> Dict[str, List[str]]: """ - Creates a dictionary with step names as keys and a list of workers - associated with each step as values. The inverse of get_worker_step_map(). - - :returns: A dict mapping step names to workers + Create a mapping of step names to associated workers. + + This method constructs a dictionary where each key is a step name + and the corresponding value is a list of workers assigned to that + step. Workers can either be associated with all steps or with + specific steps. This method serves as the inverse of the + [`get_worker_step_map`][spec.specification.MerlinSpec.get_worker_step_map] + method. + + Returns: + A dictionary mapping step names to lists of worker names + associated with each step. """ steps = self.get_study_step_names() step_worker_map = {step_name: [] for step_name in steps} @@ -590,10 +890,18 @@ def get_step_worker_map(self) -> Dict[str, List[str]]: def get_worker_step_map(self) -> Dict[str, List[str]]: """ - Creates a dictionary with worker names as keys and a list of steps - associated with each worker as values. The inverse of get_step_worker_map(). - - :returns: A dict mapping workers to the steps they watch + Create a mapping of worker names to associated steps. + + This method constructs a dictionary where each key is a worker name + and the corresponding value is a list of steps that the worker is + assigned to monitor. Workers can either be assigned to all steps or + to specific steps. It serves as the inverse of the + [`get_step_worker_map`][spec.specification.MerlinSpec.get_step_worker_map] + method. + + Returns: + A dictionary mapping worker names to lists of step names that each + worker monitors. """ worker_step_map = {} steps = self.get_study_step_names() @@ -608,13 +916,23 @@ def get_worker_step_map(self) -> Dict[str, List[str]]: worker_step_map[worker_name].append(step) return worker_step_map - def get_task_queues(self, omit_tag=False): + def get_task_queues(self, omit_tag: bool = False) -> Dict[str, str]: """ - Creates a dictionary of steps and their corresponding task queues. - This is the inverse of get_queue_step_relationship() + Create a mapping of steps to their corresponding task queues. - :param `omit_tag`: If True, omit the celery queue tag. - :returns: A dict of steps and their corresponding task queues + This method constructs a dictionary where each key is a step name + and the corresponding value is the associated task queue. The + `omit_tag` parameter allows for the optional exclusion of the Celery + queue tag from the queue names. It serves as the inverse of the + [`get_queue_step_relationship`][spec.specification.MerlinSpec.get_queue_step_relationship] + method. + + Args: + omit_tag: If True, the Celery queue tag will be omitted + from the task queue names. Default is False. + + Returns: + A dictionary mapping step names to their corresponding task queues. """ from merlin.config.configfile import CONFIG # pylint: disable=C0415 @@ -629,10 +947,17 @@ def get_task_queues(self, omit_tag=False): def get_queue_step_relationship(self) -> Dict[str, List[str]]: """ - Builds a dictionary of task queues and their associated steps. - This returns the inverse of get_task_queues(). + Build a mapping of task queues to their associated steps. - :returns: A dict of task queues and their associated steps + This method constructs a dictionary where each key is a task queue + name and the corresponding value is a list of steps that are + associated with that queue. It serves as the inverse of the + [`get_task_queues`][spec.specification.MerlinSpec.get_task_queues] + method. + + Returns: + A dictionary mapping task queue names to lists of step names + associated with each queue. """ from merlin.config.configfile import CONFIG # pylint: disable=C0415 @@ -654,13 +979,28 @@ def get_queue_step_relationship(self) -> Dict[str, List[str]]: return relationship_tracker - def get_queue_list(self, steps, omit_tag=False) -> set: + def get_queue_list(self, steps: Union[List[str], str], omit_tag: bool = False) -> Set[str]: """ - Return a sorted set of queues corresponding to spec steps - - :param `steps`: a list of step names or ['all'] - :param `omit_tag`: If True, omit the celery queue tag. - :returns: A sorted set of queues corresponding to spec steps + Return a sorted set of queues corresponding to specified steps. + + This method retrieves a list of task queues associated with the + given steps. If the `steps` parameter is set to ['all'], it will + return all available queues. The `omit_tag` parameter allows for + the optional exclusion of the Celery queue tag from the queue names. + + Args: + steps: A list of step names or a list containing the string 'all' + to represent all steps, or the name of a single step. + omit_tag: If True, the Celery queue tag will be omitted from the + task queue names. + + Returns: + A sorted set of unique task queues corresponding to the specified + steps. + + Raises: + KeyError: If any of the specified steps do not exist in the + task queues. """ queues = self.get_task_queues(omit_tag=omit_tag) if steps[0] == "all": @@ -677,17 +1017,34 @@ def get_queue_list(self, steps, omit_tag=False) -> set: raise return sorted(set(task_queues)) - def make_queue_string(self, steps): + def make_queue_string(self, steps: List[str]) -> str: """ - Return a unique queue string for the steps + Return a unique queue string for the specified steps. - param steps: a list of step names + This method constructs a comma-separated string of unique task + queues associated with the provided steps. The resulting string + is suitable for use in command-line contexts. + + Args: + steps: A list of step names for which to generate the + queue string. + + Returns: + A quoted string of unique task queues, separated by commas. """ queues = ",".join(set(self.get_queue_list(steps))) return shlex.quote(queues) - def get_worker_names(self): - """Builds a list of workers""" + def get_worker_names(self) -> List[str]: + """ + Build a list of worker names. + + This method retrieves the names of all workers defined in the + Merlin resources and returns them as a list. + + Returns: + A list of worker names. + """ result = [] for worker in self.merlin["resources"]["workers"]: result.append(worker) @@ -695,8 +1052,17 @@ def get_worker_names(self): def get_tasks_per_step(self) -> Dict[str, int]: """ - Get the number of tasks needed to complete each step, formatted as a dictionary. - :returns: A dict where the keys are the step names and the values are the number of tasks required for that step + Get the number of tasks needed to complete each step. + + This method calculates the number of tasks required for each + step in the study based on the number of samples and parameters. + It returns a dictionary where the keys are the step names and + the values are the corresponding number of tasks required for + that step. + + Returns: + A dictionary mapping step names to the number of tasks + required for each step. """ # Get the number of samples used samples = [] @@ -729,20 +1095,33 @@ def get_tasks_per_step(self) -> Dict[str, int]: return tasks_per_step - def _create_param_maps(self, param_gen: "ParameterGenerator", expanded_labels: Dict, label_param_map: Dict): # noqa: F821 + def _create_param_maps(self, param_gen: ParameterGenerator, expanded_labels: Dict, label_param_map: Dict): """ - Given a parameters block like so: + Create mappings of tokens to expanded labels and labels to parameter values. + + This private method processes a parameter generator to create two mappings: + + 1. `expanded_labels`: Maps tokens to their expanded labels based on the + provided parameter values. + 2. `label_param_map`: Maps expanded labels to their corresponding parameter + values. + + The expected structure for the parameter block is: + + ``` global.parameters: TOKEN: values: [param_val_1, param_val_2] label: label.%% - Expanded labels will map tokens to their expanded labels (e.g. {'TOKEN': ['label.param_val_1', 'label.param_val_2']}) - Label param map will map labels to parameter values - (e.g. {'label.param_val_1': {'TOKEN': 'param_val_1'}, 'label.param_val_2': {'TOKEN': 'param_val_2'}}) - - :param `param_gen`: A ParameterGenerator object from Maestro - :param `expanded_labels`: A dict to store the map from tokens to expanded labels - :param `label_param_map`: A dict to store the map from labels to parameter values + ``` + + Args: + param_gen: A `ParameterGenerator` object from Maestro containing the + parameter definitions. + expanded_labels: A dictionary to store the mapping from tokens to their + expanded labels. + label_param_map: A dictionary to store the mapping from labels to their + corresponding parameter values. """ for token, orig_label in param_gen.labels.items(): for param in param_gen.parameters[token]: @@ -755,9 +1134,13 @@ def _create_param_maps(self, param_gen: "ParameterGenerator", expanded_labels: D def get_step_param_map(self) -> Dict: # pylint: disable=R0914 """ - Create a mapping of parameters used for each step. Each step will have a cmd - to search for parameters in and could also have a restart cmd to check, too. - This creates a mapping of the form: + Create a mapping of parameters used for each step in the study. + + This method generates a mapping of parameters for each step, where each + step may have a command (`cmd`) and a restart command (`restart_cmd`). + The resulting mapping has a structure similar to the following: + + ```python step_name_with_parameters: { "cmd": { TOKEN_1: param_1_value_1, @@ -768,8 +1151,11 @@ def get_step_param_map(self) -> Dict: # pylint: disable=R0914 TOKEN_3: param_3_value_1, } } + ``` - :returns: A dict mapping between steps and params of the form shown above + Returns: + A dictionary mapping step names (with parameters) to their + respective command and restart command parameter mappings. """ # Get the steps and the parameters in the study study_steps = self.get_study_steps() diff --git a/merlin/study/__init__.py b/merlin/study/__init__.py index 57477ea1f..a1c508a55 100644 --- a/merlin/study/__init__.py +++ b/merlin/study/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -27,3 +27,26 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### + +""" +The `study` package contains functionality for defining, managing, and monitoring studies +in Merlin. A study represents a collection of tasks, steps, and workflows that can be +executed on distributed systems using various schedulers. + +Modules: + batch.py: Parses the batch section of the YAML specification, supporting worker + launches for schedulers like Slurm, LSF, and Flux. + celeryadapter.py: Provides an adapter for integrating with the Celery Distributed + Task Queue, enabling distributed task execution. + dag.py: Defines the Merlin `DAG` class, which represents the Directed Acyclic Graph + structure of a study's workflow. + script_adapter.py: Contains functionality for adapting bash scripts to work with + supported schedulers, including Flux, LSF, and Slurm. + status_constants.py: Defines constants used by the `status` module and its renderers, + helping to avoid circular import issues. + status_renderers.py: Handles the creation of formatted, task-by-task status displays + for studies. + status.py: Implements functionality for retrieving and displaying the statuses of studies. + step.py: Contains the logic for representing and managing individual steps in a study. + study.py: Implements the core logic for defining and managing a study as a whole. +""" diff --git a/merlin/study/batch.py b/merlin/study/batch.py index 16482f399..6c3dfa9c3 100644 --- a/merlin/study/batch.py +++ b/merlin/study/batch.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -33,30 +33,49 @@ Currently only the batch worker launch for slurm, lsf or flux are implemented. - """ import logging import os import subprocess -from typing import Dict, Optional, Union +from typing import Dict, Union +from merlin.spec.specification import MerlinSpec from merlin.utils import convert_timestring, get_flux_alloc, get_flux_version, get_yaml_var LOG = logging.getLogger(__name__) -def batch_check_parallel(spec): +def batch_check_parallel(spec: MerlinSpec) -> bool: """ - Check for a parallel batch section in the yaml file. + Check for a parallel batch section in the provided MerlinSpec object. + + This function examines the 'batch' section of the given specification to determine + whether it is configured for parallel execution. It checks the 'type' attribute + within the batch section, defaulting to 'local' if not specified. If the type + is anything other than 'local', the function will return True, indicating that + parallel processing is enabled. + + Args: + spec (spec.specification.MerlinSpec): An instance of the + [`MerlinSpec`][spec.specification.MerlinSpec] class that contains the + configuration details, including the batch section. + + Returns: + Returns True if the batch type is set to a value other than 'local', + indicating that parallel processing is enabled; otherwise, returns False. + + Raises: + AttributeError: If the 'batch' section is not present in the specification, + an error is logged and an AttributeError is raised. """ parallel = False try: batch = spec.batch - except AttributeError: + except AttributeError as exc: LOG.error("The batch section is required in the specification file.") - raise + raise exc btype = get_yaml_var(batch, "type", "local") if btype != "local": @@ -65,14 +84,29 @@ def batch_check_parallel(spec): return parallel -def check_for_scheduler(scheduler, scheduler_legend): +def check_for_scheduler(scheduler: str, scheduler_legend: Dict[str, str]) -> bool: """ - Check which scheduler (Flux, Slurm, LSF, or PBS) is the main - scheduler for the cluster. - :param `scheduler`: A string representing the scheduler to check for - Options: flux, slurm, lsf, or pbs - :param `scheduler_legend`: A dict of information related to each scheduler - :returns: A bool representing whether `scheduler` is the main scheduler for the cluster + Check which scheduler (Flux, Slurm, LSF, or PBS) is the main scheduler for the cluster. + + This function verifies if the specified scheduler is the main scheduler by executing + a command associated with it from the provided scheduler legend. It returns a boolean + indicating whether the specified scheduler is active. + + Args: + scheduler: A string representing the scheduler to check for. Options include 'flux', + 'slurm', 'lsf', or 'pbs'. + scheduler_legend: A dictionary containing information related to each scheduler, + including the command to check its status and the expected output. See + [`construct_scheduler_legend`][study.batch.construct_scheduler_legend] + for more information on all the settings this dict contains. + + Returns: + Returns True if the specified scheduler is the main scheduler for the + cluster, otherwise returns False. + + Raises: + FileNotFoundError: If the command associated with the scheduler cannot be found. + PermissionError: If there are insufficient permissions to execute the command. """ # Check for invalid scheduler if scheduler not in ("flux", "slurm", "lsf", "pbs"): @@ -96,14 +130,26 @@ def check_for_scheduler(scheduler, scheduler_legend): return False -def get_batch_type(scheduler_legend, default=None): +def get_batch_type(scheduler_legend: Dict[str, str], default: str = None) -> str: """ Determine which batch scheduler to use. - :param scheduler_legend: A dict storing info related to each scheduler - :param default: (str) The default batch scheduler to use if a scheduler - can't be determined. The default is None. - :returns: (str) The batch name (available options: slurm, flux, lsf, pbs). + This function checks a predefined list of batch schedulers in a specific order + to determine which one is available for use. If none of the schedulers are found, + it checks the system type environment variable to suggest a default scheduler. + If no suitable scheduler is determined, it returns the specified default value. + + Args: + scheduler_legend: A dictionary storing information related to each + scheduler, including commands and expected outputs for checking their + availability. See [`construct_scheduler_legend`][study.batch.construct_scheduler_legend] + for more information on all the settings this dict contains. + default: The default batch scheduler to use if a scheduler cannot be determined. + + Returns: + The name of the available batch scheduler. Possible options include + 'slurm', 'flux', 'lsf', or 'pbs'. If no scheduler is found, returns + the specified default value. """ # These schedulers are listed in order of which should be checked for first # 1. Flux should be checked first due to slurm emulation scripts @@ -126,13 +172,29 @@ def get_batch_type(scheduler_legend, default=None): return default -def get_node_count(parsed_batch: Dict, default=1): +def get_node_count(parsed_batch: Dict, default: int = 1) -> int: """ Determine a default node count based on the environment. - :param default: (int) The number of nodes to return if a node count from - the environment cannot be determined. - :param returns: (int) The number of nodes to use. + This function checks the environment and the Flux version to determine the + appropriate number of nodes to use for batch processing. It first verifies + the Flux version, then attempts to retrieve the node count from the Flux + allocation or environment variables specific to Slurm or LSF. If no valid + node count can be determined, it returns a specified default value. + + Args: + parsed_batch: A dictionary containing parsed batch configurations. + See [`parse_batch_block`][study.batch.parse_batch_block] for more + information on all the settings in this dictionary. + default: The number of nodes to return if a node count from the + environment cannot be determined. + + Returns: + The number of nodes to use for the batch job. This value is determined + based on the environment and scheduler specifics. + + Raises: + ValueError: If the Flux version is too old (below 0.17.0). """ # Flux version check @@ -166,9 +228,38 @@ def get_node_count(parsed_batch: Dict, default=1): def parse_batch_block(batch: Dict) -> Dict: """ - A function to parse the batch block of the yaml file. - :param `batch`: The batch block to read in - :returns: A dict with all the info (or defaults) from the batch block + Parse the batch block of a YAML configuration file. + + This function extracts relevant information from the provided batch block + dictionary, including paths, execution options, and defaults. It retrieves + the Flux executable path and allocation details, and populates a dictionary + with the parsed values. + + Args: + batch: A dictionary representing the batch block from the YAML + configuration file. + + Returns: + A dictionary containing parsed information from the batch block, + including:\n + - `btype`: The type of batch job (default is 'local'). + - `nodes`: The number of nodes to use (default is None). + - `shell`: The shell to use (default is 'bash'). + - `bank`: The bank to charge for the job (default is an empty string). + - `queue`: The queue to submit the job to (default is an empty string). + - `walltime`: The maximum wall time for the job (default is an empty string). + - `launch pre`: Any commands to run before launching (default is an empty string). + - `launch args`: Arguments for the launch command (default is an empty string). + - `launch command`: Custom command to launch workers. This will override the + default launch command (default is an empty string). + - `flux path`: Optional path to flux bin. + - `flux exe`: The full path to the Flux executable. + - `flux exec`: Optional flux exec command to launch workers on all nodes if + `flux_exec_workers` is True (default is None). + - `flux alloc`: The Flux allocation retrieved from the executable. + - `flux opts`: Optional flux start options (default is an empty string). + - `flux exec workers`: Optional flux argument to launch workers + on all nodes (default is True). """ flux_path: str = get_yaml_var(batch, "flux_path", "") if "/" in flux_path: @@ -204,9 +295,20 @@ def parse_batch_block(batch: Dict) -> Dict: def get_flux_launch(parsed_batch: Dict) -> str: """ - Build the flux launch command based on the batch section of the yaml. - :param `parsed_batch`: A dict of batch configurations - :returns: The flux launch command + Build the Flux launch command based on the batch section of the YAML configuration. + + This function constructs the command to launch a Flux job using the parameters + specified in the parsed batch configuration. It determines the appropriate + execution command for Flux workers and integrates it with the launch command + provided in the batch configuration. + + Args: + parsed_batch: A dictionary containing batch configuration parameters. + See [`parse_batch_block`][study.batch.parse_batch_block] for more information + on all the settings in this dictionary. + + Returns: + The constructed Flux launch command, ready to be executed. """ default_flux_exec = "flux exec" if parsed_batch["launch command"] else f"{parsed_batch['flux exe']} exec" flux_exec: str = "" @@ -225,20 +327,37 @@ def get_flux_launch(parsed_batch: Dict) -> str: def batch_worker_launch( - spec: Dict, + spec: MerlinSpec, com: str, - nodes: Optional[Union[str, int]] = None, - batch: Optional[Dict] = None, + nodes: Union[str, int] = None, + batch: Dict = None, ) -> str: """ - The configuration in the batch section of the merlin spec - is used to create the worker launch line, which may be - different from a simulation launch. - - : param spec : (Dict) workflow specification - : param com : (str): The command to launch with batch configuration - : param nodes : (Optional[Union[str, int]]): The number of nodes to use in the batch launch - : param batch : (Optional[Dict]): An optional batch override from the worker config + Create the worker launch command based on the batch configuration in the + workflow specification. + + This function constructs a command to launch a worker process using the + specified batch configuration. It handles different batch types and + integrates any necessary pre-launch commands, launch arguments, and + node specifications. + + Args: + spec (spec.specification.MerlinSpec): An instance of the + [`MerlinSpec`][spec.specification.MerlinSpec] class that contains the + configuration details, including the batch section. + com: The command to launch with the batch configuration. + nodes: The number of nodes to use in the batch launch. If not specified, + it will default to the value in the batch configuration. + batch: An optional batch override from the worker configuration. If not + provided, the function will attempt to retrieve the batch section from + the specification. + + Returns: + The constructed worker launch command, ready to be executed. + + Raises: + AttributeError: If the batch section is missing in the specification. + TypeError: If the `nodes` parameter is of an invalid type. """ if batch is None: try: @@ -289,17 +408,34 @@ def batch_worker_launch( def construct_scheduler_legend(parsed_batch: Dict, nodes: int) -> Dict: """ - Constructs a legend of relevant information needed for each scheduler. This includes: - - bank (str): The flag to add a bank to the launch command - - check cmd (list): The command to run to check if this is the main scheduler for the cluster - - expected check output (str): The expected output from running the check cmd - - launch (str): The initial launch command for the scheduler - - queue (str): The flag to add a queue to the launch command - - walltime (str): The flag to add a walltime to the launch command - - :param `parsed_batch`: A dict of batch configurations - :param `nodes`: An int representing the number of nodes to use in a launch command - :returns: A dict of scheduler related information + Constructs a legend of relevant information needed for each scheduler. + + This function generates a dictionary containing configuration details for various + job schedulers based on the provided batch configuration. The returned dictionary + includes flags for bank, queue, and walltime, as well as commands to check the + scheduler and the initial launch command. + + Args: + parsed_batch: A dictionary of batch configurations, which must include `bank`, + `queue`, `walltime`, and `flux alloc`. See + [`parse_batch_block`][study.batch.parse_batch_block] for more information on + all the settings in this dictionary. + nodes: The number of nodes to use in the launch command. + + Returns: + A dictionary containing scheduler-related information, structured as + follows:\n + - For each scheduler (e.g., 'flux', 'lsf', 'pbs', 'slurm'):\n + - `bank` (str): The flag to add a bank to the launch command. + - `check cmd` (List[str]): The command to run to check if this is the main + scheduler for the cluster. + - `expected check output` (bytes): The expected output from running + the check command. + - `launch` (str): The initial launch command for the scheduler. + - `queue` (str): The flag to add a queue to the launch command (if + applicable). + - `walltime` (str): The flag to add a walltime to the launch command + (if applicable). """ scheduler_legend = { "flux": { @@ -338,11 +474,26 @@ def construct_scheduler_legend(parsed_batch: Dict, nodes: int) -> Dict: def construct_worker_launch_command(parsed_batch: Dict, nodes: int) -> str: """ - If no 'worker_launch' is found in the batch yaml, this method constructs the needed launch command. - - :param `parsed_batch`: A dict of batch configurations - :param `nodes`:: The number of nodes to use in the batch launch - :returns: The launch command + Constructs the worker launch command based on the provided batch configuration. + + This function generates a launch command for a worker process when no + 'worker_launch' command is specified in the batch configuration. It + utilizes the scheduler legend to incorporate necessary flags such as + bank, queue, and walltime, depending on the workload manager. + + Args: + parsed_batch: A dictionary of batch configurations, which must include + `btype`, `bank`, `queue`, and `walltime`. See + [`parse_batch_block`][study.batch.parse_batch_block] for more information + on all the settings in this dictionary. + nodes: The number of nodes to use in the batch launch. + + Returns: + The constructed launch command for the worker process. + + Raises: + TypeError: If the PBS scheduler is enabled for a batch type other than 'flux'. + KeyError: If the workload manager is not found in the scheduler legend. """ # Initialize launch_command and get the scheduler_legend and workload_manager launch_command: str = "" diff --git a/merlin/study/celeryadapter.py b/merlin/study/celeryadapter.py index 5b5bdd419..fbc3a7488 100644 --- a/merlin/study/celeryadapter.py +++ b/merlin/study/celeryadapter.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -38,7 +38,8 @@ import time from contextlib import suppress from datetime import datetime -from typing import Dict, List, Optional, Tuple +from types import SimpleNamespace +from typing import Dict, List, Set, Tuple from amqp.exceptions import ChannelError from celery import Celery @@ -46,7 +47,9 @@ from merlin.common.dumper import dump_handler from merlin.config import Config +from merlin.spec.specification import MerlinSpec from merlin.study.batch import batch_check_parallel, batch_worker_launch +from merlin.study.study import MerlinStudy from merlin.utils import apply_list_of_regex, check_machines, get_procs, get_yaml_var, is_running @@ -55,10 +58,20 @@ # TODO figure out a better way to handle the import of celery app and CONFIG -def run_celery(study, run_mode=None): +def run_celery(study: MerlinStudy, run_mode: str = None): """ - Run the given MerlinStudy object. If the run mode is set to "local" - configure Celery to run locally (without workers). + Run the given [`MerlinStudy`][study.study.MerlinStudy] object with optional + Celery configuration. + + This function executes the provided [`MerlinStudy`][study.study.MerlinStudy] + object. If the `run_mode` is set to "local", it configures Celery to run in + local mode (without utilizing workers). Otherwise, it connects to the Celery + server to queue tasks. + + Args: + study (study.study.MerlinStudy): The study object to be executed. + run_mode: The mode in which to run the study. If set to "local", + Celery runs locally. """ # Only import celery stuff if we want celery in charge # Pylint complains about circular import between merlin.common.tasks -> merlin.router -> merlin.study.celeryadapter @@ -81,14 +94,25 @@ def run_celery(study, run_mode=None): def get_running_queues(celery_app_name: str, test_mode: bool = False) -> List[str]: """ - Check for running celery workers by looking at the currently running processes. - If there are running celery workers, we'll pull the queues from the -Q tag in the - process command. The list returned here will contain only unique celery queue names. - This must be run on the allocation where the workers are running. + Check for running Celery workers and retrieve their associated queues. - :param `celery_app_name`: The name of the celery app (typically merlin here unless testing) - :param `test_mode`: If True, run this function in test mode - :returns: A unique list of celery queues with workers attached to them + This function inspects currently running processes to identify active + Celery workers. It extracts queue names from the `-Q` tag in the + command line of the worker processes. The returned list contains + only unique Celery queue names. This function must be executed + on the allocation where the workers are running. + + Note: + Unlike [`get_active_celery_queues`][study.celeryadapter.get_active_celery_queues], + this function does _not_ go through the application's server. + + Args: + celery_app_name: The name of the Celery app (typically "merlin" + unless in test mode). + test_mode: If True, the function runs in test mode. + + Returns: + A unique list of Celery queue names with workers attached to them. """ running_queues = [] @@ -111,26 +135,36 @@ def get_running_queues(celery_app_name: str, test_mode: bool = False) -> List[st return running_queues -def get_active_celery_queues(app): - """Get all active queues and workers for a celery application. - - Unlike get_running_queues, this goes through the application's server. - Also returns a dictionary with entries for each worker attached to - the given queues. - - :param `celery.Celery` app: the celery application - - :return: queues dictionary with connected workers, all workers - :rtype: (dict of lists of strings, list of strings) - - :example: - - >>> from merlin.celery import app - >>> queues, workers = get_active_celery_queues(app) - >>> queue_names = [*queues] - >>> workers_on_q0 = queues[queue_names[0]] - >>> workers_not_on_q0 = [worker for worker in workers - if worker not in workers_on_q0] +def get_active_celery_queues(app: Celery) -> Tuple[Dict[str, List[str]], List[str]]: + """ + Retrieve all active queues and their associated workers for a Celery application. + + This function queries the application's server to obtain a comprehensive + view of active queues and the workers connected to them. It returns a + dictionary where each key is a queue name and the value is a list of + workers attached to that queue. Additionally, it provides a list of all + active workers in the application. + + Note: + Unlike [`get_running_queues`][study.celeryadapter.get_running_queues], + this function goes through the application's server. + + Args: + app: The Celery application instance. + + Returns: + A tuple containing:\n + - A dictionary mapping queue names to lists of workers connected to them. + - A list of all active workers in the application. + + Example: + ```python + from merlin.celery import app + queues, workers = get_active_celery_queues(app) + queue_names = list(queues) + workers_on_q0 = queues[queue_names[0]] + workers_not_on_q0 = [worker for worker in workers if worker not in workers_on_q0] + ``` """ i = app.control.inspect() active_workers = i.active_queues() @@ -146,14 +180,22 @@ def get_active_celery_queues(app): return queues, [*active_workers] -def get_active_workers(app): +def get_active_workers(app: Celery) -> Dict[str, List[str]]: """ - This is the inverse of get_active_celery_queues() defined above. This function - builds a dict where the keys are worker names and the values are lists - of queues attached to the worker. + Retrieve a mapping of active workers to their associated queues for a Celery application. + + This function serves as the inverse of + [`get_active_celery_queues()`][study.celeryadapter.get_active_celery_queues]. It constructs + a dictionary where each key is a worker's name and the corresponding value is a + list of queues that the worker is connected to. This allows for easy identification + of which queues are being handled by each worker. + + Args: + app: The Celery application instance. - :param `app`: The celery application - :returns: A dict mapping active workers to queues + Returns: + A dictionary mapping active worker names to lists of queue names they are + attached to. If no active workers are found, an empty dictionary is returned. """ # Get the information we need from celery i = app.control.inspect() @@ -173,14 +215,19 @@ def get_active_workers(app): return worker_queue_map -def celerize_queues(queues: List[str], config: Optional[Dict] = None): +def celerize_queues(queues: List[str], config: SimpleNamespace = None): """ - Celery requires a queue tag to be prepended to their - queues so this function will 'celerize' every queue in - a list you provide it by prepending the queue tag. + Prepend a queue tag to each queue in the provided list to conform to Celery's + queue naming requirements. - :param `queues`: A list of queues that need the queue tag prepended. - :param `config`: A dict of configuration settings + This function modifies the input list of queues by adding a specified queue tag + from the configuration. If no configuration is provided, it defaults to using + the global configuration settings. + + Args: + queues: A list of queue names that need the queue tag prepended. + config: A SimpleNamespace of configuration settings. If not provided, the + function will use the default configuration. """ if config is None: from merlin.config.configfile import CONFIG as config # pylint: disable=C0415 @@ -189,14 +236,18 @@ def celerize_queues(queues: List[str], config: Optional[Dict] = None): queues[i] = f"{config.celery.queue_tag}{queue}" -def _build_output_table(worker_list, output_table): +def _build_output_table(worker_list: List[str], output_table: List[Tuple[str, str]]): """ - Helper function for query-status that will build a table - that we'll use as output. + Construct an output table for displaying the status of workers and their associated queues. + + This helper function populates the provided output table with entries for each worker + in the given worker list. It retrieves the mapping of active workers to their queues + and formats the data accordingly. - :param `worker_list`: A list of workers to add to the table - :param `output_table`: A list of tuples where each entry is - of the form (worker name, associated queues) + Args: + worker_list: A list of worker names to be included in the output table. + output_table: A list of tuples where each entry will be of the form + (worker name, associated queues). """ from merlin.celery import app # pylint: disable=C0415 @@ -211,15 +262,22 @@ def _build_output_table(worker_list, output_table): output_table.append((worker, ", ".join(worker_queue_map[worker]))) -def query_celery_workers(spec_worker_names, queues, workers_regex): +def query_celery_workers(spec_worker_names: List[str], queues: List[str], workers_regex: List[str]): """ - Look for existing celery workers. Filter by spec, queues, or - worker names if provided by user. At the end, print a table - of workers and their associated queues. - - :param `spec_worker_names`: The worker names defined in a spec file - :param `queues`: A list of queues to filter by - :param `workers_regex`: A list of regexs to filter by + Query and filter existing Celery workers based on specified criteria, + and print a table of the workers along with their associated queues. + + This function retrieves the list of active Celery workers and filters them + according to the provided specifications, including worker names from a + spec file, specific queues, and regular expressions for worker names. + It then constructs and displays a table of the matching workers and their + associated queues. + + Args: + spec_worker_names: A list of worker names defined in a spec file + to filter the workers. + queues: A list of queues to filter the workers by. + workers_regex: A list of regular expressions to filter the worker names. """ from merlin.celery import app # pylint: disable=C0415 @@ -280,11 +338,21 @@ def query_celery_workers(spec_worker_names, queues, workers_regex): def build_csv_queue_info(query_return: List[Tuple[str, int, int]], date: str) -> Dict[str, List]: """ - Build the lists of column labels and queue info to write to the csv file. + Construct a dictionary containing queue information and column labels + for writing to a CSV file. + + This function processes the output from the [`query_queues`][router.query_queues] + function and organizes the data into a format suitable for CSV export. It includes + a timestamp to indicate when the status was recorded. + + Args: + query_return: The output from the [`query_queues`][router.query_queues] function, + containing queue names and their associated statistics. + date: A timestamp indicating when the queue status was recorded. - :param query_return: The output of `query_queues` - :param date: A timestamp for us to mark when this status occurred - :returns: A dict of queue information to dump to a csv file + Returns: + A dictionary where keys are column labels and values are lists containing the + corresponding queue information, formatted for CSV output. """ # Build the list of labels if necessary csv_to_dump = {"time": [date]} @@ -297,11 +365,22 @@ def build_csv_queue_info(query_return: List[Tuple[str, int, int]], date: str) -> def build_json_queue_info(query_return: List[Tuple[str, int, int]], date: str) -> Dict: """ - Build the dict of queue info to dump to the json file. - - :param query_return: The output of `query_queues` - :param date: A timestamp for us to mark when this status occurred - :returns: A dictionary that's ready to dump to a json outfile + Construct a dictionary containing queue information for JSON export. + + This function processes the output from the [`query_queues`][router.query_queues] + function and organizes the data into a structured format suitable for JSON + serialization. It includes a timestamp to indicate when the queue status was + recorded. + + Args: + query_return: The output from the [`query_queues`][router.query_queues] + function, containing queue names and their associated statistics. + date: A timestamp indicating when the queue status was recorded. + + Returns: + A dictionary structured for JSON output, where the keys are timestamps + and the values are dictionaries containing queue names and their + corresponding statistics (tasks and consumers). """ # Get the datetime so we can track different entries and initalize a new json entry json_to_dump = {date: {}} @@ -315,11 +394,17 @@ def build_json_queue_info(query_return: List[Tuple[str, int, int]], date: str) - def dump_celery_queue_info(query_return: List[Tuple[str, int, int]], dump_file: str): """ - Format the information we're going to dump in a way that the Dumper class can - understand and add a timestamp to the info. + Format and dump Celery queue information to a specified file. + + This function processes the output from the `query_queues` function, formats + the data according to the file type (CSV or JSON), and adds a timestamp + to the information before writing it to the specified file. - :param query_return: The output of `query_queues` - :param dump_file: The filepath of the file we'll dump queue info to + Args: + query_return: The output from the [`query_queues`][router.query_queues] + function, containing queue names and their associated statistics. + dump_file: The filepath of the file where the queue information + will be written. The file extension determines the format (CSV or JSON). """ # Get a timestamp for this dump date = datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -336,16 +421,25 @@ def dump_celery_queue_info(query_return: List[Tuple[str, int, int]], dump_file: dump_handler(dump_file, dump_info) -def _get_specific_queues(queues: set, specific_queues: List[str], spec: "MerlinSpec", verbose=True) -> set: # noqa: F821 +def _get_specific_queues(queues: Set[str], specific_queues: List[str], spec: MerlinSpec, verbose=True) -> Set[str]: """ - Search for specific queues that the user asked for. The queues that cannot be found will not - be returned. The queues that can be found will be added to a set and returned. - - :param queues: Either an empty set or a set of queues from `spec` - :param specific_queues: The list of queues that we're going to search for - :param spec: A `MerlinSpec` object or None - :param verbose: If True, display log messages. Otherwise, don't. - :returns: A set of the specific queues that were found to exist. + Retrieve a set of specific queues requested by the user, filtering out those that do not exist. + + This function checks a provided list of specific queues against a set of existing queues + (from a [`MerlinSpec`][spec.specification.MerlinSpec] object) and returns a set of queues + that are found. If a queue is not found in the existing set, it will be excluded from the + results. The function also logs messages based on the verbosity setting. + + Args: + queues: A set of existing queues, which may be empty or populated from the `spec` + object. + specific_queues: A list of specific queue names to search for. + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object that may provide context for the search. Can be None. + verbose: If True, log messages will be displayed. + + Returns: + A set containing the specific queues that were found in the existing queues. """ if verbose: LOG.info(f"Filtering queues to query by these specific queues: {specific_queues}") @@ -376,21 +470,30 @@ def _get_specific_queues(queues: set, specific_queues: List[str], spec: "MerlinS def build_set_of_queues( - spec: "MerlinSpec", # noqa: F821 + spec: MerlinSpec, steps: List[str], specific_queues: List[str], - verbose: Optional[bool] = True, - app: Optional["Celery"] = None, # noqa: F821 -) -> set: + verbose: bool = True, + app: Celery = None, +) -> Set[str]: """ - Build a set of queues to query based on the parameters given here. - - :param spec: A `MerlinSpec` object or None - :param steps: Spaced-separated list of stepnames to query. Default is all - :param specific_queues: A list of queue names to query or None - :param verbose: A bool to determine whether to output log statements or not - :param app: A celery app object, if left out we'll just import it - :returns: A set of queues to investigate + Construct a set of queues to query based on the provided parameters. + + This function builds a set of queues by querying a [`MerlinSpec`][spec.specification.MerlinSpec] + object for queues associated with specified steps and/or filtering for specific queue names. + If no spec or specific queues are provided, it defaults to querying active queues from the Celery + application. + + Args: + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object that defines the context for the query. Can be None. + steps: A list of step names to query. If empty, all steps are considered. + specific_queues: A list of specific queue names to filter. Can be None. + verbose: If True, log statements will be output. Defaults to True. + app: A Celery application instance. If None, it will be imported. + + Returns: + A set of queue names to investigate based on the provided parameters. """ if app is None: from merlin.celery import app # pylint: disable=C0415 @@ -424,15 +527,32 @@ def build_set_of_queues( return queues -def query_celery_queues(queues: List[str], app: Celery = None, config: Config = None) -> Dict[str, List[str]]: +def query_celery_queues(queues: List[str], app: Celery = None, config: Config = None) -> Dict[str, Dict[str, int]]: """ - Build a dict of information about the number of jobs and consumers attached - to specific queues that we want information on. - - :param queues: A list of the queues we want to know about - :param app: The celery application (this will be none unless testing) - :param config: The configuration object that has the broker name (this will be none unless testing) - :returns: A dict of info on the number of jobs and consumers for each queue in `queues` + Retrieve information about the number of jobs and consumers for specified Celery queues. + + This function constructs a dictionary containing details about the number of jobs + and consumers associated with each queue provided in the input list. It connects + to the Celery application to gather this information, handling both Redis and + RabbitMQ brokers. + + Notes: + - If the specified queue does not exist or has no jobs, it will be handled gracefully. + - For Redis brokers, the function counts consumers by inspecting active queues + since Redis does not track consumers like RabbitMQ does. + + Args: + queues: A list of queue names for which to gather information. + app: The Celery application instance. Defaults to None, which triggers an import + for testing purposes. + config (config.Config): A configuration object containing broker details. + Defaults to None, which also triggers an import for testing. + + Returns: + A dictionary where each key is a queue name and the value is another dictionary + containing:\n + - `jobs`: The number of jobs in the queue. + - `consumers`: The number of consumers attached to the queue. """ if app is None: from merlin.celery import app # pylint: disable=C0415 @@ -474,12 +594,17 @@ def query_celery_queues(queues: List[str], app: Celery = None, config: Config = return queue_info -def get_workers_from_app(): - """Get all workers connected to a celery application. +def get_workers_from_app() -> List[str]: + """ + Retrieve a list of all workers connected to the Celery application. + + This function uses the Celery control interface to inspect the current state + of the application and returns a list of workers that are currently connected. + If no workers are found, an empty list is returned. - :param `celery.Celery` app: the celery application - :return: A list of all connected workers - :rtype: list + Returns: + A list of worker names that are currently connected to the Celery application. + If no workers are connected, an empty list is returned. """ from merlin.celery import app # pylint: disable=C0415 @@ -492,11 +617,19 @@ def get_workers_from_app(): def check_celery_workers_processing(queues_in_spec: List[str], app: Celery) -> bool: """ - Query celery to see if any workers are still processing tasks. + Check if any Celery workers are currently processing tasks from specified queues. + + This function queries the Celery application to determine if there are any active + tasks being processed by workers for the given list of queues. It returns a boolean + indicating whether any tasks are currently active. - :param queues_in_spec: A list of queues to check if tasks are still active in - :param app: The celery app that we're querying - :returns: True if workers are still processing tasks, False otherwise + Args: + queues_in_spec: A list of queue names to check for active tasks. + app: The Celery application instance used for querying. + + Returns: + True if any workers are processing tasks in the specified queues; False + otherwise. """ # Query celery for active tasks active_tasks = app.control.inspect().active() @@ -511,15 +644,23 @@ def check_celery_workers_processing(queues_in_spec: List[str], app: Celery) -> b return False -def _get_workers_to_start(spec, steps): +def _get_workers_to_start(spec: MerlinSpec, steps: List[str]) -> Set[str]: """ - Helper function to return a set of workers to start based on - the steps provided by the user. + Determine the set of workers to start based on the specified steps. + + This helper function retrieves a mapping of steps to their corresponding workers + from a [`MerlinSpec`][spec.specification.MerlinSpec] object and returns a unique + set of workers that should be started for the provided list of steps. If a step + is not found in the mapping, a warning is logged. - :param `spec`: A MerlinSpec object - :param `steps`: A list of steps to start workers for + Args: + spec (spec.specification.MerlinSpec): An instance of the + [`MerlinSpec`][spec.specification.MerlinSpec] class that contains the + mapping of steps to workers. + steps: A list of steps for which workers need to be started. - :returns: A set of workers to start + Returns: + A set of unique workers to be started based on the specified steps. """ workers_to_start = [] step_worker_map = spec.get_step_worker_map() @@ -535,14 +676,25 @@ def _get_workers_to_start(spec, steps): return workers_to_start -def _create_kwargs(spec): +def _create_kwargs(spec: MerlinSpec) -> Tuple[Dict[str, str], Dict]: """ - Helper function to handle creating the kwargs dict that - we'll pass to subprocess.Popen when we launch the worker. - - :param `spec`: A MerlinSpec object - :returns: A tuple where the first entry is the kwargs and - the second entry is variables defined in the spec + Construct the keyword arguments for launching a worker process. + + This helper function creates a dictionary of keyword arguments that will be + passed to `subprocess.Popen` when launching a worker. It retrieves the + environment variables defined in a [`MerlinSpec`][spec.specification.MerlinSpec] + object and updates the shell environment accordingly. + + Args: + spec (spec.specification.MerlinSpec): An instance of the MerlinSpec class + that contains environment specifications. + + Returns: + A tuple containing: + - A dictionary of keyword arguments for `subprocess.Popen`, including + the updated environment. + - A dictionary of variables defined in the spec, or None if no variables + were defined. """ # Get the environment from the spec and the shell spec_env = spec.environment @@ -563,15 +715,24 @@ def _create_kwargs(spec): return kwargs, yaml_vars -def _get_steps_to_start(wsteps, steps, steps_provided): +def _get_steps_to_start(wsteps: List[str], steps: List[str], steps_provided: bool) -> List[str]: """ - Determine which steps to start workers for. - - :param `wsteps`: A list of steps associated with a worker - :param `steps`: A list of steps to start provided by the user - :param `steps`: A bool representing whether the user gave specific - steps to start or not - :returns: A list of steps to start workers for + Identify the steps for which workers should be started. + + This function determines which steps to initiate based on the steps + associated with a worker and the user-provided steps. If specific steps + are provided by the user, only those steps that match the worker's steps + will be included. If no specific steps are provided, all worker-associated + steps will be returned. + + Args: + wsteps: A list of steps that are associated with a worker. + steps: A list of steps specified by the user to start workers for. + steps_provided: A boolean indicating whether the user provided + specific steps to start. + + Returns: + A list of steps for which workers should be started. """ steps_to_start = [] if steps_provided: @@ -584,31 +745,48 @@ def _get_steps_to_start(wsteps, steps, steps_provided): return steps_to_start -def start_celery_workers(spec, steps, celery_args, disable_logs, just_return_command): # pylint: disable=R0914,R0915 - """Start the celery workers on the allocation - - :param MerlinSpec spec: A MerlinSpec object representing our study - :param list steps: A list of steps to start workers for - :param str celery_args: A string of arguments to provide to the celery workers - :param bool disable_logs: A boolean flag to turn off the celery logs for the workers - :param bool just_return_command: When True, workers aren't started and just the launch command(s) - are returned - :side effect: Starts subprocesses for each worker we launch - :returns: A string of all the worker launch commands - ... - - example config: - - merlin: - resources: - task_server: celery - overlap: False - workers: - simworkers: - args: -O fair --prefetch-multiplier 1 -E -l info --concurrency 4 - steps: [run, data] - nodes: 1 - machine: [hostA, hostB] +def start_celery_workers( + spec: MerlinSpec, steps: List[str], celery_args: str, disable_logs: bool, just_return_command: bool +) -> str: # pylint: disable=R0914,R0915 + """ + Start Celery workers based on the provided specifications and steps. + + This function initializes and starts Celery workers for the specified steps + in the given [`MerlinSpec`][spec.specification.MerlinSpec]. It constructs + the necessary command-line arguments and handles the launching of subprocesses + for each worker. If the `just_return_command` flag is set to `True`, it will + return the command(s) to start the workers without actually launching them. + + Args: + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object representing the study configuration. + steps: A list of steps for which to start workers. + celery_args: A string of additional arguments to pass to the Celery workers. + disable_logs: A flag to disable logging for the Celery workers. + just_return_command: If `True`, returns the launch command(s) without starting the workers. + + Returns: + A string containing all the worker launch commands. + + Side Effects: + - Starts subprocesses for each worker that is launched, so long as `just_return_command` + is not True. + + Example: + Below is an example configuration for Merlin workers: + + ```yaml + merlin: + resources: + task_server: celery + overlap: False + workers: + simworkers: + args: -O fair --prefetch-multiplier 1 -E -l info --concurrency 4 + steps: [run, data] + nodes: 1 + machine: [hostA, hostB] + ``` """ if not just_return_command: LOG.info("Starting workers") @@ -701,10 +879,25 @@ def start_celery_workers(spec, steps, celery_args, disable_logs, just_return_com return str(worker_list) -def examine_and_log_machines(worker_val, yenv) -> bool: +def examine_and_log_machines(worker_val: Dict, yenv: Dict[str, str]) -> bool: """ - Examines whether a worker should be skipped in a step of start_celery_workers(), logs errors in output path for a celery - worker. + Determine if a worker should be skipped based on machine availability and log any errors. + + This function checks the specified machines for a worker and determines + whether the worker can be started. If the machines are not available, + it logs an error message regarding the output path for the Celery worker. + If the environment variables (`yenv`) are not provided or do not specify + an output path, a warning is logged. + + Args: + worker_val: A dictionary containing worker configuration, including + the list of machines associated with the worker. + yenv: A dictionary of environment variables that may include the + output path for logging. + + Returns: + Returns `True` if the worker should be skipped (i.e., machines are + unavailable), otherwise returns `False`. """ worker_machines = get_yaml_var(worker_val, "machines", None) if worker_machines: @@ -725,8 +918,27 @@ def examine_and_log_machines(worker_val, yenv) -> bool: return False -def verify_args(spec, worker_args, worker_name, overlap, disable_logs=False): - """Examines the args passed to a worker for completeness.""" +def verify_args(spec: MerlinSpec, worker_args: str, worker_name: str, overlap: bool, disable_logs: bool = False) -> str: + """ + Validate and enhance the arguments passed to a Celery worker for completeness. + + This function checks the provided worker arguments to ensure that they include + recommended settings for running parallel tasks. It adds default values for + concurrency, prefetch multiplier, and logging level if they are not specified. + Additionally, it generates a unique worker name based on the current time if + the `-n` argument is not provided. + + Args: + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object containing the study configuration. + worker_args: A string of arguments passed to the worker that may need validation. + worker_name: The name of the worker, used for generating a unique worker identifier. + overlap: A flag indicating whether multiple workers can overlap in their queue processing. + disable_logs: A flag to disable logging configuration for the worker. + + Returns: + The validated and potentially modified worker arguments string. + """ parallel = batch_check_parallel(spec) if parallel: if "--concurrency" not in worker_args: @@ -750,15 +962,29 @@ def verify_args(spec, worker_args, worker_name, overlap, disable_logs=False): return worker_args -def launch_celery_worker(worker_cmd, worker_list, kwargs): +def launch_celery_worker(worker_cmd: str, worker_list: List[str], kwargs: Dict): """ - Using the worker launch command provided, launch a celery worker. - :param str worker_cmd: The celery command to launch a worker - :param list worker_list: A list of worker launch commands - :param dict kwargs: A dictionary containing additional keyword args to provide - to subprocess.Popen - - :side effect: Launches a celery worker via a subprocess + Launch a Celery worker using the specified command and parameters. + + This function executes the provided Celery command to start a worker as a + subprocess. It appends the command to the given list of worker commands + for tracking purposes. If the worker fails to start, an error is logged. + + Args: + worker_cmd: The command string used to launch the Celery worker. + worker_list: A list that will be updated to include the launched + worker command for tracking active workers. + kwargs: A dictionary of additional keyword arguments to pass to + `subprocess.Popen`, allowing for customization of the subprocess + behavior. + + Raises: + Exception: If the worker fails to start, an error is logged, and the + exception is re-raised. + + Side Effects: + - Launches a Celery worker process in the background. + - Modifies the `worker_list` by appending the launched worker command. """ try: _ = subprocess.Popen(worker_cmd, **kwargs) # pylint: disable=R1732 @@ -768,12 +994,25 @@ def launch_celery_worker(worker_cmd, worker_list, kwargs): raise -def get_celery_cmd(queue_names, worker_args="", just_return_command=False): +def get_celery_cmd(queue_names: str, worker_args: str = "", just_return_command: bool = False) -> str: """ - Get the appropriate command to launch celery workers for the specified MerlinStudy. - queue_names The name(s) of the queue(s) to associate a worker with - worker_args Optional celery arguments for the workers - just_return_command Don't execute, just return the command + Construct the command to launch Celery workers for the specified queues. + + This function generates a command string that can be used to start Celery + workers associated with the provided queue names. It allows for optional + worker arguments to be included and can return the command without executing it. + + Args: + queue_names: A comma-separated string of the queue name(s) to which the worker + will be associated. + worker_args: Additional command-line arguments for the Celery worker. + just_return_command: If True, the function will return the constructed command + without executing it. + + Returns: + The constructed command string for launching the Celery worker. If + `just_return_command` is True, returns the command; otherwise, returns an + empty string. """ worker_command = " ".join(["celery -A merlin worker", worker_args, "-Q", queue_names]) if just_return_command: @@ -783,12 +1022,24 @@ def get_celery_cmd(queue_names, worker_args="", just_return_command=False): return "" -def purge_celery_tasks(queues, force): +def purge_celery_tasks(queues: str, force: bool) -> int: """ - Purge celery tasks for the specified spec file. - - queues Which queues to purge - force Purge without asking for confirmation + Purge Celery tasks from the specified queues. + + This function constructs and executes a command to purge tasks from the + specified Celery queues. If the `force` parameter is set to True, the + purge operation will be executed without prompting for confirmation. + + Args: + queues: A comma-separated string of the queue name(s) from which + tasks should be purged. + force: If True, the purge operation will be executed without asking + for user confirmation. + + Returns: + The return code from the subprocess execution. A return code of + 0 indicates success, while any non-zero value indicates an error + occurred during the purge operation. """ # This version will purge all queues. # from merlin.celery import app @@ -801,24 +1052,33 @@ def purge_celery_tasks(queues, force): return subprocess.run(purge_command, shell=True).returncode -def stop_celery_workers(queues=None, spec_worker_names=None, worker_regex=None): # pylint: disable=R0912 - """Send a stop command to celery workers. - - Default behavior is to stop all connected workers. - As options can downselect to only workers on certain queues and/or that - match a regular expression. - - :param list queues: The queues to send stop signals to. If None: stop all - :param list spec_worker_names: Worker names read from a spec to stop, in addition to worker_regex matches. - :param str worker_regex: The regex string to match worker names. If None: - :return: Return code from stop command - - :example: - - >>> stop_celery_workers(queues=['hello'], worker_regex='celery@*my_machine*') - - >>> stop_celery_workers() - +def stop_celery_workers( + queues: List[str] = None, spec_worker_names: List[str] = None, worker_regex: List[str] = None +): # pylint: disable=R0912 + """ + Send a stop command to Celery workers. + + This function sends a shutdown command to Celery workers associated with + specified queues. By default, it stops all connected workers, but it can + be configured to target specific workers based on queue names or regular + expression patterns. + + Args: + queues: A list of queue names to which the stop command will be sent. + If None, all connected workers across all queues will be stopped. + spec_worker_names: A list of specific worker names to stop, in addition + to those matching the `worker_regex`. + worker_regex: A regular expression string used to match worker names. + If None, no regex filtering will be applied. + + Side Effects: + - Broadcasts a shutdown signal to Celery workers + + Example: + ```python + stop_celery_workers(queues=['hello'], worker_regex='celery@*my_machine*') + stop_celery_workers() + ``` """ from merlin.celery import app # pylint: disable=C0415 @@ -870,13 +1130,24 @@ def stop_celery_workers(queues=None, spec_worker_names=None, worker_regex=None): LOG.warning("No workers found to stop") -def create_celery_config(config_dir, data_file_name, data_file_path): +def create_celery_config(config_dir: str, data_file_name: str, data_file_path: str): """ - Command to setup default celery merlin config. - - :param `config_dir`: The directory to create the config file. - :param `data_file_name`: The name of the config file. - :param `data_file_path`: The full data file path. + Set up the default Celery configuration for Merlin. + + This function creates a configuration file for Celery in the specified + directory. If the configuration file already exists, it logs an + informational message and does not overwrite the existing file. If the + file does not exist, it reads from a specified data file and writes + its contents to the new configuration file. + + Args: + config_dir: The directory where the configuration file will be created. + data_file_name: The name of the configuration file to be created. + data_file_path: The full path to the data file from which the + configuration content will be read. + + Side Effects: + - Creates a configuration file if one does not already exist """ # This will need to come from the server interface MERLIN_CONFIG = os.path.join(config_dir, data_file_name) # pylint: disable=C0103 diff --git a/merlin/study/dag.py b/merlin/study/dag.py index c1b9dff78..71b1cffda 100644 --- a/merlin/study/dag.py +++ b/merlin/study/dag.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,57 +29,107 @@ ############################################################################### """ -Holds DAG class. TODO make this an interface, separate from Maestro. +Holds the Merlin Directed Acyclic Graph (DAG) class. """ from collections import OrderedDict +from typing import Dict, List from merlin.study.step import Step +# TODO make this an interface, separate from Maestro. class DAG: """ This class provides methods on a task graph that Merlin needs for staging - tasks in celery. It is initialized from am maestro ExecutionGraph, and the + tasks in Celery. It is initialized from a Maestro `ExecutionGraph`, and the major entry point is the group_tasks method, which provides groups of independent chains of tasks. + + Attributes: + backwards_adjacency (Dict): A dictionary mapping each task to its parent tasks for reverse + traversal. + column_labels (List[str]): A list of column labels provided in the spec file. + maestro_adjacency_table (OrderedDict): An ordered dict showing adjacency of nodes. Comes from + a maestrowf `ExecutionGraph`. + maestro_values (OrderedDict): An ordered dict of the values at each node. Comes from a maestrowf + `ExecutionGraph`. + parameter_info (Dict): A dict containing information about parameters in the study. + study_name (str): The name of the study. + + Methods: + calc_backwards_adjacency: Initializes the backwards adjacency table. + calc_depth: Calculate the depth of the given node and its children. + children: Return the children of the task. + compatible_merlin_expansion: Check if two tasks are compatible for Merlin expansion. + find_chain: Find the chain containing the task. + find_independent_chains: Finds independent chains and adjusts with the groups of chains + to maximize parallelism. + group_by_depth: Group Directed Acyclic Graph (DAG) tasks by depth. + group_tasks: Group independent tasks in a DAG. + num_children: Find the number of children for the given task in the DAG. + num_parents: Find the number of parents for the given task in the DAG. + parents: Return the parents of the task. + step: Return a [`Step`][study.step.Step] object for the given task name. """ def __init__( - self, maestro_adjacency_table, maestro_values, column_labels, study_name, parameter_info + self, + maestro_adjacency_table: OrderedDict, + maestro_values: OrderedDict, + column_labels: List[str], + study_name: str, + parameter_info: Dict, ): # pylint: disable=R0913 """ - :param `maestro_adjacency_table`: An ordered dict showing adjacency of nodes. Comes from a maestrowf ExecutionGraph. - :param `maestro_values`: An ordered dict of the values at each node. Comes from a maestrowf ExecutionGraph. - :param `column_labels`: A list of column labels provided in the spec file. - :param `study_name`: The name of the study - :param `parameter_info`: A dict containing information about parameters in the study + Initializes a Directed Acyclic Graph (DAG) object, which represents a task graph used by Merlin + for staging tasks in Celery. The DAG is initialized from a Maestro `ExecutionGraph` by unpacking + its adjacency table and node values. + + Args: + maestro_adjacency_table: An ordered dictionary representing the adjacency + relationships between tasks in the graph. This comes from a Maestro `ExecutionGraph`. + maestro_values: An ordered dictionary containing the values or metadata + associated with each task in the graph. This also comes from a Maestro `ExecutionGraph`. + column_labels: A list of column labels provided in the specification file, + typically used to identify parameters or task attributes. + study_name: The name of the study to which this DAG belongs. + parameter_info: A dictionary containing information about the parameters in the study, + such as their names and values. """ # We used to store the entire maestro ExecutionGraph here but now it's # unpacked so we're only storing the 2 attributes from it that we use: # the adjacency table and the values. This had to happen to get pickle # to work for Celery. - self.maestro_adjacency_table = maestro_adjacency_table - self.maestro_values = maestro_values - self.column_labels = column_labels - self.study_name = study_name - self.parameter_info = parameter_info - self.backwards_adjacency = {} + self.maestro_adjacency_table: OrderedDict = maestro_adjacency_table + self.maestro_values: OrderedDict = maestro_values + self.column_labels: List[str] = column_labels + self.study_name: str = study_name + self.parameter_info: Dict = parameter_info + self.backwards_adjacency: Dict = {} self.calc_backwards_adjacency() - def step(self, task_name): - """Return a Step object for the given task name + def step(self, task_name: str) -> Step: + """ + Return a Step object for the given task name. + + Args: + task_name: The task name. - :param `task_name`: The task name. - :return: A Merlin Step object. + Returns: + A Merlin [`Step`][study.step.Step] object representing the + task's configuration and parameters. """ return Step(self.maestro_values[task_name], self.study_name, self.parameter_info) - def calc_depth(self, node, depths, current_depth=0): - """Calculate the depth of the given node and its children. + def calc_depth(self, node: str, depths: Dict, current_depth: int = 0): + """ + Calculate the depth of the given node and its children. This recursive + method will update `depths` in place. - :param `node`: The node (str) to start at. - :param `depths`: the dictionary of depths to update. - :param `current_depth`: the current depth in the graph traversal. + Args: + node: The node to start at. + depths: The dictionary of depths to update. + current_depth: The current depth in the graph traversal. """ if node not in depths: depths[node] = current_depth @@ -90,22 +140,30 @@ def calc_depth(self, node, depths, current_depth=0): self.calc_depth(child, depths, current_depth=depths[node] + 1) @staticmethod - def group_by_depth(depths): - """Group DAG tasks by depth. + def group_by_depth(depths: Dict) -> List[List[List]]: + """ + Group Directed Acyclic Graph (DAG) tasks by depth. + + This method only groups by depth, and has one task in every chain. + [`find_independent_chains`][study.dag.DAG.find_independent_chains] is used + to figure out how to coalesce chains across depths. - :param `depths`: the dictionary of depths to group by + Args: + depths: The dictionary of depths to group by. - :return: a list of lists of lists ordered by depth + Returns: + A list of lists of lists ordered by depth. - ([[["tasks"],["with"],["Depth 0"]],[["tasks"],["with"],["Depth 1"]]]) + Example: + This method will return a list that could look something like this: - The outer index of this list is the depth, the middle index is which - chain of tasks in that depth, and the inner index is the task id in - that chain. + ```python + [[["tasks"], ["with"], ["Depth 0"]], [["tasks"], ["with"], ["Depth 1"]]] + ``` - This method only groups by depth, and has one task in every chain. - find_independent_chains is used to figure out how to coalesce chains - across depths. + Here, the outer index of this list is the depth, the middle index is + which chain of tasks in that depth, and the inner index is the task + id in that chain. """ groups = {} for node in depths: @@ -123,44 +181,66 @@ def group_by_depth(depths): return list_of_groups_of_chains - def children(self, task_name): - """Return the children of the task. - :param `task_name`: The name of the task to get the children of. + def children(self, task_name: str) -> List: + """ + Return the children of the task. + + Args: + task_name: The name of the task to get the children of. - :return: list of children of this task. + Returns: + List of children of this task. """ return self.maestro_adjacency_table[task_name] - def num_children(self, task_name): - """Find the number of children for the given task in the dag. - :param `task_name`: The name of the task to count the children of. + def num_children(self, task_name: str) -> int: + """ + Find the number of children for the given task in the Directed Acyclic Graph (DAG). + + Args: + task_name: The name of the task to count the children of. - :return : number of children this task has + Returns: + Number of children this task has. """ return len(self.children(task_name)) - def parents(self, task_name): - """Return the parents of the task. - :param `task_name` : The name of the task to get the parents of. + def parents(self, task_name: str) -> List: + """ + Return the parents of the task. + + Args: + task_name: The name of the task to get the parents of. - :return : list of parents of this task""" + Returns: + List of parents of this task. + """ return self.backwards_adjacency[task_name] - def num_parents(self, task_name): - """find the number of parents for the given task in the dag - :param `task_name` : The name of the task to count the parents of + def num_parents(self, task_name: str) -> int: + """ + Find the number of parents for the given task in the Directed Acyclic Graph (DAG). + + Args: + task_name: The name of the task to count the parents of. - :return : number of parents this task has""" + Returns: + Number of parents this task has. + """ return len(self.parents(task_name)) @staticmethod - def find_chain(task_name, list_of_groups_of_chains): - """find the chain containing the task - :param `task_name` : The task to search for. - :param `list_of_groups_of_chains` : list of groups of chains to search - for the task + def find_chain(task_name: str, list_of_groups_of_chains: List[List[List]]) -> List: + """ + Find the chain containing the task. - :return : the list representing the chain containing task_name""" + Args: + task_name: The task to search for. + list_of_groups_of_chains: List of groups of chains to search for the task. + + Returns: + The list representing the chain containing task_name, or None if not found. + """ for group in list_of_groups_of_chains: for chain in group: if task_name in chain: @@ -168,7 +248,42 @@ def find_chain(task_name, list_of_groups_of_chains): return None def calc_backwards_adjacency(self): - """initializes our backwards adjacency table""" + """ + Initializes the backwards adjacency table. + + This method constructs a mapping of each task to its parent tasks in the Directed + Acyclic Graph (DAG). The backwards adjacency table allows for reverse traversal + of the graph, enabling the identification of dependencies for each task. + + The method iterates through each parent task in the `maestro_adjacency_table` + and updates the `backwards_adjacency` dictionary. For each task that is a child + of a parent, it adds the parent to the list of that task's parents in the + `backwards_adjacency` table. + + This is essential for operations that require knowledge of a task's dependencies, + such as determining the order of execution or identifying independent tasks. + + Example: + If the `maestro_adjacency_table` is structured as follows: + + ```python + { + 'A': ['B', 'C'], + 'B': ['D'], + 'C': ['D'] + } + ``` + + After calling this method, the `backwards_adjacency` will be: + + ```python + { + 'B': ['A'], + 'C': ['A'], + 'D': ['B', 'C'] + } + ``` + """ for parent in self.maestro_adjacency_table: for task_name in self.maestro_adjacency_table[parent]: if task_name in self.backwards_adjacency: @@ -176,34 +291,52 @@ def calc_backwards_adjacency(self): else: self.backwards_adjacency[task_name] = [parent] - def compatible_merlin_expansion(self, task1, task2): + def compatible_merlin_expansion(self, task1: str, task2: str) -> bool: """ - TODO + Check if two tasks are compatible for Merlin expansion. + + This method compares the expansion needs of two tasks to determine + if they can be expanded together. + + Args: + task1: The first task. + task2: The second task. + + Returns: + True if compatible, False otherwise. """ step1 = self.step(task1) step2 = self.step(task2) return step1.check_if_expansion_needed(self.column_labels) == step2.check_if_expansion_needed(self.column_labels) - def find_independent_chains(self, list_of_groups_of_chains): + def find_independent_chains(self, list_of_groups_of_chains: List[List[List]]) -> List[List[List]]: """ - Finds independent chains and adjusts with the groups of chains to - maximalize parallelism + Finds independent chains and adjusts with the groups of chains to maximize parallelism. - :param list_of_groups_of_chains: List of list of lists, as returned by - self.group_by_depth + This method looks for opportunities to move tasks in deeper groups of chains + into chains in shallower groups, thus increasing available parallelism in execution. - e.g., + Args: + list_of_groups_of_chains: List of list of lists, as returned by + [`group_by_depth`][study.dag.DAG.group_by_depth]. - ([[["task1"],["with"],["Depth 0"]],[["task2"],["has"],["Depth 1"]]]) + Returns: + Adjusted list of groups of chains to maximize parallelism. - :return : This takes the groups of chains and looks for opportunities - to move tasks in deeper groups of chains into chains in shallower - groups, thus increasing available parallelism in the execution. + Example: + Given input chains, the method may return a modified structure that allows + for more tasks to be executed in parallel. For example, we might start with + this: - Depending on the precise parental relationships between the tasks - in the graph the output may be something like: + ```python + [[["task1"], ["with"], ["Depth 0"]], [["task2"], ["has"], ["Depth 1"]]] + ``` - ([[["task1", "has"],["with","task2"],["Depth 0"]],["Depth 1"]]]) + and finish with this: + + ```python + [[["task1", "has"], ["with", "task2"], ["Depth 0"]], ["Depth 1"]]] + ``` """ for group in list_of_groups_of_chains: for chain in group: @@ -220,14 +353,18 @@ def find_independent_chains(self, list_of_groups_of_chains): return new_list_2 - def group_tasks(self, source_node): - """Group independent tasks in a directed acyclic graph (DAG). + def group_tasks(self, source_node: str) -> List[List[List]]: + """ + Group independent tasks in a Directed Acyclic Graph (DAG). + + Starts from a source node and works down, grouping tasks by depth, + then identifies independent parallel chains in those groups. - Starts from a source node and works down, grouping tasks by - depth, then identify independent parallel chains in those groups. + Args: + source_node: The source node from which to start grouping tasks. - :param dag : The DAG - :param source_node: The source node. + Returns: + A list of independent chains of tasks. """ depths = {} self.calc_depth(source_node, depths) diff --git a/merlin/study/script_adapter.py b/merlin/study/script_adapter.py index 6380b1a9b..633506ed9 100644 --- a/merlin/study/script_adapter.py +++ b/merlin/study/script_adapter.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -29,34 +29,43 @@ ############################################################################### """ -Merlin script adapter module +This module stores the functionality for adapting bash scripts to use schedulers. + +Supported schedulers are currently: Flux, LSF, and Slurm. """ import logging import os -from typing import Dict, List, Set +from typing import Dict, List, Set, Tuple, Union +from maestrowf.abstracts.enums import StepPriority +from maestrowf.abstracts.interfaces.scriptadapter import ScriptAdapter +from maestrowf.datastructures.core.study import StudyStep from maestrowf.interfaces.script import SubmissionRecord from maestrowf.interfaces.script.localscriptadapter import LocalScriptAdapter from maestrowf.interfaces.script.slurmscriptadapter import SlurmScriptAdapter from maestrowf.utils import start_process -from merlin.common.abstracts.enums import ReturnCode +from merlin.common.enums import ReturnCode from merlin.utils import convert_timestring, find_vlaunch_var LOG = logging.getLogger(__name__) -def setup_vlaunch(step_run: str, batch_type: str, gpu_config: bool) -> None: +def setup_vlaunch(step_run: str, batch_type: str, gpu_config: bool): """ - Check for the VLAUNCHER keyword int the step run string, find - the MERLIN variables and configure VLAUNCHER. + Check for the VLAUNCHER keyword in the step run string and configure VLAUNCHER settings. + + This function examines the provided step run command string for the presence of the + VLAUNCHER keyword. If found, it replaces the keyword with the LAUNCHER keyword and + extracts relevant MERLIN variables such as nodes, processes, and cores per task. + It also configures GPU settings based on the provided boolean flag. - :param `step_run`: the step.run command string - :param `batch_type`: the batch type string - :param `gpu_config`: bool to determin if gpus should be configured - :returns: None + Args: + step_run: The step.run command string that may contain the VLAUNCHER keyword. + batch_type: A string representing the type of batch processing being used. + gpu_config: A boolean indicating whether GPUs should be configured. """ if "$(VLAUNCHER)" in step_run["cmd"]: step_run["cmd"] = step_run["cmd"].replace("$(VLAUNCHER)", "$(LAUNCHER)") @@ -74,20 +83,34 @@ def setup_vlaunch(step_run: str, batch_type: str, gpu_config: bool) -> None: class MerlinLSFScriptAdapter(SlurmScriptAdapter): """ - A SchedulerScriptAdapter class for slurm blocking parallel launches, - the LSFScriptAdapter uses non-blocking submits. + A `SchedulerScriptAdapter` class for SLURM blocking parallel launches. + The `MerlinLSFScriptAdapter` uses non-blocking submits for executing LSF parallel jobs + in a Celery worker. + + Attributes: + key (str): A unique key identifier for the adapter. + _cmd_flags (Dict[str, str]): A dictionary containing command flags for LSF execution. + _unsupported (Set[str]): A set of parameters that are unsupported by this adapter. + + Methods: + get_header: Generates the header for LSF execution scripts. + get_parallelize_command: Generates the LSF parallelization segment of the command line. + get_priority: Overrides the abstract method to fix a pylint error. + write_script: Overwrites the write_script method from the base ScriptAdapter class. """ - key = "merlin-lsf" + key: str = "merlin-lsf" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Dict): """ - Initialize an instance of the MerinLSFScriptAdapter. - The MerlinLSFScriptAdapter is the adapter that is used for workflows that + Initialize an instance of the `MerinLSFScriptAdapter`. + + The `MerlinLSFScriptAdapter` is the adapter that is used for workflows that will execute LSF parallel jobs in a celery worker. The only configurable aspect to this adapter is the shell that scripts are executed in. - :param **kwargs: A dictionary with default settings for the adapter. + Args: + **kwargs: A dictionary with default settings for the adapter. """ super().__init__(**kwargs) @@ -123,27 +146,43 @@ def __init__(self, **kwargs): "walltime", } - def get_priority(self, priority): - """This is implemented to override the abstract method and fix a pylint error""" + def get_priority(self, priority: StepPriority): + """ + This is implemented to override the abstract method and fix a pylint error. + + Args: + priority: Float or + [`StepPriority`](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/abstracts/enums/index.html#maestrowf.abstracts.enums.StepPriority) + enum representing priorty. + """ - def get_header(self, step): + def get_header(self, step: StudyStep) -> str: """ Generate the header present at the top of LSF execution scripts. - :param step: A StudyStep instance. - :returns: A string of the header based on internal batch parameters and - the parameter step. + Args: + step: A Maestro StudyStep instance that contains parameters relevant to the execution. + + Returns: + A string of the header based on internal batch parameters and the parameter step. """ return f"#!{self._exec}" - def get_parallelize_command(self, procs, nodes=None, **kwargs): + def get_parallelize_command(self, procs: int, nodes: int = None, **kwargs: Dict) -> str: """ - Generate the LSF parallelization segement of the command line. - :param procs: Number of processors to allocate to the parallel call. - :param nodes: Number of nodes to allocate to the parallel call - (default = 1). - :returns: A string of the parallelize command configured using nodes - and procs. + Generate the LSF parallelization segment of the command line. + + This method constructs a command line segment for parallel execution in LSF. + It allows specifying the number of processors and nodes to be allocated for the parallel call, + along with additional command flags through keyword arguments. + + Args: + procs: Number of processors to allocate to the parallel call. + nodes: Number of nodes to allocate to the parallel call. Defaults to 1. + **kwargs: Additional command flags that may be supported by the LSF command. + + Returns: + A string representing the parallelization command configured using nodes and procs. """ if not nodes: nodes = 1 @@ -180,18 +219,21 @@ def get_parallelize_command(self, procs, nodes=None, **kwargs): return " ".join(args) - def write_script(self, ws_path, step): + def write_script(self, ws_path: str, step: StudyStep) -> Tuple[bool, str, str]: """ - This will overwrite the write_script in method from Maestro's base ScriptAdapter + This will overwrite the `write_script` method from Maestro's base ScriptAdapter class but will eventually call it. This is necessary for the VLAUNCHER to work. - :param `ws_path`: the path to the workspace where we'll write the scripts - :param `step`: the Maestro StudyStep object containing info for our step - :returns: a tuple containing: - - a boolean representing whether this step is to be scheduled or not - - Merlin can ignore this - - a path to the script for the cmd - - a path to the script for the restart cmd + Args: + ws_path: The path to the workspace where the scripts will be written. + step: The Maestro StudyStep object containing information for the step. + + Returns: + A tuple containing:\n + - bool: A boolean indicating whether this step is to be scheduled or not. + (Merlin can ignore this value.) + - str: The path to the script for the command. + - str: The path to the script for the restart command. """ setup_vlaunch(step.run, "lsf", False) @@ -200,23 +242,42 @@ class but will eventually call it. This is necessary for the VLAUNCHER to work. class MerlinSlurmScriptAdapter(SlurmScriptAdapter): """ - A SchedulerScriptAdapter class for slurm blocking parallel launches, - the SlurmScriptAdapter uses non-blocking submits. + A `SchedulerScriptAdapter` class for SLURM blocking parallel launches. + + This class extends the `SlurmScriptAdapter` to provide support for blocking parallel + launches in SLURM. Unlike the base class, which uses non-blocking submits, this adapter + is designed for workflows that execute SLURM parallel jobs in a Celery worker. + + Attributes: + key (str): A unique identifier for the adapter, set to "merlin-slurm". + _cmd_flags (Dict[str, str]): A dictionary containing command flags for SLURM. + _unsupported (Set[str]): A set of command flags that are not supported by this adapter. + + Methods: + get_header: Generates the header for SLURM execution scripts. + get_parallelize_command: Generates the SLURM parallelization segment of the command line. + get_priority: Overrides the abstract method to fix a pylint error. + time_format: Converts a timestring to HH:MM:SS format. + write_script: Overwrites the write_script method from the base class to ensure VLAUNCHER compatibility. """ key: str = "merlin-slurm" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Dict): """ - Initialize an instance of the MerinSlurmScriptAdapter. - The MerlinSlurmScriptAdapter is the adapter that is used for workflows that + Initialize an instance of the `MerinSlurmScriptAdapter`. + + The `MerlinSlurmScriptAdapter` is the adapter that is used for workflows that will execute SLURM parallel jobs in a celery worker. The only configurable aspect to this adapter is the shell that scripts are executed in. - :param **kwargs: A dictionary with default settings for the adapter. + Args: + **kwargs: A dictionary with default settings for the adapter. """ super().__init__(**kwargs) + self._cmd_flags: Dict[str, str] + self._cmd_flags["slurm"] = "" self._cmd_flags["walltime"] = "-t" @@ -236,33 +297,63 @@ def __init__(self, **kwargs): ] self._unsupported: Set[str] = set(list(self._unsupported) + new_unsupported) - def get_priority(self, priority): - """This is implemented to override the abstract method and fix a pylint error""" + def get_priority(self, priority: StepPriority): + """ + This is implemented to override the abstract method and fix a pylint error. + + Args: + priority: Float or + [`StepPriority`](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/abstracts/enums/index.html#maestrowf.abstracts.enums.StepPriority) + enum representing priorty. + """ - def get_header(self, step): + def get_header(self, step: StudyStep) -> str: """ Generate the header present at the top of Slurm execution scripts. - :param step: A StudyStep instance. - :returns: A string of the header based on internal batch parameters and - the parameter step. + Args: + step: A Maestro StudyStep instance that contains parameters relevant to the execution. + + Returns: + A string of the header based on internal batch parameters and the parameter step. """ return f"#!{self._exec}" - def time_format(self, val): + def time_format(self, val: Union[str, int]) -> str: """ - Convert the timestring to HH:MM:SS + Convert the input timestring or integer to HH:MM:SS format. + + This method utilizes the [`convert_timestring`][utils.convert_timestring] + function to convert a given timestring or integer (representing seconds) + into a formatted string in the 'hours:minutes:seconds' (HH:MM:SS) format. + + Args: + val: A timestring in the format '[days]:[hours]:[minutes]:seconds' or + an integer representing time in seconds. + + Returns: + A string representation of the input time formatted as 'HH:MM:SS'. """ return convert_timestring(val, format_method="HMS") - def get_parallelize_command(self, procs, nodes=None, **kwargs): + def get_parallelize_command(self, procs: int, nodes: int = None, **kwargs: Dict) -> str: """ - Generate the SLURM parallelization segement of the command line. - :param procs: Number of processors to allocate to the parallel call. - :param nodes: Number of nodes to allocate to the parallel call - (default = 1). - :returns: A string of the parallelize command configured using nodes - and procs. + Generate the SLURM parallelization segment of the command line. + + This method constructs the command line segment required for parallel execution + in SLURM, including the number of processors and nodes to allocate. It also + incorporates any additional supported command flags provided in `kwargs`. + + Args: + procs: The number of processors to allocate for the parallel call. + nodes: The number of nodes to allocate for the parallel call (default is 1). + **kwargs: Additional command flags to customize the SLURM command. + Supported flags include 'walltime' and others defined in the + `_cmd_flags` attribute, excluding those in the `_unsupported` set. + + Returns: + A string representing the SLURM parallelization command, formatted with the + specified number of processors, nodes, and any additional flags. """ args = [ # SLURM srun command @@ -297,18 +388,21 @@ def get_parallelize_command(self, procs, nodes=None, **kwargs): return " ".join(args) - def write_script(self, ws_path, step): + def write_script(self, ws_path: str, step: StudyStep) -> Tuple[bool, str, str]: """ - This will overwrite the write_script in method from Maestro's base ScriptAdapter + This will overwrite the `write_script` method from Maestro's base ScriptAdapter class but will eventually call it. This is necessary for the VLAUNCHER to work. - :param `ws_path`: the path to the workspace where we'll write the scripts - :param `step`: the Maestro StudyStep object containing info for our step - :returns: a tuple containing: - - a boolean representing whether this step is to be scheduled or not - - Merlin can ignore this - - a path to the script for the cmd - - a path to the script for the restart cmd + Args: + ws_path: The path to the workspace where the scripts will be written. + step: The Maestro `StudyStep` object containing information for the step. + + Returns: + A tuple containing:\n + - bool: A boolean indicating whether this step is to be scheduled or not. + (Merlin can ignore this value.) + - str: The path to the script for the command. + - str: The path to the script for the restart command. """ setup_vlaunch(step.run, "slurm", False) @@ -317,26 +411,41 @@ class but will eventually call it. This is necessary for the VLAUNCHER to work. class MerlinFluxScriptAdapter(MerlinSlurmScriptAdapter): """ - A SchedulerScriptAdapter class for flux blocking parallel launches, - the FluxScriptAdapter uses non-blocking submits. + A `SchedulerScriptAdapter` class for flux blocking parallel launches. + + The `MerlinFluxScriptAdapter` is designed for workflows that execute flux parallel jobs + in a Celery worker. It utilizes non-blocking submits and allows for configuration of the + shell in which scripts are executed. + + Attributes: + key (str): A unique identifier for the adapter, set to "merlin-flux". + _cmd_flags (Dict[str, str]): A dictionary containing command-line flags for the flux command. + _unsupported (Set[str]): A set of command flags that are not supported by this adapter. + + Methods: + get_priority: Retrieves the priority of the step. + time_format: Converts a time format to flux standard designation. + write_script: Writes the script for the specified step and returns relevant paths. """ - key = "merlin-flux" + key: str = "merlin-flux" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Dict): """ - Initialize an instance of the MerinFluxScriptAdapter. - The MerlinFluxScriptAdapter is the adapter that is used for workflows that + Initialize an instance of the `MerinFluxScriptAdapter`. + + The `MerlinFluxScriptAdapter` is the adapter that is used for workflows that will execute flux parallel jobs in a celery worker. The only configurable aspect to this adapter is the shell that scripts are executed in. - :param **kwargs: A dictionary with default settings for the adapter. + Args: + **kwargs: A dictionary with default settings for the adapter. """ # The flux_command should always be overriden by the study object's flux_command property flux_command = kwargs.pop("flux_command", "flux run") super().__init__(**kwargs) - self._cmd_flags = { + self._cmd_flags: Dict[str, str] = { "cmd": flux_command, "ntasks": "-n", "nodes": "-N", @@ -366,29 +475,51 @@ def __init__(self, **kwargs): "lsf", "slurm", ] - self._unsupported = set(new_unsupported) # noqa + self._unsupported: Set[str] = set(new_unsupported) # noqa + + def get_priority(self, priority: StepPriority): + """ + This is implemented to override the abstract method and fix a pylint error. - def get_priority(self, priority): - """This is implemented to override the abstract method and fix a pylint error""" + Args: + priority: Float or + [`StepPriority`](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/abstracts/enums/index.html#maestrowf.abstracts.enums.StepPriority) + enum representing priorty. + """ - def time_format(self, val): + def time_format(self, val: Union[str, int]) -> str: """ - Convert a time format to flux standard designation. + Convert a time format to Flux Standard Duration (FSD). + + This method takes a time value and converts it into a format that is compatible + with Flux's standard time representation. The conversion is performed using the + [`convert_timestring`][utils.convert_timestring] function with the specified format + method. + + Args: + val: The time value to be converted. This can be a string representing a time + duration or an integer representing a time value. + + Returns: + The time formatted according to Flux Standard Duration (FSD). """ return convert_timestring(val, format_method="FSD") - def write_script(self, ws_path, step): + def write_script(self, ws_path: str, step: StudyStep) -> Tuple[bool, str, str]: """ - This will overwrite the write_script in method from Maestro's base ScriptAdapter + This will overwrite the `write_script` method from Maestro's base ScriptAdapter class but will eventually call it. This is necessary for the VLAUNCHER to work. - :param `ws_path`: the path to the workspace where we'll write the scripts - :param `step`: the Maestro StudyStep object containing info for our step - :returns: a tuple containing: - - a boolean representing whether this step is to be scheduled or not - - Merlin can ignore this - - a path to the script for the cmd - - a path to the script for the restart cmd + Args: + ws_path: The path to the workspace where the scripts will be written. + step: The Maestro `StudyStep` object containing information for the step. + + Returns: + A tuple containing:\n + - bool: A boolean indicating whether this step is to be scheduled or not. + (Merlin can ignore this value.) + - str: The path to the script for the command. + - str: The path to the script for the restart command. """ setup_vlaunch(step.run, "flux", True) @@ -397,23 +528,40 @@ class but will eventually call it. This is necessary for the VLAUNCHER to work. class MerlinScriptAdapter(LocalScriptAdapter): """ - A ScriptAdapter class for interfacing for execution in Merlin + A `ScriptAdapter` class for interfacing with execution in Merlin. + + This class serves as an adapter for executing scripts in a Celery worker + environment. It allows for configuration of the execution environment and + manages the execution of scripts with appropriate logging and error handling. + + Attributes: + batch_adapter (ScriptAdapter): An instance of a batch adapter used for executing scripts + based on the specified batch type. + batch_type (str): The type of batch processing to be used, derived from + the provided keyword arguments. + key (str): A unique identifier for the adapter, set to "merlin-local". + + Methods: + submit: Executes a workflow step locally. + write_script: Writes a script using the batch adapter. """ - key = "merlin-local" + key: str = "merlin-local" - def __init__(self, **kwargs): + def __init__(self, **kwargs: Dict): """ - Initialize an instance of the MerinScriptAdapter. - The MerlinScriptAdapter is the adapter that is used for workflows that + Initialize an instance of the `MerinScriptAdapter`. + + The `MerlinScriptAdapter` is the adapter that is used for workflows that will execute in a celery worker. The only configurable aspect to this adapter is the shell that scripts are executed in. - :param **kwargs: A dictionary with default settings for the adapter. + Args: + **kwargs: A dictionary with default settings for the adapter. """ super().__init__(**kwargs) - self.batch_type = "merlin-" + kwargs.get("batch_type", "local") + self.batch_type: str = "merlin-" + kwargs.get("batch_type", "local") if "host" not in kwargs: kwargs["host"] = "None" @@ -423,33 +571,51 @@ def __init__(self, **kwargs): kwargs["queue"] = "None" # Using super prevents recursion. - self.batch_adapter = super() + self.batch_adapter: ScriptAdapter = super() if self.batch_type != "merlin-local": self.batch_adapter = MerlinScriptAdapterFactory.get_adapter(self.batch_type)(**kwargs) - def write_script(self, *args, **kwargs): + def write_script(self, *args, **kwargs) -> Tuple[bool, str, str]: """ - TODO + Generate a script for execution using the batch adapter. + + This method delegates the script writing process to the associated + batch adapter and returns the generated script along with a restart + script if applicable. + + Returns: + A tuple containing:\n + - bool: A boolean indicating whether this step is to be scheduled or not. + (Merlin can ignore this value.) + - str: The path to the script for the command. + - str: The path to the script for the restart command. """ _, script, restart_script = self.batch_adapter.write_script(*args, **kwargs) return True, script, restart_script # Pylint complains that there's too many arguments but it's fine in this case - def submit(self, step, path, cwd, job_map=None, env=None): # pylint: disable=R0913 - """ - Execute the step locally. - If cwd is specified, the submit method will operate outside of the path - specified by the 'cwd' parameter. - If env is specified, the submit method will set the environment - variables for submission to the specified values. The 'env' parameter - should be a dictionary of environment variables. - - :param step: An instance of a StudyStep. - :param path: Path to the script to be executed. - :param cwd: Path to the current working directory. - :param job_map: A map of workflow step names to their job identifiers. - :param env: A dict containing a modified environment for execution. - :returns: The return code of the command and processID of the command. + def submit( + self, step: StudyStep, path: str, cwd: str, job_map: Dict = None, env: Dict = None + ) -> SubmissionRecord: # pylint: disable=R0913 + """ + Execute a workflow step locally. + + This method runs a specified script in the local environment, allowing for + customization of the working directory and environment variables. It handles + the execution of the script and logs the results, including any errors or + specific return codes. + + Args: + step: An instance of the StudyStep that contains information about the + workflow step being executed. + path: The file path to the script that is to be executed. + cwd: The current working directory from which the script will be executed. + job_map: A mapping of workflow step names to their job identifiers. + env: A dictionary containing environment variables to be set for the execution. + + Returns: + An object containing the return code of the command, the process ID of the + command, and any additional information about the execution. """ LOG.debug("cwd = %s", cwd) LOG.debug("Script to execute: %s", path) @@ -488,21 +654,28 @@ def submit(self, step, path, cwd, job_map=None, env=None): # pylint: disable=R0 # TODO is there currently ever a scenario where join output is True? We should look into this # Pylint is complaining there's too many local variables and args but it makes this function cleaner so ignore - def _execute_subprocess(self, output_name, script_path, cwd, env=None, join_output=False): # pylint: disable=R0913,R0914 - """ - Execute the subprocess script locally. - If cwd is specified, the submit method will operate outside of the path - specified by the 'cwd' parameter. - If env is specified, the submit method will set the environment - variables for submission to the specified values. The 'env' parameter - should be a dictionary of environment variables. - - :param output_name: Output name for stdout and stderr (output_name.out). If None, don't write. - :param script_path: Path to the script to be executed. - :param cwd: Path to the current working directory. - :param env: A dict containing a modified environment for execution. - :param join_output: If True, append stderr to stdout - :returns: The return code of the submission command and job identifier (SubmissionRecord). + def _execute_subprocess( + self, output_name: str, script_path: str, cwd: str, env: Dict = None, join_output: bool = False + ) -> SubmissionRecord: # pylint: disable=R0913,R0914 + """ + Execute a subprocess script locally and manage output. + + This method runs a specified script in a subprocess, capturing its output + and error streams. It allows for customization of the working directory, + environment variables, and output handling. The output can be saved to + files, and error messages can be appended to the standard output if desired. + + Args: + output_name: The base name for the output files (stdout and stderr). + If None, no output files will be created. + script_path: The file path to the script that is to be executed. + cwd: The current working directory from which the script will be executed. + env: A dictionary containing environment variables to be set for the execution. + join_output: If True, appends stderr to stdout in the output file. + + Returns: + An object containing the return code of the command the process ID of the + command, and any additional information about the execution. """ script_bn = os.path.basename(script_path) new_output_name = os.path.splitext(script_bn)[0] @@ -537,9 +710,23 @@ def _execute_subprocess(self, output_name, script_path, cwd, env=None, join_outp class MerlinScriptAdapterFactory: - """This class routes to the correct ScriptAdapter""" + """ + This class routes to the correct `ScriptAdapter`. + + The `MerlinScriptAdapterFactory` is responsible for providing the appropriate + `ScriptAdapter` based on the specified adapter ID. It maintains a mapping of + available adapters and offers methods to retrieve them. + + Attributes: + factories: A dictionary mapping adapter IDs (str) to their corresponding + `ScriptAdapter` classes. + + Methods: + get_adapter: Returns the appropriate `ScriptAdapter` class for the given adapter ID. + get_valid_adapters: Returns a list of valid adapter IDs that can be used with this factory. + """ - factories = { + factories: Dict[str, ScriptAdapter] = { "merlin-flux": MerlinFluxScriptAdapter, "merlin-lsf": MerlinLSFScriptAdapter, "merlin-lsf-srun": MerlinSlurmScriptAdapter, @@ -548,8 +735,23 @@ class MerlinScriptAdapterFactory: } @classmethod - def get_adapter(cls, adapter_id): - """Returns the appropriate ScriptAdapter to use""" + def get_adapter(cls, adapter_id: str) -> ScriptAdapter: + """ + Returns the appropriate `ScriptAdapter` to use. + + This method retrieves the `ScriptAdapter` class associated with the given + adapter ID. If the adapter ID is not found in the factory's mapping, + a ValueError is raised. + + Args: + adapter_id: The ID of the desired `ScriptAdapter`. + + Returns: + The corresponding `ScriptAdapter` class. + + Raises: + ValueError: If the specified adapter_id is not found in the factories. + """ if adapter_id.lower() not in cls.factories: msg = f"""Adapter '{str(adapter_id)}' not found. Specify an adapter that exists or implement a new one mapping to the '{str(adapter_id)}'""" @@ -559,6 +761,15 @@ def get_adapter(cls, adapter_id): return cls.factories[adapter_id] @classmethod - def get_valid_adapters(cls): - """Returns the valid ScriptAdapters""" + def get_valid_adapters(cls) -> List[str]: + """ + Returns the valid ScriptAdapters. + + This method provides a list of all valid adapter IDs that can be used + with this factory. The IDs are derived from the keys of the factories + dictionary. + + Returns: + A list of valid adapter IDs. + """ return cls.factories.keys() diff --git a/merlin/study/status.py b/merlin/study/status.py index fbeb4d46d..2649d57da 100644 --- a/merlin/study/status.py +++ b/merlin/study/status.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1 +# This file is part of Merlin, Version: 1.12.2 # # For details, see https://github.com/LLNL/merlin. # @@ -27,7 +27,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. ############################################################################### -"""This module handles all the functionality of getting the statuses of studies""" +"""This module handles all the functionality of getting the statuses of studies.""" import json import logging import os @@ -37,7 +37,7 @@ from datetime import datetime from glob import glob from traceback import print_exception -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Dict, List, Set, Tuple, Union import numpy as np from filelock import FileLock, Timeout @@ -47,6 +47,7 @@ from merlin.common.dumper import dump_handler from merlin.display import ANSI_COLORS, display_status_summary, display_status_task_by_task from merlin.spec.expansion import get_spec_with_expansion +from merlin.spec.specification import MerlinSpec from merlin.study.status_constants import ( ALL_VALID_FILTERS, CELERY_KEYS, @@ -72,15 +73,58 @@ class Status: """ - This class handles everything to do with status besides displaying it. - Display functionality is handled in display.py. + Handles the management and retrieval of status information for studies. + + This class is responsible for loading specifications, tracking the status of steps, + calculating runtime statistics, and formatting status information for output in + various formats (JSON, CSV). It interacts with the file system to read status files + and provides methods to display and dump status information. + + Attributes: + args (Namespace): Command-line arguments provided by the user. + full_step_name_map (Dict[str, Set[str]]): A mapping of overall step names to full step names. + num_requested_statuses (int): Counts the number of task statuses in the `requested_statuses` + dictionary. + requested_statuses (Dict): A dictionary storing the statuses that the user wants to view. + run_time_info (Dict[str, Dict]): A dictionary storing runtime statistics for each step. + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object loaded from the workspace or spec file. + step_tracker (Dict[str, List[str]]): A dictionary tracking started and unstarted steps. + tasks_per_step (Dict[str, int]): A mapping of tasks per step for accurate totals. + workspace (str): The path to the workspace containing study data. + + Methods: + display: Displays a high-level summary of the status. + dump: Dumps the status information to a specified file. + format_csv_dump: Prepares the dictionary of statuses for CSV output. + format_json_dump: Prepares the dictionary of statuses for JSON output. + format_status_for_csv: Reformats statuses into a dictionary suitable for CSV output. + get_runtime_avg_std_dev: Calculates and stores the average and standard deviation of + runtimes for a step. + get_step_statuses: Reads and returns the statuses for a given step. + get_steps_to_display: Generates a list of steps to display the status for. + load_requested_statuses: Populates the `requested_statuses` dictionary with statuses + from the study. """ def __init__(self, args: Namespace, spec_display: bool, file_or_ws: str): + """ + Initializes the `Status` object, which manages and retrieves status information for studies. + + Args: + args: Command-line arguments provided by the user, including filters and options + for displaying or dumping status information. + spec_display: A flag indicating whether the status should be loaded from a specification + file (`True`) or from a workspace (`False`). + file_or_ws: The path to the specification file or workspace, depending on the value of + `spec_display`. + """ # Save the args to this class instance and check if the steps filter was given - self.args = args + self.args: Namespace = args # Load in the workspace path and spec object + self.workspace: str + self.spec: MerlinSpec if spec_display: self.workspace, self.spec = self._load_from_spec(file_or_ws) else: @@ -91,26 +135,31 @@ def __init__(self, args: Namespace, spec_display: bool, file_or_ws: str): self._verify_filter_args() # Create a step tracker that will tell us which steps have started/not started - self.step_tracker = self.get_steps_to_display() + self.step_tracker: Dict[str, List[str]] = self.get_steps_to_display() # Create a tasks per step mapping in order to give accurate totals for each step - self.tasks_per_step = self.spec.get_tasks_per_step() + self.tasks_per_step: Dict[str, int] = self.spec.get_tasks_per_step() # This attribute will store a map between the overall step name and the full step names # that are created with parameters (e.g. step name is hello and uses a "GREET: hello" parameter # so the real step name is hello_GREET.hello) - self.full_step_name_map = {} + self.full_step_name_map: Dict[str, Set[str]] = {} # Variable to store run time information for each step - self.run_time_info = {} + self.run_time_info: Dict[str, Dict] = {} # Variable to store the statuses that the user wants - self.requested_statuses = {} + self.requested_statuses: Dict = {} self.load_requested_statuses() def _print_requested_statuses(self): """ - Helper method to print out the requested statuses dict. + Print the requested statuses stored in the `requested_statuses` dictionary. + + This helper method iterates through the `requested_statuses` attribute, which contains + information about the statuses of various steps. It prints the step names along with + their corresponding status information. Non-workspace keys are printed directly, while + workspace-related keys are further detailed by their status keys and values. """ print("self.requested_statuses:") for step_name, overall_step_info in self.requested_statuses.items(): @@ -125,16 +174,38 @@ def _print_requested_statuses(self): def _verify_filter_args(self): """ - This is an abstract method since we'll need to verify filter args for DetailedStatus - but not for Status. + Verify the filter arguments for the status retrieval. + + This is an abstract method intended to be implemented in subclasses, such as + [`DetailedStatus`][study.status.DetailedStatus]. The method will ensure that + the filter arguments provided for retrieving statuses are valid and meet the + necessary criteria. The implementation details will depend on the specific + requirements of the subclass. """ def _get_latest_study(self, studies: List[str]) -> str: """ - Given a list of studies, get the latest one. + Retrieve the latest study from a list of studies. + + This method examines a list of study identifiers and determines which one is the latest + based on the timestamp embedded in the study names. It assumes that the newest study is + represented by the last entry in the list but verifies this assumption by comparing the + timestamps of all studies. + + The method extracts the timestamp from the last 15 characters of each study identifier, + converts it to a datetime object, and compares it to find the most recent study. + + Args: + studies: A list of study identifiers to evaluate. - :param `studies`: A list of studies to sort through - :returns: The latest study in the list provided + Returns: + The identifier of the latest study. + + Example: + ```python + >>> self._get_latest_study(["study_20231101-174102", "study_20231101-182044", "study_20231101-163327"]) + 'study_20231101-182044' + ``` """ # We can assume the newest study is the last one to be added to the list of potential studies newest_study = studies[-1] @@ -155,12 +226,22 @@ def _obtain_study(self, study_output_dir: str, num_studies: int, potential_studi """ Grab the study that the user wants to view the status of based on a list of potential studies provided. - :param `study_output_dir`: A string representing the output path of a study; equivalent to $(OUTPUT_PATH) - :param `num_studies`: The number of potential studies we found - :param `potential_studies`: The list of potential studies we found; - Each entry is of the form (index, potential_study_name) - :returns: A directory path to the study that the user wants - to view the status of ("study_output_dir/selected_potential_study") + This method checks the number of potential studies found and either selects the latest study + automatically or prompts the user to choose from the available options. It constructs the + directory path to the selected study. + + Args: + study_output_dir: A string representing the output path of a study; equivalent to $(OUTPUT_PATH). + num_studies: The number of potential studies found. + potential_studies: A list of potential studies found, where each entry is of the form (index, + potential_study_name). + + Returns: + A directory path to the study that the user wants to view the status of, formatted as + "study_output_dir/selected_potential_study". + + Raises: + ValueError: If no potential studies are found or if the user input is invalid. """ study_to_check = f"{study_output_dir}/" if num_studies == 0: @@ -197,14 +278,25 @@ def _obtain_study(self, study_output_dir: str, num_studies: int, potential_studi return study_to_check - def _load_from_spec(self, filepath: str) -> Tuple[str, "MerlinSpec"]: # noqa: F821 pylint: disable=R0914 + def _load_from_spec(self, filepath: str) -> Tuple[str, MerlinSpec]: # pylint: disable=R0914 """ - Get the desired workspace from the user and load up it's yaml spec - for further processing. + Get the desired workspace from the user and load its YAML spec for further processing. + + This method verifies the output path based on user input or the spec file and builds a list + of potential study output directories. It then calls another method to obtain the study to + check the status for and loads the corresponding spec. + + Args: + filepath: The filepath to a spec provided by the user. - :param `filepath`: The filepath to a spec given by the user - :returns: The workspace of the study we'll check the status for and a MerlinSpec - object loaded in from the workspace's merlin_info subdirectory. + Returns: + A tuple containing the workspace of the study to check the status for and a + [`MerlinSpec`][spec.specification.MerlinSpec] object loaded from the workspace's + merlin_info subdirectory. + + Raises: + ValueError: If the specified output directory does not contain a merlin_info subdirectory, + or if multiple or no expanded spec options are found in the directory. """ # If the user provided a new output path to look in, use that if self.args.output_path is not None: @@ -263,11 +355,17 @@ def _load_from_spec(self, filepath: str) -> Tuple[str, "MerlinSpec"]: # noqa: F return study_to_check, actual_spec - def _load_from_workspace(self) -> "MerlinSpec": # noqa: F821 + def _load_from_workspace(self) -> MerlinSpec: """ - Create a MerlinSpec object based on the spec file in the workspace. + Create a [`MerlinSpec`][spec.specification.MerlinSpec] object based on the expanded spec file + in the workspace. + + Returns: + spec.specification.MerlinSpec: A [`MerlinSpec`][spec.specification.MerlinSpec] object loaded + from the workspace provided by the user. - :returns: A MerlinSpec object loaded from the workspace provided by the user + Raises: + ValueError: If multiple or no expanded spec options are found in the workspace's merlin_info directory. """ # Grab the spec file from the directory provided expanded_spec_options = glob(f"{self.workspace}/merlin_info/*.expanded.yaml") @@ -285,11 +383,19 @@ def _load_from_workspace(self) -> "MerlinSpec": # noqa: F821 def _create_step_tracker(self, steps_to_check: List[str]) -> Dict[str, List[str]]: """ - Creates a dictionary of started and unstarted steps that we - will display the status for. + Creates a dictionary of started and unstarted steps to display their status. + + This method checks the workspace for steps that have been started and compares them + against a provided list of steps to determine which steps are started and which are + unstarted. It returns a dictionary categorizing the steps accordingly. + + Args: + steps_to_check: A list of step names to check the status of. - :param `steps_to_check`: A list of steps to view the status of - :returns: A dictionary mapping of started and unstarted steps. Values are lists of step names. + Returns: + A dictionary with two keys:\n + - "started_steps": A list of steps that have been started. + - "unstarted_steps": A list of steps that have not been started. """ step_tracker = {"started_steps": [], "unstarted_steps": []} started_steps = next(os.walk(self.workspace))[1] @@ -310,10 +416,22 @@ def _create_step_tracker(self, steps_to_check: List[str]) -> Dict[str, List[str] def get_steps_to_display(self) -> Dict[str, List[str]]: """ - Generates a list of steps to display the status for based on information - provided to the merlin status command by the user. - - :returns: A dictionary of started and unstarted steps for us to display the status of + Generates a dictionary of steps to display their status based on user input + provided to the merlin status command. + + This method retrieves the names of existing steps from the study specification + and creates a step tracker to categorize them into started and unstarted steps. + + Returns: + A dictionary with two keys:\n + - `started_steps`: A list of steps that have been started. + - `unstarted_steps`: A list of steps that have not been started. + + Example: + ```python + >>> self.get_steps_to_display() + {"started_steps": ["step1"], "unstarted_steps": ["step2", "step3"]} + ``` """ existing_steps = self.spec.get_study_step_names() @@ -328,10 +446,13 @@ def get_steps_to_display(self) -> Dict[str, List[str]]: return step_tracker @property - def num_requested_statuses(self): + def num_requested_statuses(self) -> int: """ - Count the number of task statuses in a the requested_statuses dict. - We need to ignore non workspace keys when we count. + Counts the number of task statuses in the requested_statuses dictionary, + excluding non-workspace keys. + + Returns: + The count of requested task statuses that are not non-workspace keys. """ num_statuses = 0 for overall_step_info in self.requested_statuses.values(): @@ -341,12 +462,20 @@ def num_requested_statuses(self): def get_step_statuses(self, step_workspace: str, started_step_name: str) -> Dict[str, List[str]]: """ - Given a step workspace and the name of the step, read in all the statuses - for the step and return them in a dict. + Reads the statuses for a specified step from the given step workspace. + + This method traverses the specified step workspace directory to locate + `MERLIN_STATUS.json` files, reads their contents, and aggregates the statuses + into a dictionary. It also tracks the full names of the steps and counts + the number of statuses read. - :param step_workspace: The path to the step we're going to read statuses from - :param started_step_name: The name of the step that we're gathering statuses for - :returns: A dict of statuses for the given step + Args: + step_workspace: The path to the step directory from which to read statuses. + started_step_name: The name of the step for which statuses are being gathered. + + Returns: + A dictionary containing the statuses for the specified step, where each key is a full + step name and the value is a list of status information. """ step_statuses = {} num_statuses_read = 0 @@ -390,7 +519,13 @@ def get_step_statuses(self, step_workspace: str, started_step_name: str) -> Dict def load_requested_statuses(self): """ - Populate the requested_statuses dict with the statuses from the study. + Populates the `requested_statuses` dictionary with statuses from the study. + + This method iterates through the started steps in the step tracker, + retrieves their statuses using the + [`get_step_statuses`][study.status.Status.get_step_statuses] method, and merges + these statuses into the `requested_statuses` dictionary. It also calculates + the average and standard deviation of the run times for each step. """ LOG.info(f"Reading task statuses from {self.workspace}") @@ -406,14 +541,20 @@ def load_requested_statuses(self): # Count how many statuses in total that we just read in LOG.info(f"Read in {self.num_requested_statuses} statuses total.") - def get_runtime_avg_std_dev(self, step_statuses: Dict, step_name: str) -> Dict: + def get_runtime_avg_std_dev(self, step_statuses: Dict, step_name: str): """ - Calculate the mean and standard deviation for the runtime of each step. - Add this to the state information once calculated. - - :param `step_statuses`: A dict of step status information that we'll parse for run times - :param `step_name`: The name of the step - :returns: An updated dict of step status info with run time avg and std dev + Calculates the average and standard deviation of the runtime for a specified step. + + This method parses the provided step status information to extract runtime values, + computes the mean and standard deviation of these runtimes, and updates the state + information with the calculated values. The runtimes are expected to be in a specific + format (e.g., "1h30m15s") and are converted to seconds for the calculations. + + Args: + step_statuses: A dictionary containing step status information, where each + entry includes runtime data to be parsed. + step_name: The name of the step for which the average and standard deviation + of the runtime are being calculated. """ # Initialize a list to track all existing runtimes run_times_in_seconds = [] @@ -454,32 +595,53 @@ def get_runtime_avg_std_dev(self, step_statuses: Dict, step_name: str) -> Dict: self.run_time_info[step_name]["run_time_std_dev"] = f"±{pretty_format_hms(convert_timestring(run_time_std_dev))}" LOG.debug(f"Run time avg and std dev for step '{step_name}' calculated.") - def display(self, test_mode: Optional[bool] = False) -> Dict: + def display(self, test_mode: bool = False) -> Dict: """ - Displays the high level summary of the status. + Displays a high-level summary of the status. - :param `test_mode`: If true, run this in testing mode and don't print any output - :returns: A dict that will be empty if test_mode is False. Otherwise, the dict will - contain the status info that would be displayed. + This method provides an overview of the current status of the workflow. If + `test_mode` is enabled, it will not print any output but will return the + status information in a dictionary. + + Args: + test_mode: If true, run this in testing mode and don't print any output. + + Returns: + An empty dictionary if `test_mode` is False; otherwise, a dictionary containing + the status information that would be displayed. """ return display_status_summary(self, NON_WORKSPACE_KEYS, test_mode=test_mode) def format_json_dump(self, date: datetime) -> Dict: """ - Build the dict of statuses to dump to the json file. + Builds a dictionary of statuses to dump to a JSON file. + + This method prepares the status information for serialization by adding a timestamp + to the existing status data. - :param `date`: A timestamp for us to mark when this status occurred - :returns: A dictionary that's ready to dump to a json outfile + Args: + date: A timestamp marking when this status occurred. + + Returns: + A dictionary ready to be dumped to a JSON file, containing the timestamp + and the requested statuses. """ # Statuses are already in json format so we'll just add a timestamp for the dump here return {date: self.requested_statuses} def format_csv_dump(self, date: datetime) -> Dict: """ - Add the timestamp to the statuses to write. + Adds a timestamp to the statuses for CSV output. + + This method reformats the status information into a structure suitable for CSV + output, including a timestamp entry as the first column. - :param `date`: A timestamp for us to mark when this status occurred - :returns: A dict equivalent of formatted statuses with a timestamp entry at the start of the dict. + Args: + date: A timestamp marking when this status occurred. + + Returns: + A dictionary equivalent of formatted statuses with a timestamp entry + at the start of the dictionary. """ # Reformat the statuses to a new dict where the keys are the column labels and rows are the values LOG.debug("Formatting statuses for csv dump...") @@ -494,7 +656,11 @@ def format_csv_dump(self, date: datetime) -> Dict: def dump(self): """ - Dump the status information to a file. + Dumps the status information to a file. + + This method handles the creation of a timestamp and determines the appropriate + file format (CSV or JSON) for dumping the status information. It then calls + the appropriate formatting method and writes the data to the specified file. """ # Get a timestamp for this dump date = datetime.now().strftime("%Y-%m-%d %H:%M:%S") @@ -512,10 +678,16 @@ def dump(self): def format_status_for_csv(self) -> Dict: """ - Reformat our statuses to csv format so they can use Maestro's status renderer layouts. + Reformats statuses for CSV output to comply with + [Maestro's status renderer layouts](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/index.html). - :returns: A formatted dictionary where each key is a column and the values are the rows - of information to display for that column. + This method transforms the status information into a dictionary format where each + key represents a column label and the corresponding values are the rows of information + to display for that column. + + Returns: + A formatted dictionary where each key is a column and the values are the + rows of information to display for that column. """ reformatted_statuses = { "step_name": [], @@ -591,10 +763,42 @@ def format_status_for_csv(self) -> Dict: class DetailedStatus(Status): """ This class handles obtaining and filtering requested statuses from the user. - This class shares similar methodology to the Status class it inherits from. + It inherits from the [`Status`][study.status.Status] class and provides + additional functionality for filtering and displaying task statuses based on + user-defined criteria. + + Attributes: + args (Namespace): A namespace containing user-defined arguments for filtering. + num_requested_statuses (int): The number of task statuses in the `requested_statuses` dictionary. + requested_statuses (Dict): A dictionary holding the statuses requested by the user. + spec (spec.specification.MerlinSpec): A [`MerlinSpec`][spec.specification.MerlinSpec] + object loaded from the workspace or spec file. + steps_filter_provided (bool): Indicates if a specific steps filter was provided. + + Methods: + apply_filters: Applies user-defined filters to the requested statuses. + apply_max_tasks_limit: Limits the number of tasks displayed based on the user-defined maximum. + display: Displays a task-by-task view of the status based on user filters. + filter_via_prompts: Interacts with the user to manage task display filters. + get_steps_to_display: Generates a list of steps to display the status for. + get_user_filters: Prompts the user for filters to apply to the statuses. + get_user_max_tasks: Prompts the user for a maximum task limit to display. + load_requested_statuses: Populates the requested statuses dictionary based on user-defined filters. """ def __init__(self, args: Namespace, spec_display: bool, file_or_ws: str): + """ + Initializes the `DetailedStatus` object, extending the functionality of the `Status` class + to include filtering and detailed task-by-task status handling. + + Args: + args: Command-line arguments provided by the user, including options + for filtering, displaying, or dumping detailed task statuses. + spec_display: A flag indicating whether the status should be loaded from a + specification file (`True`) or from a workspace (`False`). + file_or_ws: The path to the specification file or workspace, depending on the + value of `spec_display`. + """ args_copy = Namespace(**vars(args)) super().__init__(args, spec_display, file_or_ws) @@ -603,23 +807,30 @@ def __init__(self, args: Namespace, spec_display: bool, file_or_ws: str): os.environ["MANPAGER"] = "less -r" # Check if the steps filter was given - self.steps_filter_provided = "all" not in args_copy.steps + self.steps_filter_provided: bool = "all" not in args_copy.steps def _verify_filters( self, filters_to_check: List[str], valid_options: Union[List, Tuple], suppress_warnings: bool, - warning_msg: Optional[str] = "", + warning_msg: str = "", ): """ - Check each filter in a list of filters provided by the user against a list of valid options. - If the filter is invalid, remove it from the list of filters. - - :param `filters_to_check`: A list of filters provided by the user - :param `valid_options`: A list of valid options for this particular filter - :param `suppress_warnings`: If True, don't log warnings. Otherwise, log them - :param `warning_msg`: An optional warning message to attach to output + Verify and validate a list of user-provided filters against a set of valid options. + + This method checks each filter in the `filters_to_check` list to determine if it is present + in the `valid_options`. If a filter is found to be invalid (i.e., not in `valid_options`), + it is removed from the `filters_to_check` list. Depending on the value of `suppress_warnings`, + a warning message may be logged for each invalid filter. + + Args: + filters_to_check: A list of filters provided by the user that need to be validated. + valid_options: A list or tuple of valid options against which the filters will be checked. + suppress_warnings: A boolean flag indicating whether to suppress warning messages. + If True, no warnings will be logged for invalid filters. + warning_msg: An optional string that provides additional context for the warning message + logged when an invalid filter is detected. Default is an empty string. """ for filter_arg in filters_to_check[:]: if filter_arg not in valid_options: @@ -627,11 +838,17 @@ def _verify_filters( LOG.warning(f"The filter '{filter_arg}' is invalid. {warning_msg}") filters_to_check.remove(filter_arg) - def _verify_filter_args(self, suppress_warnings: Optional[bool] = False): + def _verify_filter_args(self, suppress_warnings: bool = False): """ - Verify that our filters are all valid and able to be used. + Verify the validity of filter arguments used in the current context. + + This method checks various filter arguments, including steps, max_tasks, task_status, + return_code, task_queues, and workers, to ensure they are valid and can be used. + Invalid filters are removed from their respective lists, and warnings may be logged + based on the `suppress_warnings` flag. - :param `suppress_warnings`: If True, don't log warnings. Otherwise, log them. + Args: + suppress_warnings: If True, suppress logging of warnings for invalid filters. """ # Ensure the steps are valid if "all" not in self.args.steps: @@ -717,8 +934,11 @@ def _verify_filter_args(self, suppress_warnings: Optional[bool] = False): def _process_task_queue(self): """ - Modifies the list of steps to display status for based on - the list of task queues provided by the user. + Modify the list of steps to display status for based on the provided task queues. + + This method processes the task queues specified by the user, removing any duplicates + and checking for their validity. It updates the list of steps to include those associated + with the valid task queues. If a provided task queue does not exist, a warning is logged. """ from merlin.config.configfile import CONFIG # pylint: disable=C0415 @@ -744,11 +964,16 @@ def _process_task_queue(self): def get_steps_to_display(self) -> Dict[str, List[str]]: """ - Generates a list of steps to display the status for based on information - provided to the merlin detailed-status command by the user. This function - will handle the --steps and --task-queues filter options. + Generate a dictionary of steps to display the status for based on user-provided filters. - :returns: A dictionary of started and unstarted steps for us to display the status of + This method processes the `--steps` and `--task-queues` options from the `merlin + detailed-status` command. It determines which steps should be included in the status + display based on the existing steps in the study and the specified filters. + + Returns: + A dictionary containing two lists:\n + - `started`: A list of steps that have been started. + - `unstarted`: A list of steps that have not yet been started. """ existing_steps = self.spec.get_study_step_names() @@ -781,9 +1006,16 @@ def get_steps_to_display(self) -> Dict[str, List[str]]: def _remove_steps_without_statuses(self): """ - After applying filters, there's a chance that certain steps will still exist - in self.requested_statuses but won't have any tasks to view the status of so - we'll remove those here. + Remove steps from the requested statuses that do not have any associated tasks. + + This method iterates through the `requested_statuses` dictionary and checks each step + for associated sub-steps. If a step does not have any valid sub-step workspaces (i.e., + it has no tasks to view the status of), it is removed from the `requested_statuses`. + + Note: + After applying filters, there's a chance that certain steps will still exist + in self.requested_statuses but won't have any tasks to view the status of. That's + why this method is necessary. """ result = deepcopy(self.requested_statuses) @@ -797,11 +1029,16 @@ def _remove_steps_without_statuses(self): def _search_for_filter(self, filter_to_apply: List[str], entry_to_search: Union[List[str], str]) -> bool: """ - Search an entry to see if our filter(s) apply to this entry. If they do, return True. Otherwise, False. + Search an entry to see if the specified filters apply to it. - :param filter_to_apply: A list of filters to search for - :param entry_to_search: A list or string of entries to search for our filters in - :returns: True if a filter was found in the entry. False otherwise. + This method checks if any of the provided filters match the given entry or entries. + + Args: + filter_to_apply: A list of filters to search for. + entry_to_search: A list or string of entries to search for the filters in. + + Returns: + True if a filter was found in the entry; False otherwise. """ if not isinstance(entry_to_search, list): entry_to_search = [entry_to_search] @@ -814,10 +1051,12 @@ def _search_for_filter(self, filter_to_apply: List[str], entry_to_search: Union[ def apply_filters(self): """ - Apply any filters given by the --workers, --return-code, and/or --task-status arguments. - This function will also apply the --max-tasks limit if it was set by a user. We apply this - limit here so it can be done in-place; if we called apply_max_tasks_limit instead, this - would become a two-pass algorithm and can be really slow with lots of statuses. + Apply filters based on the provided command-line arguments for workers, return code, + and task status, as well as enforce a maximum task limit if specified. + + This method processes the `requested_statuses` to filter out entries that do not match + the specified criteria. It ensures that the filtering is done in-place to optimize performance + and avoid a two-pass algorithm, which can be inefficient with a large number of statuses. """ if self.args.max_tasks is not None: # Make sure the max_tasks variable is set to a reasonable number and store that value @@ -897,8 +1136,13 @@ def apply_filters(self): def apply_max_tasks_limit(self): """ - Given a number representing the maximum amount of tasks to display, filter the dict of statuses - so that there are at most a max_tasks amount of tasks. + Filter the dictionary of statuses to ensure that the number of displayed tasks does not exceed + the specified maximum limit. + + This method checks the current value of `max_tasks` and adjusts it if it exceeds the number + of available statuses. It then iterates through the `requested_statuses`, removing excess + entries to comply with the `max_tasks` limit. The method also merges the allowed task statuses + into a new dictionary and updates the `requested_statuses` accordingly. """ # Make sure the max_tasks variable is set to a reasonable number and store that value if self.args.max_tasks > self.num_requested_statuses: @@ -959,15 +1203,24 @@ def load_requested_statuses(self): def get_user_filters(self) -> bool: """ - Get a filter on the statuses to display from the user. Possible options - for filtering: - - A str MAX_TASKS -> will ask the user for another input that's equivalent to the --max-tasks flag - - A list of statuses -> equivalent to the --task-status flag - - A list of return codes -> equivalent to the --return-code flag - - A list of workers -> equivalent to the --workers flag - - An exit keyword to leave the filter prompt without filtering - - :returns: True if we need to exit without filtering. False otherwise. + Prompt the user to specify filters for the statuses to display. The user can choose from + several filtering options, including setting a maximum number of tasks, filtering by status, + return code, or worker, or exiting the filter prompt without applying any filters. + + The method displays available filter options and their descriptions, then collects and + validates the user's input. If the user provides valid filters, they are stored in the + corresponding attributes. If the user opts to exit, the method returns True; otherwise, + it returns False. + + Possible filtering options include:\n + - A string "MAX_TASKS" to request a limit on the number of tasks. + - A list of statuses to filter by, corresponding to the `--task-status` flag. + - A list of return codes to filter by, corresponding to the `--return-code` flag. + - A list of workers to filter by, corresponding to the `--workers` flag. + - An exit keyword to leave the filter prompt without applying any filters. + + Returns: + True if the user chooses to exit without filtering; False otherwise. """ valid_workers = tuple(self.spec.get_worker_names()) @@ -1050,7 +1303,17 @@ def get_user_filters(self) -> bool: def get_user_max_tasks(self): """ - Get a limit for the amount of tasks to display from the user. + Prompt the user to specify a maximum limit for the number of tasks to display. + + The method repeatedly requests input from the user until a valid integer greater than 0 + is provided. Once a valid input is received, it sets the `max_tasks` attribute in the + `args` object to the specified limit. + + This method ensures that the user input is validated and handles any exceptions + related to invalid input types or values. + + Raises: + ValueError: If the input is not a valid integer greater than 0. """ invalid_input = True @@ -1069,8 +1332,16 @@ def get_user_max_tasks(self): def filter_via_prompts(self): """ - Interact with the user to manage how many/which tasks are displayed. This helps to - prevent us from overloading the terminal by displaying a bazillion tasks at once. + Interact with the user to determine how many and which tasks should be displayed, + preventing terminal overload by limiting the output to a manageable number of tasks. + + This method prompts the user for filtering options, including task statuses, return codes, + and worker specifications. It also handles the case where the user opts to exit without + applying any filters. If filters are provided, it applies them accordingly. + + Warning: + The method includes specific handling for the "RESTART" and "RETRY" return codes, + which are currently not implemented, and issues warnings if these filters are selected. """ # Get the filters from the user exit_without_filtering = self.get_user_filters() @@ -1095,11 +1366,18 @@ def filter_via_prompts(self): elif self.args.max_tasks is not None: self.apply_max_task_limit() - def display(self, test_mode: Optional[bool] = False): + def display(self, test_mode: bool = False): """ - Displays a task-by-task view of the status based on user filter(s). + Displays a task-by-task view of the statuses based on the user-defined filters. - :param `test_mode`: If true, run this in testing mode and don't print any output + This method checks for any requested statuses and, if found, invokes the + `display_status_task_by_task` function to present the tasks accordingly. + If no statuses are available to display, it logs a warning message. + + Args: + test_mode: If set to True, the method runs in testing mode, suppressing + any output to the terminal. This is useful for unit testing or debugging + without cluttering the output. """ # Check that there's statuses found and display them if self.requested_statuses: @@ -1111,26 +1389,31 @@ def display(self, test_mode: Optional[bool] = False): # Pylint complains that args is unused but we can ignore that def status_conflict_handler(*args, **kwargs) -> Any: # pylint: disable=W0613 """ - The conflict handler function to apply to any status entries that have conflicting - values while merging two status files together. - - kwargs should include: - - dict_a_val: The conflicting value from the dictionary that we're merging into - - dict_b_val: The conflicting value from the dictionary that we're pulling from - - key: The key into each dictionary that has a conflict - - path: The path down the dictionary tree that `dict_deep_merge` is currently at - - When we're reading in status files, we're merging all of the statuses into one dictionary. - This function defines the merge rules in case there is a merge conflict. We ignore the list - and dictionary entries since `dict_deep_merge` from `utils.py` handles these scenarios already. - - There are currently 4 rules: - - string-concatenate: take the two conflicting values and concatenate them in a string - - use-dict_b-and-log-debug: use the value from dict_b and log a debug message - - use-longest-time: use the longest time between the two conflicting values - - use-max: use the larger integer between the two conflicting values - - :returns: The value to merge into dict_a at `key` + Handles conflicts that arise when merging two status files by applying specific merge rules + to conflicting values. + + This function is designed to be used during the merging process of status entries, where + conflicting values may exist. It defines how to resolve these conflicts based on predefined + rules, ensuring that the merged dictionary maintains integrity and clarity. + + The merge rules currently implemented are:\n + - **string-concatenate**: Concatenates the two conflicting string values. + - **use-dict_b-and-log-debug**: Uses the value from dict_b and logs a debug message indicating + the conflict. + - **use-longest-time**: Chooses the longest time value between the two conflicting entries, + converting them to a timedelta for comparison. + - **use-max**: Selects the maximum integer value from the two conflicting entries. + + If a key does not have a defined merge rule, a warning is logged, and the function returns None. + + The function expects the following keyword arguments:\n + - `dict_a_val`: The conflicting value from the dictionary that we are merging into (dict_a). + - `dict_b_val`: The conflicting value from the dictionary that we are merging from (dict_b). + - `key`: The key in each dictionary that has a conflict. + - `path`: The current path in the dictionary tree during the merge process. + + Returns: + The resolved value to merge into dict_a at the specified key. """ # Grab the arguments passed into this function dict_a_val = kwargs.get("dict_a_val", None) @@ -1203,12 +1486,25 @@ def read_status( """ Locks the status file for reading and returns its contents. - :param status_filepath: The path to the status file that we'll read from. - :param lock_file: The path to the lock file that we'll use to create a FileLock. - :param display_fnf_message: If True, display the file not found warning. Otherwise don't. - :param raise_errors: A boolean indicating whether to ignore errors or raise them. - :param timeout: An integer representing how long to hold a lock for before timing out. - :returns: A dict of the contents in the status file + This function attempts to read the contents of a status file while ensuring that the file is + locked to prevent race conditions. It handles various exceptions that may occur during the + reading process, including file not found errors and JSON decoding errors. + + Args: + status_filepath: The path to the status file that will be read. + lock_file: The path to the lock file used to create a FileLock. + display_fnf_message: If True, displays a warning message if the file is not found. + raise_errors: If True, raises exceptions when errors occur. + timeout: The maximum time (in seconds) to hold the lock before timing out. + + Returns: + A dictionary containing the contents of the status file. + + Raises: + Timeout: If the lock acquisition times out. + FileNotFoundError: If the status file does not exist and `raise_errors` is True. + json.decoder.JSONDecodeError: If the status file is empty or contains invalid JSON and `raise_errors` is True. + Exception: Any other exceptions that occur during the reading process if `raise_errors` is True. """ statuses_read = {} @@ -1249,13 +1545,20 @@ def read_status( def write_status(status_to_write: Dict, status_filepath: str, lock_file: str, timeout: int = 10): """ - Locks the status file for writing. We're not catching any errors here since we likely want to - know if something went wrong in this process. + Locks the status file for writing and writes the provided status to the file. + + This function ensures that the status file is locked during the write operation to prevent + race conditions. It does not catch errors during the writing process, as it is important to + be aware of any issues that may arise. + + Args: + status_to_write: The status data to write to the status file. + status_filepath: The path to the status file where the status will be written. + lock_file: The path to the lock file used to create a FileLock for the write operation. + timeout: The maximum time (in seconds) to hold the lock before timing out. - :param status_to_write: The status to write to the status file - :param status_filepath: The path to the status file that we'll write the status to - :param lock_file: The path to the lock file we'll use for this status write - :param timeout: A timeout value for the lock so it's always released eventually + Raises: + Exception: Any exceptions that occur during the writing process will be logged, but not caught. """ # Pylint complains that we're instantiating an abstract class but this is correct usage try: diff --git a/merlin/study/status_constants.py b/merlin/study/status_constants.py index b7dfe7fa3..e5dee8c36 100644 --- a/merlin/study/status_constants.py +++ b/merlin/study/status_constants.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1 +# This file is part of Merlin, Version: 1.12.2 # # For details, see https://github.com/LLNL/merlin. # @@ -29,8 +29,8 @@ ############################################################################### """ This file contains all of the constants used for the status command. -Separating this from status.py and status_renderers.py helps with circular -import issues. +Separating this from [`status.py`][study.status] and [`status_renderers.py`][study.status_renderers] +helps with circular import issues. """ VALID_STATUS_FILTERS = ("INITIALIZED", "RUNNING", "FINISHED", "FAILED", "CANCELLED", "DRY_RUN", "UNKNOWN") diff --git a/merlin/study/status_renderers.py b/merlin/study/status_renderers.py index 02d6ab948..2dee300ab 100644 --- a/merlin/study/status_renderers.py +++ b/merlin/study/status_renderers.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1 +# This file is part of Merlin, Version: 1.12.2 # # For details, see https://github.com/LLNL/merlin. # @@ -29,7 +29,7 @@ ############################################################################### """This module handles creating a formatted task-by-task status display""" import logging -from typing import Dict, List, Optional, Union +from typing import Dict, List, Union from maestrowf import BaseStatusRenderer, FlatStatusRenderer, StatusRendererFactory from rich import box @@ -45,14 +45,20 @@ LOG = logging.getLogger(__name__) -def format_label(label_to_format: str, delimiter: Optional[str] = "_") -> str: +def format_label(label_to_format: str, delimiter: str = "_") -> str: """ - Take a string of the format 'word1_word2_...' and format it so it's prettier. - This would turn the string above to 'Word1 Word2 ...'. + Format a string by replacing a specified delimiter with spaces and capitalizing each word. - :param `label_to_format`: The string we want to format - :param `delimiter`: The character separating words in `label_to_format` - :returns: A formatted string based on `label_to_format` + This function takes a string that uses a specific delimiter to separate words and returns a + more readable version of that string, where the words are separated by spaces and each word + is capitalized. + + Args: + label_to_format: The string to format. + delimiter: The character that separates words in `label_to_format`. + + Returns: + A formatted string where the delimiter is replaced with spaces and each word is capitalized. """ return label_to_format.replace(delimiter, " ").title() @@ -60,20 +66,41 @@ def format_label(label_to_format: str, delimiter: Optional[str] = "_") -> str: class MerlinDefaultRenderer(BaseStatusRenderer): """ This class handles the default status formatting for task-by-task display. - It will separate the display on a step-by-step basis. - - Similar to Maestro's 'narrow' status display. + It will separate the display on a step-by-step basis, similar to Maestro's 'narrow' status display. + + Attributes: + disable_theme (bool): Flag to disable theming for the display. + disable_pager (bool): Flag to disable pager functionality for the display. + _theme_dict (Dict[str, str]): A dictionary containing the theme settings for various status types. + _status_table (Table): A Table object that contains the formatted status information. + + Methods: + create_param_table: Creates the parameter section of the display. + create_step_table: Creates each step entry in the display. + create_task_details_table: Creates the task details section of the display. + layout: Sets up the overall layout of the display. + render: Performs the actual printing of the status table with optional theme customization. """ - def __init__(self, *args, **kwargs): + def __init__(self, *args: List, **kwargs: Dict): + """ + Initializes the `MerlinDefaultRenderer` instance, which handles the default status formatting + for task-by-task display, with optional theming and pager functionality. + + Args: + *args: Positional arguments passed to the superclass (`BaseStatusRenderer`). + **kwargs: Keyword arguments used to configure the renderer. Supported keys include:\n + - disable_theme (bool, optional): If `True`, disables theming for the display. Defaults to `False`. + - disable_pager (bool, optional): If `True`, disables pager functionality for the display. Defaults to `False`. + """ super().__init__(*args, **kwargs) - self.disable_theme = kwargs.pop("disable_theme", False) - self.disable_pager = kwargs.pop("disable_pager", False) + self.disable_theme: bool = kwargs.pop("disable_theme", False) + self.disable_pager: bool = kwargs.pop("disable_pager", False) # Setup default theme # TODO modify this theme to add more colors - self._theme_dict = { + self._theme_dict: Dict[str, str] = { "INITIALIZED": "blue", "RUNNING": "blue", "DRY_RUN": "green", @@ -92,14 +119,23 @@ def __init__(self, *args, **kwargs): } # Setup the status table that will contain our formatted status - self._status_table = Table.grid(padding=0) + self._status_table: Table = Table.grid(padding=0) def create_param_table(self, parameters: Dict[str, Dict[str, str]]) -> Columns: """ - Create the parameter section of the display + Create the parameter section of the display. - :param `parameters`: A dict of the form {"cmd": {"TOKEN1": "value1"}, "restart": {"TOKEN2": "value1"}} - :returns: A rich Columns object with the parameter info formatted appropriately + This method generates a formatted table for the parameters associated with each command type. + Each command type (e.g., "cmd", "restart") will have its own sub-table displaying the tokens + and their corresponding values. + + Args: + parameters: A dictionary where each key is a command type (e.g., "cmd", "restart") + and each value is another dictionary containing token-value pairs. + Example format: `{"cmd": {"TOKEN1": "value1"}, "restart": {"TOKEN2": "value1"}}`. + + Returns: + A rich Columns object containing the formatted parameter tables, arranged side-by-side. """ param_table = [] # Loop through cmd and restart entries @@ -135,18 +171,26 @@ def create_step_table( self, step_name: str, parameters: Dict[str, Dict[str, str]], - task_queue: Optional[str] = None, - workers: Optional[str] = None, + task_queue: str = None, + workers: str = None, ) -> Table: """ - Create each step entry in the display - - :param `step_name`: The name of the step that we're setting the layout for - :param `parameters`: The parameters dict for this step - :param `task_queue`: The name of the task queue associated with this step if one was provided - :param `workers`: The name of the worker(s) that ran this step if one was provided - :returns: A rich Table object with info for one sub step (here a 'sub step' is referencing a step - with multiple parameters; each parameter set will have it's own entry in the output) + Create each step entry in the display. + + This method constructs a formatted table entry for a specific step in the process, including + relevant details such as the step name, associated task queue, worker(s), and any parameters + related to the step. Each parameter set will be displayed in a sub-table format. + + Args: + step_name: The name of the step for which the layout is being created. + parameters: A dictionary of parameters associated with the step, where each key is a + parameter type and each value is a dictionary of token-value pairs. + task_queue: The name of the task queue associated with this step, if provided. + workers: The name(s) of the worker(s) that executed this step, if provided. + + Returns: + A rich Table object containing the formatted information for the specified step, + including its parameters and any associated task queue or worker details. """ # Initialize the table that will have our step entry information step_table = Table(box=box.SIMPLE_HEAVY, show_header=False) @@ -172,10 +216,20 @@ def create_step_table( def create_task_details_table(self, task_statuses: Dict) -> Table: """ - Create the task details section of the display + Create the task details section of the display. + + This method constructs a formatted table that displays detailed information about various tasks, + including their statuses, return codes, elapsed times, run times, restarts, and associated workers. + Each task is represented as a row in the table, with specific styling applied based on the task's status. - :param `task_statuses`: A dict of task statuses to format into our layout - :returns: A rich Table with the formatted task info for a sub step + Args: + task_statuses: A dictionary containing task statuses, where each key represents a step workspace + and each value is another dictionary with details such as status, return code, elapsed time, + run time, restarts, and workers. + + Returns: + A rich Table object containing the formatted task details, structured for easy readability + and visual distinction based on task status. """ # Initialize the task details table task_details = Table(title="Task Details") @@ -222,15 +276,23 @@ def create_task_details_table(self, task_statuses: Dict) -> Table: return task_details - def layout( - self, status_data, study_title: Optional[str] = None, status_time: Optional[str] = None - ): # pylint: disable=W0237 + def layout(self, status_data: Dict, study_title: str = None, status_time: str = None): # pylint: disable=W0237 """ - Setup the overall layout of the display + Setup the overall layout of the display. + + This method configures the main display layout for the status data, including setting up + the title with optional study information and timestamp. It organizes the status data into + a structured table format, displaying each step's details along with associated task information. - :param `status_data`: A dict of status data to display - :param `study_title`: A title for the study to display at the top of the output - :param `status_time`: A timestamp to add to the title + Args: + status_data: A dictionary containing status data to be displayed, where each key + represents a step and its associated information. + study_title: A title for the study to be displayed at the top of the output. + status_time: A timestamp to be included in the title, indicating when the status + data was captured. + + Raises: + ValueError: If `status_data` is not a dictionary or is empty. """ if isinstance(status_data, dict) and status_data: self._status_data = status_data @@ -280,11 +342,17 @@ def layout( # Add this step to the full status table self._status_table.add_row(step_table, end_section=True) - def render(self, theme: Optional[Dict[str, str]] = None): + def render(self, theme: Dict[str, str] = None): """ - Do the actual printing + Do the actual printing of the status table. + + This method is responsible for rendering the status table to the console, applying any specified + theme settings for visual customization. It handles the enabling or disabling of themes and + manages the output display, either using a pager for long outputs or printing directly to the console. - :param `theme`: A dict of theme settings (see self._theme_dict for the appropriate layout) + Args: + theme: A dictionary of theme settings that define the appearance of the output. The keys and + values should correspond to the layout defined in `self._theme_dict`. """ # Apply any theme customization if theme: @@ -314,24 +382,37 @@ class MerlinFlatRenderer(FlatStatusRenderer): """ This class handles the flat status formatting for task-by-task display. It will not separate the display on a step-by-step basis and instead group - all statuses together in a single table. + all statuses together in a single table, similar to Maestro's 'flat' status display. + + Attributes: + disable_theme (bool): A flag indicating whether to disable theme customization for the output. + disable_pager (bool): A flag indicating whether to disable the use of a pager for long outputs. - Similar to Maestro's 'flat' status display. + Methods: + layout: Sets up the layout of the display, formatting the status data and study title. + render: Renders the status table to the console, applying any specified theme settings and + managing the output display. """ def __init__(self, *args, **kwargs): super().__init__(args, kwargs) - self.disable_theme = kwargs.pop("disable_theme", False) - self.disable_pager = kwargs.pop("disable_pager", False) + self.disable_theme: bool = kwargs.pop("disable_theme", False) + self.disable_pager: bool = kwargs.pop("disable_pager", False) - def layout( - self, status_data: Dict[str, List[Union[str, int]]], study_title: Optional[str] = None - ): # pylint: disable=W0221 + def layout(self, status_data: Dict[str, List[Union[str, int]]], study_title: str = None): # pylint: disable=W0221 """ - Setup the layout of the display - - :param `status_data`: A dict of status information that we'll display - :param `study_title`: The title of the study to display at the top of the output + Set up the layout of the display for the status information. + + This method processes the provided status data by removing unnecessary parameters, + capitalizing the column labels, and preparing the data for display. It also allows + for an optional study title to be displayed at the top of the output. + + Args: + status_data: A dictionary containing status information to be displayed. The + keys represent the status categories, and the values are lists of + corresponding status values. + study_title: The title of the study to display at the top of the output. + If provided, it will be included in the layout. """ if "cmd_parameters" in status_data: del status_data["cmd_parameters"] @@ -344,11 +425,22 @@ def layout( super().layout(status_data, study_title=study_title) - def render(self, theme: Optional[Dict[str, str]] = None): + def render(self, theme: Dict[str, str] = None): """ - Do the actual printing - - :param `theme`: A dict of theme settings (see self._theme_dict for the appropriate layout) + Render the status table to the console. + + This method is responsible for displaying the formatted status information + in the console. It applies any specified theme settings to customize the + appearance of the output. If the theme is disabled, it sets all theme + attributes to 'none'. The method also handles the output display, either + printing directly to the console or using a pager for long outputs based + on the `disable_pager` attribute. + + Args: + theme (Dict[str, str], optional): A dictionary of theme settings that + customize the appearance of the output. The keys represent the + theme attributes, and the values are the corresponding settings. + If not provided, the default theme settings will be used. """ # Apply any theme customization if theme: @@ -377,6 +469,21 @@ def render(self, theme: Optional[Dict[str, str]] = None): class MerlinStatusRendererFactory(StatusRendererFactory): """ This class keeps track of all available status layouts for Merlin. + + The `MerlinStatusRendererFactory` is responsible for managing different + status layout renderers used in the Merlin application. It provides a + method to retrieve the appropriate renderer based on the specified layout + type and user preferences regarding theme and pager usage. + + Attributes: + _layouts (Dict[str, BaseStatusRenderer]): A dictionary mapping layout names to their corresponding renderer + classes. Currently includes "table" for + [`MerlinFlatRenderer`][study.status_renderers.MerlinFlatRenderer] and + "default" for [`MerlinDefaultRenderer`][study.status_renderers.MerlinDefaultRenderer]. + + Methods: + get_renderer: Retrieves an instance of the specified layout renderer, applying + user preferences for theme and pager settings. """ # TODO: when maestro releases the pager changes: @@ -387,21 +494,29 @@ class MerlinStatusRendererFactory(StatusRendererFactory): # - remove render method in MerlinDefaultRenderer # - this will also be in BaseStatusRenderer in Maestro def __init__(self): # pylint: disable=W0231 - self._layouts = { + self._layouts: Dict[str, BaseStatusRenderer] = { "table": MerlinFlatRenderer, "default": MerlinDefaultRenderer, } - def get_renderer(self, layout: str, disable_theme: bool, disable_pager: bool): # pylint: disable=W0221 - """Get handle for specific layout renderer to instantiate + def get_renderer( + self, layout: str, disable_theme: bool, disable_pager: bool + ) -> BaseStatusRenderer: # pylint: disable=W0221 + """ + Get handle for specific layout renderer to instantiate. + + Args: + layout: A string denoting the name of the layout renderer to use. + disable_theme: True if the user wants to disable themes when displaying + status; False otherwise. + disable_pager: True if the user wants to disable the pager when displaying + status; False otherwise. - :param `layout`: A string denoting the name of the layout renderer to use - :param `disable_theme`: True if the user wants to disable themes when displaying status. - False otherwise. - :param `disable_pager`: True if the user wants to disable the pager when displaying status. - False otherwise. + Returns: + The status renderer class to use for displaying the output. - :returns: The status renderer class to use for displaying the output + Raises: + ValueError: If the specified layout is not found in the available layouts. """ renderer = self._layouts.get(layout) diff --git a/merlin/study/step.py b/merlin/study/step.py index 26d737e14..8cf9d9ca5 100644 --- a/merlin/study/step.py +++ b/merlin/study/step.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -34,14 +34,16 @@ import re from contextlib import suppress from copy import deepcopy -from typing import Dict, Optional, Tuple +from typing import Dict, List, Tuple from celery import current_task from maestrowf.abstracts.enums import State +from maestrowf.abstracts.interfaces.scriptadapter import ScriptAdapter from maestrowf.datastructures.core.executiongraph import _StepRecord from maestrowf.datastructures.core.study import StudyStep +from maestrowf.interfaces.script import SubmissionRecord -from merlin.common.abstracts.enums import ReturnCode +from merlin.common.enums import ReturnCode from merlin.study.script_adapter import MerlinScriptAdapter from merlin.study.status import read_status, write_status from merlin.utils import needs_merlin_expansion @@ -50,15 +52,34 @@ LOG = logging.getLogger(__name__) -def get_current_worker(): - """Get the worker on the current running task from celery""" +def get_current_worker() -> str: + """ + Get the worker on the current running task from Celery. + + This function retrieves the name of the worker that is currently + executing the task. It extracts the worker's name from the task's + request hostname. + + Returns: + The name of the current worker. + """ worker = re.search(r"@.+\.", current_task.request.hostname).group() worker = worker[1 : len(worker) - 1] return worker -def get_current_queue(): - """Get the queue on the current running task from celery""" +def get_current_queue() -> str: + """ + Get the queue on the current running task from Celery. + + This function retrieves the name of the queue that the current + task is associated with. It extracts the routing key from the + task's delivery information and removes the queue tag defined + in the configuration. + + Returns: + The name of the current queue. + """ from merlin.config.configfile import CONFIG # pylint: disable=C0415 queue = current_task.request.delivery_info["routing_key"] @@ -68,24 +89,55 @@ def get_current_queue(): class MerlinStepRecord(_StepRecord): """ - This class is a wrapper for the Maestro _StepRecord to remove + This class is a wrapper for the Maestro `_StepRecord` to remove a re-submit message and handle status updates. + + Attributes: + condensed_workspace (str): A condensed version of the workspace path. + elapsed_time (str): The total elapsed time for the step execution. + jobid (List[int]): A list of job identifiers assigned by the scheduler. + maestro_step (StudyStep): The StudyStep object associated with this step. + merlin_step (Step): The Step object associated with this step. + restart_limit (int): Upper limit on the number of restart attempts. + restart_script (str): Script to resume record execution (if applicable). + run_time (str): The run time for the step execution. + status (State): The current status of the step. + to_be_scheduled (bool): Indicates if the record needs scheduling. + workspace (Variable): The output workspace for this step, represented as a Variable. + + Methods: + mark_end: Marks the end of the step with the given state. + mark_restart: Increments the restart count for the step. + mark_running: Marks the step as running and updates the status file. + setup_workspace: Initializes the workspace and status file for the step. """ def __init__(self, workspace: str, maestro_step: StudyStep, merlin_step: "Step", **kwargs): """ - :param `workspace`: The output workspace for this step - :param `maestro_step`: The StudyStep object associated with this step - :param `merlin_step`: The Step object associated with this step + Initializes the `MerlinStepRecord` class which helps track the status of a step. + + Args: + workspace: The output workspace for this step. + maestro_step: The + [StudyStep](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/datastructures/core/index.html#maestrowf.datastructures.core.StudyStep) + object associated with this step. + merlin_step: The [Step][study.step.Step] object associated with this step. """ _StepRecord.__init__(self, workspace, maestro_step, status=State.INITIALIZED, **kwargs) - self.merlin_step = merlin_step + self.merlin_step: Step = merlin_step @property def condensed_workspace(self) -> str: """ - Put together a smaller version of the workspace path to display. - :returns: A condensed workspace name + Generate a condensed version of the workspace path for display purposes. + + This property constructs a shorter representation of the workspace path by extracting relevant + components based on the study name and a timestamp pattern. If a match is found using a regular + expression, the workspace path is split to isolate the condensed portion. If no match is found, + a fallback method is used to manually create a condensed path based on the step name. + + Returns: + A string representing the condensed workspace path, which is easier to read and display. """ timestamp_regex = r"\d{8}-\d{6}/" match = re.search(rf"{self.merlin_step.study_name}_{timestamp_regex}", self.workspace.value) @@ -102,15 +154,21 @@ def condensed_workspace(self) -> str: LOG.debug(f"Condense workspace '{condensed_workspace}'") return condensed_workspace - def _execute(self, adapter: "ScriptAdapter", script: str) -> Tuple["SubmissionRecord", int]: # noqa: F821 + def _execute(self, adapter: ScriptAdapter, script: str) -> Tuple[SubmissionRecord, int]: """ - Overwrites _StepRecord's _execute method from Maestro since self.to_be_scheduled is - always true here. Also, if we didn't overwrite this we wouldn't be able to call - self.mark_running() for status updates. + Executes the script using the provided adapter, overriding the default behavior to ensure + that the step is marked as running and to facilitate job submission. - :param `adapter`: The script adapter to submit jobs to - :param `script`: The script to send to the script adapter - :returns: A tuple of a return code and the jobid from the execution of `script` + This method overrides the `_execute` method from the base class `_StepRecord` in Maestro. + It ensures that `self.to_be_scheduled` is always true, allowing for the invocation of + `self.mark_running()` to update the status of the step. + + Args: + adapter: The script adapter used to submit jobs. + script: The script to be submitted to the script adapter. + + Returns: + A tuple containing the return code and the job identifier from the execution of the script. """ self.mark_running() @@ -129,11 +187,17 @@ def mark_running(self): def mark_end(self, state: ReturnCode, max_retries: bool = False): """ - Mark the end time of the record with associated termination state - and update the status file. + Marks the end time of the record with the associated termination state + and updates the status file. + + This method logs the action of marking the end of the step, maps the provided + termination state to a corresponding Maestro state and result, and updates + the status file accordingly. If the maximum number of retries has been reached + for a soft failure, it appends a message to the result. - :param `state`: A merlin ReturnCode object representing the end state of a task - :param `max_retries`: A bool representing whether we hit the max number of retries or not + Args: + state: A ReturnCode object representing the end state of the task. + max_retries: A flag indicating whether the maximum number of retries has been reached. """ LOG.debug(f"Marking end for {self.name}") @@ -189,7 +253,7 @@ def mark_end(self, state: ReturnCode, max_retries: bool = False): self._update_status_file(result=step_result) def mark_restart(self): - """Increment the number of restarts we've had for this step and update the status file""" + """Increment the number of restarts we've had for this step and update the status file.""" LOG.debug(f"Marking restart for {self.name}") if self.restart_limit == 0 or self._num_restarts < self.restart_limit: self._num_restarts += 1 @@ -203,16 +267,22 @@ def setup_workspace(self): def _update_status_file( self, - result: Optional[str] = None, - task_server: Optional[str] = "celery", + result: str = None, + task_server: str = "celery", ): """ - Puts together a dictionary full of status info and creates a signature - for the update_status celery task. This signature is ran here as well. + Constructs a dictionary containing status information and creates a signature + for the update_status Celery task. This signature is executed within the method. - :param `result`: Optional parameter only applied when we've finished running - this step. String representation of a ReturnCode value. - :param `task_server`: Optional parameter to define the task server we're using. + This method checks if a status file already exists; if it does, it updates the + existing file with the current status information. If not, it initializes a new + status dictionary. The method also includes optional parameters for the result + of the task and the task server being used. + + Args: + result: An optional string representation of a ReturnCode value, applied + when the step has finished running. + task_server: An optional parameter to specify the task server being used. """ # This dict is used for converting an enum value to a string for readability @@ -294,43 +364,91 @@ def _update_status_file( class Step: """ This class provides an abstraction for an execution step, which can be - executed by calling execute. + executed by calling the [`execute`][study.step.Step.execute] method. + + Attributes: + max_retries (int): Returns the maximum number of retries for this step. + mstep (_StepRecord): The Maestro StepRecord object associated with this step. + parameter_info (dict): A dictionary containing information about parameters in the study. + params (Dict): A dictionary containing command parameters for the step, including 'cmd' and 'restart_cmd'. + restart (bool): Property to get or set the restart status of the step. + retry_delay (int): Returns the retry delay for the step (default is 1). + study_name (str): The name of the study this step belongs to. + + Methods: + check_if_expansion_needed: Checks if command expansion is needed based on specified labels. + clone_changing_workspace_and_cmd: Produces a deep copy of the current step, with optional command + and workspace modifications. + establish_params: Pulls parameters from the step parameter map if applicable. + execute: Executes the step using the provided adapter configuration. + get_cmd: Retrieves the run command text body. + get_restart_cmd: Retrieves the restart command text body, or None if not available. + get_task_queue: Retrieves the task queue for the step. + get_task_queue_from_dict: Static method to get the task queue from a step dictionary. + get_workspace: Retrieves the workspace where this step is to be executed. + name: Retrieves the name of the step. + name_no_params: Gets the original name of the step without parameters or sample labels. """ - def __init__(self, maestro_step_record, study_name, parameter_info): + def __init__(self, maestro_step_record: _StepRecord, study_name: str, parameter_info: Dict): """ - :param maestro_step_record: The StepRecord object. - :param `study_name`: The name of the study - :param `parameter_info`: A dict containing information about parameters in the study + Initializes the `Step` object which acts as a way to track everything about a step. + + Args: + maestro_step_record: The `StepRecord` object. + study_name: The name of the study + parameter_info: A dict containing information about parameters in the study """ - self.mstep = maestro_step_record - self.study_name = study_name - self.parameter_info = parameter_info - self.__restart = False - self.params = {"cmd": {}, "restart_cmd": {}} + self.mstep: _StepRecord = maestro_step_record + self.study_name: str = study_name + self.parameter_info: Dict = parameter_info + self.__restart: bool = False + self.params: Dict = {"cmd": {}, "restart_cmd": {}} self.establish_params() - def get_cmd(self): + def get_cmd(self) -> str: """ - get the run command text body" + Retrieve the run command text body for the step. + + Returns: + The run command text body for the step. """ return self.mstep.step.__dict__["run"]["cmd"] - def get_restart_cmd(self): + def get_restart_cmd(self) -> str: """ - get the restart command text body, else return None" - """ - return self.mstep.step.__dict__["run"]["restart"] + Retrieve the restart command text body for the step. - def clone_changing_workspace_and_cmd(self, new_cmd=None, cmd_replacement_pairs=None, new_workspace=None): + Returns: + The restart command text body for the step, or None if no restart command is available. """ - Produces a deep copy of the current step, performing variable - substitutions as we go + return self.mstep.step.__dict__["run"]["restart"] - :param new_cmd : (Optional) replace the existing cmd with the new_cmd. - :param cmd_replacement_pairs : (Optional) replaces strings in the cmd - according to the list of pairs in cmd_replacement_pairs - :param new_workspace : (Optional) the workspace for the new step. + def clone_changing_workspace_and_cmd( + self, + new_cmd: str = None, + cmd_replacement_pairs: List[Tuple[str]] = None, + new_workspace: str = None, + ) -> "Step": + """ + Produces a deep copy of the current step, with optional modifications to + the command and workspace, performing variable substitutions as we go. + + This method creates a new instance of the Step class by cloning the + current step and allowing for modifications to the command text and + workspace. It performs variable substitutions in the command based on + the provided replacement pairs. + + Args: + new_cmd: If provided, replaces the existing command with this new command. + cmd_replacement_pairs: A list of pairs where each pair contains a string to + be replaced and its replacement. The method will perform replacements in + both the run command and the restart command. + new_workspace: If provided, sets this as the workspace for the new step. If + not specified, the current workspace will be used. + + Returns: + A new Step instance with the modified command and workspace. """ LOG.debug(f"clone called with new_workspace {new_workspace}") step_dict = deepcopy(self.mstep.step.__dict__) @@ -356,13 +474,35 @@ def clone_changing_workspace_and_cmd(self, new_cmd=None, cmd_replacement_pairs=N study_step.run = step_dict["run"] return Step(MerlinStepRecord(new_workspace, study_step, self), self.study_name, self.parameter_info) - def get_task_queue(self): - """Retrieve the task queue for the Step.""" + def get_task_queue(self) -> str: + """ + Retrieve the task queue for the current Step. + + Returns: + The name of the task queue for the Step, which may be influenced + by the configuration settings. + """ return self.get_task_queue_from_dict(self.mstep.step.__dict__) @staticmethod - def get_task_queue_from_dict(step_dict): - """given a maestro step dict, get the task queue""" + def get_task_queue_from_dict(step_dict: Dict) -> str: + """ + Get the task queue from a given Maestro step dictionary. + + This static method extracts the task queue information from the + provided step dictionary. It considers the configuration settings + to determine the appropriate queue name, including handling cases + where the task queue may be omitted. + + Args: + step_dict: A dictionary representation of a Maestro step, expected + to contain a "run" key with a "task_queue" entry. + + Returns: + The name of the task queue. If the task queue is not specified + or is set to "none", it returns the default queue name based + on the configuration. + """ from merlin.config.configfile import CONFIG # pylint: disable=C0415 queue_tag = CONFIG.celery.queue_tag @@ -382,34 +522,56 @@ def get_task_queue_from_dict(step_dict): return queue @property - def retry_delay(self): - """Returns the retry delay (default 1)""" + def retry_delay(self) -> int: + """ + Get the retry delay for the step. + + Returns: + The retry delay in seconds. Defaults to 1 if not specified. + """ default_retry_delay = 1 return self.mstep.step.__dict__["run"].get("retry_delay", default_retry_delay) @property - def max_retries(self): + def max_retries(self) -> int: """ - Returns the max number of retries for this step. + Get the maximum number of retries for this step. + + Returns: + The maximum number of retries for the step. """ return self.mstep.step.__dict__["run"]["max_retries"] @property - def restart(self): + def restart(self) -> bool: """ - Get the restart property + Get the restart property. + + Returns: + True if the step is set to restart, False otherwise. """ return self.__restart @restart.setter - def restart(self, val): + def restart(self, val: bool): """ - Set the restart property ensuring that restart is false + Set the restart property. + + Args: + val: The new value for the restart property. It should be + a boolean value indicating whether the step should restart. """ self.__restart = val def establish_params(self): - """If this step uses parameters, pull them from the step param map.""" + """ + Establish parameters for the step from the parameter map. + + This method checks if the current step uses parameters by accessing + the `step_param_map` from `parameter_info`. If parameters are found + for the current step, it updates the `params` dictionary with the + corresponding values. + """ try: step_params = self.parameter_info["step_param_map"][self.name()] for cmd_type in step_params: @@ -417,29 +579,52 @@ def establish_params(self): except KeyError: pass - def check_if_expansion_needed(self, labels): + def check_if_expansion_needed(self, labels: List[str]) -> bool: """ - :return : True if the cmd has any of the default keywords or spec - specified sample column labels. + Check if expansion is needed based on commands and labels. + + This method determines whether the command associated with the + current step requires expansion. It checks for the presence of + default keywords or specified sample column labels. + + Args: + labels: A list of labels to check against the commands. + + Returns: + True if the command requires expansion, False otherwise. """ return needs_merlin_expansion(self.get_cmd(), self.get_restart_cmd(), labels) - def get_workspace(self): + def get_workspace(self) -> str: """ - :return : The workspace this step is to be executed in. + Get the workspace for the current step. + + Returns: + The workspace associated with this step. """ return self.mstep.workspace.value - def name(self): + def name(self) -> str: """ - :return : The step name. + Get the name of the current step. + + Returns: + The name of the step. """ return self.mstep.step.__dict__["_name"] - def name_no_params(self): + def name_no_params(self) -> str: """ - Get the original name of the step without any parameters/samples in the name. - :returns: A string representing the name of the step + Get the original name of the step without parameters or sample labels. + + This method retrieves the name of the step and removes any + parameter labels or sample identifiers that may be included + in the name. It ensures that the returned name is clean and + free from extraneous characters, such as trailing periods or + underscores. + + Returns: + The cleaned name of the step, free from parameters and sample labels. """ # Get the name with everything still in it name = self.name() @@ -459,13 +644,28 @@ def name_no_params(self): return name - def execute(self, adapter_config): + def execute(self, adapter_config: Dict) -> ReturnCode: """ - Execute the step. + Execute the step with the provided adapter configuration. + + This method performs the execution of the step by configuring + the necessary parameters and invoking the appropriate adapter. + It updates the adapter configuration based on the step's + requirements, sets up the workspace, and generates the script + for execution. If a dry run is specified, it prepares the + workspace without executing any tasks. + + Args: + adapter_config (dict): A dictionary containing configuration + for the maestro script adapter, including:\n + - `shell`: The shell to use for execution. + - `batch_type`: The type of batch processing to use. + - `dry_run`: A boolean indicating whether to perform a + dry run (setup only, no execution). - :param adapter_config : A dictionary containing configuration for - the maestro script adapter, as well as which sort of adapter - to use. + Returns: + (common.enums.ReturnCode): A [`ReturnCode`][common.enums.ReturnCode] object representing + the result of the execution. """ # Update shell if the task overrides the default value from the batch section default_shell = adapter_config.get("shell") diff --git a/merlin/study/study.py b/merlin/study/study.py index f30e36058..f37882fe7 100644 --- a/merlin/study/study.py +++ b/merlin/study/study.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -37,13 +37,16 @@ from contextlib import suppress from copy import deepcopy from pathlib import Path +from typing import Dict, List, Union +import numpy as np from cached_property import cached_property from maestrowf.datastructures.core import Study +from maestrowf.datastructures.core.studyenvironment import StudyEnvironment from maestrowf.maestro import load_parameter_generator from maestrowf.utils import create_dictionary -from merlin.common.abstracts.enums import ReturnCode +from merlin.common.enums import ReturnCode from merlin.spec import defaults from merlin.spec.expansion import determine_user_variables, expand_by_line, expand_env_vars, expand_line from merlin.spec.override import error_override_vars, replace_override_vars @@ -63,37 +66,83 @@ class MerlinStudy: # pylint: disable=R0902,R0904 """ Represents a Merlin study run on a specification. Used for 'merlin run'. - :param `filepath`: path to the desired specification file. - :param `override_vars`: Dictionary (keyword-variable name, value-variable - value) to override in the spec. - :param `restart_dir`: Filepath to restart study. If None, study runs - normally. - :param `samples_file`: File to load samples from. Ignores sample lookup - and generation in the spec if set. - :param `dry_run`: Flag to dry-run a workflow, which sets up the workspace but does not launch tasks. - :param `no_errors`: Flag to ignore some errors for testing. + This class manages the execution of a study based on a provided specification file, + handling sample data, output paths, workspace management, and the generation of a + Directed Acyclic Graph (DAG) for execution. + + Attributes: + dag (study.dag.DAG): Directed acyclic graph representing the execution flow of the study. + dry_run (bool): Flag indicating whether to perform a dry run of the workflow. + expanded_spec (spec.specification.MerlinSpec): The expanded specification after applying overrides. + filepath (str): Path to the desired specification file. + flux_command (str): Command for running flux jobs, if applicable. + info (str): Path to the 'merlin_info' directory within the workspace. + level_max_dirs (int): The number of directories at each level of the sample hierarchy. + no_errors (bool): Flag to ignore some errors for testing purposes. + original_spec (spec.specification.MerlinSpec): The original specification loaded + from the filepath. + output_path (str): Path to the output directory for the study. + override_vars (Dict[str, Union[str, int]]): Dictionary of variables to override in the specification. + parameter_labels (List[str]): List of parameter labels used in the study. + pargs (List[str]): Arguments for the parameter generator. + pgen_file (str): Filepath for the parameter generator, if applicable. + restart_dir (str): Filepath to restart the study, if applicable. + sample_labels (List[str]): The column labels of the samples. + samples (np.ndarray): The samples in the study. + samples_file (str): File to load samples from, if specified. + special_vars (Dict[str, str]): Dictionary of special variables used in the study. + timestamp (str): Timestamp representing the start time of the study. + user_vars (Dict[str, str]): The user-defined variables in the study. + workspace (str): Path to the workspace directory for the study. + + Methods: + generate_samples: Executes a command to generate sample data if the sample file is missing. + get_adapter_config: Builds and returns the adapter configuration dictionary. + get_expanded_spec: Returns a new YAML spec file with defaults, CLI overrides, and variable expansions. + get_sample_labels: Retrieves the column labels for the samples. + get_user_vars: Returns a dictionary of expanded user-defined variables from the specification. + label_clash_error: Checks for clashes between sample and parameter names. + load_dag: Generates a Directed Acyclic Graph (DAG) for the study's execution. + load_pgen: Executes a parameter generator script. + load_samples: Loads samples from disk or generates them if the file does not exist. + write_original_spec: Copies the original specification to the 'merlin_info' directory. """ def __init__( # pylint: disable=R0913 self, - filepath, - override_vars=None, - restart_dir=None, - samples_file=None, - dry_run=False, - no_errors=False, - pgen_file=None, - pargs=None, + filepath: str, + override_vars: Dict[str, Union[str, int]] = None, + restart_dir: str = None, + samples_file: str = None, + dry_run: bool = False, + no_errors: bool = False, + pgen_file: str = None, + pargs: List[str] = None, ): - self.filepath = filepath - self.original_spec = MerlinSpec.load_specification(filepath) - self.override_vars = override_vars + """ + Initializes a MerlinStudy object, which represents a study run based on a specification file. + + Args: + filepath: Path to the specification file for the study. + override_vars: Dictionary of variables to override in the specification. + restart_dir: Path to the directory for restarting the study. + samples_file: Path to a file containing sample data. If specified, the samples + will be loaded from this file. + dry_run: Flag indicating whether to perform a dry run of the workflow + without executing tasks. + no_errors: Flag to suppress certain errors for testing purposes. + pgen_file: Path to a parameter generator file. + pargs: Arguments for the parameter generator. + """ + self.filepath: str = filepath + self.original_spec: MerlinSpec = MerlinSpec.load_specification(filepath) + self.override_vars: Dict = override_vars error_override_vars(self.override_vars, self.original_spec.path) - self.samples_file = samples_file + self.samples_file: str = samples_file self.label_clash_error() - self.dry_run = dry_run - self.no_errors = no_errors + self.dry_run: bool = dry_run + self.no_errors: bool = no_errors # If we load from a file, record that in the object for provenance # downstream @@ -101,9 +150,9 @@ def __init__( # pylint: disable=R0913 self.original_spec.merlin["samples"]["file"] = self.samples_file self.original_spec.merlin["samples"]["generate"]["cmd"] = "" - self.restart_dir = restart_dir + self.restart_dir: str = restart_dir - self.special_vars = { + self.special_vars: Dict[str, str] = { "SPECROOT": self.original_spec.specroot, "MERLIN_TIMESTAMP": self.timestamp, "MERLIN_INFO": self.info, @@ -120,14 +169,21 @@ def __init__( # pylint: disable=R0913 } self._set_special_file_vars() - self.pgen_file = pgen_file - self.pargs = pargs + self.pgen_file: str = pgen_file + self.pargs: List[str] = pargs - self.dag = None + self.dag: DAG = None self.load_dag() def _set_special_file_vars(self): - """Setter for the orig, partial, and expanded file paths of a study.""" + """ + Sets the original, partial, and expanded file paths for a study. + + This method constructs file paths for three special variables + related to the study's specifications. It generates paths for + the original template, the executed run, and the archived copy + of the specification files based on the study's base file name. + """ shortened_filepath = self.filepath.replace(".out", "").replace(".partial", "").replace(".expanded", "") base_name = Path(shortened_filepath).stem self.special_vars["MERLIN_SPEC_ORIGINAL_TEMPLATE"] = os.path.join( @@ -145,16 +201,29 @@ def _set_special_file_vars(self): def write_original_spec(self): """ - Copy the original spec into merlin_info/ as '.orig.yaml'. + Copies the original specification file to the designated directory. + + This method copies the original specification file from its + current location to the `merlin_info/` directory, renaming it + to '.orig.yaml'. The base file name is derived + from the original specification's path. """ shutil.copyfile(self.original_spec.path, self.special_vars["MERLIN_SPEC_ORIGINAL_TEMPLATE"]) def label_clash_error(self): """ - Detect any illegal clashes between merlin's - merlin -> samples -> column_labels and Maestro's - global.parameters. Raises an error if any such - clash exists. + Detects illegal clashes between Merlin's sample column labels and + [Maestro's global parameters](https://maestrowf.readthedocs.io/en/latest/Maestro/specification.html#parameters-globalparameters). + + This method checks for any conflicts between the column labels + defined in the `merlin` section of the original specification and + the global parameters defined in the same specification. If a + column label is found to also exist in the global parameters, + a ValueError is raised to indicate the clash. + + Raises: + ValueError: If any column label in `merlin.samples.column_labels` + is also found in `merlin.globals`, indicating an illegal clash. """ if self.original_spec.merlin["samples"]: for label in self.original_spec.merlin["samples"]["column_labels"]: @@ -165,10 +234,24 @@ def label_clash_error(self): # to not use the MerlinStudy object so we disable this pylint error # pylint: disable=duplicate-code @staticmethod - def get_user_vars(spec): + def get_user_vars(spec: MerlinSpec) -> Dict[str, str]: """ - Using the spec environment, return a dictionary - of expanded user-defined variables. + Retrieves and expands user-defined variables from the specification environment. + + This static method examines the provided specification's environment + to collect user-defined variables and labels. It constructs a list + of these variables and passes them to the `determine_user_variables` + function to obtain a dictionary of expanded variables. + + Args: + spec (spec.specification.MerlinSpec): The specification object containing the environment from which + to extract user-defined variables. The environment should have keys + "variables" and/or "labels" that contain the relevant data. + + Returns: + A dictionary of expanded user-defined variables, where the keys + are variable names and the values are their corresponding + expanded values. """ uvars = [] if "variables" in spec.environment: @@ -180,14 +263,35 @@ def get_user_vars(spec): # pylint: enable=duplicate-code @property - def user_vars(self): - """Get the user defined variables""" + def user_vars(self) -> Dict[str, str]: + """ + Retrieves the user-defined variables for the study. + + This property accesses the original specification of the study and + retrieves the user-defined variables using the `get_user_vars` + method from this class. + + Returns: + A dictionary containing the user-defined variables + associated with the study. + """ return MerlinStudy.get_user_vars(self.original_spec) - def get_expanded_spec(self): + def get_expanded_spec(self) -> MerlinSpec: """ - Get a new yaml spec file with defaults, cli overrides, and variable expansions. - Useful for provenance. + Generates a new YAML specification file with applied defaults, + command-line interface (CLI) overrides, and variable expansions. + + This method creates a modified version of the original specification + by incorporating default values and user-defined overrides from the + command line. It also expands user-defined variables and reserved + words to produce a fully resolved specification. This is particularly + useful for tracking provenance and ensuring that the specification + accurately reflects all applied configurations. + + Returns: + (spec.specification.MerlinSpec): A new instance of the [`MerlinSpec`][spec.specification.MerlinSpec] + class that contains the fully expanded specification. """ # get specification including defaults and cli-overridden user variables new_env = replace_override_vars(self.original_spec.environment, self.override_vars) @@ -204,63 +308,98 @@ def get_expanded_spec(self): return expand_env_vars(result) @property - def samples(self): + def samples(self) -> np.ndarray: """ - Return this study's corresponding samples. + Retrieves the samples associated with this study. + + This property checks if there are any samples defined in the + expanded specification of the study. If samples are present, + it loads and returns them; otherwise, it returns an empty list. - :return: list of samples + Returns: + A numpy array of samples corresponding to the study. + If no samples are defined, an empty list is returned. """ if self.expanded_spec.merlin["samples"]: return self.load_samples() return [] - def get_sample_labels(self, from_spec): - """Return the column labels of the samples (if any)""" + def get_sample_labels(self, from_spec: MerlinSpec) -> List[str]: + """ + Retrieves the column labels of the samples from the provided specification. + + This method checks the specified [`MerlinSpec`][spec.specification.MerlinSpec] + object for sample information and returns the associated column labels if they + exist. If no sample labels are found, an empty list is returned. + + Args: + from_spec (spec.specification.MerlinSpec): The specification object + from which to extract sample column labels. It is expected to contain + a "samples" key within its "merlin" dictionary. + + Returns: + A list of column labels for the samples. If no sample labels are + present, an empty list is returned. + """ if from_spec.merlin["samples"]: return from_spec.merlin["samples"]["column_labels"] return [] @property - def sample_labels(self): + def sample_labels(self) -> List[str]: """ - Return this study's corresponding sample labels + Retrieves the labels of the samples associated with this study. + + This property extracts the sample labels from the study's + expanded specification. It returns a list of labels that + correspond to the samples defined in the specification. - Example spec_file contents: + Returns: + A list of sample labels. If no labels are defined, an empty list is returned. - --spec_file.yaml-- - ... - merlin: - samples: - column_labels: [X0, X1] + Example: + Given the following contents in a specification file: - :return: list of labels (e.g. ["X0", "X1"] ) + ```yaml + merlin: + samples: + column_labels: [X0, X1] + ``` + + This property would return: `["X0", "X1"]` """ return self.get_sample_labels(from_spec=self.expanded_spec) - def load_samples(self): + def load_samples(self) -> np.ndarray: """ - load this study's samples from disk, generating if the file does - not yet exist and the file is defined in the YAML file. - (no generation will occur if file is defined via __init__) + Loads the study's samples from disk, generating them if the file + does not exist and is defined in the YAML specification. - Runs the function defined in 'generate' and then loads up - the sample files defined in 'file', assigning them to the - variables in 'column_labels' + This method checks if a sample file is specified in the expanded + specification. If the file does not exist, it will invoke the + generation command defined in the 'generate' section of the + specification to create the sample file. Once the file is available, + it loads the samples into a NumPy array and assigns them to the + variables specified in 'column_labels'. - Example spec_file contents: + Returns: + A NumPy array containing the loaded samples. The shape of the + array will be (n_samples, n_features), where n_samples is + the number of samples loaded and n_features is the number + of features corresponding to the column labels. - --spec_file.yaml-- - ... - merlin: - samples: - generate: - cmd: python make_samples.py -outfile=samples.npy - file: samples.npy - column_labels: [X0, X1] + Example: + The spec file contents will look something like: - :return: numpy samples - :return: the samples loaded + ```yaml + merlin: + samples: + generate: + cmd: python make_samples.py -outfile=samples.npy + file: samples.npy + column_labels: [X0, X1] + ``` """ if self.samples_file is None: if self.expanded_spec.merlin["samples"]: @@ -287,18 +426,41 @@ def load_samples(self): return samples @property - def level_max_dirs(self): + def level_max_dirs(self) -> int: """ - Returns the maximum number of directory levels. + Retrieves the maximum number of directory levels for sample organization. + + This property checks the expanded specification for the maximum + number of directory levels defined under the 'merlin' section. + If the value is not found, it falls back to a default value + specified in the `defaults.SAMPLES` dictionary. + + Returns: + The maximum number of directory levels. If the value is + not specified in the expanded specification, the default + value from `defaults.SAMPLES["level_max_dirs"]` is returned. """ with suppress(TypeError, KeyError): return self.expanded_spec.merlin["samples"]["level_max_dirs"] return defaults.SAMPLES["level_max_dirs"] @cached_property - def output_path(self): + def output_path(self) -> str: """ Determines and creates an output directory for this study. + + This property checks if a restart directory is specified. If so, it validates + the existence of the directory and returns its absolute path. If no restart + directory is provided, it constructs the output path based on the original + specification and any override variables. The output path is expanded to + include user-defined variables and environment variables. If the directory + does not exist, it is created. + + Returns: + The absolute path to the output directory for the study. + + Raises: + ValueError: If the specified restart directory does not exist. """ if self.restart_dir is not None: output_path = self.restart_dir @@ -331,10 +493,18 @@ def output_path(self): return output_path @cached_property - def timestamp(self): + def timestamp(self) -> str: """ - Returns a timestamp string, representing the time this - study began. May be used as an id or unique identifier. + Returns a timestamp string representing the time this study began. + + This property generates a unique identifier based on the current time + when the study is initiated. If a restart directory is specified, it + extracts a substring from the directory name as the timestamp. Otherwise, + it formats the current time in the 'YYYYMMDD-HHMMSS' format. + + Returns: + A string representing the timestamp of the study's initiation, + which can be used as an identifier or unique key. """ if self.restart_dir is not None: return self.restart_dir.strip("/")[-15:] @@ -343,12 +513,22 @@ def timestamp(self): # TODO look into why pylint complains that this method is hidden # - might be because we reset self.workspace's value in the expanded_spec method @cached_property - def workspace(self): # pylint: disable=E0202 + def workspace(self) -> str: # pylint: disable=E0202 """ - Determines, makes, and returns the path to this study's - workspace directory. This directory holds workspace directories - for each step in the study, as well as 'merlin_info/'. The - name of this directory ends in a timestamp. + Determines, creates, and returns the path to this study's workspace directory. + + This property generates a unique workspace directory for the study, which + contains subdirectories for each step of the study and a 'merlin_info/' + directory. The name of the workspace directory is derived from the original + specification name and includes a timestamp to ensure uniqueness. If a + restart directory is specified, it validates the existence of the directory + and returns its absolute path. + + Returns: + The absolute path to the workspace directory for the study. + + Raises: + ValueError: If the specified restart directory does not exist. """ if self.restart_dir is not None: if not os.path.isdir(self.restart_dir): @@ -366,9 +546,17 @@ def workspace(self): # pylint: disable=E0202 # TODO look into why pylint complains that this method is hidden # - might be because we reset self.info's value in the expanded_spec method @cached_property - def info(self): # pylint: disable=E0202 + def info(self) -> str: # pylint: disable=E0202 """ - Creates the 'merlin_info' directory inside this study's workspace directory. + Creates and returns the path to the 'merlin_info' directory within the study's workspace. + + This property checks if a restart directory is specified. If not, it creates + the 'merlin_info' directory inside the study's workspace directory. This + directory is intended to store metadata and other relevant information related + to the study. + + Returns: + The absolute path to the 'merlin_info' directory. """ info_name = os.path.join(self.workspace, "merlin_info") if self.restart_dir is None: @@ -376,10 +564,22 @@ def info(self): # pylint: disable=E0202 return info_name @cached_property - def expanded_spec(self): + def expanded_spec(self) -> MerlinSpec: """ - Determines, writes to yaml, and loads into memory an expanded - specification. + Determines, writes to YAML, and loads into memory an expanded specification. + + This property handles the expansion of the study's specification based on + the original specification and any provided environment variables. If the + study is being restarted, it retrieves the previously expanded specification + without re-expanding it. Otherwise, it processes the original specification, + expands any tokens or shell references, and updates paths accordingly. + + Returns: + (spec.specification.MerlinSpec): The expanded specification object. + + Raises: + ValueError: If the expanded name for the workspace contains invalid + characters for a filename. """ # If we are restarting, we don't need to re-expand, just need to read # in the previously expanded spec @@ -448,9 +648,17 @@ def expanded_spec(self): return result @cached_property - def flux_command(self): + def flux_command(self) -> str: """ - Returns the flux command, this will include the full path, if flux_path given in the workflow. + Returns the full path to the flux command based on the specified workflow configuration. + + This property constructs the command to execute the flux binary. If a + `flux_path` is provided in the expanded specification's batch configuration, + it will use that path to create the full command. Otherwise, it defaults + to the standard 'flux' command. + + Returns: + The complete command string for executing flux. """ flux_bin = "flux" if "flux_path" in self.expanded_spec.batch.keys(): @@ -459,18 +667,24 @@ def flux_command(self): def generate_samples(self): """ - Runs the function defined in 'generate' if self.samples_file is not - yet a file. + Generates sample data by executing the command defined in the + 'generate' section of the specification if the sample file does + not already exist. - Example spec_file contents: + This method checks if the specified sample file exists. If it + does not, it retrieves the command from the YAML specification + and executes it using a subprocess. The output and error logs + from the command execution are saved to files for later review. - --spec_file.yaml-- - ... - merlin: - samples: - generate: - cmd: python make_samples.py -outfile=samples.npy + Example: + Here's an example sample generation command: + ```yaml + merlin: + samples: + generate: + cmd: python make_samples.py -outfile=samples.npy + ``` """ try: if not os.path.exists(self.samples_file): @@ -495,8 +709,33 @@ def generate_samples(self): LOG.error(f"Could not generate samples:\n{e}") return - def load_pgen(self, filepath, pargs, env): - """Creates a dict of variable names and values defined in a pgen script""" + def load_pgen(self, filepath: str, pargs: List[str], env: StudyEnvironment) -> Dict[str, Dict[str, str]]: + """ + Loads a parameter generator script and creates a dictionary of + variable names and their corresponding values. + + This method reads a parameter generator script from the specified + file path and extracts variable names and values defined within + the script. It constructs a dictionary where each key is a + variable name, and the value is another dictionary containing + the variable's label and its associated values. + + Args: + filepath: The path to the parameter generator script to be loaded. + pargs: A list of additional arguments to be passed to the parameter + generator. If None, an empty list will be used. + env: A Maestro + [`StudyEnvironment`](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/datastructures/core/studyenvironment.html) + object containing custom information. + + Returns: + A dictionary where each key is a variable name and each value + is a dictionary containing:\n + - `values`: The values associated with the variable, + or None if not defined. + - `label`: The label of the variable as defined in the + parameter generator script. + """ if filepath: if pargs is None: pargs = [] @@ -512,8 +751,20 @@ def load_pgen(self, filepath, pargs, env): def load_dag(self): """ - Generates a dag (a directed acyclic execution graph). - Assigns it to `self.dag`. + Generates a Directed Acyclic Graph (DAG) for the execution of + the study and assigns it to the `self.dag` attribute. + + This method constructs a DAG based on the specifications defined + in the expanded study specification. It retrieves the study + environment, steps, and parameters, and initializes a Maestro + [`Study`](https://maestrowf.readthedocs.io/en/latest/Maestro/reference_guide/api_reference/datastructures/core/study.html) + object. The method then sets up the workspace and environment + for the study, configures it, and generates the DAG using the + Maestro framework. + + The generated DAG contains the execution flow of the study, + ensuring that all steps are executed in the correct order + without cycles. """ environment = self.expanded_spec.get_study_environment() steps = self.expanded_spec.get_study_steps() @@ -555,8 +806,25 @@ def load_dag(self): # To avoid pickling issues with _pass_detect_cycle from maestro, we unpack the dag here self.dag = DAG(maestro_dag.adjacency_table, maestro_dag.values, column_labels, study.name, parameter_info) - def get_adapter_config(self, override_type=None): - """Builds and returns the adapter configuration dictionary""" + def get_adapter_config(self, override_type: str = None) -> Dict[str, str]: + """ + Builds and returns the adapter configuration dictionary. + + This method constructs a configuration dictionary for the adapter + based on the specifications defined in `self.expanded_spec.batch`. + It ensures that the configuration includes a type, which can be + overridden if specified. The method also checks for a dry run + flag and adds relevant commands if the batch type is set to + "flux". + + Args: + override_type: An optional string to override the default adapter + type. If not provided, the type from the expanded specification + will be used. + + Returns: + A dictionary containing the adapter configuration. + """ adapter_config = dict(self.expanded_spec.batch) if "type" not in adapter_config.keys(): @@ -579,10 +847,16 @@ def get_adapter_config(self, override_type=None): return adapter_config @property - def parameter_labels(self): + def parameter_labels(self) -> List[str]: """ - Get the parameter labels for this study. - :returns: A list of parameter labels used in this study + Retrieves the parameter labels associated with this study. + + This property extracts parameter labels from the expanded specification + of the study. It accesses the parameters and their associated metadata, + collecting all labels defined for each parameter. + + Returns: + A list of parameter labels used in this study. """ parameters = self.expanded_spec.get_parameters() metadata = parameters.get_metadata() diff --git a/merlin/utils.py b/merlin/utils.py index 2e69f5779..093b505d8 100644 --- a/merlin/utils.py +++ b/merlin/utils.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -34,6 +34,7 @@ import getpass import logging import os +import pickle import re import socket import subprocess @@ -42,7 +43,7 @@ from copy import deepcopy from datetime import datetime, timedelta from types import SimpleNamespace -from typing import Callable, List, Optional, Union +from typing import Any, Callable, Dict, Generator, List, Tuple, Union import numpy as np import pkg_resources @@ -51,24 +52,30 @@ from tabulate import tabulate -try: - import cPickle as pickle -except ImportError: - import pickle - - LOG = logging.getLogger(__name__) ARRAY_FILE_FORMATS = ".npy, .csv, .tab" DEFAULT_FLUX_VERSION = "0.48.0" -def get_user_process_info(user=None, attrs=None): +def get_user_process_info(user: str = None, attrs: List[str] = None) -> List[Dict]: """ - Return a list of process info for all of the user's running processes. + Return a list of process information for all of the user's running processes. + + This function retrieves and returns details about the currently running processes + for a specified user. If no user is specified, it defaults to the current user. + It can also return information for all users if specified. + + Args: + user: The username for which to retrieve process information. If set to + 'all_users', retrieves processes for all users. Defaults to the current + user's username if not provided. + attrs: A list of attributes to include in the process information. Defaults + to ["pid", "name", "username", "cmdline"] if None. If "username" is not + included in the list, it will be added. - :param `user`: user name (default from getpass). Option: 'all_users': get - all processes - :param `atts`: the attributes to include + Returns: + A list of dictionaries containing the specified attributes for each process + belonging to the specified user or all users if 'all_users' is specified. """ if attrs is None: attrs = ["pid", "name", "username", "cmdline"] @@ -84,13 +91,23 @@ def get_user_process_info(user=None, attrs=None): return [p.info for p in psutil.process_iter(attrs=attrs) if user in p.info["username"]] -def check_pid(pid, user=None): +def check_pid(pid: int, user: str = None) -> bool: """ - Check if pid is in process list. + Check if a given process ID (PID) is in the process list for a specified user. - :param `pid`: process id - :param `user`: user name (default from getpass). Option: 'all_users': get - all processes + This function determines whether a specific PID is currently running + for the specified user. If no user is specified, it defaults to the + current user. It can also check for all users if specified. + + Args: + pid: The process ID to check for in the process list. + user: The username for which to check the process. If set to 'all_users', + checks processes for all users. Defaults to the current user's username + if not provided. + + Returns: + True if the specified PID is found in the process list for the given user, + False otherwise. """ user_processes = get_user_process_info(user=user) for process in user_processes: @@ -99,13 +116,23 @@ def check_pid(pid, user=None): return False -def get_pid(name, user=None): +def get_pid(name: str, user: str = None) -> List[int]: """ - Return pid of process with name. + Return the process ID(s) (PID) of processes with the specified name. + + This function retrieves the PID(s) of all running processes that match + the given name for a specified user. If no user is specified, it defaults + to the current user. It can also retrieve PIDs for all users if specified. - :param `name`: process name - :param `user`: user name (default from getpass). Option: 'all_users': get - all processes + Args: + name: The name of the process to search for. + user: The username for which to retrieve the process IDs. If set to + 'all_users', retrieves processes for all users. Defaults to the + current user's username if not provided. + + Returns: + A list of PIDs for processes matching the specified name. Returns None + if no matching processes are found. """ user_processes = get_user_process_info(user=user) name_list = [p["pid"] for p in user_processes if name in p["name"]] @@ -114,37 +141,68 @@ def get_pid(name, user=None): return None -def get_procs(name, user=None): +def get_procs(name: str, user: str = None) -> List[Tuple[int, str]]: """ - Return a list of (pid, cmdline) tuples of process with name. + Return a list of tuples containing the process ID (PID) and command line + of processes with the specified name. + + This function retrieves all running processes that match the given name + for a specified user. If no user is specified, it defaults to the current + user. It can also retrieve processes for all users if specified. - :param `name`: process name - :param `user`: user name (default from getpass). Option: 'all_users': get - all processes + Args: + name: The name of the process to search for. + user: The username for which to retrieve the process information. + If set to 'all_users', retrieves processes for all users. + Defaults to the current user's username if not provided. + + Returns: + A list of tuples, each containing the PID and command line of processes + matching the specified name. Returns an empty list if no matching + processes are found. """ user_processes = get_user_process_info(user=user) procs = [(p["pid"], p["cmdline"]) for p in user_processes if name in p["name"]] return procs -def is_running_psutil(cmd, user=None): +def is_running_psutil(cmd: str, user: str = None) -> bool: """ - Determine if process with given command is running. - Uses psutil command instead of call to 'ps' + Determine if a process with the given command is currently running. + + This function checks for the existence of any running processes that + match the specified command. It uses the `psutil` library to gather + process information instead of making a call to the 'ps' command. - :param `cmd`: process cmd - :param `user`: user name (default from getpass). Option: 'all_users': get - all processes + Args: + cmd: The command or command line snippet to search for in running + processes. + user: The username for which to check running processes. If set to + 'all_users', checks processes for all users. Defaults to the + current user's username if not provided. + + Returns: + True if at least one matching process is found; otherwise, False. """ user_processes = get_user_process_info(user=user) return any(cmd in " ".join(p["cmdline"]) for p in user_processes) -def is_running(name, all_users=False): +def is_running(name: str, all_users: bool = False) -> bool: """ - Determine if process with name is running. + Determine if a process with the specified name is currently running. + + This function checks for the existence of a running process with the + provided name by executing the 'ps' command. It can be configured to + check processes for all users or just the current user. + + Args: + name: The name of the process to search for. + all_users: If True, checks for processes across all users. Defaults + to False, which checks only the current user's processes. - :param `name`: process name + Returns: + True if a process with the specified name is found; otherwise, False. """ cmd = ["ps", "ux"] @@ -164,23 +222,42 @@ def is_running(name, all_users=False): return False -def expandvars2(path): +def expandvars2(path: str) -> str: """ - Replace shell strings from the current environment variables + Replace shell variables in the given path with their corresponding + environment variable values. - :param `path`: a path + This function expands shell-style variable references (e.g., $VAR) + in the input path using the current environment variables. It also + ensures that any escaped dollar signs (e.g., \\$) are not expanded. + + Args: + path: The input path containing shell variable references to be expanded. + + Returns: + The path with shell variables replaced by their corresponding values + from the environment, with unescaped variables expanded. """ return re.sub(r"(? List[str]: """ - Apply a regex filter to a list + Apply a regex filter to a list. - :param `regex` : the regular expression - :param `list_to_filter` : the list to filter + This function filters a given list based on a specified regular expression. + Depending on the `match` parameter, it can either match the entire string + or search for the regex pattern within the strings of the list. - :return `new_list` + Args: + regex: The regular expression to use for filtering the list. + list_to_filter: The list of strings to be filtered based on the regex. + match: If True, uses re.match to filter items that match the regex from + the start. If False, uses re.search to filter items that contain the + regex pattern. + + Returns: + A new list containing the filtered items that match the regex. """ r = re.compile(regex) # pylint: disable=C0103 if match: @@ -188,16 +265,27 @@ def regex_list_filter(regex, list_to_filter, match=True): return list(filter(r.search, list_to_filter)) -def apply_list_of_regex(regex_list, list_to_filter, result_list, match=False, display_warning: bool = True): +def apply_list_of_regex( + regex_list: List[str], list_to_filter: List[str], result_list: List[str], match: bool = False, display_warning: bool = True +): """ - Take a list of regex's, apply each regex to a list we're searching through, - and append each result to a result list. + Apply a list of regex patterns to a list and accumulate the results. + + This function takes each regex from the provided list of regex patterns + and applies it to the specified list. The results of each successful + match or search are appended to a result list. Optionally, it can display + a warning if a regex does not match any item in the list. - :param `regex_list`: A list of regular expressions to apply to the list_to_filter - :param `list_to_filter`: A list that we'll apply regexs to - :param `result_list`: A list that we'll append results of the regex filters to - :param `match`: A bool where when true we use re.match for applying the regex, - when false we use re.search for applying the regex. + Args: + regex_list: A list of regular expressions to apply to the `list_to_filter`. + list_to_filter: The list of strings that the regex patterns will be applied to. + result_list: The list where results of the regex filters will be appended. + match: If True, uses re.match for applying the regex. If False, uses re.search. + display_warning: If True, displays a warning message when no matches are + found for a regex. + + Side Effect: + This function modifies the `result_list` in place. """ for regex in regex_list: filter_results = set(regex_list_filter(regex, list_to_filter, match)) @@ -209,28 +297,37 @@ def apply_list_of_regex(regex_list, list_to_filter, result_list, match=False, di result_list += filter_results -def load_yaml(filepath): +def load_yaml(filepath: str) -> Dict: """ - Safely read a yaml file. + Safely read a YAML file and return its contents. - :param `filepath`: a filepath to a yaml file - :type filepath: str + Args: + filepath: The file path to the YAML file to be read. - :returns: Python objects holding the contents of the yaml file + Returns: + A dict representing the contents of the YAML file. """ with open(filepath, "r") as _file: return yaml.safe_load(_file) -def get_yaml_var(entry, var, default): +def get_yaml_var(entry: Dict[str, Any], var: str, default: Any) -> Any: """ - Return entry[var], else return default + Retrieve the value associated with a specified key from a YAML dictionary. - :param `entry`: a yaml dict - :param `var`: a yaml key - :param `default`: default value in the absence of data - """ + This function attempts to return the value of `var` from the provided `entry` + dictionary. If the key does not exist, it will try to access it as an attribute + of the entry object. If neither is found, the function returns the specified + `default` value. + + Args: + entry: A dictionary representing the contents of a YAML file. + var: The key or attribute name to retrieve from the entry. + default: The default value to return if the key or attribute is not found. + Returns: + The value associated with `var` in the entry, or `default` if not found. + """ try: return entry[var] except (TypeError, KeyError): @@ -240,19 +337,31 @@ def get_yaml_var(entry, var, default): return default -def load_array_file(filename, ndmin=2): +def load_array_file(filename: str, ndmin: int = 2) -> np.ndarray: """ - Loads up an array stored in filename, based on extension. + Load an array from a file based on its extension. - Valid filename extensions: - '.npy' : numpy binary file - '.csv' : comma separated text file - '.tab' : whitespace (or tab) separated text file + This function reads an array stored in the specified `filename`. + It supports three file types based on their extensions: - :param `filename` : The file to load - :param `ndmin` : The minimum number of dimensions to load - """ + - `.npy` for NumPy binary files + - `.csv` for comma-separated values + - `.tab` for whitespace (or tab) separated values + + The function ensures that the loaded array has at least `ndmin` dimensions. + If the array is in binary format, it checks the dimensions without altering the data. + + Args: + filename: The path to the file to load. + ndmin: The minimum number of dimensions the array should have. + Returns: + The loaded array. + + Raises: + TypeError: If the file extension is not one of the supported types + (`.npy`, `.csv`, `.tab`). + """ protocol = determine_protocol(filename) # Don't change binary-stored numpy arrays; just check dimensions @@ -277,9 +386,18 @@ def load_array_file(filename, ndmin=2): return array -def determine_protocol(fname): +def determine_protocol(fname: str) -> str: """ - Determines a file protocol based on file name extension. + Determine the file protocol based on the file name extension. + + Args: + fname: The name of the file whose protocol is to be determined. + + Returns: + The protocol corresponding to the file extension (e.g., 'hdf5'). + + Raises: + ValueError: If the provided file name does not have a valid extension. """ _, ext = os.path.splitext(fname) if ext.startswith("."): @@ -294,13 +412,21 @@ def determine_protocol(fname): def verify_filepath(filepath: str) -> str: """ - Verify that the filepath argument is a valid - file. + Verify that the given file path is valid and return its absolute form. + + This function checks if the specified `filepath` points to an existing file. + It expands any user directory shortcuts (e.g., `~`) and environment variables + in the provided path before verifying its existence. If the file does not exist, + a ValueError is raised. + + Args: + filepath: The path of the file to verify. - :param [str] `filepath`: the path of a file + Returns: + The verified absolute file path with expanded environment variables. - :return: the verified absolute filepath with expanded environment variables. - :rtype: str + Raises: + ValueError: If the provided file path does not point to a valid file. """ filepath = os.path.abspath(os.path.expandvars(os.path.expanduser(filepath))) if not os.path.isfile(filepath): @@ -310,13 +436,21 @@ def verify_filepath(filepath: str) -> str: def verify_dirpath(dirpath: str) -> str: """ - Verify that the dirpath argument is a valid - directory. + Verify that the given directory path is valid and return its absolute form. - :param [str] `dirpath`: the path of a directory + This function checks if the specified `dirpath` points to an existing directory. + It expands any user directory shortcuts (e.g., `~`) and environment variables + in the provided path before verifying its existence. If the directory does not exist, + a ValueError is raised. - :return: returns the absolute path with expanded environment vars for a given dirpath. - :rtype: str + Args: + dirpath: The path of the directory to verify. + + Returns: + The verified absolute directory path with expanded environment variables. + + Raises: + ValueError: If the provided directory path does not point to a valid directory. """ dirpath: str = os.path.abspath(os.path.expandvars(os.path.expanduser(dirpath))) if not os.path.isdir(dirpath): @@ -325,9 +459,19 @@ def verify_dirpath(dirpath: str) -> str: @contextmanager -def cd(path): # pylint: disable=C0103 +def cd(path: str) -> Generator[None, None, None]: # pylint: disable=C0103 """ - TODO + Context manager for changing the current working directory. + + This context manager changes the current working directory to the specified `path` + while executing the block of code within the context. Once the block is exited, + it restores the original working directory. + + Args: + path: The path to the directory to change to. + + Yields: + Control is yielded back to the block of code within the context. """ old_dir = os.getcwd() os.chdir(path) @@ -337,15 +481,37 @@ def cd(path): # pylint: disable=C0103 os.chdir(old_dir) -def pickle_data(filepath, content): - """Dump content to a pickle file""" +def pickle_data(filepath: str, content: Any): + """ + Dump content to a pickle file. + + This function serializes the given `content` and writes it to a specified file + in pickle format. The file is opened in write mode, which will overwrite any + existing content in the file. + + Args: + filepath: The path to the file where the content will be saved. + content: The data to be serialized and saved to the pickle file. + """ with open(filepath, "w") as f: # pylint: disable=C0103 pickle.dump(content, f) -def get_source_root(filepath): - """Used to find the absolute project path given a sample file path from - within the project. +def get_source_root(filepath: str) -> str: + """ + Find the absolute project path given a file path from within the project. + + This function determines the root directory of a project by analyzing the given + file path. It works by traversing the directory structure upwards until it + encounters a directory name that is not an integer, which is assumed to be the + project root. + + Args: + filepath: The file path from within the project for which to find the root. + + Returns: + The absolute path to the root directory of the project. Returns None if + the path corresponds to the root directory itself. """ filepath = os.path.abspath(filepath) sep = os.path.sep @@ -367,9 +533,19 @@ def get_source_root(filepath): return root -def ensure_directory_exists(**kwargs): +def ensure_directory_exists(**kwargs: Dict[Any, Any]) -> bool: """ - TODO + Ensure that the directory for the specified aggregate file exists. + + This function checks if the directory for the given `aggregate_file` exists. + If it does not exist, the function creates the necessary directories. + + Args: + **kwargs: Keyword arguments that must include:\n + - `aggregate_file` (str): The file path for which the directory needs to be ensured. + + Returns: + True if the directory already existed. False otherwise. """ aggregate_bundle = kwargs["aggregate_file"] dirname = os.path.dirname(aggregate_bundle) @@ -381,9 +557,23 @@ def ensure_directory_exists(**kwargs): return True -def nested_dict_to_namespaces(dic): - """Code for recursively converting dictionaries of dictionaries - into SimpleNamespaces instead. +def nested_dict_to_namespaces(dic: Dict) -> SimpleNamespace: + """ + Convert a nested dictionary into a nested SimpleNamespace structure. + + This function recursively transforms a dictionary (which may contain other + dictionaries) into a structure of SimpleNamespace objects. Each key in the + dictionary becomes an attribute of a SimpleNamespace, allowing for attribute-style + access to the data. + + Args: + dic: The nested dictionary to be converted. + + Returns: + A SimpleNamespace object representing the nested structure of the input dictionary. + + Raises: + TypeError: If the input is not a dictionary. """ def recurse(dic): @@ -400,9 +590,23 @@ def recurse(dic): return recurse(new_dic) -def nested_namespace_to_dicts(namespaces): - """Code for recursively converting namespaces of namespaces - into dictionaries instead. +def nested_namespace_to_dicts(namespaces: SimpleNamespace) -> Dict: + """ + Convert a nested SimpleNamespace structure into a nested dictionary. + + This function recursively transforms a SimpleNamespace (which may contain + other SimpleNamespaces) into a dictionary structure. Each attribute of the + SimpleNamespace becomes a key in the resulting dictionary. + + Args: + namespaces: The nested SimpleNamespace to be converted. + + Returns: + A dictionary representing the nested structure of the input + SimpleNamespace. + + Raises: + TypeError: If the input is not a SimpleNamespace. """ def recurse(namespaces): @@ -419,12 +623,28 @@ def recurse(namespaces): return recurse(new_ns) -def get_flux_version(flux_path, no_errors=False): +def get_flux_version(flux_path: str, no_errors: bool = False) -> str: """ - Return the flux version as a string + Retrieve the version of Flux as a string. + + This function executes the Flux binary located at `flux_path` with the + "version" command and parses the output to return the version number. + If the command fails or the Flux binary cannot be found, it can either + raise an error or return a default version based on the `no_errors` flag. - :param `flux_path`: the full path to the flux bin - :param `no_errors`: a flag to determine if this a test run to ignore errors + Args: + flux_path: The full path to the Flux binary. + no_errors: A flag to suppress error messages and exceptions. If set to + True, errors will be logged but not raised. + + Returns: + The version of Flux as a string. + + Raises: + FileNotFoundError: If the Flux binary cannot be found and `no_errors` + is set to False. + ValueError: If the version cannot be determined from the output and + `no_errors` is set to False. """ cmd = [flux_path, "version"] @@ -451,12 +671,23 @@ def get_flux_version(flux_path, no_errors=False): return flux_ver -def get_flux_cmd(flux_path, no_errors=False): +def get_flux_cmd(flux_path: str, no_errors: bool = False) -> str: """ - Return the flux run command as string + Generate the Flux run command based on the installed version. + + This function determines the appropriate Flux command to use for + running jobs, depending on the version of Flux installed at the + specified `flux_path`. It defaults to "flux run" for versions + greater than or equal to 0.48.x. For older versions, it adjusts + the command accordingly. - :param `flux_path`: the full path to the flux bin - :param `no_errors`: a flag to determine if this a test run to ignore errors + Args: + flux_path: The full path to the Flux binary. + no_errors: A flag to suppress error messages and exceptions + if set to True. + + Returns: + The appropriate Flux run command as a string. """ # The default is for flux version >= 0.48.x # this may change in the future. @@ -474,12 +705,23 @@ def get_flux_cmd(flux_path, no_errors=False): return flux_cmd -def get_flux_alloc(flux_path, no_errors=False): +def get_flux_alloc(flux_path: str, no_errors: bool = False) -> str: """ - Return the flux alloc command as string + Generate the `flux alloc` command based on the installed version. + + This function constructs the appropriate command for allocating + resources with Flux, depending on the version of Flux installed + at the specified `flux_path`. It defaults to "{flux_path} alloc" + for versions greater than or equal to 0.48.x. For older versions, + it adjusts the command accordingly. - :param `flux_path`: the full path to the flux bin - :param `no_errors`: a flag to determine if this a test run to ignore errors + Args: + flux_path: The full path to the Flux binary. + no_errors: A flag to suppress error messages and exceptions + if set to True. + + Returns: + The appropriate Flux allocation command as a string. """ # The default is for flux version >= 0.48.x # this may change in the future. @@ -495,12 +737,21 @@ def get_flux_alloc(flux_path, no_errors=False): return flux_alloc -def check_machines(machines): +def check_machines(machines: Union[str, List[str], Tuple[str]]) -> bool: """ - Return a True if the current machine is in the list of machines. + Check if the current machine is in the list of specified machines. + + This function determines whether the hostname of the current + machine matches any entry in a provided list of machine names. + It returns True if a match is found, otherwise it returns False. - :param `machines`: A single machine or list of machines to compare - with the current machine. + Args: + machines: A single machine name or a list/tuple of machine + names to compare with the current machine's hostname. + + Returns: + True if the current machine's hostname matches any of the + specified machines; False otherwise. """ local_hostname = socket.gethostname() @@ -514,36 +765,67 @@ def check_machines(machines): return False -def contains_token(string): +def contains_token(string: str) -> bool: """ - Return True if given string contains a token of the form $(STR). + Check if the given string contains a token of the form $(STR). + + This function uses a regular expression to search for tokens + that match the pattern $(), where consists of + alphanumeric characters and underscores. It returns True if + such a token is found; otherwise, it returns False. + + Args: + string: The input string to be checked for tokens. + + Returns: + True if the input string contains a token of the form + $(STR); False otherwise. """ if re.search(r"\$\(\w+\)", string): return True return False -def contains_shell_ref(string): +def contains_shell_ref(string: str) -> bool: """ - Return True if given string contains a shell variable reference - of the form $STR or ${STR}. + Check if the given string contains a shell variable reference. + + This function searches for shell variable references in the + format of $ or ${}, where + consists of alphanumeric characters and underscores. It returns + True if a match is found; otherwise, it returns False. + + Args: + string: The input string to be checked for shell + variable references. + + Returns: + True if the input string contains a shell variable + reference of the form $STR or ${STR}; False otherwise. """ if re.search(r"\$\w+", string) or re.search(r"\$\{\w+\}", string): return True return False -def needs_merlin_expansion( - cmd: str, restart_cmd: str, labels: List[str], include_sample_keywords: Optional[bool] = True -) -> bool: +def needs_merlin_expansion(cmd: str, restart_cmd: str, labels: List[str], include_sample_keywords: bool = True) -> bool: """ - Check if the cmd or restart cmd provided have variables that need expansion. + Check if the provided command or restart command contains variables that require expansion. - :param `cmd`: The command inside a study step to check for expansion - :param `restart_cmd`: The restart command inside a study step to check for expansion - :param `labels`: A list of labels to check for inside `cmd` and `restart_cmd` - :return : True if the cmd has any of the default keywords or spec - specified sample column labels. False otherwise. + This function checks both the command (`cmd`) and the restart command (`restart_cmd`) + for the presence of specified labels or sample keywords that indicate a need for variable + expansion. + + Args: + cmd: The command inside a study step to check for variable expansion. + restart_cmd: The restart command inside a study step to check for variable expansion. + labels: A list of labels to check for inside `cmd` and `restart_cmd`. + include_sample_keywords: Flag to indicate whether to include default sample keywords + in the label check. + + Returns: + True if either `cmd` or `restart_cmd` contains any of the specified labels + or default sample keywords, indicating a need for expansion. False otherwise. """ sample_keywords = ["MERLIN_SAMPLE_ID", "MERLIN_SAMPLE_PATH", "merlin_sample_id", "merlin_sample_path"] if include_sample_keywords: @@ -560,20 +842,28 @@ def needs_merlin_expansion( return False -def dict_deep_merge(dict_a: dict, dict_b: dict, path: str = None, conflict_handler: Callable = None): +def dict_deep_merge(dict_a: Dict, dict_b: Dict, path: str = None, conflict_handler: Callable = None): """ - This function recursively merges dict_b into dict_a. The built-in - merge of dictionaries in python (dict(dict_a) | dict(dict_b)) does not do a - deep merge so this function is necessary. This will only merge in new keys, - it will NOT update existing ones, unless you specify a conflict handler function. - Credit to this stack overflow post: https://stackoverflow.com/a/7205107. + Recursively merges `dict_b` into `dict_a`, performing a deep merge. + + This function combines two dictionaries by recursively merging + the contents of `dict_b` into `dict_a`. Unlike Python's built-in + dictionary merge, this function performs a deep merge, meaning + it will merge nested dictionaries instead of just updating top-level keys. + Existing keys in `dict_a` will not be updated unless a conflict handler + is provided to resolve key conflicts. - :param `dict_a`: A dict that we'll merge dict_b into - :param `dict_b`: A dict that we want to merge into dict_a - :param `path`: The path down the dictionary tree that we're currently at - :param `conflict_handler`: An optional function to handle conflicts between values at the same key. - The function should return the value to be used in the merged dictionary. - The default behavior without this argument is to log a warning. + Credit to [this stack overflow post](https://stackoverflow.com/a/7205107). + + Args: + dict_a: The dictionary that will be merged into. + dict_b: The dictionary to merge into `dict_a`. + path: The current path in the dictionary tree. This is used for logging + purposes during recursion. + conflict_handler: A function to handle conflicts when both dictionaries + have the same key with different values. The function should return + the value to be used in the merged dictionary. If not provided, a + warning will be logged for conflicts. """ # Check to make sure we have valid dict_a and dict_b input @@ -609,15 +899,28 @@ def dict_deep_merge(dict_a: dict, dict_b: dict, path: str = None, conflict_handl dict_a[key] = dict_b[key] -def find_vlaunch_var(vlaunch_var: str, step_cmd: str, accept_no_matches=False) -> str: +def find_vlaunch_var(vlaunch_var: str, step_cmd: str, accept_no_matches: bool = False) -> str: """ - Given a variable used for VLAUNCHER and the step cmd value, find - the variable. + Find and return the specified VLAUNCHER variable from the step command. + + This function searches for a variable defined in the VLAUNCHER context + within the provided step command string. It looks for the variable in + the format `MERLIN_=`. If the variable is found, + it returns the variable in a format suitable for use in a command string. + If the variable is not found, the behavior depends on the `accept_no_matches` flag. + + Args: + vlaunch_var: The name of the VLAUNCHER variable (without the prefix 'MERLIN_'). + step_cmd: The command string of a step where the variable may be defined. + accept_no_matches: If True, returns None if the variable is not found. + If False, raises a ValueError. Defaults to False. - :param `vlaunch_var`: The name of the VLAUNCHER variable (without MERLIN_) - :param `step_cmd`: The string for the cmd of a step - :param `accept_no_matches`: If True, return None if we couldn't find the variable. Otherwise, raise an error. - :returns: the `vlaunch_var` variable or None + Returns: + The variable in the format '${MERLIN_}' if found, otherwise None + (if `accept_no_matches` is True) or raises a ValueError (if False). + + Raises: + ValueError: If the variable is not found and `accept_no_matches` is False. """ matches = list(re.findall(rf"^(?!#).*MERLIN_{vlaunch_var}=\d+", step_cmd, re.MULTILINE)) @@ -631,10 +934,25 @@ def find_vlaunch_var(vlaunch_var: str, step_cmd: str, accept_no_matches=False) - # Time utilities def convert_to_timedelta(timestr: Union[str, int]) -> timedelta: - """Convert a timestring to a timedelta object. - Timestring is given in in the format '[days]:[hours]:[minutes]:seconds' - with days, hours, minutes all optional add ons. - If passed as an int, will convert to a string first and interpreted as seconds. + """ + Convert a time string or integer to a timedelta object. + + The function takes a time string formatted as + '[days]:[hours]:[minutes]:seconds', where days, hours, and minutes + are optional. If an integer is provided, it is interpreted as the + total number of seconds. + + Args: + timestr: The time string in the specified format or an integer + representing seconds. + + Returns: + A timedelta object representing the duration specified by the input + string or integer. + + Raises: + ValueError: If the input string does not conform to the expected + format or contains more than four time fields. """ # make sure it's a string in case we get an int timestr = str(timestr) @@ -652,7 +970,20 @@ def convert_to_timedelta(timestr: Union[str, int]) -> timedelta: def _repr_timedelta_HMS(time_delta: timedelta) -> str: # pylint: disable=C0103 - """Represent a timedelta object as a string in hours:minutes:seconds""" + """ + Represent a timedelta object as a string in 'HH:MM:SS' format. + + This function converts a given timedelta object into a string that + represents the duration in hours, minutes, and seconds. The output + is formatted as 'HH:MM:SS', with leading zeros for single-digit + hours, minutes, or seconds. + + Args: + time_delta: The timedelta object to be converted. + + Returns: + A string representation of the timedelta in the format 'HH:MM:SS'. + """ hours, remainder = divmod(time_delta.total_seconds(), 3600) minutes, seconds = divmod(remainder, 60) hours, minutes, seconds = int(hours), int(minutes), int(seconds) @@ -660,20 +991,46 @@ def _repr_timedelta_HMS(time_delta: timedelta) -> str: # pylint: disable=C0103 def _repr_timedelta_FSD(time_delta: timedelta) -> str: # pylint: disable=C0103 - """Represent a timedelta as a flux standard duration string, using seconds. + """ + Represent a timedelta as a Flux Standard Duration (FSD) string in seconds. + + The FSD format represents a duration as a floating-point number followed + by a suffix indicating the time unit. This function simplifies the + representation by using seconds and appending an 's' suffix. + + Args: + time_delta: The timedelta object to be converted. - flux standard duration (FSD) is a floating point number with a single character suffix: s,m,h or d. - This uses seconds for simplicity. + Returns: + A string representation of the timedelta in FSD format, expressed + in seconds (e.g., '123.45s'). """ fsd = f"{time_delta.total_seconds()}s" return fsd def repr_timedelta(time_delta: timedelta, method: str = "HMS") -> str: - """Represent a timedelta object as a string using a particular method. + """ + Represent a timedelta object as a string using a specified format method. + + This function formats a given timedelta object according to the chosen + method. The available methods are: + + - HMS: Represents the duration in 'hours:minutes:seconds' format. + - FSD: Represents the duration in Flux Standard Duration (FSD), + expressed as a floating-point number of seconds with an 's' suffix. - method - HMS: 'hours:minutes:seconds' - method - FSD: flux standard duration: 'seconds.s'""" + Args: + time_delta: The timedelta object to be formatted. + method: The method to use for formatting. Must be either 'HMS' or 'FSD'. + + Returns: + A string representation of the timedelta formatted according + to the specified method. + + Raises: + ValueError: If an invalid method is provided. + """ if method == "HMS": return _repr_timedelta_HMS(time_delta) if method == "FSD": @@ -682,16 +1039,27 @@ def repr_timedelta(time_delta: timedelta, method: str = "HMS") -> str: def convert_timestring(timestring: Union[str, int], format_method: str = "HMS") -> str: - """Converts a timestring to a different format. + """ + Converts a timestring to a specified format. - timestring: -either- - a timestring in in the format '[days]:[hours]:[minutes]:seconds' - days, hours, minutes are all optional add ons - -or- - an integer representing seconds - format_method: HMS - 'hours:minutes:seconds' - FSD - 'seconds.s' (flux standard duration) + This function accepts a timestring in a specific format or an integer + representing seconds, and converts it to a formatted string based on + the chosen format method. The available format methods are: + - HMS: Represents the duration in 'hours:minutes:seconds' format. + - FSD: Represents the duration in Flux Standard Duration (FSD), + expressed as a floating-point number of seconds with an 's' suffix. + + Args: + timestring: A string representing time in the format + '[days]:[hours]:[minutes]:seconds' (where days, hours, and + minutes are optional), or an integer representing time in seconds. + format_method: The method to use for formatting. Must be either + 'HMS' or 'FSD'. + + Returns: + A string representation of the converted timestring formatted + according to the specified method. """ LOG.debug(f"Timestring is: {timestring}") tdelta = convert_to_timedelta(timestring) @@ -701,16 +1069,37 @@ def convert_timestring(timestring: Union[str, int], format_method: str = "HMS") def pretty_format_hms(timestring: str) -> str: """ - Given an HMS timestring, format it so it removes blank entries and adds - labels. + Format an HMS timestring to remove blank entries and add appropriate labels. + + This function takes a timestring in the 'HH:MM:SS' format and formats + it by removing any components that are zero and appending the relevant + labels (days, hours, minutes, seconds). The output is a cleaner string + representation of the time. - :param `timestring`: the HMS timestring we'll format - :returns: a formatted timestring + Args: + timestring: A timestring formatted as 'DD:HH:MM:SS'. Each component + represents days, hours, minutes, and seconds, respectively. + Only the last four components are relevant and may include + leading zeros. + + Returns: + A formatted timestring with non-zero components labeled appropriately. + + Raises: + ValueError: If the input timestring contains more than four components + or is not in the expected format. Examples: - - "00:00:34:00" -> "34m" - - "01:00:00:25" -> "01d:25s" - - "00:19:44:28" -> "19h:44m:28s" + ```python + >>> pretty_format_hms("00:00:34:00") + '34m' + >>> pretty_format_hms("01:00:00:25") + '01d:25s' + >>> pretty_format_hms("00:19:44:28") + '19h:44m:28s' + >>> pretty_format_hms("00:00:00:00") + '00s' + ``` """ # Create labels and split the timestring labels = ["d", "h", "m", "s"] @@ -735,10 +1124,23 @@ def pretty_format_hms(timestring: str) -> str: def ws_time_to_dt(ws_time: str) -> datetime: """ - Converts a workspace timestring to a datetime object. + Convert a workspace timestring to a datetime object. + + This function takes a workspace timestring formatted as 'YYYYMMDD-HHMMSS' + and converts it into a corresponding datetime object. The input string + must adhere to the specified format to ensure accurate conversion. - :param `ws_time`: A workspace timestring in the format YYYYMMDD-HHMMSS - :returns: A datetime object created from the workspace timestring + Args: + ws_time: A workspace timestring in the format 'YYYYMMDD-HHMMSS', where:\n + - YYYY is the four-digit year, + - MM is the two-digit month (01 to 12), + - DD is the two-digit day (01 to 31), + - HH is the two-digit hour (00 to 23), + - MM is the two-digit minute (00 to 59), + - SS is the two-digit second (00 to 59). + + Returns: + A datetime object constructed from the provided workspace timestring. """ year = int(ws_time[:4]) month = int(ws_time[4:6]) @@ -751,11 +1153,19 @@ def ws_time_to_dt(ws_time: str) -> datetime: def get_package_versions(package_list: List[str]) -> str: """ - Return a table of the versions and locations of installed packages, including python. - If the package is not installed says "Not installed" + Generate a formatted table of installed package versions and their locations. + + This function takes a list of package names and checks for their installed + versions and locations. If a package is not installed, it indicates that + the package is "Not installed". The output includes the Python version + and its executable location at the top of the table. + + Args: + package_list: A list of package names to check for installed versions. - :param `package_list`: A list of packages. - :returns: A string that's a formatted table. + Returns: + A formatted string representing a table of package names, their versions, + and installation locations. """ table = [] for package in package_list: diff --git a/mkdocs.yml b/mkdocs.yml index 76c123bd3..d7b9126b4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -48,6 +48,9 @@ nav: - API Reference: "api_reference/" - Contact Us: "contact.md" +exclude_docs: | + /README.md + theme: name: material language: en @@ -104,17 +107,17 @@ plugins: - search - codeinclude: title_mode: pymdownx.tabbed - # - gen-files: - # scripts: - # - docs/gen_ref_pages.py - # - mkdocstrings: - # handlers: - # python: - # paths: [merlin] - # options: - # docstring_style: sphinx - # - literate-nav: - # nav_file: SUMMARY.md + - gen-files: + scripts: + - docs/gen_ref_pages.py + - mkdocstrings: + handlers: + python: + paths: [merlin] + options: + docstring_style: google + - literate-nav: + nav_file: SUMMARY.md extra: social: diff --git a/requirements/dev.txt b/requirements/dev.txt index 3695c6164..097e419da 100644 --- a/requirements/dev.txt +++ b/requirements/dev.txt @@ -5,11 +5,13 @@ dep-license flake8 isort pytest +pytest-cov pylint twine sphinx>=2.0.0 alabaster johnnydep deepdiff +orderly-set==5.3.0; python_version == '3.8' pytest-order pytest-mock diff --git a/requirements/release.txt b/requirements/release.txt index dcdb9b81b..b2b0309ce 100644 --- a/requirements/release.txt +++ b/requirements/release.txt @@ -2,8 +2,6 @@ cached_property celery[redis,sqlalchemy]>=5.0.3 coloredlogs cryptography -importlib_metadata<5.0.0; python_version == '3.7' -importlib_resources; python_version < '3.7' maestrowf>=1.1.9dev1 numpy parse diff --git a/setup.cfg b/setup.cfg index a000df59a..77ac2d84f 100644 --- a/setup.cfg +++ b/setup.cfg @@ -26,3 +26,9 @@ max-line-length = 127 files=best_practices,test ignore_missing_imports=true + +[coverage:run] +omit = + merlin/ascii_art.py + merlin/config/celeryconfig.py + merlin/examples/examples.py diff --git a/setup.py b/setup.py index 0ee113e0a..0873efc58 100644 --- a/setup.py +++ b/setup.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -96,11 +96,12 @@ def extras_require(): long_description=readme(), long_description_content_type="text/markdown", classifiers=[ - "Programming Language :: Python :: 3.7", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", ], keywords="machine learning workflow", url="https://github.com/LLNL/merlin", @@ -111,7 +112,6 @@ def extras_require(): entry_points={ "console_scripts": [ "merlin=merlin.main:main", - "merlin-templates=merlin.merlin_templates:main", ] }, include_package_data=True, diff --git a/tests/README.md b/tests/README.md index 22efc5470..e2fa22cbf 100644 --- a/tests/README.md +++ b/tests/README.md @@ -1,17 +1,25 @@ # Tests This directory utilizes pytest to create and run our test suite. -Here we use pytest fixtures to create a local redis server and a celery app for testing. This directory is organized like so: -- `conftest.py` - The script containing all fixtures for our tests -- `unit/` - The directory containing unit tests - - `test_*.py` - The actual test scripts to run +- `conftest.py` - The script containing common fixtures for our tests +- `constants.py` - Constant values to be used throughout the test suite. +- `fixture_data_classes.py` - Dataclasses to help group pytest fixtures together, reducing the required number of imports. +- `fixture_types.py` - Aliases for type hinting fixtures. +- `context_managers/` - The directory containing context managers used for testing + - `celery_workers_manager.py` - A context manager used to manage celery workers for integration testing + - `server_manager.py` - A context manager used to manage the redis server used for integration testing +- `fixtures/` - The directory containing specific test module fixtures + - `.py` - Fixtures for specific test modules - `integration/` - The directory containing integration tests - - `definitions.py` - The test definitions - `run_tests.py` - The script to run the tests defined in `definitions.py` - `conditions.py` - The conditions to test against + - `commands/` - The directory containing tests for commands of the Merlin library. + - `workflows/` The directory containing tests for entire workflow runs. +- `unit/` - The directory containing unit tests + - `test_*.py` - The actual test scripts to run ## How to Run @@ -44,6 +52,28 @@ To run one unique test: python -m pytest /path/to/test_specific_file.py::TestCertainClass::test_unique_test ``` +## Viewing Results + +Test results will be written to `/tmp/$(whoami)/pytest-of-$(whoami)/pytest-current/python_{major}.{minor}.{micro}_current/`. + +It's good practice to set up a subdirectory in this temporary output folder for each module that you're testing. You can see an example of how this is set up in the files within the module-specific fixture directory. For instance, you can see this in the `examples_testing_dir` fixture from the `tests/fixtures/examples.py` file: + +``` +@pytest.fixture(scope="session") +def examples_testing_dir(temp_output_dir: str) -> str: + """ + Fixture to create a temporary output directory for tests related to the examples functionality. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + :returns: The path to the temporary testing directory for examples tests + """ + testing_dir = f"{temp_output_dir}/examples_testing" + if not os.path.exists(testing_dir): + os.mkdir(testing_dir) + + return testing_dir +``` + ## Killing the Test Server In case of an issue with the test suite, or if you stop the tests with `ctrl+C`, you may need to stop @@ -58,58 +88,45 @@ not connected> quit ## The Fixture Process Explained -In the world of pytest testing, fixtures are like the building blocks that create a sturdy foundation for your tests. -They ensure that every test starts from the same fresh ground, leading to reliable and consistent results. This section -will dive into the nitty-gritty of these fixtures, showing you how they're architected in this test suite, how to use -them in your tests here, how to combine them for more complex scenarios, how long they stick around during testing, and -what it means to yield a fixture. +In the world of pytest testing, fixtures are like the building blocks that create a sturdy foundation for your tests. They ensure that every test starts from the same fresh ground, leading to reliable and consistent results. This section will dive into the nitty-gritty of these fixtures, showing you how they're architected in this test suite, how to use them in your tests here, how to combine them for more complex scenarios, how long they stick around during testing, and what it means to yield a fixture. ### Fixture Architecture -Fixtures can be defined in two locations: +Fixtures can be defined in two locations within this test suite: -1. `tests/conftest.py`: This file located at the root of the test suite houses common fixtures that are utilized -across various test modules -2. `tests/fixtures/`: This directory contains specific test module fixtures. Each fixture file is named according -to the module(s) that the fixtures defined within are for. +1. `tests/conftest.py`: This file located at the root of the test suite houses common fixtures that are utilized across various test modules +2. `tests/fixtures/`: This directory contains specific test module fixtures. Each fixture file is named according to the module(s) that the fixtures defined within are for. Credit for this setup must be given to [this Medium article](https://medium.com/@nicolaikozel/modularizing-pytest-fixtures-fd40315c5a93). #### Fixture Naming Conventions -For fixtures defined within the `tests/fixtures/` directory, the fixture name should be prefixed by the name of the -fixture file they are defined in. +For fixtures defined within the `tests/fixtures/` directory, the fixture name should be prefixed by the name of the fixture file they are defined in. #### Importing Fixtures as Plugins -Fixtures located in the `tests/fixtures/` directory are technically plugins. Therefore, to use them we must -register them as plugins within the `conftest.py` file (see the top of said file for the implementation). -This allows them to be discovered and used by test modules throughout the suite. +Fixtures located in the `tests/fixtures/` directory are technically plugins. Therefore, to use them we must register them as plugins within the `conftest.py` file (see the top of said file for the implementation). This allows them to be discovered and used by test modules throughout the suite. -**You do not have to register the fixtures you define as plugins in `conftest.py` since the registration there -uses `glob` to grab everything from the `tests/fixtures/` directory automatically.** +**You do not have to register the fixtures you define as plugins in `conftest.py` since the registration there uses `glob` to grab everything from the `tests/fixtures/` directory automatically.** ### How to Integrate Fixtures Into Tests -Probably the most important part of fixtures is understanding how to use them. Luckily, this process is very -simple and can be dumbed down to just a couple steps: +Probably the most important part of fixtures is understanding how to use them. Luckily, this process is very simple and can be dumbed down to just a couple steps: -0. **[Module-specific fixtures only]** If you're creating a module-specific fixture (i.e. a fixture that won't be used throughout the entire test -suite), then create a file in the `tests/fixtures/` directory. +0. **[Module-specific fixtures only]** If you're creating a module-specific fixture (i.e. a fixture that won't be used throughout the entire test suite), then create a file in the `tests/fixtures/` directory. -1. Create a fixture in either the `conftest.py` file or the file you created in the `tests/fixtures/` directory -by using the `@pytest.fixture` decorator. For example: +1. Create a fixture in either the `conftest.py` file or the file you created in the `tests/fixtures/` directory by using the `@pytest.fixture` decorator. For example: ``` @pytest.fixture -def dummy_fixture(): +def dummy_fixture() -> str: return "hello world" ``` 2. Use it as an argument in a test function (you don't even need to import it!): ``` -def test_dummy(dummy_fixture): +def test_dummy(dummy_fixture: str): assert dummy_fixture == "hello world" ``` @@ -117,22 +134,18 @@ For more information, see [Pytest's documentation](https://docs.pytest.org/en/7. ### Fixtureception -One of the coolest and most useful aspects of fixtures that we utilize in this test suite is the ability for -fixtures to be used within other fixtures. For more info on this from pytest, see -[here](https://docs.pytest.org/en/7.1.x/how-to/fixtures.html#fixtures-can-request-other-fixtures). +One of the coolest and most useful aspects of fixtures that we utilize in this test suite is the ability for fixtures to be used within other fixtures. For more info on this from pytest, see [here](https://docs.pytest.org/en/7.1.x/how-to/fixtures.html#fixtures-can-request-other-fixtures). + +Pytest will handle fixtures within fixtures in a stack-based way. Let's look at how creating the `redis_pass` fixture from our `conftest.py` file works in order to illustrate the process. -Pytest will handle fixtures within fixtures in a stack-based way. Let's look at how creating the `redis_pass` -fixture from our `conftest.py` file works in order to illustrate the process. -1. First, we start by telling pytest that we want to use the `redis_pass` fixture by providing it as an argument -to a test/fixture: +1. First, we start by telling pytest that we want to use the `redis_pass` fixture by providing it as an argument to a test/fixture: ``` def test_example(redis_pass): ... ``` -2. Now pytest will find the `redis_pass` fixture and put it at the top of the stack to be created. However, -it'll see that this fixture requires another fixture `merlin_server_dir` as an argument: +2. Now pytest will find the `redis_pass` fixture and put it at the top of the stack to be created. However, it'll see that this fixture requires another fixture `merlin_server_dir` as an argument: ``` @pytest.fixture(scope="session") @@ -140,8 +153,7 @@ def redis_pass(merlin_server_dir): ... ``` -3. Pytest then puts the `merlin_server_dir` fixture at the top of the stack, but similarly it sees that this fixture -requires yet another fixture `temp_output_dir`: +3. Pytest then puts the `merlin_server_dir` fixture at the top of the stack, but similarly it sees that this fixture requires yet another fixture `temp_output_dir`: ``` @pytest.fixture(scope="session") @@ -149,34 +161,24 @@ def merlin_server_dir(temp_output_dir: str) -> str: ... ``` -4. This process continues until it reaches a fixture that doesn't require any more fixtures. At this point the base -fixture is created and pytest will start working its way back up the stack to the first fixture it looked at (in this -case `redis_pass`). +4. This process continues until it reaches a fixture that doesn't require any more fixtures. At this point the base fixture is created and pytest will start working its way back up the stack to the first fixture it looked at (in this case `redis_pass`). -5. Once all required fixtures are created, execution will be returned to the test which can now access the fixture -that was requested (`redis_pass`). +5. Once all required fixtures are created, execution will be returned to the test which can now access the fixture that was requested (`redis_pass`). -As you can see, if we have to re-do this process for every test it could get pretty time intensive. This is where fixture -scopes come to save the day. +As you can see, if we have to re-do this process for every test it could get pretty time intensive. This is where fixture scopes come to save the day. ### Fixture Scopes -There are several different scopes that you can set for fixtures. The majority of our fixtures in `conftest.py` -use a `session` scope so that we only have to create the fixtures one time (as some of them can take a few seconds -to set up). The goal is to create fixtures with the most general use-case in mind so that we can re-use them for -larger scopes, which helps with efficiency. +There are several different scopes that you can set for fixtures. The majority of our fixtures in `conftest.py` use a `session` scope so that we only have to create the fixtures one time (as some of them can take a few seconds to set up). The goal for fixtures defined in `conftest.py` is to create fixtures with the most general use-case in mind so that we can re-use them for larger scopes, which helps with efficiency. + +For fixtures that need to be reset on each run, we generally try to place these in the module-specific fixture directory `tests/fixtures/`. -For more info on scopes, see -[Pytest's Fixture Scope documentation](https://docs.pytest.org/en/6.2.x/fixture.html#scope-sharing-fixtures-across-classes-modules-packages-or-session). +For more info on scopes, see [Pytest's Fixture Scope documentation](https://docs.pytest.org/en/6.2.x/fixture.html#scope-sharing-fixtures-across-classes-modules-packages-or-session). ### Yielding Fixtures -In several fixtures throughout our test suite, we need to run some sort of teardown for the fixture. For example, -once we no longer need the `redis_server` fixture, we need to shut the server down so it stops using resources. -This is where yielding fixtures becomes extremely useful. +In several fixtures throughout our test suite, we need to run some sort of teardown for the fixture. For example, once we no longer need the `redis_server` fixture, we need to shut the server down so it stops using resources. This is where yielding fixtures becomes extremely useful. -Using the `yield` keyword allows execution to be returned to a test that needs the fixture once the feature has -been set up. After all tests using the fixture have been ran, execution will return to the fixture for us to run -our teardown code. +Using the `yield` keyword allows execution to be returned to a test that needs the fixture once the feature has been set up. After all tests using the fixture have been ran, execution will return to the fixture for us to run our teardown code. For more information on yielding fixtures, see [Pytest's documentation](https://docs.pytest.org/en/7.1.x/how-to/fixtures.html#teardown-cleanup-aka-fixture-finalization). \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py index bea07f64c..46fad196b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -28,46 +28,176 @@ # SOFTWARE. ############################################################################### """ -This module contains pytest fixtures to be used throughout the entire -integration test suite. +This module contains pytest fixtures to be used throughout the entire test suite. """ import os -import subprocess +import sys +from copy import copy from glob import glob from time import sleep -from typing import Dict import pytest -import redis +import yaml from _pytest.tmpdir import TempPathFactory from celery import Celery -from celery.canvas import Signature +from redis import Redis -from tests.celery_test_workers import CeleryTestWorkersManager +from merlin.config.configfile import CONFIG +from tests.constants import CERT_FILES, SERVER_PASS +from tests.context_managers.celery_workers_manager import CeleryWorkersManager +from tests.context_managers.server_manager import RedisServerManager +from tests.fixture_data_classes import RedisBrokerAndBackend +from tests.fixture_types import ( + FixtureBytes, + FixtureCallable, + FixtureCelery, + FixtureDict, + FixtureModification, + FixtureRedis, + FixtureSignature, + FixtureStr, +) +from tests.utils import create_cert_files, create_pass_file + + +# pylint: disable=redefined-outer-name ####################################### # Loading in Module Specific Fixtures # ####################################### + +fixture_glob = os.path.join("tests", "fixtures", "**", "*.py") pytest_plugins = [ - fixture_file.replace("/", ".").replace(".py", "") for fixture_file in glob("tests/fixtures/[!__]*.py", recursive=True) + fixture_file.replace(os.sep, ".").replace(".py", "") + for fixture_file in glob(fixture_glob, recursive=True) + if not fixture_file.endswith("__init__.py") ] -class RedisServerError(Exception): +####################################### +#### Helper Functions for Fixtures #### +####################################### + + +def create_encryption_file(key_filepath: str, encryption_key: bytes, app_yaml_filepath: str = None): """ - Exception to signal that the server wasn't pinged properly. + Check if an encryption file already exists (it will if the redis server has been started) + and if it hasn't then create one and write the encryption key to the file. If an app.yaml + filepath has been passed to this function then we'll need to update it so that the encryption + key points to the `key_filepath`. + + :param key_filepath: The path to the file that will store our encryption key + :param encryption_key: An encryption key to be used for testing + :param app_yaml_filepath: A path to the app.yaml file that needs to be updated """ + if not os.path.exists(key_filepath): + with open(key_filepath, "w") as key_file: + key_file.write(encryption_key.decode("utf-8")) + + if app_yaml_filepath is not None: + # Load up the app.yaml that was created by starting the server + with open(app_yaml_filepath, "r") as app_yaml_file: + app_yaml = yaml.load(app_yaml_file, yaml.Loader) + + # Modify the path to the encryption key and then save it + app_yaml["results_backend"]["encryption_key"] = key_filepath + with open(app_yaml_filepath, "w") as app_yaml_file: + yaml.dump(app_yaml, app_yaml_file) -class ServerInitError(Exception): +def setup_redis_config(config_type: str, merlin_server_dir: str): """ - Exception to signal that there was an error initializing the server. + Sets up the Redis configuration for either broker or results backend. + + Args: + config_type: The type of configuration to set up ('broker' or 'results_backend'). + merlin_server_dir: The directory to the merlin test server configuration. + """ + port = 6379 + name = "redis" + pass_file = os.path.join(merlin_server_dir, "redis.pass") + create_pass_file(pass_file) + + if config_type == "broker": + CONFIG.broker.password = pass_file + CONFIG.broker.port = port + CONFIG.broker.name = name + elif config_type == "results_backend": + CONFIG.results_backend.password = pass_file + CONFIG.results_backend.port = port + CONFIG.results_backend.name = name + else: + raise ValueError("Invalid config_type. Must be 'broker' or 'results_backend'.") + + +####################################### +######### Fixture Definitions ######### +####################################### + + +@pytest.fixture(scope="session") +def path_to_test_specs() -> FixtureStr: + """ + Fixture to provide the path to the directory containing test specifications. + + This fixture returns the absolute path to the 'test_specs' directory + within the 'integration' folder of the test directory. It expands + environment variables and user home directory as necessary. + + Returns: + The absolute path to the 'test_specs' directory. + """ + path_to_test_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(os.path.dirname(__file__)))) + return os.path.join(path_to_test_dir, "integration", "test_specs") + + +@pytest.fixture(scope="session") +def path_to_merlin_codebase() -> FixtureStr: + """ + Fixture to provide the path to the directory containing the Merlin code. + + This fixture returns the absolute path to the 'merlin' directory at the + top level of this repository. It expands environment variables and user + home directory as necessary. + + Returns: + The absolute path to the 'merlin' directory. + """ + path_to_test_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(os.path.dirname(__file__)))) + return os.path.join(path_to_test_dir, "..", "merlin") + + +@pytest.fixture(scope="session") +def create_testing_dir() -> FixtureCallable: + """ + Fixture to create a temporary testing directory. + + Returns: + A function that creates the testing directory. """ + def _create_testing_dir(base_dir: str, sub_dir: str) -> str: + """ + Helper function to create a temporary testing directory. + + Args: + base_dir: The base directory where the testing directory will be created. + sub_dir: The name of the subdirectory to create. + + Returns: + The path to the created testing directory. + """ + testing_dir = os.path.join(base_dir, sub_dir) + if not os.path.exists(testing_dir): + os.makedirs(testing_dir) # Use makedirs to create intermediate directories if needed + return testing_dir + + return _create_testing_dir + @pytest.fixture(scope="session") -def temp_output_dir(tmp_path_factory: TempPathFactory) -> str: +def temp_output_dir(tmp_path_factory: TempPathFactory) -> FixtureStr: """ This fixture will create a temporary directory to store output files of integration tests. The temporary directory will be stored at /tmp/`whoami`/pytest-of-`whoami`/. There can be at most @@ -78,7 +208,9 @@ def temp_output_dir(tmp_path_factory: TempPathFactory) -> str: """ # Log the cwd, then create and move into the temporary one cwd = os.getcwd() - temp_integration_outfile_dir = tmp_path_factory.mktemp("integration_outfiles_") + temp_integration_outfile_dir = tmp_path_factory.mktemp( + f"python_{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}_" + ) os.chdir(temp_integration_outfile_dir) yield temp_integration_outfile_dir @@ -88,88 +220,74 @@ def temp_output_dir(tmp_path_factory: TempPathFactory) -> str: @pytest.fixture(scope="session") -def redis_pass() -> str: - """ - This fixture represents the password to the merlin test server. - - :returns: The redis password for our test server - """ - return "merlin-test-server" - - -@pytest.fixture(scope="session") -def merlin_server_dir(temp_output_dir: str, redis_pass: str) -> str: # pylint: disable=redefined-outer-name +def merlin_server_dir(temp_output_dir: FixtureStr) -> FixtureStr: """ - This fixture will initialize the merlin server (i.e. create all the files we'll - need to start up a local redis server). It will return the path to the directory - containing the files needed for the server to start up. + The path to the merlin_server directory that will be created by the `redis_server` fixture. :param temp_output_dir: The path to the temporary output directory we'll be using for this test run - :param redis_pass: The password to the test redis server that we'll create here - :returns: The path to the merlin_server directory with the server configurations + :returns: The path to the merlin_server directory that will be created by the `redis_server` fixture """ - # Initialize the setup for the local redis server - # We'll also set the password to 'merlin-test-server' so it'll be easy to shutdown if there's an issue - subprocess.run(f"merlin server init; merlin server config -pwd {redis_pass}", shell=True, capture_output=True, text=True) - - # Check that the merlin server was initialized properly - server_dir = f"{temp_output_dir}/merlin_server" + server_dir = os.path.join(temp_output_dir, "merlin_server") if not os.path.exists(server_dir): - raise ServerInitError("The merlin server was not initialized properly.") - + os.mkdir(server_dir) return server_dir @pytest.fixture(scope="session") -def redis_server(merlin_server_dir: str, redis_pass: str) -> str: # pylint: disable=redefined-outer-name,unused-argument +def redis_server(merlin_server_dir: FixtureStr, test_encryption_key: FixtureBytes) -> FixtureStr: """ Start a redis server instance that runs on localhost:6379. This will yield the redis server uri that can be used to create a connection with celery. - :param merlin_server_dir: The directory to the merlin test server configuration. - This will not be used here but we need the server configurations before we can - start the server. - :param redis_pass: The raw redis password stored in the redis.pass file + :param merlin_server_dir: The directory to the merlin test server configuration + :param test_encryption_key: An encryption key to be used for testing :yields: The local redis server uri """ - # Start the local redis server - try: - # Need to set LC_ALL='C' before starting the server or else redis causes a failure - subprocess.run("export LC_ALL='C'; merlin server start", shell=True, timeout=5) - except subprocess.TimeoutExpired: - pass + os.environ["CELERY_ENV"] = "test" + with RedisServerManager(merlin_server_dir, SERVER_PASS) as redis_server_manager: + redis_server_manager.initialize_server() + redis_server_manager.start_server() + create_encryption_file( + os.path.join(merlin_server_dir, "encrypt_data_key"), + test_encryption_key, + app_yaml_filepath=os.path.join(merlin_server_dir, "app.yaml"), + ) + # Yield the redis_server uri to any fixtures/tests that may need it + yield redis_server_manager.redis_server_uri + # The server will be stopped once this context reaches the end of it's execution here - # Ensure the server started properly - host = "localhost" - port = 6379 - database = 0 - username = "default" - redis_client = redis.Redis(host=host, port=port, db=database, password=redis_pass, username=username) - if not redis_client.ping(): - raise RedisServerError("The redis server could not be pinged. Check that the server is running with 'ps ux'.") - # Hand over the redis server url to any other fixtures/tests that need it - redis_server_uri = f"redis://{username}:{redis_pass}@{host}:{port}/{database}" - yield redis_server_uri +@pytest.fixture(scope="session") +def redis_client(redis_server: FixtureStr) -> FixtureRedis: + """ + Fixture that provides a Redis client instance for the test session. + It connects to this client using the url created from the `redis_server` + fixture. + + Args: + redis_server: The redis server uri we'll use to connect to redis - # Kill the server; don't run this until all tests are done (accomplished with 'yield' above) - kill_process = subprocess.run("merlin server stop", shell=True, capture_output=True, text=True) - assert "Merlin server terminated." in kill_process.stderr + Returns: + An instance of the Redis client that can be used to interact + with the Redis server. + """ + return Redis.from_url(url=redis_server) @pytest.fixture(scope="session") -def celery_app(redis_server: str) -> Celery: # pylint: disable=redefined-outer-name +def celery_app(redis_server: FixtureStr) -> FixtureCelery: """ Create the celery app to be used throughout our integration tests. :param redis_server: The redis server uri we'll use to connect to redis :returns: The celery app object we'll use for testing """ + os.environ["CELERY_ENV"] = "test" return Celery("merlin_test_app", broker=redis_server, backend=redis_server) @pytest.fixture(scope="session") -def sleep_sig(celery_app: Celery) -> Signature: # pylint: disable=redefined-outer-name +def sleep_sig(celery_app: FixtureCelery) -> FixtureSignature: """ Create a task registered to our celery app and return a signature for it. Once requested by a test, you can set the queue you'd like to send this to @@ -191,7 +309,7 @@ def sleep_task(): @pytest.fixture(scope="session") -def worker_queue_map() -> Dict[str, str]: +def worker_queue_map() -> FixtureDict[str, str]: """ Worker and queue names to be used throughout tests @@ -201,7 +319,7 @@ def worker_queue_map() -> Dict[str, str]: @pytest.fixture(scope="class") -def launch_workers(celery_app: Celery, worker_queue_map: Dict[str, str]): # pylint: disable=redefined-outer-name +def launch_workers(celery_app: FixtureCelery, worker_queue_map: FixtureDict[str, str]): """ Launch the workers on the celery app fixture using the worker and queue names defined in the worker_queue_map fixture. @@ -213,6 +331,352 @@ def launch_workers(celery_app: Celery, worker_queue_map: Dict[str, str]): # pyl # (basically just add in concurrency value to worker_queue_map) worker_info = {worker_name: {"concurrency": 1, "queues": [queue]} for worker_name, queue in worker_queue_map.items()} - with CeleryTestWorkersManager(celery_app) as workers_manager: + with CeleryWorkersManager(celery_app) as workers_manager: workers_manager.launch_workers(worker_info) yield + + +@pytest.fixture(scope="session") +def test_encryption_key() -> FixtureBytes: + """ + An encryption key to be used for tests that need it. + + :returns: The test encryption key + """ + return b"Q3vLp07Ljm60ahfU9HwOOnfgGY91lSrUmqcTiP0v9i0=" + + +####################################### +########### CONFIG Fixtures ########### +####################################### +# These are intended to be used # +# either by themselves or together # +# For example, you can use a rabbit # +# broker config and a redis results # +# backend config together # +####################################### +############ !!!WARNING!!! ############ +# DO NOT USE THE `config` FIXTURE # +# IN A TEST; IT HAS UNSET VALUES # +####################################### + + +def _config(merlin_server_dir: FixtureStr, test_encryption_key: FixtureBytes): + """ + Sets up the configuration for testing purposes by modifying the global CONFIG object. + + This helper function prepares the broker and results backend configurations for testing + by creating necessary encryption key files and resetting the CONFIG object to its + original state after the tests are executed. + + Args: + merlin_server_dir: The directory to the merlin test server configuration + test_encryption_key: An encryption key to be used for testing + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + # Create a copy of the CONFIG option so we can reset it after the test + orig_config = copy(CONFIG) + + # Create an encryption key file (if it doesn't already exist) + key_file = os.path.join(merlin_server_dir, "encrypt_data_key") + create_encryption_file(key_file, test_encryption_key) + + # Set the broker configuration for testing + CONFIG.broker.password = None # This will be updated in `redis_broker_config_*` or `rabbit_broker_config` + CONFIG.broker.port = None # This will be updated in `redis_broker_config_*` or `rabbit_broker_config` + CONFIG.broker.name = None # This will be updated in `redis_broker_config_*` or `rabbit_broker_config` + CONFIG.broker.server = "127.0.0.1" + CONFIG.broker.username = "default" + CONFIG.broker.vhost = "host4testing" + CONFIG.broker.cert_reqs = "none" + + # Set the results_backend configuration for testing + CONFIG.results_backend.password = ( + None # This will be updated in `redis_results_backend_config_function` or `mysql_results_backend_config` + ) + CONFIG.results_backend.port = None # This will be updated in `redis_results_backend_config_function` + CONFIG.results_backend.name = ( + None # This will be updated in `redis_results_backend_config_function` or `mysql_results_backend_config` + ) + CONFIG.results_backend.dbname = None # This will be updated in `mysql_results_backend_config` + CONFIG.results_backend.server = "127.0.0.1" + CONFIG.results_backend.username = "default" + CONFIG.results_backend.cert_reqs = "none" + CONFIG.results_backend.encryption_key = key_file + CONFIG.results_backend.db_num = 0 + + # Go run the tests + yield + + # Reset the configuration + CONFIG.celery = orig_config.celery + CONFIG.broker = orig_config.broker + CONFIG.results_backend = orig_config.results_backend + + +@pytest.fixture(scope="function") +def config_function(merlin_server_dir: FixtureStr, test_encryption_key: FixtureBytes) -> FixtureModification: + """ + Sets up the configuration for testing with a function scope. + + Warning: + DO NOT USE THIS FIXTURE IN A TEST, USE ONE OF THE SERVER SPECIFIC CONFIGURATIONS + (LIKE `redis_broker_config_function`, `rabbit_broker_config`, etc.) INSTEAD. + + This fixture modifies the global CONFIG object to prepare the broker and results backend + configurations for testing. It creates necessary encryption key files and ensures that + the original configuration is restored after the tests are executed. + + Args: + merlin_server_dir: The directory to the merlin test server configuration + test_encryption_key: An encryption key to be used for testing + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + yield from _config(merlin_server_dir, test_encryption_key) + + +@pytest.fixture(scope="class") +def config_class(merlin_server_dir: FixtureStr, test_encryption_key: FixtureBytes) -> FixtureModification: + """ + Sets up the configuration for testing with a class scope. + + Warning: + DO NOT USE THIS FIXTURE IN A TEST, USE ONE OF THE SERVER SPECIFIC CONFIGURATIONS + (LIKE `redis_broker_config_class`, `rabbit_broker_config`, etc.) INSTEAD. + + This fixture modifies the global CONFIG object to prepare the broker and results backend + configurations for testing. It creates necessary encryption key files and ensures that + the original configuration is restored after the tests are executed. + + Args: + merlin_server_dir: The directory to the merlin test server configuration + test_encryption_key: An encryption key to be used for testing + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + yield from _config(merlin_server_dir, test_encryption_key) + + +@pytest.fixture(scope="function") +def redis_broker_config_function( + merlin_server_dir: FixtureStr, config_function: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + Fixture for configuring the Redis broker for testing with a function scope. + + This fixture sets up the CONFIG object to use a Redis broker for testing any functionality + in the codebase that interacts with the broker. It modifies the configuration to point + to the specified Redis broker settings. + + Args: + merlin_server_dir: The directory to the merlin test server configuration. + config_function: The fixture that sets up most of the CONFIG object for testing. + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + setup_redis_config("broker", merlin_server_dir) + yield + + +@pytest.fixture(scope="class") +def redis_broker_config_class( + merlin_server_dir: FixtureStr, config_class: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + Fixture for configuring the Redis broker for testing with a class scope. + + This fixture sets up the CONFIG object to use a Redis broker for testing any functionality + in the codebase that interacts with the broker. It modifies the configuration to point + to the specified Redis broker settings. + + Args: + merlin_server_dir: The directory to the merlin test server configuration. + config_function: The fixture that sets up most of the CONFIG object for testing. + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + setup_redis_config("broker", merlin_server_dir) + yield + + +@pytest.fixture(scope="function") +def redis_results_backend_config_function( + merlin_server_dir: FixtureStr, config_function: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + Fixture for configuring the Redis results backend for testing with a function scope. + + This fixture sets up the CONFIG object to use a Redis results backend for testing any + functionality in the codebase that interacts with the results backend. It modifies the + configuration to point to the specified Redis results backend settings. + + Args: + merlin_server_dir: The directory to the merlin test server configuration. + config_function: The fixture that sets up most of the CONFIG object for testing. + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + setup_redis_config("results_backend", merlin_server_dir) + yield + + +@pytest.fixture(scope="class") +def redis_results_backend_config_class( + merlin_server_dir: FixtureStr, config_class: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + Fixture for configuring the Redis results backend for testing with a class scope. + + This fixture sets up the CONFIG object to use a Redis results backend for testing any + functionality in the codebase that interacts with the results backend. It modifies the + configuration to point to the specified Redis results backend settings. + + Args: + merlin_server_dir: The directory to the merlin test server configuration. + config_function: The fixture that sets up most of the CONFIG object for testing. + + Yields: + This function yields control back to the test function, allowing tests to run + with the modified CONFIG settings. + """ + setup_redis_config("results_backend", merlin_server_dir) + yield + + +@pytest.fixture(scope="function") +def rabbit_broker_config( + merlin_server_dir: FixtureStr, config_function: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + This fixture is intended to be used for testing any functionality in the codebase + that uses the CONFIG object with a RabbitMQ broker. + + :param merlin_server_dir: The directory to the merlin test server configuration + :param config: The fixture that sets up most of the CONFIG object for testing + """ + pass_file = os.path.join(merlin_server_dir, "rabbit.pass") + create_pass_file(pass_file) + + CONFIG.broker.password = pass_file + CONFIG.broker.port = 5671 + CONFIG.broker.name = "rabbitmq" + + yield + + +@pytest.fixture(scope="function") +def mysql_results_backend_config( + merlin_server_dir: FixtureStr, config_function: FixtureModification # pylint: disable=redefined-outer-name,unused-argument +) -> FixtureModification: + """ + This fixture is intended to be used for testing any functionality in the codebase + that uses the CONFIG object with a MySQL results_backend. + + :param merlin_server_dir: The directory to the merlin test server configuration + :param config: The fixture that sets up most of the CONFIG object for testing + """ + pass_file = os.path.join(merlin_server_dir, "mysql.pass") + create_pass_file(pass_file) + + create_cert_files(merlin_server_dir, CERT_FILES) + + CONFIG.results_backend.password = pass_file + CONFIG.results_backend.name = "mysql" + CONFIG.results_backend.dbname = "test_mysql_db" + CONFIG.results_backend.keyfile = CERT_FILES["ssl_key"] + CONFIG.results_backend.certfile = CERT_FILES["ssl_cert"] + CONFIG.results_backend.ca_certs = CERT_FILES["ssl_ca"] + + yield + + +@pytest.fixture(scope="function") +def redis_broker_and_backend_function( + redis_client: FixtureRedis, + redis_server: FixtureStr, + redis_broker_config_function: FixtureModification, + redis_results_backend_config_function: FixtureModification, +): + """ + Fixture for setting up Redis broker and backend for function-scoped tests. + + This fixture creates an instance of `RedisBrokerAndBackend`, which + encapsulates all necessary Redis-related fixtures required for + establishing connections to Redis as both a broker and a backend + during function-scoped tests. + + Args: + redis_client: A fixture that provides a client for interacting with the + Redis server. + redis_server: A fixture providing the connection string to the Redis + server instance. + redis_broker_config_function: A fixture that modifies the configuration + to point to the Redis server used as the message broker for + function-scoped tests. + redis_results_backend_config_function: A fixture that modifies the + configuration to point to the Redis server used for storing results + in function-scoped tests. + + Returns: + An instance containing the Redis client, server connection string, and + configuration modifications for both the broker and backend. + """ + return RedisBrokerAndBackend( + client=redis_client, + server=redis_server, + broker_config=redis_broker_config_function, + results_backend_config=redis_results_backend_config_function, + ) + + +@pytest.fixture(scope="class") +def redis_broker_and_backend_class( + redis_client: FixtureRedis, + redis_server: FixtureStr, + redis_broker_config_class: FixtureModification, + redis_results_backend_config_class: FixtureModification, +) -> RedisBrokerAndBackend: + """ + Fixture for setting up Redis broker and backend for class-scoped tests. + + This fixture creates an instance of `RedisBrokerAndBackend`, which + encapsulates all necessary Redis-related fixtures required for + establishing connections to Redis as both a broker and a backend + during class-scoped tests. + + Args: + redis_client: A fixture that provides a client for interacting with the + Redis server. + redis_server: A fixture providing the connection string to the Redis + server instance. + redis_broker_config_function: A fixture that modifies the configuration + to point to the Redis server used as the message broker for + class-scoped tests. + redis_results_backend_config_function: A fixture that modifies the + configuration to point to the Redis server used for storing results + in class-scoped tests. + + Returns: + An instance containing the Redis client, server connection string, and + configuration modifications for both the broker and backend. + """ + return RedisBrokerAndBackend( + client=redis_client, + server=redis_server, + broker_config=redis_broker_config_class, + results_backend_config=redis_results_backend_config_class, + ) diff --git a/tests/constants.py b/tests/constants.py new file mode 100644 index 000000000..26cfe4c0a --- /dev/null +++ b/tests/constants.py @@ -0,0 +1,11 @@ +""" +This module will store constants that will be used throughout our test suite. +""" + +SERVER_PASS = "merlin-test-server" + +CERT_FILES = { + "ssl_cert": "test-rabbit-client-cert.pem", + "ssl_ca": "test-mysql-ca-cert.pem", + "ssl_key": "test-rabbit-client-key.pem", +} diff --git a/tests/context_managers/__init__.py b/tests/context_managers/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/context_managers/celery_task_manager.py b/tests/context_managers/celery_task_manager.py new file mode 100644 index 000000000..b71dcd8a4 --- /dev/null +++ b/tests/context_managers/celery_task_manager.py @@ -0,0 +1,155 @@ +""" +Module to define functionality for sending tasks to the server +and ensuring they're cleared from the server when the test finishes. +""" + +from types import TracebackType +from typing import List, Type + +from celery import Celery +from celery.result import AsyncResult +from redis import Redis + + +class CeleryTaskManager: + """ + A context manager for managing Celery tasks. + + This class provides a way to send tasks to a Celery server and clean up + any tasks that were sent during its lifetime. It is designed to be used + as a context manager, ensuring that tasks are properly removed from the + server when the context is exited. + + Attributes: + celery_app: The Celery application instance. + redis_server: The Redis server connection string. + """ + + def __init__(self, app: Celery, redis_client: Redis): + self.celery_app: Celery = app + self.redis_client = redis_client + + def __enter__(self) -> "CeleryTaskManager": + """ + Enters the runtime context related to this object. + + Returns: + The current instance of the manager. + """ + return self + + def __exit__(self, exc_type: Type[Exception], exc_value: Exception, traceback: TracebackType): + """ + Exits the runtime context and performs cleanup. + + This method removes any tasks currently in the server. + + Args: + exc_type: The exception type raised, if any. + exc_value: The exception instance raised, if any. + traceback: The traceback object, if an exception was raised. + """ + self.remove_tasks() + + def send_task(self, task_name: str, *args, **kwargs) -> AsyncResult: + """ + Sends a task to the Celery server. + + This method will be used for tests that don't call + `merlin run`, allowing for isolated test functionality. + + Args: + task_name: The name of the task to send to the server. + *args: Additional positional arguments to pass to the task. + **kwargs: Additional keyword arguments to pass to the task. + + Returns: + A Celery AsyncResult object containing information about the + task that was sent to the server. + """ + valid_kwargs = [ + "add_to_parent", + "chain", + "chord", + "compression", + "connection", + "countdown", + "eta", + "exchange", + "expires", + "group_id", + "group_index", + "headers", + "ignore_result", + "link", + "link_error", + "parent_id", + "priority", + "producer", + "publisher", + "queue", + "replaced_task_nesting", + "reply_to", + "result_cls", + "retries", + "retry", + "retry_policy", + "root_id", + "route_name", + "router", + "routing_key", + "serializer", + "shadow", + "soft_time_limit", + "task_id", + "task_type", + "time_limit", + ] + send_task_kwargs = {key: kwargs.pop(key) for key in valid_kwargs if key in kwargs} + + return self.celery_app.send_task(task_name, args=args, kwargs=kwargs, **send_task_kwargs) + + def remove_tasks(self): + """ + Removes tasks from the Celery server. + + Tasks are removed in two ways: + 1. By purging the Celery app queues, which will only purge tasks + sent with `send_task`. + 2. By deleting the remaining queues in the Redis server, which will + purge any tasks that weren't sent with `send_task` (e.g., tasks + sent with `merlin run`). + """ + # Purge the tasks + self.celery_app.control.purge() + + # Purge any remaining tasks directly through redis that may have been missed + queues = self.get_queue_list() + for queue in queues: + self.redis_client.delete(queue) + + def get_queue_list(self) -> List[str]: + """ + Builds a list of Celery queues that exist on the Redis server. + + Queries the Redis server for its keys and returns the keys + that represent the Celery queues. + + Returns: + A list of Celery queue names. + """ + cursor = 0 + queues = [] + while True: + # Get the 'merlin' queue if it exists + cursor, matching_queues = self.redis_client.scan(cursor=cursor, match="merlin") + queues.extend(matching_queues) + + # Get any queues that start with '[merlin]' + cursor, matching_queues = self.redis_client.scan(cursor=cursor, match="\\[merlin\\]*") + queues.extend(matching_queues) + + if cursor == 0: + break + + return queues diff --git a/tests/celery_test_workers.py b/tests/context_managers/celery_workers_manager.py similarity index 84% rename from tests/celery_test_workers.py rename to tests/context_managers/celery_workers_manager.py index ad81d30e6..279eab325 100644 --- a/tests/celery_test_workers.py +++ b/tests/context_managers/celery_workers_manager.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -42,7 +42,7 @@ from celery import Celery -class CeleryTestWorkersManager: +class CeleryWorkersManager: """ A class to handle the setup and teardown of celery workers. This should be treated as a context and used with python's @@ -52,7 +52,8 @@ class CeleryTestWorkersManager: def __init__(self, app: Celery): self.app = app - self.running_workers = [] + self.running_workers = set() + self.run_worker_processes = set() self.worker_processes = {} self.echo_processes = {} @@ -80,8 +81,9 @@ def __exit__(self, exc_type: Type[Exception], exc_value: Exception, traceback: T try: if str(pid) in ps_proc.stdout: os.kill(pid, signal.SIGKILL) - except ProcessLookupError as exc: - raise ProcessLookupError(f"PID {pid} not found. Output of 'ps ux':\n{ps_proc.stdout}") from exc + # If the process can't be found then it doesn't exist anymore + except ProcessLookupError: + pass def _is_worker_ready(self, worker_name: str, verbose: bool = False) -> bool: """ @@ -135,7 +137,7 @@ def start_worker(self, worker_launch_cmd: List[str]): app.worker_main instead of the normal "celery -A worker" command to launch the workers since our celery app is created in a pytest fixture and is unrecognizable by the celery command. For each worker, the output of it's logs are sent to - /tmp/`whoami`/pytest-of-`whoami`/pytest-current/integration_outfiles_current/ under a file with a name + /tmp/`whoami`/pytest-of-`whoami`/pytest-current/python_{major}.{minor}.{micro}_current/ under a file with a name similar to: test_worker_*.log. NOTE: pytest-current/ will have the results of the most recent test run. If you want to see a previous run check under pytest-/. HOWEVER, only the 3 most recent test runs will be saved. @@ -144,7 +146,7 @@ def start_worker(self, worker_launch_cmd: List[str]): """ self.app.worker_main(worker_launch_cmd) - def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = 1): + def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = 1, prefetch: int = 1): """ Launch a single worker. We'll add the process that the worker is running in to the list of worker processes. We'll also create an echo process to simulate a celery worker command that will show up with 'ps ux'. @@ -158,6 +160,8 @@ def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = self.stop_all_workers() raise ValueError(f"The worker {worker_name} is already running. Choose a different name.") + queues = [f"[merlin]_{queue}" for queue in queues] + # Create the launch command for this worker worker_launch_cmd = [ "worker", @@ -167,6 +171,8 @@ def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = ",".join(queues), "--concurrency", str(concurrency), + "--prefetch-multiplier", + str(prefetch), f"--logfile={worker_name}.log", "--loglevel=DEBUG", ] @@ -174,9 +180,9 @@ def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = # Create an echo command to simulate a running celery worker since our celery worker will be spun up in # a different process and we won't be able to see it with 'ps ux' like we normally would echo_process = subprocess.Popen( # pylint: disable=consider-using-with - f"echo 'celery merlin_test_app {' '.join(worker_launch_cmd)}'; sleep inf", + f"echo 'celery -A merlin_test_app {' '.join(worker_launch_cmd)}'; sleep inf", shell=True, - preexec_fn=os.setpgrp, # Make this the parent of the group so we can kill the 'sleep inf' that's spun up + start_new_session=True, # Make this the parent of the group so we can kill the 'sleep inf' that's spun up ) self.echo_processes[worker_name] = echo_process.pid @@ -184,7 +190,7 @@ def launch_worker(self, worker_name: str, queues: List[str], concurrency: int = worker_process = multiprocessing.Process(target=self.start_worker, args=(worker_launch_cmd,)) worker_process.start() self.worker_processes[worker_name] = worker_process - self.running_workers.append(worker_name) + self.running_workers.add(worker_name) # Wait for the worker to launch properly try: @@ -204,6 +210,24 @@ def launch_workers(self, worker_info: Dict[str, Dict]): for worker_name, worker_settings in worker_info.items(): self.launch_worker(worker_name, worker_settings["queues"], worker_settings["concurrency"]) + def add_run_workers_process(self, pid: int): + """ + Add a process ID for a `merlin run-workers` process to the + set that tracks all `merlin run-workers` processes that are + currently running. + + Warning: + The process that's added here must utilize the + `start_new_session=True` setting of subprocess.Popen. This + is necessary for us to be able to terminate all the workers + that are started with it safely since they will be seen as + child processes of the `merlin run-workers` process. + + Args: + pid: The process ID running `merlin run-workers`. + """ + self.run_worker_processes.add(pid) + def stop_worker(self, worker_name: str): """ Stop a single running worker and its associated processes. @@ -223,12 +247,16 @@ def stop_worker(self, worker_name: str): self.worker_processes[worker_name].kill() # Terminate the echo process and its sleep inf subprocess - os.killpg(os.getpgid(self.echo_processes[worker_name]), signal.SIGTERM) - sleep(2) + if self.echo_processes[worker_name] is not None: + os.killpg(os.getpgid(self.echo_processes[worker_name]), signal.SIGKILL) + sleep(2) def stop_all_workers(self): """ Stop all of the running workers and the processes associated with them. """ + for run_worker_pid in self.run_worker_processes: + os.killpg(os.getpgid(run_worker_pid), signal.SIGKILL) + for worker_name in self.running_workers: self.stop_worker(worker_name) diff --git a/tests/context_managers/server_manager.py b/tests/context_managers/server_manager.py new file mode 100644 index 000000000..bb5d86036 --- /dev/null +++ b/tests/context_managers/server_manager.py @@ -0,0 +1,106 @@ +""" +Module to define functionality for managing the containerized +server used for testing. +""" + +import os +import signal +import subprocess +from types import TracebackType +from typing import Type + +import redis +import yaml + + +class RedisServerError(Exception): + """ + Exception to signal that the server wasn't pinged properly. + """ + + +class ServerInitError(Exception): + """ + Exception to signal that there was an error initializing the server. + """ + + +class RedisServerManager: + """ + A class to handle the setup and teardown of a containerized redis server. + This should be treated as a context and used with python's built-in 'with' + statement. If you use it without this statement, beware that the processes + spun up here may never be stopped. + """ + + def __init__(self, server_dir: str, redis_pass: str): + self._redis_pass = redis_pass + self.server_dir = server_dir + self.host = "localhost" + self.port = 6379 + self.database = 0 + self.username = "default" + self.redis_server_uri = f"redis://{self.username}:{self._redis_pass}@{self.host}:{self.port}/{self.database}" + + def __enter__(self): + """This magic method is necessary for allowing this class to be used as a context manager""" + return self + + def __exit__(self, exc_type: Type[Exception], exc_value: Exception, traceback: TracebackType): + """ + This will always run at the end of a context with statement, even if an error is raised. + It's a safe way to ensure all of our server gets stopped no matter what. + """ + self.stop_server() + + def initialize_server(self): + """ + Initialize the setup for the local redis server. We'll write the folder to: + /tmp/`whoami`/pytest-of-`whoami`/pytest-current/python_{major}.{minor}.{micro}_current/ + We'll set the password to be 'merlin-test-server' so it'll be easy to shutdown if necessary + """ + subprocess.run( + f"merlin server init; merlin server config -pwd {self._redis_pass}", shell=True, capture_output=True, text=True + ) + + # Check that the merlin server was initialized properly + if not os.path.exists(self.server_dir): + raise ServerInitError("The merlin server was not initialized properly.") + + def start_server(self): + """Attempt to start the local redis server.""" + try: + # Need to set LC_ALL='C' before starting the server or else redis causes a failure + subprocess.run("export LC_ALL='C'; merlin server start", shell=True, timeout=5) + except subprocess.TimeoutExpired: + pass + + # Ensure the server started properly + redis_client = redis.Redis( + host=self.host, port=self.port, db=self.database, password=self._redis_pass, username=self.username + ) + if not redis_client.ping(): + raise RedisServerError("The redis server could not be pinged. Check that the server is running with 'ps ux'.") + + def stop_server(self): + """Stop the server.""" + # Attempt to stop the server gracefully with `merlin server` + kill_process = subprocess.run("merlin server stop", shell=True, capture_output=True, text=True) + + # Check that the server was terminated + if "Merlin server terminated." not in kill_process.stderr: + # If it wasn't, try to kill the process by using the pid stored in a file created by `merlin server` + try: + with open(os.path.join(self.server_dir, "merlin_server.pf"), "r") as process_file: + server_process_info = yaml.load(process_file, yaml.Loader) + os.kill(int(server_process_info["image_pid"]), signal.SIGKILL) + # If the file can't be found then let's make sure there's even a redis-server process running + except FileNotFoundError as exc: + process_query = subprocess.run("ps ux", shell=True, text=True, capture_output=True) + # If there is a file running we didn't start it in this test run so we can't kill it + if "redis-server" in process_query.stdout: + raise RedisServerError( + "Found an active redis server but cannot stop it since there is no process file (merlin_server.pf). " + "Did you start this server before running tests?" + ) from exc + # No else here. If there's no redis-server process found then there's nothing to stop diff --git a/tests/fixture_data_classes.py b/tests/fixture_data_classes.py new file mode 100644 index 000000000..2ff7bf877 --- /dev/null +++ b/tests/fixture_data_classes.py @@ -0,0 +1,80 @@ +""" +This module houses dataclasses to be used with pytest fixtures. +""" + +from dataclasses import dataclass + +from tests.fixture_types import FixtureInt, FixtureModification, FixtureRedis, FixtureStr + + +@dataclass +class RedisBrokerAndBackend: + """ + Data class to encapsulate all Redis-related fixtures required for + establishing connections to Redis for both the broker and backend. + + This class simplifies the management of Redis fixtures by grouping + them into a single object, reducing the number of individual fixture + imports needed in tests that require Redis functionality. + + Attributes: + client: A fixture that provides a client for interacting + with the Redis server. + server: A fixture providing the connection string to the + Redis server instance. + broker_config: A fixture that modifies the configuration + to point to the Redis server used as the message broker. + results_backend_config: A fixture that modifies the + configuration to point to the Redis server used for storing + results. + """ + + client: FixtureRedis + server: FixtureStr + results_backend_config: FixtureModification + broker_config: FixtureModification + + +@dataclass +class FeatureDemoSetup: + """ + Data class to encapsulate all feature-demo-related fixtures required + for testing the feature demo workflow. + + This class simplifies the management of feature demo setup fixtures + by grouping them into a single object, reducing the number of individual + fixture imports needed in tests that require feature demo setup. + + Attributes: + testing_dir: The path to the temp output directory for feature_demo workflow tests. + num_samples: An integer representing the number of samples to use in the feature_demo + workflow. + name: A string representing the name to use for the feature_demo workflow. + path: The path to the feature demo YAML file. + """ + + testing_dir: FixtureStr + num_samples: FixtureInt + name: FixtureStr + path: FixtureStr + + +@dataclass +class ChordErrorSetup: + """ + Data class to encapsulate all chord-error-related fixtures required + for testing the chord error workflow. + + This class simplifies the management of chord error setup fixtures + by grouping them into a single object, reducing the number of individual + fixture imports needed in tests that require chord error setup. + + Attributes: + testing_dir: The path to the temp output directory for chord_err workflow tests. + name: A string representing the name to use for the chord_err workflow. + path: The path to the chord error YAML file. + """ + + testing_dir: FixtureStr + name: FixtureStr + path: FixtureStr diff --git a/tests/fixture_types.py b/tests/fixture_types.py new file mode 100644 index 000000000..b6e6839ea --- /dev/null +++ b/tests/fixture_types.py @@ -0,0 +1,74 @@ +""" +It's hard to type hint pytest fixtures in a way that makes it clear +that the variable being used is a fixture. This module will created +aliases for these fixtures in order to make it easier to track what's +happening. + +The types here will be defined as such: +- `FixtureBytes`: A fixture that returns bytes +- `FixtureCelery`: A fixture that returns a Celery app object +- `FixtureDict`: A fixture that returns a dictionary +- `FixtureInt`: A fixture that returns an integer +- `FixtureModification`: A fixture that modifies something but never actually + returns/yields a value to be used in the test. +- `FixtureRedis`: A fixture that returns a Redis client +- `FixtureSignature`: A fixture that returns a Celery Signature object +- `FixtureStr`: A fixture that returns a string +""" + +import sys +from argparse import Namespace +from collections.abc import Callable +from typing import Any, Dict, Generic, Tuple, TypeVar + +import pytest +from celery import Celery +from celery.canvas import Signature +from redis import Redis + + +# TODO convert unit test type hinting to use these +# - likely will do this when I work on API docs for test library + +K = TypeVar("K") +V = TypeVar("V") + +# TODO when we drop support for Python 3.8, remove this if/else statement +# Check Python version +if sys.version_info >= (3, 9): + from typing import Annotated + + FixtureBytes = Annotated[bytes, pytest.fixture] + FixtureCallable = Annotated[Callable, pytest.fixture] + FixtureCelery = Annotated[Celery, pytest.fixture] + FixtureDict = Annotated[Dict[K, V], pytest.fixture] + FixtureInt = Annotated[int, pytest.fixture] + FixtureModification = Annotated[Any, pytest.fixture] + FixtureNamespace = Annotated[Namespace, pytest.fixture] + FixtureRedis = Annotated[Redis, pytest.fixture] + FixtureSignature = Annotated[Signature, pytest.fixture] + FixtureStr = Annotated[str, pytest.fixture] + FixtureTuple = Annotated[Tuple[K, V], pytest.fixture] +else: + # Fallback for Python 3.8 + class FixtureDict(Generic[K, V], Dict[K, V]): + """ + This class is necessary to allow FixtureDict to be subscriptable + when using it to type hint. + """ + + class FixtureTuple(Generic[K, V], Tuple[K, V]): + """ + This class is necessary to allow FixtureTuple to be subscriptable + when using it to type hint. + """ + + FixtureBytes = pytest.fixture + FixtureCallable = pytest.fixture + FixtureCelery = pytest.fixture + FixtureInt = pytest.fixture + FixtureModification = pytest.fixture + FixtureNamespace = pytest.fixture + FixtureRedis = pytest.fixture + FixtureSignature = pytest.fixture + FixtureStr = pytest.fixture diff --git a/tests/fixtures/chord_err.py b/tests/fixtures/chord_err.py new file mode 100644 index 000000000..54c639072 --- /dev/null +++ b/tests/fixtures/chord_err.py @@ -0,0 +1,114 @@ +""" +Fixtures specifically for help testing the chord_err workflow. +""" + +import os +import subprocess + +import pytest + +from tests.fixture_data_classes import ChordErrorSetup, RedisBrokerAndBackend +from tests.fixture_types import FixtureCallable, FixtureStr +from tests.integration.helper_funcs import copy_app_yaml_to_cwd, run_workflow + + +# pylint: disable=redefined-outer-name + + +@pytest.fixture(scope="session") +def chord_err_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a temporary output directory for tests related to testing the + chord_err workflow. + + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary ouptut directory we'll be using for this test run. + + Returns: + The path to the temporary testing directory for chord_err workflow tests. + """ + return create_testing_dir(temp_output_dir, "chord_err_testing") + + +@pytest.fixture(scope="session") +def chord_err_name() -> FixtureStr: + """ + Defines a specific name to use for the chord_err workflow. This helps ensure + that even if changes were made to the chord_err workflow, tests using this fixture + should still run the same thing. + + Returns: + A string representing the name to use for the chord_err workflow. + """ + return "chord_err_test" + + +@pytest.fixture(scope="session") +def chord_err_setup( + chord_err_testing_dir: FixtureStr, + chord_err_name: FixtureStr, + path_to_test_specs: FixtureStr, +) -> ChordErrorSetup: + """ + Fixture for setting up the environment required for testing the chord error workflow. + + This fixture prepares the necessary configuration and paths for executing tests related + to the chord error workflow. It aggregates the required parameters into a single + [`ChordErrorSetup`][fixture_data_classes.ChordErrorSetup] data class instance, which + simplifies the management of these parameters in tests. + + Args: + chord_err_testing_dir: The path to the temporary output directory where chord error + workflow tests will store their results. + chord_err_name: A string representing the name to use for the chord error workflow. + path_to_test_specs: The base path to the Merlin test specs directory, which is + used to locate the chord error YAML file. + + Returns: + A [`ChordErrorSetup`][fixture_data_classes.ChordErrorSetup] instance containing + the testing directory, name, and path to the chord error YAML file, which can + be used in tests that require this setup. + """ + chord_err_path = os.path.join(path_to_test_specs, "chord_err.yaml") + return ChordErrorSetup( + testing_dir=chord_err_testing_dir, + name=chord_err_name, + path=chord_err_path, + ) + + +@pytest.fixture(scope="class") +def chord_err_run_workflow( + redis_broker_and_backend_class: RedisBrokerAndBackend, + chord_err_setup: ChordErrorSetup, + merlin_server_dir: FixtureStr, +) -> subprocess.CompletedProcess: + """ + Run the chord error workflow. + + This fixture sets up and executes the chord error workflow using the specified configurations + and parameters. It prepares the environment by modifying the CONFIG object to connect to a + Redis server and runs the workflow with the provided name and output path. + + Args: + redis_broker_and_backend_class: Fixture for setting up Redis broker and + backend for class-scoped tests. + chord_err_setup: A fixture that returns a [`ChordErrorSetup`][fixture_data_classes.ChordErrorSetup] + instance. + merlin_server_dir: A fixture to provide the path to the merlin_server directory that will be + created by the [`redis_server`][conftest.redis_server] fixture. + + Returns: + The completed process object containing information about the execution of the workflow, including + return code, stdout, and stderr. + """ + # Setup the test + copy_app_yaml_to_cwd(merlin_server_dir) + # chord_err_path = os.path.join(path_to_test_specs, "chord_err.yaml") + + # Create the variables to pass in to the workflow + vars_to_substitute = [f"NAME={chord_err_setup.name}", f"OUTPUT_PATH={chord_err_setup.testing_dir}"] + + # Run the workflow + return run_workflow(redis_broker_and_backend_class.client, chord_err_setup.path, vars_to_substitute) diff --git a/tests/fixtures/examples.py b/tests/fixtures/examples.py new file mode 100644 index 000000000..4096b0e76 --- /dev/null +++ b/tests/fixtures/examples.py @@ -0,0 +1,22 @@ +""" +Fixtures specifically for help testing the modules in the examples/ directory. +""" + +import pytest + +from tests.fixture_types import FixtureCallable, FixtureStr + + +@pytest.fixture(scope="session") +def examples_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a temporary output directory for tests related to the examples functionality. + + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary output directory we'll be using for this test run. + + Returns: + The path to the temporary testing directory for examples tests. + """ + return create_testing_dir(temp_output_dir, "examples_testing") diff --git a/tests/fixtures/feature_demo.py b/tests/fixtures/feature_demo.py new file mode 100644 index 000000000..ca77f23fb --- /dev/null +++ b/tests/fixtures/feature_demo.py @@ -0,0 +1,136 @@ +""" +Fixtures specifically for help testing the feature_demo workflow. +""" + +import os +import subprocess + +import pytest + +from tests.fixture_data_classes import FeatureDemoSetup, RedisBrokerAndBackend +from tests.fixture_types import FixtureCallable, FixtureInt, FixtureStr +from tests.integration.helper_funcs import copy_app_yaml_to_cwd, run_workflow + + +# pylint: disable=redefined-outer-name + + +@pytest.fixture(scope="session") +def feature_demo_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a temporary output directory for tests related to testing the + feature_demo workflow. + + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary ouptut directory we'll be using for this test run. + + Returns: + The path to the temporary testing directory for feature_demo workflow tests. + """ + return create_testing_dir(temp_output_dir, "feature_demo_testing") + + +@pytest.fixture(scope="session") +def feature_demo_num_samples() -> FixtureInt: + """ + Defines a specific number of samples to use for the feature_demo workflow. + This helps ensure that even if changes were made to the feature_demo workflow, + tests using this fixture should still run the same thing. + + Returns: + An integer representing the number of samples to use in the feature_demo workflow. + """ + return 8 + + +@pytest.fixture(scope="session") +def feature_demo_name() -> FixtureStr: + """ + Defines a specific name to use for the feature_demo workflow. This helps ensure + that even if changes were made to the feature_demo workflow, tests using this fixture + should still run the same thing. + + Returns: + A string representing the name to use for the feature_demo workflow. + """ + return "feature_demo_test" + + +@pytest.fixture(scope="session") +def feature_demo_setup( + feature_demo_testing_dir: FixtureStr, + feature_demo_num_samples: FixtureInt, + feature_demo_name: FixtureStr, + path_to_merlin_codebase: FixtureStr, +) -> FeatureDemoSetup: + """ + Fixture for setting up the environment required for testing the feature demo workflow. + + This fixture prepares the necessary configuration and paths for executing tests related + to the feature demo workflow. It aggregates the required parameters into a single + [`FeatureDemoSetup`][fixture_data_classes.FeatureDemoSetup] data class instance, which + simplifies the management of these parameters in tests. + + Args: + feature_demo_testing_dir: The path to the temporary output directory where + feature demo workflow tests will store their results. + feature_demo_num_samples: An integer representing the number of samples + to use in the feature demo workflow. + feature_demo_name: A string representing the name to use for the feature + demo workflow. + path_to_merlin_codebase: The base path to the Merlin codebase, which is + used to locate the feature demo YAML file. + + Returns: + A [`FeatureDemoSetup`][fixture_data_classes.FeatureDemoSetup] instance containing + the testing directory, number of samples, name, and path to the feature demo + YAML file, which can be used in tests that require this setup. + """ + demo_workflow = os.path.join("examples", "workflows", "feature_demo", "feature_demo.yaml") + feature_demo_path = os.path.join(path_to_merlin_codebase, demo_workflow) + return FeatureDemoSetup( + testing_dir=feature_demo_testing_dir, + num_samples=feature_demo_num_samples, + name=feature_demo_name, + path=feature_demo_path, + ) + + +@pytest.fixture(scope="class") +def feature_demo_run_workflow( + redis_broker_and_backend_class: RedisBrokerAndBackend, + feature_demo_setup: FeatureDemoSetup, + merlin_server_dir: FixtureStr, +) -> subprocess.CompletedProcess: + """ + Run the feature demo workflow. + + This fixture sets up and executes the feature demo workflow using the specified configurations + and parameters. It prepares the environment by modifying the CONFIG object to connect to a + Redis server and runs the demo workflow with the provided sample size and name. + + Args: + redis_broker_and_backend_class: Fixture for setting up Redis broker and + backend for class-scoped tests. + feature_demo_setup: A fixture that returns a [`FeatureDemoSetup`][fixture_data_classes.FeatureDemoSetup] + instance. + merlin_server_dir: A fixture to provide the path to the merlin_server directory that will be + created by the [`redis_server`][conftest.redis_server] fixture. + + Returns: + The completed process object containing information about the execution of the workflow, including + return code, stdout, and stderr. + """ + # Setup the test + copy_app_yaml_to_cwd(merlin_server_dir) + + # Create the variables to pass in to the workflow + vars_to_substitute = [ + f"N_SAMPLES={feature_demo_setup.num_samples}", + f"NAME={feature_demo_setup.name}", + f"OUTPUT_PATH={feature_demo_setup.testing_dir}", + ] + + # Run the workflow + return run_workflow(redis_broker_and_backend_class.client, feature_demo_setup.path, vars_to_substitute) diff --git a/tests/fixtures/run_command.py b/tests/fixtures/run_command.py new file mode 100644 index 000000000..5a666209b --- /dev/null +++ b/tests/fixtures/run_command.py @@ -0,0 +1,23 @@ +""" +Fixtures specifically for help testing the `merlin run` command. +""" + +import pytest + +from tests.fixture_types import FixtureCallable, FixtureStr + + +@pytest.fixture(scope="session") +def run_command_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a temporary output directory for tests related to testing the + `merlin run` functionality. + + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary ouptut directory we'll be using for this test run. + + Returns: + The path to the temporary testing directory for `merlin run` tests. + """ + return create_testing_dir(temp_output_dir, "run_command_testing") diff --git a/tests/fixtures/server.py b/tests/fixtures/server.py new file mode 100644 index 000000000..c2bcdc762 --- /dev/null +++ b/tests/fixtures/server.py @@ -0,0 +1,310 @@ +""" +Fixtures specifically for help testing the modules in the server/ directory. +""" + +import os +from argparse import Namespace +from typing import Dict, Union + +import pytest +import yaml + +from tests.fixture_types import FixtureCallable, FixtureDict, FixtureNamespace, FixtureStr + + +# pylint: disable=redefined-outer-name + + +@pytest.fixture(scope="session") +def server_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a temporary output directory for tests related to the server functionality. + + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary ouptut directory we'll be using for this test run. + + Returns: + The path to the temporary testing directory for server tests. + """ + return create_testing_dir(temp_output_dir, "server_testing") + + +@pytest.fixture(scope="session") +def server_redis_conf_file(server_testing_dir: FixtureStr) -> FixtureStr: + """ + Fixture to write a redis.conf file to the temporary output directory. + + If a test will modify this file with a file write, you should make a copy of + this file to modify instead. + + :param server_testing_dir: A pytest fixture that defines a path to the output directory we'll write to + :returns: The path to the redis configuration file we'll use for testing + """ + redis_conf_file = f"{server_testing_dir}/redis.conf" + file_contents = """ + # ip address + bind 127.0.0.1 + + # port + port 6379 + + # password + requirepass merlin_password + + # directory + dir ./ + + # snapshot + save 300 100 + + # db file + dbfilename dump.rdb + + # append mode + appendfsync everysec + + # append file + appendfilename appendonly.aof + + # dummy trailing comment + """.strip().replace( + " ", "" + ) + + with open(redis_conf_file, "w") as rcf: + rcf.write(file_contents) + + return redis_conf_file + + +@pytest.fixture(scope="session") +def server_redis_pass_file(server_testing_dir: FixtureStr) -> FixtureStr: + """ + Fixture to create a redis password file in the temporary output directory. + + If a test will modify this file with a file write, you should make a copy of + this file to modify instead. + + :param server_testing_dir: A pytest fixture that defines a path to the output directory we'll write to + :returns: The path to the redis password file + """ + redis_pass_file = f"{server_testing_dir}/redis.pass" + + with open(redis_pass_file, "w") as rpf: + rpf.write("server-tests-password") + + return redis_pass_file + + +@pytest.fixture(scope="session") +def server_users() -> FixtureDict[str, Dict[str, str]]: + """ + Create a dictionary of two test users with identical configuration settings. + + :returns: A dict containing the two test users and their settings + """ + users = { + "default": { + "channels": "*", + "commands": "@all", + "hash_password": "1ba9249af0c73dacb0f9a70567126624076b5bee40de811e65f57eabcdaf490a", + "keys": "*", + "status": "on", + }, + "test_user": { + "channels": "*", + "commands": "@all", + "hash_password": "1ba9249af0c73dacb0f9a70567126624076b5bee40de811e65f57eabcdaf490a", + "keys": "*", + "status": "on", + }, + } + return users + + +@pytest.fixture(scope="session") +def server_redis_users_file(server_testing_dir: FixtureStr, server_users: FixtureDict[str, Dict[str, str]]) -> FixtureStr: + """ + Fixture to write a redis.users file to the temporary output directory. + + If a test will modify this file with a file write, you should make a copy of + this file to modify instead. + + :param server_testing_dir: A pytest fixture that defines a path to the output directory we'll write to + :param server_users: A dict of test user configurations + :returns: The path to the redis user configuration file we'll use for testing + """ + redis_users_file = f"{server_testing_dir}/redis.users" + + with open(redis_users_file, "w") as ruf: + yaml.dump(server_users, ruf) + + return redis_users_file + + +@pytest.fixture(scope="class") +def server_container_config_data( + server_testing_dir: FixtureStr, + server_redis_conf_file: FixtureStr, + server_redis_pass_file: FixtureStr, + server_redis_users_file: FixtureStr, +) -> FixtureDict[str, str]: + """ + Fixture to provide sample data for ContainerConfig tests. + + :param server_testing_dir: A pytest fixture that defines a path to the output directory we'll write to + :param server_redis_conf_file: A pytest fixture that defines a path to a redis configuration file + :param server_redis_pass_file: A pytest fixture that defines a path to a redis password file + :param server_redis_users_file: A pytest fixture that defines a path to a redis users file + :returns: A dict containing the necessary key/values for the ContainerConfig object + """ + + return { + "format": "singularity", + "image_type": "redis", + "image": "redis_latest.sif", + "url": "docker://redis", + "config": server_redis_conf_file.split("/")[-1], + "config_dir": server_testing_dir, + "pfile": "merlin_server.pf", + "pass_file": server_redis_pass_file.split("/")[-1], + "user_file": server_redis_users_file.split("/")[-1], + } + + +@pytest.fixture(scope="class") +def server_container_format_config_data() -> FixtureDict[str, str]: + """ + Fixture to provide sample data for ContainerFormatConfig tests + + :returns: A dict containing the necessary key/values for the ContainerFormatConfig object + """ + return { + "command": "singularity", + "run_command": "{command} run -H {home_dir} {image} {config}", + "stop_command": "kill", + "pull_command": "{command} pull {image} {url}", + } + + +@pytest.fixture(scope="class") +def server_process_config_data() -> FixtureDict[str, str]: + """ + Fixture to provide sample data for ProcessConfig tests + + :returns: A dict containing the necessary key/values for the ProcessConfig object + """ + return { + "status": "pgrep -P {pid}", + "kill": "kill {pid}", + } + + +@pytest.fixture(scope="class") +def server_server_config( + server_container_config_data: FixtureDict[str, str], + server_process_config_data: FixtureDict[str, str], + server_container_format_config_data: FixtureDict[str, str], +) -> FixtureDict[str, FixtureDict[str, str]]: + """ + Fixture to provide sample data for ServerConfig tests + + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + :param server_process_config_data: A pytest fixture of test data to pass to the ProcessConfig class + :param server_container_format_config_data: A pytest fixture of test data to pass to the ContainerFormatConfig class + :returns: A dictionary containing each of the configuration dicts we'll need + """ + return { + "container": server_container_config_data, + "process": server_process_config_data, + "singularity": server_container_format_config_data, + } + + +@pytest.fixture(scope="function") +def server_app_yaml_contents( + server_redis_pass_file: FixtureStr, + server_container_config_data: FixtureDict[str, str], + server_process_config_data: FixtureDict[str, str], +) -> FixtureDict[str, Union[str, int]]: + """ + Fixture to create the contents of an app.yaml file. + + :param server_redis_pass_file: A pytest fixture that defines a path to a redis password file + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + :param server_process_config_data: A pytest fixture of test data to pass to the ProcessConfig class + :returns: A dict with typical app.yaml contents + """ + contents = { + "broker": { + "cert_reqs": "none", + "name": "redis", + "password": server_redis_pass_file, + "port": 6379, + "server": "127.0.0.1", + "username": "default", + "vhost": "testhost", + }, + "container": server_container_config_data, + "process": server_process_config_data, + "results_backend": { + "cert_reqs": "none", + "db_num": 0, + "name": "redis", + "password": server_redis_pass_file, + "port": 6379, + "server": "127.0.0.1", + "username": "default", + }, + } + return contents + + +@pytest.fixture(scope="function") +def server_app_yaml(server_testing_dir: FixtureStr, server_app_yaml_contents: FixtureDict[str, Union[str, int]]) -> FixtureStr: + """ + Fixture to create an app.yaml file in the temporary output directory. + + NOTE this must be function scoped since server_app_yaml_contents is function scoped. + + :param server_testing_dir: A pytest fixture that defines a path to the output directory we'll write to + :param server_app_yaml_contents: A pytest fixture that creates a dict of contents for an app.yaml file + :returns: The path to the app.yaml file + """ + app_yaml_file = f"{server_testing_dir}/app.yaml" + + if not os.path.exists(app_yaml_file): + with open(app_yaml_file, "w") as ayf: + yaml.dump(server_app_yaml_contents, ayf) + + return app_yaml_file + + +@pytest.fixture(scope="function") +def server_process_file_contents() -> FixtureDict[str, Union[str, int]]: + """Fixture to represent process file contents.""" + return {"parent_pid": 123, "image_pid": 456, "port": 6379, "hostname": "dummy_server"} + + +@pytest.fixture(scope="function") +def server_config_server_args() -> FixtureNamespace: + """ + Setup an argparse Namespace with all args that the `config_server` + function will need. These can be modified on a test-by-test basis. + + :returns: An argparse Namespace with args needed by `config_server` + """ + return Namespace( + ipaddress=None, + port=None, + password=None, + directory=None, + snapshot_seconds=None, + snapshot_changes=None, + snapshot_file=None, + append_mode=None, + append_file=None, + add_user=None, + remove_user=None, + ) diff --git a/tests/fixtures/status.py b/tests/fixtures/status.py index f26cea37c..57ff16bac 100644 --- a/tests/fixtures/status.py +++ b/tests/fixtures/status.py @@ -11,26 +11,30 @@ import pytest import yaml +from tests.fixture_types import FixtureCallable, FixtureNamespace, FixtureStr from tests.unit.study.status_test_files import status_test_variables +# pylint: disable=redefined-outer-name + + @pytest.fixture(scope="session") -def status_testing_dir(temp_output_dir: str) -> str: +def status_testing_dir(create_testing_dir: FixtureCallable, temp_output_dir: FixtureStr) -> FixtureStr: """ A pytest fixture to set up a temporary directory to write files to for testing status. - :param temp_output_dir: The path to the temporary output directory we'll be using for this test run - :returns: The path to the temporary testing directory for status testing - """ - testing_dir = f"{temp_output_dir}/status_testing/" - if not os.path.exists(testing_dir): - os.mkdir(testing_dir) + Args: + create_testing_dir: A fixture which returns a function that creates the testing directory. + temp_output_dir: The path to the temporary ouptut directory we'll be using for this test run. - return testing_dir + Returns: + The path to the temporary testing directory for status tests. + """ + return create_testing_dir(temp_output_dir, "status_testing") -@pytest.fixture(scope="session") -def status_empty_file(status_testing_dir: str) -> str: # pylint: disable=W0621 +@pytest.fixture(scope="class") +def status_empty_file(status_testing_dir: FixtureStr) -> FixtureStr: """ A pytest fixture to create an empty status file. @@ -46,7 +50,7 @@ def status_empty_file(status_testing_dir: str) -> str: # pylint: disable=W0621 @pytest.fixture(scope="session") -def status_spec_path(status_testing_dir: str) -> str: # pylint: disable=W0621 +def status_spec_path(status_testing_dir: FixtureStr) -> FixtureStr: # pylint: disable=W0621 """ Copy the test spec to the temp directory and modify the OUTPUT_PATH in the spec to point to the temp location. @@ -91,7 +95,7 @@ def set_sample_path(output_workspace: str): @pytest.fixture(scope="session") -def status_output_workspace(status_testing_dir: str) -> str: # pylint: disable=W0621 +def status_output_workspace(status_testing_dir: FixtureStr) -> FixtureStr: # pylint: disable=W0621 """ A pytest fixture to copy the test output workspace for status to the temporary status testing directory. @@ -107,7 +111,7 @@ def status_output_workspace(status_testing_dir: str) -> str: # pylint: disable= @pytest.fixture(scope="function") -def status_args(): +def status_args() -> FixtureNamespace: """ A pytest fixture to set up a namespace with all the arguments necessary for the Status object. @@ -127,7 +131,7 @@ def status_args(): @pytest.fixture(scope="session") -def status_nested_workspace(status_testing_dir: str) -> str: # pylint: disable=W0621 +def status_nested_workspace(status_testing_dir: FixtureStr) -> FixtureStr: # pylint: disable=W0621 """ Create an output workspace that contains another output workspace within one of its steps. In this case it will copy the status test workspace then within the 'just_samples' diff --git a/tests/integration/commands/__init__.py b/tests/integration/commands/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/integration/commands/pgen.py b/tests/integration/commands/pgen.py new file mode 100644 index 000000000..6973317d6 --- /dev/null +++ b/tests/integration/commands/pgen.py @@ -0,0 +1,36 @@ +""" +This file contains pgen functionality for testing purposes. +It's specifically set up to work with the feature demo example. +""" + +import random + +from maestrowf.datastructures.core import ParameterGenerator + + +# pylint complains about unused argument `env` but it's necessary for Maestro +def get_custom_generator(env, **kwargs): # pylint: disable=unused-argument + """ + Custom parameter generator that's used for testing the `--pgen` flag + of the `merlin run` command. + """ + p_gen = ParameterGenerator() + + # Unpack any pargs passed in + x2_min = int(kwargs.get("X2_MIN", "0")) + x2_max = int(kwargs.get("X2_MAX", "1")) + n_name_min = int(kwargs.get("N_NAME_MIN", "0")) + n_name_max = int(kwargs.get("N_NAME_MAX", "10")) + + # We'll only have two parameter entries each just for testing + num_points = 2 + + params = { + "X2": {"values": [random.uniform(x2_min, x2_max) for _ in range(num_points)], "label": "X2.%%"}, + "N_NEW": {"values": [random.randint(n_name_min, n_name_max) for _ in range(num_points)], "label": "N_NEW.%%"}, + } + + for key, value in params.items(): + p_gen.add_parameter(key, value["values"], value["label"]) + + return p_gen diff --git a/tests/integration/commands/test_purge.py b/tests/integration/commands/test_purge.py new file mode 100644 index 000000000..91340e946 --- /dev/null +++ b/tests/integration/commands/test_purge.py @@ -0,0 +1,363 @@ +""" +This module will contain the testing logic +for the `merlin purge` command. +""" + +import os +import subprocess +from typing import Dict, List, Tuple, Union + +from merlin.spec.expansion import get_spec_with_expansion +from tests.context_managers.celery_task_manager import CeleryTaskManager +from tests.fixture_data_classes import RedisBrokerAndBackend +from tests.fixture_types import FixtureRedis, FixtureStr +from tests.integration.conditions import HasRegex, HasReturnCode +from tests.integration.helper_funcs import check_test_conditions, copy_app_yaml_to_cwd + + +class TestPurgeCommand: + """ + Tests for the `merlin purge` command. + """ + + demo_workflow = os.path.join("examples", "workflows", "feature_demo", "feature_demo.yaml") + + def setup_test(self, path_to_merlin_codebase: FixtureStr, merlin_server_dir: FixtureStr) -> str: + """ + Setup the test environment for these tests by: + 1. Copying the app.yaml file created by the `redis_server` fixture to the cwd so that + Merlin can connect to the test server. + 2. Obtaining the path to the feature_demo spec that we'll use for these tests. + + Args: + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + + Returns: + The path to the feature_demo spec file. + """ + copy_app_yaml_to_cwd(merlin_server_dir) + return os.path.join(path_to_merlin_codebase, self.demo_workflow) + + def setup_tasks(self, celery_task_manager: CeleryTaskManager, spec_file: str) -> Tuple[Dict[str, str], int]: + """ + Helper method to setup tasks in the specified queues. + + This method sends tasks named 'task_for_{queue}' to each queue defined in the + provided spec file and returns the total number of queues that received tasks. + + Args: + celery_task_manager: + A context manager for managing Celery tasks, used to send tasks to the server. + spec_file: + The path to the spec file from which queues will be extracted. + + Returns: + A tuple with: + - A dictionary where the keys are step names and values are their associated queues. + - The number of queues that received tasks + """ + spec = get_spec_with_expansion(spec_file) + queues_in_spec = spec.get_task_queues() + + for queue in queues_in_spec.values(): + celery_task_manager.send_task(f"task_for_{queue}", queue=queue) + + return queues_in_spec, len(queues_in_spec.values()) + + def run_purge( + self, + spec_file: str, + input_value: str = None, + force: bool = False, + steps_to_purge: List[str] = None, + ) -> Dict[str, Union[str, int]]: + """ + Helper method to run the purge command. + + Args: + spec_file: The path to the spec file from which queues will be purged. + input_value: Any input we need to send to the subprocess. + force: If True, add the `-f` option to the purge command. + steps_to_purge: An optional list of steps to send to the purge command. + + Returns: + The result from executing the command in a subprocess. + """ + purge_cmd = ( + "merlin purge" + + (" -f" if force else "") + + f" {spec_file}" + + (f" --steps {' '.join(steps_to_purge)}" if steps_to_purge is not None else "") + ) + result = subprocess.run(purge_cmd, shell=True, capture_output=True, text=True, input=input_value) + return { + "stdout": result.stdout, + "stderr": result.stderr, + "return_code": result.returncode, + } + + def check_queues( + self, + redis_client: FixtureRedis, + queues_in_spec: Dict[str, str], + expected_task_count: int, + steps_to_purge: List[str] = None, + ): + """ + Check the state of queues in Redis against expected task counts. + + When `steps_to_purge` is set, the `expected_task_count` will represent the + number of expected tasks in the queues that _are not_ associated with the + steps in the `steps_to_purge` list. + + Args: + redis_client: The Redis client instance. + queues_in_spec: A dictionary of queues to check. + expected_task_count: The expected number of tasks in the queues (0 or 1). + steps_to_purge: Optional list of steps to determine which queues should be purged. + """ + for queue in queues_in_spec.values(): + # Brackets are special chars in regex so we have to add \ to make them literal + queue = queue.replace("[", "\\[").replace("]", "\\]") + matching_queues_on_server = redis_client.keys(pattern=f"{queue}*") + + for matching_queue in matching_queues_on_server: + tasks = redis_client.lrange(matching_queue, 0, -1) + if steps_to_purge and matching_queue in [queues_in_spec[step] for step in steps_to_purge]: + assert len(tasks) == 0, f"Expected 0 tasks in {matching_queue}, found {len(tasks)}." + else: + assert ( + len(tasks) == expected_task_count + ), f"Expected {expected_task_count} tasks in {matching_queue}, found {len(tasks)}." + + def test_no_options_tasks_exist_y( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + ): + """ + Test the `merlin purge` command with no options added and + tasks sent to the server. This should come up with a y/N + prompt in which we type 'y'. This should then purge the + tasks from the server. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + feature_demo = self.setup_test(path_to_merlin_codebase, merlin_server_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client) as celery_task_manager: + # Send tasks to the server for every queue in the spec + queues_in_spec, num_queues = self.setup_tasks(celery_task_manager, feature_demo) + + # Run the purge test + test_info = self.run_purge(feature_demo, input_value="y") + + # Make sure the subprocess ran and the correct output messages are given + conditions = [ + HasReturnCode(), + HasRegex("Are you sure you want to delete all tasks?"), + HasRegex(f"Purged {num_queues} messages from {num_queues} known task queues."), + ] + check_test_conditions(conditions, test_info) + + # Check on the Redis queues to ensure they were purged + self.check_queues(redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=0) + + def test_no_options_no_tasks_y( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + ): + """ + Test the `merlin purge` command with no options added and + no tasks sent to the server. This should come up with a y/N + prompt in which we type 'y'. This should then give us a "No + messages purged" log. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + feature_demo = self.setup_test(path_to_merlin_codebase, merlin_server_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client): + # Get the queues from the spec file + spec = get_spec_with_expansion(feature_demo) + queues_in_spec = spec.get_task_queues() + num_queues = len(queues_in_spec.values()) + + # Check that there are no tasks in the queues before we run the purge command + self.check_queues(redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=0) + + # Run the purge test + test_info = self.run_purge(feature_demo, input_value="y") + + # Make sure the subprocess ran and the correct output messages are given + conditions = [ + HasReturnCode(), + HasRegex("Are you sure you want to delete all tasks?"), + HasRegex(f"No messages purged from {num_queues} queues."), + ] + check_test_conditions(conditions, test_info) + + # Check that the Redis server still has no tasks + self.check_queues(redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=0) + + def test_no_options_n( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + ): + """ + Test the `merlin purge` command with no options added and + tasks sent to the server. This should come up with a y/N + prompt in which we type 'N'. This should take us out of the + command without purging the tasks. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + feature_demo = self.setup_test(path_to_merlin_codebase, merlin_server_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client) as celery_task_manager: + # Send tasks to the server for every queue in the spec + queues_in_spec, num_queues = self.setup_tasks(celery_task_manager, feature_demo) + + # Run the purge test + test_info = self.run_purge(feature_demo, input_value="N") + + # Make sure the subprocess ran and the correct output messages are given + conditions = [ + HasReturnCode(), + HasRegex("Are you sure you want to delete all tasks?"), + HasRegex(f"Purged {num_queues} messages from {num_queues} known task queues.", negate=True), + ] + check_test_conditions(conditions, test_info) + + # Check on the Redis queues to ensure they were not purged + self.check_queues(redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=1) + + def test_force_option( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + ): + """ + Test the `merlin purge` command with the `--force` option + enabled. This should not bring up a y/N prompt and should + immediately purge all tasks. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + feature_demo = self.setup_test(path_to_merlin_codebase, merlin_server_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client) as celery_task_manager: + # Send tasks to the server for every queue in the spec + queues_in_spec, num_queues = self.setup_tasks(celery_task_manager, feature_demo) + + # Run the purge test + test_info = self.run_purge(feature_demo, force=True) + + # Make sure the subprocess ran and the correct output messages are given + conditions = [ + HasReturnCode(), + HasRegex("Are you sure you want to delete all tasks?", negate=True), + HasRegex(f"Purged {num_queues} messages from {num_queues} known task queues."), + ] + check_test_conditions(conditions, test_info) + + # Check on the Redis queues to ensure they were purged + self.check_queues(redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=0) + + def test_steps_option( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + ): + """ + Test the `merlin purge` command with the `--steps` option + enabled. This should only purge the tasks in the task queues + associated with the steps provided. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + feature_demo = self.setup_test(path_to_merlin_codebase, merlin_server_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client) as celery_task_manager: + # Send tasks to the server for every queue in the spec + queues_in_spec, _ = self.setup_tasks(celery_task_manager, feature_demo) + + # Run the purge test + steps_to_purge = ["hello", "collect"] + test_info = self.run_purge(feature_demo, input_value="y", steps_to_purge=steps_to_purge) + + # Make sure the subprocess ran and the correct output messages are given + num_steps_to_purge = len(steps_to_purge) + conditions = [ + HasReturnCode(), + HasRegex("Are you sure you want to delete all tasks?"), + HasRegex(f"Purged {num_steps_to_purge} messages from {num_steps_to_purge} known task queues."), + ] + check_test_conditions(conditions, test_info) + + # Check on the Redis queues to ensure they were not purged + self.check_queues( + redis_broker_and_backend_function.client, queues_in_spec, expected_task_count=1, steps_to_purge=steps_to_purge + ) diff --git a/tests/integration/commands/test_run.py b/tests/integration/commands/test_run.py new file mode 100644 index 000000000..c7b27b3bb --- /dev/null +++ b/tests/integration/commands/test_run.py @@ -0,0 +1,468 @@ +""" +This module will contain the testing logic +for the `merlin run` command. +""" + +import csv +import os +import re +import shutil +import subprocess +from typing import Dict, Union + +from merlin.spec.expansion import get_spec_with_expansion +from tests.context_managers.celery_task_manager import CeleryTaskManager +from tests.fixture_data_classes import RedisBrokerAndBackend +from tests.fixture_types import FixtureStr +from tests.integration.conditions import HasReturnCode, PathExists, StepFinishedFilesCount +from tests.integration.helper_funcs import check_test_conditions, copy_app_yaml_to_cwd + + +# pylint: disable=import-outside-toplevel,unused-argument + + +class TestRunCommand: + """ + Base class for testing the `merlin run` command. + """ + + demo_workflow = os.path.join("examples", "workflows", "feature_demo", "feature_demo.yaml") + + def setup_test_environment( + self, path_to_merlin_codebase: FixtureStr, merlin_server_dir: FixtureStr, run_command_testing_dir: FixtureStr + ) -> str: + """ + Setup the test environment for these tests by: + 1. Moving into the temporary output directory created specifically for these tests. + 2. Copying the app.yaml file created by the `redis_server` fixture to the cwd so that + Merlin can connect to the test server. + 3. Obtaining the path to the feature_demo spec that we'll use for these tests. + + Args: + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + + Returns: + The path to the feature_demo spec file. + """ + os.chdir(run_command_testing_dir) + copy_app_yaml_to_cwd(merlin_server_dir) + return os.path.join(path_to_merlin_codebase, self.demo_workflow) + + def run_merlin_command(self, command: str) -> Dict[str, Union[str, int]]: + """ + Open a subprocess and run the command specified by the `command` parameter. + Ensure this command runs successfully and return the process results. + + Args: + command: The command to execute in a subprocess. + + Returns: + The results from executing the command in a subprocess. + + Raises: + AssertionError: If the command fails (non-zero return code). + """ + result = subprocess.run(command, shell=True, capture_output=True, text=True) + return { + "stdout": result.stdout, + "stderr": result.stderr, + "return_code": result.returncode, + } + + def get_output_workspace_from_logs(self, test_info: Dict[str, Union[str, int]]) -> str: + """ + Extracts the workspace path from the provided standard output and error logs. + + This method searches for a specific message indicating the study workspace + in the combined logs (both stdout and stderr). The expected message format + is: "Study workspace is ''". If the message is found, + the method returns the extracted workspace path. If the message is not + found, an assertion error is raised. + + Args: + test_info: The results from executing our test. + + Returns: + The extracted workspace path from the logs. + + Raises: + AssertionError: If the expected message is not found in the combined logs. + """ + workspace_pattern = re.compile(r"Study workspace is '(\S+)'") + combined_output = test_info["stdout"] + test_info["stderr"] + match = workspace_pattern.search(combined_output) + assert match, "No 'Study workspace is...' message found in command output." + return match.group(1) + + +class TestRunCommandDistributed(TestRunCommand): + """ + Tests for the `merlin run` command that are run in a distributed manner + rather than being run locally. + """ + + def test_distributed_run( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + run_command_testing_dir: FixtureStr, + ): + """ + This test verifies that tasks can be successfully sent to a Redis server + using the `merlin run` command with no flags. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + """ + from merlin.celery import app as celery_app + + # Setup the testing environment + feature_demo = self.setup_test_environment(path_to_merlin_codebase, merlin_server_dir, run_command_testing_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client): + # Send tasks to the server + test_info = self.run_merlin_command(f"merlin run {feature_demo} --vars NAME=run_command_test_distributed_run") + + # Check that the test ran properly + check_test_conditions([HasReturnCode()], test_info) + + # Get the queues we need to query + spec = get_spec_with_expansion(feature_demo) + queues_in_spec = spec.get_task_queues() + + for queue in queues_in_spec.values(): + # Brackets are special chars in regex so we have to add \ to make them literal + queue = queue.replace("[", "\\[").replace("]", "\\]") + matching_queues_on_server = redis_broker_and_backend_function.client.keys(pattern=f"{queue}*") + + # Make sure any queues that exist on the server have tasks in them + for matching_queue in matching_queues_on_server: + tasks = redis_broker_and_backend_function.client.lrange(matching_queue, 0, -1) + assert len(tasks) > 0 + + def test_samplesfile_option( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + run_command_testing_dir: FixtureStr, + ): + """ + This test verifies that passing in a samples filepath from the command line will + substitute in the file properly. It should copy the samples file that's passed + in to the merlin_info subdirectory. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + """ + from merlin.celery import app as celery_app + + # Setup the testing environment + feature_demo = self.setup_test_environment(path_to_merlin_codebase, merlin_server_dir, run_command_testing_dir) + + # Create a new samples file to pass into our test workflow + data = [ + ["X1, Value 1", "X2, Value 1"], + ["X1, Value 2", "X2, Value 2"], + ["X1, Value 3", "X2, Value 3"], + ] + sample_filename = "test_samplesfile.csv" + new_samples_file = os.path.join(run_command_testing_dir, sample_filename) + with open(new_samples_file, mode="w", newline="") as file: + writer = csv.writer(file) + writer.writerows(data) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client): + # Send tasks to the server + test_info = self.run_merlin_command( + f"merlin run {feature_demo} --vars NAME=run_command_test_samplesfile_option --samplesfile {new_samples_file}" + ) + + # Check that the test ran properly and created the correct directories/files + expected_workspace_path = self.get_output_workspace_from_logs(test_info) + conditions = [ + HasReturnCode(), + PathExists(expected_workspace_path), + PathExists(os.path.join(expected_workspace_path, "merlin_info", sample_filename)), + ] + check_test_conditions(conditions, test_info) + + def test_pgen_and_pargs_options( # pylint: disable=too-many-locals + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + run_command_testing_dir: FixtureStr, + ): + """ + Test the `--pgen` and `--pargs` options with the `merlin run` command. + This should update the parameter block of the expanded yaml file to have + 2 entries for both `X2` and `N_NEW`. The `X2` parameter should be between + `X2_MIN` and `X2_MAX`, and the `N_NEW` parameter should be between `N_NEW_MIN` + and `N_NEW_MAX`. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + """ + from merlin.celery import app as celery_app + + # Setup test vars and the testing environment + bounds = {"X2": (1, 2), "N_NEW": (5, 15)} + pgen_filepath = os.path.join( + os.path.abspath(os.path.expandvars(os.path.expanduser(os.path.dirname(__file__)))), "pgen.py" + ) + feature_demo = self.setup_test_environment(path_to_merlin_codebase, merlin_server_dir, run_command_testing_dir) + + with CeleryTaskManager(celery_app, redis_broker_and_backend_function.client): + # Send tasks to the server + test_info = self.run_merlin_command( + f"merlin run {feature_demo} " + "--vars NAME=run_command_test_pgen_and_pargs_options " + f"--pgen {pgen_filepath} " + f'--parg "X2_MIN:{bounds["X2"][0]}" ' + f'--parg "X2_MAX:{bounds["X2"][1]}" ' + f'--parg "N_NAME_MIN:{bounds["N_NEW"][0]}" ' + f'--parg "N_NAME_MAX:{bounds["N_NEW"][1]}"' + ) + + # Check that the test ran properly and created the correct directories/files + expected_workspace_path = self.get_output_workspace_from_logs(test_info) + expanded_yaml = os.path.join(expected_workspace_path, "merlin_info", "feature_demo.expanded.yaml") + conditions = [HasReturnCode(), PathExists(expected_workspace_path), PathExists(os.path.join(expanded_yaml))] + check_test_conditions(conditions, test_info) + + # Read in the parameters from the expanded yaml and ensure they're within the new bounds we provided + params = get_spec_with_expansion(expanded_yaml).get_parameters() + for param_name, (min_val, max_val) in bounds.items(): + for param in params.parameters[param_name]: + assert min_val <= param <= max_val + + +class TestRunCommandLocal(TestRunCommand): + """ + Tests for the `merlin run` command that are run in a locally rather + than in a distributed manner. + """ + + def test_dry_run( # pylint: disable=too-many-locals + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + run_command_testing_dir: FixtureStr, + ): + """ + Test the `merlin run` command's `--dry` option. This should create all the output + subdirectories for each step but it shouldn't execute anything for the steps. In + other words, the only file in each step subdirectory should be the .sh file. + + Note: + This test will run locally so that we don't have to worry about starting + & stopping workers. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + """ + # Setup the test environment + feature_demo = self.setup_test_environment(path_to_merlin_codebase, merlin_server_dir, run_command_testing_dir) + + # Run the test and grab the output workspace generated from it + test_info = self.run_merlin_command(f"merlin run {feature_demo} --vars NAME=run_command_test_dry_run --local --dry") + + # Check that the test ran properly and created the correct directories/files + expected_workspace_path = self.get_output_workspace_from_logs(test_info) + check_test_conditions([HasReturnCode(), PathExists(expected_workspace_path)], test_info) + + # Check that every step was ran by looking for an existing output workspace + for step in get_spec_with_expansion(feature_demo).get_study_steps(): + step_directory = os.path.join(expected_workspace_path, step.name) + assert os.path.exists(step_directory), f"Output directory for step '{step.name}' not found: {step_directory}" + + allowed_dry_run_files = {"MERLIN_STATUS.json", "status.lock"} + for dirpath, dirnames, filenames in os.walk(step_directory): + # Check if the current directory has no subdirectories (leaf directory) + if not dirnames: + # Check for unexpected files + unexpected_files = [ + file for file in filenames if file not in allowed_dry_run_files and not file.endswith(".sh") + ] + assert not unexpected_files, ( + f"Unexpected files found in {dirpath}: {unexpected_files}. " + f"Expected only .sh files or {allowed_dry_run_files}." + ) + + # Check that there is exactly one .sh file + sh_file_count = sum(1 for file in filenames if file.endswith(".sh")) + assert ( + sh_file_count == 1 + ), f"Expected exactly one .sh file in {dirpath} but found {sh_file_count} .sh files." + + def test_local_run( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_merlin_codebase: FixtureStr, + merlin_server_dir: FixtureStr, + run_command_testing_dir: FixtureStr, + ): + """ + This test verifies that tasks can be successfully executed locally using + the `merlin run` command with the `--local` flag. + + Args: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_merlin_codebase: + A fixture to provide the path to the directory containing Merlin's core + functionality. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + run_command_testing_dir: + The path to the the temp output directory for `merlin run` tests. + """ + # Setup the test environment + feature_demo = self.setup_test_environment(path_to_merlin_codebase, merlin_server_dir, run_command_testing_dir) + + # Run the test and grab the output workspace generated from it + study_name = "run_command_test_local_run" + num_samples = 8 + vars_dict = {"NAME": study_name, "OUTPUT_PATH": run_command_testing_dir, "N_SAMPLES": num_samples} + vars_str = " ".join(f"{key}={value}" for key, value in vars_dict.items()) + command = f"merlin run {feature_demo} --vars {vars_str} --local" + test_info = self.run_merlin_command(command) + + # Check that the test ran properly and created the correct directories/files + expected_workspace_path = self.get_output_workspace_from_logs(test_info) + conditions = [ + HasReturnCode(), + PathExists(expected_workspace_path), + StepFinishedFilesCount( # The rest of the conditions will ensure every step ran to completion + step="hello", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=num_samples, + ), + StepFinishedFilesCount( + step="python3_hello", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="collect", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="translate", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="learn", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="make_new_samples", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="predict", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="verify", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ), + ] + + # GitHub actions doesn't have a python2 path so we'll conditionally add this check + if shutil.which("python2"): + conditions.append( + StepFinishedFilesCount( + step="python2_hello", + study_name=study_name, + output_path=run_command_testing_dir, + num_parameters=1, + num_samples=0, + ) + ) + + check_test_conditions(conditions, test_info) + + # # Check that every step was ran by looking for an existing output workspace and MERLIN_FINISHED files + # for step in get_spec_with_expansion(feature_demo).get_study_steps(): + # step_directory = os.path.join(expected_workspace_path, step.name) + # assert os.path.exists(step_directory), f"Output directory for step '{step.name}' not found: {step_directory}" + # for dirpath, dirnames, filenames in os.walk(step_directory): + # # Check if the current directory has no subdirectories (leaf directory) + # if not dirnames: + # # Check for the existence of the MERLIN_FINISHED file + # assert ( + # "MERLIN_FINISHED" in filenames + # ), f"Expected a MERLIN_FINISHED file in list of files for {dirpath} but did not find one" + + +# pylint: enable=import-outside-toplevel,unused-argument diff --git a/tests/integration/commands/test_stop_and_query_workers.py b/tests/integration/commands/test_stop_and_query_workers.py new file mode 100644 index 000000000..beed6599b --- /dev/null +++ b/tests/integration/commands/test_stop_and_query_workers.py @@ -0,0 +1,380 @@ +""" +This module will contain the testing logic for +the `stop-workers` and `query-workers` commands. +""" + +import os +import subprocess +from contextlib import contextmanager +from enum import Enum +from typing import List + +import pytest + +from tests.context_managers.celery_workers_manager import CeleryWorkersManager +from tests.fixture_data_classes import RedisBrokerAndBackend +from tests.fixture_types import FixtureStr +from tests.integration.conditions import Condition, HasRegex +from tests.integration.helper_funcs import check_test_conditions, copy_app_yaml_to_cwd, load_workers_from_spec + + +# pylint: disable=unused-argument,import-outside-toplevel + + +class WorkerMessages(Enum): + """ + Enumerated strings to help keep track of the messages + that we're expecting (or not expecting) to see from the + tests in this module. + """ + + NO_WORKERS_MSG_STOP = "No workers found to stop" + NO_WORKERS_MSG_QUERY = "No workers found!" + STEP_1_WORKER = "step_1_merlin_test_worker" + STEP_2_WORKER = "step_2_merlin_test_worker" + OTHER_WORKER = "other_merlin_test_worker" + + +class TestStopAndQueryWorkersCommands: + """ + Tests for the `merlin stop-workers` and `merlin query-workers` commands. + Most of these tests will: + 1. Start workers from a spec file used for testing + - Use CeleryWorkerManager for this to ensure safe stoppage of workers + if something goes wrong + 2. Run the test command from a subprocess + """ + + @contextmanager + def run_test_with_workers( # pylint: disable=too-many-arguments + self, + path_to_test_specs: FixtureStr, + merlin_server_dir: FixtureStr, + conditions: List[Condition], + command: str, + flag: str = None, + ): + """ + Helper method to run common testing logic for tests with workers started. + This method must also be a context manager so we can check the status of the + workers prior to the CeleryWorkersManager running it's exit code that shuts down + all active workers. + + This method will: + 0. Read in the necessary fixtures as parameters. These fixtures grab paths to + our test specs and the merlin server directory created from starting the + containerized redis server. + 1. Load in the worker specifications from the `multiple_workers.yaml` file. + 2. Use a context manager to start up the workers on the celery app connected to + the containerized redis server + 3. Copy the app.yaml file for the containerized redis server to the current working + directory so that merlin will connect to it when we run our test + 4. Run the test command that's provided and check that the conditions given are + passing. + 5. Yield control back to the calling method. + 6. Safely terminate workers that may have not been stopped once the calling method + completes. + + Parameters: + path_to_test_specs: + A fixture to provide the path to the directory containing test specifications. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + conditions: + A list of `Condition` instances that need to pass in order for this test to + be successful. + command: + The command that we're testing. E.g. "merlin stop-workers" + flag: + An optional flag to add to the command that we're testing so we can test + different functionality for the command. + """ + from merlin.celery import app as celery_app + + # Grab worker configurations from the spec file + multiple_worker_spec = os.path.join(path_to_test_specs, "multiple_workers.yaml") + workers_from_spec = load_workers_from_spec(multiple_worker_spec) + + # We use a context manager to start workers so that they'll safely stop even if this test fails + with CeleryWorkersManager(celery_app) as workers_manager: + workers_manager.launch_workers(workers_from_spec) + + # Copy the app.yaml to the cwd so merlin will connect to the testing server + copy_app_yaml_to_cwd(merlin_server_dir) + + # Run the test + cmd_to_test = f"{command} {flag}" if flag else command + result = subprocess.run(cmd_to_test, capture_output=True, text=True, shell=True) + + info = { + "stdout": result.stdout, + "stderr": result.stderr, + "return_code": result.returncode, + } + + # Ensure all test conditions are satisfied + check_test_conditions(conditions, info) + + yield + + def get_no_workers_msg(self, command_to_test: str) -> WorkerMessages: + """ + Retrieve the appropriate "no workers" found message. + + This method checks the command to test and returns a corresponding + message based on whether the command is to stop workers or query for them. + + Returns: + The message indicating that no workers are available, depending on the + command being tested. + """ + no_workers_msg = None + if command_to_test == "merlin stop-workers": + no_workers_msg = WorkerMessages.NO_WORKERS_MSG_STOP.value + else: + no_workers_msg = WorkerMessages.NO_WORKERS_MSG_QUERY.value + return no_workers_msg + + @pytest.mark.parametrize("command_to_test", ["merlin stop-workers", "merlin query-workers"]) + def test_no_workers( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + merlin_server_dir: FixtureStr, + command_to_test: str, + ): + """ + Test the `merlin stop-workers` and `merlin query-workers` commands with no workers + started in the first place. + + This test will: + 0. Setup the pytest fixtures which include: + - starting a containerized Redis server + - updating the CONFIG object to point to the containerized Redis server + - obtaining the path to the merlin server directory created from starting + the containerized Redis server + 1. Copy the app.yaml file for the containerized redis server to the current working + directory so that merlin will connect to it when we run our test + 2. Run the test command that's provided and check that the conditions given are + passing. + + Parameters: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + command_to_test: + The command that we're testing, obtained from the parametrize call. + """ + conditions = [ + HasRegex(self.get_no_workers_msg(command_to_test)), + HasRegex(WorkerMessages.STEP_1_WORKER.value, negate=True), + HasRegex(WorkerMessages.STEP_2_WORKER.value, negate=True), + HasRegex(WorkerMessages.OTHER_WORKER.value, negate=True), + ] + + # Copy the app.yaml to the cwd so merlin will connect to the testing server + copy_app_yaml_to_cwd(merlin_server_dir) + + # Run the test + result = subprocess.run(command_to_test, capture_output=True, text=True, shell=True) + info = { + "stdout": result.stdout, + "stderr": result.stderr, + "return_code": result.returncode, + } + + # Ensure all test conditions are satisfied + check_test_conditions(conditions, info) + + @pytest.mark.parametrize("command_to_test", ["merlin stop-workers", "merlin query-workers"]) + def test_no_flags( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_test_specs: FixtureStr, + merlin_server_dir: FixtureStr, + command_to_test: str, + ): + """ + Test the `merlin stop-workers` and `merlin query-workers` commands with no flags. + + Run the commands referenced above and ensure the text output from Merlin is correct. + For the `stop-workers` command, we check if all workers are stopped as well. + To see more information on exactly what this test is doing, see the + `run_test_with_workers()` method. + + Parameters: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_test_specs: + A fixture to provide the path to the directory containing test specifications. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + command_to_test: + The command that we're testing, obtained from the parametrize call. + """ + conditions = [ + HasRegex(self.get_no_workers_msg(command_to_test), negate=True), + HasRegex(WorkerMessages.STEP_1_WORKER.value), + HasRegex(WorkerMessages.STEP_2_WORKER.value), + HasRegex(WorkerMessages.OTHER_WORKER.value), + ] + with self.run_test_with_workers(path_to_test_specs, merlin_server_dir, conditions, command_to_test): + if command_to_test == "merlin stop-workers": + # After the test runs and before the CeleryWorkersManager exits, ensure there are no workers on the app + from merlin.celery import app as celery_app + + active_queues = celery_app.control.inspect().active_queues() + assert active_queues is None + + @pytest.mark.parametrize("command_to_test", ["merlin stop-workers", "merlin query-workers"]) + def test_spec_flag( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_test_specs: FixtureStr, + merlin_server_dir: FixtureStr, + command_to_test: str, + ): + """ + Test the `merlin stop-workers` and `merlin query-workers` commands with the `--spec` + flag. + + Run the commands referenced above with the `--spec` flag and ensure the text output + from Merlin is correct. For the `stop-workers` command, we check if all workers defined + in the spec file are stopped as well. To see more information on exactly what this test + is doing, see the `run_test_with_workers()` method. + + Parameters: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_test_specs: + A fixture to provide the path to the directory containing test specifications. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + command_to_test: + The command that we're testing, obtained from the parametrize call. + """ + conditions = [ + HasRegex(self.get_no_workers_msg(command_to_test), negate=True), + HasRegex(WorkerMessages.STEP_1_WORKER.value), + HasRegex(WorkerMessages.STEP_2_WORKER.value), + HasRegex(WorkerMessages.OTHER_WORKER.value), + ] + with self.run_test_with_workers( + path_to_test_specs, + merlin_server_dir, + conditions, + command_to_test, + flag=f"--spec {os.path.join(path_to_test_specs, 'multiple_workers.yaml')}", + ): + if command_to_test == "merlin stop-workers": + from merlin.celery import app as celery_app + + active_queues = celery_app.control.inspect().active_queues() + assert active_queues is None + + @pytest.mark.parametrize("command_to_test", ["merlin stop-workers", "merlin query-workers"]) + def test_workers_flag( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_test_specs: FixtureStr, + merlin_server_dir: FixtureStr, + command_to_test: str, + ): + """ + Test the `merlin stop-workers` and `merlin query-workers` commands with the `--workers` + flag. + + Run the commands referenced above with the `--workers` flag and ensure the text output + from Merlin is correct. For the `stop-workers` command, we check to make sure that all + workers given with this flag are stopped. To see more information on exactly what this + test is doing, see the `run_test_with_workers()` method. + + Parameters: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_test_specs: + A fixture to provide the path to the directory containing test specifications. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + command_to_test: + The command that we're testing, obtained from the parametrize call. + """ + conditions = [ + HasRegex(self.get_no_workers_msg(command_to_test), negate=True), + HasRegex(WorkerMessages.STEP_1_WORKER.value), + HasRegex(WorkerMessages.STEP_2_WORKER.value), + HasRegex(WorkerMessages.OTHER_WORKER.value, negate=True), + ] + with self.run_test_with_workers( + path_to_test_specs, + merlin_server_dir, + conditions, + command_to_test, + flag=f"--workers {WorkerMessages.STEP_1_WORKER.value} {WorkerMessages.STEP_2_WORKER.value}", + ): + if command_to_test == "merlin stop-workers": + from merlin.celery import app as celery_app + + active_queues = celery_app.control.inspect().active_queues() + worker_name = f"celery@{WorkerMessages.OTHER_WORKER.value}" + assert worker_name in active_queues + + @pytest.mark.parametrize("command_to_test", ["merlin stop-workers", "merlin query-workers"]) + def test_queues_flag( + self, + redis_broker_and_backend_function: RedisBrokerAndBackend, + path_to_test_specs: FixtureStr, + merlin_server_dir: FixtureStr, + command_to_test: str, + ): + """ + Test the `merlin stop-workers` and `merlin query-workers` commands with the `--queues` + flag. + + Run the commands referenced above with the `--queues` flag and ensure the text output + from Merlin is correct. For the `stop-workers` command, we check that only the workers + attached to the given queues are stopped. To see more information on exactly what this + test is doing, see the `run_test_with_workers()` method. + + Parameters: + redis_broker_and_backend_function: Fixture for setting up Redis broker and + backend for function-scoped tests. + path_to_test_specs: + A fixture to provide the path to the directory containing test specifications. + merlin_server_dir: + A fixture to provide the path to the merlin_server directory that will be + created by the `redis_server` fixture. + command_to_test: + The command that we're testing, obtained from the parametrize call. + """ + conditions = [ + HasRegex(self.get_no_workers_msg(command_to_test), negate=True), + HasRegex(WorkerMessages.STEP_1_WORKER.value), + HasRegex(WorkerMessages.STEP_2_WORKER.value, negate=True), + HasRegex(WorkerMessages.OTHER_WORKER.value, negate=True), + ] + with self.run_test_with_workers( + path_to_test_specs, + merlin_server_dir, + conditions, + command_to_test, + flag="--queues hello_queue", + ): + if command_to_test == "merlin stop-workers": + from merlin.celery import app as celery_app + + active_queues = celery_app.control.inspect().active_queues() + workers_that_should_be_alive = [ + f"celery@{WorkerMessages.OTHER_WORKER.value}", + f"celery@{WorkerMessages.STEP_2_WORKER.value}", + ] + for worker_name in workers_that_should_be_alive: + assert worker_name in active_queues + + +# pylint: enable=unused-argument,import-outside-toplevel diff --git a/tests/integration/conditions.py b/tests/integration/conditions.py index 83f07aafe..2ae0aa02a 100644 --- a/tests/integration/conditions.py +++ b/tests/integration/conditions.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -34,7 +34,6 @@ from re import search -# TODO when moving command line tests to pytest, change Condition boolean returns to assertions class Condition(ABC): """Abstract Condition class that other conditions will inherit from""" @@ -131,7 +130,7 @@ def __init__(self, study_name, output_path): """ self.study_name = study_name self.output_path = output_path - self.dirpath_glob = f"{self.output_path}/{self.study_name}" f"_[0-9]*-[0-9]*" + self.dirpath_glob = os.path.join(self.output_path, f"{self.study_name}_[0-9]*-[0-9]*") def glob(self, glob_string): """ @@ -154,7 +153,7 @@ class StepFileExists(StudyOutputAware): A StudyOutputAware that checks for a particular file's existence. """ - def __init__(self, step, filename, study_name, output_path, params=False): # pylint: disable=R0913 + def __init__(self, step, filename, study_name, output_path, params=False, samples=False): # pylint: disable=R0913 """ :param `step`: the name of a step :param `filename`: name of file to search for in step's workspace directory @@ -165,6 +164,7 @@ def __init__(self, step, filename, study_name, output_path, params=False): # py self.step = step self.filename = filename self.params = params + self.samples = samples def __str__(self): return f"{__class__.__name__} expected to find file '{self.glob_string}', but file did not exist" @@ -174,10 +174,9 @@ def glob_string(self): """ Returns a regex string for the glob library to recursively find files with. """ - param_glob = "" - if self.params: - param_glob = "*/" - return f"{self.dirpath_glob}/{self.step}/{param_glob}{self.filename}" + param_glob = "*" if self.params else "" + samples_glob = "**" if self.samples else "" + return os.path.join(self.dirpath_glob, self.step, param_glob, samples_glob, self.filename) def file_exists(self): """Check if the file path created by glob_string exists""" @@ -229,7 +228,7 @@ def contains(self): with open(filename, "r") as textfile: filetext = textfile.read() return self.is_within(filetext) - except Exception: # pylint: disable=W0718 + except Exception: # pylint: disable=broad-except return False def is_within(self, text): @@ -243,6 +242,108 @@ def passes(self): return self.contains() +# TODO when writing API docs for tests make sure this looks correct and has functioning links +# - Do we want to list expected_count, glob_string, and passes as methods since they're already attributes? +class StepFinishedFilesCount(StudyOutputAware): + """ + A [`StudyOutputAware`][integration.conditions.StudyOutputAware] that checks for the + exact number of `MERLIN_FINISHED` files in a specified step's output directory based + on the number of parameters and samples. + + Attributes: + step: The name of the step to check. + study_name: The name of the study. + output_path: The output path of the study. + num_parameters: The number of parameters for the step. + num_samples: The number of samples for the step. + expected_count: The expected number of `MERLIN_FINISHED` files based on parameters and samples or explicitly set. + glob_string: The glob pattern to find `MERLIN_FINISHED` files in the specified step's output directory. + passes: Checks if the count of `MERLIN_FINISHED` files matches the expected count. + + Methods: + expected_count: Calculates the expected number of `MERLIN_FINISHED` files. + glob_string: Constructs the glob pattern for searching `MERLIN_FINISHED` files. + count_finished_files: Counts the number of `MERLIN_FINISHED` files found. + passes: Checks if the count of `MERLIN_FINISHED` files matches the expected count. + """ + + # All of these parameters are necessary for this Condition so we'll ignore pylint + def __init__( + self, + step: str, + study_name: str, + output_path: str, + num_parameters: int = 0, + num_samples: int = 0, + expected_count: int = None, + ): # pylint: disable=too-many-arguments + super().__init__(study_name, output_path) + self.step = step + self.num_parameters = num_parameters + self.num_samples = num_samples + self._expected_count = expected_count + + @property + def expected_count(self) -> int: + """ + Calculate the expected number of `MERLIN_FINISHED` files. + + Returns: + The expected number of `MERLIN_FINISHED` files. + """ + # Return the explicitly set expected count if given + if self._expected_count is not None: + return self._expected_count + + # Otherwise calculate the correct number of MERLIN_FINISHED files to expect + if self.num_parameters > 0 and self.num_samples > 0: + return self.num_parameters * self.num_samples + if self.num_parameters > 0: + return self.num_parameters + if self.num_samples > 0: + return self.num_samples + + return 1 # Default case when there are no parameters or samples + + @property + def glob_string(self) -> str: + """ + Glob pattern to find `MERLIN_FINISHED` files in the specified step's output directory. + + Returns: + A glob pattern to find `MERLIN_FINISHED` files. + """ + param_glob = "*" if self.num_parameters > 0 else "" + samples_glob = "**" if self.num_samples > 0 else "" + return os.path.join(self.dirpath_glob, self.step, param_glob, samples_glob, "MERLIN_FINISHED") + + def count_finished_files(self) -> int: + """ + Count the number of `MERLIN_FINISHED` files found. + + Returns: + The actual number of `MERLIN_FINISHED` files that exist in the step's output directory. + """ + finished_files = glob(self.glob_string) # Adjust the glob pattern as needed + return len(finished_files) + + @property + def passes(self) -> bool: + """ + Check if the count of `MERLIN_FINISHED` files matches the expected count. + + Returns: + True if the expected count matches the actual count. False otherwise. + """ + return self.count_finished_files() == self.expected_count + + def __str__(self) -> str: + return ( + f"{__class__.__name__} expected {self.expected_count} `MERLIN_FINISHED` " + f"files, but found {self.count_finished_files()}" + ) + + class ProvenanceYAMLFileHasRegex(HasRegex): """ A condition that a Merlin provenance yaml spec in the 'merlin_info' directory @@ -339,7 +440,7 @@ def contains(self) -> bool: with open(self.filename, "r") as f: # pylint: disable=C0103 filetext = f.read() return self.is_within(filetext) - except Exception: # pylint: disable=W0718 + except Exception: # pylint: disable=broad-except return False def is_within(self, text): diff --git a/tests/integration/definitions.py b/tests/integration/definitions.py index 59c1fa256..4714275f5 100644 --- a/tests/integration/definitions.py +++ b/tests/integration/definitions.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -108,16 +108,12 @@ def define_tests(): # pylint: disable=R0914,R0915 workers_lsf = get_worker_by_cmd("jsrun", workers) run = f"merlin {err_lvl} run" restart = f"merlin {err_lvl} restart" - purge = "merlin purge" - stop = "merlin stop-workers" - query = "merlin query-workers" # Shortcuts for example workflow paths examples = "merlin/examples/workflows" dev_examples = "merlin/examples/dev_workflows" test_specs = "tests/integration/test_specs" demo = f"{examples}/feature_demo/feature_demo.yaml" - remote_demo = f"{examples}/remote_feature_demo/remote_feature_demo.yaml" demo_pgen = f"{examples}/feature_demo/scripts/pgen.py" simple = f"{examples}/simple_chain/simple_chain.yaml" slurm = f"{test_specs}/slurm_test.yaml" @@ -126,9 +122,7 @@ def define_tests(): # pylint: disable=R0914,R0915 flux_restart = f"{examples}/flux/flux_par_restart.yaml" flux_native = f"{test_specs}/flux_par_native_test.yaml" lsf = f"{examples}/lsf/lsf_par.yaml" - mul_workers_demo = f"{dev_examples}/multiple_workers.yaml" cli_substitution_wf = f"{test_specs}/cli_substitution_test.yaml" - chord_err_wf = f"{test_specs}/chord_err.yaml" # Other shortcuts black = "black --check --target-version py36" @@ -650,213 +644,6 @@ def define_tests(): # pylint: disable=R0914,R0915 "run type": "local", }, } - stop_workers_tests = { - "stop workers no workers": { - "cmds": f"{stop}", - "conditions": [ - HasReturnCode(), - HasRegex("No workers found to stop"), - HasRegex("step_1_merlin_test_worker", negate=True), - HasRegex("step_2_merlin_test_worker", negate=True), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - }, - "stop workers no flags": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{stop}", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found to stop", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker"), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "stop workers with spec flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{stop} --spec {mul_workers_demo}", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found to stop", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker"), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "stop workers with workers flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{stop} --workers step_1_merlin_test_worker step_2_merlin_test_worker", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found to stop", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "stop workers with queues flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{stop} --queues hello_queue", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found to stop", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker", negate=True), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - } - query_workers_tests = { - "query workers no workers": { - "cmds": f"{query}", - "conditions": [ - HasReturnCode(), - HasRegex("No workers found!"), - HasRegex("step_1_merlin_test_worker", negate=True), - HasRegex("step_2_merlin_test_worker", negate=True), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - }, - "query workers no flags": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{query}", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found!", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker"), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "query workers with spec flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{query} --spec {mul_workers_demo}", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found!", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker"), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "query workers with workers flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{query} --workers step_1_merlin_test_worker step_2_merlin_test_worker", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found!", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker"), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - "query workers with queues flag": { - "cmds": [ - f"{workers} {mul_workers_demo}", - f"{query} --queues hello_queue", - ], - "conditions": [ - HasReturnCode(), - HasRegex("No workers found!", negate=True), - HasRegex("step_1_merlin_test_worker"), - HasRegex("step_2_merlin_test_worker", negate=True), - HasRegex("other_merlin_test_worker", negate=True), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - }, - } - distributed_tests = { # noqa: F841 - "run and purge feature_demo": { - "cmds": f"{run} {demo}; {purge} {demo} -f", - "conditions": HasReturnCode(), - "run type": "distributed", - }, - "remote feature_demo": { - "cmds": f"""{run} {remote_demo} --vars OUTPUT_PATH=./{OUTPUT_DIR} WORKER_NAME=cli_test_demo_workers; - {workers} {remote_demo} --vars OUTPUT_PATH=./{OUTPUT_DIR} WORKER_NAME=cli_test_demo_workers""", - "conditions": [ - HasReturnCode(), - ProvenanceYAMLFileHasRegex( - regex="cli_test_demo_workers:", - spec_file_name="remote_feature_demo", - study_name="feature_demo", - output_path=OUTPUT_DIR, - provenance_type="expanded", - ), - StepFileExists( - "verify", - "MERLIN_FINISHED", - "feature_demo", - OUTPUT_DIR, - params=True, - ), - ], - "run type": "distributed", - }, - } - distributed_error_checks = { - "check chord error continues wf": { - "cmds": [ - f"{workers} {chord_err_wf} --vars OUTPUT_PATH=./{OUTPUT_DIR}", - f"{run} {chord_err_wf} --vars OUTPUT_PATH=./{OUTPUT_DIR}; sleep 40; tree {OUTPUT_DIR}", - ], - "conditions": [ - HasReturnCode(), - PathExists( # Check that the sample that's supposed to raise an error actually raises an error - f"{OUTPUT_DIR}/process_samples/01/MERLIN_FINISHED", - negate=True, - ), - StepFileExists( # Check that step 3 is actually started and completes - "step_3", - "MERLIN_FINISHED", - "chord_err", - OUTPUT_DIR, - ), - ], - "run type": "distributed", - "cleanup": KILL_WORKERS, - "num procs": 2, - } - } # combine and return test dictionaries all_tests = {} @@ -876,10 +663,6 @@ def define_tests(): # pylint: disable=R0914,R0915 # provenence_equality_checks, # omitting provenance equality check because it is broken # style_checks, # omitting style checks due to different results on different machines dependency_checks, - stop_workers_tests, - query_workers_tests, - distributed_tests, - distributed_error_checks, ]: all_tests.update(test_dict) diff --git a/tests/integration/helper_funcs.py b/tests/integration/helper_funcs.py new file mode 100644 index 000000000..4837b516b --- /dev/null +++ b/tests/integration/helper_funcs.py @@ -0,0 +1,164 @@ +""" +This module contains helper functions for the integration +test suite. +""" + +import os +import re +import shutil +import subprocess +from time import sleep +from typing import Dict, List + +from merlin.spec.expansion import get_spec_with_expansion +from tests.context_managers.celery_task_manager import CeleryTaskManager +from tests.context_managers.celery_workers_manager import CeleryWorkersManager +from tests.fixture_types import FixtureRedis +from tests.integration.conditions import Condition + + +def load_workers_from_spec(spec_filepath: str) -> dict: + """ + Load worker specifications from a YAML file. + + This function reads a YAML file containing study specifications and + extracts the worker information under the "merlin" section. It + constructs a dictionary in the form that + [`CeleryWorkersManager.launch_workers`][context_managers.celery_workers_manager.CeleryWorkersManager.launch_workers] + requires. + + Args: + spec_filepath: The file path to the YAML specification file. + + Returns: + A dictionary containing the worker specifications from the + "merlin" section of the YAML file. + """ + worker_info = {} + spec = get_spec_with_expansion(spec_filepath) + steps_and_queues = spec.get_task_queues(omit_tag=True) + + for worker_name, worker_settings in spec.merlin["resources"]["workers"].items(): + match = re.search(r"--concurrency\s+(\d+)", worker_settings["args"]) + concurrency = int(match.group(1)) if match else 1 + worker_info[worker_name] = {"concurrency": concurrency} + if worker_settings["steps"] == ["all"]: + worker_info[worker_name]["queues"] = list(steps_and_queues.values()) + else: + worker_info[worker_name]["queues"] = [steps_and_queues[step] for step in worker_settings["steps"]] + + return worker_info + + +def copy_app_yaml_to_cwd(merlin_server_dir: str): + """ + Copy the app.yaml file from the directory provided to the current working + directory. + + Grab the app.yaml file from `merlin_server_dir` and copy it to the current + working directory so that Merlin will read this in as the server configuration + for whatever test is calling this. + + Args: + merlin_server_dir: The path to the `merlin_server` directory that should be created by the + [`redis_server`][conftest.redis_server] fixture. + """ + copied_app_yaml = os.path.join(os.getcwd(), "app.yaml") + if not os.path.exists(copied_app_yaml): + server_app_yaml = os.path.join(merlin_server_dir, "app.yaml") + shutil.copy(server_app_yaml, copied_app_yaml) + + +def check_test_conditions(conditions: List[Condition], info: Dict[str, str]): + """ + Ensure all specified test conditions are satisfied based on the output + from a subprocess. + + This function iterates through a list of [`Condition`][integration.conditions.Condition] + instances, ingests the provided information (stdout, stderr, and return + code) for each condition, and checks if each condition passes. If any + condition fails, an AssertionError is raised with a detailed message that + includes the condition that failed, along with the captured output and + return code. + + Args: + conditions: A list of Condition instances that define the expectations for the test. + info: A dictionary containing the output from the subprocess, which should + include the following keys:\n + - 'stdout': The standard output captured from the subprocess. + - 'stderr': The standard error output captured from the subprocess. + - 'return_code': The return code of the subprocess, indicating success + or failure of the command executed. + + Raises: + AssertionError: If any of the conditions do not pass, an AssertionError is raised with + a detailed message including the failed condition and the subprocess + output. + """ + for condition in conditions: + condition.ingest_info(info) + try: + assert condition.passes + except AssertionError as exc: + error_message = ( + f"Condition failed: {condition}\n" + f"Captured stdout: {info['stdout']}\n" + f"Captured stderr: {info['stderr']}\n" + f"Return code: {info['return_code']}\n" + ) + raise AssertionError(error_message) from exc + + +def run_workflow(redis_client: FixtureRedis, workflow_path: str, vars_to_substitute: List[str]) -> subprocess.CompletedProcess: + """ + Run a Merlin workflow using the `merlin run` and `merlin run-workers` commands. + + This function executes a Merlin workflow using a specified path to a study and variables to + configure the study with. It utilizes context managers to safely send tasks to the server + and start up workers. The tasks are given 15 seconds to be sent to the server. Once tasks + exist on the server, the workflow is given 30 seconds to run to completion, which should be + plenty of time. + + Args: + redis_client: A fixture that connects us to a redis client that we can interact with. + workflow_path: The path to the study that we're going to run here + vars_to_substitute: A list of variables in the form ["VAR_NAME=var_value"] to be modified + in the workflow. + + Returns: + The completed process object containing information about the execution of the workflow, including + return code, stdout, and stderr. + """ + from merlin.celery import app as celery_app # pylint: disable=import-outside-toplevel + + run_workers_proc = None + + with CeleryTaskManager(celery_app, redis_client): + # Send the tasks to the server + try: + subprocess.run( + f"merlin run {workflow_path} --vars {' '.join(vars_to_substitute)}", + shell=True, + capture_output=True, + text=True, + timeout=15, + ) + except subprocess.TimeoutExpired as exc: + raise TimeoutError("Could not send tasks to the server within the allotted time.") from exc + + # We use a context manager to start workers so that they'll safely stop even if this test fails + with CeleryWorkersManager(celery_app) as celery_worker_manager: + # Start the workers then add them to the context manager so they can be stopped safely later + run_workers_proc = subprocess.Popen( # pylint: disable=consider-using-with + f"merlin run-workers {workflow_path}".split(), + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) + celery_worker_manager.add_run_workers_process(run_workers_proc.pid) + + # Let the workflow try to run for 30 seconds + sleep(30) + + return run_workers_proc diff --git a/tests/integration/run_tests.py b/tests/integration/run_tests.py index ef2bef882..48d558f33 100644 --- a/tests/integration/run_tests.py +++ b/tests/integration/run_tests.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -39,10 +39,9 @@ from contextlib import suppress from subprocess import TimeoutExpired, run +from definitions import OUTPUT_DIR, define_tests # pylint: disable=E0401 from tabulate import tabulate -from tests.integration.definitions import OUTPUT_DIR, define_tests # pylint: disable=E0401 - def get_definition_issues(test): """ diff --git a/tests/unit/study/test_celeryadapter.py b/tests/integration/test_celeryadapter.py similarity index 68% rename from tests/unit/study/test_celeryadapter.py rename to tests/integration/test_celeryadapter.py index 0572d6c66..89241088b 100644 --- a/tests/unit/study/test_celeryadapter.py +++ b/tests/integration/test_celeryadapter.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # @@ -34,12 +34,9 @@ import json import os from datetime import datetime -from time import sleep from typing import Dict -import pytest from celery import Celery -from celery.canvas import Signature from deepdiff import DeepDiff from merlin.config import Config @@ -49,150 +46,153 @@ from tests.unit.study.status_test_files.status_test_variables import SPEC_PATH -@pytest.mark.order(before="TestInactive") -class TestActive: - """ - This class will test functions in the celeryadapter.py module. - It will run tests where we need active queues/workers to interact with. - - NOTE: The tests in this class must be ran before the TestInactive class or else the - Celery workers needed for this class don't start - - TODO: fix the bug noted above and then check if we still need pytest-order - """ - - def test_query_celery_queues( - self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 - ): - """ - Test the query_celery_queues function by providing it with a list of active queues. - This should return a dict where keys are queue names and values are more dicts containing - the number of jobs and consumers in that queue. - - :param `celery_app`: A pytest fixture for the test Celery app - :param launch_workers: A pytest fixture that launches celery workers for us to interact with - :param worker_queue_map: A pytest fixture that returns a dict of workers and queues - """ - # Set up a dummy configuration to use in the test - dummy_config = Config({"broker": {"name": "redis"}}) - - # Get the actual output - queues_to_query = list(worker_queue_map.values()) - actual_queue_info = celeryadapter.query_celery_queues(queues_to_query, app=celery_app, config=dummy_config) - - # Ensure all 3 queues in worker_queue_map were queried before looping - assert len(actual_queue_info) == 3 - - # Ensure each queue has a worker attached - for queue_name, queue_info in actual_queue_info.items(): - assert queue_name in worker_queue_map.values() - assert queue_info == {"consumers": 1, "jobs": 0} - - def test_get_running_queues(self, launch_workers: "Fixture", worker_queue_map: Dict[str, str]): # noqa: F821 - """ - Test the get_running_queues function with queues active. - This should return a list of active queues. - - :param `launch_workers`: A pytest fixture that launches celery workers for us to interact with - :param `worker_queue_map`: A pytest fixture that returns a dict of workers and queues - """ - result = celeryadapter.get_running_queues("merlin_test_app", test_mode=True) - assert sorted(result) == sorted(list(worker_queue_map.values())) - - def test_get_active_celery_queues( - self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 - ): - """ - Test the get_active_celery_queues function with queues active. - This should return a tuple where the first entry is a dict of queue info - and the second entry is a list of worker names. - - :param `celery_app`: A pytest fixture for the test Celery app - :param `launch_workers`: A pytest fixture that launches celery workers for us to interact with - :param `worker_queue_map`: A pytest fixture that returns a dict of workers and queues - """ - # Start the queues and run the test - queue_result, worker_result = celeryadapter.get_active_celery_queues(celery_app) - - # Ensure we got output before looping - assert len(queue_result) == len(worker_result) == 3 - - for worker, queue in worker_queue_map.items(): - # Check that the entry in the queue_result dict for this queue is correct - assert queue in queue_result - assert len(queue_result[queue]) == 1 - assert worker in queue_result[queue][0] - - # Remove this entry from the queue_result dict - del queue_result[queue] - - # Check that this worker was added to the worker_result list - worker_found = False - for worker_name in worker_result[:]: - if worker in worker_name: - worker_found = True - worker_result.remove(worker_name) - break - assert worker_found - - # Ensure there was no extra output that we weren't expecting - assert queue_result == {} - assert worker_result == [] - - def test_build_set_of_queues( - self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 - ): - """ - Test the build_set_of_queues function with queues active. - This should return a set of queues (the queues defined in setUp). - """ - # Run the test - result = celeryadapter.build_set_of_queues( - steps=["all"], spec=None, specific_queues=None, verbose=False, app=celery_app - ) - assert result == set(worker_queue_map.values()) - - @pytest.mark.order(index=1) - def test_check_celery_workers_processing_tasks( - self, - celery_app: Celery, - sleep_sig: Signature, - launch_workers: "Fixture", # noqa: F821 - ): - """ - Test the check_celery_workers_processing function with workers active and a task in a queue. - This function will query workers for any tasks they're still processing. We'll send a - a task that sleeps for 3 seconds to our workers before we run this test so that there should be - a task for this function to find. - - NOTE: the celery app fixture shows strange behavior when using app.control.inspect() calls (which - check_celery_workers_processing uses) so we have to run this test first in this class in order to - have it run properly. - - :param celery_app: A pytest fixture for the test Celery app - :param sleep_sig: A pytest fixture for a celery signature of a task that sleeps for 3 sec - :param launch_workers: A pytest fixture that launches celery workers for us to interact with - """ - # Our active workers/queues are test_worker_[0-2]/test_queue_[0-2] so we're - # sending this to test_queue_0 for test_worker_0 to process - queue_for_signature = "test_queue_0" - sleep_sig.set(queue=queue_for_signature) - result = sleep_sig.delay() - - # We need to give the task we just sent to the server a second to get picked up by the worker - sleep(1) - - # Run the test now that the task should be getting processed - active_queue_test = celeryadapter.check_celery_workers_processing([queue_for_signature], celery_app) - assert active_queue_test is True - - # Now test that a queue without any tasks returns false - # We sent the signature to task_queue_0 so task_queue_1 shouldn't have any tasks to find - non_active_queue_test = celeryadapter.check_celery_workers_processing(["test_queue_1"], celery_app) - assert non_active_queue_test is False - - # Wait for the worker to finish running the task - result.get() +# from time import sleep +# import pytest +# from celery.canvas import Signature +# @pytest.mark.order(before="TestInactive") +# class TestActive: +# """ +# This class will test functions in the celeryadapter.py module. +# It will run tests where we need active queues/workers to interact with. + +# NOTE: The tests in this class must be ran before the TestInactive class or else the +# Celery workers needed for this class don't start + +# TODO: fix the bug noted above and then check if we still need pytest-order +# """ + +# def test_query_celery_queues( +# self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 +# ): +# """ +# Test the query_celery_queues function by providing it with a list of active queues. +# This should return a dict where keys are queue names and values are more dicts containing +# the number of jobs and consumers in that queue. + +# :param `celery_app`: A pytest fixture for the test Celery app +# :param launch_workers: A pytest fixture that launches celery workers for us to interact with +# :param worker_queue_map: A pytest fixture that returns a dict of workers and queues +# """ +# # Set up a dummy configuration to use in the test +# dummy_config = Config({"broker": {"name": "redis"}}) + +# # Get the actual output +# queues_to_query = list(worker_queue_map.values()) +# actual_queue_info = celeryadapter.query_celery_queues(queues_to_query, app=celery_app, config=dummy_config) + +# # Ensure all 3 queues in worker_queue_map were queried before looping +# assert len(actual_queue_info) == 3 + +# # Ensure each queue has a worker attached +# for queue_name, queue_info in actual_queue_info.items(): +# assert queue_name in worker_queue_map.values() +# assert queue_info == {"consumers": 1, "jobs": 0} + +# def test_get_running_queues(self, launch_workers: "Fixture", worker_queue_map: Dict[str, str]): # noqa: F821 +# """ +# Test the get_running_queues function with queues active. +# This should return a list of active queues. + +# :param `launch_workers`: A pytest fixture that launches celery workers for us to interact with +# :param `worker_queue_map`: A pytest fixture that returns a dict of workers and queues +# """ +# result = celeryadapter.get_running_queues("merlin_test_app", test_mode=True) +# assert sorted(result) == sorted(list(worker_queue_map.values())) + +# def test_get_active_celery_queues( +# self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 +# ): +# """ +# Test the get_active_celery_queues function with queues active. +# This should return a tuple where the first entry is a dict of queue info +# and the second entry is a list of worker names. + +# :param `celery_app`: A pytest fixture for the test Celery app +# :param `launch_workers`: A pytest fixture that launches celery workers for us to interact with +# :param `worker_queue_map`: A pytest fixture that returns a dict of workers and queues +# """ +# # Start the queues and run the test +# queue_result, worker_result = celeryadapter.get_active_celery_queues(celery_app) + +# # Ensure we got output before looping +# assert len(queue_result) == len(worker_result) == 3 + +# for worker, queue in worker_queue_map.items(): +# # Check that the entry in the queue_result dict for this queue is correct +# assert queue in queue_result +# assert len(queue_result[queue]) == 1 +# assert worker in queue_result[queue][0] + +# # Remove this entry from the queue_result dict +# del queue_result[queue] + +# # Check that this worker was added to the worker_result list +# worker_found = False +# for worker_name in worker_result[:]: +# if worker in worker_name: +# worker_found = True +# worker_result.remove(worker_name) +# break +# assert worker_found + +# # Ensure there was no extra output that we weren't expecting +# assert queue_result == {} +# assert worker_result == [] + +# def test_build_set_of_queues( +# self, celery_app: Celery, launch_workers: "Fixture", worker_queue_map: Dict[str, str] # noqa: F821 +# ): +# """ +# Test the build_set_of_queues function with queues active. +# This should return a set of queues (the queues defined in setUp). +# """ +# # Run the test +# result = celeryadapter.build_set_of_queues( +# steps=["all"], spec=None, specific_queues=None, verbose=False, app=celery_app +# ) +# assert result == set(worker_queue_map.values()) + +# @pytest.mark.order(index=1) +# def test_check_celery_workers_processing_tasks( +# self, +# celery_app: Celery, +# sleep_sig: Signature, +# launch_workers: "Fixture", # noqa: F821 +# ): +# """ +# Test the check_celery_workers_processing function with workers active and a task in a queue. +# This function will query workers for any tasks they're still processing. We'll send a +# a task that sleeps for 3 seconds to our workers before we run this test so that there should be +# a task for this function to find. + +# NOTE: the celery app fixture shows strange behavior when using app.control.inspect() calls (which +# check_celery_workers_processing uses) so we have to run this test first in this class in order to +# have it run properly. + +# :param celery_app: A pytest fixture for the test Celery app +# :param sleep_sig: A pytest fixture for a celery signature of a task that sleeps for 3 sec +# :param launch_workers: A pytest fixture that launches celery workers for us to interact with +# """ +# # Our active workers/queues are test_worker_[0-2]/test_queue_[0-2] so we're +# # sending this to test_queue_0 for test_worker_0 to process +# queue_for_signature = "test_queue_0" +# sleep_sig.set(queue=queue_for_signature) +# result = sleep_sig.delay() + +# # We need to give the task we just sent to the server a second to get picked up by the worker +# sleep(1) + +# # Run the test now that the task should be getting processed +# active_queue_test = celeryadapter.check_celery_workers_processing([queue_for_signature], celery_app) +# assert active_queue_test is True + +# # Now test that a queue without any tasks returns false +# # We sent the signature to task_queue_0 so task_queue_1 shouldn't have any tasks to find +# non_active_queue_test = celeryadapter.check_celery_workers_processing(["test_queue_1"], celery_app) +# assert non_active_queue_test is False + +# # Wait for the worker to finish running the task +# result.get() class TestInactive: @@ -250,7 +250,7 @@ def test_get_running_queues(self): This should return an empty list. """ result = celeryadapter.get_running_queues("merlin_test_app", test_mode=True) - assert result == [] + assert not result def test_get_active_celery_queues(self, celery_app: Celery): """ @@ -261,8 +261,8 @@ def test_get_active_celery_queues(self, celery_app: Celery): :param `celery_app`: A pytest fixture for the test Celery app """ queue_result, worker_result = celeryadapter.get_active_celery_queues(celery_app) - assert queue_result == {} - assert worker_result == [] + assert not queue_result + assert not worker_result def test_check_celery_workers_processing_tasks(self, celery_app: Celery, worker_queue_map: Dict[str, str]): """ @@ -476,7 +476,7 @@ def test_dump_celery_queue_info_csv(self, worker_queue_map: Dict[str, str]): # Make sure the rest of the csv file was created as expected dump_diff = DeepDiff(csv_dump_output, expected_output) - assert dump_diff == {} + assert not dump_diff finally: try: os.remove(outfile) @@ -513,7 +513,7 @@ def test_dump_celery_queue_info_json(self, worker_queue_map: Dict[str, str]): # There should only be one entry in the json dump file so this will only 'loop' once for dump_entry in json_df_contents.values(): json_dump_diff = DeepDiff(dump_entry, expected_output) - assert json_dump_diff == {} + assert not json_dump_diff finally: try: os.remove(outfile) diff --git a/tests/integration/test_specs/chord_err.yaml b/tests/integration/test_specs/chord_err.yaml index 3da99ae03..9fe7d55ea 100644 --- a/tests/integration/test_specs/chord_err.yaml +++ b/tests/integration/test_specs/chord_err.yaml @@ -1,10 +1,11 @@ description: - name: chord_err + name: $(NAME) description: test the chord err problem env: variables: OUTPUT_PATH: ./studies + NAME: chord_err global.parameters: TEST_PARAM: diff --git a/tests/integration/test_specs/multiple_workers.yaml b/tests/integration/test_specs/multiple_workers.yaml new file mode 100644 index 000000000..967582a53 --- /dev/null +++ b/tests/integration/test_specs/multiple_workers.yaml @@ -0,0 +1,56 @@ +description: + name: multiple_workers + description: a very simple merlin workflow with multiple workers + +global.parameters: + GREET: + values : ["hello","hola"] + label : GREET.%% + WORLD: + values : ["world","mundo"] + label : WORLD.%% + +study: + - name: step_1 + description: say hello + run: + cmd: | + echo "$(GREET), $(WORLD)!" + task_queue: hello_queue + + - name: step_2 + description: step 2 + run: + cmd: | + echo "step_2" + depends: [step_1_*] + task_queue: echo_queue + + - name: step_3 + description: stop workers + run: + cmd: | + echo "stop workers" + depends: [step_2] + task_queue: other_queue + + - name: step_4 + description: another step + run: + cmd: | + echo "another step" + depends: [step_3] + task_queue: other_queue + +merlin: + resources: + workers: + step_1_merlin_test_worker: + args: -l INFO --concurrency 1 + steps: [step_1] + step_2_merlin_test_worker: + args: -l INFO --concurrency 1 + steps: [step_2] + other_merlin_test_worker: + args: -l INFO --concurrency 1 + steps: [step_3, step_4] diff --git a/tests/integration/workflows/test_chord_error.py b/tests/integration/workflows/test_chord_error.py new file mode 100644 index 000000000..e1f9fad0b --- /dev/null +++ b/tests/integration/workflows/test_chord_error.py @@ -0,0 +1,63 @@ +""" +This module contains tests for the feature_demo workflow. +""" + +import subprocess + +from tests.fixture_data_classes import ChordErrorSetup +from tests.integration.conditions import HasRegex, StepFinishedFilesCount +from tests.integration.helper_funcs import check_test_conditions + + +class TestChordError: + """ + Tests for the chord error workflow. + """ + + def test_chord_error_continues( + self, + chord_err_setup: ChordErrorSetup, + chord_err_run_workflow: subprocess.CompletedProcess, + ): + """ + Test that this workflow continues through to the end of its execution, even + though a ChordError will be raised. + + Args: + chord_err_setup: A fixture that returns a [`ChordErrorSetup`][fixture_data_classes.ChordErrorSetup] + instance. + chord_err_run_workflow: A fixture to run the chord error study. + """ + + conditions = [ + HasRegex("Exception raised by request from the user"), + StepFinishedFilesCount( # Check that the `process_samples` step has only 2 MERLIN_FINISHED files + step="process_samples", + study_name=chord_err_setup.name, + output_path=chord_err_setup.testing_dir, + expected_count=2, + num_samples=3, + ), + StepFinishedFilesCount( # Check that the `samples_and_params` step has all of its MERLIN_FINISHED files + step="samples_and_params", + study_name=chord_err_setup.name, + output_path=chord_err_setup.testing_dir, + num_parameters=2, + num_samples=3, + ), + StepFinishedFilesCount( # Check that the final step has a MERLIN_FINISHED file + step="step_3", + study_name=chord_err_setup.name, + output_path=chord_err_setup.testing_dir, + num_parameters=0, + num_samples=0, + ), + ] + + info = { + "return_code": chord_err_run_workflow.returncode, + "stdout": chord_err_run_workflow.stdout.read(), + "stderr": chord_err_run_workflow.stderr.read(), + } + + check_test_conditions(conditions, info) diff --git a/tests/integration/workflows/test_feature_demo.py b/tests/integration/workflows/test_feature_demo.py new file mode 100644 index 000000000..7974ebecc --- /dev/null +++ b/tests/integration/workflows/test_feature_demo.py @@ -0,0 +1,132 @@ +""" +This module contains tests for the feature_demo workflow. +""" + +import shutil +import subprocess + +from tests.fixture_data_classes import FeatureDemoSetup +from tests.integration.conditions import ProvenanceYAMLFileHasRegex, StepFinishedFilesCount + + +class TestFeatureDemo: + """ + Tests for the feature_demo workflow. + """ + + def test_end_to_end_run( + self, feature_demo_setup: FeatureDemoSetup, feature_demo_run_workflow: subprocess.CompletedProcess + ): + """ + Test that the workflow runs from start to finish with no problems. + + This will check that each step has the proper amount of `MERLIN_FINISHED` files. + The workflow will be run via the + [`feature_demo_run_workflow`][fixtures.feature_demo.feature_demo_run_workflow] + fixture. + + Args: + feature_demo_setup: A fixture that returns a + [`FeatureDemoSetup`][fixture_data_classes.FeatureDemoSetup] instance. + feature_demo_run_workflow: A fixture to run the feature demo study. + """ + conditions = [ + ProvenanceYAMLFileHasRegex( # This condition will check that variable substitution worked + regex=f"N_SAMPLES: {feature_demo_setup.num_samples}", + spec_file_name="feature_demo", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + provenance_type="expanded", + ), + StepFinishedFilesCount( # The rest of the conditions will ensure every step ran to completion + step="hello", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=feature_demo_setup.num_samples, + ), + StepFinishedFilesCount( + step="python3_hello", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="collect", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="translate", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="learn", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="make_new_samples", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="predict", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + StepFinishedFilesCount( + step="verify", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ), + ] + + # GitHub actions doesn't have a python2 path so we'll conditionally add this check + if shutil.which("python2"): + conditions.append( + StepFinishedFilesCount( + step="python2_hello", + study_name=feature_demo_setup.name, + output_path=feature_demo_setup.testing_dir, + num_parameters=1, + num_samples=0, + ) + ) + + for condition in conditions: + assert condition.passes + + # TODO implement the below tests + # def test_step_execution_order(self): + # """ + # Test that steps are executed in the correct order. + # """ + # # TODO build a list with the correct order that steps should be ran + # # TODO compare the list against the logs from the worker + + # def test_workflow_error_handling(self): + # """ + # Test the behavior when errors arise during the worfklow. + + # TODO should this test both soft and hard fails? should this test all return codes? + # """ + + # def test_data_passing(self): + # """ + # Test that data can be successfully passed between steps using built-in Merlin variables. + # """ diff --git a/tests/unit/common/test_dumper.py b/tests/unit/common/test_dumper.py new file mode 100644 index 000000000..c52e9fe90 --- /dev/null +++ b/tests/unit/common/test_dumper.py @@ -0,0 +1,168 @@ +""" +Tests for the `dumper.py` file. +""" + +import csv +import json +import os +from datetime import datetime +from time import sleep + +import pytest + +from merlin.common.dumper import dump_handler + + +NUM_ROWS = 5 +CSV_INFO_TO_DUMP = { + "row_num": [i for i in range(1, NUM_ROWS + 1)], + "other_info": [f"test_info_{i}" for i in range(1, NUM_ROWS + 1)], +} +JSON_INFO_TO_DUMP = {str(i): {f"other_info_{i}": f"test_info_{i}"} for i in range(1, NUM_ROWS + 1)} +DUMP_HANDLER_DIR = "{temp_output_dir}/dump_handler" + + +def test_dump_handler_invalid_dump_file(): + """ + This is really testing the initialization of the Dumper class with an invalid file type. + This should raise a ValueError. + """ + with pytest.raises(ValueError) as excinfo: + dump_handler("bad_file.txt", CSV_INFO_TO_DUMP) + assert "Invalid file type for bad_file.txt. Supported file types are: ['csv', 'json']" in str(excinfo.value) + + +def get_output_file(temp_dir: str, file_name: str): + """ + Helper function to get a full path to the temporary output file. + + :param temp_dir: The path to the temporary output directory that pytest gives us + :param file_name: The name of the file + """ + dump_dir = DUMP_HANDLER_DIR.format(temp_output_dir=temp_dir) + if not os.path.exists(dump_dir): + os.mkdir(dump_dir) + dump_file = f"{dump_dir}/{file_name}" + return dump_file + + +def run_csv_dump_test(dump_file: str, fmode: str): + """ + Run the test for csv dump. + + :param dump_file: The file that the dump was written to + :param fmode: The type of write that we're testing ("w" for write, "a" for append) + """ + + # Check that the file exists and that read in the contents of the file + assert os.path.exists(dump_file) + with open(dump_file, "r") as df: + reader = csv.reader(df) + written_data = list(reader) + + expected_rows = NUM_ROWS * 2 if fmode == "a" else NUM_ROWS + assert len(written_data) == expected_rows + 1 # Adding one because of the header row + for i, row in enumerate(written_data): + assert len(row) == 2 # Check number of columns + if i == 0: # Checking the header row + assert row[0] == "row_num" + assert row[1] == "other_info" + else: # Checking the data rows + assert row[0] == str(CSV_INFO_TO_DUMP["row_num"][(i % NUM_ROWS) - 1]) + assert row[1] == str(CSV_INFO_TO_DUMP["other_info"][(i % NUM_ROWS) - 1]) + + +def test_dump_handler_csv_write(temp_output_dir: str): + """ + This is really testing the write method of the Dumper class. + This should create a csv file and write to it. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the path to the file we'll write to + dump_file = get_output_file(temp_output_dir, "csv_write.csv") + + # Run the actual call to dump to the file + dump_handler(dump_file, CSV_INFO_TO_DUMP) + + # Assert that everything ran properly + run_csv_dump_test(dump_file, "w") + + +def test_dump_handler_csv_append(temp_output_dir: str): + """ + This is really testing the write method of the Dumper class with the file write mode set to append. + We'll write to a csv file first and then run again to make sure we can append to it properly. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the path to the file we'll write to + dump_file = get_output_file(temp_output_dir, "csv_append.csv") + + # Run the first call to create the csv file + dump_handler(dump_file, CSV_INFO_TO_DUMP) + + # Run the second call to append to the csv file + dump_handler(dump_file, CSV_INFO_TO_DUMP) + + # Assert that everything ran properly + run_csv_dump_test(dump_file, "a") + + +def test_dump_handler_json_write(temp_output_dir: str): + """ + This is really testing the write method of the Dumper class. + This should create a json file and write to it. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the path to the file we'll write to + dump_file = get_output_file(temp_output_dir, "json_write.json") + + # Run the actual call to dump to the file + dump_handler(dump_file, JSON_INFO_TO_DUMP) + + # Check that the file exists and that the contents are correct + assert os.path.exists(dump_file) + with open(dump_file, "r") as df: + contents = json.load(df) + assert contents == JSON_INFO_TO_DUMP + + +def test_dump_handler_json_append(temp_output_dir: str): + """ + This is really testing the write method of the Dumper class with the file write mode set to append. + We'll write to a json file first and then run again to make sure we can append to it properly. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the path to the file we'll write to + dump_file = get_output_file(temp_output_dir, "json_append.json") + + # Run the first call to create the file + timestamp_1 = str(datetime.now()) + first_dump = {timestamp_1: JSON_INFO_TO_DUMP} + dump_handler(dump_file, first_dump) + + # Sleep so we don't accidentally get the same timestamp + sleep(0.5) + + # Run the second call to append to the file + timestamp_2 = str(datetime.now()) + second_dump = {timestamp_2: JSON_INFO_TO_DUMP} + dump_handler(dump_file, second_dump) + + # Check that the file exists and that the contents are correct + assert os.path.exists(dump_file) + with open(dump_file, "r") as df: + contents = json.load(df) + keys = contents.keys() + assert len(keys) == 2 + assert timestamp_1 in keys + assert timestamp_2 in keys + assert contents[timestamp_1] == JSON_INFO_TO_DUMP + assert contents[timestamp_2] == JSON_INFO_TO_DUMP diff --git a/tests/unit/common/test_encryption.py b/tests/unit/common/test_encryption.py new file mode 100644 index 000000000..3a06c0ab9 --- /dev/null +++ b/tests/unit/common/test_encryption.py @@ -0,0 +1,139 @@ +""" +Tests for the `encrypt.py` and `encrypt_backend_traffic.py` files. +""" + +import os + +import celery +import pytest + +from merlin.common.security.encrypt import _gen_key, _get_key, _get_key_path, decrypt, encrypt +from merlin.common.security.encrypt_backend_traffic import _decrypt_decode, _encrypt_encode, set_backend_funcs +from merlin.config.configfile import CONFIG + + +class TestEncryption: + """ + This class will house all tests necessary for our encryption modules. + """ + + def test_encrypt(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test that our encryption function is encrypting the bytes that we're + passing to it. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + str_to_encrypt = b"super secret string shhh" + encrypted_str = encrypt(str_to_encrypt) + for word in str_to_encrypt.decode("utf-8").split(" "): + assert word not in encrypted_str.decode("utf-8") + + def test_decrypt(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test that our decryption function is decrypting the bytes that we're passing to it. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # This is the output of the bytes from the encrypt test + str_to_decrypt = b"gAAAAABld6k-jEncgCW5AePgrwn-C30dhr7dzGVhqzcqskPqFyA2Hdg3VWmo0qQnLklccaUYzAGlB4PMxyp4T-1gAYlAOf_7sC_bJOEcYOIkhZFoH6cX4Uw=" + decrypted_str = decrypt(str_to_decrypt) + assert decrypted_str == b"super secret string shhh" + + def test_get_key_path(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `_get_key_path` function. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Test the default behavior (`_get_key_path` will pull from CONFIG.results_backend which + # will be set to the temporary output path for our tests in the `use_fake_encrypt_data_key` fixture) + actual_default = _get_key_path() + assert actual_default.startswith("/tmp/") and actual_default.endswith("/encrypt_data_key") + + # Test with having the encryption key set to None + temp = CONFIG.results_backend.encryption_key + CONFIG.results_backend.encryption_key = None + with pytest.raises(ValueError) as excinfo: + _get_key_path() + assert "Error! No password provided for RabbitMQ" in str(excinfo.value) + CONFIG.results_backend.encryption_key = temp + + # Test with having the entire results_backend wiped from CONFIG + orig_results_backend = CONFIG.results_backend + CONFIG.results_backend = None + actual_no_results_backend = _get_key_path() + assert actual_no_results_backend == os.path.abspath(os.path.expanduser("~/.merlin/encrypt_data_key")) + CONFIG.results_backend = orig_results_backend + + def test_gen_key(self, temp_output_dir: str): + """ + Test the `_gen_key` function. + + :param temp_output_dir: The path to the temporary output directory for this test run + """ + # Create the file but don't put anything in it + key_gen_test_file = f"{temp_output_dir}/key_gen_test" + with open(key_gen_test_file, "w"): + pass + + # Ensure nothing is in the file + with open(key_gen_test_file, "r") as key_gen_file: + key_gen_contents = key_gen_file.read() + assert key_gen_contents == "" + + # Run the test and then check to make sure the file is now populated + _gen_key(key_gen_test_file) + with open(key_gen_test_file, "r") as key_gen_file: + key_gen_contents = key_gen_file.read() + assert key_gen_contents != "" + + def test_get_key( + self, + merlin_server_dir: str, + test_encryption_key: bytes, + redis_results_backend_config_function: "fixture", # noqa: F821 + ): + """ + Test the `_get_key` function. + + :param merlin_server_dir: The directory to the merlin test server configuration + :param test_encryption_key: A fixture to establish a fixed encryption key for testing + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Test the default functionality + actual_default = _get_key() + assert actual_default == test_encryption_key + + # Modify the permission of the key file so that it can't be read by anyone + # (we're purposefully trying to raise an IOError) + key_path = f"{merlin_server_dir}/encrypt_data_key" + orig_file_permissions = os.stat(key_path).st_mode + os.chmod(key_path, 0o222) + with pytest.raises(IOError): + _get_key() + os.chmod(key_path, orig_file_permissions) + + # Reset the key value to our test value since the IOError test will rewrite the key + with open(key_path, "w") as key_file: + key_file.write(test_encryption_key.decode("utf-8")) + + def test_set_backend_funcs(self): + """ + Test the `set_backend_funcs` function. + """ + orig_encode = celery.backends.base.Backend.encode + orig_decode = celery.backends.base.Backend.decode + + # Make sure these values haven't been set yet + assert celery.backends.base.Backend.encode != _encrypt_encode + assert celery.backends.base.Backend.decode != _decrypt_decode + + set_backend_funcs() + + # Ensure the new functions have been set + assert celery.backends.base.Backend.encode == _encrypt_encode + assert celery.backends.base.Backend.decode == _decrypt_decode + + celery.backends.base.Backend.encode = orig_encode + celery.backends.base.Backend.decode = orig_decode diff --git a/tests/unit/common/test_sample_index.py b/tests/unit/common/test_sample_index.py index c693827f0..c9cd108ee 100644 --- a/tests/unit/common/test_sample_index.py +++ b/tests/unit/common/test_sample_index.py @@ -1,178 +1,519 @@ +""" +Tests for the `sample_index.py` and `sample_index_factory.py` files. +""" + import os -import shutil -from contextlib import suppress +import pytest + +from merlin.common.sample_index import SampleIndex, new_dir, uniform_directories from merlin.common.sample_index_factory import create_hierarchy, read_hierarchy -TEST_DIR = "UNIT_TEST_SPACE" +def test_uniform_directories(): + """ + Test the `uniform_directories` function with different inputs. + """ + # Create the tests and the expected outputs + tests = [ + # SMALL SAMPLE SIZE + (10, 1, 100), # Bundle size of 1 and max dir level of 100 is default + (10, 1, 2), + (10, 2, 100), + (10, 2, 2), + # MEDIUM SAMPLE SIZE + (10000, 1, 100), # Bundle size of 1 and max dir level of 100 is default + (10000, 1, 5), + (10000, 5, 100), + (10000, 5, 10), + # LARGE SAMPLE SIZE + (1000000000, 1, 100), # Bundle size of 1 and max dir level of 100 is default + (1000000000, 1, 5), + (1000000000, 5, 100), + (1000000000, 5, 10), + ] + expected_outputs = [ + # SMALL SAMPLE SIZE + [1], + [8, 4, 2, 1], + [2], + [8, 4, 2], + # MEDIUM SAMPLE SIZE + [100, 1], + [3125, 625, 125, 25, 5, 1], + [500, 5], + [5000, 500, 50, 5], + # LARGE SAMPLE SIZE + [100000000, 1000000, 10000, 100, 1], + [244140625, 48828125, 9765625, 1953125, 390625, 78125, 15625, 3125, 625, 125, 25, 5, 1], + [500000000, 5000000, 50000, 500, 5], + [500000000, 50000000, 5000000, 500000, 50000, 5000, 500, 50, 5], + ] + # Run the tests and compare outputs + for i, test in enumerate(tests): + actual = uniform_directories(num_samples=test[0], bundle_size=test[1], level_max_dirs=test[2]) + assert actual == expected_outputs[i] -def clear_test_tree(): - with suppress(FileNotFoundError): - shutil.rmtree(TEST_DIR) +def test_new_dir(temp_output_dir: str): + """ + Test the `new_dir` function. This will test a valid path and also raising an OSError during + creation. -def clear(func): - def wrapper(): - clear_test_tree() - func() - clear_test_tree() + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Test basic functionality + test_path = f"{os.getcwd()}/test_sample_index/test_new_dir" + new_dir(test_path) + assert os.path.exists(test_path) - return wrapper + # Test OSError functionality + new_dir(test_path) -@clear -def test_index_file_writing(): - indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=TEST_DIR) - indx.write_directories() - indx.write_multiple_sample_index_files() - indx2 = read_hierarchy(TEST_DIR) - assert indx2.get_path_to_sample(123000123) == indx.get_path_to_sample(123000123) +class TestSampleIndex: + """ + These tests focus on testing the SampleIndex class used for creating the + sample hierarchy. + NOTE to see output of creating any hierarchy, change `write_all_hierarchies` to True. + The results of each hierarchy will be written to: + /tmp/`whoami`/pytest/pytest-of-`whoami`/pytest-current/python_{major}.{minor}.{micro}_current/test_sample_index/ + """ -def test_bundle_retrieval(): - indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=TEST_DIR) - expected = f"{TEST_DIR}/0/0/0/samples0-10000.ext" - result = indx.get_path_to_sample(123) - assert expected == result + write_all_hierarchies = False - expected = f"{TEST_DIR}/0/0/0/samples10000-20000.ext" - result = indx.get_path_to_sample(10000) - assert expected == result + def get_working_dir(self, test_workspace: str): + """ + This method is called for every test to get a unique workspace in the temporary + directory for the test output. - expected = f"{TEST_DIR}/1/2/3/samples123000000-123010000.ext" - result = indx.get_path_to_sample(123000123) - assert expected == result + :param test_workspace: The unique name for this workspace + (all tests use their unique test name for this value usually) + """ + return f"{os.getcwd()}/test_sample_index/{test_workspace}" + def write_hierarchy_for_debug(self, indx: SampleIndex): + """ + This method is for debugging purposes. It will cause all tests that don't write + hierarchies to write them so the output can be investigated. -def test_start_sample_id(): - expected = """: DIRECTORY MIN 203 MAX 303 NUM_BUNDLES 10 - 0: BUNDLE 0 MIN 203 MAX 213 - 1: BUNDLE 1 MIN 213 MAX 223 - 2: BUNDLE 2 MIN 223 MAX 233 - 3: BUNDLE 3 MIN 233 MAX 243 - 4: BUNDLE 4 MIN 243 MAX 253 - 5: BUNDLE 5 MIN 253 MAX 263 - 6: BUNDLE 6 MIN 263 MAX 273 - 7: BUNDLE 7 MIN 273 MAX 283 - 8: BUNDLE 8 MIN 283 MAX 293 - 9: BUNDLE 9 MIN 293 MAX 303 -""" - idx203 = create_hierarchy(100, 10, start_sample_id=203) - assert expected == str(idx203) - - -@clear -def test_directory_writing(): - path = os.path.join(TEST_DIR) - indx = create_hierarchy(2, 1, [1], root=path) - expected = """: DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2 - 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1 - 0.0: BUNDLE 0 MIN 0 MAX 1 - 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1 - 1.0: BUNDLE 1 MIN 1 MAX 2 -""" - assert expected == str(indx) - indx.write_directories() - assert os.path.isdir(f"{TEST_DIR}/0") - assert os.path.isdir(f"{TEST_DIR}/1") - indx.write_multiple_sample_index_files() - - clear_test_tree() - - path = os.path.join(TEST_DIR) - indx = create_hierarchy(1000000000, 10000, [100000000, 10000000], root=path) - indx.write_directories() - path = indx.get_path_to_sample(123000123) - assert os.path.exists(os.path.dirname(path)) - assert path != TEST_DIR - path = indx.get_path_to_sample(10000000000) - assert path == TEST_DIR - - clear_test_tree() - - path = os.path.join(TEST_DIR) - indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=path) - indx.write_directories() - - -def test_directory_path(): - indx = create_hierarchy(20, 1, [20, 5, 1], root="") - leaves = indx.make_directory_string() - expected_leaves = "0/0/0 0/0/1 0/0/2 0/0/3 0/0/4 0/1/0 0/1/1 0/1/2 0/1/3 0/1/4 0/2/0 0/2/1 0/2/2 0/2/3 0/2/4 0/3/0 0/3/1 0/3/2 0/3/3 0/3/4" - assert leaves == expected_leaves - all_dirs = indx.make_directory_string(just_leaf_directories=False) - expected_all_dirs = " 0 0/0 0/0/0 0/0/1 0/0/2 0/0/3 0/0/4 0/1 0/1/0 0/1/1 0/1/2 0/1/3 0/1/4 0/2 0/2/0 0/2/1 0/2/2 0/2/3 0/2/4 0/3 0/3/0 0/3/1 0/3/2 0/3/3 0/3/4" - assert all_dirs == expected_all_dirs - - -@clear -def test_subhierarchy_insertion(): - indx = create_hierarchy(2, 1, [1], root=TEST_DIR) - print("Writing directories") - indx.write_directories() - indx.write_multiple_sample_index_files() - print("reading heirarchy") - top = read_hierarchy(os.path.abspath(TEST_DIR)) - expected = """: DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2 - 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1 - 0.0: BUNDLE -1 MIN 0 MAX 1 - 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1 - 1.0: BUNDLE -1 MIN 1 MAX 2 -""" - assert str(top) == expected - print("creating sub_heirarchy") - sub_h = create_hierarchy(100, 10, address="1.0") - print("inserting sub_heirarchy") - top["1.0"] = sub_h - print(str(indx)) - print("after insertion") - print(str(top)) - expected = """: DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2 - 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1 - 0.0: BUNDLE -1 MIN 0 MAX 1 - 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1 - 1.0: DIRECTORY MIN 0 MAX 100 NUM_BUNDLES 10 - 1.0.0: BUNDLE 0 MIN 0 MAX 10 - 1.0.1: BUNDLE 1 MIN 10 MAX 20 - 1.0.2: BUNDLE 2 MIN 20 MAX 30 - 1.0.3: BUNDLE 3 MIN 30 MAX 40 - 1.0.4: BUNDLE 4 MIN 40 MAX 50 - 1.0.5: BUNDLE 5 MIN 50 MAX 60 - 1.0.6: BUNDLE 6 MIN 60 MAX 70 - 1.0.7: BUNDLE 7 MIN 70 MAX 80 - 1.0.8: BUNDLE 8 MIN 80 MAX 90 - 1.0.9: BUNDLE 9 MIN 90 MAX 100 -""" - assert str(top) == expected + :param indx: The `SampleIndex` object to write the hierarchy for + """ + if self.write_all_hierarchies: + indx.write_directories() + indx.write_multiple_sample_index_files() + def test_invalid_children(self): + """ + This will test that an invalid type for the `children` argument will raise + an error. + """ + tests = [ + ["a", "b", "c"], + True, + "a b c", + ] + for test in tests: + with pytest.raises(TypeError): + SampleIndex(0, 10, test, "name") -def test_sample_index(): - """Run through some basic testing of the SampleIndex class.""" - tests = [ - (10, 1, []), - (10, 3, []), - (11, 2, [5]), - (10, 3, [3]), - (10, 3, [1]), - (10, 1, [3]), - (10, 3, [1, 3]), - (10, 1, [2]), - (1000, 100, [500]), - (1000, 50, [500, 100]), - (1000000000, 100000132, []), - ] + def test_is_parent_of_leaf(self, temp_output_dir: str): + """ + Test the `is_parent_of_leaf` property. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_parent_of_leaf") + indx = create_hierarchy(10, 1, [2], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Test to see if parent of leaf is recognized + assert indx.is_parent_of_leaf is False + assert indx.children["0"].is_parent_of_leaf is True + + # Test to see if leaf is recognized + leaf_node = indx.children["0"].children["0.0"] + assert leaf_node.is_parent_of_leaf is False + + def test_is_grandparent_of_leaf(self, temp_output_dir: str): + """ + Test the `is_grandparent_of_leaf` property. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_grandparent_of_leaf") + indx = create_hierarchy(10, 1, [2], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Test to see if grandparent of leaf is recognized + assert indx.is_grandparent_of_leaf is True + assert indx.children["0"].is_grandparent_of_leaf is False + + # Test to see if leaf is recognized + leaf_node = indx.children["0"].children["0.0"] + assert leaf_node.is_grandparent_of_leaf is False + + def test_is_great_grandparent_of_leaf(self, temp_output_dir: str): + """ + Test the `is_great_grandparent_of_leaf` property. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_great_grandparent_of_leaf") + indx = create_hierarchy(10, 1, [5, 1], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Test to see if great grandparent of leaf is recognized + assert indx.is_great_grandparent_of_leaf is True + assert indx.children["0"].is_great_grandparent_of_leaf is False + assert indx.children["0"].children["0.0"].is_great_grandparent_of_leaf is False + + # Test to see if leaf is recognized + leaf_node = indx.children["0"].children["0.0"].children["0.0.0"] + assert leaf_node.is_great_grandparent_of_leaf is False + + def test_traverse_bundle(self, temp_output_dir: str): + """ + Test the `traverse_bundle` method to make sure it's just returning leaves. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_grandparent_of_leaf") + indx = create_hierarchy(10, 1, [2], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Ensure all nodes in the traversal are leaves + for _, node in indx.traverse_bundles(): + assert node.is_leaf + + def test_getitem(self, temp_output_dir: str): + """ + Test the `__getitem__` magic method. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_grandparent_of_leaf") + indx = create_hierarchy(10, 1, [2], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Test getting that requesting the root returns itself + assert indx[""] == indx + + # Test a valid address + assert indx["0"] == indx.children["0"] + + # Test an invalid address + with pytest.raises(KeyError): + indx["10"] + + def test_setitem(self, temp_output_dir: str): + """ + Test the `__setitem__` magic method. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create a hierarchy to test + working_dir = self.get_working_dir("test_is_grandparent_of_leaf") + indx = create_hierarchy(10, 1, [2], root=working_dir) + self.write_hierarchy_for_debug(indx) + + invalid_indx = SampleIndex(1, 3, {}, "invalid_indx") + + # Ensure that trying to change the root raises an error + with pytest.raises(KeyError): + indx[""] = invalid_indx + + # Ensure we can't just add a new subtree to a level + with pytest.raises(KeyError): + indx["10"] = invalid_indx + + # Test that invalid subtrees are caught + with pytest.raises(TypeError): + indx["0"] = invalid_indx + + # Test a valid set operation + dummy_indx = SampleIndex(0, 1, {}, "dummy_indx", leafid=0, address="0.0") + indx["0"]["0.0"] = dummy_indx + + def test_index_file_writing(self, temp_output_dir: str): + """ + Test the functionality of writing multiple index files. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + working_dir = self.get_working_dir("test_index_file_writing") + indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=working_dir) + indx.write_directories() + indx.write_multiple_sample_index_files() + indx2 = read_hierarchy(working_dir) + assert indx2.get_path_to_sample(123000123) == indx.get_path_to_sample(123000123) + + def test_directory_writing_small(self, temp_output_dir: str): + """ + Test that writing a small directory functions properly. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create the directory and ensure it has the correct format + working_dir = self.get_working_dir("test_directory_writing_small/") + indx = create_hierarchy(2, 1, [1], root=working_dir) + expected = ( + ": DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2\n" + " 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1\n" + " 0.0: BUNDLE 0 MIN 0 MAX 1\n" + " 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1\n" + " 1.0: BUNDLE 1 MIN 1 MAX 2\n" + ) + assert expected == str(indx) + + # Write the directories and ensure the paths are actually written + indx.write_directories() + assert os.path.isdir(f"{working_dir}/0") + assert os.path.isdir(f"{working_dir}/1") + indx.write_multiple_sample_index_files() + + def test_directory_writing_large(self, temp_output_dir: str): + """ + Test that writing a large directory functions properly. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + working_dir = self.get_working_dir("test_directory_writing_large") + indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=working_dir) + indx.write_directories() + path = indx.get_path_to_sample(123000123) + assert os.path.exists(os.path.dirname(path)) + assert path != working_dir + path = indx.get_path_to_sample(10000000000) + assert path == working_dir + + def test_bundle_retrieval(self, temp_output_dir: str): + """ + Test the functionality to get a bundle of samples when providing a sample id to find. + This will test a large sample hierarchy to ensure this scales properly. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create the hierarchy + working_dir = self.get_working_dir("test_bundle_retrieval") + indx = create_hierarchy(1000000000, 10000, [100000000, 10000000, 1000000], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Test for a small sample id + expected = f"{working_dir}/0/0/0/samples0-10000.ext" + result = indx.get_path_to_sample(123) + assert expected == result + + # Test for a mid size sample id + expected = f"{working_dir}/0/0/0/samples10000-20000.ext" + result = indx.get_path_to_sample(10000) + assert expected == result + + # Test for a large sample id + expected = f"{working_dir}/1/2/3/samples123000000-123010000.ext" + result = indx.get_path_to_sample(123000123) + assert expected == result + + def test_start_sample_id(self, temp_output_dir: str): + """ + Test creating a hierarchy using a starting sample id. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + working_dir = self.get_working_dir("test_start_sample_id") + expected = ( + ": DIRECTORY MIN 203 MAX 303 NUM_BUNDLES 10\n" + " 0: BUNDLE 0 MIN 203 MAX 213\n" + " 1: BUNDLE 1 MIN 213 MAX 223\n" + " 2: BUNDLE 2 MIN 223 MAX 233\n" + " 3: BUNDLE 3 MIN 233 MAX 243\n" + " 4: BUNDLE 4 MIN 243 MAX 253\n" + " 5: BUNDLE 5 MIN 253 MAX 263\n" + " 6: BUNDLE 6 MIN 263 MAX 273\n" + " 7: BUNDLE 7 MIN 273 MAX 283\n" + " 8: BUNDLE 8 MIN 283 MAX 293\n" + " 9: BUNDLE 9 MIN 293 MAX 303\n" + ) + idx203 = create_hierarchy(100, 10, start_sample_id=203, root=working_dir) + self.write_hierarchy_for_debug(idx203) + + assert expected == str(idx203) + + def test_make_directory_string(self, temp_output_dir: str): + """ + Test the `make_directory_string` method of `SampleIndex`. This will check + both the normal functionality where we just request paths to the leaves and + also the inverse functionality where we request all paths that are not leaves. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Creating the hierarchy + working_dir = self.get_working_dir("test_make_directory_string") + indx = create_hierarchy(20, 1, [20, 5, 1], root=working_dir) + self.write_hierarchy_for_debug(indx) + + # Testing normal functionality (just leaf directories) + leaves = indx.make_directory_string() + expected_leaves_list = [ + f"{working_dir}/0/0/0", + f"{working_dir}/0/0/1", + f"{working_dir}/0/0/2", + f"{working_dir}/0/0/3", + f"{working_dir}/0/0/4", + f"{working_dir}/0/1/0", + f"{working_dir}/0/1/1", + f"{working_dir}/0/1/2", + f"{working_dir}/0/1/3", + f"{working_dir}/0/1/4", + f"{working_dir}/0/2/0", + f"{working_dir}/0/2/1", + f"{working_dir}/0/2/2", + f"{working_dir}/0/2/3", + f"{working_dir}/0/2/4", + f"{working_dir}/0/3/0", + f"{working_dir}/0/3/1", + f"{working_dir}/0/3/2", + f"{working_dir}/0/3/3", + f"{working_dir}/0/3/4", + ] + expected_leaves = " ".join(expected_leaves_list) + assert leaves == expected_leaves + + # Testing no leaf functionality + all_dirs = indx.make_directory_string(just_leaf_directories=False) + expected_all_dirs_list = [ + working_dir, + f"{working_dir}/0", + f"{working_dir}/0/0", + f"{working_dir}/0/0/0", + f"{working_dir}/0/0/1", + f"{working_dir}/0/0/2", + f"{working_dir}/0/0/3", + f"{working_dir}/0/0/4", + f"{working_dir}/0/1", + f"{working_dir}/0/1/0", + f"{working_dir}/0/1/1", + f"{working_dir}/0/1/2", + f"{working_dir}/0/1/3", + f"{working_dir}/0/1/4", + f"{working_dir}/0/2", + f"{working_dir}/0/2/0", + f"{working_dir}/0/2/1", + f"{working_dir}/0/2/2", + f"{working_dir}/0/2/3", + f"{working_dir}/0/2/4", + f"{working_dir}/0/3", + f"{working_dir}/0/3/0", + f"{working_dir}/0/3/1", + f"{working_dir}/0/3/2", + f"{working_dir}/0/3/3", + f"{working_dir}/0/3/4", + ] + expected_all_dirs = " ".join(expected_all_dirs_list) + assert all_dirs == expected_all_dirs + + def test_subhierarchy_insertion(self, temp_output_dir: str): + """ + Test that a subhierarchy can be inserted into our `SampleIndex` properly. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Create the hierarchy and read it + working_dir = self.get_working_dir("test_subhierarchy_insertion") + indx = create_hierarchy(2, 1, [1], root=working_dir) + indx.write_directories() + indx.write_multiple_sample_index_files() + top = read_hierarchy(os.path.abspath(working_dir)) + + # Compare results + expected = ( + ": DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2\n" + " 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1\n" + " 0.0: BUNDLE -1 MIN 0 MAX 1\n" + " 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1\n" + " 1.0: BUNDLE -1 MIN 1 MAX 2\n" + ) + assert str(top) == expected + + # Create and insert the sub hierarchy + sub_h = create_hierarchy(100, 10, address="1.0") + top["1.0"] = sub_h + + # Compare results + expected = ( + ": DIRECTORY MIN 0 MAX 2 NUM_BUNDLES 2\n" + " 0: DIRECTORY MIN 0 MAX 1 NUM_BUNDLES 1\n" + " 0.0: BUNDLE -1 MIN 0 MAX 1\n" + " 1: DIRECTORY MIN 1 MAX 2 NUM_BUNDLES 1\n" + " 1.0: DIRECTORY MIN 0 MAX 100 NUM_BUNDLES 10\n" + " 1.0.0: BUNDLE 0 MIN 0 MAX 10\n" + " 1.0.1: BUNDLE 1 MIN 10 MAX 20\n" + " 1.0.2: BUNDLE 2 MIN 20 MAX 30\n" + " 1.0.3: BUNDLE 3 MIN 30 MAX 40\n" + " 1.0.4: BUNDLE 4 MIN 40 MAX 50\n" + " 1.0.5: BUNDLE 5 MIN 50 MAX 60\n" + " 1.0.6: BUNDLE 6 MIN 60 MAX 70\n" + " 1.0.7: BUNDLE 7 MIN 70 MAX 80\n" + " 1.0.8: BUNDLE 8 MIN 80 MAX 90\n" + " 1.0.9: BUNDLE 9 MIN 90 MAX 100\n" + ) + assert str(top) == expected + + def test_sample_index_creation_and_insertion(self, temp_output_dir: str): + """ + Run through some basic testing of the SampleIndex class. This will try + creating hierarchies of different sizes and inserting subhierarchies of + different sizes as well. + + :param temp_output_dir: A pytest fixture defined in conftest.py that creates a + temporary output path for our tests + """ + # Define the tests for hierarchies of varying sizes + tests = [ + (10, 1, []), + (10, 3, []), + (11, 2, [5]), + (10, 3, [3]), + (10, 3, [1]), + (10, 1, [3]), + (10, 3, [1, 3]), + (10, 1, [2]), + (1000, 100, [500]), + (1000, 50, [500, 100]), + (1000000000, 100000132, []), + ] + + # Run all the tests we defined above + for i, args in enumerate(tests): + working_dir = self.get_working_dir(f"test_sample_index_creation_and_insertion/{i}") + + # Put at root address of "0" to guarantee insertion at "0.1" later is valid + idx = create_hierarchy(args[0], args[1], args[2], address="0", root=working_dir) + self.write_hierarchy_for_debug(idx) - for args in tests: - print(f"############ TEST {args[0]} {args[1]} {args[2]} ###########") - # put at root address of "0" to guarantee insertion at "0.1" later is valid - idx = create_hierarchy(args[0], args[1], args[2], address="0") - print(str(idx)) - try: - idx["0.1"] = create_hierarchy(args[0], args[1], args[2], address="0.1") - print("successful set") - print(str(idx)) - except KeyError as error: - print(error) - assert False + # Inserting hierarchy at 0.1 + try: + idx["0.1"] = create_hierarchy(args[0], args[1], args[2], address="0.1") + except KeyError: + assert False diff --git a/tests/unit/common/test_util_sampling.py b/tests/unit/common/test_util_sampling.py new file mode 100644 index 000000000..b4cc252d5 --- /dev/null +++ b/tests/unit/common/test_util_sampling.py @@ -0,0 +1,45 @@ +""" +Tests for the `util_sampling.py` file. +""" + +import numpy as np +import pytest + +from merlin.common.util_sampling import scale_samples + + +class TestUtilSampling: + """ + This class will hold all of the tests for the `util_sampling.py` file. + """ + + def test_scale_samples_basic(self): + """Test basic functionality""" + samples_norm = np.array([[0.2, 0.4], [0.6, 0.8]]) + limits = [(-1, 1), (2, 6)] + result = scale_samples(samples_norm, limits) + expected_result = np.array([[-0.6, 3.6], [0.2, 5.2]]) + np.testing.assert_array_almost_equal(result, expected_result) + + def test_scale_samples_logarithmic(self): + """Test functionality with log enabled""" + samples_norm = np.array([[0.2, 0.4], [0.6, 0.8]]) + limits = [(1, 5), (1, 100)] + result = scale_samples(samples_norm, limits, do_log=[False, True]) + expected_result = np.array([[1.8, 6.309573], [3.4, 39.810717]]) + np.testing.assert_array_almost_equal(result, expected_result) + + def test_scale_samples_invalid_input(self): + """Test that function raises ValueError for invalid input""" + with pytest.raises(ValueError): + # Invalid input: samples_norm should be a 2D array + scale_samples([0.2, 0.4, 0.6], [(1, 5), (2, 6)]) + + def test_scale_samples_with_custom_limits_norm(self): + """Test functionality with custom limits_norm""" + samples_norm = np.array([[0.2, 0.4], [0.6, 0.8]]) + limits = [(1, 5), (2, 6)] + limits_norm = (-1, 1) + result = scale_samples(samples_norm, limits, limits_norm=limits_norm) + expected_result = np.array([[3.4, 4.8], [4.2, 5.6]]) + np.testing.assert_array_almost_equal(result, expected_result) diff --git a/tests/unit/config/dummy_app.yaml b/tests/unit/config/dummy_app.yaml new file mode 100644 index 000000000..966156566 --- /dev/null +++ b/tests/unit/config/dummy_app.yaml @@ -0,0 +1,33 @@ +broker: + cert_reqs: none + name: redis + password: redis.pass + port: '6379' + server: 127.0.0.1 + username: default + vhost: host4gunny +celery: + override: + visibility_timeout: 86400 +container: + config: redis.conf + config_dir: ./merlin_server/ + format: singularity + image: redis_latest.sif + image_type: redis + pass_file: redis.pass + pfile: merlin_server.pf + url: docker://redis + user_file: redis.users +process: + kill: kill {pid} + status: pgrep -P {pid} +results_backend: + cert_reqs: none + db_num: 0 + encryption_key: encrypt_data_key + name: redis + password: redis.pass + port: '6379' + server: 127.0.0.1 + username: default \ No newline at end of file diff --git a/tests/unit/config/old_test_configfile.py b/tests/unit/config/old_test_configfile.py deleted file mode 100644 index 39139ec11..000000000 --- a/tests/unit/config/old_test_configfile.py +++ /dev/null @@ -1,97 +0,0 @@ -"""Tests for the configfile module.""" - -import os -import shutil -import tempfile -import unittest -from getpass import getuser - -from merlin.config import configfile - -from .utils import mkfile - - -CONFIG_FILE_CONTENTS = """ -celery: - certs: path/to/celery/config/files - -broker: - name: rabbitmq - username: testuser - password: rabbit.password # The filename that contains the password. - server: jackalope.llnl.gov - -results_backend: - name: mysql - dbname: testuser - username: mlsi - password: mysql.password # The filename that contains the password. - server: rabbit.llnl.gov - -""" - - -class TestFindConfigFile(unittest.TestCase): - def setUp(self): - self.tmpdir = tempfile.mkdtemp() - self.appfile = mkfile(self.tmpdir, "app.yaml") - - def tearDown(self): - shutil.rmtree(self.tmpdir, ignore_errors=True) - - def test_tempdir(self): - self.assertTrue(os.path.isdir(self.tmpdir)) - - def test_find_config_file(self): - """ - Given the path to a vaild config file, find and return the full - filepath. - """ - path = configfile.find_config_file(path=self.tmpdir) - expected = os.path.join(self.tmpdir, self.appfile) - self.assertEqual(path, expected) - - def test_find_config_file_error(self): - """Given an invalid path, return None.""" - invalid = "invalid/path" - expected = None - - path = configfile.find_config_file(path=invalid) - self.assertEqual(path, expected) - - -class TestConfigFile(unittest.TestCase): - """Unit tests for loading the config file.""" - - def setUp(self): - self.tmpdir = tempfile.mkdtemp() - self.configfile = mkfile(self.tmpdir, "app.yaml", content=CONFIG_FILE_CONTENTS) - - def tearDown(self): - shutil.rmtree(self.tmpdir, ignore_errors=True) - - def test_get_config(self): - """ - Given the directory path to a valid merlin config file, then - `get_config` should find the merlin config file and load the YAML - contents to a dictionary. - """ - expected = { - "broker": { - "name": "rabbitmq", - "password": "rabbit.password", - "server": "jackalope.llnl.gov", - "username": "testuser", - "vhost": getuser(), - }, - "celery": {"certs": "path/to/celery/config/files"}, - "results_backend": { - "dbname": "testuser", - "name": "mysql", - "password": "mysql.password", - "server": "rabbit.llnl.gov", - "username": "mlsi", - }, - } - - self.assertDictEqual(configfile.get_config(self.tmpdir), expected) diff --git a/tests/unit/config/old_test_results_backend.py b/tests/unit/config/old_test_results_backend.py deleted file mode 100644 index 638f13eb8..000000000 --- a/tests/unit/config/old_test_results_backend.py +++ /dev/null @@ -1,67 +0,0 @@ -"""Tests for the results_backend module.""" - -import os -import shutil -import tempfile -import unittest - -from merlin.config import results_backend - -from .utils import mkfile - - -class TestResultsBackend(unittest.TestCase): - def setUp(self): - self.tmpdir = tempfile.mkdtemp() - - # Create test files. - self.tmpfile1 = mkfile(self.tmpdir, "mysql_test1.txt") - self.tmpfile2 = mkfile(self.tmpdir, "mysql_test2.txt") - - def tearDown(self): - shutil.rmtree(self.tmpdir, ignore_errors=True) - - def test_mysql_config(self): - """ - Given the path to a directory containing the MySQL cert files and a - dictionary of files to look for, then find and return the full path to - all the certs. - """ - certs = {"test1": "mysql_test1.txt", "test2": "mysql_test2.txt"} - - # This will just be the above dictionary with the full file paths. - expected = { - "test1": os.path.join(self.tmpdir, certs["test1"]), - "test2": os.path.join(self.tmpdir, certs["test2"]), - } - results = results_backend.get_mysql_config(self.tmpdir, certs) - self.assertDictEqual(results, expected) - - def test_mysql_config_no_files(self): - """ - Given the path to a directory containing the MySQL cert files and - an empty dictionary, then `get_mysql_config` should return an empty - dictionary. - """ - files = {} - result = results_backend.get_mysql_config(self.tmpdir, files) - self.assertEqual(result, {}) - - -class TestConfingMysqlErrorPath(unittest.TestCase): - """ - Test `get_mysql_config` against cases were the given path does not exist. - """ - - def test_mysql_config_false(self): - """ - Given a path that does not exist, then `get_mysql_config` should return - False. - """ - path = "invalid/path" - - # We don't need the dictionary populated for this test. The function - # should return False before trying to process the dictionary. - certs = {} - result = results_backend.get_mysql_config(path, certs) - self.assertFalse(result) diff --git a/tests/unit/config/test_broker.py b/tests/unit/config/test_broker.py new file mode 100644 index 000000000..3b44c5261 --- /dev/null +++ b/tests/unit/config/test_broker.py @@ -0,0 +1,551 @@ +""" +Tests for the `broker.py` file. +""" + +import os +from ssl import CERT_NONE +from typing import Any, Dict + +import pytest + +from merlin.config.broker import ( + RABBITMQ_CONNECTION, + REDISSOCK_CONNECTION, + get_connection_string, + get_rabbit_connection, + get_redis_connection, + get_redissock_connection, + get_ssl_config, + read_file, +) +from merlin.config.configfile import CONFIG +from tests.constants import SERVER_PASS +from tests.utils import create_pass_file + + +def test_read_file(merlin_server_dir: str): + """ + Test the `read_file` function. We'll start up our containerized redis server + so that we have a password file to read here. + + :param merlin_server_dir: The directory to the merlin test server configuration + """ + pass_file = f"{merlin_server_dir}/redis.pass" + create_pass_file(pass_file) + actual = read_file(pass_file) + assert actual == SERVER_PASS + + +def test_get_connection_string_invalid_broker(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with an invalid broker (a broker that isn't one of: + ["rabbitmq", "redis", "rediss", "redis+socket", "amqps", "amqp"]). + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.name = "invalid_broker" + with pytest.raises(ValueError): + get_connection_string() + + +def test_get_connection_string_no_broker(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function without a broker name value in the CONFIG object. This + should raise a ValueError just like the `test_get_connection_string_invalid_broker` does. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.name + with pytest.raises(ValueError): + get_connection_string() + + +def test_get_connection_string_simple(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function in the simplest way that we can. This function + will automatically check for a broker url and if it finds one in the CONFIG object it will just + return the value it finds. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + test_url = "test_url" + CONFIG.broker.url = test_url + actual = get_connection_string() + assert actual == test_url + + +def test_get_ssl_config_no_broker(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function without a broker. This should return False. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.name + assert not get_ssl_config() + + +class TestRabbitBroker: + """ + This class will house all tests necessary for our broker module when using a + rabbit broker. + """ + + def run_get_rabbit_connection(self, expected_vals: Dict[str, Any], include_password: bool, conn: str): + """ + Helper method to run the tests for the `get_rabbit_connection`. + + :param expected_vals: A dict of expected values for this test. Format: + {"conn": "", + "vhost": "host4testing", + "username": "default", + "password": "", + "server": "127.0.0.1", + "port": } + :param include_password: If True, include the password in the output. Otherwise don't. + :param conn: The connection type to pass in (either amqp or amqps) + """ + expected = RABBITMQ_CONNECTION.format(**expected_vals) + actual = get_rabbit_connection(include_password=include_password, conn=conn) + assert actual == expected + + def test_get_rabbit_connection(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + conn = "amqps" + expected_vals = { + "conn": conn, + "vhost": "host4testing", + "username": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + "port": 5671, + } + self.run_get_rabbit_connection(expected_vals=expected_vals, include_password=True, conn=conn) + + def test_get_rabbit_connection_dont_include_password(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function but set include_password to False. This should * out the + password + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + conn = "amqps" + expected_vals = { + "conn": conn, + "vhost": "host4testing", + "username": "default", + "password": "******", + "server": "127.0.0.1", + "port": 5671, + } + self.run_get_rabbit_connection(expected_vals=expected_vals, include_password=False, conn=conn) + + def test_get_rabbit_connection_no_port_amqp(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function with no port in the CONFIG object. This should use + 5672 as the port since we're using amqp as the connection. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.port + CONFIG.broker.name = "amqp" + conn = "amqp" + expected_vals = { + "conn": conn, + "vhost": "host4testing", + "username": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + "port": 5672, + } + self.run_get_rabbit_connection(expected_vals=expected_vals, include_password=True, conn=conn) + + def test_get_rabbit_connection_no_port_amqps(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function with no port in the CONFIG object. This should use + 5671 as the port since we're using amqps as the connection. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.port + conn = "amqps" + expected_vals = { + "conn": conn, + "vhost": "host4testing", + "username": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + "port": 5671, + } + self.run_get_rabbit_connection(expected_vals=expected_vals, include_password=True, conn=conn) + + def test_get_rabbit_connection_no_password(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function with no password file set. This should raise a ValueError. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.password + with pytest.raises(ValueError) as excinfo: + get_rabbit_connection(True) + assert "Broker: No password provided for RabbitMQ" in str(excinfo.value) + + def test_get_rabbit_connection_invalid_pass_filepath(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_rabbit_connection` function with an invalid password filepath. + This should raise a ValueError. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.password = "invalid_filepath" + expanded_filepath = os.path.abspath(os.path.expanduser(CONFIG.broker.password)) + with pytest.raises(ValueError) as excinfo: + get_rabbit_connection(True) + assert f"Broker: RabbitMQ password file {expanded_filepath} does not exist" in str(excinfo.value) + + def run_get_connection_string(self, expected_vals: Dict[str, Any]): + """ + Helper method to run the tests for the `get_connection_string`. + + :param expected_vals: A dict of expected values for this test. Format: + {"conn": "", + "vhost": "host4testing", + "username": "default", + "password": "", + "server": "127.0.0.1", + "port": } + """ + expected = RABBITMQ_CONNECTION.format(**expected_vals) + actual = get_connection_string() + assert actual == expected + + def test_get_connection_string_rabbitmq(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with rabbitmq as the broker. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "conn": "amqps", + "vhost": "host4testing", + "username": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + "port": 5671, + } + self.run_get_connection_string(expected_vals) + + def test_get_connection_string_amqp(self, rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with amqp as the broker. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.port + CONFIG.broker.name = "amqp" + expected_vals = { + "conn": "amqp", + "vhost": "host4testing", + "username": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + "port": 5672, + } + self.run_get_connection_string(expected_vals) + + +class TestRedisBroker: + """ + This class will house all tests necessary for our broker module when using a + redis broker. + """ + + def run_get_redissock_connection(self, expected_vals: Dict[str, str]): + """ + Helper method to run the tests for the `get_redissock_connection`. + + :param expected_vals: A dict of expected values for this test. Format: + {"db_num": "", "path": ""} + """ + expected = REDISSOCK_CONNECTION.format(**expected_vals) + actual = get_redissock_connection() + assert actual == expected + + def test_get_redissock_connection(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redissock_connection` function with both a db_num and a broker path set. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Create and store a fake path and db_num for testing + test_path = "/fake/path/to/broker" + test_db_num = "45" + CONFIG.broker.path = test_path + CONFIG.broker.db_num = test_db_num + + # Set up our expected vals and compare against the actual result + expected_vals = {"db_num": test_db_num, "path": test_path} + self.run_get_redissock_connection(expected_vals) + + def test_get_redissock_connection_no_db(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redissock_connection` function with a broker path set but no db num. + This should default the db_num to 0. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Create and store a fake path for testing + test_path = "/fake/path/to/broker" + CONFIG.broker.path = test_path + + # Set up our expected vals and compare against the actual result + expected_vals = {"db_num": 0, "path": test_path} + self.run_get_redissock_connection(expected_vals) + + def test_get_redissock_connection_no_path(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redissock_connection` function with a db num set but no broker path. + This should raise an AttributeError since there will be no path value to read from + in `CONFIG.broker`. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.db_num = "45" + with pytest.raises(AttributeError): + get_redissock_connection() + + def test_get_redissock_connection_no_path_nor_db(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redissock_connection` function with neither a broker path nor a db num set. + This should raise an AttributeError since there will be no path value to read from + in `CONFIG.broker`. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + with pytest.raises(AttributeError): + get_redissock_connection() + + def run_get_redis_connection(self, expected_vals: Dict[str, Any], include_password: bool, use_ssl: bool): + """ + Helper method to run the tests for the `get_redis_connection`. + + :param expected_vals: A dict of expected values for this test. Format: + {"urlbase": "", "spass": "", "server": "127.0.0.1", "port": , "db_num": } + :param include_password: If True, include the password in the output. Otherwise don't. + :param use_ssl: If True, use ssl for the connection. Otherwise don't. + """ + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_redis_connection(include_password=include_password, use_ssl=use_ssl) + assert expected == actual + + def test_get_redis_connection(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "redis", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + + def test_get_redis_connection_no_port(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + We'll run this after deleting the port setting from the CONFIG object. This should still run and give us + port = 6379. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.port + expected_vals = { + "urlbase": "redis", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + + def test_get_redis_connection_with_db(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + We'll run this after adding the db_num setting to the CONFIG object. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + test_db_num = "45" + CONFIG.broker.db_num = test_db_num + expected_vals = { + "urlbase": "redis", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": test_db_num, + } + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + + def test_get_redis_connection_no_username(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + We'll run this after deleting the username setting from the CONFIG object. This should still run and give us + username = ''. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.username + expected_vals = {"urlbase": "redis", "spass": ":merlin-test-server@", "server": "127.0.0.1", "port": 6379, "db_num": 0} + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + + def test_get_redis_connection_invalid_pass_file(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + We'll run this after changing the permissions of the password file so it can't be opened. This should still + run and give us password = CONFIG.broker.password. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Capture the initial permissions of the password file so we can reset them + orig_file_permissions = os.stat(CONFIG.broker.password).st_mode + + # Change the permissions of the password file so it can't be read + os.chmod(CONFIG.broker.password, 0o222) + + try: + # Run the test + expected_vals = { + "urlbase": "redis", + "spass": f"default:{CONFIG.broker.password}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + except AssertionError as exc: + # If this test failed, make sure to reset the permissions in case other tests need to read this file + os.chmod(CONFIG.broker.password, orig_file_permissions) + raise AssertionError from exc + + os.chmod(CONFIG.broker.password, orig_file_permissions) + + def test_get_redis_connection_dont_include_password(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function without including the password. This should place 6 *s + where the password would normally be placed in spass. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = {"urlbase": "redis", "spass": "default:******@", "server": "127.0.0.1", "port": 6379, "db_num": 0} + self.run_get_redis_connection(expected_vals=expected_vals, include_password=False, use_ssl=False) + + def test_get_redis_connection_use_ssl(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with using ssl. This should change the urlbase to rediss (with two 's'). + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "rediss", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=True) + + def test_get_redis_connection_no_password(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis_connection` function with default functionality (including password and not using ssl). + We'll run this after deleting the password setting from the CONFIG object. This should still run and give us + spass = ''. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.broker.password + expected_vals = {"urlbase": "redis", "spass": "", "server": "127.0.0.1", "port": 6379, "db_num": 0} + self.run_get_redis_connection(expected_vals=expected_vals, include_password=True, use_ssl=False) + + def test_get_connection_string_redis(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with redis as the broker (this is what our CONFIG + is set to by default with the redis_broker_config_function fixture). + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "redis", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_connection_string() + assert expected == actual + + def test_get_connection_string_rediss(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with rediss (with two 's') as the broker. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.name = "rediss" + expected_vals = { + "urlbase": "rediss", + "spass": "default:merlin-test-server@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_connection_string() + assert expected == actual + + def test_get_connection_string_redis_socket(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with redis+socket as the broker. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + # Change our broker + CONFIG.broker.name = "redis+socket" + + # Create and store a fake path and db_num for testing + test_path = "/fake/path/to/broker" + test_db_num = "45" + CONFIG.broker.path = test_path + CONFIG.broker.db_num = test_db_num + + # Set up our expected vals and compare against the actual result + expected_vals = {"db_num": test_db_num, "path": test_path} + expected = REDISSOCK_CONNECTION.format(**expected_vals) + actual = get_connection_string() + assert actual == expected + + def test_get_ssl_config_redis(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with redis as the broker (this is the default in our tests). + This should return False. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + assert not get_ssl_config() + + def test_get_ssl_config_rediss(self, redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with rediss (with two 's') as the broker. + This should return a dict of cert reqs with ssl.CERT_NONE as the value. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.name = "rediss" + expected = {"ssl_cert_reqs": CERT_NONE} + actual = get_ssl_config() + assert actual == expected diff --git a/tests/unit/config/test_config_object.py b/tests/unit/config/test_config_object.py new file mode 100644 index 000000000..64e56b7d9 --- /dev/null +++ b/tests/unit/config/test_config_object.py @@ -0,0 +1,150 @@ +""" +Test the functionality of the Config object. +""" + +from copy import copy, deepcopy +from types import SimpleNamespace + +from merlin.config import Config + + +class TestConfig: + """ + Class for testing the Config object. We'll store a valid `app_dict` + as an attribute here so that each test doesn't have to redefine it + each time. + """ + + app_dict = { + "celery": {"override": {"visibility_timeout": 86400}}, + "broker": { + "cert_reqs": "none", + "name": "rabbitmq", + "password": "/path/to/pass_file", + "port": 5671, + "server": "127.0.0.1", + "username": "default", + "vhost": "host4testing", + }, + "results_backend": { + "cert_reqs": "none", + "db_num": 0, + "name": "rediss", + "password": "/path/to/pass_file", + "port": 6379, + "server": "127.0.0.1", + "username": "default", + "vhost": "host4testing", + "encryption_key": "/path/to/encryption_key", + }, + } + + def test_config_creation(self): + """ + Test the creation of the Config object. This should create nested namespaces + for each key in the `app_dict` variable and save them to their respective + attributes in the object. + """ + config = Config(self.app_dict) + + # Create the nested namespace for celery and compare result + override_namespace = SimpleNamespace(**self.app_dict["celery"]["override"]) + updated_celery_dict = deepcopy(self.app_dict) + updated_celery_dict["celery"]["override"] = override_namespace + celery_namespace = SimpleNamespace(**updated_celery_dict["celery"]) + assert config.celery == celery_namespace + + # Broker and Results Backend are easier since there's no nested namespace here + assert config.broker == SimpleNamespace(**self.app_dict["broker"]) + assert config.results_backend == SimpleNamespace(**self.app_dict["results_backend"]) + + def test_config_creation_no_celery(self): + """ + Test the creation of the Config object without the celery key. This should still + work and just not set anything for the celery attribute. + """ + + # Copy the celery section so we can restore it later and then delete it + celery_section = copy(self.app_dict["celery"]) + del self.app_dict["celery"] + config = Config(self.app_dict) + + # Broker and Results Backend are the only things loaded here + assert config.broker == SimpleNamespace(**self.app_dict["broker"]) + assert config.results_backend == SimpleNamespace(**self.app_dict["results_backend"]) + + # Ensure the celery attribute is not loaded + assert "celery" not in dir(config) + + # Reset celery section in case other tests use it after this + self.app_dict["celery"] = celery_section + + def test_config_copy(self): + """ + Test the `__copy__` magic method of the Config object. Here we'll make sure + each attribute was copied properly but the ids should be different. + """ + orig_config = Config(self.app_dict) + copied_config = copy(orig_config) + + assert orig_config.celery == copied_config.celery + assert orig_config.broker == copied_config.broker + assert orig_config.results_backend == copied_config.results_backend + + assert id(orig_config) != id(copied_config) + + def test_config_str(self): + """ + Test the `__str__` magic method of the Config object. This should just give us + a formatted string of the attributes in the object. + """ + config = Config(self.app_dict) + + # Test normal printing + actual = config.__str__() + expected = ( + "config:\n" + " celery:\n" + " override: namespace(visibility_timeout=86400)\n" + " broker:\n" + " cert_reqs: 'none'\n" + " name: 'rabbitmq'\n" + " password: '/path/to/pass_file'\n" + " port: 5671\n" + " server: '127.0.0.1'\n" + " username: 'default'\n" + " vhost: 'host4testing'\n" + " results_backend:\n" + " cert_reqs: 'none'\n" + " db_num: 0\n" + " name: 'rediss'\n" + " password: '/path/to/pass_file'\n" + " port: 6379\n" + " server: '127.0.0.1'\n" + " username: 'default'\n" + " vhost: 'host4testing'\n" + " encryption_key: '/path/to/encryption_key'" + ) + + assert actual == expected + + # Test printing with one section set to None + config.results_backend = None + actual_with_none = config.__str__() + expected_with_none = ( + "config:\n" + " celery:\n" + " override: namespace(visibility_timeout=86400)\n" + " broker:\n" + " cert_reqs: 'none'\n" + " name: 'rabbitmq'\n" + " password: '/path/to/pass_file'\n" + " port: 5671\n" + " server: '127.0.0.1'\n" + " username: 'default'\n" + " vhost: 'host4testing'\n" + " results_backend:\n" + " None" + ) + + assert actual_with_none == expected_with_none diff --git a/tests/unit/config/test_configfile.py b/tests/unit/config/test_configfile.py new file mode 100644 index 000000000..49b12ba86 --- /dev/null +++ b/tests/unit/config/test_configfile.py @@ -0,0 +1,696 @@ +""" +Tests for the configfile.py module. +""" + +import getpass +import os +import shutil +import ssl +from copy import copy, deepcopy + +import pytest +import yaml + +from merlin.config.configfile import ( + CONFIG, + default_config_info, + find_config_file, + get_cert_file, + get_config, + get_ssl_entries, + is_debug, + load_config, + load_default_celery, + load_defaults, + merge_sslmap, + process_ssl_map, + set_username_and_vhost, +) +from tests.constants import CERT_FILES +from tests.utils import create_dir + + +CONFIGFILE_DIR = "{temp_output_dir}/test_configfile" +COPIED_APP_FILENAME = "app_copy.yaml" +DUMMY_APP_FILEPATH = f"{os.path.dirname(__file__)}/dummy_app.yaml" + + +def create_app_yaml(app_yaml_filepath: str): + """ + Create a dummy app.yaml file at `app_yaml_filepath`. + + :param app_yaml_filepath: The location to create an app.yaml file at + """ + full_app_yaml_filepath = f"{app_yaml_filepath}/app.yaml" + if not os.path.exists(full_app_yaml_filepath): + shutil.copy(DUMMY_APP_FILEPATH, full_app_yaml_filepath) + + +def test_load_config(temp_output_dir: str): + """ + Test the `load_config` function. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + configfile_dir = CONFIGFILE_DIR.format(temp_output_dir=temp_output_dir) + create_dir(configfile_dir) + create_app_yaml(configfile_dir) + + with open(DUMMY_APP_FILEPATH, "r") as dummy_app_file: + expected = yaml.load(dummy_app_file, yaml.Loader) + + actual = load_config(f"{configfile_dir}/app.yaml") + assert actual == expected + + +def test_load_config_invalid_file(): + """ + Test the `load_config` function with an invalid filepath. + """ + assert load_config("invalid/filepath") is None + + +def test_find_config_file_valid_path(temp_output_dir: str): + """ + Test the `find_config_file` function with passing a valid path in. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + configfile_dir = CONFIGFILE_DIR.format(temp_output_dir=temp_output_dir) + create_dir(configfile_dir) + create_app_yaml(configfile_dir) + + assert find_config_file(configfile_dir) == f"{configfile_dir}/app.yaml" + + +def test_find_config_file_invalid_path(): + """ + Test the `find_config_file` function with passing an invalid path in. + """ + assert find_config_file("invalid/path") is None + + +def test_find_config_file_local_path(temp_output_dir: str): + """ + Test the `find_config_file` function by having it find a local (in our cwd) app.yaml file. + We'll use the `temp_output_dir` fixture so that our current working directory is in a temp + location. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the configfile directory and put an app.yaml file there + configfile_dir = CONFIGFILE_DIR.format(temp_output_dir=temp_output_dir) + create_dir(configfile_dir) + create_app_yaml(configfile_dir) + + # Move into the configfile directory and run the test + os.chdir(configfile_dir) + try: + assert find_config_file() == f"{os.getcwd()}/app.yaml" + except AssertionError as exc: + # Move back to the temp output directory even if the test fails + os.chdir(temp_output_dir) + raise AssertionError from exc + + # Move back to the temp output directory + os.chdir(temp_output_dir) + + +def test_find_config_file_merlin_home_path(temp_output_dir: str): + """ + Test the `find_config_file` function by having it find an app.yaml file in our merlin directory. + We'll use the `temp_output_dir` fixture so that our current working directory is in a temp + location. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + merlin_home = os.path.expanduser("~/.merlin") + if not os.path.exists(merlin_home): + os.mkdir(merlin_home) + create_app_yaml(merlin_home) + assert find_config_file() == f"{merlin_home}/app.yaml" + + +def check_for_and_move_app_yaml(dir_to_check: str) -> bool: + """ + Check for any app.yaml files in `dir_to_check`. If one is found, rename it. + Return True if an app.yaml was found, false otherwise. + + :param dir_to_check: The directory to search for an app.yaml in + :returns: True if an app.yaml was found. False otherwise. + """ + for filename in os.listdir(dir_to_check): + full_path = os.path.join(dir_to_check, filename) + if os.path.isfile(full_path) and filename == "app.yaml": + os.rename(full_path, f"{dir_to_check}/{COPIED_APP_FILENAME}") + return True + return False + + +def test_find_config_file_no_path(temp_output_dir: str): + """ + Test the `find_config_file` function by making it unable to find any app.yaml path. + We'll use the `temp_output_dir` fixture so that our current working directory is in a temp + location. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Rename any app.yaml in the cwd + cwd_path = os.getcwd() + cwd_had_app_yaml = check_for_and_move_app_yaml(cwd_path) + + # Rename any app.yaml in the merlin home directory + merlin_home_dir = os.path.expanduser("~/.merlin") + merlin_home_had_app_yaml = check_for_and_move_app_yaml(merlin_home_dir) + + try: + assert find_config_file() is None + except AssertionError as exc: + # Reset the cwd app.yaml even if the test fails + if cwd_had_app_yaml: + os.rename(f"{cwd_path}/{COPIED_APP_FILENAME}", f"{cwd_path}/app.yaml") + + # Reset the merlin home app.yaml even if the test fails + if merlin_home_had_app_yaml: + os.rename(f"{merlin_home_dir}/{COPIED_APP_FILENAME}", f"{merlin_home_dir}/app.yaml") + + raise AssertionError from exc + + # Reset the cwd app.yaml + if cwd_had_app_yaml: + os.rename(f"{cwd_path}/{COPIED_APP_FILENAME}", f"{cwd_path}/app.yaml") + + # Reset the merlin home app.yaml + if merlin_home_had_app_yaml: + os.rename(f"{merlin_home_dir}/{COPIED_APP_FILENAME}", f"{merlin_home_dir}/app.yaml") + + +def test_set_username_and_vhost_nothing_to_load(): + """ + Test the `set_username_and_vhost` function with nothing to load. In other words, in this + test the config dict will have a username and vhost already set for the broker. We'll + create the dict then make a copy of it to test against after calling the function. + """ + actual_config = {"broker": {"username": "default", "vhost": "host4testing"}} + expected_config = deepcopy(actual_config) + assert actual_config is not expected_config + + set_username_and_vhost(actual_config) + + # Ensure that nothing was modified after our call to set_username_and_vhost + assert actual_config == expected_config + + +def test_set_username_and_vhost_no_username(): + """ + Test the `set_username_and_vhost` function with no username. In other words, in this + test the config dict will have vhost already set for the broker but not a username. + """ + expected_config = {"broker": {"username": getpass.getuser(), "vhost": "host4testing"}} + actual_config = {"broker": {"vhost": "host4testing"}} + set_username_and_vhost(actual_config) + + # Ensure that the username was set in the call to set_username_and_vhost + assert actual_config == expected_config + + +def test_set_username_and_vhost_no_vhost(): + """ + Test the `set_username_and_vhost` function with no vhost. In other words, in this + test the config dict will have username already set for the broker but not a vhost. + """ + expected_config = {"broker": {"username": "default", "vhost": getpass.getuser()}} + actual_config = {"broker": {"username": "default"}} + set_username_and_vhost(actual_config) + + # Ensure that the vhost was set in the call to set_username_and_vhost + assert actual_config == expected_config + + +def test_load_default_celery_nothing_to_load(): + """ + Test the `load_default_celery` function with nothing to load. In other words, in this + test the config dict will have a celery entry containing omit_queue_tag, queue_tag, and + override. We'll create the dict then make a copy of it to test against after calling + the function. + """ + actual_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} + expected_config = deepcopy(actual_config) + assert actual_config is not expected_config + + load_default_celery(actual_config) + + # Ensure that nothing was modified after our call to load_default_celery + assert actual_config == expected_config + + +def test_load_default_celery_no_omit_queue_tag(): + """ + Test the `load_default_celery` function with no omit_queue_tag. The function should + create a default entry of False for this. + """ + actual_config = {"celery": {"queue_tag": "[merlin]_", "override": None}} + expected_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} + load_default_celery(actual_config) + + # Ensure that the omit_queue_tag was set in the call to load_default_celery + assert actual_config == expected_config + + +def test_load_default_celery_no_queue_tag(): + """ + Test the `load_default_celery` function with no queue_tag. The function should + create a default entry of '[merlin]_' for this. + """ + actual_config = {"celery": {"omit_queue_tag": False, "override": None}} + expected_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} + load_default_celery(actual_config) + + # Ensure that the queue_tag was set in the call to load_default_celery + assert actual_config == expected_config + + +def test_load_default_celery_no_override(): + """ + Test the `load_default_celery` function with no override. The function should + create a default entry of None for this. + """ + actual_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_"}} + expected_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} + load_default_celery(actual_config) + + # Ensure that the override was set in the call to load_default_celery + assert actual_config == expected_config + + +def test_load_default_celery_no_celery_block(): + """ + Test the `load_default_celery` function with no celery block. The function should + create a default entry of + {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} for this. + """ + actual_config = {} + expected_config = {"celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}} + load_default_celery(actual_config) + + # Ensure that the celery block was set in the call to load_default_celery + assert actual_config == expected_config + + +def test_load_defaults(): + """ + Test that the `load_defaults` function loads the user names and the celery block properly. + """ + actual_config = {"broker": {}} + expected_config = { + "broker": {"username": getpass.getuser(), "vhost": getpass.getuser()}, + "celery": {"omit_queue_tag": False, "queue_tag": "[merlin]_", "override": None}, + } + load_defaults(actual_config) + + assert actual_config == expected_config + + +def test_get_config(temp_output_dir: str): + """ + Test the `get_config` function. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the configfile directory and put an app.yaml file there + configfile_dir = CONFIGFILE_DIR.format(temp_output_dir=temp_output_dir) + create_dir(configfile_dir) + create_app_yaml(configfile_dir) + + # Load up the contents of the dummy app.yaml file that we copied + with open(DUMMY_APP_FILEPATH, "r") as dummy_app_file: + expected = yaml.load(dummy_app_file, yaml.Loader) + + # Add in default settings that should be added + expected["celery"]["omit_queue_tag"] = False + expected["celery"]["queue_tag"] = "[merlin]_" + + actual = get_config(configfile_dir) + + assert actual == expected + + +def test_get_config_invalid_path(): + """ + Test the `get_config` function with an invalid path. This should raise a ValueError. + """ + with pytest.raises(ValueError) as excinfo: + get_config("invalid/path") + + assert "Cannot find a merlin config file!" in str(excinfo.value) + + +def test_is_debug_no_merlin_debug(): + """ + Test the `is_debug` function without having MERLIN_DEBUG in the environment. + This should return False. + """ + + # Delete the current val of MERLIN_DEBUG and store it (if there is one) + reset_merlin_debug = False + debug_val = None + if "MERLIN_DEBUG" in os.environ: + debug_val = copy(os.environ["MERLIN_DEBUG"]) + del os.environ["MERLIN_DEBUG"] + reset_merlin_debug = True + + # Run the test + try: + assert is_debug() is False + except AssertionError as exc: + # Make sure to reset the value of MERLIN_DEBUG even if the test fails + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + raise AssertionError from exc + + # Reset the value of MERLIN_DEBUG + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + + +def test_is_debug_with_merlin_debug(): + """ + Test the `is_debug` function with having MERLIN_DEBUG in the environment. + This should return True. + """ + + # Grab the current value of MERLIN_DEBUG if there is one + reset_merlin_debug = False + debug_val = None + if "MERLIN_DEBUG" in os.environ and int(os.environ["MERLIN_DEBUG"]) != 1: + debug_val = copy(os.environ["MERLIN_DEBUG"]) + reset_merlin_debug = True + + # Set the MERLIN_DEBUG value to be 1 + os.environ["MERLIN_DEBUG"] = "1" + + try: + assert is_debug() is True + except AssertionError as exc: + # Make sure to reset the value of MERLIN_DEBUG even if the test fails + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + raise AssertionError from exc + + # Reset the value of MERLIN_DEBUG + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + + +def test_default_config_info(temp_output_dir: str): + """ + Test the `default_config_info` function. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + + # Create the configfile directory and put an app.yaml file there + configfile_dir = CONFIGFILE_DIR.format(temp_output_dir=temp_output_dir) + create_dir(configfile_dir) + create_app_yaml(configfile_dir) + cwd = os.getcwd() + os.chdir(configfile_dir) + + # Delete the current val of MERLIN_DEBUG and store it (if there is one) + reset_merlin_debug = False + debug_val = None + if "MERLIN_DEBUG" in os.environ: + debug_val = copy(os.environ["MERLIN_DEBUG"]) + del os.environ["MERLIN_DEBUG"] + reset_merlin_debug = True + + # Create the merlin home directory if it doesn't already exist + merlin_home = f"{os.path.expanduser('~')}/.merlin" + remove_merlin_home = False + if not os.path.exists(merlin_home): + os.mkdir(merlin_home) + remove_merlin_home = True + + # Run the test + try: + expected = { + "config_file": f"{configfile_dir}/app.yaml", + "is_debug": False, + "merlin_home": merlin_home, + "merlin_home_exists": True, + } + actual = default_config_info() + assert actual == expected + except AssertionError as exc: + # Make sure to reset values even if the test fails + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + if remove_merlin_home: + os.rmdir(merlin_home) + raise AssertionError from exc + + # Reset values if necessary + if reset_merlin_debug: + os.environ["MERLIN_DEBUG"] = debug_val + if remove_merlin_home: + os.rmdir(merlin_home) + + os.chdir(cwd) + + +def test_get_cert_file_all_valid_args(mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_cert_file` function with all valid arguments. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The path to the temporary merlin server directory that's housing our cert files + """ + expected = f"{merlin_server_dir}/{CERT_FILES['ssl_key']}" + actual = get_cert_file( + server_type="Results Backend", config=CONFIG.results_backend, cert_name="keyfile", cert_path=merlin_server_dir + ) + assert actual == expected + + +def test_get_cert_file_invalid_cert_name(mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_cert_file` function with an invalid cert_name argument. This should just return None. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The path to the temporary merlin server directory that's housing our cert files + """ + actual = get_cert_file( + server_type="Results Backend", config=CONFIG.results_backend, cert_name="invalid", cert_path=merlin_server_dir + ) + assert actual is None + + +def test_get_cert_file_nonexistent_cert_path( + mysql_results_backend_config: "fixture", temp_output_dir: str, merlin_server_dir: str # noqa: F821 +): + """ + Test the `get_cert_file` function with cert_path argument that doesn't exist. + This should still return the nonexistent path at the root of our temporary directory for testing. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + :param merlin_server_dir: The path to the temporary merlin server directory that's housing our cert files + """ + CONFIG.results_backend.certfile = "new_certfile.pem" + expected = f"{temp_output_dir}/new_certfile.pem" + actual = get_cert_file( + server_type="Results Backend", config=CONFIG.results_backend, cert_name="certfile", cert_path=merlin_server_dir + ) + assert actual == expected + + +def test_get_ssl_entries_required_certs(mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_entries` function with mysql as the results_backend. For this test we'll make + cert reqs be required. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + CONFIG.results_backend.cert_reqs = "required" + + expected = { + "ssl_key": f"{temp_output_dir}/{CERT_FILES['ssl_key']}", + "ssl_cert": f"{temp_output_dir}/{CERT_FILES['ssl_cert']}", + "ssl_ca": f"{temp_output_dir}/{CERT_FILES['ssl_ca']}", + "cert_reqs": ssl.CERT_REQUIRED, + } + actual = get_ssl_entries( + server_type="Results Backend", server_name="mysql", server_config=CONFIG.results_backend, cert_path=temp_output_dir + ) + assert expected == actual + + +def test_get_ssl_entries_optional_certs(mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_entries` function with mysql as the results_backend. For this test we'll make + cert reqs be optional. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + CONFIG.results_backend.cert_reqs = "optional" + + expected = { + "ssl_key": f"{temp_output_dir}/{CERT_FILES['ssl_key']}", + "ssl_cert": f"{temp_output_dir}/{CERT_FILES['ssl_cert']}", + "ssl_ca": f"{temp_output_dir}/{CERT_FILES['ssl_ca']}", + "cert_reqs": ssl.CERT_OPTIONAL, + } + actual = get_ssl_entries( + server_type="Results Backend", server_name="mysql", server_config=CONFIG.results_backend, cert_path=temp_output_dir + ) + assert expected == actual + + +def test_get_ssl_entries_none_certs(mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_entries` function with mysql as the results_backend. For this test we won't require + any cert reqs. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + CONFIG.results_backend.cert_reqs = "none" + + expected = { + "ssl_key": f"{temp_output_dir}/{CERT_FILES['ssl_key']}", + "ssl_cert": f"{temp_output_dir}/{CERT_FILES['ssl_cert']}", + "ssl_ca": f"{temp_output_dir}/{CERT_FILES['ssl_ca']}", + "cert_reqs": ssl.CERT_NONE, + } + actual = get_ssl_entries( + server_type="Results Backend", server_name="mysql", server_config=CONFIG.results_backend, cert_path=temp_output_dir + ) + assert expected == actual + + +def test_get_ssl_entries_omit_certs(mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_entries` function with mysql as the results_backend. For this test we'll completely + omit the cert_reqs option + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + del CONFIG.results_backend.cert_reqs + + expected = { + "ssl_key": f"{temp_output_dir}/{CERT_FILES['ssl_key']}", + "ssl_cert": f"{temp_output_dir}/{CERT_FILES['ssl_cert']}", + "ssl_ca": f"{temp_output_dir}/{CERT_FILES['ssl_ca']}", + "cert_reqs": ssl.CERT_REQUIRED, + } + actual = get_ssl_entries( + server_type="Results Backend", server_name="mysql", server_config=CONFIG.results_backend, cert_path=temp_output_dir + ) + assert expected == actual + + +def test_get_ssl_entries_with_ssl_protocol(mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_entries` function with mysql as the results_backend. For this test we'll add in a + dummy ssl_protocol value that should get added to the dict that's output. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + protocol = "test_protocol" + CONFIG.results_backend.ssl_protocol = protocol + + expected = { + "ssl_key": f"{temp_output_dir}/{CERT_FILES['ssl_key']}", + "ssl_cert": f"{temp_output_dir}/{CERT_FILES['ssl_cert']}", + "ssl_ca": f"{temp_output_dir}/{CERT_FILES['ssl_ca']}", + "cert_reqs": ssl.CERT_NONE, + "ssl_protocol": protocol, + } + actual = get_ssl_entries( + server_type="Results Backend", server_name="mysql", server_config=CONFIG.results_backend, cert_path=temp_output_dir + ) + assert expected == actual + + +def test_process_ssl_map_mysql(): + """Test the `process_ssl_map` function with mysql as the server name.""" + expected = {"keyfile": "ssl_key", "certfile": "ssl_cert", "ca_certs": "ssl_ca"} + actual = process_ssl_map("mysql") + assert actual == expected + + +def test_process_ssl_map_rediss(): + """Test the `process_ssl_map` function with rediss as the server name.""" + expected = { + "keyfile": "ssl_keyfile", + "certfile": "ssl_certfile", + "ca_certs": "ssl_ca_certs", + "cert_reqs": "ssl_cert_reqs", + } + actual = process_ssl_map("rediss") + assert actual == expected + + +def test_merge_sslmap_all_keys_present(): + """ + Test the `merge_sslmap` function with all keys from server_ssl in ssl_map. + We'll assume we're using a rediss server for this. + """ + expected = { + "ssl_keyfile": "/path/to/keyfile", + "ssl_certfile": "/path/to/certfile", + "ssl_ca_certs": "/path/to/ca_file", + "ssl_cert_reqs": ssl.CERT_NONE, + } + test_server_ssl = { + "keyfile": "/path/to/keyfile", + "certfile": "/path/to/certfile", + "ca_certs": "/path/to/ca_file", + "cert_reqs": ssl.CERT_NONE, + } + test_ssl_map = { + "keyfile": "ssl_keyfile", + "certfile": "ssl_certfile", + "ca_certs": "ssl_ca_certs", + "cert_reqs": "ssl_cert_reqs", + } + actual = merge_sslmap(test_server_ssl, test_ssl_map) + assert actual == expected + + +def test_merge_sslmap_some_keys_present(): + """ + Test the `merge_sslmap` function with some keys from server_ssl in ssl_map and others not. + We'll assume we're using a rediss server for this. + """ + expected = { + "ssl_keyfile": "/path/to/keyfile", + "ssl_certfile": "/path/to/certfile", + "ssl_ca_certs": "/path/to/ca_file", + "ssl_cert_reqs": ssl.CERT_NONE, + "new_key": "new_val", + "second_new_key": "second_new_val", + } + test_server_ssl = { + "keyfile": "/path/to/keyfile", + "certfile": "/path/to/certfile", + "ca_certs": "/path/to/ca_file", + "cert_reqs": ssl.CERT_NONE, + "new_key": "new_val", + "second_new_key": "second_new_val", + } + test_ssl_map = { + "keyfile": "ssl_keyfile", + "certfile": "ssl_certfile", + "ca_certs": "ssl_ca_certs", + "cert_reqs": "ssl_cert_reqs", + } + actual = merge_sslmap(test_server_ssl, test_ssl_map) + assert actual == expected diff --git a/tests/unit/config/test_results_backend.py b/tests/unit/config/test_results_backend.py new file mode 100644 index 000000000..55459f3ec --- /dev/null +++ b/tests/unit/config/test_results_backend.py @@ -0,0 +1,593 @@ +""" +Tests for the `results_backend.py` file. +""" + +import os +from ssl import CERT_NONE +from typing import Any, Dict + +import pytest + +from merlin.config.configfile import CONFIG +from merlin.config.results_backend import ( + MYSQL_CONFIG_FILENAMES, + MYSQL_CONNECTION_STRING, + SQLITE_CONNECTION_STRING, + get_backend_password, + get_connection_string, + get_mysql, + get_mysql_config, + get_redis, + get_ssl_config, +) +from tests.constants import CERT_FILES, SERVER_PASS +from tests.utils import create_cert_files, create_pass_file + + +RESULTS_BACKEND_DIR = "{temp_output_dir}/test_results_backend" + + +def test_get_backend_password_pass_file_in_merlin(): + """ + Test the `get_backend_password` function with the password file in the ~/.merlin/ + directory. We'll create a dummy file in this directory and delete it once the test + is done. + """ + + # Check if the .merlin directory exists and create it if it doesn't + remove_merlin_dir_after_test = False + path_to_merlin_dir = os.path.expanduser("~/.merlin") + if not os.path.exists(path_to_merlin_dir): + remove_merlin_dir_after_test = True + os.mkdir(path_to_merlin_dir) + + # Create the test password file + pass_filename = "test.pass" + full_pass_filepath = f"{path_to_merlin_dir}/{pass_filename}" + create_pass_file(full_pass_filepath) + + try: + # Run the test + assert get_backend_password(pass_filename) == SERVER_PASS + # Cleanup + os.remove(full_pass_filepath) + if remove_merlin_dir_after_test: + os.rmdir(path_to_merlin_dir) + except AssertionError as exc: + # If the test fails, make sure we clean up the files/dirs created + os.remove(full_pass_filepath) + if remove_merlin_dir_after_test: + os.rmdir(path_to_merlin_dir) + raise AssertionError from exc + + +def test_get_backend_password_pass_file_not_in_merlin(temp_output_dir: str): + """ + Test the `get_backend_password` function with the password file not in the ~/.merlin/ + directory. By using the `temp_output_dir` fixture, our cwd will be the temporary directory. + We'll create a password file in the this directory for this test and have `get_backend_password` + read from that. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + pass_file = "test.pass" + create_pass_file(pass_file) + + assert get_backend_password(pass_file) == SERVER_PASS + + +def test_get_backend_password_directly_pass_password(): + """ + Test the `get_backend_password` function by passing the password directly to this + function instead of a password file. + """ + assert get_backend_password(SERVER_PASS) == SERVER_PASS + + +def test_get_backend_password_using_certs_path(temp_output_dir: str): + """ + Test the `get_backend_password` function with certs_path set to our temporary testing path. + We'll create a password file in the temporary directory for this test and have `get_backend_password` + read from that. + + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + pass_filename = "test_certs.pass" + test_dir = RESULTS_BACKEND_DIR.format(temp_output_dir=temp_output_dir) + if not os.path.exists(test_dir): + os.mkdir(test_dir) + full_pass_filepath = f"{test_dir}/{pass_filename}" + create_pass_file(full_pass_filepath) + + assert get_backend_password(pass_filename, certs_path=test_dir) == SERVER_PASS + + +def test_get_ssl_config_no_results_backend(config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with no results_backend set. This should return False. + NOTE: we're using the config fixture here to make sure values are reset after this test finishes. + We won't actually use anything from the config fixture. + + :param config: A fixture to set up the CONFIG object for us + """ + del CONFIG.results_backend.name + assert get_ssl_config() is False + + +def test_get_connection_string_no_results_backend(config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with no results_backend set. + This should raise a ValueError. + NOTE: we're using the config fixture here to make sure values are reset after this test finishes. + We won't actually use anything from the config fixture. + + :param config: A fixture to set up the CONFIG object for us + """ + del CONFIG.results_backend.name + with pytest.raises(ValueError) as excinfo: + get_connection_string() + + assert "'' is not a supported results backend" in str(excinfo.value) + + +class TestRedisResultsBackend: + """ + This class will house all tests necessary for our results_backend module when using a + redis results_backend. + """ + + def run_get_redis( + self, + expected_vals: Dict[str, Any], + certs_path: str = None, + include_password: bool = True, + ssl: bool = False, + ): + """ + Helper method for running tests for the `get_redis` function. + + :param expected_vals: A dict of expected values for this test. Format: + {"urlbase": "redis", + "spass": "", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0} + :param certs_path: A string denoting the path to the certification files + :param include_password: If True, include the password in the output. Otherwise don't. + :param ssl: If True, use ssl. Otherwise, don't. + """ + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_redis(certs_path=certs_path, include_password=include_password, ssl=ssl) + assert actual == expected + + def test_get_redis(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with default functionality. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "redis", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + + def test_get_redis_dont_include_password(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with the password hidden. This should * out the password. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "redis", + "spass": "default:******@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=False, ssl=False) + + def test_get_redis_using_ssl(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with ssl enabled. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "rediss", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=True) + + def test_get_redis_no_port(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with no port in our CONFIG object. This should default to port=6379. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.results_backend.port + expected_vals = { + "urlbase": "redis", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + + def test_get_redis_no_db_num(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with no db_num in our CONFIG object. This should default to db_num=0. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.results_backend.db_num + expected_vals = { + "urlbase": "redis", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + + def test_get_redis_no_username(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with no username in our CONFIG object. This should default to username=''. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.results_backend.username + expected_vals = { + "urlbase": "redis", + "spass": f":{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + + def test_get_redis_no_password_file(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function with no password filepath in our CONFIG object. This should default to spass=''. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.results_backend.password + expected_vals = { + "urlbase": "redis", + "spass": "", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + + def test_get_redis_invalid_pass_file(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_redis` function. We'll run this after changing the permissions of the password file so it + can't be opened. This should still run and give us password=CONFIG.results_backend.password. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + + # Capture the initial permissions of the password file so we can reset them + orig_file_permissions = os.stat(CONFIG.results_backend.password).st_mode + + # Change the permissions of the password file so it can't be read + os.chmod(CONFIG.results_backend.password, 0o222) + + try: + # Run the test + expected_vals = { + "urlbase": "redis", + "spass": f"default:{CONFIG.results_backend.password}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + self.run_get_redis(expected_vals=expected_vals, certs_path=None, include_password=True, ssl=False) + os.chmod(CONFIG.results_backend.password, orig_file_permissions) + except AssertionError as exc: + # If this test failed, make sure to reset the permissions in case other tests need to read this file + os.chmod(CONFIG.results_backend.password, orig_file_permissions) + raise AssertionError from exc + + def test_get_ssl_config_redis(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with redis as the results_backend. This should return False since + ssl requires using rediss (with two 's'). + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + assert get_ssl_config() is False + + def test_get_ssl_config_rediss(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with rediss as the results_backend. + This should return a dict of cert reqs with ssl.CERT_NONE as the value. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.name = "rediss" + assert get_ssl_config() == {"ssl_cert_reqs": CERT_NONE} + + def test_get_ssl_config_rediss_no_cert_reqs(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with rediss as the results_backend and no cert_reqs set. + This should return True. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + del CONFIG.results_backend.cert_reqs + CONFIG.results_backend.name = "rediss" + assert get_ssl_config() is True + + def test_get_connection_string_redis(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with redis as the results_backend. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + expected_vals = { + "urlbase": "redis", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_connection_string() + assert actual == expected + + def test_get_connection_string_rediss(self, redis_results_backend_config_function: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with rediss as the results_backend. + + :param redis_results_backend_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.name = "rediss" + expected_vals = { + "urlbase": "rediss", + "spass": f"default:{SERVER_PASS}@", + "server": "127.0.0.1", + "port": 6379, + "db_num": 0, + } + expected = "{urlbase}://{spass}{server}:{port}/{db_num}".format(**expected_vals) + actual = get_connection_string() + assert actual == expected + + +class TestMySQLResultsBackend: + """ + This class will house all tests necessary for our results_backend module when using a + MySQL results_backend. + NOTE: You'll notice a lot of these tests are setting CONFIG.results_backend.name to be + "invalid". This is so that we can get by the first if statement in the `get_mysql_config` + function. + """ + + def test_get_mysql_config_certs_set(self, mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_mysql_config` function with the certs dict getting set and returned. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.results_backend.name = "invalid" + expected = {} + for key, cert_file in CERT_FILES.items(): + expected[key] = f"{merlin_server_dir}/{cert_file}" + actual = get_mysql_config(merlin_server_dir, CERT_FILES) + assert actual == expected + + def test_get_mysql_config_ssl_exists(self, mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_mysql_config` function with mysql_ssl being found. This should just return the ssl value that's found. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + expected = {key: f"{temp_output_dir}/{cert_file}" for key, cert_file in CERT_FILES.items()} + expected["cert_reqs"] = CERT_NONE + assert get_mysql_config(None, None) == expected + + def test_get_mysql_config_no_mysql_certs( + self, mysql_results_backend_config: "fixture", merlin_server_dir: str # noqa: F821 + ): + """ + Test the `get_mysql_config` function with no mysql certs dict. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.results_backend.name = "invalid" + assert get_mysql_config(merlin_server_dir, {}) == {} + + def test_get_mysql_config_invalid_certs_path(self, mysql_results_backend_config: "fixture"): # noqa: F821 + """ + Test the `get_mysql_config` function with an invalid certs path. This should return False. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.name = "invalid" + assert get_mysql_config("invalid/path", CERT_FILES) is False + + def run_get_mysql( + self, expected_vals: Dict[str, Any], certs_path: str, mysql_certs: Dict[str, str], include_password: bool + ): + """ + Helper method for running tests for the `get_mysql` function. + + :param expected_vals: A dict of expected values for this test. Format: + {"cert_reqs": cert reqs dict, + "user": "default", + "password": "", + "server": "127.0.0.1", + "ssl_cert": "test-rabbit-client-cert.pem", + "ssl_ca": "test-mysql-ca-cert.pem", + "ssl_key": "test-rabbit-client-key.pem"} + :param certs_path: A string denoting the path to the certification files + :param mysql_certs: A dict of cert files + :param include_password: If True, include the password in the output. Otherwise don't. + """ + expected = MYSQL_CONNECTION_STRING.format(**expected_vals) + actual = get_mysql(certs_path=certs_path, mysql_certs=mysql_certs, include_password=include_password) + assert actual == expected + + def test_get_mysql(self, mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_mysql` function with default behavior. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.results_backend.name = "invalid" + expected_vals = { + "cert_reqs": CERT_NONE, + "user": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + } + for key, cert_file in CERT_FILES.items(): + expected_vals[key] = f"{merlin_server_dir}/{cert_file}" + self.run_get_mysql( + expected_vals=expected_vals, certs_path=merlin_server_dir, mysql_certs=CERT_FILES, include_password=True + ) + + def test_get_mysql_dont_include_password( + self, mysql_results_backend_config: "fixture", merlin_server_dir: str # noqa: F821 + ): + """ + Test the `get_mysql` function but set include_password to False. This should * out the password. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.results_backend.name = "invalid" + expected_vals = { + "cert_reqs": CERT_NONE, + "user": "default", + "password": "******", + "server": "127.0.0.1", + } + for key, cert_file in CERT_FILES.items(): + expected_vals[key] = f"{merlin_server_dir}/{cert_file}" + self.run_get_mysql( + expected_vals=expected_vals, certs_path=merlin_server_dir, mysql_certs=CERT_FILES, include_password=False + ) + + def test_get_mysql_no_mysql_certs(self, mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_mysql` function with no mysql_certs passed in. This should use default config filenames so we'll + have to create these default files. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.results_backend.name = "invalid" + expected_vals = { + "cert_reqs": CERT_NONE, + "user": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + } + + create_cert_files(merlin_server_dir, MYSQL_CONFIG_FILENAMES) + + for key, cert_file in MYSQL_CONFIG_FILENAMES.items(): + # Password file is already is already set in expected_vals dict + if key == "password": + continue + expected_vals[key] = f"{merlin_server_dir}/{cert_file}" + + self.run_get_mysql(expected_vals=expected_vals, certs_path=merlin_server_dir, mysql_certs=None, include_password=True) + + def test_get_mysql_no_server(self, mysql_results_backend_config: "fixture"): # noqa: F821 + """ + Test the `get_mysql` function with no server set. This should raise a TypeError. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.server = False + with pytest.raises(TypeError) as excinfo: + get_mysql() + assert "Results backend: server False does not have a configuration" in str(excinfo.value) + + def test_get_mysql_invalid_certs_path(self, mysql_results_backend_config: "fixture"): # noqa: F821 + """ + Test the `get_mysql` function with an invalid certs_path. This should raise a TypeError. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.name = "invalid" + with pytest.raises(TypeError) as excinfo: + get_mysql(certs_path="invalid_path", mysql_certs=CERT_FILES) + err_msg = f"""The connection information for MySQL could not be set, cannot find:\n + {CERT_FILES}\ncheck the celery/certs path or set the ssl information in the app.yaml file.""" + assert err_msg in str(excinfo.value) + + def test_get_ssl_config_mysql(self, mysql_results_backend_config: "fixture", temp_output_dir: str): # noqa: F821 + """ + Test the `get_ssl_config` function with mysql as the results_backend. + This should return a dict of cert reqs with ssl.CERT_NONE as the value. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param temp_output_dir: The path to the temporary output directory we'll be using for this test run + """ + expected = {key: f"{temp_output_dir}/{cert_file}" for key, cert_file in CERT_FILES.items()} + expected["cert_reqs"] = CERT_NONE + assert get_ssl_config() == expected + + def test_get_ssl_config_mysql_celery_check(self, mysql_results_backend_config: "fixture"): # noqa: F821 + """ + Test the `get_ssl_config` function with mysql as the results_backend and celery_check set. + This should return False. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + assert get_ssl_config(celery_check=True) is False + + def test_get_connection_string_mysql(self, mysql_results_backend_config: "fixture", merlin_server_dir: str): # noqa: F821 + """ + Test the `get_connection_string` function with MySQL as the results_backend. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + :param merlin_server_dir: The directory that has the test certification files + """ + CONFIG.celery.certs = merlin_server_dir + + create_cert_files(merlin_server_dir, MYSQL_CONFIG_FILENAMES) + CONFIG.results_backend.keyfile = MYSQL_CONFIG_FILENAMES["ssl_key"] + CONFIG.results_backend.certfile = MYSQL_CONFIG_FILENAMES["ssl_cert"] + CONFIG.results_backend.ca_certs = MYSQL_CONFIG_FILENAMES["ssl_ca"] + + expected_vals = { + "cert_reqs": CERT_NONE, + "user": "default", + "password": SERVER_PASS, + "server": "127.0.0.1", + } + for key, cert_file in MYSQL_CONFIG_FILENAMES.items(): + # Password file is already is already set in expected_vals dict + if key == "password": + continue + expected_vals[key] = f"{merlin_server_dir}/{cert_file}" + + assert MYSQL_CONNECTION_STRING.format(**expected_vals) == get_connection_string(include_password=True) + + def test_get_connection_string_sqlite(self, mysql_results_backend_config: "fixture"): # noqa: F821 + """ + Test the `get_connection_string` function with sqlite as the results_backend. + + :param mysql_results_backend_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.results_backend.name = "sqlite" + assert get_connection_string() == SQLITE_CONNECTION_STRING diff --git a/tests/unit/config/test_utils.py b/tests/unit/config/test_utils.py new file mode 100644 index 000000000..b4f7c52fa --- /dev/null +++ b/tests/unit/config/test_utils.py @@ -0,0 +1,117 @@ +""" +Tests for the merlin/config/utils.py module. +""" + +import pytest + +from merlin.config.configfile import CONFIG +from merlin.config.utils import Priority, determine_priority_map, get_priority, is_rabbit_broker, is_redis_broker + + +def test_is_rabbit_broker(): + """Test the `is_rabbit_broker` by passing in rabbit as the broker""" + assert is_rabbit_broker("rabbitmq") is True + assert is_rabbit_broker("amqp") is True + assert is_rabbit_broker("amqps") is True + + +def test_is_rabbit_broker_invalid(): + """Test the `is_rabbit_broker` by passing in an invalid broker""" + assert is_rabbit_broker("redis") is False + assert is_rabbit_broker("") is False + + +def test_is_redis_broker(): + """Test the `is_redis_broker` by passing in redis as the broker""" + assert is_redis_broker("redis") is True + assert is_redis_broker("rediss") is True + assert is_redis_broker("redis+socket") is True + + +def test_is_redis_broker_invalid(): + """Test the `is_redis_broker` by passing in an invalid broker""" + assert is_redis_broker("rabbitmq") is False + assert is_redis_broker("") is False + + +def test_get_priority_rabbit_broker(rabbit_broker_config: "fixture"): # noqa: F821 + """ + Test the `get_priority` function with rabbit as the broker. + Low priority for rabbit is 1 and high is 9. + + :param rabbit_broker_config: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + assert get_priority(Priority.LOW) == 1 + assert get_priority(Priority.MID) == 5 + assert get_priority(Priority.HIGH) == 9 + assert get_priority(Priority.RETRY) == 10 + + +def test_get_priority_redis_broker(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_priority` function with redis as the broker. + Low priority for redis is 10 and high is 2. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + assert get_priority(Priority.LOW) == 10 + assert get_priority(Priority.MID) == 5 + assert get_priority(Priority.HIGH) == 2 + assert get_priority(Priority.RETRY) == 1 + + +def test_get_priority_invalid_broker(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_priority` function with an invalid broker. + This should raise a ValueError. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + CONFIG.broker.name = "invalid" + with pytest.raises(ValueError) as excinfo: + get_priority(Priority.LOW) + assert "Unsupported broker name: invalid" in str(excinfo.value) + + +def test_get_priority_invalid_priority(redis_broker_config_function: "fixture"): # noqa: F821 + """ + Test the `get_priority` function with an invalid priority. + This should raise a TypeError. + + :param redis_broker_config_function: A fixture to set the CONFIG object to a test configuration that we'll use here + """ + with pytest.raises(ValueError) as excinfo: + get_priority("invalid_priority") + assert "Invalid priority: invalid_priority" in str(excinfo.value) + + +def test_determine_priority_map_rabbit(): + """ + Test the `determine_priority_map` function with rabbit as the broker. + This should return the following map: + {Priority.LOW: 1, Priority.MID: 5, Priority.HIGH: 9, Priority.RETRY: 10} + """ + expected = {Priority.LOW: 1, Priority.MID: 5, Priority.HIGH: 9, Priority.RETRY: 10} + actual = determine_priority_map("rabbitmq") + assert actual == expected + + +def test_determine_priority_map_redis(): + """ + Test the `determine_priority_map` function with redis as the broker. + This should return the following map: + {Priority.LOW: 10, Priority.MID: 5, Priority.HIGH: 2, Priority.RETRY: 1} + """ + expected = {Priority.LOW: 10, Priority.MID: 5, Priority.HIGH: 2, Priority.RETRY: 1} + actual = determine_priority_map("redis") + assert actual == expected + + +def test_determine_priority_map_invalid(): + """ + Test the `determine_priority_map` function with an invalid broker. + This should raise a ValueError. + """ + with pytest.raises(ValueError) as excinfo: + determine_priority_map("invalid_broker") + assert "Unsupported broker name: invalid_broker" in str(excinfo.value) diff --git a/tests/unit/config/utils.py b/tests/unit/config/utils.py deleted file mode 100644 index 1765e8478..000000000 --- a/tests/unit/config/utils.py +++ /dev/null @@ -1,24 +0,0 @@ -""" -Utils module for common test functionality. -""" - -import os - - -def mkfile(tmpdir, filename, content=""): - """ - A simple function for creating a file and returning the path. This is to - abstract out file creation logic in the tests. - - :param tmpdir: (str) The path to the temp directory. - :param filename: (str) The name of the file. - :param contents: (str) Optional contents to write to the file. Defaults to - an empty string. - :returns: (str) The appended path of the given tempdir and filename. - """ - filepath = os.path.join(tmpdir, filename) - - with open(filepath, "w") as f: - f.write(content) - - return filepath diff --git a/tests/unit/server/__init__.py b/tests/unit/server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/server/test_RedisConfig.py b/tests/unit/server/test_RedisConfig.py new file mode 100644 index 000000000..321d2f38a --- /dev/null +++ b/tests/unit/server/test_RedisConfig.py @@ -0,0 +1,556 @@ +""" +Tests for the RedisConfig class of the `server_util.py` module. + +This class is especially large so that's why these tests have been +moved to their own file. +""" + +import filecmp +import logging +from typing import Any + +import pytest + +from merlin.server.server_util import RedisConfig + + +class TestRedisConfig: + """Tests for the RedisConfig class.""" + + def test_initialization(self, server_redis_conf_file: str): + """ + Using a dummy redis configuration file, test that the initialization + of the RedisConfig class behaves as expected. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + expected_entries = { + "bind": "127.0.0.1", + "port": "6379", + "requirepass": "merlin_password", + "dir": "./", + "save": "300 100", + "dbfilename": "dump.rdb", + "appendfsync": "everysec", + "appendfilename": "appendonly.aof", + } + expected_comments = { + "bind": "# ip address\n", + "port": "\n# port\n", + "requirepass": "\n# password\n", + "dir": "\n# directory\n", + "save": "\n# snapshot\n", + "dbfilename": "\n# db file\n", + "appendfsync": "\n# append mode\n", + "appendfilename": "\n# append file\n", + } + expected_trailing_comment = "\n# dummy trailing comment" + expected_entry_order = list(expected_entries.keys()) + redis_config = RedisConfig(server_redis_conf_file) + assert redis_config.filename == server_redis_conf_file + assert not redis_config.changed + assert redis_config.entries == expected_entries + assert redis_config.entry_order == expected_entry_order + assert redis_config.comments == expected_comments + assert redis_config.trailing_comments == expected_trailing_comment + + def test_write(self, server_redis_conf_file: str, server_testing_dir: str): + """ + Test that the write functionality works by writing the contents of a dummy + configuration file to a blank configuration file. + + :param server_redis_conf_file: The path to a dummy redis configuration file + :param server_testing_dir: The path to the the temp output directory for server tests + """ + copy_redis_conf_file = f"{server_testing_dir}/redis_copy.conf" + + # Create a RedisConf object with the basic redis conf file + redis_config = RedisConfig(server_redis_conf_file) + + # Change the filepath of the redis config file to be the copy that we'll write to + redis_config.set_filename(copy_redis_conf_file) + + # Run the test + redis_config.write() + + # Check that the contents of the copied file match the contents of the basic file + assert filecmp.cmp(server_redis_conf_file, copy_redis_conf_file) + + @pytest.mark.parametrize("key, val, expected_return", [("port", 1234, True), ("invalid_key", "dummy_val", False)]) + def test_set_config_value(self, server_redis_conf_file: str, key: str, val: Any, expected_return: bool): + """ + Test the `set_config_value` method with valid and invalid keys. + + :param server_redis_conf_file: The path to a dummy redis configuration file + :param key: The key value to modify with `set_config_value` + :param val: The value to set `key` to + :param expected_return: The expected return from `set_config_value` + """ + redis_config = RedisConfig(server_redis_conf_file) + actual_return = redis_config.set_config_value(key, val) + assert actual_return == expected_return + if expected_return: + assert redis_config.entries[key] == val + assert redis_config.changes_made() + else: + assert not redis_config.changes_made() + + @pytest.mark.parametrize( + "key, expected_val", + [ + ("bind", "127.0.0.1"), + ("port", "6379"), + ("requirepass", "merlin_password"), + ("dir", "./"), + ("save", "300 100"), + ("dbfilename", "dump.rdb"), + ("appendfsync", "everysec"), + ("appendfilename", "appendonly.aof"), + ("invalid_key", None), + ], + ) + def test_get_config_value(self, server_redis_conf_file: str, key: str, expected_val: str): + """ + Test the `get_config_value` method with valid and invalid keys. + + :param server_redis_conf_file: The path to a dummy redis configuration file + :param key: The key value to modify with `set_config_value` + :param expected_val: The value we're expecting to get by querying `key` + """ + redis_conf = RedisConfig(server_redis_conf_file) + assert redis_conf.get_config_value(key) == expected_val + + @pytest.mark.parametrize( + "ip_to_set", + [ + "127.0.0.1", # Most common IP + "0.0.0.0", # Edge case (low) + "255.255.255.255", # Edge case (high) + "123.222.199.20", # Random valid IP + ], + ) + def test_set_ip_address_valid(self, caplog: "Fixture", server_redis_conf_file: str, ip_to_set: str): # noqa: F821 + """ + Test the `set_ip_address` method with valid ips. These should all return True + and set the 'bind' value to whatever `ip_to_set` is. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param ip_to_set: The ip address to set + """ + caplog.set_level(logging.INFO) + redis_config = RedisConfig(server_redis_conf_file) + assert redis_config.set_ip_address(ip_to_set) + assert f"Ipaddress is set to {ip_to_set}" in caplog.text, "Missing expected log message" + assert redis_config.get_ip_address() == ip_to_set + + @pytest.mark.parametrize( + "ip_to_set, expected_log", + [ + (None, None), # No IP + ("0.0.0", "Invalid IPv4 address given."), # Invalid IPv4 + ("bind-unset", "Unable to set ip address for redis config"), # Special invalid case where bind doesn't exist + ], + ) + def test_set_ip_address_invalid( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + ip_to_set: str, + expected_log: str, + ): + """ + Test the `set_ip_address` method with invalid ips. These should all return False. + and not modify the 'bind' setting. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param ip_to_set: The ip address to set + :param expected_log: The string we're expecting the logger to log + """ + redis_config = RedisConfig(server_redis_conf_file) + # For the test where bind is unset, delete bind from dict and set new ip val to a valid value + if ip_to_set == "bind-unset": + del redis_config.entries["bind"] + ip_to_set = "127.0.0.1" + assert not redis_config.set_ip_address(ip_to_set) + assert redis_config.get_ip_address() != ip_to_set + if expected_log is not None: + assert expected_log in caplog.text, "Missing expected log message" + + @pytest.mark.parametrize( + "port_to_set", + [ + 6379, # Most common port + 1, # Edge case (low) + 65535, # Edge case (high) + 12345, # Random valid port + ], + ) + def test_set_port_valid( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + port_to_set: str, + ): + """ + Test the `set_port` method with valid ports. These should all return True + and set the 'port' value to whatever `port_to_set` is. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param port_to_set: The port to set + """ + caplog.set_level(logging.INFO) + redis_config = RedisConfig(server_redis_conf_file) + assert redis_config.set_port(port_to_set) + assert redis_config.get_port() == port_to_set + assert f"Port is set to {port_to_set}" in caplog.text, "Missing expected log message" + + @pytest.mark.parametrize( + "port_to_set, expected_log", + [ + (None, None), # No port + (0, "Invalid port given."), # Edge case (low) + (65536, "Invalid port given."), # Edge case (high) + ("port-unset", "Unable to set port for redis config"), # Special invalid case where port doesn't exist + ], + ) + def test_set_port_invalid( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + port_to_set: str, + expected_log: str, + ): + """ + Test the `set_port` method with invalid inputs. These should all return False + and not modify the 'port' setting. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param port_to_set: The port to set + :param expected_log: The string we're expecting the logger to log + """ + redis_config = RedisConfig(server_redis_conf_file) + # For the test where port is unset, delete port from dict and set port val to a valid value + if port_to_set == "port-unset": + del redis_config.entries["port"] + port_to_set = 5 + assert not redis_config.set_port(port_to_set) + assert redis_config.get_port() != port_to_set + if expected_log is not None: + assert expected_log in caplog.text, "Missing expected log message" + + @pytest.mark.parametrize( + "pass_to_set, expected_return", + [ + ("valid_password", True), # Valid password + (None, False), # Invalid password + ], + ) + def test_set_password( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + pass_to_set: str, + expected_return: bool, + ): + """ + Test the `set_password` method with both valid and invalid input. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param pass_to_set: The password to set + :param expected_return: The expected return value + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + assert redis_conf.set_password(pass_to_set) == expected_return + if expected_return: + assert redis_conf.get_password() == pass_to_set + assert "New password set" in caplog.text, "Missing expected log message" + + def test_set_directory_valid( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + server_testing_dir: str, + ): + """ + Test the `set_directory` method with valid input. This should return True, modify the + 'dir' value, and log some messages about creating/setting the directory. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param server_testing_dir: The path to the the temp output directory for server tests + """ + caplog.set_level(logging.INFO) + redis_config = RedisConfig(server_redis_conf_file) + dir_to_set = f"{server_testing_dir}/dummy_dir" + assert redis_config.set_directory(dir_to_set) + assert redis_config.get_config_value("dir") == dir_to_set + assert f"Created directory {dir_to_set}" in caplog.text, "Missing created log message" + assert f"Directory is set to {dir_to_set}" in caplog.text, "Missing set log message" + + def test_set_directory_none(self, server_redis_conf_file: str): + """ + Test the `set_directory` method with None as the input. This should return False + and not modify the 'dir' setting. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_config = RedisConfig(server_redis_conf_file) + assert not redis_config.set_directory(None) + assert redis_config.get_config_value("dir") is not None + + def test_set_directory_dir_unset( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + server_testing_dir: str, + ): + """ + Test the `set_directory` method with the 'dir' setting not existing. This should + return False and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param server_testing_dir: The path to the the temp output directory for server tests + """ + redis_config = RedisConfig(server_redis_conf_file) + del redis_config.entries["dir"] + dir_to_set = f"{server_testing_dir}/dummy_dir" + assert not redis_config.set_directory(dir_to_set) + assert "Unable to set directory for redis config" in caplog.text, "Missing expected log message" + + def test_set_snapshot_valid(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot` method with a valid input for 'seconds' and 'changes'. + This should return True and modify both values of 'save'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + snap_sec_to_set = 20 + snap_changes_to_set = 30 + assert redis_conf.set_snapshot(seconds=snap_sec_to_set, changes=snap_changes_to_set) + save_val = redis_conf.get_config_value("save").split() + assert save_val[0] == str(snap_sec_to_set) + assert save_val[1] == str(snap_changes_to_set) + expected_log = ( + f"Snapshot wait time is set to {snap_sec_to_set} seconds. " + f"Snapshot threshold is set to {snap_changes_to_set} changes" + ) + assert expected_log in caplog.text, "Missing expected log message" + + def test_set_snapshot_just_seconds(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot` method with a valid input for 'seconds'. This should + return True and modify the first value of 'save'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + orig_save = redis_conf.get_config_value("save").split() + snap_sec_to_set = 20 + assert redis_conf.set_snapshot(seconds=snap_sec_to_set) + save_val = redis_conf.get_config_value("save").split() + assert save_val[0] == str(snap_sec_to_set) + assert save_val[1] == orig_save[1] + expected_log = f"Snapshot wait time is set to {snap_sec_to_set} seconds. " + assert expected_log in caplog.text, "Missing expected log message" + + def test_set_snapshot_just_changes(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot` method with a valid input for 'changes'. This should + return True and modify the second value of 'save'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + orig_save = redis_conf.get_config_value("save").split() + snap_changes_to_set = 30 + assert redis_conf.set_snapshot(changes=snap_changes_to_set) + save_val = redis_conf.get_config_value("save").split() + assert save_val[0] == orig_save[0] + assert save_val[1] == str(snap_changes_to_set) + expected_log = f"Snapshot threshold is set to {snap_changes_to_set} changes" + assert expected_log in caplog.text, "Missing expected log message" + + def test_set_snapshot_none(self, server_redis_conf_file: str): + """ + Test the `set_snapshot` method with None as the input for both seconds + and changes. This should return False. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + assert not redis_conf.set_snapshot(seconds=None, changes=None) + + def test_set_snapshot_save_unset(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot` method with the 'save' setting not existing. This should + return False and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + del redis_conf.entries["save"] + assert not redis_conf.set_snapshot(seconds=20) + assert "Unable to get exisiting parameter values for snapshot" in caplog.text, "Missing expected log message" + + def test_set_snapshot_file_valid(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot_file` method with a valid input. This should + return True and modify the value of 'dbfilename'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + filename = "dummy_file.rdb" + assert redis_conf.set_snapshot_file(filename) + assert redis_conf.get_config_value("dbfilename") == filename + assert f"Snapshot file is set to {filename}" in caplog.text, "Missing expected log message" + + def test_set_snapshot_file_none(self, server_redis_conf_file: str): + """ + Test the `set_snapshot_file` method with None as the input. + This should return False. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + assert not redis_conf.set_snapshot_file(None) + + def test_set_snapshot_file_dbfilename_unset(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_snapshot` method with the 'dbfilename' setting not existing. This should + return False and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + del redis_conf.entries["dbfilename"] + filename = "dummy_file.rdb" + assert not redis_conf.set_snapshot_file(filename) + assert redis_conf.get_config_value("dbfilename") != filename + assert "Unable to set snapshot_file name" in caplog.text, "Missing expected log message" + + @pytest.mark.parametrize( + "mode_to_set", + [ + "always", + "everysec", + "no", + ], + ) + def test_set_append_mode_valid( + self, + caplog: "Fixture", # noqa: F821 + server_redis_conf_file: str, + mode_to_set: str, + ): + """ + Test the `set_append_mode` method with valid modes. These should all return True + and modify the value of 'appendfsync'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + :param mode_to_set: The mode to set + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + assert redis_conf.set_append_mode(mode_to_set) + assert redis_conf.get_config_value("appendfsync") == mode_to_set + assert f"Append mode is set to {mode_to_set}" in caplog.text, "Missing expected log message" + + def test_set_append_mode_invalid(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_append_mode` method with an invalid mode. This should return False + and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + invalid_mode = "invalid" + assert not redis_conf.set_append_mode(invalid_mode) + assert redis_conf.get_config_value("appendfsync") != invalid_mode + expected_log = "Not a valid append_mode (Only valid modes are always, everysec, no)" + assert expected_log in caplog.text, "Missing expected log message" + + def test_set_append_mode_none(self, server_redis_conf_file: str): + """ + Test the `set_append_mode` method with None as the input. + This should return False. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + assert not redis_conf.set_append_mode(None) + + def test_set_append_mode_appendfsync_unset(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_append_mode` method with the 'appendfsync' setting not existing. This should + return False and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + del redis_conf.entries["appendfsync"] + mode = "no" + assert not redis_conf.set_append_mode(mode) + assert redis_conf.get_config_value("appendfsync") != mode + assert "Unable to set append_mode in redis config" in caplog.text, "Missing expected log message" + + def test_set_append_file_valid(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_append_file` method with a valid file. This should return True + and modify the value of 'appendfilename'. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + caplog.set_level(logging.INFO) + redis_conf = RedisConfig(server_redis_conf_file) + valid_file = "valid" + assert redis_conf.set_append_file(valid_file) + assert redis_conf.get_config_value("appendfilename") == f'"{valid_file}"' + assert f"Append file is set to {valid_file}" in caplog.text, "Missing expected log message" + + def test_set_append_file_none(self, server_redis_conf_file: str): + """ + Test the `set_append_file` method with None as the input. + This should return False. + + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + assert not redis_conf.set_append_file(None) + + def test_set_append_file_appendfilename_unset(self, caplog: "Fixture", server_redis_conf_file: str): # noqa: F821 + """ + Test the `set_append_file` method with the 'appendfilename' setting not existing. This should + return False and log an error message. + + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + redis_conf = RedisConfig(server_redis_conf_file) + del redis_conf.entries["appendfilename"] + filename = "valid_filename" + assert not redis_conf.set_append_file(filename) + assert redis_conf.get_config_value("appendfilename") != filename + assert "Unable to set append filename." in caplog.text, "Missing expected log message" diff --git a/tests/unit/server/test_server_commands.py b/tests/unit/server/test_server_commands.py new file mode 100644 index 000000000..ec52df2a0 --- /dev/null +++ b/tests/unit/server/test_server_commands.py @@ -0,0 +1,647 @@ +""" +Tests for the `server_commands.py` module. +""" + +import logging +import os +import subprocess +from argparse import Namespace +from typing import Dict, List + +import pytest + +from merlin.server.server_commands import ( + check_for_not_running_server, + config_server, + init_server, + restart_server, + server_started, + start_container, + start_server, + status_server, + stop_server, +) +from merlin.server.server_config import ServerStatus +from merlin.server.server_util import ServerConfig + + +def test_init_server_create_server_fail(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `init_server` function with `create_server_config` returning False. + This should log a failure message. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + caplog.set_level(logging.INFO) + create_server_mock = mocker.patch("merlin.server.server_commands.create_server_config", return_value=False) + init_server() + create_server_mock.assert_called_once() + assert "Merlin server initialization failed." in caplog.text + + +def test_init_server_create_server_success(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `init_server` function with `create_server_config` returning True. + This should log a sucess message. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + caplog.set_level(logging.INFO) + create_server_mock = mocker.patch("merlin.server.server_commands.create_server_config", return_value=True) + pull_server_mock = mocker.patch("merlin.server.server_commands.pull_server_image", return_value=True) + config_merlin_mock = mocker.patch("merlin.server.server_commands.config_merlin_server", return_value=True) + init_server() + create_server_mock.assert_called_once() + pull_server_mock.assert_called_once() + config_merlin_mock.assert_called_once() + assert "Merlin server initialization successful." in caplog.text + + +def test_config_server_no_server_config( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_config_server_args: Namespace, +): + """ + Test the `config_server` function with no server config. This should log an error + and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_config_server_args: An argparse Namespace with args needed by `config_server` + """ + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=None) + assert not config_server(server_config_server_args) + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +@pytest.mark.parametrize( + "server_status, status_name", + [ + (ServerStatus.RUNNING, "running"), + (ServerStatus.NOT_RUNNING, "not_running"), + ], +) +def test_config_server_add_user_remove_user_success( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_config_server_args: Namespace, + server_server_config: Dict[str, Dict[str, str]], + server_status: ServerStatus, + status_name: str, +): + """ + Test the `config_server` function by adding and removing a user. This will be ran with and without + the server status being set to RUNNING. For each scenario we should expect: + - RUNNING -> RedisUsers.write and RedisUsers.apply_to_redis are both called twice + - NOT_RUNNING -> RedisUsers.write is called twice and RedisUsers.apply_to_redis is not called at all + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_config_server_args: An argparse Namespace with args needed by `config_server` + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param server_status: The server status for this test (either RUNNING or NOT_RUNNING) + :param status_name: The name of the status in string form so we can have unique users for each test + """ + caplog.set_level(logging.INFO) + + # Set up the add_user and remove_user calls to test + user_to_add_and_remove = f"test_config_server_modification_user_{status_name}" + server_config_server_args.add_user = [user_to_add_and_remove, "test_config_server_modification_password"] + server_config_server_args.remove_user = user_to_add_and_remove + + # Create mocks of the necessary calls for this function + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.apply_config_changes") + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + write_mock = mocker.patch("merlin.server.server_util.RedisUsers.write") + apply_to_redis_mock = mocker.patch("merlin.server.server_util.RedisUsers.apply_to_redis") + + # Run the test + expected_apply_calls = 2 if server_status == ServerStatus.RUNNING else 0 + assert config_server(server_config_server_args) is None + assert write_mock.call_count == 2 + assert apply_to_redis_mock.call_count == expected_apply_calls + assert f"Added user {user_to_add_and_remove} to merlin server" in caplog.text + assert f"Removed user {user_to_add_and_remove} to merlin server" in caplog.text + + +def test_config_server_add_user_remove_user_failure( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_config_server_args: Namespace, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `config_server` function by attempting to add a user that already exists (we do this through mock) + and removing a user that doesn't exist. This should run to completion but never call RedisUsers.write. It + should also log two error messages. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_config_server_args: An argparse Namespace with args needed by `config_server` + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + # Set up the add_user and remove_user calls to test (these users should never actually be added/removed) + user_to_add_and_remove = "test_config_user_not_ever_added" + server_config_server_args.add_user = [user_to_add_and_remove, "test_config_server_modification_password"] + server_config_server_args.remove_user = user_to_add_and_remove + + # Create mocks of the necessary calls for this function + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.apply_config_changes") + mocker.patch("merlin.server.server_util.RedisUsers.add_user", return_value=False) + write_mock = mocker.patch("merlin.server.server_util.RedisUsers.write") + + # Run the test + assert config_server(server_config_server_args) is None + assert write_mock.call_count == 0 + assert f"User '{user_to_add_and_remove}' already exisits within current users" in caplog.text + assert f"User '{user_to_add_and_remove}' doesn't exist within current users." in caplog.text + + +@pytest.mark.parametrize( + "server_status, expected_log_msgs", + [ + ( + ServerStatus.NOT_INITIALIZED, + ["Merlin server has not been initialized.", "Please initalize server by running 'merlin server init'"], + ), + ( + ServerStatus.MISSING_CONTAINER, + ["Unable to find server image.", "Ensure there is a .sif file in merlin server directory."], + ), + (ServerStatus.NOT_RUNNING, ["Merlin server is not running."]), + (ServerStatus.RUNNING, ["Merlin server is running."]), + ], +) +def test_status_server( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_status: ServerStatus, + expected_log_msgs: List[str], +): + """ + Test the `status_server` function to make sure it produces the correct logs for each status. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_status: The server status for this test + :param expected_log_msgs: The logs we're expecting from this test + """ + caplog.set_level(logging.INFO) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + status_server() + for expected_log_msg in expected_log_msgs: + assert expected_log_msg in caplog.text + + +@pytest.mark.parametrize( + "server_status, expected_result, expected_log_msg", + [ + ( + ServerStatus.NOT_INITIALIZED, + False, + "Merlin server has not been intitialized. Please run 'merlin server init' first.", + ), + ( + ServerStatus.MISSING_CONTAINER, + False, + "Merlin server has not been intitialized. Please run 'merlin server init' first.", + ), + (ServerStatus.NOT_RUNNING, True, None), + ( + ServerStatus.RUNNING, + False, + """Merlin server already running. + Stop current server with 'merlin server stop' before attempting to start a new server.""", + ), + ], +) +def test_check_for_not_running_server( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_status: ServerStatus, + expected_result: bool, + expected_log_msg: str, +): + """ + Test the `check_for_not_running_server` function with different server statuses. + There should be a logged message for each status and the results we should expect are as + follows: + - NOT_RUNNING status should return True + - any other status should return False + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_status: The server status for this test + :param expected_result: The expected result (T/F) for this test + :param expected_log_msg: The log we're expecting from this test + """ + caplog.set_level(logging.INFO) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + assert check_for_not_running_server() == expected_result + if expected_log_msg is not None: + assert expected_log_msg in caplog.text + + +def test_start_container_invalid_image_path( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `start_container` function with a nonexistent image path. + This should log an error and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + image_file = "nonexistent.image" + server_server_config["container"]["image"] = image_file + server_server_config["container"]["config"] = "start_container.config" + + # Create the config path so we ensure it exists + config_path = f"{server_testing_dir}/{server_server_config['container']['config']}" + with open(config_path, "w"): + pass + + assert start_container(ServerConfig(server_server_config)) is None + assert f"Unable to find image at {os.path.join(server_testing_dir, image_file)}" in caplog.text + + +def test_start_container_invalid_config_path( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `start_container` function with a nonexistent config path. + This should log an error and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + config_file = "nonexistent.config" + server_server_config["container"]["image"] = "start_container.image" + server_server_config["container"]["config"] = config_file + + # Create the config path so we ensure it exists + image_path = f"{server_testing_dir}/{server_server_config['container']['image']}" + with open(image_path, "w"): + pass + + assert start_container(ServerConfig(server_server_config)) is None + assert f"Unable to find config file at {os.path.join(server_testing_dir, config_file)}" in caplog.text + + +def test_start_container_valid_paths(mocker: "Fixture", server_server_config: Dict[str, Dict[str, str]]): # noqa: F821 + """ + Test the `start_container` function with valid image and config paths. + This should return a subprocess. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + expected_return = "fake subprocess" + mocker.patch("subprocess.Popen", return_value=expected_return) + mocker.patch("os.path.exists", return_value=True) + assert start_container(ServerConfig(server_server_config)) == expected_return + + +def test_server_started_no_redis_start(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `server_started` function with redis not starting. This should log errors and + return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + mock_process = mocker.Mock() + mock_process.stdout = mocker.Mock() + + expected_redis_out_msg = "Reached end of redis output without seeing 'Ready to accept connections'" + mocker.patch("merlin.server.server_commands.parse_redis_output", return_value=(False, expected_redis_out_msg)) + + assert not server_started(mock_process, "unecessary_config") + assert "Redis is unable to start" in caplog.text + assert 'Check to see if there is an unresponsive instance of redis with "ps -e"' in caplog.text + assert expected_redis_out_msg in caplog.text + + +def test_server_started_process_file_dump_fail( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `server_started` function with the dump to the process file failing. + This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mock_process = mocker.Mock() + mock_process.pid = 1234 + mock_process.stdout = mocker.Mock() + + image_pid = 5678 + mocker.patch("merlin.server.server_commands.parse_redis_output", return_value=(True, {"pid": image_pid})) + mocker.patch("merlin.server.server_commands.dump_process_file", return_value=False) + + assert not server_started(mock_process, ServerConfig(server_server_config)) + assert "Unable to create process file for container." in caplog.text + + +@pytest.mark.parametrize( + "server_status", + [ + ServerStatus.NOT_RUNNING, + ServerStatus.MISSING_CONTAINER, + ServerStatus.NOT_INITIALIZED, + ], +) +def test_server_started_server_not_running( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_server_config: Dict[str, Dict[str, str]], + server_status: ServerStatus, +): + """ + Test the `server_started` function with the server status returning a non-running status. + This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param server_status: The server status for this test + """ + mock_process = mocker.Mock() + mock_process.pid = 1234 + mock_process.stdout = mocker.Mock() + + image_pid = 5678 + mocker.patch("merlin.server.server_commands.parse_redis_output", return_value=(True, {"pid": image_pid})) + mocker.patch("merlin.server.server_commands.dump_process_file", return_value=True) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + + assert not server_started(mock_process, ServerConfig(server_server_config)) + assert "Unable to start merlin server." in caplog.text + + +def test_server_started_no_issues(mocker: "Fixture", server_server_config: Dict[str, Dict[str, str]]): # noqa: F821 + """ + Test the `server_started` function with no issues starting the server. + This should return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mock_process = mocker.Mock() + mock_process.pid = 1234 + mock_process.stdout = mocker.Mock() + + image_pid = 5678 + mocker.patch("merlin.server.server_commands.parse_redis_output", return_value=(True, {"pid": image_pid, "port": 6379})) + mocker.patch("merlin.server.server_commands.dump_process_file", return_value=True) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=ServerStatus.RUNNING) + + assert server_started(mock_process, ServerConfig(server_server_config)) + + +def test_start_server_no_running_server(mocker: "Fixture"): # noqa: F821 + """ + Test the `start_server` function with no running server. This should return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + """ + mocker.patch("merlin.server.server_commands.check_for_not_running_server", return_value=False) + assert not start_server() + + +def test_start_server_no_server_config(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `start_server` function with no running server. This should return False + and log an error. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + mocker.patch("merlin.server.server_commands.check_for_not_running_server", return_value=True) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=None) + assert not start_server() + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +def test_start_server_redis_container_startup_fail(mocker: "Fixture"): # noqa: F821 + """ + Test the `start_server` function with the redis container startup failing. This should return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + """ + mocker.patch("merlin.server.server_commands.check_for_not_running_server", return_value=True) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=True) + mocker.patch("merlin.server.server_commands.start_container", return_value=None) + assert not start_server() + + +def test_start_server_server_did_not_start(mocker: "Fixture"): # noqa: F821 + """ + Test the `start_server` function with the server startup failing. This should return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + """ + mocker.patch("merlin.server.server_commands.check_for_not_running_server", return_value=True) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=True) + mocker.patch("merlin.server.server_commands.start_container", return_value=True) + mocker.patch("merlin.server.server_commands.server_started", return_value=False) + assert not start_server() + + +def test_start_server_successful_start( + mocker: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], + server_redis_conf_file: str, +): + """ + Test the `start_server` function with a successful start. This should return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param server_redis_conf_file: The path to a dummy redis configuration file + """ + mocker.patch("merlin.server.server_commands.check_for_not_running_server", return_value=True) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.start_container", return_value=True) + mocker.patch("merlin.server.server_commands.server_started", return_value=True) + mocker.patch("merlin.server.server_commands.RedisUsers") + mocker.patch("merlin.server.server_commands.RedisConfig") + mocker.patch("merlin.server.server_commands.AppYaml") + mocker.patch("merlin.server.server_util.ContainerConfig.get_config_path", return_value=server_redis_conf_file) + mocker.patch("os.path.join", return_value=f"{server_testing_dir}/start_server_app.yaml") + + assert start_server() + + +@pytest.mark.parametrize( + "server_status", + [ + ServerStatus.NOT_RUNNING, + ServerStatus.MISSING_CONTAINER, + ServerStatus.NOT_INITIALIZED, + ], +) +def test_stop_server_server_not_running( + mocker: "Fixture", caplog: "Fixture", server_status: ServerStatus # noqa: F821 # noqa: F821 +): + """ + Test the `stop_server` function with a server that's not running. This should log two messages + and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_status: The server status for this test + """ + caplog.set_level(logging.INFO) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + assert not stop_server() + assert "There is no instance of merlin server running." in caplog.text + assert "Start a merlin server first with 'merlin server start'" in caplog.text + + +def test_stop_server_no_server_config(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `stop_server` function with no server config being pulled. This should log a message + and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + mocker.patch("merlin.server.server_commands.get_server_status", return_value=ServerStatus.RUNNING) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=None) + assert not stop_server() + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +def test_stop_server_empty_stdout( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_server_config: ServerConfig, +): + """ + Test the `stop_server` function with no server config being pulled. This should log a message + and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mocker.patch("merlin.server.server_commands.get_server_status", return_value=ServerStatus.RUNNING) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.pull_process_file", return_value={"parent_pid": 123}) + mock_run = mocker.patch("subprocess.run") + mock_run.return_value.stdout = b"" + assert not stop_server() + assert "Unable to get the PID for the current merlin server." in caplog.text + + +def test_stop_server_unable_to_stop_server( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_server_config: ServerConfig, +): + """ + Test the `stop_server` function with the server status still RUNNING after trying + to kill the server. This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mocker.patch("merlin.server.server_commands.get_server_status", return_value=ServerStatus.RUNNING) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.pull_process_file", return_value={"parent_pid": 123}) + mock_run = mocker.patch("subprocess.run") + mock_run.return_value.stdout = b"some output from status check" + assert not stop_server() + assert "Unable to kill process." in caplog.text + + +def test_stop_server_stop_command_is_not_kill( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_server_config: ServerConfig, +): + """ + Test the `stop_server` function with a stop command that's not 'kill'. + This should run through the command successfully and return True. The subprocess + should run the command we provide in this test. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mocker.patch( + "merlin.server.server_commands.get_server_status", side_effect=[ServerStatus.RUNNING, ServerStatus.NOT_RUNNING] + ) + mocker.patch("merlin.server.server_commands.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_commands.pull_process_file", return_value={"parent_pid": 123}) + custom_stop_command = "not a kill command" + mocker.patch("merlin.server.server_util.ContainerFormatConfig.get_stop_command", return_value=custom_stop_command) + mock_run = mocker.patch("subprocess.run") + mock_run.return_value.stdout = b"some output from status check" + assert stop_server() + mock_run.assert_called_with(custom_stop_command.split(), stdout=subprocess.PIPE) + + +@pytest.mark.parametrize( + "server_status", + [ + ServerStatus.NOT_RUNNING, + ServerStatus.MISSING_CONTAINER, + ServerStatus.NOT_INITIALIZED, + ], +) +def test_restart_server_server_not_running(mocker: "Fixture", caplog: "Fixture", server_status: ServerStatus): # noqa: F821 + """ + Test the `restart_server` function with a server that's not running. + This should log two messages and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_status: The server status for this test + """ + caplog.set_level(logging.INFO) + mocker.patch("merlin.server.server_commands.get_server_status", return_value=server_status) + assert not restart_server() + assert "Merlin server is not currently running." in caplog.text + assert "Please start a merlin server instance first with 'merlin server start'" in caplog.text + + +def test_restart_server_successful_restart(mocker: "Fixture"): # noqa: F821 + """ + Test the `restart_server` function with a successful restart. This should call + `stop_server` and `start_server`, and return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + """ + mocker.patch("merlin.server.server_commands.get_server_status", return_value=ServerStatus.RUNNING) + stop_server_mock = mocker.patch("merlin.server.server_commands.stop_server") + start_server_mock = mocker.patch("merlin.server.server_commands.start_server") + assert restart_server() + stop_server_mock.assert_called_once() + start_server_mock.assert_called_once() diff --git a/tests/unit/server/test_server_config.py b/tests/unit/server/test_server_config.py new file mode 100644 index 000000000..3b337d90e --- /dev/null +++ b/tests/unit/server/test_server_config.py @@ -0,0 +1,853 @@ +""" +Tests for the `server_config.py` module. +""" + +import io +import logging +import os +import string +from importlib import resources +from typing import Dict, Tuple, Union + +import pytest +import yaml + +from merlin.server.server_config import ( + MERLIN_CONFIG_DIR, + PASSWORD_LENGTH, + ServerStatus, + check_process_file_format, + config_merlin_server, + copy_container_command_files, + create_server_config, + dump_process_file, + generate_password, + get_server_status, + parse_redis_output, + pull_process_file, + pull_server_config, + pull_server_image, +) +from merlin.server.server_util import CONTAINER_TYPES, MERLIN_SERVER_SUBDIR, ServerConfig + + +def test_generate_password_no_pass_command(): + """ + Test the `generate_password` function with no password command. + This should generate a password of 256 (PASSWORD_LENGTH) random ASCII characters. + """ + generated_password = generate_password(PASSWORD_LENGTH) + assert len(generated_password) == PASSWORD_LENGTH + valid_ascii_chars = string.ascii_letters + string.digits + "!@#$%^&*()" + for ch in generated_password: + assert ch in valid_ascii_chars + + +def test_generate_password_with_pass_command(): + """ + Test the `generate_password` function with no password command. + This should generate a password of 256 (PASSWORD_LENGTH) random ASCII characters. + """ + test_pass = "test-password" + generated_password = generate_password(0, pass_command=f"echo {test_pass}") + assert generated_password == test_pass + + +@pytest.mark.parametrize( + "line, expected_return", + [ + (None, (False, "None passed as redis output")), + (b"", (False, "Reached end of redis output without seeing 'Ready to accept connections'")), + (b"Ready to accept connections", (True, {})), + (b"aborting", (False, "aborting")), + (b"Fatal error", (False, "Fatal error")), + ], +) +def test_parse_redis_output_with_basic_input(line: Union[None, bytes], expected_return: Tuple[bool, Union[str, Dict]]): + """ + Test the `parse_redis_output` function with basic input. + Here "basic input" means single line input or None as input. + + :param line: The value to pass in as input to `parse_redis_output` + :param expected_return: The expected return value based on what was passed in for `line` + """ + if line is None: + reader_input = None + else: + buffer = io.BytesIO(line) + reader_input = io.BufferedReader(buffer) + actual_return = parse_redis_output(reader_input) + assert expected_return == actual_return + + +@pytest.mark.parametrize( + "lines, expected_config", + [ + ( # Testing setting vars before initialized message + b"port=6379 blah blah server=127.0.0.1\nServer initialized\nReady to accept connections", + {"port": "6379", "server": "127.0.0.1"}, + ), + ( # Testing setting vars after initialized message + b"Server initialized\nport=6379 blah blah server=127.0.0.1\nReady to accept connections", + {}, + ), + ( # Testing setting vars before + after initialized message + b"blah blah max_connections=100 blah\n" + b"Server initialized\n" + b"port=6379 blah blah server=127.0.0.1\n" + b"Ready to accept connections", + {"max_connections": "100"}, + ), + ], +) +def test_parse_redis_output_with_vars(lines: bytes, expected_config: Tuple[bool, Union[str, Dict]]): + """ + Test the `parse_redis_output` function with input that has variables in lines. + This should set any variable given before the "Server initialized" message is provided. + + We'll test setting vars before the initialized message, after, and both before and after. + + :param lines: The lines to pass in as input to `parse_redis_output` + :param expected_config: The expected config dict based on what was passed in for `lines` + """ + buffer = io.BytesIO(lines) + reader_input = io.BufferedReader(buffer) + _, actual_vars = parse_redis_output(reader_input) + assert expected_config == actual_vars + + +def test_copy_container_command_files_with_existing_files( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Test the `copy_container_command_files` function with files that already exist. + This should skip trying to create the files, log 3 "file already exists" messages, + and return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + caplog.set_level(logging.INFO) + mocker.patch("os.path.exists", return_value=True) + assert copy_container_command_files(server_testing_dir) + file_names = [f"{container}.yaml" for container in CONTAINER_TYPES] + for file in file_names: + assert f"{file} already exists." in caplog.text + + +def test_copy_container_command_files_with_nonexisting_files( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Test the `copy_container_command_files` function with files that don't already exist. + This should create the files, log messages for each file, and return True + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + caplog.set_level(logging.INFO) + + # Mock the os.path.exists function so it returns False + mocker.patch("os.path.exists", return_value=False) + + # Mock the resources.path context manager + mock_path = mocker.patch("merlin.server.server_config.resources.path") + mock_path.return_value.__enter__.return_value = "mocked_file_path" + + # Mock the open builtin + mock_data = mocker.mock_open(read_data="mocked data") + mocker.patch("builtins.open", mock_data) + + assert copy_container_command_files(server_testing_dir) + file_names = [f"{container}.yaml" for container in CONTAINER_TYPES] + for file in file_names: + assert f"Copying file {file} to configuration directory." in caplog.text + + +def test_copy_container_command_files_with_oserror( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Test the `copy_container_command_files` function with an OSError being raised. + This should log an error message and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + # Mock the open function to raise an OSError + mocker.patch("builtins.open", side_effect=OSError("File not writeable")) + + assert not copy_container_command_files(server_testing_dir) + assert f"Destination location {server_testing_dir} is not writable." in caplog.text + + +def test_create_server_config_merlin_config_dir_nonexistent( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Tests the `create_server_config` function with MERLIN_CONFIG_DIR not existing. + This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + nonexistent_dir = f"{server_testing_dir}/merlin_config_dir" + mocker.patch("merlin.server.server_config.MERLIN_CONFIG_DIR", nonexistent_dir) + assert not create_server_config() + assert f"Unable to find main merlin configuration directory at {nonexistent_dir}" in caplog.text + + +def test_create_server_config_server_subdir_nonexistent_oserror( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Tests the `create_server_config` function with MERLIN_CONFIG_DIR/MERLIN_SERVER_SUBDIR + not existing and an OSError being raised. This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + + # Mock MERLIN_CONFIG_DIR and MERLIN_SERVER_SUBDIR + nonexistent_server_subdir = "test_create_server_config_server_subdir_nonexistent" + mocker.patch("merlin.server.server_config.MERLIN_CONFIG_DIR", server_testing_dir) + mocker.patch("merlin.server.server_config.MERLIN_SERVER_SUBDIR", nonexistent_server_subdir) + + # Mock os.mkdir so it raises an OSError + err_msg = "File not writeable" + mocker.patch("os.mkdir", side_effect=OSError(err_msg)) + assert not create_server_config() + assert err_msg in caplog.text + + +def test_create_server_config_no_server_config( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, +): + """ + Tests the `create_server_config` function with the call to `pull_server_config()` + returning None. This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + """ + + # Mock the necessary variables/functions to get us to the pull_server_config call + mocker.patch("merlin.server.server_config.MERLIN_CONFIG_DIR", server_testing_dir) + mocker.patch("merlin.server.server_config.copy_container_command_files", return_value=True) + mock_open_func = mocker.mock_open(read_data="key: value") + mocker.patch("builtins.open", mock_open_func) + + # Mock the pull_server_config call (what we're actually testing) and run the test + mocker.patch("merlin.server.server_config.pull_server_config", return_value=None) + assert not create_server_config() + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +def test_create_server_config_no_server_dir( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, str], +): + """ + Tests the `create_server_config` function with the call to + `server_config.container.get_config_dir()` returning a non-existent path. This should + log a message and create the directory, then return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + caplog.set_level(logging.INFO) + + # Mock the necessary variables/functions to get us to the get_config_dir call + mocker.patch("merlin.server.server_config.MERLIN_CONFIG_DIR", server_testing_dir) + mocker.patch("merlin.server.server_config.copy_container_command_files", return_value=True) + mock_open_func = mocker.mock_open(read_data="key: value") + mocker.patch("builtins.open", mock_open_func) + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + + # Mock the get_config_dir call to return a directory that doesn't exist yet + nonexistent_dir = f"{server_testing_dir}/merlin_server" + mocker.patch("merlin.server.server_util.ContainerConfig.get_config_dir", return_value=nonexistent_dir) + + assert create_server_config() + assert os.path.exists(nonexistent_dir) + assert "Creating merlin server directory." in caplog.text + + +def test_config_merlin_server_no_server_config(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `config_merlin_server` function with the call to `pull_server_config()` + returning None. This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + mocker.patch("merlin.server.server_config.pull_server_config", return_value=None) + assert not config_merlin_server() + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +def test_config_merlin_server_pass_user_exist( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, str], +): + """ + Tests the `config_merlin_server` function with a password file and user file already + existing. This should log 2 messages and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + caplog.set_level(logging.INFO) + + # Create the password file and user file + pass_file = f"{server_testing_dir}/existent_pass_file.txt" + user_file = f"{server_testing_dir}/existent_user_file.txt" + with open(pass_file, "w"), open(user_file, "w"): + pass + + # Mock necessary calls + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_util.ContainerConfig.get_pass_file_path", return_value=pass_file) + mocker.patch("merlin.server.server_util.ContainerConfig.get_user_file_path", return_value=user_file) + + assert config_merlin_server() is None + assert "Password file already exists. Skipping password generation step." in caplog.text + assert "User file already exists." in caplog.text + + +def test_config_merlin_server_pass_user_dont_exist( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, str], +): + """ + Tests the `config_merlin_server` function with a password file and user file that don't + already exist. This should log 2 messages, create the files, and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + caplog.set_level(logging.INFO) + + # Create vars for the nonexistent password file and user file + pass_file = f"{server_testing_dir}/nonexistent_pass_file.txt" + user_file = f"{server_testing_dir}/nonexistent_user_file.txt" + + # Mock necessary calls + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_util.ContainerConfig.get_pass_file_path", return_value=pass_file) + mocker.patch("merlin.server.server_util.ContainerConfig.get_user_file_path", return_value=user_file) + + assert config_merlin_server() is None + assert os.path.exists(pass_file) + assert os.path.exists(user_file) + assert "Creating password file for merlin server container." in caplog.text + assert f"User {os.environ.get('USER')} created in user file for merlin server container" in caplog.text + + +def setup_pull_server_config_mock( + mocker: "Fixture", # noqa: F821 + server_testing_dir: str, + server_app_yaml_contents: Dict[str, Union[str, int]], + server_server_config: Dict[str, Dict[str, str]], +): + """ + Setup the necessary mocker calls for the `pull_server_config` function. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_app_yaml_contents: A dict of app.yaml configurations + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mocker.patch("merlin.server.server_util.AppYaml.get_data", return_value=server_app_yaml_contents) + mocker.patch("merlin.server.server_config.MERLIN_CONFIG_DIR", server_testing_dir) + mock_data = mocker.mock_open(read_data=str(server_server_config)) + mocker.patch("builtins.open", mock_data) + + +@pytest.mark.parametrize( + "key_to_delete, expected_log_message", + [ + ("container", 'Unable to find "container" object in {default_app_yaml}'), + ("container.format", 'Unable to find "format" in {default_app_yaml}'), + ("process", "Process config not found in {default_app_yaml}"), + ], +) +def test_pull_server_config_missing_config_keys( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_app_yaml_contents: Dict[str, Union[str, int]], + server_server_config: Dict[str, Dict[str, str]], + key_to_delete: str, + expected_log_message: str, +): + """ + Test the `pull_server_config` function with missing container-related keys in the + app.yaml file contents. This should log an error message and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_app_yaml_contents: A dict of app.yaml configurations + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param key_to_delete: The key to delete from the app.yaml contents + :param expected_log_message: The expected log message when the key is missing + """ + # Handle nested key deletion + keys = key_to_delete.split(".") + temp_app_yaml = server_app_yaml_contents + for key in keys[:-1]: + temp_app_yaml = temp_app_yaml[key] + del temp_app_yaml[keys[-1]] + + setup_pull_server_config_mock(mocker, server_testing_dir, server_app_yaml_contents, server_server_config) + + assert pull_server_config() is None + default_app_yaml = os.path.join(MERLIN_CONFIG_DIR, "app.yaml") + assert expected_log_message.format(default_app_yaml=default_app_yaml) in caplog.text + + +@pytest.mark.parametrize("key_to_delete", ["command", "run_command", "stop_command", "pull_command"]) +def test_pull_server_config_missing_format_needed_keys( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_app_yaml_contents: Dict[str, Union[str, int]], + server_container_format_config_data: Dict[str, str], + server_server_config: Dict[str, Dict[str, str]], + key_to_delete: str, +): + """ + Test the `pull_server_config` function with necessary format keys missing in the + singularity.yaml file contents. This should log an error message and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_app_yaml_contents: A dict of app.yaml configurations + :param server_container_format_config_data: A pytest fixture of test data to pass to the ContainerFormatConfig class + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param key_to_delete: The key to delete from the singularity.yaml contents + """ + del server_container_format_config_data[key_to_delete] + setup_pull_server_config_mock(mocker, server_testing_dir, server_app_yaml_contents, server_server_config) + + assert pull_server_config() is None + format_file_basename = server_app_yaml_contents["container"]["format"] + ".yaml" + format_file = os.path.join(server_testing_dir, MERLIN_SERVER_SUBDIR) + format_file = os.path.join(format_file, format_file_basename) + assert f'Unable to find necessary "{key_to_delete}" value in format config file {format_file}' in caplog.text + + +@pytest.mark.parametrize("key_to_delete", ["status", "kill"]) +def test_pull_server_config_missing_process_needed_key( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_app_yaml_contents: Dict[str, Union[str, int]], + server_process_config_data: Dict[str, str], + server_server_config: Dict[str, Dict[str, str]], + key_to_delete: str, +): + """ + Test the `pull_server_config` function with necessary process keys missing. + This should log an error message and return None. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_app_yaml_contents: A dict of app.yaml configurations + :param server_process_config_data: A pytest fixture of test data to pass to the ProcessConfig class + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param key_to_delete: The key to delete from the process config entry + """ + del server_process_config_data[key_to_delete] + setup_pull_server_config_mock(mocker, server_testing_dir, server_app_yaml_contents, server_server_config) + + assert pull_server_config() is None + default_app_yaml = os.path.join(MERLIN_CONFIG_DIR, "app.yaml") + assert f'Process necessary "{key_to_delete}" command configuration not found in {default_app_yaml}' in caplog.text + + +def test_pull_server_config_no_issues( + mocker: "Fixture", # noqa: F821 + server_testing_dir: str, + server_app_yaml_contents: Dict[str, Union[str, int]], + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `pull_server_config` function without any problems. This should + return a ServerConfig object. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_app_yaml_contents: A dict of app.yaml configurations + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + setup_pull_server_config_mock(mocker, server_testing_dir, server_app_yaml_contents, server_server_config) + assert isinstance(pull_server_config(), ServerConfig) + + +def test_pull_server_image_no_server_config(mocker: "Fixture", caplog: "Fixture"): # noqa: F821 + """ + Test the `pull_server_image` function with no server config being found. + This should return False and log an error message. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + """ + mocker.patch("merlin.server.server_config.pull_server_config", return_value=None) + assert not pull_server_image() + assert 'Try to run "merlin server init" again to reinitialize values.' in caplog.text + + +def setup_pull_server_image_mock( + mocker: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], + config_dir: str, + config_file: str, + image_file: str, + create_config_file: bool = False, + create_image_file: bool = False, +): + """ + Set up the necessary mock calls for the `pull_server_image` function. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + image_url = "docker://redis" + image_path = f"{server_testing_dir}/{image_file}" + + os.makedirs(config_dir, exist_ok=True) + + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_util.ContainerConfig.get_config_dir", return_value=config_dir) + mocker.patch("merlin.server.server_util.ContainerConfig.get_config_name", return_value=config_file) + mocker.patch("merlin.server.server_util.ContainerConfig.get_image_url", return_value=image_url) + mocker.patch("merlin.server.server_util.ContainerConfig.get_image_path", return_value=image_path) + + if create_config_file: + with open(os.path.join(config_dir, config_file), "w"): + pass + + if create_image_file: + with open(image_path, "w"): + pass + + +def test_pull_server_image_no_image_path_no_config_path( + mocker: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `pull_server_image` function with no image path and no configuration + path. This should run a subprocess for the image path and create the configuration + file. It should also return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + # Set up mock calls to simulate the setup of this function + config_dir = f"{server_testing_dir}/config_dir" + config_file = "redis.conf" + image_file = "pull_server_image_no_image_path_no_config_path_image_nonexistent.sif" + setup_pull_server_image_mock(mocker, server_testing_dir, server_server_config, config_dir, config_file, image_file) + mocked_subprocess = mocker.patch("subprocess.run") + + # Mock the open function + read_data = "Mocked file content" + mocked_open = mocker.mock_open(read_data=read_data) + mocked_open.write = mocker.Mock() + mocker.patch("builtins.open", mocked_open) + + # Call the function + assert pull_server_image() + + # Assert that the subprocess call to pull the image was called + mocked_subprocess.assert_called_once() + + # Assert that open was called with the correct arguments + mocked_open.assert_any_call(os.path.join(config_dir, config_file), "w") + with resources.path("merlin.server", config_file) as file: + mocked_open.assert_any_call(file, "r") + assert mocked_open.call_count == 2 + + # Assert that the write method was called with the expected content + mocked_open().write.assert_called_once_with(read_data) + + +def test_pull_server_image_both_paths_exist( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `pull_server_image` function with both an image path and a configuration + path that both exist. This should log two messages and return True. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + caplog.set_level(logging.INFO) + + # Set up mock calls to simulate the setup of this function + config_dir = f"{server_testing_dir}/config_dir" + config_file = "pull_server_image_both_paths_exist_config.yaml" + image_file = "pull_server_image_both_paths_exist_image.sif" + setup_pull_server_image_mock( + mocker, + server_testing_dir, + server_server_config, + config_dir, + config_file, + image_file, + create_config_file=True, + create_image_file=True, + ) + + assert pull_server_image() + assert f"{image_file} already exists." in caplog.text + assert "Redis configuration file already exist." in caplog.text + + +def test_pull_server_image_os_error( + mocker: "Fixture", # noqa: F821 + caplog: "Fixture", # noqa: F821 + server_testing_dir: str, + server_server_config: Dict[str, Dict[str, str]], +): + """ + Test the `pull_server_image` function with an image path but no configuration + path. We'll force this to raise an OSError when writing to the configuration file + to ensure it's handled properly. This should log an error and return False. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param caplog: A built-in fixture from the pytest library to capture logs + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + # Set up mock calls to simulate the setup of this function + config_dir = f"{server_testing_dir}/config_dir" + config_file = "pull_server_image_os_error_config.yaml" + image_file = "pull_server_image_os_error_config_nonexistent.sif" + setup_pull_server_image_mock( + mocker, + server_testing_dir, + server_server_config, + config_dir, + config_file, + image_file, + create_image_file=True, + ) + + # Mock the open function + mocker.patch("builtins.open", side_effect=OSError) + + # Run the test + assert not pull_server_image() + assert f"Destination location {config_dir} is not writable." in caplog.text + + +@pytest.mark.parametrize( + "server_config_exists, config_exists, image_exists, pfile_exists, expected_status", + [ + (False, True, True, True, ServerStatus.NOT_INITIALIZED), # No server config + (True, False, True, True, ServerStatus.NOT_INITIALIZED), # Config dir does not exist + (True, True, False, True, ServerStatus.MISSING_CONTAINER), # Image path does not exist + (True, True, True, False, ServerStatus.NOT_RUNNING), # Pfile path does not exist + ], +) +def test_get_server_status_initial_checks( + mocker: "Fixture", # noqa: F821 + server_server_config: Dict[str, Dict[str, str]], + server_config_exists: bool, + config_exists: bool, + image_exists: bool, + pfile_exists: bool, + expected_status: ServerStatus, +): + """ + Test the `get_server_status` function for the initial conditional checks that it looks for. + These checks include: + - no server configuration -> should return NOT_INITIALIZED + - no config directory path -> should return NOT_INITIALIZED + - no image path -> should return MISSING_CONTAINER + - no password file -> should return NOT_RUNNING + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + :param server_config_exists: A boolean to denote whether the server config exists in this test or not + :param config_exists: A boolean to denote whether the config dir exists in this test or not + :param image_exists: A boolean to denote whether the image path exists in this test or not + :param pfile_exists: A boolean to denote whether the password file exists in this test or not + :param expected_status: The status we're expecting `get_server_status` to return for this test + """ + # Mock the necessary calls + if server_config_exists: + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("merlin.server.server_util.ContainerConfig.get_config_dir", return_value="config_dir") + mocker.patch("merlin.server.server_util.ContainerConfig.get_image_path", return_value="image_path") + mocker.patch("merlin.server.server_util.ContainerConfig.get_pfile_path", return_value="pfile_path") + + # Mock os.path.exists to return the desired values + mocker.patch( + "os.path.exists", + side_effect=lambda path: {"config_dir": config_exists, "image_path": image_exists, "pfile_path": pfile_exists}.get( + path, False + ), + ) + else: + mocker.patch("merlin.server.server_config.pull_server_config", return_value=None) + + # Call the function and assert the expected status + assert get_server_status() == expected_status + + +@pytest.mark.parametrize( + "stdout_val, expected_status", + [ + (b"", ServerStatus.NOT_RUNNING), # No stdout from subprocess + (b"Successfully started", ServerStatus.RUNNING), # Stdout from subprocess exists + ], +) +def test_get_server_status_subprocess_check( + mocker: "Fixture", # noqa: F821 + server_server_config: Dict[str, Dict[str, str]], + stdout_val: bytes, + expected_status: ServerStatus, +): + """ + Test the `get_server_status` function with empty stdout return from the subprocess run. + This should return a NOT_RUNNING status. + + :param mocker: A built-in fixture from the pytest-mock library to create a Mock object + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + mocker.patch("merlin.server.server_config.pull_server_config", return_value=ServerConfig(server_server_config)) + mocker.patch("os.path.exists", return_value=True) + mocker.patch("merlin.server.server_config.pull_process_file", return_value={"parent_pid": 123}) + mock_run = mocker.patch("subprocess.run") + mock_run.return_value.stdout = stdout_val + + assert get_server_status() == expected_status + + +@pytest.mark.parametrize( + "data_to_test, expected_result", + [ + ({"image_pid": 123, "port": 6379, "hostname": "dummy_server"}, False), # No parent_pid entry + ({"parent_pid": 123, "port": 6379, "hostname": "dummy_server"}, False), # No image_pid entry + ({"parent_pid": 123, "image_pid": 456, "hostname": "dummy_server"}, False), # No port entry + ({"parent_pid": 123, "image_pid": 123, "port": 6379}, False), # No hostname entry + ({"parent_pid": 123, "image_pid": 123, "port": 6379, "hostname": "dummy_server"}, True), # All required entries exist + ], +) +def test_check_process_file_format(data_to_test: Dict[str, Union[int, str]], expected_result: bool): + """ + Test the `check_process_file_format` function. The first 4 parametrized tests above should all + return False as they're all missing a required key. The final parametrized test above should return + True since it has every required key. + + :param data_to_test: The data dict that we'll pass in to the `check_process_file_format` function + :param expected_result: The return value we expect based on `data_to_test` + """ + assert check_process_file_format(data_to_test) == expected_result + + +def test_pull_process_file_valid_file(server_testing_dir: str, server_process_file_contents: Dict[str, Union[int, str]]): + """ + Test the `pull_process_file` function with a valid process file. This test will create a test + process file with valid contents that `pull_process_file` will read in and return. + + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_process_file_contents: A fixture representing process file contents + """ + # Create the valid process file in our temp testing directory + process_filepath = f"{server_testing_dir}/valid_process_file.yaml" + with open(process_filepath, "w") as process_file: + yaml.dump(server_process_file_contents, process_file) + + # Run the test + assert pull_process_file(process_filepath) == server_process_file_contents + + +def test_pull_process_file_invalid_file(server_testing_dir: str, server_process_file_contents: Dict[str, Union[int, str]]): + """ + Test the `pull_process_file` function with an invalid process file. This test will create a test + process file with invalid contents that `pull_process_file` will try to read in. Once it sees + that the file is invalid it will return None. + + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_process_file_contents: A fixture representing process file contents + """ + # Remove a key from the process file contents so that it's no longer valid + del server_process_file_contents["hostname"] + + # Create the invalid process file in our temp testing directory + process_filepath = f"{server_testing_dir}/invalid_process_file.yaml" + with open(process_filepath, "w") as process_file: + yaml.dump(server_process_file_contents, process_file) + + # Run the test + assert pull_process_file(process_filepath) is None + + +def test_dump_process_file_invalid_file(server_process_file_contents: Dict[str, Union[int, str]]): + """ + Test the `dump_process_file` function with invalid process file data. This should return False. + + :param server_process_file_contents: A fixture representing process file contents + """ + # Remove a key from the process file contents so that it's no longer valid and run the test + del server_process_file_contents["parent_pid"] + assert not dump_process_file(server_process_file_contents, "some_filepath.yaml") + + +def test_dump_process_file_valid_file(server_testing_dir: str, server_process_file_contents: Dict[str, Union[int, str]]): + """ + Test the `dump_process_file` function with invalid process file data. This should return False. + + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_process_file_contents: A fixture representing process file contents + """ + process_filepath = f"{server_testing_dir}/dumped_process_file.yaml" + assert dump_process_file(server_process_file_contents, process_filepath) + assert os.path.exists(process_filepath) diff --git a/tests/unit/server/test_server_util.py b/tests/unit/server/test_server_util.py new file mode 100644 index 000000000..909cb7cdf --- /dev/null +++ b/tests/unit/server/test_server_util.py @@ -0,0 +1,565 @@ +""" +Tests for the `server_util.py` module. +""" + +import filecmp +import hashlib +import os +from typing import Dict, Union + +import pytest + +from merlin.server.server_util import ( + AppYaml, + ContainerConfig, + ContainerFormatConfig, + ProcessConfig, + RedisConfig, + RedisUsers, + ServerConfig, + valid_ipv4, + valid_port, +) + + +@pytest.mark.parametrize( + "valid_ip", + [ + "0.0.0.0", + "127.0.0.1", + "14.105.200.58", + "255.255.255.255", + ], +) +def test_valid_ipv4_valid_ip(valid_ip: str): + """ + Test the `valid_ipv4` function with valid IPs. + This should return True. + + :param valid_ip: A valid port to test. These are pulled from the parametrized + list defined above this test. + """ + assert valid_ipv4(valid_ip) + + +@pytest.mark.parametrize( + "invalid_ip", + [ + "256.0.0.1", + "-1.0.0.1", + None, + "127.0.01", + ], +) +def test_valid_ipv4_invalid_ip(invalid_ip: Union[str, None]): + """ + Test the `valid_ipv4` function with invalid IPs. + An IP is valid if every integer separated by the '.' delimiter are between 0 and 255. + This should return False for both IPs tested here. + + :param invalid_ip: An invalid port to test. These are pulled from the parametrized + list defined above this test. + """ + assert not valid_ipv4(invalid_ip) + + +@pytest.mark.parametrize( + "valid_input", + [ + 1, + 433, + 65535, + ], +) +def test_valid_port_valid_input(valid_input: int): + """ + Test the `valid_port` function with valid port numbers. + Valid ports are ports between 1 and 65535. + This should return True. + + :param valid_input: A valid input value to test. These are pulled from the parametrized + list defined above this test. + """ + assert valid_port(valid_input) + + +@pytest.mark.parametrize( + "invalid_input", + [ + -1, + 0, + 65536, + ], +) +def test_valid_port_invalid_input(invalid_input: int): + """ + Test the `valid_port` function with invalid inputs. + Valid ports are ports between 1 and 65535. + This should return False for each invalid input tested. + + :param invalid_input: An invalid input value to test. These are pulled from the parametrized + list defined above this test. + """ + assert not valid_port(invalid_input) + + +class TestContainerConfig: + """Tests for the ContainerConfig class.""" + + def test_init_with_complete_data(self, server_container_config_data: Dict[str, str]): + """ + Tests that __init__ populates attributes correctly with complete data. + + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + """ + config = ContainerConfig(server_container_config_data) + assert config.format == server_container_config_data["format"] + assert config.image_type == server_container_config_data["image_type"] + assert config.image == server_container_config_data["image"] + assert config.url == server_container_config_data["url"] + assert config.config == server_container_config_data["config"] + assert config.config_dir == server_container_config_data["config_dir"] + assert config.pfile == server_container_config_data["pfile"] + assert config.pass_file == server_container_config_data["pass_file"] + assert config.user_file == server_container_config_data["user_file"] + + def test_init_with_missing_data(self): + """ + Tests that __init__ uses defaults for missing data. + """ + incomplete_data = {"format": "docker"} + config = ContainerConfig(incomplete_data) + assert config.format == incomplete_data["format"] + assert config.image_type == ContainerConfig.IMAGE_TYPE + assert config.image == ContainerConfig.IMAGE_NAME + assert config.url == ContainerConfig.REDIS_URL + assert config.config == ContainerConfig.CONFIG_FILE + assert config.config_dir == ContainerConfig.CONFIG_DIR + assert config.pfile == ContainerConfig.PROCESS_FILE + assert config.pass_file == ContainerConfig.PASSWORD_FILE + assert config.user_file == ContainerConfig.USERS_FILE + + @pytest.mark.parametrize( + "attr_name", + [ + "image", + "config", + "pfile", + "pass_file", + "user_file", + ], + ) + def test_get_path_methods(self, server_container_config_data: Dict[str, str], attr_name: str): + """ + Tests that get_*_path methods construct the correct path. + + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + :param attr_name: Name of the attribute to be tested. These are pulled from the parametrized list defined above this test. + """ + config = ContainerConfig(server_container_config_data) + get_path_method = getattr(config, f"get_{attr_name}_path") # Dynamically get the method based on attr_name + expected_path = os.path.join(server_container_config_data["config_dir"], server_container_config_data[attr_name]) + assert get_path_method() == expected_path + + @pytest.mark.parametrize( + "getter_name, expected_attr", + [ + ("get_format", "format"), + ("get_image_type", "image_type"), + ("get_image_name", "image"), + ("get_image_url", "url"), + ("get_config_name", "config"), + ("get_config_dir", "config_dir"), + ("get_pfile_name", "pfile"), + ("get_pass_file_name", "pass_file"), + ("get_user_file_name", "user_file"), + ], + ) + def test_getter_methods(self, server_container_config_data: Dict[str, str], getter_name: str, expected_attr: str): + """ + Tests that all getter methods return the correct attribute values. + + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + :param getter_name: Name of the getter method to test. This is pulled from the parametrized list defined above this test. + :param expected_attr: Name of the corresponding attribute. This is pulled from the parametrized list defined above this test. + """ + config = ContainerConfig(server_container_config_data) + getter = getattr(config, getter_name) + assert getter() == server_container_config_data[expected_attr] + + def test_get_container_password(self, server_testing_dir: str, server_container_config_data: Dict[str, str]): + """ + Test that the `get_container_password` method is reading the password file properly. + + :param server_testing_dir: The path to the the temp output directory for server tests + :param server_container_config_data: A pytest fixture of test data to pass to the ContainerConfig class + """ + # Write a fake password to the password file + test_password = "super-secret-password" + temp_pass_file = f"{server_testing_dir}/temp.pass" + with open(temp_pass_file, "w") as pass_file: + pass_file.write(test_password) + + # Use temp pass file + orig_pass_file = server_container_config_data["pass_file"] + server_container_config_data["pass_file"] = temp_pass_file + + try: + # Run the test + config = ContainerConfig(server_container_config_data) + assert config.get_container_password() == test_password + except Exception as exc: + # If there was a problem, reset to the original password file + server_container_config_data["pass_file"] = orig_pass_file + raise exc + + +class TestContainerFormatConfig: + """Tests for the ContainerFormatConfig class.""" + + def test_init_with_complete_data(self, server_container_format_config_data: Dict[str, str]): + """ + Tests that __init__ populates attributes correctly with complete data. + + :param server_container_format_config_data: A pytest fixture of test data to pass to the ContainerFormatConfig class + """ + config = ContainerFormatConfig(server_container_format_config_data) + assert config.command == server_container_format_config_data["command"] + assert config.run_command == server_container_format_config_data["run_command"] + assert config.stop_command == server_container_format_config_data["stop_command"] + assert config.pull_command == server_container_format_config_data["pull_command"] + + def test_init_with_missing_data(self): + """ + Tests that __init__ uses defaults for missing data. + """ + incomplete_data = {"command": "docker"} + config = ContainerFormatConfig(incomplete_data) + assert config.command == incomplete_data["command"] + assert config.run_command == config.RUN_COMMAND + assert config.stop_command == config.STOP_COMMAND + assert config.pull_command == config.PULL_COMMAND + + @pytest.mark.parametrize( + "getter_name, expected_attr", + [ + ("get_command", "command"), + ("get_run_command", "run_command"), + ("get_stop_command", "stop_command"), + ("get_pull_command", "pull_command"), + ], + ) + def test_getter_methods(self, server_container_format_config_data: Dict[str, str], getter_name: str, expected_attr: str): + """ + Tests that all getter methods return the correct attribute values. + + :param server_container_format_config_data: A pytest fixture of test data to pass to the ContainerFormatConfig class + :param getter_name: Name of the getter method to test. This is pulled from the parametrized list defined above this test. + :param expected_attr: Name of the corresponding attribute. This is pulled from the parametrized list defined above this test. + """ + config = ContainerFormatConfig(server_container_format_config_data) + getter = getattr(config, getter_name) + assert getter() == server_container_format_config_data[expected_attr] + + +class TestProcessConfig: + """Tests for the ProcessConfig class.""" + + def test_init_with_complete_data(self, server_process_config_data: Dict[str, str]): + """ + Tests that __init__ populates attributes correctly with complete data. + + :param server_process_config_data: A pytest fixture of test data to pass to the ProcessConfig class + """ + config = ProcessConfig(server_process_config_data) + assert config.status == server_process_config_data["status"] + assert config.kill == server_process_config_data["kill"] + + def test_init_with_missing_data(self): + """ + Tests that __init__ uses defaults for missing data. + """ + incomplete_data = {"status": "status {pid}"} + config = ProcessConfig(incomplete_data) + assert config.status == incomplete_data["status"] + assert config.kill == config.KILL_COMMAND + + @pytest.mark.parametrize( + "getter_name, expected_attr", + [ + ("get_status_command", "status"), + ("get_kill_command", "kill"), + ], + ) + def test_getter_methods(self, server_process_config_data: Dict[str, str], getter_name: str, expected_attr: str): + """ + Tests that all getter methods return the correct attribute values. + + :param server_process_config_data: A pytest fixture of test data to pass to the ProcessConfig class + :param getter_name: Name of the getter method to test. This is pulled from the parametrized list defined above this test. + :param expected_attr: Name of the corresponding attribute. This is pulled from the parametrized list defined above this test. + """ + config = ProcessConfig(server_process_config_data) + getter = getattr(config, getter_name) + assert getter() == server_process_config_data[expected_attr] + + +class TestServerConfig: + """Tests for the ServerConfig class.""" + + def test_init_with_complete_data(self, server_server_config: Dict[str, str]): + """ + Tests that __init__ populates attributes correctly with complete data. + + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + config = ServerConfig(server_server_config) + assert config.container == ContainerConfig(server_server_config["container"]) + assert config.process == ProcessConfig(server_server_config["process"]) + assert config.container_format == ContainerFormatConfig(server_server_config["singularity"]) + + def test_init_with_missing_data(self, server_process_config_data: Dict[str, str]): + """ + Tests that __init__ uses None for missing data. + + :param server_process_config_data: A pytest fixture of test data to pass to the ContainerConfig class + """ + incomplete_data = {"process": server_process_config_data} + config = ServerConfig(incomplete_data) + assert config.process == ProcessConfig(server_process_config_data) + assert config.container is None + assert config.container_format is None + + +class TestRedisUsers: + """ + Tests for the RedisUsers class. + + TODO add integration test(s) for `apply_to_redis` method of this class. + """ + + class TestUser: + """Tests for the RedisUsers.User class.""" + + def test_initializaiton(self): + """Test the initialization process of the User class.""" + user = RedisUsers.User() + assert user.status == "on" + assert user.hash_password == hashlib.sha256(b"password").hexdigest() + assert user.keys == "*" + assert user.channels == "*" + assert user.commands == "@all" + + def test_parse_dict(self): + """Test the `parse_dict` method of the User class.""" + test_dict = { + "status": "test_status", + "hash_password": "test_password", + "keys": "test_keys", + "channels": "test_channels", + "commands": "test_commands", + } + user = RedisUsers.User() + user.parse_dict(test_dict) + assert user.status == test_dict["status"] + assert user.hash_password == test_dict["hash_password"] + assert user.keys == test_dict["keys"] + assert user.channels == test_dict["channels"] + assert user.commands == test_dict["commands"] + + def test_get_user_dict(self): + """Test the `get_user_dict` method of the User class.""" + test_dict = { + "status": "test_status", + "hash_password": "test_password", + "keys": "test_keys", + "channels": "test_channels", + "commands": "test_commands", + "invalid_key": "invalid_val", + } + user = RedisUsers.User() + user.parse_dict(test_dict) # Set the test values + actual_dict = user.get_user_dict() + assert "invalid_key" not in actual_dict # Check that the invalid key isn't parsed + + # Check that the values are as expected + for key, val in actual_dict.items(): + if key == "status": + assert val == "on" + else: + assert val == test_dict[key] + + def test_set_password(self): + """Test the `set_password` method of the User class.""" + user = RedisUsers.User() + pass_to_set = "dummy_password" + user.set_password(pass_to_set) + assert user.hash_password == hashlib.sha256(bytes(pass_to_set, "utf-8")).hexdigest() + + def test_initialization(self, server_redis_users_file: str, server_users: dict): + """ + Test the initialization process of the RedisUsers class. + + :param server_redis_users_file: The path to a dummy redis users file + :param server_users: A dict of test user configurations + """ + redis_users = RedisUsers(server_redis_users_file) + assert redis_users.filename == server_redis_users_file + assert len(redis_users.users) == len(server_users) + + def test_write(self, server_redis_users_file: str, server_testing_dir: str): + """ + Test that the write functionality works by writing the contents of a dummy + users file to a blank users file. + + :param server_redis_users_file: The path to a dummy redis users file + :param server_testing_dir: The path to the the temp output directory for server tests + """ + copy_redis_users_file = f"{server_testing_dir}/redis_copy.users" + + # Create a RedisUsers object with the basic redis users file + redis_users = RedisUsers(server_redis_users_file) + + # Change the filepath of the redis users file to be the copy that we'll write to + redis_users.filename = copy_redis_users_file + + # Run the test + redis_users.write() + + # Check that the contents of the copied file match the contents of the basic file + assert filecmp.cmp(server_redis_users_file, copy_redis_users_file) + + def test_add_user_nonexistent(self, server_redis_users_file: str): + """ + Test the `add_user` method with a user that doesn't exists. + This should return True and add the user to the list of users. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + num_users_before = len(redis_users.users) + assert redis_users.add_user("new_user") + assert len(redis_users.users) == num_users_before + 1 + + def test_add_user_exists(self, server_redis_users_file: str): + """ + Test the `add_user` method with a user that already exists. + This should return False. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + assert not redis_users.add_user("test_user") + + def test_set_password_valid(self, server_redis_users_file: str): + """ + Test the `set_password` method with a user that exists. + This should return True and change the password for the user. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + pass_to_set = "new_password" + assert redis_users.set_password("test_user", pass_to_set) + expected_hash_pass = hashlib.sha256(bytes(pass_to_set, "utf-8")).hexdigest() + assert redis_users.users["test_user"].hash_password == expected_hash_pass + + def test_set_password_invalid(self, server_redis_users_file: str): + """ + Test the `set_password` method with a user that doesn't exist. + This should return False. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + assert not redis_users.set_password("nonexistent_user", "new_password") + + def test_remove_user_valid(self, server_redis_users_file: str): + """ + Test the `remove_user` method with a user that exists. + This should return True and remove the user from the list of users. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + num_users_before = len(redis_users.users) + assert redis_users.remove_user("test_user") + assert len(redis_users.users) == num_users_before - 1 + + def test_remove_user_invalid(self, server_redis_users_file: str): + """ + Test the `remove_user` method with a user that doesn't exist. + This should return False and not modify the user list. + + :param server_redis_users_file: The path to a dummy redis users file + """ + redis_users = RedisUsers(server_redis_users_file) + assert not redis_users.remove_user("nonexistent_user") + + +class TestAppYaml: + """Tests for the AppYaml class.""" + + def test_initialization(self, server_app_yaml: str, server_app_yaml_contents: dict): + """ + Test the initialization process of the AppYaml class. + + :param server_app_yaml: The path to an app.yaml file + :param server_app_yaml_contents: A dict of app.yaml configurations + """ + app_yaml = AppYaml(server_app_yaml) + assert app_yaml.get_data() == server_app_yaml_contents + + def test_apply_server_config(self, server_app_yaml: str, server_server_config: Dict[str, str]): + """ + Test the `apply_server_config` method. This should update the data attribute. + + :param server_app_yaml: The path to an app.yaml file + :param server_server_config: A pytest fixture of test data to pass to the ServerConfig class + """ + app_yaml = AppYaml(server_app_yaml) + server_config = ServerConfig(server_server_config) + redis_config = RedisConfig(server_config.container.get_config_path()) + app_yaml.apply_server_config(server_config) + + assert app_yaml.data[app_yaml.broker_name]["name"] == server_config.container.get_image_type() + assert app_yaml.data[app_yaml.broker_name]["username"] == "default" + assert app_yaml.data[app_yaml.broker_name]["password"] == server_config.container.get_pass_file_path() + assert app_yaml.data[app_yaml.broker_name]["server"] == redis_config.get_ip_address() + assert app_yaml.data[app_yaml.broker_name]["port"] == redis_config.get_port() + + assert app_yaml.data[app_yaml.results_name]["name"] == server_config.container.get_image_type() + assert app_yaml.data[app_yaml.results_name]["username"] == "default" + assert app_yaml.data[app_yaml.results_name]["password"] == server_config.container.get_pass_file_path() + assert app_yaml.data[app_yaml.results_name]["server"] == redis_config.get_ip_address() + assert app_yaml.data[app_yaml.results_name]["port"] == redis_config.get_port() + + def test_update_data(self, server_app_yaml: str): + """ + Test the `update_data` method. This should update the data attribute. + + :param server_app_yaml: The path to an app.yaml file + """ + app_yaml = AppYaml(server_app_yaml) + new_data = {app_yaml.broker_name: {"username": "new_user"}} + app_yaml.update_data(new_data) + + assert app_yaml.data[app_yaml.broker_name]["username"] == "new_user" + + def test_write(self, server_app_yaml: str, server_testing_dir: str): + """ + Test the `write` method. This should write data to a file. + + :param server_app_yaml: The path to an app.yaml file + :param server_testing_dir: The path to the the temp output directory for server tests + """ + copy_app_yaml = f"{server_testing_dir}/app_copy.yaml" + + # Create a AppYaml object with the basic app.yaml file + app_yaml = AppYaml(server_app_yaml) + + # Run the test + app_yaml.write(copy_app_yaml) + + # Check that the contents of the copied file match the contents of the basic file + assert filecmp.cmp(server_app_yaml, copy_app_yaml) diff --git a/tests/unit/study/__init__.py b/tests/unit/study/__init__.py index 57477ea1f..37cabcad1 100644 --- a/tests/unit/study/__init__.py +++ b/tests/unit/study/__init__.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/study/status_test_files/combine_status_files.py b/tests/unit/study/status_test_files/combine_status_files.py index f7021a97a..b52b35f1d 100644 --- a/tests/unit/study/status_test_files/combine_status_files.py +++ b/tests/unit/study/status_test_files/combine_status_files.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/study/status_test_files/shared_tests.py b/tests/unit/study/status_test_files/shared_tests.py index fb31b96a7..3e6b0fde8 100644 --- a/tests/unit/study/status_test_files/shared_tests.py +++ b/tests/unit/study/status_test_files/shared_tests.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/study/status_test_files/status_test_variables.py b/tests/unit/study/status_test_files/status_test_variables.py index ffb3cba31..6a7f7eb28 100644 --- a/tests/unit/study/status_test_files/status_test_variables.py +++ b/tests/unit/study/status_test_files/status_test_variables.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/study/test_detailed_status.py b/tests/unit/study/test_detailed_status.py index ae278c975..e6eea4748 100644 --- a/tests/unit/study/test_detailed_status.py +++ b/tests/unit/study/test_detailed_status.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/study/test_status.py b/tests/unit/study/test_status.py index 9d602848f..6786a1230 100644 --- a/tests/unit/study/test_status.py +++ b/tests/unit/study/test_status.py @@ -6,7 +6,7 @@ # # LLNL-CODE-797170 # All rights reserved. -# This file is part of Merlin, Version: 1.12.2b1. +# This file is part of Merlin, Version: 1.12.2. # # For details, see https://github.com/LLNL/merlin. # diff --git a/tests/unit/test_examples_generator.py b/tests/unit/test_examples_generator.py new file mode 100644 index 000000000..7d4d879fb --- /dev/null +++ b/tests/unit/test_examples_generator.py @@ -0,0 +1,475 @@ +""" +Tests for the `merlin/examples/generator.py` module. +""" + +import os +from typing import List + +import pytest +from tabulate import tabulate + +from merlin.examples.generator import ( + EXAMPLES_DIR, + gather_all_examples, + gather_example_dirs, + list_examples, + setup_example, + write_example, +) +from tests.utils import create_dir + + +EXAMPLES_GENERATOR_DIR = "{temp_output_dir}/examples_generator" + + +def test_gather_example_dirs(): + """Test the `gather_example_dirs` function.""" + example_workflows = [ + "feature_demo", + "flux", + "hello", + "hpc_demo", + "iterative_demo", + "lsf", + "null_spec", + "openfoam_wf", + "openfoam_wf_no_docker", + "openfoam_wf_singularity", + "optimization", + "remote_feature_demo", + "restart", + "restart_delay", + "simple_chain", + "slurm", + ] + expected = {} + for wf_dir in example_workflows: + expected[wf_dir] = wf_dir + actual = gather_example_dirs() + assert actual == expected + + +def test_gather_all_examples(): + """Test the `gather_all_examples` function.""" + expected = [ + f"{EXAMPLES_DIR}/feature_demo/feature_demo.yaml", + f"{EXAMPLES_DIR}/flux/flux_local.yaml", + f"{EXAMPLES_DIR}/flux/flux_par_restart.yaml", + f"{EXAMPLES_DIR}/flux/flux_par.yaml", + f"{EXAMPLES_DIR}/flux/paper.yaml", + f"{EXAMPLES_DIR}/hello/hello_samples.yaml", + f"{EXAMPLES_DIR}/hello/hello.yaml", + f"{EXAMPLES_DIR}/hello/my_hello.yaml", + f"{EXAMPLES_DIR}/hpc_demo/hpc_demo.yaml", + f"{EXAMPLES_DIR}/iterative_demo/iterative_demo.yaml", + f"{EXAMPLES_DIR}/lsf/lsf_par_srun.yaml", + f"{EXAMPLES_DIR}/lsf/lsf_par.yaml", + f"{EXAMPLES_DIR}/null_spec/null_chain.yaml", + f"{EXAMPLES_DIR}/null_spec/null_spec.yaml", + f"{EXAMPLES_DIR}/openfoam_wf/openfoam_wf_docker_template.yaml", + f"{EXAMPLES_DIR}/openfoam_wf/openfoam_wf.yaml", + f"{EXAMPLES_DIR}/openfoam_wf_no_docker/openfoam_wf_no_docker_template.yaml", + f"{EXAMPLES_DIR}/openfoam_wf_no_docker/openfoam_wf_no_docker.yaml", + f"{EXAMPLES_DIR}/openfoam_wf_singularity/openfoam_wf_singularity.yaml", + f"{EXAMPLES_DIR}/openfoam_wf_singularity/openfoam_wf_singularity_template.yaml", + f"{EXAMPLES_DIR}/optimization/optimization_basic.yaml", + f"{EXAMPLES_DIR}/remote_feature_demo/remote_feature_demo.yaml", + f"{EXAMPLES_DIR}/restart/restart.yaml", + f"{EXAMPLES_DIR}/restart_delay/restart_delay.yaml", + f"{EXAMPLES_DIR}/simple_chain/simple_chain.yaml", + f"{EXAMPLES_DIR}/slurm/slurm_par_restart.yaml", + f"{EXAMPLES_DIR}/slurm/slurm_par.yaml", + ] + actual = gather_all_examples() + assert sorted(actual) == sorted(expected) + + +def test_write_example_dir(examples_testing_dir: str): + """ + Test the `write_example` function with the src_path as a directory. + + :param examples_testing_dir: The path to the the temp output directory for examples tests + """ + dir_to_copy = f"{EXAMPLES_DIR}/feature_demo/" + dst_dir = f"{examples_testing_dir}/write_example_dir" + write_example(dir_to_copy, dst_dir) + assert sorted(os.listdir(dir_to_copy)) == sorted(os.listdir(dst_dir)) + + +def test_write_example_file(examples_testing_dir: str): + """ + Test the `write_example` function with the src_path as a file. + + :param examples_testing_dir: The path to the the temp output directory for examples tests + """ + file_to_copy = f"{EXAMPLES_DIR}/flux/flux_par.yaml" + dst_path = f"{examples_testing_dir}/flux_par.yaml" + write_example(file_to_copy, dst_path) + assert os.path.exists(dst_path) + + +def test_list_examples(): + """Test the `list_examples` function to see if it gives us all of the examples that we want.""" + expected_headers = ["name", "description"] + expected_rows = [ + ["feature_demo", "Run 10 hello worlds."], + ["flux_local", "Run a scan through Merlin/Maestro"], + ["flux_par", "A simple ensemble of parallel MPI jobs run by flux."], + ["flux_par_restart", "A simple ensemble of parallel MPI jobs run by flux."], + ["paper_flux", "Use flux to run single core MPI jobs and record timings."], + ["hello", "a very simple merlin workflow"], + ["hello_samples", "a very simple merlin workflow, with samples"], + ["hpc_demo", "Demo running a workflow on HPC machines"], + ["iterative_demo", "Demo of a workflow with self driven iteration/looping"], + ["lsf_par", "A simple ensemble of parallel MPI jobs run by lsf (jsrun)."], + ["lsf_par_srun", "A simple ensemble of parallel MPI jobs run by lsf using the srun wrapper (srun)."], + [ + "null_chain", + "Run N_SAMPLES steps of TIME seconds each at CONC concurrency.\n" + "May be used to measure overhead in merlin.\n" + "Iterates thru a chain of workflows.", + ], + [ + "null_spec", + "run N_SAMPLES null steps at CONC concurrency for TIME seconds each. May be used to measure overhead in merlin.", + ], + [ + "openfoam_wf", + "A parameter study that includes initializing, running,\n" + "post-processing, collecting, learning and visualizing OpenFOAM runs\n" + "using docker.", + ], + [ + "openfoam_wf_no_docker", + "A parameter study that includes initializing, running,\n" + "post-processing, collecting, learning and vizualizing OpenFOAM runs\n" + "without using docker.", + ], + [ + "openfoam_wf_singularity", + "A parameter study that includes initializing, running,\n" + "post-processing, collecting, learning and visualizing OpenFOAM runs\n" + "using singularity.", + ], + [ + "optimization_basic", + "Design Optimization Template\n" + "To use,\n" + "1. Specify the first three variables here (N_DIMS, TEST_FUNCTION, DEBUG)\n" + "2. Run the template_config file in current directory using `python template_config.py`\n" + "3. Merlin run as usual (merlin run optimization.yaml)\n" + "* MAX_ITER and the N_SAMPLES options use default values unless using DEBUG mode\n" + "* BOUNDS_X and UNCERTS_X are configured using the template_config.py scripts", + ], + ["remote_feature_demo", "Run 10 hello worlds."], + ["restart", "A simple ensemble of with restarts."], + ["restart_delay", "A simple ensemble of with restart delay times."], + ["simple_chain", "test to see that chains are not run in parallel"], + ["slurm_par", "A simple ensemble of parallel MPI jobs run by slurm (srun)."], + ["slurm_par_restart", "A simple ensemble of parallel MPI jobs run by slurm (srun)."], + ] + expected = "\n" + tabulate(expected_rows, expected_headers) + "\n" + actual = list_examples() + print(f"actual:\n{actual}") + print(f"expected:\n{expected}") + assert actual == expected + + +def test_setup_example_invalid_name(): + """ + Test the `setup_example` function with an invalid example name. + This should just return None. + """ + assert setup_example("invalid_example_name", None) is None + + +def test_setup_example_no_outdir(examples_testing_dir: str): + """ + Test the `setup_example` function with an invalid example name. + This should create a directory with the example name (in this case hello) + and copy all of the example contents to this folder. + We'll create a directory specifically for this test and move into it so that + the `setup_example` function creates the hello/ subdirectory in a directory with + the name of this test (setup_no_outdir). + + :param examples_testing_dir: The path to the the temp output directory for examples tests + """ + cwd = os.getcwd() + + # Create the temp path to store this setup and move into that directory + setup_example_dir = os.path.join(examples_testing_dir, "setup_no_outdir") + create_dir(setup_example_dir) + os.chdir(setup_example_dir) + + # This should still work and return to us the name of the example + try: + assert setup_example("hello", None) == "hello" + except AssertionError as exc: + os.chdir(cwd) + raise AssertionError from exc + + # All files from this example should be written to a directory with the example name + full_output_path = os.path.join(setup_example_dir, "hello") + expected_files = [ + os.path.join(full_output_path, "hello_samples.yaml"), + os.path.join(full_output_path, "hello.yaml"), + os.path.join(full_output_path, "my_hello.yaml"), + os.path.join(full_output_path, "requirements.txt"), + os.path.join(full_output_path, "make_samples.py"), + ] + try: + for file in expected_files: + assert os.path.exists(file) + except AssertionError as exc: + os.chdir(cwd) + raise AssertionError from exc + finally: + os.chdir(cwd) + + +def test_setup_example_outdir_exists(examples_testing_dir: str): + """ + Test the `setup_example` function with an output directory that already exists. + This should just return None. + + :param examples_testing_dir: The path to the the temp output directory for examples tests + """ + assert setup_example("hello", examples_testing_dir) is None + + +@pytest.mark.parametrize( + "example_name, example_files, expected_return", + [ + ( + "feature_demo", + [ + ".gitignore", + "feature_demo.yaml", + "requirements.txt", + "scripts/features.json", + "scripts/hello_world.py", + "scripts/pgen.py", + ], + "feature_demo", + ), + ( + "flux_local", + [ + "flux_local.yaml", + "flux_par_restart.yaml", + "flux_par.yaml", + "paper.yaml", + "requirements.txt", + "scripts/flux_info.py", + "scripts/hello_sleep.c", + "scripts/hello.c", + "scripts/make_samples.py", + "scripts/paper_workers.sbatch", + "scripts/test_workers.sbatch", + "scripts/workers.sbatch", + "scripts/workers.bsub", + ], + "flux", + ), + ( + "lsf_par", + [ + "lsf_par_srun.yaml", + "lsf_par.yaml", + "scripts/hello.c", + "scripts/make_samples.py", + ], + "lsf", + ), + ( + "slurm_par", + [ + "slurm_par.yaml", + "slurm_par_restart.yaml", + "requirements.txt", + "scripts/hello.c", + "scripts/make_samples.py", + "scripts/test_workers.sbatch", + "scripts/workers.sbatch", + ], + "slurm", + ), + ( + "hello", + [ + "hello_samples.yaml", + "hello.yaml", + "my_hello.yaml", + "requirements.txt", + "make_samples.py", + ], + "hello", + ), + ( + "hpc_demo", + [ + "hpc_demo.yaml", + "cumulative_sample_processor.py", + "faker_sample.py", + "sample_collector.py", + "sample_processor.py", + "requirements.txt", + ], + "hpc_demo", + ), + ( + "iterative_demo", + [ + "iterative_demo.yaml", + "cumulative_sample_processor.py", + "faker_sample.py", + "sample_collector.py", + "sample_processor.py", + "requirements.txt", + ], + "iterative_demo", + ), + ( + "null_spec", + [ + "null_spec.yaml", + "null_chain.yaml", + ".gitignore", + "Makefile", + "requirements.txt", + "scripts/aggregate_chain_output.sh", + "scripts/aggregate_output.sh", + "scripts/check_completion.sh", + "scripts/kill_all.sh", + "scripts/launch_chain_job.py", + "scripts/launch_jobs.py", + "scripts/make_samples.py", + "scripts/read_output_chain.py", + "scripts/read_output.py", + "scripts/search.sh", + "scripts/submit_chain.sbatch", + "scripts/submit.sbatch", + ], + "null_spec", + ), + ( + "openfoam_wf", + [ + "openfoam_wf.yaml", + "openfoam_wf_docker_template.yaml", + "README.md", + "requirements.txt", + "scripts/make_samples.py", + "scripts/blockMesh_template.txt", + "scripts/cavity_setup.sh", + "scripts/combine_outputs.py", + "scripts/learn.py", + "scripts/mesh_param_script.py", + "scripts/run_openfoam", + ], + "openfoam_wf", + ), + ( + "openfoam_wf_no_docker", + [ + "openfoam_wf_no_docker.yaml", + "openfoam_wf_no_docker_template.yaml", + "requirements.txt", + "scripts/make_samples.py", + "scripts/blockMesh_template.txt", + "scripts/cavity_setup.sh", + "scripts/combine_outputs.py", + "scripts/learn.py", + "scripts/mesh_param_script.py", + "scripts/run_openfoam", + ], + "openfoam_wf_no_docker", + ), + ( + "openfoam_wf_singularity", + [ + "openfoam_wf_singularity.yaml", + "openfoam_wf_singularity_template.yaml", + "requirements.txt", + "scripts/make_samples.py", + "scripts/blockMesh_template.txt", + "scripts/cavity_setup.sh", + "scripts/combine_outputs.py", + "scripts/learn.py", + "scripts/mesh_param_script.py", + "scripts/run_openfoam", + ], + "openfoam_wf_singularity", + ), + ( + "optimization_basic", + [ + "optimization_basic.yaml", + "requirements.txt", + "template_config.py", + "template_optimization.temp", + "scripts/collector.py", + "scripts/optimizer.py", + "scripts/test_functions.py", + "scripts/visualizer.py", + ], + "optimization", + ), + ( + "remote_feature_demo", + [ + ".gitignore", + "remote_feature_demo.yaml", + "requirements.txt", + "scripts/features.json", + "scripts/hello_world.py", + "scripts/pgen.py", + ], + "remote_feature_demo", + ), + ("restart", ["restart.yaml", "scripts/make_samples.py"], "restart"), + ("restart_delay", ["restart_delay.yaml", "scripts/make_samples.py"], "restart_delay"), + ], +) +def test_setup_example(examples_testing_dir: str, example_name: str, example_files: List[str], expected_return: str): + """ + Run tests for the `setup_example` function. + Each test will consist of: + 1. The name of the example to setup + 2. A list of files that we're expecting to be setup + 3. The expected return value + Each test is a tuple in the parametrize decorator above this test function. + + :param examples_testing_dir: The path to the the temp output directory for examples tests + :param example_name: The name of the example to setup + :param example_files: A list of filenames that should be copied by setup_example + :param expected_return: The expected return value from `setup_example` + """ + # Create the temp path to store this setup + setup_example_dir = os.path.join(examples_testing_dir, f"setup_{example_name}") + + # Ensure that the example name is returned + actual = setup_example(example_name, setup_example_dir) + assert actual == expected_return + + # Ensure all of the files that should've been copied were copied + expected_files = [os.path.join(setup_example_dir, expected_file) for expected_file in example_files] + for file in expected_files: + assert os.path.exists(file) + + +def test_setup_example_simple_chain(examples_testing_dir: str): + """ + Test the `setup_example` function for the simple_chain example. + This example just writes a single file so we can't run it in the `test_setup_example` test. + + :param examples_testing_dir: The path to the the temp output directory for examples tests + """ + + # Create the temp path to store this setup + output_file = os.path.join(examples_testing_dir, "simple_chain.yaml") + + # Ensure that the example name is returned + actual = setup_example("simple_chain", output_file) + assert actual == "simple_chain" + assert os.path.exists(output_file) diff --git a/tests/utils.py b/tests/utils.py new file mode 100644 index 000000000..0b408db54 --- /dev/null +++ b/tests/utils.py @@ -0,0 +1,44 @@ +""" +Utility functions for our test suite. +""" + +import os +from typing import Dict + +from tests.constants import SERVER_PASS + + +def create_pass_file(pass_filepath: str): + """ + Check if a password file already exists (it will if the redis server has been started) + and if it hasn't then create one and write the password to the file. + + :param pass_filepath: The path to the password file that we need to check for/create + """ + if not os.path.exists(pass_filepath): + with open(pass_filepath, "w") as pass_file: + pass_file.write(SERVER_PASS) + + +def create_cert_files(cert_filepath: str, cert_files: Dict[str, str]): + """ + Check if cert files already exist and if they don't then create them. + + :param cert_filepath: The path to the cert files + :param cert_files: A dict of certification files to create + """ + for cert_file in cert_files.values(): + full_cert_filepath = f"{cert_filepath}/{cert_file}" + if not os.path.exists(full_cert_filepath): + with open(full_cert_filepath, "w"): + pass + + +def create_dir(dirpath: str): + """ + Check if `dirpath` exists and if it doesn't then create it. + + :param dirpath: The directory to create + """ + if not os.path.exists(dirpath): + os.mkdir(dirpath)