{"action":"created","comment":{"url":"https://api.github.com/repos/amplab/keystone/pulls/comments/56698714","id":56698714,"diff_hunk":"@@ -5,96 +5,341 @@ import org.apache.spark.rdd.RDD\n /**\n  * Node-level optimization, such as selecting a Linear Solver\n  *\n- * @param sampleFraction The fraction of the RDD to use for operations\n+ * @param samplePerPartition The number of items per partition to look at when optimizing nodes\n  */\n-class NodeOptimizationRule(sampleFraction: Double = 0.01, seed: Long = 0) extends Rule {\n+class NodeOptimizationRule(samplePerPartition: Int = 3) extends Rule {\n+  case class OptimizationState(\n+    optimizedInstructions: Seq[Instruction],\n+    unoptimizedInstructionIndexToOptimizedIndex: Int => Int,\n+    registers: Map[Int, InstructionOutput],\n+    numPerPartitionPerNode: Map[Int, Map[Int, Int]]\n+  )\n+\n+  /**\n+   * Get the set of instructions that we need to check for optimizations at.\n+   *\n+   * This is any instruction that either applies an optimizable transformer,\n+   * or fits an optimizable estimator,\n+   * and is not supposed to be executed on test data.\n+   */\n+  private def getInstructionsToOptimize(instructions: Seq[Instruction]): Set[Int] = {\n+    instructions.zipWithIndex.collect {\n+      case (TransformerApplyNode(tIndex, _), index)\n+        if instructions(tIndex).isInstanceOf[OptimizableTransformer[_, _]] => index\n+\n+      case (EstimatorFitNode(estIndex, _), index)\n+        if instructions(estIndex).isInstanceOf[OptimizableEstimator[_, _]] => index\n+\n+      case (EstimatorFitNode(estIndex, _), index)\n+        if instructions(estIndex).isInstanceOf[OptimizableLabelEstimator[_, _, _]] => index\n+\n+    }.toSet -- WorkflowUtils.getChildren(Pipeline.SOURCE, instructions)\n+  }\n+\n+  /**\n+   * Execute the instruction at the given index, and output the new state of the optimization process\n+   *\n+   * @param optimizationState\n+   * @param instruction\n+   * @param index\n+   * @return\n+   */\n+  private def executeInstruction(optimizationState: OptimizationState, instruction: Instruction, index: Int)\n+  : OptimizationState = {\n+    // Get the dependencies from the registers and run the instruction on them (caching RDD outputs)\n+    val registers = optimizationState.registers\n+    val deps = instruction.getDependencies.map(registers)\n+    val instructionOutput = instruction.execute(deps) match {\n+      case RDDOutput(rdd) => RDDOutput(rdd.cache())\n+      case out => out\n+    }\n+    val newRegisters = registers + (index -> instructionOutput)\n+\n+    // If any of the dependencies is an RDD and has a numPerPartition, make sure to feed it forward\n+    val numPerPartitionPerNode = optimizationState.numPerPartitionPerNode\n+    val numPerPartition = instruction.getDependencies.collectFirst {\n+      case dep if numPerPartitionPerNode.contains(dep) => numPerPartitionPerNode(dep)\n+    }\n+    val newNumPerPartitionPerNode = if (numPerPartition.nonEmpty) {\n+      numPerPartitionPerNode + (index -> numPerPartition.get)\n+    } else {\n+      numPerPartitionPerNode\n+    }\n+\n+    // Create the new state\n+    optimizationState.copy(registers = newRegisters, numPerPartitionPerNode = newNumPerPartitionPerNode)\n+  }\n+\n+  /**\n+   * Optimizes the transformer being applied by the instruction at the given index.\n+   * Outputs a new optimization state.\n+   *\n+   * @param os The previous optimization state\n+   * @param transformerApplyNode\n+   * @param index\n+   * @return\n+   */\n+  private def optimizeTransformer(os: OptimizationState, transformerApplyNode: TransformerApplyNode, index: Int)\n+  : OptimizationState = {\n+    val tIndex = transformerApplyNode.transformer\n+    val inputIndices = transformerApplyNode.inputs\n+\n+    val numPerPartition = os.numPerPartitionPerNode(inputIndices.head)\n+    val inputs = inputIndices.map(os.registers).collect {\n+      case RDDOutput(rdd) => rdd\n+    }\n+\n+    // Optimize the transformer if possible\n+    os.registers(tIndex) match {\n+      case TransformerOutput(ot: OptimizableTransformer[a, b]) => {","path":"src/main/scala/workflow/NodeOptimizationRule.scala","position":92,"original_position":92,"commit_id":"a57a1a2d044ce978118246bb8b31c98e9986540f","original_commit_id":"a57a1a2d044ce978118246bb8b31c98e9986540f","user":{"login":"etrain","id":1326181,"avatar_url":"https://avatars.githubusercontent.com/u/1326181?v=3","gravatar_id":"","url":"https://api.github.com/users/etrain","html_url":"https://github.com/etrain","followers_url":"https://api.github.com/users/etrain/followers","following_url":"https://api.github.com/users/etrain/following{/other_user}","gists_url":"https://api.github.com/users/etrain/gists{/gist_id}","starred_url":"https://api.github.com/users/etrain/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/etrain/subscriptions","organizations_url":"https://api.github.com/users/etrain/orgs","repos_url":"https://api.github.com/users/etrain/repos","events_url":"https://api.github.com/users/etrain/events{/privacy}","received_events_url":"https://api.github.com/users/etrain/received_events","type":"User","site_admin":false},"body":"Is there a reason you're using lower-case type names here? Makes it a little surprising/hard to read.","created_at":"2016-03-18T18:06:48Z","updated_at":"2016-03-18T18:06:48Z","html_url":"https://github.com/amplab/keystone/pull/227#discussion_r56698714","pull_request_url":"https://api.github.com/repos/amplab/keystone/pulls/227","_links":{"self":{"href":"https://api.github.com/repos/amplab/keystone/pulls/comments/56698714"},"html":{"href":"https://github.com/amplab/keystone/pull/227#discussion_r56698714"},"pull_request":{"href":"https://api.github.com/repos/amplab/keystone/pulls/227"}}},"pull_request":{"url":"https://api.github.com/repos/amplab/keystone/pulls/227","id":58974971,"html_url":"https://github.com/amplab/keystone/pull/227","diff_url":"https://github.com/amplab/keystone/pull/227.diff","patch_url":"https://github.com/amplab/keystone/pull/227.patch","issue_url":"https://api.github.com/repos/amplab/keystone/issues/227","number":227,"state":"open","locked":false,"title":"Node-level optimization generalization","user":{"login":"tomerk","id":676357,"avatar_url":"https://avatars.githubusercontent.com/u/676357?v=3","gravatar_id":"","url":"https://api.github.com/users/tomerk","html_url":"https://github.com/tomerk","followers_url":"https://api.github.com/users/tomerk/followers","following_url":"https://api.github.com/users/tomerk/following{/other_user}","gists_url":"https://api.github.com/users/tomerk/gists{/gist_id}","starred_url":"https://api.github.com/users/tomerk/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/tomerk/subscriptions","organizations_url":"https://api.github.com/users/tomerk/orgs","repos_url":"https://api.github.com/users/tomerk/repos","events_url":"https://api.github.com/users/tomerk/events{/privacy}","received_events_url":"https://api.github.com/users/tomerk/received_events","type":"User","site_admin":false},"body":"Allows node-level optimization rules to construct pipelines instead of just returning a single node.\r\n\r\nThe code in it's current form is very convoluted but the unit tests pass. Suggestions for how to improve it would be appreciated.","created_at":"2016-02-11T00:03:30Z","updated_at":"2016-03-18T18:06:48Z","closed_at":null,"merged_at":null,"merge_commit_sha":"3bc759a83346eb64f47b8b14b16b75583b6a2571","assignee":null,"milestone":null,"commits_url":"https://api.github.com/repos/amplab/keystone/pulls/227/commits","review_comments_url":"https://api.github.com/repos/amplab/keystone/pulls/227/comments","review_comment_url":"https://api.github.com/repos/amplab/keystone/pulls/comments{/number}","comments_url":"https://api.github.com/repos/amplab/keystone/issues/227/comments","statuses_url":"https://api.github.com/repos/amplab/keystone/statuses/a57a1a2d044ce978118246bb8b31c98e9986540f","head":{"label":"tomerk:node-optimization-generalization","ref":"node-optimization-generalization","sha":"a57a1a2d044ce978118246bb8b31c98e9986540f","user":{"login":"tomerk","id":676357,"avatar_url":"https://avatars.githubusercontent.com/u/676357?v=3","gravatar_id":"","url":"https://api.github.com/users/tomerk","html_url":"https://github.com/tomerk","followers_url":"https://api.github.com/users/tomerk/followers","following_url":"https://api.github.com/users/tomerk/following{/other_user}","gists_url":"https://api.github.com/users/tomerk/gists{/gist_id}","starred_url":"https://api.github.com/users/tomerk/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/tomerk/subscriptions","organizations_url":"https://api.github.com/users/tomerk/orgs","repos_url":"https://api.github.com/users/tomerk/repos","events_url":"https://api.github.com/users/tomerk/events{/privacy}","received_events_url":"https://api.github.com/users/tomerk/received_events","type":"User","site_admin":false},"repo":{"id":36905160,"name":"keystone","full_name":"tomerk/keystone","owner":{"login":"tomerk","id":676357,"avatar_url":"https://avatars.githubusercontent.com/u/676357?v=3","gravatar_id":"","url":"https://api.github.com/users/tomerk","html_url":"https://github.com/tomerk","followers_url":"https://api.github.com/users/tomerk/followers","following_url":"https://api.github.com/users/tomerk/following{/other_user}","gists_url":"https://api.github.com/users/tomerk/gists{/gist_id}","starred_url":"https://api.github.com/users/tomerk/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/tomerk/subscriptions","organizations_url":"https://api.github.com/users/tomerk/orgs","repos_url":"https://api.github.com/users/tomerk/repos","events_url":"https://api.github.com/users/tomerk/events{/privacy}","received_events_url":"https://api.github.com/users/tomerk/received_events","type":"User","site_admin":false},"private":false,"html_url":"https://github.com/tomerk/keystone","description":"Simplifying robust end-to-end machine learning on Apache Spark.","fork":true,"url":"https://api.github.com/repos/tomerk/keystone","forks_url":"https://api.github.com/repos/tomerk/keystone/forks","keys_url":"https://api.github.com/repos/tomerk/keystone/keys{/key_id}","collaborators_url":"https://api.github.com/repos/tomerk/keystone/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/tomerk/keystone/teams","hooks_url":"https://api.github.com/repos/tomerk/keystone/hooks","issue_events_url":"https://api.github.com/repos/tomerk/keystone/issues/events{/number}","events_url":"https://api.github.com/repos/tomerk/keystone/events","assignees_url":"https://api.github.com/repos/tomerk/keystone/assignees{/user}","branches_url":"https://api.github.com/repos/tomerk/keystone/branches{/branch}","tags_url":"https://api.github.com/repos/tomerk/keystone/tags","blobs_url":"https://api.github.com/repos/tomerk/keystone/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/tomerk/keystone/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/tomerk/keystone/git/refs{/sha}","trees_url":"https://api.github.com/repos/tomerk/keystone/git/trees{/sha}","statuses_url":"https://api.github.com/repos/tomerk/keystone/statuses/{sha}","languages_url":"https://api.github.com/repos/tomerk/keystone/languages","stargazers_url":"https://api.github.com/repos/tomerk/keystone/stargazers","contributors_url":"https://api.github.com/repos/tomerk/keystone/contributors","subscribers_url":"https://api.github.com/repos/tomerk/keystone/subscribers","subscription_url":"https://api.github.com/repos/tomerk/keystone/subscription","commits_url":"https://api.github.com/repos/tomerk/keystone/commits{/sha}","git_commits_url":"https://api.github.com/repos/tomerk/keystone/git/commits{/sha}","comments_url":"https://api.github.com/repos/tomerk/keystone/comments{/number}","issue_comment_url":"https://api.github.com/repos/tomerk/keystone/issues/comments{/number}","contents_url":"https://api.github.com/repos/tomerk/keystone/contents/{+path}","compare_url":"https://api.github.com/repos/tomerk/keystone/compare/{base}...{head}","merges_url":"https://api.github.com/repos/tomerk/keystone/merges","archive_url":"https://api.github.com/repos/tomerk/keystone/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/tomerk/keystone/downloads","issues_url":"https://api.github.com/repos/tomerk/keystone/issues{/number}","pulls_url":"https://api.github.com/repos/tomerk/keystone/pulls{/number}","milestones_url":"https://api.github.com/repos/tomerk/keystone/milestones{/number}","notifications_url":"https://api.github.com/repos/tomerk/keystone/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/tomerk/keystone/labels{/name}","releases_url":"https://api.github.com/repos/tomerk/keystone/releases{/id}","deployments_url":"https://api.github.com/repos/tomerk/keystone/deployments","created_at":"2015-06-05T01:29:58Z","updated_at":"2016-01-21T20:52:37Z","pushed_at":"2016-03-18T05:18:51Z","git_url":"git://github.com/tomerk/keystone.git","ssh_url":"git@github.com:tomerk/keystone.git","clone_url":"https://github.com/tomerk/keystone.git","svn_url":"https://github.com/tomerk/keystone","homepage":"http://keystone-ml.org/","size":23222,"stargazers_count":0,"watchers_count":0,"language":"Scala","has_issues":false,"has_downloads":true,"has_wiki":true,"has_pages":true,"forks_count":1,"mirror_url":null,"open_issues_count":1,"forks":1,"open_issues":1,"watchers":0,"default_branch":"master"}},"base":{"label":"amplab:master","ref":"master","sha":"8f3fdb1bc6134d8033a9338c71ada065b2f1cab5","user":{"login":"amplab","id":1438454,"avatar_url":"https://avatars.githubusercontent.com/u/1438454?v=3","gravatar_id":"","url":"https://api.github.com/users/amplab","html_url":"https://github.com/amplab","followers_url":"https://api.github.com/users/amplab/followers","following_url":"https://api.github.com/users/amplab/following{/other_user}","gists_url":"https://api.github.com/users/amplab/gists{/gist_id}","starred_url":"https://api.github.com/users/amplab/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/amplab/subscriptions","organizations_url":"https://api.github.com/users/amplab/orgs","repos_url":"https://api.github.com/users/amplab/repos","events_url":"https://api.github.com/users/amplab/events{/privacy}","received_events_url":"https://api.github.com/users/amplab/received_events","type":"Organization","site_admin":false},"repo":{"id":33748726,"name":"keystone","full_name":"amplab/keystone","owner":{"login":"amplab","id":1438454,"avatar_url":"https://avatars.githubusercontent.com/u/1438454?v=3","gravatar_id":"","url":"https://api.github.com/users/amplab","html_url":"https://github.com/amplab","followers_url":"https://api.github.com/users/amplab/followers","following_url":"https://api.github.com/users/amplab/following{/other_user}","gists_url":"https://api.github.com/users/amplab/gists{/gist_id}","starred_url":"https://api.github.com/users/amplab/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/amplab/subscriptions","organizations_url":"https://api.github.com/users/amplab/orgs","repos_url":"https://api.github.com/users/amplab/repos","events_url":"https://api.github.com/users/amplab/events{/privacy}","received_events_url":"https://api.github.com/users/amplab/received_events","type":"Organization","site_admin":false},"private":false,"html_url":"https://github.com/amplab/keystone","description":"Simplifying robust end-to-end machine learning on Apache Spark.","fork":false,"url":"https://api.github.com/repos/amplab/keystone","forks_url":"https://api.github.com/repos/amplab/keystone/forks","keys_url":"https://api.github.com/repos/amplab/keystone/keys{/key_id}","collaborators_url":"https://api.github.com/repos/amplab/keystone/collaborators{/collaborator}","teams_url":"https://api.github.com/repos/amplab/keystone/teams","hooks_url":"https://api.github.com/repos/amplab/keystone/hooks","issue_events_url":"https://api.github.com/repos/amplab/keystone/issues/events{/number}","events_url":"https://api.github.com/repos/amplab/keystone/events","assignees_url":"https://api.github.com/repos/amplab/keystone/assignees{/user}","branches_url":"https://api.github.com/repos/amplab/keystone/branches{/branch}","tags_url":"https://api.github.com/repos/amplab/keystone/tags","blobs_url":"https://api.github.com/repos/amplab/keystone/git/blobs{/sha}","git_tags_url":"https://api.github.com/repos/amplab/keystone/git/tags{/sha}","git_refs_url":"https://api.github.com/repos/amplab/keystone/git/refs{/sha}","trees_url":"https://api.github.com/repos/amplab/keystone/git/trees{/sha}","statuses_url":"https://api.github.com/repos/amplab/keystone/statuses/{sha}","languages_url":"https://api.github.com/repos/amplab/keystone/languages","stargazers_url":"https://api.github.com/repos/amplab/keystone/stargazers","contributors_url":"https://api.github.com/repos/amplab/keystone/contributors","subscribers_url":"https://api.github.com/repos/amplab/keystone/subscribers","subscription_url":"https://api.github.com/repos/amplab/keystone/subscription","commits_url":"https://api.github.com/repos/amplab/keystone/commits{/sha}","git_commits_url":"https://api.github.com/repos/amplab/keystone/git/commits{/sha}","comments_url":"https://api.github.com/repos/amplab/keystone/comments{/number}","issue_comment_url":"https://api.github.com/repos/amplab/keystone/issues/comments{/number}","contents_url":"https://api.github.com/repos/amplab/keystone/contents/{+path}","compare_url":"https://api.github.com/repos/amplab/keystone/compare/{base}...{head}","merges_url":"https://api.github.com/repos/amplab/keystone/merges","archive_url":"https://api.github.com/repos/amplab/keystone/{archive_format}{/ref}","downloads_url":"https://api.github.com/repos/amplab/keystone/downloads","issues_url":"https://api.github.com/repos/amplab/keystone/issues{/number}","pulls_url":"https://api.github.com/repos/amplab/keystone/pulls{/number}","milestones_url":"https://api.github.com/repos/amplab/keystone/milestones{/number}","notifications_url":"https://api.github.com/repos/amplab/keystone/notifications{?since,all,participating}","labels_url":"https://api.github.com/repos/amplab/keystone/labels{/name}","releases_url":"https://api.github.com/repos/amplab/keystone/releases{/id}","deployments_url":"https://api.github.com/repos/amplab/keystone/deployments","created_at":"2015-04-10T20:16:11Z","updated_at":"2016-03-17T14:11:20Z","pushed_at":"2016-03-18T05:18:51Z","git_url":"git://github.com/amplab/keystone.git","ssh_url":"git@github.com:amplab/keystone.git","clone_url":"https://github.com/amplab/keystone.git","svn_url":"https://github.com/amplab/keystone","homepage":"http://keystone-ml.org/","size":17808,"stargazers_count":297,"watchers_count":297,"language":"Scala","has_issues":true,"has_downloads":true,"has_wiki":true,"has_pages":true,"forks_count":63,"mirror_url":null,"open_issues_count":47,"forks":63,"open_issues":47,"watchers":297,"default_branch":"master"}},"_links":{"self":{"href":"https://api.github.com/repos/amplab/keystone/pulls/227"},"html":{"href":"https://github.com/amplab/keystone/pull/227"},"issue":{"href":"https://api.github.com/repos/amplab/keystone/issues/227"},"comments":{"href":"https://api.github.com/repos/amplab/keystone/issues/227/comments"},"review_comments":{"href":"https://api.github.com/repos/amplab/keystone/pulls/227/comments"},"review_comment":{"href":"https://api.github.com/repos/amplab/keystone/pulls/comments{/number}"},"commits":{"href":"https://api.github.com/repos/amplab/keystone/pulls/227/commits"},"statuses":{"href":"https://api.github.com/repos/amplab/keystone/statuses/a57a1a2d044ce978118246bb8b31c98e9986540f"}}}}