[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).