feat: add Transport abstraction for SSH connections - #724
Merged
ArangoGutierrez merged 3 commits intoMar 13, 2026
Conversation
Introduce Transport interface so each provider controls how SSH connections are established. DirectTransport for single-node and SSH provider. SSMTransport for AWS cluster mode with retry-based dial (D1). Provisioner unchanged — just receives working connections. - Transport interface with Dial() and Target() methods - DirectTransport: TCP dial to host:22 (existing behavior) - SSMTransport: AWS SSM port forwarding with exponential backoff retry - Provisioner gets functional option WithTransport() - NodeInfo gets InstanceID and Transport fields - Backwards compatible: nil transport falls back to DirectTransport Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
There was a problem hiding this comment.
Pull request overview
This PR introduces a new SSH connection “Transport” abstraction in the provisioner layer to support multiple connection mechanisms (direct TCP and AWS SSM port-forwarding) and begins wiring it through cluster provisioning.
Changes:
- Added
Transportinterface plusDirectTransportandSSMTransportimplementations. - Updated
Provisioner.New(...)to accept functional options, includingWithTransport(...). - Extended
NodeInfowithInstanceIDandTransport, and updated cluster provisioning call sites to pass transport options.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| pkg/provisioner/transport.go | Adds Transport abstraction, Direct/SSM transports, retry dial helper, and WithTransport option. |
| pkg/provisioner/transport_test.go | Adds unit tests for transports and retry dialing behavior. |
| pkg/provisioner/provisioner.go | Adds transport Transport field and variadic options to New(...). |
| pkg/provisioner/cluster.go | Adds InstanceID/Transport to NodeInfo and passes transport options into provisioner creation. |
You can also share your feedback on Copilot code review. Take the survey.
Comment on lines
+87
to
+91
| if p.transport == nil { | ||
| p.transport = NewDirectTransport(hostUrl) | ||
| } | ||
|
|
||
| client, err := connectOrDie(keyPath, userName, hostUrl) |
Comment on lines
186
to
189
| cp.log.Info("Provisioning base dependencies on %s (%s)", node.Name, node.PublicIP) | ||
|
|
||
| provisioner, err := New(cp.log, cp.KeyPath, cp.getUsernameForNode(node), node.PublicIP) | ||
| provisioner, err := New(cp.log, cp.KeyPath, cp.getUsernameForNode(node), node.PublicIP, cp.transportOptsForNode(node)...) | ||
| if err != nil { |
Comment on lines
+88
to
+92
| // transportOptsForNode returns functional options for New() based on the node's transport. | ||
| // If the node has a Transport configured, it is passed via WithTransport; otherwise | ||
| // the default DirectTransport(PublicIP) is used automatically by New(). | ||
| func (cp *ClusterProvisioner) transportOptsForNode(node NodeInfo) []Option { | ||
| if node.Transport != nil { |
Comment on lines
+118
to
+122
| s.cmd = exec.Command("aws", args...) //nolint:gosec // args are constructed from validated fields | ||
| if err := s.cmd.Start(); err != nil { | ||
| return nil, fmt.Errorf("ssm transport: start session for %s: %w", s.InstanceID, err) | ||
| } | ||
|
|
Comment on lines
+156
to
+162
| for attempt := range maxAttempts { | ||
| conn, err := net.DialTimeout("tcp", addr, ssmDialTimeout) | ||
| if err == nil { | ||
| return conn, nil | ||
| } | ||
| lastErr = err | ||
| time.Sleep(baseDelay * (1 << attempt)) |
Address Distinguished Engineer review findings for PR NVIDIA#724: B1 (BLOCKING): connectOrDie() now uses transport.Dial() + ssh.NewClientConn() when a Transport is provided, instead of ignoring the transport and calling ssh.Dial("tcp", ...) directly. All reconnection paths (waitForNodeReboot, provision) pass p.transport through. B2 (BLOCKING): installK8sPrereqs() and GetClusterHealth() now use transportOptsForNode() to pass transport options to New(). R1: SSMTransport.Dial() is now idempotent — closes any existing session before starting a new one. R2: Transport interface includes Close() method. DirectTransport has a no-op Close(). SSMTransport.Close() terminates the SSM process. R3: SSM stderr is captured in a bytes.Buffer and included in error messages for diagnostics. R4: GetClusterHealth() builds transport options from node info in environment status. Signed-off-by: Eduardo Arango <earango@nvidia.com> Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
ArangoGutierrez
marked this pull request as ready for review
March 13, 2026 12:52
Fix golangci-lint errcheck violations for ln.Close() and conn.Close() in transport_test.go. Signed-off-by: Eduardo Arango <earango@nvidia.com> Signed-off-by: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
Pull Request Test Coverage Report for Build 23051945581Details
💛 - Coveralls |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Test plan