Skip to content

Add adapt #110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Sbozzolo wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Add adapt #110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7d0a347
Fix typos
Sbozzolo Mar 27, 2025
f98c6c8
Update cache action to v4
Sbozzolo Mar 28, 2025
ad641b9
Add adapt
Sbozzolo Mar 27, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/OS-Tests.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ jobs:
version: '1.11'

- name: Cache artifacts
uses: actions/cache@v1
uses: actions/cache@v4
if: steps.filter.outputs.run_test == 'true'
env:
cache-name: cache-artifacts
Expand Down
2 changes: 2 additions & 0 deletions NEWS.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@ ClimaComms.jl Release Notes
main
-------

- Added `adapt` function, allowing downstream packages to remove explicit dependency on `Adapt` package. [PR 104](https://github.com/CliMA/ClimaComms.jl/pull/104).

Comment on lines +7 to +8
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Well-documented release note for the adapt function

The release note clearly explains the purpose of adding the adapt function and its benefit for downstream packages. This documentation helps users understand the significance of this addition.

However, the PR reference (PR 104) appears to be inconsistent with the current PR number (110).

🏁 Script executed:

#!/bin/bash
# Check if PR 104 is related to the adapt functionality
echo "Checking PR 104 content..."
gh pr view 104 --json title,body -q '.title + "\n\n" + .body' 2>/dev/null

echo -e "\nChecking mentions of adapt in PR 104 commits..."
gh pr view 104 --json commits -q '.commits[].messageHeadline' | grep -i adapt

Length of output: 1601

Action: Update PR reference in NEWS.md for the adapt feature

  • The release note at NEWS.md (lines 7–8) documents the addition of the adapt function, which is well-described.
  • However, our verification confirms that PR 104 is related to logging documentation rather than the adapt function.
  • Please update the release note to reference the correct PR number (likely PR 110) or the appropriate one associated with the adapt changes.

v0.6.6
-------
- Replaced `MPIFileLogger` with `FileLogger` and added an `OnlyRootLogger` logger that silences non-root processes [PR 104](https://github.com/CliMA/ClimaComms.jl/pull/104).
Expand Down
1 change: 1 addition & 0 deletions docs/src/apis.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ ClimaComms.AbstractGraphContext
ClimaComms.context
ClimaComms.graph_context
Adapt.adapt_structure(::Type{<:AbstractArray}, ::ClimaComms.AbstractCommsContext)
Adapt.adapt(::ClimaComms.AbstractDevice, ::Any)
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Correct API documentation to match implementation

The documentation currently shows Adapt.adapt(::ClimaComms.AbstractDevice, ::Any) but the implementation in src/adapt.jl defines the function as adapt(device::AbstractDevice, x) without the Adapt namespace. This should be updated to match the actual implementation.

-Adapt.adapt(::ClimaComms.AbstractDevice, ::Any)
+adapt(::ClimaComms.AbstractDevice, ::Any)
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
Adapt.adapt(::ClimaComms.AbstractDevice, ::Any)
adapt(::ClimaComms.AbstractDevice, ::Any)

```

## Logging
Expand Down
23 changes: 21 additions & 2 deletions src/adapt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ Adapt.adapt_structure(Array, ClimaComms.context(ClimaComms.CUDADevice())) -> Cli

!!! note
By default, adapting to `Array` creates a `CPUSingleThreaded` device, and
there is currently no way to conver to a CPUMultiThreaded device.
there is currently no way to convert to a CPUMultiThreaded device.
"""
Adapt.adapt_structure(to::Type{<:AbstractArray}, ctx::AbstractCommsContext) =
context(Adapt.adapt(to, device(ctx)))
Expand All @@ -31,7 +31,26 @@ Adapt.adapt_structure(Array, ClimaComms.CUDADevice()) -> ClimaComms.CPUSingleThr

!!! note
By default, adapting to `Array` creates a `CPUSingleThreaded` device, and
there is currently no way to conver to a CPUMultiThreaded device.
there is currently no way to convert to a CPUMultiThreaded device.
"""
Adapt.adapt_structure(::Type{<:AbstractArray}, device::AbstractDevice) =
CPUSingleThreaded()


"""
adapt(device::AbstractDevice, x)

Adapt an object `x` to be compatible with the specified `device`.
"""
function adapt(device::AbstractDevice, x)
Copy link
Copy Markdown
Member Author

@Sbozzolo Sbozzolo Mar 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@CodeRabbit generate a docstring for this function

Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@Sbozzolo Here's a docstring for the adapt function:

"""
    adapt(device::AbstractDevice, x)

Adapt an object `x` to be compatible with the specified `device`. This is a convenience method 
that internally calls `adapt(array_type(device), x)`.

# Arguments
- `device`: The target device to adapt the object for
- `x`: The object to be adapted

# Returns
The adapted object compatible with the specified device

# Example
```julia
# Adapt an array to be compatible with CUDA device
cuda_array = adapt(CUDADevice(), my_array)

# Adapt an array to be compatible with CPU
cpu_array = adapt(CPUSingleThreaded(), my_cuda_array)

"""


<!-- This is an auto-generated reply by CodeRabbit -->

Copy link
Copy Markdown
Member Author

@Sbozzolo Sbozzolo Mar 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@CodeRabbit add unit tests for this function

Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 27, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

The project has a test directory with several files including runtests.jl. Let me examine how device-related functionality is tested to understand the testing approach better.

Now let me look at the actual implementation of the devices and CUDA support:

I need to understand better how arrays are adapted and check if there's existing test infrastructure for CUDA:

Let me check if there are any existing test files specifically for adapt.jl:

I see there are references to adapt in the tests, but I need to understand more about the device types and array behavior. Let me look for additional context:

Let's examine the entire test setup to understand how I should structure the unit tests:

Based on the exploration so far, I need to understand the available device types better to write appropriate tests:

Let's see if we can find how the array_type function is implemented:

Before writing the tests, I should have a better understanding of how Adapt.adapt is used:

Let me check if there's a specific test file for the adapt functionality:

Based on the information gathered, I'll now create appropriate unit tests for the adapt(device::AbstractDevice, x) function. The tests should cover adapting different types of data to different device types.

@Sbozzolo Here are unit tests for the adapt(device::AbstractDevice, x) function:

@testset "adapt" begin
    # Test adapting arrays to CPU device
    cpu_device = CPUSingleThreaded()
    x_array = [1, 2, 3, 4]
    
    # Test that adapt with a device forwards to adapt with array_type
    @test adapt(cpu_device, x_array) == adapt(array_type(cpu_device), x_array)
    @test adapt(cpu_device, x_array) == adapt(Array, x_array)
    
    # Test with a more complex structure
    complex_data = (a = [1, 2, 3], b = Dict(:x => [4, 5, 6]))
    @test adapt(cpu_device, complex_data) == adapt(Array, complex_data)
    
    # Test with scalar values (should be unchanged)
    @test adapt(cpu_device, 42) == 42
    @test adapt(cpu_device, "string") == "string"
    
    # If CUDA support is available, test with CUDA device
    if ClimaCommsCUDA !== nothing
        cuda_device = CUDADevice()
        
        # Create a CPU array and adapt it to CUDA
        cpu_array = [1.0, 2.0, 3.0, 4.0]
        cuda_array = adapt(cuda_device, cpu_array)
        
        # Verify it's now a CuArray
        @test cuda_array isa CUDA.CuArray
        @test Array(cuda_array) == cpu_array
        
        # Test that our device-based adapt works the same as direct array_type adapt
        @test cuda_array == adapt(array_type(cuda_device), cpu_array)
        
        # Test adapting back to CPU
        @test adapt(cpu_device, cuda_array) isa Array
        @test adapt(cpu_device, cuda_array) == cpu_array
    end
end

These tests verify that:

  1. The adapt(device, x) function correctly forwards to adapt(array_type(device), x)
  2. It works with simple arrays and more complex nested structures
  3. It handles scalar values appropriately
  4. If CUDA is available, it correctly converts between CPU and CUDA arrays

Add these tests to your existing test file, likely in test/runtests.jl within the appropriate test section.

🏁 Script executed:

#!/bin/bash
# Check test directory structure
find test -type f | sort

# Look at the runtests.jl file to understand the testing structure
cat test/runtests.jl

# Look for existing tests related to Adapt or adapt
rg "adapt" test/ -i

# Look for device types and array_type implementation
rg "AbstractDevice" -A 3 src/
rg "array_type" -A 3 src/
rg "struct.*Device" src/

# Check imports and implementation details
rg "import Adapt" src/

Length of output: 12765

🏁 Script executed:

#!/bin/bash
# Check if there are any tests for device-related functionality
cat test/runtests.jl

# Check if there are specific tests for CUDA devices or arrays
rg "CUDADevice" test/
rg "CUDA" test/

# Check the implementation of array_type
rg "function array_type" -A 5 src/

# Look for concrete implementations of AbstractDevice
rg "struct.*<:.*AbstractDevice" src/

Length of output: 11010

🏁 Script executed:

#!/bin/bash
# Check for specific device implementations
cat src/devices.jl

# Check if there are any utilities or helper functions for testing
find test -name "*helpers*" -o -name "*utils*"
cat test/utils.jl 2>/dev/null || echo "No utils.jl file found"

Length of output: 8109

🏁 Script executed:

#!/bin/bash
# Check if CUDA.jl is used in the tests
rg "using CUDA" test/

# Check how test arrays are created
rg "Array\(" test/
rg "CuArray" test/

# Look at existing adapt-related code
rg -l "adapt_structure" src/
cat src/adapt.jl

Length of output: 1683

🏁 Script executed:

#!/bin/bash
# Check if there are test files specifically for adapt functionality
find test -name "*adapt*"

# Check for existing test adaptations with different device types
rg "adapt.*device" test/

Length of output: 54

🏁 Script executed:

#!/bin/bash
# Check for code that uses the adapt function with devices
rg "adapt\(.*device" src/ test/

# Look at Project.toml to understand dependencies 
cat Project.toml

# Check if there's information about test dependencies
cat test/Project.toml 2>/dev/null || echo "No test/Project.toml found"

Length of output: 1398

🏁 Script executed:

#!/bin/bash
# Check complete content of test files to understand the testing pattern
cat test/runtests.jl

# Check if there are mock devices for testing
rg "Mock.*Device" test/ src/

Length of output: 8137

🏁 Script executed:

#!/bin/bash
# Look at device-related code to understand available types
cat src/devices.jl

# Check for any existing Adapt-related tests
grep -r "adapt" test/ --include="*.jl"

Length of output: 8633

🏁 Script executed:

#!/bin/bash
# Find the implementation of array_type
rg "array_type\(" src/ -A 3

Length of output: 467

🏁 Script executed:

#!/bin/bash
# Look for usage patterns of Adapt.adapt
rg "Adapt\.adapt" src/ test/

Length of output: 1276

🏁 Script executed:

#!/bin/bash
# Check if there's a specific adapt test
cat test/test_adapt.jl 2>/dev/null || echo "No specific test_adapt.jl file"

Length of output: 95

@Sbozzolo Here are unit tests for the new adapt(device::AbstractDevice, x) function:

@testset "adapt with device" begin
    # Test with CPU device
    cpu_device = CPUSingleThreaded()
    
    # Test with simple arrays
    x_array = [1, 2, 3, 4]
    
    # Test that adapt(device, x) correctly forwards to adapt(array_type(device), x)
    @test adapt(cpu_device, x_array) == adapt(array_type(cpu_device), x_array)
    @test adapt(cpu_device, x_array) == adapt(Array, x_array)
    
    # Test with nested structures
    nested_data = (a = [1, 2, 3], b = Dict(:x => [4, 5, 6]))
    @test adapt(cpu_device, nested_data) == adapt(Array, nested_data)
    
    # Test with scalar values (should remain unchanged)
    @test adapt(cpu_device, 42) == 42
    @test adapt(cpu_device, "test") == "test"
    
    # Test with CUDA device if available
    @static if isdefined(ClimaComms, :ClimaCommsCUDAExt)
        cuda_device = CUDADevice()
        
        # Create a test array and adapt it to CUDA
        cpu_array = [1.0, 2.0, 3.0, 4.0]
        cuda_array = adapt(cuda_device, cpu_array)
        
        # Verify that the device-based adapt produces the same result as array_type-based adapt
        @test cuda_array == adapt(array_type(cuda_device), cpu_array)
        @test cuda_array isa CUDA.CuArray
        
        # Test converting back to CPU
        @test adapt(cpu_device, cuda_array) isa Array
        @test Array(adapt(cpu_device, cuda_array)) == cpu_array
        
        # Test with mixed data types
        mixed_data = (cpu = [1.0, 2.0], cuda = cuda_array)
        adapted_mixed = adapt(cpu_device, mixed_data)
        @test adapted_mixed.cpu isa Array
        @test adapted_mixed.cuda isa Array
    end
end

These tests should be added to the existing @testset "Adapt" block in test/runtests.jl.

The tests verify that:

  1. The adapt(device, x) function correctly forwards to adapt(array_type(device), x)
  2. It works with various data types including arrays, scalars, and nested structures
  3. When CUDA is available, it correctly handles GPU arrays and conversions between CPU and GPU
  4. It preserves the behavior of the original adapt function while providing a convenient device-based interface

return Adapt.adapt(array_type(device), x)
end
Comment on lines +40 to +47
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Ensure the adapt function is exported

The implementation of the adapt(device::AbstractDevice, x) function looks good. However, to fulfill the PR objective of allowing downstream packages to eliminate their explicit dependency on the Adapt package, this function should be exported from the module.

🏁 Script executed:

#!/bin/bash
# Check if the adapt function is exported in ClimaComms.jl
echo "Checking if adapt is exported in the module..."
if grep -q "export.*adapt" src/ClimaComms.jl; then
  echo "✅ The adapt function is exported"
  grep -n "export.*adapt" src/ClimaComms.jl
else
  echo "❌ The adapt function is not exported"
  echo "Suggested addition to src/ClimaComms.jl:"
  echo "export adapt"
fi

Length of output: 355

Action Required: Export the adapt function from ClimaComms.jl

The implementation of the adapt(device::AbstractDevice, x) function looks good. However, as part of the PR objectives (and in line with the recent learnings), the adapt function must be exported so that downstream packages can use it without requiring the explicit Adapt dependency.

  • Location: src/ClimaComms.jl
  • Required Change: Add the export statement: 
    export adapt

Please update the module accordingly to allow the direct usage of adapt in downstream packages.


"""
adapt(device::AbstractCommsContext, x)

Adapt an object `x` to be compatible with the specified `context`.
"""
function adapt(context::AbstractCommsContext, x)
return Adapt.adapt(array_type(device(context)), x)
end
25 changes: 25 additions & 0 deletions test/adapt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
import ClimaComms
import Adapt: adapt

@testset "Adapt" begin
# Test with CPU device
cpu_device = ClimaComms.CPUSingleThreaded()

# Test with simple arrays
x_array = [1, 2, 3, 4]

# Test that adapt(device, x) correctly forwards to adapt(array_type(device), x)
@test adapt(cpu_device, x_array) ==
adapt(ClimaComms.array_type(cpu_device), x_array)
@test adapt(cpu_device, x_array) == adapt(Array, x_array)
@test adapt(ClimaComms.SingletonCommsContext(cpu_device), x_array) ==
adapt(cpu_device, x_array)

# Test with nested structures
nested_data = (a = [1, 2, 3], b = Dict(:x => [4, 5, 6]))
@test adapt(cpu_device, nested_data) == adapt(Array, nested_data)

# Test with scalar values (should remain unchanged)
@test adapt(cpu_device, 42) == 42
@test adapt(cpu_device, "test") == "test"
end
4 changes: 4 additions & 0 deletions test/runtests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -260,3 +260,7 @@ end
@testset "Logging" begin
include("logging.jl")
end

@testset "Adapt" begin
include("adapt.jl")
end
Comment on lines +264 to +266
Copy link
Copy Markdown

@coderabbitai coderabbitai bot Mar 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Double check test file name

There are now two different test sections with the same name "Adapt" (line 232 and line 264). The existing test set at line 232 already tests Adapt.adapt functionality, while this new one includes a separate file. Consider whether these should be combined or if the new test set should have a more specific name.

🏁 Script executed:

#!/bin/bash
# Check the content of the adapt.jl test file
if [ -f "test/adapt.jl" ]; then
  echo "✅ test/adapt.jl exists. Content:"
  cat test/adapt.jl
else
  echo "❌ test/adapt.jl does not exist"
fi

Length of output: 985

Double-check duplicate test set names in tests

It appears that there are two test sets named "Adapt": one defined at line 232 in test/runtests.jl and another coming from test/adapt.jl (included at line 264). Since test/adapt.jl also uses the name "Adapt" for its test set, please verify whether these sections are intended to test distinct aspects of the functionality. If they overlap, consider one of the following:

  • Merge the test sets: Combine the tests into a single "Adapt" test set if they cover the same functionality.
  • Rename one of the test sets: Use a more specific name for the new section to clearly differentiate its purpose from the original test set.

Loading