You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@celeborn.apache.org by et...@apache.org on 2024/05/23 13:10:35 UTC
(celeborn) branch main updated: [CELEBORN-914] Support memory file storage
This is an automated email from the ASF dual-hosted git repository.
ethanfeng pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/celeborn.git
The following commit(s) were added to refs/heads/main by this push:
new 89d56c9bb [CELEBORN-914] Support memory file storage
89d56c9bb is described below
commit 89d56c9bbcef3249ff92d5e727ab1cbfd4a0057b
Author: mingji <fe...@alibaba-inc.com>
AuthorDate: Thu May 23 21:05:52 2024 +0800
[CELEBORN-914] Support memory file storage
### What changes were proposed in this pull request?
To support memory file storage.
### Why are the changes needed?
To improve shuffle performance for small shuffle files.
Design doc: https://docs.google.com/document/d/1SM-oOM0JHEIoRHTYhE9PYH60_1D3NMxDR50LZIM7uW0/edit?usp=sharing
### Does this PR introduce _any_ user-facing change?
No.
### How was this patch tested?
Pass GA and manually test on a cluster.
Closes #2300 from FMX/B914.
Authored-by: mingji <fe...@alibaba-inc.com>
Signed-off-by: mingji <fe...@alibaba-inc.com>
---
assets/grafana/celeborn-dashboard.json | 587 ++++++++++++++++----
.../celeborn/client/read/CelebornInputStream.java | 5 +-
.../celeborn/client/read/DfsPartitionReader.java | 3 +-
.../apache/celeborn/common/meta/DiskFileInfo.java | 57 +-
.../org/apache/celeborn/common/meta/FileInfo.java | 55 +-
.../celeborn/common/meta/MemoryFileInfo.java | 46 +-
.../celeborn/common/meta/ReduceFileMeta.java | 28 +-
.../buffer/ChunkBuffers.java} | 40 +-
.../buffer/FileChunkBuffers.java} | 37 +-
.../buffer/MemoryChunkBuffers.java} | 36 +-
.../celeborn/common/protocol/StorageInfo.java | 11 +-
.../common/util/ShuffleBlockInfoUtils.java | 78 ++-
.../org/apache/celeborn/common/CelebornConf.scala | 31 +-
.../common/network/CelebornRackResolverSuite.scala | 1 -
.../celeborn/common/util/PbSerDeUtilsTest.scala | 4 +-
docs/configuration/worker.md | 3 +-
.../tests/spark/memory/MemorySkewJoinSuite.scala | 113 ++++
.../tests/spark/memory/MemorySparkTestBase.scala | 62 +++
.../memory/MemoryStorageReusedExchangeSuite.scala | 83 +++
.../deploy/worker/memory/MemoryManager.java | 144 +++--
.../deploy/worker/storage/ChunkStreamManager.java | 31 +-
.../worker/storage/MapPartitionDataWriter.java | 30 +-
.../deploy/worker/storage/PartitionDataWriter.java | 380 +++++++++----
.../worker/storage/PartitionDataWriterContext.java | 99 ++++
.../worker/storage/PartitionFilesSorter.java | 238 +++++---
.../worker/storage/ReducePartitionDataWriter.java | 122 ++---
.../service/deploy/worker/Controller.scala | 42 +-
.../service/deploy/worker/FetchHandler.scala | 96 ++--
.../service/deploy/worker/PushDataHandler.scala | 52 +-
.../celeborn/service/deploy/worker/Worker.scala | 19 +-
.../service/deploy/worker/WorkerSource.scala | 5 +
.../service/deploy/worker/storage/FlushTask.scala | 9 +-
.../service/deploy/worker/storage/Flusher.scala | 35 +-
.../deploy/worker/storage/StorageManager.scala | 282 +++++++---
.../service/deploy/worker/FetchHandlerSuiteJ.java | 14 +-
.../worker/storage/ChunkStreamManagerSuiteJ.java | 10 +-
.../worker/storage/CreditStreamManagerSuiteJ.java | 6 +-
.../storage/PartitionDataWriterSuiteUtils.java | 113 +++-
.../SSLReducePartitionDataWriterSuiteJ.java | 7 +-
.../DiskMapPartitionDataWriterSuiteJ.java} | 44 +-
.../DiskPartitionFilesSorterSuiteJ.java} | 29 +-
.../DiskReducePartitionDataWriterSuiteJ.java} | 367 +++++++------
.../memory/MemoryPartitionFilesSorterSuiteJ.java | 159 ++++++
.../MemoryReducePartitionDataWriterSuiteJ.java} | 600 +++++++++++----------
.../service/deploy/memory/MemoryManagerSuite.scala | 4 +-
45 files changed, 2954 insertions(+), 1263 deletions(-)
diff --git a/assets/grafana/celeborn-dashboard.json b/assets/grafana/celeborn-dashboard.json
index f7b4de9cf..f914886cd 100644
--- a/assets/grafana/celeborn-dashboard.json
+++ b/assets/grafana/celeborn-dashboard.json
@@ -130,7 +130,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 1
+ "y": 33
},
"id": 2,
"options": {
@@ -220,7 +220,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 1
+ "y": 33
},
"id": 94,
"options": {
@@ -595,7 +595,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 20
+ "y": 52
},
"id": 121,
"options": {
@@ -688,7 +688,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 20
+ "y": 52
},
"id": 120,
"options": {
@@ -781,7 +781,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 28
+ "y": 60
},
"id": 122,
"options": {
@@ -873,7 +873,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 28
+ "y": 60
},
"id": 124,
"options": {
@@ -965,7 +965,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 36
+ "y": 68
},
"id": 100,
"options": {
@@ -1055,7 +1055,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 36
+ "y": 68
},
"id": 102,
"options": {
@@ -1145,7 +1145,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 44
+ "y": 76
},
"id": 36,
"options": {
@@ -1236,7 +1236,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 44
+ "y": 76
},
"id": 117,
"options": {
@@ -1432,7 +1432,7 @@
"h": 10,
"w": 24,
"x": 0,
- "y": 61
+ "y": 85
},
"id": 84,
"options": {
@@ -1612,7 +1612,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 69
+ "y": 95
},
"id": 60,
"options": {
@@ -1702,7 +1702,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 69
+ "y": 95
},
"id": 62,
"options": {
@@ -1791,7 +1791,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 77
+ "y": 104
},
"id": 90,
"options": {
@@ -1881,7 +1881,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 77
+ "y": 104
},
"id": 92,
"options": {
@@ -1972,7 +1972,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 85
+ "y": 112
},
"id": 182,
"options": {
@@ -2063,7 +2063,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 85
+ "y": 112
},
"id": 184,
"options": {
@@ -2155,7 +2155,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 93
+ "y": 120
},
"id": 181,
"options": {
@@ -2178,8 +2178,8 @@
},
"editorMode": "builder",
"expr": "metrics_ActiveShuffleSize_Value{role=\"Worker\"}",
- "legendFormat": "${baseLegend}",
"instant": false,
+ "legendFormat": "${baseLegend}",
"range": true,
"refId": "A"
}
@@ -2248,7 +2248,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 93
+ "y": 120
},
"id": 183,
"options": {
@@ -2271,8 +2271,8 @@
},
"editorMode": "builder",
"expr": "metrics_ActiveShuffleFileCount_Value{role=\"Worker\"}",
- "legendFormat": "${baseLegend}",
"instant": false,
+ "legendFormat": "${baseLegend}",
"range": true,
"refId": "A"
}
@@ -2340,7 +2340,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 101
+ "y": 128
},
"id": 179,
"options": {
@@ -2445,7 +2445,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 110
+ "y": 137
},
"id": 68,
"options": {
@@ -2535,7 +2535,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 110
+ "y": 137
},
"id": 70,
"options": {
@@ -2625,7 +2625,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 118
+ "y": 146
},
"id": 72,
"options": {
@@ -2715,7 +2715,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 118
+ "y": 146
},
"id": 74,
"options": {
@@ -2804,7 +2804,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 126
+ "y": 155
},
"id": 83,
"options": {
@@ -2895,7 +2895,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 126
+ "y": 155
},
"id": 76,
"options": {
@@ -2986,7 +2986,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 134
+ "y": 163
},
"id": 128,
"options": {
@@ -3077,7 +3077,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 134
+ "y": 163
},
"id": 129,
"options": {
@@ -3168,7 +3168,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 142
+ "y": 171
},
"id": 130,
"options": {
@@ -3259,7 +3259,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 142
+ "y": 171
},
"id": 132,
"options": {
@@ -3350,7 +3350,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 150
+ "y": 179
},
"id": 131,
"options": {
@@ -3441,7 +3441,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 150
+ "y": 179
},
"id": 133,
"options": {
@@ -3638,7 +3638,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 167
+ "y": 188
},
"id": 66,
"options": {
@@ -3728,7 +3728,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 167
+ "y": 188
},
"id": 96,
"options": {
@@ -3818,7 +3818,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 175
+ "y": 196
},
"id": 17,
"options": {
@@ -3908,7 +3908,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 175
+ "y": 196
},
"id": 18,
"options": {
@@ -3997,7 +3997,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 183
+ "y": 204
},
"id": 81,
"options": {
@@ -4088,7 +4088,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 183
+ "y": 204
},
"id": 77,
"options": {
@@ -4179,7 +4179,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 191
+ "y": 212
},
"id": 82,
"options": {
@@ -4270,7 +4270,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 191
+ "y": 212
},
"id": 75,
"options": {
@@ -4361,7 +4361,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 199
+ "y": 220
},
"id": 73,
"options": {
@@ -4467,7 +4467,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 208
+ "y": 155
},
"id": 78,
"options": {
@@ -4557,7 +4557,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 208
+ "y": 155
},
"id": 80,
"options": {
@@ -4647,7 +4647,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 216
+ "y": 163
},
"id": 4,
"options": {
@@ -4737,7 +4737,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 216
+ "y": 163
},
"id": 6,
"options": {
@@ -4827,7 +4827,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 224
+ "y": 172
},
"id": 56,
"options": {
@@ -4917,7 +4917,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 224
+ "y": 172
},
"id": 58,
"options": {
@@ -5005,7 +5005,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5021,7 +5022,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 233
+ "y": 7
},
"id": 19,
"options": {
@@ -5096,7 +5097,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5104,7 +5106,192 @@
}
]
},
- "unit": "bytes"
+ "unit": "percentunit"
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 12,
+ "y": 7
+ },
+ "id": 185,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "builder",
+ "expr": "metrics_DirectMemoryUsageRatio_Value",
+ "instant": false,
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "DirectMemoryUsageRatio",
+ "type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green",
+ "value": null
+ },
+ {
+ "color": "red",
+ "value": 80
+ }
+ ]
+ },
+ "unit": "decbytes"
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 0,
+ "y": 15
+ },
+ "id": 186,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "builder",
+ "expr": "metrics_MemoryFileStorageSize_Value",
+ "instant": false,
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "MemoryFileStorageSize",
+ "type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green",
+ "value": null
+ },
+ {
+ "color": "red",
+ "value": 80
+ }
+ ]
+ }
},
"overrides": []
},
@@ -5112,7 +5299,100 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 233
+ "y": 15
+ },
+ "id": 188,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "builder",
+ "expr": "metrics_MemoryStorageFileCount_Value",
+ "instant": false,
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "MemoryStorageFileCount",
+ "type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green",
+ "value": null
+ },
+ {
+ "color": "red",
+ "value": 80
+ }
+ ]
+ },
+ "unit": "bytes"
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 0,
+ "y": 23
},
"id": 20,
"options": {
@@ -5187,7 +5467,100 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
+ },
+ {
+ "color": "red",
+ "value": 80
+ }
+ ]
+ }
+ },
+ "overrides": []
+ },
+ "gridPos": {
+ "h": 8,
+ "w": 12,
+ "x": 12,
+ "y": 23
+ },
+ "id": 187,
+ "options": {
+ "legend": {
+ "calcs": [],
+ "displayMode": "list",
+ "placement": "bottom",
+ "showLegend": true
+ },
+ "tooltip": {
+ "mode": "single",
+ "sort": "none"
+ }
+ },
+ "targets": [
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "editorMode": "builder",
+ "expr": "metrics_EvictedFileCount_Value",
+ "instant": false,
+ "range": true,
+ "refId": "A"
+ }
+ ],
+ "title": "EvictedFileCount",
+ "type": "timeseries"
+ },
+ {
+ "datasource": {
+ "type": "prometheus",
+ "uid": "${DS_PROMETHEUS}"
+ },
+ "fieldConfig": {
+ "defaults": {
+ "color": {
+ "mode": "palette-classic"
+ },
+ "custom": {
+ "axisCenteredZero": false,
+ "axisColorMode": "text",
+ "axisLabel": "",
+ "axisPlacement": "auto",
+ "barAlignment": 0,
+ "drawStyle": "line",
+ "fillOpacity": 0,
+ "gradientMode": "none",
+ "hideFrom": {
+ "legend": false,
+ "tooltip": false,
+ "viz": false
+ },
+ "lineInterpolation": "linear",
+ "lineWidth": 1,
+ "pointSize": 5,
+ "scaleDistribution": {
+ "type": "linear"
+ },
+ "showPoints": "auto",
+ "spanNulls": false,
+ "stacking": {
+ "group": "A",
+ "mode": "none"
+ },
+ "thresholdsStyle": {
+ "mode": "off"
+ }
+ },
+ "mappings": [],
+ "thresholds": {
+ "mode": "absolute",
+ "steps": [
+ {
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5203,7 +5576,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 241
+ "y": 31
},
"id": 165,
"options": {
@@ -5279,7 +5652,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5295,7 +5669,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 241
+ "y": 31
},
"id": 166,
"options": {
@@ -5371,7 +5745,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5387,7 +5762,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 249
+ "y": 39
},
"id": 167,
"options": {
@@ -5463,7 +5838,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5479,7 +5855,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 249
+ "y": 39
},
"id": 168,
"options": {
@@ -5555,7 +5931,8 @@
"mode": "absolute",
"steps": [
{
- "color": "green"
+ "color": "green",
+ "value": null
},
{
"color": "red",
@@ -5571,7 +5948,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 257
+ "y": 47
},
"id": 169,
"options": {
@@ -5663,7 +6040,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 257
+ "y": 47
},
"id": 170,
"options": {
@@ -5754,7 +6131,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 265
+ "y": 55
},
"id": 108,
"options": {
@@ -5846,7 +6223,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 265
+ "y": 55
},
"id": 104,
"options": {
@@ -5937,7 +6314,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 273
+ "y": 63
},
"id": 106,
"options": {
@@ -6043,7 +6420,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 282
+ "y": 230
},
"id": 44,
"options": {
@@ -6133,7 +6510,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 282
+ "y": 230
},
"id": 46,
"options": {
@@ -6222,7 +6599,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 290
+ "y": 238
},
"id": 48,
"options": {
@@ -6312,7 +6689,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 290
+ "y": 238
},
"id": 180,
"options": {
@@ -6403,7 +6780,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 298
+ "y": 246
},
"id": 88,
"options": {
@@ -6493,7 +6870,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 298
+ "y": 246
},
"id": 135,
"options": {
@@ -6600,7 +6977,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 307
+ "y": 255
},
"id": 159,
"options": {
@@ -6693,7 +7070,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 307
+ "y": 255
},
"id": 160,
"options": {
@@ -6786,7 +7163,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 315
+ "y": 263
},
"id": 161,
"options": {
@@ -6892,7 +7269,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 324
+ "y": 272
},
"id": 139,
"options": {
@@ -6984,7 +7361,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 324
+ "y": 272
},
"id": 141,
"options": {
@@ -7076,7 +7453,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 332
+ "y": 281
},
"id": 142,
"options": {
@@ -7168,7 +7545,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 332
+ "y": 281
},
"id": 143,
"options": {
@@ -7260,7 +7637,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 340
+ "y": 290
},
"id": 144,
"options": {
@@ -7352,7 +7729,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 340
+ "y": 290
},
"id": 145,
"options": {
@@ -7444,7 +7821,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 348
+ "y": 299
},
"id": 146,
"options": {
@@ -7536,7 +7913,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 348
+ "y": 299
},
"id": 147,
"options": {
@@ -7628,7 +8005,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 356
+ "y": 308
},
"id": 148,
"options": {
@@ -7720,7 +8097,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 356
+ "y": 308
},
"id": 149,
"options": {
@@ -7812,7 +8189,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 364
+ "y": 317
},
"id": 150,
"options": {
@@ -7904,7 +8281,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 364
+ "y": 317
},
"id": 151,
"options": {
@@ -7995,7 +8372,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 372
+ "y": 326
},
"id": 153,
"options": {
@@ -8086,7 +8463,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 372
+ "y": 326
},
"id": 154,
"options": {
@@ -8177,7 +8554,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 380
+ "y": 334
},
"id": 155,
"options": {
@@ -8282,7 +8659,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 389
+ "y": 343
},
"id": 112,
"options": {
@@ -8373,7 +8750,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 389
+ "y": 343
},
"id": 116,
"options": {
@@ -8479,7 +8856,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 398
+ "y": 352
},
"id": 125,
"options": {
@@ -8572,7 +8949,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 398
+ "y": 352
},
"id": 126,
"options": {
@@ -8665,7 +9042,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 406
+ "y": 360
},
"id": 163,
"options": {
@@ -8758,7 +9135,7 @@
"h": 8,
"w": 12,
"x": 12,
- "y": 406
+ "y": 360
},
"id": 162,
"options": {
@@ -8851,7 +9228,7 @@
"h": 8,
"w": 12,
"x": 0,
- "y": 414
+ "y": 368
},
"id": 127,
"options": {
@@ -8957,7 +9334,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 423
+ "y": 377
},
"id": 174,
"options": {
@@ -9050,7 +9427,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 423
+ "y": 377
},
"id": 176,
"options": {
@@ -9142,7 +9519,7 @@
"h": 9,
"w": 12,
"x": 0,
- "y": 431
+ "y": 386
},
"id": 175,
"options": {
@@ -9235,7 +9612,7 @@
"h": 9,
"w": 12,
"x": 12,
- "y": 431
+ "y": 386
},
"id": 177,
"options": {
@@ -9305,7 +9682,7 @@
]
},
"time": {
- "from": "now-5m",
+ "from": "now-30m",
"to": "now"
},
"timepicker": {},
@@ -9314,4 +9691,4 @@
"uid": "U_qgru_7z",
"version": 1,
"weekStart": ""
-}
+}
\ No newline at end of file
diff --git a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
index 99b131a48..6b4337929 100644
--- a/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
+++ b/client/src/main/java/org/apache/celeborn/client/read/CelebornInputStream.java
@@ -444,7 +444,10 @@ public abstract class CelebornInputStream extends InputStream {
switch (storageInfo.getType()) {
case HDD:
case SSD:
- if (enabledReadLocalShuffle && location.getHost().equals(localHostAddress)) {
+ case MEMORY:
+ if (enabledReadLocalShuffle
+ && location.getHost().equals(localHostAddress)
+ && storageInfo.getType() != StorageInfo.Type.MEMORY) {
logger.debug("Read local shuffle file {}", localHostAddress);
containLocalRead = true;
return new LocalPartitionReader(
diff --git a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
index 38306f6c2..712bd82e4 100644
--- a/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
+++ b/client/src/main/java/org/apache/celeborn/client/read/DfsPartitionReader.java
@@ -167,7 +167,8 @@ public class DfsPartitionReader implements PartitionReader {
startMapIndex,
endMapIndex,
shuffleChunkSize,
- ShuffleBlockInfoUtils.parseShuffleBlockInfosFromByteBuffer(indexBuffer)));
+ ShuffleBlockInfoUtils.parseShuffleBlockInfosFromByteBuffer(indexBuffer),
+ false));
}
return offsets;
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java b/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
index ab10df2ec..ce99192c6 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/DiskFileInfo.java
@@ -20,14 +20,14 @@ package org.apache.celeborn.common.meta;
import java.io.File;
import java.util.ArrayList;
import java.util.Arrays;
-import java.util.Set;
-import java.util.concurrent.ConcurrentHashMap;
+import com.google.common.annotations.VisibleForTesting;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.protocol.StorageInfo;
import org.apache.celeborn.common.util.Utils;
@@ -36,11 +36,9 @@ import org.apache.celeborn.common.util.Utils;
* This class will describe shuffle files located on local disks, HDFS and other remote storage.
* */
public class DiskFileInfo extends FileInfo {
- private static Logger logger = LoggerFactory.getLogger(DiskFileInfo.class);
- private final Set<Long> streams = ConcurrentHashMap.newKeySet();
- private String filePath;
- private StorageInfo.Type storageType;
- private volatile long bytesFlushed;
+ private static final Logger logger = LoggerFactory.getLogger(DiskFileInfo.class);
+ private final String filePath;
+ private final StorageInfo.Type storageType;
public DiskFileInfo(
UserIdentifier userIdentifier,
@@ -66,11 +64,12 @@ public class DiskFileInfo extends FileInfo {
this.bytesFlushed = bytesFlushed;
}
- public DiskFileInfo(File file, UserIdentifier userIdentifier) {
+ @VisibleForTesting
+ public DiskFileInfo(File file, UserIdentifier userIdentifier, CelebornConf conf) {
this(
userIdentifier,
true,
- new ReduceFileMeta(new ArrayList(Arrays.asList(0L))),
+ new ReduceFileMeta(new ArrayList<>(Arrays.asList(0L)), conf.shuffleChunkSize()),
file.getAbsolutePath(),
StorageInfo.Type.HDD);
}
@@ -81,42 +80,6 @@ public class DiskFileInfo extends FileInfo {
this.storageType = StorageInfo.Type.HDD;
}
- public boolean addStream(long streamId) {
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) fileMeta;
- synchronized (reduceFileMeta.getSorted()) {
- if (reduceFileMeta.getSorted().get()) {
- return false;
- } else {
- streams.add(streamId);
- return true;
- }
- }
- }
-
- public void closeStream(long streamId) {
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) fileMeta;
- synchronized (reduceFileMeta.getSorted()) {
- streams.remove(streamId);
- }
- }
-
- public boolean isStreamsEmpty() {
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) fileMeta;
- synchronized (reduceFileMeta.getSorted()) {
- return streams.isEmpty();
- }
- }
-
- @Override
- public long getFileLength() {
- return bytesFlushed;
- }
-
- public long updateBytesFlushed(long numBytes) {
- bytesFlushed += numBytes;
- return bytesFlushed;
- }
-
public File getFile() {
return new File(filePath);
}
@@ -185,10 +148,6 @@ public class DiskFileInfo extends FileInfo {
((MapFileMeta) fileMeta).setMountPoint(mountPoint);
}
- public long getBytesFlushed() {
- return bytesFlushed;
- }
-
public boolean isHdfs() {
return Utils.isHdfsPath(filePath);
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java b/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
index 34ba7cb13..c21c19847 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/FileInfo.java
@@ -17,6 +17,9 @@
package org.apache.celeborn.common.meta;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
import org.apache.celeborn.common.identity.UserIdentifier;
public abstract class FileInfo {
@@ -25,22 +28,40 @@ public abstract class FileInfo {
// now it's just used for mappartition to compatible with old client which can't support split
private boolean partitionSplitEnabled;
protected FileMeta fileMeta;
+ protected final Set<Long> streams = ConcurrentHashMap.newKeySet();
+ protected volatile long bytesFlushed;
+ private boolean isReduceFileMeta;
public FileInfo(UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
this.userIdentifier = userIdentifier;
this.partitionSplitEnabled = partitionSplitEnabled;
this.fileMeta = fileMeta;
+ this.isReduceFileMeta = fileMeta instanceof ReduceFileMeta;
}
public void replaceFileMeta(FileMeta meta) {
this.fileMeta = meta;
+ this.isReduceFileMeta = meta instanceof ReduceFileMeta;
}
public FileMeta getFileMeta() {
return fileMeta;
}
- public abstract long getFileLength();
+ public ReduceFileMeta getReduceFileMeta() {
+ return (ReduceFileMeta) fileMeta;
+ }
+
+ public long getFileLength() {
+ return bytesFlushed;
+ }
+
+ public void updateBytesFlushed(long bytes) {
+ bytesFlushed += bytes;
+ if (isReduceFileMeta) {
+ getReduceFileMeta().updateChunkOffset(bytesFlushed, false);
+ }
+ }
public UserIdentifier getUserIdentifier() {
return userIdentifier;
@@ -53,4 +74,36 @@ public abstract class FileInfo {
public void setPartitionSplitEnabled(boolean partitionSplitEnabled) {
this.partitionSplitEnabled = partitionSplitEnabled;
}
+
+ public boolean addStream(long streamId) {
+ if (!isReduceFileMeta) {
+ throw new IllegalStateException("In addStream, filemeta cannot be MapFileMeta");
+ }
+ synchronized (getReduceFileMeta().getSorted()) {
+ if (getReduceFileMeta().getSorted().get()) {
+ return false;
+ } else {
+ streams.add(streamId);
+ return true;
+ }
+ }
+ }
+
+ public void closeStream(long streamId) {
+ if (!isReduceFileMeta) {
+ throw new IllegalStateException("In closeStream, filemeta cannot be MapFileMeta");
+ }
+ synchronized (getReduceFileMeta().getSorted()) {
+ streams.remove(streamId);
+ }
+ }
+
+ public boolean isStreamsEmpty() {
+ if (!isReduceFileMeta) {
+ throw new IllegalStateException("In isStreamsEmpty, filemeta cannot be MapFileMeta");
+ }
+ synchronized (getReduceFileMeta().getSorted()) {
+ return streams.isEmpty();
+ }
+ }
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java b/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
index e58263bdf..f75826b84 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
@@ -17,19 +17,36 @@
package org.apache.celeborn.common.meta;
+import java.util.*;
+
import io.netty.buffer.CompositeByteBuf;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.identity.UserIdentifier;
+import org.apache.celeborn.common.util.ShuffleBlockInfoUtils.ShuffleBlockInfo;
public class MemoryFileInfo extends FileInfo {
+ Logger logger = LoggerFactory.getLogger(MemoryFileInfo.class);
private CompositeByteBuf buffer;
- private long length;
+ private CompositeByteBuf sortedBuffer;
+ private Map<Integer, List<ShuffleBlockInfo>> sortedIndexes;
public MemoryFileInfo(
UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
super(userIdentifier, partitionSplitEnabled, fileMeta);
}
+ // This constructor is only used in partition sorter for temp new memory file
+ public MemoryFileInfo(
+ UserIdentifier userIdentifier,
+ boolean partitionSplitEnabled,
+ FileMeta fileMeta,
+ CompositeByteBuf buffer) {
+ super(userIdentifier, partitionSplitEnabled, fileMeta);
+ this.buffer = buffer;
+ }
+
public CompositeByteBuf getBuffer() {
return buffer;
}
@@ -38,12 +55,29 @@ public class MemoryFileInfo extends FileInfo {
this.buffer = buffer;
}
- public void setBufferSize(int bufferSize) {
- this.length = bufferSize;
+ public CompositeByteBuf getSortedBuffer() {
+ return sortedBuffer;
+ }
+
+ public void setSortedBuffer(CompositeByteBuf sortedBuffer) {
+ this.sortedBuffer = sortedBuffer;
+ }
+
+ public Map<Integer, List<ShuffleBlockInfo>> getSortedIndexes() {
+ return sortedIndexes;
+ }
+
+ public void setSortedIndexes(Map<Integer, List<ShuffleBlockInfo>> sortedIndexes) {
+ this.sortedIndexes = sortedIndexes;
}
- @Override
- public long getFileLength() {
- return length;
+ public int releaseMemoryBuffers() {
+ int bufferSize = 0;
+ if (buffer != null) {
+ bufferSize = buffer.writerIndex();
+ buffer.release();
+ }
+ logger.info("Memory File Info {} expire, removed {}", this, bufferSize);
+ return bufferSize;
}
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/ReduceFileMeta.java b/common/src/main/java/org/apache/celeborn/common/meta/ReduceFileMeta.java
index ab782cd9c..abc449881 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/ReduceFileMeta.java
+++ b/common/src/main/java/org/apache/celeborn/common/meta/ReduceFileMeta.java
@@ -22,12 +22,25 @@ import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;
public class ReduceFileMeta implements FileMeta {
- private final List<Long> chunkOffsets;
private final AtomicBoolean sorted = new AtomicBoolean(false);
+ private final List<Long> chunkOffsets;
+ private long chunkSize;
+ private long nextBoundary;
- public ReduceFileMeta() {
+ public ReduceFileMeta(long chunkSize) {
this.chunkOffsets = new ArrayList<>();
chunkOffsets.add(0L);
+ this.chunkSize = chunkSize;
+ nextBoundary = chunkSize;
+ }
+
+ public ReduceFileMeta(List<Long> chunkOffsets, long chunkSize) {
+ this.chunkOffsets = chunkOffsets;
+ nextBoundary = chunkSize;
+ if (!chunkOffsets.isEmpty()) {
+ nextBoundary += chunkOffsets.get(chunkOffsets.size() - 1);
+ }
+ this.chunkSize = chunkSize;
}
public ReduceFileMeta(List<Long> chunkOffsets) {
@@ -39,7 +52,16 @@ public class ReduceFileMeta implements FileMeta {
}
public synchronized void addChunkOffset(long offset) {
- chunkOffsets.add(offset);
+ nextBoundary = offset + chunkSize;
+ if (chunkOffsets.isEmpty() || chunkOffsets.get(chunkOffsets.size() - 1) != offset) {
+ chunkOffsets.add(offset);
+ }
+ }
+
+ public void updateChunkOffset(long bytesFlushed, boolean force) {
+ if (bytesFlushed >= nextBoundary || force) {
+ addChunkOffset(bytesFlushed);
+ }
}
public synchronized long getLastChunkOffset() {
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/FileManagedBuffers.java b/common/src/main/java/org/apache/celeborn/common/network/buffer/ChunkBuffers.java
similarity index 61%
rename from common/src/main/java/org/apache/celeborn/common/meta/FileManagedBuffers.java
rename to common/src/main/java/org/apache/celeborn/common/network/buffer/ChunkBuffers.java
index 52c91ee4f..f4049953b 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/FileManagedBuffers.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/buffer/ChunkBuffers.java
@@ -15,25 +15,25 @@
* limitations under the License.
*/
-package org.apache.celeborn.common.meta;
+package org.apache.celeborn.common.network.buffer;
-import java.io.File;
import java.util.List;
-import org.apache.celeborn.common.network.buffer.FileSegmentManagedBuffer;
-import org.apache.celeborn.common.network.buffer.ManagedBuffer;
-import org.apache.celeborn.common.network.util.TransportConf;
+import scala.Tuple2;
-public class FileManagedBuffers {
- private final File file;
- private final long[] offsets;
- private final int numChunks;
+import com.google.common.base.Preconditions;
- private final TransportConf conf;
+import org.apache.celeborn.common.meta.ReduceFileMeta;
- public FileManagedBuffers(DiskFileInfo fileInfo, TransportConf conf) {
- file = fileInfo.getFile();
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) fileInfo.getFileMeta();
+public abstract class ChunkBuffers {
+ protected long[] offsets;
+ protected int numChunks;
+
+ public int numChunks() {
+ return numChunks;
+ }
+
+ public ChunkBuffers(ReduceFileMeta reduceFileMeta) {
numChunks = reduceFileMeta.getNumChunks();
if (numChunks > 0) {
offsets = new long[numChunks + 1];
@@ -44,19 +44,15 @@ public class FileManagedBuffers {
} else {
offsets = new long[] {0};
}
- this.conf = conf;
}
- public int numChunks() {
- return numChunks;
- }
-
- public ManagedBuffer chunk(int chunkIndex, int offset, int len) {
- // offset of the beginning of the chunk in the file
+ public Tuple2<Long, Long> getChunkOffsetLength(int chunkIndex, int offset, int len) {
final long chunkOffset = offsets[chunkIndex];
final long chunkLength = offsets[chunkIndex + 1] - chunkOffset;
- assert offset < chunkLength;
+ Preconditions.checkArgument(offset < chunkLength);
long length = Math.min(chunkLength - offset, len);
- return new FileSegmentManagedBuffer(conf, file, chunkOffset + offset, length);
+ return new Tuple2<>(chunkOffset + offset, length);
}
+
+ public abstract ManagedBuffer chunk(int chunkIndex, int offset, int len);
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java b/common/src/main/java/org/apache/celeborn/common/network/buffer/FileChunkBuffers.java
similarity index 53%
copy from common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
copy to common/src/main/java/org/apache/celeborn/common/network/buffer/FileChunkBuffers.java
index e58263bdf..ed4969d3a 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/buffer/FileChunkBuffers.java
@@ -15,35 +15,28 @@
* limitations under the License.
*/
-package org.apache.celeborn.common.meta;
+package org.apache.celeborn.common.network.buffer;
-import io.netty.buffer.CompositeByteBuf;
+import java.io.File;
-import org.apache.celeborn.common.identity.UserIdentifier;
+import scala.Tuple2;
-public class MemoryFileInfo extends FileInfo {
- private CompositeByteBuf buffer;
- private long length;
+import org.apache.celeborn.common.meta.DiskFileInfo;
+import org.apache.celeborn.common.network.util.TransportConf;
- public MemoryFileInfo(
- UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
- super(userIdentifier, partitionSplitEnabled, fileMeta);
- }
-
- public CompositeByteBuf getBuffer() {
- return buffer;
- }
-
- public void setBuffer(CompositeByteBuf buffer) {
- this.buffer = buffer;
- }
+public class FileChunkBuffers extends ChunkBuffers {
+ private final File file;
+ private final TransportConf conf;
- public void setBufferSize(int bufferSize) {
- this.length = bufferSize;
+ public FileChunkBuffers(DiskFileInfo fileInfo, TransportConf conf) {
+ super(fileInfo.getReduceFileMeta());
+ file = fileInfo.getFile();
+ this.conf = conf;
}
@Override
- public long getFileLength() {
- return length;
+ public ManagedBuffer chunk(int chunkIndex, int offset, int len) {
+ Tuple2<Long, Long> offsetLen = getChunkOffsetLength(chunkIndex, offset, len);
+ return new FileSegmentManagedBuffer(conf, file, offsetLen._1, offsetLen._2);
}
}
diff --git a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java b/common/src/main/java/org/apache/celeborn/common/network/buffer/MemoryChunkBuffers.java
similarity index 55%
copy from common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
copy to common/src/main/java/org/apache/celeborn/common/network/buffer/MemoryChunkBuffers.java
index e58263bdf..ec2d7286f 100644
--- a/common/src/main/java/org/apache/celeborn/common/meta/MemoryFileInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/network/buffer/MemoryChunkBuffers.java
@@ -15,35 +15,27 @@
* limitations under the License.
*/
-package org.apache.celeborn.common.meta;
+package org.apache.celeborn.common.network.buffer;
-import io.netty.buffer.CompositeByteBuf;
-
-import org.apache.celeborn.common.identity.UserIdentifier;
+import scala.Tuple2;
-public class MemoryFileInfo extends FileInfo {
- private CompositeByteBuf buffer;
- private long length;
-
- public MemoryFileInfo(
- UserIdentifier userIdentifier, boolean partitionSplitEnabled, FileMeta fileMeta) {
- super(userIdentifier, partitionSplitEnabled, fileMeta);
- }
+import io.netty.buffer.CompositeByteBuf;
- public CompositeByteBuf getBuffer() {
- return buffer;
- }
+import org.apache.celeborn.common.meta.MemoryFileInfo;
+import org.apache.celeborn.common.meta.ReduceFileMeta;
- public void setBuffer(CompositeByteBuf buffer) {
- this.buffer = buffer;
- }
+public class MemoryChunkBuffers extends ChunkBuffers {
+ private final CompositeByteBuf buffer;
- public void setBufferSize(int bufferSize) {
- this.length = bufferSize;
+ public MemoryChunkBuffers(MemoryFileInfo memoryFileInfo) {
+ super((ReduceFileMeta) memoryFileInfo.getFileMeta());
+ buffer = memoryFileInfo.getBuffer();
}
@Override
- public long getFileLength() {
- return length;
+ public ManagedBuffer chunk(int chunkIndex, int offset, int len) {
+ Tuple2<Long, Long> offsetLen = getChunkOffsetLength(chunkIndex, offset, len);
+ return new NettyManagedBuffer(
+ buffer.retainedSlice(offsetLen._1.intValue(), offsetLen._2.intValue()));
}
}
diff --git a/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java b/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
index 5b59a10ed..d92909128 100644
--- a/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
+++ b/common/src/main/java/org/apache/celeborn/common/protocol/StorageInfo.java
@@ -38,8 +38,8 @@ public class StorageInfo implements Serializable {
}
}
- public static Map<Integer, Type> typesMap = new HashMap<>();
- public static Set<String> typeNames = new HashSet<>();
+ public static final Map<Integer, Type> typesMap = new HashMap<>();
+ public static final Set<String> typeNames = new HashSet<>();
static {
for (Type type : Type.values()) {
@@ -129,6 +129,11 @@ public class StorageInfo implements Serializable {
+ '}';
}
+ public boolean memoryAvailable() {
+ return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
+ || (availableStorageTypes & MEMORY_MASK) > 0;
+ }
+
public static boolean localDiskAvailable(int availableStorageTypes) {
return availableStorageTypes == ALL_TYPES_AVAILABLE_MASK
|| (availableStorageTypes & LOCAL_DISK_MASK) > 0;
@@ -181,7 +186,7 @@ public class StorageInfo implements Serializable {
return Objects.hash(type, mountPoint, finalResult, filePath, availableStorageTypes);
}
- public static final boolean validate(String typeStr) {
+ public static boolean validate(String typeStr) {
return typeNames.contains(typeStr);
}
diff --git a/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java b/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
index 193b8d3bc..a7f25cc48 100644
--- a/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
+++ b/common/src/main/java/org/apache/celeborn/common/util/ShuffleBlockInfoUtils.java
@@ -23,6 +23,9 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
+
public class ShuffleBlockInfoUtils {
public static class ShuffleBlockInfo {
@@ -34,7 +37,8 @@ public class ShuffleBlockInfoUtils {
int startMapIndex,
int endMapIndex,
long fetchChunkSize,
- Map<Integer, List<ShuffleBlockInfo>> indexMap) {
+ Map<Integer, List<ShuffleBlockInfo>> indexMap,
+ boolean isInMemory) {
List<Long> sortedChunkOffset = new ArrayList<>();
ShuffleBlockInfo lastBlock = null;
int maxMapIndex = endMapIndex;
@@ -43,24 +47,49 @@ public class ShuffleBlockInfoUtils {
maxMapIndex = indexMap.keySet().stream().max(Integer::compareTo).get() + 1;
}
- for (int i = startMapIndex; i < maxMapIndex; i++) {
- List<ShuffleBlockInfo> blockInfos = indexMap.get(i);
- if (blockInfos != null) {
- for (ShuffleBlockInfo info : blockInfos) {
- if (sortedChunkOffset.size() == 0) {
- sortedChunkOffset.add(info.offset);
+ if (isInMemory) {
+ long currentChunkOffset = 0;
+ long lastChunkOffset = 0;
+ // This sorted chunk offsets are used for fetch handler.
+ // Sorted byte buf is a new composite byte buf containing the required data.
+ // It will not reuse the old buffer of memory file, so the offset starts from 0.
+ sortedChunkOffset.add(0l);
+ for (int i = startMapIndex; i < maxMapIndex; i++) {
+ List<ShuffleBlockInfo> blockInfos = indexMap.get(i);
+ if (blockInfos != null) {
+ for (ShuffleBlockInfo info : blockInfos) {
+ currentChunkOffset += info.length;
+ if (currentChunkOffset - lastChunkOffset > fetchChunkSize) {
+ lastChunkOffset = currentChunkOffset;
+ sortedChunkOffset.add(currentChunkOffset);
+ }
}
- if (info.offset - sortedChunkOffset.get(sortedChunkOffset.size() - 1) >= fetchChunkSize) {
- sortedChunkOffset.add(info.offset);
+ }
+ }
+ if (lastChunkOffset != currentChunkOffset) {
+ sortedChunkOffset.add(currentChunkOffset);
+ }
+ } else {
+ for (int i = startMapIndex; i < maxMapIndex; i++) {
+ List<ShuffleBlockInfo> blockInfos = indexMap.get(i);
+ if (blockInfos != null) {
+ for (ShuffleBlockInfo info : blockInfos) {
+ if (sortedChunkOffset.size() == 0) {
+ sortedChunkOffset.add(info.offset);
+ }
+ if (info.offset - sortedChunkOffset.get(sortedChunkOffset.size() - 1)
+ >= fetchChunkSize) {
+ sortedChunkOffset.add(info.offset);
+ }
+ lastBlock = info;
}
- lastBlock = info;
}
}
- }
- if (lastBlock != null) {
- long endChunkOffset = lastBlock.length + lastBlock.offset;
- if (!sortedChunkOffset.contains(endChunkOffset)) {
- sortedChunkOffset.add(endChunkOffset);
+ if (lastBlock != null) {
+ long endChunkOffset = lastBlock.length + lastBlock.offset;
+ if (!sortedChunkOffset.contains(endChunkOffset)) {
+ sortedChunkOffset.add(endChunkOffset);
+ }
}
}
return sortedChunkOffset;
@@ -90,4 +119,23 @@ public class ShuffleBlockInfoUtils {
}
return indexMap;
}
+
+ public static void sliceSortedBufferByMapRange(
+ int startMapIndex,
+ int endMapIndex,
+ Map<Integer, List<ShuffleBlockInfo>> indexMap,
+ CompositeByteBuf sortedByteBuf,
+ CompositeByteBuf targetByteBuf) {
+ for (int i = startMapIndex; i < endMapIndex; i++) {
+ List<ShuffleBlockInfo> blockInfos = indexMap.get(i);
+ if (blockInfos != null) {
+ for (ShuffleBlockInfo blockInfo : blockInfos) {
+ ByteBuf slice = sortedByteBuf.slice((int) blockInfo.offset, (int) blockInfo.length);
+ // Do not retain this buffer because this buffer
+ // will be released when the fileinfo is released
+ targetByteBuf.addComponent(slice);
+ }
+ }
+ }
+ }
}
diff --git a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
index 13eaed0b9..62cb30b4d 100644
--- a/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
+++ b/common/src/main/scala/org/apache/celeborn/common/CelebornConf.scala
@@ -636,7 +636,7 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
def masterSlotAssignPolicy: SlotsAssignPolicy =
SlotsAssignPolicy.valueOf(get(MASTER_SLOT_ASSIGN_POLICY))
def availableStorageTypes: Int = {
- val types = get(ACTIVE_STORAGE_TYPES).split(",").map(StorageInfo.Type.valueOf(_)).toList
+ val types = get(ACTIVE_STORAGE_TYPES).split(",").map(StorageInfo.Type.valueOf).toList
StorageInfo.getAvailableTypes(types.asJava)
}
def hasHDFSStorage: Boolean =
@@ -1200,8 +1200,10 @@ class CelebornConf(loadDefaults: Boolean) extends Cloneable with Logging with Se
get(WORKER_DIRECT_MEMORY_TRIM_CHANNEL_WAIT_INTERVAL)
def workerDirectMemoryTrimFlushWaitInterval: Long =
get(WORKER_DIRECT_MEMORY_TRIM_FLUSH_WAIT_INTERVAL)
- def workerDirectMemoryRatioForShuffleStorage: Double =
- get(WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE)
+ def workerDirectMemoryRatioForMemoryFilesStorage: Double =
+ get(WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE)
+ def workerMemoryFileStorageMaxFileSize: Long =
+ get(WORKER_MEMORY_FILE_STORAGE_MAX_FILE_SIZE)
// //////////////////////////////////////////////////////
// Rate Limit controller //
@@ -2583,6 +2585,7 @@ object CelebornConf extends Logging {
.doc("Max chunk size of reducer's merged shuffle data. For example, if a reducer's " +
"shuffle data is 128M and the data will need 16 fetch chunk requests to fetch.")
.bytesConf(ByteUnit.BYTE)
+ .checkValue(v => v < Integer.MAX_VALUE, "Chunk size can not be larger than 2GB")
.createWithDefaultString("8m")
val CLIENT_FETCH_DFS_READ_CHUNK_SIZE: ConfigEntry[Long] =
@@ -3192,13 +3195,14 @@ object CelebornConf extends Logging {
.doubleConf
.createWithDefault(0.1)
- val WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE: ConfigEntry[Double] =
- buildConf("celeborn.worker.directMemoryRatioForMemoryShuffleStorage")
+ val WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE: ConfigEntry[Double] =
+ buildConf("celeborn.worker.directMemoryRatioForMemoryFileStorage")
.categories("worker")
- .doc("Max ratio of direct memory to store shuffle data")
- .version("0.2.0")
+ .doc("Max ratio of direct memory to store shuffle data. " +
+ "This feature is experimental and disabled by default.")
+ .version("0.5.0")
.doubleConf
- .createWithDefault(0.0)
+ .createWithDefault(0)
val WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE: ConfigEntry[Double] =
buildConf("celeborn.worker.directMemoryRatioToPauseReceive")
@@ -3225,6 +3229,15 @@ object CelebornConf extends Logging {
.doubleConf
.createWithDefault(0.7)
+ val WORKER_MEMORY_FILE_STORAGE_MAX_FILE_SIZE: ConfigEntry[Long] =
+ buildConf("celeborn.worker.memoryFileStorage.maxFileSize")
+ .categories("worker")
+ .doc("Max size for a memory storage file. It must be less than 2GB.")
+ .version("0.5.0")
+ .bytesConf(ByteUnit.BYTE)
+ .checkValue(v => v < Int.MaxValue, "A single memory storage file can not be larger than 2GB")
+ .createWithDefaultString("8MB")
+
val WORKER_CONGESTION_CONTROL_ENABLED: ConfigEntry[Boolean] =
buildConf("celeborn.worker.congestionControl.enabled")
.categories("worker")
@@ -4793,7 +4806,7 @@ object CelebornConf extends Logging {
"Enabled storages. Available options: MEMORY,HDD,SSD,HDFS. Note: HDD and SSD would be treated as identical.")
.stringConf
.transform(_.toUpperCase(Locale.ROOT))
- .checkValue(p => p.split(",").map(StorageInfo.validate(_)).reduce(_ && _), "")
+ .checkValue(p => p.split(",").map(StorageInfo.validate).reduce(_ && _), "")
.createWithDefault("HDD")
val READ_LOCAL_SHUFFLE_FILE: ConfigEntry[Boolean] =
diff --git a/common/src/test/scala/org/apache/celeborn/common/network/CelebornRackResolverSuite.scala b/common/src/test/scala/org/apache/celeborn/common/network/CelebornRackResolverSuite.scala
index 4b46b2602..e7a0f8e40 100644
--- a/common/src/test/scala/org/apache/celeborn/common/network/CelebornRackResolverSuite.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/network/CelebornRackResolverSuite.scala
@@ -135,7 +135,6 @@ class CelebornRackResolverSuite extends AnyFunSuite {
|$hostName2 /default/rack2
|$hostName3 /default/rack1
|""".stripMargin)
-
Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS)
assertEquals("/default/rack1", resolver.resolve(hostName3).getNetworkLocation)
assertEquals(false, resolver.isOnSameRack(hostName1, hostName2))
diff --git a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
index 840b9896d..a9ad63e27 100644
--- a/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
+++ b/common/src/test/scala/org/apache/celeborn/common/util/PbSerDeUtilsTest.scala
@@ -66,13 +66,13 @@ class PbSerDeUtilsTest extends CelebornFunSuite {
val fileInfo1 = new DiskFileInfo(
userIdentifier1,
true,
- new ReduceFileMeta(chunkOffsets1),
+ new ReduceFileMeta(chunkOffsets1, 123),
file1.getAbsolutePath,
3000L)
val fileInfo2 = new DiskFileInfo(
userIdentifier2,
true,
- new ReduceFileMeta(chunkOffsets2),
+ new ReduceFileMeta(chunkOffsets2, 123),
file2.getAbsolutePath,
6000L)
val fileInfoMap = JavaUtils.newConcurrentHashMap[String, DiskFileInfo]()
diff --git a/docs/configuration/worker.md b/docs/configuration/worker.md
index 93cbae431..42d4ec6ed 100644
--- a/docs/configuration/worker.md
+++ b/docs/configuration/worker.md
@@ -58,7 +58,7 @@ license: |
| celeborn.worker.congestionControl.user.inactive.interval | 10min | false | How long will consider this user is inactive if it doesn't send data | 0.3.0 | |
| celeborn.worker.decommission.checkInterval | 30s | false | The wait interval of checking whether all the shuffle expired during worker decommission | 0.4.0 | |
| celeborn.worker.decommission.forceExitTimeout | 6h | false | The wait time of waiting for all the shuffle expire during worker decommission. | 0.4.0 | |
-| celeborn.worker.directMemoryRatioForMemoryShuffleStorage | 0.0 | false | Max ratio of direct memory to store shuffle data | 0.2.0 | |
+| celeborn.worker.directMemoryRatioForMemoryFileStorage | 0.0 | false | Max ratio of direct memory to store shuffle data. This feature is experimental and disabled by default. | 0.5.0 | |
| celeborn.worker.directMemoryRatioForReadBuffer | 0.1 | false | Max ratio of direct memory for read buffer | 0.2.0 | |
| celeborn.worker.directMemoryRatioToPauseReceive | 0.85 | false | If direct memory usage reaches this limit, the worker will stop to receive data from Celeborn shuffle clients. | 0.2.0 | |
| celeborn.worker.directMemoryRatioToPauseReplicate | 0.95 | false | If direct memory usage reaches this limit, the worker will stop to receive replication data from other workers. This value should be higher than celeborn.worker.directMemoryRatioToPauseReceive. | 0.2.0 | |
@@ -101,6 +101,7 @@ license: |
| celeborn.worker.jvmQuake.exitCode | 502 | false | The exit code of system kill for the maximum GC 'deficit' during worker jvm quake. | 0.4.0 | |
| celeborn.worker.jvmQuake.kill.threshold | 60s | false | The threshold of system kill for the maximum GC 'deficit' which can be accumulated before jvmquake takes action. | 0.4.0 | |
| celeborn.worker.jvmQuake.runtimeWeight | 5.0 | false | The factor by which to multiply running JVM time, when weighing it against GCing time. 'Deficit' is accumulated as `gc_time - runtime * runtime_weight`, and is compared against threshold to determine whether to take action. | 0.4.0 | |
+| celeborn.worker.memoryFileStorage.maxFileSize | 8MB | false | Max size for a memory storage file. It must be less than 2GB. | 0.5.0 | |
| celeborn.worker.monitor.disk.check.interval | 30s | false | Intervals between device monitor to check disk. | 0.3.0 | celeborn.worker.monitor.disk.checkInterval |
| celeborn.worker.monitor.disk.check.timeout | 30s | false | Timeout time for worker check device status. | 0.3.0 | celeborn.worker.disk.check.timeout |
| celeborn.worker.monitor.disk.checklist | readwrite,diskusage | false | Monitor type for disk, available items are: iohang, readwrite and diskusage. | 0.2.0 | |
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySkewJoinSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySkewJoinSuite.scala
new file mode 100644
index 000000000..d8ae35144
--- /dev/null
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySkewJoinSuite.scala
@@ -0,0 +1,113 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.celeborn.tests.spark.memory
+
+import scala.util.Random
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.internal.SQLConf
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.common.CelebornConf
+import org.apache.celeborn.common.protocol.CompressionCodec
+
+class MemorySkewJoinSuite extends AnyFunSuite
+ with MemorySparkTestBase
+ with BeforeAndAfterEach {
+
+ override def beforeEach(): Unit = {
+ ShuffleClient.reset()
+ }
+
+ override def afterEach(): Unit = {
+ System.gc()
+ }
+
+ private def enableCeleborn(conf: SparkConf) = {
+ conf.set("spark.shuffle.manager", "org.apache.spark.shuffle.celeborn.SparkShuffleManager")
+ .set(s"spark.${CelebornConf.MASTER_ENDPOINTS.key}", masterInfo._1.rpcEnv.address.toString)
+ .set(s"spark.${CelebornConf.SHUFFLE_PARTITION_SPLIT_THRESHOLD.key}", "10MB")
+ }
+
+ CompressionCodec.values.foreach { codec =>
+ test(s"celeborn spark integration test - skew join - $codec") {
+ val sparkConf = new SparkConf().setAppName("celeborn-demo")
+ .setMaster("local[2]")
+ .set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "true")
+ .set("spark.sql.adaptive.skewJoin.enabled", "true")
+ .set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "16MB")
+ .set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "12MB")
+ .set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
+ .set("spark.sql.adaptive.autoBroadcastJoinThreshold", "-1")
+ .set(s"spark.celeborn.storage.availableTypes", "HDD,MEMORY")
+ .set(SQLConf.PARQUET_COMPRESSION.key, "gzip")
+ .set(s"spark.${CelebornConf.SHUFFLE_COMPRESSION_CODEC.key}", codec.name)
+ .set(s"spark.${CelebornConf.SHUFFLE_RANGE_READ_FILTER_ENABLED.key}", "true")
+
+ enableCeleborn(sparkConf)
+
+ val sparkSession = SparkSession.builder().config(sparkConf).getOrCreate()
+ if (sparkSession.version.startsWith("3")) {
+ import sparkSession.implicits._
+ val df = sparkSession.sparkContext.parallelize(1 to 120000, 8)
+ .map(i => {
+ val random = new Random()
+ val oriKey = random.nextInt(64)
+ val key = if (oriKey < 32) 1 else oriKey
+ val fas = random.nextInt(1200000)
+ val fa = Range(fas, fas + 100).mkString(",")
+ val fbs = random.nextInt(1200000)
+ val fb = Range(fbs, fbs + 100).mkString(",")
+ val fcs = random.nextInt(1200000)
+ val fc = Range(fcs, fcs + 100).mkString(",")
+ val fds = random.nextInt(1200000)
+ val fd = Range(fds, fds + 100).mkString(",")
+
+ (key, fa, fb, fc, fd)
+ })
+ .toDF("fa", "f1", "f2", "f3", "f4")
+ df.createOrReplaceTempView("view1")
+ val df2 = sparkSession.sparkContext.parallelize(1 to 8, 8)
+ .map(i => {
+ val random = new Random()
+ val oriKey = random.nextInt(64)
+ val key = if (oriKey < 32) 1 else oriKey
+ val fas = random.nextInt(1200000)
+ val fa = Range(fas, fas + 100).mkString(",")
+ val fbs = random.nextInt(1200000)
+ val fb = Range(fbs, fbs + 100).mkString(",")
+ val fcs = random.nextInt(1200000)
+ val fc = Range(fcs, fcs + 100).mkString(",")
+ val fds = random.nextInt(1200000)
+ val fd = Range(fds, fds + 100).mkString(",")
+ (key, fa, fb, fc, fd)
+ })
+ .toDF("fb", "f6", "f7", "f8", "f9")
+ df2.createOrReplaceTempView("view2")
+ sparkSession.sql("drop table if exists fres")
+ sparkSession.sql("create table fres using parquet as select * from view1 a inner join view2 b on a.fa=b.fb where a.fa=1 ")
+ sparkSession.sql("drop table fres")
+ sparkSession.stop()
+ }
+ }
+ }
+
+}
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySparkTestBase.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySparkTestBase.scala
new file mode 100644
index 000000000..d38d84561
--- /dev/null
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemorySparkTestBase.scala
@@ -0,0 +1,62 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.celeborn.tests.spark.memory
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.internal.SQLConf
+import org.scalatest.{BeforeAndAfterAll, BeforeAndAfterEach}
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.common.CelebornConf._
+import org.apache.celeborn.common.internal.Logging
+import org.apache.celeborn.common.protocol.ShuffleMode
+import org.apache.celeborn.service.deploy.MiniClusterFeature
+import org.apache.celeborn.tests.spark.SparkTestBase
+
+trait MemorySparkTestBase extends AnyFunSuite
+ with Logging with MiniClusterFeature with BeforeAndAfterAll with BeforeAndAfterEach
+ with SparkTestBase {
+
+ override def beforeAll(): Unit = {
+ logInfo("test initialized , setup Celeborn mini cluster")
+ val workerConfs = Map("celeborn.worker.directMemoryRatioForMemoryFileStorage" -> "0.2")
+ setupMiniClusterWithRandomPorts(workerConf = workerConfs, workerNum = 5)
+ }
+
+ override def afterAll(): Unit = {
+ logInfo("all test complete , stop Celeborn mini cluster")
+ shutdownMiniCluster()
+ }
+
+ override def updateSparkConf(sparkConf: SparkConf, mode: ShuffleMode): SparkConf = {
+ sparkConf.set("spark.serializer", "org.apache.spark.serializer.KryoSerializer")
+ sparkConf.set(
+ "spark.shuffle.manager",
+ "org.apache.spark.shuffle.celeborn.SparkShuffleManager")
+ sparkConf.set("spark.shuffle.useOldFetchProtocol", "true")
+ sparkConf.set(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key, "false")
+ sparkConf.set("spark.shuffle.service.enabled", "false")
+ sparkConf.set("spark.sql.adaptive.skewJoin.enabled", "false")
+ sparkConf.set("spark.sql.adaptive.localShuffleReader.enabled", "false")
+ sparkConf.set(s"spark.${MASTER_ENDPOINTS.key}", masterInfo._1.rpcEnv.address.toString)
+ sparkConf.set(s"spark.${SPARK_SHUFFLE_WRITER_MODE.key}", mode.toString)
+ sparkConf.set(s"spark.celeborn.storage.availableTypes", "HDD,MEMORY")
+
+ sparkConf
+ }
+}
diff --git a/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemoryStorageReusedExchangeSuite.scala b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemoryStorageReusedExchangeSuite.scala
new file mode 100644
index 000000000..5304eaede
--- /dev/null
+++ b/tests/spark-it/src/test/scala/org/apache/celeborn/tests/spark/memory/MemoryStorageReusedExchangeSuite.scala
@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.celeborn.tests.spark.memory
+
+import org.apache.spark.SparkConf
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.functions._
+import org.scalatest.BeforeAndAfterEach
+import org.scalatest.funsuite.AnyFunSuite
+
+import org.apache.celeborn.client.ShuffleClient
+import org.apache.celeborn.common.CelebornConf
+
+class MemoryStorageReusedExchangeSuite extends AnyFunSuite
+ with MemorySparkTestBase
+ with BeforeAndAfterEach {
+
+ override def beforeEach(): Unit = {
+ ShuffleClient.reset()
+ }
+
+ test("[CELEBORN-980] Asynchronously delete original files to fix ReusedExchange bug") {
+ testReusedExchange(false)
+ }
+
+ def testReusedExchange(readLocalShuffle: Boolean): Unit = {
+ val sparkConf = new SparkConf().setAppName("celeborn-test").setMaster("local[2]")
+ .set("spark.shuffle.manager", "org.apache.spark.shuffle.celeborn.SparkShuffleManager")
+ .set(s"spark.${CelebornConf.MASTER_ENDPOINTS.key}", masterInfo._1.rpcEnv.address.toString)
+ .set(s"spark.${CelebornConf.READ_LOCAL_SHUFFLE_FILE.key}", readLocalShuffle.toString)
+ .set("spark.sql.autoBroadcastJoinThreshold", "-1")
+ .set("spark.sql.adaptive.skewJoin.skewedPartitionThresholdInBytes", "100")
+ .set("spark.sql.adaptive.advisoryPartitionSizeInBytes", "100")
+ .set(s"spark.celeborn.storage.availableTypes", "HDD,MEMORY")
+ val spark = SparkSession.builder()
+ .config(sparkConf)
+ .getOrCreate()
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id as k1", "id as v1")
+ .createOrReplaceTempView("ta")
+ spark.range(0, 1000, 1, 10)
+ .selectExpr("id % 1 as k21", "id % 1 as k22", "id as v2")
+ .createOrReplaceTempView("tb")
+ spark.range(140)
+ .select(
+ col("id").cast("long").as("k3"),
+ concat(col("id").cast("string"), lit("a")).as("v3"))
+ .createOrReplaceTempView("tc")
+
+ spark.sql(
+ """
+ |SELECT *
+ |FROM ta
+ |LEFT JOIN tb ON ta.k1 = tb.k21
+ |LEFT JOIN tc ON tb.k22 = tc.k3
+ |""".stripMargin)
+ .createOrReplaceTempView("v1")
+
+ spark.sql(
+ """
+ |SELECT * FROM v1 WHERE v3 IS NOT NULL
+ |UNION
+ |SELECT * FROM v1
+ |""".stripMargin)
+ .collect()
+ spark.stop
+ }
+}
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
index c58ccc55e..ae25060c3 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/memory/MemoryManager.java
@@ -17,8 +17,8 @@
package org.apache.celeborn.service.deploy.worker.memory;
-import java.util.ArrayList;
-import java.util.List;
+import java.io.IOException;
+import java.util.*;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
@@ -29,6 +29,7 @@ import java.util.concurrent.atomic.LongAdder;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
+import io.netty.buffer.PooledByteBufAllocator;
import io.netty.util.internal.PlatformDependent;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -39,14 +40,16 @@ import org.apache.celeborn.common.util.ThreadUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.reflect.DynMethods;
import org.apache.celeborn.service.deploy.worker.storage.CreditStreamManager;
+import org.apache.celeborn.service.deploy.worker.storage.PartitionDataWriter;
+import org.apache.celeborn.service.deploy.worker.storage.StorageManager;
public class MemoryManager {
private static final Logger logger = LoggerFactory.getLogger(MemoryManager.class);
private static volatile MemoryManager _INSTANCE = null;
- @VisibleForTesting public long maxDirectorMemory = 0;
+ @VisibleForTesting public long maxDirectMemory;
private final long pausePushDataThreshold;
private final long pauseReplicateThreshold;
- private final long resumeThreshold;
+ private final double resumeRatio;
private final long maxSortMemory;
private final int forceAppendPauseSpentTimeThreshold;
private final List<MemoryPressureListener> memoryPressureListeners = new ArrayList<>();
@@ -75,8 +78,8 @@ public class MemoryManager {
private volatile boolean isPaused = false;
// For credit stream
private final AtomicLong readBufferCounter = new AtomicLong(0);
- private long readBufferThreshold = 0;
- private long readBufferTarget = 0;
+ private long readBufferThreshold;
+ private long readBufferTarget;
private ReadBufferDispatcher readBufferDispatcher;
private List<ReadBufferTargetChangeListener> readBufferTargetChangeListeners;
private long lastNotifiedTarget = 0;
@@ -85,11 +88,18 @@ public class MemoryManager {
"worker-memory-manager-read-buffer-target-updater");
private CreditStreamManager creditStreamManager = null;
- private long memoryShuffleStorageThreshold = 0;
+ private long memoryFileStorageThreshold;
+ private final LongAdder memoryFileStorageCounter = new LongAdder();
+ private final StorageManager storageManager;
+ @VisibleForTesting
public static MemoryManager initialize(CelebornConf conf) {
+ return initialize(conf, null);
+ }
+
+ public static MemoryManager initialize(CelebornConf conf, StorageManager storageManager) {
if (_INSTANCE == null) {
- _INSTANCE = new MemoryManager(conf);
+ _INSTANCE = new MemoryManager(conf, storageManager);
}
return _INSTANCE;
}
@@ -104,28 +114,27 @@ public class MemoryManager {
return _INSTANCE;
}
- private MemoryManager(CelebornConf conf) {
+ private MemoryManager(CelebornConf conf, StorageManager storageManager) {
double pausePushDataRatio = conf.workerDirectMemoryRatioToPauseReceive();
double pauseReplicateRatio = conf.workerDirectMemoryRatioToPauseReplicate();
- double resumeRatio = conf.workerDirectMemoryRatioToResume();
+ this.resumeRatio = conf.workerDirectMemoryRatioToResume();
double maxSortMemRatio = conf.workerPartitionSorterDirectMemoryRatioThreshold();
double readBufferRatio = conf.workerDirectMemoryRatioForReadBuffer();
- double shuffleStorageRatio = conf.workerDirectMemoryRatioForShuffleStorage();
+ double memoryFileStorageRatio = conf.workerDirectMemoryRatioForMemoryFilesStorage();
long checkInterval = conf.workerDirectMemoryPressureCheckIntervalMs();
long reportInterval = conf.workerDirectMemoryReportIntervalSecond();
double readBufferTargetRatio = conf.readBufferTargetRatio();
long readBufferTargetUpdateInterval = conf.readBufferTargetUpdateInterval();
long readBufferTargetNotifyThreshold = conf.readBufferTargetNotifyThreshold();
forceAppendPauseSpentTimeThreshold = conf.metricsWorkerForceAppendPauseSpentTimeThreshold();
-
- maxDirectorMemory =
+ maxDirectMemory =
DynMethods.builder("maxDirectMemory")
.impl("jdk.internal.misc.VM") // for Java 10 and above
.impl("sun.misc.VM") // for Java 9 and previous
.buildStatic()
.<Long>invoke();
- Preconditions.checkArgument(maxDirectorMemory > 0);
+ Preconditions.checkArgument(maxDirectMemory > 0);
Preconditions.checkArgument(
pauseReplicateRatio > pausePushDataRatio,
String.format(
@@ -135,15 +144,16 @@ public class MemoryManager {
CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(),
pausePushDataRatio));
Preconditions.checkArgument(pausePushDataRatio > resumeRatio);
- Preconditions.checkArgument(resumeRatio > (readBufferRatio + shuffleStorageRatio));
+ if (memoryFileStorageRatio > 0) {
+ Preconditions.checkArgument(resumeRatio > (readBufferRatio + memoryFileStorageRatio));
+ }
- maxSortMemory = ((long) (maxDirectorMemory * maxSortMemRatio));
- pausePushDataThreshold = (long) (maxDirectorMemory * pausePushDataRatio);
- pauseReplicateThreshold = (long) (maxDirectorMemory * pauseReplicateRatio);
- resumeThreshold = (long) (maxDirectorMemory * resumeRatio);
- readBufferThreshold = (long) (maxDirectorMemory * readBufferRatio);
+ maxSortMemory = ((long) (maxDirectMemory * maxSortMemRatio));
+ pausePushDataThreshold = (long) (maxDirectMemory * pausePushDataRatio);
+ pauseReplicateThreshold = (long) (maxDirectMemory * pauseReplicateRatio);
+ readBufferThreshold = (long) (maxDirectMemory * readBufferRatio);
readBufferTarget = (long) (readBufferThreshold * readBufferTargetRatio);
- memoryShuffleStorageThreshold = (long) (maxDirectorMemory * shuffleStorageRatio);
+ memoryFileStorageThreshold = (long) (maxDirectMemory * memoryFileStorageRatio);
checkService.scheduleWithFixedDelay(
() -> {
@@ -160,12 +170,17 @@ public class MemoryManager {
reportService.scheduleWithFixedDelay(
() ->
logger.info(
- "Direct memory usage: {}/{}, disk buffer size: {}, sort memory size: {}, read buffer size: {}",
+ "Direct memory usage: {}/{}, "
+ + "disk buffer size: {}, "
+ + "sort memory size: {}, "
+ + "read buffer size: {}, "
+ + "memory file storage size : {}",
Utils.bytesToString(getNettyUsedDirectMemory()),
- Utils.bytesToString(maxDirectorMemory),
+ Utils.bytesToString(maxDirectMemory),
Utils.bytesToString(diskBufferCounter.get()),
Utils.bytesToString(sortMemoryCounter.get()),
- Utils.bytesToString(readBufferCounter.get())),
+ Utils.bytesToString(readBufferCounter.get()),
+ Utils.bytesToString(memoryFileStorageCounter.sum())),
reportInterval,
reportInterval,
TimeUnit.SECONDS);
@@ -208,19 +223,64 @@ public class MemoryManager {
TimeUnit.MILLISECONDS);
}
+ this.storageManager = storageManager;
+ if (memoryFileStorageThreshold > 0
+ && storageManager != null
+ && storageManager.localOrHdfsStorageAvailable()) {
+ ScheduledExecutorService memoryFileStorageService =
+ ThreadUtils.newDaemonSingleThreadScheduledExecutor("memory-file-storage-checker");
+ memoryFileStorageService.scheduleWithFixedDelay(
+ () -> {
+ try {
+ if ((memoryFileStorageCounter.sum() >= 0.5 * memoryFileStorageThreshold)
+ && currentServingState() != ServingState.NONE_PAUSED) {
+ List<PartitionDataWriter> memoryWriters =
+ new ArrayList<>(storageManager.memoryWriters().values());
+ if (memoryWriters.isEmpty()) {
+ return;
+ }
+ logger.info("Start evicting {} memory file infos", memoryWriters.size());
+ // always evict the largest memory file info first
+ memoryWriters.sort(
+ Comparator.comparingLong(o -> o.getMemoryFileInfo().getFileLength()));
+ Collections.reverse(memoryWriters);
+ try {
+ for (PartitionDataWriter writer : memoryWriters) {
+ // this branch means that there is no memory pressure
+ if ((memoryFileStorageCounter.sum() < 0.5 * memoryFileStorageThreshold)
+ || currentServingState() == ServingState.NONE_PAUSED) {
+ break;
+ }
+ logger.debug("Evict writer {}", writer);
+ writer.evict(true);
+ }
+ } catch (IOException e) {
+ logger.warn("Partition data writer evict failed", e);
+ }
+ }
+ } catch (Exception e) {
+ logger.error("Evict thread encounter error", e);
+ }
+ },
+ checkInterval,
+ checkInterval,
+ TimeUnit.MILLISECONDS);
+ }
+
logger.info(
"Memory tracker initialized with: "
+ "max direct memory: {}, pause push memory: {}, "
- + "pause replication memory: {}, resume memory: {}, "
+ + "pause replication memory: {}, "
+ "read buffer memory limit: {} target: {}, "
- + "memory shuffle storage limit: {}",
- Utils.bytesToString(maxDirectorMemory),
+ + "memory shuffle storage limit: {}，"
+ + "resume memory ratio: {},",
+ Utils.bytesToString(maxDirectMemory),
Utils.bytesToString(pausePushDataThreshold),
Utils.bytesToString(pauseReplicateThreshold),
- Utils.bytesToString(resumeThreshold),
Utils.bytesToString(readBufferThreshold),
Utils.bytesToString(readBufferTarget),
- Utils.bytesToString(memoryShuffleStorageThreshold));
+ Utils.bytesToString(memoryFileStorageThreshold),
+ resumeRatio);
}
public ServingState currentServingState() {
@@ -237,7 +297,7 @@ public class MemoryManager {
return ServingState.PUSH_PAUSED;
}
// trigger resume
- if (memoryUsage < resumeThreshold) {
+ if (memoryUsage / (double) (maxDirectMemory) < resumeRatio) {
isPaused = false;
return ServingState.NONE_PAUSED;
}
@@ -449,6 +509,26 @@ public class MemoryManager {
this.creditStreamManager = creditStreamManager;
}
+ public double workerMemoryUsageRatio() {
+ return getMemoryUsage() / (double) (maxDirectMemory);
+ }
+
+ public long getMemoryFileStorageCounter() {
+ return memoryFileStorageCounter.sum();
+ }
+
+ public boolean memoryFileStorageAvailable() {
+ return memoryFileStorageCounter.sum() < memoryFileStorageThreshold;
+ }
+
+ public void increaseMemoryFileStorage(int bytes) {
+ memoryFileStorageCounter.add(bytes);
+ }
+
+ public void releaseMemoryFileStorage(int bytes) {
+ memoryFileStorageCounter.add(-1 * bytes);
+ }
+
public void close() {
checkService.shutdown();
reportService.shutdown();
@@ -459,6 +539,10 @@ public class MemoryManager {
readBufferDispatcher.close();
}
+ public PooledByteBufAllocator getStoragePooledByteBufAllocator() {
+ return storageManager.storageBufferAllocator();
+ }
+
@VisibleForTesting
public static void reset() {
_INSTANCE = null;
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManager.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManager.java
index 006ecbe44..7ea390168 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManager.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManager.java
@@ -28,8 +28,8 @@ import com.google.common.annotations.VisibleForTesting;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
-import org.apache.celeborn.common.meta.FileManagedBuffers;
import org.apache.celeborn.common.meta.TimeWindow;
+import org.apache.celeborn.common.network.buffer.ChunkBuffers;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
import org.apache.celeborn.common.util.JavaUtils;
@@ -47,20 +47,16 @@ public class ChunkStreamManager {
protected final ConcurrentHashMap<String, Set<Long>> shuffleStreamIds;
/** State of a single stream. */
- protected static class StreamState {
- final FileManagedBuffers buffers;
- final String shuffleKey;
- final String fileName;
- final TimeWindow fetchTimeMetric;
-
+ public static class StreamState {
+ public final ChunkBuffers buffers;
+ public final String shuffleKey;
+ public final String fileName;
+ public final TimeWindow fetchTimeMetric;
// Used to keep track of the number of chunks being transferred and not finished yet.
volatile long chunksBeingTransferred = 0L;
StreamState(
- String shuffleKey,
- FileManagedBuffers buffers,
- String fileName,
- TimeWindow fetchTimeMetric) {
+ String shuffleKey, ChunkBuffers buffers, String fileName, TimeWindow fetchTimeMetric) {
this.buffers = buffers;
this.shuffleKey = shuffleKey;
this.fileName = fileName;
@@ -87,7 +83,7 @@ public class ChunkStreamManager {
String.format("Requested chunk index beyond end %s", chunkIndex));
}
- FileManagedBuffers buffers = state.buffers;
+ ChunkBuffers buffers = state.buffers;
return buffers.chunk(chunkIndex, offset, len);
}
@@ -146,7 +142,7 @@ public class ChunkStreamManager {
* stream is not properly closed, it will eventually be cleaned up by `cleanupExpiredShuffleKey`.
*/
public long registerStream(
- String shuffleKey, FileManagedBuffers buffers, String fileName, TimeWindow fetchTimeMetric) {
+ String shuffleKey, ChunkBuffers buffers, String fileName, TimeWindow fetchTimeMetric) {
long myStreamId = nextStreamId.getAndIncrement();
return registerStream(myStreamId, shuffleKey, buffers, fileName, fetchTimeMetric);
}
@@ -154,10 +150,11 @@ public class ChunkStreamManager {
public long registerStream(
long streamId,
String shuffleKey,
- FileManagedBuffers buffers,
+ ChunkBuffers buffers,
String fileName,
TimeWindow fetchTimeMetric) {
- streams.put(streamId, new StreamState(shuffleKey, buffers, fileName, fetchTimeMetric));
+ StreamState streamState = new StreamState(shuffleKey, buffers, fileName, fetchTimeMetric);
+ streams.put(streamId, streamState);
shuffleStreamIds.compute(
shuffleKey,
(key, value) -> {
@@ -193,6 +190,10 @@ public class ChunkStreamManager {
streams.size());
}
+ public StreamState getStreamState(long streamId) {
+ return streams.get(streamId);
+ }
+
public Tuple2<String, String> getShuffleKeyAndFileName(long streamId) {
StreamState state = streams.get(streamId);
return new Tuple2<>(state.shuffleKey, state.fileName);
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriter.java
index 2ce0cbbc1..4eab8f0b1 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriter.java
@@ -23,6 +23,7 @@ import java.nio.ByteOrder;
import java.nio.channels.FileChannel;
import java.util.Arrays;
+import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import org.apache.hadoop.fs.FSDataOutputStream;
import org.apache.hadoop.fs.Path;
@@ -30,11 +31,8 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.MapFileMeta;
import org.apache.celeborn.common.metrics.source.AbstractSource;
-import org.apache.celeborn.common.protocol.PartitionSplitMode;
-import org.apache.celeborn.common.protocol.PartitionType;
import org.apache.celeborn.common.util.FileChannelUtils;
import org.apache.celeborn.common.util.Utils;
@@ -57,28 +55,14 @@ public final class MapPartitionDataWriter extends PartitionDataWriter {
public MapPartitionDataWriter(
StorageManager storageManager,
- DiskFileInfo diskFileInfo,
- Flusher flusher,
AbstractSource workerSource,
CelebornConf conf,
DeviceMonitor deviceMonitor,
- long splitThreshold,
- PartitionSplitMode splitMode,
- boolean rangeReadFilter,
- String shuffleKey)
+ PartitionDataWriterContext writerContext)
throws IOException {
- super(
- storageManager,
- diskFileInfo,
- flusher,
- workerSource,
- conf,
- deviceMonitor,
- splitThreshold,
- splitMode,
- PartitionType.MAP,
- rangeReadFilter,
- shuffleKey);
+ super(storageManager, workerSource, conf, deviceMonitor, writerContext, false);
+
+ Preconditions.checkState(diskFileInfo != null);
if (!diskFileInfo.isHdfs()) {
indexChannel = FileChannelUtils.createWritableFileChannel(diskFileInfo.getIndexPath());
} else {
@@ -134,9 +118,7 @@ public final class MapPartitionDataWriter extends PartitionDataWriter {
@Override
public synchronized long close() throws IOException {
return super.close(
- () -> {
- flushIndex();
- },
+ this::flushIndex,
() -> {
if (diskFileInfo.isHdfs()) {
if (StorageManager.hadoopFs().exists(diskFileInfo.getHdfsPeerWriterSuccessPath())) {
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
index 2058dd649..a2f00b0a8 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriter.java
@@ -19,26 +19,33 @@ package org.apache.celeborn.service.deploy.worker.storage;
import java.io.File;
import java.io.IOException;
+import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
-import javax.annotation.concurrent.GuardedBy;
+import scala.Tuple4;
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.PooledByteBufAllocator;
import org.roaringbitmap.RoaringBitmap;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.exception.AlreadyClosedException;
+import org.apache.celeborn.common.exception.CelebornIOException;
import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.DiskStatus;
+import org.apache.celeborn.common.meta.FileInfo;
+import org.apache.celeborn.common.meta.MemoryFileInfo;
import org.apache.celeborn.common.metrics.source.AbstractSource;
import org.apache.celeborn.common.protocol.PartitionSplitMode;
-import org.apache.celeborn.common.protocol.PartitionType;
import org.apache.celeborn.common.protocol.StorageInfo;
import org.apache.celeborn.common.unsafe.Platform;
import org.apache.celeborn.common.util.FileChannelUtils;
@@ -54,69 +61,106 @@ public abstract class PartitionDataWriter implements DeviceObserver {
private static final Logger logger = LoggerFactory.getLogger(PartitionDataWriter.class);
private static final long WAIT_INTERVAL_MS = 5;
- protected final DiskFileInfo diskFileInfo;
+ // After commit file, there will be only 1 fileinfo left.
+ protected DiskFileInfo diskFileInfo = null;
+ protected MemoryFileInfo memoryFileInfo = null;
private FileChannel channel;
private volatile boolean closed;
private volatile boolean destroyed;
protected final AtomicInteger numPendingWrites = new AtomicInteger();
- public final Flusher flusher;
- private final int flushWorkerIndex;
+ public Flusher flusher;
+ private int flushWorkerIndex;
- @GuardedBy("flushLock")
- private CompositeByteBuf flushBuffer;
+ protected CompositeByteBuf flushBuffer;
- private final Object flushLock = new Object();
+ protected final Object flushLock = new Object();
private final long writerCloseTimeoutMs;
- protected final long flusherBufferSize;
+ protected long flusherBufferSize;
protected final DeviceMonitor deviceMonitor;
protected final AbstractSource source; // metrics
- private long splitThreshold = 0;
+ private final long splitThreshold;
private final PartitionSplitMode splitMode;
- private final PartitionType partitionType;
private final boolean rangeReadFilter;
protected boolean deleted = false;
private RoaringBitmap mapIdBitMap = null;
protected final FlushNotifier notifier = new FlushNotifier();
// It's only needed when graceful shutdown is enabled
- private String shuffleKey;
- private final StorageManager storageManager;
+ private final String shuffleKey;
+ protected final StorageManager storageManager;
private final boolean workerGracefulShutdown;
+ protected final long memoryFileStorageMaxFileSize;
+ protected AtomicBoolean isMemoryShuffleFile = new AtomicBoolean();
+ protected final String filename;
+ protected PooledByteBufAllocator pooledByteBufAllocator;
+ private final PartitionDataWriterContext writerContext;
+ private final long localFlusherBufferSize;
+ private final long hdfsFlusherBufferSize;
+ private Exception exception = null;
+ private boolean metricsCollectCriticalEnabled;
public PartitionDataWriter(
StorageManager storageManager,
- DiskFileInfo diskFileInfo,
- Flusher flusher,
AbstractSource workerSource,
CelebornConf conf,
DeviceMonitor deviceMonitor,
- long splitThreshold,
- PartitionSplitMode splitMode,
- PartitionType partitionType,
- boolean rangeReadFilter,
- String shuffleKey)
+ PartitionDataWriterContext writerContext,
+ boolean supportInMemory)
throws IOException {
this.storageManager = storageManager;
- this.diskFileInfo = diskFileInfo;
- this.flusher = flusher;
- this.flushWorkerIndex = flusher.getWorkerIndex();
this.writerCloseTimeoutMs = conf.workerWriterCloseTimeoutMs();
this.workerGracefulShutdown = conf.workerGracefulShutdown();
- this.splitThreshold = splitThreshold;
+ this.splitThreshold = writerContext.getSplitThreshold();
this.deviceMonitor = deviceMonitor;
- this.splitMode = splitMode;
- this.partitionType = partitionType;
- this.rangeReadFilter = rangeReadFilter;
- this.shuffleKey = shuffleKey;
+ this.splitMode = writerContext.getPartitionSplitMode();
+ this.rangeReadFilter = writerContext.isRangeReadFilter();
+ this.shuffleKey = writerContext.getShuffleKey();
+ this.memoryFileStorageMaxFileSize = conf.workerMemoryFileStorageMaxFileSize();
+ this.filename = writerContext.getPartitionLocation().getFileName();
+ this.writerContext = writerContext;
+ this.localFlusherBufferSize = conf.workerFlusherBufferSize();
+ this.hdfsFlusherBufferSize = conf.workerHdfsFlusherBufferSize();
+ this.metricsCollectCriticalEnabled = conf.metricsCollectCriticalEnabled();
+
+ Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File> createFileResult =
+ storageManager.createFile(writerContext, supportInMemory);
+
+ // Reduce partition data writers support memory storage now
+ if (supportInMemory && createFileResult._1() != null) {
+ this.memoryFileInfo = createFileResult._1();
+ this.pooledByteBufAllocator = storageManager.storageBufferAllocator();
+ this.isMemoryShuffleFile.set(true);
+ storageManager.registerMemoryPartitionWriter(this, createFileResult._1());
+ } else if (createFileResult._2() != null) {
+ this.diskFileInfo = createFileResult._3();
+ this.flusher = createFileResult._2();
+ File workingDir = createFileResult._4();
+ this.isMemoryShuffleFile.set(false);
+ initFileChannelsForDiskFile();
+ storageManager.registerDiskFilePartitionWriter(this, workingDir, diskFileInfo);
+ } else {
+ throw new CelebornIOException(
+ "Create file failed for location:" + writerContext.getPartitionLocation().toString());
+ }
+
+ source = workerSource;
+ logger.debug("FileWriter {} split threshold {} mode {}", this, splitThreshold, splitMode);
+ if (rangeReadFilter) {
+ this.mapIdBitMap = new RoaringBitmap();
+ }
+ takeBuffer();
+ }
+
+ public void initFileChannelsForDiskFile() throws IOException {
if (!this.diskFileInfo.isHdfs()) {
- this.flusherBufferSize = conf.workerFlusherBufferSize();
+ this.flusherBufferSize = localFlusherBufferSize;
channel = FileChannelUtils.createWritableFileChannel(this.diskFileInfo.getFilePath());
} else {
- this.flusherBufferSize = conf.workerHdfsFlusherBufferSize();
+ this.flusherBufferSize = hdfsFlusherBufferSize;
// We open the stream and close immediately because HDFS output stream will
// create a DataStreamer that is a thread.
// If we reuse HDFS output stream, we will exhaust the memory soon.
@@ -132,12 +176,6 @@ public abstract class PartitionDataWriter implements DeviceObserver {
StorageManager.hadoopFs().create(this.diskFileInfo.getHdfsPath(), true).close();
}
}
- source = workerSource;
- logger.debug("FileWriter {} split threshold {} mode {}", this, splitThreshold, splitMode);
- if (rangeReadFilter) {
- this.mapIdBitMap = new RoaringBitmap();
- }
- takeBuffer();
}
public DiskFileInfo getDiskFileInfo() {
@@ -156,23 +194,56 @@ public abstract class PartitionDataWriter implements DeviceObserver {
numPendingWrites.decrementAndGet();
}
- protected void flush(boolean finalFlush) throws IOException {
- synchronized (flushLock) {
- // flushBuffer == null here means writer already closed
- if (flushBuffer != null) {
- int numBytes = flushBuffer.readableBytes();
- if (numBytes != 0) {
- notifier.checkException();
+ @VisibleForTesting
+ public void flush(boolean finalFlush, boolean fromEvict) throws IOException {
+ // flushBuffer == null here means this writer is already closed
+ if (flushBuffer != null) {
+ int numBytes = flushBuffer.readableBytes();
+ if (numBytes != 0) {
+ notifier.checkException();
+ FlushTask task = null;
+ if (fromEvict) {
notifier.numPendingFlushes.incrementAndGet();
- FlushTask task = null;
+ // duplicate buffer before its released
+ ByteBuf dupBuf = flushBuffer.retainedDuplicate();
+ // flush task will release the buffer of memory shuffle file
if (channel != null) {
- task = new LocalFlushTask(flushBuffer, channel, notifier);
+ task = new LocalFlushTask(flushBuffer, channel, notifier, false);
} else if (diskFileInfo.isHdfs()) {
- task = new HdfsFlushTask(flushBuffer, diskFileInfo.getHdfsPath(), notifier);
+ task = new HdfsFlushTask(flushBuffer, diskFileInfo.getHdfsPath(), notifier, false);
+ }
+ MemoryManager.instance().releaseMemoryFileStorage(numBytes);
+ MemoryManager.instance().incrementDiskBuffer(numBytes);
+ // read flush buffer to generate correct chunk offsets
+ // data header layout (mapId, attemptId, nextBatchId, length)
+ ByteBuffer headerBuf = ByteBuffer.allocate(16);
+ while (dupBuf.isReadable()) {
+ headerBuf.rewind();
+ dupBuf.readBytes(headerBuf);
+ byte[] batchHeader = headerBuf.array();
+ int compressedSize = Platform.getInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 12);
+ dupBuf.skipBytes(compressedSize);
+ diskFileInfo.updateBytesFlushed(compressedSize + 16);
}
+ dupBuf.release();
+ } else {
+ if (!isMemoryShuffleFile.get()) {
+ notifier.numPendingFlushes.incrementAndGet();
+ if (channel != null) {
+ task = new LocalFlushTask(flushBuffer, channel, notifier, true);
+ } else if (diskFileInfo.isHdfs()) {
+ task = new HdfsFlushTask(flushBuffer, diskFileInfo.getHdfsPath(), notifier, true);
+ }
+ }
+ }
+ // task won't be null in real workloads
+ // task will be null in UT to check chunk size and offset
+ if (task != null) {
addTask(task);
flushBuffer = null;
- diskFileInfo.updateBytesFlushed(numBytes);
+ if (!fromEvict) {
+ diskFileInfo.updateBytesFlushed(numBytes);
+ }
if (!finalFlush) {
takeBuffer();
}
@@ -181,10 +252,20 @@ public abstract class PartitionDataWriter implements DeviceObserver {
}
}
+ public boolean needHardSplitForMemoryShuffleStorage() {
+ if (!isMemoryShuffleFile.get()) {
+ return false;
+ } else {
+ return !storageManager.localOrHdfsStorageAvailable()
+ && (memoryFileInfo.getFileLength() > memoryFileStorageMaxFileSize
+ || !MemoryManager.instance().memoryFileStorageAvailable());
+ }
+ }
+
/** assume data size is less than chunk capacity */
public void write(ByteBuf data) throws IOException {
if (closed) {
- String msg = "FileWriter has already closed!, fileName " + diskFileInfo.getFilePath();
+ String msg = getFileAlreadyClosedMsg();
logger.warn(msg);
throw new AlreadyClosedException(msg);
}
@@ -203,49 +284,95 @@ public abstract class PartitionDataWriter implements DeviceObserver {
}
final int numBytes = data.readableBytes();
- MemoryManager.instance().incrementDiskBuffer(numBytes);
-
- Optional.ofNullable(CongestionController.instance())
- .ifPresent(
- congestionController ->
- congestionController.produceBytes(diskFileInfo.getUserIdentifier(), numBytes));
+ if (isMemoryShuffleFile.get()) {
+ MemoryManager.instance().increaseMemoryFileStorage(numBytes);
+ } else {
+ MemoryManager.instance().incrementDiskBuffer(numBytes);
+ Optional.ofNullable(CongestionController.instance())
+ .ifPresent(
+ congestionController ->
+ congestionController.produceBytes(diskFileInfo.getUserIdentifier(), numBytes));
+ }
synchronized (flushLock) {
if (closed) {
- String msg = "FileWriter has already closed!, fileName " + diskFileInfo.getFilePath();
+ String msg = getFileAlreadyClosedMsg();
logger.warn(msg);
throw new AlreadyClosedException(msg);
}
if (rangeReadFilter) {
mapIdBitMap.add(mapId);
}
- if (flushBuffer.readableBytes() != 0
- && flushBuffer.readableBytes() + numBytes >= flusherBufferSize) {
- flush(false);
+ int flushBufferReadableBytes = flushBuffer.readableBytes();
+ if (!isMemoryShuffleFile.get()) {
+ if (flushBufferReadableBytes != 0
+ && flushBufferReadableBytes + numBytes >= flusherBufferSize) {
+ flush(false, false);
+ }
+ } else {
+ if (flushBufferReadableBytes > memoryFileStorageMaxFileSize
+ && storageManager.localOrHdfsStorageAvailable()) {
+ logger.debug(
+ "{} Evict, memory buffer is {}",
+ writerContext.getPartitionLocation().getFileName(),
+ flushBufferReadableBytes);
+ evict(false);
+ }
}
data.retain();
flushBuffer.addComponent(true, data);
+ if (isMemoryShuffleFile.get()) {
+ memoryFileInfo.updateBytesFlushed(numBytes);
+ }
}
numPendingWrites.decrementAndGet();
}
+ public void evictInternal() throws IOException {
+ if (exception != null) {
+ return;
+ }
+ Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File> createFileResult =
+ storageManager.createFile(writerContext, false);
+ if (createFileResult._4() != null) {
+ this.diskFileInfo = createFileResult._3();
+ this.flusher = createFileResult._2();
+
+ isMemoryShuffleFile.set(false);
+ initFileChannelsForDiskFile();
+ flush(closed, true);
+
+ logger.debug("evict {} {}", shuffleKey, filename);
+ storageManager.unregisterMemoryPartitionWriterAndFileInfo(
+ memoryFileInfo, shuffleKey, filename);
+ storageManager.evictedFileCount().incrementAndGet();
+ memoryFileInfo = null;
+ } else {
+ exception = new CelebornIOException("PartitionDataWriter create disk-related file failed");
+ throw (CelebornIOException) exception;
+ }
+ }
+
public RoaringBitmap getMapIdBitMap() {
return mapIdBitMap;
}
public StorageInfo getStorageInfo() {
- if (flusher instanceof LocalFlusher) {
- LocalFlusher localFlusher = (LocalFlusher) flusher;
- // do not write file path to reduce rpc size
- return new StorageInfo(localFlusher.diskType(), true, "");
- } else {
- if (deleted) {
- return null;
+ if (diskFileInfo != null) {
+ if (diskFileInfo.isHdfs()) {
+ if (deleted) {
+ return null;
+ } else {
+ return new StorageInfo(StorageInfo.Type.HDFS, true, diskFileInfo.getFilePath());
+ }
} else {
- return new StorageInfo(StorageInfo.Type.HDFS, true, diskFileInfo.getFilePath());
+ return new StorageInfo(((LocalFlusher) flusher).diskType(), true, "");
}
+ } else {
+ Preconditions.checkArgument(memoryFileInfo != null);
+ return new StorageInfo(StorageInfo.Type.MEMORY, true, "");
}
}
@@ -266,7 +393,7 @@ public abstract class PartitionDataWriter implements DeviceObserver {
RunnableWithIOException finalClose)
throws IOException {
if (closed) {
- String msg = "FileWriter has already closed! fileName " + diskFileInfo.getFilePath();
+ String msg = getFileAlreadyClosedMsg();
logger.error(msg);
throw new AlreadyClosedException(msg);
}
@@ -276,15 +403,18 @@ public abstract class PartitionDataWriter implements DeviceObserver {
closed = true;
synchronized (flushLock) {
- if (flushBuffer.readableBytes() > 0) {
- flush(true);
+ if (!isMemoryShuffleFile.get()) {
+ // memory shuffle file doesn't need final flush
+ if (flushBuffer != null && flushBuffer.readableBytes() > 0) {
+ flush(true, false);
+ }
}
}
tryClose.run();
waitOnNoPending(notifier.numPendingFlushes);
} finally {
- returnBuffer();
+ returnBuffer(false);
try {
if (channel != null) {
channel.close();
@@ -297,15 +427,53 @@ public abstract class PartitionDataWriter implements DeviceObserver {
finalClose.run();
// unregister from DeviceMonitor
- if (!diskFileInfo.isHdfs()) {
+ if (diskFileInfo != null && !diskFileInfo.isHdfs()) {
logger.debug("file info {} unregister from device monitor", diskFileInfo);
deviceMonitor.unregisterFileWriter(this);
}
}
if (workerGracefulShutdown) {
- storageManager.notifyFileInfoCommitted(shuffleKey, getFile().getName(), diskFileInfo);
+ if (diskFileInfo != null) {
+ storageManager.notifyFileInfoCommitted(shuffleKey, getFile().getName(), diskFileInfo);
+ }
+ }
+ if (diskFileInfo != null) {
+ return diskFileInfo.getFileLength();
+ } else {
+ return memoryFileInfo.getFileLength();
+ }
+ }
+
+ private String getFileAlreadyClosedMsg() {
+ String msg = "PartitionDataWriter has already closed! ";
+ if (isMemoryShuffleFile.get()) {
+ msg += "In memory file name:" + filename;
+ } else {
+ msg += "Disk file name:" + diskFileInfo.getFilePath();
+ }
+ return msg;
+ }
+
+ public void evict(boolean checkClose) throws IOException {
+ // this lock is used to make sure that
+ // memory manager won't evict with writer thread concurrently
+ synchronized (flushLock) {
+ if (checkClose) {
+ // close and evict might be invoked concurrently
+ // do not evict committed files from memory manager
+ // evict memory file info if worker is shutdown gracefully
+ if (isClosed()) {
+ return;
+ }
+ }
+ if (memoryFileInfo != null) {
+ evictInternal();
+ if (isClosed()) {
+ waitOnNoPending(notifier.numPendingFlushes);
+ storageManager.notifyFileInfoCommitted(shuffleKey, getFile().getName(), diskFileInfo);
+ }
+ }
}
- return diskFileInfo.getFileLength();
}
public synchronized void destroy(IOException ioException) {
@@ -314,7 +482,7 @@ public abstract class PartitionDataWriter implements DeviceObserver {
if (!notifier.hasException()) {
notifier.setException(ioException);
}
- returnBuffer();
+ returnBuffer(true);
try {
if (channel != null) {
channel.close();
@@ -329,15 +497,25 @@ public abstract class PartitionDataWriter implements DeviceObserver {
if (!destroyed) {
destroyed = true;
- diskFileInfo.deleteAllFiles(StorageManager.hadoopFs());
+ if (diskFileInfo != null) {
+ diskFileInfo.deleteAllFiles(StorageManager.hadoopFs());
- // unregister from DeviceMonitor
- if (!diskFileInfo.isHdfs()) {
- deviceMonitor.unregisterFileWriter(this);
+ // unregister from DeviceMonitor
+ if (!diskFileInfo.isHdfs()) {
+ deviceMonitor.unregisterFileWriter(this);
+ }
}
}
}
+ protected FileInfo getCurrentFileInfo() {
+ if (!isMemoryShuffleFile.get()) {
+ return diskFileInfo;
+ } else {
+ return memoryFileInfo;
+ }
+ }
+
public IOException getException() {
if (notifier.hasException()) {
return notifier.exception.get();
@@ -368,18 +546,26 @@ public abstract class PartitionDataWriter implements DeviceObserver {
}
protected void takeBuffer() {
- if (source.metricsCollectCriticalEnabled()) {
- String metricsName = WorkerSource.TAKE_BUFFER_TIME();
- String fileAbsPath = diskFileInfo.getFilePath();
+ String metricsName = null;
+ String fileAbsPath = null;
+ if (metricsCollectCriticalEnabled) {
+ metricsName = WorkerSource.TAKE_BUFFER_TIME();
+ fileAbsPath = diskFileInfo.getFilePath();
source.startTimer(metricsName, fileAbsPath);
- synchronized (flushLock) {
+ }
+
+ synchronized (flushLock) {
+ if (diskFileInfo != null) {
flushBuffer = flusher.takeBuffer();
+ } else {
+ if (flushBuffer == null) {
+ flushBuffer = pooledByteBufAllocator.compositeBuffer(Integer.MAX_VALUE);
+ }
}
+ }
+
+ if (metricsCollectCriticalEnabled) {
source.stopTimer(metricsName, fileAbsPath);
- } else {
- synchronized (flushLock) {
- flushBuffer = flusher.takeBuffer();
- }
}
}
@@ -391,11 +577,18 @@ public abstract class PartitionDataWriter implements DeviceObserver {
}
}
- protected void returnBuffer() {
+ protected void returnBuffer(boolean destroy) {
synchronized (flushLock) {
if (flushBuffer != null) {
- flusher.returnBuffer(flushBuffer);
- flushBuffer = null;
+ if (flusher != null) {
+ flusher.returnBuffer(flushBuffer, true);
+ flushBuffer = null;
+ } else {
+ if (destroy) {
+ flushBuffer.removeComponents(0, flushBuffer.numComponents());
+ flushBuffer.release();
+ }
+ }
}
}
}
@@ -415,12 +608,13 @@ public abstract class PartitionDataWriter implements DeviceObserver {
@Override
public String toString() {
- return diskFileInfo.getFilePath();
+ return shuffleKey + "-" + filename + " partition-writer";
}
public void flushOnMemoryPressure() throws IOException {
synchronized (flushLock) {
- flush(false);
+ // this won't happen if this writer is in memory
+ flush(false, false);
}
}
@@ -456,7 +650,7 @@ public abstract class PartitionDataWriter implements DeviceObserver {
@Override
public void notifyNonCriticalError(String mountPoint, DiskStatus diskStatus) {}
- public PartitionType getPartitionType() {
- return partitionType;
+ public MemoryFileInfo getMemoryFileInfo() {
+ return memoryFileInfo;
}
}
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterContext.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterContext.java
new file mode 100644
index 000000000..b73b901ad
--- /dev/null
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterContext.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.celeborn.service.deploy.worker.storage;
+
+import org.apache.celeborn.common.identity.UserIdentifier;
+import org.apache.celeborn.common.protocol.PartitionLocation;
+import org.apache.celeborn.common.protocol.PartitionSplitMode;
+import org.apache.celeborn.common.protocol.PartitionType;
+import org.apache.celeborn.common.util.Utils;
+
+public class PartitionDataWriterContext {
+ private final long splitThreshold;
+ private final PartitionSplitMode partitionSplitMode;
+ private final boolean rangeReadFilter;
+ private final PartitionLocation partitionLocation;
+ private final String appId;
+ private final int shuffleId;
+ private final UserIdentifier userIdentifier;
+ private final boolean partitionSplitEnabled;
+ private final String shuffleKey;
+ private final PartitionType partitionType;
+
+ public PartitionDataWriterContext(
+ long splitThreshold,
+ PartitionSplitMode partitionSplitMode,
+ boolean rangeReadFilter,
+ PartitionLocation partitionLocation,
+ String appId,
+ int shuffleId,
+ UserIdentifier userIdentifier,
+ PartitionType partitionType,
+ boolean partitionSplitEnabled) {
+ this.splitThreshold = splitThreshold;
+ this.partitionSplitMode = partitionSplitMode;
+ this.rangeReadFilter = rangeReadFilter;
+ this.partitionLocation = partitionLocation;
+ this.appId = appId;
+ this.shuffleId = shuffleId;
+ this.userIdentifier = userIdentifier;
+ this.partitionSplitEnabled = partitionSplitEnabled;
+ this.partitionType = partitionType;
+ this.shuffleKey = Utils.makeShuffleKey(appId, shuffleId);
+ }
+
+ public long getSplitThreshold() {
+ return splitThreshold;
+ }
+
+ public PartitionSplitMode getPartitionSplitMode() {
+ return partitionSplitMode;
+ }
+
+ public boolean isRangeReadFilter() {
+ return rangeReadFilter;
+ }
+
+ public PartitionLocation getPartitionLocation() {
+ return partitionLocation;
+ }
+
+ public String getAppId() {
+ return appId;
+ }
+
+ public int getShuffleId() {
+ return shuffleId;
+ }
+
+ public UserIdentifier getUserIdentifier() {
+ return userIdentifier;
+ }
+
+ public boolean isPartitionSplitEnabled() {
+ return partitionSplitEnabled;
+ }
+
+ public String getShuffleKey() {
+ return shuffleKey;
+ }
+
+ public PartitionType getPartitionType() {
+ return partitionType;
+ }
+}
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
index 0025e6482..481ea6c96 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorter.java
@@ -40,6 +40,8 @@ import java.util.concurrent.locks.ReentrantLock;
import com.google.common.annotations.VisibleForTesting;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
+import io.netty.buffer.ByteBuf;
+import io.netty.buffer.CompositeByteBuf;
import org.apache.commons.io.IOUtils;
import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FSDataOutputStream;
@@ -49,9 +51,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
-import org.apache.celeborn.common.meta.DiskFileInfo;
-import org.apache.celeborn.common.meta.FileInfo;
-import org.apache.celeborn.common.meta.ReduceFileMeta;
+import org.apache.celeborn.common.meta.*;
import org.apache.celeborn.common.metrics.source.AbstractSource;
import org.apache.celeborn.common.unsafe.Platform;
import org.apache.celeborn.common.util.*;
@@ -88,7 +88,6 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
protected final long shuffleChunkSize;
protected final boolean prefetchEnabled;
protected final long reservedMemoryPerPartition;
- private final boolean gracefulShutdown;
private final long partitionSorterShutdownAwaitTime;
private DB sortedFilesDb;
@@ -96,7 +95,6 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
private final ExecutorService fileSorterExecutors;
private final ExecutorService fileSorterSchedulerThread;
- private final long indexCacheMaxWeight;
public PartitionFilesSorter(
MemoryManager memoryManager, CelebornConf conf, AbstractSource source) {
@@ -106,10 +104,11 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
this.reservedMemoryPerPartition = conf.workerPartitionSorterReservedMemoryPerPartition();
this.partitionSorterShutdownAwaitTime =
conf.workerGracefulShutdownPartitionSorterCloseAwaitTimeMs();
- this.indexCacheMaxWeight = conf.workerPartitionSorterIndexCacheMaxWeight();
+ long indexCacheMaxWeight = conf.workerPartitionSorterIndexCacheMaxWeight();
this.source = source;
this.cleaner = new PartitionFilesCleaner(this);
- this.gracefulShutdown = conf.workerGracefulShutdown();
+ boolean gracefulShutdown = conf.workerGracefulShutdown();
+ // Assume a chunk won't be larger than 2GB
// ShuffleClient can fetch shuffle data from a restarted worker only
// when the worker's fetching port is stable and enables graceful shutdown.
if (gracefulShutdown) {
@@ -161,8 +160,6 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
() -> {
try {
task.sort();
- } catch (InterruptedException e) {
- logger.warn("File sorter thread was interrupted.");
} finally {
if (task.isPrefetch) {
memoryManager.releaseSortMemory(reservedMemoryPerPartition);
@@ -194,74 +191,166 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
// 2. If the FileSorter task is already in the sorting queue but the sorted file has not been
// generated, it awaits until a timeout occurs (default 220 seconds).
// 3. If the sorted file is generated, it returns the sorted FileInfo.
- public DiskFileInfo getSortedFileInfo(
- String shuffleKey, String fileName, DiskFileInfo fileInfo, int startMapIndex, int endMapIndex)
+ // This method will generate temporary file info for this shuffle read
+ public FileInfo getSortedFileInfo(
+ String shuffleKey, String fileName, FileInfo fileInfo, int startMapIndex, int endMapIndex)
throws IOException {
- String fileId = shuffleKey + "-" + fileName;
- UserIdentifier userIdentifier = fileInfo.getUserIdentifier();
- Set<String> sorted =
- sortedShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
- Set<String> sorting =
- sortingShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
-
- String sortedFilePath = Utils.getSortedFilePath(fileInfo.getFilePath());
- String indexFilePath = Utils.getIndexFilePath(fileInfo.getFilePath());
- boolean fileSorting = true;
- synchronized (sorting) {
- if (sorted.contains(fileId)) {
- fileSorting = false;
- } else if (!sorting.contains(fileId)) {
- try {
- FileSorter fileSorter = new FileSorter(fileInfo, fileId, shuffleKey);
- sorting.add(fileId);
- logger.debug(
- "Adding sorter to sort queue shuffle key {}, file name {}", shuffleKey, fileName);
- shuffleSortTaskDeque.put(fileSorter);
- } catch (InterruptedException e) {
- logger.error("Sorter scheduler thread is interrupted means worker is shutting down.", e);
- throw new IOException(
- "Sort scheduler thread is interrupted means worker is shutting down.", e);
- } catch (IOException e) {
- logger.error("File sorter access HDFS failed.", e);
- throw new IOException("File sorter access HDFS failed.", e);
+ if (fileInfo instanceof MemoryFileInfo) {
+ MemoryFileInfo memoryFileInfo = ((MemoryFileInfo) fileInfo);
+ Map<Integer, List<ShuffleBlockInfo>> indexesMap;
+ sortMemoryShuffleFile(memoryFileInfo);
+ indexesMap = memoryFileInfo.getSortedIndexes();
+
+ ReduceFileMeta reduceFileMeta =
+ new ReduceFileMeta(
+ ShuffleBlockInfoUtils.getChunkOffsetsFromShuffleBlockInfos(
+ startMapIndex, endMapIndex, shuffleChunkSize, indexesMap, true),
+ shuffleChunkSize);
+ CompositeByteBuf targetBuffer =
+ MemoryManager.instance()
+ .getStoragePooledByteBufAllocator()
+ .compositeBuffer(Integer.MAX_VALUE);
+ ShuffleBlockInfoUtils.sliceSortedBufferByMapRange(
+ startMapIndex, endMapIndex, indexesMap, memoryFileInfo.getSortedBuffer(), targetBuffer);
+ return new MemoryFileInfo(
+ memoryFileInfo.getUserIdentifier(),
+ memoryFileInfo.isPartitionSplitEnabled(),
+ reduceFileMeta,
+ targetBuffer);
+ } else {
+ DiskFileInfo diskFileInfo = ((DiskFileInfo) fileInfo);
+ String fileId = shuffleKey + "-" + fileName;
+ UserIdentifier userIdentifier = diskFileInfo.getUserIdentifier();
+ Set<String> sorted =
+ sortedShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
+ Set<String> sorting =
+ sortingShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
+
+ String sortedFilePath = Utils.getSortedFilePath(diskFileInfo.getFilePath());
+ String indexFilePath = Utils.getIndexFilePath(diskFileInfo.getFilePath());
+ boolean fileSorting = true;
+ synchronized (sorting) {
+ if (sorted.contains(fileId)) {
+ fileSorting = false;
+ } else if (!sorting.contains(fileId)) {
+ try {
+ FileSorter fileSorter = new FileSorter(diskFileInfo, fileId, shuffleKey);
+ sorting.add(fileId);
+ logger.debug(
+ "Adding sorter to sort queue shuffle key {}, file name {}", shuffleKey, fileName);
+ shuffleSortTaskDeque.put(fileSorter);
+ } catch (InterruptedException e) {
+ logger.error(
+ "Sorter scheduler thread is interrupted means worker is shutting down.", e);
+ throw new IOException(
+ "Sort scheduler thread is interrupted means worker is shutting down.", e);
+ } catch (IOException e) {
+ logger.error("File sorter access HDFS failed.", e);
+ throw new IOException("File sorter access HDFS failed.", e);
+ }
}
}
- }
- if (fileSorting) {
- long sortStartTime = System.currentTimeMillis();
- while (!sorted.contains(fileId)) {
- if (sorting.contains(fileId)) {
- try {
- Thread.sleep(50);
- if (System.currentTimeMillis() - sortStartTime > sortTimeout) {
- logger.error("Sorting file {} timeout after {}ms", fileId, sortTimeout);
+ if (fileSorting) {
+ long sortStartTime = System.currentTimeMillis();
+ while (!sorted.contains(fileId)) {
+ if (sorting.contains(fileId)) {
+ try {
+ Thread.sleep(50);
+ if (System.currentTimeMillis() - sortStartTime > sortTimeout) {
+ logger.error("Sorting file {} timeout after {}ms", fileId, sortTimeout);
+ throw new IOException(
+ "Sort file " + diskFileInfo.getFilePath() + " timeout after " + sortTimeout);
+ }
+ } catch (InterruptedException e) {
+ logger.error(
+ "Sorter scheduler thread is interrupted means worker is shutting down.", e);
throw new IOException(
- "Sort file " + fileInfo.getFilePath() + " timeout after " + sortTimeout);
+ "Sorter scheduler thread is interrupted means worker is shutting down.", e);
}
- } catch (InterruptedException e) {
- logger.error(
- "Sorter scheduler thread is interrupted means worker is shutting down.", e);
+ } else {
+ logger.debug(
+ "Sorting shuffle file for {} {} failed.", shuffleKey, diskFileInfo.getFilePath());
throw new IOException(
- "Sorter scheduler thread is interrupted means worker is shutting down.", e);
+ "Sorting shuffle file for "
+ + shuffleKey
+ + " "
+ + diskFileInfo.getFilePath()
+ + " failed.");
}
- } else {
- logger.debug(
- "Sorting shuffle file for {} {} failed.", shuffleKey, fileInfo.getFilePath());
- throw new IOException(
- "Sorting shuffle file for " + shuffleKey + " " + fileInfo.getFilePath() + " failed.");
}
}
+
+ return resolve(
+ shuffleKey,
+ fileId,
+ userIdentifier,
+ sortedFilePath,
+ indexFilePath,
+ startMapIndex,
+ endMapIndex);
}
+ }
- return resolve(
- shuffleKey,
- fileId,
- userIdentifier,
- sortedFilePath,
- indexFilePath,
- startMapIndex,
- endMapIndex);
+ public static void sortMemoryShuffleFile(MemoryFileInfo memoryFileInfo) {
+ ReduceFileMeta reduceFileMeta = ((ReduceFileMeta) memoryFileInfo.getFileMeta());
+ synchronized (reduceFileMeta.getSorted()) {
+ if (!reduceFileMeta.getSorted().get()) {
+ CompositeByteBuf originBuffer = memoryFileInfo.getBuffer();
+ Map<Integer, List<ShuffleBlockInfo>> blocksMap = new TreeMap<>();
+ int originReaderIndex = originBuffer.readerIndex();
+ int originWriterIndex = originBuffer.writerIndex();
+ int bufLength = originBuffer.readableBytes();
+ int index = 0;
+ ByteBuffer headerBuf = ByteBuffer.allocate(16);
+
+ while (index != bufLength) {
+ headerBuf.rewind();
+ originBuffer.readerIndex(index);
+ originBuffer.readBytes(headerBuf);
+ byte[] batchHeader = headerBuf.array();
+ int mapId = Platform.getInt(batchHeader, Platform.BYTE_ARRAY_OFFSET);
+ int compressedSize = Platform.getInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 12);
+ ShuffleBlockInfo shuffleBlockInfo = new ShuffleBlockInfo();
+ shuffleBlockInfo.offset = index;
+ shuffleBlockInfo.length = 16 + compressedSize;
+ List<ShuffleBlockInfo> singleMapIdShuffleBlockList =
+ blocksMap.computeIfAbsent(mapId, v -> new ArrayList<>());
+ singleMapIdShuffleBlockList.add(shuffleBlockInfo);
+ index += 16 + compressedSize;
+ }
+ originBuffer.setIndex(originReaderIndex, originWriterIndex);
+
+ // sorted buffer should not consolidate
+ // because this will affect origin buffer's reference count
+ CompositeByteBuf sortedBuffer =
+ MemoryManager.instance()
+ .getStoragePooledByteBufAllocator()
+ .compositeBuffer(Integer.MAX_VALUE - 1);
+ Map<Integer, List<ShuffleBlockInfo>> sortedBlocks = new TreeMap<>();
+ int sortedBufferIndex = 0;
+ for (Map.Entry<Integer, List<ShuffleBlockInfo>> entry : blocksMap.entrySet()) {
+ int mapId = entry.getKey();
+ List<ShuffleBlockInfo> blockInfos = entry.getValue();
+ List<ShuffleBlockInfo> sortedMapBlocks = new ArrayList<>(blockInfos.size());
+ sortedBlocks.put(mapId, sortedMapBlocks);
+ for (ShuffleBlockInfo blockInfo : blockInfos) {
+ int offset = (int) blockInfo.offset;
+ int length = (int) blockInfo.length;
+ ByteBuf slice = originBuffer.slice(offset, length);
+ sortedBuffer.addComponent(true, slice);
+ ShuffleBlockInfo shuffleBlockInfo = new ShuffleBlockInfo();
+ shuffleBlockInfo.offset = sortedBufferIndex;
+ shuffleBlockInfo.length = blockInfo.length;
+ sortedBufferIndex += (int) blockInfo.length;
+ sortedMapBlocks.add(shuffleBlockInfo);
+ }
+ }
+ memoryFileInfo.setSortedBuffer(sortedBuffer);
+ memoryFileInfo.setSortedIndexes(sortedBlocks);
+ reduceFileMeta.setSorted();
+ }
+ }
}
public void cleanup(HashSet<String> expiredShuffleKeys) {
@@ -301,7 +390,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
}
}
long end = System.currentTimeMillis();
- logger.info("Await partition sorter executor complete cost " + (end - start) + "ms");
+ logger.info("Await partition sorter executor complete cost {}ms", end - start);
} else {
fileSorterSchedulerThread.shutdownNow();
fileSorterExecutors.shutdownNow();
@@ -355,8 +444,8 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
}
@VisibleForTesting
- public Set<String> initSortedShuffleFiles(String shuffleKey) {
- return sortedShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
+ public void initSortedShuffleFiles(String shuffleKey) {
+ sortedShuffleFiles.computeIfAbsent(shuffleKey, v -> ConcurrentHashMap.newKeySet());
}
@VisibleForTesting
@@ -494,7 +583,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
FileChannel indexChannel = null;
FSDataInputStream hdfsIndexStream = null;
boolean isHdfs = Utils.isHdfsPath(indexFilePath);
- int indexSize = 0;
+ int indexSize;
try {
if (isHdfs) {
hdfsIndexStream = StorageManager.hadoopFs().open(new Path(indexFilePath));
@@ -536,7 +625,8 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
ReduceFileMeta reduceFileMeta =
new ReduceFileMeta(
ShuffleBlockInfoUtils.getChunkOffsetsFromShuffleBlockInfos(
- startMapIndex, endMapIndex, shuffleChunkSize, indexMap));
+ startMapIndex, endMapIndex, shuffleChunkSize, indexMap, false),
+ shuffleChunkSize);
return new DiskFileInfo(userIdentifier, reduceFileMeta, sortedFilePath);
}
@@ -585,7 +675,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
}
}
- public void sort() throws InterruptedException {
+ public void sort() {
source.startTimer(WorkerSource.SORT_TIME(), fileId);
try {
@@ -640,7 +730,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
writeIndex(sortedBlockInfoMap, indexFilePath, isHdfs);
updateSortedShuffleFiles(shuffleKey, fileId, originFileLen);
- ((ReduceFileMeta) originFileInfo.getFileMeta()).setSorted();
+ originFileInfo.getReduceFileMeta().setSorted();
cleaner.add(this);
logger.debug("sort complete for {} {}", shuffleKey, originFilePath);
} catch (Exception e) {
@@ -699,7 +789,7 @@ public class PartitionFilesSorter extends ShuffleRecoverHelper {
}
public void deleteOriginFiles() throws IOException {
- boolean deleteSuccess = false;
+ boolean deleteSuccess;
if (isHdfs) {
deleteSuccess = StorageManager.hadoopFs().delete(new Path(originFilePath), false);
} else {
@@ -785,7 +875,7 @@ class PartitionFilesCleaner {
while (it.hasNext()) {
PartitionFilesSorter.FileSorter sorter = it.next();
try {
- if (((DiskFileInfo) sorter.getOriginFileInfo()).isStreamsEmpty()) {
+ if (sorter.getOriginFileInfo().isStreamsEmpty()) {
logger.debug(
"Deleting the original files for shuffle key {}: {}",
sorter.getShuffleKey(),
diff --git a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriter.java b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriter.java
index 04d4d46d3..20910ab88 100644
--- a/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriter.java
+++ b/worker/src/main/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriter.java
@@ -24,11 +24,8 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.meta.DiskFileInfo;
-import org.apache.celeborn.common.meta.ReduceFileMeta;
+import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.metrics.source.AbstractSource;
-import org.apache.celeborn.common.protocol.PartitionSplitMode;
-import org.apache.celeborn.common.protocol.PartitionType;
/*
* reduce partition file writer, it will create chunk index
@@ -36,49 +33,27 @@ import org.apache.celeborn.common.protocol.PartitionType;
public final class ReducePartitionDataWriter extends PartitionDataWriter {
private static final Logger logger = LoggerFactory.getLogger(ReducePartitionDataWriter.class);
- private long nextBoundary;
- private final long shuffleChunkSize;
-
public ReducePartitionDataWriter(
StorageManager storageManager,
- DiskFileInfo diskFileInfo,
- Flusher flusher,
AbstractSource workerSource,
CelebornConf conf,
DeviceMonitor deviceMonitor,
- long splitThreshold,
- PartitionSplitMode splitMode,
- boolean rangeReadFilter,
- String shuffleKey)
+ PartitionDataWriterContext writerContext)
throws IOException {
- super(
- storageManager,
- diskFileInfo,
- flusher,
- workerSource,
- conf,
- deviceMonitor,
- splitThreshold,
- splitMode,
- PartitionType.REDUCE,
- rangeReadFilter,
- shuffleKey);
- this.shuffleChunkSize = conf.shuffleChunkSize();
- this.nextBoundary = this.shuffleChunkSize;
+ super(storageManager, workerSource, conf, deviceMonitor, writerContext, true);
}
- @Override
- protected void flush(boolean finalFlush) throws IOException {
- super.flush(finalFlush);
- maybeSetChunkOffsets(finalFlush);
+ private void updateLastChunkOffset() {
+ FileInfo fileInfo = getCurrentFileInfo();
+ fileInfo.getReduceFileMeta().updateChunkOffset(fileInfo.getFileLength(), true);
}
- private void maybeSetChunkOffsets(boolean forceSet) {
- long bytesFlushed = diskFileInfo.getFileLength();
- if (bytesFlushed >= nextBoundary || forceSet) {
- ((ReduceFileMeta) diskFileInfo.getFileMeta()).addChunkOffset(bytesFlushed);
- nextBoundary = bytesFlushed + shuffleChunkSize;
- }
+ @Override
+ public void flush(boolean finalFlush, boolean fromEvict) throws IOException {
+ super.flush(finalFlush, fromEvict);
+ getCurrentFileInfo()
+ .getReduceFileMeta()
+ .updateChunkOffset(getCurrentFileInfo().getFileLength(), finalFlush);
}
private boolean isChunkOffsetValid() {
@@ -90,36 +65,57 @@ public final class ReducePartitionDataWriter extends PartitionDataWriter {
// but its size is smaller than the nextBoundary, then the
// chunk offset will not be set after flushing. we should
// set it during FileWriter close.
- return ((ReduceFileMeta) diskFileInfo.getFileMeta()).getLastChunkOffset()
- == diskFileInfo.getFileLength();
+ if (diskFileInfo != null) {
+ return (diskFileInfo.getReduceFileMeta()).getLastChunkOffset()
+ == diskFileInfo.getFileLength();
+ }
+ if (memoryFileInfo != null) {
+ return (memoryFileInfo.getReduceFileMeta()).getLastChunkOffset()
+ == memoryFileInfo.getFileLength();
+ }
+ // this should not happen
+ return false;
}
@Override
public synchronized long close() throws IOException {
- return super.close(
- () -> {
- if (!isChunkOffsetValid()) {
- maybeSetChunkOffsets(true);
- }
- },
- () -> {
- if (diskFileInfo.isHdfs()) {
- if (StorageManager.hadoopFs().exists(diskFileInfo.getHdfsPeerWriterSuccessPath())) {
- StorageManager.hadoopFs().delete(diskFileInfo.getHdfsPath(), false);
- deleted = true;
- } else {
- StorageManager.hadoopFs().create(diskFileInfo.getHdfsWriterSuccessPath()).close();
- FSDataOutputStream indexOutputStream =
- StorageManager.hadoopFs().create(diskFileInfo.getHdfsIndexPath());
- indexOutputStream.writeInt(
- ((ReduceFileMeta) diskFileInfo.getFileMeta()).getChunkOffsets().size());
- for (Long offset : ((ReduceFileMeta) diskFileInfo.getFileMeta()).getChunkOffsets()) {
- indexOutputStream.writeLong(offset);
+ long streamId =
+ super.close(
+ () -> {
+ if (!isChunkOffsetValid()) {
+ updateLastChunkOffset();
+ }
+ },
+ () -> {
+ if (diskFileInfo != null) {
+ if (diskFileInfo.isHdfs()) {
+ if (StorageManager.hadoopFs()
+ .exists(diskFileInfo.getHdfsPeerWriterSuccessPath())) {
+ StorageManager.hadoopFs().delete(diskFileInfo.getHdfsPath(), false);
+ deleted = true;
+ } else {
+ StorageManager.hadoopFs()
+ .create(diskFileInfo.getHdfsWriterSuccessPath())
+ .close();
+ FSDataOutputStream indexOutputStream =
+ StorageManager.hadoopFs().create(diskFileInfo.getHdfsIndexPath());
+ indexOutputStream.writeInt(
+ (diskFileInfo.getReduceFileMeta()).getChunkOffsets().size());
+ for (Long offset : (diskFileInfo.getReduceFileMeta()).getChunkOffsets()) {
+ indexOutputStream.writeLong(offset);
+ }
+ indexOutputStream.close();
+ }
+ }
+ } else {
+ synchronized (flushLock) {
+ // merge and free small components
+ flushBuffer.consolidate();
+ memoryFileInfo.setBuffer(flushBuffer);
+ }
}
- indexOutputStream.close();
- }
- }
- },
- () -> {});
+ },
+ () -> {});
+ return streamId;
}
}
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
index 7024ca889..6d9a44df1 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Controller.scala
@@ -31,7 +31,7 @@ import org.roaringbitmap.RoaringBitmap
import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.identity.UserIdentifier
import org.apache.celeborn.common.internal.Logging
-import org.apache.celeborn.common.meta.{WorkerInfo, WorkerPartitionLocationInfo}
+import org.apache.celeborn.common.meta.{ReduceFileMeta, WorkerInfo, WorkerPartitionLocationInfo}
import org.apache.celeborn.common.metrics.MetricsSystem
import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType, StorageInfo}
import org.apache.celeborn.common.protocol.message.ControlMessages._
@@ -200,8 +200,8 @@ private[deploy] class Controller(
logWarning(s"[handleReserveSlots] $msg, will destroy writers.")
primaryLocs.asScala.foreach { partitionLocation =>
val fileWriter = partitionLocation.asInstanceOf[WorkingPartition].getFileWriter
- fileWriter.destroy(new IOException(s"Destroy FileWriter ${fileWriter} caused by " +
- s"reserving slots failed for ${shuffleKey}."))
+ fileWriter.destroy(new IOException(s"Destroy FileWriter $fileWriter caused by " +
+ s"reserving slots failed for $shuffleKey."))
}
context.reply(ReserveSlotsResponse(StatusCode.RESERVE_SLOTS_FAILED, msg))
return
@@ -240,13 +240,13 @@ private[deploy] class Controller(
logWarning(s"[handleReserveSlots] $msg, destroy writers.")
primaryLocs.asScala.foreach { partitionLocation =>
val fileWriter = partitionLocation.asInstanceOf[WorkingPartition].getFileWriter
- fileWriter.destroy(new IOException(s"Destroy FileWriter ${fileWriter} caused by " +
- s"reserving slots failed for ${shuffleKey}."))
+ fileWriter.destroy(new IOException(s"Destroy FileWriter $fileWriter caused by " +
+ s"reserving slots failed for $shuffleKey."))
}
replicaLocs.asScala.foreach { partitionLocation =>
val fileWriter = partitionLocation.asInstanceOf[WorkingPartition].getFileWriter
- fileWriter.destroy(new IOException(s"Destroy FileWriter ${fileWriter} caused by " +
- s"reserving slots failed for ${shuffleKey}."))
+ fileWriter.destroy(new IOException(s"Destroy FileWriter $fileWriter caused by " +
+ s"reserving slots failed for $shuffleKey."))
}
context.reply(ReserveSlotsResponse(StatusCode.RESERVE_SLOTS_FAILED, msg))
return
@@ -346,14 +346,16 @@ private[deploy] class Controller(
private def waitMapPartitionRegionFinished(
fileWriter: PartitionDataWriter,
waitTimeout: Long): Unit = {
- if (fileWriter.isInstanceOf[MapPartitionDataWriter]) {
- if (fileWriter.asInstanceOf[MapPartitionDataWriter].checkPartitionRegionFinished(
- waitTimeout)) {
- logDebug(s"CommitFile succeed to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
- } else {
- logWarning(
- s"CommitFile failed to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
- }
+ fileWriter match {
+ case writer: MapPartitionDataWriter =>
+ if (writer.checkPartitionRegionFinished(
+ waitTimeout)) {
+ logDebug(s"CommitFile succeed to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
+ } else {
+ logWarning(
+ s"CommitFile failed to waitMapPartitionRegionFinished ${fileWriter.getFile.getAbsolutePath}")
+ }
+ case _ =>
}
}
@@ -424,11 +426,11 @@ private[deploy] class Controller(
commitInfo.synchronized {
if (commitInfo.status == CommitInfo.COMMIT_FINISHED) {
- logInfo(s"${shuffleKey} CommitFinished, just return the response")
+ logInfo(s"$shuffleKey CommitFinished, just return the response")
context.reply(commitInfo.response)
return
} else if (commitInfo.status == CommitInfo.COMMIT_INPROCESS) {
- logInfo(s"${shuffleKey} CommitFiles inprogress, wait for finish")
+ logInfo(s"$shuffleKey CommitFiles inprogress, wait for finish")
commitThreadPool.submit(new Runnable {
override def run(): Unit = {
waitForCommitFinish()
@@ -436,7 +438,7 @@ private[deploy] class Controller(
})
return
} else {
- logInfo(s"Start commitFiles for ${shuffleKey}")
+ logInfo(s"Start commitFiles for $shuffleKey")
commitInfo.status = CommitInfo.COMMIT_INPROCESS
workerSource.startTimer(WorkerSource.COMMIT_FILES_TIME, shuffleKey)
}
@@ -529,10 +531,10 @@ private[deploy] class Controller(
logInfo(
s"CommitFiles for $shuffleKey success with " +
s"${committedPrimaryIds.size()} committed primary partitions, " +
- s"${emptyFilePrimaryIds.size()} empty primary partitions, " +
+ s"${emptyFilePrimaryIds.size()} empty primary partitions ${emptyFilePrimaryIds.asScala.mkString(",")}, " +
s"${failedPrimaryIds.size()} failed primary partitions, " +
s"${committedReplicaIds.size()} committed replica partitions, " +
- s"${emptyFileReplicaIds.size()} empty replica partitions, " +
+ s"${emptyFileReplicaIds.size()} empty replica partitions ${emptyFileReplicaIds.asScala.mkString(",")} , " +
s"${failedReplicaIds.size()} failed replica partitions.")
CommitFilesResponse(
StatusCode.SUCCESS,
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
index e9b1a6607..934a4e967 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/FetchHandler.scala
@@ -34,15 +34,15 @@ import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.CelebornConf.MAX_CHUNKS_BEING_TRANSFERRED
import org.apache.celeborn.common.exception.CelebornIOException
import org.apache.celeborn.common.internal.Logging
-import org.apache.celeborn.common.meta.{DiskFileInfo, FileManagedBuffers, MapFileMeta, ReduceFileMeta}
-import org.apache.celeborn.common.network.buffer.NioManagedBuffer
+import org.apache.celeborn.common.meta.{DiskFileInfo, FileInfo, MapFileMeta, MemoryFileInfo, ReduceFileMeta}
+import org.apache.celeborn.common.network.buffer.{FileChunkBuffers, MemoryChunkBuffers, NioManagedBuffer}
import org.apache.celeborn.common.network.client.{RpcResponseCallback, TransportClient}
import org.apache.celeborn.common.network.protocol._
import org.apache.celeborn.common.network.server.BaseMessageHandler
import org.apache.celeborn.common.network.util.{NettyUtils, TransportConf}
import org.apache.celeborn.common.protocol.{MessageType, PbBufferStreamEnd, PbChunkFetchRequest, PbOpenStream, PbOpenStreamList, PbOpenStreamListResponse, PbReadAddCredit, PbStreamHandler, PbStreamHandlerOpt, StreamType}
import org.apache.celeborn.common.protocol.message.StatusCode
-import org.apache.celeborn.common.util.{ExceptionUtils, ThreadUtils, Utils}
+import org.apache.celeborn.common.util.{ExceptionUtils, Utils}
import org.apache.celeborn.service.deploy.worker.storage.{ChunkStreamManager, CreditStreamManager, PartitionFilesSorter, StorageManager}
class FetchHandler(
@@ -81,11 +81,11 @@ class FetchHandler(
this.registered = Some(worker.registered)
}
- def getRawDiskFileInfo(
+ def getRawFileInfo(
shuffleKey: String,
- fileName: String): DiskFileInfo = {
+ fileName: String): FileInfo = {
// find FileWriter responsible for the data
- val fileInfo = storageManager.getDiskFileInfo(shuffleKey, fileName)
+ val fileInfo = storageManager.getFileInfo(shuffleKey, fileName)
if (fileInfo == null) {
val errMsg = s"Could not find file $fileName for $shuffleKey."
logWarning(errMsg)
@@ -240,14 +240,14 @@ class FetchHandler(
s"$endIndex get file name $fileName from client channel " +
s"${NettyUtils.getRemoteAddress(client.getChannel)}")
- var fileInfo = getRawDiskFileInfo(shuffleKey, fileName)
+ var fileInfo = getRawFileInfo(shuffleKey, fileName)
val streamId = chunkStreamManager.nextStreamId()
// we must get sorted fileInfo for the following cases.
// 1. when the current request is a non-range openStream, but the original unsorted file
// has been deleted by another range's openStream request.
// 2. when the current request is a range openStream request.
- if ((endIndex != Int.MaxValue) || (endIndex == Int.MaxValue && !fileInfo.addStream(
- streamId))) {
+ if ((endIndex != Int.MaxValue) || (endIndex == Int.MaxValue
+ && !fileInfo.addStream(streamId))) {
fileInfo = partitionsSorter.getSortedFileInfo(
shuffleKey,
fileName,
@@ -255,9 +255,9 @@ class FetchHandler(
startIndex,
endIndex)
}
- val meta = fileInfo.getFileMeta.asInstanceOf[ReduceFileMeta]
+ val meta = fileInfo.getReduceFileMeta
val streamHandler =
- if (readLocalShuffle) {
+ if (readLocalShuffle && !fileInfo.isInstanceOf[MemoryFileInfo]) {
chunkStreamManager.registerStream(
streamId,
shuffleKey,
@@ -266,31 +266,45 @@ class FetchHandler(
streamId,
meta.getNumChunks,
meta.getChunkOffsets,
- fileInfo.getFilePath)
- } else if (fileInfo.isHdfs) {
- chunkStreamManager.registerStream(
- streamId,
- shuffleKey,
- fileName)
- makeStreamHandler(streamId, numChunks = 0)
- } else {
- chunkStreamManager.registerStream(
- streamId,
- shuffleKey,
- new FileManagedBuffers(fileInfo, transportConf),
- fileName,
- storageManager.getFetchTimeMetric(fileInfo.getFile))
- if (meta.getNumChunks == 0)
- logDebug(s"StreamId $streamId, fileName $fileName, mapRange " +
- s"[$startIndex-$endIndex] is empty. Received from client channel " +
- s"${NettyUtils.getRemoteAddress(client.getChannel)}")
- else logDebug(
- s"StreamId $streamId, fileName $fileName, numChunks ${meta.getNumChunks}, " +
- s"mapRange [$startIndex-$endIndex]. Received from client channel " +
- s"${NettyUtils.getRemoteAddress(client.getChannel)}")
- makeStreamHandler(
- streamId,
- meta.getNumChunks)
+ fileInfo.asInstanceOf[DiskFileInfo].getFilePath)
+ } else fileInfo match {
+ case info: DiskFileInfo if info.isHdfs =>
+ chunkStreamManager.registerStream(
+ streamId,
+ shuffleKey,
+ fileName)
+ makeStreamHandler(streamId, numChunks = 0)
+ case _ =>
+ val managedBuffer = fileInfo match {
+ case df: DiskFileInfo =>
+ new FileChunkBuffers(df, transportConf)
+ case mf: MemoryFileInfo =>
+ new MemoryChunkBuffers(mf)
+ }
+ val fetchTimeMetric =
+ fileInfo match {
+ case info: DiskFileInfo =>
+ storageManager.getFetchTimeMetric(info.getFile)
+ case _ =>
+ null
+ }
+ chunkStreamManager.registerStream(
+ streamId,
+ shuffleKey,
+ managedBuffer,
+ fileName,
+ fetchTimeMetric)
+ if (meta.getNumChunks == 0)
+ logDebug(s"StreamId $streamId, fileName $fileName, mapRange " +
+ s"[$startIndex-$endIndex] is empty. Received from client channel " +
+ s"${NettyUtils.getRemoteAddress(client.getChannel)}")
+ else logDebug(
+ s"StreamId $streamId, fileName $fileName, numChunks ${meta.getNumChunks}, " +
+ s"mapRange [$startIndex-$endIndex]. Received from client channel " +
+ s"${NettyUtils.getRemoteAddress(client.getChannel)}")
+ makeStreamHandler(
+ streamId,
+ meta.getNumChunks)
}
workerSource.incCounter(WorkerSource.OPEN_STREAM_SUCCESS_COUNT)
PbStreamHandlerOpt.newBuilder().setStreamHandler(streamHandler)
@@ -323,7 +337,7 @@ class FetchHandler(
workerSource.recordAppActiveConnection(client, shuffleKey)
workerSource.startTimer(WorkerSource.OPEN_STREAM_TIME, shuffleKey)
try {
- val fileInfo = getRawDiskFileInfo(shuffleKey, fileName)
+ val fileInfo = getRawFileInfo(shuffleKey, fileName)
fileInfo.getFileMeta match {
case _: ReduceFileMeta =>
val pbStreamHandlerOpt =
@@ -358,7 +372,7 @@ class FetchHandler(
initialCredit,
startIndex,
endIndex,
- fileInfo)
+ fileInfo.asInstanceOf[DiskFileInfo])
}
workerSource.incCounter(WorkerSource.OPEN_STREAM_SUCCESS_COUNT)
} catch {
@@ -440,10 +454,10 @@ class FetchHandler(
streamType: StreamType): Unit = {
streamType match {
case StreamType.ChunkStream =>
- val (shuffleKey, fileName) = chunkStreamManager.getShuffleKeyAndFileName(streamId)
+ val streamState = chunkStreamManager.getStreamState(streamId)
+ val (shuffleKey, fileName) = (streamState.shuffleKey, streamState.fileName)
workerSource.recordAppActiveConnection(client, shuffleKey)
- getRawDiskFileInfo(shuffleKey, fileName).closeStream(
- streamId)
+ getRawFileInfo(shuffleKey, fileName).closeStream(streamId)
case StreamType.CreditStream =>
val shuffleKey = creditStreamManager.getStreamShuffleKey(streamId)
if (shuffleKey != null) {
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
index d21a1fdea..0fe387f58 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/PushDataHandler.scala
@@ -269,7 +269,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
s"[handlePushData] FileWriter is already closed! File path ${fileWriter.getDiskFileInfo.getFilePath}")
callbackWithTimer.onFailure(new CelebornIOException("File already closed!"))
fileWriter.decrementPendingWrites()
- return;
+ return
}
val writePromise = Promise[Unit]()
// for primary, send data to replica
@@ -816,7 +816,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
s"[handleMapPartitionPushData] FileWriter is already closed! File path ${fileWriter.getDiskFileInfo.getFilePath}")
callback.onFailure(new CelebornIOException("File already closed!"))
fileWriter.decrementPendingWrites()
- return;
+ return
}
val writePromise = Promise[Unit]()
writeLocalData(Seq(fileWriter), body, shuffleKey, isPrimary, None, writePromise)
@@ -867,7 +867,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
}
private def mapPartitionRpcRequest(rpcRequest: RpcRequest)
- : Tuple8[GeneratedMessageV3, Message, Boolean, Type, Mode, String, String, Boolean] = {
+ : (GeneratedMessageV3, Message, Boolean, Type, Mode, String, String, Boolean) = {
try {
val msg = TransportMessage.fromByteBuffer(
rpcRequest.body().nioByteBuffer()).getParsedPayload.asInstanceOf[GeneratedMessageV3]
@@ -1185,7 +1185,7 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
}
private def checkDiskFull(fileWriter: PartitionDataWriter): Boolean = {
- if (fileWriter.flusher.isInstanceOf[HdfsFlusher]) {
+ if (fileWriter.flusher == null || fileWriter.flusher.isInstanceOf[HdfsFlusher]) {
return false
}
val mountPoint = fileWriter.flusher.asInstanceOf[LocalFlusher].mountPoint
@@ -1206,28 +1206,34 @@ class PushDataHandler(val workerSource: WorkerSource) extends BaseMessageHandler
|CheckDiskFullAndSplit in
|diskFull:$diskFull,
|partitionSplitMinimumSize:$partitionSplitMinimumSize,
- |splitThreshold:${fileWriter.getSplitThreshold()},
+ |splitThreshold:${fileWriter.getSplitThreshold},
|fileLength:${fileWriter.getDiskFileInfo.getFileLength}
|fileName:${fileWriter.getDiskFileInfo.getFilePath}
|""".stripMargin)
- if (workerPartitionSplitEnabled && ((diskFull && fileWriter.getDiskFileInfo.getFileLength > partitionSplitMinimumSize) ||
- (isPrimary && fileWriter.getDiskFileInfo.getFileLength > fileWriter.getSplitThreshold()))) {
- if (softSplit != null && fileWriter.getSplitMode == PartitionSplitMode.SOFT &&
- (fileWriter.getDiskFileInfo.getFileLength < partitionSplitMaximumSize)) {
- softSplit.set(true)
- } else {
- workerSource.incCounter(WorkerSource.WRITE_DATA_HARD_SPLIT_COUNT)
- callback.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
- logTrace(
- s"""
- |CheckDiskFullAndSplit hardSplit
- |diskFull:$diskFull,
- |partitionSplitMinimumSize:$partitionSplitMinimumSize,
- |splitThreshold:${fileWriter.getSplitThreshold()},
- |fileLength:${fileWriter.getDiskFileInfo.getFileLength},
- |fileName:${fileWriter.getDiskFileInfo.getFilePath}
- |""".stripMargin)
- return true
+ if (fileWriter.needHardSplitForMemoryShuffleStorage()) {
+ return true
+ }
+ val diskFileInfo = fileWriter.getDiskFileInfo
+ if (diskFileInfo != null) {
+ if (workerPartitionSplitEnabled && ((diskFull && diskFileInfo.getFileLength > partitionSplitMinimumSize) ||
+ (isPrimary && diskFileInfo.getFileLength > fileWriter.getSplitThreshold))) {
+ if (softSplit != null && fileWriter.getSplitMode == PartitionSplitMode.SOFT &&
+ (fileWriter.getDiskFileInfo.getFileLength < partitionSplitMaximumSize)) {
+ softSplit.set(true)
+ } else {
+ workerSource.incCounter(WorkerSource.WRITE_DATA_HARD_SPLIT_COUNT)
+ callback.onSuccess(ByteBuffer.wrap(Array[Byte](StatusCode.HARD_SPLIT.getValue)))
+ logTrace(
+ s"""
+ |CheckDiskFullAndSplit hardSplit
+ |diskFull:$diskFull,
+ |partitionSplitMinimumSize:$partitionSplitMinimumSize,
+ |splitThreshold:${fileWriter.getSplitThreshold},
+ |fileLength:${diskFileInfo.getFileLength},
+ |fileName:${diskFileInfo.getFilePath}
+ |""".stripMargin)
+ return true
+ }
}
}
false
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
index 07d29e465..f3039fc20 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/Worker.scala
@@ -170,7 +170,7 @@ private[celeborn] class Worker(
val storageManager = new StorageManager(conf, workerSource)
- val memoryManager: MemoryManager = MemoryManager.initialize(conf)
+ val memoryManager: MemoryManager = MemoryManager.initialize(conf, storageManager)
memoryManager.registerMemoryListener(storageManager)
val partitionsSorter = new PartitionFilesSorter(memoryManager, conf, workerSource)
@@ -386,11 +386,24 @@ private[celeborn] class Worker(
workerSource.addGauge(WorkerSource.READ_BUFFER_ALLOCATED_COUNT) { () =>
memoryManager.getAllocatedReadBuffers
}
+ workerSource.addGauge(WorkerSource.MEMORY_FILE_STORAGE_SIZE) { () =>
+ memoryManager.getMemoryFileStorageCounter
+ }
+ workerSource.addGauge(WorkerSource.DIRECT_MEMORY_USAGE_RATIO) { () =>
+ memoryManager.workerMemoryUsageRatio()
+ }
+ workerSource.addGauge(WorkerSource.EVICTED_FILE_COUNT) { () =>
+ storageManager.evictedFileCount.get()
+ }
+ workerSource.addGauge(WorkerSource.MEMORY_STORAGE_FILE_COUNT) { () =>
+ storageManager.memoryWriters.size()
+ }
+
workerSource.addGauge(WorkerSource.ACTIVE_SHUFFLE_SIZE) { () =>
- storageManager.getActiveShuffleSize()
+ storageManager.getActiveShuffleSize
}
workerSource.addGauge(WorkerSource.ACTIVE_SHUFFLE_FILE_COUNT) { () =>
- storageManager.getActiveShuffleFileCount()
+ storageManager.getActiveShuffleFileCount
}
workerSource.addGauge(WorkerSource.PAUSE_PUSH_DATA_TIME) { () =>
memoryManager.getPausePushDataTime
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
index 6be2a9256..726f7acec 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/WorkerSource.scala
@@ -187,6 +187,11 @@ object WorkerSource {
val BUFFER_STREAM_READ_BUFFER = "BufferStreamReadBuffer"
val READ_BUFFER_DISPATCHER_REQUESTS_LENGTH = "ReadBufferDispatcherRequestsLength"
val READ_BUFFER_ALLOCATED_COUNT = "ReadBufferAllocatedCount"
+ val MEMORY_FILE_STORAGE_SIZE = "MemoryFileStorageSize"
+ val DIRECT_MEMORY_USAGE_RATIO = "DirectMemoryUsageRatio"
+ val EVICTED_FILE_COUNT = "EvictedFileCount"
+
+ val MEMORY_STORAGE_FILE_COUNT = "MemoryStorageFileCount"
// credit
val ACTIVE_CREDIT_STREAM_COUNT = "ActiveCreditStreamCount"
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
index 6c6bc64aa..607e47ddc 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/FlushTask.scala
@@ -24,14 +24,16 @@ import org.apache.hadoop.fs.Path
abstract private[worker] class FlushTask(
val buffer: CompositeByteBuf,
- val notifier: FlushNotifier) {
+ val notifier: FlushNotifier,
+ val keepBuffer: Boolean) {
def flush(): Unit
}
private[worker] class LocalFlushTask(
buffer: CompositeByteBuf,
fileChannel: FileChannel,
- notifier: FlushNotifier) extends FlushTask(buffer, notifier) {
+ notifier: FlushNotifier,
+ keepBuffer: Boolean) extends FlushTask(buffer, notifier, keepBuffer) {
override def flush(): Unit = {
val buffers = buffer.nioBuffers()
for (buffer <- buffers) {
@@ -45,7 +47,8 @@ private[worker] class LocalFlushTask(
private[worker] class HdfsFlushTask(
buffer: CompositeByteBuf,
val path: Path,
- notifier: FlushNotifier) extends FlushTask(buffer, notifier) {
+ notifier: FlushNotifier,
+ keepBuffer: Boolean) extends FlushTask(buffer, notifier, keepBuffer) {
override def flush(): Unit = {
val hdfsStream = StorageManager.hadoopFs.append(path, 256 * 1024)
hdfsStream.write(ByteBufUtil.getBytes(buffer))
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala
index ffee36f0c..2ac06b829 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/Flusher.scala
@@ -77,17 +77,19 @@ abstract private[worker] class Flusher(
}
} catch {
case t: Throwable =>
- if (t.isInstanceOf[IOException]) {
- task.notifier.setException(t.asInstanceOf[IOException])
- processIOException(
- t.asInstanceOf[IOException],
- DiskStatus.READ_OR_WRITE_FAILURE)
+ t match {
+ case exception: IOException =>
+ task.notifier.setException(exception)
+ processIOException(
+ exception,
+ DiskStatus.READ_OR_WRITE_FAILURE)
+ case _ =>
}
- logWarning(s"Flusher-${this}-thread-${index} encounter exception.", t)
+ logWarning(s"Flusher-$this-thread-$index encounter exception.", t)
}
lastBeginFlushTime.set(index, -1)
}
- Utils.tryLogNonFatalError(returnBuffer(task.buffer))
+ Utils.tryLogNonFatalError(returnBuffer(task.buffer, task.keepBuffer))
task.notifier.numPendingFlushes.decrementAndGet()
}
}
@@ -97,11 +99,6 @@ abstract private[worker] class Flusher(
ThreadPoolSource.registerSource(s"$this", workers)
}
- def getWorkerIndex: Int = synchronized {
- nextWorkerIndex = (nextWorkerIndex + 1) % threadCount
- nextWorkerIndex
- }
-
def takeBuffer(): CompositeByteBuf = {
var buffer = bufferQueue.poll()
if (buffer == null) {
@@ -110,15 +107,19 @@ abstract private[worker] class Flusher(
buffer
}
- def returnBuffer(buffer: CompositeByteBuf): Unit = {
- MemoryManager.instance().releaseDiskBuffer(buffer.readableBytes())
+ def returnBuffer(buffer: CompositeByteBuf, keepBuffer: Boolean = false): Unit = {
+ val bufferSize = buffer.readableBytes()
+ MemoryManager.instance().releaseDiskBuffer(bufferSize)
Option(CongestionController.instance())
.foreach(
- _.consumeBytes(buffer.readableBytes()))
+ _.consumeBytes(bufferSize))
buffer.removeComponents(0, buffer.numComponents())
buffer.clear()
-
- bufferQueue.put(buffer)
+ if (keepBuffer) {
+ bufferQueue.put(buffer)
+ } else {
+ buffer.release()
+ }
}
def addTask(task: FlushTask, timeoutMs: Long, workerIndex: Int): Boolean = {
diff --git a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
index 82e0f2ef1..857939929 100644
--- a/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
+++ b/worker/src/main/scala/org/apache/celeborn/service/deploy/worker/storage/StorageManager.scala
@@ -22,12 +22,13 @@ import java.nio.charset.StandardCharsets
import java.nio.file.{FileAlreadyExistsException, Files, Paths}
import java.util
import java.util.concurrent.{ConcurrentHashMap, ScheduledExecutorService, ThreadPoolExecutor, TimeUnit}
-import java.util.concurrent.atomic.AtomicInteger
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
import java.util.function.{BiConsumer, IntUnaryOperator}
import scala.collection.JavaConverters._
import scala.concurrent.duration._
+import com.google.common.annotations.VisibleForTesting
import io.netty.buffer.PooledByteBufAllocator
import org.apache.commons.io.FileUtils
import org.apache.hadoop.fs.{FileSystem, Path}
@@ -37,22 +38,32 @@ import org.apache.celeborn.common.CelebornConf
import org.apache.celeborn.common.exception.CelebornException
import org.apache.celeborn.common.identity.UserIdentifier
import org.apache.celeborn.common.internal.Logging
-import org.apache.celeborn.common.meta.{DeviceInfo, DiskFileInfo, DiskInfo, DiskStatus, MapFileMeta, ReduceFileMeta, TimeWindow}
+import org.apache.celeborn.common.meta.{DeviceInfo, DiskFileInfo, DiskInfo, DiskStatus, FileInfo, MapFileMeta, MemoryFileInfo, ReduceFileMeta, TimeWindow}
import org.apache.celeborn.common.metrics.source.{AbstractSource, ThreadPoolSource}
import org.apache.celeborn.common.network.util.{NettyUtils, TransportConf}
import org.apache.celeborn.common.protocol.{PartitionLocation, PartitionSplitMode, PartitionType, StorageInfo}
import org.apache.celeborn.common.quota.ResourceConsumption
import org.apache.celeborn.common.util.{CelebornExitKind, CelebornHadoopUtils, JavaUtils, PbSerDeUtils, ThreadUtils, Utils}
import org.apache.celeborn.service.deploy.worker._
+import org.apache.celeborn.service.deploy.worker.memory.MemoryManager
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager.MemoryPressureListener
import org.apache.celeborn.service.deploy.worker.shuffledb.{DB, DBBackend, DBProvider}
import org.apache.celeborn.service.deploy.worker.storage.StorageManager.hadoopFs
final private[worker] class StorageManager(conf: CelebornConf, workerSource: AbstractSource)
extends ShuffleRecoverHelper with DeviceObserver with Logging with MemoryPressureListener {
+ // fileInfos and partitionDataWriters are one to one mapping
// mount point -> file writer
val workingDirWriters =
JavaUtils.newConcurrentHashMap[File, ConcurrentHashMap[String, PartitionDataWriter]]()
+ val hdfsWriters = JavaUtils.newConcurrentHashMap[String, PartitionDataWriter]()
+ val memoryWriters = JavaUtils.newConcurrentHashMap[MemoryFileInfo, PartitionDataWriter]()
+ // (shuffleKey->(filename->DiskFileInfo))
+ private val diskFileInfos =
+ JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[String, DiskFileInfo]]()
+ // (shuffleKey->(filename->MemoryFileInfo))
+ val memoryFileInfos =
+ JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[String, MemoryFileInfo]]()
val hasHDFSStorage = conf.hasHDFSStorage
@@ -107,21 +118,22 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
private val deviceMonitor =
DeviceMonitor.createDeviceMonitor(conf, this, deviceInfos, tmpDiskInfos, workerSource)
- private val byteBufAllocator: PooledByteBufAllocator =
+ val storageBufferAllocator: PooledByteBufAllocator =
NettyUtils.getPooledByteBufAllocator(new TransportConf("StorageManager", conf), null, true)
+
// (mountPoint -> LocalFlusher)
private val (
localFlushers: ConcurrentHashMap[String, LocalFlusher],
_totalLocalFlusherThread: Int) = {
val flushers = JavaUtils.newConcurrentHashMap[String, LocalFlusher]()
- var totalThread = 0;
+ var totalThread = 0
disksSnapshot().foreach { diskInfo =>
if (!flushers.containsKey(diskInfo.mountPoint)) {
val flusher = new LocalFlusher(
workerSource,
deviceMonitor,
diskInfo.threadCount,
- byteBufAllocator,
+ storageBufferAllocator,
conf.workerPushMaxComponents,
diskInfo.mountPoint,
diskInfo.storageType,
@@ -137,10 +149,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
val hdfsDir = conf.hdfsDir
val hdfsPermission = new FsPermission("755")
- val hdfsWriters = JavaUtils.newConcurrentHashMap[String, PartitionDataWriter]()
val (hdfsFlusher, _totalHdfsFlusherThread) =
if (hasHDFSStorage) {
- logInfo(s"Initialize HDFS support with path ${hdfsDir}")
+ logInfo(s"Initialize HDFS support with path $hdfsDir")
try {
StorageManager.hadoopFs = CelebornHadoopUtils.getHadoopFS(conf)
} catch {
@@ -152,7 +163,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
Some(new HdfsFlusher(
workerSource,
conf.workerHdfsFlusherThreads,
- byteBufAllocator,
+ storageBufferAllocator,
conf.workerPushMaxComponents)),
conf.workerHdfsFlusherThreads)
} else {
@@ -160,6 +171,12 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
def totalFlusherThread: Int = _totalLocalFlusherThread + _totalHdfsFlusherThread
+ val activeTypes = conf.availableStorageTypes
+
+ def localOrHdfsStorageAvailable(): Boolean = {
+ StorageInfo.HDFSAvailable(activeTypes) || StorageInfo.localDiskAvailable(
+ activeTypes) || hdfsDir.nonEmpty || !diskInfos.isEmpty
+ }
override def notifyError(mountPoint: String, diskStatus: DiskStatus): Unit = this.synchronized {
if (diskStatus == DiskStatus.CRITICAL_ERROR) {
@@ -193,11 +210,9 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
// shuffleKey -> (fileName -> file info)
- private val diskFileInfos =
- JavaUtils.newConcurrentHashMap[String, ConcurrentHashMap[String, DiskFileInfo]]()
private val RECOVERY_FILE_NAME_PREFIX = "recovery"
private var RECOVERY_FILE_NAME = "recovery.ldb"
- private var db: DB = null
+ private var db: DB = _
private var saveCommittedFileInfosExecutor: ScheduledExecutorService = _
private val saveCommittedFileInfoBySyncMethod =
conf.workerGracefulShutdownSaveCommittedFileInfoSync
@@ -293,12 +308,18 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
logDebug(s"Update FileInfos into DB: $shuffleKey -> $files")
} catch {
case exception: Exception =>
- logError(s"Update FileInfos into DB: ${shuffleKey} failed.", exception)
+ logError(s"Update FileInfos into DB: $shuffleKey failed.", exception)
}
}
}
- private def getNextIndex() = counter.getAndUpdate(counterOperator)
+ def saveAllMemoryStorageFileInfosToDB(): Unit = {
+ for (writer <- memoryWriters.asScala) {
+ Utils.tryLogNonFatalError(writer._2.evict(false))
+ }
+ }
+
+ private def getNextIndex = counter.getAndUpdate(counterOperator)
private val newMapFunc =
new java.util.function.Function[String, ConcurrentHashMap[String, DiskFileInfo]]() {
@@ -312,12 +333,21 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
JavaUtils.newConcurrentHashMap()
}
+ private val memoryFileInfoMapFunc =
+ new java.util.function.Function[String, ConcurrentHashMap[String, MemoryFileInfo]]() {
+ override def apply(key: String): ConcurrentHashMap[String, MemoryFileInfo] =
+ JavaUtils.newConcurrentHashMap()
+ }
+
private val workingDirWriterListFunc =
new java.util.function.Function[File, ConcurrentHashMap[String, PartitionDataWriter]]() {
override def apply(t: File): ConcurrentHashMap[String, PartitionDataWriter] =
JavaUtils.newConcurrentHashMap[String, PartitionDataWriter]()
}
+ @VisibleForTesting
+ val evictedFileCount = new AtomicLong
+
@throws[IOException]
def createPartitionDataWriter(
appId: String,
@@ -354,40 +384,32 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
if (healthyWorkingDirs().size <= 0 && !hasHDFSStorage) {
throw new IOException("No available working dirs!")
}
- val shuffleKey = Utils.makeShuffleKey(appId, shuffleId)
- val (flusher, diskFileInfo, workingDir) = createFile(
+ val partitionDataWriterContext = new PartitionDataWriterContext(
+ splitThreshold,
+ splitMode,
+ rangeReadFilter,
location,
appId,
shuffleId,
- location.getFileName,
userIdentifier,
partitionType,
partitionSplitEnabled)
+
val writer =
try {
partitionType match {
case PartitionType.MAP => new MapPartitionDataWriter(
this,
- diskFileInfo,
- flusher,
workerSource,
conf,
deviceMonitor,
- splitThreshold,
- splitMode,
- rangeReadFilter,
- shuffleKey)
+ partitionDataWriterContext)
case PartitionType.REDUCE => new ReducePartitionDataWriter(
this,
- diskFileInfo,
- flusher,
workerSource,
conf,
deviceMonitor,
- splitThreshold,
- splitMode,
- rangeReadFilter,
- shuffleKey)
+ partitionDataWriterContext)
case _ => throw new UnsupportedOperationException(s"Not support $partitionType yet")
}
} catch {
@@ -395,24 +417,58 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
logError("Create partition data writer failed", e)
throw e
}
- if (!(writer.getDiskFileInfo.isHdfs)) {
- deviceMonitor.registerFileWriter(writer)
- workingDirWriters.computeIfAbsent(workingDir, workingDirWriterListFunc).put(
- diskFileInfo.getFilePath,
- writer)
- } else {
- hdfsWriters.put(diskFileInfo.getFilePath, writer)
- }
writer
}
- def getDiskFileInfo(shuffleKey: String, fileName: String): DiskFileInfo = {
- val shuffleMap = diskFileInfos.get(shuffleKey)
- if (shuffleMap ne null) {
- shuffleMap.get(fileName)
- } else {
- null
+ def registerMemoryPartitionWriter(writer: PartitionDataWriter, fileInfo: MemoryFileInfo): Unit = {
+ memoryWriters.put(fileInfo, writer)
+ }
+
+ def unregisterMemoryPartitionWriterAndFileInfo(
+ fileInfo: MemoryFileInfo,
+ shuffleKey: String,
+ fileName: String): Unit = {
+ memoryWriters.remove(fileInfo)
+ val map = memoryFileInfos.get(shuffleKey)
+ if (map != null) {
+ map.remove(fileName)
+ }
+ }
+
+ def registerDiskFilePartitionWriter(
+ writer: PartitionDataWriter,
+ workingDir: File,
+ fileInfo: DiskFileInfo): Unit = {
+ if (writer.getDiskFileInfo.isHdfs) {
+ hdfsWriters.put(fileInfo.getFilePath, writer)
+ return
+ }
+ deviceMonitor.registerFileWriter(writer)
+ workingDirWriters.computeIfAbsent(workingDir, workingDirWriterListFunc).put(
+ fileInfo.getFilePath,
+ writer)
+ }
+
+ def getFileInfo(
+ shuffleKey: String,
+ fileName: String): FileInfo = {
+ val memoryShuffleMap = memoryFileInfos.get(shuffleKey)
+ if (memoryShuffleMap != null) {
+ val memoryFileInfo = memoryShuffleMap.get(fileName)
+ if (memoryFileInfo != null) {
+ return memoryFileInfo
+ }
}
+
+ val diskShuffleMap = diskFileInfos.get(shuffleKey)
+ if (diskShuffleMap != null) {
+ val diskFileInfo = diskShuffleMap.get(fileName)
+ if (diskFileInfo != null) {
+ return diskFileInfo
+ }
+ }
+
+ null
}
def getFetchTimeMetric(file: File): TimeWindow = {
@@ -443,34 +499,40 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
def cleanFile(shuffleKey: String, fileName: String): Unit = {
- val fileInfo = getDiskFileInfo(shuffleKey, fileName)
- if (fileInfo != null) {
- cleanFileInternal(shuffleKey, fileInfo)
- }
+ cleanFileInternal(shuffleKey, getFileInfo(shuffleKey, fileName))
}
- def cleanFileInternal(shuffleKey: String, fileInfo: DiskFileInfo): Boolean = {
+ def cleanFileInternal(shuffleKey: String, fileInfo: FileInfo): Boolean = {
+ if (fileInfo == null) return false
var isHdfsExpired = false
- if (fileInfo.isHdfs) {
- isHdfsExpired = true
- val hdfsFileWriter = hdfsWriters.get(fileInfo.getFilePath)
- if (hdfsFileWriter != null) {
- hdfsFileWriter.destroy(new IOException(
- s"Destroy FileWriter $hdfsFileWriter caused by shuffle $shuffleKey expired."))
- hdfsWriters.remove(fileInfo.getFilePath)
- }
- } else {
- val workingDir =
- fileInfo.getFile.getParentFile.getParentFile.getParentFile
- val writers = workingDirWriters.get(workingDir)
- if (writers != null) {
- val fileWriter = writers.get(fileInfo.getFilePath)
- if (fileWriter != null) {
- fileWriter.destroy(new IOException(
- s"Destroy FileWriter $fileWriter caused by shuffle $shuffleKey expired."))
- writers.remove(fileInfo.getFilePath)
+ fileInfo match {
+ case info: DiskFileInfo =>
+ if (info.isHdfs) {
+ isHdfsExpired = true
+ val hdfsFileWriter = hdfsWriters.get(info.getFilePath)
+ if (hdfsFileWriter != null) {
+ hdfsFileWriter.destroy(new IOException(
+ s"Destroy FileWriter $hdfsFileWriter caused by shuffle $shuffleKey expired."))
+ hdfsWriters.remove(info.getFilePath)
+ }
+ } else {
+ val workingDir =
+ info.getFile.getParentFile.getParentFile.getParentFile
+ val writers = workingDirWriters.get(workingDir)
+ if (writers != null) {
+ val fileWriter = writers.get(info.getFilePath)
+ if (fileWriter != null) {
+ fileWriter.destroy(new IOException(
+ s"Destroy FileWriter $fileWriter caused by shuffle $shuffleKey expired."))
+ writers.remove(info.getFilePath)
+ }
+ }
}
- }
+ case mInfo: MemoryFileInfo =>
+ val memoryWriter = memoryWriters.remove(mInfo)
+ memoryWriter.destroy(new IOException(
+ s"Destroy FileWriter $memoryWriter caused by shuffle $shuffleKey expired."))
+ case _ =>
}
isHdfsExpired
@@ -517,6 +579,13 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
}
}
}
+ if (memoryFileInfos.containsKey(shuffleKey)) {
+ val memoryFileMaps = memoryFileInfos.remove(shuffleKey)
+ memoryFileMaps.asScala.foreach(u => {
+ cleanFileInternal(shuffleKey, u._2)
+ MemoryManager.instance().releaseMemoryFileStorage(u._2.releaseMemoryBuffers())
+ })
+ }
}
}
@@ -630,6 +699,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
if (db != null) {
if (exitKind == CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN) {
try {
+ saveAllMemoryStorageFileInfosToDB()
saveAllCommittedFileInfosToDB()
db.close()
} catch {
@@ -783,15 +853,83 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
committedFileInfos.computeIfAbsent(shuffleKey, newMapFunc).put(fileName, fileInfo)
}
- def getActiveShuffleSize(): Long = {
- diskFileInfos.values().asScala.map(_.values().asScala.map(_.getBytesFlushed).sum).sum
+ def getActiveShuffleSize: Long = {
+ diskFileInfos.values().asScala.map(_.values().asScala.map(_.getFileLength).sum).sum
}
- def getActiveShuffleFileCount(): Long = {
+ def getActiveShuffleFileCount: Long = {
diskFileInfos.asScala.values.map(_.size()).sum
}
def createFile(
+ partitionDataWriterContext: PartitionDataWriterContext,
+ useMemoryShuffle: Boolean): (MemoryFileInfo, Flusher, DiskFileInfo, File) = {
+ logDebug(
+ s"create file for ${partitionDataWriterContext.getShuffleKey} ${partitionDataWriterContext.getPartitionLocation.getFileName}")
+ val location = partitionDataWriterContext.getPartitionLocation
+ if (useMemoryShuffle
+ && location.getStorageInfo.memoryAvailable()
+ && MemoryManager.instance().memoryFileStorageAvailable()) {
+ logDebug(s"Create memory file for ${partitionDataWriterContext.getShuffleKey} ${partitionDataWriterContext.getPartitionLocation.getFileName}")
+ (
+ createMemoryFile(
+ partitionDataWriterContext.getAppId,
+ partitionDataWriterContext.getShuffleId,
+ location.getFileName,
+ partitionDataWriterContext.getUserIdentifier,
+ partitionDataWriterContext.getPartitionType,
+ partitionDataWriterContext.isPartitionSplitEnabled),
+ null,
+ null,
+ null)
+ } else if (location.getStorageInfo.localDiskAvailable() || location.getStorageInfo.HDFSAvailable()) {
+ logDebug(s"create non-memory file for ${partitionDataWriterContext.getShuffleKey} ${partitionDataWriterContext.getPartitionLocation.getFileName}")
+ val createDiskFileResult = createDiskFile(
+ location,
+ partitionDataWriterContext.getAppId,
+ partitionDataWriterContext.getShuffleId,
+ location.getFileName,
+ partitionDataWriterContext.getUserIdentifier,
+ partitionDataWriterContext.getPartitionType,
+ partitionDataWriterContext.isPartitionSplitEnabled)
+ (null, createDiskFileResult._1, createDiskFileResult._2, createDiskFileResult._3)
+ } else {
+ (null, null, null, null)
+ }
+ }
+
+ def createMemoryFile(
+ appId: String,
+ shuffleId: Int,
+ fileName: String,
+ userIdentifier: UserIdentifier,
+ partitionType: PartitionType,
+ partitionSplitEnabled: Boolean): MemoryFileInfo = {
+ val fileMeta = partitionType match {
+ case PartitionType.REDUCE =>
+ new ReduceFileMeta(conf.shuffleChunkSize)
+ case PartitionType.MAP =>
+ new MapFileMeta()
+ case PartitionType.MAPGROUP =>
+ throw new NotImplementedError("Map group is not implemented")
+ }
+ val shuffleKey = Utils.makeShuffleKey(appId, shuffleId)
+ val memoryFileInfo =
+ new MemoryFileInfo(
+ userIdentifier,
+ partitionSplitEnabled,
+ fileMeta)
+ logDebug(s"create memory file for ${shuffleKey} ${fileName} and put it int memoryFileInfos")
+ memoryFileInfos.computeIfAbsent(shuffleKey, memoryFileInfoMapFunc).put(
+ fileName,
+ memoryFileInfo)
+ memoryFileInfo
+ }
+
+ /**
+ * @return (Flusher,DiskFileInfo,workingDir)
+ */
+ def createDiskFile(
location: PartitionLocation,
appId: String,
shuffleId: Int,
@@ -825,7 +963,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
val hdfsFileInfo = new DiskFileInfo(
userIdentifier,
partitionSplitEnabled,
- new ReduceFileMeta(),
+ new ReduceFileMeta(conf.shuffleChunkSize),
hdfsFilePath,
StorageInfo.Type.HDFS)
diskFileInfos.computeIfAbsent(shuffleKey, diskFileInfoMapFunc).put(
@@ -833,7 +971,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
hdfsFileInfo)
return (hdfsFlusher.get, hdfsFileInfo, null)
} else if (dirs.nonEmpty && location.getStorageInfo.localDiskAvailable()) {
- val dir = dirs(getNextIndex() % dirs.size)
+ val dir = dirs(getNextIndex % dirs.size)
val mountPoint = DeviceInfo.getMountPoint(dir.getAbsolutePath, mountPoints)
val shuffleDir = new File(dir, s"$appId/$shuffleId")
shuffleDir.mkdirs()
@@ -852,7 +990,7 @@ final private[worker] class StorageManager(conf: CelebornConf, workerSource: Abs
val filePath = file.getAbsolutePath
val fileMeta = partitionType match {
case PartitionType.REDUCE =>
- new ReduceFileMeta()
+ new ReduceFileMeta(conf.shuffleChunkSize)
case PartitionType.MAP =>
new MapFileMeta()
case PartitionType.MAPGROUP =>
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
index ef4f172bf..b54a0b99a 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/FetchHandlerSuiteJ.java
@@ -18,7 +18,6 @@
package org.apache.celeborn.service.deploy.worker;
import static org.junit.Assert.*;
-import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.*;
import java.io.File;
@@ -47,7 +46,6 @@ import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.FileInfo;
-import org.apache.celeborn.common.meta.ReduceFileMeta;
import org.apache.celeborn.common.network.buffer.NioManagedBuffer;
import org.apache.celeborn.common.network.client.RpcResponseCallback;
import org.apache.celeborn.common.network.client.TransportClient;
@@ -80,7 +78,7 @@ public class FetchHandlerSuiteJ {
private static final Logger LOG = LoggerFactory.getLogger(FetchHandlerSuiteJ.class);
- private static CelebornConf conf = new CelebornConf();
+ private static final CelebornConf conf = new CelebornConf();
private static final Random random = new Random();
@@ -93,7 +91,7 @@ public class FetchHandlerSuiteJ {
byte[] batchHeader = new byte[16];
File shuffleFile = File.createTempFile("celeborn", UUID.randomUUID().toString());
- FileInfo fileInfo = new DiskFileInfo(shuffleFile, userIdentifier);
+ DiskFileInfo fileInfo = new DiskFileInfo(shuffleFile, userIdentifier, conf);
FileOutputStream fileOutputStream = new FileOutputStream(shuffleFile);
FileChannel channel = fileOutputStream.getChannel();
Map<Integer, Integer> batchIds = new HashMap<>();
@@ -132,10 +130,10 @@ public class FetchHandlerSuiteJ {
for (long offset = conf.shuffleChunkSize();
offset <= originFileLen;
offset += conf.shuffleChunkSize()) {
- ((ReduceFileMeta) fileInfo.getFileMeta()).getChunkOffsets().add(offset);
+ (fileInfo.getReduceFileMeta()).addChunkOffset(offset);
}
// update sorted fileInfo chunk offsets
- ((DiskFileInfo) fileInfo).updateBytesFlushed(originFileLen);
+ fileInfo.updateBytesFlushed(originFileLen);
return fileInfo;
}
@@ -319,7 +317,7 @@ public class FetchHandlerSuiteJ {
Mockito.doReturn(partitionFilesSorter).when(worker).partitionsSorter();
fetchHandler0.init(worker);
FetchHandler fetchHandler = spy(fetchHandler0);
- Mockito.doReturn(fileInfo).when(fetchHandler).getRawDiskFileInfo(anyString(), anyString());
+ Mockito.doReturn(fileInfo).when(fetchHandler).getRawFileInfo(anyString(), anyString());
return fetchHandler;
}
@@ -443,7 +441,7 @@ public class FetchHandlerSuiteJ {
}
private void checkOriginFileBeDeleted(FileInfo fileInfo) {
- assertTrue(!((DiskFileInfo) fileInfo).getFilePath().endsWith(".sorted"));
+ assertFalse(((DiskFileInfo) fileInfo).getFilePath().endsWith(".sorted"));
long startTs = System.currentTimeMillis();
boolean deleted = false;
long timeout = 5 * 1000; // 5s
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManagerSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManagerSuiteJ.java
index 515b176d5..5a055b7b4 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManagerSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ChunkStreamManagerSuiteJ.java
@@ -24,7 +24,7 @@ import org.junit.Assert;
import org.junit.Test;
import org.mockito.Mockito;
-import org.apache.celeborn.common.meta.FileManagedBuffers;
+import org.apache.celeborn.common.network.buffer.FileChunkBuffers;
public class ChunkStreamManagerSuiteJ {
@Test
@@ -32,12 +32,12 @@ public class ChunkStreamManagerSuiteJ {
ChunkStreamManager manager = new ChunkStreamManager();
@SuppressWarnings("unchecked")
- FileManagedBuffers buffers = Mockito.mock(FileManagedBuffers.class);
+ FileChunkBuffers buffers = Mockito.mock(FileChunkBuffers.class);
@SuppressWarnings("unchecked")
- FileManagedBuffers buffers2 = Mockito.mock(FileManagedBuffers.class);
- FileManagedBuffers buffers3 = Mockito.mock(FileManagedBuffers.class);
- FileManagedBuffers buffers4 = Mockito.mock(FileManagedBuffers.class);
+ FileChunkBuffers buffers2 = Mockito.mock(FileChunkBuffers.class);
+ FileChunkBuffers buffers3 = Mockito.mock(FileChunkBuffers.class);
+ FileChunkBuffers buffers4 = Mockito.mock(FileChunkBuffers.class);
manager.registerStream("shuffleKey1", buffers, "shuffleFile1", null);
manager.registerStream("shuffleKey1", buffers2, "shuffleFile1", null);
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java
index 0896f2462..95064b90e 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/CreditStreamManagerSuiteJ.java
@@ -46,16 +46,16 @@ public class CreditStreamManagerSuiteJ {
private static final Logger LOG = LoggerFactory.getLogger(CreditStreamManagerSuiteJ.class);
private static File tempDir =
Utils.createTempDir(System.getProperty("java.io.tmpdir"), "celeborn");
+ private static CelebornConf conf = new CelebornConf();
@BeforeClass
public static void beforeAll() {
- CelebornConf conf = new CelebornConf();
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(), "0.8");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_REPLICATE().key(), "0.9");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
- conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
@@ -77,7 +77,7 @@ public class CreditStreamManagerSuiteJ {
Channel channel = Mockito.mock(Channel.class);
DiskFileInfo diskFileInfo =
new DiskFileInfo(
- createTemporaryFileWithIndexFile(), new UserIdentifier("default", "default"));
+ createTemporaryFileWithIndexFile(), new UserIdentifier("default", "default"), conf);
MapFileMeta mapFileMeta = new MapFileMeta(1024, 10);
diskFileInfo.replaceFileMeta(mapFileMeta);
Consumer<Long> streamIdConsumer = streamId -> Assert.assertTrue(streamId > 0);
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterSuiteUtils.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterSuiteUtils.java
index e5926ad10..44baa691e 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterSuiteUtils.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionDataWriterSuiteUtils.java
@@ -20,15 +20,25 @@ package org.apache.celeborn.service.deploy.worker.storage;
import java.io.File;
import java.io.IOException;
import java.util.UUID;
+import java.util.concurrent.atomic.AtomicInteger;
-import scala.Tuple3;
+import scala.Function0;
import scala.Tuple4;
+import org.apache.hadoop.shaded.com.google.common.io.Files;
import org.mockito.Mockito;
+import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.meta.MapFileMeta;
+import org.apache.celeborn.common.meta.MemoryFileInfo;
+import org.apache.celeborn.common.meta.ReduceFileMeta;
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+import org.apache.celeborn.common.network.util.NettyUtils;
+import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.protocol.StorageInfo;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
public class PartitionDataWriterSuiteUtils {
public static File getTemporaryFile(File tempDir) throws IOException {
@@ -38,25 +48,100 @@ public class PartitionDataWriterSuiteUtils {
return temporaryFile;
}
- public static Tuple4<StorageManager, Flusher, DiskFileInfo, File> prepareTestFileContext(
- File tempDir, UserIdentifier userIdentifier, Flusher flusher, boolean reduceMeta)
+ public static StorageManager prepareDiskFileTestEnvironment(
+ File tempDir,
+ UserIdentifier userIdentifier,
+ Flusher flusher,
+ boolean reduceMeta,
+ CelebornConf conf)
throws IOException {
File file = getTemporaryFile(tempDir);
- DiskFileInfo fileInfo = new DiskFileInfo(file, userIdentifier);
+ DiskFileInfo fileInfo = new DiskFileInfo(file, userIdentifier, conf);
if (!reduceMeta) {
fileInfo.replaceFileMeta(new MapFileMeta(32 * 1024, 10));
}
StorageManager storageManager = Mockito.mock(StorageManager.class);
- Mockito.doAnswer(invocation -> new Tuple3<Flusher, DiskFileInfo, File>(flusher, fileInfo, null))
+ Mockito.doAnswer(
+ invocation ->
+ new Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File>(
+ null, flusher, fileInfo, file))
.when(storageManager)
- .createFile(
- Mockito.any(),
- Mockito.any(),
- Mockito.anyInt(),
- Mockito.any(),
- Mockito.any(),
- Mockito.any(),
- Mockito.anyBoolean());
- return new Tuple4<>(storageManager, flusher, fileInfo, null);
+ .createFile(Mockito.any(), Mockito.anyBoolean());
+ return storageManager;
+ }
+
+ public static StorageManager prepareMemoryFileTestEnvironment(
+ UserIdentifier userIdentifier,
+ boolean reduceMeta,
+ StorageManager storageManager,
+ CelebornConf celebornConf) {
+ ReduceFileMeta reduceFileMeta = new ReduceFileMeta(celebornConf.shuffleChunkSize());
+ MemoryFileInfo memoryFileInfo = new MemoryFileInfo(userIdentifier, false, reduceFileMeta);
+ if (!reduceMeta) {
+ memoryFileInfo.replaceFileMeta(new MapFileMeta(32 * 1024, 10));
+ }
+
+ Mockito.doAnswer(
+ invocation ->
+ new Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File>(
+ memoryFileInfo, null, null, null))
+ .when(storageManager)
+ .createFile(Mockito.any(), Mockito.anyBoolean());
+ return storageManager;
+ }
+
+ public static StorageManager prepareMemoryEvictEnvironment(
+ UserIdentifier userIdentifier,
+ boolean reduceMeta,
+ StorageManager storageManager,
+ CelebornConf celebornConf)
+ throws IOException {
+ ReduceFileMeta reduceFileMeta = new ReduceFileMeta(celebornConf.shuffleChunkSize());
+ MemoryFileInfo memoryFileInfo = new MemoryFileInfo(userIdentifier, false, reduceFileMeta);
+ if (!reduceMeta) {
+ memoryFileInfo.replaceFileMeta(new MapFileMeta(32 * 1024, 10));
+ }
+
+ File tempDir = Files.createTempDir();
+ tempDir.deleteOnExit();
+ File file = getTemporaryFile(tempDir);
+ DiskFileInfo fileInfo = new DiskFileInfo(file, userIdentifier, celebornConf);
+ if (!reduceMeta) {
+ fileInfo.replaceFileMeta(new MapFileMeta(32 * 1024, 10));
+ }
+
+ AbstractSource source = Mockito.mock(WorkerSource.class);
+ Mockito.doAnswer(
+ invocationOnMock -> {
+ Function0<?> function = (Function0<?>) invocationOnMock.getArguments()[2];
+ return function.apply();
+ })
+ .when(source)
+ .sample(Mockito.anyString(), Mockito.anyString(), Mockito.any(Function0.class));
+ AtomicInteger callCounter = new AtomicInteger(0);
+ LocalFlusher flusher =
+ new LocalFlusher(
+ source,
+ DeviceMonitor$.MODULE$.EmptyMonitor(),
+ 1,
+ NettyUtils.getPooledByteBufAllocator(
+ new TransportConf("test", celebornConf), null, true),
+ 256,
+ "disk1",
+ StorageInfo.Type.HDD,
+ null);
+ Mockito.doAnswer(
+ invocation -> {
+ if (callCounter.getAndIncrement() == 0) {
+ return new Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File>(
+ memoryFileInfo, null, null, null);
+ } else {
+ return new Tuple4<MemoryFileInfo, Flusher, DiskFileInfo, File>(
+ null, flusher, fileInfo, file);
+ }
+ })
+ .when(storageManager)
+ .createFile(Mockito.any(), Mockito.anyBoolean());
+ return storageManager;
}
}
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/SSLReducePartitionDataWriterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/SSLReducePartitionDataWriterSuiteJ.java
index e6e37ed51..4fa2d54ee 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/SSLReducePartitionDataWriterSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/SSLReducePartitionDataWriterSuiteJ.java
@@ -24,8 +24,9 @@ import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.network.TestHelper;
import org.apache.celeborn.common.network.ssl.SslSampleConfigs;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.service.deploy.worker.storage.local.DiskReducePartitionDataWriterSuiteJ;
-public class SSLReducePartitionDataWriterSuiteJ extends ReducePartitionDataWriterSuiteJ {
+public class SSLReducePartitionDataWriterSuiteJ extends DiskReducePartitionDataWriterSuiteJ {
protected TransportConf createModuleTransportConf(String module) {
CelebornConf conf =
TestHelper.updateCelebornConfWithMap(
@@ -35,11 +36,11 @@ public class SSLReducePartitionDataWriterSuiteJ extends ReducePartitionDataWrite
@BeforeClass
public static void beforeAll() {
- ReducePartitionDataWriterSuiteJ.beforeAll();
+ DiskReducePartitionDataWriterSuiteJ.beforeAll();
}
@AfterClass
public static void afterAll() {
- ReducePartitionDataWriterSuiteJ.afterAll();
+ DiskReducePartitionDataWriterSuiteJ.afterAll();
}
}
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskMapPartitionDataWriterSuiteJ.java
similarity index 85%
rename from worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriterSuiteJ.java
rename to worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskMapPartitionDataWriterSuiteJ.java
index bc7772001..f3fb4030d 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/MapPartitionDataWriterSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskMapPartitionDataWriterSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.worker.storage;
+package org.apache.celeborn.service.deploy.worker.storage.local;
import static org.junit.Assert.assertEquals;
@@ -24,12 +24,10 @@ import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.charset.StandardCharsets;
-import java.util.UUID;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import scala.Function0;
-import scala.Tuple4;
import scala.collection.mutable.ListBuffer;
import io.netty.buffer.Unpooled;
@@ -42,19 +40,21 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
-import org.apache.celeborn.common.meta.DiskFileInfo;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.network.util.TransportConf;
+import org.apache.celeborn.common.protocol.PartitionLocation;
import org.apache.celeborn.common.protocol.PartitionSplitMode;
+import org.apache.celeborn.common.protocol.PartitionType;
import org.apache.celeborn.common.protocol.StorageInfo;
import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
+import org.apache.celeborn.service.deploy.worker.storage.*;
-public class MapPartitionDataWriterSuiteJ {
+public class DiskMapPartitionDataWriterSuiteJ {
- private static final Logger LOG = LoggerFactory.getLogger(MapPartitionDataWriterSuiteJ.class);
+ private static final Logger LOG = LoggerFactory.getLogger(DiskMapPartitionDataWriterSuiteJ.class);
private static final CelebornConf CONF = new CelebornConf();
public static final Long SPLIT_THRESHOLD = 256 * 1024 * 1024L;
@@ -100,7 +100,7 @@ public class MapPartitionDataWriterSuiteJ {
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
- conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
@@ -121,21 +121,24 @@ public class MapPartitionDataWriterSuiteJ {
@Test
public void testMultiThreadWrite() throws IOException {
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, false);
+ PartitionLocation partitionLocation = Mockito.mock(PartitionLocation.class);
MapPartitionDataWriter fileWriter =
new MapPartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, false, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ partitionLocation,
+ "app1",
+ 1,
+ userIdentifier,
+ PartitionType.MAP,
+ false));
fileWriter.pushDataHandShake(2, 32 * 1024);
fileWriter.regionStart(0, false);
byte[] partData0 = generateData(0);
@@ -157,13 +160,6 @@ public class MapPartitionDataWriterSuiteJ {
assertEquals(fileWriter.getFile().length(), bytesWritten);
}
- private File getTemporaryFile() throws IOException {
- String filename = UUID.randomUUID().toString();
- File temporaryFile = new File(tempDir, filename);
- temporaryFile.createNewFile();
- return temporaryFile;
- }
-
private byte[] generateData(int partitionId) {
ThreadLocalRandom rand = ThreadLocalRandom.current();
byte[] hello = "hello, world".getBytes(StandardCharsets.UTF_8);
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskPartitionFilesSorterSuiteJ.java
similarity index 89%
rename from worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorterSuiteJ.java
rename to worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskPartitionFilesSorterSuiteJ.java
index cce7e1f29..d8546ed49 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/PartitionFilesSorterSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskPartitionFilesSorterSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.worker.storage;
+package org.apache.celeborn.service.deploy.worker.storage.local;
import static org.mockito.Mockito.when;
@@ -37,6 +37,7 @@ import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.DiskFileInfo;
+import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.meta.ReduceFileMeta;
import org.apache.celeborn.common.unsafe.Platform;
import org.apache.celeborn.common.util.CelebornExitKind;
@@ -44,18 +45,20 @@ import org.apache.celeborn.common.util.JavaUtils;
import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
+import org.apache.celeborn.service.deploy.worker.storage.PartitionDataWriter;
+import org.apache.celeborn.service.deploy.worker.storage.PartitionFilesSorter;
-public class PartitionFilesSorterSuiteJ {
+public class DiskPartitionFilesSorterSuiteJ {
- private static Logger logger = LoggerFactory.getLogger(PartitionFilesSorterSuiteJ.class);
+ private static final Logger logger =
+ LoggerFactory.getLogger(DiskPartitionFilesSorterSuiteJ.class);
- private Random random = new Random();
+ private final Random random = new Random();
private File shuffleFile;
private DiskFileInfo fileInfo;
private String originFileName;
- private long originFileLen;
private PartitionDataWriter partitionDataWriter;
- private UserIdentifier userIdentifier = new UserIdentifier("mock-tenantId", "mock-name");
+ private final UserIdentifier userIdentifier = new UserIdentifier("mock-tenantId", "mock-name");
private static final int MAX_MAP_ID = 50;
@@ -65,7 +68,7 @@ public class PartitionFilesSorterSuiteJ {
shuffleFile = File.createTempFile("Celeborn", "sort-suite");
originFileName = shuffleFile.getAbsolutePath();
- fileInfo = new DiskFileInfo(shuffleFile, userIdentifier);
+ fileInfo = new DiskFileInfo(shuffleFile, userIdentifier, new CelebornConf());
FileOutputStream fileOutputStream = new FileOutputStream(shuffleFile);
FileChannel channel = fileOutputStream.getChannel();
Map<Integer, Integer> batchIds = new HashMap<>();
@@ -102,8 +105,8 @@ public class PartitionFilesSorterSuiteJ {
}
partitionSize[mapId] = partitionSize[mapId] + batchHeader.length + mockedData.length;
}
- originFileLen = channel.size();
- ((ReduceFileMeta) fileInfo.getFileMeta()).getChunkOffsets().add(originFileLen);
+ long originFileLen = channel.size();
+ fileInfo.getReduceFileMeta().getChunkOffsets().add(originFileLen);
fileInfo.updateBytesFlushed(originFileLen);
logger.info(shuffleFile.getAbsolutePath() + " filelen: " + Utils.bytesToString(originFileLen));
@@ -113,7 +116,7 @@ public class PartitionFilesSorterSuiteJ {
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
- conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
@@ -141,7 +144,7 @@ public class PartitionFilesSorterSuiteJ {
conf.set(CelebornConf.SHUFFLE_CHUNK_SIZE().key(), "8m");
PartitionFilesSorter partitionFilesSorter =
new PartitionFilesSorter(MemoryManager.instance(), conf, new WorkerSource(conf));
- DiskFileInfo info =
+ FileInfo info =
partitionFilesSorter.getSortedFileInfo(
"application-1",
originFileName,
@@ -159,7 +162,7 @@ public class PartitionFilesSorterSuiteJ {
long actualTotalChunkSize =
((ReduceFileMeta) info.getFileMeta()).getLastChunkOffset()
- ((ReduceFileMeta) info.getFileMeta()).getChunkOffsets().get(0);
- Assert.assertTrue(totalSizeToFetch == actualTotalChunkSize);
+ Assert.assertEquals(totalSizeToFetch, actualTotalChunkSize);
} finally {
clean();
}
@@ -207,7 +210,7 @@ public class PartitionFilesSorterSuiteJ {
Assert.assertEquals(
partitionFilesSorter2.getSortedShuffleFiles("application-1-1").toString(),
"[0-0-3, 0-0-2, 0-0-1]");
- Assert.assertEquals(partitionFilesSorter2.getSortedShuffleFiles("application-2-1"), null);
+ Assert.assertNull(partitionFilesSorter2.getSortedShuffleFiles("application-2-1"));
Assert.assertEquals(
partitionFilesSorter2.getSortedShuffleFiles("application-3-1").toString(), "[0-0-1]");
partitionFilesSorter2.close(CelebornExitKind.WORKER_GRACEFUL_SHUTDOWN());
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskReducePartitionDataWriterSuiteJ.java
similarity index 72%
copy from worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java
copy to worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskReducePartitionDataWriterSuiteJ.java
index 6417e2aa0..748cd9d83 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/local/DiskReducePartitionDataWriterSuiteJ.java
@@ -15,7 +15,7 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.worker.storage;
+package org.apache.celeborn.service.deploy.worker.storage.local;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotEquals;
@@ -27,13 +27,7 @@ import java.io.IOException;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Set;
-import java.util.UUID;
+import java.util.*;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
@@ -43,7 +37,6 @@ import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;
import scala.Function0;
-import scala.Tuple4;
import scala.collection.mutable.ListBuffer;
import io.netty.buffer.ByteBuf;
@@ -57,9 +50,9 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.exception.CelebornException;
import org.apache.celeborn.common.identity.UserIdentifier;
import org.apache.celeborn.common.meta.DiskFileInfo;
+import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.meta.ReduceFileMeta;
import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
@@ -77,10 +70,12 @@ import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.FetchHandler;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
+import org.apache.celeborn.service.deploy.worker.storage.*;
-public class ReducePartitionDataWriterSuiteJ {
+public class DiskReducePartitionDataWriterSuiteJ {
- private static final Logger LOG = LoggerFactory.getLogger(ReducePartitionDataWriterSuiteJ.class);
+ private static final Logger LOG =
+ LoggerFactory.getLogger(DiskReducePartitionDataWriterSuiteJ.class);
private static final CelebornConf CONF = new CelebornConf();
public static final Long SPLIT_THRESHOLD = 256 * 1024 * 1024L;
@@ -130,7 +125,7 @@ public class ReducePartitionDataWriterSuiteJ {
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
- conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
@@ -152,7 +147,7 @@ public class ReducePartitionDataWriterSuiteJ {
}
@Override
- public DiskFileInfo getRawDiskFileInfo(String shuffleKey, String fileName) {
+ public FileInfo getRawFileInfo(String shuffleKey, String fileName) {
return info;
}
@@ -222,7 +217,7 @@ public class ReducePartitionDataWriterSuiteJ {
return message.toByteBuffer();
}
- private void setUpConn(TransportClient client) throws IOException, CelebornException {
+ private void setUpConn(TransportClient client) throws IOException {
ByteBuffer resp = client.sendRpcSync(createOpenMessage(), 10000);
PbStreamHandler streamHandler = TransportMessage.fromByteBuffer(resp).getParsedPayload();
streamId = streamHandler.getStreamId();
@@ -269,22 +264,25 @@ public class ReducePartitionDataWriterSuiteJ {
@Test
public void testMultiThreadWrite() throws IOException, ExecutionException, InterruptedException {
final int threadsNum = 8;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-1");
@@ -318,22 +316,24 @@ public class ReducePartitionDataWriterSuiteJ {
public void testMultiThreadWriteDuringClose()
throws IOException, ExecutionException, InterruptedException {
final int threadsNum = 8;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-1");
@@ -361,29 +361,31 @@ public class ReducePartitionDataWriterSuiteJ {
assertEquals(length.get(), bytesWritten.get());
assertEquals(partitionDataWriter.getFile().length(), bytesWritten.get());
- assertEquals(context._3().getFileLength(), bytesWritten.get());
+ assertEquals(partitionDataWriter.getDiskFileInfo().getFileLength(), bytesWritten.get());
}
@Test
public void testAfterStressfulWriteWillReadCorrect()
throws IOException, ExecutionException, InterruptedException {
final int threadsNum = Runtime.getRuntime().availableProcessors();
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-2");
@@ -432,23 +434,24 @@ public class ReducePartitionDataWriterSuiteJ {
@Test
public void testWriteAndChunkRead() throws Exception {
final int threadsNum = 16;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
- DiskFileInfo fileInfo = context._3();
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-2");
@@ -477,7 +480,7 @@ public class ReducePartitionDataWriterSuiteJ {
long bytesWritten = partitionDataWriter.close();
assertEquals(length.get(), bytesWritten);
- setupChunkServer(fileInfo);
+ setupChunkServer(partitionDataWriter.getDiskFileInfo());
TransportClient client =
clientFactory.createClient(InetAddress.getLocalHost().getHostAddress(), server.getPort());
@@ -547,24 +550,28 @@ public class ReducePartitionDataWriterSuiteJ {
conf.set(CelebornConf.WORKER_FLUSHER_BUFFER_SIZE().key(), "128B");
// case 1: write 8MiB
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
partitionDataWriter.write(generateData(8 * 1024 * 1024));
partitionDataWriter.close();
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ ReduceFileMeta reduceFileMeta =
+ (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 8 * 1024 * 1024);
assertEquals(
@@ -572,224 +579,242 @@ public class ReducePartitionDataWriterSuiteJ {
8 * 1024 * 1024);
// case 2: write 1024B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 8; i++) {
partitionDataWriter.write(generateData(128));
}
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1024);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 3: write 1023B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
partitionDataWriter.write(generateData(1020));
partitionDataWriter.write(generateData(3));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1023);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1023);
// case 4: write 1025B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 8; i++) {
partitionDataWriter.write(generateData(128));
}
partitionDataWriter.write(generateData(1));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1025);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 5: write 2048B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
partitionDataWriter.write(generateData(128));
}
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2048);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 5.1: write 2048B with trim; without PR #1702 this case will fail
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
partitionDataWriter.write(generateData(128));
}
// mock trim
- partitionDataWriter.flush(false);
+ partitionDataWriter.flush(false, false);
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2048);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 6: write 2049B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
partitionDataWriter.write(generateData(128));
}
partitionDataWriter.write(generateData(1));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2049);
assertEquals(
reduceFileMeta.getChunkOffsets().get(2) - reduceFileMeta.getChunkOffsets().get(1), 1024);
// case 7: write 4097B with 3 chunks
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
partitionDataWriter.write(generateData(1024));
for (int i = 0; i < 9; i++) {
partitionDataWriter.write(generateData(128));
}
partitionDataWriter.write(generateData(1920));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 4096);
assertEquals(
reduceFileMeta.getChunkOffsets().get(3) - reduceFileMeta.getChunkOffsets().get(2), 2048);
// case 7.2: write 4097B with 3 chunks with trim; without PR #1702 this case will fail
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareDiskFileTestEnvironment(
+ tempDir, userIdentifier, localFlusher, true, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
partitionDataWriter.write(generateData(1024));
for (int i = 0; i < 9; i++) {
partitionDataWriter.write(generateData(128));
- partitionDataWriter.flush(false);
+ partitionDataWriter.flush(false, false);
}
partitionDataWriter.write(generateData(1920));
// mock trim
- partitionDataWriter.flush(false);
+ partitionDataWriter.flush(false, false);
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getDiskFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 4096);
assertEquals(
@@ -818,9 +843,7 @@ public class ReducePartitionDataWriterSuiteJ {
private ByteBuf generateData(int len) {
byte[] data = new byte[len];
- for (int i = 0; i < len; i++) {
- data[i] = 'a';
- }
+ Arrays.fill(data, (byte) 'a');
return Unpooled.wrappedBuffer(data);
}
}
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryPartitionFilesSorterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryPartitionFilesSorterSuiteJ.java
new file mode 100644
index 000000000..1d39839b0
--- /dev/null
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryPartitionFilesSorterSuiteJ.java
@@ -0,0 +1,159 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 org.apache.celeborn.service.deploy.worker.storage.memory;
+
+import static org.mockito.Mockito.when;
+
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.Random;
+
+import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.PooledByteBufAllocator;
+import org.junit.Assert;
+import org.junit.Test;
+import org.mockito.Mockito;
+
+import org.apache.celeborn.common.CelebornConf;
+import org.apache.celeborn.common.identity.UserIdentifier;
+import org.apache.celeborn.common.meta.FileInfo;
+import org.apache.celeborn.common.meta.MemoryFileInfo;
+import org.apache.celeborn.common.meta.ReduceFileMeta;
+import org.apache.celeborn.common.metrics.source.AbstractSource;
+import org.apache.celeborn.common.network.util.NettyUtils;
+import org.apache.celeborn.common.unsafe.Platform;
+import org.apache.celeborn.service.deploy.worker.WorkerSource;
+import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
+import org.apache.celeborn.service.deploy.worker.storage.PartitionDataWriter;
+import org.apache.celeborn.service.deploy.worker.storage.PartitionFilesSorter;
+import org.apache.celeborn.service.deploy.worker.storage.StorageManager;
+
+public class MemoryPartitionFilesSorterSuiteJ {
+
+ private final Random random = new Random();
+ private MemoryFileInfo fileInfo;
+ private PartitionDataWriter partitionDataWriter;
+ private final UserIdentifier userIdentifier = new UserIdentifier("mock-tenantId", "mock-name");
+
+ private static final int MAX_MAP_ID = 50;
+
+ public long[] prepare(int mapCount) {
+ long[] partitionSize = new long[MAX_MAP_ID];
+ byte[] batchHeader = new byte[16];
+ fileInfo = new MemoryFileInfo(userIdentifier, true, new ReduceFileMeta(8 * 1024 * 1024));
+
+ AbstractSource source = Mockito.mock(AbstractSource.class);
+ PooledByteBufAllocator allocator =
+ NettyUtils.getSharedPooledByteBufAllocator(new CelebornConf(), source, false);
+ CompositeByteBuf buffer = allocator.compositeBuffer();
+ fileInfo.setBuffer(buffer);
+
+ Map<Integer, Integer> batchIds = new HashMap<>();
+
+ for (int i = 0; i < mapCount; i++) {
+ int mapId = random.nextInt(MAX_MAP_ID);
+ int currentAttemptId = 0;
+ int batchId =
+ batchIds.compute(
+ mapId,
+ (k, v) -> {
+ if (v == null) {
+ v = 0;
+ } else {
+ v++;
+ }
+ return v;
+ });
+ // [63.9k, 192k + 63.9k]
+ int dataSize = random.nextInt(192 * 1024) + 65525;
+ byte[] mockedData = new byte[dataSize];
+ Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET, mapId);
+ Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 4, currentAttemptId);
+ Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 8, batchId);
+ Platform.putInt(batchHeader, Platform.BYTE_ARRAY_OFFSET + 12, dataSize);
+ ByteBuffer buf1 = ByteBuffer.wrap(batchHeader);
+ while (buf1.hasRemaining()) {
+ buffer.writeBytes(buf1);
+ }
+ random.nextBytes(mockedData);
+ ByteBuffer buf2 = ByteBuffer.wrap(mockedData);
+ while (buf2.hasRemaining()) {
+ buffer.writeBytes(buf2);
+ }
+ partitionSize[mapId] = partitionSize[mapId] + batchHeader.length + mockedData.length;
+ }
+ long originFileLen = buffer.readableBytes();
+ fileInfo.getReduceFileMeta().addChunkOffset(originFileLen);
+ fileInfo.updateBytesFlushed((int) originFileLen);
+
+ CelebornConf conf = new CelebornConf();
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(), "0.8");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_REPLICATE().key(), "0.9");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
+ conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
+
+ StorageManager storageManager = Mockito.mock(StorageManager.class);
+ Mockito.when(storageManager.storageBufferAllocator()).thenAnswer(a -> allocator);
+ MemoryManager.initialize(conf, storageManager);
+ partitionDataWriter = Mockito.mock(PartitionDataWriter.class);
+ when(partitionDataWriter.getMemoryFileInfo()).thenAnswer(i -> fileInfo);
+
+ return partitionSize;
+ }
+
+ private void check(int mapCount, int startMapIndex, int endMapIndex) throws IOException {
+ long[] partitionSize = prepare(mapCount);
+ CelebornConf conf = new CelebornConf();
+ conf.set(CelebornConf.SHUFFLE_CHUNK_SIZE().key(), "8m");
+ PartitionFilesSorter partitionFilesSorter =
+ new PartitionFilesSorter(MemoryManager.instance(), conf, new WorkerSource(conf));
+ FileInfo info =
+ partitionFilesSorter.getSortedFileInfo(
+ "application-1",
+ "",
+ partitionDataWriter.getMemoryFileInfo(),
+ startMapIndex,
+ endMapIndex);
+ long totalSizeToFetch = 0;
+ for (int i = startMapIndex; i < endMapIndex; i++) {
+ totalSizeToFetch += partitionSize[i];
+ }
+ long numChunks = totalSizeToFetch / conf.shuffleChunkSize() + 1;
+ Assert.assertTrue(
+ 0 < ((ReduceFileMeta) info.getFileMeta()).getNumChunks()
+ && ((ReduceFileMeta) info.getFileMeta()).getNumChunks() <= numChunks);
+ long actualTotalChunkSize =
+ ((ReduceFileMeta) info.getFileMeta()).getLastChunkOffset()
+ - ((ReduceFileMeta) info.getFileMeta()).getChunkOffsets().get(0);
+ Assert.assertEquals(totalSizeToFetch, actualTotalChunkSize);
+ fileInfo.getBuffer().release();
+ }
+
+ @Test
+ public void testSortMemoryShuffleFile() throws IOException {
+ int startMapIndex = random.nextInt(5);
+ int endMapIndex = startMapIndex + random.nextInt(5) + 5;
+ check(1000, startMapIndex, endMapIndex);
+ }
+}
diff --git a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryReducePartitionDataWriterSuiteJ.java
similarity index 57%
rename from worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java
rename to worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryReducePartitionDataWriterSuiteJ.java
index 6417e2aa0..cf71200dc 100644
--- a/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/ReducePartitionDataWriterSuiteJ.java
+++ b/worker/src/test/java/org/apache/celeborn/service/deploy/worker/storage/memory/MemoryReducePartitionDataWriterSuiteJ.java
@@ -15,41 +15,25 @@
* limitations under the License.
*/
-package org.apache.celeborn.service.deploy.worker.storage;
+package org.apache.celeborn.service.deploy.worker.storage.memory;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.fail;
+import static org.junit.Assert.*;
import static org.mockito.Mockito.*;
-import java.io.File;
import java.io.IOException;
import java.net.InetAddress;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.LinkedList;
-import java.util.List;
-import java.util.Set;
-import java.util.UUID;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Future;
-import java.util.concurrent.Semaphore;
-import java.util.concurrent.ThreadLocalRandom;
-import java.util.concurrent.TimeUnit;
+import java.util.*;
+import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicLong;
import scala.Function0;
-import scala.Tuple4;
-import scala.collection.mutable.ListBuffer;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.CompositeByteBuf;
+import io.netty.buffer.PooledByteBufAllocator;
import io.netty.buffer.Unpooled;
-import org.junit.AfterClass;
import org.junit.BeforeClass;
import org.junit.Test;
import org.mockito.Mockito;
@@ -57,9 +41,8 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.apache.celeborn.common.CelebornConf;
-import org.apache.celeborn.common.exception.CelebornException;
import org.apache.celeborn.common.identity.UserIdentifier;
-import org.apache.celeborn.common.meta.DiskFileInfo;
+import org.apache.celeborn.common.meta.FileInfo;
import org.apache.celeborn.common.meta.ReduceFileMeta;
import org.apache.celeborn.common.network.TransportContext;
import org.apache.celeborn.common.network.buffer.ManagedBuffer;
@@ -71,36 +54,37 @@ import org.apache.celeborn.common.network.server.TransportServer;
import org.apache.celeborn.common.network.util.NettyUtils;
import org.apache.celeborn.common.network.util.TransportConf;
import org.apache.celeborn.common.protocol.*;
-import org.apache.celeborn.common.util.JavaUtils;
+import org.apache.celeborn.common.unsafe.Platform;
import org.apache.celeborn.common.util.ThreadUtils;
-import org.apache.celeborn.common.util.Utils;
import org.apache.celeborn.service.deploy.worker.FetchHandler;
import org.apache.celeborn.service.deploy.worker.WorkerSource;
import org.apache.celeborn.service.deploy.worker.memory.MemoryManager;
+import org.apache.celeborn.service.deploy.worker.storage.*;
-public class ReducePartitionDataWriterSuiteJ {
+public class MemoryReducePartitionDataWriterSuiteJ {
- private static final Logger LOG = LoggerFactory.getLogger(ReducePartitionDataWriterSuiteJ.class);
+ private static final Logger LOG =
+ LoggerFactory.getLogger(MemoryReducePartitionDataWriterSuiteJ.class);
private static final CelebornConf CONF = new CelebornConf();
public static final Long SPLIT_THRESHOLD = 256 * 1024 * 1024L;
public static final PartitionSplitMode splitMode = PartitionSplitMode.HARD;
- private static File tempDir = null;
- private static LocalFlusher localFlusher = null;
private static WorkerSource source = null;
- private TransportContext transportContext;
- private TransportServer server;
- private TransportClientFactory clientFactory;
+ private static TransportServer server;
+ private static TransportClientFactory clientFactory;
private static long streamId;
private static int numChunks;
private final UserIdentifier userIdentifier = new UserIdentifier("mock-tenantId", "mock-name");
+ private static final TransportConf transConf = new TransportConf("shuffle", new CelebornConf());
+ private static StorageManager storageManager;
+
@BeforeClass
public static void beforeAll() {
- tempDir = Utils.createTempDir(System.getProperty("java.io.tmpdir"), "celeborn");
CONF.set(CelebornConf.SHUFFLE_CHUNK_SIZE().key(), "1k");
+ CONF.set(CelebornConf.ACTIVE_STORAGE_TYPES().key(), "MEMORY,HDD");
source = Mockito.mock(WorkerSource.class);
Mockito.doAnswer(
@@ -111,39 +95,27 @@ public class ReducePartitionDataWriterSuiteJ {
.when(source)
.sample(Mockito.anyString(), Mockito.anyString(), Mockito.any(Function0.class));
- ListBuffer<File> dirs = new ListBuffer<>();
- dirs.$plus$eq(tempDir);
- localFlusher =
- new LocalFlusher(
- source,
- DeviceMonitor$.MODULE$.EmptyMonitor(),
- 1,
- NettyUtils.getPooledByteBufAllocator(new TransportConf("test", CONF), null, true),
- 256,
- "disk1",
- StorageInfo.Type.HDD,
- null);
-
CelebornConf conf = new CelebornConf();
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_RECEIVE().key(), "0.8");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_PAUSE_REPLICATE().key(), "0.9");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_RESUME().key(), "0.5");
conf.set(CelebornConf.WORKER_PARTITION_SORTER_DIRECT_MEMORY_RATIO_THRESHOLD().key(), "0.6");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_READ_BUFFER().key(), "0.1");
- conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_SHUFFLE_STORAGE().key(), "0.1");
+ conf.set(CelebornConf.WORKER_DIRECT_MEMORY_RATIO_FOR_MEMORY_FILE_STORAGE().key(), "0.1");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_CHECK_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_DIRECT_MEMORY_REPORT_INTERVAL().key(), "10");
conf.set(CelebornConf.WORKER_READBUFFER_ALLOCATIONWAIT().key(), "10ms");
- MemoryManager.initialize(conf);
- }
-
- protected TransportConf createModuleTransportConf(String module) {
- return new TransportConf(module, new CelebornConf());
+ PooledByteBufAllocator allocator =
+ NettyUtils.getPooledByteBufAllocator(transConf, source, false);
+ storageManager = Mockito.mock(StorageManager.class);
+ AtomicLong evictCount = new AtomicLong();
+ Mockito.when(storageManager.evictedFileCount()).thenAnswer(a -> evictCount);
+ Mockito.when(storageManager.localOrHdfsStorageAvailable()).thenAnswer(a -> true);
+ Mockito.when(storageManager.storageBufferAllocator()).thenAnswer(a -> allocator);
+ MemoryManager.initialize(conf, storageManager);
}
- public void setupChunkServer(DiskFileInfo info) throws IOException {
- TransportConf transConf = createModuleTransportConf("shuffle");
-
+ public static void setupChunkServer(FileInfo info) throws IOException {
FetchHandler handler =
new FetchHandler(transConf.getCelebornConf(), transConf, mock(WorkerSource.class)) {
@Override
@@ -152,7 +124,7 @@ public class ReducePartitionDataWriterSuiteJ {
}
@Override
- public DiskFileInfo getRawDiskFileInfo(String shuffleKey, String fileName) {
+ public FileInfo getRawFileInfo(String shuffleKey, String fileName) {
return info;
}
@@ -171,28 +143,15 @@ public class ReducePartitionDataWriterSuiteJ {
.when(sorter)
.getSortedFileInfo(anyString(), anyString(), eq(info), anyInt(), anyInt());
handler.setPartitionsSorter(sorter);
- transportContext = new TransportContext(transConf, handler);
- server = transportContext.createServer();
+ TransportContext context = new TransportContext(transConf, handler);
+ server = context.createServer();
- clientFactory = transportContext.createClientFactory();
- }
-
- @AfterClass
- public static void afterAll() {
- if (tempDir != null) {
- try {
- JavaUtils.deleteRecursively(tempDir);
- tempDir = null;
- } catch (IOException e) {
- LOG.error("Failed to delete temp dir.", e);
- }
- }
+ clientFactory = context.createClientFactory();
}
- public void closeChunkServer() {
+ public static void closeChunkServer() {
server.close();
clientFactory.close();
- transportContext.close();
}
static class FetchResult {
@@ -212,7 +171,7 @@ public class ReducePartitionDataWriterSuiteJ {
new TransportMessage(
MessageType.OPEN_STREAM,
PbOpenStream.newBuilder()
- .setShuffleKey("shuffleKey-123")
+ .setShuffleKey("app-1")
.setFileName("location")
.setStartIndex(0)
.setEndIndex(Integer.MAX_VALUE)
@@ -222,7 +181,7 @@ public class ReducePartitionDataWriterSuiteJ {
return message.toByteBuffer();
}
- private void setUpConn(TransportClient client) throws IOException, CelebornException {
+ private void setUpConn(TransportClient client) throws IOException {
ByteBuffer resp = client.sendRpcSync(createOpenMessage(), 10000);
PbStreamHandler streamHandler = TransportMessage.fromByteBuffer(resp).getParsedPayload();
streamId = streamHandler.getStreamId();
@@ -251,6 +210,7 @@ public class ReducePartitionDataWriterSuiteJ {
@Override
public void onFailure(int chunkIndex, Throwable e) {
res.failedChunks.add(chunkIndex);
+ LOG.error("Fetch chunk failed", e);
sem.release();
}
};
@@ -269,22 +229,25 @@ public class ReducePartitionDataWriterSuiteJ {
@Test
public void testMultiThreadWrite() throws IOException, ExecutionException, InterruptedException {
final int threadsNum = 8;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-1");
@@ -294,7 +257,7 @@ public class ReducePartitionDataWriterSuiteJ {
futures.add(
es.submit(
() -> {
- byte[] bytes = generateData();
+ byte[] bytes = generateDataWithHeader();
length.addAndGet(bytes.length);
ByteBuf buf = Unpooled.wrappedBuffer(bytes);
try {
@@ -311,29 +274,31 @@ public class ReducePartitionDataWriterSuiteJ {
long bytesWritten = partitionDataWriter.close();
assertEquals(length.get(), bytesWritten);
- assertEquals(partitionDataWriter.getFile().length(), bytesWritten);
+ assertEquals(partitionDataWriter.getMemoryFileInfo().getFileLength(), bytesWritten);
}
@Test
public void testMultiThreadWriteDuringClose()
throws IOException, ExecutionException, InterruptedException {
final int threadsNum = 8;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-1");
@@ -344,7 +309,7 @@ public class ReducePartitionDataWriterSuiteJ {
futures.add(
es.submit(
() -> {
- byte[] bytes = generateData();
+ byte[] bytes = generateDataWithHeader();
ByteBuf buf = Unpooled.wrappedBuffer(bytes);
try {
partitionDataWriter.write(buf);
@@ -360,30 +325,31 @@ public class ReducePartitionDataWriterSuiteJ {
}
assertEquals(length.get(), bytesWritten.get());
- assertEquals(partitionDataWriter.getFile().length(), bytesWritten.get());
- assertEquals(context._3().getFileLength(), bytesWritten.get());
+ assertEquals(partitionDataWriter.getMemoryFileInfo().getFileLength(), bytesWritten.get());
}
@Test
public void testAfterStressfulWriteWillReadCorrect()
throws IOException, ExecutionException, InterruptedException {
final int threadsNum = Runtime.getRuntime().availableProcessors();
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryEvictEnvironment(
+ userIdentifier, true, storageManager, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-2");
@@ -393,7 +359,7 @@ public class ReducePartitionDataWriterSuiteJ {
es.submit(
() -> {
for (int j = 0; j < 100; ++j) {
- byte[] bytes = generateData();
+ byte[] bytes = generateDataWithHeader();
length.addAndGet(bytes.length);
ByteBuf buf = Unpooled.wrappedBuffer(bytes);
try {
@@ -413,42 +379,101 @@ public class ReducePartitionDataWriterSuiteJ {
}
@Test
- public void testHugeBufferQueueSize() throws IOException {
- File file = getTemporaryFile();
- ListBuffer<File> dirs = new ListBuffer<>();
- dirs.$plus$eq(file);
- localFlusher =
- new LocalFlusher(
+ public void testWriteAndChunkRead() throws Exception {
+ final int threadsNum = 2;
+ PartitionDataWriter partitionDataWriter =
+ new ReducePartitionDataWriter(
+ PartitionDataWriterSuiteUtils.prepareMemoryEvictEnvironment(
+ userIdentifier, true, storageManager, CONF),
source,
+ CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- 1,
- NettyUtils.getPooledByteBufAllocator(new TransportConf("test", CONF), null, true),
- 256,
- "disk2",
- StorageInfo.Type.HDD,
- null);
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
+
+ List<Future<?>> futures = new ArrayList<>();
+ ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-2");
+ AtomicLong length = new AtomicLong(0);
+ futures.add(
+ es.submit(
+ () -> {
+ for (int j = 0; j < 10; ++j) {
+ byte[] bytes = generateDataWithHeader();
+ length.addAndGet(bytes.length);
+ ByteBuf buf = Unpooled.wrappedBuffer(bytes);
+ buf.retain();
+ try {
+ partitionDataWriter.incrementPendingWrites();
+ partitionDataWriter.write(buf);
+ } catch (IOException e) {
+ LOG.error("Failed to write buffer.", e);
+ }
+ buf.release();
+ }
+ }));
+ for (Future<?> future : futures) {
+ future.get();
+ }
+
+ long bytesWritten = partitionDataWriter.close();
+ assertEquals(length.get(), bytesWritten);
+
+ setupChunkServer(partitionDataWriter.getMemoryFileInfo());
+
+ TransportClient client =
+ clientFactory.createClient(InetAddress.getLocalHost().getHostAddress(), server.getPort());
+
+ setUpConn(client);
+
+ int chunkNums = numChunks;
+ List<Integer> indices = new ArrayList<>();
+
+ for (int i = 0; i < chunkNums; i++) {
+ indices.add(i);
+ }
+
+ FetchResult result = fetchChunks(client, indices);
+
+ Thread.sleep(3000);
+
+ assertEquals(result.buffers.size(), chunkNums);
+ assertEquals(result.successChunks.size(), chunkNums);
+
+ result.releaseBuffers();
+
+ closeChunkServer();
}
@Test
- public void testWriteAndChunkRead() throws Exception {
+ public void testEvictAndChunkRead() throws Exception {
final int threadsNum = 16;
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> context =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
- DiskFileInfo fileInfo = context._3();
-
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- context._1(),
- context._3(),
- context._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryEvictEnvironment(
+ userIdentifier, true, storageManager, CONF),
source,
CONF,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
List<Future<?>> futures = new ArrayList<>();
ExecutorService es = ThreadUtils.newDaemonFixedThreadPool(threadsNum, "FileWriter-UT-2");
@@ -457,7 +482,7 @@ public class ReducePartitionDataWriterSuiteJ {
es.submit(
() -> {
for (int j = 0; j < 1000; ++j) {
- byte[] bytes = generateData();
+ byte[] bytes = generateDataWithHeader();
length.addAndGet(bytes.length);
ByteBuf buf = Unpooled.wrappedBuffer(bytes);
buf.retain();
@@ -477,7 +502,7 @@ public class ReducePartitionDataWriterSuiteJ {
long bytesWritten = partitionDataWriter.close();
assertEquals(length.get(), bytesWritten);
- setupChunkServer(fileInfo);
+ setupChunkServer(partitionDataWriter.getDiskFileInfo());
TransportClient client =
clientFactory.createClient(InetAddress.getLocalHost().getHostAddress(), server.getPort());
@@ -547,24 +572,28 @@ public class ReducePartitionDataWriterSuiteJ {
conf.set(CelebornConf.WORKER_FLUSHER_BUFFER_SIZE().key(), "128B");
// case 1: write 8MiB
- Tuple4<StorageManager, Flusher, DiskFileInfo, File> pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
PartitionDataWriter partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
- partitionDataWriter.write(generateData(8 * 1024 * 1024));
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
+ partitionDataWriter.write(generateDataWithHeader(8 * 1024 * 1024));
partitionDataWriter.close();
- ReduceFileMeta reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ ReduceFileMeta reduceFileMeta =
+ (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 8 * 1024 * 1024);
assertEquals(
@@ -572,255 +601,270 @@ public class ReducePartitionDataWriterSuiteJ {
8 * 1024 * 1024);
// case 2: write 1024B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 8; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1024);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 3: write 1023B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
- partitionDataWriter.write(generateData(1020));
- partitionDataWriter.write(generateData(3));
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
+ partitionDataWriter.write(generateDataWithHeader(1020));
+ partitionDataWriter.write(generateDataWithHeader(3));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 1);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1023);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1023);
// case 4: write 1025B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 8; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
- partitionDataWriter.write(generateData(1));
+ partitionDataWriter.write(generateDataWithHeader(1));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 1025);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 5: write 2048B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2048);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 5.1: write 2048B with trim; without PR #1702 this case will fail
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
// mock trim
- partitionDataWriter.flush(false);
+ partitionDataWriter.flush(false, false);
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 2);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2048);
assertEquals(
reduceFileMeta.getChunkOffsets().get(1) - reduceFileMeta.getChunkOffsets().get(0), 1024);
// case 6: write 2049B
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
for (int i = 0; i < 16; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
- partitionDataWriter.write(generateData(1));
+ partitionDataWriter.write(generateDataWithHeader(1));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 2049);
assertEquals(
reduceFileMeta.getChunkOffsets().get(2) - reduceFileMeta.getChunkOffsets().get(1), 1024);
// case 7: write 4097B with 3 chunks
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
- partitionDataWriter.write(generateData(1024));
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
+ partitionDataWriter.write(generateDataWithHeader(1024));
for (int i = 0; i < 9; i++) {
- partitionDataWriter.write(generateData(128));
+ partitionDataWriter.write(generateDataWithHeader(128));
}
- partitionDataWriter.write(generateData(1920));
+ partitionDataWriter.write(generateDataWithHeader(1920));
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 4096);
assertEquals(
reduceFileMeta.getChunkOffsets().get(3) - reduceFileMeta.getChunkOffsets().get(2), 2048);
// case 7.2: write 4097B with 3 chunks with trim; without PR #1702 this case will fail
- pctx =
- PartitionDataWriterSuiteUtils.prepareTestFileContext(
- tempDir, userIdentifier, localFlusher, true);
-
partitionDataWriter =
new ReducePartitionDataWriter(
- pctx._1(),
- pctx._3(),
- pctx._2(),
+ PartitionDataWriterSuiteUtils.prepareMemoryFileTestEnvironment(
+ userIdentifier, true, storageManager, conf),
source,
conf,
DeviceMonitor$.MODULE$.EmptyMonitor(),
- SPLIT_THRESHOLD,
- splitMode,
- false,
- "app1-1");
- partitionDataWriter.write(generateData(1024));
+ new PartitionDataWriterContext(
+ SPLIT_THRESHOLD,
+ splitMode,
+ false,
+ new PartitionLocation(
+ 1, 0, "host", 1111, 1112, 1113, 1114, PartitionLocation.Mode.PRIMARY, null),
+ "app1-1",
+ 1,
+ userIdentifier,
+ PartitionType.REDUCE,
+ false));
+ partitionDataWriter.write(generateDataWithHeader(1024));
for (int i = 0; i < 9; i++) {
- partitionDataWriter.write(generateData(128));
- partitionDataWriter.flush(false);
+ partitionDataWriter.write(generateDataWithHeader(128));
+ partitionDataWriter.flush(false, false);
}
- partitionDataWriter.write(generateData(1920));
+ partitionDataWriter.write(generateDataWithHeader(1920));
// mock trim
- partitionDataWriter.flush(false);
+ partitionDataWriter.flush(false, false);
partitionDataWriter.close();
- reduceFileMeta = (ReduceFileMeta) pctx._3().getFileMeta();
+ reduceFileMeta = (ReduceFileMeta) partitionDataWriter.getMemoryFileInfo().getFileMeta();
assertEquals(reduceFileMeta.getNumChunks(), 3);
assertEquals(reduceFileMeta.getLastChunkOffset(), 4096);
assertEquals(
reduceFileMeta.getChunkOffsets().get(3) - reduceFileMeta.getChunkOffsets().get(2), 2048);
}
- private File getTemporaryFile() throws IOException {
- String filename = UUID.randomUUID().toString();
- File temporaryFile = new File(tempDir, filename);
- temporaryFile.createNewFile();
- return temporaryFile;
- }
-
- private byte[] generateData() {
+ private byte[] generateDataWithHeader() {
ThreadLocalRandom rand = ThreadLocalRandom.current();
+ byte[] header = new byte[16];
+ Platform.putInt(header, Platform.BYTE_ARRAY_OFFSET, 1);
+ Platform.putInt(header, Platform.BYTE_ARRAY_OFFSET + 4, 1);
+ Platform.putInt(header, Platform.BYTE_ARRAY_OFFSET + 8, 1);
byte[] hello = "hello, world".getBytes(StandardCharsets.UTF_8);
int tempLen = rand.nextInt(256 * 1024) + 128 * 1024;
int len = (int) (Math.ceil(1.0 * tempLen / hello.length) * hello.length);
+ Platform.putInt(header, Platform.BYTE_ARRAY_OFFSET + 12, len);
- byte[] data = new byte[len];
+ byte[] data = new byte[16 + len];
+ System.arraycopy(header, 0, data, 0, 16);
for (int i = 0; i < len; i += hello.length) {
- System.arraycopy(hello, 0, data, i, hello.length);
+ System.arraycopy(hello, 0, data, 16 + i, hello.length);
}
return data;
}
- private ByteBuf generateData(int len) {
+ private ByteBuf generateDataWithHeader(int len) {
byte[] data = new byte[len];
- for (int i = 0; i < len; i++) {
- data[i] = 'a';
- }
+ Arrays.fill(data, (byte) 'a');
return Unpooled.wrappedBuffer(data);
}
}
diff --git a/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala b/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
index 6f91b4374..78fc43670 100644
--- a/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
+++ b/worker/src/test/scala/org/apache/celeborn/service/deploy/memory/MemoryManagerSuite.scala
@@ -53,7 +53,7 @@ class MemoryManagerSuite extends CelebornFunSuite {
val conf = new CelebornConf()
try {
val memoryManager = MemoryManager.initialize(conf)
- val maxDirectorMemory = memoryManager.maxDirectorMemory
+ val maxDirectorMemory = memoryManager.maxDirectMemory
val pushThreshold =
(conf.workerDirectMemoryRatioToPauseReceive * maxDirectorMemory).longValue()
val replicateThreshold =
@@ -90,7 +90,7 @@ class MemoryManagerSuite extends CelebornFunSuite {
test("[CELEBORN-882] Test MemoryManager check memory thread logic") {
val conf = new CelebornConf()
val memoryManager = MemoryManager.initialize(conf)
- val maxDirectorMemory = memoryManager.maxDirectorMemory
+ val maxDirectorMemory = memoryManager.maxDirectMemory
val pushThreshold =
(conf.workerDirectMemoryRatioToPauseReceive * maxDirectorMemory).longValue()
val replicateThreshold =