Refactoring

.gitignore → .gitignogre

@@ -1,3 +1,4 @@
+.idea/
 aws_config/
 
 # Byte-compiled / optimized / DLL files

README.md

@@ -25,6 +25,11 @@ python scripts/ec2_setup.py
 
 - (Optional) Set up [Docker](https://docs.docker.com/engine/installation/). This is required on the target machine if running in a Docker-enabled mode.
 
+- (Optional) Set up GCP
+  - https://cloud.google.com/sdk/docs/quickstart-debian-ubuntu
+  - https://cloud.google.com/compute/docs/tutorials/python-guide
+  - https://cloud.google.com/storage/docs/reference/libraries#client-libraries-install-python
+
 
 ## Example
 

doodad/arg_parse.py

@@ -9,7 +9,7 @@
 
 
 __ARGS = None
-def __get_arg_config():
+def __get_unpickled_arg_config():
     """
     global __ARGS
     if __ARGS is not None:
@@ -34,7 +34,7 @@ def __get_arg_config():
 
 
 def get_args(key=None, default=None):
-    args = __get_arg_config()
+    args = __get_unpickled_arg_config()
 
     if args.args_data:
         if args.use_cloudpickle:

doodad/gcp/gcp_util.py

@@ -0,0 +1,55 @@
+import os
+
+from doodad.utils import hash_file, call_and_wait, CommandBuilder, REPO_DIR
+
+GCP_STARTUP_SCRIPT_PATH = os.path.join(REPO_DIR, "scripts/gcp/gcp_startup_script.sh")
+GCP_SHUTDOWN_SCRIPT_PATH = os.path.join(REPO_DIR, "scripts/gcp/gcp_shutdown_script.sh")
+
+def upload_file_to_gcp_storage(
+    bucket_name,
+    file_name,
+    remote_filename=None,
+    dry=False,
+    check_exists=True
+):
+    from google.cloud import storage
+    storage_client = storage.Client()
+    bucket = storage_client.get_bucket(bucket_name)
+
+    if remote_filename is None:
+        remote_filename = os.path.basename(file_name)
+    remote_path = 'doodad/mount/' + remote_filename
+    blob = bucket.blob(remote_path)
+    if check_exists and blob.exists(storage_client):
+        print("{remote_path} already exists".format(remote_path=remote_path))
+        return remote_path
+    blob.upload_from_filename(file_name)
+    return remote_path
+
+def get_machine_type(zone, instance_type):
+    return "zones/{zone}/machineTypes/{instance_type}".format(
+            zone=zone,
+            instance_type=instance_type,
+    )
+
+def get_gpu_type(project, zone, gpu_model):
+    """
+    Check the available gpu models for each zone
+    https://cloud.google.com/compute/docs/gpus/
+    """
+    assert gpu_model in [
+        'nvidia-tesla-p4',
+        'nvidia-tesla-k80',
+        'nvidia-tesla-v100',
+        'nvidia-tesla-p100'
+    ]
+
+    return (
+        "https://www.googleapis.com/compute/v1/"
+        "projects/{project}/zones/{zone}/acceleratorTypes/{gpu_model}".format(
+            project=project,
+            zone=zone,
+            gpu_model=gpu_model
+        )
+    )
+

doodad/launch_tools.py

@@ -22,13 +22,12 @@ def launch_python(
         mode=LOCAL,
         mount_points=None,
         args=None,
-        env=None,
-        dry=False,
         fake_display=False,
         target_mount_dir='target',
-        verbose=False,
         use_cloudpickle=False,
         target_mount=None,
+        launch_locally=None,
+        **launch_command_kwargs
 ):
     """
 
@@ -37,7 +36,6 @@ def launch_python(
     :param mode:
     :param mount_points:
     :param args:
-    :param env:
     :param dry:
     :param fake_display:
     :param target_mount_dir:
@@ -50,19 +48,21 @@ def launch_python(
         args = {}
     if mount_points is None:
         mount_points = []
+    if launch_locally is None:
+        launch_locally = isinstance(mode, Local)
 
     if target_mount is None:
         # mount
         target_dir = os.path.dirname(target)
         if not target_mount_dir:
             target_mount_dir = target_dir
         target_mount_dir = os.path.join(target_mount_dir, os.path.basename(target_dir))
-        if isinstance(mode, Local):
+        if launch_locally:
             target_mount = MountLocal(local_dir=target_dir)
         else:
             target_mount = MountLocal(local_dir=target_dir, mount_point=target_mount_dir)
     mount_points = mount_points + [target_mount]
-    target_full_path = os.path.join(target_mount.docker_mount_dir(), os.path.basename(target))
+    target_full_path = os.path.join(target_mount.mount_dir(), os.path.basename(target))
 
     command = make_python_command(
         target_full_path,
@@ -71,7 +71,8 @@ def launch_python(
         fake_display=fake_display,
         use_cloudpickle=use_cloudpickle,
     )
-    mode.launch_command(command, mount_points=mount_points, dry=dry, verbose=verbose)
+    mode.launch_command(command, mount_points=mount_points,
+                        **launch_command_kwargs)
     return target_mount
 
 HEADLESS = 'xvfb-run -a -s "-ac -screen 0 1400x900x24 +extension RANDR"'
@@ -90,8 +91,8 @@ def make_python_command(
 
     args_encoded, cp_version = encode_args(args, cloudpickle=use_cloudpickle)
     if args:
-        cmd = '%s=%s %s=%s %s=%s %s' % (ARGS_DATA, args_encoded, 
-                USE_CLOUDPICKLE, str(int(use_cloudpickle)), 
+        cmd = '%s=%s %s=%s %s=%s %s' % (ARGS_DATA, args_encoded,
+                USE_CLOUDPICKLE, str(int(use_cloudpickle)),
                 CLOUDPICKLE_VERSION, cp_version,
                 cmd)