feat: abstract worker

Author: justinmerrell

README.md

@@ -10,6 +10,7 @@
 - [Installation](#installation)
 - [API Language Library](#api-language-library)
 - [SDK - Serverless Worker](#sdk---serverless-worker)
+  - [Quick Start](#quick-start)
 
 ## Installation
 
@@ -30,3 +31,15 @@ runpod.api_key = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"
 ## SDK - Serverless Worker
 
 This python package can also be used to create a serverless worker that can be deployed to RunPod.
+
+### Quick Start
+
+Create an executable file called 'worker' in the root of your project that contains the following:
+
+```python
+#!/usr/bin/env python
+
+import runpod
+
+runpod.serverless.pod_worker()
+```

src/runpod/serverless/pod_worker.py

@@ -11,37 +11,44 @@
 from .modules.logging import log
 
 
-worker_life = lifecycle.LifecycleManager()
+def start_worker():
+    '''
+    Starts the worker.
+    '''
+    worker_life = lifecycle.LifecycleManager()
+
+    if not worker_life.is_worker_zero:
+        log("Not worker zero, starting TTL timer thread.")
+        threading.Thread(target=worker_life.check_worker_ttl_thread).start()
+    else:
+        log("Worker zero, not starting TTL timer thread.")
 
-if not worker_life.is_worker_zero:
-    log("Not worker zero, starting TTL timer thread.")
-    threading.Thread(target=worker_life.check_worker_ttl_thread).start()
-else:
-    log("Worker zero, not starting TTL timer thread.")
+    while True:
+        if os.environ.get('TEST_LOCAL', 'false') != 'true':
+            next_job = job.get(worker_life.worker_id)
+        else:
+            next_job = job.get_local()
 
+        if next_job is not None:
+            worker_life.work_in_progress = True  # Rests when "reset_worker_ttl" is called
 
-while True:
-    if os.environ.get('TEST_LOCAL', 'false') != 'true':
-        next_job = job.get(worker_life.worker_id)
-    else:
-        next_job = job.get_local()
+            try:
+                output_urls, job_duration_ms = job.run(next_job['id'], next_job['input'])
+                job.post(worker_life.worker_id, next_job['id'], output_urls, job_duration_ms)
+            except ValueError as err:
+                job.error(worker_life.worker_id, next_job['id'], str(err))
 
-    if next_job is not None:
-        worker_life.work_in_progress = True  # Rests when "reset_worker_ttl" is called
+            # -------------------------------- Job Cleanup ------------------------------- #
+            shutil.rmtree("input_objects", ignore_errors=True)
+            shutil.rmtree("output_objects", ignore_errors=True)
+            os.remove('output.zip')
 
-        try:
-            output_urls, job_duration_ms = job.run(next_job['id'], next_job['input'])
-            job.post(worker_life.worker_id, next_job['id'], output_urls, job_duration_ms)
-        except ValueError as err:
-            job.error(worker_life.worker_id, next_job['id'], str(err))
+            worker_life.reset_worker_ttl()
 
-        # -------------------------------- Job Cleanup ------------------------------- #
-        shutil.rmtree("input_objects", ignore_errors=True)
-        shutil.rmtree("output_objects", ignore_errors=True)
-        os.remove('output.zip')
+        if os.environ.get('TEST_LOCAL', 'false') == 'true':
+            log("Local testing complete, exiting.")
+            break
 
-        worker_life.reset_worker_ttl()
 
-    if os.environ.get('TEST_LOCAL', 'false') == 'true':
-        log("Local testing complete, exiting.")
-        break
+if __name__ == '__main__':
+    start_worker()

tests/test_serverless_module_download.py

@@ -30,9 +30,10 @@ def __init__(self, content, status_code):
 class TestDownloadInputObjects(unittest.TestCase):
     ''' Tests for download_input_objects '''
 
+    @patch('os.makedirs', return_value=None)
     @patch('requests.get', side_effect=mock_requests_get)
     @patch('builtins.open', new_callable=mock_open)
-    def test_download_input_objects(self, mock_open_file, mock_get):
+    def test_download_input_objects(self, mock_open_file, mock_get, mock_makedirs):
         '''
         Tests download_input_objects
         '''
@@ -43,3 +44,4 @@ def test_download_input_objects(self, mock_open_file, mock_get):
         self.assertEqual(len(objects), 1)
         self.assertIn('https://example.com/picture.jpg', mock_get.call_args_list[0][0])
         mock_open_file.assert_called_once_with(objects[0], 'wb')
+        mock_makedirs.assert_called_once_with('input_objects', exist_ok=True)