├── .github ├── ISSUE_TEMPLATE │ ├── bug_report.md │ └── enhancement-proposal.md ├── dependabot.yml └── workflows │ ├── ci-docker.yml │ ├── ci-docs.yml │ ├── ci-format.yml │ ├── ci-pypi.yml │ ├── ci-test.yml │ ├── release-please.yml │ ├── sync-gitee.yml │ └── sync-github.yml ├── .gitignore ├── CHANGELOG.md ├── Dockerfile ├── LICENSE ├── MANIFEST.in ├── README.md ├── VERSION ├── docs ├── Makefile ├── developer.md ├── imgs │ └── dflow_architecture.png ├── make.bat ├── requirements.txt └── source │ ├── conf.py │ └── index.rst ├── examples ├── test_async.py ├── test_big_parameter.py ├── test_bohrium_datasets.py ├── test_conditional_outputs.py ├── test_dag.py ├── test_dag_slices.py ├── test_datasets.py ├── test_datasets_local.py ├── test_dispatcher.py ├── test_error_handling.py ├── test_func_op.py ├── test_launching.py ├── test_launching_2.py ├── test_lineage.py ├── test_loop_slices.py ├── test_makevasp.py ├── test_metadata.py ├── test_op_parameter.py ├── test_progress.py ├── test_python.py ├── test_ray.py ├── test_ray_after_ray.py ├── test_recurse.py ├── test_reuse.py ├── test_slices.py ├── test_slurm.py ├── test_steps.py ├── test_subpath_slices.py ├── test_success_ratio.py ├── test_utils.py └── test_wlm.py ├── manifests ├── configurator.yaml ├── crds │ ├── wlm_v1alpha1_slurmjob.yaml │ └── wlm_v1alpha1_wlmjob.yaml ├── operator-rbac.yaml ├── operator.yaml ├── quick-start-postgres-3.4.1-deepmodeling.yaml ├── quick-start-postgres-stable-cn.yaml └── quick-start-postgres.yaml ├── requirements.txt ├── scripts ├── docker-compose.yml ├── install-linux-cn.sh ├── install-linux.sh ├── install-mac-cn.sh ├── install-mac.sh ├── register_cluster.sh └── start-slurm.sh ├── setup.py ├── src └── dflow │ ├── __init__.py │ ├── __main__.py │ ├── argo_objects.py │ ├── client │ ├── __init__.py │ ├── v1alpha1_artifact.py │ ├── v1alpha1_dag_task.py │ ├── v1alpha1_lifecycle_hook.py │ ├── v1alpha1_parameter.py │ ├── v1alpha1_retry_strategy.py │ ├── v1alpha1_sequence.py │ ├── v1alpha1_template.py │ ├── v1alpha1_value_from.py │ └── v1alpha1_workflow_step.py │ ├── code_gen.py │ ├── common.py │ ├── config.py │ ├── context.py │ ├── context_syntax.py │ ├── dag.py │ ├── executor.py │ ├── io.py │ ├── main.py │ ├── op_template.py │ ├── plugins │ ├── __init__.py │ ├── bohrium.py │ ├── datasets.py │ ├── dispatcher.py │ ├── launching.py │ ├── lebesgue.py │ ├── metadata.py │ ├── oss.py │ └── ray.py │ ├── py.typed │ ├── python │ ├── __init__.py │ ├── op.py │ ├── opio.py │ ├── python_op_template.py │ ├── utils.py │ └── vendor │ │ └── typeguard │ │ ├── __init__.py │ │ ├── importhook.py │ │ ├── py.typed │ │ └── pytest_plugin.py │ ├── resource.py │ ├── slurm.py │ ├── step.py │ ├── steps.py │ ├── task.py │ ├── util_ops.py │ ├── utils.py │ └── workflow.py ├── tests ├── test_big_parameter.py ├── test_conditional_outputs.py ├── test_dag.py ├── test_group_size.py ├── test_makevasp.py ├── test_python.py ├── test_recurse.py ├── test_reuse.py ├── test_slices.py └── test_subpath_slices.py └── tutorials ├── cp2k_opt └── input.inp ├── dflow-conditional.ipynb ├── dflow-function.ipynb ├── dflow-python.ipynb ├── dflow-recurse.ipynb ├── dflow-reuse.ipynb ├── dflow-slices.ipynb ├── dflow-slurm.ipynb ├── imgs ├── access_one_node.png ├── argoui_main_page.png ├── connection_warning.png ├── minikube_image_bug.png ├── minikube_start_fail_bug.png └── workflow_overview.png ├── install_manual_linuxos.md ├── install_manual_macos.md ├── install_manual_windowsos.md └── readme.md /.github/ISSUE_TEMPLATE/bug_report.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Bug report 3 | about: Create a report to help us improve 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | ## Checklist 11 | 12 | 13 | 14 | * [ ] Checked the syntax is legal within dflow. 15 | * [ ] Tested using the latest version. 16 | 17 | ## Summary 18 | 19 | What happened/what you expected to happen? 20 | 21 | What version are you running? 22 | 23 | ## Diagnostics 24 | 25 | Paste the smallest python script that reproduces the bug 26 | 27 | ```python 28 | 29 | ``` 30 | -------------------------------------------------------------------------------- /.github/ISSUE_TEMPLATE/enhancement-proposal.md: -------------------------------------------------------------------------------- 1 | --- 2 | name: Enhancement proposal 3 | about: Propose an enhancement for this project. 4 | title: '' 5 | labels: '' 6 | assignees: '' 7 | 8 | --- 9 | 10 | # Summary 11 | 12 | What change needs making? 13 | 14 | # Use Cases 15 | 16 | When would you use this? 17 | -------------------------------------------------------------------------------- /.github/dependabot.yml: -------------------------------------------------------------------------------- 1 | version: 2 2 | updates: 3 | - package-ecosystem: "github-actions" 4 | directory: "/" 5 | schedule: 6 | interval: "weekly" 7 | -------------------------------------------------------------------------------- /.github/workflows/ci-docker.yml: -------------------------------------------------------------------------------- 1 | name: Build docker image and push to Docker Hub 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | paths: 8 | - 'VERSION' 9 | 10 | jobs: 11 | build-n-push: 12 | if: github.repository_owner == 'deepmodeling' 13 | name: Build docker image and push to Docker Hub 14 | runs-on: ubuntu-latest 15 | steps: 16 | - name: Check out the repo 17 | uses: actions/checkout@v3 18 | 19 | - name: Log in to Docker Hub 20 | uses: docker/login-action@f054a8b539a109f9f41c372932f1ae047eff08c9 21 | with: 22 | username: ${{ secrets.DOCKER_USERNAME }} 23 | password: ${{ secrets.DOCKER_PASSWORD }} 24 | 25 | - name: Set up QEMU 26 | uses: docker/setup-qemu-action@v2 27 | 28 | - name: Set up Docker Buildx 29 | id: buildx 30 | uses: docker/setup-buildx-action@v2 31 | 32 | - name: read version 33 | id: read 34 | uses: juliangruber/read-file-action@v1 35 | with: 36 | path: ./VERSION 37 | 38 | - name: Build and push 39 | run: | 40 | docker buildx build . --push --platform linux/arm64,linux/amd64 -t dptechnology/dflow:v${{ steps.read.outputs.content }} 41 | docker buildx build . --push --platform linux/arm64,linux/amd64 -t dptechnology/dflow:latest 42 | 43 | - name: send email 44 | uses: dawidd6/action-send-mail@v3 45 | with: 46 | server_address: smtp.feishu.cn 47 | server_port: 465 48 | username: ${{ secrets.MAILUSERNAME }} 49 | password: ${{ secrets.MAILPASSWORD }} 50 | subject: Docker Auto Build For Dflow 51 | body: Docker image has been pushed to dptechnology/dflow:v${{ steps.read.outputs.content }} 52 | to: ${{ secrets.MAIL_RECEIVER_LIST }} 53 | from: Github Actions 54 | content_type: text 55 | -------------------------------------------------------------------------------- /.github/workflows/ci-docs.yml: -------------------------------------------------------------------------------- 1 | name: Build docs 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | 8 | jobs: 9 | build-n-push: 10 | if: github.repository_owner == 'deepmodeling' 11 | name: Build docs 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: actions/checkout@v1 15 | 16 | - name: Sphinx apidoc 17 | run: | 18 | python -m pip install sphinx jsonpickle typeguard 19 | cd docs 20 | sphinx-apidoc -o ./source ../src && make html 21 | cd .. 22 | 23 | - name: Commit documentation changes 24 | run: | 25 | git clone https://github.com/deepmodeling/dflow.git --branch docs --single-branch dflow-docs 26 | cp -r docs/build/html/* dflow-docs 27 | cd dflow-docs 28 | git config --local user.email "action@github.com" 29 | git config --local user.name "GitHub Action" 30 | git add . 31 | git diff-index --quiet HEAD || git commit -m "Update documentation" 32 | 33 | - name: Push changes 34 | uses: ad-m/github-push-action@master 35 | with: 36 | branch: docs 37 | directory: dflow-docs 38 | github_token: ${{ secrets.GITHUB_TOKEN }} 39 | 40 | - name: send email 41 | uses: dawidd6/action-send-mail@v3 42 | with: 43 | server_address: smtp.feishu.cn 44 | server_port: 465 45 | username: ${{ secrets.MAILUSERNAME }} 46 | password: ${{ secrets.MAILPASSWORD }} 47 | subject: Documentation Auto Build For Dflow 48 | body: Docs has been published to https://deepmodeling.com/dflow/ 49 | to: ${{ secrets.MAIL_RECEIVER_LIST }} 50 | from: Github Actions 51 | content_type: text 52 | -------------------------------------------------------------------------------- /.github/workflows/ci-format.yml: -------------------------------------------------------------------------------- 1 | name: auto-formating the commit 2 | 3 | on: 4 | pull_request: 5 | branches: 6 | - master 7 | 8 | jobs: 9 | build: 10 | 11 | runs-on: ubuntu-latest 12 | 13 | strategy: 14 | matrix: 15 | python-version: ['3.9'] 16 | 17 | steps: 18 | - uses: actions/checkout@v2 19 | - name: Set up Python ${{ matrix.python-version }} 20 | uses: actions/setup-python@v2 21 | with: 22 | python-version: ${{ matrix.python-version }} 23 | - name: Install flake8 24 | run: | 25 | python -m pip install flake8==3.8.3 26 | - name: Run formating 27 | run: | 28 | flake8 src -------------------------------------------------------------------------------- /.github/workflows/ci-pypi.yml: -------------------------------------------------------------------------------- 1 | name: Publish Python distributions to PyPI 2 | 3 | on: 4 | push: 5 | branches: 6 | - master 7 | paths: 8 | - 'VERSION' 9 | 10 | jobs: 11 | build-n-publish: 12 | if: github.repository_owner == 'deepmodeling' 13 | name: Build and publish Python distributions to PyPI 14 | runs-on: ubuntu-latest 15 | steps: 16 | - name: Check out the repo 17 | uses: actions/checkout@master 18 | 19 | - name: Set up Python 3.9 20 | uses: actions/setup-python@master 21 | with: 22 | python-version: 3.9 23 | 24 | - name: Install pypa/build 25 | run: >- 26 | python -m 27 | pip install 28 | build 29 | --user 30 | 31 | - name: Build a binary wheel and a source tarball 32 | run: >- 33 | python -m 34 | build 35 | --sdist 36 | --wheel 37 | --outdir dist/ 38 | . 39 | 40 | - name: Publish distribution to PyPI 41 | uses: pypa/gh-action-pypi-publish@master 42 | with: 43 | user: __token__ 44 | password: ${{ secrets.PYPI_API_TOKEN }} 45 | verbose: true 46 | 47 | - name: read version 48 | id: read 49 | uses: juliangruber/read-file-action@v1 50 | with: 51 | path: ./VERSION 52 | 53 | - name: send email 54 | uses: dawidd6/action-send-mail@v3 55 | with: 56 | server_address: smtp.feishu.cn 57 | server_port: 465 58 | username: ${{ secrets.MAILUSERNAME }} 59 | password: ${{ secrets.MAILPASSWORD }} 60 | subject: PyPI Auto Build For Dflow 61 | body: Distribution has been published to https://pypi.org/project/pydflow/${{ steps.read.outputs.content }}/ 62 | to: ${{ secrets.MAIL_RECEIVER_LIST }} 63 | from: Github Actions 64 | content_type: text 65 | -------------------------------------------------------------------------------- /.github/workflows/ci-test.yml: -------------------------------------------------------------------------------- 1 | name: Python unit-tests 2 | 3 | on: 4 | - pull_request 5 | 6 | jobs: 7 | build: 8 | runs-on: ubuntu-latest 9 | 10 | steps: 11 | - uses: actions/checkout@master 12 | - name: Set up Python 3.8 13 | uses: actions/setup-python@master 14 | with: 15 | python-version: 3.8 16 | - name: Install dependencies 17 | run: | 18 | pip install . 19 | pip install pytest 20 | - name: Test 21 | run: DFLOW_MODE=debug pytest -vs tests 22 | -------------------------------------------------------------------------------- /.github/workflows/release-please.yml: -------------------------------------------------------------------------------- 1 | on: 2 | push: 3 | branches: 4 | - master 5 | name: release-please 6 | jobs: 7 | release-please: 8 | if: github.repository_owner == 'deepmodeling' 9 | runs-on: ubuntu-latest 10 | steps: 11 | - uses: google-github-actions/release-please-action@v3 12 | with: 13 | release-type: simple 14 | package-name: release-please-action 15 | version-file: "VERSION" 16 | -------------------------------------------------------------------------------- /.github/workflows/sync-gitee.yml: -------------------------------------------------------------------------------- 1 | name: Sync to Gitee Repo 2 | 3 | on: [ push, delete, create ] 4 | 5 | # Ensures that only one mirror task will run at a time. 6 | concurrency: 7 | group: git-mirror 8 | 9 | jobs: 10 | git-mirror: 11 | if: github.repository_owner == 'deepmodeling' 12 | runs-on: ubuntu-latest 13 | steps: 14 | - uses: wearerequired/git-mirror-action@v1 15 | env: 16 | ORGANIZATION: deepmodeling 17 | SSH_PRIVATE_KEY: ${{ secrets.SYNC_GITEE_PRIVATE_KEY }} 18 | with: 19 | source-repo: "https://github.com/deepmodeling/dflow.git" 20 | destination-repo: "git@gitee.com:deepmodeling/dflow.git" -------------------------------------------------------------------------------- /.github/workflows/sync-github.yml: -------------------------------------------------------------------------------- 1 | name: Sync to Github Repo 2 | 3 | on: [ push, delete, create ] 4 | 5 | jobs: 6 | git-mirror: 7 | if: github.repository_owner == 'deepmodeling' 8 | runs-on: ubuntu-latest 9 | steps: 10 | - uses: actions/checkout@v2 11 | with: 12 | fetch-depth: 0 13 | - uses: pixta-dev/repository-mirroring-action@v1 14 | with: 15 | target_repo_url: 16 | git@github.com:dptech-corp/dflow.git 17 | ssh_private_key: 18 | ${{ secrets.DPTECH_CORP_SSH_PRIVATE_KEY }} 19 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *~ 2 | *.bk* 3 | # Byte-compiled / optimized / DLL files 4 | __pycache__/ 5 | *.py[cod] 6 | *$py.class 7 | 8 | # C extensions 9 | *.so 10 | 11 | # Distribution / packaging 12 | .Python 13 | build/ 14 | develop-eggs/ 15 | dist/ 16 | downloads/ 17 | eggs/ 18 | .eggs/ 19 | lib/ 20 | lib64/ 21 | parts/ 22 | sdist/ 23 | var/ 24 | wheels/ 25 | pip-wheel-metadata/ 26 | share/python-wheels/ 27 | *.egg-info/ 28 | .installed.cfg 29 | *.egg 30 | MANIFEST 31 | 32 | # PyInstaller 33 | # Usually these files are written by a python script from a template 34 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 35 | *.manifest 36 | *.spec 37 | 38 | # Installer logs 39 | pip-log.txt 40 | pip-delete-this-directory.txt 41 | 42 | # Unit test / coverage reports 43 | htmlcov/ 44 | .tox/ 45 | .nox/ 46 | .coverage 47 | .coverage.* 48 | .cache 49 | nosetests.xml 50 | coverage.xml 51 | *.cover 52 | *.py,cover 53 | .hypothesis/ 54 | .pytest_cache/ 55 | 56 | # Translations 57 | *.mo 58 | *.pot 59 | 60 | # Django stuff: 61 | *.log 62 | local_settings.py 63 | db.sqlite3 64 | db.sqlite3-journal 65 | 66 | # Flask stuff: 67 | instance/ 68 | .webassets-cache 69 | 70 | # Scrapy stuff: 71 | .scrapy 72 | 73 | # Sphinx documentation 74 | docs/_build/ 75 | 76 | # PyBuilder 77 | target/ 78 | 79 | # Jupyter Notebook 80 | .ipynb_checkpoints 81 | 82 | # IPython 83 | profile_default/ 84 | ipython_config.py 85 | 86 | # pyenv 87 | .python-version 88 | 89 | # pipenv 90 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 91 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 92 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 93 | # install all needed dependencies. 94 | #Pipfile.lock 95 | 96 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 97 | __pypackages__/ 98 | 99 | # Celery stuff 100 | celerybeat-schedule 101 | celerybeat.pid 102 | 103 | # SageMath parsed files 104 | *.sage.py 105 | 106 | # Environments 107 | .env 108 | .venv 109 | env/ 110 | venv/ 111 | ENV/ 112 | env.bak/ 113 | venv.bak/ 114 | 115 | # Spyder project settings 116 | .spyderproject 117 | .spyproject 118 | 119 | # Rope project settings 120 | .ropeproject 121 | 122 | # mkdocs documentation 123 | /site 124 | 125 | # mypy 126 | .mypy_cache/ 127 | .dmypy.json 128 | dmypy.json 129 | 130 | # Pyre type checker 131 | .pyre/ 132 | 133 | # cached files 134 | _date.py 135 | _version.py 136 | 137 | # vscode 138 | .vscode/ 139 | .idea/ 140 | 141 | # Mac OS 142 | .DS_store 143 | -------------------------------------------------------------------------------- /Dockerfile: -------------------------------------------------------------------------------- 1 | FROM python:3.8 2 | 3 | WORKDIR /data/dflow 4 | ADD requirements.txt ./ 5 | RUN pip install -r requirements.txt 6 | COPY ./ ./ 7 | RUN pip install . 8 | 9 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include VERSION 2 | -------------------------------------------------------------------------------- /VERSION: -------------------------------------------------------------------------------- 1 | 1.8.127 2 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- 1 | # Minimal makefile for Sphinx documentation 2 | # 3 | 4 | # You can set these variables from the command line, and also 5 | # from the environment for the first two. 6 | SPHINXOPTS ?= 7 | SPHINXBUILD ?= sphinx-build 8 | SOURCEDIR = source 9 | BUILDDIR = build 10 | 11 | # Put it first so that "make" without argument is like "make help". 12 | help: 13 | @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 14 | 15 | .PHONY: help Makefile 16 | 17 | # Catch-all target: route all unknown targets to Sphinx using the new 18 | # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). 19 | %: Makefile 20 | @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) 21 | -------------------------------------------------------------------------------- /docs/developer.md: -------------------------------------------------------------------------------- 1 | # Developers' guide 2 | 3 | - [How to contribute](#how-to-contribute) 4 | 5 | ## How to contribute 6 | 7 | Dflow uses release-please action to manage versions and CHANGELOG. Developers and contributors MUST write [conventionalcommits](https://www.conventionalcommits.org/) to make release-please work. 8 | -------------------------------------------------------------------------------- /docs/imgs/dflow_architecture.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/docs/imgs/dflow_architecture.png -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- 1 | @ECHO OFF 2 | 3 | pushd %~dp0 4 | 5 | REM Command file for Sphinx documentation 6 | 7 | if "%SPHINXBUILD%" == "" ( 8 | set SPHINXBUILD=sphinx-build 9 | ) 10 | set SOURCEDIR=source 11 | set BUILDDIR=build 12 | 13 | if "%1" == "" goto help 14 | 15 | %SPHINXBUILD% >NUL 2>NUL 16 | if errorlevel 9009 ( 17 | echo. 18 | echo.The 'sphinx-build' command was not found. Make sure you have Sphinx 19 | echo.installed, then set the SPHINXBUILD environment variable to point 20 | echo.to the full path of the 'sphinx-build' executable. Alternatively you 21 | echo.may add the Sphinx directory to PATH. 22 | echo. 23 | echo.If you don't have Sphinx installed, grab it from 24 | echo.https://www.sphinx-doc.org/ 25 | exit /b 1 26 | ) 27 | 28 | %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 29 | goto end 30 | 31 | :help 32 | %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% 33 | 34 | :end 35 | popd 36 | -------------------------------------------------------------------------------- /docs/requirements.txt: -------------------------------------------------------------------------------- 1 | ../requirements.txt -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- 1 | # Configuration file for the Sphinx documentation builder. 2 | # 3 | # This file only contains a selection of the most common options. For a full 4 | # list see the documentation: 5 | # https://www.sphinx-doc.org/en/master/usage/configuration.html 6 | 7 | # -- Path setup -------------------------------------------------------------- 8 | 9 | # If extensions (or modules to document with autodoc) are in another directory, 10 | # add these directories to sys.path here. If the directory is relative to the 11 | # documentation root, use os.path.abspath to make it absolute, like shown here. 12 | # 13 | import os 14 | import sys 15 | sys.path.insert(0, os.path.abspath('../../src')) 16 | 17 | 18 | # -- Project information ----------------------------------------------------- 19 | 20 | project = 'dflow' 21 | copyright = '2022, zjgemi' 22 | author = 'zjgemi' 23 | 24 | # The full version, including alpha/beta/rc tags 25 | with open("../../VERSION", "r") as f: 26 | release = f.read() 27 | 28 | 29 | # -- General configuration --------------------------------------------------- 30 | 31 | # Add any Sphinx extension module names here, as strings. They can be 32 | # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom 33 | # ones. 34 | extensions = [ 35 | 'sphinx.ext.autodoc', 36 | 'sphinx.ext.doctest', 37 | 'sphinx.ext.intersphinx', 38 | 'sphinx.ext.todo', 39 | 'sphinx.ext.coverage', 40 | 'sphinx.ext.mathjax', 41 | 'sphinx.ext.napoleon', 42 | ] 43 | 44 | # Add any paths that contain templates here, relative to this directory. 45 | templates_path = ['_templates'] 46 | 47 | # List of patterns, relative to source directory, that match files and 48 | # directories to ignore when looking for source files. 49 | # This pattern also affects html_static_path and html_extra_path. 50 | exclude_patterns = [] 51 | 52 | 53 | # -- Options for HTML output ------------------------------------------------- 54 | 55 | # The theme to use for HTML and HTML Help pages. See the documentation for 56 | # a list of builtin themes. 57 | # 58 | html_theme = 'default' 59 | 60 | # Add any paths that contain custom static files (such as style sheets) here, 61 | # relative to this directory. They are copied after the builtin static files, 62 | # so a file named "default.css" will overwrite the builtin "default.css". 63 | html_static_path = ['_static'] 64 | -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- 1 | .. dflow documentation master file, created by 2 | sphinx-quickstart on Sun May 29 08:42:35 2022. 3 | You can adapt this file completely to your liking, but it should at least 4 | contain the root `toctree` directive. 5 | 6 | Welcome to dflow's documentation! 7 | ================================= 8 | 9 | .. toctree:: 10 | :maxdepth: 2 11 | :caption: Contents: 12 | 13 | 14 | 15 | Indices and tables 16 | ================== 17 | 18 | * :ref:`genindex` 19 | * :ref:`modindex` 20 | * :ref:`search` 21 | -------------------------------------------------------------------------------- /examples/test_async.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import (InputArtifact, OutputArtifact, ShellOPTemplate, Step, 4 | Workflow) 5 | 6 | if __name__ == "__main__": 7 | hello = ShellOPTemplate(name='Hello', 8 | image="alpine:latest", 9 | script="sleep 60 && echo Hello > /tmp/bar.txt") 10 | hello.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")} 11 | 12 | echo = ShellOPTemplate(name='Echo', 13 | image="alpine:latest", 14 | script="while [ 1 ]; do echo 'waiting...';" 15 | " if [ -f /tmp/foo.txt ]; then cat /tmp/foo.txt;" 16 | " break; fi; sleep 1; done") 17 | echo.inputs.artifacts = {"foo": InputArtifact(path="/tmp/foo.txt")} 18 | 19 | wf = Workflow(name="async") 20 | step0 = Step(name="hello", template=hello) 21 | # This step will give output artifact "bar" which contains "Hello" after 22 | # 60 seconds 23 | step1 = Step(name="echo", template=echo, artifacts={ 24 | "foo": step0.outputs.artifacts["bar"].pvc()}) 25 | # This step will wait the last step to finish and then print its output 26 | # artifact 27 | wf.add([step0, step1]) 28 | wf.submit() 29 | 30 | while wf.query_status() in ["Pending", "Running"]: 31 | time.sleep(1) 32 | 33 | assert(wf.query_status() == "Succeeded") 34 | -------------------------------------------------------------------------------- /examples/test_big_parameter.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import InputParameter, OutputParameter, Step, Steps, Workflow 4 | from dflow.python import (OP, OPIO, BigParameter, OPIOSign, PythonOPTemplate, 5 | upload_packages) 6 | 7 | if "__file__" in locals(): 8 | upload_packages.append(__file__) 9 | 10 | 11 | class Hello: 12 | def __init__(self, msg): 13 | self.msg = msg 14 | 15 | 16 | class Duplicate(OP): 17 | def __init__(self): 18 | pass 19 | 20 | @classmethod 21 | def get_input_sign(cls): 22 | return OPIOSign({ 23 | 'foo': BigParameter(Hello) 24 | }) 25 | 26 | @classmethod 27 | def get_output_sign(cls): 28 | return OPIOSign({ 29 | 'foo': BigParameter(Hello) 30 | }) 31 | 32 | @OP.exec_sign_check 33 | def execute( 34 | self, 35 | op_in: OPIO, 36 | ) -> OPIO: 37 | foo = op_in["foo"] 38 | print(foo.msg) 39 | foo.msg = foo.msg * 2 40 | op_out = OPIO({ 41 | "foo": foo 42 | }) 43 | return op_out 44 | 45 | 46 | def test_big_parameter(): 47 | wf = Workflow(name="big-param") 48 | 49 | steps = Steps(name="hello-steps") 50 | steps.inputs.parameters["foo"] = InputParameter() 51 | steps.outputs.parameters["foo"] = OutputParameter() 52 | 53 | step1 = Step( 54 | name="step1", 55 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 56 | parameters={"foo": steps.inputs.parameters["foo"]}, 57 | key="step1" 58 | ) 59 | steps.add(step1) 60 | 61 | step2 = Step( 62 | name="step2", 63 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 64 | parameters={"foo": step1.outputs.parameters["foo"]}, 65 | key="step2" 66 | ) 67 | steps.add(step2) 68 | 69 | steps.outputs.parameters["foo"].value_from_parameter = \ 70 | step2.outputs.parameters["foo"] 71 | 72 | big_step = Step(name="big-step", template=steps, 73 | parameters={"foo": Hello("hello")}) 74 | wf.add(big_step) 75 | wf.submit() 76 | 77 | while wf.query_status() in ["Pending", "Running"]: 78 | time.sleep(1) 79 | 80 | assert(wf.query_status() == "Succeeded") 81 | step = wf.query_step(name="step1")[0] 82 | assert(step.phase == "Succeeded") 83 | print(step.outputs.parameters["foo"].value) 84 | 85 | step.modify_output_parameter("foo", Hello("byebye")) 86 | wf = Workflow(name="big-param-resubmit") 87 | wf.add(big_step) 88 | wf.submit(reuse_step=[step]) 89 | 90 | 91 | if __name__ == "__main__": 92 | test_big_parameter() 93 | -------------------------------------------------------------------------------- /examples/test_bohrium_datasets.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow 5 | from dflow.plugins.bohrium import BohriumDatasetsArtifact 6 | from dflow.plugins.dispatcher import DispatcherExecutor 7 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate 8 | 9 | 10 | class Hello(OP): 11 | def __init__(self): 12 | pass 13 | 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign({ 17 | 'foo': Artifact(Path), 18 | }) 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign() 23 | 24 | @OP.exec_sign_check 25 | def execute( 26 | self, 27 | op_in: OPIO, 28 | ) -> OPIO: 29 | os.system("ls -l %s" % op_in["foo"]) 30 | return OPIO() 31 | 32 | 33 | if __name__ == "__main__": 34 | dispatcher_executor = DispatcherExecutor( 35 | machine_dict={ 36 | "batch_type": "Bohrium", 37 | "context_type": "Bohrium", 38 | }, 39 | ) 40 | 41 | wf = Workflow(name="datasets") 42 | art = BohriumDatasetsArtifact("/bohr/water-example-m7kb/v1") 43 | step = Step( 44 | name="step", 45 | template=PythonOPTemplate(Hello, image="python:3.8"), 46 | artifacts={"foo": art}, 47 | executor=dispatcher_executor, 48 | ) 49 | wf.add(step) 50 | wf.submit() 51 | -------------------------------------------------------------------------------- /examples/test_conditional_outputs.py: -------------------------------------------------------------------------------- 1 | import random 2 | import time 3 | 4 | from dflow import (OutputArtifact, OutputParameter, Outputs, Step, Steps, 5 | Workflow, if_expression) 6 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 7 | upload_packages) 8 | 9 | if "__file__" in locals(): 10 | upload_packages.append(__file__) 11 | 12 | 13 | class Random(OP): 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign() 17 | 18 | @classmethod 19 | def get_output_sign(cls): 20 | return OPIOSign({ 21 | "is_head": bool, 22 | "msg1": str, 23 | "msg2": str, 24 | "foo": Artifact(str), 25 | "bar": Artifact(str) 26 | }) 27 | 28 | @OP.exec_sign_check 29 | def execute( 30 | self, 31 | op_in: OPIO, 32 | ) -> OPIO: 33 | open("foo.txt", "w").write("head") 34 | open("bar.txt", "w").write("tail") 35 | if random.random() < 0.5: 36 | is_head = True 37 | else: 38 | is_head = False 39 | return OPIO({ 40 | "is_head": is_head, 41 | "msg1": "head", 42 | "msg2": "tail", 43 | "foo": "foo.txt", 44 | "bar": "bar.txt" 45 | }) 46 | 47 | 48 | def test_conditional_outputs(): 49 | steps = Steps("conditional-steps", outputs=Outputs( 50 | parameters={"msg": OutputParameter()}, 51 | artifacts={"res": OutputArtifact()})) 52 | 53 | random_step = Step( 54 | name="random", 55 | template=PythonOPTemplate(Random, image="python:3.8") 56 | ) 57 | steps.add(random_step) 58 | 59 | steps.outputs.parameters["msg"].value_from_expression = if_expression( 60 | _if=random_step.outputs.parameters["is_head"], 61 | _then=random_step.outputs.parameters["msg1"], 62 | _else=random_step.outputs.parameters["msg2"]) 63 | 64 | steps.outputs.artifacts["res"].from_expression = if_expression( 65 | _if=random_step.outputs.parameters["is_head"], 66 | _then=random_step.outputs.artifacts["foo"], 67 | _else=random_step.outputs.artifacts["bar"]) 68 | 69 | wf = Workflow(name="conditional", steps=steps) 70 | 71 | wf.submit() 72 | while wf.query_status() in ["Pending", "Running"]: 73 | time.sleep(1) 74 | 75 | assert(wf.query_status() == "Succeeded") 76 | 77 | 78 | if __name__ == "__main__": 79 | test_conditional_outputs() 80 | -------------------------------------------------------------------------------- /examples/test_dag.py: -------------------------------------------------------------------------------- 1 | from dflow import (InputArtifact, InputParameter, OutputArtifact, 2 | OutputParameter, ShellOPTemplate, Task, Workflow) 3 | 4 | if __name__ == "__main__": 5 | hello = ShellOPTemplate( 6 | name='Hello', 7 | image="alpine:latest", 8 | script="echo Hello > /tmp/bar.txt && echo 1 > /tmp/result.txt") 9 | hello.outputs.parameters = {"msg": OutputParameter( 10 | value_from_path="/tmp/result.txt")} 11 | hello.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")} 12 | 13 | duplicate = ShellOPTemplate( 14 | name='Duplicate', 15 | image="alpine:latest", 16 | script="cat /tmp/foo.txt /tmp/foo.txt > /tmp/bar.txt && " 17 | "echo $(({{inputs.parameters.msg}}*2)) > /tmp/result.txt") 18 | duplicate.inputs.parameters = {"msg": InputParameter()} 19 | duplicate.outputs.parameters = {"msg": OutputParameter( 20 | value_from_path="/tmp/result.txt")} 21 | duplicate.inputs.artifacts = {"foo": InputArtifact(path="/tmp/foo.txt")} 22 | duplicate.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")} 23 | 24 | wf = Workflow(name="dag") 25 | hello0 = Task(name="hello0", template=hello) 26 | wf.add(hello0) 27 | hello1 = Task(name="hello1", 28 | template=duplicate, 29 | parameters={"msg": hello0.outputs.parameters["msg"]}, 30 | artifacts={"foo": hello0.outputs.artifacts["bar"]}) 31 | wf.add(hello1) 32 | wf.submit() 33 | -------------------------------------------------------------------------------- /examples/test_dag_slices.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import DAG, OutputArtifact, Task, Workflow, argo_range 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Hello(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | 'filename': str 20 | }) 21 | 22 | @classmethod 23 | def get_output_sign(cls): 24 | return OPIOSign({ 25 | 'foo': Artifact(str) 26 | }) 27 | 28 | @OP.exec_sign_check 29 | def execute( 30 | self, 31 | op_in: OPIO, 32 | ) -> OPIO: 33 | open(op_in["filename"], "w").write("foo") 34 | op_out = OPIO({ 35 | 'foo': op_in["filename"] 36 | }) 37 | return op_out 38 | 39 | 40 | class Check(OP): 41 | def __init__(self): 42 | pass 43 | 44 | @classmethod 45 | def get_input_sign(cls): 46 | return OPIOSign({ 47 | 'foo': Artifact(List[str]) 48 | }) 49 | 50 | @classmethod 51 | def get_output_sign(cls): 52 | return OPIOSign() 53 | 54 | @OP.exec_sign_check 55 | def execute( 56 | self, 57 | op_in: OPIO, 58 | ) -> OPIO: 59 | print(op_in["foo"]) 60 | return OPIO() 61 | 62 | 63 | def test_dag_slices(): 64 | dag = DAG("slices-dag") 65 | hello0 = Task("hello0", 66 | PythonOPTemplate(Hello, image="python:3.8", 67 | slices=Slices("{{item}}", 68 | input_parameter=["filename"], 69 | output_artifact=["foo"] 70 | ) 71 | ), 72 | parameters={"filename": ["f1.txt", "f2.txt"]}, 73 | with_param=argo_range(2), 74 | key="hello-0-{{item}}") 75 | dag.add(hello0) 76 | check0 = Task("check0", 77 | PythonOPTemplate(Check, image="python:3.8"), 78 | artifacts={"foo": hello0.outputs.artifacts["foo"]}) 79 | dag.add(check0) 80 | hello1 = Task("hello1", 81 | PythonOPTemplate(Hello, image="python:3.8", 82 | slices=Slices("{{item}}", 83 | input_parameter=["filename"], 84 | output_artifact=["foo"] 85 | ) 86 | ), 87 | parameters={"filename": []}, 88 | with_param=argo_range(0), 89 | key="hello-1-{{item}}") 90 | dag.add(hello1) 91 | check1 = Task("check1", 92 | PythonOPTemplate(Check, image="python:3.8"), 93 | artifacts={"foo": hello1.outputs.artifacts["foo"]}) 94 | dag.add(check1) 95 | dag.outputs.artifacts["foo"] = OutputArtifact( 96 | _from=hello1.outputs.artifacts["foo"]) 97 | 98 | wf = Workflow("dag-slices", dag=dag) 99 | wf.submit() 100 | 101 | while wf.query_status() in ["Pending", "Running"]: 102 | time.sleep(1) 103 | 104 | assert(wf.query_status() == "Succeeded") 105 | step0 = wf.query_step(key="hello-0-0")[0] 106 | 107 | wf2 = Workflow("dag-slices-resubmit", dag=dag) 108 | wf2.submit(reuse_step=[step0]) 109 | 110 | 111 | if __name__ == "__main__": 112 | test_dag_slices() 113 | -------------------------------------------------------------------------------- /examples/test_datasets.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow 5 | from dflow.plugins.datasets import DatasetsArtifact 6 | from dflow.plugins.dispatcher import DispatcherExecutor 7 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate 8 | 9 | 10 | class Hello(OP): 11 | def __init__(self): 12 | pass 13 | 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign({ 17 | 'foo': Artifact(Path), 18 | }) 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign() 23 | 24 | @OP.exec_sign_check 25 | def execute( 26 | self, 27 | op_in: OPIO, 28 | ) -> OPIO: 29 | print(os.listdir(op_in["foo"])) 30 | return OPIO() 31 | 32 | 33 | if __name__ == "__main__": 34 | dispatcher_executor = DispatcherExecutor( 35 | machine_dict={ 36 | "batch_type": "Bohrium", 37 | "context_type": "Bohrium", 38 | }, 39 | ) 40 | 41 | wf = Workflow(name="datasets") 42 | art = DatasetsArtifact.from_urn("launching+datasets://pdbbind@0.1.2") 43 | step = Step( 44 | name="step", 45 | template=PythonOPTemplate(Hello, image="python:3.8"), 46 | artifacts={"foo": art}, 47 | executor=dispatcher_executor, 48 | ) 49 | wf.add(step) 50 | wf.submit() 51 | -------------------------------------------------------------------------------- /examples/test_datasets_local.py: -------------------------------------------------------------------------------- 1 | import os 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow 5 | from dflow.plugins.datasets import DatasetsArtifact, config 6 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate 7 | 8 | 9 | class Hello(OP): 10 | def __init__(self): 11 | pass 12 | 13 | @classmethod 14 | def get_input_sign(cls): 15 | return OPIOSign({ 16 | 'foo': Artifact(Path), 17 | }) 18 | 19 | @classmethod 20 | def get_output_sign(cls): 21 | return OPIOSign() 22 | 23 | @OP.exec_sign_check 24 | def execute( 25 | self, 26 | op_in: OPIO, 27 | ) -> OPIO: 28 | print(os.listdir(op_in["foo"])) 29 | return OPIO() 30 | 31 | 32 | if __name__ == "__main__": 33 | wf = Workflow(name="datasets") 34 | # art = DatasetsArtifact.from_rclone_config(""" 35 | # [pdbbind@0.1.2] 36 | # type = ftp 37 | # disable_tls13 = true 38 | # concurrency = 3 39 | # host = uftp.mlops-passthrough.dp.tech 40 | # user = xxx 41 | # pass = xxx 42 | # explicit_tls = true""") 43 | config["user"] = "xxx" 44 | config["password"] = "xxx" 45 | art = DatasetsArtifact.from_urn("launching+datasets://pdbbind@0.1.2") 46 | step = Step( 47 | name="step", 48 | template=PythonOPTemplate(Hello, image="python:3.8"), 49 | artifacts={"foo": art}, 50 | ) 51 | wf.add(step) 52 | wf.submit() 53 | -------------------------------------------------------------------------------- /examples/test_dispatcher.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | from typing import List 4 | 5 | from dflow import Step, Workflow, download_artifact, upload_artifact 6 | from dflow.plugins.dispatcher import DispatcherExecutor 7 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate 8 | 9 | 10 | class Duplicate(OP): 11 | def __init__(self): 12 | pass 13 | 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign({ 17 | 'msg': str, 18 | 'num': int, 19 | 'foo': Artifact(Path), 20 | 'idir': Artifact(Path), 21 | }) 22 | 23 | @classmethod 24 | def get_output_sign(cls): 25 | return OPIOSign({ 26 | 'msg': List[str], 27 | 'bar': Artifact(Path), 28 | 'odir': Artifact(Path), 29 | }) 30 | 31 | @OP.exec_sign_check 32 | def execute( 33 | self, 34 | op_in: OPIO, 35 | ) -> OPIO: 36 | op_out = OPIO({ 37 | "msg": [op_in["msg"] * op_in["num"]], 38 | "bar": Path("output.txt"), 39 | "odir": Path("todir"), 40 | }) 41 | 42 | content = open(op_in['foo'], "r").read() 43 | open("output.txt", "w").write(content * op_in["num"]) 44 | 45 | Path(op_out['odir']).mkdir() 46 | for ii in ['f1', 'f2']: 47 | (op_out['odir']/ii).write_text(op_in['num'] 48 | * (op_in['idir']/ii).read_text()) 49 | 50 | return op_out 51 | 52 | 53 | def make_idir(): 54 | idir = Path("tidir") 55 | idir.mkdir(exist_ok=True) 56 | (idir / "f1").write_text("foo") 57 | (idir / "f2").write_text("bar") 58 | 59 | 60 | if __name__ == "__main__": 61 | wf = Workflow(name="dispatcher") 62 | 63 | with open("foo.txt", "w") as f: 64 | f.write("Hi") 65 | make_idir() 66 | 67 | artifact0 = upload_artifact("foo.txt") 68 | artifact1 = upload_artifact("tidir") 69 | print(artifact0) 70 | print(artifact1) 71 | 72 | # run ../scripts/start-slurm.sh first to start up a slurm cluster 73 | import socket 74 | 75 | def get_my_ip_address(): 76 | with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: 77 | s.connect(("8.8.8.8", 80)) 78 | return s.getsockname()[0] 79 | 80 | dispatcher_executor = DispatcherExecutor( 81 | host=get_my_ip_address(), 82 | username="root", 83 | port=31129, 84 | queue_name="normal", 85 | remote_root="/data", 86 | password="password", 87 | ) 88 | step = Step( 89 | name="step", 90 | template=PythonOPTemplate(Duplicate), 91 | parameters={"msg": "Hello", "num": 3}, 92 | artifacts={"foo": artifact0, "idir": artifact1}, 93 | executor=dispatcher_executor 94 | ) 95 | wf.add(step) 96 | wf.submit() 97 | 98 | while wf.query_status() in ["Pending", "Running"]: 99 | time.sleep(1) 100 | 101 | assert(wf.query_status() == "Succeeded") 102 | step = wf.query_step(name="step")[0] 103 | assert(step.phase == "Succeeded") 104 | 105 | print(download_artifact(step.outputs.artifacts["bar"])) 106 | print(download_artifact(step.outputs.artifacts["odir"])) 107 | 108 | print(step.outputs.parameters["msg"].value, 109 | type(step.outputs.parameters["msg"].value)) 110 | -------------------------------------------------------------------------------- /examples/test_error_handling.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import Step, Workflow 4 | from dflow.python import (OP, OPIO, FatalError, OPIOSign, PythonOPTemplate, 5 | TransientError) 6 | 7 | 8 | class Hello(OP): 9 | def __init__(self): 10 | pass 11 | 12 | @classmethod 13 | def get_input_sign(cls): 14 | return OPIOSign() 15 | 16 | @classmethod 17 | def get_output_sign(cls): 18 | return OPIOSign() 19 | 20 | @OP.exec_sign_check 21 | def execute( 22 | self, 23 | op_in: OPIO, 24 | ) -> OPIO: 25 | raise TransientError("Hello") 26 | return OPIO() 27 | 28 | 29 | class Timeout(OP): 30 | def __init__(self): 31 | pass 32 | 33 | @classmethod 34 | def get_input_sign(cls): 35 | return OPIOSign() 36 | 37 | @classmethod 38 | def get_output_sign(cls): 39 | return OPIOSign() 40 | 41 | @OP.exec_sign_check 42 | def execute( 43 | self, 44 | op_in: OPIO, 45 | ) -> OPIO: 46 | time.sleep(100) 47 | return OPIO() 48 | 49 | 50 | class Goodbye(OP): 51 | def __init__(self): 52 | pass 53 | 54 | @classmethod 55 | def get_input_sign(cls): 56 | return OPIOSign() 57 | 58 | @classmethod 59 | def get_output_sign(cls): 60 | return OPIOSign() 61 | 62 | @OP.exec_sign_check 63 | def execute( 64 | self, 65 | op_in: OPIO, 66 | ) -> OPIO: 67 | raise FatalError("Goodbye") 68 | return OPIO() 69 | 70 | 71 | if __name__ == "__main__": 72 | wf = Workflow(name="error-handling") 73 | 74 | step = Step( 75 | name="hello0", 76 | template=PythonOPTemplate( 77 | Hello, image="python:3.8", retry_on_transient_error=1), 78 | continue_on_failed=True 79 | ) 80 | wf.add(step) 81 | step = Step( 82 | name="hello1", 83 | template=PythonOPTemplate(Timeout, image="python:3.8", timeout=10, 84 | retry_on_transient_error=1, 85 | timeout_as_transient_error=True), 86 | continue_on_failed=True 87 | ) 88 | wf.add(step) 89 | step = Step( 90 | name="hello2", 91 | template=PythonOPTemplate( 92 | Goodbye, image="python:3.8", retry_on_transient_error=1) 93 | ) 94 | wf.add(step) 95 | wf.submit() 96 | -------------------------------------------------------------------------------- /examples/test_func_op.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow, download_artifact, upload_artifact 5 | from dflow.python import OP, Artifact, PythonOPTemplate, upload_packages 6 | 7 | if '__file__' in locals(): 8 | upload_packages.append(__file__) 9 | 10 | 11 | @OP.function 12 | def Duplicate(msg: str, num: int, foo: Artifact(Path), idir: Artifact(Path) 13 | ) -> {'msg': str, 'bar': Artifact(Path), 'odir': Artifact(Path)}: 14 | msg_out = msg * num 15 | bar = Path('output.txt') 16 | odir = Path('todir') 17 | content = open(foo, 'r').read() 18 | open('output.txt', 'w').write(content * num) 19 | odir.mkdir() 20 | for ii in ['f1', 'f2']: 21 | (odir / ii).write_text(num * (idir / ii).read_text()) 22 | return {'msg': msg_out, 'bar': Path(bar), 'odir': Path(odir), } 23 | 24 | 25 | def make_idir(): 26 | idir = Path('tidir') 27 | idir.mkdir(exist_ok=True) 28 | (idir / 'f1').write_text('foo') 29 | (idir / 'f2').write_text('bar') 30 | 31 | 32 | def test_python(): 33 | with open('foo.txt', 'w') as f: 34 | f.write('Hi') 35 | make_idir() 36 | 37 | artifact0 = upload_artifact('foo.txt') 38 | artifact1 = upload_artifact('tidir') 39 | print(artifact0) 40 | print(artifact1) 41 | 42 | with Workflow(name='python-sugar') as wf: 43 | step = Step( 44 | name='step', template=PythonOPTemplate( 45 | Duplicate, image='python:3.8'), parameters={ 46 | 'msg': 'Hello', 'num': 3}, artifacts={ 47 | 'foo': artifact0, 'idir': artifact1}, ) 48 | 49 | while wf.query_status() in ['Pending', 'Running']: 50 | time.sleep(1) 51 | 52 | assert (wf.query_status() == 'Succeeded') 53 | step = wf.query_step(name='step')[0] 54 | assert (step.phase == 'Succeeded') 55 | 56 | print(download_artifact(step.outputs.artifacts['bar'])) 57 | print(download_artifact(step.outputs.artifacts['odir'])) 58 | 59 | 60 | if __name__ == '__main__': 61 | test_python() 62 | -------------------------------------------------------------------------------- /examples/test_launching.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | 4 | from dflow.python import OP, OPIO, Artifact, OPIOSign, Parameter 5 | 6 | 7 | """@OP.function 8 | def Duplicate( 9 | foo: Artifact(Path, description="input file"), 10 | num: Parameter(int, default=2, description="number"), 11 | ) -> {"bar": Artifact(Path)}: 12 | with open(foo, "r") as f: 13 | content = f.read() 14 | with open("bar.txt", "w") as f: 15 | f.write(content * num) 16 | return {"bar": Path("bar.txt")}""" 17 | 18 | 19 | class Duplicate(OP): 20 | def __init__(self): 21 | pass 22 | 23 | @classmethod 24 | def get_input_sign(cls): 25 | return OPIOSign({ 26 | "foo": Artifact(Path, description="input file"), 27 | "num": Parameter(int, default=2, description="number"), 28 | }) 29 | 30 | @classmethod 31 | def get_output_sign(cls): 32 | return OPIOSign({ 33 | "bar": Artifact(Path), 34 | }) 35 | 36 | @OP.exec_sign_check 37 | def execute( 38 | self, 39 | op_in: OPIO, 40 | ) -> OPIO: 41 | with open(op_in["foo"], "r") as f: 42 | content = f.read() 43 | with open("bar.txt", "w") as f: 44 | f.write(content * op_in["num"]) 45 | return OPIO({"bar": Path("bar.txt")}) 46 | 47 | 48 | if __name__ == '__main__': 49 | from dflow.plugins.launching import OP_to_parser 50 | to_parser = OP_to_parser(Duplicate) 51 | to_parser()(sys.argv[1:]) 52 | -------------------------------------------------------------------------------- /examples/test_launching_2.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from pathlib import Path 3 | 4 | from dflow import InputArtifact, InputParameter, OutputArtifact, Step, Steps 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, Parameter, 6 | PythonOPTemplate) 7 | 8 | 9 | class Duplicate(OP): 10 | def __init__(self): 11 | pass 12 | 13 | @classmethod 14 | def get_input_sign(cls): 15 | return OPIOSign({ 16 | "foo": Artifact(Path, description="input file"), 17 | "num": Parameter(int, default=2, description="number"), 18 | }) 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign({ 23 | "bar": Artifact(Path), 24 | }) 25 | 26 | @OP.exec_sign_check 27 | def execute( 28 | self, 29 | op_in: OPIO, 30 | ) -> OPIO: 31 | with open(op_in["foo"], "r") as f: 32 | content = f.read() 33 | with open("bar.txt", "w") as f: 34 | f.write(content * op_in["num"]) 35 | return OPIO({"bar": Path("bar.txt")}) 36 | 37 | 38 | if __name__ == '__main__': 39 | from dflow.plugins.launching import OP_to_parser 40 | steps = Steps(name="duplicate-steps") 41 | steps.inputs.parameters["num"] = InputParameter(value=2, 42 | description="number") 43 | steps.inputs.artifacts["foo"] = InputArtifact(description="input file") 44 | step = Step( 45 | name="duplicate", 46 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 47 | parameters={"num": steps.inputs.parameters["num"]}, 48 | artifacts={"foo": steps.inputs.artifacts["foo"]}) 49 | steps.add(step) 50 | steps.outputs.artifacts["bar"] = OutputArtifact( 51 | _from=step.outputs.artifacts["bar"]) 52 | to_parser = OP_to_parser(steps) 53 | to_parser()(sys.argv[1:]) 54 | -------------------------------------------------------------------------------- /examples/test_lineage.py: -------------------------------------------------------------------------------- 1 | from pathlib import Path 2 | 3 | from dflow import Step, Workflow, config, upload_artifact 4 | from dflow.plugins.metadata import MetadataClient 5 | from dflow.python import OP, Artifact, PythonOPTemplate 6 | 7 | 8 | @OP.function 9 | def duplicate(foo: Artifact(Path)) -> {'bar': Artifact(Path)}: 10 | with open(foo, 'r') as f: 11 | content = f.read() 12 | with open('output.txt', 'w') as f: 13 | f.write(content * 2) 14 | return {'bar': Path('output.txt')} 15 | 16 | 17 | if __name__ == "__main__": 18 | config["lineage"] = MetadataClient( 19 | project="", 20 | token="", 21 | ) 22 | config["register_tasks"] = True 23 | 24 | wf = Workflow(name="lineage") 25 | 26 | with open("foo.txt", "w") as f: 27 | f.write("Hi") 28 | 29 | art = upload_artifact("foo.txt", namespace="lineage-test", name="foo") 30 | step1 = Step( 31 | name="step1", 32 | template=PythonOPTemplate( 33 | duplicate, image="registry.dp.tech/dptech/dp-metadata-sdk:latest"), 34 | artifacts={"foo": art}, 35 | ) 36 | wf.add(step1) 37 | 38 | step2 = Step( 39 | name="step2", 40 | template=PythonOPTemplate( 41 | duplicate, image="registry.dp.tech/dptech/dp-metadata-sdk:latest"), 42 | artifacts={"foo": step1.outputs.artifacts["bar"]}, 43 | ) 44 | wf.add(step2) 45 | wf.submit() 46 | -------------------------------------------------------------------------------- /examples/test_loop_slices.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import OutputArtifact, Step, Steps, Workflow, argo_range 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Hello(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | 'filename': str 20 | }) 21 | 22 | @classmethod 23 | def get_output_sign(cls): 24 | return OPIOSign({ 25 | 'foo': Artifact(str) 26 | }) 27 | 28 | @OP.exec_sign_check 29 | def execute( 30 | self, 31 | op_in: OPIO, 32 | ) -> OPIO: 33 | open(op_in["filename"], "w").write("foo") 34 | op_out = OPIO({ 35 | 'foo': op_in["filename"] 36 | }) 37 | return op_out 38 | 39 | 40 | class Check(OP): 41 | def __init__(self): 42 | pass 43 | 44 | @classmethod 45 | def get_input_sign(cls): 46 | return OPIOSign({ 47 | 'foo': Artifact(List[str]) 48 | }) 49 | 50 | @classmethod 51 | def get_output_sign(cls): 52 | return OPIOSign({ 53 | 'path': List[str] 54 | }) 55 | 56 | @OP.exec_sign_check 57 | def execute( 58 | self, 59 | op_in: OPIO, 60 | ) -> OPIO: 61 | print(op_in["foo"]) 62 | return OPIO({ 63 | 'path': op_in["foo"] 64 | }) 65 | 66 | 67 | def test_loop_slices(): 68 | steps = Steps("slices-steps") 69 | hello0 = Step("hello0", 70 | PythonOPTemplate(Hello, image="python:3.8", 71 | slices=Slices("{{item}}", 72 | input_parameter=["filename"], 73 | output_artifact=["foo"] 74 | ) 75 | ), 76 | parameters={"filename": ["f1.txt", "f2.txt"]}, 77 | with_param=argo_range(2), 78 | key="hello-0-{{item}}") 79 | steps.add(hello0) 80 | check0 = Step("check0", 81 | PythonOPTemplate(Check, image="python:3.8"), 82 | artifacts={"foo": hello0.outputs.artifacts["foo"]}) 83 | steps.add(check0) 84 | hello1 = Step("hello1", 85 | PythonOPTemplate(Hello, image="python:3.8", 86 | slices=Slices("{{item}}", 87 | input_parameter=["filename"], 88 | output_artifact=["foo"] 89 | ) 90 | ), 91 | parameters={"filename": []}, 92 | with_param=argo_range(0), 93 | key="hello-1-{{item}}") 94 | steps.add(hello1) 95 | check1 = Step("check1", 96 | PythonOPTemplate(Check, image="python:3.8"), 97 | artifacts={"foo": hello1.outputs.artifacts["foo"]}) 98 | steps.add(check1) 99 | steps.outputs.artifacts["foo"] = OutputArtifact( 100 | _from=hello1.outputs.artifacts["foo"]) 101 | 102 | wf = Workflow("loop-slices", steps=steps) 103 | wf.submit() 104 | 105 | while wf.query_status() in ["Pending", "Running"]: 106 | time.sleep(1) 107 | 108 | assert(wf.query_status() == "Succeeded") 109 | step0 = wf.query_step(key="hello-0-0")[0] 110 | 111 | wf2 = Workflow("loop-slices-resubmit", steps=steps) 112 | wf2.submit(reuse_step=[step0]) 113 | 114 | while wf2.query_status() in ["Pending", "Running"]: 115 | time.sleep(1) 116 | 117 | assert(wf2.query_status() == "Succeeded") 118 | check0 = wf2.query_step("check0")[0] 119 | assert len(check0.outputs.parameters["path"].value) == 2 120 | 121 | 122 | if __name__ == "__main__": 123 | test_loop_slices() 124 | -------------------------------------------------------------------------------- /examples/test_metadata.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow, config, upload_artifact # , S3Artifact 5 | from dflow.plugins.metadata import MetadataClient 6 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate 7 | 8 | 9 | class Duplicate(OP): 10 | def __init__(self): 11 | pass 12 | 13 | @classmethod 14 | def get_input_sign(cls): 15 | return OPIOSign({ 16 | 'foo': Artifact(Path), 17 | }) 18 | 19 | @classmethod 20 | def get_output_sign(cls): 21 | return OPIOSign({ 22 | 'bar': Artifact(Path), 23 | }) 24 | 25 | @OP.exec_sign_check 26 | def execute( 27 | self, 28 | op_in: OPIO, 29 | ) -> OPIO: 30 | with open(op_in['foo'], "r") as f: 31 | content = f.read() 32 | with open("output.txt", "w") as f: 33 | f.write(content*2) 34 | self.register_output_artifact( 35 | "bar", 36 | namespace="tiefblue", 37 | dataset_name="dflow-test-%s" % datetime.datetime.now().strftime( 38 | "%Y%m%d%H%M%S"), 39 | description="Output artifact", 40 | tags=["test"], 41 | properties={"length": "4"}, 42 | ) 43 | return OPIO({ 44 | "bar": Path("output.txt"), 45 | }) 46 | 47 | 48 | if __name__ == "__main__": 49 | config["lineage"] = MetadataClient( 50 | project="", 51 | token="", 52 | ) 53 | wf = Workflow(name="metadata") 54 | 55 | with open("foo.txt", "w") as f: 56 | f.write("Hi") 57 | 58 | # artifact0 = S3Artifact(urn="") 59 | artifact0 = upload_artifact( 60 | "foo.txt", 61 | namespace="tiefblue", 62 | dataset_name="dflow-test-%s" % datetime.datetime.now().strftime( 63 | "%Y%m%d%H%M%S"), 64 | description="Uploaded artifact", 65 | tags=["test"], 66 | properties={"length": "2"}) 67 | print(artifact0.urn) 68 | step = Step( 69 | name="step", 70 | template=PythonOPTemplate( 71 | Duplicate, image="registry.dp.tech/dptech/dp-metadata-sdk:latest"), 72 | artifacts={"foo": artifact0}, 73 | ) 74 | wf.add(step) 75 | wf.submit() 76 | -------------------------------------------------------------------------------- /examples/test_op_parameter.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import Step, Workflow 4 | from dflow.python import (OP, OPIO, OPIOSign, Parameter, PythonOPTemplate, 5 | upload_packages) 6 | 7 | if "__file__" in locals(): 8 | upload_packages.append(__file__) 9 | 10 | 11 | class Hello(OP): 12 | def __init__(self): 13 | pass 14 | 15 | @classmethod 16 | def get_input_sign(cls): 17 | return OPIOSign({ 18 | 'msg': Parameter(str, default="hello"), 19 | 'num': Parameter(int, default=5), 20 | }) 21 | 22 | @classmethod 23 | def get_output_sign(cls): 24 | return OPIOSign({ 25 | 'msg': Parameter(str, default="Hello dflow!"), 26 | }) 27 | 28 | @OP.exec_sign_check 29 | def execute( 30 | self, 31 | op_in: OPIO, 32 | ) -> OPIO: 33 | print(op_in["msg"] * op_in["num"]) 34 | return OPIO() 35 | 36 | 37 | def test_op_parameter(): 38 | wf = Workflow(name="op-parameter") 39 | 40 | step = Step( 41 | name="step", 42 | template=PythonOPTemplate(Hello, image="python:3.8") 43 | ) 44 | wf.add(step) 45 | wf.submit() 46 | while wf.query_status() in ["Pending", "Running"]: 47 | time.sleep(1) 48 | 49 | assert(wf.query_status() == "Succeeded") 50 | 51 | 52 | if __name__ == "__main__": 53 | test_op_parameter() 54 | -------------------------------------------------------------------------------- /examples/test_progress.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import Step, Workflow 4 | from dflow.python import OP, OPIO, OPIOSign, PythonOPTemplate 5 | 6 | 7 | class Progress(OP): 8 | progress_total = 100 9 | 10 | @classmethod 11 | def get_input_sign(cls): 12 | return OPIOSign() 13 | 14 | @classmethod 15 | def get_output_sign(cls): 16 | return OPIOSign() 17 | 18 | @OP.exec_sign_check 19 | def execute( 20 | self, 21 | op_in: OPIO, 22 | ) -> OPIO: 23 | for i in range(10): 24 | time.sleep(10) 25 | self.progress_current = 10 * (i + 1) 26 | return OPIO() 27 | 28 | 29 | if __name__ == "__main__": 30 | wf = Workflow(name="progress") 31 | step = Step(name="step", template=PythonOPTemplate(Progress, 32 | image="python:3.8")) 33 | # This step will report progress every 10 seconds 34 | wf.add(step) 35 | wf.submit() 36 | 37 | while wf.query_status() in ["Pending", "Running"]: 38 | time.sleep(1) 39 | step = wf.query_step(name="step") 40 | if len(step) > 0 and hasattr(step[0], "progress"): 41 | print(step[0].progress) 42 | 43 | assert(wf.query_status() == "Succeeded") 44 | step = wf.query_step(name="step")[0] 45 | assert(step.progress == "100/100") 46 | -------------------------------------------------------------------------------- /examples/test_python.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow, download_artifact, upload_artifact 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Duplicate(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | 'msg': str, 20 | 'num': int, 21 | 'foo': Artifact(Path), 22 | 'idir': Artifact(Path), 23 | }) 24 | 25 | @classmethod 26 | def get_output_sign(cls): 27 | return OPIOSign({ 28 | 'msg': str, 29 | 'bar': Artifact(Path), 30 | 'odir': Artifact(Path), 31 | }) 32 | 33 | @OP.exec_sign_check 34 | def execute( 35 | self, 36 | op_in: OPIO, 37 | ) -> OPIO: 38 | op_out = OPIO({ 39 | "msg": op_in["msg"] * op_in["num"], 40 | "bar": Path("output.txt"), 41 | "odir": Path("todir"), 42 | }) 43 | 44 | content = open(op_in['foo'], "r").read() 45 | open("output.txt", "w").write(content * op_in["num"]) 46 | 47 | Path(op_out['odir']).mkdir() 48 | for ii in ['f1', 'f2']: 49 | (op_out['odir']/ii).write_text(op_in['num'] 50 | * (op_in['idir']/ii).read_text()) 51 | 52 | return op_out 53 | 54 | 55 | def make_idir(): 56 | idir = Path("tidir") 57 | idir.mkdir(exist_ok=True) 58 | (idir / "f1").write_text("foo") 59 | (idir / "f2").write_text("bar") 60 | 61 | 62 | def test_python(): 63 | wf = Workflow(name="python") 64 | 65 | with open("foo.txt", "w") as f: 66 | f.write("Hi") 67 | make_idir() 68 | 69 | artifact0 = upload_artifact("foo.txt") 70 | artifact1 = upload_artifact("tidir") 71 | print(artifact0) 72 | print(artifact1) 73 | step = Step( 74 | name="step", 75 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 76 | parameters={"msg": "Hello", "num": 3}, 77 | artifacts={"foo": artifact0, "idir": artifact1}, 78 | ) 79 | # This step will give output parameter "msg" with value "HelloHelloHello", 80 | # and output artifact "bar" which contains "HiHiHi" 81 | # output artifact "odir/f1" contains foofoofoo and "odir/f2" contains 82 | # barbarbar 83 | wf.add(step) 84 | wf.submit() 85 | 86 | while wf.query_status() in ["Pending", "Running"]: 87 | time.sleep(1) 88 | 89 | assert(wf.query_status() == "Succeeded") 90 | step = wf.query_step(name="step")[0] 91 | assert(step.phase == "Succeeded") 92 | 93 | print(download_artifact(step.outputs.artifacts["bar"])) 94 | print(download_artifact(step.outputs.artifacts["odir"])) 95 | 96 | 97 | if __name__ == "__main__": 98 | test_python() 99 | -------------------------------------------------------------------------------- /examples/test_ray.py: -------------------------------------------------------------------------------- 1 | import platform 2 | import time 3 | from pathlib import Path 4 | 5 | from dflow import Step, Workflow, download_artifact 6 | from dflow.plugins.ray import RayClusterExecutor 7 | from dflow.python import (OP, OPIO, OPIOSign, PythonOPTemplate, 8 | upload_packages, Artifact) 9 | 10 | if '__file__' in locals(): 11 | upload_packages.append(__file__) 12 | 13 | # You need ray to define remote functions 14 | import ray 15 | 16 | 17 | @ray.remote 18 | def return_value(input_value): 19 | time.sleep(1) 20 | return {'input_value': input_value, 'node_name': platform.node()} 21 | 22 | 23 | @ray.remote 24 | def add_one(inputs): 25 | input_value = inputs['input_value'] 26 | node_from = inputs["node_name"] 27 | time.sleep(1) 28 | return { 29 | 'result': f"From {node_from} " 30 | f"with value {str(input_value)} to {platform.node()}", 31 | 'node_name': platform.node()} 32 | 33 | 34 | class Duplicate(OP): 35 | def __init__(self): 36 | resources = ray.cluster_resources() 37 | node_keys = [key for key in resources if 'node' in key] 38 | num_nodes = sum(resources[node_key] for node_key in node_keys) 39 | print('num of nodes in cluster:', num_nodes) 40 | 41 | @classmethod 42 | def get_input_sign(cls): 43 | return OPIOSign({ 44 | 'msg': str, 45 | 'num': int 46 | }) 47 | 48 | @classmethod 49 | def get_output_sign(cls): 50 | return OPIOSign({ 51 | 'nodes': Artifact(Path), 52 | 'results': Artifact(Path) 53 | }) 54 | 55 | @OP.exec_sign_check 56 | def execute( 57 | self, 58 | op_in: OPIO, 59 | ) -> OPIO: 60 | # Ray need list generators to get ObjectRef and dispatch tasks. 61 | remote_fun_return_true = [add_one.remote(return_value.remote(i)) for i 62 | in 63 | range(op_in['num'])] 64 | op_out = OPIO({ 65 | 'nodes': Path('nodes.txt'), 66 | 'results': Path('results.txt') 67 | }) 68 | # ray.get() call on remote obj to gather results. 69 | results = ray.get(remote_fun_return_true) 70 | for result in results: 71 | op_out['nodes'].open('a').write(str(result["node_name"]) + "\n") 72 | op_out['results'].open('a').write(result["result"] + "\n") 73 | 74 | return op_out 75 | 76 | 77 | def run_ray(): 78 | wf = Workflow(name='ray-test') 79 | 80 | # 1. To run a workflow 81 | # !!! Please make sure you set the correct ray header url as `ray_host`. 82 | # !!! The ray_host could be connected INSIDE the kubernetes argoflow on. 83 | # (often 10.*.*.*:10001) 84 | 85 | # 2. choose an image 86 | # RayClusterExecutor will exam your image, if it has no ray package with 87 | # default python (make sure you are not working on some virtual 88 | # environmental by pip or conda), init container will try to install 89 | # ray with `pip install ray`. 90 | 91 | # 3. set up mirror if install is needed (optional) 92 | # For users with special package installation settings, 93 | # such as private package servers, 94 | # please set ray_install_mirror to your package server mirror. 95 | raycluster = RayClusterExecutor( 96 | ray_host='ray://【】【】【】【】【】:10001', # !!! change this 97 | ray_install_mirror='https://pypi.tuna.tsinghua.edu.cn/simple' 98 | # comment this if you don't need 99 | ) 100 | 101 | # !!! change this to same python minor version as your ray cluster. 102 | IMAGE_NAME = 'python:3.9' 103 | 104 | step = Step( 105 | name='step', 106 | template=PythonOPTemplate(Duplicate, image=IMAGE_NAME), 107 | parameters={'msg': 'Hello', 'num': 10}, 108 | executor=raycluster 109 | ) 110 | wf.add(step) 111 | wf.submit() 112 | while wf.query_status() in ['Pending', 'Running']: 113 | time.sleep(1) 114 | assert (wf.query_status() == 'Succeeded') 115 | step = wf.query_step(name='step')[0] 116 | assert (step.phase == 'Succeeded') 117 | download_artifact(step.outputs.artifacts['nodes']) 118 | download_artifact(step.outputs.artifacts['results']) 119 | 120 | 121 | if __name__ == '__main__': 122 | run_ray() 123 | -------------------------------------------------------------------------------- /examples/test_ray_after_ray.py: -------------------------------------------------------------------------------- 1 | import platform 2 | import time 3 | from pathlib import Path 4 | 5 | from dflow import Step, Workflow 6 | from dflow.plugins.ray import RayClusterExecutor 7 | from dflow.python import (OP, OPIO, OPIOSign, PythonOPTemplate, 8 | upload_packages, Artifact) 9 | 10 | if '__file__' in locals(): 11 | upload_packages.append(__file__) 12 | 13 | # You need ray to define remote functions 14 | import ray 15 | 16 | 17 | @ray.remote 18 | def return_value(input_value): 19 | time.sleep(1) 20 | return {'input_value': input_value, 'node_name': platform.node()} 21 | 22 | 23 | @ray.remote 24 | def add_one(inputs): 25 | input_value = inputs['input_value'] 26 | node_from = inputs["node_name"] 27 | time.sleep(1) 28 | return { 29 | 'result': f"From {node_from} " 30 | f"with value {str(input_value)} to {platform.node()}", 31 | 'node_name': platform.node()} 32 | 33 | 34 | @ray.remote 35 | def count_work_time(one_result): 36 | if len(one_result) > 0: 37 | result_list = one_result.split() 38 | from_node = result_list[1] 39 | to_node = result_list[-1] 40 | return { 41 | "from_node": from_node, 42 | "to_node": to_node, 43 | "on_node": platform.node() 44 | } 45 | else: 46 | return { 47 | "on_node": platform.node() 48 | } 49 | 50 | 51 | class Sleeps(OP): 52 | def __init__(self): 53 | resources = ray.cluster_resources() 54 | node_keys = [key for key in resources if 'node' in key] 55 | num_nodes = sum(resources[node_key] for node_key in node_keys) 56 | print('num of nodes in cluster:', num_nodes) 57 | 58 | @classmethod 59 | def get_input_sign(cls): 60 | return OPIOSign({ 61 | 'msg': str, 62 | 'num': int 63 | }) 64 | 65 | @classmethod 66 | def get_output_sign(cls): 67 | return OPIOSign({ 68 | 'nodes': Artifact(Path), 69 | 'results': Artifact(Path) 70 | }) 71 | 72 | @OP.exec_sign_check 73 | def execute( 74 | self, 75 | op_in: OPIO, 76 | ) -> OPIO: 77 | # Ray need list generators to get ObjectRef and dispatch tasks. 78 | remote_fun_return_true = [add_one.remote(return_value.remote(i)) for i 79 | in 80 | range(op_in['num'])] 81 | op_out = OPIO({ 82 | 'nodes': Path('nodes.txt'), 83 | 'results': Path('results.txt') 84 | }) 85 | # ray.get() call on remote obj to gather results. 86 | results = ray.get(remote_fun_return_true) 87 | for result in results: 88 | op_out['nodes'].open('a').write(str(result["node_name"]) + "\n") 89 | op_out['results'].open('a').write(result["result"] + "\n") 90 | 91 | return op_out 92 | 93 | 94 | class CountTheQuotes(OP): 95 | def __init__(self): 96 | resources = ray.cluster_resources() 97 | node_keys = [key for key in resources if 'node' in key] 98 | num_nodes = sum(resources[node_key] for node_key in node_keys) 99 | print('num of nodes in cluster:', num_nodes) 100 | 101 | @classmethod 102 | def get_input_sign(cls): 103 | return OPIOSign({ 104 | 'results': Artifact(Path), 105 | }) 106 | 107 | @classmethod 108 | def get_output_sign(cls): 109 | return OPIOSign({ 110 | "count": list 111 | }) 112 | 113 | @OP.exec_sign_check 114 | def execute( 115 | self, 116 | op_in: OPIO, 117 | ) -> OPIO: 118 | results = op_in["results"].read_text().split("\n") 119 | # Ray need list generators to get ObjectRef and dispatch tasks. 120 | remote_fun_call = [count_work_time.remote(result_one) 121 | for result_one in results] 122 | op_out = OPIO({ 123 | 'count': ray.get(remote_fun_call) 124 | }) 125 | return op_out 126 | 127 | 128 | def run_ray(): 129 | wf = Workflow(name='ray-after-ray-test') 130 | 131 | # 1. To run a workflow 132 | # !!! Please make sure you set the correct ray header url as `ray_host`. 133 | # !!! The ray_host could be connected INSIDE the kubernetes argoflow on. 134 | # (often 10.*.*.*:10001) 135 | 136 | # 2. choose an image 137 | # RayClusterExecutor will exam your image, if it has no ray package with 138 | # default python (make sure you are not working on some virtual 139 | # environmental by pip or conda), init container will try to install 140 | # ray with `pip install ray`. 141 | 142 | # 3. set up mirror if install is needed (optional) 143 | # For users with special package installation settings, 144 | # such as private package servers, 145 | # please set ray_install_mirror to your package server mirror. 146 | raycluster = RayClusterExecutor( 147 | ray_host='ray://【】【】【】【】【】【】:10001', # !!! change this 148 | ray_install_mirror='https://pypi.tuna.tsinghua.edu.cn/simple' 149 | # comment this if you don't need 150 | ) 151 | 152 | # !!! change this to same python minor version as your ray cluster. 153 | IMAGE_NAME = 'python:3.9' 154 | 155 | step1 = Step( 156 | name='Sleeps', 157 | template=PythonOPTemplate(Sleeps, image=IMAGE_NAME), 158 | parameters={'msg': 'Hello', 'num': 10}, 159 | executor=raycluster 160 | ) 161 | wf.add(step1) 162 | step2 = Step( 163 | name='CountTheQuotes', 164 | template=PythonOPTemplate(CountTheQuotes, image=IMAGE_NAME), 165 | artifacts={'results': step1.outputs.artifacts["results"]}, 166 | executor=raycluster 167 | ) 168 | wf.add(step2) 169 | wf.submit() 170 | 171 | 172 | if __name__ == '__main__': 173 | run_ray() 174 | -------------------------------------------------------------------------------- /examples/test_recurse.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import (InputParameter, Inputs, OutputParameter, ShellOPTemplate, 4 | Step, Steps, Workflow) 5 | 6 | 7 | def test_recurse(): 8 | plus1 = ShellOPTemplate( 9 | name='plus1', 10 | image="alpine:3.15", 11 | script="echo 'This is iter {{inputs.parameters.iter}}' && " 12 | "echo $(({{inputs.parameters.iter}}+1)) > /tmp/result.txt") 13 | plus1.inputs.parameters = {"iter": InputParameter()} 14 | plus1.outputs.parameters = {"iter": OutputParameter( 15 | value_from_path="/tmp/result.txt")} 16 | 17 | steps = Steps(name="iter", inputs=Inputs( 18 | parameters={"iter": InputParameter(value=0), 19 | "limit": InputParameter(value=3)})) 20 | hello = Step(name="hello", template=plus1, parameters={ 21 | "iter": steps.inputs.parameters["iter"]}) 22 | steps.add(hello) 23 | next = Step(name="next", template=steps, 24 | parameters={"iter": hello.outputs.parameters["iter"]}, 25 | when="%s < %s" % ( 26 | hello.outputs.parameters["iter"], 27 | steps.inputs.parameters["limit"])) 28 | # This step use steps as its template (note that Steps is a subclass of 29 | # OPTemplate), meanwhile the steps it used contains this step, 30 | # which gives a recursion. The recursion will stop when the "when" 31 | # condition is satisfied (after 10 loops in this example) 32 | steps.add(next) 33 | 34 | wf = Workflow("recurse", steps=steps) 35 | wf.submit() 36 | 37 | while wf.query_status() in ["Pending", "Running"]: 38 | time.sleep(1) 39 | 40 | assert(wf.query_status() == "Succeeded") 41 | 42 | 43 | if __name__ == "__main__": 44 | test_recurse() 45 | -------------------------------------------------------------------------------- /examples/test_reuse.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import InputParameter, Inputs, Step, Steps, Workflow 4 | from dflow.python import OP, OPIO, OPIOSign, PythonOPTemplate 5 | 6 | 7 | class Plus1(OP): 8 | def __init__(self): 9 | pass 10 | 11 | @classmethod 12 | def get_input_sign(cls): 13 | return OPIOSign({ 14 | 'iter': int 15 | }) 16 | 17 | @classmethod 18 | def get_output_sign(cls): 19 | return OPIOSign({ 20 | 'iter': int 21 | }) 22 | 23 | @OP.exec_sign_check 24 | def execute( 25 | self, 26 | op_in: OPIO, 27 | ) -> OPIO: 28 | return OPIO({ 29 | 'iter': op_in['iter'] + 1 30 | }) 31 | 32 | 33 | if __name__ == "__main__": 34 | steps = Steps(name="iter", inputs=Inputs( 35 | parameters={"iter": InputParameter(value=0), 36 | "limit": InputParameter(value=5)})) 37 | plus1 = Step(name="plus1", 38 | template=PythonOPTemplate(Plus1, 39 | image="python:3.8"), 40 | parameters={"iter": steps.inputs.parameters["iter"]}, 41 | key="iter-%s" % steps.inputs.parameters["iter"]) 42 | steps.add(plus1) 43 | next = Step(name="next", template=steps, 44 | parameters={"iter": plus1.outputs.parameters["iter"]}, 45 | when="%s < %s" % ( 46 | plus1.outputs.parameters["iter"], 47 | steps.inputs.parameters["limit"])) 48 | steps.add(next) 49 | 50 | wf = Workflow("recurse", steps=steps) 51 | wf.submit() 52 | 53 | while wf.query_status() in ["Pending", "Running"]: 54 | time.sleep(1) 55 | 56 | assert(wf.query_status() == "Succeeded") 57 | 58 | step0 = wf.query_step(key="iter-0")[0] 59 | step1 = wf.query_step(key="iter-1")[0] 60 | step1.modify_output_parameter("iter", 3) 61 | 62 | wf = Workflow("recurse-resubmit", steps=steps) 63 | wf.submit(reuse_step=[step0, step1]) 64 | -------------------------------------------------------------------------------- /examples/test_slices.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from dflow import Step, Workflow, argo_range 4 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, Slices 5 | 6 | 7 | class Hello(OP): 8 | def __init__(self): 9 | pass 10 | 11 | @classmethod 12 | def get_input_sign(cls): 13 | return OPIOSign({ 14 | 'filename': str 15 | }) 16 | 17 | @classmethod 18 | def get_output_sign(cls): 19 | return OPIOSign({ 20 | 'foo': Artifact(str) 21 | }) 22 | 23 | @OP.exec_sign_check 24 | def execute( 25 | self, 26 | op_in: OPIO, 27 | ) -> OPIO: 28 | open(op_in["filename"], "w").write("foo") 29 | op_out = OPIO({ 30 | 'foo': op_in["filename"] 31 | }) 32 | return op_out 33 | 34 | 35 | class Check(OP): 36 | def __init__(self): 37 | pass 38 | 39 | @classmethod 40 | def get_input_sign(cls): 41 | return OPIOSign({ 42 | 'foo': Artifact(List[str]) 43 | }) 44 | 45 | @classmethod 46 | def get_output_sign(cls): 47 | return OPIOSign() 48 | 49 | @OP.exec_sign_check 50 | def execute( 51 | self, 52 | op_in: OPIO, 53 | ) -> OPIO: 54 | print(op_in["foo"]) 55 | return OPIO() 56 | 57 | 58 | if __name__ == "__main__": 59 | wf = Workflow("slices") 60 | hello = Step("hello", 61 | PythonOPTemplate(Hello, image="python:3.8", 62 | slices=Slices("{{item}}", 63 | input_parameter=["filename"], 64 | output_artifact=["foo"] 65 | ) 66 | ), 67 | parameters={"filename": ["f1.txt", "f2.txt"]}, 68 | with_param=argo_range(2)) 69 | wf.add(hello) 70 | check = Step("check", 71 | PythonOPTemplate(Check, image="python:3.8"), 72 | artifacts={"foo": hello.outputs.artifacts["foo"]}) 73 | wf.add(check) 74 | wf.submit() 75 | -------------------------------------------------------------------------------- /examples/test_slurm.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from dflow import SlurmRemoteExecutor, Step, Workflow, argo_range 4 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, Slices 5 | 6 | 7 | class Hello(OP): 8 | def __init__(self): 9 | pass 10 | 11 | @classmethod 12 | def get_input_sign(cls): 13 | return OPIOSign({ 14 | 'filename': str 15 | }) 16 | 17 | @classmethod 18 | def get_output_sign(cls): 19 | return OPIOSign({ 20 | 'foo': Artifact(str) 21 | }) 22 | 23 | @OP.exec_sign_check 24 | def execute( 25 | self, 26 | op_in: OPIO, 27 | ) -> OPIO: 28 | open(op_in["filename"], "w").write("foo") 29 | op_out = OPIO({ 30 | 'foo': op_in["filename"] 31 | }) 32 | return op_out 33 | 34 | 35 | class Check(OP): 36 | def __init__(self): 37 | pass 38 | 39 | @classmethod 40 | def get_input_sign(cls): 41 | return OPIOSign({ 42 | 'foo': Artifact(List[str]) 43 | }) 44 | 45 | @classmethod 46 | def get_output_sign(cls): 47 | return OPIOSign() 48 | 49 | @OP.exec_sign_check 50 | def execute( 51 | self, 52 | op_in: OPIO, 53 | ) -> OPIO: 54 | print(op_in["foo"]) 55 | for filename in op_in["foo"]: 56 | with open(filename, "r") as f: 57 | print(f.read()) 58 | return OPIO() 59 | 60 | 61 | if __name__ == "__main__": 62 | # run ../scripts/start-slurm.sh first to start up a slurm cluster 63 | import socket 64 | 65 | def get_my_ip_address(): 66 | with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s: 67 | s.connect(("8.8.8.8", 80)) 68 | return s.getsockname()[0] 69 | 70 | slurm_remote_executor = SlurmRemoteExecutor( 71 | host=get_my_ip_address(), 72 | port=31129, 73 | username="root", 74 | password="password", 75 | header="#!/bin/bash\n#SBATCH -N 1\n#SBATCH -n 1\n", 76 | workdir="/data/dflow/workflows/{{workflow.name}}/{{pod.name}}", 77 | ) 78 | 79 | wf = Workflow("slurm-slices") 80 | hello = Step("hello", 81 | PythonOPTemplate(Hello, 82 | slices=Slices("{{item}}", 83 | input_parameter=["filename"], 84 | output_artifact=["foo"] 85 | ) 86 | ), 87 | parameters={"filename": ["f1.txt", "f2.txt"]}, 88 | with_param=argo_range(2), 89 | key="hello-{{item}}", 90 | executor=slurm_remote_executor) 91 | wf.add(hello) 92 | check = Step("check", 93 | PythonOPTemplate(Check, image="python:3.8"), 94 | artifacts={"foo": hello.outputs.artifacts["foo"]}) 95 | wf.add(check) 96 | wf.submit() 97 | -------------------------------------------------------------------------------- /examples/test_steps.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import (InputArtifact, InputParameter, OutputArtifact, 4 | OutputParameter, ShellOPTemplate, Step, Workflow, 5 | download_artifact) 6 | 7 | 8 | def test_steps(): 9 | hello = ShellOPTemplate( 10 | name='Hello', 11 | image="alpine:3.15", 12 | script="echo Hello > /tmp/bar.txt && echo 1 > /tmp/result.txt") 13 | hello.outputs.parameters = {"msg": OutputParameter( 14 | value_from_path="/tmp/result.txt")} 15 | hello.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")} 16 | 17 | duplicate = ShellOPTemplate( 18 | name='Duplicate', 19 | image="alpine:3.15", 20 | script="cat /tmp/foo.txt /tmp/foo.txt > /tmp/bar.txt && " 21 | "echo $(({{inputs.parameters.msg}}*2)) > /tmp/result.txt") 22 | duplicate.inputs.parameters = {"msg": InputParameter()} 23 | duplicate.outputs.parameters = { 24 | "msg": OutputParameter(value_from_path="/tmp/result.txt")} 25 | duplicate.inputs.artifacts = {"foo": InputArtifact(path="/tmp/foo.txt")} 26 | duplicate.outputs.artifacts = {"bar": OutputArtifact(path="/tmp/bar.txt")} 27 | 28 | wf = Workflow(name="steps") 29 | hello0 = Step(name="hello0", template=hello) 30 | # This step will give output parameter "msg" with value 1, and output 31 | # artifact "bar" which contains "Hello" 32 | wf.add(hello0) 33 | # This step use the output parameter "msg" of step "hello0" as input 34 | # parameter "msg", and the output artifact "bar" of step "hello0" as input 35 | # artifact "foo" 36 | hello1 = Step(name="hello1", template=duplicate, parameters={ 37 | "msg": hello0.outputs.parameters["msg"]}, 38 | artifacts={"foo": hello0.outputs.artifacts["bar"]}) 39 | # This step will give output parameter "msg" with value 2, and output 40 | # artifact "bar" which contains "HelloHello" 41 | wf.add(hello1) 42 | hello2 = Step(name="hello2", template=duplicate, parameters={ 43 | "msg": hello1.outputs.parameters["msg"]}, 44 | artifacts={"foo": hello1.outputs.artifacts["bar"]}) 45 | hello3 = Step(name="hello3", template=duplicate, parameters={ 46 | "msg": hello1.outputs.parameters["msg"]}, 47 | artifacts={"foo": hello1.outputs.artifacts["bar"]}) 48 | wf.add([hello2, hello3]) # parallel steps 49 | # These steps will give output parameter "msg" with value 4, and output 50 | # artifact "bar" which contains "HelloHelloHelloHello" 51 | wf.submit() 52 | 53 | while wf.query_status() in ["Pending", "Running"]: 54 | time.sleep(1) 55 | 56 | assert(wf.query_status() == "Succeeded") 57 | step = wf.query_step(name="hello3")[0] 58 | assert(step.phase == "Succeeded") 59 | assert(step.outputs.parameters["msg"].value == "4") 60 | download_artifact(step.outputs.artifacts["bar"]) 61 | assert(open("bar.txt", "r").read() == "Hello\nHello\nHello\nHello\n") 62 | 63 | 64 | if __name__ == "__main__": 65 | test_steps() 66 | -------------------------------------------------------------------------------- /examples/test_subpath_slices.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import Step, Workflow # , upload_artifact 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Prepare(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | }) 20 | 21 | @classmethod 22 | def get_output_sign(cls): 23 | return OPIOSign({ 24 | 'foo': Artifact(List[str], archive=None) 25 | }) 26 | 27 | @OP.exec_sign_check 28 | def execute( 29 | self, 30 | op_in: OPIO, 31 | ) -> OPIO: 32 | with open("foo1.txt", "w") as f: 33 | f.write("foo1") 34 | with open("foo2.txt", "w") as f: 35 | f.write("foo2") 36 | op_out = OPIO({ 37 | 'foo': ["foo1.txt", "foo2.txt"] 38 | }) 39 | return op_out 40 | 41 | 42 | class Hello(OP): 43 | def __init__(self): 44 | pass 45 | 46 | @classmethod 47 | def get_input_sign(cls): 48 | return OPIOSign({ 49 | 'foo': Artifact(str) 50 | }) 51 | 52 | @classmethod 53 | def get_output_sign(cls): 54 | return OPIOSign() 55 | 56 | @OP.exec_sign_check 57 | def execute( 58 | self, 59 | op_in: OPIO, 60 | ) -> OPIO: 61 | with open(op_in["foo"], "r") as f: 62 | print(f.read()) 63 | return OPIO() 64 | 65 | 66 | def test_subpath_slices(): 67 | wf = Workflow("subpath-slices") 68 | # with open("foo1.txt", "w") as f: 69 | # f.write("foo1") 70 | # with open("foo2.txt", "w") as f: 71 | # f.write("foo2") 72 | # artifact = upload_artifact(["foo1.txt", "foo2.txt"], archive=None) 73 | prepare = Step("prepare", 74 | PythonOPTemplate(Prepare, image="python:3.8")) 75 | wf.add(prepare) 76 | 77 | hello = Step("hello", 78 | PythonOPTemplate(Hello, image="python:3.8", 79 | slices=Slices(sub_path=True, 80 | input_artifact=["foo"] 81 | ) 82 | ), 83 | artifacts={"foo": prepare.outputs.artifacts["foo"]}) 84 | # artifacts={"foo": artifact}) 85 | wf.add(hello) 86 | wf.submit() 87 | while wf.query_status() in ["Pending", "Running"]: 88 | time.sleep(1) 89 | 90 | assert(wf.query_status() == "Succeeded") 91 | 92 | 93 | if __name__ == "__main__": 94 | test_subpath_slices() 95 | -------------------------------------------------------------------------------- /examples/test_success_ratio.py: -------------------------------------------------------------------------------- 1 | import random 2 | 3 | from dflow import Step, Workflow, argo_range 4 | from dflow.python import OP, OPIO, OPIOSign, PythonOPTemplate, TransientError 5 | 6 | 7 | class Hello(OP): 8 | @classmethod 9 | def get_input_sign(cls): 10 | return OPIOSign() 11 | 12 | @classmethod 13 | def get_output_sign(cls): 14 | return OPIOSign() 15 | 16 | @OP.exec_sign_check 17 | def execute( 18 | self, 19 | op_in: OPIO, 20 | ) -> OPIO: 21 | if random.random() < 0.5: 22 | raise TransientError("Hello") 23 | return OPIO() 24 | 25 | 26 | if __name__ == "__main__": 27 | wf = Workflow(name="success-ratio") 28 | 29 | step = Step( 30 | name="hello0", 31 | template=PythonOPTemplate(Hello, image="python:3.8"), 32 | continue_on_success_ratio=0.6, 33 | # continue_on_num_success=3, 34 | with_param=argo_range(5) 35 | ) 36 | wf.add(step) 37 | wf.submit() 38 | -------------------------------------------------------------------------------- /examples/test_utils.py: -------------------------------------------------------------------------------- 1 | import os 2 | import shutil 3 | from pathlib import Path 4 | 5 | import pytest 6 | from dflow import copy_artifact, upload_artifact 7 | from dflow.utils import catalog_of_artifact, run_command, set_directory 8 | 9 | 10 | def test_set_directory(): 11 | pwd = Path.cwd() 12 | with set_directory("test_dir", mkdir=True) as wdir: 13 | assert str(wdir) == str(pwd / "test_dir") 14 | assert os.getcwd() == str(wdir) 15 | shutil.rmtree(wdir) 16 | 17 | 18 | def test_run_command(): 19 | code, out, err = run_command("echo test") 20 | assert code == 0 21 | assert out == "test\n" 22 | 23 | 24 | def test_run_command_err(): 25 | with pytest.raises(AssertionError): 26 | run_command(["python", "-c", "raise ValueError('error')"]) 27 | 28 | 29 | def test_run_command_input(): 30 | code, out, err = run_command(["sh"], input="echo test\nexit") 31 | assert code == 0 32 | assert out == "test\n" 33 | 34 | 35 | def test_copy_artifact(): 36 | with open("foo.txt", "w"): 37 | pass 38 | with open("bar.txt", "w"): 39 | pass 40 | art_1 = upload_artifact(["foo.txt"], archive=None) 41 | art_2 = upload_artifact(["bar.txt"], archive=None) 42 | copy_artifact(art_1, art_2, sort=True) 43 | catalog = catalog_of_artifact(art_2) 44 | catalog.sort(key=lambda x: x["order"]) 45 | assert catalog == [{'dflow_list_item': 'bar.txt', 'order': 0}, 46 | {'dflow_list_item': 'foo.txt', 'order': 1}] 47 | 48 | art_1 = upload_artifact(["foo.txt"], archive=None) 49 | art_2 = upload_artifact(["bar.txt"], archive=None) 50 | copy_artifact(art_2, art_1, sort=True) 51 | catalog = catalog_of_artifact(art_1) 52 | catalog.sort(key=lambda x: x["order"]) 53 | assert catalog == [{'dflow_list_item': 'foo.txt', 'order': 0}, 54 | {'dflow_list_item': 'bar.txt', 'order': 1}] 55 | 56 | os.remove("foo.txt") 57 | os.remove("bar.txt") 58 | -------------------------------------------------------------------------------- /examples/test_wlm.py: -------------------------------------------------------------------------------- 1 | from typing import List 2 | 3 | from dflow import SlurmJobTemplate, Step, Workflow, argo_range 4 | from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, Slices 5 | 6 | 7 | class Hello(OP): 8 | def __init__(self): 9 | pass 10 | 11 | @classmethod 12 | def get_input_sign(cls): 13 | return OPIOSign({ 14 | 'filename': str 15 | }) 16 | 17 | @classmethod 18 | def get_output_sign(cls): 19 | return OPIOSign({ 20 | 'foo': Artifact(str) 21 | }) 22 | 23 | @OP.exec_sign_check 24 | def execute( 25 | self, 26 | op_in: OPIO, 27 | ) -> OPIO: 28 | open(op_in["filename"], "w").write("foo") 29 | op_out = OPIO({ 30 | 'foo': op_in["filename"] 31 | }) 32 | return op_out 33 | 34 | 35 | class Check(OP): 36 | def __init__(self): 37 | pass 38 | 39 | @classmethod 40 | def get_input_sign(cls): 41 | return OPIOSign({ 42 | 'foo': Artifact(List[str]) 43 | }) 44 | 45 | @classmethod 46 | def get_output_sign(cls): 47 | return OPIOSign() 48 | 49 | @OP.exec_sign_check 50 | def execute( 51 | self, 52 | op_in: OPIO, 53 | ) -> OPIO: 54 | print(op_in["foo"]) 55 | for filename in op_in["foo"]: 56 | with open(filename, "r") as f: 57 | print(f.read()) 58 | return OPIO() 59 | 60 | 61 | if __name__ == "__main__": 62 | wf = Workflow(name="wlm") 63 | 64 | hello = Step("hello", 65 | PythonOPTemplate(Hello, image="python:3.8", 66 | slices=Slices("{{item}}", 67 | input_parameter=["filename"], 68 | output_artifact=["foo"] 69 | ) 70 | ), 71 | parameters={"filename": ["f1.txt", "f2.txt"]}, 72 | with_param=argo_range(2), 73 | key="hello-{{item}}", 74 | executor=SlurmJobTemplate( 75 | header="#!/bin/sh\n#SBATCH --nodes=1", 76 | node_selector={ 77 | "kubernetes.io/hostname": "slurm-minikube-v100"})) 78 | wf.add(hello) 79 | check = Step("check", 80 | PythonOPTemplate(Check, image="python:3.8"), 81 | artifacts={"foo": hello.outputs.artifacts["foo"]}) 82 | wf.add(check) 83 | wf.submit() 84 | -------------------------------------------------------------------------------- /manifests/configurator.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: rbac.authorization.k8s.io/v1 2 | kind: ClusterRole 3 | metadata: 4 | name: virtual-kubelet 5 | rules: 6 | - apiGroups: 7 | - "" 8 | resources: 9 | - configmaps 10 | - secrets 11 | - services 12 | verbs: 13 | - get 14 | - list 15 | - watch 16 | - apiGroups: 17 | - "" 18 | resources: 19 | - pods 20 | verbs: 21 | - delete 22 | - get 23 | - list 24 | - watch 25 | - patch 26 | - create 27 | - apiGroups: 28 | - "" 29 | resources: 30 | - nodes 31 | verbs: 32 | - create 33 | - get 34 | - patch 35 | - list 36 | - delete 37 | - apiGroups: 38 | - "" 39 | resources: 40 | - nodes/status 41 | verbs: 42 | - update 43 | - patch 44 | - apiGroups: 45 | - "" 46 | resources: 47 | - pods/status 48 | verbs: 49 | - update 50 | - patch 51 | - apiGroups: 52 | - "" 53 | resources: 54 | - events 55 | verbs: 56 | - create 57 | - patch 58 | - apiGroups: 59 | - "" 60 | - wlm.sylabs.io 61 | resources: 62 | - slurmjobs 63 | - wlmjobs 64 | verbs: 65 | - get 66 | --- 67 | apiVersion: rbac.authorization.k8s.io/v1 68 | kind: ClusterRoleBinding 69 | metadata: 70 | name: argo-virtual-kubelet 71 | subjects: 72 | - kind: ServiceAccount 73 | name: argo 74 | namespace: argo 75 | roleRef: 76 | apiGroup: rbac.authorization.k8s.io 77 | kind: ClusterRole 78 | name: virtual-kubelet 79 | --- 80 | apiVersion: apps/v1 81 | kind: DaemonSet 82 | metadata: 83 | name: configurator 84 | namespace: argo 85 | spec: 86 | selector: 87 | matchLabels: 88 | name: configurator 89 | template: 90 | metadata: 91 | labels: 92 | name: configurator 93 | spec: 94 | serviceAccountName: argo 95 | containers: 96 | - name: configurator 97 | image: dptechnology/hpc-configurator:latest 98 | imagePullPolicy: IfNotPresent 99 | args: 100 | - --sock 101 | - "/syslurm/red-box.sock" 102 | volumeMounts: 103 | - name: syslurm-mount 104 | mountPath: /syslurm 105 | env: 106 | - name: HOST_NAME 107 | valueFrom: 108 | fieldRef: 109 | fieldPath: spec.nodeName 110 | - name: SERVICE_ACCOUNT 111 | valueFrom: 112 | fieldRef: 113 | fieldPath: spec.serviceAccountName 114 | - name: KUBELET_IMAGE 115 | value: "dptechnology/hpc-vk:latest" 116 | - name: RESULTS_IMAGE 117 | value: "dptechnology/hpc-results:latest" 118 | - name: NAMESPACE 119 | value: argo 120 | volumes: 121 | - name: syslurm-mount 122 | hostPath: 123 | path: /var/run/syslurm 124 | type: Directory 125 | -------------------------------------------------------------------------------- /manifests/crds/wlm_v1alpha1_slurmjob.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: apiextensions.k8s.io/v1 2 | kind: CustomResourceDefinition 3 | metadata: 4 | creationTimestamp: null 5 | labels: 6 | controller-tools.k8s.io: "1.0" 7 | name: slurmjobs.wlm.sylabs.io 8 | spec: 9 | versions: 10 | - name: v1alpha1 11 | additionalPrinterColumns: 12 | - jsonPath: .metadata.creationTimestamp 13 | name: Age 14 | type: date 15 | - jsonPath: .status.status 16 | description: status of the kind 17 | name: Status 18 | type: string 19 | schema: 20 | openAPIV3Schema: 21 | type: object 22 | properties: 23 | apiVersion: 24 | description: 'APIVersion defines the versioned schema of this representation 25 | of an object. Servers should convert recognized schemas to the latest 26 | internal value, and may reject unrecognized values. More info: https://git.k8s.io/community/contributors/devel/api-conventions.md#resources' 27 | type: string 28 | kind: 29 | description: 'Kind is a string value representing the REST resource this 30 | object represents. Servers may infer this from the endpoint the client 31 | submits requests to. Cannot be updated. In CamelCase. More info: https://git.k8s.io/community/contributors/devel/api-conventions.md#types-kinds' 32 | type: string 33 | metadata: 34 | type: object 35 | spec: 36 | properties: 37 | batch: 38 | description: Batch is a script that will be submitted to a Slurm cluster 39 | as a batch job. 40 | minLength: 1 41 | type: string 42 | nodeSelector: 43 | description: 'NodeSelector is a selector which must be true for the 44 | SlurmJob to fit on a node. Selector which must match a node''s labels 45 | for the SlurmJob to be scheduled on that node. More info: https://kubernetes.io/docs/concepts/configuration/assign-pod-node/.' 46 | type: object 47 | x-kubernetes-map-type: atomic 48 | x-kubernetes-preserve-unknown-fields: true 49 | prepare: 50 | description: Prepare may be specified for an optional data preparation step. 51 | When specified, before job is started required data will be uploaded to Slurm 52 | cluster with respect to this configuration. 53 | properties: 54 | to: 55 | description: To is a path to the data to be uploaded to a Slurm cluster. 56 | type: string 57 | mount: 58 | description: Mount is a directory where input data will be given. 59 | type: object 60 | x-kubernetes-map-type: atomic 61 | x-kubernetes-preserve-unknown-fields: true 62 | required: 63 | - mount 64 | - to 65 | type: object 66 | results: 67 | description: Results may be specified for an optional results collection 68 | step. When specified, after job is completed all results will be downloaded 69 | from Slurm cluster with respect to this configuration. 70 | properties: 71 | from: 72 | description: From is a path to the results to be collected from 73 | a Slurm cluster. 74 | type: string 75 | mount: 76 | description: Mount is a directory where job results will be stored. 77 | After results collection all job generated files can be found 78 | in Mount/ directory. 79 | type: object 80 | x-kubernetes-map-type: atomic 81 | x-kubernetes-preserve-unknown-fields: true 82 | required: 83 | - mount 84 | - from 85 | type: object 86 | required: 87 | - batch 88 | type: object 89 | status: 90 | properties: 91 | status: 92 | description: Status reflects job status, e.g running, succeeded. 93 | type: string 94 | required: 95 | - status 96 | type: object 97 | served: true 98 | storage: true 99 | subresources: 100 | status: {} 101 | group: wlm.sylabs.io 102 | names: 103 | kind: SlurmJob 104 | plural: slurmjobs 105 | shortNames: 106 | - sj 107 | scope: Namespaced 108 | status: 109 | acceptedNames: 110 | kind: "" 111 | plural: "" 112 | conditions: [] 113 | storedVersions: [] 114 | -------------------------------------------------------------------------------- /manifests/operator-rbac.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: rbac.authorization.k8s.io/v1 2 | kind: ClusterRoleBinding 3 | metadata: 4 | name: argo-wlm-operator 5 | subjects: 6 | - kind: ServiceAccount 7 | name: argo 8 | namespace: argo 9 | roleRef: 10 | kind: ClusterRole 11 | name: wlm-operator 12 | apiGroup: rbac.authorization.k8s.io 13 | --- 14 | apiVersion: rbac.authorization.k8s.io/v1 15 | kind: ClusterRole 16 | metadata: 17 | name: wlm-operator 18 | rules: 19 | - apiGroups: 20 | - "" 21 | resources: 22 | - pods 23 | - services 24 | - endpoints 25 | - persistentvolumeclaims 26 | - events 27 | - configmaps 28 | - secrets 29 | verbs: 30 | - '*' 31 | - apiGroups: 32 | - "" 33 | resources: 34 | - namespaces 35 | verbs: 36 | - get 37 | - apiGroups: 38 | - apps 39 | resources: 40 | - deployments 41 | - daemonsets 42 | - replicasets 43 | - statefulsets 44 | verbs: 45 | - '*' 46 | - apiGroups: 47 | - monitoring.coreos.com 48 | resources: 49 | - servicemonitors 50 | verbs: 51 | - get 52 | - create 53 | - apiGroups: 54 | - wlm.sylabs.io 55 | resources: 56 | - '*' 57 | verbs: 58 | - '*' 59 | -------------------------------------------------------------------------------- /manifests/operator.yaml: -------------------------------------------------------------------------------- 1 | apiVersion: apps/v1 2 | kind: Deployment 3 | metadata: 4 | name: wlm-operator 5 | namespace: argo 6 | spec: 7 | replicas: 1 8 | selector: 9 | matchLabels: 10 | name: wlm-operator 11 | template: 12 | metadata: 13 | labels: 14 | name: wlm-operator 15 | spec: 16 | serviceAccountName: argo 17 | containers: 18 | - name: wlm-operator 19 | image: dptechnology/hpc-operator:latest 20 | imagePullPolicy: IfNotPresent 21 | env: 22 | - name: WATCH_NAMESPACE 23 | valueFrom: 24 | fieldRef: 25 | fieldPath: metadata.namespace 26 | - name: POD_NAME 27 | valueFrom: 28 | fieldRef: 29 | fieldPath: metadata.name 30 | - name: OPERATOR_NAME 31 | value: "wlm-operator" 32 | securityContext: 33 | runAsUser: 1000 34 | runAsGroup: 1000 35 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | six 2 | python-dateutil 3 | urllib3 4 | certifi 5 | argo-workflows==5.0.0 6 | jsonpickle 7 | minio 8 | kubernetes 9 | pyyaml 10 | cloudpickle==2.2.0 11 | requests 12 | tqdm 13 | psutil 14 | filelock 15 | -------------------------------------------------------------------------------- /scripts/docker-compose.yml: -------------------------------------------------------------------------------- 1 | version: "2.2" 2 | 3 | services: 4 | mysql: 5 | image: mysql:5.7.29 6 | hostname: mysql 7 | container_name: mysql 8 | environment: 9 | MYSQL_RANDOM_ROOT_PASSWORD: "yes" 10 | MYSQL_DATABASE: slurm_acct_db 11 | MYSQL_USER: slurm 12 | MYSQL_PASSWORD: password 13 | volumes: 14 | - var_lib_mysql:/var/lib/mysql 15 | networks: 16 | common-network: 17 | 18 | slurmdbd: 19 | image: dptechnology/slurm-docker-cluster:latest 20 | command: ["slurmdbd"] 21 | container_name: slurmdbd 22 | hostname: slurmdbd 23 | volumes: 24 | - etc_munge:/etc/munge 25 | - etc_slurm:/etc/slurm 26 | - var_log_slurm:/var/log/slurm 27 | expose: 28 | - "6819" 29 | depends_on: 30 | - mysql 31 | networks: 32 | common-network: 33 | 34 | slurmctld: 35 | image: dptechnology/slurm-docker-cluster:latest 36 | command: ["slurmctld"] 37 | container_name: slurmctld 38 | hostname: slurmctld 39 | volumes: 40 | - etc_munge:/etc/munge 41 | - etc_slurm:/etc/slurm 42 | - slurm_jobdir:/data 43 | - var_log_slurm:/var/log/slurm 44 | ports: 45 | - "31129:22" 46 | expose: 47 | - "6817" 48 | depends_on: 49 | - "slurmdbd" 50 | networks: 51 | common-network: 52 | ipv4_address: 10.1.1.10 53 | cap_add: 54 | - NET_ADMIN 55 | 56 | c1: 57 | image: dptechnology/slurm-docker-cluster:latest 58 | command: ["slurmd"] 59 | hostname: c1 60 | container_name: c1 61 | volumes: 62 | - etc_munge:/etc/munge 63 | - etc_slurm:/etc/slurm 64 | - slurm_jobdir:/data 65 | - var_log_slurm:/var/log/slurm 66 | expose: 67 | - "6818" 68 | depends_on: 69 | - "slurmctld" 70 | networks: 71 | common-network: 72 | ipv4_address: 10.1.1.11 73 | cap_add: 74 | - NET_ADMIN 75 | 76 | c2: 77 | image: dptechnology/slurm-docker-cluster:latest 78 | command: ["slurmd"] 79 | hostname: c2 80 | container_name: c2 81 | volumes: 82 | - etc_munge:/etc/munge 83 | - etc_slurm:/etc/slurm 84 | - slurm_jobdir:/data 85 | - var_log_slurm:/var/log/slurm 86 | expose: 87 | - "6818" 88 | depends_on: 89 | - "slurmctld" 90 | networks: 91 | common-network: 92 | ipv4_address: 10.1.1.12 93 | cap_add: 94 | - NET_ADMIN 95 | 96 | volumes: 97 | etc_munge: 98 | etc_slurm: 99 | slurm_jobdir: 100 | var_lib_mysql: 101 | var_log_slurm: 102 | 103 | networks: 104 | common-network: 105 | driver: bridge 106 | ipam: 107 | driver: default 108 | config: 109 | - subnet: 10.1.1.0/24 110 | -------------------------------------------------------------------------------- /scripts/install-linux-cn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function INFO() { 4 | echo "[INFO] $@" 5 | } 6 | 7 | function WARNING() { 8 | echo >&2 "[WARNING] $@" 9 | } 10 | 11 | function ERROR() { 12 | echo >&2 "[ERROR] $@" 13 | } 14 | 15 | docker_path=$(which docker) 16 | if [[ -n "$docker_path" ]]; then 17 | INFO "Found docker executable at $docker_path" 18 | else 19 | INFO "Docker not found, installing docker..." 20 | sh -c "$(curl -fsSL https://get.docker.com/)" 21 | if [[ $? != 0 ]]; then 22 | ERROR "Fail to install docker" 23 | exit 1 24 | fi 25 | fi 26 | 27 | minikube_path=$(which minikube) 28 | if [[ -n "$minikube_path" ]]; then 29 | INFO "Found minikube binary at $minikube_path" 30 | else 31 | INFO "Minikube not found, installing minikube 1.25.2 ..." 32 | curl -o minikube -L https://registry.npmmirror.com/-/binary/minikube/v1.25.2/minikube-linux-amd64 33 | if [[ $? != 0 ]]; then 34 | ERROR "Fail to download minikube" 35 | exit 1 36 | fi 37 | sudo install minikube /usr/local/bin/minikube 38 | if [[ $? != 0 ]]; then 39 | ERROR "Fail to install minikube" 40 | exit 1 41 | fi 42 | fi 43 | 44 | kubectl=$(which kubectl) 45 | if [[ $? != 0 ]]; then 46 | curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl" 47 | if [[ $? != 0 ]]; then 48 | ERROR "Fail to download kubectl" 49 | exit 1 50 | fi 51 | sudo install kubectl /usr/local/bin/kubectl 52 | if [[ $? != 0 ]]; then 53 | ERROR "Fail to install kubectl" 54 | exit 1 55 | fi 56 | fi 57 | 58 | 59 | # if [[ $EUID > 0 ]]; then 60 | # minikube start --image-mirror-country=cn $@ 61 | # else 62 | # INFO "minikube can not start with root user" 63 | # INFO "Creating new user" 64 | # adduser developer 65 | # usermod -aG sudo developer 66 | # INFO "Login to the newly created user" 67 | # su - developer 68 | # INFO "Add user to the Docker Group" 69 | # sudo groupadd docker 70 | # sudo usermod -aG docker $USER 71 | # exit 72 | # su - developer 73 | # fi 74 | 75 | 76 | 77 | 78 | minikube status 1>/dev/null 2>/dev/null 79 | if [[ $? < 2 ]]; then 80 | INFO "Minikube has been started" 81 | else 82 | INFO "Starting minikube..." 83 | minikube start --image-mirror-country=cn $@ 84 | if [[ $? != 0 ]]; then 85 | ERROR "Fail to start minikube" 86 | exit 1 87 | fi 88 | fi 89 | 90 | kubectl create ns argo 1>/dev/null 2>/dev/null 91 | wget https://gitee.com/deepmodeling/dflow/raw/master/manifests/quick-start-postgres-3.4.1-deepmodeling.yaml 92 | kubectl apply -n argo -f quick-start-postgres-3.4.1-deepmodeling.yaml 1>/dev/null 93 | if [[ $? != 0 ]]; then 94 | ERROR "Fail to apply argo yaml" 95 | exit 1 96 | fi 97 | 98 | function waitForReady() { 99 | while true; do 100 | ready=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.readyReplicas}') 101 | replicas=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.replicas}') 102 | if [[ $? != 0 ]]; then 103 | ERROR "Fail to get status of $1" 104 | exit 1 105 | fi 106 | if [[ $replicas > 0 && $ready == $replicas ]]; then 107 | INFO "$1 has been ready..." 108 | break 109 | fi 110 | INFO "Waiting for $1 ready..." 111 | sleep 3 112 | done 113 | } 114 | 115 | waitForReady argo-server 116 | waitForReady httpbin 117 | waitForReady minio 118 | waitForReady postgres 119 | waitForReady workflow-controller 120 | 121 | function forward() { 122 | pid=`ps -ef | grep port-forward | grep $1 | grep $2 | awk '{print $2}'` 123 | if [[ -n "$pid" ]]; then 124 | kill -9 $pid 125 | fi 126 | INFO "Forwarding $1:$2 to localhost:$2" 127 | nohup kubectl -n argo port-forward deployment/$1 $2:$2 --address 0.0.0.0 & 128 | } 129 | forward argo-server 2746 130 | forward minio 9000 131 | forward minio 9001 132 | 133 | sleep 3 134 | INFO "dflow server has been installed successfully!" 135 | -------------------------------------------------------------------------------- /scripts/install-linux.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function INFO() { 4 | echo "[INFO] $@" 5 | } 6 | 7 | function WARNING() { 8 | echo >&2 "[WARNING] $@" 9 | } 10 | 11 | function ERROR() { 12 | echo >&2 "[ERROR] $@" 13 | } 14 | 15 | docker_path=$(which docker) 16 | if [[ -n "$docker_path" ]]; then 17 | INFO "Found docker executable at $docker_path" 18 | else 19 | INFO "Docker not found, installing docker..." 20 | sh -c "$(curl -fsSL https://get.docker.com/)" 21 | if [[ $? != 0 ]]; then 22 | ERROR "Fail to install docker" 23 | exit 1 24 | fi 25 | fi 26 | 27 | minikube_path=$(which minikube) 28 | if [[ -n "$minikube_path" ]]; then 29 | INFO "Found minikube binary at $minikube_path" 30 | else 31 | INFO "Minikube not found, installing minikube..." 32 | curl -o minikube -L https://storage.googleapis.com/minikube/releases/latest/minikube-linux-amd64 33 | if [[ $? != 0 ]]; then 34 | ERROR "Fail to download minikube" 35 | exit 1 36 | fi 37 | sudo install minikube /usr/local/bin/minikube 38 | if [[ $? != 0 ]]; then 39 | ERROR "Fail to install minikube" 40 | exit 1 41 | fi 42 | fi 43 | 44 | kubectl=$(which kubectl) 45 | if [[ $? != 0 ]]; then 46 | curl -LO "https://dl.k8s.io/release/$(curl -L -s https://dl.k8s.io/release/stable.txt)/bin/linux/amd64/kubectl" 47 | if [[ $? != 0 ]]; then 48 | ERROR "Fail to download kubectl" 49 | exit 1 50 | fi 51 | sudo install kubectl /usr/local/bin/kubectl 52 | if [[ $? != 0 ]]; then 53 | ERROR "Fail to install kubectl" 54 | exit 1 55 | fi 56 | fi 57 | 58 | minikube status 1>/dev/null 2>/dev/null 59 | if [[ $? < 2 ]]; then 60 | INFO "Minikube has been started" 61 | else 62 | INFO "Starting minikube..." 63 | minikube start $@ 64 | if [[ $? != 0 ]]; then 65 | ERROR "Fail to start minikube" 66 | exit 1 67 | fi 68 | fi 69 | 70 | kubectl create ns argo 1>/dev/null 2>/dev/null 71 | kubectl apply -n argo -f https://raw.githubusercontent.com/deepmodeling/dflow/master/manifests/quick-start-postgres-3.4.1-deepmodeling.yaml 1>/dev/null 72 | if [[ $? != 0 ]]; then 73 | ERROR "Fail to apply argo yaml" 74 | exit 1 75 | fi 76 | 77 | function waitForReady() { 78 | while true; do 79 | ready=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.readyReplicas}') 80 | replicas=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.replicas}') 81 | if [[ $? != 0 ]]; then 82 | ERROR "Fail to get status of $1" 83 | exit 1 84 | fi 85 | if [[ $replicas > 0 && $ready == $replicas ]]; then 86 | INFO "$1 has been ready..." 87 | break 88 | fi 89 | INFO "Waiting for $1 ready..." 90 | sleep 3 91 | done 92 | } 93 | 94 | waitForReady argo-server 95 | waitForReady minio 96 | waitForReady postgres 97 | waitForReady workflow-controller 98 | 99 | function forward() { 100 | pid=`ps -ef | grep port-forward | grep $1 | grep $2 | awk '{print $2}'` 101 | if [[ -n "$pid" ]]; then 102 | kill -9 $pid 103 | fi 104 | INFO "Forwarding $1:$2 to localhost:$2" 105 | nohup kubectl -n argo port-forward deployment/$1 $2:$2 --address 0.0.0.0 & 106 | } 107 | forward argo-server 2746 108 | forward minio 9000 109 | forward minio 9001 110 | 111 | sleep 3 112 | INFO "dflow server has been installed successfully!" 113 | -------------------------------------------------------------------------------- /scripts/install-mac-cn.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function INFO() { 4 | echo "[INFO] $@" 5 | } 6 | 7 | function WARNING() { 8 | echo >&2 "[WARNING] $@" 9 | } 10 | 11 | function ERROR() { 12 | echo >&2 "[ERROR] $@" 13 | } 14 | 15 | docker_path=$(which docker) 16 | if [[ -n "$docker_path" ]]; then 17 | INFO "Found docker client at $docker_path" 18 | else 19 | INFO "Docker client not found, installing docker client..." 20 | brew install docker 21 | if [[ $? != 0 ]]; then 22 | ERROR "Fail to install docker client" 23 | exit 1 24 | fi 25 | fi 26 | 27 | docker info 1>/dev/null 2>/dev/null 28 | if [[ $? == 0 ]]; then 29 | INFO "Docker server has been running" 30 | elif [[ $? == 1 ]]; then 31 | if [[ -d "/Applications/Docker.app" ]]; then 32 | INFO "Found docker app at /Applications/Docker.app" 33 | else 34 | INFO "Downloading docker app..." 35 | machine_type=$(uname -m) 36 | if [[ "$machine_type" == "amd64" || "$machine_type" == "x86_64" ]]; then 37 | curl -o Docker.dmg -L "https://desktop.docker.com/mac/main/amd64/Docker.dmg?utm_source=docker&utm_medium=webreferral&utm_campaign=docs-driven-download-mac-amd64" 38 | elif [[ "$machine_type" == "arm64" ]]; then 39 | curl -o Docker.dmg -L "https://desktop.docker.com/mac/main/arm64/Docker.dmg?utm_source=docker&utm_medium=webreferral&utm_campaign=docs-driven-download-mac-arm64" 40 | else 41 | ERROR "Unrecognized machine architecture" 42 | fi 43 | if [[ $? != 0 ]]; then 44 | ERROR "Fail to download docker" 45 | exit 1 46 | fi 47 | sudo hdiutil attach Docker.dmg 48 | sudo /Volumes/Docker/Docker.app/Contents/MacOS/install 49 | if [[ $? != 0 ]]; then 50 | ERROR "Fail to install docker" 51 | exit 1 52 | fi 53 | sudo hdiutil detach /Volumes/Docker 54 | fi 55 | 56 | INFO "Starting docker server" 57 | sudo open /Applications/Docker.app 58 | if [[ $? != 0 ]]; then 59 | ERROR "Fail to start docker app" 60 | exit 1 61 | fi 62 | while true; do 63 | docker info 1>/dev/null 2>/dev/null 64 | if [[ $? == 0 ]]; then 65 | INFO "Docker server has been running..." 66 | break 67 | fi 68 | INFO "Waiting docker server running..." 69 | sleep 3 70 | done 71 | fi 72 | 73 | minikube_path=$(which minikube) 74 | if [[ -n "$minikube_path" ]]; then 75 | INFO "Found minikube binary at $minikube_path" 76 | else 77 | INFO "Minikube not found, installing minikube 1.25.2 ..." 78 | machine_type=$(uname -m) 79 | if [[ "$machine_type" == "amd64" || "$machine_type" == "x86_64" ]]; then 80 | curl -o minikube -L "https://registry.npmmirror.com/-/binary/minikube/v1.25.2/minikube-darwin-amd64" 81 | elif [[ "$machine_type" == "arm64" ]]; then 82 | curl -o minikube -L "https://registry.npmmirror.com/-/binary/minikube/v1.25.2/minikube-darwin-arm64" 83 | else 84 | ERROR "Unrecognized machine architecture" 85 | fi 86 | if [[ $? != 0 ]]; then 87 | ERROR "Fail to download minikube" 88 | exit 1 89 | fi 90 | sudo install minikube /usr/local/bin/minikube 91 | if [[ $? != 0 ]]; then 92 | ERROR "Fail to install minikube" 93 | exit 1 94 | fi 95 | fi 96 | 97 | kubectl_path=$(which kubectl) 98 | if [[ -z "$kubectl_path" ]]; then 99 | echo "alias kubectl=\"minikube kubectl --\"" >> ~/.bash_profile 100 | source ~/.bash_profile 101 | fi 102 | 103 | minikube status 1>/dev/null 2>/dev/null 104 | if [[ $? < 2 ]]; then 105 | INFO "Minikube has been started" 106 | else 107 | INFO "Starting minikube..." 108 | minikube start --image-mirror-country=cn $@ 109 | if [[ $? != 0 ]]; then 110 | ERROR "Fail to start minikube" 111 | exit 1 112 | fi 113 | fi 114 | 115 | kubectl create ns argo 1>/dev/null 2>/dev/null 116 | wget https://gitee.com/deepmodeling/dflow/raw/master/manifests/quick-start-postgres-3.4.1-deepmodeling.yaml 117 | kubectl apply -n argo -f quick-start-postgres-3.4.1-deepmodeling.yaml 1>/dev/null 118 | if [[ $? != 0 ]]; then 119 | ERROR "Fail to apply argo yaml" 120 | exit 1 121 | fi 122 | 123 | function waitForReady() { 124 | while true; do 125 | ready=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.readyReplicas}') 126 | replicas=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.replicas}') 127 | if [[ $? != 0 ]]; then 128 | ERROR "Fail to get status of $1" 129 | exit 1 130 | fi 131 | if [[ $replicas > 0 && $ready == $replicas ]]; then 132 | INFO "$1 has been ready..." 133 | break 134 | fi 135 | INFO "Waiting for $1 ready..." 136 | sleep 3 137 | done 138 | } 139 | 140 | waitForReady argo-server 141 | waitForReady httpbin 142 | waitForReady minio 143 | waitForReady postgres 144 | waitForReady workflow-controller 145 | 146 | function forward() { 147 | pid=`ps -ef | grep port-forward | grep $1 | grep $2 | awk '{print $2}'` 148 | if [[ -n "$pid" ]]; then 149 | kill -9 $pid 150 | fi 151 | INFO "Forwarding $1:$2 to localhost:$2" 152 | nohup kubectl -n argo port-forward deployment/$1 $2:$2 --address 0.0.0.0 & 153 | } 154 | forward argo-server 2746 155 | forward minio 9000 156 | forward minio 9001 157 | 158 | sleep 3 159 | INFO "dflow server has been installed successfully!" 160 | -------------------------------------------------------------------------------- /scripts/install-mac.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | function INFO() { 4 | echo "[INFO] $@" 5 | } 6 | 7 | function WARNING() { 8 | echo >&2 "[WARNING] $@" 9 | } 10 | 11 | function ERROR() { 12 | echo >&2 "[ERROR] $@" 13 | } 14 | 15 | docker_path=$(which docker) 16 | if [[ -n "$docker_path" ]]; then 17 | INFO "Found docker client at $docker_path" 18 | else 19 | INFO "Docker client not found, installing docker client..." 20 | brew install docker 21 | if [[ $? != 0 ]]; then 22 | ERROR "Fail to install docker client" 23 | exit 1 24 | fi 25 | fi 26 | 27 | docker info 1>/dev/null 2>/dev/null 28 | if [[ $? == 0 ]]; then 29 | INFO "Docker server has been running" 30 | elif [[ $? == 1 ]]; then 31 | if [[ -d "/Applications/Docker.app" ]]; then 32 | INFO "Found docker app at /Applications/Docker.app" 33 | else 34 | INFO "Downloading docker app..." 35 | machine_type=$(uname -m) 36 | if [[ "$machine_type" == "amd64" || "$machine_type" == "x86_64" ]]; then 37 | curl -o Docker.dmg -L "https://desktop.docker.com/mac/main/amd64/Docker.dmg?utm_source=docker&utm_medium=webreferral&utm_campaign=docs-driven-download-mac-amd64" 38 | elif [[ "$machine_type" == "arm64" ]]; then 39 | curl -o Docker.dmg -L "https://desktop.docker.com/mac/main/arm64/Docker.dmg?utm_source=docker&utm_medium=webreferral&utm_campaign=docs-driven-download-mac-arm64" 40 | else 41 | ERROR "Unrecognized machine architecture" 42 | fi 43 | if [[ $? != 0 ]]; then 44 | ERROR "Fail to download docker" 45 | exit 1 46 | fi 47 | sudo hdiutil attach Docker.dmg 48 | sudo /Volumes/Docker/Docker.app/Contents/MacOS/install 49 | if [[ $? != 0 ]]; then 50 | ERROR "Fail to install docker" 51 | exit 1 52 | fi 53 | sudo hdiutil detach /Volumes/Docker 54 | fi 55 | 56 | INFO "Starting docker server" 57 | sudo open /Applications/Docker.app 58 | if [[ $? != 0 ]]; then 59 | ERROR "Fail to start docker app" 60 | exit 1 61 | fi 62 | while true; do 63 | docker info 1>/dev/null 2>/dev/null 64 | if [[ $? == 0 ]]; then 65 | INFO "Docker server has been running..." 66 | break 67 | fi 68 | INFO "Waiting docker server running..." 69 | sleep 3 70 | done 71 | fi 72 | 73 | minikube_path=$(which minikube) 74 | if [[ -n "$minikube_path" ]]; then 75 | INFO "Found minikube binary at $minikube_path" 76 | else 77 | INFO "Minikube not found, installing minikube..." 78 | machine_type=$(uname -m) 79 | if [[ "$machine_type" == "amd64" || "$machine_type" == "x86_64" ]]; then 80 | curl -o minikube -L "https://storage.googleapis.com/minikube/releases/latest/minikube-darwin-amd64" 81 | elif [[ "$machine_type" == "arm64" ]]; then 82 | curl -o minikube -L "https://storage.googleapis.com/minikube/releases/latest/minikube-darwin-arm64" 83 | else 84 | ERROR "Unrecognized machine architecture" 85 | fi 86 | if [[ $? != 0 ]]; then 87 | ERROR "Fail to download minikube" 88 | exit 1 89 | fi 90 | sudo install minikube /usr/local/bin/minikube 91 | if [[ $? != 0 ]]; then 92 | ERROR "Fail to install minikube" 93 | exit 1 94 | fi 95 | fi 96 | 97 | kubectl_path=$(which kubectl) 98 | if [[ -z "$kubectl_path" ]]; then 99 | echo "alias kubectl=\"minikube kubectl --\"" >> ~/.bash_profile 100 | source ~/.bash_profile 101 | fi 102 | 103 | minikube status 1>/dev/null 2>/dev/null 104 | if [[ $? < 2 ]]; then 105 | INFO "Minikube has been started" 106 | else 107 | INFO "Starting minikube..." 108 | minikube start $@ 109 | if [[ $? != 0 ]]; then 110 | ERROR "Fail to start minikube" 111 | exit 1 112 | fi 113 | fi 114 | 115 | kubectl create ns argo 1>/dev/null 2>/dev/null 116 | kubectl apply -n argo -f https://raw.githubusercontent.com/deepmodeling/dflow/master/manifests/quick-start-postgres-3.4.1-deepmodeling.yaml 1>/dev/null 117 | if [[ $? != 0 ]]; then 118 | ERROR "Fail to apply argo yaml" 119 | exit 1 120 | fi 121 | 122 | function waitForReady() { 123 | while true; do 124 | ready=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.readyReplicas}') 125 | replicas=$(kubectl get deployment $1 -n argo -o jsonpath='{.status.replicas}') 126 | if [[ $? != 0 ]]; then 127 | ERROR "Fail to get status of $1" 128 | exit 1 129 | fi 130 | if [[ $replicas > 0 && $ready == $replicas ]]; then 131 | INFO "$1 has been ready..." 132 | break 133 | fi 134 | INFO "Waiting for $1 ready..." 135 | sleep 3 136 | done 137 | } 138 | 139 | waitForReady argo-server 140 | waitForReady minio 141 | waitForReady postgres 142 | waitForReady workflow-controller 143 | 144 | function forward() { 145 | pid=`ps -ef | grep port-forward | grep $1 | grep $2 | awk '{print $2}'` 146 | if [[ -n "$pid" ]]; then 147 | kill -9 $pid 148 | fi 149 | INFO "Forwarding $1:$2 to localhost:$2" 150 | nohup kubectl -n argo port-forward deployment/$1 $2:$2 --address 0.0.0.0 & 151 | } 152 | forward argo-server 2746 153 | forward minio 9000 154 | forward minio 9001 155 | 156 | sleep 3 157 | INFO "dflow server has been installed successfully!" 158 | -------------------------------------------------------------------------------- /scripts/register_cluster.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | set -e 3 | 4 | docker exec slurmctld bash -c "/usr/bin/sacctmgr --immediate add cluster name=linux" && \ 5 | docker-compose restart slurmdbd slurmctld 6 | -------------------------------------------------------------------------------- /scripts/start-slurm.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | docker-compose up -d --no-build 4 | 5 | while [ `./register_cluster.sh 2>&1 | grep "sacctmgr: error" | wc -l` -ne 0 ] 6 | do 7 | echo "Waiting for SLURM cluster to become ready"; 8 | sleep 2 9 | done 10 | echo "SLURM properly configured" 11 | 12 | docker exec slurmctld /usr/sbin/sshd 13 | docker exec slurmctld rm -f /etc/locale.conf 14 | # On some clusters the login node does not have the same interface as the 15 | # compute nodes. The next three lines allow to test this edge case by adding 16 | # separate interfaces on the worker and on the scheduler nodes. 17 | docker exec slurmctld ip addr add 10.1.1.20/24 dev eth0 label eth0:scheduler 18 | docker exec c1 ip addr add 10.1.1.21/24 dev eth0 label eth0:worker 19 | docker exec c2 ip addr add 10.1.1.22/24 dev eth0 label eth0:worker 20 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup 2 | 3 | with open('VERSION', 'r') as f: 4 | VERSION = f.read() 5 | 6 | with open('README.md', 'r') as f: 7 | LONG_DESCRIPTION = f.read() 8 | 9 | setup( 10 | name='pydflow', 11 | version=VERSION, 12 | description='Dflow is a Python framework for constructing scientific ' 13 | 'computing workflows employing Argo Workflows as the workflow engine.', 14 | long_description=LONG_DESCRIPTION, 15 | long_description_content_type="text/markdown", 16 | author="Xinzijian Liu", 17 | author_email="liuxzj@dp.tech", 18 | url="https://github.com/deepmodeling/dflow", 19 | license="LGPLv3", 20 | package_dir={'': 'src'}, 21 | packages=[ 22 | "dflow", 23 | "dflow/python", 24 | "dflow/python/vendor", 25 | "dflow/python/vendor/typeguard", 26 | "dflow/client", 27 | "dflow/plugins", 28 | ], 29 | include_package_data=True, 30 | python_requires='>=3.6', 31 | install_requires=[ 32 | "six", 33 | "python-dateutil", 34 | "urllib3", 35 | "certifi", 36 | "argo-workflows==5.0.0", 37 | "jsonpickle", 38 | "minio", 39 | "kubernetes", 40 | "pyyaml", 41 | "cloudpickle==2.2.0", 42 | "requests", 43 | "tqdm", 44 | "psutil", 45 | "filelock", 46 | ], 47 | entry_points={ 48 | 'console_scripts': ['dflow=dflow.main:main'], 49 | }, 50 | ) 51 | -------------------------------------------------------------------------------- /src/dflow/__init__.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import os 3 | 4 | from .argo_objects import ArgoStep, ArgoWorkflow 5 | from .code_gen import gen_code 6 | from .common import (CustomArtifact, HTTPArtifact, LineageClient, 7 | LocalArtifact, S3Artifact, import_func, jsonpickle) 8 | from .config import config, s3_config, set_config, set_s3_config 9 | from .context import Context 10 | from .dag import DAG 11 | from .executor import ContainerExecutor, Executor, RemoteExecutor 12 | from .io import (AutonamedDict, IfExpression, InputArtifact, InputParameter, 13 | Inputs, OutputArtifact, OutputParameter, Outputs, 14 | if_expression) 15 | from .op_template import (HTTPOPTemplate, OPTemplate, PythonScriptOPTemplate, 16 | Secret, ShellOPTemplate) 17 | from .resource import Resource 18 | from .slurm import SlurmJob, SlurmJobTemplate, SlurmRemoteExecutor 19 | from .step import (HookStep, Step, argo_concat, argo_enumerate, argo_len, 20 | argo_range, argo_sequence, argo_sum) 21 | from .steps import Steps 22 | from .task import Task 23 | from .utils import (copy_artifact, copy_s3, download_artifact, download_s3, 24 | path_list_of_artifact, path_object_of_artifact, randstr, 25 | upload_artifact, upload_s3) 26 | from .workflow import (DockerSecret, Workflow, parse_repo, 27 | query_archived_workflows, query_workflows) 28 | 29 | log_level = os.environ.get('LOG_LEVEL') 30 | if log_level: 31 | logging.basicConfig(level=getattr(logging, log_level.upper(), None)) 32 | 33 | __all__ = ["S3Artifact", "DAG", "Executor", "RemoteExecutor", "AutonamedDict", 34 | "IfExpression", "InputArtifact", "InputParameter", "Inputs", 35 | "OutputArtifact", "OutputParameter", "Outputs", 36 | "if_expression", "OPTemplate", "PythonScriptOPTemplate", 37 | "ShellOPTemplate", "Resource", "SlurmJob", "SlurmJobTemplate", 38 | "SlurmRemoteExecutor", "Step", "argo_len", "argo_range", 39 | "argo_sequence", "Steps", "Task", "copy_artifact", "copy_s3", 40 | "download_artifact", "download_s3", "path_list_of_artifact", 41 | "s3_config", "upload_artifact", "upload_s3", "Workflow", "config", 42 | "Context", "randstr", "LocalArtifact", "set_config", 43 | "set_s3_config", "DockerSecret", "argo_sum", "argo_concat", 44 | "LineageClient", "Secret", "query_workflows", 45 | "query_archived_workflows", "ContainerExecutor", "ArgoStep", 46 | "ArgoWorkflow", "argo_enumerate", "path_object_of_artifact", 47 | "CustomArtifact", "gen_code", "jsonpickle", "HTTPArtifact", 48 | "HookStep", "HTTPOPTemplate"] 49 | 50 | 51 | if os.environ.get("DFLOW_LINEAGE"): 52 | config["lineage"] = import_func(os.environ.get("DFLOW_LINEAGE"))() 53 | if os.environ.get("DFLOW_S3_STORAGE_CLIENT"): 54 | s3_config["storage_client"] = import_func(os.environ.get( 55 | "DFLOW_S3_STORAGE_CLIENT"))() 56 | parse_repo() 57 | -------------------------------------------------------------------------------- /src/dflow/__main__.py: -------------------------------------------------------------------------------- 1 | from .main import main 2 | 3 | if __name__ == "__main__": 4 | main() 5 | -------------------------------------------------------------------------------- /src/dflow/client/__init__.py: -------------------------------------------------------------------------------- 1 | from .v1alpha1_artifact import V1alpha1Artifact 2 | from .v1alpha1_dag_task import V1alpha1DAGTask 3 | from .v1alpha1_lifecycle_hook import V1alpha1LifecycleHook 4 | from .v1alpha1_parameter import V1alpha1Parameter 5 | from .v1alpha1_retry_strategy import V1alpha1RetryStrategy 6 | from .v1alpha1_sequence import V1alpha1Sequence 7 | from .v1alpha1_template import V1alpha1Template 8 | from .v1alpha1_value_from import V1alpha1ValueFrom 9 | from .v1alpha1_workflow_step import V1alpha1WorkflowStep 10 | 11 | __all__ = ["V1alpha1Artifact", "V1alpha1LifecycleHook", "V1alpha1Parameter", 12 | "V1alpha1RetryStrategy", "V1alpha1Sequence", "V1alpha1ValueFrom", 13 | "V1alpha1WorkflowStep", "V1alpha1DAGTask", "V1alpha1Template"] 14 | -------------------------------------------------------------------------------- /src/dflow/context.py: -------------------------------------------------------------------------------- 1 | import abc 2 | from abc import ABC 3 | 4 | from .op_template import OPTemplate 5 | 6 | 7 | class Context(ABC): 8 | """ 9 | Context 10 | """ 11 | @abc.abstractmethod 12 | def render( 13 | self, 14 | template: OPTemplate, 15 | ) -> OPTemplate: 16 | """ 17 | render original template and return a new template 18 | """ 19 | return template 20 | -------------------------------------------------------------------------------- /src/dflow/context_syntax.py: -------------------------------------------------------------------------------- 1 | import warnings 2 | from typing import Any, List, Union 3 | 4 | dflow = Any 5 | 6 | 7 | class _Context(object): 8 | """Global Context Manager.""" 9 | 10 | def __init__(self): 11 | """Init context with default to false.""" 12 | self._in_context = False 13 | self.current_workflow = None 14 | 15 | @property 16 | def in_context(self) -> bool: 17 | """whether it is in context environment or not.""" 18 | return self._in_context 19 | 20 | @in_context.setter 21 | def in_context(self, value: bool): 22 | if not isinstance(value, bool): 23 | raise ValueError(f'Unsupported in_context value: {value}.' 24 | f'Expected bool, got {type(value)}') 25 | if not value: 26 | self._reset() 27 | self._in_context = value 28 | 29 | def _reset(self): 30 | """Reset context to origin condition. 31 | 32 | (for context exit usage only.) 33 | """ 34 | self._in_context = False 35 | self.current_workflow = None 36 | 37 | def to_in_context(self): 38 | """Switch to be in context environment.""" 39 | if self.in_context: 40 | warnings.warning( 41 | 'Already in context. But call `to_in_context()` again.') 42 | self._in_context = True 43 | 44 | def to_out_context(self): 45 | """Switch to be not in context environment.""" 46 | if not self.in_context: 47 | warnings.warning( 48 | 'Not in context. But call `to_out_context()` again.') 49 | self._in_context = False 50 | 51 | def registry_step(self, 52 | step: Union['dflow.Step', 53 | List['dflow.Step'], 54 | 'dflow.Task', 55 | List['dflow.Task']]): 56 | """Registry step to context.""" 57 | self.current_workflow.add(step) 58 | 59 | 60 | GLOBAL_CONTEXT = _Context() 61 | 62 | 63 | class Range_Context(object): 64 | """Local context for range.""" 65 | 66 | def __init__(self): 67 | """Init context with default to false.""" 68 | self.range_param_name = None 69 | self.range_target_name = None 70 | self._in_context = False 71 | self.current_step = None 72 | self.range_param_len = 0 73 | 74 | @property 75 | def in_context(self) -> bool: 76 | """whether it is in context environment or not.""" 77 | return self._in_context 78 | 79 | @in_context.setter 80 | def in_context(self, value: bool): 81 | if not isinstance(value, bool): 82 | raise ValueError(f'Unsupported in_context value: {value}.' 83 | f'Expected bool, got {type(value)}') 84 | if not value: 85 | self._reset() 86 | self._in_context = value 87 | 88 | def _reset(self): 89 | """Reset context to origin condition. 90 | 91 | (for context exit usage only.) 92 | """ 93 | self._in_context = False 94 | self.current_workflow = None 95 | self.range_param_name = None 96 | self.range_target_name = None 97 | self.range_param_len = 0 98 | 99 | def to_in_context(self): 100 | """Switch to be in context environment.""" 101 | if self.in_context: 102 | warnings.warning( 103 | 'Already in context. But call `to_in_context()` again.') 104 | self._in_context = True 105 | 106 | def to_out_context(self): 107 | """Switch to be not in context environment.""" 108 | if not self.in_context: 109 | warnings.warning( 110 | 'Not in context. But call `to_out_context()` again.') 111 | self._in_context = False 112 | 113 | def get_current_range_param_name(self): 114 | if self.range_param_name: 115 | return self.range_param_name 116 | else: 117 | raise SyntaxError("It seems not right.") 118 | 119 | 120 | GLOBAL_RANGE_CONTEXT = Range_Context() 121 | -------------------------------------------------------------------------------- /src/dflow/plugins/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/src/dflow/plugins/__init__.py -------------------------------------------------------------------------------- /src/dflow/plugins/launching.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import os 3 | import shutil 4 | from typing import Union 5 | 6 | from dp.launching.cli import to_runner 7 | from dp.launching.typing import BaseModel, Field 8 | from dp.launching.typing.io import InputFilePath, OutputDirectory 9 | from pydantic.main import ModelMetaclass 10 | 11 | from dflow import (OPTemplate, S3Artifact, Step, Workflow, download_artifact, 12 | upload_s3) 13 | from dflow.python import OP, Artifact, BigParameter, Parameter 14 | from dflow.python.utils import (handle_input_artifact, handle_output_artifact, 15 | handle_output_parameter) 16 | 17 | 18 | def OP_template_to_parser(templ: OPTemplate, version: str = "1.0.0"): 19 | ns = {"output_dir": Field(default='./outputs', description='output dir'), 20 | "__annotations__": {"output_dir": OutputDirectory}} 21 | for name, par in templ.inputs.parameters.items(): 22 | ns["__annotations__"][name] = par.type 23 | kwargs = {} 24 | if hasattr(par, "description"): 25 | kwargs["description"] = par.description 26 | if hasattr(par, "value"): 27 | kwargs["default"] = par.value 28 | ns[name] = Field(**kwargs) 29 | for name, art in templ.inputs.artifacts.items(): 30 | ns["__annotations__"][name] = InputFilePath 31 | kwargs = {} 32 | if hasattr(art, "description"): 33 | kwargs["description"] = art.description 34 | ns[name] = Field(**kwargs) 35 | 36 | OPOptions = ModelMetaclass.__new__(ModelMetaclass, "OPOptions", 37 | (BaseModel,), ns) 38 | 39 | def op_runner(opts: OPOptions) -> int: 40 | step = Step("dflow-main", templ, parameters={ 41 | name: getattr(opts, name) 42 | for name in templ.inputs.parameters.keys() 43 | }, artifacts={ 44 | name: S3Artifact(key=upload_s3(getattr(opts, name).get_path())) 45 | for name in templ.inputs.artifacts.keys() 46 | }) 47 | wf = Workflow(templ.name) 48 | wf.add(step) 49 | wf.submit() 50 | wf.wait() 51 | assert wf.query_status() == "Succeeded" 52 | step = wf.query_step("dflow-main")[0] 53 | if hasattr(step, "outputs") and hasattr(step.outputs, "parameters"): 54 | os.makedirs('outputs/parameters', exist_ok=True) 55 | for name, par in step.outputs.parameters.items(): 56 | with open('outputs/parameters/%s' % name, 'w') as f: 57 | f.write(par["value"]) 58 | if hasattr(step, "outputs") and hasattr(step.outputs, "artifacts"): 59 | for name, art in step.outputs.artifacts.items(): 60 | download_artifact(art, path='outputs/artifact/%s' % name, 61 | remove_catalog=False) 62 | return 0 63 | 64 | def to_parser(): 65 | return to_runner( 66 | OPOptions, 67 | op_runner, 68 | version=version, 69 | ) 70 | return to_parser 71 | 72 | 73 | def python_OP_to_parser(op: OP, version: str = "1.0.0"): 74 | ns = {"output_dir": Field(default='./outputs', description='output dir'), 75 | "__annotations__": {"output_dir": OutputDirectory}} 76 | for name, sign in op.get_input_sign().items(): 77 | if isinstance(sign, Artifact): 78 | ns["__annotations__"][name] = InputFilePath 79 | kwargs = {} 80 | if hasattr(sign, "description"): 81 | kwargs["description"] = sign.description 82 | ns[name] = Field(**kwargs) 83 | elif isinstance(sign, (Parameter, BigParameter)): 84 | ns["__annotations__"][name] = sign.type 85 | kwargs = {} 86 | if hasattr(sign, "description"): 87 | kwargs["description"] = sign.description 88 | if hasattr(sign, "default"): 89 | kwargs["default"] = sign.default 90 | ns[name] = Field(**kwargs) 91 | else: 92 | ns["__annotations__"][name] = sign 93 | ns[name] = Field() 94 | 95 | OPOptions = ModelMetaclass.__new__(ModelMetaclass, "OPOptions", 96 | (BaseModel,), ns) 97 | 98 | def op_runner(opts: OPOptions) -> int: 99 | op_in = {} 100 | for name, sign in op.get_input_sign().items(): 101 | if isinstance(sign, Artifact): 102 | op_in[name] = handle_input_artifact( 103 | name, sign, 104 | path=os.path.abspath(getattr(opts, name).get_path())) 105 | else: 106 | op_in[name] = getattr(opts, name) 107 | op_out = op.execute(op_in) 108 | os.makedirs('outputs/parameters', exist_ok=True) 109 | os.makedirs('outputs/artifacts', exist_ok=True) 110 | for name, sign in op.get_output_sign().items(): 111 | value = op_out[name] 112 | if isinstance(sign, Artifact): 113 | if os.path.isdir('outputs/artifacts/%s' % name): 114 | shutil.rmtree('outputs/artifacts/%s' % name) 115 | handle_output_artifact(name, value, sign, data_root=".") 116 | else: 117 | handle_output_parameter(name, value, sign, data_root=".") 118 | return 0 119 | 120 | def to_parser(): 121 | return to_runner( 122 | OPOptions, 123 | op_runner, 124 | version=version, 125 | ) 126 | return to_parser 127 | 128 | 129 | def OP_to_parser(op: Union[OP, OPTemplate], version: str = "1.0.0"): 130 | if inspect.isclass(op) and issubclass(op, OP): 131 | return python_OP_to_parser(op(), version) 132 | elif isinstance(op, OP): 133 | return python_OP_to_parser(op, version) 134 | elif isinstance(op, OPTemplate): 135 | return OP_template_to_parser(op, version) 136 | else: 137 | raise ValueError("Only Python OP/OP template supported") 138 | -------------------------------------------------------------------------------- /src/dflow/plugins/metadata.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Dict, List, Optional, Union 3 | 4 | from dp.metadata import Dataset, MetadataContext 5 | from dp.metadata.entity.task import Task 6 | from dp.metadata.entity.workflow import WorkFlow 7 | 8 | from .. import LineageClient 9 | 10 | config = { 11 | "gms_endpoint": os.environ.get("METADATA_GMS_ENDPOINT", 12 | "https://datahub-gms.dp.tech"), 13 | "project": os.environ.get("METADATA_PROJECT", None), 14 | "token": os.environ.get("METADATA_TOKEN", None), 15 | } 16 | 17 | 18 | class MetadataClient(LineageClient): 19 | def __init__(self, gms_endpoint=None, project=None, token=None): 20 | if gms_endpoint is None: 21 | gms_endpoint = config["gms_endpoint"] 22 | if project is None: 23 | project = config["project"] 24 | if token is None: 25 | token = config["token"] 26 | self.gms_endpoint = gms_endpoint 27 | self.project = project 28 | self.token = token 29 | 30 | def register_workflow( 31 | self, 32 | workflow_name: str) -> str: 33 | with MetadataContext(project=self.project, endpoint=self.gms_endpoint, 34 | token=self.token) as context: 35 | client = context.client 36 | workflow = WorkFlow(workflow_name) 37 | job = client.prepare_workflow(workflow) 38 | client.begin_workflow(job) 39 | return str(job.urn) 40 | 41 | def register_artifact( 42 | self, 43 | namespace: str, 44 | name: str, 45 | uri: str, 46 | description: str = "", 47 | tags: Optional[List[str]] = None, 48 | properties: Optional[Dict[str, str]] = None, 49 | **kwargs) -> str: 50 | if tags is None: 51 | tags = [] 52 | if properties is None: 53 | properties = {} 54 | with MetadataContext(project=self.project, endpoint=self.gms_endpoint, 55 | token=self.token) as context: 56 | client = context.client 57 | urn = Dataset.gen_urn(context, namespace, name) 58 | ds = Dataset(urn=urn, display_name=name, uri=uri, 59 | description=description, tags=tags, 60 | properties=properties) 61 | client.update_dataset(ds) 62 | return urn 63 | 64 | def register_task( 65 | self, 66 | task_name: str, 67 | input_urns: Dict[str, Union[str, List[str]]], 68 | output_uris: Dict[str, str], 69 | workflow_urn: str) -> Dict[str, str]: 70 | with MetadataContext(project=self.project, endpoint=self.gms_endpoint, 71 | token=self.token) as context: 72 | client = context.client 73 | task = Task(task_name, workflow_urn) 74 | job = client.prepare_job(task) 75 | inputs = [] 76 | for urn in input_urns.values(): 77 | if isinstance(urn, list): 78 | inputs += urn 79 | elif isinstance(urn, str): 80 | inputs.append(urn) 81 | inputs = list(filter(lambda x: x != "", inputs)) 82 | run = client.begin_job(job, inputs=inputs) 83 | output_urns = {} 84 | for name, uri in output_uris.items(): 85 | urn = Dataset.gen_urn(context, task_name, name) 86 | ds = Dataset(urn=urn, display_name=name, uri=uri) 87 | client.update_dataset(ds) 88 | output_urns[name] = urn 89 | client.end_job(run, outputs=list(output_urns.values())) 90 | return output_urns 91 | 92 | def get_artifact_metadata(self, urn: str) -> object: 93 | with MetadataContext(project=self.project, endpoint=self.gms_endpoint, 94 | token=self.token) as context: 95 | client = context.client 96 | ds = client.get_dataset(urn) 97 | return ds 98 | -------------------------------------------------------------------------------- /src/dflow/plugins/oss.py: -------------------------------------------------------------------------------- 1 | import os 2 | from typing import Optional 3 | 4 | import oss2 5 | 6 | from ..config import s3_config 7 | from ..utils import StorageClient 8 | 9 | 10 | class OSSClient(StorageClient): 11 | def __init__( 12 | self, 13 | endpoint: Optional[str] = None, 14 | bucket_name: Optional[str] = None, 15 | access_key_id: Optional[str] = None, 16 | access_key_secret: Optional[str] = None, 17 | ) -> None: 18 | if endpoint is None: 19 | endpoint = os.environ.get("OSS_ENDPOINT") 20 | if bucket_name is None: 21 | bucket_name = os.environ.get("OSS_BUCKET_NAME") 22 | if access_key_id is None: 23 | access_key_id = os.environ.get("OSS_ACCESS_KEY_ID") 24 | if access_key_secret is None: 25 | access_key_secret = os.environ.get("OSS_ACCESS_KEY_SECRET") 26 | self.endpoint = endpoint 27 | self.bucket_name = bucket_name 28 | self.access_key_id = access_key_id 29 | self.access_key_secret = access_key_secret 30 | auth = oss2.Auth(access_key_id, access_key_secret) 31 | bucket = oss2.Bucket(auth, endpoint, bucket_name) 32 | self.bucket = bucket 33 | 34 | def to_dict(self): 35 | retained_keys = ["endpoint", "bucket_name"] 36 | return {k: self.__dict__[k] for k in retained_keys} 37 | 38 | def prefixing(self, key): 39 | if not key.startswith(s3_config["repo_prefix"]): 40 | return s3_config["repo_prefix"] + key 41 | return key 42 | 43 | def unprefixing(self, key): 44 | if key.startswith(s3_config["repo_prefix"]): 45 | return key[len(s3_config["repo_prefix"]):] 46 | return key 47 | 48 | def upload(self, key, path): 49 | self.bucket.put_object_from_file(self.prefixing(key), path) 50 | 51 | def download(self, key, path): 52 | if os.path.dirname(path): 53 | os.makedirs(os.path.dirname(path), exist_ok=True) 54 | self.bucket.get_object_to_file(self.prefixing(key), path) 55 | 56 | def list(self, prefix, recursive=False): 57 | prefix = self.prefixing(prefix) 58 | keys = [] 59 | if recursive: 60 | marker = "" 61 | while True: 62 | r = self.bucket.list_objects(prefix, marker=marker) 63 | for obj in r.object_list: 64 | if not obj.key.endswith("/"): 65 | keys.append(self.unprefixing(obj.key)) 66 | if not r.is_truncated: 67 | break 68 | marker = r.next_marker 69 | else: 70 | marker = "" 71 | while True: 72 | r = self.bucket.list_objects(prefix, delimiter="/", 73 | marker=marker) 74 | for obj in r.object_list: 75 | if obj.key == prefix and obj.key.endswith("/"): 76 | continue 77 | keys.append(self.unprefixing(obj.key)) 78 | for key in r.prefix_list: 79 | keys.append(self.unprefixing(key)) 80 | if not r.is_truncated: 81 | break 82 | marker = r.next_marker 83 | return keys 84 | 85 | def copy(self, src, dst): 86 | self.bucket.copy_object(self.bucket_name, self.prefixing(src), 87 | self.prefixing(dst)) 88 | 89 | def get_md5(self, key): 90 | return self.bucket.get_object_meta(self.prefixing(key).etag) 91 | -------------------------------------------------------------------------------- /src/dflow/plugins/ray.py: -------------------------------------------------------------------------------- 1 | from copy import deepcopy 2 | from typing import Any, List, Optional, Union 3 | 4 | from dflow.executor import Executor 5 | from dflow.op_template import ScriptOPTemplate 6 | from dflow.utils import randstr 7 | 8 | try: 9 | from argo.workflows.client import (V1alpha1UserContainer, 10 | V1EmptyDirVolumeSource, V1Volume, 11 | V1VolumeMount) 12 | except ImportError: 13 | pass 14 | 15 | 16 | def _ray_init_container(main_image: str, 17 | install_mirror: Union[bool, str, None] = None, 18 | ray_remote_path_prefix: str = '/tmp/ray', 19 | container_name: str = 'ray-init', 20 | ): 21 | """Get init_container with the same image as main_image do. If the 22 | container have python with ray (import ray), skip the install; else run pip 23 | install ray. 24 | 25 | Args: 26 | main_image: image name of main container 27 | install_mirror: mirror url used by pip install, 28 | such as tuna: "https://pypi.tuna.tsinghua.edu.cn/simple" 29 | ray_remote_path_prefix: ray install path (default: "/tmp/ray") 30 | 31 | Returns: 32 | """ 33 | RAY_INSTALL_STATEMENT = 'pip install ray ' \ 34 | f'--target={ray_remote_path_prefix}' 35 | if install_mirror: 36 | RAY_INSTALL_STATEMENT += f' -i {install_mirror}' 37 | return V1alpha1UserContainer( 38 | name=container_name, 39 | image=main_image, 40 | command=['sh'], 41 | args=['-c', 42 | f"python -c 'import ray' >/dev/null 2>&1 ||" 43 | f' {{ {RAY_INSTALL_STATEMENT}; }}'], 44 | volume_mounts=[V1VolumeMount( 45 | name='ray-python-packages', mount_path=ray_remote_path_prefix)], 46 | ) 47 | 48 | 49 | class RayClusterExecutor(Executor): 50 | 51 | def __init__( 52 | self, 53 | ray_host: str, 54 | ray_remote_path: str = '/tmp/ray', 55 | workdir: str = '~/dflow/workflows/{{workflow.name}}/{{pod.name}}', 56 | ray_install_mirror=None, 57 | ray_dependencies: Optional[List[Any]] = None, 58 | ) -> None: 59 | """Ray cluster executor. 60 | 61 | Args: 62 | ray_host: 63 | ray_remote_path: 64 | workdir: 65 | ray_install_mirror: 66 | ray_dependencies: `py_modules` of ray.init(runtime_env={}) 67 | """ 68 | if ray_dependencies is None: 69 | ray_dependencies = [] 70 | self.ray_host = ray_host 71 | self.ray_remote_path = ray_remote_path 72 | self.workdir = workdir 73 | self.ray_install_mirror = ray_install_mirror 74 | self.ray_dependencies = ray_dependencies 75 | 76 | def render(self, template: ScriptOPTemplate): 77 | new_template = deepcopy(template) 78 | new_template.name += '-' + randstr() 79 | if new_template.init_containers is None: 80 | new_template.init_containers = [ 81 | _ray_init_container(template.image, self.ray_install_mirror, 82 | self.ray_remote_path)] 83 | else: 84 | assert isinstance(new_template.init_containers, 85 | list), f'Type of current init_containers is ' \ 86 | f'not list, but ' \ 87 | f'{type(new_template.init_containers)}, '\ 88 | f'check your other ' \ 89 | f'init_container settings.' 90 | new_template.init_containers.append( 91 | _ray_init_container(template.image, self.ray_install_mirror, 92 | self.ray_remote_path)) 93 | new_template.script = 'import sys,os ' \ 94 | '\n' \ 95 | f"sys.path.append('{self.ray_remote_path}') " \ 96 | f"\n" \ 97 | "if os.environ.get('RAY_ADDRESS') is None: " \ 98 | "\n" \ 99 | f" os.environ['RAY_ADDRESS']=" \ 100 | f"'{self.ray_host}' " \ 101 | f"\n" \ 102 | 'else: ' \ 103 | '\n' \ 104 | " print(f\"Not use input ray_host address," \ 105 | " use {os.environ['RAY_ADDRESS']} instead.\") " \ 106 | "\n" \ 107 | + template.script 108 | # To locate the initialization of package path in `python_op_template`. 109 | insert_index = new_template.script.find('from dflow import config') 110 | new_script = list(new_template.script) 111 | _dependencies_str = ','.join(item.__name__ 112 | for item in self.ray_dependencies) 113 | ray_dependencies_import = \ 114 | f"import {_dependencies_str}\n" if len( 115 | self.ray_dependencies) > 0 else '\n' 116 | new_script.insert( 117 | insert_index, 118 | 'import ray\n' + 119 | ray_dependencies_import + 120 | "ray.init(os.environ['RAY_ADDRESS'], runtime_env={" 121 | f"'py_modules':[" 122 | f"{','.join(item.__name__ for item in self.ray_dependencies)}" 123 | ']})\n') 124 | new_template.script = ''.join(new_script) 125 | new_template.volumes.append(V1Volume( 126 | name='ray-python-packages', empty_dir=V1EmptyDirVolumeSource())) 127 | new_template.mounts.append(V1VolumeMount( 128 | name='ray-python-packages', mount_path=self.ray_remote_path)) 129 | return new_template 130 | -------------------------------------------------------------------------------- /src/dflow/py.typed: -------------------------------------------------------------------------------- 1 | 2 | -------------------------------------------------------------------------------- /src/dflow/python/__init__.py: -------------------------------------------------------------------------------- 1 | from .op import OP 2 | from .opio import (OPIO, Artifact, BigParameter, HDF5Datasets, NestedDict, 3 | OPIOSign, Parameter) 4 | from .python_op_template import (FatalError, PythonOPTemplate, Slices, 5 | TransientError, upload_packages) 6 | 7 | __all__ = ["OP", "OPIO", "Artifact", "BigParameter", "OPIOSign", "Parameter", 8 | "FatalError", "PythonOPTemplate", "Slices", "TransientError", 9 | "upload_packages", "NestedDict", "HDF5Datasets"] 10 | -------------------------------------------------------------------------------- /src/dflow/python/vendor/typeguard/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/src/dflow/python/vendor/typeguard/py.typed -------------------------------------------------------------------------------- /src/dflow/python/vendor/typeguard/pytest_plugin.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | from typeguard.importhook import install_import_hook 4 | 5 | 6 | def pytest_addoption(parser): 7 | group = parser.getgroup('typeguard') 8 | group.addoption('--typeguard-packages', action='store', 9 | help='comma separated name list of packages and modules to instrument for ' # noqa: E501 10 | 'type checking') 11 | 12 | 13 | def pytest_configure(config): 14 | value = config.getoption("typeguard_packages") 15 | if not value: 16 | return 17 | 18 | packages = [pkg.strip() for pkg in value.split(",")] 19 | 20 | already_imported_packages = sorted( 21 | package for package in packages if package in sys.modules 22 | ) 23 | if already_imported_packages: 24 | message = ( 25 | "typeguard cannot check these packages because they " 26 | "are already imported: {}" 27 | ) 28 | raise RuntimeError(message.format(", ".join(already_imported_packages))) # noqa: E501 29 | 30 | install_import_hook(packages=packages) 31 | -------------------------------------------------------------------------------- /src/dflow/resource.py: -------------------------------------------------------------------------------- 1 | import abc 2 | from abc import ABC 3 | from typing import Optional 4 | 5 | from .op_template import OPTemplate 6 | 7 | 8 | class Resource(ABC): 9 | """ 10 | Resource 11 | 12 | Args: 13 | action: action on the Kubernetes resource 14 | success_condition: expression representing success 15 | failure_condition: expression representing failure 16 | """ 17 | action: Optional[str] = None 18 | success_condition: Optional[str] = None 19 | failure_condition: Optional[str] = None 20 | 21 | @abc.abstractmethod 22 | def get_manifest( 23 | self, 24 | template: OPTemplate, 25 | ) -> OPTemplate: 26 | """ 27 | The method to get the manifest (str) 28 | """ 29 | raise NotImplementedError() 30 | -------------------------------------------------------------------------------- /src/dflow/task.py: -------------------------------------------------------------------------------- 1 | from typing import List, Optional, Union 2 | 3 | from .io import OutputArtifact, OutputParameter 4 | from .op_template import OPTemplate 5 | from .step import Step 6 | 7 | try: 8 | from argo.workflows.client import V1alpha1Arguments, V1alpha1ContinueOn 9 | from .client import V1alpha1DAGTask 10 | except Exception: 11 | pass 12 | 13 | 14 | class Task(Step): 15 | """ 16 | Task 17 | 18 | Args: 19 | name: the name of the task 20 | template: OP template the task uses 21 | parameters: input parameters passed to the task as arguments 22 | artifacts: input artifacts passed to the task as arguments 23 | when: conditional task if the condition is satisfied 24 | with_param: generate parallel tasks with respect to a list as 25 | a parameter 26 | continue_on_failed: continue if the task fails 27 | continue_on_num_success: continue if the success number of the 28 | generated parallel tasks greater than certain number 29 | continue_on_success_ratio: continue if the success ratio of the 30 | generated parallel tasks greater than certain number 31 | with_sequence: generate parallel tasks with respect to a sequence 32 | key: the key of the task 33 | executor: define the executor to execute the script 34 | use_resource: use k8s resource 35 | util_image: image for utility step 36 | util_image_pull_policy: image pull policy for utility step 37 | util_command: command for utility step 38 | dependencies: extra dependencies of the task 39 | """ 40 | 41 | def __init__( 42 | self, 43 | name: str, 44 | template: OPTemplate, 45 | dependencies: Optional[List[Union["Task", str]]] = None, 46 | **kwargs, 47 | ) -> None: 48 | self.dependencies = [] 49 | super().__init__(name=name, template=template, **kwargs) 50 | # override inferred dependencies if specified explicitly 51 | if dependencies is not None: 52 | self.dependencies = dependencies 53 | if self.prepare_step is not None: 54 | self.dependencies.append(self.prepare_step) 55 | if self.check_step is not None: 56 | self.check_step.dependencies = [ 57 | "%s.Succeeded || %s.Failed || %s.Errored" % (self, self, self)] 58 | 59 | @property 60 | def expr(self): 61 | return "tasks['%s']" % self.id 62 | 63 | @classmethod 64 | def from_dict(cls, d, templates): 65 | task = super().from_dict(d, templates) 66 | task.dependencies = d.get("dependencies", []) 67 | if d.get("depends"): 68 | for dep in d["depends"].split("&&"): 69 | dep = dep.strip() 70 | # removeprefix and removesuffix only supported for python>=3.9 71 | dep = dep[1:] if dep.startswith("(") else dep 72 | dep = dep[:-1] if dep.endswith(")") else dep 73 | dep = dep[:-10] if dep.endswith(".Succeeded") else dep 74 | task.dependencies.append(dep) 75 | return task 76 | 77 | def set_parameters(self, parameters): 78 | super().set_parameters(parameters) 79 | for v in self.inputs.parameters.values(): 80 | if hasattr(v, "value"): 81 | def handle(obj): 82 | # TODO: Only support output parameter, dict and list 83 | if isinstance(obj, OutputParameter): 84 | if obj.step not in self.dependencies: 85 | self.dependencies.append(obj.step) 86 | elif isinstance(obj, dict): 87 | for v in obj.values(): 88 | handle(v) 89 | elif isinstance(obj, list): 90 | for v in obj: 91 | handle(v) 92 | handle(v.value) 93 | 94 | def set_artifacts(self, artifacts): 95 | super().set_artifacts(artifacts) 96 | for v in self.inputs.artifacts.values(): 97 | if isinstance(v.source, OutputArtifact) and v.source.step not in \ 98 | self.dependencies: 99 | self.dependencies.append(v.source.step) 100 | 101 | def convert_to_argo(self, context=None): 102 | self.prepare_argo_arguments(context) 103 | depends = [] 104 | for task in self.dependencies: 105 | if isinstance(task, Task): 106 | if task.check_step is not None: 107 | depends.append("(%s.Succeeded)" % task.check_step) 108 | else: 109 | depends.append("(%s.Succeeded)" % task) 110 | else: 111 | depends.append("(%s)" % task) 112 | if self.continue_on_failed or self.continue_on_error: 113 | continue_on = V1alpha1ContinueOn( 114 | failed=self.continue_on_failed, error=self.continue_on_error) 115 | else: 116 | continue_on = None 117 | return V1alpha1DAGTask( 118 | name=self.name, template=self.template.name, 119 | arguments=V1alpha1Arguments( 120 | parameters=self.argo_parameters, 121 | artifacts=self.argo_artifacts 122 | ), when=self.when, with_param=self.with_param, 123 | continue_on=continue_on, 124 | with_sequence=self.with_sequence, 125 | depends=" && ".join(depends), 126 | hooks={ 127 | name: hook.convert_to_argo(context) 128 | for name, hook in self.hooks.items() 129 | }, 130 | ) 131 | 132 | def convert_to_graph(self): 133 | g = super().convert_to_graph() 134 | g["dependencies"] = [str(d) for d in self.dependencies] 135 | return g 136 | -------------------------------------------------------------------------------- /tests/test_big_parameter.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import InputParameter, OutputParameter, Step, Steps, Workflow 4 | from dflow.python import (OP, OPIO, BigParameter, OPIOSign, PythonOPTemplate, 5 | upload_packages) 6 | 7 | if "__file__" in locals(): 8 | upload_packages.append(__file__) 9 | 10 | 11 | class Hello: 12 | def __init__(self, msg): 13 | self.msg = msg 14 | 15 | 16 | class Duplicate(OP): 17 | def __init__(self): 18 | pass 19 | 20 | @classmethod 21 | def get_input_sign(cls): 22 | return OPIOSign({ 23 | 'foo': BigParameter(Hello) 24 | }) 25 | 26 | @classmethod 27 | def get_output_sign(cls): 28 | return OPIOSign({ 29 | 'foo': BigParameter(Hello) 30 | }) 31 | 32 | @OP.exec_sign_check 33 | def execute( 34 | self, 35 | op_in: OPIO, 36 | ) -> OPIO: 37 | foo = op_in["foo"] 38 | print(foo.msg) 39 | foo.msg = foo.msg * 2 40 | op_out = OPIO({ 41 | "foo": foo 42 | }) 43 | return op_out 44 | 45 | 46 | def test_big_parameter(): 47 | wf = Workflow(name="big-param") 48 | 49 | steps = Steps(name="hello-steps") 50 | steps.inputs.parameters["foo"] = InputParameter() 51 | steps.outputs.parameters["foo"] = OutputParameter() 52 | 53 | step1 = Step( 54 | name="step1", 55 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 56 | parameters={"foo": steps.inputs.parameters["foo"]}, 57 | key="step1" 58 | ) 59 | steps.add(step1) 60 | 61 | step2 = Step( 62 | name="step2", 63 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 64 | parameters={"foo": step1.outputs.parameters["foo"]}, 65 | key="step2" 66 | ) 67 | steps.add(step2) 68 | 69 | steps.outputs.parameters["foo"].value_from_parameter = \ 70 | step2.outputs.parameters["foo"] 71 | 72 | big_step = Step(name="big-step", template=steps, 73 | parameters={"foo": Hello("hello")}) 74 | wf.add(big_step) 75 | wf.submit() 76 | 77 | while wf.query_status() in ["Pending", "Running"]: 78 | time.sleep(1) 79 | 80 | assert(wf.query_status() == "Succeeded") 81 | step = wf.query_step(key="step1")[0] 82 | assert(step.phase == "Succeeded") 83 | assert(isinstance(step.outputs.parameters["foo"].value, Hello)) 84 | 85 | 86 | if __name__ == "__main__": 87 | test_big_parameter() 88 | -------------------------------------------------------------------------------- /tests/test_conditional_outputs.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import (OutputArtifact, OutputParameter, Outputs, Step, Steps, 4 | Workflow, if_expression) 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Random(OP): 13 | @classmethod 14 | def get_input_sign(cls): 15 | return OPIOSign() 16 | 17 | @classmethod 18 | def get_output_sign(cls): 19 | return OPIOSign({ 20 | "is_head": bool, 21 | "msg1": str, 22 | "msg2": str, 23 | "foo": Artifact(str), 24 | "bar": Artifact(str) 25 | }) 26 | 27 | @OP.exec_sign_check 28 | def execute( 29 | self, 30 | op_in: OPIO, 31 | ) -> OPIO: 32 | open("foo.txt", "w").write("head") 33 | open("bar.txt", "w").write("tail") 34 | is_head = False 35 | return OPIO({ 36 | "is_head": is_head, 37 | "msg1": "head", 38 | "msg2": "tail", 39 | "foo": "foo.txt", 40 | "bar": "bar.txt" 41 | }) 42 | 43 | 44 | def test_conditional_outputs(): 45 | steps = Steps("steps", outputs=Outputs( 46 | parameters={"msg": OutputParameter()}, 47 | artifacts={"res": OutputArtifact()})) 48 | 49 | random_step = Step( 50 | name="random", 51 | template=PythonOPTemplate(Random, image="python:3.8") 52 | ) 53 | steps.add(random_step) 54 | 55 | steps.outputs.parameters["msg"].value_from_expression = if_expression( 56 | _if=random_step.outputs.parameters["is_head"], 57 | _then=random_step.outputs.parameters["msg1"], 58 | _else=random_step.outputs.parameters["msg2"]) 59 | 60 | steps.outputs.artifacts["res"].from_expression = if_expression( 61 | _if=random_step.outputs.parameters["is_head"], 62 | _then=random_step.outputs.artifacts["foo"], 63 | _else=random_step.outputs.artifacts["bar"]) 64 | 65 | wf = Workflow(name="conditional") 66 | step = Step("conditional", template=steps) 67 | wf.add(step) 68 | wf.submit() 69 | 70 | while wf.query_status() in ["Pending", "Running"]: 71 | time.sleep(1) 72 | 73 | assert(wf.query_status() == "Succeeded") 74 | step = wf.query_step(name="conditional")[0] 75 | assert(step.phase == "Succeeded") 76 | assert step.outputs.parameters["msg"].value == "tail" 77 | 78 | 79 | if __name__ == "__main__": 80 | test_conditional_outputs() 81 | -------------------------------------------------------------------------------- /tests/test_dag.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | 4 | from dflow import DAG, Task, Workflow, download_artifact 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Hello(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign() 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign({ 23 | 'msg': int, 24 | 'bar': Artifact(Path), 25 | }) 26 | 27 | @OP.exec_sign_check 28 | def execute( 29 | self, 30 | op_in: OPIO, 31 | ) -> OPIO: 32 | with open("output.txt", "w") as f: 33 | f.write("Hello") 34 | 35 | return OPIO({ 36 | "msg": 1, 37 | "bar": Path("output.txt"), 38 | }) 39 | 40 | 41 | class Duplicate(OP): 42 | def __init__(self): 43 | pass 44 | 45 | @classmethod 46 | def get_input_sign(cls): 47 | return OPIOSign({ 48 | 'msg': int, 49 | 'foo': Artifact(Path), 50 | }) 51 | 52 | @classmethod 53 | def get_output_sign(cls): 54 | return OPIOSign({ 55 | 'msg': int, 56 | 'bar': Artifact(Path), 57 | }) 58 | 59 | @OP.exec_sign_check 60 | def execute( 61 | self, 62 | op_in: OPIO, 63 | ) -> OPIO: 64 | with open("output.txt", "w") as f: 65 | with open(op_in["foo"], "r") as f2: 66 | f.write(f2.read()*2) 67 | return OPIO({ 68 | "msg": op_in["msg"]*2, 69 | "bar": Path("output.txt"), 70 | }) 71 | 72 | 73 | def test_dag(): 74 | dag = DAG() 75 | hello0 = Task(name="hello0", 76 | template=PythonOPTemplate(Hello, image="python:3.8")) 77 | dag.add(hello0) 78 | hello1 = Task(name="hello1", 79 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 80 | parameters={"msg": hello0.outputs.parameters["msg"]}, 81 | artifacts={"foo": hello0.outputs.artifacts["bar"]}) 82 | dag.add(hello1) 83 | 84 | wf = Workflow(name="dag", dag=dag) 85 | wf.submit() 86 | 87 | while wf.query_status() in ["Pending", "Running"]: 88 | time.sleep(1) 89 | 90 | assert(wf.query_status() == "Succeeded") 91 | step = wf.query_step(name="hello1")[0] 92 | assert(step.phase == "Succeeded") 93 | 94 | assert step.outputs.parameters["msg"].value == 2 95 | bar = download_artifact(step.outputs.artifacts["bar"])[0] 96 | with open(bar, "r") as f: 97 | content = f.read() 98 | assert content == "HelloHello" 99 | 100 | 101 | if __name__ == "__main__": 102 | test_dag() 103 | -------------------------------------------------------------------------------- /tests/test_group_size.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import Step, Workflow, argo_len, argo_range 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Prepare(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign() 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign({ 23 | 'data': List[int] 24 | }) 25 | 26 | @OP.exec_sign_check 27 | def execute( 28 | self, 29 | op_in: OPIO, 30 | ) -> OPIO: 31 | op_out = OPIO({ 32 | 'data': list(range(1, 18)) 33 | }) 34 | return op_out 35 | 36 | 37 | class Process(OP): 38 | def __init__(self): 39 | pass 40 | 41 | @classmethod 42 | def get_input_sign(cls): 43 | return OPIOSign({ 44 | 'data': int 45 | }) 46 | 47 | @classmethod 48 | def get_output_sign(cls): 49 | return OPIOSign({ 50 | "output": Artifact(str) 51 | }) 52 | 53 | @OP.exec_sign_check 54 | def execute( 55 | self, 56 | op_in: OPIO, 57 | ) -> OPIO: 58 | with open("%s.txt" % op_in["data"], "w") as f: 59 | f.write("Hello") 60 | return OPIO({ 61 | "output": "%s.txt" % op_in["data"] 62 | }) 63 | 64 | 65 | def test_group_size(): 66 | wf = Workflow("slices") 67 | prepare = Step("prepare", 68 | PythonOPTemplate(Prepare, image="python:3.8")) 69 | wf.add(prepare) 70 | process = Step( 71 | "process", 72 | PythonOPTemplate( 73 | Process, 74 | image="python:3.8", 75 | slices=Slices( 76 | "{{item}}", 77 | input_parameter=["data"], 78 | output_artifact=["output"], 79 | pool_size=1, 80 | group_size=10, 81 | ), 82 | ), 83 | parameters={"data": prepare.outputs.parameters["data"]}, 84 | with_param=argo_range(argo_len(prepare.outputs.parameters["data"])), 85 | ) 86 | wf.add(process) 87 | wf.submit() 88 | 89 | while wf.query_status() in ["Pending", "Running"]: 90 | time.sleep(1) 91 | 92 | assert(wf.query_status() == "Succeeded") 93 | 94 | 95 | if __name__ == "__main__": 96 | test_group_size() 97 | -------------------------------------------------------------------------------- /tests/test_python.py: -------------------------------------------------------------------------------- 1 | import time 2 | from pathlib import Path 3 | 4 | from dflow import Step, Workflow, download_artifact, upload_artifact 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Duplicate(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | 'msg': str, 20 | 'num': int, 21 | 'foo': Artifact(Path), 22 | 'idir': Artifact(Path), 23 | }) 24 | 25 | @classmethod 26 | def get_output_sign(cls): 27 | return OPIOSign({ 28 | 'msg': str, 29 | 'bar': Artifact(Path), 30 | 'odir': Artifact(Path), 31 | }) 32 | 33 | @OP.exec_sign_check 34 | def execute( 35 | self, 36 | op_in: OPIO, 37 | ) -> OPIO: 38 | op_out = OPIO({ 39 | "msg": op_in["msg"] * op_in["num"], 40 | "bar": Path("output.txt"), 41 | "odir": Path("todir"), 42 | }) 43 | 44 | content = open(op_in['foo'], "r").read() 45 | open("output.txt", "w").write(content * op_in["num"]) 46 | 47 | Path(op_out['odir']).mkdir() 48 | for ii in ['f1', 'f2']: 49 | (op_out['odir']/ii).write_text(op_in['num'] 50 | * (op_in['idir']/ii).read_text()) 51 | 52 | return op_out 53 | 54 | 55 | def make_idir(): 56 | idir = Path("tidir") 57 | idir.mkdir(exist_ok=True) 58 | (idir / "f1").write_text("foo") 59 | (idir / "f2").write_text("bar") 60 | 61 | 62 | def test_python(): 63 | wf = Workflow(name="python") 64 | 65 | with open("foo.txt", "w") as f: 66 | f.write("Hi") 67 | make_idir() 68 | 69 | artifact0 = upload_artifact("foo.txt") 70 | artifact1 = upload_artifact("tidir") 71 | print(artifact0) 72 | print(artifact1) 73 | step = Step( 74 | name="step", 75 | template=PythonOPTemplate(Duplicate, image="python:3.8"), 76 | parameters={"msg": "Hello", "num": 3}, 77 | artifacts={"foo": artifact0, "idir": artifact1}, 78 | ) 79 | # This step will give output parameter "msg" with value "HelloHelloHello", 80 | # and output artifact "bar" which contains "HiHiHi" 81 | # output artifact "odir/f1" contains foofoofoo and "odir/f2" contains 82 | # barbarbar 83 | wf.add(step) 84 | wf.submit() 85 | 86 | while wf.query_status() in ["Pending", "Running"]: 87 | time.sleep(1) 88 | 89 | assert(wf.query_status() == "Succeeded") 90 | step = wf.query_step(name="step")[0] 91 | assert(step.phase == "Succeeded") 92 | 93 | print(download_artifact(step.outputs.artifacts["bar"])) 94 | print(download_artifact(step.outputs.artifacts["odir"])) 95 | 96 | 97 | if __name__ == "__main__": 98 | test_python() 99 | -------------------------------------------------------------------------------- /tests/test_recurse.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import InputParameter, Inputs, Step, Steps, Workflow 4 | from dflow.python import OP, OPIO, OPIOSign, PythonOPTemplate, upload_packages 5 | 6 | if "__file__" in locals(): 7 | upload_packages.append(__file__) 8 | 9 | 10 | class Plus1(OP): 11 | def __init__(self): 12 | pass 13 | 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign({ 17 | 'iter': int 18 | }) 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign({ 23 | 'iter': int 24 | }) 25 | 26 | @OP.exec_sign_check 27 | def execute( 28 | self, 29 | op_in: OPIO, 30 | ) -> OPIO: 31 | print("This is iter %i" % op_in["iter"]) 32 | return OPIO({ 33 | 'iter': op_in['iter'] + 1 34 | }) 35 | 36 | 37 | def test_recurse(): 38 | steps = Steps(name="iter", inputs=Inputs( 39 | parameters={"iter": InputParameter(value=0), 40 | "limit": InputParameter(value=5)})) 41 | plus1 = Step(name="plus1", 42 | template=PythonOPTemplate(Plus1, 43 | image="python:3.8"), 44 | parameters={"iter": steps.inputs.parameters["iter"]}, 45 | key="iter-%s" % steps.inputs.parameters["iter"]) 46 | steps.add(plus1) 47 | next = Step(name="next", template=steps, 48 | parameters={"iter": plus1.outputs.parameters["iter"]}, 49 | when="%s < %s" % ( 50 | plus1.outputs.parameters["iter"], 51 | steps.inputs.parameters["limit"])) 52 | steps.add(next) 53 | 54 | wf = Workflow("recurse", steps=steps) 55 | wf.submit() 56 | 57 | while wf.query_status() in ["Pending", "Running"]: 58 | time.sleep(1) 59 | 60 | assert(wf.query_status() == "Succeeded") 61 | step = wf.query_step(key="iter-4")[0] 62 | assert(step.phase == "Succeeded") 63 | assert step.outputs.parameters["iter"].value == 5 64 | 65 | 66 | if __name__ == "__main__": 67 | test_recurse() 68 | -------------------------------------------------------------------------------- /tests/test_reuse.py: -------------------------------------------------------------------------------- 1 | import time 2 | 3 | from dflow import InputParameter, Inputs, Step, Steps, Workflow 4 | from dflow.python import OP, OPIO, OPIOSign, PythonOPTemplate, upload_packages 5 | 6 | if "__file__" in locals(): 7 | upload_packages.append(__file__) 8 | 9 | 10 | class Plus1(OP): 11 | def __init__(self): 12 | pass 13 | 14 | @classmethod 15 | def get_input_sign(cls): 16 | return OPIOSign({ 17 | 'iter': int 18 | }) 19 | 20 | @classmethod 21 | def get_output_sign(cls): 22 | return OPIOSign({ 23 | 'iter': int 24 | }) 25 | 26 | @OP.exec_sign_check 27 | def execute( 28 | self, 29 | op_in: OPIO, 30 | ) -> OPIO: 31 | print("This is iter %i" % op_in["iter"]) 32 | return OPIO({ 33 | 'iter': op_in['iter'] + 1 34 | }) 35 | 36 | 37 | def test_reuse(): 38 | steps = Steps(name="iter", inputs=Inputs( 39 | parameters={"iter": InputParameter(value=0), 40 | "limit": InputParameter(value=5)})) 41 | plus1 = Step(name="plus1", 42 | template=PythonOPTemplate(Plus1, 43 | image="python:3.8"), 44 | parameters={"iter": steps.inputs.parameters["iter"]}, 45 | key="iter-%s" % steps.inputs.parameters["iter"]) 46 | steps.add(plus1) 47 | next = Step(name="next", template=steps, 48 | parameters={"iter": plus1.outputs.parameters["iter"]}, 49 | when="%s < %s" % ( 50 | plus1.outputs.parameters["iter"], 51 | steps.inputs.parameters["limit"])) 52 | steps.add(next) 53 | 54 | wf = Workflow("recurse", steps=steps) 55 | wf.submit() 56 | 57 | while wf.query_status() in ["Pending", "Running"]: 58 | time.sleep(1) 59 | 60 | assert(wf.query_status() == "Succeeded") 61 | 62 | step0 = wf.query_step(key="iter-0")[0] 63 | step1 = wf.query_step(key="iter-1")[0] 64 | step1.modify_output_parameter("iter", 3) 65 | 66 | wf = Workflow("recurse-resubmit", steps=steps) 67 | wf.submit(reuse_step=[step0, step1]) 68 | 69 | 70 | if __name__ == "__main__": 71 | test_reuse() 72 | -------------------------------------------------------------------------------- /tests/test_slices.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import Step, Workflow, argo_range 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Hello(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | 'filename': str 20 | }) 21 | 22 | @classmethod 23 | def get_output_sign(cls): 24 | return OPIOSign({ 25 | 'foo': Artifact(str) 26 | }) 27 | 28 | @OP.exec_sign_check 29 | def execute( 30 | self, 31 | op_in: OPIO, 32 | ) -> OPIO: 33 | open(op_in["filename"], "w").write("foo") 34 | op_out = OPIO({ 35 | 'foo': op_in["filename"] 36 | }) 37 | return op_out 38 | 39 | 40 | class Check(OP): 41 | def __init__(self): 42 | pass 43 | 44 | @classmethod 45 | def get_input_sign(cls): 46 | return OPIOSign({ 47 | 'foo': Artifact(List[str]) 48 | }) 49 | 50 | @classmethod 51 | def get_output_sign(cls): 52 | return OPIOSign() 53 | 54 | @OP.exec_sign_check 55 | def execute( 56 | self, 57 | op_in: OPIO, 58 | ) -> OPIO: 59 | print(op_in["foo"]) 60 | assert len(op_in["foo"]) == 2 61 | return OPIO() 62 | 63 | 64 | def test_slices(): 65 | wf = Workflow("slices") 66 | hello = Step("hello", 67 | PythonOPTemplate(Hello, image="python:3.8", 68 | slices=Slices("{{item}}", 69 | input_parameter=["filename"], 70 | output_artifact=["foo"] 71 | ) 72 | ), 73 | parameters={"filename": ["f1.txt", "f2.txt"]}, 74 | with_param=argo_range(2)) 75 | wf.add(hello) 76 | check = Step("check", 77 | PythonOPTemplate(Check, image="python:3.8"), 78 | artifacts={"foo": hello.outputs.artifacts["foo"]}) 79 | wf.add(check) 80 | wf.submit() 81 | 82 | while wf.query_status() in ["Pending", "Running"]: 83 | time.sleep(1) 84 | 85 | assert(wf.query_status() == "Succeeded") 86 | 87 | 88 | if __name__ == "__main__": 89 | test_slices() 90 | -------------------------------------------------------------------------------- /tests/test_subpath_slices.py: -------------------------------------------------------------------------------- 1 | import time 2 | from typing import List 3 | 4 | from dflow import Step, Workflow # , upload_artifact 5 | from dflow.python import (OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, 6 | Slices, upload_packages) 7 | 8 | if "__file__" in locals(): 9 | upload_packages.append(__file__) 10 | 11 | 12 | class Prepare(OP): 13 | def __init__(self): 14 | pass 15 | 16 | @classmethod 17 | def get_input_sign(cls): 18 | return OPIOSign({ 19 | }) 20 | 21 | @classmethod 22 | def get_output_sign(cls): 23 | return OPIOSign({ 24 | 'foo': Artifact(List[str], archive=None) 25 | }) 26 | 27 | @OP.exec_sign_check 28 | def execute( 29 | self, 30 | op_in: OPIO, 31 | ) -> OPIO: 32 | with open("foo1.txt", "w") as f: 33 | f.write("foo1") 34 | with open("foo2.txt", "w") as f: 35 | f.write("foo2") 36 | op_out = OPIO({ 37 | 'foo': ["foo1.txt", "foo2.txt"] 38 | }) 39 | return op_out 40 | 41 | 42 | class Hello(OP): 43 | def __init__(self): 44 | pass 45 | 46 | @classmethod 47 | def get_input_sign(cls): 48 | return OPIOSign({ 49 | 'foo': Artifact(str) 50 | }) 51 | 52 | @classmethod 53 | def get_output_sign(cls): 54 | return OPIOSign() 55 | 56 | @OP.exec_sign_check 57 | def execute( 58 | self, 59 | op_in: OPIO, 60 | ) -> OPIO: 61 | with open(op_in["foo"], "r") as f: 62 | print(f.read()) 63 | return OPIO() 64 | 65 | 66 | def test_subpath_slices(): 67 | wf = Workflow("subpath-slices") 68 | # with open("foo1.txt", "w") as f: 69 | # f.write("foo1") 70 | # with open("foo2.txt", "w") as f: 71 | # f.write("foo2") 72 | # artifact = upload_artifact(["foo1.txt", "foo2.txt"], archive=None) 73 | prepare = Step("prepare", 74 | PythonOPTemplate(Prepare, image="python:3.8")) 75 | wf.add(prepare) 76 | 77 | hello = Step("hello", 78 | PythonOPTemplate(Hello, image="python:3.8", 79 | slices=Slices(sub_path=True, 80 | input_artifact=["foo"] 81 | ) 82 | ), 83 | artifacts={"foo": prepare.outputs.artifacts["foo"]}) 84 | # artifacts={"foo": artifact}) 85 | wf.add(hello) 86 | wf.submit() 87 | 88 | while wf.query_status() in ["Pending", "Running"]: 89 | time.sleep(1) 90 | 91 | assert(wf.query_status() == "Succeeded") 92 | 93 | 94 | if __name__ == "__main__": 95 | test_subpath_slices() 96 | -------------------------------------------------------------------------------- /tutorials/cp2k_opt/input.inp: -------------------------------------------------------------------------------- 1 | &GLOBAL 2 | PROJECT input 3 | PRINT_LEVEL LOW 4 | RUN_TYPE CELL_OPT 5 | &END GLOBAL 6 | 7 | &FORCE_EVAL 8 | METHOD Quickstep 9 | &SUBSYS 10 | &CELL 11 | A 5.43069750 0.00000000 0.00000000 12 | B 0.00000000 5.43069750 0.00000000 13 | C 0.00000000 0.00000000 5.43069750 14 | PERIODIC XYZ #Direction of applied PBC (geometry aspect) 15 | &END CELL 16 | &COORD 17 | Si 0.00000000 0.00000000 0.00000000 18 | Si 0.00000000 2.71534870 2.71534870 19 | Si 2.71534870 2.71534870 0.00000000 20 | Si 2.71534870 0.00000000 2.71534870 21 | Si 4.07302310 1.35767440 4.07302310 22 | Si 1.35767440 1.35767440 1.35767440 23 | Si 1.35767440 4.07302310 4.07302310 24 | Si 4.07302310 4.07302310 1.35767440 25 | &END COORD 26 | &KIND Si 27 | ELEMENT Si 28 | BASIS_SET DZVP-MOLOPT-SR-GTH-q4 29 | POTENTIAL GTH-PBE 30 | &END KIND 31 | &END SUBSYS 32 | 33 | &DFT 34 | BASIS_SET_FILE_NAME BASIS_MOLOPT 35 | POTENTIAL_FILE_NAME POTENTIAL 36 | # WFN_RESTART_FILE_NAME input-RESTART.wfn 37 | CHARGE 0 #Net charge 38 | MULTIPLICITY 1 #Spin multiplicity 39 | &KPOINTS 40 | SCHEME MONKHORST-PACK 2 2 2 41 | &END KPOINTS 42 | &QS 43 | EPS_DEFAULT 1E-8 #This is default. Set all EPS_xxx to values such that the energy will be correct up to this value 44 | &END QS 45 | &POISSON 46 | PERIODIC XYZ #Direction(s) of PBC for calculating electrostatics 47 | PSOLVER PERIODIC #The way to solve Poisson equation 48 | &END POISSON 49 | &XC 50 | &XC_FUNCTIONAL PBE 51 | &END XC_FUNCTIONAL 52 | &END XC 53 | &MGRID 54 | CUTOFF 300 55 | REL_CUTOFF 40 56 | &END MGRID 57 | &SCF 58 | MAX_SCF 128 59 | EPS_SCF 2.0E-06 #Convergence threshold of density matrix during SCF 60 | # SCF_GUESS RESTART #Use wavefunction from WFN_RESTART_FILE_NAME file as initial guess 61 | &DIAGONALIZATION 62 | ALGORITHM STANDARD #Algorithm for diagonalization. DAVIDSON is faster for large systems 63 | &END DIAGONALIZATION 64 | &MIXING #How to mix old and new density matrices 65 | METHOD BROYDEN_MIXING #PULAY_MIXING is also a good alternative 66 | ALPHA 0.4 #Default. Mixing 40% of new density matrix with the old one 67 | NBROYDEN 8 #Default is 4. Number of previous steps stored for the actual mixing scheme 68 | &END MIXING 69 | &PRINT 70 | &RESTART #Note: Use "&RESTART OFF" can prevent generating wfn file 71 | BACKUP_COPIES 0 #Maximum number of backup copies of wfn file. 0 means never 72 | &END RESTART 73 | &END PRINT 74 | &END SCF 75 | &END DFT 76 | STRESS_TENSOR ANALYTICAL #Compute full stress tensor analytically 77 | &END FORCE_EVAL 78 | 79 | &MOTION 80 | &CELL_OPT 81 | MAX_ITER 250 #Maximum number of geometry optimization 82 | EXTERNAL_PRESSURE 1.01325E+00 #External pressure for cell optimization (bar) 83 | CONSTRAINT NONE #Constraint of cell length, can be: NONE, X, Y, Z, XY, XZ, YZ 84 | KEEP_ANGLES F #If T, then cell angles will be kepted 85 | KEEP_SYMMETRY F #If T, then cell symmetry specified by &CELL / SYMMETRY will be kepted 86 | KEEP_SPACE_GROUP F #If T, then space group will be detected and preserved 87 | TYPE DIRECT_CELL_OPT #Geometry and cell are optimized at the same time. Can also be GEO_OPT, MD 88 | #The following thresholds of optimization convergence are the default ones 89 | MAX_DR 3E-3 #Maximum geometry change 90 | RMS_DR 1.5E-3 #RMS geometry change 91 | MAX_FORCE 4.5E-4 #Maximum force 92 | RMS_FORCE 3E-4 #RMS force 93 | PRESSURE_TOLERANCE 100 #Pressure tolerance (w.r.t EXTERNAL_PRESSURE) 94 | OPTIMIZER BFGS #Can also be CG (more robust for difficult cases) or LBFGS 95 | &BFGS 96 | TRUST_RADIUS 0.2 #Trust radius (maximum stepsize) in Angstrom 97 | # RESTART_HESSIAN T #If read initial Hessian, uncomment this line and specify the file in the next line 98 | # RESTART_FILE_NAME to_be_specified 99 | &END BFGS 100 | &END CELL_OPT 101 | &PRINT 102 | &TRAJECTORY 103 | FORMAT pdb 104 | &END TRAJECTORY 105 | &RESTART 106 | BACKUP_COPIES 0 #Maximum number of backing up restart file, 0 means never 107 | &END RESTART 108 | &RESTART_HISTORY 109 | &EACH 110 | CELL_OPT 0 #How often a history .restart file is generated, 0 means never 111 | &END EACH 112 | &END RESTART_HISTORY 113 | &END PRINT 114 | &END MOTION 115 | -------------------------------------------------------------------------------- /tutorials/dflow-conditional.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "attachments": {}, 5 | "cell_type": "markdown", 6 | "metadata": {}, 7 | "source": [ 8 | "# Task\n", 9 | "In this tutorial, you will learn how to use if_expression." 10 | ] 11 | }, 12 | { 13 | "cell_type": "code", 14 | "execution_count": 1, 15 | "metadata": {}, 16 | "outputs": [], 17 | "source": [ 18 | "import random\n", 19 | "\n", 20 | "from dflow import (OutputArtifact, OutputParameter, Outputs, Step, Steps,\n", 21 | " Workflow, if_expression)\n", 22 | "from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate" 23 | ] 24 | }, 25 | { 26 | "attachments": {}, 27 | "cell_type": "markdown", 28 | "metadata": {}, 29 | "source": [ 30 | "**Random OP**\n", 31 | "This OP's outputs include five parameters, as shown below.\n", 32 | "- `is_header` will get a bool type value through random function, true or false.\n", 33 | "- `msg1` is \"head\",\n", 34 | "- `msg2` is \"tail\",\n", 35 | "- `foo` is a text file with content \"head\"\n", 36 | "- `bar` is a text file with content \"tail\"\n" 37 | ] 38 | }, 39 | { 40 | "cell_type": "code", 41 | "execution_count": 2, 42 | "metadata": {}, 43 | "outputs": [], 44 | "source": [ 45 | "class Random(OP):\n", 46 | " @classmethod\n", 47 | " def get_input_sign(cls):\n", 48 | " return OPIOSign()\n", 49 | "\n", 50 | " @classmethod\n", 51 | " def get_output_sign(cls):\n", 52 | " return OPIOSign({\n", 53 | " \"is_head\": bool,\n", 54 | " \"msg1\": str,\n", 55 | " \"msg2\": str,\n", 56 | " \"foo\": Artifact(str),\n", 57 | " \"bar\": Artifact(str)\n", 58 | " })\n", 59 | "\n", 60 | " @OP.exec_sign_check\n", 61 | " def execute(\n", 62 | " self,\n", 63 | " op_in: OPIO,\n", 64 | " ) -> OPIO:\n", 65 | " open(\"foo.txt\", \"w\").write(\"head\")\n", 66 | " open(\"bar.txt\", \"w\").write(\"tail\")\n", 67 | " if random.random() < 0.5:\n", 68 | " is_head = True\n", 69 | " else:\n", 70 | " is_head = False\n", 71 | " return OPIO({\n", 72 | " \"is_head\": is_head,\n", 73 | " \"msg1\": \"head\",\n", 74 | " \"msg2\": \"tail\",\n", 75 | " \"foo\": \"foo.txt\",\n", 76 | " \"bar\": \"bar.txt\"\n", 77 | " })" 78 | ] 79 | }, 80 | { 81 | "attachments": {}, 82 | "cell_type": "markdown", 83 | "metadata": {}, 84 | "source": [ 85 | "Believe you remember what you have learned from `dflow-reuse.py` about `Steps`. There, `Steps` has a output parameter and artifact.\n", 86 | "\n", 87 | "The Random defined above is used here" 88 | ] 89 | }, 90 | { 91 | "cell_type": "code", 92 | "execution_count": 3, 93 | "metadata": {}, 94 | "outputs": [], 95 | "source": [ 96 | "steps = Steps(\n", 97 | " name=\"conditional-steps\", \n", 98 | " outputs=Outputs(\n", 99 | " parameters={\"msg\": OutputParameter()},\n", 100 | " artifacts={\"res\": OutputArtifact()}\n", 101 | " ),\n", 102 | ")\n", 103 | "\n", 104 | "import sys\n", 105 | "random_step = Step(\n", 106 | " name=\"random\",\n", 107 | " template=PythonOPTemplate(Random, image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\")\n", 108 | ")\n", 109 | "\n", 110 | "steps.add(random_step)" 111 | ] 112 | }, 113 | { 114 | "attachments": {}, 115 | "cell_type": "markdown", 116 | "metadata": {}, 117 | "source": [ 118 | "# if_expression\n", 119 | "\n", 120 | "In this part, `if_expression` is used to determine which message and artifact will be used. To illustrate: \n", 121 | "- if `is_head` is true, then `steps.outputs.parameters`'s `msg` will be `msg1` and `res` will be `foo` that includes the content of head.\n", 122 | "- else, `steps.outputs.parameters`'s `msg` will be `msg2` and `res` will be `bar` that includes the content of tail." 123 | ] 124 | }, 125 | { 126 | "cell_type": "code", 127 | "execution_count": 4, 128 | "metadata": {}, 129 | "outputs": [], 130 | "source": [ 131 | "steps.outputs.parameters[\"msg\"].value_from_expression = if_expression(\n", 132 | " _if=random_step.outputs.parameters[\"is_head\"],\n", 133 | " _then=random_step.outputs.parameters[\"msg1\"],\n", 134 | " _else=random_step.outputs.parameters[\"msg2\"])\n", 135 | "\n", 136 | "steps.outputs.artifacts[\"res\"].from_expression = if_expression(\n", 137 | " _if=random_step.outputs.parameters[\"is_head\"],\n", 138 | " _then=random_step.outputs.artifacts[\"foo\"],\n", 139 | " _else=random_step.outputs.artifacts[\"bar\"])" 140 | ] 141 | }, 142 | { 143 | "cell_type": "code", 144 | "execution_count": 5, 145 | "metadata": {}, 146 | "outputs": [ 147 | { 148 | "name": "stdout", 149 | "output_type": "stream", 150 | "text": [ 151 | "Workflow has been submitted (ID: conditional-crb56, UID: 552b993e-d5e0-4bda-9e80-76691b044112)\n" 152 | ] 153 | } 154 | ], 155 | "source": [ 156 | "wf = Workflow(name=\"conditional\", steps=steps)\n", 157 | "wf.submit();" 158 | ] 159 | } 160 | ], 161 | "metadata": { 162 | "kernelspec": { 163 | "display_name": "base", 164 | "language": "python", 165 | "name": "python3" 166 | }, 167 | "language_info": { 168 | "codemirror_mode": { 169 | "name": "ipython", 170 | "version": 3 171 | }, 172 | "file_extension": ".py", 173 | "mimetype": "text/x-python", 174 | "name": "python", 175 | "nbconvert_exporter": "python", 176 | "pygments_lexer": "ipython3", 177 | "version": "3.9.7" 178 | }, 179 | "orig_nbformat": 4, 180 | "vscode": { 181 | "interpreter": { 182 | "hash": "0d3b56f35093c43ef3a807ec55a8177d3d51ef411c9a162a01ec53961f392e60" 183 | } 184 | } 185 | }, 186 | "nbformat": 4, 187 | "nbformat_minor": 2 188 | } 189 | -------------------------------------------------------------------------------- /tutorials/dflow-function.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "attachments": {}, 5 | "cell_type": "markdown", 6 | "metadata": {}, 7 | "source": [ 8 | "# Task\n", 9 | "We want to achieve the following workflow:\n", 10 | "\n", 11 | "Step 1. \n", 12 | "1. Echo a string to msg.txt \n", 13 | " \n", 14 | "2. Output the length of the string as a number\n", 15 | " \n", 16 | "Step 2.\n", 17 | "1. Duplicate the content in the file from step 1 two times and redirect it to a new file\n", 18 | " \n", 19 | "2. Get the number from step 1 and times the number by 2 and output it" 20 | ] 21 | }, 22 | { 23 | "attachments": {}, 24 | "cell_type": "markdown", 25 | "metadata": {}, 26 | "source": [ 27 | "##### In this tutorial, we will complete the task using function OP" 28 | ] 29 | }, 30 | { 31 | "attachments": {}, 32 | "cell_type": "markdown", 33 | "metadata": {}, 34 | "source": [ 35 | "import necessary packages" 36 | ] 37 | }, 38 | { 39 | "cell_type": "code", 40 | "execution_count": null, 41 | "metadata": {}, 42 | "outputs": [], 43 | "source": [ 44 | "import sys\n", 45 | "from pathlib import Path\n", 46 | "\n", 47 | "from dflow import Workflow\n", 48 | "from dflow.python import OP, Artifact" 49 | ] 50 | }, 51 | { 52 | "attachments": {}, 53 | "cell_type": "markdown", 54 | "metadata": {}, 55 | "source": [ 56 | "For step 1: \n", 57 | "\n", 58 | "This OP is to write files. In the example, we try to write a file containing message of string format, and output a number of int format.\n", 59 | "- input:\n", 60 | " - \"msg\": the input message\n", 61 | "- output:\n", 62 | " - \"out_art\": file containing the input message\n", 63 | " - \"length\": length of the input message" 64 | ] 65 | }, 66 | { 67 | "cell_type": "code", 68 | "execution_count": null, 69 | "metadata": {}, 70 | "outputs": [], 71 | "source": [ 72 | "@OP.function(image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\")\n", 73 | "def write_file(msg: str) -> {\"out_art\": Artifact(Path), \"length\": int}:\n", 74 | " with open(\"msg.txt\",\"w\") as f:\n", 75 | " f.write(msg)\n", 76 | " \n", 77 | " return {\n", 78 | " \"out_art\": Path(\"msg.txt\"),\n", 79 | " \"length\": len(msg),\n", 80 | " }" 81 | ] 82 | }, 83 | { 84 | "attachments": {}, 85 | "cell_type": "markdown", 86 | "metadata": {}, 87 | "source": [ 88 | "For Step 2:\n", 89 | "\n", 90 | "This OP is to duplicate the content in the file from in_art, and to multiply the input number by 2." 91 | ] 92 | }, 93 | { 94 | "cell_type": "code", 95 | "execution_count": null, 96 | "metadata": {}, 97 | "outputs": [], 98 | "source": [ 99 | "@OP.function(image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\")\n", 100 | "def duplicate(in_art: Artifact(Path), in_num: int) -> {\"out_art\": Artifact(Path), \"out_num\": int}:\n", 101 | " with open(in_art, \"r\") as f:\n", 102 | " content = f.read()\n", 103 | " with open(\"bar.txt\", \"w\") as f:\n", 104 | " f.write(content * 2)\n", 105 | "\n", 106 | " return {\n", 107 | " \"out_art\": Path(\"bar.txt\"),\n", 108 | " \"out_num\": in_num * 2,\n", 109 | " }" 110 | ] 111 | }, 112 | { 113 | "attachments": {}, 114 | "cell_type": "markdown", 115 | "metadata": {}, 116 | "source": [ 117 | "After defining OPs, call the OPs in series in the context of a workflow, which will connect them together as a workflow and submit it finally." 118 | ] 119 | }, 120 | { 121 | "cell_type": "code", 122 | "execution_count": null, 123 | "metadata": {}, 124 | "outputs": [], 125 | "source": [ 126 | "with Workflow(name=\"python\") as wf:\n", 127 | " out = write_file(msg=\"HelloWorld!\")\n", 128 | " duplicate(in_num=out[\"length\"], in_art=out[\"out_art\"])" 129 | ] 130 | } 131 | ], 132 | "metadata": { 133 | "kernelspec": { 134 | "display_name": "base", 135 | "language": "python", 136 | "name": "python3" 137 | }, 138 | "language_info": { 139 | "codemirror_mode": { 140 | "name": "ipython", 141 | "version": 3 142 | }, 143 | "file_extension": ".py", 144 | "mimetype": "text/x-python", 145 | "name": "python", 146 | "nbconvert_exporter": "python", 147 | "pygments_lexer": "ipython3", 148 | "version": "3.9.7" 149 | }, 150 | "orig_nbformat": 4, 151 | "vscode": { 152 | "interpreter": { 153 | "hash": "0d3b56f35093c43ef3a807ec55a8177d3d51ef411c9a162a01ec53961f392e60" 154 | } 155 | } 156 | }, 157 | "nbformat": 4, 158 | "nbformat_minor": 2 159 | } 160 | -------------------------------------------------------------------------------- /tutorials/dflow-slices.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "attachments": {}, 5 | "cell_type": "markdown", 6 | "metadata": {}, 7 | "source": [ 8 | "# Task\n", 9 | "We want to achieve the following workflow:\n", 10 | "\n", 11 | "Step 1. \n", 12 | "1. Write \"Hello\" in a file \n", 13 | "\n", 14 | "Step 2:\n", 15 | "1. Check the file from step 1\n", 16 | "\n", 17 | "**However, this is too simple for dflow. We want to write \"Hello\" to 10 different files in parallel.**" 18 | ] 19 | }, 20 | { 21 | "cell_type": "code", 22 | "execution_count": 1, 23 | "metadata": {}, 24 | "outputs": [], 25 | "source": [ 26 | "from typing import List\n", 27 | "\n", 28 | "from dflow import Step, Workflow\n", 29 | "from dflow.python import OP, OPIO, Artifact, OPIOSign, PythonOPTemplate, Slices" 30 | ] 31 | }, 32 | { 33 | "attachments": {}, 34 | "cell_type": "markdown", 35 | "metadata": {}, 36 | "source": [ 37 | "** Hello OP **\n", 38 | "- input: \n", 39 | " - 'filename': str that is the name of the file to generate" 40 | ] 41 | }, 42 | { 43 | "cell_type": "code", 44 | "execution_count": 2, 45 | "metadata": {}, 46 | "outputs": [], 47 | "source": [ 48 | "class Hello(OP):\n", 49 | " def __init__(self):\n", 50 | " pass\n", 51 | "\n", 52 | " @classmethod\n", 53 | " def get_input_sign(cls):\n", 54 | " return OPIOSign({\n", 55 | " 'filename': str\n", 56 | " })\n", 57 | "\n", 58 | " @classmethod\n", 59 | " def get_output_sign(cls):\n", 60 | " return OPIOSign({\n", 61 | " 'out_art': Artifact(str)\n", 62 | " })\n", 63 | "\n", 64 | " @OP.exec_sign_check\n", 65 | " def execute(\n", 66 | " self,\n", 67 | " op_in: OPIO,\n", 68 | " ) -> OPIO:\n", 69 | " file_num = int(op_in[\"filename\"].split('.')[0][1:])\n", 70 | " with open(op_in[\"filename\"], \"w\") as f:\n", 71 | " f.write(\"Hello\" * file_num)\n", 72 | " op_out = OPIO({\n", 73 | " 'out_art': op_in[\"filename\"]\n", 74 | " })\n", 75 | " return op_out" 76 | ] 77 | }, 78 | { 79 | "attachments": {}, 80 | "cell_type": "markdown", 81 | "metadata": {}, 82 | "source": [ 83 | "**Check OP**\n", 84 | "It is to check the files that have been written and print the files' pathes. `print(op_in[\"files\"])` " 85 | ] 86 | }, 87 | { 88 | "cell_type": "code", 89 | "execution_count": 3, 90 | "metadata": {}, 91 | "outputs": [], 92 | "source": [ 93 | "class Check(OP):\n", 94 | " def __init__(self):\n", 95 | " pass\n", 96 | "\n", 97 | " @classmethod\n", 98 | " def get_input_sign(cls):\n", 99 | " return OPIOSign({\n", 100 | " 'files': Artifact(List[str])\n", 101 | " })\n", 102 | "\n", 103 | " @classmethod\n", 104 | " def get_output_sign(cls):\n", 105 | " return OPIOSign()\n", 106 | "\n", 107 | " @OP.exec_sign_check\n", 108 | " def execute(\n", 109 | " self,\n", 110 | " op_in: OPIO,\n", 111 | " ) -> OPIO:\n", 112 | " print(op_in[\"files\"])\n", 113 | " return OPIO()" 114 | ] 115 | }, 116 | { 117 | "attachments": {}, 118 | "cell_type": "markdown", 119 | "metadata": {}, 120 | "source": [ 121 | "And then we can use OP above to write files repeatedly using Slices. This step is using Hello OP and python: 3.8 image.\n", 122 | "\n", 123 | "**Slices**\n", 124 | "We need to define input_parameter to be sliced and output_artifact to be stacked." 125 | ] 126 | }, 127 | { 128 | "cell_type": "code", 129 | "execution_count": 4, 130 | "metadata": {}, 131 | "outputs": [ 132 | { 133 | "name": "stdout", 134 | "output_type": "stream", 135 | "text": [ 136 | "Workflow has been submitted (ID: slices-ldtjd, UID: 2842d9b9-9b1d-4319-aaff-d59866213278)\n" 137 | ] 138 | } 139 | ], 140 | "source": [ 141 | "import sys\n", 142 | "wf = Workflow(\"slices\")\n", 143 | "hello = Step(\"hello\",\n", 144 | " PythonOPTemplate(Hello, image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\",\n", 145 | " slices=Slices(\"{{item}}\",\n", 146 | " input_parameter=[\"filename\"],\n", 147 | " output_artifact=[\"out_art\"]\n", 148 | " )\n", 149 | " ),\n", 150 | " parameters={\"filename\": [f\"f{x}.txt\" for x in range(3)]},\n", 151 | " with_param=range(3))\n", 152 | "check = Step(\"check\",\n", 153 | " PythonOPTemplate(Check, image=f\"python:{sys.version_info.major}.{sys.version_info.minor}\"),\n", 154 | " artifacts={\"files\": hello.outputs.artifacts[\"out_art\"]},\n", 155 | " )\n", 156 | "wf.add(hello)\n", 157 | "wf.add(check)\n", 158 | "wf.submit();" 159 | ] 160 | } 161 | ], 162 | "metadata": { 163 | "kernelspec": { 164 | "display_name": "base", 165 | "language": "python", 166 | "name": "python3" 167 | }, 168 | "language_info": { 169 | "codemirror_mode": { 170 | "name": "ipython", 171 | "version": 3 172 | }, 173 | "file_extension": ".py", 174 | "mimetype": "text/x-python", 175 | "name": "python", 176 | "nbconvert_exporter": "python", 177 | "pygments_lexer": "ipython3", 178 | "version": "3.9.7" 179 | }, 180 | "orig_nbformat": 4, 181 | "vscode": { 182 | "interpreter": { 183 | "hash": "0d3b56f35093c43ef3a807ec55a8177d3d51ef411c9a162a01ec53961f392e60" 184 | } 185 | } 186 | }, 187 | "nbformat": 4, 188 | "nbformat_minor": 2 189 | } 190 | -------------------------------------------------------------------------------- /tutorials/dflow-slurm.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "code", 5 | "execution_count": null, 6 | "metadata": {}, 7 | "outputs": [], 8 | "source": [ 9 | "from dflow import upload_artifact, Workflow, Step\n", 10 | "from dflow.python import PythonOPTemplate, OP, OPIO, OPIOSign, Artifact\n", 11 | "from dflow.plugins.dispatcher import DispatcherExecutor\n", 12 | "from pathlib import Path" 13 | ] 14 | }, 15 | { 16 | "attachments": {}, 17 | "cell_type": "markdown", 18 | "metadata": {}, 19 | "source": [ 20 | "**Duplicate OP**\n", 21 | "\n", 22 | "This OP duplicate input message and content of input file by a specified number of times as output message and output file, respectively.\n", 23 | "- input:\n", 24 | " - \"msg\": the input message\n", 25 | " - \"num\": the number of times\n", 26 | " - \"foo\": the input file\n", 27 | "- output:\n", 28 | " - \"msg\": the output message\n", 29 | " - \"bar\": the output file" 30 | ] 31 | }, 32 | { 33 | "cell_type": "code", 34 | "execution_count": null, 35 | "metadata": {}, 36 | "outputs": [], 37 | "source": [ 38 | "class Duplicate(OP):\n", 39 | " def __init__(self):\n", 40 | " pass\n", 41 | "\n", 42 | " @classmethod\n", 43 | " def get_input_sign(cls):\n", 44 | " return OPIOSign({\n", 45 | " \"msg\": str,\n", 46 | " \"num\": int,\n", 47 | " \"foo\": Artifact(Path),\n", 48 | " })\n", 49 | "\n", 50 | " @classmethod\n", 51 | " def get_output_sign(cls):\n", 52 | " return OPIOSign({\n", 53 | " \"msg\": str,\n", 54 | " \"bar\": Artifact(Path),\n", 55 | " })\n", 56 | "\n", 57 | " @OP.exec_sign_check\n", 58 | " def execute(\n", 59 | " self,\n", 60 | " op_in: OPIO,\n", 61 | " ) -> OPIO:\n", 62 | " with open(op_in[\"foo\"], \"r\") as f:\n", 63 | " content = f.read()\n", 64 | " with open(\"bar.txt\", \"w\") as f:\n", 65 | " f.write(content * op_in[\"num\"])\n", 66 | "\n", 67 | " op_out = OPIO({\n", 68 | " \"msg\": op_in[\"msg\"] * op_in[\"num\"],\n", 69 | " \"bar\": Path(\"bar.txt\"),\n", 70 | " })\n", 71 | " return op_out" 72 | ] 73 | }, 74 | { 75 | "attachments": {}, 76 | "cell_type": "markdown", 77 | "metadata": {}, 78 | "source": [ 79 | "[DPDispatcher](https://github.com/deepmodeling/dpdispatcher) is a python package used to generate HPC scheduler systems (Slurm/PBS/LSF) or Bohrium jobs input scripts and submit these scripts and poke until they finish. Dflow provides `DispatcherExecutor` plugin to invoke dispatcher as executor to complete a certain step\n", 80 | "\n", 81 | "For SSH authentication, one can use password, specify path of private key file locally, or upload authorized private key to each node (or equivalently add each node to the authorized host list). In this example, password is used.\n", 82 | "\n", 83 | "For configuring extra machine, resources or task parameters for dispatcher, use `DispatcherExecutor(..., machine_dict=m, resources_dict=r, task_dict=t)`." 84 | ] 85 | }, 86 | { 87 | "cell_type": "code", 88 | "execution_count": null, 89 | "metadata": {}, 90 | "outputs": [], 91 | "source": [ 92 | "dispatcher_executor = DispatcherExecutor(\n", 93 | " host=\"your-cluster-address\",\n", 94 | " username=\"your-login-username\",\n", 95 | " password=\"your-login-password\",\n", 96 | ")\n", 97 | "\n", 98 | "with open(\"foo.txt\", \"w\") as f:\n", 99 | " f.write(\"Hello world!\")\n", 100 | "\n", 101 | "step = Step(\n", 102 | " \"duplicate\",\n", 103 | " PythonOPTemplate(Duplicate),\n", 104 | " parameters={\"msg\": \"Hello\", \"num\": 2},\n", 105 | " artifacts={\"foo\": upload_artifact(\"foo.txt\")},\n", 106 | " executor=dispatcher_executor,\n", 107 | ")\n", 108 | "\n", 109 | "wf = Workflow(name=\"slurm\")\n", 110 | "wf.add(step)\n", 111 | "wf.submit()" 112 | ] 113 | } 114 | ], 115 | "metadata": { 116 | "kernelspec": { 117 | "display_name": "Python 3.9.12 ('base')", 118 | "language": "python", 119 | "name": "python3" 120 | }, 121 | "language_info": { 122 | "codemirror_mode": { 123 | "name": "ipython", 124 | "version": 3 125 | }, 126 | "file_extension": ".py", 127 | "mimetype": "text/x-python", 128 | "name": "python", 129 | "nbconvert_exporter": "python", 130 | "pygments_lexer": "ipython3", 131 | "version": "3.9.7" 132 | }, 133 | "orig_nbformat": 4, 134 | "vscode": { 135 | "interpreter": { 136 | "hash": "65b01cda8a5255d697b7c650722434fd8759cb966fc0703c59c131e9aaea8cdf" 137 | } 138 | } 139 | }, 140 | "nbformat": 4, 141 | "nbformat_minor": 2 142 | } 143 | -------------------------------------------------------------------------------- /tutorials/imgs/access_one_node.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/access_one_node.png -------------------------------------------------------------------------------- /tutorials/imgs/argoui_main_page.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/argoui_main_page.png -------------------------------------------------------------------------------- /tutorials/imgs/connection_warning.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/connection_warning.png -------------------------------------------------------------------------------- /tutorials/imgs/minikube_image_bug.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/minikube_image_bug.png -------------------------------------------------------------------------------- /tutorials/imgs/minikube_start_fail_bug.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/minikube_start_fail_bug.png -------------------------------------------------------------------------------- /tutorials/imgs/workflow_overview.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/deepmodeling/dflow/1ec8cdfc34c9e1b28654289fb544bef2f4c12181/tutorials/imgs/workflow_overview.png -------------------------------------------------------------------------------- /tutorials/install_manual_linuxos.md: -------------------------------------------------------------------------------- 1 | # Install on LinuxOS 2 | ## Container Engine 3 | There are different container engine. But we recommend to use Docker. 4 | This guide shows your how to install Docker Engine (not Docker Desktop). This is a useful guide if you are using dflow on a cloud. 5 | ### Download 6 | Go to Docker official website to follow the official download guide: [Get Docker](https://docs.docker.com/engine/install/#server) 7 | 8 | We will show you how to install Docker Engine on a Ubuntu platform: [Install Docker Engine on Ubuntu](https://docs.docker.com/engine/install/ubuntu/) 9 | 10 | ### Install using the repository 11 | This guide is copied from: https://docs.docker.com/engine/install/ubuntu/#install-using-the-repository. 12 | 13 | #### Set up the repository 14 | 1. Update the `apt` package index and install packages to allow `apt` to use a 15 | repository over HTTPS: 16 | 17 | ```console 18 | sudo apt-get update 19 | 20 | sudo apt-get install \ 21 | ca-certificates \ 22 | curl \ 23 | gnupg \ 24 | lsb-release 25 | ``` 26 | 27 | 2. Add Docker's official GPG key: 28 | 29 | ```console 30 | sudo mkdir -p /etc/apt/keyrings 31 | curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg 32 | ``` 33 | 34 | 3. Use the following command to set up the repository: 35 | 36 | ```console 37 | echo \ 38 | "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \ 39 | $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null 40 | ``` 41 | 42 | #### Install Docker Engine 43 | 44 | 1. Update the `apt` package index, and install the _latest version_ of Docker 45 | Engine, containerd, and Docker Compose, or go to the next step to install a specific version: 46 | 47 | ```console 48 | sudo apt-get update 49 | sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin 50 | ``` 51 | 52 | > Receiving a GPG error when running `apt-get update`? 53 | > 54 | > Your default umask may not be set correctly, causing the public key file 55 | > for the repo to not be detected. Run the following command and then try to 56 | > update your repo again: `sudo chmod a+r /etc/apt/keyrings/docker.gpg`. 57 | 58 | 2. To install a _specific version_ of Docker Engine, list the available versions 59 | in the repo, then select and install: 60 | 61 | a. List the versions available in your repo: 62 | 63 | ```console 64 | apt-cache madison docker-ce 65 | 66 | docker-ce | 5:20.10.16~3-0~ubuntu-jammy | https://download.docker.com/linux/ubuntu jammy/stable amd64 Packages 67 | docker-ce | 5:20.10.15~3-0~ubuntu-jammy | https://download.docker.com/linux/ubuntu jammy/stable amd64 Packages 68 | docker-ce | 5:20.10.14~3-0~ubuntu-jammy | https://download.docker.com/linux/ubuntu jammy/stable amd64 Packages 69 | docker-ce | 5:20.10.13~3-0~ubuntu-jammy | https://download.docker.com/linux/ubuntu jammy/stable amd64 Packages 70 | ``` 71 | 72 | b. Install a specific version using the version string from the second column, 73 | for example, `5:20.10.16~3-0~ubuntu-jammy`. 74 | 75 | ```console 76 | sudo apt-get install docker-ce= docker-ce-cli= containerd.io docker-compose-plugin 77 | ``` 78 | 79 | ### Verify the installation 80 | Run the `hello-world` official image to verify the installation. 81 | ```bash 82 | docker run hello-world 83 | ``` 84 | This command downloads a test image and runs it in a container. When the 85 | container runs, it prints a message and exits. 86 | 87 | ## Kubernetes 88 | To use kubernetes, we can install minikube, a local kubernetes. 89 | 90 | ### Before installation 91 | Check if you are the root user. Minikube does not allow root user to start. Follow the following guide to add a user. See the original issue here: https://github.com/kubernetes/minikube/issues/7903 92 | #### Add new user 93 | ```bash 94 | adduser developer 95 | #add the user. Follow the prompt 96 | usermod -aG sudo developer 97 | ``` 98 | #### Login to the newly created user 99 | ```bash 100 | su - developer 101 | ``` 102 | #### Add user to the Docker Group 103 | ```bash 104 | sudo groupadd docker 105 | sudo usermod -aG docker $USER 106 | - Re-Login or Restart the Server 107 | ``` 108 | 109 | ### Download 110 | #### Up-to-date version 111 | Go to minikube official website to download: [Minikube](https://minikube.sigs.k8s.io/docs/start/) 112 | #### Older version 113 | **NOTE: this is the recommended route for users in China**. The minikube image repository for the up-to-date version is not correct, which leads to minikube start failure. 114 | - Install older version from the official github repository: [Releases · kubernetes/minikube](https://github.com/kubernetes/minikube/releases) 115 | - Install older version using [mirror in China](https://npmmirror.com/): [minikube](https://registry.npmmirror.com/binary.html?path=minikube/) (The latest version that works is 1.25.2) 116 | 117 | Command to install minikube older version (1.25.2) 118 | ```bash 119 | curl -o minikube -L https://registry.npmmirror.com/-/binary/minikube/v1.25.2/minikube-linux-amd64 120 | ``` 121 | #### Install 122 | Install you downloaded minikube to `/usr/local/bin`: 123 | ```bash 124 | sudo install minikube /usr/local/bin/minikube 125 | ``` 126 | 127 | #### Verify the installation 128 | ```bash 129 | minikube version 130 | ``` 131 | 132 | 133 | -------------------------------------------------------------------------------- /tutorials/install_manual_macos.md: -------------------------------------------------------------------------------- 1 | # Install on MacOS 2 | ## Container Engine 3 | There are different container engine. But we recommend to use Docker. 4 | ### Download 5 | Go to Docker official website to download: [Get Docker](https://docs.docker.com/desktop/install/mac-install/) 6 | ### Install 7 | Find the `Docker.dmg` on the laptop. Install it as a common Mac software. 8 | ### Verify the installation 9 | Run the `hello-world` official image to verify the installation. 10 | ```bash 11 | docker run hello-world 12 | ``` 13 | This command downloads a test image and runs it in a container. When the 14 | container runs, it prints a message and exits. 15 | 16 | ### Start Docker 17 | Click Docker Desktop app to start. 18 | 19 | ## Kubernetes 20 | To use kubernetes on the laptop, we can install minikube, a local kubernetes. 21 | ### Download 22 | 23 | #### Up-to-date version 24 | Go to minikube official website to download: [Minikube](https://minikube.sigs.k8s.io/docs/start/) 25 | #### Older version 26 | **NOTE: this is the recommended route for users in China**. The minikube image repository for the up-to-date version is not correct, which leads to minikube start failure. 27 | - Install older version from the official github repository: [Releases · kubernetes/minikube](https://github.com/kubernetes/minikube/releases) 28 | - Install older version using [mirror in China](https://npmmirror.com/): [minikube](https://registry.npmmirror.com/binary.html?path=minikube/) (The latest version that works is 1.25.2) 29 | 30 | Command to install minikube older version (1.25.2) 31 | ```bash 32 | curl -o minikube -L https://registry.npmmirror.com/-/binary/minikube/v1.25.2/minikube-linux-amd64 33 | ``` 34 | 35 | #### Install 36 | Install you downloaded minikube to `/usr/local/bin`: 37 | ```bash 38 | sudo install minikube /usr/local/bin/minikube 39 | ``` 40 | 41 | #### Verify the installation 42 | ```bash 43 | minikube version 44 | ``` -------------------------------------------------------------------------------- /tutorials/install_manual_windowsos.md: -------------------------------------------------------------------------------- 1 | # Install and set up on WindowsOS 2 | ## Container engine: Docker 3 | ### Download 4 | Go to the [Docker website](https://docs.docker.com/get-docker/) and download the **Docker Desktop for Windows**. Then you can follow the instructions on installer step by step. 5 | ### WSL Linux kernel 6 | When you click and run docker desktop, you will see ***"Docker Desktop stopping"*** on the window. Wait a few seconds and then you will be informed that WSL 2 installation is incomplete with a [link](http://aka.ms/wsl2kernel) for the kernel update installation. After completing the instructions, you can restart Docker Desktop and see ***Docker Desktop running*** 7 | 8 | ## Kubernetes 9 | To use kubernetes on the laptop, we can install minikube, a local kubernetes. 10 | ### Download 11 | 1. **Latest version**: Go to minikube official website to download: [Minikube](https://minikube.sigs.k8s.io/docs/start/) 12 | 13 | 2. **Older version**: 14 | **this is the recommended version for users in China to avoid some conficts**. The minikube image repository for the up-to-date version is not correct, which leads to minikube start failure. 15 | 16 | You can find the releases from the official github repository: [Releases · kubernetes/minikube](https://github.com/kubernetes/minikube/releases). The latest version that works is [1.25.2](https://github.com/kubernetes/minikube/releases/tag/v1.25.2) and we can click it and download minikube-installer.exe of this version. 17 | 18 | ### Verify the installation 19 | ```bash 20 | minikube version 21 | ``` 22 | 23 | --------------------------------------------------------------------------------