Add Wave containers reports (preview)

This commit adds the ability to create a contaiens
report file used by Wave when running a pipeline execution.

The contains reports can be used to precisily replicate the
pipeline execution with the puntcual list of containers build
by Wave.

Signed-off-by: Paolo Di Tommaso <[email protected]>

Author: pditommaso

docs/wave.md

@@ -196,3 +196,8 @@ The following configuration options are available:
 `wave.strategy`
 : The strategy to be used when resolving ambiguous Wave container requirements (default: `'container,dockerfile,conda,spack'`).
 
+`wave.report.enabled` (preview)
+: Enable the reporting of the Wave containers used during the pipeline execution (default: `false`, requires version `23.06.0-edge` or later).
+
+`wave.report.file` (preview)
+: The name of the containers report file (default: `containers-<timestamp>.config` requires version `23.06.0-edge` or later).

plugins/nf-wave/src/main/io/seqera/wave/plugin/DescribeContainerResponse.groovy

@@ -0,0 +1,57 @@
+package io.seqera.wave.plugin
+
+import java.time.Instant
+
+import groovy.transform.Canonical
+import groovy.transform.CompileStatic
+/**
+ * Model a container request record
+ *
+ * @author Paolo Di Tommaso <[email protected]>
+ */
+@Canonical
+@CompileStatic
+class DescribeContainerResponse {
+
+    static class User {
+        Long id
+        String userName
+        String email
+    }
+
+    @Canonical
+    static class RequestInfo {
+        final User user
+        final Long workspaceId
+        final String containerImage
+        final ContainerConfig containerConfig
+        final String platform
+        final String towerEndpoint
+        final String fingerprint
+        final Instant timestamp
+        final String zoneId
+        final String ipAddress
+    }
+
+    @Canonical
+    static class BuildInfo {
+        final String containerFile
+        final String condaFile
+        final String buildRepository
+        final String cacheRepository
+    }
+
+    @Canonical
+    static class ContainerInfo {
+        String image
+        String digest
+    }
+
+    final String token
+    final Instant expiration
+    final RequestInfo request
+    final BuildInfo build
+    final ContainerInfo source
+    final ContainerInfo wave
+
+}

plugins/nf-wave/src/main/io/seqera/wave/plugin/WaveClient.groovy

@@ -22,18 +22,22 @@ import java.net.http.HttpRequest
 import java.net.http.HttpResponse
 import java.nio.file.Path
 import java.time.Duration
+import java.time.Instant
 import java.time.OffsetDateTime
 import java.util.concurrent.Callable
 import java.util.concurrent.TimeUnit
+import java.util.regex.Pattern
 
 import com.google.common.cache.Cache
 import com.google.common.cache.CacheBuilder
 import com.google.common.util.concurrent.UncheckedExecutionException
 import com.google.gson.Gson
+import com.google.gson.GsonBuilder
 import com.google.gson.reflect.TypeToken
 import groovy.json.JsonOutput
 import groovy.transform.CompileStatic
 import groovy.transform.Memoized
+import io.seqera.wave.plugin.adapter.InstantAdapter
 import io.seqera.wave.plugin.config.TowerConfig
 import io.seqera.wave.plugin.config.WaveConfig
 import io.seqera.wave.plugin.exception.BadResponseException
@@ -59,6 +63,8 @@ class WaveClient {
 
     private static Logger log = LoggerFactory.getLogger(WaveClient)
 
+    private static final Pattern CONTAINER_PATH = ~/(\S+)\/wt\/([a-z0-9]+)\/\S+/
+
     private static final List<String> DEFAULT_CONDA_CHANNELS = ['conda-forge','defaults']
 
     private static final String DEFAULT_SPACK_ARCH = 'x86_64'
@@ -89,6 +95,8 @@ class WaveClient {
 
     private List<String> condaChannels
 
+    final private String waveRegistry
+
     WaveClient(Session session) {
         this.session = session
         this.config = new WaveConfig(session.config.wave as Map ?: Collections.emptyMap(), SysEnv.get())
@@ -98,6 +106,7 @@ class WaveClient {
         this.condaChannels = session.getCondaConfig()?.getChannels() ?: DEFAULT_CONDA_CHANNELS
         log.debug "Wave server endpoint: ${endpoint}"
         this.packer = new Packer()
+        this.waveRegistry = new URI(endpoint).getAuthority()
         // create cache
         cache = CacheBuilder<String, SubmitContainerTokenResponse>
             .newBuilder()
@@ -594,4 +603,61 @@ class WaveClient {
         return null
     }
 
+    String resolveSourceContainer(String container) {
+        final token = getWaveToken(container)
+        if( !token )
+            return container
+        final resp = fetchContainerInfo(token)
+        final describe = jsonToDescribeContainerResponse(resp)
+        return describe.source.digest==describe.wave.digest
+                ? digestImage(describe.source)      // when the digest are equals, return the source because it's a stable name
+                : digestImage(describe.wave)        // otherwise returns the wave container name
+    }
+
+    protected String digestImage(DescribeContainerResponse.ContainerInfo info) {
+        if( !info.digest )
+            return info.image
+        final p = info.image.lastIndexOf(':')
+        return p!=-1
+            ? info.image.substring(0,p) + '@' + info.digest
+            : info.image + '@' + info.digest
+    }
+
+    protected String getWaveToken(String name) {
+        if( !name )
+            return null
+        final matcher = CONTAINER_PATH.matcher(name)
+        if( !matcher.find() )
+            return null
+        return matcher.group(1)==waveRegistry
+                ? matcher.group(2)
+                : null
+    }
+
+    @Memoized
+    protected String fetchContainerInfo(String token) {
+        final uri = new URI("$endpoint/container-token/$token")
+        log.trace "Wave request container info: $uri"
+        final req = HttpRequest.newBuilder()
+                .uri(uri)
+                .headers('Content-Type','application/json')
+                .GET()
+                .build()
+
+        final resp = httpClient.send(req, HttpResponse.BodyHandlers.ofString())
+        final code = resp.statusCode()
+        if( code>=200 && code<400 ) {
+            log.debug "Wave container config info: [$code] ${resp.body()}"
+            return resp.body()
+        }
+        throw new BadResponseException("Unexpected response for \'$uri\': [${resp.statusCode()}] ${resp.body()}")
+    }
+
+    protected DescribeContainerResponse jsonToDescribeContainerResponse(String json) {
+        final gson = new GsonBuilder()
+                .registerTypeAdapter(Instant.class, new InstantAdapter())
+                .create();
+        final type = new TypeToken<DescribeContainerResponse>(){}.getType()
+        return gson.fromJson(json, type)
+    }
 }

plugins/nf-wave/src/main/io/seqera/wave/plugin/WaveFactory.groovy

@@ -43,6 +43,10 @@ class WaveFactory implements TraceObserverFactory {
             wave.bundleProjectResources = true
             session.disableRemoteBinDir = true
         }
-        return Collections.emptyList()
+
+        final observer = new WaveObserver(session)
+        return wave?.enabled && observer.reportOpts().enabled()
+                ? List.<TraceObserver>of(observer)
+                : List.<TraceObserver>of()
     }
 }

plugins/nf-wave/src/main/io/seqera/wave/plugin/WaveObserver.groovy

@@ -0,0 +1,80 @@
+/*
+ * Copyright 2020-2022, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.wave.plugin
+
+import java.util.concurrent.ConcurrentHashMap
+
+import groovy.transform.CompileStatic
+import io.seqera.wave.plugin.config.ReportOpts
+import nextflow.Session
+import nextflow.file.FileHelper
+import nextflow.processor.TaskHandler
+import nextflow.trace.TraceObserver
+import nextflow.trace.TraceRecord
+/**
+ *
+ * @author Paolo Di Tommaso <[email protected]>
+ */
+@CompileStatic
+class WaveObserver implements TraceObserver {
+
+    private WaveClient client
+
+    private ConcurrentHashMap<String,String> containers = new ConcurrentHashMap<>()
+
+    WaveObserver(Session session) {
+        this.client = new WaveClient(session)
+    }
+
+    protected void apply(TaskHandler handler) {
+        final process = handler.task.getProcessor().getName()
+        containers.computeIfAbsent(process, (String it) -> {
+            final container = handler.task.getContainer()
+            return client.resolveSourceContainer(container)
+        })
+    }
+
+    void onProcessComplete(TaskHandler handler, TraceRecord trace){
+        apply(handler)
+    }
+
+    void onProcessCached(TaskHandler handler, TraceRecord trace){
+        apply(handler)
+    }
+
+    @Override
+    void onFlowComplete() {
+        final result = renderContainersConfig(containers)
+        // save the report file
+        FileHelper
+                .asPath(reportOpts().file())
+                .text = result.toString()
+    }
+
+    protected String renderContainersConfig(Map<String,String> containers) {
+        final result = new StringBuilder()
+        for( Map.Entry<String,String> entry : containers ) {
+            result.append("process { withName: '${entry.key}' { container='$entry.value' }}\n")
+        }
+        return result.toString()
+    }
+
+    ReportOpts reportOpts() {
+        client.config().reportOpts()
+    }
+}

plugins/nf-wave/src/main/io/seqera/wave/plugin/adapter/InstantAdapter.groovy

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2020-2022, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.wave.plugin.adapter
+
+import java.lang.reflect.Type
+import java.time.Instant
+
+import com.google.gson.JsonDeserializationContext
+import com.google.gson.JsonDeserializer
+import com.google.gson.JsonElement
+import com.google.gson.JsonParseException
+import com.google.gson.JsonPrimitive
+import com.google.gson.JsonSerializationContext
+import com.google.gson.JsonSerializer
+
+/**
+ * Gson adapter for java Instant class
+ * 
+ * @author Paolo Di Tommaso <[email protected]>
+ */
+class InstantAdapter implements JsonSerializer<Instant>, JsonDeserializer<Instant> {
+
+    @Override
+    JsonElement serialize(Instant instant, Type type, JsonSerializationContext context) {
+        return new JsonPrimitive(instant.toString());
+    }
+
+    @Override
+    Instant deserialize(JsonElement json, Type type, JsonDeserializationContext context) throws JsonParseException {
+        String instantString = json.getAsString();
+        return Instant.parse(instantString);
+    }
+}

plugins/nf-wave/src/main/io/seqera/wave/plugin/config/ReportOpts.groovy

@@ -0,0 +1,50 @@
+/*
+ * Copyright 2020-2022, Seqera Labs
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package io.seqera.wave.plugin.config
+
+import groovy.transform.CompileStatic
+import nextflow.trace.TraceHelper
+
+/**
+ *
+ * @author Paolo Di Tommaso <[email protected]>
+ */
+@CompileStatic
+class ReportOpts {
+
+    final private Boolean enabled
+
+    final private String file
+
+    ReportOpts(Map opts) {
+        this.enabled = opts.enabled as Boolean
+        this.file = opts.file
+    }
+
+    boolean enabled() {
+        enabled || file != null
+    }
+
+    String file() {
+        if( file )
+            return file
+        return enabled
+            ? "containers-${TraceHelper.launchTimestampFmt()}.config"
+            : null
+    }
+}

plugins/nf-wave/src/main/io/seqera/wave/plugin/config/WaveConfig.groovy

@@ -40,6 +40,7 @@ class WaveConfig {
     final private Boolean bundleProjectResources
     final private String buildRepository
     final private String cacheRepository
+    final private ReportOpts reportOpts
 
     WaveConfig(Map opts, Map<String,String> env=System.getenv()) {
         this.enabled = opts.enabled
@@ -52,6 +53,7 @@ class WaveConfig {
         this.cacheRepository = opts.navigate('build.cacheRepository') as String
         this.strategy = parseStrategy(opts.strategy)
         this.bundleProjectResources = opts.bundleProjectResources
+        this.reportOpts = new ReportOpts(opts.report as Map ?: Map.of())
         if( !endpoint.startsWith('http://') && !endpoint.startsWith('https://') )
             throw new IllegalArgumentException("Endpoint URL should start with 'http:' or 'https:' protocol prefix - offending value: $endpoint")
     }
@@ -121,4 +123,6 @@ class WaveConfig {
     Duration tokensCacheMaxDuration() { 
         return tokensCacheMaxDuration 
     }
+
+    ReportOpts reportOpts() { reportOpts }
 }