proximl 1.0.0__tar.gz → 1.0.1__tar.gz

{proximl-1.0.0/proximl.egg-info → proximl-1.0.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proximl
-Version: 1.0.0
+Version: 1.0.1
 Summary: proxiML client SDK and command line utilities
 Home-page: https://github.com/proxiML/python-sdk
 Author: proxiML

{proximl-1.0.0 → proximl-1.0.1}/proximl/__init__.py

@@ -13,5 +13,5 @@ logging.basicConfig(
 logger = logging.getLogger(__name__)
 
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"
 __all__ = "ProxiML"

{proximl-1.0.0 → proximl-1.0.1}/proximl/cli/job/__init__.py

@@ -1,3 +1,4 @@
+import asyncio
 import click
 from webbrowser import open as browse
 from proximl.cli import cli, pass_config, search_by_id_name
@@ -28,6 +29,96 @@ def attach(config, job):
     config.proximl.run(found.attach())
 
 
+async def _connect_job(job, attach, config):
+    """
+    Async helper function to handle job connection with proper
+    handling of local input/output types and attach task management.
+    """
+    # Get job properties
+    model = job._job.get("model", {})
+    data = job._job.get("data", {})
+    model_local = model.get("source_type") == "local"
+    data_local = data.get("input_type") == "local"
+    output_local = data.get("output_type") == "local"
+    early_statuses = [
+        "new",
+        "waiting for data/model download",
+        "waiting for GPUs",
+        "waiting for resources",
+    ]
+
+    # Check if we need to wait for data/model download
+    # Only wait if status is early AND (data or model is local)
+    needs_upload_wait = job.status in early_statuses and (
+        model_local or data_local
+    )
+
+    if needs_upload_wait:
+        # Wait for job to reach data/model download status
+        await job.wait_for("waiting for data/model download", 3600)
+        await job.refresh()
+
+    # Start attach task early if requested
+    attach_task = None
+    if attach:
+        attach_task = asyncio.create_task(job.attach())
+
+    # Run first connect (upload if needed)
+    await job.connect()
+
+    # For notebook jobs, handle opening
+    if job.type == "notebook":
+        # Refresh to get latest status after connect
+        await job.refresh()
+
+        if job.status in early_statuses:
+            if attach_task:
+                await attach_task
+            click.echo("Launching...", file=config.stdout)
+            browse(job.notebook_url)
+            return
+        elif job.status not in [
+            "starting",
+            "running",
+            "reinitializing",
+            "copying",
+        ]:
+            if attach_task:
+                attach_task.cancel()
+            raise click.UsageError("Notebook job not running.")
+        else:
+            await job.wait_for("running")
+            if attach_task:
+                await attach_task
+            click.echo("Launching...", file=config.stdout)
+            browse(job.notebook_url)
+            return
+
+    # For non-notebook jobs, check if we need second connect (download)
+    # Refresh to get latest status after first connect
+    await job.refresh()
+
+    # Run second connect if output_type is local
+    # (as per user's requirement: "if the output_type is 'local'")
+    if output_local:
+        # Always wait for running status before second connect
+        # (as shown in user's example)
+        await job.wait_for("running", 3600)
+        await job.refresh()
+
+        # Create second connect task (download)
+        connect_task = asyncio.create_task(job.connect())
+
+        # Gather both attach and second connect tasks
+        if attach_task:
+            await asyncio.gather(attach_task, connect_task)
+        else:
+            await connect_task
+    elif attach_task:
+        # Just wait for attach if no second connect needed
+        await attach_task
+
+
 @job.command()
 @click.option(
     "--attach/--no-attach",
@@ -49,34 +140,7 @@ def connect(config, job, attach):
     if None is found:
         raise click.UsageError("Cannot find specified job.")
 
-    if found.type != "notebook":
-        if attach:
-            config.proximl.run(found.connect(), found.attach())
-        else:
-            config.proximl.run(found.connect())
-    else:
-        if found.status in [
-            "new",
-            "waiting for data/model download",
-            "waiting for GPUs",
-        ]:
-            if attach:
-                config.proximl.run(found.connect(), found.attach())
-                click.echo("Launching...", file=config.stdout)
-                browse(found.notebook_url)
-            else:
-                config.proximl.run(found.connect())
-        elif found.status not in [
-            "starting",
-            "running",
-            "reinitializing",
-            "copying",
-        ]:
-            raise click.UsageError("Notebook job not running.")
-        else:
-            config.proximl.run(found.wait_for("running"))
-            click.echo("Launching...", file=config.stdout)
-            browse(found.notebook_url)
+    config.proximl.run(_connect_job(found, attach, config))
 
 
 @job.command()

{proximl-1.0.0 → proximl-1.0.1}/proximl/jobs.py

@@ -321,10 +321,29 @@ class Job:
         ]:
             return self.url
 
-        # Check for invalid statuses
+        # Refresh to get latest job data first, so we can check worker statuses
+        await self.refresh()
+
+        # Check worker statuses - if any worker is uploading, allow connection
+        # This handles the case where job status might be "finished" but workers are still uploading
+        workers = self._job.get("workers", [])
+        has_uploading_workers = any(
+            worker.get("status") == "uploading" for worker in workers
+        ) if workers else False
+
+        # Log worker statuses for debugging
+        if workers:
+            worker_statuses = [
+                f"Worker {i+1}: {worker.get('status')}"
+                for i, worker in enumerate(workers)
+            ]
+            logging.debug(
+                f"Job status: {self.status}, Worker statuses: {', '.join(worker_statuses)}, Has uploading workers: {has_uploading_workers}"
+            )
+
+        # Check for invalid statuses (but allow "finished" if workers are still uploading)
         if self.status in [
             "failed",
-            "finished",
             "canceled",
             "archived",
             "removed",
@@ -335,34 +354,34 @@ class Job:
                 f"You can only connect to active jobs.",
             )
 
+        # Allow "finished" status if there are workers still uploading
+        # This handles reconnection scenarios where some workers are done but others are still uploading
+        if self.status == "finished":
+            if not has_uploading_workers:
+                raise SpecificationError(
+                    "status",
+                    f"You can only connect to active jobs.",
+                )
+            logging.info(
+                f"Job status is 'finished' but has {sum(1 for w in workers if w.get('status') == 'uploading')} worker(s) still uploading. Allowing connection to download remaining workers."
+            )
+            # If we have uploading workers, fall through to download logic
+
         # Only allow specific statuses for connect
         if self.status not in [
             "waiting for data/model download",
             "uploading",
             "running",
+            "finished",  # Allow finished if workers are still uploading
         ]:
             if self.status == "new":
                 await self.wait_for("waiting for data/model download")
             else:
                 raise SpecificationError(
                     "status",
-                    f"You can only connect to jobs in 'waiting for data/model download', 'uploading', or 'running' status.",
+                    f"You can only connect to jobs in 'waiting for data/model download', 'uploading', 'running', or 'finished' (with uploading workers) status.",
                 )
 
-        # Refresh to get latest job data
-        await self.refresh()
-
-        # Re-check status after refresh (status may have changed if attach() is running in parallel)
-        if self.status not in [
-            "waiting for data/model download",
-            "uploading",
-            "running",
-        ]:
-            raise SpecificationError(
-                "status",
-                f"Job status changed to '{self.status}'. You can only connect to jobs in 'waiting for data/model download', 'uploading', or 'running' status.",
-            )
-
         if self.status == "waiting for data/model download":
             # Upload model and/or data if local
             model = self._job.get("model", {})
@@ -421,7 +440,7 @@ class Job:
             if upload_tasks:
                 await asyncio.gather(*upload_tasks)
 
-        elif self.status in ["uploading", "running"]:
+        elif self.status in ["uploading", "running", "finished"]:
             # Download output if local
             data = self._job.get("data", {})
 
@@ -455,8 +474,15 @@ class Job:
                         f"Job has no workers.",
                     )
 
-                # Check if job is finished
-                if self.status in ["finished", "canceled", "failed"]:
+                # Check if job is in a terminal state AND all workers are finished
+                # Allow "finished" status if workers are still uploading
+                all_workers_finished = all(
+                    worker.get("status") in ["finished", "removed"]
+                    for worker in workers
+                )
+                if self.status in ["canceled", "failed"]:
+                    break
+                if self.status == "finished" and all_workers_finished:
                     break
 
                 # Check all workers for uploading status
@@ -467,6 +493,7 @@ class Job:
                     worker_status = worker.get("status")
 
                     # Start download for any worker that enters uploading status
+                    # This handles both new connections and reconnections where some workers are already uploading
                     if (
                         worker_status == "uploading"
                         and worker_id not in downloading_workers
@@ -478,6 +505,8 @@ class Job:
                             logging.warning(
                                 f"Worker {worker_id} in uploading status missing output_auth_token or output_hostname, skipping."
                             )
+                            # Mark as downloading to avoid retrying
+                            downloading_workers.add(worker_id)
                             continue
 
                         downloading_workers.add(worker_id)

{proximl-1.0.0 → proximl-1.0.1}/proximl/proximl.py

@@ -214,118 +214,129 @@ class ProxiML(object):
             "Content-Type": "application/json",
         }
         try:
-            tokens = self.auth.get_tokens()
-        except ProxiMLException as e:
-            raise e
-        except Exception:
-            raise ProxiMLException(
-                f"Error getting authorization tokens.  Verify configured credentials. Error: {traceback.format_exc()}"
-            )
-        async with aiohttp.ClientSession() as session:
-            done = False
-            async with session.ws_connect(
-                f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
-                headers=headers,
-                heartbeat=30,
-            ) as ws:
-                asyncio.create_task(
-                    ws.send_json(
-                        dict(
-                            action="getlogs",
-                            data=dict(
-                                type="init",
-                                entity=entity,
-                                id=id,
-                                project_uuid=project_uuid,
-                            ),
-                        )
-                    )
+            try:
+                tokens = self.auth.get_tokens()
+            except ProxiMLException as e:
+                raise e
+            except Exception:
+                raise ProxiMLException(
+                    f"Error getting authorization tokens.  Verify configured credentials. Error: {traceback.format_exc()}"
                 )
-                asyncio.create_task(
-                    ws.send_json(
-                        dict(
-                            action="subscribe",
-                            data=dict(
-                                type="logs",
-                                entity=entity,
-                                id=id,
-                                project_uuid=project_uuid,
-                            ),
+            async with aiohttp.ClientSession() as session:
+                done = False
+                async with session.ws_connect(
+                    f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
+                    headers=headers,
+                    heartbeat=30,
+                ) as ws:
+                    asyncio.create_task(
+                        ws.send_json(
+                            dict(
+                                action="getlogs",
+                                data=dict(
+                                    type="init",
+                                    entity=entity,
+                                    id=id,
+                                    project_uuid=project_uuid,
+                                ),
+                            )
                         )
                     )
-                )
-                async for msg in ws:
-                    if msg.type in (
-                        aiohttp.WSMsgType.CLOSED,
-                        aiohttp.WSMsgType.ERROR,
-                        aiohttp.WSMsgType.CLOSE,
-                    ):
-                        logging.debug(
-                            f"Websocket Received Closed Message.  Done? {done}"
+                    asyncio.create_task(
+                        ws.send_json(
+                            dict(
+                                action="subscribe",
+                                data=dict(
+                                    type="logs",
+                                    entity=entity,
+                                    id=id,
+                                    project_uuid=project_uuid,
+                                ),
+                            )
                         )
-                        await ws.close()
-                        break
-                    data = json.loads(msg.data)
-                    if data.get("type") == "end":
-                        done = True
-                        asyncio.create_task(delayed_close(ws))
-                    else:
-                        msg_handler(data)
-            logging.debug(f"Websocket Disconnected.  Done? {done}")
+                    )
+                    async for msg in ws:
+                        if msg.type in (
+                            aiohttp.WSMsgType.CLOSED,
+                            aiohttp.WSMsgType.ERROR,
+                            aiohttp.WSMsgType.CLOSE,
+                        ):
+                            logging.debug(
+                                f"Websocket Received Closed Message.  Done? {done}"
+                            )
+                            await ws.close()
+                            break
+                        data = json.loads(msg.data)
+                        if data.get("type") == "end":
+                            done = True
+                            asyncio.create_task(delayed_close(ws))
+                        else:
+                            msg_handler(data)
+                logging.debug(f"Websocket Disconnected.  Done? {done}")
 
-            connection_tries = 0
-            while not done:
-                tokens = self.auth.get_tokens()
-                try:
-                    async with session.ws_connect(
-                        f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
-                        headers=headers,
-                        heartbeat=30,
-                    ) as ws:
-                        asyncio.create_task(
-                            ws.send_json(
-                                dict(
-                                    action="subscribe",
-                                    data=dict(
-                                        type="logs",
-                                        entity=entity,
-                                        id=id,
-                                        project_uuid=project_uuid,
-                                    ),
+                connection_tries = 0
+                while not done:
+                    tokens = self.auth.get_tokens()
+                    try:
+                        async with session.ws_connect(
+                            f"wss://{self.ws_url}?Authorization={tokens.get('id_token')}",
+                            headers=headers,
+                            heartbeat=30,
+                        ) as ws:
+                            asyncio.create_task(
+                                ws.send_json(
+                                    dict(
+                                        action="subscribe",
+                                        data=dict(
+                                            type="logs",
+                                            entity=entity,
+                                            id=id,
+                                            project_uuid=project_uuid,
+                                        ),
+                                    )
                                 )
                             )
+                            async for msg in ws:
+                                if msg.type in (
+                                    aiohttp.WSMsgType.CLOSED,
+                                    aiohttp.WSMsgType.ERROR,
+                                    aiohttp.WSMsgType.CLOSE,
+                                ):
+                                    logging.debug(
+                                        f"Websocket Received Closed Message.  Done? {done}"
+                                    )
+                                    await ws.close()
+                                    break
+                                data = json.loads(msg.data)
+                                if data.get("type") == "end":
+                                    done = True
+                                    asyncio.create_task(delayed_close(ws))
+                                else:
+                                    msg_handler(data)
+                        connection_tries = 0
+                        logging.debug(f"Websocket Disconnected.  Done? {done}")
+                    except Exception as e:
+                        connection_tries += 1
+                        logging.debug(
+                            f"Connection error: {traceback.format_exc()}"
                         )
-                        async for msg in ws:
-                            if msg.type in (
-                                aiohttp.WSMsgType.CLOSED,
-                                aiohttp.WSMsgType.ERROR,
-                                aiohttp.WSMsgType.CLOSE,
-                            ):
-                                logging.debug(
-                                    f"Websocket Received Closed Message.  Done? {done}"
-                                )
-                                await ws.close()
-                                break
-                            data = json.loads(msg.data)
-                            if data.get("type") == "end":
-                                done = True
-                                asyncio.create_task(delayed_close(ws))
-                            else:
-                                msg_handler(data)
-                    connection_tries = 0
-                    logging.debug(f"Websocket Disconnected.  Done? {done}")
-                except Exception as e:
-                    connection_tries += 1
-                    logging.debug(
-                        f"Connection error: {traceback.format_exc()}"
-                    )
-                    if connection_tries == 5:
-                        raise ApiError(
-                            500,
-                            {
-                                "message": f"Connection error: {traceback.format_exc()}"
-                            },
-                        )
+                        if connection_tries == 5:
+                            raise ApiError(
+                                500,
+                                {
+                                    "message": f"Connection error: {traceback.format_exc()}"
+                                },
+                            )
+        except GeneratorExit:
+            # Handle graceful shutdown - GeneratorExit is raised during
+            # event loop cleanup. Don't re-raise to avoid "coroutine ignored"
+            # warnings.
+            logging.debug("Websocket subscription cancelled during shutdown")
+            return
+        except asyncio.CancelledError:
+            # Re-raise CancelledError to properly propagate task cancellation
+            logging.debug("Websocket subscription task cancelled")
+            raise
 
     def set_active_project(self, project_uuid):
         CONFIG_DIR = os.path.expanduser(

{proximl-1.0.0 → proximl-1.0.1}/proximl/utils/transfer.py

@@ -1,5 +1,6 @@
 import os
 import re
+import math
 import asyncio
 import aiohttp
 import aiofiles
@@ -28,6 +29,8 @@ RETRY_STATUSES = {
 # Additional retries for DNS/connection errors (ClientConnectorError)
 DNS_MAX_RETRIES = 7  # More retries for DNS resolution issues
 DNS_INITIAL_DELAY = 1  # Initial delay in seconds before first DNS retry
+# Ping warmup timeout: calculate retries so last retry is this many seconds after first try
+PING_WARMUP_TIMEOUT = 8 * 60  # 8 minutes in seconds
 
 
 def normalize_endpoint(endpoint):
@@ -56,6 +59,38 @@ def normalize_endpoint(endpoint):
     return endpoint
 
 
+def calculate_ping_retries(timeout_seconds, backoff_base):
+    """
+    Calculate the number of retries needed for ping warmup to reach timeout.
+
+    With exponential backoff, if we have n retries, the total wait time is:
+    backoff_base^1 + backoff_base^2 + ... + backoff_base^(n-1) = (backoff_base^n - backoff_base) / (backoff_base - 1)
+
+    For backoff_base = 2, this simplifies to: 2^n - 2
+
+    Args:
+        timeout_seconds: Total timeout in seconds
+        backoff_base: Exponential backoff base
+
+    Returns:
+        Number of retries needed to reach or exceed the timeout
+    """
+    if backoff_base == 2:
+        # Simplified calculation for base 2
+        # 2^n - 2 >= timeout_seconds
+        # 2^n >= timeout_seconds + 2
+        # n >= log2(timeout_seconds + 2)
+        n = math.ceil(math.log2(timeout_seconds + 2))
+    else:
+        # General case: (backoff_base^n - backoff_base) / (backoff_base - 1) >= timeout_seconds
+        # backoff_base^n >= timeout_seconds * (backoff_base - 1) + backoff_base
+        # n >= log_base(timeout_seconds * (backoff_base - 1) + backoff_base)
+        target = timeout_seconds * (backoff_base - 1) + backoff_base
+        n = math.ceil(math.log(target, backoff_base))
+
+    return max(1, int(n))
+
+
 async def ping_endpoint(
     endpoint, auth_token, max_retries=MAX_RETRIES, retry_backoff=RETRY_BACKOFF
 ):
@@ -68,10 +103,13 @@ async def ping_endpoint(
     Creates a fresh TCPConnector for each attempt to force fresh DNS resolution
     and avoid stale DNS cache issues.
 
+    For ping warmup, calculates retries dynamically to ensure the last retry
+    occurs PING_WARMUP_TIMEOUT seconds after the first try.
+
     Args:
         endpoint: Server endpoint URL
         auth_token: Authentication token
-        max_retries: Maximum number of retry attempts
+        max_retries: Maximum number of retry attempts (ignored for ping warmup)
         retry_backoff: Exponential backoff base
 
     Raises:
@@ -80,7 +118,18 @@ async def ping_endpoint(
     """
     endpoint = normalize_endpoint(endpoint)
     attempt = 1
-    effective_max_retries = max_retries
+    # Calculate retries for ping warmup to reach PING_WARMUP_TIMEOUT
+    # Allow max_retries to override when explicitly provided (for testing)
+    if max_retries == MAX_RETRIES:
+        # Use default, calculate retries for ping warmup
+        ping_max_retries = calculate_ping_retries(
+            PING_WARMUP_TIMEOUT, retry_backoff
+        )
+        effective_max_retries = ping_max_retries
+    else:
+        # Use explicitly provided max_retries (for testing)
+        effective_max_retries = max_retries
+        ping_max_retries = max_retries  # For DNS error handling below
 
     while attempt <= effective_max_retries:
         # Create a fresh connector for each attempt to force DNS re-resolution
@@ -122,8 +171,9 @@ async def ping_endpoint(
             )
         except ClientConnectorError as e:
             # DNS resolution errors need more retries and initial delay
-            if effective_max_retries == max_retries:
-                effective_max_retries = max(max_retries, DNS_MAX_RETRIES)
+            # Use the higher of DNS_MAX_RETRIES or calculated ping retries
+            if effective_max_retries == ping_max_retries:
+                effective_max_retries = max(ping_max_retries, DNS_MAX_RETRIES)
 
             if attempt < effective_max_retries:
                 # Use initial delay for first retry, then exponential backoff
@@ -419,8 +469,11 @@ async def download(endpoint, auth_token, target_directory, file_name=None):
                             error_text = await response.text()
                         except Exception:
                             error_text = f"Unable to read response body (status: {response.status})"
-                        raise ConnectionError(
-                            f"Failed to get server info (status {response.status}): {error_text}"
+                        raise ClientResponseError(
+                            request_info=response.request_info,
+                            history=response.history,
+                            status=response.status,
+                            message=error_text,
                         )
                     return await response.json()
 
@@ -444,7 +497,14 @@ async def download(endpoint, auth_token, target_directory, file_name=None):
                     "Warning: /info endpoint not available, defaulting to TAR stream mode"
                 )
             else:
-                # For other errors, re-raise
+                # For other errors, convert ClientResponseError to ConnectionError
+                # to maintain backward compatibility
+                if isinstance(e, ClientResponseError):
+                    error_msg = getattr(e, "message", str(e))
+                    raise ConnectionError(
+                        f"Failed to get server info (status {e.status}): {error_msg}"
+                    )
+                # For ConnectionError, re-raise as-is
                 raise
 
         # Download the archive

{proximl-1.0.0 → proximl-1.0.1/proximl.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: proximl
-Version: 1.0.0
+Version: 1.0.1
 Summary: proxiML client SDK and command line utilities
 Home-page: https://github.com/proxiML/python-sdk
 Author: proxiML