Merge pull request #165 from /issues/163-add-vpc-option

Add --vpc and --subnet to `cgcloud create` (resolves #163)

core/src/cgcloud/core/box.py

@@ -295,22 +295,30 @@ def _security_group_name( self ):
         """
         return self.role( )
 
-    def __setup_security_groups( self ):
+    def __setup_security_groups( self, vpc_id=None ):
         log.info( 'Setting up security group ...' )
         name = self.ctx.to_aws_name( self._security_group_name( ) )
         try:
             sg = self.ctx.ec2.create_security_group(
                 name=name,
+                vpc_id=vpc_id,
                 description="Security group for box of role %s in namespace %s" % (
                     self.role( ), self.ctx.namespace) )
         except EC2ResponseError as e:
             if e.error_code == 'InvalidGroup.Duplicate':
+                filters = { 'group-name': name }
+                if vpc_id is not None:
+                    filters[ 'vpc-id' ] = vpc_id
                 for attempt in retry_ec2( retry_while=inconsistencies_detected,
                                           retry_for=10 * 60 ):
                     with attempt:
-                        sg = self.ctx.ec2.get_all_security_groups( groupnames=[ name ] )[ 0 ]
+                        sgs = self.ctx.ec2.get_all_security_groups( filters=filters )
+                        assert len( sgs ) == 1
+                        sg = sgs[ 0 ]
             else:
                 raise
+        # It's OK to have two security groups of the same name as long as their VPC is distinct.
+        assert vpc_id is None or sg.vpc_id == vpc_id
         rules = self._populate_security_group( sg.name )
         for rule in rules:
             try:
@@ -326,7 +334,7 @@ def __setup_security_groups( self ):
         # FIXME: What about stale rules? I tried writing code that removes them but gave up. The
         # API in both boto and EC2 is just too brain-dead.
         log.info( '... finished setting up %s.', sg.id )
-        return [ sg.name ]
+        return [ sg.id ]
 
     def _populate_security_group( self, group_name ):
         """
@@ -394,6 +402,7 @@ def __get_image( self, virtualization_types, image_ref=None ):
     def prepare( self, ec2_keypair_globs,
                  instance_type=None, image_ref=None, virtualization_type=None,
                  spot_bid=None, spot_launch_group=None, spot_auto_zone=False,
+                 vpc_id=None, subnet_id=None,
                  **options ):
         """
         Prepare to create an EC2 instance represented by this box. Return a dictionary with
@@ -421,6 +430,23 @@ def prepare( self, ec2_keypair_globs,
         Amazon EC2 to launch a set of Spot instances only if it can launch them all. In addition,
         if the Spot service must terminate one of the instances in a launch group (for example,
         if the Spot price rises above your bid price), it must terminate them all.
+        
+        :param bool spot_auto_zone: Use heuristic to automatically choose the "best" availability 
+        zone to launch spot instances in. Can't be combined with subnet_id. Overrides the 
+        availability zone in the context.
+        
+        :param: str vpc_id: The ID of a VPC to create the instance and associated security group 
+        in. If this argument is None or absent and the AWS account has a default VPC, the default 
+        VPC will be used. This is the most common case. If this argument is None or absent and 
+        the AWS account has EC2 Classic enabled and the selected instance type supports EC2 
+        classic mode, no VPC will be used. If this argument is None or absent and the AWS account 
+        has no default VPC and an instance type that only supports VPC is used, an exception will 
+        be raised.
+        
+        :param: str subnet_id: The ID of a subnet to allocate instance's private IP address from. 
+        Can't be combined with spot_auto_zone. The specified subnet must belong to the specified 
+        VPC (or the default VPC if none was specified) and reside in the context's availability 
+        zone. If this argument is None or absent, a subnet will be chosen automatically.
 
         :param dict options: Additional, role-specific options can be specified. These options
         augment the options associated with the givem image.
@@ -431,6 +457,11 @@ def prepare( self, ec2_keypair_globs,
         if spot_auto_zone and spot_bid is None:
             raise UserError( 'Need a spot bid for automatically chosing a zone for spot instances' )
 
+        if subnet_id is not None and spot_auto_zone:
+            raise UserError( 'Cannot automatically choose an availability zone for spot instances '
+                             'while placing them in an explicitly defined subnet since the subnet '
+                             'implies a specific availability zone.' )
+
         if self.instance_id is not None:
             raise AssertionError( 'Instance already bound or created' )
 
@@ -441,7 +472,19 @@ def prepare( self, ec2_keypair_globs,
         image = self.__get_image( virtualization_types, image_ref )
         self.image_id = image.id
 
-        security_groups = self.__setup_security_groups( )
+        zone = self.ctx.availability_zone
+
+        security_group_ids = self.__setup_security_groups( vpc_id=vpc_id )
+        if vpc_id is not None and subnet_id is None:
+            log.info( 'Looking up suitable subnet for VPC %s in zone %s.', vpc_id, zone )
+            subnets = self.ctx.vpc.get_all_subnets( filters={ 'vpc-id': vpc_id,
+                                                              'availability-zone': zone } )
+            if subnets:
+                subnet_id = subnets[ 0 ].id
+            else:
+                raise UserError( 'There is no subnet belonging to VPC %s in availability zone %s. '
+                                 'Please create a subnet manually using the VPC console.'
+                                 % (vpc_id, zone) )
 
         options = dict( image.tags, **options )
         self._set_instance_options( options )
@@ -457,8 +500,9 @@ def prepare( self, ec2_keypair_globs,
 
         spec = Expando( instance_type=instance_type,
                         key_name=ec2_keypairs[ 0 ].name,
-                        placement=self.ctx.availability_zone,
-                        security_groups=security_groups,
+                        placement=zone,
+                        security_group_ids=security_group_ids,
+                        subnet_id=subnet_id,
                         instance_profile_arn=self.get_instance_profile_arn( ) )
         self._spec_block_device_mapping( spec, image )
         self._spec_spot_market( spec,

core/src/cgcloud/core/commands.py

@@ -454,6 +454,22 @@ def __init__( self, application ):
                      accepted. By default on-demand instances are used. Note that some instance
                      types are not available on the spot market!""" ) )
 
+        self.option( '--vpc', metavar='VPC_ID', type=str, dest='vpc_id',
+                     help=heredoc( """The ID of a VPC to create the instance and associated 
+                     security group in. If this option is absent and the AWS account has a 
+                     default VPC, the default VPC will be used. This is the most common case. If 
+                     this option is absent and the AWS account has EC2 Classic enabled and the 
+                     selected instance type supports EC2 classic mode, no VPC will be used. If 
+                     this option is absent and the AWS account has no default VPC and an instance 
+                     type that only supports VPC is used, an exception will be raised.""" ) )
+
+        self.option( '--subnet', metavar='SUBNET_ID', type=str, dest='subnet_id',
+                     help=heredoc( """The ID of a subnet to allocate the instance's private IP 
+                     address from. Can't be combined with --spot-auto-zone. The specified subnet 
+                     must belong to the specified VPC (or the default VPC if none was given) and 
+                     reside in the availability zone given via CGCLOUD_ZONE or --zone. If this 
+                     option is absent, cgcloud will attempt to choose a subnet automatically.""" ) ) 
+
         self.option( '--spot-launch-group', metavar='NAME',
                      help=heredoc( """The name of an EC2 spot instance launch group. If
                      specified, the spot request will only be fullfilled once all instances in
@@ -529,6 +545,8 @@ def preparation_kwargs( self, options, box ):
                      ec2_keypair_globs=map( resolve_me, options.ec2_keypair_names ),
                      instance_type=options.instance_type,
                      virtualization_type=options.virtualization_type,
+                     vpc_id=options.vpc_id,
+                     subnet_id=options.subnet_id,
                      spot_bid=options.spot_bid,
                      spot_launch_group=options.spot_launch_group,
                      spot_auto_zone=options.spot_auto_zone )

lib/src/cgcloud/lib/context.py

@@ -123,11 +123,10 @@ def __init__( self, availability_zone, namespace ):
         super( Context, self ).__init__( )
 
         self.__iam = None
-        self.__ec2 = None
+        self.__vpc = None
         self.__s3 = None
         self.__sns = None
         self.__sqs = None
-        self.__vpc = None
 
         self.availability_zone = availability_zone
         m = self.availability_zone_re.match( availability_zone )
@@ -159,14 +158,7 @@ def iam( self ):
             self.__iam = self.__aws_connect( iam, 'universal' )
         return self.__iam
 
-    @property
-    def ec2( self ):
-        """
-        :rtype: EC2Connection
-        """
-        if self.__ec2 is None:
-            self.__ec2 = self.__aws_connect( ec2 )
-        return self.__ec2
+    # VPCConnection extends EC2Connection so we can use one instance of the former for both 
 
     @property
     def vpc( self ):
@@ -177,6 +169,8 @@ def vpc( self ):
             self.__vpc = self.__aws_connect( vpc )
         return self.__vpc
 
+    ec2 = vpc
+
     @property
     def s3( self ):
         """
@@ -223,7 +217,7 @@ def __exit__( self, exc_type, exc_val, exc_tb ):
         self.close( )
 
     def close( self ):
-        if self.__ec2 is not None: self.__ec2.close( )
+        if self.__vpc is not None: self.__vpc.close( )
         if self.__s3 is not None: self.__s3.close( )
         if self.__iam is not None: self.__iam.close( )
         if self.__sns is not None: self.__sns.close( )