Merge pull request #41 from deflaux/master

Add Genomics API counters for Dataflow UI display.

Author: deflaux

pom.xml

@@ -114,7 +114,7 @@
     <dependency>
       <groupId>com.google.apis</groupId>
       <artifactId>google-api-services-genomics</artifactId>
-      <version>v1beta2-rev9-1.19.0</version>
+      <version>v1beta2-rev25-1.19.1</version>
       <exclusions>
         <!-- Exclude an old version of guava which is being pulled
              in by a transitive dependency google-api-client 1.19.0 -->

src/main/java/com/google/cloud/genomics/dataflow/functions/JoinNonVariantSegmentsWithVariants.java

@@ -88,9 +88,9 @@ public static PCollection<Variant> joinVariantsTransform(
       PCollection<SearchVariantsRequest> input, GenomicsFactory.OfflineAuth auth, String fields) {
     for (String field : REQUIRED_FIELDS) {
       Preconditions
-          .checkState(
+          .checkArgument(
               fields.contains(field),
-              "Required field missing: %s Add this field to the list of Variants fields returned in the partial response.",
+              "Required field missing: '%s' Add this field to the list of Variants fields returned in the partial response.",
               field);
     }
     return joinVariants(input, auth, fields);

src/main/java/com/google/cloud/genomics/dataflow/pipelines/ConvertNonVariantSegmentsToVariants.java

@@ -142,7 +142,7 @@ private static PCollection<Read> getReads() throws IOException {
   private static PCollection<Read> getReadsFromAPI() {
     List<SearchReadsRequest> requests = getReadRequests(options);
     PCollection<SearchReadsRequest> readRequests =
-        DataflowWorkarounds.getPCollection(requests, p, options.getNumWorkers());
+        DataflowWorkarounds.getPCollection(requests, p);
     PCollection<Read> reads =
         readRequests.apply(
             ParDo.of(

src/main/java/com/google/cloud/genomics/dataflow/pipelines/CountReads.java

@@ -48,7 +48,7 @@
 
 public class IdentityByState {
   private static final String VARIANT_FIELDS =
-      "nextPageToken,variants(id,calls(genotype,callSetName))";
+      "nextPageToken,variants(start,calls(genotype,callSetName))";
 
   public static void main(String[] args) throws IOException, GeneralSecurityException,
       InstantiationException, IllegalAccessException {
@@ -64,7 +64,7 @@ public static void main(String[] args) throws IOException, GeneralSecurityExcept
     Pipeline p = Pipeline.create(options);
     DataflowWorkarounds.registerGenomicsCoders(p);
     PCollection<SearchVariantsRequest> input =
-        DataflowWorkarounds.getPCollection(requests, p, options.getNumWorkers());
+        DataflowWorkarounds.getPCollection(requests, p);
 
     PCollection<Variant> variants =
         options.getHasNonVariantSegments()

src/main/java/com/google/cloud/genomics/dataflow/pipelines/IdentityByState.java

@@ -44,7 +44,7 @@
  */
 public class TransmissionProbability {
   private static final String VARIANT_FIELDS
-      = "nextPageToken,variants(id,names,calls(info,callSetName))";
+      = "nextPageToken,variants(id,start,names,calls(info,callSetName))";
 
   public static void main(String[] args) throws IOException, GeneralSecurityException {
     GenomicsDatasetOptions options = PipelineOptionsFactory.fromArgs(args)
@@ -65,7 +65,7 @@ public static void main(String[] args) throws IOException, GeneralSecurityExcept
     //    - Groups Transmission sources by Variant,
     //    - Calculate transmission Probability for each variant
     //    - Print calculated values to a file.
-    DataflowWorkarounds.getPCollection(requests, p, options.getNumWorkers())
+    DataflowWorkarounds.getPCollection(requests, p)
         .apply(ParDo.named("VariantReader")
             .of(new VariantReader(auth, ShardBoundary.STRICT, VARIANT_FIELDS)))
         .apply(ParDo.named("ExtractFamilyVariantStatus")

src/main/java/com/google/cloud/genomics/dataflow/pipelines/TransmissionProbability.java

@@ -39,7 +39,7 @@
  */
 public class VariantSimilarity {
   private static final String VARIANT_FIELDS
-      = "nextPageToken,variants(id,calls(genotype,callSetName))";
+      = "nextPageToken,variants(start,calls(genotype,callSetName))";
 
   public static void main(String[] args) throws IOException, GeneralSecurityException {
     GenomicsDatasetOptions options = PipelineOptionsFactory.fromArgs(args)
@@ -54,7 +54,7 @@ public static void main(String[] args) throws IOException, GeneralSecurityExcept
     Pipeline p = Pipeline.create(options);
     DataflowWorkarounds.registerGenomicsCoders(p);
     DataflowWorkarounds
-        .getPCollection(requests, p, options.getNumWorkers())
+        .getPCollection(requests, p)
         .apply(
             ParDo.named("VariantReader").of(
                 new VariantReader(auth, ShardBoundary.STRICT, VARIANT_FIELDS)))

src/main/java/com/google/cloud/genomics/dataflow/pipelines/VariantSimilarity.java

@@ -17,7 +17,9 @@
 
 import com.google.api.client.json.GenericJson;
 import com.google.api.services.genomics.Genomics;
+import com.google.cloud.dataflow.sdk.transforms.Aggregator;
 import com.google.cloud.dataflow.sdk.transforms.DoFn;
+import com.google.cloud.dataflow.sdk.transforms.Sum;
 import com.google.cloud.genomics.utils.GenomicsFactory;
 
 import java.io.IOException;
@@ -31,18 +33,33 @@ public abstract class GenomicsApiReader<I extends GenericJson, O extends Generic
   // Used for access to the genomics API
   protected final GenomicsFactory.OfflineAuth auth;
   protected final String fields;
-
+  protected Aggregator<Integer> initializedRequestsCount;
+  protected Aggregator<Integer> unsuccessfulResponsesCount;
+  protected Aggregator<Integer> ioExceptionsCount;
+  protected Aggregator<Long> itemCount;
+  
   public GenomicsApiReader(GenomicsFactory.OfflineAuth auth, String fields) {
     this.auth = auth;
     this.fields = fields;
   }
 
+  @Override
+  public void startBundle(Context c) {
+    initializedRequestsCount = c.createAggregator("Genomics API Initialized Request Count", new Sum.SumIntegerFn());
+    unsuccessfulResponsesCount = c.createAggregator("Genomics API Unsuccessful Response Count", new Sum.SumIntegerFn());
+    ioExceptionsCount = c.createAggregator("Genomics API IOException Response Count", new Sum.SumIntegerFn());
+    itemCount = c.createAggregator("Genomics API Item Count", new Sum.SumLongFn());
+  }
+  
   @Override
   public void processElement(ProcessContext c) {
     try {
       GenomicsFactory factory = auth.getDefaultFactory();
       processApiCall(auth.getGenomics(factory), c, c.element());
 
+      initializedRequestsCount.addValue(factory.initializedRequestsCount());
+      unsuccessfulResponsesCount.addValue(factory.unsuccessfulResponsesCount());
+      ioExceptionsCount.addValue(factory.ioExceptionsCount());
       LOG.info("ApiReader processed " + factory.initializedRequestsCount() + " requests ("
           + factory.unsuccessfulResponsesCount() + " server errors and "
           + factory.ioExceptionsCount() + " IO exceptions)");

src/main/java/com/google/cloud/genomics/dataflow/readers/GenomicsApiReader.java

@@ -57,6 +57,7 @@ protected void processApiCall(Genomics genomics, ProcessContext c, SearchReadsRe
 
     for (Read read : Paginator.Reads.create(genomics, shardBoundary).search(request, fields)) {
       c.output(read);
+      itemCount.addValue(1L);
     }
   }
 }

src/main/java/com/google/cloud/genomics/dataflow/readers/ReadReader.java

@@ -31,7 +31,7 @@ public class VariantReader extends GenomicsApiReader<SearchVariantsRequest, Vari
 
   /**
    * Create a VariantReader using a auth and fields parameter. All fields not specified under 
-   * readFields will not be returned in the API response.
+   * variantFields will not be returned in the API response.
    * 
    * @param auth Auth class containing credentials.
    * @param variantFields Fields to return in responses.
@@ -52,12 +52,13 @@ public VariantReader(GenomicsFactory.OfflineAuth auth, ShardBoundary shardBounda
   @Override
   protected void processApiCall(Genomics genomics, ProcessContext c, SearchVariantsRequest request)
       throws IOException {
-    LOG.info("Starting Variants read loop");
+    LOG.info("Starting Variants read loop: " + request);
 
     int numberOfVariants = 0;
     for (Variant variant : Paginator.Variants.create(genomics, shardBoundary).search(request, fields)) {
       c.output(variant);
       ++numberOfVariants;
+      itemCount.addValue(1L);
     }
 
     LOG.info("Read " + numberOfVariants + " variants at: " + request.getReferenceName() + "-" + "["