openweights 0.2.0__tar.gz → 0.4.0__tar.gz

{openweights-0.2.0 → openweights-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openweights
-Version: 0.2.0
+Version: 0.4.0
 Summary: An openai-like sdk for finetuning and batch inference
 License: MIT
 License-File: LICENSE.txt

{openweights-0.2.0 → openweights-0.4.0}/entrypoint.sh

@@ -1,14 +1,23 @@
 #!/bin/bash
 
+echo "[$(date)] Starting entrypoint script"
+
 # Add public keys to authorized_keys
+echo "[$(date)] Checking for PUBLIC_KEY environment variable"
 if [ -n "$PUBLIC_KEY" ]; then
+    echo "[$(date)] Setting up SSH public key"
     mkdir -p /root/.ssh
     echo "$PUBLIC_KEY" > /root/.ssh/authorized_keys
     chmod 600 /root/.ssh/authorized_keys
+    echo "[$(date)] SSH public key setup completed"
+else
+    echo "[$(date)] No PUBLIC_KEY provided, skipping SSH key setup"
 fi
 
 # if OW_COMMIT is set, checkout the commit
+echo "[$(date)] Checking for OW_COMMIT environment variable"
 if [ -n "$OW_COMMIT" ]; then
+    echo "[$(date)] Starting repository checkout for commit: $OW_COMMIT"
     rm -rf openweights
     git clone https://github.com/longtermrisk/openweights.git openweights_dev
     cd openweights_dev
@@ -16,15 +25,25 @@ if [ -n "$OW_COMMIT" ]; then
     mv openweights ../openweights
     cd ..
     rm -rf openweights_dev
+    echo "[$(date)] Repository checkout completed"
+else
+    echo "[$(date)] No OW_COMMIT specified, skipping repository checkout"
 fi
 
 # Login to huggingface
+echo "[$(date)] Attempting to login to Hugging Face"
 python3 -c "from huggingface_hub.hf_api import HfFolder; import os; HfFolder.save_token(os.environ['HF_TOKEN'])"
+echo "[$(date)] Hugging Face login completed"
 
 # Generate SSH host keys
+echo "[$(date)] Generating SSH host keys"
 ssh-keygen -A
+echo "[$(date)] SSH host keys generated"
+
 # Start SSH service
+echo "[$(date)] Starting SSH service"
 service ssh start
+echo "[$(date)] SSH service started"
 
 # Print sshd logs to stdout
 tail -f /var/log/auth.log &
@@ -34,6 +53,7 @@ tail -f /var/log/auth.log &
 touch main.log
 
 # Start a simple Python HTTP server to serve files from logs/
+echo "[$(date)] Starting HTTP log server on port 10101"
 python3 -c '
 import http.server
 import socketserver
@@ -63,12 +83,13 @@ with socketserver.TCPServer(("", 10101), LogHandler) as httpd:
     httpd.serve_forever()
 ' &
 
-
-mkdir logs
+echo "[$(date)] HTTP log server started"
 
 # Execute the main application or run in dev mode
 if [ "$OW_DEV" = "true" ]; then
+    echo "[$(date)] Starting in development mode"
     exec tail -f /dev/null
 else
-    exec python3 openweights/worker/main.py > logs/main 2>&1
+    echo "[$(date)] Starting main application"
+    exec python3 openweights/worker/main.py \ > >(tee logs/main) \ 2> >(tee -a logs/main >&2)
 fi

{openweights-0.2.0 → openweights-0.4.0}/openweights/client/jobs.py

@@ -153,9 +153,19 @@ class Jobs:
     
     def compute_id(self, data: Dict[str, Any]) -> str:
         """Compute job ID from data"""
-        return f"{self.id_predix}-{hashlib.sha256(json.dumps(data).encode() + self._org_id.encode()).hexdigest()[:12]}"
-    
-    @backoff.on_exception(backoff.constant, Exception, interval=1, max_time=60, max_tries=60, on_backoff=lambda details: print(f"Retrying... {details['exception']}"))
+        job_id = f"{self.id_predix}-{hashlib.sha256(json.dumps(data).encode() + self._org_id.encode()).hexdigest()[:12]}"
+        if "job_id_suffix" in data.get("validated_params", {}):
+            job_id += f"-{data['validated_params']['job_id_suffix']}"
+        return job_id
+
+    @backoff.on_exception(
+        backoff.constant,
+        Exception,
+        interval=1,
+        max_time=60,
+        max_tries=60,
+        on_backoff=lambda details: print(f"Retrying... {details['exception']}"),
+    )
     def get_or_create_or_reset(self, data: Dict[str, Any]) -> Dict[str, Any]:
         """If job exists and is [pending, in_progress, completed] return it.
         If job exists and is [failed, canceled] reset it to pending and return it.

{openweights-0.2.0 → openweights-0.4.0}/openweights/cluster/start_runpod.py

@@ -25,13 +25,44 @@ IMAGES = {
 }
 
 GPUs = {
-    'A6000': 'NVIDIA RTX 6000 Ada Generation',
-    'A100': 'NVIDIA A100 80GB PCIe',
-    'A100S': 'NVIDIA A100-SXM4-80GB',
-    'H100': 'NVIDIA H100 PCIe',
-    'H100N': 'NVIDIA H100 NVL',
-    'H100S': 'NVIDIA H100 80GB HBM3',
-    "H200": "NVIDIA H200"
+    # References found at https://rest.runpod.io/v1/docs#v-0-106
+    "6000Ada": "NVIDIA RTX 6000 Ada Generation",
+    "4000Ada": "NVIDIA RTX 4000 Ada Generation",
+    "2000Ada": "NVIDIA RTX 2000 Ada Generation",
+    "A6000": "NVIDIA RTX A6000",
+    "A4000": "NVIDIA RTX A4000",
+    "A2000": "NVIDIA RTX A2000",
+    "A100": "NVIDIA A100 80GB PCIe",  # Default A100 - 80GB
+    "A100_80": "NVIDIA A100 80GB PCIe",
+    "A100_40": "NVIDIA A100-SXM4-40GB",
+    "A100S": "NVIDIA A100-SXM4-80GB",
+    "H100": "NVIDIA H100 PCIe",
+    "H100N": "NVIDIA H100 NVL",
+    "H100S": "NVIDIA H100 80GB HBM3",
+    "H200": "NVIDIA H200",
+    "RTX4090": "NVIDIA GeForce RTX 4090",
+    "A5000": "NVIDIA RTX A5000",
+    "A40": "NVIDIA A40",
+    "A4500": "NVIDIA RTX A4500",
+    "L4": "NVIDIA L4",
+    "L40": "NVIDIA L40",
+    "L40S": "NVIDIA L40S",
+    "RTX3080": "NVIDIA GeForce RTX 3080",
+    "RTX3070": "NVIDIA GeForce RTX 3070",
+    "RTX3080Ti": "NVIDIA GeForce RTX 3080 Ti",
+    "A30": "NVIDIA A30",
+    "RTX4080": "NVIDIA GeForce RTX 4080",
+    "RTX3090": "NVIDIA GeForce RTX 3090",
+    "RTX3090Ti": "NVIDIA GeForce RTX 3090 Ti",
+    "V100": "Tesla V100-SXM2-32GB",  # Default V100 - 32GB
+    "V100_32": "Tesla V100-SXM2-32GB",
+    "V100_16": "Tesla V100-SXM2-16GB",
+    "V100_16_FHHL": "Tesla V100-FHHL-16GB",
+    "V100_16_PCIE": "Tesla V100-PCIE-16GB",
+    "RTX4070Ti": "NVIDIA GeForce RTX 4070 Ti",
+    "A4000_SFF": "NVIDIA RTX 4000 SFF Ada Generation",
+    "A5000_ADA": "NVIDIA RTX 5000 Ada Generation",
+    # "MI300X": "AMD Instinct MI300X OAM",
 }
 GPU_COUNT = 1
 allowed_cuda_versions = ['12.4']

{openweights-0.2.0 → openweights-0.4.0}/openweights/jobs/unsloth/__init__.py

@@ -13,6 +13,8 @@ from validate import TrainingConfig, MCQCallbackModel, MultipleChoiceEvalModel,
 from mc_question import MultipleChoiceEvalABC, MultipleChoiceEvalFreeform, MultipleChoiceEval, Question, Choice
 from logprobs import get_logprobs
 from mcq_callback import MCQCallback
+from huggingface_hub.utils import validate_repo_id
+from huggingface_hub.errors import HFValidationError
 
 
 @register("fine_tuning")        
@@ -39,28 +41,35 @@ class FineTuning(Jobs):
         
         params = TrainingConfig(**params).model_dump()
         mounted_files = self._upload_mounted_files()
-        job_id = self.compute_id({
-            'validated_params': params,
-            'mounted_files': mounted_files
-        })
-        model_name = params['model'].split('/')[-1]
-        params['finetuned_model_id'] = params['finetuned_model_id'].format(job_id=job_id, org_id=self.client.hf_org, model_name=model_name)
+        job_id = self.compute_id(
+            {"validated_params": params, "mounted_files": mounted_files}
+        )
+        model_name = params["model"].split("/")[-1]
+        params["finetuned_model_id"] = params["finetuned_model_id"].format(
+            job_id=job_id, org_id=self.client.hf_org, model_name=model_name
+        )
+        if params.get("ft_id_suffix", None):
+            params["finetuned_model_id"] += f"-{params['ft_id_suffix']}"
+
+        try:
+            validate_repo_id(params["finetuned_model_id"])
+        except HFValidationError as e:
+            raise ValueError(
+                f"Invalid finetuned_model_id: {params['finetuned_model_id']}. Error: {e}"
+            )
 
         data = {
-            'id': job_id,
-            'type': 'fine-tuning',
-            'model': params['model'],
-            'params': {
-                'validated_params': params,
-                'mounted_files': mounted_files
-            },
-            'status': 'pending',
-            'requires_vram_gb': requires_vram_gb,
-            'allowed_hardware': allowed_hardware,
-            'docker_image': self.base_image,
-            'script': f"python training.py {job_id}"
+            "id": job_id,
+            "type": "fine-tuning",
+            "model": params["model"],
+            "params": {"validated_params": params, "mounted_files": mounted_files},
+            "status": "pending",
+            "requires_vram_gb": requires_vram_gb,
+            "allowed_hardware": allowed_hardware,
+            "docker_image": self.base_image,
+            "script": f"python training.py {job_id}",
         }
-        
+
         return self.get_or_create_or_reset(data)
 
     def get_training_config(self, **params) -> Dict[str, Any]:

{openweights-0.2.0 → openweights-0.4.0}/openweights/jobs/unsloth/validate.py

@@ -17,8 +17,14 @@ class TrainingConfig(BaseModel):
     test_file: Optional[str] = Field(None, description="File ID of the test dataset")
 
     # Output model
-    finetuned_model_id: str = Field('{org_id}/{model_name}-{job_id}', description="File ID of the finetuned model")
-    
+    finetuned_model_id: str = Field(
+        "{org_id}/{model_name}-{job_id}", description="File ID of the finetuned model"
+    )
+    ft_id_suffix: Optional[str] = Field(
+        None, description="Suffix for the finetuned model id"
+    )
+    job_id_suffix: Optional[str] = Field(None, description="Suffix for the job id")
+
     # Model configuration
     max_seq_length: int = Field(2048, description="Maximum sequence length for training")
     load_in_4bit: bool = Field(False, description="Whether to load model in 4-bit quantization")