Merge pull request #40 from jeffreylwang/feature/enable-public-ip-checkbox

Feature/enable public ip checkbox

Author: RanJiang233

.yarnrc.yml

@@ -1 +1,5 @@
+compressionLevel: mixed
+
+enableGlobalCache: false
+
 nodeLinker: node-modules

package.json

@@ -1,6 +1,6 @@
 {
     "name": "scheduler-jupyter-plugin",
-    "version": "0.1.5",
+    "version": "0.1.6",
     "description": "A JupyterLab extension.",
     "keywords": [
         "jupyter",

scheduler_jupyter_plugin/models/models.py

@@ -85,6 +85,7 @@ class DescribeVertexJob(BaseModel):
     disk_type: str = None
     disk_size: str = None
     kms_key_name: Optional[str] = None
+    enable_public_ip: Optional[bool] = True
 
     @classmethod
     def from_dict(cls, data):
@@ -121,6 +122,7 @@ class DescribeUpdateVertexJob(BaseModel):
     disk_type: Optional[str] = None
     disk_size: Optional[str] = None
     kms_key_name: Optional[str] = None
+    enable_public_ip: Optional[bool] = True
 
     @classmethod
     def from_dict(cls, data):

scheduler_jupyter_plugin/services/vertex.py

@@ -170,7 +170,7 @@ async def create_schedule(self, job, file_path, bucket_name):
                 ]["networkSpec"]["network"] = job.network
                 payload["createNotebookExecutionJobRequest"]["notebookExecutionJob"][
                     "customEnvironmentSpec"
-                ]["networkSpec"]["enableInternetAccess"] = "TRUE"
+                ]["networkSpec"]["enableInternetAccess"] = "TRUE" if job.enable_public_ip else "FALSE"
             if job.subnetwork and job.network:
                 payload["createNotebookExecutionJobRequest"]["notebookExecutionJob"][
                     "customEnvironmentSpec"
@@ -524,6 +524,7 @@ async def update_schedule(self, region_id, schedule_id, input_data):
                 custom_environment_spec["networkSpec"] = {
                     "network": data.network,
                     "subnetwork": data.subnetwork,
+                    "enableInternetAccess": "TRUE" if data.enable_public_ip else "FALSE",
                 }
             if data.disk_size or data.disk_type:
                 custom_environment_spec["persistentDiskSpec"] = {

src/scheduler/vertex/CreateVertexScheduler.tsx

@@ -28,7 +28,8 @@ import {
   FormControlLabel,
   Typography,
   Box,
-  CircularProgress
+  CircularProgress,
+  Checkbox
 } from '@mui/material';
 import CalendarMonthIcon from '@mui/icons-material/CalendarMonth';
 import { LocalizationProvider } from '@mui/x-date-pickers/LocalizationProvider';
@@ -195,6 +196,7 @@ const CreateVertexScheduler = ({
     network: string;
     subnetwork: string;
   } | null>(null);
+  const [enablePublicIp, setEnablePublicIp] = useState<boolean>(true);
   const [maxRuns, setMaxRuns] = useState<string>('');
   const [scheduleField, setScheduleField] = useState<string>('');
   const [scheduleMode, setScheduleMode] = useState<scheduleMode>('runNow');
@@ -897,7 +899,8 @@ const CreateVertexScheduler = ({
       end_time: endDate,
       disk_type: diskTypeSelected,
       disk_size: diskSize,
-      parameters: [] // Parameters for future scope
+      parameters: [], // Parameters for future scope
+      enable_public_ip: enablePublicIp
     };
 
     if (acceleratorType && acceleratedCount) {
@@ -1070,6 +1073,7 @@ const CreateVertexScheduler = ({
       setDiskTypeSelected(vertexSchedulerDetails.disk_type);
       setDiskSize(vertexSchedulerDetails.disk_size);
       setGcsPath(vertexSchedulerDetails.gcs_notebook_source ?? '');
+      setEnablePublicIp(vertexSchedulerDetails.enable_public_ip ?? true);
 
       if ('kms_key_name' in vertexSchedulerDetails) {
         if (vertexSchedulerDetails.kms_key_name) {
@@ -1869,6 +1873,27 @@ const CreateVertexScheduler = ({
               )}
             </>
           )}
+
+          {/* Enable public IP checkbox - shown when custom network is configured */}
+          {(primaryNetworkSelected || sharedNetworkSelected) && (
+            <div className="create-scheduler-form-element">
+              <FormControlLabel
+                control={
+                  <Checkbox
+                    checked={enablePublicIp}
+                    onChange={(e) => setEnablePublicIp(e.target.checked)}
+                    disabled={editMode}
+                  />
+                }
+                label={
+                  <Typography variant="body2">
+                    Enable public IP (allows internet access)
+                  </Typography>
+                }
+              />
+            </div>
+          )}
+
           <div className="create-scheduler-label">Schedule</div>
           <div className="create-scheduler-form-element">
             <FormControl>

src/scheduler/vertex/VertexInterfaces.ts

@@ -54,6 +54,7 @@ export interface ICreatePayload {
   disk_size: string;
   gcs_notebook_source?: string;
   kms_key_name?: string;
+  enable_public_ip?: boolean;
 }
 export interface IVertexScheduleList {
   displayName: string;

src/services/Vertex.tsx

@@ -680,7 +680,13 @@ export class VertexServices {
               .diskSizeGb,
           gcs_notebook_source:
             formattedResponse.createNotebookExecutionJobRequest
-              .notebookExecutionJob.gcsNotebookSource.uri
+              .notebookExecutionJob.gcsNotebookSource.uri,
+          enable_public_ip:
+            formattedResponse.createNotebookExecutionJobRequest
+              .notebookExecutionJob.customEnvironmentSpec?.networkSpec
+              ?.enableInternetAccess === 'TRUE'
+              ? true
+              : false
         };
         setCreateCompleted(false);
         setRegion(region);