Skip to content

Commit

Permalink
Refactor SSHTunnel (#1669)
Browse files Browse the repository at this point in the history
Replace 4 different `*SSHTunnel` classes with a
common `SSHTunnel` class. The class implements
both remote-to-local and local-to-remote port
forwarding where both IP and Unix sockets can be
used on either end, which will be useful for
future features.

This commit does not change `dstack` behavior for
end users.
  • Loading branch information
jvstme authored Sep 9, 2024
1 parent 25e1c8c commit 3f01327
Show file tree
Hide file tree
Showing 12 changed files with 439 additions and 289 deletions.
18 changes: 12 additions & 6 deletions src/dstack/_internal/core/services/ssh/attach.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,11 @@
from dstack._internal.core.models.instances import SSHConnectionParams
from dstack._internal.core.services.configs import ConfigManager
from dstack._internal.core.services.ssh.ports import PortsLock
from dstack._internal.core.services.ssh.tunnel import ClientTunnel
from dstack._internal.core.services.ssh.tunnel import (
FilePath,
SSHTunnel,
ports_to_forwarded_sockets,
)
from dstack._internal.utils.path import PathLike
from dstack._internal.utils.ssh import get_ssh_config, include_ssh_config, update_ssh_config

Expand Down Expand Up @@ -52,13 +56,15 @@ def __init__(
self.ports = ports_lock.dict()
self.run_name = run_name
self.ssh_config_path = str(ConfigManager().dstack_ssh_config_path)
self.tunnel = ClientTunnel(
run_name,
self.ports,
id_rsa_path=id_rsa_path,
self.tunnel = SSHTunnel(
destination=run_name,
identity=FilePath(id_rsa_path),
forwarded_sockets=ports_to_forwarded_sockets(
ports=self.ports,
bind_local=bind_address or "localhost",
),
control_sock_path=control_sock_path,
ssh_config_path=self.ssh_config_path,
bind_address=bind_address,
)
self.ssh_proxy = ssh_proxy
if ssh_proxy is None:
Expand Down
Loading

0 comments on commit 3f01327

Please sign in to comment.