diff --git a/docs/applications/development/workflows-api.md b/docs/applications/development/workflows-api.md index bf2e428d..5ccc50e2 100644 --- a/docs/applications/development/workflows-api.md +++ b/docs/applications/development/workflows-api.md @@ -140,7 +140,25 @@ op = operations.SingleTaskOperation('my-op-', my_task, shared_directory=shared_d op.execute() ``` +## Specify resources +Resources can be directly specified in the task as: + +```python +from cloudharness.workflows import operations, tasks + +my_task = tasks.CustomTask('my-gpu', 'myapp-mytask', resources={"requests": {"cpu": "50m", "memory": "128Mi"}, "limits": {"memory": "256Mi"}}) +op = operations.PipelineOperation('my-op-gpu-', [my_task]) +``` + +To use a gpu specify the resource like: + +```python +from cloudharness.workflows import operations, tasks + +my_task = tasks.CustomTask('my-gpu', 'myapp-mytask', resources={"limits": {"nvidia.com/gpu": 1}}) +op = operations.PipelineOperation('my-op-gpu-', [my_task]) +``` ## Pod execution context / affinity diff --git a/libraries/cloudharness-common/tests/test_workflow.py b/libraries/cloudharness-common/tests/test_workflow.py index 35a1933f..6b007b96 100644 --- a/libraries/cloudharness-common/tests/test_workflow.py +++ b/libraries/cloudharness-common/tests/test_workflow.py @@ -15,6 +15,7 @@ set_debug() execute = False +verbose = True assert 'registry' in CloudharnessConfig.get_configuration() @@ -286,3 +287,17 @@ def f(): 'matchExpressions'][0] assert affinity_expr['key'] == 'c' assert affinity_expr['values'][0] == 'd' + + +def test_gpu_workflow(): + + + from cloudharness.workflows import operations, tasks + + my_task = tasks.CustomTask('my-gpu', 'myapp-mytask', resources={"limits": {"nvidia.com/gpu": 1}}) + op = operations.PipelineOperation('my-op-gpu-', [my_task]) + wf = op.to_workflow() + + if verbose: + print('\n', yaml.dump(wf)) + assert "nvidia.com/gpu" in wf['spec']['templates'][1]["container"]["resources"]["limits"] \ No newline at end of file