This is an automated email from the ASF dual-hosted git repository. angerszhuuuu pushed a commit to branch branch-0.2 in repository https://gitbox.apache.org/repos/asf/incubator-celeborn.git
commit 942128f11394f7b4ec3fd592b92e752a14d218ea Author: Angerszhuuuu <[email protected]> AuthorDate: Wed Dec 14 18:24:19 2022 +0800 [CELEBORN-140][IMPROVEMENT] Quota yaml should support byte string (#1085) --- .../org/apache/celeborn/common/quota/DefaultQuotaManager.scala | 7 ++++++- common/src/test/resources/test-quota.yaml | 8 ++------ .../apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala | 5 +++-- 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala b/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala index 300b3f41..fb327d33 100644 --- a/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala +++ b/common/src/main/scala/org/apache/celeborn/common/quota/DefaultQuotaManager.scala @@ -49,7 +49,12 @@ class DefaultQuotaManager(conf: CelebornConf) extends QuotaManager(conf) { .asInstanceOf[java.util.HashMap[String, Object]] .asScala .foreach { case (key, value) => - quota.update(userIdentifier, key, value.toString.toLong) + key match { + case "diskBytesWritten" | "hdfsBytesWritten" => + quota.update(userIdentifier, key, Utils.byteStringAsBytes(value.toString)) + case _ => + quota.update(userIdentifier, key, value.toString.toLong) + } } userQuotas.put(userIdentifier, quota) } diff --git a/common/src/test/resources/test-quota.yaml b/common/src/test/resources/test-quota.yaml index 45d2cdc3..95935188 100644 --- a/common/src/test/resources/test-quota.yaml +++ b/common/src/test/resources/test-quota.yaml @@ -18,15 +18,11 @@ - tenantId: AAA name: Tom quota: - diskBytesWritten: 10000 + diskBytesWritten: 100m diskFileCount: 200 - hdfsBytesWritten: -1 - hdfsFileCount: -1 - tenantId: BBB name: Jerry quota: - diskBytesWritten: -1 - diskFileCount: -1 - hdfsBytesWritten: 10000 + hdfsBytesWritten: 200m hdfsFileCount: 200 diff --git a/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala b/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala index e7d43ef0..a1219ef6 100644 --- a/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala +++ b/common/src/test/scala/org/apache/celeborn/common/quota/DefaultQuotaManagerSuite.scala @@ -21,6 +21,7 @@ import org.junit.Assert.assertEquals import org.apache.celeborn.common.CelebornConf import org.apache.celeborn.common.identity.UserIdentifier +import org.apache.celeborn.common.util.Utils class DefaultQuotaManagerSuite extends BaseQuotaManagerSuite { @@ -39,9 +40,9 @@ class DefaultQuotaManagerSuite extends BaseQuotaManagerSuite { test("test rss quota conf") { assertEquals( quotaManager.getQuota(UserIdentifier("AAA", "Tom")), - Quota(10000, 200, -1, -1)) + Quota(Utils.byteStringAsBytes("100m"), 200, -1, -1)) assertEquals( quotaManager.getQuota(UserIdentifier("BBB", "Jerry")), - Quota(-1, -1, 10000, 200)) + Quota(-1, -1, Utils.byteStringAsBytes("200m"), 200)) } }
