[machinelearningservices] Add retry logic to az ml ssh client websocket connection

[lavakumarrepala]

🤖 PR Validation — ️✔️ All clear

Breaking Changes
️✔️ None

Summary

�z ml job connect-ssh (and connect-ssh for compute instances) proxies SSH traffic over a websocket tunnel via _ssh_connector.py. Today, if the initial websocket handshake to the NBIP proxy/broker is refused or times out, the client fails immediately with no retry.

SRE Agent investigation of SSH connection instability (ADO work item https://msdata.visualstudio.com/Vienna/_workitems/edit/5533560/) identified this as a remaining client-side gap: backend forwarder/broker stability fixes already improved connection duration stats, but the az ml SSH client still lacks the retry behavior that SSH clients such as VS Code Remote-SSH have (which show meaningfully better long-duration connection stats as a result).

Change

• Added _connect_with_retry(), an async context manager that retries the initial websocket connection with exponential backoff + jitter on transient failures (\OSError, \TimeoutError, websocket handshake/connection-closed errors during connect).
• Retry behavior is configurable via env vars: \AZUREML_SSH_CONNECT_MAX_RETRIES\ (default 5), \AZUREML_SSH_CONNECT_RETRY_BASE_DELAY_SECONDS\ (default 1), \AZUREML_SSH_CONNECT_RETRY_MAX_DELAY_SECONDS\ (default 30).
• Once retries are exhausted, the original exception is re-raised so behavior is unchanged for persistent failures.
• Deliberately out of scope: reconnecting once the tunnel is open and actively relaying SSH bytes. This tunnel carries the raw SSH transport, so a new websocket connection cannot resume an in-flight SSH session — a silent mid-stream reconnect would corrupt the SSH session rather than heal it. Recovering from mid-session drops requires server-side session continuity, tracked separately under the NBIP rearchitecture work item (https://msdata.visualstudio.com/Vienna/_workitems/edit/5533562).

Testing

• \python -m py_compile\ on the modified file.
• \pyflakes\ shows no new warnings (pre-existing unused TYPE_CHECKING import only).
• Manually verified retry behavior with a mocked \websockets.client.connect:
  • Transient failures (e.g. simulated \TimeoutError) are retried with backoff and the connection succeeds once the mock stops failing.
  • When failures persist past \AZUREML_SSH_CONNECT_MAX_RETRIES, the original exception propagates as before.

Opening as draft for review/discussion before merge.

[azure-client-tools-bot-prd]

Hi lavakumarrepala,
Please write the description of changes which can be perceived by customers into HISTORY.rst.
If you want to release a new extension version, please update the version in pyproject.toml (or setup.py, if the extension has not migrated yet) as well.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Aditya Pujara (a0x1ab)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 2 pipeline(s).

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

The new retry delay computation and env var handling can yield invalid sleep durations (and exceed the configured max delay), which can break retries under certain configurations.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR improves resiliency of az ml job connect-ssh (and compute instance connect-ssh) by adding retry logic around the initial WebSocket tunnel establishment in _ssh_connector.py, addressing transient proxy/broker handshake failures without attempting mid-session reconnects.

Changes:

• Introduces an async context manager _connect_with_retry() to retry initial WebSocket connections with exponential backoff + jitter on retryable connect-time exceptions.
• Adds env var configuration for retry count and backoff parameters (AZUREML_SSH_CONNECT_MAX_RETRIES, ...BASE_DELAY_SECONDS, ...MAX_DELAY_SECONDS).
• Switches the SSH tunnel setup to use _connect_with_retry() instead of calling websockets.client.connect() directly.

File summaries

File	Description
src/machinelearningservices/azext_mlv2/manual/custom/_ssh_connector.py	Adds a retrying WebSocket connect wrapper and uses it for the SSH-over-WebSocket tunnel.

Review details

Suppressed comments (1)

src/machinelearningservices/azext_mlv2/manual/custom/_ssh_connector.py:103

• delay is clamped to max_delay before adding jitter, so the actual sleep can exceed AZUREML_SSH_CONNECT_RETRY_MAX_DELAY_SECONDS by up to base_delay. If max_delay is meant to be a hard cap, clamp after applying jitter.

            delay = min(max_delay, base_delay * (2 ** (attempt - 1))) + random.uniform(0, base_delay)

• Files reviewed: 1/1 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.