Get skypilot status working

Author: deep1401

src/lattice/providers/example.py

@@ -44,37 +44,37 @@ def example_skypilot():
     status = provider.get_cluster_status("my-cluster")
     print(f"   Status: {status.state}, Message: {status.status_message}\n")
 
-    # Get cluster resources
-    print("3. Getting cluster resources...")
-    resources = provider.get_cluster_resources("my-cluster")
-    print(f"   Resources: {resources.num_nodes} nodes, GPUs: {resources.gpus}\n")
-
-    # Submit a job
-    print("4. Submitting job...")
-    job_config = JobConfig(
-        command="python train.py",
-        job_name="training-job",
-        env_vars={"CUDA_VISIBLE_DEVICES": "0"},
-    )
-    job_result = provider.submit_job("my-cluster", job_config)
-    print(f"   Job ID: {job_result.get('job_id')}\n")
-
-    # List jobs
-    print("5. Listing jobs...")
-    jobs = provider.list_jobs("my-cluster")
-    for job in jobs:
-        print(f"   Job {job.job_id}: {job.state} - {job.job_name}")
-
-    # Get job logs
-    if jobs:
-        print(f"\n6. Getting logs for job {jobs[0].job_id}...")
-        logs = provider.get_job_logs("my-cluster", jobs[0].job_id, tail_lines=50)
-        print(f"   Logs (first 200 chars): {str(logs)[:200]}...\n")
-
-    # Stop cluster
-    print("7. Stopping cluster...")
-    stop_result = provider.stop_cluster("my-cluster")
-    print(f"   Result: {stop_result}\n")
+    # # Get cluster resources
+    # print("3. Getting cluster resources...")
+    # resources = provider.get_cluster_resources("my-cluster")
+    # print(f"   Resources: {resources.num_nodes} nodes, GPUs: {resources.gpus}\n")
+
+    # # Submit a job
+    # print("4. Submitting job...")
+    # job_config = JobConfig(
+    #     command="python train.py",
+    #     job_name="training-job",
+    #     env_vars={"CUDA_VISIBLE_DEVICES": "0"},
+    # )
+    # job_result = provider.submit_job("my-cluster", job_config)
+    # print(f"   Job ID: {job_result.get('job_id')}\n")
+
+    # # List jobs
+    # print("5. Listing jobs...")
+    # jobs = provider.list_jobs("my-cluster")
+    # for job in jobs:
+    #     print(f"   Job {job.job_id}: {job.state} - {job.job_name}")
+
+    # # Get job logs
+    # if jobs:
+    #     print(f"\n6. Getting logs for job {jobs[0].job_id}...")
+    #     logs = provider.get_job_logs("my-cluster", jobs[0].job_id, tail_lines=50)
+    #     print(f"   Logs (first 200 chars): {str(logs)[:200]}...\n")
+
+    # # Stop cluster
+    # print("7. Stopping cluster...")
+    # stop_result = provider.stop_cluster("my-cluster")
+    # print(f"   Result: {stop_result}\n")
 
 
 def example_slurm():

src/lattice/providers/skypilot.py

@@ -349,45 +349,106 @@ def get_cluster_status(self, cluster_name: str) -> ClusterStatus:
         if self.default_entrypoint_command:
             body_json.setdefault("entrypoint_command", self.default_entrypoint_command)
         body_json.setdefault("using_remote_api_server", False)
-        body_json.setdefault("override_skypilot_config", {})
-        
+        body_json.setdefault("override_skypilot_config", {})        
         # Use SkyPilot's make_authenticated_request (matches SDK exactly)
         response = self._make_authenticated_request(
-            'POST', '/status', json_data=body_json, timeout=5
+            'POST', '/status', json_data=body_json, timeout=10
         )
 
+        # Check response status
+        if hasattr(response, 'status_code'):
+            if response.status_code != 200:
+                return ClusterStatus(
+                    cluster_name=cluster_name,
+                    state=ClusterState.UNKNOWN,
+                    status_message=f"API returned status code {response.status_code}",
+                )
+        
+        # Parse response content
+        response_content = None
+        is_null_response = False
+        if hasattr(response, 'content'):
+            if response.content == b'null' or response.content == b'':
+                is_null_response = True
+            else:
+                try:
+                    response_content = response.json() if hasattr(response, 'json') else None
+                except Exception as e:
+                    print(f"Warning: Could not parse response as JSON: {e}")
+                    response_content = None
+        
         # Get request ID using SkyPilot's method (matches SDK exactly)
         request_id = None
         if self._server_common:
             try:
                 request_id = self._server_common.get_request_id(response)
-            except Exception:
-                pass
+                # For debugging: return early after getting request_id
+                print(f"Got request_id: {request_id}")
+                return ClusterStatus(
+                    cluster_name=cluster_name,
+                    state=ClusterState.UNKNOWN,
+                    status_message=f"Debug: Got request_id {request_id}",
+                )
+            except Exception as e:
+                # If get_request_id fails, try to extract from response directly
+                if response_content and isinstance(response_content, dict):
+                    request_id = response_content.get("request_id")
+                # Also check headers
+                if not request_id and hasattr(response, 'headers'):
+                    request_id = (
+                        response.headers.get('X-Request-ID') or
+                        response.headers.get('Request-ID') or
+                        response.headers.get('X-Request-Id')
+                    )
+                if not request_id:
+                    print(f"Warning: Could not extract request_id: {e}")
+                    print(f"Response status: {getattr(response, 'status_code', 'unknown')}")
+                    print(f"Response headers: {getattr(response, 'headers', {})}")
+                    print(f"Response content: {getattr(response, 'content', b'')[:200]}")
+        
+        # If response is null and we don't have a request ID, the cluster likely doesn't exist
+        if is_null_response and not request_id:
+            return ClusterStatus(
+                cluster_name=cluster_name,
+                state=ClusterState.UNKNOWN,
+                status_message="API returned null response - cluster may not exist or request format is incorrect",
+            )
 
         # Get the actual result from the request ID
+        clusters = []
         if request_id and self._server_common:
             try:
                 # Use server_common.get() to get the actual response
-                clusters = self._server_common.get(request_id)
-            except Exception:
+                request_payload = self._server_common.get(request_id)
+                if request_payload and hasattr(request_payload, 'return_value'):
+                    # return_value is a JSON string of the list of clusters
+                    if request_payload.return_value:
+                        clusters = json.loads(request_payload.return_value)
+                elif request_payload and isinstance(request_payload, (list, dict)):
+                    # Sometimes the response is directly the clusters list
+                    clusters = request_payload if isinstance(request_payload, list) else [request_payload]
+            except Exception as e:
+                print(f"Warning: Could not get clusters from request_id {request_id}: {e}")
                 # Fallback: try to parse response directly
                 try:
-                    if hasattr(response, 'json'):
+                    if response_content:
+                        clusters = response_content if isinstance(response_content, list) else response_content.get("clusters", [])
+                    elif hasattr(response, 'json'):
                         result = response.json()
                         clusters = result if isinstance(result, list) else result.get("clusters", [])
-                    else:
-                        clusters = []
-                except Exception:
+                except Exception as parse_error:
+                    print(f"Warning: Could not parse response: {parse_error}")
                     clusters = []
         else:
             # Fallback: try to parse response directly
             try:
-                if hasattr(response, 'json'):
+                if response_content:
+                    clusters = response_content if isinstance(response_content, list) else response_content.get("clusters", [])
+                elif hasattr(response, 'json'):
                     result = response.json()
                     clusters = result if isinstance(result, list) else result.get("clusters", [])
-                else:
-                    clusters = []
-            except Exception:
+            except Exception as e:
+                print(f"Warning: Could not parse response directly: {e}")
                 clusters = []
 
         # Handle empty or invalid responses