Create a job schedule using vertex scheduler - DO NOT MERGE, FOR REVIEW ONLY.

dataproc_jupyter_plugin/models/models.py

@@ -106,7 +106,7 @@ class DescribeUpdateVertexJob(BaseModel):
     max_run_count: str = None
     region: Optional[str] = None
     cloud_storage_bucket: Optional[str] = None
-    parameters: Optional[dict] = None
+    parameters: Optional[List[str]] = None
     service_account: Optional[str] = None
     network: Optional[str] = None
     subnetwork: Optional[str] = None

dataproc_jupyter_plugin/services/vertex.py

@@ -432,6 +432,10 @@ async def update_schedule(self, region_id, schedule_id, input_data):
                 if data.time_zone == "UTC"
                 else f"TZ={data.time_zone} {schedule_value}"
             )
+            # getting list of strings from UI, the api accepts dictionary, so converting it
+            labels = {
+                param.split(":")[0]: param.split(":")[1] for param in data.parameters
+            }
 
             if data.kernel_name:
                 notebook_execution_job["kernelName"] = data.kernel_name
@@ -440,7 +444,7 @@ async def update_schedule(self, region_id, schedule_id, input_data):
             if data.cloud_storage_bucket:
                 notebook_execution_job["gcsOutputUri"] = data.cloud_storage_bucket
             if data.parameters:
-                notebook_execution_job["labels"] = data.parameters
+                notebook_execution_job["labels"] = labels
             if data.machine_type:
                 custom_environment_spec["machineSpec"] = {
                     "machineType": data.machine_type.split(" ", 1)[0],

src/Services/Storage.tsx

@@ -80,6 +80,7 @@ export class StorageServices {
         fileName: string | undefined,
         jobRunId: string | undefined,
         setJobDownloadLoading: (value: boolean) => void,
+        scheduleName: string
     ) => {
         try {
             const bucketName = gcsUrl?.split('//')[1];
@@ -89,7 +90,7 @@ export class StorageServices {
             });
             if (formattedResponse.status === 0) {
                 toast.success(
-                    `Job history downloaded successfully`,
+                    `${scheduleName} job history downloaded successfully`,
                     toastifyCustomStyle
                 );
             } else {

src/Services/Vertex.tsx

@@ -86,8 +86,13 @@ export class VertexServices {
         payload: ICreatePayload,
         setCreateCompleted: (value: boolean) => void,
         setCreatingVertexScheduler: (value: boolean) => void,
+        gcsPath: string
     ) => {
         setCreatingVertexScheduler(true);
+        if(gcsPath) {
+            payload.gcs_notebook_source = gcsPath;
+        }
+
         try {
             const data: any = await requestAPI(`api/vertex/updateSchedule?region_id=${region}&schedule_id=${jobId}`, {
                 body: JSON.stringify(payload),
@@ -335,6 +340,7 @@ export class VertexServices {
         setMaxRuns: (value: string) => void,
         setEditMode: (value: boolean) => void,
         setJobNameSelected: (value: string) => void,
+        setGcsPath: (value: string) => void
     ) => {
         setEditDagLoading(job_id);
         try {
@@ -348,6 +354,7 @@ export class VertexServices {
                 setCreateCompleted(false);
                 setRegion(region);
                 setJobNameSelected(formattedResponse.displayName);
+                setGcsPath(formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.gcsNotebookSource.uri);
 
                 // Machine type selection
                 setMachineTypeSelected(formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.customEnvironmentSpec.machineSpec.machineType)
@@ -370,9 +377,9 @@ export class VertexServices {
 
                 // Network
                 const primaryNetwork = formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.customEnvironmentSpec.networkSpec.network.split('/');
-                setPrimaryNetworkSelected({ name: primaryNetwork[primaryNetwork.length - 1], link: primaryNetwork[primaryNetwork.length - 1] });
+                setPrimaryNetworkSelected({ name: primaryNetwork[primaryNetwork.length - 1], link: formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.customEnvironmentSpec.networkSpec.network });
                 const subnetwork = formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.customEnvironmentSpec.networkSpec.subnetwork.split('/');
-                setSubNetworkSelected({ name: subnetwork[subnetwork.length - 1], link: subnetwork[subnetwork.length - 1] })
+                setSubNetworkSelected({ name: subnetwork[subnetwork.length - 1], link: formattedResponse.createNotebookExecutionJobRequest.notebookExecutionJob.customEnvironmentSpec.networkSpec.subnetwork })
 
                 setSubNetworkList([{ name: subnetwork[subnetwork.length - 1], link: subnetwork[subnetwork.length - 1] }])
                 if (formattedResponse.cron === '* * * * *' && formattedResponse.maxRunCount === '1') {

src/scheduler/VertexScheduler/CreateVertexScheduler.tsx

@@ -130,7 +130,8 @@ const CreateVertexScheduler = ({
     const timezones = Object.keys(tzdata.zones).sort();
     const [startDate, setStartDate] = useState<dayjs.Dayjs | null>(dayjs());
     const [endDate, setEndDate] = useState<dayjs.Dayjs | null>(dayjs());
-    const [endDateError, setEndDateError] = useState<boolean>(false)
+    const [endDateError, setEndDateError] = useState<boolean>(false);
+    const [gcsPath, setGcsPath] = useState('');
 
     /**
     * Changing the region value and empyting the value of machineType, accelratorType and accelratorCount
@@ -560,6 +561,7 @@ const CreateVertexScheduler = ({
                 payload,
                 setCreateCompleted,
                 setCreatingVertexScheduler,
+                gcsPath
             );
         } else {
             await VertexServices.createVertexSchedulerService(
@@ -627,7 +629,7 @@ const CreateVertexScheduler = ({
     }, [editMode]);
 
     useEffect(() => {
-        if(editMode) {
+        if(!editMode) {
             setStartDate(null);
             setEndDate(null);
         }
@@ -666,12 +668,9 @@ const CreateVertexScheduler = ({
                         setStartDate={setStartDate}
                         setEndDate={setEndDate}
                         setMaxRuns={setMaxRuns}
-                        setTimeZoneSelected={setTimeZoneSelected}
                         setEditMode={setEditMode}
                         setJobNameSelected={setJobNameSelected}
-                        setServiceAccountList={setServiceAccountList}
-                        setPrimaryNetworkList={setPrimaryNetworkList}
-                        setNetworkSelected={setNetworkSelected}
+                        setGcsPath = {setGcsPath}
                     />
                     :
                     <div className='submit-job-container'>

src/scheduler/VertexScheduler/ListVertexScheduler.tsx

@@ -66,12 +66,9 @@ function ListVertexScheduler({
   setStartDate,
   setEndDate,
   setMaxRuns,
-  setTimeZoneSelected,
   setEditMode,
   setJobNameSelected,
-  setServiceAccountList,
-  setPrimaryNetworkList,
-  setNetworkSelected
+  setGcsPath
 }: {
   region: string;
   setRegion: (value: string) => void;
@@ -92,7 +89,6 @@ function ListVertexScheduler({
   setParameterDetailUpdated: (value: string[]) => void;
   setServiceAccountSelected: (value: { displayName: string; email: string } | null) => void;
   setPrimaryNetworkSelected: (value: { name: string; link: string } | null) => void;
-  setPrimaryNetworkList: (value: { name: string; link: string }[]) => void;
   setSubNetworkSelected: (value: { name: string; link: string } | null) => void;
   setSubNetworkList: (value: { name: string; link: string }[]) => void;
   setSharedNetworkSelected: (value: { name: string; network: string, subnetwork: string } | null) => void;
@@ -101,11 +97,9 @@ function ListVertexScheduler({
   setStartDate: (value: dayjs.Dayjs | null) => void;
   setEndDate: (value: dayjs.Dayjs | null) => void;
   setMaxRuns: (value: string) => void;
-  setTimeZoneSelected: (value: string) => void;
   setEditMode: (value: boolean) => void;
   setJobNameSelected: (value: string) => void;
-  setServiceAccountList: (value: { displayName: string; email: string }[]) => void;
-  setNetworkSelected: (value: string) => void;
+  setGcsPath: (value: string) => void;
 }) {
   const [isLoading, setIsLoading] = useState<boolean>(true);
   const [dagList, setDagList] = useState<IDagList[]>([]);
@@ -286,6 +280,7 @@ function ListVertexScheduler({
         setMaxRuns,
         setEditMode,
         setJobNameSelected,
+        setGcsPath
       );
     }
   };
@@ -469,11 +464,19 @@ function ListVertexScheduler({
               <>
                 <div className='execution-history-main-wrapper'>
                   {cell.row.original.lastScheduledRunResponse === null ? 
-                    <IconActive.react
-                        tag="div"
-                        title=''
-                        className="icon-white logo-alignment-style success_icon icon-size-status"
-                      /> 
+                    (cell.row.original.status === 'ACTIVE' ? 
+                      <IconActive.react
+                          tag="div"
+                          title='ACTIVE'
+                          className="icon-white logo-alignment-style success_icon icon-size-status"
+                        />  
+                        :
+                        <IconListPause.react
+                          tag="div"
+                          title='PAUSE'
+                          className="icon-white logo-alignment-style success_icon icon-size"
+                        />
+                      )
                   : 
                   (cell.row.original.lastScheduledRunResponse && cell.row.original.lastScheduledRunResponse.runResponse ? (cell.row.original.status === 'COMPLETED' ? (cell.row.original.lastScheduledRunResponse.runResponse === 'OK' ? <div>
                     <IconSuccess.react

src/scheduler/VertexScheduler/VertexExecutionHistory.tsx

@@ -283,6 +283,7 @@ const VertexExecutionHistory = ({
                         <VertexJobRuns
                             region={region}
                             schedulerData={schedulerData}
+                            scheduleName={scheduleName}
                             dagId={scheduleName}
                             setJobRunsData={setJobRunsData}
                             setJobRunId={setJobRunId}

src/scheduler/VertexScheduler/VertexInterfaces.tsx

@@ -45,6 +45,7 @@ export interface ICreatePayload {
     subnetwork: string | undefined;
     start_time: dayjs.Dayjs | null;
     end_time: dayjs.Dayjs | null;
+    gcs_notebook_source?: string;
 }
 
 export interface IUpdateSchedulerAPIResponse {

src/scheduler/VertexScheduler/VertexJobRuns.tsx

@@ -29,6 +29,7 @@ import { StorageServices } from '../../Services/Storage';
 const VertexJobRuns = ({
     region,
     schedulerData,
+    scheduleName,
     dagId,
     setJobRunsData,
     setJobRunId,
@@ -47,6 +48,7 @@ const VertexJobRuns = ({
 }: {
     region: string;
     schedulerData: ISchedulerData | undefined;
+    scheduleName: string;
     dagId: string;
     setJobRunsData: React.Dispatch<React.SetStateAction<IDagRunList | undefined>>;
     setJobRunId: (value: string) => void;
@@ -238,7 +240,8 @@ const VertexJobRuns = ({
             data.gcsUrl,
             data.fileName,
             data.jobRunId,
-            setJobDownloadLoading
+            setJobDownloadLoading,
+            scheduleName
         );
     };
 

src/scheduler/VertexScheduler/VertexScheduleJobs.tsx

@@ -55,12 +55,9 @@ const VertexScheduleJobs = ({
     setStartDate,
     setEndDate,
     setMaxRuns,
-    setTimeZoneSelected,
     setEditMode,
     setJobNameSelected,
-    setServiceAccountList,
-    setPrimaryNetworkList,
-    setNetworkSelected
+    setGcsPath
 }: {
     app: JupyterLab;
     themeManager: IThemeManager;
@@ -82,7 +79,6 @@ const VertexScheduleJobs = ({
     setParameterDetailUpdated: (value: string[]) => void;
     setServiceAccountSelected: (value: { displayName: string; email: string } | null) => void;
     setPrimaryNetworkSelected: (value: { name: string; link: string } | null) => void;
-    setPrimaryNetworkList: (value: { name: string; link: string }[]) => void;
     setSubNetworkSelected: (value: { name: string; link: string } | null) => void;
     setSubNetworkList: (value: { name: string; link: string }[]) => void;
     setSharedNetworkSelected: (value: { name: string; network: string, subnetwork: string } | null) => void;
@@ -91,11 +87,9 @@ const VertexScheduleJobs = ({
     setStartDate: (value: dayjs.Dayjs | null) => void;
     setEndDate: (value: dayjs.Dayjs | null) => void;
     setMaxRuns: (value: string) => void;
-    setTimeZoneSelected: (value: string) => void;
     setEditMode: (value: boolean) => void;
     setJobNameSelected?: (value: string) => void;
-    setServiceAccountList: (value: { displayName: string; email: string }[]) => void;
-    setNetworkSelected: (value: string) => void;
+    setGcsPath: (value: string) => void;
 }): React.JSX.Element => {
     const [showExecutionHistory, setShowExecutionHistory] = useState<boolean>(false);
     const [schedulerData, setScheduleData] = useState<ISchedulerData>();
@@ -161,12 +155,9 @@ const VertexScheduleJobs = ({
                         setStartDate={setStartDate}
                         setEndDate={setEndDate}
                         setMaxRuns={setMaxRuns}
-                        setTimeZoneSelected={setTimeZoneSelected}
                         setEditMode={setEditMode}
                         setJobNameSelected={setJobNameSelected!}
-                        setServiceAccountList={setServiceAccountList}
-                        setPrimaryNetworkList={setPrimaryNetworkList}
-                        setNetworkSelected={setNetworkSelected}
+                        setGcsPath={setGcsPath}
                     />
                 </div>
             )}
@@ -222,12 +213,8 @@ export class NotebookJobs extends DataprocWidget {
                     throw new Error('Function not implemented.');
                 } } setServiceAccountSelected={function (value: { displayName: string; email: string; } | null): void {
                     throw new Error('Function not implemented.');
-                } } setNetworkSelected={function (value: string): void {
-                    throw new Error('Function not implemented.');
                 } } setPrimaryNetworkSelected={function (value: { name: string; link: string; } | null): void {
                     throw new Error('Function not implemented.');
-                } } setPrimaryNetworkList={function (value: { name: string; link: string; }[]): void {
-                    throw new Error('Function not implemented.');
                 } } setSubNetworkSelected={function (value: { name: string; link: string; } | null): void {
                     throw new Error('Function not implemented.');
                 } } setSubNetworkList={function (value: { name: string; link: string; }[]): void {
@@ -244,14 +231,12 @@ export class NotebookJobs extends DataprocWidget {
                     throw new Error('Function not implemented.');
                 } } setMaxRuns={function (value: string): void {
                     throw new Error('Function not implemented.');
-                } } setTimeZoneSelected={function (value: string): void {
-                    throw new Error('Function not implemented.');
                 } } setEditMode={function (value: boolean): void {
                     throw new Error('Function not implemented.');
-                } } setServiceAccountList={function (value: { displayName: string; email: string; }[]): void {
-                    throw new Error('Function not implemented.');
                 } } setJobId={function (value: string): void {
                     throw new Error('Function not implemented.');
+                } } setGcsPath={function (value: string): void {
+                    throw new Error('Function not implemented.');
                 } }/>
         );
     }

style/notebookScheduler.css

@@ -439,4 +439,12 @@
   display: flex;
   justify-content: flex-end;
   padding: 15px;
+}
+
+.icon-size-status {
+  width: 12px;
+}
+
+.text-icon {
+  line-height: 2;
 }