Pytorch

A Pytorch task type’s example and dive into information of PyDolphinScheduler.

Example

"""A example workflow for task pytorch."""

from pydolphinscheduler.core.workflow import Workflow
from pydolphinscheduler.tasks.pytorch import Pytorch

with Workflow(
    name="task_pytorch_example",
) as workflow:
    # run project with existing environment
    task_existing_env = Pytorch(
        name="task_existing_env",
        script="main.py",
        script_params="--dry-run --no-cuda",
        project_path="https://github.com/pytorch/examples#mnist",
        python_command="/home/anaconda3/envs/pytorch/bin/python3",
    )

    # run project with creating conda environment
    task_conda_env = Pytorch(
        name="task_conda_env",
        script="main.py",
        script_params="--dry-run --no-cuda",
        project_path="https://github.com/pytorch/examples#mnist",
        is_create_environment=True,
        python_env_tool="conda",
        requirements="requirements.txt",
        conda_python_version="3.7",
    )

    # run project with creating virtualenv environment
    task_virtualenv_env = Pytorch(
        name="task_virtualenv_env",
        script="main.py",
        script_params="--dry-run --no-cuda",
        project_path="https://github.com/pytorch/examples#mnist",
        is_create_environment=True,
        python_env_tool="virtualenv",
        requirements="requirements.txt",
    )

    # [start resource_limit]
    pytorch_resources_limit = Pytorch(
        name="pytorch_resources_limit",
        script="main.py",
        script_params="--dry-run --no-cuda",
        project_path="https://github.com/pytorch/examples#mnist",
        python_command="/home/anaconda3/envs/pytorch/bin/python3",
        cpu_quota=1,
        memory_max=100,
    )
    # [end resource_limit]

    workflow.submit()

Resource Limit Example

We can add resource limit like CPU quota and max memory by passing parameters when declaring tasks.

    pytorch_resources_limit = Pytorch(
        name="pytorch_resources_limit",
        script="main.py",
        script_params="--dry-run --no-cuda",
        project_path="https://github.com/pytorch/examples#mnist",
        python_command="/home/anaconda3/envs/pytorch/bin/python3",
        cpu_quota=1,
        memory_max=100,
    )

Dive Into

Task Pytorch.

class pydolphinscheduler.tasks.pytorch.DEFAULT[source]

Bases: object

Default values for Pytorch.

is_create_environment = False

project_path = '.'

python_command = '${PYTHON_HOME}'

class pydolphinscheduler.tasks.pytorch.Pytorch(name: str, script: str, script_params: str = '', project_path: str | None = '.', is_create_environment: bool | None = False, python_command: str | None = '${PYTHON_HOME}', python_env_tool: str | None = 'conda', requirements: str | None = 'requirements.txt', conda_python_version: str | None = '3.7', *args, **kwargs)[source]

Bases: WorkerResourceMixin, BatchTask

Task Pytorch object, declare behavior for Pytorch task to dolphinscheduler.

Parameters:

name – task name
script – Entry to the Python script file that you want to run.
script_params – Input parameters at run time.
project_path – The path to the project. Default “.” .
is_create_environment – is create environment. Default False.
python_command – The path to the python command. Default “${PYTHON_HOME}”.
python_env_tool – The python environment tool. Default “conda”.
requirements – The path to the requirements.txt file. Default “requirements.txt”.
conda_python_version – The python version of conda environment. Default “3.7”.

_downstream_task_codes: set[int]

_task_custom_attr: set = {'conda_python_version', 'is_create_environment', 'other_params', 'python_command', 'python_env_tool', 'python_path', 'requirements', 'script', 'script_params'}

_task_relation: set[TaskRelation]

_timeout: timedelta | int

_upstream_task_codes: set[int]

property other_params: Return other params.

YAML file example

# Define the workflow
workflow:
  name: "Pytorch"

# Define the tasks within the workflow
tasks:

  # run project with existing environment
  - name: task_existing_env
    task_type: pytorch
    script: main.py
    script_params: --dry-run --no-cuda
    project_path: https://github.com/pytorch/examples#mnist
    python_command: /home/anaconda3/envs/pytorch/bin/python3


  # run project with creating conda environment
  - name: task_conda_env
    task_type: pytorch
    script: main.py
    script_params: --dry-run --no-cuda
    project_path: https://github.com/pytorch/examples#mnist
    is_create_environment: True
    python_env_tool: conda
    requirements: requirements.txt
    conda_python_version: 3.7

  # run project with creating virtualenv environment
  - name: task_virtualenv_env
    task_type: pytorch
    script: main.py
    script_params: --dry-run --no-cuda
    project_path: https://github.com/pytorch/examples#mnist
    is_create_environment: True
    python_env_tool: virtualenv
    requirements: requirements.txt