feat: pass user_group for job submissions and retries (#92)

Author: vanyae-cqc

examples/data/add_one.wasm

@@ -30,7 +30,16 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "qnx.login()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -70,9 +79,10 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "circuit = Circuit(0, 2)\n",
+    "circuit = Circuit(1)\n",
     "# Very minimal WASM example\n",
-    "circuit.add_wasm(\"add_one\", wfh, [1], [1], [0, 1])"
+    "a = circuit.add_c_register(\"a\", 8)\n",
+    "circuit.add_wasm_to_reg(\"add_one\", wfh, [a], [a])"
    ]
   },
   {
@@ -84,7 +94,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -96,7 +106,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "metadata": {},
    "outputs": [],
    "source": [

examples/data/rus.wasm

@@ -277,3 +277,62 @@ def test_results_not_available_error(
     assert isinstance(execute_results[0].download_result(), BackendResult)
 
     assert isinstance(execute_results[0].download_backend_info(), BackendInfo)
+
+
+def test_submit_under_user_group(
+    _authenticated_nexus_circuit_ref: CircuitRef,
+    qa_project_name: str,
+    qa_circuit_name: str,
+) -> None:
+    """Test that a user can submit jobs under a user_group that
+    they belong to.
+
+    Requires that the test user is a member of a group called:
+    'QA_IntegrationTestGroup',
+    and not a member of a group called:
+    'made_up_group'.
+    """
+
+    fake_group = "made_up_group"
+
+    my_proj = qnx.projects.get(name_like=qa_project_name)
+
+    with pytest.raises(qnx_exc.ResourceCreateFailed) as exc:
+        qnx.start_compile_job(
+            circuits=[_authenticated_nexus_circuit_ref],
+            name=f"qnexus_integration_test_compile_job_{datetime.now()}",
+            project=my_proj,
+            backend_config=qnx.AerConfig(),
+            user_group=fake_group,
+        )
+        assert exc.value == f"Not a member of any group with name: {fake_group}"
+
+    qnx.start_compile_job(
+        circuits=[_authenticated_nexus_circuit_ref],
+        name=f"qnexus_integration_test_compile_job_{datetime.now()}",
+        project=my_proj,
+        backend_config=qnx.AerConfig(),
+        user_group="QA_IntegrationTestGroup",
+    )
+
+    my_circ = qnx.circuits.get(name_like=qa_circuit_name, project=my_proj)
+
+    with pytest.raises(qnx_exc.ResourceCreateFailed):
+        qnx.start_execute_job(
+            circuits=[my_circ],
+            name=f"qnexus_integration_test_execute_job_{datetime.now()}",
+            project=my_proj,
+            backend_config=qnx.AerConfig(),
+            n_shots=[10],
+            user_group="made_up_group",
+        )
+        assert exc.value == f"Not a member of any group with name: {fake_group}"
+
+    qnx.start_execute_job(
+        circuits=[my_circ],
+        name=f"qnexus_integration_test_execute_job_{datetime.now()}",
+        project=my_proj,
+        backend_config=qnx.AerConfig(),
+        n_shots=[10],
+        user_group="QA_IntegrationTestGroup",
+    )

examples/wasm_examples.ipynb

@@ -10,7 +10,7 @@ def test_user_get(_authenticated_nexus: None) -> None:
     my_user = qnx.users.get_self()
     assert isinstance(my_user, UserRef)
 
-    my_user_again = qnx.users._fetch(  # pylint: disable=protected-access
+    my_user_again = qnx.users._fetch_by_id(  # pylint: disable=protected-access
         user_id=my_user.id
     )
     assert isinstance(my_user_again, UserRef)

integration/test_jobs.py

@@ -38,8 +38,9 @@ def test_wasm_flow(
     wasm_ref_2 = qnx.wasm_modules.get(id=wasm_ref.id)
     assert wasm_ref == wasm_ref_2
 
-    circuit = Circuit(0, 2)
-    circuit.add_wasm("add_one", wfh, [1], [1], [0, 1])
+    circuit = Circuit(1)
+    a = circuit.add_c_register("a", 8)
+    circuit.add_wasm_to_reg("add_one", wfh, [a], [a])
     qa_wasm_circuit_name_fixture = (
         f"qnexus_integration_test_wasm_circuit_{datetime.now()}"
     )

integration/test_user.py

@@ -12,7 +12,7 @@
 #     pass
 
 
-# def _fetch():
+# def _fetch_by_id():
 #     pass
 
 

integration/test_wasm_modules.py

@@ -130,7 +130,7 @@ def get(
     not match exactly one object.
     """
     if id:
-        return _fetch(circuit_id=id)
+        return _fetch_by_id(circuit_id=id)
 
     return get_all(
         name_like=name_like,
@@ -237,7 +237,7 @@ def update(
     )
 
 
-def _fetch(circuit_id: UUID | str) -> CircuitRef:
+def _fetch_by_id(circuit_id: UUID | str) -> CircuitRef:
     """Utility method for fetching directly by a unique identifier."""
 
     res = get_nexus_client().get(f"/api/circuits/v1beta/{circuit_id}")

qnexus/client/backend_snapshots.py

@@ -206,7 +206,7 @@ def get(  # pylint: disable=too-many-positional-arguments
     not match exactly one object.
     """
     if id:
-        return _fetch(job_id=id)
+        return _fetch_by_id(job_id=id)
 
     return get_all(
         name_like=name_like,
@@ -225,7 +225,7 @@ def get(  # pylint: disable=too-many-positional-arguments
     ).try_unique_match()
 
 
-def _fetch(job_id: UUID | str) -> JobRef:
+def _fetch_by_id(job_id: UUID | str) -> JobRef:
     """Utility method for fetching directly by a unique identifier."""
     res = get_nexus_client().get(f"/api/jobs/v1beta/{job_id}")
 
@@ -376,13 +376,17 @@ def retry_submission(
     job: JobRef,
     retry_status: list[StatusEnum] | None = None,
     remote_retry_strategy: RemoteRetryStrategy = RemoteRetryStrategy.DEFAULT,
+    user_group: str | None = None,
 ):
     """Retry a job in Nexus according to status(es) or retry strategy.
 
     By default, jobs with the ERROR status will be retried.
     """
     body: dict[str, str | list[str]] = {"remote_retry_strategy": remote_retry_strategy}
 
+    if user_group is not None:
+        body["user_group"] = user_group
+
     if retry_status is not None:
         body["retry_status"] = [status.name for status in retry_status]
 
@@ -418,6 +422,7 @@ def compile(  # pylint: disable=redefined-builtin, too-many-positional-arguments
     properties: PropertiesDict | None = None,
     optimisation_level: int = 2,
     credential_name: str | None = None,
+    user_group: str | None = None,
     hypertket_config: HyperTketConfig | None = None,
     timeout: float | None = 300.0,
 ) -> DataframableList[CircuitRef]:
@@ -437,6 +442,7 @@ def compile(  # pylint: disable=redefined-builtin, too-many-positional-arguments
         properties=properties,
         optimisation_level=optimisation_level,
         credential_name=credential_name,
+        user_group=user_group,
         hypertket_config=hypertket_config,
     )
 
@@ -465,6 +471,7 @@ def execute(  # pylint: disable=too-many-locals, too-many-positional-arguments
     language: Language = Language.AUTO,
     seed: int | None = None,
     credential_name: str | None = None,
+    user_group: str | None = None,
     timeout: float | None = 300.0,
 ) -> list[BackendResult]:
     """
@@ -489,6 +496,7 @@ def execute(  # pylint: disable=too-many-locals, too-many-positional-arguments
         language=language,
         seed=seed,
         credential_name=credential_name,
+        user_group=user_group,
     )
 
     wait_for(job=execute_job_ref, timeout=timeout)

qnexus/client/circuits.py

@@ -23,7 +23,7 @@
 
 
 @merge_properties_from_context
-def start_compile_job(  # pylint: disable=too-many-arguments, too-many-positional-arguments
+def start_compile_job(  # pylint: disable=too-many-arguments, too-many-locals, too-many-positional-arguments
     circuits: Union[CircuitRef, list[CircuitRef]],
     backend_config: BackendConfig,
     name: str,
@@ -32,6 +32,7 @@ def start_compile_job(  # pylint: disable=too-many-arguments, too-many-positiona
     properties: PropertiesDict | None = None,
     optimisation_level: int = 2,
     credential_name: str | None = None,
+    user_group: str | None = None,
     hypertket_config: HyperTketConfig | None = None,
 ) -> CompileJobRef:
     """Submit a compile job to be run in Nexus."""
@@ -55,11 +56,10 @@ def start_compile_job(  # pylint: disable=too-many-arguments, too-many-positiona
             "definition": {
                 "job_definition_type": "compile_job_definition",
                 "backend_config": backend_config.model_dump(),
-                "hypertket_config": (
-                    hypertket_config.model_dump()
-                    if hypertket_config is not None
-                    else None
-                ),
+                "user_group": user_group,
+                "hypertket_config": hypertket_config.model_dump()
+                if hypertket_config is not None
+                else None,
                 "optimisation_level": optimisation_level,
                 "credential_name": credential_name,
                 "items": [
@@ -245,14 +245,18 @@ def _fetch_compilation_passes(
         pass_input_circuit_id = pass_info["relationships"]["original_circuit"]["data"][
             "id"
         ]
-        pass_input_circuit = circuit_api._fetch(  # pylint: disable=protected-access
-            pass_input_circuit_id
+        pass_input_circuit = (
+            circuit_api._fetch_by_id(  # pylint: disable=protected-access
+                pass_input_circuit_id
+            )
         )
         pass_output_circuit_id = pass_info["relationships"]["compiled_circuit"]["data"][
             "id"
         ]
-        pass_output_circuit = circuit_api._fetch(  # pylint: disable=protected-access
-            pass_output_circuit_id
+        pass_output_circuit = (
+            circuit_api._fetch_by_id(  # pylint: disable=protected-access
+                pass_output_circuit_id
+            )
         )
 
         pass_list.append(

qnexus/client/jobs/__init__.py

@@ -40,6 +40,7 @@ def start_execute_job(  # pylint: disable=too-many-arguments, too-many-locals, t
     seed: int | None = None,
     credential_name: str | None = None,
     wasm_module: WasmModuleRef | None = None,
+    user_group: str | None = None,
 ) -> ExecuteJobRef:
     """
     Submit an execute job to be run in Nexus. Returns an ``ExecuteJobRef``
@@ -69,6 +70,7 @@ def start_execute_job(  # pylint: disable=too-many-arguments, too-many-locals, t
             "definition": {
                 "job_definition_type": "execute_job_definition",
                 "backend_config": backend_config.model_dump(),
+                "user_group": user_group,
                 "valid_check": valid_check,
                 "postprocess": postprocess,
                 "noisy_simulator": noisy_simulator,
@@ -164,7 +166,7 @@ def _fetch_execution_result(
 
     input_circuit_id = res_dict["data"]["relationships"]["circuit"]["data"]["id"]
 
-    input_circuit = circuit_api._fetch(  # pylint: disable=protected-access
+    input_circuit = circuit_api._fetch_by_id(  # pylint: disable=protected-access
         input_circuit_id
     )
 

qnexus/client/jobs/_compile.py

@@ -116,7 +116,7 @@ def get(
     not match exactly one object.
     """
     if id:
-        return _fetch(id)
+        return _fetch_by_id(id)
 
     return get_all(
         name_like=name_like,
@@ -154,7 +154,7 @@ def get_or_create(
         )
 
 
-def _fetch(project_id: UUID | str) -> ProjectRef:
+def _fetch_by_id(project_id: UUID | str) -> ProjectRef:
     """Utility method for fetching directly by a unique identifier."""
     res = get_nexus_client().get(f"/api/projects/v1beta/{project_id}")
 

qnexus/client/jobs/_execute.py

@@ -11,7 +11,7 @@
 #     pass
 
 
-# def _fetch():
+# def _fetch_by_id():
 #     pass
 
 

qnexus/client/projects.py

@@ -67,18 +67,19 @@ def assignments(resource_ref: BaseRef) -> DataframableList[RoleInfo]:
         role_infos.append(
             RoleInfo(
                 assignment_type="user",
-                assignee=user_client._fetch(  # pylint: disable=protected-access
+                assignee=user_client._fetch_by_id(  # pylint: disable=protected-access
                     user_id=user_role_assignment["user_id"]
                 ),
                 role=roles_dict[user_role_assignment["role_id"]],
             )
         )
-
     for team_role_assignment in res_assignments["team_role_assignments"]:
         role_infos.append(
             RoleInfo(
                 assignment_type="team",
-                assignee=team_client.get(name=team_role_assignment["team_id"]),
+                assignee=team_client._fetch_by_id(  # pylint: disable=protected-access
+                    team_id=team_role_assignment["team_id"]
+                ),
                 role=roles_dict[team_role_assignment["role_id"]],
             )
         )

qnexus/client/results.py

@@ -33,7 +33,7 @@ def get(name: str) -> TeamRef:
     """
     res = get_nexus_client().get("/api/v5/user/teams", params={"name": name})
 
-    if res.status_code == 404:
+    if res.status_code == 404 or res.json() == []:
         raise qnx_exc.ZeroMatches
 
     if res.status_code != 200:
@@ -54,6 +54,27 @@ def get(name: str) -> TeamRef:
     return teams_list[0]
 
 
+def _fetch_by_id(team_id: str) -> TeamRef:  # pylint: disable=redefined-builtin
+    """
+    Get a single team by id.
+    """
+    res = get_nexus_client().get(f"/api/v5/user/teams/{team_id}")
+
+    if res.status_code == 404:
+        raise qnx_exc.ZeroMatches
+
+    if res.status_code != 200:
+        raise qnx_exc.ResourceFetchFailed(message=res.text, status_code=res.status_code)
+
+    team_dict = res.json()
+
+    return TeamRef(
+        id=team_dict["id"],
+        name=team_dict["team_name"],
+        description=team_dict["description"],
+    )
+
+
 def create(name: str, description: str | None = None) -> TeamRef:
     """Create a team in Nexus."""