Repository: firecracker-microvm/firecracker-containerd
Pull request #182 (2019-05-10 16:37:57)

Integrate fccontrol plugin with shim. 
This is a fairly significant reorganization that combines the work done on the
fccontrol plugin with the shim. Significant aspects include:
1. The shim implements the fccontrol interface (over ttrpc) and handles the
   actual VM operations. The plugin is now only responsible for starting shims
   and forwarding fccontrol requests from clients to the correct shim (based on
   VMID).
2. The fccontrol plugin is now completely stateless (as a result of the above
   change).
3. A shim and its VM are now started together, so the shim start routine no
   longer directly starts a shim process; it just gets a VMID and issues a
   CreateVM call to the plugin. This lets the plugin take care of looking to see
   if the shim for the incoming VM already exists (and starting it if not).

This change was biased towards keeping the behavior and interface unchanged, but
there are a few significant updates from the previous state of things:
1. The shim no longer by default shuts down when all tasks have exited; instead
   it continues to run until either StopVM is called or Shutdown is called when
   no tasks are left. Its lifecycle and the VM lifecycle are now the same
   essentially.
   * This results in less surprising behavior; a shim is created via a CreateVM
     call and can exist without any tasks, so it makes sense that it should be
     able to continue to exist after each task has exited too.
   * Otherwise, it's easy for users to, for example, create a VM and issue
     several CreateTask calls for that VM, but the first Task is created and
     exits quickly before the other tasks were setup, which results in the shim
     exiting and the rest of the tasks failing to be created.
   * For situations where the previous behavior is more desirable, a new
     CreateVM parameter, ExitAfterAllTasksGone, can be set to true. This is
     used by the shim start routine when the client did not specify any VMID
     (and thus don't care about VMs, just their task).
2. You can no longer pass Firecracker machine options as part of the
   CreateTask extra data field. If you create a task with a VMID that doesn't
   exist yet or no VMID at all, the VM will be created using the default
   options as configured in the runtime's config file.
   * This greatly simplifies the amount of branches our code has to handle and
     also greatly simplifies the interface we are giving to users without
     reducing the functionality. If users want custom VM options, they just
     need to use the CreateVM call. If they don't care about custom VM options,
     they can continue to just issue a CreateTask request and get the defaults
     specified in our runtime config file.

Signed-off-by: Erik Sipsma <sipsma@amazon.com>

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

- [x] Working e2e (passing e2e tests)
- [x] ~~Use actual plugin in CI system instead of the special containerd binary with plugin baked in~~ This is not going to be pursued right now, we are just going to use the baked containerd binary for now
- [x] ~~Cleanup some parts of the shim logging (it may just be cleanest for the plugin to create a FIFO for the shim)~~ For now, went with just setting the stdout/stderr of the shim process and forwarding the logs to containerd. See the commit message for the new commit for more details.
- [x] Finish cleanup of old way of passing VM parameters via Task (runc options doesn't need to be embedded in firecracker options anymore)



Patch:
diff --git a/agent/service.go b/agent/service.go
index 19df5c1..fea7853 100644
--- a/agent/service.go
+++ b/agent/service.go
@@ -103,25 +103,25 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 
 	extraData, err := unmarshalExtraData(req.Options)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to unmarshal extra data")
 	}
 
 	bundleDir := bundle.Dir(filepath.Join(containerRootDir, req.ID))
 	err = bundleDir.Create()
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to create bundle dir")
 	}
 
 	err = bundleDir.OCIConfig().Write(extraData.JsonSpec)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to write oci config file")
 	}
 
 	// TODO replace with proper drive mounting once that PR is merged. Right now, all containers in
 	// this VM start up with the same rootfs image no matter their configuration
 	err = bundleDir.MountRootfs("/dev/vdb", "ext4", nil)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to mount rootfs device")
 	}
 
 	// Create a runc shim to manage this task
@@ -130,7 +130,7 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 	taskCtx, taskCancel := context.WithCancel(ts.shimCtx)
 	runcService, err := runc.New(taskCtx, req.ID, ts.publisher, taskCancel)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to create runc shim for task")
 	}
 
 	defer func() {
@@ -143,7 +143,7 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 	fifoSet, err := cio.NewFIFOSetInDir(bundleDir.RootPath(), req.ID, req.Terminal)
 	if err != nil {
 		logger.WithError(err).Error("failed opening stdio FIFOs")
-		return nil, err
+		return nil, errors.Wrap(err, "failed to open stdio FIFOs")
 	}
 
 	// Don't try to connect any io streams that weren't requested by the client
@@ -161,7 +161,7 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 
 	task, err := ts.taskManager.AddTask(req.ID, runcService, bundleDir, extraData, fifoSet, taskCtx.Done(), taskCancel)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to add task")
 	}
 
 	logger.Debug("calling runc create")
@@ -180,13 +180,13 @@ func (ts *TaskService) Create(requestCtx context.Context, req *taskAPI.CreateTas
 	// the task to ensure we capture all task output
 	err = <-task.StartStdioProxy(taskCtx, vm.VSockToFIFO, acceptVSock)
 	if err != nil {
-		return nil, err
+		return nil, errors.Wrap(err, "failed to initialize stdio proxy")
 	}
 
 	resp, err := task.Create(requestCtx, req)
 	if err != nil {
 		logger.WithError(err).Error("error creating container")
-		return nil, err
+		return nil, errors.Wrap(err, "failed create container")
 	}
 
 	logger.WithField("pid", resp.Pid).Debugf("create succeeded")
diff --git a/examples/Makefile b/examples/Makefile
index a15aaaa..6c1132a 100644
--- a/examples/Makefile
+++ b/examples/Makefile
@@ -11,8 +11,13 @@
 # express or implied. See the License for the specific language governing
 # permissions and limitations under the License.
 
+# Set this to pass additional commandline flags to the go compiler, e.g. "make test EXTRAGOARGS=-v"
+EXTRAGOARGS?=
+
+SOURCES:=$(shell find . -name '*.go')
 GOMOD := $(shell go env GOMOD)
 GOSUM := $(GOMOD:.mod=.sum)
+DOCKER_IMAGE_TAG?=latest
 
 all: examples
 
@@ -21,6 +26,19 @@ examples: taskworkflow
 taskworkflow: taskworkflow.go $(GOMOD) $(GOSUM)
 	go build -o taskworkflow taskworkflow.go
 
+test-docker: $(SOURCES) $(GOMOD) $(GOSUM)
+	docker run --rm \
+		--privileged \
+		--ipc=host \
+		--volume /dev:/dev \
+		--volume /sys:/sys \
+		--volume /run/udev/control:/run/udev/control \
+		--volume $(shell pwd)/logs:/var/log/firecracker-containerd-test \
+		--workdir="/firecracker-containerd/examples" \
+		--env EXTRAGOARGS="${EXTRAGOARGS}" \
+		localhost/firecracker-containerd-e2etest-naive:${DOCKER_IMAGE_TAG} \
+		"make examples && ./taskworkflow"
+
 clean:
 	- rm -f taskworkflow
 
diff --git a/examples/taskworkflow.go b/examples/taskworkflow.go
index baf2a0c..5e81b32 100644
--- a/examples/taskworkflow.go
+++ b/examples/taskworkflow.go
@@ -28,6 +28,8 @@ import (
 	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/oci"
 	"github.com/firecracker-microvm/firecracker-containerd/proto"
+	fccontrol "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/grpc"
+	"github.com/firecracker-microvm/firecracker-containerd/runtime/firecrackeroci"
 	specs "github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/pkg/errors"
 )
@@ -52,7 +54,7 @@ func main() {
 	}
 }
 
-func taskWorkflow(containerIP string, gateway string, netMask string) error {
+func taskWorkflow(containerIP string, gateway string, netMask string) (err error) {
 	log.Println("Creating containerd client")
 	client, err := containerd.New("/run/containerd/containerd.sock")
 	if err != nil {
@@ -71,6 +73,39 @@ func taskWorkflow(containerIP string, gateway string, netMask string) error {
 		return errors.Wrapf(err, "creating container")
 
 	}
+
+	fcClient := fccontrol.NewFirecrackerClient(client.Conn())
+
+	vmID := "fc-example"
+	createVMRequest := &proto.CreateVMRequest{
+		VMID: vmID,
+	}
+
+	if containerIP != "" {
+		createVMRequest.NetworkInterfaces = []*proto.FirecrackerNetworkInterface{
+			{
+				MacAddress:  macAddress,
+				HostDevName: hostDevName,
+			},
+		}
+		createVMRequest.KernelArgs = fmt.Sprintf(kernelArgsFormat, containerIP, gateway, netMask)
+	}
+
+	_, err = fcClient.CreateVM(ctx, createVMRequest)
+	if err != nil {
+		return errors.Wrap(err, "failed to create VM")
+	}
+
+	defer func() {
+		_, stopErr := fcClient.StopVM(ctx, &proto.StopVMRequest{VMID: vmID})
+		if stopErr != nil {
+			log.Printf("failed to stop VM, err: %v\n", stopErr)
+		}
+		if err == nil {
+			err = stopErr
+		}
+	}()
+
 	log.Printf("Successfully pulled %s image\n", image.Name())
 	container, err := client.NewContainer(
 		ctx,
@@ -82,6 +117,7 @@ func taskWorkflow(containerIP string, gateway string, netMask string) error {
 			oci.WithHostNamespace(specs.NetworkNamespace),
 			oci.WithHostHostsFile,
 			oci.WithHostResolvconf,
+			firecrackeroci.WithVMID(vmID),
 		),
 		containerd.WithRuntime("aws.firecracker", nil),
 	)
@@ -90,26 +126,7 @@ func taskWorkflow(containerIP string, gateway string, netMask string) error {
 	}
 	defer container.Delete(ctx, containerd.WithSnapshotCleanup)
 
-	task, err := container.NewTask(ctx,
-		cio.NewCreator(cio.WithStdio),
-		func(ctx context.Context, _ *containerd.Client, ti *containerd.TaskInfo) error {
-			if containerIP == "" {
-				return nil
-			}
-			// An IP address for the container has been provided. Configure
-			// the VM opts accordingly.
-			firecrackerConfig := &proto.FirecrackerConfig{
-				NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
-					{
-						MacAddress:  macAddress,
-						HostDevName: hostDevName,
-					},
-				},
-				KernelArgs: fmt.Sprintf(kernelArgsFormat, containerIP, gateway, netMask),
-			}
-			ti.Options = firecrackerConfig
-			return nil
-		})
+	task, err := container.NewTask(ctx, cio.NewCreator(cio.WithStdio))
 	if err != nil {
 		return errors.Wrapf(err, "creating task")
 
diff --git a/firecracker-control/common.go b/firecracker-control/common.go
index 0eef3a6..1f056b3 100644
--- a/firecracker-control/common.go
+++ b/firecracker-control/common.go
@@ -13,30 +13,7 @@
 
 package service
 
-import (
-	"time"
-
-	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
-	"github.com/pkg/errors"
-)
-
 const (
 	localPluginID = "fc-control"
 	grpcPluginID  = "fc-control-service"
 )
-
-var (
-	// ErrVMNotFound means that a VM with the given ID not found
-	ErrVMNotFound = errors.New("vm not found")
-)
-
-const (
-	firecrackerStartTimeout = 5 * time.Second
-	defaultCPUTemplate      = models.CPUTemplateT2
-	defaultCPUCount         = 1
-	defaultMemSizeMb        = 128
-	defaultKernelArgs       = "console=ttyS0 noapic reboot=k panic=1 pci=off nomodules rw"
-	defaultFilesPath        = "/var/lib/firecracker-containerd/runtime/"
-	defaultKernelPath       = defaultFilesPath + "default-vmlinux.bin"
-	defaultRootfsPath       = defaultFilesPath + "default-rootfs.img"
-)
diff --git a/firecracker-control/local.go b/firecracker-control/local.go
index 509c3e9..7ce8e98 100644
--- a/firecracker-control/local.go
+++ b/firecracker-control/local.go
@@ -16,25 +16,34 @@ package service
 import (
 	"context"
 	"fmt"
+	"net"
 	"os"
-	"path/filepath"
-	"sync"
+	"os/exec"
+	"strconv"
+	"strings"
+	"syscall"
 
-	"github.com/containerd/containerd/events"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/plugin"
-	"github.com/firecracker-microvm/firecracker-go-sdk"
+	"github.com/containerd/containerd/runtime/v2/shim"
+	"github.com/containerd/ttrpc"
 	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
+	"github.com/firecracker-microvm/firecracker-containerd/internal"
+	fcShim "github.com/firecracker-microvm/firecracker-containerd/internal/shim"
+	"github.com/firecracker-microvm/firecracker-containerd/internal/vm"
 	"github.com/firecracker-microvm/firecracker-containerd/proto"
-	fccontrol "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/grpc"
+	fccontrolGrpc "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/grpc"
+	fccontrolTtrpc "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/ttrpc"
 )
 
 var (
-	_ fccontrol.FirecrackerServer = (*local)(nil)
+	_ fccontrolGrpc.FirecrackerServer = (*local)(nil)
 )
 
 func init() {
@@ -48,31 +57,9 @@ func init() {
 	})
 }
 
-const (
-	startEventName = "/firecracker-vm/start"
-	stopEventName  = "/firecracker-vm/stop"
-)
-
-// Wrap interface in order to properly generate mock with mockgen
-type machine interface {
-	firecracker.MachineIface
-}
-
-type publisher interface {
-	events.Publisher
-}
-
-type instance struct {
-	cfg     *firecracker.Config
-	machine machine
-}
-
 type local struct {
-	vm          map[string]*instance
-	vmLock      sync.RWMutex
-	rootPath    string
-	publisher   publisher
-	findVsockFn func(context.Context) (*os.File, uint32, error)
+	containerdAddress string
+	logger            *logrus.Entry
 }
 
 func newLocal(ic *plugin.InitContext) (*local, error) {
@@ -81,263 +68,256 @@ func newLocal(ic *plugin.InitContext) (*local, error) {
 	}
 
 	return &local{
-		vm:          make(map[string]*instance),
-		rootPath:    ic.Root,
-		publisher:   ic.Events,
-		findVsockFn: findNextAvailableVsockCID,
+		containerdAddress: ic.Address,
+		logger:            log.G(ic.Context),
 	}, nil
 }
 
-// CreateVM creates new Firecracker VM instance
-func (s *local) CreateVM(ctx context.Context, req *proto.CreateVMRequest) (*empty.Empty, error) {
+// CreateVM creates new Firecracker VM instance. It creates a runtime shim for the VM and the forwards
+// the CreateVM request to that shim. If there is already a VM created with the provided VMID, then
+// AlreadyExists is returned.
+func (s *local) CreateVM(requestCtx context.Context, req *proto.CreateVMRequest) (*empty.Empty, error) {
+	var err error
+
 	id := req.GetVMID()
 	if id == "" {
-		return nil, errors.New("invalid VM ID")
+		err = errors.New("invalid VM ID")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	s.vmLock.Lock()
-
-	// Make sure VM ID unique
-	if _, ok := s.vm[id]; ok {
-		s.vmLock.Unlock()
-		return nil, errors.Errorf("VM with id %q already exists", id)
+	ns, err := namespaces.NamespaceRequired(requestCtx)
+	if err != nil {
+		err = errors.Wrap(err, "error retrieving namespace of request")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	// Reserve this ID while we're creating an instance in order to avoid race conditions
-	s.vm[id] = &instance{}
-	s.vmLock.Unlock()
-
-	instance, err := s.newInstance(ctx, id, req)
-
-	s.vmLock.Lock()
-	defer s.vmLock.Unlock()
+	s.logger.Debugf("using namespace: %s", ns)
 
+	// We determine if there is already a shim managing a VM with the current VMID by attempting
+	// to listen on the abstract socket address (which is parameterized by VMID). If we get
+	// EADDRINUSE, then we assume there is already a shim for the VM and return an AlreadyExists error.
+	shimSocketAddress, err := fcShim.SocketAddress(requestCtx, id)
 	if err != nil {
-		log.G(ctx).WithError(err).WithField("vm_id", id).Error("failed to create VM")
-		delete(s.vm, id)
-
+		err = errors.Wrap(err, "failed to obtain shim socket address")
+		s.logger.WithError(err).Error()
 		return nil, err
 	}
 
-	s.vm[id] = instance
-	return &empty.Empty{}, nil
-}
+	shimSocket, err := shim.NewSocket(shimSocketAddress)
+	if isEADDRINUSE(err) {
+		return nil, status.Errorf(codes.AlreadyExists, "VM with ID %q already exists", id)
+	} else if err != nil {
+		err = errors.Wrapf(err, "failed to open shim socket at address %q", shimSocketAddress)
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
 
-func (s *local) newInstance(ctx context.Context, id string, req *proto.CreateVMRequest) (*instance, error) {
-	logger := log.G(ctx).WithField("vm_id", id)
+	// If we're here, there is no pre-existing shim for this VMID, so we spawn a new one
+	defer shimSocket.Close()
 
-	vsockDescriptor, cid, err := s.findVsockFn(ctx)
+	shimDir := vm.ShimDir(ns, id)
+	err = shimDir.Create()
 	if err != nil {
-		return nil, errors.Wrap(err, "failed to find available cid for vsock")
+		err = errors.Wrapf(err, "failed to create VM dir %q", shimDir.RootPath())
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	defer vsockDescriptor.Close()
+	defer func() {
+		if err != nil {
+			removeErr := os.RemoveAll(shimDir.RootPath())
+			if removeErr != nil {
+				s.logger.WithError(removeErr).WithField("path", shimDir.RootPath()).Error("failed to cleanup VM dir")
+			}
+		}
+	}()
 
-	cfg, err := s.buildVMConfiguration(ctx, id, cid, req)
+	// TODO we have to create separate listeners for the fccontrol service and shim service because
+	// containerd does not currently expose the shim server for us to register the fccontrol service with too.
+	// This is likely addressable through some relatively small upstream contributions; the following is a stop-gap
+	// solution until that time.
+	fcSocketAddress, err := fcShim.FCControlSocketAddress(requestCtx, id)
 	if err != nil {
+		err = errors.Wrap(err, "failed to obtain shim socket address")
+		s.logger.WithError(err).Error()
 		return nil, err
 	}
 
-	logger.Info("creating new VM")
-
-	machine, err := firecracker.NewMachine(ctx, *cfg, firecracker.WithLogger(logger))
+	fcSocket, err := shim.NewSocket(fcSocketAddress)
 	if err != nil {
-		logger.WithError(err).Error("failed to create new machine instance")
+		err = errors.Wrapf(err, "failed to open fccontrol socket at address %q", fcSocketAddress)
+		s.logger.WithError(err).Error()
 		return nil, err
 	}
 
-	// Release vsock CID so Firecracker instance can reacquire it.
-	vsockDescriptor.Close()
+	defer fcSocket.Close()
 
-	if err := s.startMachine(ctx, id, machine); err != nil {
-		logger.WithError(err).Error("failed to start VM instance")
+	cmd, err := s.newShim(ns, id, s.containerdAddress, shimSocket, fcSocket)
+	if err != nil {
 		return nil, err
 	}
 
-	logger.Info("successfully started the VM")
+	defer func() {
+		if err != nil {
+			cmd.Process.Kill()
+		}
+	}()
 
-	newInstance := &instance{
-		cfg:     cfg,
-		machine: machine,
+	client, err := s.shimFirecrackerClient(requestCtx, id)
+	if err != nil {
+		err = errors.Wrap(err, "failed to create firecracker shim client")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	return newInstance, nil
-}
-
-func (s *local) buildVMConfiguration(ctx context.Context, id string, cid uint32, req *proto.CreateVMRequest) (*firecracker.Config, error) {
-	logger := log.G(ctx).WithFields(logrus.Fields{
-		"vm_id":  id,
-		"vm_cid": cid,
-	})
-
-	cfg := &firecracker.Config{
-		SocketPath:   filepath.Join(s.rootPath, fmt.Sprintf("vm_%s.socket", id)),
-		LogFifo:      filepath.Join(s.rootPath, fmt.Sprintf("vm_%s_log.fifo", id)),
-		MetricsFifo:  filepath.Join(s.rootPath, fmt.Sprintf("vm_%s_metrics.fifo", id)),
-		VsockDevices: []firecracker.VsockDevice{{Path: "root", CID: cid}},
-		MachineCfg:   machineConfigurationFromProto(req.GetMachineCfg()),
+	resp, err := client.CreateVM(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim CreateVM returned error")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	logger.Debugf("using socket path: %s", cfg.SocketPath)
-
-	// Kernel configuration
+	return resp, nil
+}
 
-	if val := req.GetKernelArgs(); val != "" {
-		cfg.KernelArgs = val
-	} else {
-		cfg.KernelArgs = defaultKernelArgs
+func (s *local) shimFirecrackerClient(requestCtx context.Context, vmID string) (fccontrolTtrpc.FirecrackerService, error) {
+	socketAddr, err := fcShim.FCControlSocketAddress(requestCtx, vmID)
+	if err != nil {
+		err = errors.Wrap(err, "failed to get shim's fccontrol socket address")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	if val := req.GetKernelImagePath(); val != "" {
-		cfg.KernelImagePath = val
-	} else {
-		cfg.KernelImagePath = defaultKernelPath
+	conn, err := shim.Connect(socketAddr, shim.AnonDialer)
+	if err != nil {
+		err = errors.Wrap(err, "failed to connect to shim's fccontrol endpoint")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	// Drives configuration
-
-	var driveBuilder firecracker.DrivesBuilder
-	if root := req.GetRootDrive(); root != nil {
-		driveBuilder = firecracker.NewDrivesBuilder(root.GetPathOnHost())
-	} else {
-		driveBuilder = firecracker.NewDrivesBuilder(defaultRootfsPath)
-	}
+	return fccontrolTtrpc.NewFirecrackerClient(ttrpc.NewClient(conn)), nil
+}
 
-	for _, drive := range req.GetAdditionalDrives() {
-		driveBuilder = addDriveFromProto(driveBuilder, drive)
+// StopVM stops running VM instance by VM ID. This stops the VM, all tasks within the VM and the runtime shim
+// managing the VM.
+func (s *local) StopVM(requestCtx context.Context, req *proto.StopVMRequest) (*empty.Empty, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
 	}
 
-	// TODO: Reserve fake drives here (https://github.com/firecracker-microvm/firecracker-containerd/pull/154)
-
-	cfg.Drives = driveBuilder.Build()
-
-	// Setup network interfaces
-
-	for _, ni := range req.GetNetworkInterfaces() {
-		cfg.NetworkInterfaces = append(cfg.NetworkInterfaces, networkConfigFromProto(ni))
+	resp, err := client.StopVM(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to stop VM")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	return cfg, nil
+	return resp, err
 }
 
-// startMachine attempts to start an instance within timeout and runs a goroutine to monitor VM exit event
-func (s *local) startMachine(ctx context.Context, id string, machine machine) error {
-	logger := log.G(ctx).WithField("vm_id", id)
-
-	logger.Info("starting VM instance")
-
-	ctx, cancel := context.WithTimeout(ctx, firecrackerStartTimeout)
-	defer cancel()
+// GetVMInfo returns metadata for the VM with the given VMID.
+func (s *local) GetVMInfo(requestCtx context.Context, req *proto.GetVMInfoRequest) (*proto.GetVMInfoResponse, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
+	}
 
-	ns, err := namespaces.NamespaceRequired(ctx)
+	resp, err := client.GetVMInfo(requestCtx, req)
 	if err != nil {
-		return err
+		err = errors.Wrap(err, "shim client failed to get vm info")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	logger.Debugf("using namespace: %s", ns)
+	return resp, nil
+}
 
-	if err := machine.Start(ctx); err != nil {
-		logger.WithError(err).Error("failed to start the VM")
-		return err
+// SetVMMetadata sets Firecracker instance metadata for the VM with the given VMID.
+func (s *local) SetVMMetadata(requestCtx context.Context, req *proto.SetVMMetadataRequest) (*empty.Empty, error) {
+	client, err := s.shimFirecrackerClient(requestCtx, req.VMID)
+	if err != nil {
+		return nil, err
 	}
 
-	if err := s.publisher.Publish(ctx, startEventName, &proto.VMStart{VMID: id}); err != nil {
-		logger.WithError(err).Error("failed to publish start event")
-		return err
+	resp, err := client.SetVMMetadata(requestCtx, req)
+	if err != nil {
+		err = errors.Wrap(err, "shim client failed to set vm info")
+		s.logger.WithError(err).Error()
+		return nil, err
 	}
 
-	go func() {
-		ctx := namespaces.WithNamespace(context.Background(), ns)
-
-		if err := machine.Wait(ctx); err != nil {
-			logger.WithError(err).Error("failed to wait the VM")
-		}
-
-		if err := s.publisher.Publish(ctx, stopEventName, &proto.VMStop{VMID: id}); err != nil {
-			logger.WithError(err).Error("failed to publish stop event")
-		}
-
-		logger.Debug("removing machine")
-
-		s.vmLock.Lock()
-		delete(s.vm, id)
-		s.vmLock.Unlock()
-	}()
-
-	return nil
+	return resp, nil
 }
 
-// StopVM stops running VM instance by VM ID
-func (s *local) StopVM(ctx context.Context, req *proto.StopVMRequest) (*empty.Empty, error) {
-	logger := log.G(ctx).WithField("vm_id", req.GetVMID())
+func (s *local) newShim(ns, vmID, containerdAddress string, shimSocket *net.UnixListener, fcSocket *net.UnixListener) (*exec.Cmd, error) {
+	logger := s.logger.WithField("vmID", vmID)
 
-	instance, err := s.getVM(req.GetVMID())
-	if err != nil {
-		return nil, err
+	args := []string{
+		"-namespace", ns,
+		"-address", containerdAddress,
 	}
 
-	logger.Infof("stopping the VM")
+	cmd := exec.Command(internal.ShimBinaryName, args...)
+
+	cmd.Dir = vm.ShimDir(ns, vmID).RootPath()
 
-	if err := instance.machine.StopVMM(); err != nil {
-		logger.WithError(err).Error("failed to stop VM")
+	shimSocketFile, err := shimSocket.File()
+	if err != nil {
+		err = errors.Wrap(err, "failed to get shim socket fd")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
-	logger.Info("successfully stopped the VM")
-	return &empty.Empty{}, nil
-}
-
-func (s *local) GetVMInfo(ctx context.Context, req *proto.GetVMInfoRequest) (*proto.GetVMInfoResponse, error) {
-	id := req.GetVMID()
-	instance, err := s.getVM(id)
+	fcSocketFile, err := fcSocket.File()
 	if err != nil {
+		err = errors.Wrap(err, "failed to get shim fccontrol socket fd")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
-	var cid uint32
-	if len(instance.cfg.VsockDevices) > 0 {
-		cid = instance.cfg.VsockDevices[0].CID
-	}
+	cmd.ExtraFiles = append(cmd.ExtraFiles, shimSocketFile, fcSocketFile)
+	fcSocketFDNum := 2 + len(cmd.ExtraFiles) // "2 +" because ExtraFiles come after stderr (fd #2)
 
-	resp := &proto.GetVMInfoResponse{
-		VMID:            id,
-		ContextID:       cid,
-		SocketPath:      instance.cfg.SocketPath,
-		LogFifoPath:     instance.cfg.LogFifo,
-		MetricsFifoPath: instance.cfg.MetricsFifo,
-	}
+	cmd.Env = append(os.Environ(),
+		fmt.Sprintf("%s=%s", internal.VMIDEnvVarKey, vmID),
+		fmt.Sprintf("%s=%s", internal.FCSocketFDEnvKey, strconv.Itoa(fcSocketFDNum))) // TODO remove after containerd is updated to expose ttrpc server to shim
 
-	return resp, nil
-}
+	cmd.SysProcAttr = &syscall.SysProcAttr{
+		Setpgid: true,
+	}
 
-// SetVMMetadata sets Firecracker instance metadata
-func (s *local) SetVMMetadata(ctx context.Context, req *proto.SetVMMetadataRequest) (*empty.Empty, error) {
-	var (
-		id       = req.GetVMID()
-		metadata = req.GetMetadata()
-	)
+	// shim stderr is just raw text, so pass it through our logrus formatter first
+	cmd.Stderr = logger.WithField("shim_stream", "stderr").WriterLevel(logrus.ErrorLevel)
+	// shim stdout on the other hand is already formatted by logrus, so pass that transparently through to containerd logs
+	cmd.Stdout = logger.Logger.Out
 
-	instance, err := s.getVM(id)
+	err = cmd.Start()
 	if err != nil {
+		err = errors.Wrap(err, "failed to start shim child process")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
-	log.G(ctx).WithField("vm_id", id).Info("updating VM metadata")
-	if err := instance.machine.SetMetadata(ctx, metadata); err != nil {
-		return nil, errors.Wrapf(err, "failed to set metadata for VM %q", id)
+	// make sure to wait after start
+	go func() {
+		logger.Debug("waiting on shim process")
+		waitErr := cmd.Wait()
+		logger.WithError(waitErr).Debug("completed waiting on shim process")
+	}()
+
+	err = shim.SetScore(cmd.Process.Pid)
+	if err != nil {
+		err = errors.Wrap(err, "failed to set OOM Score on shim")
+		logger.WithError(err).Error()
+		return nil, err
 	}
 
-	return &empty.Empty{}, nil
+	return cmd, nil
 }
 
-func (s *local) getVM(id string) (*instance, error) {
-	s.vmLock.RLock()
-	instance, ok := s.vm[id]
-	s.vmLock.RUnlock()
-
-	if !ok {
-		return nil, ErrVMNotFound
-	}
-
-	return instance, nil
+func isEADDRINUSE(err error) bool {
+	return err != nil && strings.Contains(err.Error(), "address already in use")
 }
diff --git a/firecracker-control/vsock.go b/firecracker-control/vsock.go
deleted file mode 100644
index 8848a4d..0000000
--- a/firecracker-control/vsock.go
+++ /dev/null
@@ -1,79 +0,0 @@
-// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//	http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-
-package service
-
-import (
-	"context"
-	"math"
-	"os"
-	"syscall"
-	"unsafe"
-
-	"github.com/pkg/errors"
-	"golang.org/x/sys/unix"
-)
-
-var (
-	sysCall = syscall.Syscall
-)
-
-// findNextAvailableVsockCID finds first available vsock context ID.
-// It uses VHOST_VSOCK_SET_GUEST_CID ioctl which allows some CID ranges to be statically reserved in advance.
-// The ioctl fails with EADDRINUSE if cid is already taken and with EINVAL if the CID is invalid.
-// Taken from https://bugzilla.redhat.com/show_bug.cgi?id=1291851
-//
-// findNextAvailableVsockCID returns a file descriptor to keep context ID acquired.
-// Caller is responsible for closing file (and releasing CID)
-func findNextAvailableVsockCID(ctx context.Context) (*os.File, uint32, error) {
-	const (
-		// Corresponds to VHOST_VSOCK_SET_GUEST_CID in vhost.h
-		ioctlVsockSetGuestCID = uintptr(0x4008AF60)
-		// 0, 1 and 2 are reserved CIDs, see http://man7.org/linux/man-pages/man7/vsock.7.html
-		startCID        = 3
-		maxCID          = math.MaxUint32
-		vsockDevicePath = "/dev/vhost-vsock"
-	)
-
-	file, err := os.OpenFile(vsockDevicePath, syscall.O_RDWR, 0600)
-	if err != nil {
-		return nil, 0, errors.Wrap(err, "failed to open vsock device")
-	}
-
-	for contextID := startCID; contextID < maxCID; contextID++ {
-		select {
-		case <-ctx.Done():
-			return nil, 0, ctx.Err()
-		default:
-			cid := contextID
-			_, _, err = sysCall(
-				unix.SYS_IOCTL,
-				file.Fd(),
-				ioctlVsockSetGuestCID,
-				uintptr(unsafe.Pointer(&cid)))
-
-			switch err {
-			case unix.Errno(0):
-				return file, uint32(contextID), nil
-			case unix.EADDRINUSE:
-				// ID is already taken, try next one
-				continue
-			default:
-				// Fail if we get an error we don't expect
-				return nil, 0, err
-			}
-		}
-	}
-
-	return nil, 0, errors.New("couldn't find any available vsock context id")
-}
diff --git a/internal/common.go b/internal/common.go
index be68769..b7cdfdb 100644
--- a/internal/common.go
+++ b/internal/common.go
@@ -46,4 +46,13 @@ const (
 
 	// BundleRootfsName is the name of the bundle's directory for holding the container's rootfs
 	BundleRootfsName = "rootfs"
+
+	// VMIDEnvVarKey is the environment variable key used to provide a VMID to a shim process
+	VMIDEnvVarKey = "FIRECRACKER_VM_ID"
+
+	// FCSocketFDEnvKey is the environment variable key used to provide the FD of the fccontrol listening socket to a shim
+	FCSocketFDEnvKey = "FCCONTROL_SOCKET_FD"
+
+	// ShimBinaryName is the name of the runtime shim binary
+	ShimBinaryName = "containerd-shim-aws-firecracker"
 )
diff --git a/internal/shim/shim.go b/internal/shim/shim.go
new file mode 100644
index 0000000..84d59cf
--- /dev/null
+++ b/internal/shim/shim.go
@@ -0,0 +1,38 @@
+// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"). You may
+// not use this file except in compliance with the License. A copy of the
+// License is located at
+//
+//	http://aws.amazon.com/apache2.0/
+//
+// or in the "license" file accompanying this file. This file is distributed
+// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+// express or implied. See the License for the specific language governing
+// permissions and limitations under the License.
+
+package shim
+
+import (
+	"context"
+	"path/filepath"
+
+	"github.com/containerd/containerd/runtime/v2/shim"
+)
+
+// SocketAddress is the abstract unix socket address at which a shim
+// with the given namespace and vmID will listen and serve the shim api
+func SocketAddress(namespacedCtx context.Context, vmID string) (string, error) {
+	return shim.SocketAddress(namespacedCtx, vmID)
+}
+
+// FCControlSocketAddress is the abstract unix socket address at which a shim
+// with the given namespace and vmID will listen and serve the fccontrol api
+func FCControlSocketAddress(namespacedCtx context.Context, vmID string) (string, error) {
+	shimSocketAddr, err := SocketAddress(namespacedCtx, vmID)
+	if err != nil {
+		return "", err
+	}
+
+	return filepath.Join(shimSocketAddr, "fccontrol"), nil
+}
diff --git a/internal/vm/dir.go b/internal/vm/dir.go
index decf96a..4ff33d7 100644
--- a/internal/vm/dir.go
+++ b/internal/vm/dir.go
@@ -14,15 +14,29 @@
 package vm
 
 import (
+	"context"
+	"io"
 	"os"
 	"path/filepath"
 
 	"github.com/containerd/containerd/runtime/v2/shim"
+	"github.com/containerd/fifo"
 	"github.com/firecracker-microvm/firecracker-containerd/internal"
 	"github.com/firecracker-microvm/firecracker-containerd/internal/bundle"
 	"github.com/pkg/errors"
+	"golang.org/x/sys/unix"
 )
 
+const (
+	varRunDir = "/var/run/firecracker-containerd/"
+)
+
+// ShimDir holds files, sockets and FIFOs scoped to a single shim managing the
+// VM with the given VMID. It is unique per-VM and containerd namespace.
+func ShimDir(namespace, vmID string) Dir {
+	return Dir(filepath.Join(varRunDir, namespace, vmID))
+}
+
 // Dir represents the root of a firecracker-containerd VM directory, which
 // holds various files, sockets and FIFOs used during VM runtime.
 type Dir string
@@ -49,6 +63,11 @@ func (d Dir) LogFifoPath() string {
 	return filepath.Join(d.RootPath(), internal.ShimLogFifoName)
 }
 
+// OpenLogFifo opens the shim's log fifo as WriteOnly
+func (d Dir) OpenLogFifo(requestCtx context.Context) (io.ReadWriteCloser, error) {
+	return fifo.OpenFifo(requestCtx, d.LogFifoPath(), unix.O_WRONLY|unix.O_NONBLOCK, 0200)
+}
+
 // FirecrackerSockPath returns the path to the unix socket at which the firecracker VMM
 // services its API
 func (d Dir) FirecrackerSockPath() string {
diff --git a/proto/firecracker.pb.go b/proto/firecracker.pb.go
index e9afdf6..2d6f2e5 100644
--- a/proto/firecracker.pb.go
+++ b/proto/firecracker.pb.go
@@ -37,10 +37,12 @@ type CreateVMRequest struct {
 	// Specifies the networking configuration for a VM
 	NetworkInterfaces []*FirecrackerNetworkInterface `protobuf:"bytes,7,rep,name=NetworkInterfaces,json=networkInterfaces,proto3" json:"NetworkInterfaces,omitempty"`
 	// The number of dummy drives to reserve in advance before running FC instance.
-	ContainerCount       int32    `protobuf:"varint,8,opt,name=ContainerCount,json=containerCount,proto3" json:"ContainerCount,omitempty"`
-	XXX_NoUnkeyedLiteral struct{} `json:"-"`
-	XXX_unrecognized     []byte   `json:"-"`
-	XXX_sizecache        int32    `json:"-"`
+	ContainerCount int32 `protobuf:"varint,8,opt,name=ContainerCount,json=containerCount,proto3" json:"ContainerCount,omitempty"`
+	// Whether the VM should exit after all tasks running in it have exited
+	ExitAfterAllTasksGone bool     `protobuf:"varint,9,opt,name=ExitAfterAllTasksGone,json=exitAfterAllTasksGone,proto3" json:"ExitAfterAllTasksGone,omitempty"`
+	XXX_NoUnkeyedLiteral  struct{} `json:"-"`
+	XXX_unrecognized      []byte   `json:"-"`
+	XXX_sizecache         int32    `json:"-"`
 }
 
 func (m *CreateVMRequest) Reset()         { *m = CreateVMRequest{} }
@@ -123,6 +125,13 @@ func (m *CreateVMRequest) GetContainerCount() int32 {
 	return 0
 }
 
+func (m *CreateVMRequest) GetExitAfterAllTasksGone() bool {
+	if m != nil {
+		return m.ExitAfterAllTasksGone
+	}
+	return false
+}
+
 type StopVMRequest struct {
 	VMID                 string   `protobuf:"bytes,1,opt,name=VMID,json=vMID,proto3" json:"VMID,omitempty"`
 	XXX_NoUnkeyedLiteral struct{} `json:"-"`
@@ -326,31 +335,33 @@ func init() {
 func init() { proto.RegisterFile("firecracker.proto", fileDescriptor_a73317e9fb8da571) }
 
 var fileDescriptor_a73317e9fb8da571 = []byte{
-	// 416 bytes of a gzipped FileDescriptorProto
-	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x7c, 0x92, 0xd1, 0x8a, 0x13, 0x31,
-	0x18, 0x85, 0x19, 0xdb, 0xee, 0x76, 0xfe, 0xb2, 0xdb, 0x4e, 0xf0, 0x22, 0x2c, 0x8b, 0x0c, 0x15,
-	0x96, 0x5e, 0xad, 0xa0, 0xd7, 0x82, 0x6b, 0x4b, 0x65, 0xd4, 0x11, 0x49, 0xa1, 0x17, 0xde, 0xc5,
-	0xf4, 0x9f, 0xd9, 0x30, 0x6d, 0x52, 0x93, 0xbf, 0xab, 0x3e, 0x88, 0x0f, 0xe2, 0x1b, 0xca, 0xa4,
-	0xb6, 0x3b, 0x2d, 0x4b, 0xaf, 0x86, 0x9c, 0x7c, 0x67, 0x72, 0xe0, 0x1c, 0x48, 0x0a, 0xed, 0x50,
-	0x39, 0xa9, 0x2a, 0x74, 0xb7, 0x6b, 0x67, 0xc9, 0x5e, 0xf5, 0xe8, 0xf7, 0x1a, 0xfd, 0xf6, 0x30,
-	0xfc, 0xd3, 0x82, 0xfe, 0xd8, 0xa1, 0x24, 0x9c, 0xe7, 0x02, 0x7f, 0x6c, 0xd0, 0x13, 0x63, 0xd0,
-	0x9e, 0xe7, 0xd9, 0x84, 0x47, 0x69, 0x34, 0x8a, 0x45, 0xfb, 0x21, 0xcf, 0x26, 0xec, 0x1d, 0x40,
-	0x2e, 0xd5, 0xbd, 0x36, 0x38, 0x2e, 0x4a, 0xfe, 0x2c, 0x8d, 0x46, 0xbd, 0xd7, 0xe9, 0xed, 0xf4,
-	0xf1, 0xe7, 0xbb, 0x5b, 0x6b, 0x0a, 0x5d, 0x6e, 0x9c, 0x24, 0x6d, 0x8d, 0x80, 0xd5, 0xde, 0xc3,
-	0x46, 0xd0, 0xff, 0x84, 0xce, 0xe0, 0x32, 0x5b, 0xc9, 0x12, 0xbf, 0x4a, 0xba, 0xe7, 0xad, 0xf0,
-	0x40, 0xbf, 0x3a, 0x94, 0xd9, 0x0b, 0x80, 0x2d, 0x79, 0xe7, 0x4a, 0xcf, 0xdb, 0x01, 0x82, 0x6a,
-	0xaf, 0xb0, 0x57, 0x10, 0x0b, 0x6b, 0x69, 0xe2, 0xf4, 0x03, 0xf2, 0x4e, 0x88, 0x92, 0x34, 0xa3,
-	0x84, 0x0b, 0x11, 0xbb, 0x1d, 0xc3, 0xde, 0xc2, 0xe0, 0x6e, 0xb1, 0xd0, 0x75, 0x24, 0xb9, 0x0c,
-	0x92, 0xe7, 0x67, 0x69, 0xeb, 0x69, 0xdf, 0x40, 0x1e, 0xa1, 0xec, 0x23, 0x24, 0x5f, 0x90, 0x7e,
-	0x5a, 0x57, 0x65, 0x86, 0xd0, 0x15, 0x52, 0xa1, 0xe7, 0xe7, 0xc1, 0x7f, 0xdd, 0xf4, 0x1f, 0x43,
-	0x22, 0x31, 0xc7, 0x36, 0x76, 0x03, 0x97, 0x63, 0x6b, 0x48, 0x6a, 0x83, 0x6e, 0x6c, 0x37, 0x86,
-	0x78, 0x37, 0x8d, 0x46, 0x1d, 0x71, 0xa9, 0x0e, 0xd4, 0xe1, 0x4b, 0xb8, 0x98, 0x91, 0x5d, 0x9f,
-	0x2c, 0x65, 0x78, 0x03, 0x83, 0x0f, 0x48, 0xf3, 0x3c, 0x33, 0x85, 0x3d, 0xc5, 0xfd, 0x8d, 0x20,
-	0x69, 0x80, 0x7e, 0x6d, 0x8d, 0xc7, 0x27, 0x6b, 0xbe, 0x86, 0xb8, 0x8e, 0x87, 0xbf, 0x28, 0x9b,
-	0x84, 0x96, 0x2f, 0x44, 0xac, 0x76, 0x42, 0x5d, 0xcc, 0xcc, 0xaa, 0x0a, 0xa9, 0xd1, 0x1e, 0xf8,
-	0xbd, 0xc2, 0x52, 0xe8, 0x7d, 0xb6, 0xe5, 0x54, 0x17, 0x36, 0x00, 0xdb, 0xe6, 0x7a, 0xcb, 0x47,
-	0xa9, 0x1e, 0x41, 0x8e, 0xe4, 0xb4, 0xf2, 0x7b, 0xaa, 0xb3, 0x1d, 0xc1, 0xea, 0x50, 0x1e, 0x4e,
-	0xe1, 0xf9, 0xac, 0x8e, 0x9c, 0x23, 0xc9, 0x85, 0x24, 0x79, 0x6a, 0x9c, 0x57, 0xd0, 0xdd, 0x61,
-	0x21, 0x74, 0x2c, 0xba, 0xab, 0xff, 0xe7, 0xf7, 0xe7, 0xdf, 0x3a, 0x61, 0xe9, 0xdf, 0xcf, 0xc2,
-	0xe7, 0xcd, 0xbf, 0x00, 0x00, 0x00, 0xff, 0xff, 0x58, 0x21, 0x2d, 0x79, 0x12, 0x03, 0x00, 0x00,
+	// 446 bytes of a gzipped FileDescriptorProto
+	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x7c, 0x93, 0x51, 0x8b, 0xd3, 0x40,
+	0x14, 0x85, 0x89, 0x6d, 0x77, 0x9b, 0x5b, 0x76, 0xdb, 0x0c, 0x2e, 0x84, 0x65, 0x91, 0x50, 0x61,
+	0xc9, 0xd3, 0x0a, 0xea, 0xab, 0x60, 0x6d, 0xed, 0x12, 0x35, 0x22, 0x53, 0xe9, 0x83, 0x6f, 0x63,
+	0x7a, 0x93, 0x1d, 0x92, 0xce, 0xd4, 0x99, 0xdb, 0x75, 0xfd, 0x59, 0xfe, 0x0e, 0xff, 0x94, 0x64,
+	0xba, 0xed, 0xb6, 0xa5, 0xf4, 0x29, 0xcc, 0xb9, 0xdf, 0x49, 0x4e, 0xb8, 0x67, 0x20, 0xc8, 0xa5,
+	0xc1, 0xcc, 0x88, 0xac, 0x44, 0x73, 0xb3, 0x30, 0x9a, 0xf4, 0x65, 0x87, 0xfe, 0x2c, 0xd0, 0xae,
+	0x0e, 0xfd, 0x7f, 0x0d, 0xe8, 0x0e, 0x0d, 0x0a, 0xc2, 0x69, 0xca, 0xf1, 0xd7, 0x12, 0x2d, 0x31,
+	0x06, 0xcd, 0x69, 0x9a, 0x8c, 0x42, 0x2f, 0xf2, 0x62, 0x9f, 0x37, 0xef, 0xd3, 0x64, 0xc4, 0xde,
+	0x03, 0xa4, 0x22, 0xbb, 0x93, 0x0a, 0x87, 0x79, 0x11, 0x3e, 0x8b, 0xbc, 0xb8, 0xf3, 0x3a, 0xba,
+	0x19, 0x3f, 0xbd, 0x7c, 0x3d, 0xd5, 0x2a, 0x97, 0xc5, 0xd2, 0x08, 0x92, 0x5a, 0x71, 0x98, 0x6f,
+	0x3c, 0x2c, 0x86, 0xee, 0x67, 0x34, 0x0a, 0xab, 0x64, 0x2e, 0x0a, 0xfc, 0x26, 0xe8, 0x2e, 0x6c,
+	0xb8, 0x0f, 0x74, 0xcb, 0x5d, 0x99, 0xbd, 0x00, 0x58, 0x91, 0x03, 0x53, 0xd8, 0xb0, 0xe9, 0x20,
+	0x28, 0x37, 0x0a, 0x7b, 0x05, 0x3e, 0xd7, 0x9a, 0x46, 0x46, 0xde, 0x63, 0xd8, 0x72, 0x51, 0x82,
+	0xed, 0x28, 0x6e, 0xc0, 0x7d, 0xb3, 0x66, 0xd8, 0x3b, 0xe8, 0x0d, 0x66, 0x33, 0x59, 0x47, 0x12,
+	0x95, 0x93, 0x6c, 0x78, 0x12, 0x35, 0x0e, 0xfb, 0x7a, 0x62, 0x0f, 0x65, 0x9f, 0x20, 0xf8, 0x8a,
+	0xf4, 0x5b, 0x9b, 0x32, 0x51, 0x84, 0x26, 0x17, 0x19, 0xda, 0xf0, 0xd4, 0xf9, 0xaf, 0xb6, 0xfd,
+	0xfb, 0x10, 0x0f, 0xd4, 0xbe, 0x8d, 0x5d, 0xc3, 0xf9, 0x50, 0x2b, 0x12, 0x52, 0xa1, 0x19, 0xea,
+	0xa5, 0xa2, 0xb0, 0x1d, 0x79, 0x71, 0x8b, 0x9f, 0x67, 0x3b, 0x2a, 0x7b, 0x0b, 0x17, 0x1f, 0x1f,
+	0x24, 0x0d, 0x72, 0x42, 0x33, 0xa8, 0xaa, 0xef, 0xc2, 0x96, 0xf6, 0x56, 0x2b, 0x0c, 0xfd, 0xc8,
+	0x8b, 0xdb, 0xfc, 0x02, 0x0f, 0x0d, 0xfb, 0x2f, 0xe1, 0x6c, 0x42, 0x7a, 0x71, 0x74, 0x95, 0xfd,
+	0x6b, 0xe8, 0xdd, 0x22, 0x4d, 0xd3, 0x44, 0xe5, 0xfa, 0x18, 0xf7, 0xd7, 0x83, 0x60, 0x0b, 0xb4,
+	0x0b, 0xad, 0x2c, 0x1e, 0x2c, 0xc7, 0x15, 0xf8, 0xf5, 0x4f, 0xe1, 0x03, 0x25, 0x23, 0xd7, 0x8d,
+	0x33, 0xee, 0x67, 0x6b, 0xa1, 0x5e, 0xe7, 0x44, 0x67, 0x25, 0xd2, 0xd6, 0xce, 0xc1, 0x6e, 0x14,
+	0x16, 0x41, 0xe7, 0x8b, 0x2e, 0xc6, 0x32, 0xd7, 0x0e, 0x58, 0xed, 0xbb, 0x53, 0x3d, 0x49, 0x75,
+	0x75, 0x52, 0x24, 0x23, 0x33, 0xbb, 0xa1, 0x5a, 0xab, 0xea, 0xcc, 0x77, 0xe5, 0xfe, 0x18, 0x9e,
+	0x4f, 0xea, 0xc8, 0x29, 0x92, 0x98, 0x09, 0x12, 0xc7, 0x2a, 0x7d, 0x09, 0xed, 0x35, 0xe6, 0x42,
+	0xfb, 0xbc, 0x3d, 0x7f, 0x3c, 0x7f, 0x38, 0xfd, 0xd1, 0x72, 0xf7, 0xe3, 0xe7, 0x89, 0x7b, 0xbc,
+	0xf9, 0x1f, 0x00, 0x00, 0xff, 0xff, 0xc3, 0xf2, 0x2f, 0xc4, 0x48, 0x03, 0x00, 0x00,
 }
diff --git a/proto/firecracker.proto b/proto/firecracker.proto
index 3da4b0f..f5ffceb 100644
--- a/proto/firecracker.proto
+++ b/proto/firecracker.proto
@@ -29,6 +29,9 @@ message CreateVMRequest {
 
     // The number of dummy drives to reserve in advance before running FC instance.
     int32 ContainerCount = 8;
+
+    // Whether the VM should exit after all tasks running in it have exited
+    bool ExitAfterAllTasksGone = 9;
 }
 
 message StopVMRequest {
diff --git a/proto/types.pb.go b/proto/types.pb.go
index a4266df..a823907 100644
--- a/proto/types.pb.go
+++ b/proto/types.pb.go
@@ -92,97 +92,6 @@ func (m *ExtraData) GetStderrPort() uint32 {
 	return 0
 }
 
-// Message to specify the various Firecracker VM configurations for a task
-type FirecrackerConfig struct {
-	RuncOptions       *types.Any                       `protobuf:"bytes,1,opt,name=RuncOptions,json=runcOptions,proto3" json:"RuncOptions,omitempty"`
-	NetworkInterfaces []*FirecrackerNetworkInterface   `protobuf:"bytes,2,rep,name=NetworkInterfaces,json=networkInterfaces,proto3" json:"NetworkInterfaces,omitempty"`
-	KernelImagePath   string                           `protobuf:"bytes,3,opt,name=KernelImagePath,json=kernelImagePath,proto3" json:"KernelImagePath,omitempty"`
-	KernelArgs        string                           `protobuf:"bytes,4,opt,name=KernelArgs,json=kernelArgs,proto3" json:"KernelArgs,omitempty"`
-	MachineCfg        *FirecrackerMachineConfiguration `protobuf:"bytes,5,opt,name=MachineCfg,json=machineCfg,proto3" json:"MachineCfg,omitempty"`
-	RootDrive         *FirecrackerDrive                `protobuf:"bytes,6,opt,name=RootDrive,json=rootDrive,proto3" json:"RootDrive,omitempty"`
-	// Specifies the additional block device config for the VM. Note that
-	// the snapshotter used might itself present block devices to the be
-	// attached to the VM. This fields allows additional block devices to
-	// that list.
-	AdditionalDrives     []*FirecrackerDrive `protobuf:"bytes,7,rep,name=AdditionalDrives,json=additionalDrives,proto3" json:"AdditionalDrives,omitempty"`
-	XXX_NoUnkeyedLiteral struct{}            `json:"-"`
-	XXX_unrecognized     []byte              `json:"-"`
-	XXX_sizecache        int32               `json:"-"`
-}
-
-func (m *FirecrackerConfig) Reset()         { *m = FirecrackerConfig{} }
-func (m *FirecrackerConfig) String() string { return proto.CompactTextString(m) }
-func (*FirecrackerConfig) ProtoMessage()    {}
-func (*FirecrackerConfig) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{1}
-}
-func (m *FirecrackerConfig) XXX_Unmarshal(b []byte) error {
-	return xxx_messageInfo_FirecrackerConfig.Unmarshal(m, b)
-}
-func (m *FirecrackerConfig) XXX_Marshal(b []byte, deterministic bool) ([]byte, error) {
-	return xxx_messageInfo_FirecrackerConfig.Marshal(b, m, deterministic)
-}
-func (m *FirecrackerConfig) XXX_Merge(src proto.Message) {
-	xxx_messageInfo_FirecrackerConfig.Merge(m, src)
-}
-func (m *FirecrackerConfig) XXX_Size() int {
-	return xxx_messageInfo_FirecrackerConfig.Size(m)
-}
-func (m *FirecrackerConfig) XXX_DiscardUnknown() {
-	xxx_messageInfo_FirecrackerConfig.DiscardUnknown(m)
-}
-
-var xxx_messageInfo_FirecrackerConfig proto.InternalMessageInfo
-
-func (m *FirecrackerConfig) GetRuncOptions() *types.Any {
-	if m != nil {
-		return m.RuncOptions
-	}
-	return nil
-}
-
-func (m *FirecrackerConfig) GetNetworkInterfaces() []*FirecrackerNetworkInterface {
-	if m != nil {
-		return m.NetworkInterfaces
-	}
-	return nil
-}
-
-func (m *FirecrackerConfig) GetKernelImagePath() string {
-	if m != nil {
-		return m.KernelImagePath
-	}
-	return ""
-}
-
-func (m *FirecrackerConfig) GetKernelArgs() string {
-	if m != nil {
-		return m.KernelArgs
-	}
-	return ""
-}
-
-func (m *FirecrackerConfig) GetMachineCfg() *FirecrackerMachineConfiguration {
-	if m != nil {
-		return m.MachineCfg
-	}
-	return nil
-}
-
-func (m *FirecrackerConfig) GetRootDrive() *FirecrackerDrive {
-	if m != nil {
-		return m.RootDrive
-	}
-	return nil
-}
-
-func (m *FirecrackerConfig) GetAdditionalDrives() []*FirecrackerDrive {
-	if m != nil {
-		return m.AdditionalDrives
-	}
-	return nil
-}
-
 // Message to specify network config for a Firecracker VM
 type FirecrackerNetworkInterface struct {
 	MacAddress           string                  `protobuf:"bytes,1,opt,name=MacAddress,json=macAddress,proto3" json:"MacAddress,omitempty"`
@@ -199,7 +108,7 @@ func (m *FirecrackerNetworkInterface) Reset()         { *m = FirecrackerNetworkI
 func (m *FirecrackerNetworkInterface) String() string { return proto.CompactTextString(m) }
 func (*FirecrackerNetworkInterface) ProtoMessage()    {}
 func (*FirecrackerNetworkInterface) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{2}
+	return fileDescriptor_d938547f84707355, []int{1}
 }
 func (m *FirecrackerNetworkInterface) XXX_Unmarshal(b []byte) error {
 	return xxx_messageInfo_FirecrackerNetworkInterface.Unmarshal(m, b)
@@ -272,7 +181,7 @@ func (m *FirecrackerMachineConfiguration) Reset()         { *m = FirecrackerMach
 func (m *FirecrackerMachineConfiguration) String() string { return proto.CompactTextString(m) }
 func (*FirecrackerMachineConfiguration) ProtoMessage()    {}
 func (*FirecrackerMachineConfiguration) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{3}
+	return fileDescriptor_d938547f84707355, []int{2}
 }
 func (m *FirecrackerMachineConfiguration) XXX_Unmarshal(b []byte) error {
 	return xxx_messageInfo_FirecrackerMachineConfiguration.Unmarshal(m, b)
@@ -336,7 +245,7 @@ func (m *FirecrackerDrive) Reset()         { *m = FirecrackerDrive{} }
 func (m *FirecrackerDrive) String() string { return proto.CompactTextString(m) }
 func (*FirecrackerDrive) ProtoMessage()    {}
 func (*FirecrackerDrive) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{4}
+	return fileDescriptor_d938547f84707355, []int{3}
 }
 func (m *FirecrackerDrive) XXX_Unmarshal(b []byte) error {
 	return xxx_messageInfo_FirecrackerDrive.Unmarshal(m, b)
@@ -404,7 +313,7 @@ func (m *FirecrackerRateLimiter) Reset()         { *m = FirecrackerRateLimiter{}
 func (m *FirecrackerRateLimiter) String() string { return proto.CompactTextString(m) }
 func (*FirecrackerRateLimiter) ProtoMessage()    {}
 func (*FirecrackerRateLimiter) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{5}
+	return fileDescriptor_d938547f84707355, []int{4}
 }
 func (m *FirecrackerRateLimiter) XXX_Unmarshal(b []byte) error {
 	return xxx_messageInfo_FirecrackerRateLimiter.Unmarshal(m, b)
@@ -452,7 +361,7 @@ func (m *FirecrackerTokenBucket) Reset()         { *m = FirecrackerTokenBucket{}
 func (m *FirecrackerTokenBucket) String() string { return proto.CompactTextString(m) }
 func (*FirecrackerTokenBucket) ProtoMessage()    {}
 func (*FirecrackerTokenBucket) Descriptor() ([]byte, []int) {
-	return fileDescriptor_d938547f84707355, []int{6}
+	return fileDescriptor_d938547f84707355, []int{5}
 }
 func (m *FirecrackerTokenBucket) XXX_Unmarshal(b []byte) error {
 	return xxx_messageInfo_FirecrackerTokenBucket.Unmarshal(m, b)
@@ -495,7 +404,6 @@ func (m *FirecrackerTokenBucket) GetCapacity() int64 {
 
 func init() {
 	proto.RegisterType((*ExtraData)(nil), "ExtraData")
-	proto.RegisterType((*FirecrackerConfig)(nil), "FirecrackerConfig")
 	proto.RegisterType((*FirecrackerNetworkInterface)(nil), "FirecrackerNetworkInterface")
 	proto.RegisterType((*FirecrackerMachineConfiguration)(nil), "FirecrackerMachineConfiguration")
 	proto.RegisterType((*FirecrackerDrive)(nil), "FirecrackerDrive")
@@ -506,50 +414,42 @@ func init() {
 func init() { proto.RegisterFile("types.proto", fileDescriptor_d938547f84707355) }
 
 var fileDescriptor_d938547f84707355 = []byte{
-	// 718 bytes of a gzipped FileDescriptorProto
-	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x94, 0x94, 0xcd, 0x6e, 0xe3, 0x36,
-	0x10, 0xc7, 0xa1, 0xaa, 0x49, 0xac, 0x91, 0xf3, 0x61, 0xa1, 0x68, 0xdd, 0x34, 0x68, 0x0d, 0x9f,
-	0xdc, 0x8b, 0x03, 0xb8, 0x68, 0x81, 0x1e, 0x82, 0xd6, 0x1f, 0x29, 0xe2, 0xb4, 0x8e, 0x0d, 0x3a,
-	0xed, 0xa1, 0x37, 0x5a, 0x1a, 0xdb, 0xac, 0x25, 0x52, 0xa0, 0x28, 0x27, 0xce, 0x6b, 0xf4, 0x01,
-	0xfa, 0x24, 0x7d, 0x84, 0x7d, 0xa4, 0x05, 0x16, 0xa4, 0x25, 0x45, 0x36, 0xb2, 0xd9, 0xdd, 0x93,
-	0xc0, 0xdf, 0x7f, 0x66, 0x38, 0xf3, 0x27, 0x45, 0x70, 0xd5, 0x26, 0xc6, 0xa4, 0x1d, 0x4b, 0xa1,
-	0xc4, 0xf9, 0xd7, 0x0b, 0x21, 0x16, 0x21, 0x5e, 0x9a, 0xd5, 0x2c, 0x9d, 0x5f, 0x52, 0xbe, 0xd9,
-	0x4a, 0xcd, 0xff, 0x2d, 0x70, 0xae, 0x1f, 0x95, 0xa4, 0x03, 0xaa, 0xa8, 0x77, 0x0e, 0x95, 0xdb,
-	0x44, 0xf0, 0x69, 0x8c, 0x7e, 0xdd, 0x6a, 0x58, 0xad, 0x2a, 0xa9, 0xfc, 0x93, 0xad, 0xbd, 0x9f,
-	0xc0, 0x25, 0x29, 0xf7, 0xc7, 0xb1, 0x62, 0x82, 0x27, 0xf5, 0xcf, 0x1a, 0x56, 0xcb, 0xed, 0x7c,
-	0xd1, 0xde, 0x96, 0x6e, 0xe7, 0xa5, 0xdb, 0x5d, 0xbe, 0x21, 0xae, 0x7c, 0x0e, 0xf4, 0x2e, 0xc0,
-	0x99, 0xaa, 0x80, 0xf1, 0x89, 0x90, 0xaa, 0x6e, 0x37, 0xac, 0xd6, 0x31, 0x71, 0x92, 0x1c, 0x78,
-	0xdf, 0x02, 0x4c, 0x55, 0x20, 0x52, 0x65, 0xe4, 0xcf, 0x8d, 0x0c, 0x49, 0x41, 0x32, 0x1d, 0xa5,
-	0x34, 0xfa, 0x41, 0xa1, 0x67, 0xa4, 0xf9, 0xaf, 0x0d, 0xb5, 0xdf, 0x98, 0x44, 0x5f, 0x52, 0x7f,
-	0x85, 0xb2, 0x2f, 0xf8, 0x9c, 0x2d, 0xf6, 0x7b, 0xb5, 0x3e, 0xb6, 0xd7, 0x5b, 0xa8, 0xdd, 0xa1,
-	0x7a, 0x10, 0x72, 0x35, 0xe4, 0x0a, 0xe5, 0x9c, 0xfa, 0xa8, 0x27, 0xb5, 0x5b, 0x6e, 0xe7, 0xa2,
-	0x5d, 0xda, 0x66, 0x3f, 0x88, 0xd4, 0xf8, 0x7e, 0x9a, 0xd7, 0x82, 0xd3, 0xdf, 0x51, 0x72, 0x0c,
-	0x87, 0x11, 0x5d, 0xe0, 0x84, 0xaa, 0xa5, 0x99, 0xde, 0x21, 0xa7, 0xab, 0x5d, 0xac, 0x67, 0xdc,
-	0x46, 0x76, 0xe5, 0x22, 0x31, 0x1e, 0x38, 0x04, 0x56, 0x05, 0xf1, 0x7e, 0x05, 0x18, 0x51, 0x7f,
-	0xc9, 0x38, 0xf6, 0xe7, 0x0b, 0xe3, 0x81, 0xdb, 0x69, 0x94, 0xdb, 0xc9, 0x55, 0x33, 0x7c, 0x2a,
-	0xa9, 0x1e, 0x86, 0x40, 0x54, 0xe4, 0x78, 0x97, 0xe0, 0x10, 0x21, 0xd4, 0x40, 0xb2, 0x35, 0xd6,
-	0x0f, 0x4d, 0x81, 0x5a, 0xb9, 0x80, 0x11, 0x88, 0x23, 0xf3, 0x18, 0xef, 0x0a, 0xce, 0xba, 0x41,
-	0xc0, 0x74, 0x21, 0x1a, 0x1a, 0x94, 0xd4, 0x8f, 0x8c, 0x0f, 0x2f, 0xe4, 0x9d, 0xd1, 0xbd, 0xd0,
-	0xe6, 0x5b, 0x0b, 0xbe, 0x79, 0xc5, 0x2e, 0x3d, 0xf1, 0x88, 0xfa, 0xdd, 0x20, 0x90, 0x98, 0x6c,
-	0x8f, 0xc7, 0x31, 0xfd, 0x66, 0xc4, 0x6b, 0x80, 0x7b, 0x23, 0x12, 0x35, 0xc0, 0xf5, 0x1d, 0x8d,
-	0xd0, 0xdc, 0x35, 0x87, 0xb8, 0xcb, 0x67, 0xa4, 0x6f, 0x55, 0x37, 0x0c, 0xc5, 0xc3, 0x68, 0x34,
-	0x98, 0x1a, 0x5f, 0x2b, 0xc4, 0xa1, 0x39, 0xf0, 0xae, 0xe0, 0x78, 0xc8, 0x09, 0x55, 0xf8, 0x07,
-	0x8b, 0x98, 0x42, 0x69, 0x4c, 0x75, 0x3b, 0x5f, 0x95, 0x7b, 0x2f, 0xc9, 0xe4, 0x98, 0x95, 0xa3,
-	0xbd, 0x5f, 0xe0, 0x64, 0x9c, 0xaa, 0x72, 0xfe, 0xc1, 0xeb, 0xf9, 0x27, 0x62, 0x27, 0xbc, 0xf9,
-	0x9f, 0x05, 0xdf, 0x7d, 0xe0, 0x7c, 0xf4, 0x8c, 0xfd, 0xc9, 0x9f, 0xf7, 0x18, 0xc5, 0x21, 0x55,
-	0x98, 0x99, 0xe0, 0xfa, 0xcf, 0x48, 0xcf, 0x78, 0xa3, 0xae, 0x39, 0x9d, 0x85, 0x18, 0x18, 0x0f,
-	0x2a, 0xc4, 0x59, 0xe6, 0xc0, 0x78, 0x88, 0xd1, 0x94, 0x3d, 0xe1, 0x88, 0xcd, 0xb2, 0x1f, 0x0b,
-	0xa2, 0x82, 0xe8, 0xec, 0xbf, 0xfc, 0x38, 0xed, 0x8b, 0x94, 0xe7, 0x3f, 0x96, 0xb3, 0xce, 0x41,
-	0xf3, 0x8d, 0x05, 0x67, 0xfb, 0x07, 0xa9, 0x4b, 0x0e, 0x13, 0x82, 0x34, 0x18, 0xf3, 0x70, 0x63,
-	0x3a, 0xaa, 0x10, 0x60, 0x05, 0xf1, 0x9a, 0x50, 0x1d, 0x26, 0xe6, 0x22, 0xe1, 0x9a, 0xf9, 0x98,
-	0xf5, 0x54, 0x65, 0x25, 0xa6, 0x9f, 0x90, 0x09, 0x95, 0x2a, 0x4d, 0x59, 0x90, 0xdd, 0xf7, 0x4a,
-	0x9c, 0xad, 0x75, 0x7d, 0x7d, 0xe1, 0xc7, 0x5c, 0x1f, 0x6e, 0x7e, 0xd1, 0xe3, 0x82, 0x78, 0x3f,
-	0x83, 0xfb, 0x09, 0xa6, 0xbb, 0xb2, 0xe4, 0xf8, 0x13, 0x7c, 0xf9, 0x72, 0x98, 0xf7, 0x23, 0x38,
-	0x3d, 0xca, 0x83, 0x07, 0x16, 0xa8, 0x65, 0xf6, 0x12, 0xec, 0x94, 0xbc, 0x17, 0x2b, 0xe4, 0xbd,
-	0xd4, 0x5f, 0xa1, 0x22, 0xce, 0x2c, 0x8f, 0xf4, 0xbe, 0x07, 0x7b, 0x1c, 0xe7, 0xcf, 0xdc, 0x7b,
-	0x13, 0x6c, 0x11, 0x27, 0xcd, 0xc7, 0x9d, 0xbd, 0x4b, 0xb2, 0x36, 0x6c, 0xcc, 0xf1, 0x9e, 0x45,
-	0xd8, 0x4b, 0x65, 0xa2, 0xcc, 0xf6, 0x36, 0xa9, 0x8a, 0x12, 0xd3, 0xa6, 0x10, 0x9c, 0xb3, 0x30,
-	0xd4, 0xc8, 0xec, 0x67, 0x13, 0x90, 0x05, 0xd1, 0x86, 0xf6, 0x69, 0x4c, 0x7d, 0xa6, 0x36, 0xc6,
-	0x50, 0x9b, 0x54, 0xfc, 0x6c, 0xdd, 0x3b, 0xfa, 0xfb, 0x60, 0xfb, 0x98, 0x1d, 0x9a, 0xcf, 0x0f,
-	0xef, 0x02, 0x00, 0x00, 0xff, 0xff, 0xfc, 0x4d, 0xe2, 0x2e, 0xf7, 0x05, 0x00, 0x00,
+	// 589 bytes of a gzipped FileDescriptorProto
+	0x1f, 0x8b, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x02, 0xff, 0x94, 0x93, 0xd1, 0x6e, 0xd3, 0x3e,
+	0x14, 0xc6, 0x95, 0x7f, 0xff, 0x83, 0xf8, 0xa4, 0x9b, 0x50, 0x84, 0x60, 0x0c, 0x04, 0x55, 0xaf,
+	0xc6, 0x4d, 0x27, 0x0d, 0x81, 0xc4, 0x05, 0x42, 0x6b, 0x3b, 0xb4, 0x22, 0xba, 0x56, 0xee, 0xe0,
+	0x82, 0x3b, 0xd7, 0x39, 0x5d, 0x4d, 0x13, 0x3b, 0xb2, 0x9d, 0x6d, 0xdd, 0xcb, 0xf0, 0x24, 0x3c,
+	0x02, 0x8f, 0x84, 0x84, 0xec, 0x25, 0x69, 0x86, 0x00, 0x89, 0xab, 0xca, 0xbf, 0xef, 0x1c, 0xfb,
+	0x3b, 0xdf, 0x69, 0x20, 0xb2, 0xeb, 0x1c, 0x4d, 0x2f, 0xd7, 0xca, 0xaa, 0xbd, 0x47, 0xe7, 0x4a,
+	0x9d, 0xa7, 0x78, 0xe0, 0x4f, 0xf3, 0x62, 0x71, 0xc0, 0xe4, 0xfa, 0x46, 0xea, 0x7e, 0x0b, 0x80,
+	0x1c, 0x5f, 0x59, 0xcd, 0x86, 0xcc, 0xb2, 0x78, 0x0f, 0xc2, 0xf7, 0x46, 0xc9, 0x59, 0x8e, 0x7c,
+	0x37, 0xe8, 0x04, 0xfb, 0x6d, 0x1a, 0x7e, 0x29, 0xcf, 0xf1, 0x2b, 0x88, 0x68, 0x21, 0xf9, 0x24,
+	0xb7, 0x42, 0x49, 0xb3, 0xfb, 0x5f, 0x27, 0xd8, 0x8f, 0x0e, 0xef, 0xf7, 0x6e, 0xae, 0xee, 0x55,
+	0x57, 0xf7, 0x8e, 0xe4, 0x9a, 0x46, 0x7a, 0x53, 0x18, 0x3f, 0x01, 0x32, 0xb3, 0x89, 0x90, 0x53,
+	0xa5, 0xed, 0x6e, 0xab, 0x13, 0xec, 0x6f, 0x53, 0x62, 0x2a, 0x10, 0x3f, 0x05, 0x98, 0xd9, 0x44,
+	0x15, 0xd6, 0xcb, 0xff, 0x7b, 0x19, 0x4c, 0x4d, 0x4a, 0x1d, 0xb5, 0xf6, 0xfa, 0x56, 0xad, 0x97,
+	0xa4, 0xfb, 0x23, 0x80, 0xc7, 0xef, 0x84, 0x46, 0xae, 0x19, 0x5f, 0xa1, 0x3e, 0x45, 0x7b, 0xa9,
+	0xf4, 0x6a, 0x24, 0x2d, 0xea, 0x05, 0xe3, 0xe8, 0xfa, 0xc7, 0x8c, 0x1f, 0x25, 0x89, 0x46, 0x63,
+	0xfc, 0x4c, 0x84, 0x42, 0x56, 0x93, 0xb8, 0x03, 0xd1, 0x89, 0x32, 0x76, 0x88, 0x17, 0xa7, 0x2c,
+	0x43, 0x3f, 0x15, 0xa1, 0xd1, 0x72, 0x83, 0x9c, 0xff, 0xa3, 0x34, 0x55, 0x97, 0xe3, 0xf1, 0x70,
+	0xe6, 0xfd, 0x87, 0x94, 0xb0, 0x0a, 0xc4, 0x6f, 0x60, 0x7b, 0x24, 0x29, 0xb3, 0xf8, 0x41, 0x64,
+	0xc2, 0xa2, 0xf6, 0x23, 0x44, 0x87, 0x0f, 0x7b, 0x0d, 0x53, 0x0d, 0x99, 0x6e, 0x8b, 0x66, 0x75,
+	0xfc, 0x16, 0x76, 0x26, 0x85, 0x6d, 0xf6, 0x6f, 0xfd, 0xbd, 0x7f, 0x47, 0xdd, 0x2a, 0xef, 0x7e,
+	0x0d, 0xe0, 0x59, 0xa3, 0x74, 0xcc, 0xf8, 0x52, 0x48, 0x1c, 0x28, 0xb9, 0x10, 0xe7, 0x85, 0x66,
+	0x6e, 0x05, 0x6e, 0xc6, 0xc1, 0xf4, 0xe3, 0x19, 0x66, 0x79, 0xca, 0x2c, 0x96, 0x21, 0x44, 0x7c,
+	0x83, 0xdc, 0x8c, 0x27, 0xf6, 0x58, 0xb2, 0x79, 0x8a, 0x89, 0xcf, 0x20, 0xa4, 0x64, 0x59, 0x01,
+	0x9f, 0x21, 0x66, 0x33, 0x71, 0x8d, 0x63, 0x31, 0x2f, 0x57, 0x08, 0x59, 0x4d, 0x5c, 0xf7, 0x27,
+	0x9e, 0x17, 0x03, 0x55, 0xc8, 0x6a, 0x85, 0xe4, 0xa2, 0x02, 0xdd, 0xef, 0x01, 0xdc, 0x6b, 0x38,
+	0x1c, 0x6a, 0x71, 0xe1, 0xd7, 0x32, 0x32, 0x14, 0x59, 0x32, 0x91, 0xe9, 0xda, 0x3b, 0x0a, 0x29,
+	0x88, 0x9a, 0xc4, 0x5d, 0x68, 0x8f, 0x0c, 0x55, 0xca, 0x6d, 0x41, 0x70, 0x2c, 0x3d, 0xb5, 0x45,
+	0x83, 0xb9, 0x3f, 0xeb, 0x94, 0x69, 0x5b, 0x14, 0x22, 0xf1, 0xa6, 0x08, 0x0d, 0xf3, 0xf2, 0xec,
+	0xee, 0x9f, 0x32, 0xbb, 0x9c, 0x48, 0xb7, 0x5c, 0xef, 0x89, 0x50, 0xc8, 0x6b, 0x12, 0xbf, 0x86,
+	0xe8, 0x1f, 0x42, 0x8f, 0x74, 0x23, 0xf1, 0x6b, 0x78, 0xf0, 0xfb, 0xb2, 0xf8, 0x25, 0x90, 0x3e,
+	0x93, 0xc9, 0xa5, 0x48, 0xec, 0xd2, 0xcf, 0xf4, 0xcb, 0x95, 0x67, 0x6a, 0x85, 0xb2, 0x5f, 0xf0,
+	0x15, 0x5a, 0x4a, 0xe6, 0x55, 0x65, 0xfc, 0x1c, 0x5a, 0x93, 0xbc, 0xfa, 0xa0, 0xfe, 0xd8, 0xd0,
+	0x52, 0xb9, 0xe9, 0x5e, 0xdd, 0x7a, 0xbb, 0x21, 0xbb, 0xc0, 0x26, 0x12, 0xcf, 0x44, 0x86, 0xfd,
+	0x42, 0x1b, 0xeb, 0x9f, 0x6f, 0xd1, 0xb6, 0x6a, 0x30, 0x17, 0x0a, 0xc5, 0x85, 0x48, 0x53, 0x87,
+	0xfc, 0x7b, 0x2d, 0x0a, 0xba, 0x26, 0x2e, 0xd0, 0x01, 0xcb, 0x19, 0x17, 0x76, 0xed, 0x03, 0x6d,
+	0xd1, 0x90, 0x97, 0xe7, 0xfe, 0xdd, 0xcf, 0x5b, 0x37, 0x9f, 0xf8, 0x1d, 0xff, 0xf3, 0xe2, 0x67,
+	0x00, 0x00, 0x00, 0xff, 0xff, 0x8c, 0x86, 0x88, 0xb2, 0x61, 0x04, 0x00, 0x00,
 }
diff --git a/proto/types.proto b/proto/types.proto
index 6fab535..fc9438f 100644
--- a/proto/types.proto
+++ b/proto/types.proto
@@ -13,21 +13,6 @@ message ExtraData {
 	uint32 StderrPort = 5;
 }
 
-// Message to specify the various Firecracker VM configurations for a task
-message FirecrackerConfig {
-	google.protobuf.Any RuncOptions = 1; // Specifies Runc options that need to relayed to the CreateTask request
-	repeated FirecrackerNetworkInterface NetworkInterfaces = 2; // Specifies the networking configuration for a VM
-	string KernelImagePath = 3; // Specifies the file path where the kernel image is located
-	string KernelArgs = 4; // Specifies the commandline arguments that should be passed to the kernel
-	FirecrackerMachineConfiguration MachineCfg = 5; // Specifies the machine configuration for the VM
-	FirecrackerDrive RootDrive = 6; // Specifies the root block device for the VM
-	// Specifies the additional block device config for the VM. Note that
-	// the snapshotter used might itself present block devices to the be
-	// attached to the VM. This fields allows additional block devices to
-	// that list.
-	repeated FirecrackerDrive AdditionalDrives = 7;
-}
-
 // Message to specify network config for a Firecracker VM
 message FirecrackerNetworkInterface {
 	string MacAddress = 1; // Specifies the mac address for the the device
diff --git a/runtime/config.go b/runtime/config.go
index 929bfe8..4a4912c 100644
--- a/runtime/config.go
+++ b/runtime/config.go
@@ -19,16 +19,20 @@ import (
 	"os"
 
 	"github.com/pkg/errors"
+
+	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 )
 
 const (
-	configPathEnvName = "FIRECRACKER_CONTAINERD_RUNTIME_CONFIG_PATH"
-	defaultConfigPath = "/etc/containerd/firecracker-runtime.json"
-	defaultKernelArgs = "console=ttyS0 noapic reboot=k panic=1 pci=off nomodules rw"
-	defaultFilesPath  = "/var/lib/firecracker-containerd/runtime/"
-	defaultKernelPath = defaultFilesPath + "default-vmlinux.bin"
-	defaultRootfsPath = defaultFilesPath + "default-rootfs.img"
-	defaultCPUCount   = 1
+	configPathEnvName  = "FIRECRACKER_CONTAINERD_RUNTIME_CONFIG_PATH"
+	defaultConfigPath  = "/etc/containerd/firecracker-runtime.json"
+	defaultKernelArgs  = "console=ttyS0 noapic reboot=k panic=1 pci=off nomodules rw"
+	defaultFilesPath   = "/var/lib/firecracker-containerd/runtime/"
+	defaultKernelPath  = defaultFilesPath + "default-vmlinux.bin"
+	defaultRootfsPath  = defaultFilesPath + "default-rootfs.img"
+	defaultCPUCount    = 1
+	defaultCPUTemplate = models.CPUTemplateT2
+	defaultMemSizeMb   = 128
 )
 
 // Config represents runtime configuration parameters
diff --git a/firecracker-control/helpers.go b/runtime/helpers.go
similarity index 87%
rename from firecracker-control/helpers.go
rename to runtime/helpers.go
index 2a4783d..bba42fc 100644
--- a/firecracker-control/helpers.go
+++ b/runtime/helpers.go
@@ -11,7 +11,7 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
-package service
+package main
 
 import (
 	"time"
@@ -29,19 +29,23 @@ func machineConfigurationFromProto(req *proto.FirecrackerMachineConfiguration) m
 		MemSizeMib:  defaultMemSizeMb,
 	}
 
-	if name := req.GetCPUTemplate(); name != "" {
+	if req == nil {
+		return config
+	}
+
+	if name := req.CPUTemplate; name != "" {
 		config.CPUTemplate = models.CPUTemplate(name)
 	}
 
-	if count := req.GetVcpuCount(); count > 0 {
+	if count := req.VcpuCount; count > 0 {
 		config.VcpuCount = int64(count)
 	}
 
-	if size := req.GetMemSizeMib(); size > 0 {
+	if size := req.MemSizeMib; size > 0 {
 		config.MemSizeMib = int64(size)
 	}
 
-	config.HtEnabled = req.GetHtEnabled()
+	config.HtEnabled = req.HtEnabled
 
 	return config
 }
@@ -68,15 +72,15 @@ func networkConfigFromProto(nwIface *proto.FirecrackerNetworkInterface) firecrac
 
 func addDriveFromProto(builder firecracker.DrivesBuilder, drive *proto.FirecrackerDrive) firecracker.DrivesBuilder {
 	opt := func(d *models.Drive) {
-		d.IsRootDevice = firecracker.Bool(drive.GetIsRootDevice())
-		d.Partuuid = drive.GetPartuuid()
+		d.IsRootDevice = firecracker.Bool(drive.IsRootDevice)
+		d.Partuuid = drive.Partuuid
 
-		if limiter := drive.GetRateLimiter(); limiter != nil {
+		if limiter := drive.RateLimiter; limiter != nil {
 			d.RateLimiter = rateLimiterFromProto(limiter)
 		}
 	}
 
-	return builder.AddDrive(drive.GetPathOnHost(), drive.GetIsReadOnly(), opt)
+	return builder.AddDrive(drive.PathOnHost, drive.IsReadOnly, opt)
 }
 
 // rateLimiterFromProto creates a firecracker RateLimiter object from the
diff --git a/runtime/main.go b/runtime/main.go
index 9d9af1e..cbeaa46 100644
--- a/runtime/main.go
+++ b/runtime/main.go
@@ -14,11 +14,26 @@
 package main
 
 import (
+	"os"
+
+	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/runtime/v2/shim"
+	"github.com/sirupsen/logrus"
 )
 
 const shimID = "aws.firecracker"
 
+func init() {
+	logrus.SetFormatter(&logrus.TextFormatter{
+		TimestampFormat: log.RFC3339NanoFixed,
+		FullTimestamp:   true,
+	})
+
+	logrus.SetOutput(os.Stdout)
+}
+
 func main() {
-	shim.Run(shimID, NewService)
+	shim.Run(shimID, NewService, func(cfg *shim.Config) {
+		cfg.NoSetupLogger = true
+	})
 }
diff --git a/runtime/service.go b/runtime/service.go
index dcf8b0b..ef1902e 100644
--- a/runtime/service.go
+++ b/runtime/service.go
@@ -19,9 +19,8 @@ import (
 	"math"
 	"net"
 	"os"
-	"os/exec"
-	"path/filepath"
 	"runtime/debug"
+	"strconv"
 	"strings"
 	"sync"
 	"syscall"
@@ -33,29 +32,39 @@ import (
 	// secure randomness
 	"math/rand" // #nosec
 
+	"github.com/containerd/containerd"
 	eventsAPI "github.com/containerd/containerd/api/events"
 	"github.com/containerd/containerd/cio"
+	"github.com/containerd/containerd/events"
 	"github.com/containerd/containerd/events/exchange"
 	"github.com/containerd/containerd/log"
 	"github.com/containerd/containerd/namespaces"
 	"github.com/containerd/containerd/runtime"
 	"github.com/containerd/containerd/runtime/v2/shim"
 	taskAPI "github.com/containerd/containerd/runtime/v2/task"
+	"github.com/containerd/fifo"
 	"github.com/containerd/ttrpc"
 	"github.com/containerd/typeurl"
 	"github.com/firecracker-microvm/firecracker-go-sdk"
 	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 	"github.com/gofrs/uuid"
 	ptypes "github.com/gogo/protobuf/types"
+	"github.com/golang/protobuf/ptypes/empty"
 	"github.com/mdlayher/vsock"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/sys/unix"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 
 	"github.com/firecracker-microvm/firecracker-containerd/eventbridge"
+	"github.com/firecracker-microvm/firecracker-containerd/internal"
 	"github.com/firecracker-microvm/firecracker-containerd/internal/bundle"
+	fcShim "github.com/firecracker-microvm/firecracker-containerd/internal/shim"
 	"github.com/firecracker-microvm/firecracker-containerd/internal/vm"
 	"github.com/firecracker-microvm/firecracker-containerd/proto"
+	fccontrolGrpc "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/grpc"
+	fccontrolTtrpc "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/ttrpc"
 )
 
 func init() {
@@ -63,13 +72,15 @@ func init() {
 }
 
 const (
-	defaultVsockPort     = 10789
-	minVsockIOPort       = uint32(11000)
-	supportedMountFSType = "ext4"
+	defaultVsockPort        = 10789
+	minVsockIOPort          = uint32(11000)
+	firecrackerStartTimeout = 5 * time.Second
 
-	vmIDEnvVarKey = "FIRECRACKER_VM_ID"
+	// StartEventName is the topic published to when a VM starts
+	StartEventName = "/firecracker-vm/start"
 
-	varRunDir = "/var/run/firecracker-containerd"
+	// StopEventName is the topic published to when a VM stops
+	StopEventName = "/firecracker-vm/stop"
 )
 
 var (
@@ -84,8 +95,6 @@ var (
 type service struct {
 	taskManager vm.TaskManager
 
-	server        *ttrpc.Server
-	publish       shim.Publisher
 	eventExchange *exchange.Exchange
 	namespace     string
 
@@ -113,11 +122,15 @@ type service struct {
 	vmID   string
 	config *Config
 
-	startVMMutex sync.Mutex
-	agentStarted bool
-	agentClient  taskAPI.TaskService
+	// vmReady is closed once CreateVM has been successfully called
+	vmReady           chan struct{}
+	vmStartOnce       sync.Once
+	rootfsOnce        sync.Once // TODO remove once stub drives are merged
+	agentClient       taskAPI.TaskService
+	eventBridgeClient eventbridge.Getter
 
 	machine          *firecracker.Machine
+	machineConfig    *firecracker.Config
 	machineCID       uint32
 	vsockIOPortCount uint32
 	vsockPortMu      sync.Mutex
@@ -134,11 +147,6 @@ func shimOpts(shimCtx context.Context) (*shim.Opts, error) {
 
 // NewService creates new runtime shim.
 func NewService(shimCtx context.Context, id string, remotePublisher shim.Publisher, shimCancel func()) (shim.Shim, error) {
-	server, err := newServer()
-	if err != nil {
-		return nil, err
-	}
-
 	config, err := LoadConfig("")
 	if err != nil {
 		return nil, err
@@ -158,110 +166,116 @@ func NewService(shimCtx context.Context, id string, remotePublisher shim.Publish
 		namespace = namespaces.Default
 	}
 
-	eventExchange := exchange.NewExchange()
-
-	// Republish each event received on our exchange to the provided remote publisher.
-	// TODO ideally we would be forwarding events instead of re-publishing them, which would
-	// preserve the events' original timestamps and namespaces. However, as of this writing,
-	// the containerd v2 runtime model only provides a shim with a publisher, not a forwarder,
-	// so we have to republish for now.
-	go func() {
-		if err := <-eventbridge.Republish(shimCtx, eventExchange, remotePublisher); err != nil && err != context.Canceled {
-			log.G(shimCtx).WithError(err).Error("error while republishing events")
-		}
-	}()
+	vmID := os.Getenv(internal.VMIDEnvVarKey)
+	logger := log.G(shimCtx)
+	if vmID != "" {
+		logger = logger.WithField("vmID", vmID)
+	}
 
 	s := &service{
-		taskManager: vm.NewTaskManager(log.G(shimCtx)),
+		taskManager: vm.NewTaskManager(logger),
 
-		server:        server,
-		publish:       remotePublisher,
-		eventExchange: eventExchange,
+		eventExchange: exchange.NewExchange(),
 		namespace:     namespace,
 
-		logger:     log.G(shimCtx),
+		logger:     logger,
 		shimCtx:    shimCtx,
 		shimCancel: shimCancel,
 
-		vmID:   os.Getenv(vmIDEnvVarKey),
+		vmID:   vmID,
 		config: config,
-	}
 
-	return s, nil
-}
+		vmReady: make(chan struct{}),
+	}
 
-// vmDir holds files, sockets and FIFOs scoped to a single given VM.
-// It is unique per-VM and containerd namespace.
-func (s *service) vmDir() vm.Dir {
-	return vm.Dir(filepath.Join(varRunDir, s.namespace, s.vmID))
-}
+	s.startEventForwarders(remotePublisher)
 
-func (s *service) newShim(namespace string, containerdBinary, containerdAddress string, shimSocket *net.UnixListener) (*exec.Cmd, error) {
-	self, err := os.Executable()
+	err = s.serveFCControl()
 	if err != nil {
+		err = errors.Wrap(err, "failed to start fccontrol server")
+		s.logger.WithError(err).Error()
 		return nil, err
 	}
 
-	args := []string{
-		"-namespace", namespace,
-		"-address", containerdAddress,
-		"-publish-binary", containerdBinary,
-	}
-
-	if s.config.Debug {
-		args = append(args, "-debug")
-	}
+	return s, nil
+}
 
-	cmd := exec.Command(self, args...)
+func (s *service) startEventForwarders(remotePublisher events.Publisher) {
+	// Republish each event received on our exchange to the provided remote publisher.
+	// TODO ideally we would be forwarding events instead of re-publishing them, which would
+	// preserve the events' original timestamps and namespaces. However, as of this writing,
+	// the containerd v2 runtime model only provides a shim with a publisher, not a forwarder.
+	go func() {
+		err := <-eventbridge.Republish(s.shimCtx, s.eventExchange, remotePublisher)
+		if err != nil && err != context.Canceled {
+			s.logger.WithError(err).Error("error while republishing events")
+		}
+	}()
 
-	cmd.Dir = s.vmDir().RootPath()
+	// Once the VM is ready, also start forwarding events from it to our exchange
+	go func() {
+		<-s.vmReady
+		err := <-eventbridge.Attach(s.shimCtx, s.eventBridgeClient, s.eventExchange)
+		if err != nil && err != context.Canceled {
+			s.logger.WithError(err).Error("error while forwarding events from VM agent")
+		}
+	}()
+}
 
-	cmd.Env = append(os.Environ(),
-		fmt.Sprintf("%s=%s", vmIDEnvVarKey, s.vmID))
+// TODO we have to create separate listeners for the fccontrol service and shim service because
+// containerd does not currently expose the shim server for us to register the fccontrol service with too.
+// This is likely addressable through some relatively small upstream contributions; the following is a stop-gap
+// solution until that time.
+func (s *service) serveFCControl() error {
+	// If the fccontrol socket was configured, setup the fccontrol server
+	fcSocketFDEnvVal := os.Getenv(internal.FCSocketFDEnvKey)
+	if fcSocketFDEnvVal == "" {
+		// if there's no socket, we don't need to serve the API (this must be a shim start or shim delete call)
+		return nil
+	}
 
-	shimSocketFile, err := shimSocket.File()
+	fcServer, err := ttrpc.NewServer(ttrpc.WithServerHandshaker(ttrpc.UnixSocketRequireSameUser()))
 	if err != nil {
-		return nil, err
+		return err
 	}
-	cmd.ExtraFiles = append(cmd.ExtraFiles, shimSocketFile)
 
-	cmd.SysProcAttr = &syscall.SysProcAttr{
-		Setpgid: true,
+	socketFD, err := strconv.Atoi(fcSocketFDEnvVal)
+	if err != nil {
+		err = errors.Wrap(err, "failed to parse fccontrol socket FD value")
+		s.logger.WithError(err).Error()
+		return err
 	}
 
-	err = cmd.Start()
+	fccontrolTtrpc.RegisterFirecrackerService(fcServer, s)
+	fcListener, err := net.FileListener(os.NewFile(uintptr(socketFD), "fccontrol"))
 	if err != nil {
-		return nil, err
+		return err
 	}
 
-	// make sure to wait after start
 	go func() {
-		s.logger.WithField("vmID", s.vmID).Debug("waiting on shim process")
-		waitErr := cmd.Wait()
-		s.logger.WithError(waitErr).WithField("vmID", s.vmID).Debug("completed waiting on shim process")
+		defer fcListener.Close()
+		err := fcServer.Serve(s.shimCtx, fcListener)
+		if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
+			s.logger.WithError(err).Error("fccontrol ttrpc server error")
+		}
 	}()
 
-	err = shim.SetScore(cmd.Process.Pid)
-	if err != nil {
-		s.logger.WithError(err).WithField("vmID", s.vmID).Error("failed to set OOM score on shim")
-		return nil, errors.Wrap(err, "failed to set OOM Score on shim")
-	}
-
-	return cmd, nil
-}
-
-func isEADDRINUSE(err error) bool {
-	return err != nil && strings.Contains(err.Error(), "address already in use")
+	return nil
 }
 
 func (s *service) StartShim(shimCtx context.Context, containerID, containerdBinary, containerdAddress string) (string, error) {
-	s.logger.WithField("id", containerID).Debug("StartShim")
-
-	namespace, err := namespaces.NamespaceRequired(shimCtx)
+	// In the shim start routine, we can assume that containerd provided a "log" FIFO in the current working dir.
+	// We have to use that instead of stdout/stderr because containerd reads the stdio pipes of shim start to get
+	// either the shim address or the error returned here.
+	logFifo, err := fifo.OpenFifo(shimCtx, "log", unix.O_WRONLY, 0200)
 	if err != nil {
-		return "", errors.Wrap(err, "failed to extract namespace from context")
+		return "", err
 	}
 
+	logrus.SetOutput(logFifo)
+
+	log.G(shimCtx).WithField("id", containerID).Debug("StartShim")
+
 	// If we are running a shim start routine, we can safely assume our current working
 	// directory is the bundle directory
 	cwd, err := os.Getwd()
@@ -270,6 +284,8 @@ func (s *service) StartShim(shimCtx context.Context, containerID, containerdBina
 	}
 	bundleDir := bundle.Dir(cwd)
 
+	var exitAfterAllTasksGone bool
+
 	// Since we're running a shim start routine, we need to determine the vmID for the incoming
 	// container. Start by looking at the container's OCI annotations
 	s.vmID, err = bundleDir.OCIConfig().VMID()
@@ -287,83 +303,31 @@ func (s *service) StartShim(shimCtx context.Context, containerID, containerdBina
 		}
 
 		s.vmID = uuid.String()
-	}
-
-	// We determine if there is already a shim managing a VM with the current VMID by attempting
-	// to listen on the abstract socket address (which is parameterized by VMID). If we get
-	// EADDRINUSE, then we assume there is already a shim for the VM and just hand off to that
-	// one. This is in line with the approach used by containerd's reference runC shim v2
-	// implementation (which is also designed to manage multiple containers from a single shim
-	// process)
-	shimSocketAddress, err := shim.SocketAddress(shimCtx, s.vmID)
-	if err != nil {
-		return "", err
-	}
-
-	shimSocket, err := shim.NewSocket(shimSocketAddress)
-	if isEADDRINUSE(err) {
-		// There's already a shim for this VMID, so just hand off to it
-		err = s.vmDir().CreateBundleLink(containerID, bundleDir)
-		if err != nil {
-			return "", err
-		}
 
-		err = s.vmDir().CreateAddressLink(containerID)
-		if err != nil {
-			return "", err
-		}
-
-		return shimSocketAddress, nil
-	} else if err != nil {
-		return "", errors.Wrapf(err, "failed to create new shim socket at address \"%s\"", shimSocketAddress)
+		// If the client didn't specify a VMID, the VM should exit after this task is gone
+		exitAfterAllTasksGone = true
 	}
 
-	// If we're here, there is no pre-existing shim for this VMID, so we spawn a new one
-	defer func() {
-		s.logger.WithField("id", containerID).Debug("closing shim socket")
-		shimSocket.Close()
-	}()
-
-	err = s.vmDir().Create()
+	client, err := containerd.New(containerdAddress)
 	if err != nil {
 		return "", err
 	}
 
-	err = s.vmDir().CreateBundleLink(containerID, bundleDir)
-	if err != nil {
-		return "", err
-	}
+	fcControlClient := fccontrolGrpc.NewFirecrackerClient(client.Conn())
 
-	cmd, err := s.newShim(namespace, containerdBinary, containerdAddress, shimSocket)
+	_, err = fcControlClient.CreateVM(shimCtx, &proto.CreateVMRequest{
+		VMID:                  s.vmID,
+		ExitAfterAllTasksGone: exitAfterAllTasksGone,
+	})
 	if err != nil {
-		s.logger.WithError(err).WithField("id", containerID).Error("Failed to start new shim")
-		return "", err
-	}
-
-	defer func() {
-		if err != nil {
-			s.logger.WithError(err).WithField("id", containerID).Error("killing shim process after error")
-			cmd.Process.Kill()
+		errStatus, ok := status.FromError(err)
+		// ignore AlreadyExists errors, that just means the shim is already up and running
+		if !ok || errStatus.Code() != codes.AlreadyExists {
+			return "", errors.Wrap(err, "unexpected error from CreateVM")
 		}
-	}()
-
-	err = s.vmDir().WriteAddress(shimSocketAddress)
-	if err != nil {
-		s.logger.WithError(err).WithField("id", containerID).Error("failed to write address")
-		return "", err
-	}
-
-	err = s.vmDir().CreateAddressLink(containerID)
-	if err != nil {
-		return "", err
-	}
-
-	err = s.vmDir().CreateShimLogFifoLink(containerID)
-	if err != nil {
-		return "", err
 	}
 
-	return shimSocketAddress, nil
+	return fcShim.SocketAddress(shimCtx, s.vmID)
 }
 
 func logPanicAndDie(logger *logrus.Entry) {
@@ -372,24 +336,6 @@ func logPanicAndDie(logger *logrus.Entry) {
 	}
 }
 
-func parseCreateTaskOpts(opts *ptypes.Any) (*proto.FirecrackerConfig, *ptypes.Any, error) {
-	cfg, err := typeurl.UnmarshalAny(opts)
-	if err != nil {
-		return nil, nil, errors.Wrapf(err, "unmarshaling task create request options")
-	}
-	// We've verified that this is a valid prototype at this point of time.
-	// Check if it's the FirecrackerConfig type
-	firecrackerConfig, ok := cfg.(*proto.FirecrackerConfig)
-	if ok {
-		// We've verified that the proto message was of FirecrackerConfig type.
-		// Get runc options based on what was set in FirecrackerConfig.
-		return firecrackerConfig, firecrackerConfig.RuncOptions, nil
-	}
-	// This is a valid proto message, but is not FirecrackerConfig type.
-	// Treat the message as runc opts
-	return nil, opts, nil
-}
-
 func (s *service) generateExtraData(bundleDir bundle.Dir, options *ptypes.Any) (*proto.ExtraData, error) {
 	// Add the bundle/config.json to the request so it can be recreated
 	// inside the vm:
@@ -436,10 +382,268 @@ func (s *service) nextVSockPort() uint32 {
 	return port
 }
 
+func (s *service) waitVMReady() error {
+	select {
+	case <-s.vmReady:
+		return nil
+	case <-time.After(firecrackerStartTimeout):
+		return status.Error(codes.DeadlineExceeded, "timed out waiting for VM start")
+	}
+}
+
+// CreateVM will attempt to create the VM as specified in the provided request, but only on the first request
+// received. Any subsequent requests will be ignored and get an AlreadyExists error response.
+func (s *service) CreateVM(requestCtx context.Context, request *proto.CreateVMRequest) (*empty.Empty, error) {
+	defer logPanicAndDie(s.logger)
+
+	var (
+		err       error
+		createRan bool
+	)
+
+	s.vmStartOnce.Do(func() {
+		err = s.createVM(requestCtx, request)
+		createRan = true
+	})
+
+	if !createRan {
+		return nil, status.Error(codes.AlreadyExists, "shim cannot create VM more than once")
+	}
+
+	// If we failed to create the VM, we have no point in existing anymore, so shutdown
+	if err != nil {
+		s.shimCancel()
+		err = errors.Wrap(err, "failed to create VM")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	// creating the VM succeeded, setup monitors and publish events to celebrate
+	err = s.publishVMStart()
+	if err != nil {
+		s.logger.WithError(err).Error("failed to publish start VM event")
+	}
+
+	go s.monitorVMExit()
+	go s.monitorTaskExit(request.ExitAfterAllTasksGone)
+
+	// let all the other methods know that the VM is ready for tasks
+	close(s.vmReady)
+	return &empty.Empty{}, nil
+}
+
+func (s *service) publishVMStart() error {
+	return s.eventExchange.Publish(s.shimCtx, StartEventName, &proto.VMStart{VMID: s.vmID})
+}
+
+func (s *service) publishVMStop() error {
+	return s.eventExchange.Publish(s.shimCtx, StopEventName, &proto.VMStop{VMID: s.vmID})
+}
+
+func (s *service) createVM(requestCtx context.Context, request *proto.CreateVMRequest) (err error) {
+	var vsockFd *os.File
+	defer func() {
+		if vsockFd != nil {
+			vsockFd.Close()
+		}
+	}()
+
+	s.logger.Info("creating new VM")
+
+	s.machineCID, vsockFd, err = findNextAvailableVsockCID(requestCtx)
+	if err != nil {
+		return errors.Wrapf(err, "failed to find vsock for VM")
+	}
+
+	s.machineConfig, err = s.buildVMConfiguration(request)
+	if err != nil {
+		return errors.Wrapf(err, "failed to build VM configuration")
+	}
+
+	cmd := firecracker.VMCommandBuilder{}.
+		WithBin(s.config.FirecrackerBinaryPath).
+		WithSocketPath(s.shimDir().FirecrackerSockPath()).
+		Build(s.shimCtx) // shimCtx so the VM process is only killed when the shim shuts down
+
+	// use shimCtx so the VM is killed when the shim shuts down
+	s.machine, err = firecracker.NewMachine(s.shimCtx, *s.machineConfig,
+		firecracker.WithLogger(s.logger), firecracker.WithProcessRunner(cmd))
+	if err != nil {
+		return errors.Wrapf(err, "failed to create new machine instance")
+	}
+
+	// Close the vsock FD before starting the machine so firecracker doesn't get EADDRINUSE.
+	// This technically leaves us vulnerable to race conditions, but given the small time
+	// window and the fact that we choose a random CID from a 32bit range, the chances are
+	// infinitesimal. This will also no longer be an issue when firecracker implements its
+	// new vsock model.
+	err = vsockFd.Close()
+	if err != nil {
+		return errors.Wrapf(err, "failed to close vsock")
+	}
+
+	if err = s.machine.Start(requestCtx); err != nil {
+		return errors.Wrapf(err, "failed to start the VM")
+	}
+
+	s.logger.Info("calling agent")
+	conn, err := dialVsock(requestCtx, s.machineCID, defaultVsockPort)
+	if err != nil {
+		return errors.Wrapf(err, "failed to dial the VM over vsock")
+	}
+
+	rpcClient := ttrpc.NewClient(conn, ttrpc.WithOnClose(func() { _ = conn.Close() }))
+	s.agentClient = taskAPI.NewTaskClient(rpcClient)
+	s.eventBridgeClient = eventbridge.NewGetterClient(rpcClient)
+
+	s.logger.Info("successfully started the VM")
+
+	return nil
+}
+
+// StopVM will shutdown the firecracker VM and start this shim's shutdown procedure. If the VM has not been
+// created yet, an error will be returned but the shim will continue to shutdown.
+func (s *service) StopVM(requestCtx context.Context, request *proto.StopVMRequest) (*empty.Empty, error) {
+	defer logPanicAndDie(s.logger)
+
+	// Shutdown the shim after this method is called, whether the clean StopVMM was successful or not.
+	// If StopVMM didn't work, the cancel results in the VM cmd context being cancelled, sending SIGKILL
+	// to the VM process.
+	defer s.shimCancel()
+
+	s.logger.Info("stopping the VM")
+
+	err := s.machine.StopVMM()
+	if err != nil {
+		err = errors.Wrap(err, "failed to gracefully stop VM")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	err = os.RemoveAll(s.shimDir().RootPath())
+	if err != nil {
+		err = errors.Wrap(err, "failed to remove VM dir during shutdown")
+		s.logger.WithField("path", s.shimDir().RootPath()).WithError(err).Error()
+		return nil, err
+	}
+
+	s.logger.Info("successfully stopped the VM")
+	return &empty.Empty{}, nil
+}
+
+// GetVMInfo returns metadata for the VM being managed by this shim. If the VM has not been created yet, this
+// method will wait for up to a hardcoded timeout for it to exist, returning an error if the timeout is reached.
+func (s *service) GetVMInfo(requestCtx context.Context, request *proto.GetVMInfoRequest) (*proto.GetVMInfoResponse, error) {
+	defer logPanicAndDie(s.logger)
+
+	err := s.waitVMReady()
+	if err != nil {
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return &proto.GetVMInfoResponse{
+		VMID:            s.vmID,
+		ContextID:       s.machineCID,
+		SocketPath:      s.machineConfig.SocketPath,
+		LogFifoPath:     s.machineConfig.LogFifo,
+		MetricsFifoPath: s.machineConfig.MetricsFifo,
+	}, nil
+}
+
+// SetVMMetadata will update the VM being managed by this shim with the provided metadata. If the VM has not been created yet, this
+// method will wait for up to a hardcoded timeout for it to exist, returning an error if the timeout is reached.
+func (s *service) SetVMMetadata(requestCtx context.Context, request *proto.SetVMMetadataRequest) (*empty.Empty, error) {
+	defer logPanicAndDie(s.logger)
+
+	err := s.waitVMReady()
+	if err != nil {
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	s.logger.Info("updating VM metadata")
+	if err := s.machine.SetMetadata(requestCtx, request.Metadata); err != nil {
+		err = errors.Wrap(err, "failed to set VM metadata")
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
+	return &empty.Empty{}, nil
+}
+
+func (s *service) shimDir() vm.Dir {
+	return vm.ShimDir(s.namespace, s.vmID)
+}
+
+func (s *service) buildVMConfiguration(req *proto.CreateVMRequest) (*firecracker.Config, error) {
+	logger := s.logger.WithField("cid", s.machineCID)
+
+	cfg := firecracker.Config{
+		SocketPath:   s.shimDir().FirecrackerSockPath(),
+		VsockDevices: []firecracker.VsockDevice{{Path: "root", CID: s.machineCID}},
+		LogFifo:      s.shimDir().FirecrackerLogFifoPath(),
+		MetricsFifo:  s.shimDir().FirecrackerMetricsFifoPath(),
+		MachineCfg:   machineConfigurationFromProto(req.MachineCfg),
+	}
+
+	logger.Debugf("using socket path: %s", cfg.SocketPath)
+
+	// Kernel configuration
+
+	if val := req.KernelArgs; val != "" {
+		cfg.KernelArgs = val
+	} else {
+		cfg.KernelArgs = defaultKernelArgs
+	}
+
+	if val := req.KernelImagePath; val != "" {
+		cfg.KernelImagePath = val
+	} else {
+		cfg.KernelImagePath = defaultKernelPath
+	}
+
+	// Drives configuration
+
+	var driveBuilder firecracker.DrivesBuilder
+	if root := req.RootDrive; root != nil {
+		driveBuilder = firecracker.NewDrivesBuilder(root.PathOnHost)
+	} else {
+		driveBuilder = firecracker.NewDrivesBuilder(defaultRootfsPath)
+	}
+
+	// TODO: Reserve fake drives here (https://github.com/firecracker-microvm/firecracker-containerd/pull/154)
+	// Right now, just a single hardcoded stub drive is allocated
+	driveBuilder = driveBuilder.AddDrive("/dev/null", false, func(drive *models.Drive) {
+		drive.IsRootDevice = firecracker.Bool(false)
+		drive.DriveID = firecracker.String("containerRootfs")
+	})
+
+	for _, drive := range req.AdditionalDrives {
+		driveBuilder = addDriveFromProto(driveBuilder, drive)
+	}
+
+	cfg.Drives = driveBuilder.Build()
+
+	// Setup network interfaces
+
+	for _, ni := range req.NetworkInterfaces {
+		cfg.NetworkInterfaces = append(cfg.NetworkInterfaces, networkConfigFromProto(ni))
+	}
+
+	return &cfg, nil
+}
+
 func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTaskRequest) (*taskAPI.CreateTaskResponse, error) {
-	logger := log.G(requestCtx).WithField("id", request.ID)
+	logger := s.logger.WithField("containerID", request.ID)
 	defer logPanicAndDie(logger)
 
+	err := s.waitVMReady()
+	if err != nil {
+		logger.WithError(err).Error()
+		return nil, err
+	}
+
 	logger.WithFields(logrus.Fields{
 		"bundle":     request.Bundle,
 		"terminal":   request.Terminal,
@@ -449,41 +653,37 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 		"checkpoint": request.Checkpoint,
 	}).Debug("creating task")
 
-	var err error
-	var firecrackerConfig *proto.FirecrackerConfig
-	var runcOpts *ptypes.Any
+	bundleDir := bundle.Dir(request.Bundle)
+	err = s.shimDir().CreateBundleLink(request.ID, bundleDir)
+	if err != nil {
+		err = errors.Wrap(err, "failed to create VM dir bundle link")
+		logger.WithError(err).Error()
+		return nil, err
+	}
 
-	if request.Options != nil {
-		firecrackerConfig, runcOpts, err = parseCreateTaskOpts(request.Options)
-		if err != nil {
-			logger.WithFields(logrus.Fields{
-				"options": request.Options,
-				"error":   err,
-			}).Error("failed to unmarshal task create request options")
-			return nil, errors.Wrapf(err, "unmarshaling task create request options")
-		}
+	err = s.shimDir().CreateAddressLink(request.ID)
+	if err != nil {
+		err = errors.Wrap(err, "failed to create shim address symlink")
+		logger.WithError(err).Error()
+		return nil, err
 	}
 
-	s.startVMMutex.Lock()
-	if !s.agentStarted {
-		logger.Debug("calling startVM")
-		client, err := s.startVM(requestCtx, request, firecrackerConfig)
+	// TODO replace with proper drive mounting after that PR is merged
+	s.rootfsOnce.Do(func() {
+		defer logPanicAndDie(logger)
+
+		err := s.machine.UpdateGuestDrive(requestCtx, "containerRootfs", request.Rootfs[0].Source)
 		if err != nil {
-			s.startVMMutex.Unlock()
-			logger.WithError(err).Error("failed to start VM")
-			return nil, err
+			panic(err)
 		}
-
-		s.agentClient = client
-		s.agentStarted = true
-	}
-	s.startVMMutex.Unlock()
+	})
 
 	logger.Info("creating task")
 
-	bundleDir := bundle.Dir(request.Bundle)
-	extraData, err := s.generateExtraData(bundleDir, runcOpts)
+	extraData, err := s.generateExtraData(bundleDir, request.Options)
 	if err != nil {
+		err = errors.Wrap(err, "failed to generate extra data")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
@@ -495,6 +695,8 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 		Terminal: request.Terminal,
 	}, nil), taskCtx.Done(), taskCancel)
 	if err != nil {
+		err = errors.Wrap(err, "failed to add task")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
@@ -512,12 +714,15 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 	// Override the original request options with ExtraData needed by the VM Agent before sending it off
 	request.Options, err = task.MarshalExtraData()
 	if err != nil {
+		err = errors.Wrap(err, "failed to marshal extra data")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
 	resp, err := task.Create(requestCtx, request)
 	if err != nil {
-		logger.WithError(err).Error("create failed")
+		err = errors.Wrap(err, "failed to create task")
+		logger.WithError(err).Error()
 		return nil, err
 	}
 
@@ -527,7 +732,7 @@ func (s *service) Create(requestCtx context.Context, request *taskAPI.CreateTask
 		return nil, err
 	}
 
-	logger.WithField("pid", resp.Pid).Info("successfully created task")
+	logger.WithField("pid_in_vm", resp.Pid).Info("successfully created task")
 	return resp, nil
 }
 
@@ -750,32 +955,31 @@ func (s *service) Connect(requestCtx context.Context, req *taskAPI.ConnectReques
 func (s *service) Shutdown(requestCtx context.Context, req *taskAPI.ShutdownRequest) (*ptypes.Empty, error) {
 	defer logPanicAndDie(log.G(requestCtx))
 
-	log.G(requestCtx).WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
+	s.logger.WithFields(logrus.Fields{"id": req.ID, "now": req.Now}).Debug("shutdown")
 
 	// If we are still managing containers, don't shutdown
 	if s.taskManager.TaskCount() > 0 {
 		return &ptypes.Empty{}, nil
 	}
 
-	_, err := s.agentClient.Shutdown(requestCtx, req)
+	defer s.shimCancel()
+
+	err := s.waitVMReady()
 	if err != nil {
+		s.logger.WithError(err).Error()
 		return nil, err
 	}
 
-	log.G(requestCtx).Debug("stopping VM")
-	if err := s.stopVM(); err != nil {
-		log.G(requestCtx).WithError(err).Error("failed to stop VM")
+	_, err = s.agentClient.Shutdown(requestCtx, req)
+	if err != nil {
 		return nil, err
 	}
 
-	err = os.RemoveAll(s.vmDir().RootPath())
-	if err != nil {
-		log.G(requestCtx).WithField("path", s.vmDir().RootPath()).WithError(err).Error("failed to remove VM dir during shutdown")
+	if _, err := s.StopVM(requestCtx, &proto.StopVMRequest{VMID: s.vmID}); err != nil {
+		s.logger.WithError(err).Error("failed to stop VM")
+		return nil, err
 	}
 
-	defer s.shimCancel()
-
-	log.G(requestCtx).Debug("stopping runtime")
 	return &ptypes.Empty{}, nil
 }
 
@@ -835,6 +1039,12 @@ func (s *service) Wait(requestCtx context.Context, req *taskAPI.WaitRequest) (*t
 func (s *service) Cleanup(requestCtx context.Context) (*taskAPI.DeleteResponse, error) {
 	defer logPanicAndDie(log.G(requestCtx))
 
+	err := s.waitVMReady()
+	if err != nil {
+		s.logger.WithError(err).Error()
+		return nil, err
+	}
+
 	log.G(requestCtx).Debug("cleanup")
 	// Destroy VM/etc here?
 	// copied from runcs impl, nothing to cleanup atm
@@ -844,10 +1054,6 @@ func (s *service) Cleanup(requestCtx context.Context) (*taskAPI.DeleteResponse,
 	}, nil
 }
 
-func newServer() (*ttrpc.Server, error) {
-	return ttrpc.NewServer(ttrpc.WithServerHandshaker(ttrpc.UnixSocketRequireSameUser()))
-}
-
 func dialVsock(requestCtx context.Context, contextID, port uint32) (net.Conn, error) {
 	// VM should start within 200ms, vsock dial will make retries at 100ms, 200ms, 400ms, 800ms, 1.6s, 3.2s, 6.4s
 	const (
@@ -892,31 +1098,36 @@ func cidDialer(cid uint32) vm.VSockConnector {
 // It uses VHOST_VSOCK_SET_GUEST_CID ioctl which allows some CID ranges to be statically reserved in advance.
 // The ioctl fails with EADDRINUSE if cid is already taken and with EINVAL if the CID is invalid.
 // Taken from https://bugzilla.redhat.com/show_bug.cgi?id=1291851
-func findNextAvailableVsockCID(requestCtx context.Context) (uint32, error) {
+func findNextAvailableVsockCID(requestCtx context.Context) (cid uint32, file *os.File, err error) {
 	const (
 		// Corresponds to VHOST_VSOCK_SET_GUEST_CID in vhost.h
 		ioctlVsockSetGuestCID = uintptr(0x4008AF60)
 		// 0, 1 and 2 are reserved CIDs, see http://man7.org/linux/man-pages/man7/vsock.7.html
 		minCID          = 3
 		maxCID          = math.MaxUint32
+		cidRange        = maxCID - minCID
 		vsockDevicePath = "/dev/vhost-vsock"
 	)
 
-	file, err := os.OpenFile(vsockDevicePath, syscall.O_RDWR, 0600)
+	defer func() {
+		// close the vsock file if anything bad happens
+		if err != nil && file != nil {
+			file.Close()
+		}
+	}()
+
+	file, err = os.OpenFile(vsockDevicePath, syscall.O_RDWR, 0600)
 	if err != nil {
-		return 0, errors.Wrap(err, "failed to open vsock device")
+		return 0, nil, errors.Wrap(err, "failed to open vsock device")
 	}
-	defer file.Close()
 
 	// Start at a random ID to minimize chances of conflicts when shims are racing each other here
-	// TODO the actual fix to this problem will come when the shim is updated to just use the FC-control plugin,
-	// this current implementation is a very short-term bandaid.
-	start := rand.Intn(maxCID - minCID)
-	for n := 0; n < maxCID-minCID; n++ {
-		cid := minCID + ((start + n) % (maxCID - minCID))
+	start := rand.Intn(cidRange)
+	for n := 0; n < cidRange; n++ {
+		cid := minCID + ((start + n) % cidRange)
 		select {
 		case <-requestCtx.Done():
-			return 0, requestCtx.Err()
+			return 0, nil, requestCtx.Err()
 		default:
 			_, _, err = sysCall(
 				unix.SYS_IOCTL,
@@ -926,119 +1137,21 @@ func findNextAvailableVsockCID(requestCtx context.Context) (uint32, error) {
 
 			switch err {
 			case unix.Errno(0):
-				return uint32(cid), nil
+				return uint32(cid), file, nil
 			case unix.EADDRINUSE:
 				// ID is already taken, try next one
 				continue
 			default:
 				// Fail if we get an error we don't expect
-				return 0, err
+				return 0, nil, err
 			}
 		}
 	}
 
-	return 0, errors.New("couldn't find any available vsock context id")
+	return 0, nil, errors.New("couldn't find any available vsock context id")
 }
 
-// TODO: replace startVM with calls to the FC-control plugin
-func (s *service) startVM(requestCtx context.Context,
-	request *taskAPI.CreateTaskRequest,
-	vmConfig *proto.FirecrackerConfig,
-) (taskAPI.TaskService, error) {
-	log.G(requestCtx).Info("starting VM")
-
-	cid, err := findNextAvailableVsockCID(requestCtx)
-	if err != nil {
-		return nil, err
-	}
-
-	cfg := firecracker.Config{
-		SocketPath:      s.vmDir().FirecrackerSockPath(),
-		VsockDevices:    []firecracker.VsockDevice{{Path: "root", CID: cid}},
-		KernelImagePath: s.config.KernelImagePath,
-		KernelArgs:      s.config.KernelArgs,
-		MachineCfg: models.MachineConfiguration{
-			VcpuCount:   int64(s.config.CPUCount),
-			CPUTemplate: models.CPUTemplate(s.config.CPUTemplate),
-			MemSizeMib:  256,
-		},
-		LogFifo:     s.vmDir().FirecrackerLogFifoPath(),
-		LogLevel:    s.config.LogLevel,
-		MetricsFifo: s.vmDir().FirecrackerMetricsFifoPath(),
-		Debug:       s.config.Debug,
-	}
-
-	driveBuilder := firecracker.NewDrivesBuilder(s.config.RootDrive)
-	// Attach block devices passed from snapshotter
-	for _, mnt := range request.Rootfs {
-		if mnt.Type != supportedMountFSType {
-			return nil, errors.Errorf("unsupported mount type '%s', expected '%s'", mnt.Type, supportedMountFSType)
-		}
-
-		driveBuilder = driveBuilder.AddDrive(mnt.Source, false)
-	}
-	// Override config provided in task's create opts if any.
-	// Note: We've chosen to override here instead of merging in order to
-	// provide a cleaner, simpler interface to reason about for clients.
-	// Any config provided by clients for create task opts will always override
-	// the default config generated by the runtime.
-	cfg, driveBuilder, err = overrideVMConfigFromTaskOpts(cfg, vmConfig, driveBuilder)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to build VM options")
-	}
-
-	cfg.Drives = driveBuilder.Build()
-	cmd := firecracker.VMCommandBuilder{}.
-		WithBin(s.config.FirecrackerBinaryPath).
-		WithSocketPath(s.vmDir().FirecrackerSockPath()).
-		Build(s.shimCtx) // use shimCtx to ensure the cmd is only killed once the shim shutsdown
-	machineOpts := []firecracker.Opt{
-		firecracker.WithProcessRunner(cmd),
-		firecracker.WithLogger(s.logger),
-	}
-
-	s.machine, err = firecracker.NewMachine(requestCtx, cfg, machineOpts...)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to create firecracker config")
-	}
-	s.machineCID = cid
-
-	log.G(requestCtx).Info("starting instance")
-	if err := s.machine.Start(requestCtx); err != nil {
-		return nil, errors.Wrap(err, "failed to start firecracker VM")
-	}
-
-	log.G(requestCtx).Info("calling agent")
-	conn, err := dialVsock(requestCtx, cid, defaultVsockPort)
-	if err != nil {
-		s.stopVM()
-		return nil, errors.Wrap(err, "failed to dial vsock")
-	}
-
-	log.G(requestCtx).Info("creating clients")
-	rpcClient := ttrpc.NewClient(conn, ttrpc.WithOnClose(func() { _ = conn.Close() }))
-	apiClient := taskAPI.NewTaskClient(rpcClient)
-	eventBridgeClient := eventbridge.NewGetterClient(rpcClient)
-
-	go func() {
-		// Connect the agent's event exchange to our own own event exchange
-		// using the eventbridge client. All events that are published on the
-		// agent's exchange will also be published on our own
-		if err := <-eventbridge.Attach(s.shimCtx, eventBridgeClient, s.eventExchange); err != nil && err != context.Canceled {
-			log.G(requestCtx).WithError(err).Error("error while forwarding events from VM agent")
-		}
-	}()
-
-	go s.monitorTaskExit()
-
-	return apiClient, nil
-}
-
-func (s *service) stopVM() error {
-	return s.machine.StopVMM()
-}
-
-func (s *service) monitorTaskExit() {
+func (s *service) monitorTaskExit(exitAfterAllTasksGone bool) {
 	exitEvents, exitEventErrs := s.eventExchange.Subscribe(s.shimCtx, fmt.Sprintf(`topic=="%s"`, runtime.TaskExitEventTopic))
 
 	var err error
@@ -1064,11 +1177,13 @@ func (s *service) monitorTaskExit() {
 
 				s.taskManager.Remove(event.ContainerID)
 
-				// If we have no more containers, shutdown. If we still have containers left,
-				// this will be a no-op
-				_, err = s.Shutdown(s.shimCtx, &taskAPI.ShutdownRequest{})
-				if err != nil {
-					logger.WithError(err).Fatal("failed to shutdown after container exit")
+				if exitAfterAllTasksGone {
+					// If we have no more containers, shutdown. If we still have containers left,
+					// this will be a no-op
+					_, err = s.Shutdown(s.shimCtx, &taskAPI.ShutdownRequest{})
+					if err != nil {
+						logger.WithError(err).Fatal("failed to shutdown after container exit")
+					}
 				}
 
 			default:
@@ -1087,3 +1202,19 @@ func (s *service) monitorTaskExit() {
 		}
 	}
 }
+
+func (s *service) monitorVMExit() {
+	// once the VM shuts down, the shim should too
+	defer s.shimCancel()
+
+	// Block until the VM exits
+	waitErr := s.machine.Wait(s.shimCtx)
+	if waitErr != nil && waitErr != context.Canceled {
+		s.logger.WithError(waitErr).Error("error returned from VM wait")
+	}
+
+	publishErr := s.publishVMStop()
+	if publishErr != nil {
+		s.logger.WithError(publishErr).Error("failed to publish stop VM event")
+	}
+}
diff --git a/runtime/taskfirecrackeropts.go b/runtime/taskfirecrackeropts.go
deleted file mode 100644
index 8d70e4d..0000000
--- a/runtime/taskfirecrackeropts.go
+++ /dev/null
@@ -1,161 +0,0 @@
-// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//	http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-
-package main
-
-import (
-	"time"
-
-	"github.com/firecracker-microvm/firecracker-containerd/proto"
-	firecracker "github.com/firecracker-microvm/firecracker-go-sdk"
-	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
-	"github.com/pkg/errors"
-)
-
-// overrideVMConfigFromTaskOpts overrides the firecracker VM config for the
-// task with values set in the protobuf message.
-func overrideVMConfigFromTaskOpts(
-	cfg firecracker.Config,
-	vmConfig *proto.FirecrackerConfig,
-	drivesBuilder firecracker.DrivesBuilder,
-) (firecracker.Config, firecracker.DrivesBuilder, error) {
-	if vmConfig == nil {
-		return cfg, drivesBuilder, nil
-	}
-	var err error
-	if vmConfig.KernelImagePath != "" {
-		// If kernel image path is set, override it.
-		cfg.KernelImagePath = vmConfig.KernelImagePath
-	}
-	if vmConfig.KernelArgs != "" {
-		// If kernel args are set, override it.
-		cfg.KernelArgs = vmConfig.KernelArgs
-	}
-	// Attach network interface specified in the config.
-	if len(vmConfig.NetworkInterfaces) > 0 {
-		nwIfaces := make([]firecracker.NetworkInterface, len(vmConfig.NetworkInterfaces))
-		for i, nw := range vmConfig.NetworkInterfaces {
-			nwIfaces[i] = networkConfigFromProto(nw)
-		}
-		cfg.NetworkInterfaces = nwIfaces
-	}
-	if vmConfig.MachineCfg != nil {
-		cfg.MachineCfg, err = machineConfigFromProto(vmConfig.MachineCfg)
-		if err != nil {
-			return firecracker.Config{}, drivesBuilder, err
-		}
-	}
-	// Overwrite the root drive config if specified.
-	if vmConfig.RootDrive != nil {
-		drivesBuilder = drivesBuilderFromProto(vmConfig.RootDrive, drivesBuilder, true)
-	}
-	// Add additional drives if specified.
-	if len(vmConfig.AdditionalDrives) > 0 {
-		for _, drive := range vmConfig.AdditionalDrives {
-			drivesBuilder = drivesBuilderFromProto(drive, drivesBuilder, false)
-		}
-	}
-
-	return cfg, drivesBuilder, nil
-}
-
-// networkConfigFromProto creates a firecracker NetworkInterface object from
-// the protobuf FirecrackerNetworkInterface message.
-func networkConfigFromProto(nwIface *proto.FirecrackerNetworkInterface) firecracker.NetworkInterface {
-	result := firecracker.NetworkInterface{
-		MacAddress:  nwIface.MacAddress,
-		HostDevName: nwIface.HostDevName,
-		AllowMMDS:   nwIface.AllowMMDS,
-	}
-	if nwIface.InRateLimiter != nil {
-		result.InRateLimiter = rateLimiterFromProto(nwIface.InRateLimiter)
-	}
-	if nwIface.OutRateLimiter != nil {
-		result.OutRateLimiter = rateLimiterFromProto(nwIface.OutRateLimiter)
-	}
-
-	return result
-}
-
-// machineConfigFromProto creates a firecracker MachineConfiguration object
-// from the protobuf FirecrackerMachineConfiguration message.
-func machineConfigFromProto(mCfg *proto.FirecrackerMachineConfiguration) (models.MachineConfiguration, error) {
-	result := models.MachineConfiguration{}
-	if mCfg.MemSizeMib > 0 && mCfg.VcpuCount == 0 {
-		// Only memory is being overridden.
-		return result, errors.New("both vm memory and vcpu count need to be overridden")
-	}
-	if mCfg.VcpuCount > 0 && mCfg.MemSizeMib == 0 {
-		// Only cpu is being overridden.
-		return result, errors.New("both vm memory and vcpu count need to be overridden")
-	}
-	if mCfg.CPUTemplate != "" {
-		result.CPUTemplate = models.CPUTemplate(mCfg.CPUTemplate)
-	}
-	result.HtEnabled = mCfg.HtEnabled
-	result.MemSizeMib = int64(mCfg.MemSizeMib)
-	result.VcpuCount = int64(mCfg.VcpuCount)
-
-	return result, nil
-}
-
-// drivesBuilderFromProto returns a new DrivesBuilder that can be used to
-// build the firecracker Drives config. The new DrivesBuilder object is
-// constructed using the information present in the protobuf message.
-func drivesBuilderFromProto(drive *proto.FirecrackerDrive,
-	drivesBuilder firecracker.DrivesBuilder,
-	isRoot bool,
-) firecracker.DrivesBuilder {
-	opt := func(d *models.Drive) {
-		d.Partuuid = drive.Partuuid
-		if drive.RateLimiter != nil {
-			d.RateLimiter = rateLimiterFromProto(drive.RateLimiter)
-		}
-	}
-	if isRoot {
-		return drivesBuilder.WithRootDrive(drive.PathOnHost, opt)
-	}
-	return drivesBuilder.AddDrive(drive.PathOnHost, drive.IsReadOnly, opt)
-}
-
-// rateLimiterFromProto creates a firecracker RateLimiter object from the
-// protobuf message.
-func rateLimiterFromProto(rl *proto.FirecrackerRateLimiter) *models.RateLimiter {
-	result := models.RateLimiter{}
-	if rl.Bandwidth != nil {
-		result.Bandwidth = tokenBucketFromProto(rl.Bandwidth)
-	}
-	if rl.Ops != nil {
-		result.Ops = tokenBucketFromProto(rl.Ops)
-	}
-
-	return &result
-}
-
-// tokenBucketFromProto creates a firecracker TokenBucket object from the
-// protobuf message.
-func tokenBucketFromProto(bucket *proto.FirecrackerTokenBucket) *models.TokenBucket {
-	builder := firecracker.TokenBucketBuilder{}
-	if bucket.OneTimeBurst > 0 {
-		builder = builder.WithInitialSize(bucket.OneTimeBurst)
-	}
-	if bucket.RefillTime > 0 {
-		builder = builder.WithRefillDuration(time.Duration(bucket.RefillTime) * time.Millisecond)
-	}
-	if bucket.Capacity > 0 {
-		builder = builder.WithBucketSize(bucket.Capacity)
-	}
-
-	res := builder.Build()
-	return &res
-}
diff --git a/tools/docker/Dockerfile b/tools/docker/Dockerfile
index 8c9d9e9..cdce66e 100644
--- a/tools/docker/Dockerfile
+++ b/tools/docker/Dockerfile
@@ -101,30 +101,6 @@ RUN --mount=type=cache,from=build-base,source=/home/builder/cargo/registry,targe
 
 
 
-# Build containerd
-FROM build-base as containerd-build
-# we don't need their btrfs snapshotter, which requires more dependencies
-ENV BUILDTAGS='no_btrfs'
-# The commit ID being used ensures we have the following commit IDs (which currently are not in a tagged release):
-# * 82f2ac7edec03dfc301d99d8910c9e53a9d71d90 - Stop logging error on v2 multi shim log failure
-# * 810b3c37bb13662a289b9f6b028b5b5435be72b7 - Support disabling default setup of shim logger.
-# * 665715b8a6c595a1a96e6b60cbf8b544ad63cc78 - Allow reuse client for external GRPC services
-# Once those commits are in a tagged release it should make sense to go back to using a release branch instead of
-# a commit ID off the master containerd branch.
-RUN --mount=type=cache,from=build-base,source=/home/builder/go/pkg/mod,target=/home/builder/go/pkg/mod,sharing=locked \
-	mkdir -p /home/builder/go/src/github.com/containerd/containerd \
-	&& git clone https://github.com/containerd/containerd.git /home/builder/go/src/github.com/containerd/containerd \
-	&& cd /home/builder/go/src/github.com/containerd/containerd \
-  && git checkout 82f2ac7edec03dfc301d99d8910c9e53a9d71d90 \
-	&& go mod verify || go mod download
-RUN --mount=type=cache,from=build-base,source=/home/builder/go/pkg/mod,target=/home/builder/go/pkg/mod \
-	cd /home/builder/go/src/github.com/containerd/containerd \
-	&& make \
-	&& cp -R /home/builder/go/src/github.com/containerd/containerd/bin/* /output
-
-
-
-
 # Build runc
 FROM build-base as runc-build
 ENV BUILDTAGS='seccomp'
@@ -165,6 +141,12 @@ RUN --mount=type=cache,from=build-base,source=/home/builder/go/pkg/mod,target=/h
 		snapshotter/cmd/devmapper/devmapper_snapshotter \
 		snapshotter/cmd/naive/naive_snapshotter \
 		/output
+RUN --mount=type=cache,from=build-base,source=/home/builder/go/pkg/mod,target=/home/builder/go/pkg/mod \
+  cd firecracker-containerd/firecracker-control/cmd/containerd/ \
+  && make \
+  && cp \
+  firecracker-containerd \
+  /output/containerd
 
 
 
@@ -233,16 +215,15 @@ RUN mkdir -p /etc/apt/sources.list.d \
 
 COPY --from=firecracker-containerd-build /home/builder/firecracker-containerd /firecracker-containerd
 COPY --from=firecracker-build /output/* /usr/local/bin/
-COPY --from=containerd-build /output/* /usr/local/bin/
-COPY --from=firecracker-vm-root-builder /output/vm.ext4 /var/lib/firecracker-containerd/runtime/hello-rootfs.ext4
+COPY --from=firecracker-vm-root-builder /output/vm.ext4 /var/lib/firecracker-containerd/runtime/default-rootfs.img
 COPY --from=firecracker-containerd-build /output/* /usr/local/bin/
 COPY --from=runc-build /output/* /usr/local/bin/
 COPY tools/docker/firecracker-runtime.json /etc/containerd/firecracker-runtime.json
 
-RUN curl --silent --show-error --retry 3 --max-time 30 --output hello-vmlinux.bin \
+RUN curl --silent --show-error --retry 3 --max-time 30 --output default-vmlinux.bin \
 	"https://s3.amazonaws.com/spec.ccfc.min/img/hello/kernel/hello-vmlinux.bin" \
-	&& echo "882fa465c43ab7d92e31bd4167da3ad6a82cb9230f9b0016176df597c6014cef hello-vmlinux.bin" | sha256sum -c - \
-	&& mv hello-vmlinux.bin /var/lib/firecracker-containerd/runtime/hello-vmlinux.bin
+	&& echo "882fa465c43ab7d92e31bd4167da3ad6a82cb9230f9b0016176df597c6014cef default-vmlinux.bin" | sha256sum -c - \
+	&& mv default-vmlinux.bin /var/lib/firecracker-containerd/runtime/default-vmlinux.bin
 
 RUN --mount=type=cache,from=build-base,source=/home/builder/go/pkg/mod,target=/tmp/go/pkg/mod,readonly \
 	mkdir -p /root/go/pkg/mod \


Test patch:
diff --git a/firecracker-control/local_mock_test.go b/firecracker-control/local_mock_test.go
deleted file mode 100644
index 5778cdd..0000000
--- a/firecracker-control/local_mock_test.go
+++ /dev/null
@@ -1,148 +0,0 @@
-// Code generated by MockGen. DO NOT EDIT.
-// Source: local.go
-
-// Package service is a generated GoMock package.
-package service
-
-import (
-	context "context"
-	events "github.com/containerd/containerd/events"
-	firecracker_go_sdk "github.com/firecracker-microvm/firecracker-go-sdk"
-	gomock "github.com/golang/mock/gomock"
-	reflect "reflect"
-)
-
-// Mockmachine is a mock of machine interface
-type Mockmachine struct {
-	ctrl     *gomock.Controller
-	recorder *MockmachineMockRecorder
-}
-
-// MockmachineMockRecorder is the mock recorder for Mockmachine
-type MockmachineMockRecorder struct {
-	mock *Mockmachine
-}
-
-// NewMockmachine creates a new mock instance
-func NewMockmachine(ctrl *gomock.Controller) *Mockmachine {
-	mock := &Mockmachine{ctrl: ctrl}
-	mock.recorder = &MockmachineMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use
-func (m *Mockmachine) EXPECT() *MockmachineMockRecorder {
-	return m.recorder
-}
-
-// Start mocks base method
-func (m *Mockmachine) Start(arg0 context.Context) error {
-	ret := m.ctrl.Call(m, "Start", arg0)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Start indicates an expected call of Start
-func (mr *MockmachineMockRecorder) Start(arg0 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Start", reflect.TypeOf((*Mockmachine)(nil).Start), arg0)
-}
-
-// StopVMM mocks base method
-func (m *Mockmachine) StopVMM() error {
-	ret := m.ctrl.Call(m, "StopVMM")
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// StopVMM indicates an expected call of StopVMM
-func (mr *MockmachineMockRecorder) StopVMM() *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "StopVMM", reflect.TypeOf((*Mockmachine)(nil).StopVMM))
-}
-
-// Shutdown mocks base method
-func (m *Mockmachine) Shutdown(arg0 context.Context) error {
-	ret := m.ctrl.Call(m, "Shutdown", arg0)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Shutdown indicates an expected call of Shutdown
-func (mr *MockmachineMockRecorder) Shutdown(arg0 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Shutdown", reflect.TypeOf((*Mockmachine)(nil).Shutdown), arg0)
-}
-
-// Wait mocks base method
-func (m *Mockmachine) Wait(arg0 context.Context) error {
-	ret := m.ctrl.Call(m, "Wait", arg0)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Wait indicates an expected call of Wait
-func (mr *MockmachineMockRecorder) Wait(arg0 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Wait", reflect.TypeOf((*Mockmachine)(nil).Wait), arg0)
-}
-
-// SetMetadata mocks base method
-func (m *Mockmachine) SetMetadata(arg0 context.Context, arg1 interface{}) error {
-	ret := m.ctrl.Call(m, "SetMetadata", arg0, arg1)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// SetMetadata indicates an expected call of SetMetadata
-func (mr *MockmachineMockRecorder) SetMetadata(arg0, arg1 interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "SetMetadata", reflect.TypeOf((*Mockmachine)(nil).SetMetadata), arg0, arg1)
-}
-
-// UpdateGuestDrive mocks base method
-func (m *Mockmachine) UpdateGuestDrive(arg0 context.Context, arg1, arg2 string, arg3 ...firecracker_go_sdk.PatchGuestDriveByIDOpt) error {
-	varargs := []interface{}{arg0, arg1, arg2}
-	for _, a := range arg3 {
-		varargs = append(varargs, a)
-	}
-	ret := m.ctrl.Call(m, "UpdateGuestDrive", varargs...)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// UpdateGuestDrive indicates an expected call of UpdateGuestDrive
-func (mr *MockmachineMockRecorder) UpdateGuestDrive(arg0, arg1, arg2 interface{}, arg3 ...interface{}) *gomock.Call {
-	varargs := append([]interface{}{arg0, arg1, arg2}, arg3...)
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateGuestDrive", reflect.TypeOf((*Mockmachine)(nil).UpdateGuestDrive), varargs...)
-}
-
-// Mockpublisher is a mock of publisher interface
-type Mockpublisher struct {
-	ctrl     *gomock.Controller
-	recorder *MockpublisherMockRecorder
-}
-
-// MockpublisherMockRecorder is the mock recorder for Mockpublisher
-type MockpublisherMockRecorder struct {
-	mock *Mockpublisher
-}
-
-// NewMockpublisher creates a new mock instance
-func NewMockpublisher(ctrl *gomock.Controller) *Mockpublisher {
-	mock := &Mockpublisher{ctrl: ctrl}
-	mock.recorder = &MockpublisherMockRecorder{mock}
-	return mock
-}
-
-// EXPECT returns an object that allows the caller to indicate expected use
-func (m *Mockpublisher) EXPECT() *MockpublisherMockRecorder {
-	return m.recorder
-}
-
-// Publish mocks base method
-func (m *Mockpublisher) Publish(ctx context.Context, topic string, event events.Event) error {
-	ret := m.ctrl.Call(m, "Publish", ctx, topic, event)
-	ret0, _ := ret[0].(error)
-	return ret0
-}
-
-// Publish indicates an expected call of Publish
-func (mr *MockpublisherMockRecorder) Publish(ctx, topic, event interface{}) *gomock.Call {
-	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Publish", reflect.TypeOf((*Mockpublisher)(nil).Publish), ctx, topic, event)
-}
diff --git a/firecracker-control/local_test.go b/firecracker-control/local_test.go
deleted file mode 100644
index 5bf35f4..0000000
--- a/firecracker-control/local_test.go
+++ /dev/null
@@ -1,158 +0,0 @@
-// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//	http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-
-//go:generate mockgen -source=local.go -destination=local_mock_test.go -package=service
-
-package service
-
-import (
-	"context"
-	"os"
-	"testing"
-	"time"
-
-	"github.com/firecracker-microvm/firecracker-go-sdk"
-	"github.com/golang/mock/gomock"
-	"github.com/stretchr/testify/assert"
-
-	"github.com/firecracker-microvm/firecracker-containerd/proto"
-)
-
-var testCtx = context.Background()
-
-func TestLocal_buildVMConfiguration(t *testing.T) {
-	request := &proto.CreateVMRequest{
-		KernelArgs:      "kernel_args",
-		KernelImagePath: "kernel_path",
-		MachineCfg:      &proto.FirecrackerMachineConfiguration{},
-		RootDrive: &proto.FirecrackerDrive{
-			PathOnHost: "/",
-		},
-		NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
-			{MacAddress: "mac", HostDevName: "host", AllowMMDS: true},
-		},
-	}
-
-	obj := &local{
-		rootPath:    "/",
-		findVsockFn: func(context.Context) (*os.File, uint32, error) { return nil, 3, nil },
-	}
-
-	config, err := obj.buildVMConfiguration(testCtx, "1", 3, request)
-	assert.NoError(t, err)
-	assert.NotNil(t, config)
-
-	assert.Len(t, config.VsockDevices, 1)
-	assert.EqualValues(t, 3, config.VsockDevices[0].CID)
-
-	assert.Equal(t, "kernel_args", config.KernelArgs)
-	assert.Equal(t, "kernel_path", config.KernelImagePath)
-
-	assert.Equal(t, "/vm_1.socket", config.SocketPath)
-	assert.Equal(t, "/vm_1_log.fifo", config.LogFifo)
-	assert.Equal(t, "/vm_1_metrics.fifo", config.MetricsFifo)
-
-	assert.Len(t, config.Drives, 1)
-
-	rootDrive := config.Drives[0]
-	assert.True(t, *rootDrive.IsRootDevice)
-	assert.Equal(t, "/", *rootDrive.PathOnHost)
-
-	assert.Len(t, config.NetworkInterfaces, 1)
-
-	ni := config.NetworkInterfaces[0]
-	assert.Equal(t, "mac", ni.MacAddress)
-	assert.Equal(t, "host", ni.HostDevName)
-	assert.True(t, ni.AllowMMDS)
-}
-
-func TestLocal_startMachine(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	machine := NewMockmachine(ctrl)
-	machine.EXPECT().Start(gomock.Any()).Times(1).Return(nil)
-	machine.EXPECT().Wait(gomock.Any()).Times(1).Return(nil)
-
-	publisher := NewMockpublisher(ctrl)
-	publisher.EXPECT().Publish(gomock.Any(), startEventName, gomock.Eq(&proto.VMStart{VMID: "1"})).Return(nil)
-	publisher.EXPECT().Publish(gomock.Any(), stopEventName, gomock.Eq(&proto.VMStop{VMID: "1"})).Return(nil)
-
-	obj := &local{
-		vm:        map[string]*instance{"1": {}},
-		publisher: publisher,
-	}
-
-	err := obj.startMachine(testCtx, "1", machine)
-	time.Sleep(100 * time.Microsecond) // Give it some time to spawn the goroutine
-	assert.NoError(t, err)
-	assert.Empty(t, obj.vm) // Make sure "1" removed from the VM list
-}
-
-func TestLocal_StopVM(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	machine := NewMockmachine(ctrl)
-	machine.EXPECT().StopVMM().Times(1).Return(nil)
-
-	obj := &local{
-		vm: map[string]*instance{"1": {machine: machine}},
-	}
-
-	_, err := obj.StopVM(testCtx, &proto.StopVMRequest{VMID: "1"})
-	assert.NoError(t, err)
-}
-
-func TestLocal_StopInvalidVM(t *testing.T) {
-	obj := local{}
-	_, err := obj.StopVM(testCtx, &proto.StopVMRequest{VMID: "2"})
-	assert.Equal(t, ErrVMNotFound, err)
-}
-
-func TestLocal_GetVMInfo(t *testing.T) {
-	cfg := &firecracker.Config{
-		VsockDevices: []firecracker.VsockDevice{{CID: 123}},
-		SocketPath:   "socket",
-		LogFifo:      "logs",
-		MetricsFifo:  "metrics",
-	}
-
-	obj := &local{
-		vm: map[string]*instance{"1": {cfg: cfg}},
-	}
-
-	response, err := obj.GetVMInfo(testCtx, &proto.GetVMInfoRequest{VMID: "1"})
-	assert.NoError(t, err)
-	assert.NotNil(t, response)
-
-	assert.EqualValues(t, 123, response.ContextID)
-	assert.Equal(t, "socket", response.SocketPath)
-	assert.Equal(t, "logs", response.LogFifoPath)
-	assert.Equal(t, "metrics", response.MetricsFifoPath)
-}
-
-func TestLocal_SetVMMetadata(t *testing.T) {
-	ctrl := gomock.NewController(t)
-	defer ctrl.Finish()
-
-	machine := NewMockmachine(ctrl)
-	machine.EXPECT().SetMetadata(gomock.Any(), gomock.Eq("test")).Times(1).Return(nil)
-
-	obj := &local{
-		vm: map[string]*instance{"1": {machine: machine}},
-	}
-
-	_, err := obj.SetVMMetadata(testCtx, &proto.SetVMMetadataRequest{VMID: "1", Metadata: "test"})
-	assert.NoError(t, err)
-}
diff --git a/firecracker-control/vsock_test.go b/firecracker-control/vsock_test.go
deleted file mode 100644
index 4ab1ee0..0000000
--- a/firecracker-control/vsock_test.go
+++ /dev/null
@@ -1,45 +0,0 @@
-// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//	http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-
-package service
-
-import (
-	"context"
-	"syscall"
-	"testing"
-
-	"github.com/stretchr/testify/require"
-)
-
-func TestFindNextAvailableVsockCID(t *testing.T) {
-	sysCall = func(trap, a1, a2, a3 uintptr) (r1, r2 uintptr, err syscall.Errno) {
-		return 0, 0, 0
-	}
-
-	defer func() {
-		sysCall = syscall.Syscall
-	}()
-
-	file, cid, err := findNextAvailableVsockCID(context.Background())
-	require.NoError(t, err,
-		"Do you have permission to interact with /dev/vhost-vsock?\n"+
-			"Grant yourself permission with `sudo setfacl -m u:${USER}:rw /dev/vhost-vsock`")
-	defer file.Close()
-	require.EqualValues(t, uint32(3), cid)
-
-	ctx, cancel := context.WithCancel(context.Background())
-	cancel()
-
-	_, _, err = findNextAvailableVsockCID(ctx)
-	require.Equal(t, context.Canceled, err)
-}
diff --git a/firecracker-control/helpers_test.go b/runtime/helpers_test.go
similarity index 97%
rename from firecracker-control/helpers_test.go
rename to runtime/helpers_test.go
index fcb5a6f..fab379b 100644
--- a/firecracker-control/helpers_test.go
+++ b/runtime/helpers_test.go
@@ -11,7 +11,7 @@
 // express or implied. See the License for the specific language governing
 // permissions and limitations under the License.
 
-package service
+package main
 
 import (
 	"testing"
@@ -29,8 +29,6 @@ const (
 	capacity     = 400
 	memSize      = 4096
 	vcpuCount    = 2
-	mac          = "AA:FC:00:00:00:01"
-	hostDevName  = "tap0"
 )
 
 func TestMachineConfigurationFromProto(t *testing.T) {
diff --git a/runtime/service_integ_test.go b/runtime/service_integ_test.go
index b48197b..e9d5b6f 100644
--- a/runtime/service_integ_test.go
+++ b/runtime/service_integ_test.go
@@ -42,10 +42,13 @@ import (
 
 	"github.com/firecracker-microvm/firecracker-containerd/internal"
 	"github.com/firecracker-microvm/firecracker-containerd/proto"
+	fccontrol "github.com/firecracker-microvm/firecracker-containerd/proto/service/fccontrol/grpc"
 	"github.com/firecracker-microvm/firecracker-containerd/runtime/firecrackeroci"
 )
 
 const (
+	defaultNamespace = namespaces.Default
+
 	containerdSockPath = "/run/containerd/containerd.sock"
 	debianDockerImage  = "docker.io/library/debian:latest"
 
@@ -53,14 +56,15 @@ const (
 	naiveSnapshotterName = "firecracker-naive"
 	shimProcessName      = "containerd-shim-aws-firecracker"
 
-	defaultVMRootfsPath = "/var/lib/firecracker-containerd/runtime/hello-rootfs.ext4"
+	defaultVMRootfsPath = "/var/lib/firecracker-containerd/runtime/default-rootfs.img"
 	defaultVMNetDevName = "eth0"
+	varRunDir           = "/var/run/firecracker-containerd"
 )
 
 func TestShimExitsUponContainerKill_Isolated(t *testing.T) {
 	internal.RequiresIsolation(t)
 
-	ctx := namespaces.WithNamespace(context.Background(), namespaces.Default)
+	ctx := namespaces.WithNamespace(context.Background(), defaultNamespace)
 
 	client, err := containerd.New(containerdSockPath)
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", containerdSockPath)
@@ -108,7 +112,7 @@ func TestShimExitsUponContainerKill_Isolated(t *testing.T) {
 			return filepath.Base(processExecutable) == shimProcessName, nil
 		},
 	)
-	require.NoError(t, err, "failed waiting for expected shim process \"%s\" to come up", shimProcessName)
+	require.NoError(t, err, "failed waiting for expected shim process %q to come up", shimProcessName)
 	require.Len(t, shimProcesses, 1, "expected only one shim process to exist")
 	shimProcess := shimProcesses[0]
 
@@ -183,7 +187,7 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 		containersPerVM = 3
 	)
 
-	ctx := namespaces.WithNamespace(context.Background(), "default")
+	ctx := namespaces.WithNamespace(context.Background(), defaultNamespace)
 
 	client, err := containerd.New(containerdSockPath, containerd.WithDefaultRuntime(firecrackerRuntime))
 	require.NoError(t, err, "unable to create client to containerd service at %s, is containerd running?", containerdSockPath)
@@ -214,6 +218,24 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 			err = ioutil.WriteFile(rootfsPath, rootfsBytes, 0600)
 			require.NoError(t, err, "failed to copy vm rootfs to %s", rootfsPath)
 
+			fcClient := fccontrol.NewFirecrackerClient(client.Conn())
+			_, err = fcClient.CreateVM(ctx, &proto.CreateVMRequest{
+				VMID: strconv.Itoa(vmID),
+				RootDrive: &proto.FirecrackerDrive{
+					PathOnHost:   rootfsPath,
+					IsReadOnly:   false,
+					IsRootDevice: true,
+				},
+				NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
+					{
+						HostDevName: tapName,
+						MacAddress:  vmIDtoMacAddr(uint(vmID)),
+						AllowMMDS:   true,
+					},
+				},
+			})
+			require.NoError(t, err, "failed to create vm")
+
 			var containerWg sync.WaitGroup
 			for containerID := 0; containerID < containersPerVM; containerID++ {
 				containerWg.Add(1)
@@ -228,10 +250,7 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 						containerd.WithSnapshotter(naiveSnapshotterName),
 						containerd.WithNewSnapshot(snapshotName, image),
 						containerd.WithNewSpec(
-							oci.WithProcessArgs("sh", "-c", fmt.Sprintf("%s && %s",
-								fmt.Sprintf("cat /sys/class/net/%s/address", defaultVMNetDevName),
-								fmt.Sprintf("sleep %d", 10),
-							)),
+							oci.WithProcessArgs("cat", fmt.Sprintf("/sys/class/net/%s/address", defaultVMNetDevName)),
 							oci.WithHostNamespace(specs.NetworkNamespace),
 							firecrackeroci.WithVMID(strconv.Itoa(vmID)),
 						),
@@ -242,23 +261,7 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 					var stderr bytes.Buffer
 
 					newTask, err := newContainer.NewTask(ctx,
-						cio.NewCreator(cio.WithStreams(nil, bufio.NewWriter(&stdout), bufio.NewWriter(&stderr))),
-						func(ctx context.Context, _ *containerd.Client, ti *containerd.TaskInfo) error {
-							ti.Options = &proto.FirecrackerConfig{
-								RootDrive: &proto.FirecrackerDrive{
-									PathOnHost:   rootfsPath,
-									IsReadOnly:   false,
-									IsRootDevice: true,
-								},
-								NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
-									{
-										HostDevName: tapName,
-										MacAddress:  vmIDtoMacAddr(uint(vmID)),
-									},
-								},
-							}
-							return nil
-						})
+						cio.NewCreator(cio.WithStreams(nil, bufio.NewWriter(&stdout), bufio.NewWriter(&stderr))))
 					require.NoError(t, err, "failed to create task for container %s", containerName)
 
 					exitCh, err := newTask.Wait(ctx)
@@ -284,7 +287,30 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 				}(containerID)
 			}
 
+			// verify duplicate CreateVM call fails with right error
+			_, err = fcClient.CreateVM(ctx, &proto.CreateVMRequest{VMID: strconv.Itoa(vmID)})
+			require.Error(t, err, "did not receive expected error for duplicate CreateVM call")
+
+			// verify GetVMInfo returns expected data
+			vmInfoResp, err := fcClient.GetVMInfo(ctx, &proto.GetVMInfoRequest{VMID: strconv.Itoa(vmID)})
+			require.NoError(t, err, "failed to get VM Info for VM %d", vmID)
+			require.Equal(t, vmInfoResp.VMID, strconv.Itoa(vmID))
+			require.Equal(t, vmInfoResp.SocketPath, filepath.Join(varRunDir, defaultNamespace, strconv.Itoa(vmID), "firecracker.sock"))
+			require.Equal(t, vmInfoResp.LogFifoPath, filepath.Join(varRunDir, defaultNamespace, strconv.Itoa(vmID), "fc-logs.fifo"))
+			require.Equal(t, vmInfoResp.MetricsFifoPath, filepath.Join(varRunDir, defaultNamespace, strconv.Itoa(vmID), "fc-metrics.fifo"))
+
+			// just verify that updating the metadata doesn't return an error, a separate test case is needed
+			// to very the MMDS update propagates to the container correctly
+			_, err = fcClient.SetVMMetadata(ctx, &proto.SetVMMetadataRequest{
+				VMID:     strconv.Itoa(vmID),
+				Metadata: "{}",
+			})
+			require.NoError(t, err, "failed to set VM Metadata for VM %d", vmID)
+
 			containerWg.Wait()
+
+			_, err = fcClient.StopVM(ctx, &proto.StopVMRequest{VMID: strconv.Itoa(vmID)})
+			require.NoError(t, err, "failed to stop VM %d", vmID)
 		}(vmID)
 	}
 
diff --git a/runtime/service_test.go b/runtime/service_test.go
index 05a64bc..c07b3fd 100644
--- a/runtime/service_test.go
+++ b/runtime/service_test.go
@@ -18,9 +18,6 @@ import (
 	"syscall"
 	"testing"
 
-	"github.com/containerd/typeurl"
-	"github.com/firecracker-microvm/firecracker-containerd/proto"
-	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -38,7 +35,7 @@ func TestFindNextAvailableVsockCID(t *testing.T) {
 		sysCall = syscall.Syscall
 	}()
 
-	_, err := findNextAvailableVsockCID(context.Background())
+	_, _, err := findNextAvailableVsockCID(context.Background())
 	require.NoError(t, err,
 		"Do you have permission to interact with /dev/vhost-vsock?\n"+
 			"Grant yourself permission with `sudo setfacl -m u:${USER}:rw /dev/vhost-vsock`")
@@ -47,47 +44,6 @@ func TestFindNextAvailableVsockCID(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	cancel()
 
-	_, err = findNextAvailableVsockCID(ctx)
+	_, _, err = findNextAvailableVsockCID(ctx)
 	require.Equal(t, context.Canceled, err)
 }
-
-func TestParseCreateTaskOptsParsesFirecrackerConfig(t *testing.T) {
-	inFirecrackerConfig := &proto.FirecrackerConfig{
-		NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
-			{
-				MacAddress:  mac,
-				HostDevName: hostDevName,
-			},
-		},
-	}
-
-	protoFirecrackerConfig, err := typeurl.MarshalAny(inFirecrackerConfig)
-	require.NoError(t, err, "unable to marshal firecracker config proto message")
-	outFirecrackerConfig, outRuncOpts, err := parseCreateTaskOpts(protoFirecrackerConfig)
-	require.NoError(t, err, "unable to parse firecracker config from proto message")
-	if outRuncOpts != nil {
-		// assert.Equal is insufficient here as the nil comparison fails for
-		// empty Any protobuf message
-		t.Error("unexpected value parsed for runc options")
-	}
-	assert.Equal(t, 1, len(outFirecrackerConfig.NetworkInterfaces))
-	assert.Equal(t, hostDevName, outFirecrackerConfig.NetworkInterfaces[0].HostDevName)
-	assert.Equal(t, mac, outFirecrackerConfig.NetworkInterfaces[0].MacAddress)
-}
-
-func TestParseCreateTaskLeavesNonFirecrackerConfigAlong(t *testing.T) {
-	in := &proto.FirecrackerNetworkInterface{
-		MacAddress:  mac,
-		HostDevName: hostDevName,
-	}
-	protoIn, err := typeurl.MarshalAny(in)
-	require.NoError(t, err, "unable to marshal proto message")
-	outFirecrackerConfig, outOpts, err := parseCreateTaskOpts(protoIn)
-	require.NoError(t, err, "unable to parse firecracker config from proto message")
-	if outFirecrackerConfig != nil {
-		// assert.Equal is insufficient here as the nil comparison fails for
-		// empty Any protobuf message
-		t.Error("unexpected value parsed for firecracker config")
-	}
-	assert.Equal(t, protoIn, outOpts)
-}
diff --git a/runtime/taskfirecrackeropts_test.go b/runtime/taskfirecrackeropts_test.go
deleted file mode 100644
index ec00232..0000000
--- a/runtime/taskfirecrackeropts_test.go
+++ /dev/null
@@ -1,221 +0,0 @@
-// Copyright 2018-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
-//
-// Licensed under the Apache License, Version 2.0 (the "License"). You may
-// not use this file except in compliance with the License. A copy of the
-// License is located at
-//
-//	http://aws.amazon.com/apache2.0/
-//
-// or in the "license" file accompanying this file. This file is distributed
-// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
-// express or implied. See the License for the specific language governing
-// permissions and limitations under the License.
-
-package main
-
-import (
-	"testing"
-
-	"github.com/firecracker-microvm/firecracker-containerd/proto"
-	firecracker "github.com/firecracker-microvm/firecracker-go-sdk"
-	models "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
-	"github.com/stretchr/testify/assert"
-)
-
-const (
-	oneTimeBurst  = 800
-	refillTime    = 1000
-	capacity      = 400
-	rootDrivePath = "/path/to/root/drive"
-	loopDrivePath = "/path/to/loop/drive"
-	loopDriveRO   = false
-	memSize       = 4096
-	vcpuCount     = 2
-)
-
-var (
-	defaultDrivesBuilder = firecracker.NewDrivesBuilder(rootDrivePath).AddDrive(
-		loopDrivePath, loopDriveRO)
-	defaultDrives             = defaultDrivesBuilder.Build()
-	overrideRootDrivesBuilder = firecracker.NewDrivesBuilder(loopDrivePath)
-	additionalDrivesBuilder   = firecracker.NewDrivesBuilder(rootDrivePath).AddDrive(
-		loopDrivePath, loopDriveRO).AddDrive(
-		loopDrivePath, loopDriveRO, func(d *models.Drive) {
-			d.RateLimiter = &models.RateLimiter{
-				Ops: &models.TokenBucket{
-					OneTimeBurst: firecracker.Int64(oneTimeBurst),
-					RefillTime:   firecracker.Int64(refillTime),
-					Size:         firecracker.Int64(capacity),
-				},
-			}
-		})
-)
-
-func TestOverrideVMConfigFromTaskOpts(t *testing.T) {
-	var testCases = []struct {
-		name            string
-		in              firecracker.Config
-		inDrivesBuilder firecracker.DrivesBuilder
-		vmConfig        *proto.FirecrackerConfig
-		expectedOut     firecracker.Config
-		expectedDrives  []models.Drive
-		expectedErr     bool
-	}{
-		{
-			name:            "network config",
-			in:              firecracker.Config{},
-			inDrivesBuilder: defaultDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				NetworkInterfaces: []*proto.FirecrackerNetworkInterface{
-					{
-						MacAddress:  mac,
-						HostDevName: hostDevName,
-						AllowMMDS:   true,
-						InRateLimiter: &proto.FirecrackerRateLimiter{
-							Bandwidth: &proto.FirecrackerTokenBucket{
-								OneTimeBurst: oneTimeBurst,
-								RefillTime:   refillTime,
-								Capacity:     capacity,
-							},
-						},
-						OutRateLimiter: &proto.FirecrackerRateLimiter{
-							Bandwidth: &proto.FirecrackerTokenBucket{
-								OneTimeBurst: oneTimeBurst,
-								RefillTime:   refillTime,
-								Capacity:     capacity,
-							},
-						},
-					},
-					{
-						MacAddress:  mac,
-						HostDevName: hostDevName,
-					},
-				},
-			},
-			expectedErr: false,
-			expectedOut: firecracker.Config{
-				NetworkInterfaces: []firecracker.NetworkInterface{
-					{
-						MacAddress:  mac,
-						HostDevName: hostDevName,
-						AllowMMDS:   true,
-						InRateLimiter: &models.RateLimiter{
-							Bandwidth: &models.TokenBucket{
-								OneTimeBurst: firecracker.Int64(oneTimeBurst),
-								RefillTime:   firecracker.Int64(refillTime),
-								Size:         firecracker.Int64(capacity),
-							},
-						},
-						OutRateLimiter: &models.RateLimiter{
-							Bandwidth: &models.TokenBucket{
-								OneTimeBurst: firecracker.Int64(oneTimeBurst),
-								RefillTime:   firecracker.Int64(refillTime),
-								Size:         firecracker.Int64(capacity),
-							},
-						},
-					},
-					{
-						MacAddress:  mac,
-						HostDevName: hostDevName,
-					},
-				},
-			},
-			expectedDrives: defaultDrives,
-		},
-		{
-			name:            "machine config",
-			in:              firecracker.Config{},
-			inDrivesBuilder: defaultDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				MachineCfg: &proto.FirecrackerMachineConfiguration{
-					MemSizeMib: memSize,
-					VcpuCount:  vcpuCount,
-				},
-			},
-			expectedErr: false,
-			expectedOut: firecracker.Config{
-				MachineCfg: models.MachineConfiguration{
-					MemSizeMib: memSize,
-					VcpuCount:  vcpuCount,
-				},
-			},
-			expectedDrives: defaultDrivesBuilder.Build(),
-		},
-		{
-			name:            "invalid memory in machine config",
-			in:              firecracker.Config{},
-			inDrivesBuilder: defaultDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				MachineCfg: &proto.FirecrackerMachineConfiguration{
-					VcpuCount: vcpuCount,
-				},
-			},
-			expectedErr:    true,
-			expectedOut:    firecracker.Config{},
-			expectedDrives: defaultDrivesBuilder.Build(),
-		},
-		{
-			name:            "invalid vcpu count in machine config",
-			in:              firecracker.Config{},
-			inDrivesBuilder: defaultDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				MachineCfg: &proto.FirecrackerMachineConfiguration{
-					MemSizeMib: memSize,
-				},
-			},
-			expectedErr:    true,
-			expectedOut:    firecracker.Config{},
-			expectedDrives: defaultDrivesBuilder.Build(),
-		},
-		{
-			name:            "storage config overrides root drive",
-			in:              firecracker.Config{},
-			inDrivesBuilder: overrideRootDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				RootDrive: &proto.FirecrackerDrive{
-					PathOnHost: loopDrivePath,
-				},
-			},
-			expectedErr:    false,
-			expectedOut:    firecracker.Config{},
-			expectedDrives: overrideRootDrivesBuilder.Build(),
-		},
-		{
-			name:            "storage config adds additional drives",
-			in:              firecracker.Config{},
-			inDrivesBuilder: defaultDrivesBuilder,
-			vmConfig: &proto.FirecrackerConfig{
-				AdditionalDrives: []*proto.FirecrackerDrive{
-					{
-						PathOnHost: loopDrivePath,
-						RateLimiter: &proto.FirecrackerRateLimiter{
-							Ops: &proto.FirecrackerTokenBucket{
-								OneTimeBurst: oneTimeBurst,
-								RefillTime:   refillTime,
-								Capacity:     capacity,
-							},
-						},
-					},
-				},
-			},
-			expectedErr:    false,
-			expectedOut:    firecracker.Config{},
-			expectedDrives: additionalDrivesBuilder.Build(),
-		},
-	}
-	for _, tc := range testCases {
-		// scopelint gets sad if we use tc directly. The specific complaint is
-		// "Using the variable on range scope in function literal".
-		test := tc
-		t.Run(test.name, func(t *testing.T) {
-			out, drivesBuilder, err := overrideVMConfigFromTaskOpts(test.in, test.vmConfig, test.inDrivesBuilder)
-			if test.expectedErr {
-				assert.Error(t, err)
-			} else {
-				assert.NoError(t, err)
-				assert.Equal(t, test.expectedOut, out)
-				assert.Equal(t, test.expectedDrives, drivesBuilder.Build())
-			}
-		})
-	}
-}