使用 Amazon SageMaker 构建代码伴侣应用

引言

代码伴侣是一类基于大型语言模型（LLMs）的智能化编程工具，它可以通过自然语言理解和代码生成能力，为程序员提供高效、准确的代码编写辅助，从而使整个软件开发过程更加流畅高效。然而，主流商业代码助手通常基于封闭的专有模型，不仅订阅费用昂贵，还可能带来隐私、安全、版权等方面的潜在风险。

近期涌现出一批优秀的开源代码大模型，如 DeepSeek Coder、CodeLlama、StarCoder 等，它们基于海量开源代码训练，性能接近商业模型，为开发者提供了一条前景广阔的新路径。本文旨在利用 SageMaker Endpoint 部署开源代码模型，并通过开源插件 Continue 完成整个代码助手的开发。

SageMaker Endpoint 部署模型

Large Model Inference（LMI）

LMI 容器是一组专门为基础模型（FM）推理而构建的高性能 Docker 容器。使用这些容器，可以利用高性能的开源推理库，如 vLLM、TensorRT-LLM 和 Transformers NeuronX，在Amazon SageMaker Endpoints 上快速部署 LLM。LMI 容器将模型服务器与开源推理库捆绑在一起，提供一体化的 LLM 部署解决方案，可以在几分钟内部署流行的开源模型。更多详细信息请看官方文档。

vLLM

vLLM 是一个开源推理框架，它通过使用 PagedAttention 注意力算法，可以有效地管理 LLM 中的 Key 和 Value，优化 KVCache 的显存使用，从而提高显存利用率和模型服务吞吐量。另外 vLLM 也支持 Continuous batching ，以及不同的量化方式 GPTQ, AWQ, SqueezeLLM, FP8 KV Cache 来部署大语言模型。

应用 PagedAttention 处理不同 request 的示例（图片来源于 vLLM Blog）

DeepSeek Coder

DeepSeek Coder 提供了一系列的代码语言模型，每个模型都是从头开始经过了 2T 的 token 训练，包含 87% 的代码和 13% 的英文和中文自然语言。并且提供了不同大小的代码模型，从 1B 到 33B 版本不等。每个模型都是在项目级代码语料库上进行预训练，支持 16K 的窗口大小和额外的 FIM 任务以支持项目级代码补全和代码插入。在编码能力方面，DeepSeek Coder在多种编程语言和各种基准测试中都有不错的表现。

Deepseek coder 不同语言 benchmark 效果（图片来源于 DeepSeek Coder github）

DeepSeek Coder 训练过程

DeepSeeK Coder 系列模型训练主要分为以下三个步骤，更详细的训练过程请移步阅读 DeepSeek Coder 的论文：

第一步，使用包含 87% 代码、10% 与代码相关的语言数据和 3% 与非代码相关的中文语言的数据集进行预训练。在这一步骤中，使用了 1.8T 个 tokens 和 4K 的窗口大小对模型进行预训练。

第二步，在另外 200B 个 token 上进行进一步预训练，将模型扩展到 16K 的窗口大小，得到基础模型（DeepSeek-Coder-Base）。

第三步，在 2B 个 token 的指令数据集上进行指令微调，得到指令微调模型（DeepSeek-Coder-Instruct）。让模型拥有指令遵循或者对话的能力。

部署模型

现在我们使用 LMI 容器中的 vLLM 作为推理后端，选择 DeepSeek Coder base 6.7B 模型作为代码补全和代码插入的基础模型，使用 DeepSeek Coder instruct 6.7B 来作为 chat 模型，支持输入指令生成完整代码的功能。部署使用的机器实例我们使用 2 台 g5.2xlarge 分别进行部署。本文旨在介绍如何使用 SageMaker Endpoint 结合 Continue 一起实现一个代码助手，实际应用中，可以结合自己的编程语言，编程场景选择不同的模型进行部署。

我们这里以部署 DeepSeek Coder base 6.7B 这个模型为例，介绍如何使用 LMI 在 SageMaker 部署并调用大语言模型。完整的代码可以使用 GitHub 上的分步教程 Notebooks，重现所有步骤。

步骤一：引入需要的包，以及获取基础的 IAM role 等

import boto3
import sagemaker
from sagemaker import Model, image_uris, serializers, deserializers

role = sagemaker.get_execution_role()  # execution role for the endpoint
sess = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs
region = sess._region_name  # region name of the current SageMaker Studio environment
account_id = sess.account_id()  # account_id of the current SageMaker Studio environment
sagemaker_default_bucket = sess.default_bucket()

步骤二：打包模型基础文件

注意：如果是在中国区，需要将 Hugging face 的 HF_ENDPOINT 替换为 ‘https://hf-mirror.com’，完整代码参考上方的 deployment guide。

基本配置信息写入 serving.properties

model_name="deepseek-ai/deepseek-coder-6.7b-instruct"
model_lineage=model_name.split("/")[0]
model_specific_name = model_name.split("/")[1]

with open("serving.properties", "w") as wf:
    wf.write("engine=Python\n")
    wf.write(f"option.model_id={model_name}\n")
    wf.write("option.rolling_batch=vllm\n")
    wf.write("option.max_model_len=8192\n")

打包模型配置信息

%%sh
mkdir mymodel
mv serving.properties mymodel/
tar czvf mymodel.tar.gz mymodel/
rm -rf mymodel

步骤三：开始构建 SageMaker Endpoint

指定 LMI 容器以及对应的版本

image_uri = image_uris.retrieve(
    framework="djl-lmi",
    region=sess.boto_session.region_name,
    version="0.29.0"
    )

上传模型配置文件到 S3 桶

s3_code_prefix = f"large-model-lmi/code-{model_lineage}-{model_specific_name}"
bucket = sess.default_bucket()  # bucket to house artifacts
code_artifact = sess.upload_data("mymodel.tar.gz", bucket, s3_code_prefix)
print(f"S3 Code or Model tar ball uploaded to --- > {code_artifact}")

model = Model(image_uri=image_uri, model_data=code_artifact, role=role)

步骤四：开始部署模型到 SageMaker Endpoint

这里我们选择 g5 实例来进行部署，g5.2xlarge 拥有 1 张 A10g 的显卡，24G 显存，其他更多信息可以移步 EC2 实例类型网页获取。

instance_type = "ml.g5.2xlarge"
endpoint_name = sagemaker.utils.name_from_base(f"{model_lineage}-{model_specific_name.replace('.','-')}")

model.deploy(initial_instance_count=1,
             instance_type=instance_type,
             endpoint_name=endpoint_name,
             # container_startup_health_check_timeout=3600
            )

# our requests and responses will be in json format so we specify the serializer and the deserializer
predictor = sagemaker.Predictor(
    endpoint_name=endpoint_name,
    sagemaker_session=sess,
    serializer=serializers.JSONSerializer(),
    deserializer=sagemaker.deserializers.JSONDeserializer(),
)

步骤五：调用测试

模型部署成功以后，可直接进行测试调用。以下代码展示了使用常规 Prompt 方式调用模型，如果需要使用 streaming 方式，可以参考提供的完整代码示例里的脚本进行调用。如果需要使用 Message API （Continue 里面的 chat 接口需要使用 message api 的方式进行调用）的方式，可以参考 deepseek-ai/deepseek-coder-6.7b-instruct LMI deployment guide 里提供的代码示例进行调用。

prompt = """<｜fim▁begin｜>def quick_sort(arr):
    if len(arr) <= 1:
        return arr
    pivot = arr[0]
    left = []
    right = []
<｜fim▁hole｜>
        if arr[i] < pivot:
            left.append(arr[i])
        else:
            right.append(arr[i])
    return quick_sort(left) + [pivot] + quick_sort(right)<｜fim▁end｜>"""
parameters = {
        "max_new_tokens":1024, 
    }
response = predictor.predict(
    {"inputs": prompt, "parameters": parameters}
)

print(response["generated_text"])

步骤六：资源清理

如果不再需要调用模型服务，即可使用以下命令释放资源，避免导致服务在启用状态被收费：

sess.delete_endpoint(endpoint_name)
sess.delete_endpoint_config(endpoint_name)
model.delete_model()

IDE 插件集成

当部署完成 SageMaker 推理节点后，就可以使用第三方插件调用 SageMaker 接口来实现。

例如 https://github.com/Continuedev/Continue 是一款较为知名的第三方开源插件，支持在 Visual Studio Code 和 IntelliJ IDE 环境中，调用 Amazon Bedrock 等 API provider 来构建您自己的代码开发插件。

本文在 Continue 开源版本基础上，增加了使用 SageMaker 作为 model provider，代码已合入 Continue 项目。以下将以 VSCode 为例，介绍如何在开发 IDE 中如何构建 LLM 辅助代码伴侣应用。

配置模型鉴权

首先我们需要创建具备 IAM SageMaker 调用权限 user 的 access key 用来鉴权，如果当前电脑本地已配置过，可以跳过这一步骤。

创建 IAM user 和 access key

创建 SageMaker endpoint 调用权限 Policy

浏览器打开亚马逊云科技控制台 IAM 创建 Policy 页面，点击“JSON”，并在 Policy editor 编辑器中输入以下 Policy

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "sagemaker:InvokeEndpointWithResponseStream",
                "sagemaker:InvokeEndpoint"
            ],
            "Resource": ["arn:aws:sagemaker:*:*:endpoint/*"]
        }
    ]
}

点击 Next 下一步后，在 Policy name 填入“sagemaker-invoke-endpoint” ，然后“Create policy”创建权限

创建 SageMaker endpoint 调用用户

浏览器打开亚马逊云科技控制台 IAM 创建用户页面，然后填写用户名，用户名可以自己指定，然后 Next 下一步

然后选择“Attach policies directly”，在搜索框搜索刚刚已经创建好的权限“sagemaker-invoke-endpoint”，并选中，之后 Next 下一步

然后点击“Create user”创建用户

添加 Access key / Secret key

在创建完成用户后的页面（或者打开 IAM User 页面），点击刚刚创建的用户名

点击“Security credentials”标签页，然后点击“Create access key”

选择“Local code”本地代码，然后点击下一步

然后勾选“I understand”然后 Next 下一步

然后点击“Create access key”

创建成功后，可以将 Access key / Secret access key 复制保存下来，注意这个 key 只能够显示一次，如果忘记需要重新创建。

编辑本地鉴权文件：

Mac / Linux 用户可以编辑 ~/.aws/credentials 文件
Windows 用户可以编辑在文件管理器打开 %UserProfile%\.aws\credentials ，例如：C:\Users\YourUsername\.aws\credentials，如果 .aws 文件夹不存在，可以手动创建它

使用 VSCode 编辑器创建或编辑 credentials 文件（注意没有 .txt 等任何后缀），在这里分别填入上一步创建的 access key 和 secret access key，然后保存，即可完成全部权限配置。

[default]
aws_access_key_id = [YOUR_ACCESS_KEY]
aws_secret_access_key = [YOUR_SECRET_ACCESS_KEY]

插件安装

VSCode 应用市场安装

打开 VSCode 扩展商店，搜索 Continue，点击“安装”安装插件（在 0.8.46 版本后支持已 SageMaker 模型调用，如果之前已经安装该插件，建议更新到最新版）

手动编译安装

如果您想定制自己专属代码伴侣应用，可以参考开发文档进行更进一步定制开发。Continue 插件开发以及编译发布可以参考：https://github.com/Continuedev/Continue/blob/main/CONTRIBUTING.md。下面将介绍如何通过源码形式安装打包 VSCode 插件。

下载代码：

git clone https://github.com/Continuedev/Continue.git

使用 VSCode 打开下载的 Continue 项目文件夹，VSCode 命令面板（cmd/ctrl+shift+p），选择 Tasks: Run Task，然后选择 install-all-dependencies。
如果要打包扩展，请在 extensions/vscode 目录下运行 npm run package。这将生成 extensions/vscode/build/Continue-{VERSION}.vsix，您可以通过右键单击并选择“安装扩展 VSIX”来安装它。

插件参数配置

插件安装完成后，如果提示配置，可以点击 skip 跳过配置，然后手动修改配置文件，来实现访问 SageMaker 上部署的语言模型。

依次点击图中示意 1，2 的位置，打开 Continue 配置文件，然后修改 3 中的具体配置文件，配置文件填写可以参考下方说明
对于代码问答助手应用，需要使用针对 instruct 形式训练的对话模型，这里使用前文部署的 deepseek-6.7b-instruct 模型，需要在配置文件”models”中添加相应模型，具体填写参数说明：
- title：在插件中显示的模型名称，可以填写任何名称
- provider：这里固定为 sagemaker
- model: 前文部署好的对应 SageMaker endpoint 的名称，这里注意 endpoint 名称为 instruct 版的模型
- region：项需要填写部署模型所在的 region

完整示例如下：

 {
  "models": [
    {
      "title": "deepseek-6.7b-instruct",
      "provider": "sagemaker",
      "model": "deepseek-ai-deepseek-coder-6-7b-instruct...",
      "region": "us-west-2"
    }
  ],
  ...
 }

对于代码自动补全应用，需要使用在代码补全任务训练的预训练模型，这里使用前文部署的 deepseek-6.7b-base 模型，需要在配置文件”tabAutocompleteModel”中改为前文部署的模型，具体填写参数说明：
- title：在插件中显示的模型名称，可以填写任何名称
- provider：这里固定为 sagemaker
- model：前文部署好的对应 SageMaker endpoint 的名称，这里注意 endpoint 名称为 base 版的模型
- region：项需要填写部署模型所在的 region

完整示例如下：

{
  ...
  "tabAutocompleteModel": {
    "title": "deepseek-6.7b-base",
    "provider": "sagemaker",
    "model": "deepseek-coder-6-7b-base...",
    "region": "us-west-2"
  },
  ...
 }

代码问答

在 Continue 插件代码问答标签页，选择配置好的模型名称 deepseek-6.7b-instruct，就可以与 SageMaker 上的大语言模型进行问答。

代码补全

当开启 Continue 后，代码补全功会默认开启，当用户输入代码光标等待时自动提示补全代码。Continue 插件已经适配了不同的代码补全模型例如 Deepseek Coder、StarCoder、CodeLlama 等，会根据配置文件中的 model 参数构造代码补全的模板。

例如下图示意为当光标输入等待时，模型会在输入光标处自动给出灰色的提示补全代码，按 Tab 键采纳提示补全结果。

以上为使用 Continue 插件辅助开发的简单介绍，详细可以参考 Continue 文档查看更多使用技巧。

结语

本文详细介绍了如何利用 Amazon SageMaker 部署开源代码模型 DeepSeek Coder，并通过开源插件 Continue 实现一个功能强大的代码助手应用。通过 SageMaker Endpoint 部署 LMI 容器和 vLLM 推理框架，可以快速高效地将开源代码模型投入生产环境。同时结合 Continue 插件的代码问答和自动补全功能，为开发者提供了一个智能、便捷的编程辅助工具。

这种基于开源模型的代码助手方案，不仅能够避免商业模型的高昂费用和潜在的隐私风险，而且还能充分利用社区贡献的优质代码资源，持续提升模型性能。相信随着开源代码模型的不断发展和完善，这种基于云端部署的代码助手应用将会为广大开发者带来更加智能、高效的编程体验。

亚马逊AWS官方博客