feat(glue): support Ray jobs: Add tests and update README

aws · mergify · Feb 9, 2023 · Jan 25, 2023 · Jan 25, 2023 · Jan 25, 2023
commit 13983df012ba576e4a480f34514e89237f3ccda3
diff --git a/packages/@aws-cdk/aws-glue/README.md b/packages/@aws-cdk/aws-glue/README.md
@@ -100,7 +100,7 @@ new glue.Job(stack, 'RayJob', {
   executable: glue.JobExecutable.pythonRay({
     glueVersion: glue.GlueVersion.V4_0,
     pythonVersion: glue.PythonVersion.THREE_NINE,
-    script,
+    script: glue.Code.fromAsset(path.join(__dirname, 'job-script/hello_world.py')),
   }),
   workerType: glue.WorkerType.Z_2X,
   workerCount: 2,

diff --git a/packages/@aws-cdk/aws-glue/lib/job-executable.ts b/packages/@aws-cdk/aws-glue/lib/job-executable.ts
@@ -320,10 +320,18 @@ export class JobExecutable {
       if (config.language !== JobLanguage.PYTHON) {
         throw new Error('Python shell requires the language to be set to Python');
       }
-      if ([GlueVersion.V0_9, GlueVersion.V2_0, GlueVersion.V3_0, GlueVersion.V4_0].includes(config.glueVersion)) {
+      if ([GlueVersion.V0_9, GlueVersion.V3_0, GlueVersion.V4_0].includes(config.glueVersion)) {
         throw new Error(`Specified GlueVersion ${config.glueVersion.name} does not support Python Shell`);
       }
     }
+    if (JobType.RAY === config.type) {
+      if (config.language !== JobLanguage.PYTHON) {
+        throw new Error('Ray requires the language to be set to Python');
+      }
+      if ([GlueVersion.V0_9, GlueVersion.V1_0, GlueVersion.V2_0, GlueVersion.V3_0].includes(config.glueVersion)) {
+        throw new Error(`Specified GlueVersion ${config.glueVersion.name} does not support Ray`);
+      }
+    }
     if (config.extraJarsFirst && [GlueVersion.V0_9, GlueVersion.V1_0].includes(config.glueVersion)) {
       throw new Error(`Specified GlueVersion ${config.glueVersion.name} does not support extraJarsFirst`);
     }

diff --git a/packages/@aws-cdk/aws-glue/lib/job.ts b/packages/@aws-cdk/aws-glue/lib/job.ts
@@ -736,6 +736,8 @@ export class Job extends JobBase {
   private setupSparkUI(executable: JobExecutableConfig, role: iam.IRole, props: SparkUIProps) {
     if (JobType.PYTHON_SHELL === executable.type) {
       throw new Error('Spark UI is not available for JobType.PYTHON_SHELL jobs');
+    } else if (JobType.RAY === executable.type) {
+      throw new Error('Spark UI is not available for JobType.RAY jobs');
     }
 
     const bucket = props.bucket ?? new s3.Bucket(this, 'SparkUIBucket');

diff --git a/packages/@aws-cdk/aws-glue/test/job-executable.test.ts b/packages/@aws-cdk/aws-glue/test/job-executable.test.ts
@@ -31,6 +31,8 @@ describe('JobType', () => {
 
   test('.PYTHON_SHELL should set the name correctly', () => expect(glue.JobType.PYTHON_SHELL.name).toEqual('pythonshell'));
 
+  test('.RAY should set the name correctly', () => expect(glue.JobType.RAY.name).toEqual('glueray'));
+
   test('of(customName) should set the name correctly', () => expect(glue.JobType.of('CustomName').name).toEqual('CustomName'));
 });
 
@@ -65,6 +67,15 @@ describe('JobExecutable', () => {
       })).toThrow(/Python shell requires the language to be set to Python/);
     });
 
+    test('with JobType.RAY and a language other than JobLanguage.PYTHON should throw', () => {
+      expect(() => glue.JobExecutable.of({
+        glueVersion: glue.GlueVersion.V4_0,
+        type: glue.JobType.RAY,
+        language: glue.JobLanguage.SCALA,
+        script,
+      })).toThrow(/Ray requires the language to be set to Python/);
+    });
+
     test('with a non JobLanguage.PYTHON and extraPythonFiles set should throw', () => {
       expect(() => glue.JobExecutable.of({
         glueVersion: glue.GlueVersion.V3_0,
@@ -76,7 +87,7 @@ describe('JobExecutable', () => {
       })).toThrow(/extraPythonFiles is not supported for languages other than JobLanguage.PYTHON/);
     });
 
-    [glue.GlueVersion.V0_9, glue.GlueVersion.V2_0, glue.GlueVersion.V3_0, glue.GlueVersion.V4_0].forEach((glueVersion) => {
+    [glue.GlueVersion.V0_9, glue.GlueVersion.V3_0, glue.GlueVersion.V4_0].forEach((glueVersion) => {
       test(`with JobType.PYTHON_SHELL and GlueVersion ${glueVersion} should throw`, () => {
         expect(() => glue.JobExecutable.of({
           type: glue.JobType.PYTHON_SHELL,
@@ -113,7 +124,7 @@ describe('JobExecutable', () => {
       });
     });
 
-    test('with PythonVersion set to PythonVersion.THREE_NINE and JobType not pythonshell should throw', () => {
+    test('with PythonVersion set to PythonVersion.THREE_NINE and JobType etl should throw', () => {
       expect(() => glue.JobExecutable.of({
         type: glue.JobType.ETL,
         language: glue.JobLanguage.PYTHON,
@@ -132,5 +143,15 @@ describe('JobExecutable', () => {
         script,
       })).toBeDefined();
     });
+
+    test('with PythonVersion PythonVersion.THREE_NINE and JobType ray should succeed', () => {
+      expect(glue.JobExecutable.of({
+        type: glue.JobType.RAY,
+        glueVersion: glue.GlueVersion.V4_0,
+        language: glue.JobLanguage.PYTHON,
+        pythonVersion: glue.PythonVersion.THREE_NINE,
+        script,
+      })).toBeDefined();
+    });
   });
 });
diff --git a/packages/@aws-cdk/aws-glue/test/job.test.ts b/packages/@aws-cdk/aws-glue/test/job.test.ts
@@ -14,6 +14,10 @@ describe('WorkerType', () => {
 
   test('.G_2X should set the name correctly', () => expect(glue.WorkerType.G_2X.name).toEqual('G.2X'));
 
+  test('.G_025X should set the name correctly', () => expect(glue.WorkerType.G_025X.name).toEqual('G.025X'));
+
+  test('.Z_2X should set the name correctly', () => expect(glue.WorkerType.Z_2X.name).toEqual('Z.2X'));
+
   test('of(customType) should set name correctly', () => expect(glue.WorkerType.of('CustomType').name).toEqual('CustomType'));
 });
 
@@ -604,6 +608,33 @@ describe('Job', () => {
       });
     });
 
+    describe('ray job', () => {
+      test('with unsupported glue version should throw', () => {
+        expect(() => new glue.Job(stack, 'Job', {
+          executable: glue.JobExecutable.pythonRay({
+            glueVersion: glue.GlueVersion.V3_0,
+            pythonVersion: glue.PythonVersion.THREE_NINE,
+            script,
+          }),
+          workerType: glue.WorkerType.Z_2X,
+          workerCount: 2,
+        })).toThrow('Specified GlueVersion 3.0 does not support Ray');
+      });
+
+      test('with unsupported Spark UI prop should throw', () => {
+        expect(() => new glue.Job(stack, 'Job', {
+          executable: glue.JobExecutable.pythonRay({
+            glueVersion: glue.GlueVersion.V4_0,
+            pythonVersion: glue.PythonVersion.THREE_NINE,
+            script,
+          }),
+          workerType: glue.WorkerType.Z_2X,
+          workerCount: 2,
+          sparkUI: { enabled: true },
+        })).toThrow('Spark UI is not available for JobType.RAY');
+      });
+    });
+
 
     test('etl job with all props should synthesize correctly', () => {
       new glue.Job(stack, 'Job', {