Added the ability to add command line arguments when submitting a job.

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()

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()

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

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)

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: