diff --git a/src/sfapi_client/_async/compute.py b/src/sfapi_client/_async/compute.py index 2bcf495..875b58a 100644 --- a/src/sfapi_client/_async/compute.py +++ b/src/sfapi_client/_async/compute.py @@ -126,10 +126,13 @@ async def _wait_for_task(self, task_id) -> AsyncCommandTaskResult: return task.result @check_auth - async def submit_job(self, script: Union[str, AsyncRemotePath]) -> AsyncJobSqueue: + async def submit_job( + self, script: Union[str, AsyncRemotePath], args: Optional[List[str]] = None + ) -> AsyncJobSqueue: """Submit a job to the compute resource :param script: Path to file on the compute system, or script to run beginning with `#!`. + :param args: An optional list of command line arguments to pass to the script. :return: Object containing information about the job, its job id, and status on the system. """ @@ -150,6 +153,12 @@ async def submit_job(self, script: Union[str, AsyncRemotePath]) -> AsyncJobSqueu raise SfApiError(f"Script path not present or is not a file, {script}") data = {"job": script, "isPath": is_path} + if args is not None: + if not is_path: + raise ValueError( + "Command line arguments cannot be passed when the script is not a file." + ) + data["args"] = args r = await self.client.post(f"compute/jobs/{self.name}", data) r.raise_for_status() diff --git a/src/sfapi_client/_sync/compute.py b/src/sfapi_client/_sync/compute.py index 148454b..2943fd8 100644 --- a/src/sfapi_client/_sync/compute.py +++ b/src/sfapi_client/_sync/compute.py @@ -127,10 +127,11 @@ def _wait_for_task(self, task_id) -> CommandTaskResult: return task.result @check_auth - def submit_job(self, script: Union[str, RemotePath]) -> JobSqueue: + def submit_job(self, script: Union[str, RemotePath], args: Optional[List[str]] = None) -> JobSqueue: """Submit a job to the compute resource :param script: Path to file on the compute system, or script to run beginning with `#!`. + :param args: An optional list of command line arguments to pass to the script. :return: Object containing information about the job, its job id, and status on the system. """ @@ -151,6 +152,10 @@ def submit_job(self, script: Union[str, RemotePath]) -> JobSqueue: raise SfApiError(f"Script path not present or is not a file, {script}") data = {"job": script, "isPath": is_path} + if args is not None: + if not is_path: + raise ValueError("Command line arguments cannot be passed when the script is not a file.") + data["args"] = args r = self.client.post(f"compute/jobs/{self.name}", data) r.raise_for_status() diff --git a/tests/conftest.py b/tests/conftest.py index 9e8acd4..d0e1ab8 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import json import random import string +from io import BytesIO from typing import Optional, Union, Dict from pathlib import Path from cryptography.hazmat.primitives.asymmetric import rsa @@ -23,6 +24,7 @@ class Settings(BaseSettings): SFAPI_DEV_CLIENT_ID: Optional[str] = None SFAPI_DEV_CLIENT_SECRET: Optional[Union[str, Dict]] = None TEST_JOB_PATH: Optional[str] = None + TEST_JOB_ACCOUNT: Optional[str] = None TEST_MACHINE: Machine = Machine.perlmutter TEST_RESOURCE: Resource = Resource.spin TEST_USERNAME: Optional[str] = None @@ -75,6 +77,11 @@ def test_job_path(): return settings.TEST_JOB_PATH +@pytest.fixture +def test_job_account(): + return settings.TEST_JOB_ACCOUNT + + @pytest.fixture def test_machine(): return settings.TEST_MACHINE @@ -214,3 +221,21 @@ def empty_key_file(tmp_path_factory): if temp_path.exists(): (temp_path / "nokey.pem").unlink(missing_ok=True) temp_path.rmdir() + + +@pytest.fixture +def test_arg_job_script(test_job_account, test_tmp_dir): + script = BytesIO(f"""#!/usr/bin/env bash +#SBATCH --nodes=1 +#SBATCH --constraint=cpu +#SBATCH --account={test_job_account} +#SBATCH --qos=debug +#SBATCH --time=00:01:00 +#SBATCH --output={test_tmp_dir}/slurm-%j.out + +set -euo pipefail + +echo "$@" +""".encode()) + script.filename = "echo-args.sh" + return script diff --git a/tests/test_jobs.py b/tests/test_jobs.py index 497f971..fe8c48c 100755 --- a/tests/test_jobs.py +++ b/tests/test_jobs.py @@ -17,6 +17,30 @@ def test_submit(authenticated_client, test_job_path, test_machine): assert state == JobState.COMPLETED +def test_submit_with_args( + authenticated_client, test_machine, test_tmp_dir, test_arg_job_script +): + with authenticated_client as client: + machine = client.compute(test_machine) + + [remote_dir] = machine.ls(test_tmp_dir, directory=True) + remote_script = remote_dir.upload(test_arg_job_script) + try: + args = ["1", "2", "3"] + + job = machine.submit_job(remote_script, args) + + state = job.complete() + + assert state == JobState.COMPLETED + + [output_path] = machine.ls(f"{test_tmp_dir}/slurm-{job.jobid}.out") + output = output_path.download().read() + assert "1 2 3" in output + finally: + machine.run(["rm", "-f", str(remote_script)]) + + def test_cancel(authenticated_client, test_job_path, test_machine): with authenticated_client as client: machine = client.compute(test_machine) diff --git a/tests/test_jobs_async.py b/tests/test_jobs_async.py index 3c1b5ab..3d8ceee 100755 --- a/tests/test_jobs_async.py +++ b/tests/test_jobs_async.py @@ -16,6 +16,31 @@ async def test_submit(async_authenticated_client, test_job_path, test_machine): assert state == JobState.COMPLETED +@pytest.mark.asyncio +async def test_submit_with_args( + async_authenticated_client, test_machine, test_tmp_dir, test_arg_job_script +): + async with async_authenticated_client as client: + machine = await client.compute(test_machine) + + [remote_dir] = await machine.ls(test_tmp_dir, directory=True) + remote_script = await remote_dir.upload(test_arg_job_script) + try: + args = ["1", "2", "3"] + + job = await machine.submit_job(remote_script, args) + + state = await job.complete() + + assert state == JobState.COMPLETED + + [output_path] = await machine.ls(f"{test_tmp_dir}/slurm-{job.jobid}.out") + output = (await output_path.download()).read() + assert "1 2 3" in output + finally: + await machine.run(["rm", "-f", str(remote_script)]) + + @pytest.mark.asyncio async def test_cancel(async_authenticated_client, test_job_path, test_machine): async with async_authenticated_client as client: