Skip to content

Commit ebefcc0

Browse files
committed
Add SSH and IDE connection info to runs API
`/api/project/:project/runs/get` now returns a new structure with connection info for each running job within a run. The info includes: * SSH command to connect (all configurations) * IDE name and URL (dev-environments only) Both command and IDE URL have two flavors: * attached – for use with CLI attached to the run * proxied – for use with sshproxy component (optional) Part-of: #3644
1 parent 1568525 commit ebefcc0

27 files changed

Lines changed: 688 additions & 194 deletions

File tree

src/dstack/_internal/cli/services/configurators/run.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
get_repo_creds_and_default_branch,
5858
)
5959
from dstack._internal.core.services.ssh.ports import PortUsedError
60+
from dstack._internal.settings import FeatureFlags
6061
from dstack._internal.utils.common import local_time
6162
from dstack._internal.utils.interpolator import InterpolatorError, VariablesInterpolator
6263
from dstack._internal.utils.logging import get_logger
@@ -215,6 +216,7 @@ def apply_configuration(
215216
current_job_submission = run._run.latest_job_submission
216217
if run.status in (RunStatus.RUNNING, RunStatus.DONE):
217218
_print_service_urls(run)
219+
_print_dev_environment_connection_info(run)
218220
bind_address: Optional[str] = getattr(
219221
configurator_args, _BIND_ADDRESS_ARG, None
220222
)
@@ -806,6 +808,30 @@ def _print_service_urls(run: Run) -> None:
806808
console.print()
807809

808810

811+
def _print_dev_environment_connection_info(run: Run) -> None:
812+
if not FeatureFlags.CLI_PRINT_JOB_CONNECTION_INFO:
813+
return
814+
if run._run.run_spec.configuration.type != RunConfigurationType.DEV_ENVIRONMENT.value:
815+
return
816+
jci = run._run.jobs[0].job_connection_info
817+
if jci is None:
818+
return
819+
if jci.ide_name:
820+
urls = [u for u in (jci.attached_ide_url, jci.proxied_ide_url) if u]
821+
if urls:
822+
console.print(
823+
f"To open in {jci.ide_name}, use link{'s' if len(urls) > 1 else ''} below:\n"
824+
)
825+
for link in urls:
826+
console.print(f" [link={link}]{link}[/]\n")
827+
ssh_commands = [" ".join(c) for c in (jci.attached_ssh_command, jci.proxied_ssh_command) if c]
828+
if ssh_commands:
829+
console.print(
830+
f"To connect via SSH, use: {' or '.join(f'[code]{c}[/]' for c in ssh_commands)}\n"
831+
)
832+
console.print()
833+
834+
809835
def print_finished_message(run: Run):
810836
status_message = (
811837
run._run.latest_job_submission.status_message

src/dstack/_internal/core/compatibility/runs.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,8 @@ def get_apply_plan_excludes(plan: ApplyRunPlanInput) -> Optional[IncludeExcludeD
4646
]
4747
),
4848
},
49+
# Contains only informational computed fields, safe to exclude unconditionally
50+
"job_connection_info": True,
4951
}
5052
}
5153
if current_resource.latest_job_submission is not None:

src/dstack/_internal/core/models/runs.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -431,9 +431,53 @@ def duration(self) -> timedelta:
431431
return end_time - self.submitted_at
432432

433433

434+
class JobConnectionInfo(CoreModel):
435+
ide_name: Annotated[
436+
Optional[str], Field(description="Dev environment IDE name for UI, human-readable.")
437+
]
438+
attached_ide_url: Annotated[
439+
Optional[str],
440+
Field(
441+
description=(
442+
"Dev environment IDE URL."
443+
" Not set if the job has not started yet."
444+
" Only works if the user is attached to the run via CLI or Python API."
445+
)
446+
),
447+
]
448+
proxied_ide_url: Annotated[
449+
Optional[str],
450+
Field(
451+
description=(
452+
"Dev environment IDE URL."
453+
" Not set if the job has hot started yet or sshproxy is not configured."
454+
)
455+
),
456+
]
457+
attached_ssh_command: Annotated[
458+
Optional[list[str]],
459+
Field(
460+
description=(
461+
"SSH command to connect to the job, list of command line arguments."
462+
" Only works if the user is attached to the run via CLI or Python API."
463+
)
464+
),
465+
]
466+
proxied_ssh_command: Annotated[
467+
Optional[list[str]],
468+
Field(
469+
description=(
470+
"SSH command to connect to the job, list of command line arguments."
471+
" Not set if sshproxy is not configured."
472+
)
473+
),
474+
]
475+
476+
434477
class Job(CoreModel):
435478
job_spec: JobSpec
436479
job_submissions: List[JobSubmission]
480+
job_connection_info: Optional[JobConnectionInfo] = None
437481

438482

439483
class RunSpecConfig(CoreConfig):

src/dstack/_internal/server/routers/sshproxy.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,21 @@
1-
import os
21
from typing import Annotated
32

43
from fastapi import APIRouter, Depends
54
from sqlalchemy.ext.asyncio import AsyncSession
65

76
from dstack._internal.core.errors import ResourceNotExistsError
7+
from dstack._internal.server import settings
88
from dstack._internal.server.db import get_session
99
from dstack._internal.server.schemas.sshproxy import GetUpstreamRequest, GetUpstreamResponse
1010
from dstack._internal.server.security.permissions import AlwaysForbidden, ServiceAccount
11-
from dstack._internal.server.services.sshproxy import get_upstream_response
11+
from dstack._internal.server.services.sshproxy.handlers import get_upstream_response
1212
from dstack._internal.server.utils.routers import (
1313
CustomORJSONResponse,
1414
get_base_api_additional_responses,
1515
)
1616

17-
if _token := os.getenv("DSTACK_SSHPROXY_API_TOKEN"):
18-
_auth = ServiceAccount(_token)
17+
if settings.SSHPROXY_API_TOKEN is not None:
18+
_auth = ServiceAccount(settings.SSHPROXY_API_TOKEN)
1919
else:
2020
_auth = AlwaysForbidden()
2121

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
from typing import Literal, Optional
2+
3+
from dstack._internal.server.services.ides.base import IDE
4+
from dstack._internal.server.services.ides.cursor import CursorDesktop
5+
from dstack._internal.server.services.ides.vscode import VSCodeDesktop
6+
from dstack._internal.server.services.ides.windsurf import WindsurfDesktop
7+
8+
_IDELiteral = Literal["vscode", "cursor", "windsurf"]
9+
10+
_ide_literal_to_ide_class_map: dict[_IDELiteral, type[IDE]] = {
11+
"vscode": VSCodeDesktop,
12+
"cursor": CursorDesktop,
13+
"windsurf": WindsurfDesktop,
14+
}
15+
16+
17+
def get_ide(ide_literal: _IDELiteral) -> Optional[IDE]:
18+
ide_class = _ide_literal_to_ide_class_map.get(ide_literal)
19+
if ide_class is None:
20+
return None
21+
return ide_class()
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
from abc import ABC, abstractmethod
2+
from typing import ClassVar, Optional
3+
4+
5+
class IDE(ABC):
6+
name: ClassVar[str]
7+
url_scheme: ClassVar[str]
8+
9+
@abstractmethod
10+
def get_install_commands(
11+
self, version: Optional[str] = None, extensions: Optional[list[str]] = None
12+
) -> list[str]:
13+
pass
14+
15+
def get_url(self, authority: str, working_dir: str) -> str:
16+
return f"{self.url_scheme}://vscode-remote/ssh-remote+{authority}{working_dir}"
17+
18+
def get_print_readme_commands(self, authority: str) -> list[str]:
19+
url = self.get_url(authority, working_dir="$DSTACK_WORKING_DIR")
20+
return [
21+
f"echo 'To open in {self.name}, use link below:'",
22+
"echo",
23+
f'echo " {url}"',
24+
"echo",
25+
]
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
from typing import Optional
2+
3+
from dstack._internal.server.services.ides.base import IDE
4+
5+
6+
class CursorDesktop(IDE):
7+
name = "Cursor"
8+
url_scheme = "cursor"
9+
10+
def get_install_commands(
11+
self, version: Optional[str] = None, extensions: Optional[list[str]] = None
12+
) -> list[str]:
13+
commands = []
14+
if version is not None:
15+
url = f"https://cursor.blob.core.windows.net/remote-releases/{version}/vscode-reh-linux-$arch.tar.gz"
16+
archive = "vscode-reh-linux-$arch.tar.gz"
17+
target = f'~/.cursor-server/cli/servers/"Stable-{version}"/server'
18+
commands.extend(
19+
[
20+
'if [ $(uname -m) = "aarch64" ]; then arch="arm64"; else arch="x64"; fi',
21+
"mkdir -p /tmp",
22+
f'wget -q --show-progress "{url}" -O "/tmp/{archive}"',
23+
f"mkdir -vp {target}",
24+
f'tar --no-same-owner -xz --strip-components=1 -C {target} -f "/tmp/{archive}"',
25+
f'rm "/tmp/{archive}"',
26+
]
27+
)
28+
if extensions:
29+
_extensions = " ".join(f'--install-extension "{name}"' for name in extensions)
30+
commands.append(f'PATH="$PATH":{target}/bin cursor-server {_extensions}')
31+
return commands
Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
from typing import Optional
2+
3+
from dstack._internal.server.services.ides.base import IDE
4+
5+
6+
class VSCodeDesktop(IDE):
7+
name = "VS Code"
8+
url_scheme = "vscode"
9+
10+
def get_install_commands(
11+
self, version: Optional[str] = None, extensions: Optional[list[str]] = None
12+
) -> list[str]:
13+
commands = []
14+
if version is not None:
15+
url = (
16+
f"https://update.code.visualstudio.com/commit:{version}/server-linux-$arch/stable"
17+
)
18+
archive = "vscode-server-linux-$arch.tar.gz"
19+
target = f'~/.vscode-server/bin/"{version}"'
20+
commands.extend(
21+
[
22+
'if [ $(uname -m) = "aarch64" ]; then arch="arm64"; else arch="x64"; fi',
23+
"mkdir -p /tmp",
24+
f'wget -q --show-progress "{url}" -O "/tmp/{archive}"',
25+
f"mkdir -vp {target}",
26+
f'tar --no-same-owner -xz --strip-components=1 -C {target} -f "/tmp/{archive}"',
27+
f'rm "/tmp/{archive}"',
28+
]
29+
)
30+
if extensions:
31+
_extensions = " ".join(f'--install-extension "{name}"' for name in extensions)
32+
commands.append(f'PATH="$PATH":{target}/bin code-server {_extensions}')
33+
return commands
Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
from typing import Optional
2+
3+
from dstack._internal.server.services.ides.base import IDE
4+
5+
6+
class WindsurfDesktop(IDE):
7+
name = "Windsurf"
8+
url_scheme = "windsurf"
9+
10+
def get_install_commands(
11+
self, version: Optional[str] = None, extensions: Optional[list[str]] = None
12+
) -> list[str]:
13+
commands = []
14+
if version is not None:
15+
version, commit = version.split("@")
16+
url = f"https://windsurf-stable.codeiumdata.com/linux-reh-$arch/stable/{commit}/windsurf-reh-linux-$arch-{version}.tar.gz"
17+
archive = "windsurf-reh-linux-$arch.tar.gz"
18+
target = f'~/.windsurf-server/bin/"{commit}"'
19+
commands.extend(
20+
[
21+
'if [ $(uname -m) = "aarch64" ]; then arch="arm64"; else arch="x64"; fi',
22+
"mkdir -p /tmp",
23+
f'wget -q --show-progress "{url}" -O "/tmp/{archive}"',
24+
f"mkdir -vp {target}",
25+
f'tar --no-same-owner -xz --strip-components=1 -C {target} -f "/tmp/{archive}"',
26+
f'rm "/tmp/{archive}"',
27+
]
28+
)
29+
if extensions:
30+
_extensions = " ".join(f'--install-extension "{name}"' for name in extensions)
31+
commands.append(f'PATH="$PATH":{target}/bin windsurf-server {_extensions}')
32+
return commands

src/dstack/_internal/server/services/jobs/__init__.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from dstack._internal.core.models.configurations import RunConfigurationType
2020
from dstack._internal.core.models.runs import (
2121
Job,
22+
JobConnectionInfo,
2223
JobProvisioningData,
2324
JobRuntimeData,
2425
JobSpec,
@@ -37,6 +38,7 @@
3738
)
3839
from dstack._internal.server.services import events
3940
from dstack._internal.server.services import volumes as volumes_services
41+
from dstack._internal.server.services.ides import get_ide
4042
from dstack._internal.server.services.instances import (
4143
get_instance_ssh_private_keys,
4244
)
@@ -51,9 +53,14 @@
5153
from dstack._internal.server.services.probes import probe_model_to_probe
5254
from dstack._internal.server.services.runner import client
5355
from dstack._internal.server.services.runner.ssh import runner_ssh_tunnel
56+
from dstack._internal.server.services.sshproxy import (
57+
build_proxied_job_ssh_command,
58+
build_proxied_job_ssh_url_authority,
59+
)
5460
from dstack._internal.utils import common
5561
from dstack._internal.utils.common import run_async
5662
from dstack._internal.utils.logging import get_logger
63+
from dstack._internal.utils.ssh import build_ssh_command, build_ssh_url_authority
5764

5865
logger = get_logger(__name__)
5966

@@ -490,6 +497,47 @@ def remove_job_spec_sensitive_info(spec: JobSpec):
490497
spec.ssh_key = None
491498

492499

500+
def get_job_connection_info(job_model: JobModel, run_spec: RunSpec) -> JobConnectionInfo:
501+
# Run.attach() Python API method, used internally by CLI, uses the following as the Hostname
502+
# in the SSH config:
503+
# * for the (job=0 replica=0) job - run name, e.g., `my-task`
504+
# * for other jobs - job name, e.g., `my-task-0-1`
505+
attached_hostname = run_spec.run_name
506+
if job_model.job_num != 0 or job_model.replica_num != 0:
507+
attached_hostname = job_model.job_name
508+
assert attached_hostname is not None
509+
510+
# ide_* fields are for dev-environment only
511+
ide_name: Optional[str] = None
512+
# IDE URLs are not set until the job status is switched to RUNNING,
513+
# as JobRuntimeData.working_dir, which is required to build URLs, is returned
514+
# by dstack-runner's `/api/run` method
515+
attached_ide_url: Optional[str] = None
516+
proxied_ide_url: Optional[str] = None
517+
if (
518+
run_spec.configuration.type == RunConfigurationType.DEV_ENVIRONMENT.value
519+
and run_spec.configuration.ide is not None
520+
):
521+
ide = get_ide(run_spec.configuration.ide)
522+
if ide is not None:
523+
ide_name = ide.name
524+
jrd = get_job_runtime_data(job_model)
525+
if jrd is not None and jrd.working_dir is not None:
526+
attached_url_authority = build_ssh_url_authority(hostname=attached_hostname)
527+
attached_ide_url = ide.get_url(attached_url_authority, jrd.working_dir)
528+
proxied_url_authority = build_proxied_job_ssh_url_authority(job_model)
529+
if proxied_url_authority is not None:
530+
proxied_ide_url = ide.get_url(proxied_url_authority, jrd.working_dir)
531+
532+
return JobConnectionInfo(
533+
ide_name=ide_name,
534+
attached_ide_url=attached_ide_url,
535+
proxied_ide_url=proxied_ide_url,
536+
attached_ssh_command=build_ssh_command(hostname=attached_hostname),
537+
proxied_ssh_command=build_proxied_job_ssh_command(job_model),
538+
)
539+
540+
493541
def _get_job_mount_point_attached_volume(
494542
volumes: List[Volume],
495543
job_provisioning_data: JobProvisioningData,

0 commit comments

Comments
 (0)