Skip to content

Commit c59b2ea

Browse files
Merge master into feature/cloudformation
2 parents c6ca60b + b13ad47 commit c59b2ea

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+4142
-544
lines changed

.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ packages/toolkit/package.nls.json
4747
packages/toolkit/resources
4848
packages/amazonq/package.nls.json
4949
packages/amazonq/resources
50+
packages/sagemaker-ssh-kiro/resources
5051

5152
# Icons
5253
packages/*/resources/fonts/aws-toolkit-icons.woff

LICENSE-THIRD-PARTY

Lines changed: 423 additions & 396 deletions
Large diffs are not rendered by default.

aws-toolkit-vscode.code-workspace

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
{
1313
"path": "packages/amazonq",
1414
},
15+
{
16+
"path": "packages/sagemaker-ssh-kiro",
17+
},
1518
],
1619
"settings": {
1720
"typescript.tsdk": "node_modules/typescript/lib",

package-lock.json

Lines changed: 92 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.
Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
---
2+
inclusion: always
3+
---
4+
5+
# SageMaker Unified Studio Space Context
6+
7+
This workspace is running on an Amazon SageMaker Unified Studio Space.
8+
9+
## Environment
10+
- Operating system: Ubuntu-based SageMaker Distribution
11+
- User: sagemaker-user
12+
- Home directory: /home/sagemaker-user
13+
- AWS credentials are available via the container credentials provider (AWS_CONTAINER_CREDENTIALS_RELATIVE_URI)
14+
- Do NOT hardcode AWS credentials; use the default credential chain (e.g., boto3.Session())
15+
16+
## Project Info
17+
- ~/README.md contains project-specific configuration such as connection names and available compute resources.
18+
- ~/shared/README.md contains shared project data catalog and storage information.
19+
Refer to these files when you need details about the project's connections, databases, or S3 paths.
20+
21+
## Project Library (`sagemaker_studio`)
22+
The `sagemaker_studio` package is pre-installed and provides access to project resources.
23+
24+
### Project
25+
```python
26+
from sagemaker_studio import Project
27+
project = Project()
28+
29+
project.id
30+
project.name
31+
project.iam_role # project IAM role ARN
32+
project.kms_key_arn # project KMS key ARN (if configured)
33+
project.mlflow_tracking_server_arn # MLflow ARN (if configured)
34+
project.s3.root # project S3 root path
35+
```
36+
37+
### Connections
38+
```python
39+
project.connections # list all connections
40+
project.connection() # default IAM connection
41+
project.connection("redshift") # named connection
42+
conn.name, conn.id, conn.iam_role
43+
conn.physical_endpoints[0].host # endpoint host
44+
conn.data # all connection properties
45+
conn.secret # credentials (dict or string)
46+
conn.create_client() # boto3 client with connection credentials
47+
conn.create_client("glue") # boto3 client for specific service
48+
```
49+
50+
### Catalogs, Databases, and Tables
51+
```python
52+
catalog = project.connection().catalog() # default catalog
53+
catalog = project.connection().catalog("catalog_id")
54+
catalog.databases # list databases
55+
db = catalog.database("my_db")
56+
db.tables # list tables
57+
table = db.table("my_table")
58+
table.columns # list columns (name, type)
59+
```
60+
61+
### SQL Utilities
62+
```python
63+
from sagemaker_studio import sqlutils
64+
65+
# DuckDB (local, no connection needed)
66+
result = sqlutils.sql("SELECT * FROM my_df WHERE id > 1")
67+
68+
# Athena
69+
result = sqlutils.sql("SELECT * FROM orders", connection_name="project.athena")
70+
71+
# Redshift
72+
result = sqlutils.sql("SELECT * FROM products", connection_name="project.redshift")
73+
74+
# Parameterized queries
75+
result = sqlutils.sql(
76+
"SELECT * FROM orders WHERE status = :status",
77+
parameters={"status": "completed"},
78+
connection_name="project.redshift"
79+
)
80+
81+
# Get SQLAlchemy engine
82+
engine = sqlutils.get_engine(connection_name="project.redshift")
83+
```
84+
85+
### DataFrame Utilities
86+
```python
87+
from sagemaker_studio import dataframeutils
88+
import pandas as pd
89+
90+
# Read from catalog table
91+
df = pd.read_catalog_table(database="my_db", table="my_table")
92+
93+
# Write to catalog table
94+
df.to_catalog_table(database="my_db", table="my_table")
95+
96+
# S3 Tables catalog
97+
df = pd.read_catalog_table(
98+
database="my_db", table="my_table",
99+
catalog="s3tablescatalog/my_catalog"
100+
)
101+
```
102+
103+
### Spark Utilities
104+
```python
105+
from sagemaker_studio import sparkutils
106+
107+
# Initialize Spark Connect session
108+
spark = sparkutils.init()
109+
spark = sparkutils.init(connection_name="my_spark_connection")
110+
111+
# Get Spark options for JDBC connections
112+
options = sparkutils.get_spark_options("my_redshift_connection")
113+
df = spark.read.format("jdbc").options(**options).option("dbtable", "my_table").load()
114+
```
115+
116+
## Compute Options
117+
- **Local Python**: Runs directly on the Space instance. Use for single-machine Python, ML, and AI workloads.
118+
- **Apache Spark (AWS Glue / Amazon EMR)**: Use `%%pyspark`, `%%scalaspark`, or `%%sql` cell magics in notebooks. Default Spark connection is `project.spark.compatibility`.
119+
- **SQL (Athena)**: Use `%%sql project.athena` for Trino SQL queries via Amazon Athena.
120+
- **SQL (Redshift)**: Use `%%sql project.redshift` if a Redshift connection is available.
121+
122+
## Code Patterns
123+
- Use `sagemaker_studio.Project()` for project-aware sessions and resource discovery
124+
- Reference data using S3 URIs in s3://bucket/prefix format
125+
- Write Spark DataFrames to the project catalog: `df.write.saveAsTable(f"{database}.table_name", format='parquet', mode='overwrite')`
126+
- SQL query results are available as DataFrames in subsequent cells via the `_` variable
127+
- Use `sqlutils.sql()` for programmatic SQL execution against any connection
128+
- Use `pd.read_catalog_table()` / `df.to_catalog_table()` for pandas catalog I/O
129+
130+
## MCP Server Configuration
131+
- When configuring MCP servers, pass AWS credentials via environment variable expansion:
132+
"AWS_CONTAINER_CREDENTIALS_RELATIVE_URI": "${AWS_CONTAINER_CREDENTIALS_RELATIVE_URI}"

packages/core/src/awsService/sagemaker/commands.ts

Lines changed: 29 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,13 +13,18 @@ import { getLogger } from '../../shared/logger/logger'
1313
import { SagemakerSpaceNode, tryRefreshNode } from './explorer/sagemakerSpaceNode'
1414
import { isRemoteWorkspace } from '../../shared/vscode/env'
1515
import _ from 'lodash'
16-
import { prepareDevEnvConnection, tryRemoteConnection } from './model'
16+
import {
17+
prepareDevEnvConnection,
18+
startRemoteViaSageMakerSshKiro,
19+
tryRemoteConnection,
20+
useSageMakerSshKiroExtension,
21+
} from './model'
22+
import { ensureSageMakerSshKiroExtension } from './sagemakerSshKiroUtils'
1723
import { ExtContext } from '../../shared/extensions'
1824
import { SagemakerClient } from '../../shared/clients/sagemaker'
1925
import { AccessDeniedException } from '@amzn/sagemaker-client'
2026
import { ToolkitError, isUserCancelledError } from '../../shared/errors'
2127
import { showConfirmationMessage } from '../../shared/utilities/messages'
22-
import { RemoteSessionError } from '../../shared/remoteSession'
2328
import {
2429
ConnectFromRemoteWorkspaceMessage,
2530
InstanceTypeInsufficientMemory,
@@ -113,11 +118,11 @@ export async function deeplinkConnect(
113118
)
114119

115120
getLogger().info(
116-
`sm:deeplinkConnect:
117-
domain: ${domain},
118-
appType: ${appType},
119-
workspaceName: ${workspaceName},
120-
namespace: ${namespace},
121+
`sm:deeplinkConnect:
122+
domain: ${domain},
123+
appType: ${appType},
124+
workspaceName: ${workspaceName},
125+
namespace: ${namespace},
121126
eksClusterArn: ${eksClusterArn}`
122127
)
123128

@@ -153,13 +158,21 @@ export async function deeplinkConnect(
153158
)
154159

155160
try {
156-
await startVscodeRemote(
157-
remoteEnv.SessionProcess,
158-
remoteEnv.hostname,
159-
'/home/sagemaker-user',
160-
remoteEnv.vscPath,
161-
'sagemaker-user'
162-
)
161+
const path = '/home/sagemaker-user'
162+
const username = 'sagemaker-user'
163+
164+
if (useSageMakerSshKiroExtension()) {
165+
await ensureSageMakerSshKiroExtension(ctx.extensionContext)
166+
await startRemoteViaSageMakerSshKiro(
167+
remoteEnv.SessionProcess,
168+
remoteEnv.hostname,
169+
path,
170+
remoteEnv.vscPath,
171+
username
172+
)
173+
} else {
174+
await startVscodeRemote(remoteEnv.SessionProcess, remoteEnv.hostname, path, remoteEnv.vscPath, username)
175+
}
163176
} catch (remoteErr: any) {
164177
throw new ToolkitError(
165178
`Failed to establish remote connection: ${remoteErr.message}. Check Remote-SSH logs for details.`,
@@ -174,9 +187,9 @@ export async function deeplinkConnect(
174187
isSMUS
175188
)
176189

177-
if (![RemoteSessionError.MissingExtension, RemoteSessionError.ExtensionVersionTooLow].includes(err.code)) {
190+
if (!isUserCancelledError(err)) {
178191
void vscode.window.showErrorMessage(
179-
`Remote connection failed: ${err.message || 'Unknown error'}. Check Output > Log (Window) for details.`
192+
`Remote connection failed: ${err?.message || 'Unknown error'}. Check Output > Log (Window) for details.`
180193
)
181194
throw err
182195
}

packages/core/src/awsService/sagemaker/detached-server/errorPage.ts

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -79,12 +79,12 @@ export const ErrorText = {
7979
'Your session has expired. Update the credentials associated with the IAM profile or use a valid IAM profile, then try again.',
8080
},
8181
[ExceptionType.INTERNAL_FAILURE]: {
82-
Title: 'Failed to connect remotely to VSCode',
83-
Text: 'Unable to establish remote connection to VSCode. This could be due to several factors. Please try again by clicking the VSCode button. If the problem persists, please contact your admin.',
82+
Title: 'Failed to connect remotely to the Space',
83+
Text: 'Unable to establish remote connection to the Space. This could be due to several factors. Please try again by clicking the connect button. If the problem persists, please contact your admin.',
8484
},
8585
[ExceptionType.RESOURCE_LIMIT_EXCEEDED]: {
8686
Title: 'Connection limit reached',
87-
Text: 'You have 10 active remote connections to this space. Stop an existing connection to start a new one.',
87+
Text: 'You have 10 active remote connections to this space. Stop an existing connection to start a new one. If the problem persists, try restarting the Space.',
8888
},
8989
[ExceptionType.THROTTLING]: {
9090
Title: 'Too many connection attempts',

0 commit comments

Comments
 (0)