tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py (6)

39-52:Add a constant for the MPI service name and lightweight docstring; fix prints/raise usage.

Hardcoding "my_port" in multiple places is brittle. Define a single constant (UPPER_SNAKE_CASE per guidelines) and reuse it across helpers and workers. Also, add a brief docstring and use flush/raise best practices.

Apply within this hunk:

-def mpi_publish_name():+def mpi_publish_name():+    """Open an MPI port and publish it under MPI_SERVICE_NAME.++    Returns:+        str: The MPI port name string.+    """     port_name = None     try:         port_name = MPI.Open_port()-        MPI.Publish_name('my_port', port_name)+        MPI.Publish_name(MPI_SERVICE_NAME, port_name)     except MPI.Exception as e:-        print(f"Error publishing port name: {e}")-        raise e+        print(f"Error publishing port name: {e}", flush=True)+        raise     except Exception as e:-        print(f"Unexpected error publishing port name: {e}")-        raise e+        print(f"Unexpected error publishing port name: {e}", flush=True)+        raise

And add the constant near the other MPI constants (outside this hunk):

# Single point of truth for the service name used by Publish/LookupMPI_SERVICE_NAME="disagg_test_port"

Also consider adding the NVIDIA copyright header at the top of the file as per repository guidelines. Want me to generate it with the current year?

---

68-74:Generalize termination to N workers instead of hardcoding 0 and 1.

Use Get_remote_size() to avoid fragile assumptions and ease future scaling.

-def mpi_send_termination_request(intercomm):-    if intercomm is not None:-        # Send termination requests-        intercomm.send(None, dest=0, tag=MPI_REQUEST)-        intercomm.send(None, dest=1, tag=MPI_REQUEST)-        print("Sent termination requests to the workers.")+def mpi_send_termination_request(intercomm):+    if intercomm is not None:+        for dst in range(intercomm.Get_remote_size()):+            intercomm.send(None, dest=dst, tag=MPI_REQUEST)+        print("Sent termination requests to the workers.", flush=True)

---

88-97:Fix misleading error messages and use the service-name constant in Lookup.

The current messages say “publishing” while this block performs Lookup/Connect. Use MPI_SERVICE_NAME for consistency with Publish_name.

-    try:-        port_name = MPI.Lookup_name('my_port')-        intercomm = MPI.COMM_WORLD.Connect(port_name)-    except MPI.Exception as e:-        print(f"Error publishing port name: {e}")-        raise e-    except Exception as e:-        print(f"Unexpected error publishing port name: {e}")-        raise e+    try:+        port_name = MPI.Lookup_name(MPI_SERVICE_NAME)+        intercomm = MPI.COMM_WORLD.Connect(port_name)+    except MPI.Exception as e:+        print(f"Error during Lookup/Connect: {e}", flush=True)+        raise+    except Exception as e:+        print(f"Unexpected error during Lookup/Connect: {e}", flush=True)+        raise

---

233-235:Use bareraise to preserve traceback context; include flush.

Minor, but improves debuggability.

-        except Exception as e:-            print(f"Exception encountered: {e}", flush=True)-            raise e+        except Exception as e:+            print(f"Exception encountered: {e}", flush=True)+            raise

---

380-386:Improve MPI exception logging and ensure termination is robust.

Print the exception details and flush. Also relies on intercomm being defined—apply the earlier predefinition.

-        except MPI.Exception as e:-            print(f"MPI Error")-            raise e+        except MPI.Exception as e:+            print(f"MPI error: {e}", flush=True)+            raise

After applying, please re-run the relevant tests to confirm no residual leaks or name-collision errors occur when running the whole test module.

---

481-486:Improve MPI exception logging and keep termination safe.

Same as above: print the exception object for context and use bare raise.

-        except MPI.Exception as e:-            print(f"MPI Error")-            raise e+        except MPI.Exception as e:+            print(f"MPI error: {e}", flush=True)+            raise

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/disaggregated/test_disaggregated_single_gpu.py (4)

186-186:LGTM: helper used for publishing.

Using a helper here improves readability and centralizes error handling. No issues.

---

241-241:LGTM: explicit “waiting” log with flush improves CI readability.

No action needed.

---

329-329:LGTM: consistent use of mpi_publish_name().

No issues here.

---

431-431:LGTM: consistent use of mpi_publish_name().

No issues here.