Skip to content

Commit 3571c4d

Browse files
authored
Update container image tags and GPU support methods (#699)
- Changed version tags in `OllamaContainerImageTags.cs` for Ollama and Open Web UI images. - Modified AMD GPU support in `OllamaResourceBuilderExtensions.cs` to use a new method that appends `-rocm` to the image tag and includes additional device arguments. - Updated tests in `AddOllamaTests.cs` to reflect the new GPU support methods, splitting the previous test into separate tests for Nvidia and AMD, with updated assertions. Fixes #692
1 parent b430044 commit 3571c4d

File tree

3 files changed

+61
-10
lines changed

3 files changed

+61
-10
lines changed

src/CommunityToolkit.Aspire.Hosting.Ollama/OllamaContainerImageTags.cs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ internal static class OllamaContainerImageTags
44
{
55
public const string Registry = "docker.io";
66
public const string Image = "ollama/ollama";
7-
public const string Tag = "0.6.8";
7+
public const string Tag = "0.7.1";
88

99
public const string OpenWebUIRegistry = "ghcr.io";
1010
public const string OpenWebUIImage = "open-webui/open-webui";
11-
public const string OpenWebUITag = "0.5.20";
11+
public const string OpenWebUITag = "0.6.10";
1212
}

src/CommunityToolkit.Aspire.Hosting.Ollama/OllamaResourceBuilderExtensions.cs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -103,11 +103,23 @@ public static IResourceBuilder<OllamaResource> WithGPUSupport(this IResourceBuil
103103
return vendor switch
104104
{
105105
OllamaGpuVendor.Nvidia => builder.WithContainerRuntimeArgs("--gpus", "all"),
106-
OllamaGpuVendor.AMD => builder.WithContainerRuntimeArgs("--device", "/dev/kfd"),
106+
OllamaGpuVendor.AMD => builder.WithAMDGPUSupport(),
107107
_ => throw new ArgumentException("Invalid GPU vendor", nameof(vendor))
108108
};
109109
}
110110

111+
private static IResourceBuilder<OllamaResource> WithAMDGPUSupport(this IResourceBuilder<OllamaResource> builder)
112+
{
113+
if (builder.Resource.TryGetLastAnnotation<ContainerImageAnnotation>(out var containerAnnotation))
114+
{
115+
if (containerAnnotation.Tag?.EndsWith("rocm") == false)
116+
{
117+
containerAnnotation.Tag += "-rocm";
118+
}
119+
}
120+
return builder.WithContainerRuntimeArgs("--device", "/dev/kfd", "--device", "/dev/dri");
121+
}
122+
111123
private static OllamaResource AddServerResourceCommand(
112124
this OllamaResource ollamaResource,
113125
string name,

tests/CommunityToolkit.Aspire.Hosting.Ollama.Tests/AddOllamaTests.cs

Lines changed: 46 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -581,13 +581,40 @@ public void OllamaResourceCommandsUpdateState(string commandType)
581581
Assert.Equal(ResourceCommandState.Enabled, state);
582582
}
583583

584-
[Theory]
585-
[InlineData(OllamaGpuVendor.Nvidia, "--gpus", "all")]
586-
[InlineData(OllamaGpuVendor.AMD, "--device", "/dev/kfd")]
587-
public async Task WithGPUSupport(OllamaGpuVendor vendor, string expectedArg, string expectedValue)
584+
[Fact]
585+
public async Task WithNvidiaGPUSupport()
586+
{
587+
var builder = DistributedApplication.CreateBuilder();
588+
_ = builder.AddOllama("ollama").WithGPUSupport(OllamaGpuVendor.Nvidia);
589+
590+
using var app = builder.Build();
591+
592+
var appModel = app.Services.GetRequiredService<DistributedApplicationModel>();
593+
594+
var resource = Assert.Single(appModel.Resources.OfType<OllamaResource>());
595+
596+
Assert.True(resource.TryGetLastAnnotation(out ContainerRuntimeArgsCallbackAnnotation? argsAnnotations));
597+
ContainerRuntimeArgsCallbackContext context = new([]);
598+
await argsAnnotations.Callback(context);
599+
600+
Assert.Collection(
601+
context.Args,
602+
arg =>
603+
{
604+
Assert.Equal("--gpus", arg);
605+
},
606+
arg =>
607+
{
608+
Assert.Equal("all", arg);
609+
}
610+
);
611+
}
612+
613+
[Fact]
614+
public async Task WithAMDGPUSupport()
588615
{
589616
var builder = DistributedApplication.CreateBuilder();
590-
_ = builder.AddOllama("ollama").WithGPUSupport(vendor);
617+
_ = builder.AddOllama("ollama").WithGPUSupport(OllamaGpuVendor.AMD);
591618

592619
using var app = builder.Build();
593620

@@ -603,12 +630,24 @@ public async Task WithGPUSupport(OllamaGpuVendor vendor, string expectedArg, str
603630
context.Args,
604631
arg =>
605632
{
606-
Assert.Equal(expectedArg, arg);
633+
Assert.Equal("--device", arg);
634+
},
635+
arg =>
636+
{
637+
Assert.Equal("/dev/kfd", arg);
607638
},
608639
arg =>
609640
{
610-
Assert.Equal(expectedValue, arg);
641+
Assert.Equal("--device", arg);
642+
},
643+
arg =>
644+
{
645+
Assert.Equal("/dev/dri", arg);
611646
}
612647
);
648+
649+
Assert.True(resource.TryGetLastAnnotation<ContainerImageAnnotation>(out var imageAnnotation));
650+
Assert.NotNull(imageAnnotation);
651+
Assert.EndsWith("-rocm", imageAnnotation.Tag, StringComparison.OrdinalIgnoreCase);
613652
}
614653
}

0 commit comments

Comments
 (0)