Add priority option to run the job (#1096)

* Add priority option to run the job

Companion PR of #11559

* Update src/datachain/remote/studio.py

* fix test

Author: amritghimire

docs/commands/job/run.md

@@ -29,6 +29,7 @@ This command runs a job in Studio using the specified query file. You can config
 * `--python-version PYTHON_VERSION` - Python version for the job (e.g., 3.9, 3.10, 3.11)
 * `--req-file REQ_FILE` - Python requirements file
 * `--req REQ` - Python package requirements
+* `--priority PRIORITY` - Priority for the job in range 0-5. Lower value is higher priority (default: 5)
 * `-h`, `--help` - Show the help message and exit.
 * `-v`, `--verbose` - Be verbose.
 * `-q`, `--quiet` - Be quiet.
@@ -65,6 +66,11 @@ datachain job run --env API_KEY=123 --req pandas numpy query.py
 datachain job run --repository https://github.com/iterative/datachain query.py
 ```
 
+7. Run a job with higher priority
+```bash
+datachain job run --priority 2 query.py
+```
+
 ## Notes
 
 * Closing the logs command (e.g., with Ctrl+C) will only stop displaying the logs but will not cancel the job execution

src/datachain/cli/parser/job.py

@@ -82,6 +82,13 @@ def add_jobs_parser(subparsers, parent_parser) -> None:
         nargs="+",
         help="Python package requirements",
     )
+    studio_run_parser.add_argument(
+        "--priority",
+        type=int,
+        default=5,
+        help="Priority for the job in range 0-5. "
+        "Lower value is higher priority (default: 5)",
+    )
 
     studio_ls_help = "List jobs in Studio"
     studio_ls_description = "List jobs in Studio."

src/datachain/remote/studio.py

@@ -388,6 +388,7 @@ def create_job(
         python_version: Optional[str] = None,
         requirements: Optional[str] = None,
         repository: Optional[str] = None,
+        priority: Optional[int] = None,
     ) -> Response[JobData]:
         data = {
             "query": query,
@@ -399,6 +400,7 @@ def create_job(
             "python_version": python_version,
             "requirements": requirements,
             "repository": repository,
+            "priority": priority,
         }
         return self._send_request("datachain/job", data)
 

src/datachain/studio.py

@@ -40,6 +40,7 @@ def process_jobs_args(args: "Namespace"):
             args.repository,
             args.req,
             args.req_file,
+            args.priority,
         )
 
     if args.cmd == "cancel":
@@ -266,6 +267,7 @@ def create_job(
     repository: Optional[str] = None,
     req: Optional[list[str]] = None,
     req_file: Optional[str] = None,
+    priority: Optional[int] = None,
 ):
     query_type = "PYTHON" if query_file.endswith(".py") else "SHELL"
     with open(query_file) as f:
@@ -294,6 +296,7 @@ def create_job(
         python_version=python_version,
         repository=repository,
         requirements=requirements,
+        priority=priority,
     )
     if not response.ok:
         raise DataChainError(response.message)

tests/test_cli_studio.py

@@ -424,4 +424,5 @@ def test_studio_run(capsys, mocker, tmp_dir):
         "requirements": "pyjokes\nstupidity",
         "team_name": "team_name",
         "repository": "https://github.com/iterative/datachain",
+        "priority": 5,
     }